Heisenberg model

A minimal spin-1/2 Heisenberg chain example written directly in the DSL. This tutorial is configured to match NetKet’s Heisenberg(..., sign_rule=False) exactly.

We build

\[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)\]

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.

Setup

import netket as nk
import nkdsl

L = 8
hi = nk.hilbert.Spin(s=0.5, N=L, total_sz=0)
g = nk.graph.Chain(length=L, pbc=True)
edges = g.edges()
∣NK⟩ Tip: Debug multi-node HPC? `djaxrun -np 2 python Examples/Sharding/multi_process.py`

Construct the symbolic operator

H = (
    nkdsl.SymbolicDiscreteJaxOperator(hi, "heisenberg_sym", hermitian=True)
    .for_each(("i", "j"), over=edges)
    .emit(
        nkdsl.identity(),
        matrix_element=nkdsl.site("i").value * nkdsl.site("j").value,
    )
    .for_each(("i", "j"), over=edges)
    .where(nkdsl.site("i").value * nkdsl.site("j").value < 0)
    .emit(
        nkdsl.swap("i", "j"),
        matrix_element=2.0,
    )
    .compile()
)

Construct the state

model = nk.models.RBM(alpha=1, param_dtype=float)
state = nk.vqs.FullSumState(hi, model, seed=123)
opt = nk.optimizer.Sgd(learning_rate=0.01)
driver = nk.driver.VMC(H, opt, variational_state=state)

Run the VMC

driver.run(60, out=None, show_progress=False, timeit=True)
print("Final energy statistics:", state.expect(H))

╭────────────────────────────────────────────── Timing Information ───────────────────────────────────────────────╮
│ Total: 0.546                                                                                                    │
│ └── (94.8%) | VMC._forward_and_backward : 0.518 s                                                               │
╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯

Final energy statistics: -7.819e+00 [σ²=3.3e+01]

Compare with NetKet native

Operator

heisenberg_nk = nk.operator.Heisenberg(hi, g, J=1.0, sign_rule=False)

State

model = nk.models.RBM(alpha=1, param_dtype=float)
state = nk.vqs.FullSumState(hi, model, seed=123)
opt = nk.optimizer.Sgd(learning_rate=0.01)
driver = nk.driver.VMC(heisenberg_nk, opt, variational_state=state)

VMC

driver.run(60, out=None, show_progress=False, timeit=True)
print("Final energy statistics:", state.expect(heisenberg_nk))

╭────────────────────────────────────────────── Timing Information ───────────────────────────────────────────────╮
│ Total: 0.408                                                                                                    │
│ └── (96.1%) | VMC._forward_and_backward : 0.392 s                                                               │
╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯

Final energy statistics: -7.819e+00 [σ²=3.3e+01]

What this notebook shows

This is the complete nkdsl workflow in one place, with explicit parity checks against NetKet.

  1. Define a Hilbert space and static tuples to iterate over.

  2. Build the Hamiltonian declaratively with iterators, predicates, and emissions.

  3. Compile to a NetKet-compatible JAX operator.

  4. Verify dense-operator equality against native NetKet Heisenberg.

  5. Verify matched VMC traces from identical initial parameters.