{ "cells": [ { "cell_type": "markdown", "id": "fb23010d", "metadata": {}, "source": [ "# J1-J2 model\n", "\n", "A compact J1-J2 tutorial in the same spirit as the NetKet example, but focused only on the end-to-end user workflow. The DSL makes it natural to separate nearest- and next-nearest-neighbour branches while keeping the final operator readable." ] }, { "cell_type": "markdown", "id": "2b220d25", "metadata": {}, "source": [ "We consider\n", "\n", "$$H = J_1 \\sum_i \\vec{\\sigma}_i \\cdot \\vec{\\sigma}_{i+1} + J_2 \\sum_i \\vec{\\sigma}_i \\cdot \\vec{\\sigma}_{i+2}$$\n", "\n", "This is just the Heisenberg pattern repeated for two bond families with different couplings." ] }, { "cell_type": "markdown", "id": "efee9404dbf0791e", "metadata": {}, "source": [ "## Setup" ] }, { "cell_type": "code", "id": "2a2e6ee8", "metadata": { "execution": { "iopub.execute_input": "2026-04-07T06:10:28.474076Z", "iopub.status.busy": "2026-04-07T06:10:28.473864Z", "iopub.status.idle": "2026-04-07T06:10:28.508783Z", "shell.execute_reply": "2026-04-07T06:10:28.508232Z" }, "ExecuteTime": { "end_time": "2026-05-02T11:56:27.702318Z", "start_time": "2026-05-02T11:56:26.748544Z" } }, "source": [ "import netket as nk\n", "import nkdsl\n", "\n", "L = 8\n", "J1 = 1.0\n", "J2 = 0.5\n", "\n", "hi = nk.hilbert.Spin(s=0.5, N=L, total_sz=0)\n", "nearest_edges = [(i, (i + 1) % L) for i in range(L)]\n", "next_nearest_edges = [(i, (i + 2) % L) for i in range(L)]\n" ], "outputs": [ { "data": { "text/plain": [ "\u001B[1;36m∣NK⟩ Tip: \u001B[0m\u001B[1;35mlog\u001B[0m\u001B[1m(\u001B[0mψ\u001B[1m)\u001B[0m ∈ ℜ → ψ = \u001B[1;35mexp\u001B[0m\u001B[1m(\u001B[0mlogψ\u001B[1m)\u001B[0m ∈ ℜ₊ \u001B[1m(\u001B[0mreal NN gives only positive wave-function\u001B[1m)\u001B[0m.\n" ], "text/html": [ "
∣NK⟩ Tip: log(ψ) ∈ ℜ → ψ = exp(logψ) ∈ ℜ₊ (real NN gives only positive wave-function).\n",
       "
\n" ] }, "metadata": {}, "output_type": "display_data", "jetTransient": { "display_id": null } } ], "execution_count": 1 }, { "cell_type": "markdown", "id": "645c29ed67d1aaad", "metadata": {}, "source": [ "## Construct the symbolic operator" ] }, { "cell_type": "code", "id": "dc0fdf86", "metadata": { "execution": { "iopub.execute_input": "2026-04-07T06:10:28.510333Z", "iopub.status.busy": "2026-04-07T06:10:28.510257Z", "iopub.status.idle": "2026-04-07T06:10:28.522557Z", "shell.execute_reply": "2026-04-07T06:10:28.522120Z" }, "ExecuteTime": { "end_time": "2026-05-02T11:56:30.284373Z", "start_time": "2026-05-02T11:56:27.705596Z" } }, "source": [ "H = (\n", " nkdsl.SymbolicDiscreteJaxOperator(hi, \"j1j2_sym\", hermitian=True)\n", " .for_each((\"i\", \"j\"), over=nearest_edges)\n", " .emit(\n", " nkdsl.identity(),\n", " matrix_element=J1 * nkdsl.site(\"i\").value * nkdsl.site(\"j\").value,\n", " )\n", " .for_each((\"i\", \"j\"), over=nearest_edges)\n", " .where(nkdsl.site(\"i\").value * nkdsl.site(\"j\").value < 0)\n", " .emit(\n", " nkdsl.swap(\"i\", \"j\"),\n", " matrix_element=2.0 * J1,\n", " )\n", " .for_each((\"i\", \"j\"), over=next_nearest_edges)\n", " .emit(\n", " nkdsl.identity(),\n", " matrix_element=J2 * nkdsl.site(\"i\").value * nkdsl.site(\"j\").value,\n", " )\n", " .for_each((\"i\", \"j\"), over=next_nearest_edges)\n", " .where(nkdsl.site(\"i\").value * nkdsl.site(\"j\").value < 0)\n", " .emit(\n", " nkdsl.swap(\"i\", \"j\"),\n", " matrix_element=2.0 * J2,\n", " )\n", " .compile()\n", ")" ], "outputs": [], "execution_count": 2 }, { "cell_type": "markdown", "id": "8f5724fcde1a0a17", "metadata": {}, "source": [ "## Construct the state" ] }, { "cell_type": "code", "id": "17a1284df5405c02", "metadata": { "execution": { "iopub.execute_input": "2026-04-07T06:10:28.524038Z", "iopub.status.busy": "2026-04-07T06:10:28.523975Z", "iopub.status.idle": "2026-04-07T06:10:29.110988Z", "shell.execute_reply": "2026-04-07T06:10:29.110589Z" }, "ExecuteTime": { "end_time": "2026-05-02T11:56:30.882477Z", "start_time": "2026-05-02T11:56:30.313111Z" } }, "source": [ "model = nk.models.RBM(alpha=1, param_dtype=complex)\n", "state = nk.vqs.FullSumState(hi, model, seed=7)\n", "opt = nk.optimizer.Sgd(learning_rate=0.01)\n", "driver = nk.driver.VMC_SR(H, opt, variational_state=state, diag_shift=0.01)\n" ], "outputs": [], "execution_count": 3 }, { "cell_type": "markdown", "id": "f17a6bfb2b0c6412", "metadata": {}, "source": [ "## Run the VMC" ] }, { "cell_type": "code", "id": "cd5be1db", "metadata": { "execution": { "iopub.execute_input": "2026-04-07T06:10:29.112237Z", "iopub.status.busy": "2026-04-07T06:10:29.112173Z", "iopub.status.idle": "2026-04-07T06:10:32.045531Z", "shell.execute_reply": "2026-04-07T06:10:32.045134Z" }, "ExecuteTime": { "end_time": "2026-05-02T11:56:32.126741Z", "start_time": "2026-05-02T11:56:30.886117Z" } }, "source": [ "print(\"Initial energy statistics:\", state.expect(H))\n", "driver.run(1000, out=None, show_progress=False, timeit=True)\n", "final_stats = state.expect(H)\n", "exact_energy = nk.exact.lanczos_ed(H)[0]\n", "print(\"Final energy statistics:\", final_stats)\n", "print(\"Exact ground-state energy:\", exact_energy)\n", "print(\"Absolute error:\", abs(final_stats.mean - exact_energy))\n" ], "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Initial energy statistics: 1.199e+01-1.388e-17j [σ²=9.2e-02]\n" ] }, { "data": { "text/plain": [], "text/html": [ "
\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "╭────────────────────────────────────────────── Timing Information ───────────────────────────────────────────────╮\n",
      "│ Total: 0.709                                                                                                    │\n",
      "│ └── (90.5%) | VMC_SR._forward_and_backward : 0.642 s                                                            │\n",
      "│     └── (65.9%) | _sr_srt_common : 0.423 s                                                                      │\n",
      "│         └── (5.9%) | jacobian : 0.025 s                                                                         │\n",
      "╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\n",
      "\n",
      "Final energy statistics: -1.200e+01-6.897e-17j [σ²=6.1e-05]\n",
      "Exact ground-state energy: -12.000000000000007\n",
      "Absolute error: 5.409540474232699e-06\n"
     ]
    }
   ],
   "execution_count": 4
  },
  {
   "cell_type": "markdown",
   "id": "9af8a9a6c1344e96",
   "metadata": {},
   "source": [
    "## Compare against a native NetKet reference operator\n"
   ]
  },
  {
   "cell_type": "code",
   "id": "32523631e04c518a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-04-07T06:10:32.046758Z",
     "iopub.status.busy": "2026-04-07T06:10:32.046695Z",
     "iopub.status.idle": "2026-04-07T06:10:34.515563Z",
     "shell.execute_reply": "2026-04-07T06:10:34.515146Z"
    },
    "ExecuteTime": {
     "end_time": "2026-05-02T11:56:34.369708Z",
     "start_time": "2026-05-02T11:56:32.133488Z"
    }
   },
   "source": [
    "H_ref = nk.operator.LocalOperator(hi, dtype=complex)\n",
    "\n",
    "for i, j in nearest_edges:\n",
    "    H_ref += J1 * (\n",
    "        nk.operator.spin.sigmax(hi, i) @ nk.operator.spin.sigmax(hi, j)\n",
    "        + nk.operator.spin.sigmay(hi, i) @ nk.operator.spin.sigmay(hi, j)\n",
    "        + nk.operator.spin.sigmaz(hi, i) @ nk.operator.spin.sigmaz(hi, j)\n",
    "    )\n",
    "\n",
    "for i, j in next_nearest_edges:\n",
    "    H_ref += J2 * (\n",
    "        nk.operator.spin.sigmax(hi, i) @ nk.operator.spin.sigmax(hi, j)\n",
    "        + nk.operator.spin.sigmay(hi, i) @ nk.operator.spin.sigmay(hi, j)\n",
    "        + nk.operator.spin.sigmaz(hi, i) @ nk.operator.spin.sigmaz(hi, j)\n",
    "    )\n",
    "\n",
    "max_dense_diff = abs(H.to_dense() - H_ref.to_dense()).max()\n",
    "e0_sym = nk.exact.lanczos_ed(H)[0]\n",
    "e0_ref = nk.exact.lanczos_ed(H_ref)[0]\n",
    "\n",
    "print(\"Max |H_sym - H_ref| in dense form:\", max_dense_diff)\n",
    "print(\"E0 (Symbolic):\", e0_sym)\n",
    "print(\"E0 (NetKet reference):\", e0_ref)\n"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Max |H_sym - H_ref| in dense form: 0.0\n",
      "E0 (Symbolic): -12.000000000000004\n",
      "E0 (NetKet reference): -11.999999999999986\n"
     ]
    }
   ],
   "execution_count": 5
  },
  {
   "cell_type": "markdown",
   "id": "3eb14184",
   "metadata": {},
   "source": [
    "## What this notebook shows\n",
    "\n",
    "This is the complete ``nkdsl`` workflow in one place.\n",
    "\n",
    "1. Define a Hilbert space and a static set of tuples to iterate over.\n",
    "2. Build the Hamiltonian declaratively with iterators, predicates, and emissions.\n",
    "3. Compile the symbolic description into a NetKet-compatible JAX operator.\n",
    "4. Optimise a variational state directly against that compiled operator.\n",
    "5. Verify dense/operator-level parity against a native NetKet reference construction.\n"
   ]
  },
  {
   "cell_type": "code",
   "id": "a57e1e80a9d183e1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-05-02T11:56:34.378426Z",
     "start_time": "2026-05-02T11:56:34.374523Z"
    }
   },
   "source": [],
   "outputs": [],
   "execution_count": 5
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.11"
  },
  "mystnb": {
   "execution_mode": "cache"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}