{
"cells": [
{
"cell_type": "markdown",
"id": "47f07200",
"metadata": {},
"source": [
"# Heisenberg model\n",
"\n",
"A minimal spin-1/2 Heisenberg chain example written directly in the DSL.\n",
"This tutorial is configured to match NetKet's ``Heisenberg(..., sign_rule=False)`` exactly.\n"
]
},
{
"cell_type": "markdown",
"id": "8db6615a",
"metadata": {},
"source": [
"We build\n",
"\n",
"$$H = \\sum_{\\langle i,j \\rangle}\\left(\\sigma_i^x\\sigma_j^x + \\sigma_i^y\\sigma_j^y + \\sigma_i^z\\sigma_j^z\\right)$$\n",
"\n",
"In the computational basis, the diagonal contribution is ``sigma^z_i sigma^z_j`` and the off-diagonal contribution swaps anti-aligned nearest-neighbour spins with weight ``2``."
]
},
{
"cell_type": "markdown",
"id": "180c50a01cead626",
"metadata": {},
"source": [
"## Setup"
]
},
{
"cell_type": "code",
"id": "c69cdd74",
"metadata": {
"ExecuteTime": {
"end_time": "2026-05-02T11:56:09.078871Z",
"start_time": "2026-05-02T11:56:07.695424Z"
}
},
"source": [
"import netket as nk\n",
"import nkdsl\n",
"\n",
"L = 8\n",
"hi = nk.hilbert.Spin(s=0.5, N=L, total_sz=0)\n",
"g = nk.graph.Chain(length=L, pbc=True)\n",
"edges = g.edges()\n"
],
"outputs": [
{
"data": {
"text/plain": [
"\u001B[1;36m∣NK⟩ Tip: \u001B[0mDebug multi-node HPC? `djaxrun -np \u001B[1;36m2\u001B[0m python Examples/Sharding/multi_process.py`\n"
],
"text/html": [
"∣NK⟩ Tip: Debug multi-node HPC? `djaxrun -np 2 python Examples/Sharding/multi_process.py`\n",
"\n"
]
},
"metadata": {},
"output_type": "display_data",
"jetTransient": {
"display_id": null
}
}
],
"execution_count": 1
},
{
"cell_type": "markdown",
"id": "82f68f84a7057f8c",
"metadata": {},
"source": [
"## Construct the symbolic operator"
]
},
{
"cell_type": "code",
"id": "17bdf3aa",
"metadata": {
"ExecuteTime": {
"end_time": "2026-05-02T11:56:10.866804Z",
"start_time": "2026-05-02T11:56:09.081732Z"
}
},
"source": [
"H = (\n",
" nkdsl.SymbolicDiscreteJaxOperator(hi, \"heisenberg_sym\", hermitian=True)\n",
" .for_each((\"i\", \"j\"), over=edges)\n",
" .emit(\n",
" nkdsl.identity(),\n",
" matrix_element=nkdsl.site(\"i\").value * nkdsl.site(\"j\").value,\n",
" )\n",
" .for_each((\"i\", \"j\"), over=edges)\n",
" .where(nkdsl.site(\"i\").value * nkdsl.site(\"j\").value < 0)\n",
" .emit(\n",
" nkdsl.swap(\"i\", \"j\"),\n",
" matrix_element=2.0,\n",
" )\n",
" .compile()\n",
")\n"
],
"outputs": [],
"execution_count": 2
},
{
"cell_type": "markdown",
"id": "6afbd6f2063a7091",
"metadata": {},
"source": [
"## Construct the state"
]
},
{
"cell_type": "code",
"id": "bec1a112cd504c56",
"metadata": {
"ExecuteTime": {
"end_time": "2026-05-02T11:56:11.302940Z",
"start_time": "2026-05-02T11:56:10.925257Z"
}
},
"source": [
"model = nk.models.RBM(alpha=1, param_dtype=float)\n",
"state = nk.vqs.FullSumState(hi, model, seed=123)\n",
"opt = nk.optimizer.Sgd(learning_rate=0.01)\n",
"driver = nk.driver.VMC(H, opt, variational_state=state)"
],
"outputs": [],
"execution_count": 3
},
{
"cell_type": "markdown",
"id": "40abbae1110fa5f2",
"metadata": {},
"source": [
"## Run the VMC"
]
},
{
"cell_type": "code",
"id": "e962362f",
"metadata": {
"ExecuteTime": {
"end_time": "2026-05-02T11:56:11.975228Z",
"start_time": "2026-05-02T11:56:11.304155Z"
}
},
"source": [
"driver.run(60, out=None, show_progress=False, timeit=True)\n",
"print(\"Final energy statistics:\", state.expect(H))"
],
"outputs": [
{
"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.546 │\n",
"│ └── (94.8%) | VMC._forward_and_backward : 0.518 s │\n",
"╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\n",
"\n",
"Final energy statistics: -7.819e+00 [σ²=3.3e+01]\n"
]
}
],
"execution_count": 4
},
{
"cell_type": "markdown",
"id": "56923e956c63279f",
"metadata": {},
"source": [
"## Compare with NetKet native\n"
]
},
{
"cell_type": "markdown",
"id": "af117c1a6c3d4b3e",
"metadata": {},
"source": [
"### Operator"
]
},
{
"cell_type": "code",
"id": "b226c4c60ad41ab2",
"metadata": {
"ExecuteTime": {
"end_time": "2026-05-02T11:56:11.990296Z",
"start_time": "2026-05-02T11:56:11.983609Z"
}
},
"source": [
"heisenberg_nk = nk.operator.Heisenberg(hi, g, J=1.0, sign_rule=False)"
],
"outputs": [],
"execution_count": 5
},
{
"cell_type": "markdown",
"id": "101329331a299d7",
"metadata": {},
"source": [
"### State"
]
},
{
"cell_type": "code",
"id": "8d922b5008e5f6f2",
"metadata": {
"ExecuteTime": {
"end_time": "2026-05-02T11:56:11.995465Z",
"start_time": "2026-05-02T11:56:11.990734Z"
}
},
"source": [
"model = nk.models.RBM(alpha=1, param_dtype=float)\n",
"state = nk.vqs.FullSumState(hi, model, seed=123)\n",
"opt = nk.optimizer.Sgd(learning_rate=0.01)\n",
"driver = nk.driver.VMC(heisenberg_nk, opt, variational_state=state)"
],
"outputs": [],
"execution_count": 6
},
{
"cell_type": "markdown",
"id": "b17c42608cfd938b",
"metadata": {},
"source": [
"### VMC"
]
},
{
"cell_type": "code",
"id": "8bb289f1d6fe28c2",
"metadata": {
"ExecuteTime": {
"end_time": "2026-05-02T11:56:12.459082Z",
"start_time": "2026-05-02T11:56:11.995838Z"
}
},
"source": [
"driver.run(60, out=None, show_progress=False, timeit=True)\n",
"print(\"Final energy statistics:\", state.expect(heisenberg_nk))"
],
"outputs": [
{
"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.408 │\n",
"│ └── (96.1%) | VMC._forward_and_backward : 0.392 s │\n",
"╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\n",
"\n",
"Final energy statistics: -7.819e+00 [σ²=3.3e+01]\n"
]
}
],
"execution_count": 7
},
{
"cell_type": "markdown",
"id": "1d51a61e",
"metadata": {},
"source": [
"## What this notebook shows\n",
"\n",
"This is the complete ``nkdsl`` workflow in one place, with explicit parity checks against NetKet.\n",
"\n",
"1. Define a Hilbert space and static tuples to iterate over.\n",
"2. Build the Hamiltonian declaratively with iterators, predicates, and emissions.\n",
"3. Compile to a NetKet-compatible JAX operator.\n",
"4. Verify dense-operator equality against native NetKet Heisenberg.\n",
"5. Verify matched VMC traces from identical initial parameters.\n"
]
},
{
"cell_type": "code",
"id": "8cf1e6991bccdf8b",
"metadata": {
"ExecuteTime": {
"end_time": "2026-05-02T11:56:12.501351Z",
"start_time": "2026-05-02T11:56:12.460407Z"
}
},
"source": [],
"outputs": [],
"execution_count": 7
},
{
"cell_type": "code",
"id": "fb62921c2bf5c879",
"metadata": {
"ExecuteTime": {
"end_time": "2026-05-02T11:56:12.547799Z",
"start_time": "2026-05-02T11:56:12.502382Z"
}
},
"source": [],
"outputs": [],
"execution_count": 7
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python",
"pygments_lexer": "ipython3"
},
"mystnb": {
"execution_mode": "cache"
}
},
"nbformat": 4,
"nbformat_minor": 5
}