J1-J2 model

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.

We consider

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

This is just the Heisenberg pattern repeated for two bond families with different couplings.

Setup

import netket as nk
import nkdsl

L = 8
J1 = 1.0
J2 = 0.5

hi = nk.hilbert.Spin(s=0.5, N=L, total_sz=0)
nearest_edges = [(i, (i + 1) % L) for i in range(L)]
next_nearest_edges = [(i, (i + 2) % L) for i in range(L)]

∣NK⟩ Tip: log(ψ) ∈ ℜ → ψ = exp(logψ) ∈ ℜ₊ (real NN gives only positive wave-function).

Construct the symbolic operator

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

Construct the state

model = nk.models.RBM(alpha=1, param_dtype=complex)
state = nk.vqs.FullSumState(hi, model, seed=7)
opt = nk.optimizer.Sgd(learning_rate=0.01)
driver = nk.driver.VMC_SR(H, opt, variational_state=state, diag_shift=0.01)

Run the VMC

print("Initial energy statistics:", state.expect(H))
driver.run(1000, out=None, show_progress=False, timeit=True)
final_stats = state.expect(H)
exact_energy = nk.exact.lanczos_ed(H)[0]
print("Final energy statistics:", final_stats)
print("Exact ground-state energy:", exact_energy)
print("Absolute error:", abs(final_stats.mean - exact_energy))

Initial energy statistics: 1.199e+01-1.388e-17j [σ²=9.2e-02]



╭────────────────────────────────────────────── Timing Information ───────────────────────────────────────────────╮
│ Total: 0.709                                                                                                    │
│ └── (90.5%) | VMC_SR._forward_and_backward : 0.642 s                                                            │
│     └── (65.9%) | _sr_srt_common : 0.423 s                                                                      │
│         └── (5.9%) | jacobian : 0.025 s                                                                         │
╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯

Final energy statistics: -1.200e+01-6.897e-17j [σ²=6.1e-05]
Exact ground-state energy: -12.000000000000007
Absolute error: 5.409540474232699e-06













Compare against a native NetKet reference operator






H_ref = nk.operator.LocalOperator(hi, dtype=complex)

for i, j in nearest_edges:
    H_ref += J1 * (
        nk.operator.spin.sigmax(hi, i) @ nk.operator.spin.sigmax(hi, j)
        + nk.operator.spin.sigmay(hi, i) @ nk.operator.spin.sigmay(hi, j)
        + nk.operator.spin.sigmaz(hi, i) @ nk.operator.spin.sigmaz(hi, j)
    )

for i, j in next_nearest_edges:
    H_ref += J2 * (
        nk.operator.spin.sigmax(hi, i) @ nk.operator.spin.sigmax(hi, j)
        + nk.operator.spin.sigmay(hi, i) @ nk.operator.spin.sigmay(hi, j)
        + nk.operator.spin.sigmaz(hi, i) @ nk.operator.spin.sigmaz(hi, j)
    )

max_dense_diff = abs(H.to_dense() - H_ref.to_dense()).max()
e0_sym = nk.exact.lanczos_ed(H)[0]
e0_ref = nk.exact.lanczos_ed(H_ref)[0]

print("Max |H_sym - H_ref| in dense form:", max_dense_diff)
print("E0 (Symbolic):", e0_sym)
print("E0 (NetKet reference):", e0_ref)









Max |H_sym - H_ref| in dense form: 0.0
E0 (Symbolic): -12.000000000000004
E0 (NetKet reference): -11.999999999999986













What this notebook shows


This is the complete nkdsl workflow in one place.


    

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

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

  3. Compile the symbolic description into a NetKet-compatible JAX operator.
    4. 

  4. Optimise a variational state directly against that compiled operator.
    5. 

  5. Verify dense/operator-level parity against a native NetKet reference construction.
    6.