Transaction-Science/open-standards
Watch
1
0
Fork
You've already forked open-standards
 0
Code Issues Pull requests Projects Releases 15 Packages Wiki Activity Actions 8

wai-quantum 0.3.38: neural quantum states held in their symmetry sector #180

Merged
dcharlot merged 2 commits from wai-quantum-nqs-symmetry into main 2026-10-08 10:24:11 -04:00
Conversation 0 Commits 2 Files changed 11 +506 -50
11 changed files with 506 additions and 50 deletions
Download patch file Download diff file Expand all files Collapse all files

2
wai/wai-pack/Cargo.lock generated
View file

@ -682,7 +682,7 @@ dependencies = [
[[package]]
name = "wai-quantum"
version = "0.3.36"
version = "0.3.38"
dependencies = [
"blake3",
"ed25519-dalek",

2
wai/wai-quantum-serve/Cargo.lock generated
View file

@ -589,7 +589,7 @@ dependencies = [
[[package]]
name = "wai-quantum"
version = "0.3.37"
version = "0.3.38"
dependencies = [
"blake3",
"ed25519-dalek",

2
wai/wai-quantum/Cargo.lock generated
View file

@ -415,7 +415,7 @@ checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"
[[package]]
name = "wai-quantum"
version = "0.3.37"
version = "0.3.38"
dependencies = [
"blake3",
"ed25519-dalek",

2
wai/wai-quantum/Cargo.toml
View file

@ -1,6 +1,6 @@
[package]
name = "wai-quantum"
version = "0.3.37"
version = "0.3.38"
edition = "2024"
rust-version = "1.97"
license = "Apache-2.0"

8
wai/wai-quantum/README.md
View file

@ -69,7 +69,7 @@ never creates information.
| `quantum_dmrg` | **ground states of spin chains by DMRG**: matrix-product operators for nearest-neighbour chains (XXZ Heisenberg in a field, transverse-field Ising), two-site sweeps with Lanczos and density-matrix truncation, entanglement entropies, a symmetric eigensolver, and two exact referees: full diagonalisation up to 22 sites, and the free-fermion solution of the Ising chain at any length. The critical Ising chain is reproduced to about `10⁻¹¹` per site at 200 sites |
| `quantum_tdvp` | **time evolution of spin chains by TDVP**: two-site sweeps forward and back, Lanczos exponentials that halve a step they cannot converge, bonds split by singular value decomposition, and every basis completed with null vectors so a product state's small bonds do not hold the evolution back. Measures site and bond observables, entanglement at every bond, energy and the Loschmidt echo. Held to two exact referees: Krylov evolution of all amplitudes up to 20 sites, and the free-fermion solution of an Ising quench at any length. On a 100-site critical quench at bond 64 it stays within `10⁻⁹` of the exact answer until the entanglement fills the bond, and past that its error follows the discarded weight it reports |
| `quantum_chem` | **molecules by DMRG**: FCIDUMP integrals, the second-quantised Hamiltonian on spin orbitals compiled into an exact MPO (minimum vertex covers, no numerical compression), DMRG from Hartree–Fock with complete bases, and full configuration interaction as the referee. LiH and H₂O reach full CI to `10⁻¹³` |
| `quantum_nqs` | **neural quantum states**: spin models on any graph (Heisenberg, frustrated J₁–J₂, transverse-field Ising), a restricted Boltzmann machine with real parameters (and Marshall's sign on bipartite lattices) or complex ones that learn a sign structure nobody supplied, variational Monte Carlo by Metropolis chains or exact enumeration, and stochastic reconfiguration solved in whichever of parameter or sample space is smaller. Held to exact diagonalisation and to DMRG on the same operator sum |
| `quantum_nqs` | **neural quantum states**: spin models on any graph (Heisenberg, frustrated J₁–J₂, transverse-field Ising), a restricted Boltzmann machine with real parameters (and Marshall's sign on bipartite lattices) or complex ones that learn a sign structure nobody supplied, projected onto a momentum and spin-flip sector, variational Monte Carlo by Metropolis chains or exact enumeration, and stochastic reconfiguration solved in whichever of parameter or sample space is smaller. Held to exact diagonalisation and to DMRG on the same operator sum |
| `quantum_noisy` | **noisy circuits**: quantum trajectories on the state vector (depolarizing, dephasing, amplitude damping, readout), an exact density-matrix referee up to 12 qubits that reuses the same kernels, linear cross-entropy fidelity (sampled, exact and the circuit's own), and the digital error model. Measures where that model is off: below the true fidelity by a margin that grows with the error rate, not with depth |
| `quantum_ldpc` | **qLDPC memory at circuit level**: bivariate-bicycle codes from their two polynomials (`k` by rank, logical bases), the depth-8 syndrome cycle with circuit-level noise, its detector error model, and BP-OSD decoding (min-sum with adaptive scaling, ordered statistics with a combination sweep). The error model has the same columns as the source's own decoding matrices, and on identical syndromes the decoder makes the source's decisions shot for shot |
| `quantum_arch` | **architecture-level estimation**: yoked surface-code storage (row and grid parity checks over patches, hot and cold), magic-state cultivation read off the published simulation curves, CCZ factories, and a clock set by whichever is slowest of lattice surgery, the magic-state supply and the control loop. Assembles whole-machine estimates; its storage search matches the source's footprint tool layout for layout |
@ -243,6 +243,7 @@ weight that accounts for the difference.
wai-quantum nqs --lx 8 --exact --steps 200 --lr 0.1 # every configuration enumerated: no sampling noise
wai-quantum nqs --model ising --lx 4 --ly 4 --h 3.044 --alpha 1 --steps 400
wai-quantum nqs --model j1j2 --lx 8 --periodic --complex --exact --steps 250 --lr 0.1 # frustrated: the sign is learned
wai-quantum nqs --lx 10 --periodic --complex --project --kpi --odd --exact --steps 200 --lr 0.1 # held in its sector
```
Neural-network wavefunctions are offered as evidence about systems that no
@ -252,7 +253,10 @@ lands within its own error bar of exact diagonalisation. Past exact reach,
the same Hamiltonian goes to DMRG through its exact MPO, so two independent
methods meet on one number. For frustrated models, whose signs no rule gives,
a complex network learns them. On the Majumdar–Ghosh ring it reaches the
exact energy to rounding.
exact energy to rounding. Projected onto a lattice's translations and the
global spin flip, a network is held in the sector the ground state lives in.
Rings reach their exact energies to rounding, and the frustrated 4 × 4 torus
at `J₂ = 0.5` comes within a percent.
## Noisy circuits

18
wai/wai-quantum/examples/nqs_ground.rs
View file

@ -6,7 +6,11 @@
//! cargo run --release --example nqs_ground --features quantum_nqs -- \
//! --model heisenberg|j1j2|ising --lx L --ly L [--j2 J] [--h H] [--periodic]
//! [--complex] [--marshall] [--alpha A] [--steps N] [--samples S] [--lr R]
//! [--dmrg-bond D] [--threads T]
//! [--project [--kpi] [--odd]] [--dmrg-bond D] [--threads T]
//!
//! `--project` projects onto translations of the periodic lattice at
//! momentum 0 (or π with `--kpi`) and onto even (or `--odd`) parity under a
//! global spin flip.
//! ```
use wai_quantum::quantum_dmrg::{dmrg_mpo, DmrgConfig, Mps, C};
use wai_quantum::quantum_nqs::*;
@ -60,17 +64,29 @@ fn main() {
let cfg = VmcConfig { samples: arg(&a, "--samples", 1000), chains: 16, magnetisation: sector, lr: arg(&a, "--lr", 0.02), shift: 1e-3, threads, ..VmcConfig::default() };
let steps: usize = arg(&a, "--steps", 600);
let marshall = a.iter().any(|x| x == "--marshall");
let symmetry = a.iter().any(|x| x == "--project").then(|| {
assert!(periodic, "--project needs --periodic");
let kpi = a.iter().any(|x| x == "--kpi");
let base = if ly == 1 { Symmetry::ring(n, kpi) } else { Symmetry::torus(lx, ly, kpi, kpi) };
base.with_spin_flip(a.iter().any(|x| x == "--odd"))
});
if a.iter().any(|x| x == "--complex") {
let mut rbm: Rbm<C> = Rbm::complex(n, alpha, 0.01, 7);
if marshall {
rbm = rbm.with_marshall_sign(checkerboard(lx, ly));
}
if let Some(sym) = symmetry {
rbm = rbm.with_symmetry(sym);
}
train(model, rbm, cfg, steps, reference);
} else {
let mut rbm = Rbm::new(n, alpha, 0.01, 7);
if marshall {
rbm = rbm.with_marshall_sign(checkerboard(lx, ly));
}
if let Some(sym) = symmetry {
rbm = rbm.with_symmetry(sym);
}
train(model, rbm, cfg, steps, reference);
}
}

485
wai/wai-quantum/src/quantum_nqs.rs
View file

@ -17,12 +17,24 @@
//! - A complex network carries a phase, so training can find a sign
//! structure nobody supplied. Frustrated models need that, because no
//! sign rule is known for them.
//! - Weights and biases start small and random. With the biases at zero,
//! the network would be exactly even under a global spin flip.
//! - **The sector.** [`Symmetry`] projects the network,
//! `Ψ(s) = Σ_g χ(g) ψ(g·s)`, over a lattice's translations (momentum 0 or
//! π along each direction) and, optionally, the flip of every spin (even
//! or odd). Ratios and log-derivatives become weighted averages over the
//! transformed configurations, with a cache that Metropolis moves keep
//! up to date.
//! - **Variational Monte Carlo.** [`Vmc`] samples `|ψ|²` by Metropolis
//! chains: single flips, or swaps of antiparallel spins at fixed
//! magnetisation. Or it enumerates every configuration (at most 24 spins)
//! for expectations with no Monte Carlo noise. Each step is stochastic
//! reconfiguration, `θ ← θ − η (S + λ)⁻¹ F`, solved in parameter space or
//! in sample space, whichever is smaller; the two give the same step.
//! in sample space, whichever is smaller; the two give the same step. When
//! Cholesky loses a pivot to rounding, the step comes from the
//! eigendecomposition instead, with eigenvalues below `10⁻¹²` of the
//! largest dropped. A projected network can take its metric there:
//! amplitudes that nearly cancel give log-derivatives of order `10⁸`.
//! - **Referees.** [`SpinModel::exact_ground_energy`] is Lanczos on all `2ⁿ`
//! amplitudes (at most 22 spins). [`SpinModel::opsum`] gives the model to
//! DMRG on its exact MPO, for lattices past exact diagonalisation.
@ -46,6 +58,22 @@
//! one, to `10⁻¹³`. Its ratios equal ratios of its amplitudes. Its
//! holomorphic log-derivatives agree with differences along a real step.
//! Both spaces take the same Hermitian reconfiguration step.
//! - **The projection.** It equals the explicit sum over the group, ratios
//! and log-derivatives included. It carries its sector: a translation
//! multiplies `Ψ` by its character, a global flip by its parity. Through
//! accepted moves, the cache stays equal to a fresh one.
//! - **Rings in their sector.** With every configuration enumerated, `α = 1`,
//! reaches each ring's exact energy to rounding:
//!
//! | ring (sector) | steps | relative error |
//! |----------------------------------------|-------|----------------|
//! | 6 sites, Heisenberg (π, odd) | 150 | `< 10⁻⁸` |
//! | 10 sites, Heisenberg (π) | 300 | 7·10⁻¹⁴ |
//! | 10 sites, Majumdar–Ghosh (π) | 300 | 3·10⁻¹⁵ |
//! | 10 sites, J₁–J₂ 0.3 (π) | 300 | 3·10⁻¹⁴ |
//! | 12 sites, Heisenberg (0, even) | 400 | 7·10⁻¹⁶ |
//! | 12 sites, Majumdar–Ghosh (0, even) | 400 | 4·10⁻¹⁶ |
//! | 12 sites, J₁–J₂ 0.3 (0, even) | 400 | 1·10⁻¹⁵ |
//! - **A sign nobody supplied.** Exact diagonalisation gives the
//! Majumdar–Ghosh ring (`J₂ = J₁/2`) its exact `−3n/8`. On 8-site rings
//! with no sign rule given, Heisenberg and Majumdar–Ghosh, a real network
@ -86,9 +114,13 @@
//! - Starting the biases away from zero frees it. The same 10-site runs then
//! reach `3·10⁻⁴` (Heisenberg), `1·10⁻⁶` (Majumdar–Ghosh) and `1.4·10⁻⁵`
//! (J₁–J₂ 0.3) in 400 steps.
//! - On 12-site rings, where the ground state is even, the even start is the
//! better one (the table above). Drawn biases there reach only `2·10⁻²`,
//! `6·10⁻⁷` and `3·10⁻³`.
//! - On 12-site rings, where the ground state is even, the even start did
//! better than drawn biases (`2·10⁻³` against `2·10⁻²` for Heisenberg).
//! It was a sector chosen by accident.
//!
//! Biases are now drawn, and a sector is chosen on purpose with
//! [`Symmetry`]. That is the table of rings above, where projection takes
//! every case to rounding.
//!
//! The variance cannot catch the trap, since an excited eigenstate has none.
//! Only a referee can.
@ -102,14 +134,20 @@
//!
//! # Honest boundaries
//!
//! - **The hardest point is not reached.** At `J₂ = 0.5` on the periodic
//! 4 × 4 lattice, exact diagonalisation gives `−8.45792335`. At `α = 1`,
//! 600 steps of 1000 samples leave all three networks about 3 % above it:
//! complex with Marshall's sign, complex without, and real with it. At
//! this size and budget the network is the limit, not the sampling.
//! - **Symmetry is not chosen.** A network is not projected onto a momentum
//! or spin-flip sector. Started with zero biases it stays even under a
//! global spin flip, and that settles the sector, as above.
//! - **The hardest point, within a percent.** At `J₂ = 0.5` on the periodic
//! 4 × 4 lattice, exact diagonalisation gives `−8.45792335`.
//! - Unprojected at `α = 1`, 600 steps of 1000 samples leave every network
//! about 3 % above it: complex with Marshall's sign, complex without,
//! and real with it.
//! - Projected onto momentum zero and even parity, a complex `α = 1`
//! network is 1.7 % above after 600 steps. After 2000 it is
//! `7.7·10⁻³ ± 1.2·10⁻³` above (final estimate of 8000 samples).
//! - Rotations and reflections are left out of the group, and the network
//! is small; what remains is theirs to close.
//! - **Translations and the global flip only.** Rotations and reflections
//! of the lattice are not in the group, and characters are `±1`, so the
//! momenta are 0 and π. Projection multiplies the cost of every amplitude
//! by the group's size.
//! - **Error bars ignore autocorrelation.** The reported error is the
//! standard error of uncorrelated samples. Chains thinned by two sweeps are
//! correlated, so the true error is larger.
@ -327,6 +365,7 @@ pub trait Amplitude: Scalar {
fn tanh(self) -> Self;
/// `self / d` for a real `d`.
fn div_re(self, d: f64) -> Self;
fn div(self, o: Self) -> Self;
/// An initial weight, uniform in `±scale` (each part, when complex).
fn draw(uniform: &mut dyn FnMut() -> f64, scale: f64) -> Self;
}
@ -347,6 +386,9 @@ impl Amplitude for f64 {
fn div_re(self, d: f64) -> f64 {
self / d
}
fn div(self, o: f64) -> f64 {
self / o
}
fn draw(uniform: &mut dyn FnMut() -> f64, scale: f64) -> f64 {
scale * (2.0 * uniform() - 1.0)
}
@ -386,6 +428,9 @@ impl Amplitude for C {
fn div_re(self, d: f64) -> C {
C { re: self.re / d, im: self.im / d }
}
fn div(self, o: C) -> C {
c_div(self, o)
}
fn draw(uniform: &mut dyn FnMut() -> f64, scale: f64) -> C {
let re = scale * (2.0 * uniform() - 1.0);
let im = scale * (2.0 * uniform() - 1.0);
@ -393,6 +438,99 @@ impl Amplitude for C {
}
}
/// A symmetry group of the lattice: site permutations, each possibly
/// followed by flipping every spin, each with a real character. The
/// projection `Ψ(s) = Σ_g χ(g) ψ(g·s)` picks the sector: a momentum of 0 or
/// π along each direction and, with the flip, even or odd parity.
#[derive(Clone, Debug, PartialEq)]
pub struct Symmetry {
/// `perms[g][i]`: the site whose spin lands on site `i`.
pub perms: Vec<Vec<usize>>,
/// Whether element `g` also flips every spin.
pub flips: Vec<bool>,
/// `±1`, the character of each element.
pub chars: Vec<f64>,
/// `inverse[g][j]`: where site `j`'s spin lands.
inverse: Vec<Vec<usize>>,
}
impl Symmetry {
pub fn new(perms: Vec<Vec<usize>>, chars: Vec<f64>) -> Symmetry {
let flips = vec![false; perms.len()];
Symmetry::with_flips(perms, flips, chars)
}
fn with_flips(perms: Vec<Vec<usize>>, flips: Vec<bool>, chars: Vec<f64>) -> Symmetry {
assert_eq!(perms.len(), chars.len());
assert_eq!(perms.len(), flips.len());
let inverse = perms
.iter()
.map(|p| {
let mut inv = vec![usize::MAX; p.len()];
for (i, &j) in p.iter().enumerate() {
inv[j] = i;
}
assert!(inv.iter().all(|&x| x != usize::MAX), "not a permutation");
inv
})
.collect();
Symmetry { perms, flips, chars, inverse }
}
/// Add the flip of every spin to the group, projecting onto even or odd
/// parity under it. A singlet's parity is `(−1)^{n/2}`.
pub fn with_spin_flip(self, odd: bool) -> Symmetry {
let mut perms = self.perms.clone();
let mut flips = self.flips.clone();
let mut chars = self.chars.clone();
for g in 0..self.perms.len() {
perms.push(self.perms[g].clone());
flips.push(!self.flips[g]);
chars.push(if odd { -self.chars[g] } else { self.chars[g] });
}
Symmetry::with_flips(perms, flips, chars)
}
/// The translations of a ring of `n` sites, at momentum 0 or π.
pub fn ring(n: usize, momentum_pi: bool) -> Symmetry {
let perms = (0..n).map(|t| (0..n).map(|i| (i + t) % n).collect()).collect();
let chars = (0..n).map(|t| if momentum_pi && t % 2 == 1 { -1.0 } else { 1.0 }).collect();
Symmetry::new(perms, chars)
}
/// The translations of an `lx × ly` torus (site `y·lx + x`), at momentum
/// 0 or π along each direction.
pub fn torus(lx: usize, ly: usize, kx_pi: bool, ky_pi: bool) -> Symmetry {
let mut perms = Vec::new();
let mut chars = Vec::new();
for ty in 0..ly {
for tx in 0..lx {
perms.push((0..lx * ly).map(|i| ((i / lx + ty) % ly) * lx + (i % lx + tx) % lx).collect());
let odd = (kx_pi && tx % 2 == 1) ^ (ky_pi && ty % 2 == 1);
chars.push(if odd { -1.0 } else { 1.0 });
}
}
Symmetry::new(perms, chars)
}
pub fn len(&self) -> usize {
self.perms.len()
}
pub fn is_empty(&self) -> bool {
self.perms.is_empty()
}
}
/// What a configuration's amplitude ratios and derivatives need: `θ(s)`,
/// or for a projected network `θ` of every transformed configuration and its
/// share `p_g = χ(g) ψ(s∘g) / Ψ(s)` of the amplitude.
#[derive(Clone, Debug)]
enum Cache<T> {
Plain(Vec<T>),
Projected { configs: Vec<Vec<i8>>, thetas: Vec<Vec<T>>, shares: Vec<T>, log_abs: f64 },
}
/// A restricted Boltzmann machine: `ψ(s) = sign(s) · exp(Σ aᵢsᵢ) ·
/// Π_j cosh(b_j + Σ W_jᵢ sᵢ)`, with `m = α·n` hidden units. Its parameters are
/// real (`Rbm<f64>`), so `ψ` is positive up to the optional sign, or complex
@ -409,6 +547,8 @@ pub struct Rbm<T = f64> {
/// `W[j·n + i]`.
pub w: Vec<T>,
pub sublattice: Option<Vec<bool>>,
/// Project onto a momentum: `Ψ(s) = Σ_g χ(g) ψ(s∘g)`.
pub symmetry: Option<Symmetry>,
}
/// A deterministic generator (splitmix64).
@ -456,7 +596,13 @@ impl<T: Amplitude> Rbm<T> {
let mut rng = Rng(seed);
let mut u = || rng.uniform();
let w = (0..m * n).map(|_| T::draw(&mut u, scale)).collect();
Rbm { n, m, a: vec![T::ZERO; n], b: vec![T::ZERO; m], w, sublattice: None }
// Biases drawn too: with both at zero the machine is exactly even
// under flipping every spin, and training keeps it so, which shuts
// it out of every odd state (singlets on rings whose half-length is
// odd among them).
let a = (0..n).map(|_| T::draw(&mut u, scale)).collect();
let b = (0..m).map(|_| T::draw(&mut u, scale)).collect();
Rbm { n, m, a, b, w, sublattice: None, symmetry: None }
}
/// Multiply by Marshall's sign for this sublattice.
@ -466,6 +612,122 @@ impl<T: Amplitude> Rbm<T> {
self
}
/// Project onto the momentum `symmetry` selects.
pub fn with_symmetry(mut self, symmetry: Symmetry) -> Rbm<T> {
assert!(symmetry.perms.iter().all(|p| p.len() == self.n));
self.symmetry = Some(symmetry);
self
}
/// The configuration's cache.
fn cache(&self, s: &[i8]) -> Cache<T> {
let Some(sym) = &self.symmetry else {
return Cache::Plain(self.theta(s));
};
let configs: Vec<Vec<i8>> = sym.perms.iter().zip(&sym.flips).map(|(p, &f)| p.iter().map(|&j| if f { -s[j] } else { s[j] }).collect()).collect();
let thetas: Vec<Vec<T>> = configs.iter().map(|t| self.theta(t)).collect();
let logs: Vec<T> = configs.iter().map(|t| self.log_psi(t)).collect();
let top = logs.iter().map(|l| l.re()).fold(f64::NEG_INFINITY, f64::max);
// χ(g) sign(s∘g) e^{ln ψ(s∘g) − top}, summed in order.
let terms: Vec<T> = logs.iter().zip(&configs).zip(&sym.chars).map(|((l, t), &c)| l.sub(T::ONE.scale(top)).exp().scale(c * self.sign(t))).collect();
let total = terms.iter().fold(T::ZERO, |acc, &x| acc.add(x));
let mag2 = total.norm2();
// A sum below the rounding error of its own terms is zero: the
// configuration is one the symmetry forbids, and what is left is
// noise. Its log-derivative, noise over noise, would otherwise feed
// the metric something of order one.
let scale: f64 = terms.iter().map(|x| x.norm2().sqrt()).sum();
let vanishes = mag2.sqrt() <= 1e-13 * scale;
let log_abs = if vanishes { f64::NEG_INFINITY } else { top + 0.5 * ln(mag2) };
let shares = if vanishes { vec![T::ZERO; terms.len()] } else { terms.iter().map(|&x| x.div(total)).collect() };
Cache::Projected { configs, thetas, shares, log_abs }
}
/// `Ψ(s′)/Ψ(s)` where `s′` is `s` with `flips` flipped.
fn ratio_cached(&self, s: &[i8], cache: &Cache<T>, flips: &[usize]) -> T {
match cache {
Cache::Plain(theta) => self.ratio(s, theta, flips),
Cache::Projected { configs, thetas, shares, .. } => {
let sym = self.symmetry.as_ref().expect("a projected cache has a symmetry");
let mut r = T::ZERO;
let mut moved = Vec::with_capacity(flips.len());
for g in 0..configs.len() {
if shares[g].is_zero() {
continue;
}
moved.clear();
moved.extend(flips.iter().map(|&j| sym.inverse[g][j]));
r = r.add(shares[g].mul(self.ratio(&configs[g], &thetas[g], &moved)));
}
r
}
}
}
/// `∂ ln Ψ / ∂(a, b, W)`.
fn derivatives_cached(&self, s: &[i8], cache: &Cache<T>, out: &mut [T]) {
match cache {
Cache::Plain(theta) => self.log_derivatives(s, theta, out),
Cache::Projected { configs, thetas, shares, .. } => {
out.iter_mut().for_each(|o| *o = T::ZERO);
let mut one = vec![T::ZERO; out.len()];
for g in 0..configs.len() {
if shares[g].is_zero() {
continue;
}
self.log_derivatives(&configs[g], &thetas[g], &mut one);
for (o, &d) in out.iter_mut().zip(&one) {
*o = o.add(shares[g].mul(d));
}
}
}
}
}
/// `ln |Ψ(s)|` from the cache (up to a constant).
fn log_abs_cached(&self, s: &[i8], cache: &Cache<T>) -> f64 {
match cache {
Cache::Plain(_) => self.log_amplitude(s),
Cache::Projected { log_abs, .. } => *log_abs,
}
}
/// Flip `flips` in `s` and bring the cache along, given the ratio `r`
/// already computed for the move.
fn accept(&self, s: &mut [i8], cache: &mut Cache<T>, flips: &[usize], r: T) {
let n = self.n;
match cache {
Cache::Plain(theta) => {
for &i in flips {
let ds = -2.0 * f64::from(s[i]);
for (j, t) in theta.iter_mut().enumerate() {
*t = t.add(self.w[j * n + i].scale(ds));
}
s[i] = -s[i];
}
}
Cache::Projected { configs, thetas, shares, log_abs } => {
let sym = self.symmetry.as_ref().expect("a projected cache has a symmetry");
for g in 0..configs.len() {
let moved: Vec<usize> = flips.iter().map(|&j| sym.inverse[g][j]).collect();
let rg = self.ratio(&configs[g], &thetas[g], &moved);
shares[g] = shares[g].mul(rg).div(r);
for &i in &moved {
let ds = -2.0 * f64::from(configs[g][i]);
for (j, t) in thetas[g].iter_mut().enumerate() {
*t = t.add(self.w[j * n + i].scale(ds));
}
configs[g][i] = -configs[g][i];
}
}
*log_abs += 0.5 * ln(r.norm2());
for &i in flips {
s[i] = -s[i];
}
}
}
}
/// The number of variational parameters (complex ones counted once).
pub fn params(&self) -> usize {
self.n + self.m + self.n * self.m
@ -565,7 +827,7 @@ impl<T: Amplitude> Rbm<T> {
}
/// `E_loc(s) = Σ_{s′} H_{ss′} ψ(s′)/ψ(s)`.
fn local_energy<T: Amplitude>(model: &SpinModel, rbm: &Rbm<T>, s: &[i8], theta: &[T]) -> T {
fn local_energy<T: Amplitude>(model: &SpinModel, rbm: &Rbm<T>, s: &[i8], cache: &Cache<T>) -> T {
let mut e = T::ZERO;
for &(i, j, c) in &model.zz {
e = e.add(T::ONE.scale(c * f64::from(s[i]) * f64::from(s[j])));
@ -575,11 +837,11 @@ fn local_energy<T: Amplitude>(model: &SpinModel, rbm: &Rbm<T>, s: &[i8], theta:
}
for &(i, j, c) in &model.exchange {
if s[i] != s[j] {
e = e.add(rbm.ratio(s, theta, &[i, j]).scale(2.0 * c));
e = e.add(rbm.ratio_cached(s, cache, &[i, j]).scale(2.0 * c));
}
}
for &(i, c) in &model.x {
e = e.add(rbm.ratio(s, theta, &[i]).scale(c));
e = e.add(rbm.ratio_cached(s, cache, &[i]).scale(c));
}
e
}
@ -698,11 +960,15 @@ impl<T: Amplitude> Vmc<T> {
for idx in range {
let x = configs[idx];
let s: Vec<i8> = (0..n).map(|i| if x >> (n - 1 - i) & 1 == 0 { 1 } else { -1 }).collect();
let theta = self.rbm.theta(&s);
let cache = self.rbm.cache(&s);
let r = idx - start;
logs[r] = self.rbm.log_amplitude(&s);
energies[r] = local_energy(&self.model, &self.rbm, &s, &theta);
self.rbm.log_derivatives(&s, &theta, &mut derivs[r * p..(r + 1) * p]);
logs[r] = self.rbm.log_abs_cached(&s, &cache);
if logs[r] == f64::NEG_INFINITY {
// The projection vanishes here: no weight, nothing to add.
continue;
}
energies[r] = local_energy(&self.model, &self.rbm, &s, &cache);
self.rbm.derivatives_cached(&s, &cache, &mut derivs[r * p..(r + 1) * p]);
}
};
if threads <= 1 || cfg!(target_arch = "wasm32") {
@ -752,9 +1018,9 @@ impl<T: Amplitude> Vmc<T> {
let mut accepted = 0usize;
for k in 0..per_chain {
accepted += metropolis(rbm, s, rng, thin, exchange);
let theta = rbm.theta(s);
energies.push(local_energy(model, rbm, s, &theta));
rbm.log_derivatives(s, &theta, &mut derivs[k * p..(k + 1) * p]);
let cache = rbm.cache(s);
energies.push(local_energy(model, rbm, s, &cache));
rbm.derivatives_cached(s, &cache, &mut derivs[k * p..(k + 1) * p]);
}
(energies, derivs, accepted as f64 / (per_chain * thin) as f64)
};
@ -874,7 +1140,7 @@ fn sr_parameter_space<T: Amplitude>(o: &[T], eps: &[T], k: usize, p: usize, shif
f[c] = f[c].add(o[r * p + c].conj().mul(eps[r]));
}
}
cholesky_solve(&mut s, &mut f, p);
hermitian_solve(&mut s, &mut f, p);
f
}
@ -893,7 +1159,7 @@ fn sr_sample_space<T: Amplitude>(o: &[T], eps: Vec<T>, k: usize, p: usize, shift
t[i * k + i] = t[i * k + i].add(T::ONE.scale(shift));
}
let mut x = eps;
cholesky_solve(&mut t, &mut x, k);
hermitian_solve(&mut t, &mut x, k);
let mut delta = vec![T::ZERO; p];
for r in 0..k {
for c in 0..p {
@ -903,15 +1169,52 @@ fn sr_sample_space<T: Amplitude>(o: &[T], eps: Vec<T>, k: usize, p: usize, shift
delta
}
/// Solve `A x = b` for Hermitian positive-definite `A` (overwritten by its
/// Cholesky factor `L`, `A = L L†`), `b` overwritten by `x`.
/// Solve `A x = b` for Hermitian positive-definite `A`, `b` overwritten by
/// `x`. Cholesky first. If a pivot is lost to rounding, the metric is
/// singular at the scale of its own entries (projected networks with nearly
/// cancelling amplitudes drive it there), and the step is taken from the
/// eigendecomposition instead, dropping eigenvalues below `10⁻¹²` of the
/// largest.
fn hermitian_solve<T: Amplitude>(a: &mut [T], b: &mut [T], n: usize) {
let (a0, b0) = (a.to_vec(), b.to_vec());
if cholesky_solve(a, b, n) {
return;
}
let (vals, vecs) = T::eigh(&a0, n);
let top = vals.iter().fold(0.0f64, |m, v| m.max(v.abs()));
for x in b.iter_mut() {
*x = T::ZERO;
}
for (i, &lambda) in vals.iter().enumerate() {
if lambda <= 1e-12 * top {
continue;
}
// (v_i† b₀ / λ_i) v_i
let mut c = T::ZERO;
for r in 0..n {
c = c.add(vecs[r * n + i].conj().mul(b0[r]));
}
let c = c.div_re(lambda);
for r in 0..n {
b[r] = b[r].add(vecs[r * n + i].mul(c));
}
}
}
/// `A x = b` by Cholesky (`A` overwritten by `L`, `A = L L†`; `b` by `x`).
/// Returns false, leaving both unusable, if a pivot is lost: not positive
/// or below rounding against its diagonal entry.
#[allow(clippy::needless_range_loop)]
fn cholesky_solve<T: Amplitude>(a: &mut [T], b: &mut [T], n: usize) {
fn cholesky_solve<T: Amplitude>(a: &mut [T], b: &mut [T], n: usize) -> bool {
for j in 0..n {
let mut d = a[j * n + j].re();
let diag = a[j * n + j].re();
let mut d = diag;
for k in 0..j {
d -= a[j * n + k].norm2();
}
if !d.is_finite() || d <= 1e-14 * diag.abs() {
return false;
}
let d = d.max(1e-300).sqrt();
a[j * n + j] = T::ONE.scale(d);
for i in j + 1..n {
@ -936,6 +1239,7 @@ fn cholesky_solve<T: Amplitude>(a: &mut [T], b: &mut [T], n: usize) {
}
b[i] = s.div_re(a[i * n + i].re());
}
b.iter().all(|x| x.re().is_finite() && x.norm2().is_finite())
}
/// A random configuration, at the given magnetisation if any.
@ -960,7 +1264,7 @@ fn start_configuration(n: usize, magnetisation: Option<i32>, rng: &mut Rng) -> V
/// accepted.
fn metropolis<T: Amplitude>(rbm: &Rbm<T>, s: &mut [i8], rng: &mut Rng, moves: usize, exchange: bool) -> usize {
let n = s.len();
let mut theta = rbm.theta(s);
let mut cache = rbm.cache(s);
let mut accepted = 0;
for _ in 0..moves {
let flips: Vec<usize> = if exchange {
@ -973,15 +1277,9 @@ fn metropolis<T: Amplitude>(rbm: &Rbm<T>, s: &mut [i8], rng: &mut Rng, moves: us
} else {
vec![rng.below(n)]
};
let r = rbm.ratio(s, &theta, &flips);
let r = rbm.ratio_cached(s, &cache, &flips);
if rng.uniform() < r.norm2() {
for &i in &flips {
let ds = -2.0 * f64::from(s[i]);
for (j, t) in theta.iter_mut().enumerate() {
*t = t.add(rbm.w[j * n + i].scale(ds));
}
s[i] = -s[i];
}
rbm.accept(s, &mut cache, &flips, r);
accepted += 1;
}
}
@ -1122,7 +1420,7 @@ mod tests {
fn a_complex_machine_with_real_parameters_is_the_real_one() {
let n = 6;
let real = Rbm::new(n, 2, 0.4, 5);
let cplx = Rbm { n: real.n, m: real.m, a: real.a.iter().map(|&re| C { re, im: 0.0 }).collect(), b: real.b.iter().map(|&re| C { re, im: 0.0 }).collect(), w: real.w.iter().map(|&re| C { re, im: 0.0 }).collect(), sublattice: None };
let cplx = Rbm { n: real.n, m: real.m, a: real.a.iter().map(|&re| C { re, im: 0.0 }).collect(), b: real.b.iter().map(|&re| C { re, im: 0.0 }).collect(), w: real.w.iter().map(|&re| C { re, im: 0.0 }).collect(), sublattice: None, symmetry: None };
let mut rng = Rng(8);
for _ in 0..10 {
let s = start_configuration(n, None, &mut rng);
@ -1225,4 +1523,113 @@ mod tests {
assert!((ec.energy - exact) / exact.abs() < 2e-3, "complex: {} vs {exact}", ec.energy);
}
}
/// `Ψ(s)` by summing over the group.
fn projected<T: Amplitude>(rbm: &Rbm<T>, s: &[i8]) -> T {
let sym = rbm.symmetry.as_ref().unwrap();
let mut total = T::ZERO;
for g in 0..sym.len() {
let t: Vec<i8> = sym.perms[g].iter().map(|&j| if sym.flips[g] { -s[j] } else { s[j] }).collect();
total = total.add(rbm.log_psi(&t).exp().scale(sym.chars[g] * rbm.sign(&t)));
}
total
}
#[test]
fn the_projection_is_the_group_sum() {
let n = 6;
for odd in [false, true] {
let rbm = Rbm::complex(n, 1, 0.3, 4).with_symmetry(Symmetry::ring(n, true).with_spin_flip(odd));
let mut rng = Rng(3);
for _ in 0..12 {
let s = start_configuration(n, Some(0), &mut rng);
let base = projected(&rbm, &s);
if base.norm2() < 1e-20 {
continue;
}
let cache = rbm.cache(&s);
for flips in [vec![0usize, 1], vec![2, 5], vec![3]] {
let mut t = s.clone();
for &i in &flips {
t[i] = -t[i];
}
let want = projected(&rbm, &t).div(base);
let got = rbm.ratio_cached(&s, &cache, &flips);
assert!(got.sub(want).norm2().sqrt() < 1e-10 * want.norm2().sqrt().max(1.0), "{got:?} vs {want:?}");
}
// The sector: a translation by one site multiplies by −1
// (momentum π), a global flip by the parity.
let shifted: Vec<i8> = (0..n).map(|i| s[(i + 1) % n]).collect();
let flipped: Vec<i8> = s.iter().map(|&x| -x).collect();
assert!(projected(&rbm, &shifted).add(base).norm2().sqrt() < 1e-10 * base.norm2().sqrt());
let sign = if odd { -1.0 } else { 1.0 };
assert!(projected(&rbm, &flipped).sub(base.scale(sign)).norm2().sqrt() < 1e-10 * base.norm2().sqrt());
// Derivatives of ln Ψ by central differences along a real step.
let p = rbm.params();
let mut d = vec![C::ZERO; p];
rbm.derivatives_cached(&s, &cache, &mut d);
let h = 1e-6;
for k in [0, n + 1, p - 1] {
let mut e = vec![C::ZERO; p];
e[k] = C { re: -h, im: 0.0 };
let mut up = rbm.clone();
up.shift(&e, 1.0);
e[k] = C { re: h, im: 0.0 };
let mut dn = rbm.clone();
dn.shift(&e, 1.0);
let fd = projected(&up, &s).sub(projected(&dn, &s)).div(base).scale(0.5 / h);
assert!(fd.sub(d[k]).norm2().sqrt() < 1e-6, "param {k}: {fd:?} vs {:?}", d[k]);
}
}
}
}
#[test]
fn accepted_moves_keep_the_cache_exact() {
let n = 8;
let rbm = Rbm::complex(n, 1, 0.3, 9).with_symmetry(Symmetry::ring(n, false).with_spin_flip(false));
let mut rng = Rng(21);
let mut s = start_configuration(n, Some(0), &mut rng);
let mut cache = rbm.cache(&s);
for _ in 0..200 {
let (i, j) = (rng.below(n), rng.below(n));
if s[i] == s[j] {
continue;
}
let r = rbm.ratio_cached(&s, &cache, &[i, j]);
rbm.accept(&mut s, &mut cache, &[i, j], r);
}
let fresh = rbm.cache(&s);
match (&cache, &fresh) {
(Cache::Projected { configs: c1, shares: p1, log_abs: l1, thetas: t1 }, Cache::Projected { configs: c2, shares: p2, log_abs: l2, thetas: t2 }) => {
assert_eq!(c1, c2);
assert!((l1 - l2).abs() < 1e-9, "{l1} vs {l2}");
for (a, b) in p1.iter().zip(p2) {
assert!(a.sub(*b).norm2().sqrt() < 1e-9, "{a:?} vs {b:?}");
}
for (a, b) in t1.iter().flatten().zip(t2.iter().flatten()) {
assert!(a.sub(*b).norm2().sqrt() < 1e-12);
}
}
_ => unreachable!(),
}
}
#[test]
fn projection_reaches_the_ground_state_in_its_sector() {
// A 6-site ring: half its length is odd, so the singlet ground state
// has momentum π and odd parity under a global flip.
let n = 6;
let model = SpinModel::heisenberg(n, &chain_edges(n, true), 1.0);
let exact = model.exact_ground_energy(Some(0)).unwrap();
let cfg = VmcConfig { exact: true, magnetisation: Some(0), lr: 0.1, shift: 1e-3, ..VmcConfig::default() };
let rbm = Rbm::complex(n, 1, 0.05, 3).with_symmetry(Symmetry::ring(n, true).with_spin_flip(true));
let mut v = Vmc::new(model, rbm, cfg);
let mut e = v.step();
for _ in 0..150 {
e = v.step();
}
assert!(e.energy >= exact - 1e-9, "variational: {} vs {exact}", e.energy);
assert!((e.energy - exact) / exact.abs() < 1e-8, "{} vs {exact}", e.energy);
}
}

14
wai/wai-quantum/tests/bit_identity.rs
View file

@ -555,6 +555,18 @@ fn neural_quantum_states() {
v.extend([e.energy, e.variance]);
}
v.extend(cplx.rbm.w.iter().flat_map(|c| [c.re, c.im]));
assert_eq!(digest(&v), 0x0ede_030d_b0c7_54dd);
// A projected machine on a 10-site ring (momentum π, odd), started near
// even: the metric goes singular and the eigendecomposition takes over.
let ring = SpinModel::heisenberg(10, &chain_edges(10, true), 1.0);
let rbm = Rbm::complex(10, 1, 0.01, 7).with_symmetry(Symmetry::ring(10, true).with_spin_flip(true));
let mut proj = Vmc::new(ring, rbm, VmcConfig { exact: true, magnetisation: Some(0), lr: 0.1, threads: 1, ..VmcConfig::default() });
for k in 0..60 {
let e = proj.step();
if k % 10 == 9 {
v.extend([e.energy, e.variance]);
}
}
v.extend(proj.rbm.b.iter().flat_map(|c| [c.re, c.im]));
assert_eq!(digest(&v), 0x834c_988e_b34e_814f);
}

2
wai/wai-rs/Cargo.lock generated
View file

@ -3900,7 +3900,7 @@ dependencies = [
[[package]]
name = "wai-quantum"
version = "0.3.37"
version = "0.3.38"
dependencies = [
"blake3",
"ed25519-dalek",

19
wai/wai-rs/src/bin/wai_quantum.rs
View file

@ -174,7 +174,8 @@ fn usage() -> ! {
wai-quantum chem --fcidump FILE [--bond D] [--sweeps S] [--penalty P] [--no-fci] [--json]\n\n\
NEURAL QUANTUM STATES (RBM, variational Monte Carlo, stochastic reconfiguration):\n \
wai-quantum nqs [--model heisenberg|j1j2|ising] [--lx L] [--ly L] [--j2 J] [--h H] [--periodic]\n \
[--complex] [--alpha A] [--steps N] [--samples S] [--lr R] [--exact] [--seed S] [--json]\n\n\
[--complex] [--project [--kpi] [--odd]] [--alpha A] [--steps N] [--samples S] [--lr R]\n \
[--exact] [--seed S] [--json]\n\n\
OPTIMISATION (simulated bifurcation, annealing, exact):\n \
wai-quantum qio (--gset FILE | --sk N [--instance S]) [--method dsb|bsb|hdsb|hbsb|sa|exact]\n \
[--steps N] [--trials T] [--dt D] [--c1 C] [--gamma G] [--beta0 B --beta1 B] [--descend]\n \
@ -1194,6 +1195,16 @@ fn main() {
let alpha: usize = num(args, "--alpha", "", 2);
let seed: u64 = num(args, "--seed", "", 7);
let reference = if n <= 20 { model.exact_ground_energy(sector) } else { None };
// Projection onto translations at momentum 0 or π, and a parity
// under the global flip; the lattice must be periodic.
let symmetry = flag(args, "--project").then(|| {
if !periodic {
die("--project needs --periodic");
}
let kpi = flag(args, "--kpi");
let base = if ly == 1 { Symmetry::ring(n, kpi) } else { Symmetry::torus(lx, ly, kpi, kpi) };
base.with_spin_flip(flag(args, "--odd"))
});
fn run<T: Amplitude>(model: SpinModel, rbm: Rbm<T>, cfg: VmcConfig, steps: usize) -> (usize, Estimate, Estimate) {
let params = rbm.params();
let mut vmc = Vmc::new(model, rbm, cfg);
@ -1208,12 +1219,18 @@ fn main() {
if marshall {
rbm = rbm.with_marshall_sign(checkerboard(lx, ly));
}
if let Some(sym) = symmetry {
rbm = rbm.with_symmetry(sym);
}
run(model, rbm, cfg, steps)
} else {
let mut rbm = Rbm::new(n, alpha, 0.01, seed);
if marshall {
rbm = rbm.with_marshall_sign(checkerboard(lx, ly));
}
if let Some(sym) = symmetry {
rbm = rbm.with_symmetry(sym);
}
run(model, rbm, cfg, steps)
};
if flag(args, "--json") {

2
wai/wai-web/Cargo.lock generated
View file

@ -734,7 +734,7 @@ dependencies = [
[[package]]
name = "wai-quantum"
version = "0.3.37"
version = "0.3.38"
dependencies = [
"blake3",
"ed25519-dalek",