Quantum many-body physics at functional scale.

A simulation in ten lines

Build a Hamiltonian, define an ansatz, and run a variational Monte Carlo optimisation loop.

from tachys.lattice.spins.hamiltonians.ising_transverse_field import ising_transverse_field_square_pbc
from tachys.lattice.spins.spin_state import SpinState, init_config_fixed_magn
from tachys.lattice.ansatz.rbm import SpinRBM
from tachys.lattice.operator.local_estimator import compute_expectation
from tachys.montecarlo import sample
from tachys.wavefunction import WaveFunction
from tachys.optimizer import SR

H     = ising_transverse_field_square_pbc(L=4)
model = SpinRBM(num_hidden=1, dtype=jnp.float64)
state = SpinState(spins=init_config_fixed_magn(key, N, sz=0, N_mc=16), Ns=N)
wf    = WaveFunction(params=model.init(key, state), apply_fn=model.apply)

optimizer = SR(diag_shift=1e-4)
opt_state = optimizer.init(wf.params)

for step in range(100):
    state, log_amps, _ = sample(1, state, action, keys, wf)
    E_L, e_mean, e2_mean = compute_expectation(H, wf, state, log_amps)
    updates, opt_state = optimizer(E_L, opt_state, state, wf)
    wf = wf.apply_gradients(updates, eta=0.01)

1. 1ising_transverse_field_square_pbc — builds the Hamiltonian as a callable operator sum. Adding or scaling terms returns a new operator — no matrix is allocated yet.
2. 2SpinRBM — a Flax module implementing a Restricted Boltzmann Machine ansatz for spin-½ systems. num_hidden sets the hidden layer width.
3. 3init_config_fixed_magn — samples a batch of random spin configurations with fixed total magnetization sz, returning a SpinState pytree of shape (N_mc, N).
4. 4WaveFunction — wraps parameters and apply function into a single pytree, so the wavefunction flows through jit, vmap, and grad.
5. 5SR — Stochastic Reconfiguration optimizer. Preconditions the energy gradient with the quantum geometric tensor (via NTK). diag_shift regularises the inversion.
6. 6sample — runs Metropolis-Hastings sampling and returns the updated state, its log-amplitudes, and the acceptance rate.
7. 7compute_expectation — evaluates the local energy estimator E_L, then reduces to e_mean = ⟨E⟩ and e2_mean = ⟨|E|²⟩ across all devices.
8. 8optimizer — applies SR to E_L and returns the parameter updates and the new optimizer state.
9. 9apply_gradients — performs the descent step θ ← θ − η · updates, returning a new WaveFunction.

How Tachys works

Three primitives — states, operators, and the diagonalization engine — cover the full workflow. Each is independently composable.