e7d4380eba
The full reference was computed first and then replaced by the sector's; past 20 spins that first step was a DMRG run (136 s on the 6 × 4 torus, and 3 % off the exact energy) whose answer was thrown away. Co-Authored-By: Claude Opus 5.5 <[email protected]>
149 lines
8 KiB
Rust
149 lines
8 KiB
Rust
//! A neural quantum state for a spin model's ground state, by variational
|
||
//! Monte Carlo with stochastic reconfiguration, against exact
|
||
//! diagonalisation (up to 20 spins) or DMRG on the same operator sum.
|
||
//!
|
||
//! ```text
|
||
//! 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] [--shift λ]
|
||
//! [--init W] [--exact] [--guide ε]
|
||
//! [--project [--kpi] [--point-group a1|a2|b1|b2 | --reflections a1|a2|b1|b2]
|
||
//! [--odd]] [--dmrg-bond D]
|
||
//! [--threads T] [--exact-final]
|
||
//!
|
||
//! `--project` projects onto translations of the periodic lattice at
|
||
//! momentum 0 (or π with `--kpi`), with `--point-group` on a square lattice
|
||
//! onto a representation of its rotations and reflections too, with
|
||
//! `--reflections` on any lattice onto one of its half turn and axis
|
||
//! reflections (`RectIrrep`), and onto even
|
||
//! (or `--odd`) parity under a global spin flip. Up to 20 spins, the
|
||
//! reference is then the exact ground energy in that sector. `--exact-final`
|
||
//! also evaluates the trained network with every configuration enumerated
|
||
//! (at most 24 spins), separating what training left from what sampling
|
||
//! adds. `--guide ε` samples `|ψ|²` with a floor `ε` under it, weighted
|
||
//! back (`VmcConfig::guide`); projected networks default to `0.1`, others to
|
||
//! none.
|
||
//! ```
|
||
use wai_quantum::quantum_dmrg::{dmrg_mpo, DmrgConfig, Mps, C};
|
||
use wai_quantum::quantum_nqs::*;
|
||
|
||
fn arg<T: std::str::FromStr>(a: &[String], key: &str, default: T) -> T {
|
||
a.iter().position(|x| x == key).and_then(|i| a.get(i + 1)).and_then(|v| v.parse().ok()).unwrap_or(default)
|
||
}
|
||
|
||
fn train<T: Amplitude>(model: SpinModel, rbm: Rbm<T>, cfg: VmcConfig, steps: usize, reference: f64, exact_final: bool) {
|
||
println!("RBM α={}: {} parameters, {} samples a step", rbm.m / rbm.n, rbm.params(), cfg.samples);
|
||
let mut vmc = Vmc::new(model, rbm, cfg);
|
||
let t0 = std::time::Instant::now();
|
||
for it in 1..=steps {
|
||
let e = vmc.step();
|
||
if it % (steps / 10).max(1) == 0 {
|
||
println!("step {it:>5} E {:.6} ± {:.6} variance {:.2e} acceptance {:.2} relative error {:.2e} {:.0}s", e.energy, e.error, e.variance, e.acceptance, (e.energy - reference) / reference.abs(), t0.elapsed().as_secs_f64());
|
||
}
|
||
}
|
||
vmc.cfg.samples *= 8;
|
||
let f = vmc.estimate();
|
||
println!("final E {:.6} ± {:.6} variance {:.2e} relative error {:.2e}", f.energy, f.error, f.variance, (f.energy - reference) / reference.abs());
|
||
// The trained network's own energy and variance, every configuration
|
||
// enumerated: what is left once sampling is taken out.
|
||
if exact_final {
|
||
vmc.cfg.exact = true;
|
||
let x = vmc.estimate();
|
||
println!("exact E {:.8} variance {:.2e} relative error {:.2e} (every configuration)", x.energy, x.variance, (x.energy - reference) / reference.abs());
|
||
}
|
||
}
|
||
|
||
fn main() {
|
||
let a: Vec<String> = std::env::args().collect();
|
||
let kind: String = arg(&a, "--model", "heisenberg".to_string());
|
||
let (lx, ly): (usize, usize) = (arg(&a, "--lx", 4), arg(&a, "--ly", 4));
|
||
let periodic = a.iter().any(|x| x == "--periodic");
|
||
let threads: usize = arg(&a, "--threads", VmcConfig::default().threads);
|
||
let n = lx * ly;
|
||
let nn = if ly == 1 { chain_edges(n, periodic) } else { square_edges(lx, ly, periodic) };
|
||
let nnn = if ly == 1 { chain_next_nearest(n, periodic) } else { square_diagonals(lx, ly, periodic) };
|
||
let model = match kind.as_str() {
|
||
"heisenberg" => SpinModel::heisenberg(n, &nn, 1.0),
|
||
"j1j2" => SpinModel::j1j2(n, &nn, &nnn, 1.0, arg(&a, "--j2", 0.5)),
|
||
"ising" => SpinModel::ising(n, &nn, 1.0, arg(&a, "--h", 3.044)),
|
||
other => panic!("{other}"),
|
||
};
|
||
let sector = if model.conserves_magnetisation() { Some(0) } else { None };
|
||
let point = a.iter().position(|x| x == "--point-group").map(|i| match a.get(i + 1).map(String::as_str) {
|
||
Some("a1") => PointIrrep::A1,
|
||
Some("a2") => PointIrrep::A2,
|
||
Some("b1") => PointIrrep::B1,
|
||
Some("b2") => PointIrrep::B2,
|
||
other => panic!("--point-group: a1, a2, b1 or b2, not {other:?}"),
|
||
});
|
||
let reflections = a.iter().position(|x| x == "--reflections").map(|i| match a.get(i + 1).map(String::as_str) {
|
||
Some("a1") => RectIrrep::A1,
|
||
Some("a2") => RectIrrep::A2,
|
||
Some("b1") => RectIrrep::B1,
|
||
Some("b2") => RectIrrep::B2,
|
||
other => panic!("--reflections: a1, a2, b1 or b2, not {other:?}"),
|
||
});
|
||
let symmetry = a.iter().any(|x| x == "--project").then(|| {
|
||
assert!(periodic, "--project needs --periodic");
|
||
let kpi = a.iter().any(|x| x == "--kpi");
|
||
let base = match (point, reflections) {
|
||
(Some(_), Some(_)) => panic!("--point-group or --reflections, not both"),
|
||
(Some(irrep), None) => {
|
||
assert_eq!(lx, ly, "--point-group needs a square lattice");
|
||
Symmetry::square(lx, kpi, irrep)
|
||
}
|
||
(None, Some(irrep)) => Symmetry::rectangle(lx, ly, kpi, kpi, irrep),
|
||
(None, None) if ly == 1 => Symmetry::ring(n, kpi),
|
||
(None, None) => Symmetry::torus(lx, ly, kpi, kpi),
|
||
};
|
||
base.with_spin_flip(a.iter().any(|x| x == "--odd"))
|
||
});
|
||
let t0 = std::time::Instant::now();
|
||
// A projected network is held to its own sector's ground energy.
|
||
let in_sector = match &symmetry {
|
||
Some(sym) if n <= 24 => Some(model.exact_ground_energy_in(sector, sym).expect("the model is invariant under the group, and the sector is not empty")),
|
||
_ => None,
|
||
};
|
||
let (reference, how) = match (in_sector, if n <= 20 && in_sector.is_none() { model.exact_ground_energy(sector) } else { None }) {
|
||
(Some(e), _) => (e, "exact diagonalisation in the projection's sector".to_string()),
|
||
(None, Some(e)) => (e, "exact diagonalisation".to_string()),
|
||
(None, None) => {
|
||
let bond = arg(&a, "--dmrg-bond", 128usize);
|
||
let r = dmrg_mpo(&model.opsum().mpo(), Mps::random(n, 8, 1), &DmrgConfig { max_bond: bond, sweeps: 6, threads, ..DmrgConfig::default() });
|
||
(r.energy, format!("DMRG, bond {bond}, discarded {:.1e}", r.discarded))
|
||
}
|
||
};
|
||
println!("{kind} {lx}×{ly}{}: reference {reference:.8} ({how}, {:.1}s)", if periodic { " periodic" } else { "" }, t0.elapsed().as_secs_f64());
|
||
let alpha: usize = arg(&a, "--alpha", 2);
|
||
let mut cfg = VmcConfig { samples: arg(&a, "--samples", 1000), chains: 16, magnetisation: sector, lr: arg(&a, "--lr", 0.02), shift: arg(&a, "--shift", 1e-3), exact: a.iter().any(|x| x == "--exact"), threads, ..VmcConfig::default() };
|
||
let steps: usize = arg(&a, "--steps", 600);
|
||
let marshall = a.iter().any(|x| x == "--marshall");
|
||
let exact_final = a.iter().any(|x| x == "--exact-final");
|
||
// Parameters start uniform in ±init. A projected network that starts
|
||
// near zero is nearly the projection of a constant, the ferromagnet: an
|
||
// eigenstate, which reconfiguration barely moves.
|
||
let init: f64 = arg(&a, "--init", if symmetry.is_some() { 0.1 } else { 0.01 });
|
||
// A projected network's amplitude is a sum that cancels near its nodes,
|
||
// where the local energy is large and |ψ|² rarely samples; a guide floor
|
||
// visits there, and training then corrects what it finds.
|
||
cfg.guide = arg(&a, "--guide", if symmetry.is_some() { 0.1 } else { 0.0 });
|
||
if a.iter().any(|x| x == "--complex") {
|
||
let mut rbm: Rbm<C> = Rbm::complex(n, alpha, init, 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, exact_final);
|
||
} else {
|
||
let mut rbm = Rbm::new(n, alpha, init, 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, exact_final);
|
||
}
|
||
}
|