1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
//! Chromatic block-Gibbs sampling.
//!
//! One sweep updates every node exactly once, color class by color class. Within a class, node
//! updates are conditionally independent (no two adjacent) — the parallelism a TSU exploits in
//! physics and a GPU exploits in threads; here the classes are simple loops, kept in the same
//! order so CPU, WebGPU, and device runs are cross-checkable draw for draw.
use crate::graph::Graph;
use crate::ledger::Ledger;
use crate::rng::Pcg;
// The logistic the module docs write as sigma(2 beta f_i). Only the conditional test
// exercises it, so it is dead in a non-test build — kept because it is the executable
// statement of the formula graph.rs and kernel.rs both cite in prose.
#[cfg_attr(not(test), allow(dead_code))]
#[inline]
fn sigma(x: f64) -> f64 {
1.0 / (1.0 + (-x).exp())
}
pub struct Sampler<'g> {
pub g: &'g Graph,
pub beta: f64,
pub s: Vec<i8>,
pub rng: Pcg,
/// Nodes whose value is held fixed (conditioning / "clamping"); sweeps skip them.
pub clamped: Vec<bool>,
/// Base seed for the parallel path's per-(sweep, class, chunk) RNG streams.
par_seed: u64,
/// Sweeps completed via the parallel path (advances its stream derivation).
par_sweeps: u64,
}
impl<'g> Sampler<'g> {
pub fn new(g: &'g Graph, beta: f64, seed: u64) -> Self {
let mut rng = Pcg::new(seed, 0x5EED);
let s = (0..g.n).map(|_| rng.spin(0.5)).collect();
Sampler { g, beta, s, rng, clamped: vec![false; g.n], par_seed: seed, par_sweeps: 0 }
}
/// Clamp node i to value v (observation / conditioning input).
pub fn clamp(&mut self, i: usize, v: i8) {
debug_assert!(v == 1 || v == -1);
self.s[i] = v;
self.clamped[i] = true;
}
pub fn unclamp(&mut self, i: usize) {
self.clamped[i] = false;
}
/// One full chromatic sweep (every free node updated once). If a ledger is given, it is
/// charged one Gibbs cycle per free node — the device-side price of this sweep.
pub fn sweep(&mut self, ledger: Option<&mut Ledger>) {
let mut updated = 0u64;
for class in &self.g.classes {
for &iu in class {
let i = iu as usize;
if self.clamped[i] {
continue;
}
let f = self.g.field(i, &self.s);
let p_up = crate::kernel::p_up(f, self.beta);
self.s[i] = self.rng.spin(p_up);
updated += 1;
}
}
if let Some(l) = ledger {
l.samples += updated;
}
}
/// Run `n` sweeps.
pub fn sweeps(&mut self, n: usize, mut ledger: Option<&mut Ledger>) {
for _ in 0..n {
self.sweep(ledger.as_deref_mut());
}
}
/// One full chromatic sweep across `threads` OS threads — the performance core.
///
/// Within a color class every node's conditional is independent (no two adjacent), so the
/// class is split into contiguous chunks, each updated by its own thread reading the shared
/// spin field and writing only its own chunk's nodes. Reads touch only OTHER-color nodes,
/// which no thread writes during this phase, so the access pattern is race-free by
/// construction of the coloring.
///
/// Determinism: each (sweep, class, chunk) gets its own counter-derived RNG stream, so the
/// result is bit-reproducible for a fixed (seed, threads). A different thread count is a
/// different, equally valid sample path (document the thread count next to the seed).
pub fn sweep_par(&mut self, threads: usize, ledger: Option<&mut Ledger>) {
assert!(threads >= 1);
let beta = self.beta;
let g = self.g;
let sweep_idx = self.par_sweeps;
let base = self.par_seed;
let mut updated = 0u64;
for (ci, class) in g.classes.iter().enumerate() {
let chunk = class.len().div_ceil(threads);
if chunk == 0 {
continue;
}
// SAFETY: chunks are disjoint index sets within one color class; every write target
// is unique to one thread, and every read is either a bias, an other-color neighbour
// (not written this phase), or the thread's own not-yet-updated node.
let sp = self.s.as_mut_ptr() as usize;
let clamped = &self.clamped;
std::thread::scope(|scope| {
for (ti, part) in class.chunks(chunk).enumerate() {
let part: &[u32] = part;
scope.spawn(move || {
let mut rng = Pcg::new(
base ^ sweep_idx.wrapping_mul(0x9E3779B97F4A7C15) ^ (ci as u64) << 32,
0xC0DE ^ ti as u64,
);
let s_ptr = sp as *mut i8;
for &iu in part {
let i = iu as usize;
if clamped[i] {
continue;
}
let mut f = g.h[i];
for k in g.offset[i]..g.offset[i + 1] {
f += g.w[k] * unsafe { *s_ptr.add(g.nbr[k] as usize) } as f64;
}
let p_up = crate::kernel::p_up(f, beta);
unsafe {
*s_ptr.add(i) = rng.spin(p_up);
}
}
});
}
});
updated += class.iter().filter(|&&iu| !self.clamped[iu as usize]).count() as u64;
}
self.par_sweeps += 1;
if let Some(l) = ledger {
l.samples += updated;
}
}
/// Run `n` parallel sweeps.
pub fn sweeps_par(&mut self, n: usize, threads: usize, mut ledger: Option<&mut Ledger>) {
for _ in 0..n {
self.sweep_par(threads, ledger.as_deref_mut());
}
}
/// Read the full state (device price: one read per node). Prefer [`Self::read_subset`]:
/// full-state readback is the crossings-tax regime.
pub fn read_all(&self, ledger: Option<&mut Ledger>) -> Vec<i8> {
if let Some(l) = ledger {
l.reads += self.g.n as u64;
}
self.s.clone()
}
/// Read only the named nodes (e.g. action bits).
pub fn read_subset(&self, idx: &[usize], ledger: Option<&mut Ledger>) -> Vec<i8> {
if let Some(l) = ledger {
l.reads += idx.len() as u64;
}
idx.iter().map(|&i| self.s[i]).collect()
}
}
#[cfg(test)]
mod tests {
use super::sigma;
use super::*;
use crate::graph::GraphBuilder;
/// The sampler's stationary distribution must match the exact Boltzmann distribution on an
/// enumerable system. 4-node cycle, mixed couplings and biases, TV < 0.02.
#[test]
fn matches_exact_boltzmann() {
let mut gb = GraphBuilder::new(4);
gb.couple(0, 1, 0.7);
gb.couple(1, 2, -0.4);
gb.couple(2, 3, 0.55);
gb.couple(3, 0, 0.3);
gb.bias(0, 0.2);
gb.bias(2, -0.35);
let g = gb.build();
let beta = 0.9;
// exact
let mut z = 0.0;
let mut p_exact = [0.0f64; 16];
for m in 0..16u32 {
let s: Vec<i8> = (0..4).map(|b| if m >> b & 1 == 1 { 1 } else { -1 }).collect();
let w = (-beta * g.energy(&s)).exp();
p_exact[m as usize] = w;
z += w;
}
for p in p_exact.iter_mut() {
*p /= z;
}
// sampled
let mut smp = Sampler::new(&g, beta, 0xC0FFEE);
smp.sweeps(200, None); // burn-in
let mut counts = [0u64; 16];
let n_samples = 200_000;
for _ in 0..n_samples {
smp.sweep(None);
let mut m = 0usize;
for b in 0..4 {
if smp.s[b] == 1 {
m |= 1 << b;
}
}
counts[m] += 1;
}
let tv: f64 = (0..16)
.map(|m| (counts[m] as f64 / n_samples as f64 - p_exact[m]).abs())
.sum::<f64>()
/ 2.0;
assert!(tv < 0.02, "TV distance to exact Boltzmann = {tv}");
}
/// The parallel path must satisfy the same physics standard as the sequential one: Onsager's
/// exact magnetization on the 2D lattice, and bit-reproducibility for fixed (seed, threads).
#[test]
fn parallel_sweep_physics_and_determinism() {
let g = crate::ising::lattice2d(48, 1.0);
let beta = 0.6;
let mut smp = Sampler::new(&g, beta, 0x9A7);
for s in smp.s.iter_mut() {
*s = 1;
}
smp.sweeps_par(2000, 8, None);
let mut acc = 0.0;
let reads = 2000;
for _ in 0..reads {
smp.sweep_par(8, None);
let m: i64 = smp.s.iter().map(|&v| v as i64).sum();
acc += (m as f64 / g.n as f64).abs();
}
let m = acc / reads as f64;
let exact = crate::ising::onsager_m(beta);
assert!((m - exact).abs() < 0.01, "parallel |M| {m:.4} vs Onsager {exact:.4}");
// determinism for fixed (seed, threads)
let mut a = Sampler::new(&g, beta, 0x1234);
let mut b = Sampler::new(&g, beta, 0x1234);
a.sweeps_par(50, 4, None);
b.sweeps_par(50, 4, None);
assert_eq!(a.s, b.s, "same (seed, threads) must reproduce bit-identically");
}
/// Clamped nodes must never change and must steer the conditional distribution.
#[test]
fn clamping_conditions() {
let mut gb = GraphBuilder::new(2);
gb.couple(0, 1, 1.5);
let g = gb.build();
let mut smp = Sampler::new(&g, 1.0, 7);
smp.clamp(0, 1);
let mut up = 0u64;
let n = 20_000;
for _ in 0..n {
smp.sweep(None);
assert_eq!(smp.s[0], 1);
if smp.s[1] == 1 {
up += 1;
}
}
// exact: P(s1=+1 | s0=+1) = sigma(2*beta*J) = sigma(3.0)
let want = sigma(3.0);
let got = up as f64 / n as f64;
assert!((got - want).abs() < 0.01, "got {got}, want {want}");
}
}