Skip to main content

spin_sim/statistics/
equilibration.rs

1/// Replica-averaged energy, link overlap and squared overlap over the log-binned
2/// window ending at `sweep`.
3///
4/// The window is the last `⌈sweep / 2⌉` sweeps, `(⌊sweep / 2⌋, sweep]` counting from
5/// one, i.e. `[2^(k-1), 2^k)` at the power-of-two checkpoints (Katzgraber, Palassini
6/// & Young, PRB 63, 184422 (2001)). Early non-equilibrium sweeps therefore drop out
7/// of later checkpoints instead of biasing a cumulative average.
8pub struct EquilCheckpoint {
9    pub sweep: usize,
10    pub energy_avg: Vec<f64>,
11    pub link_overlap_avg: Vec<f64>,
12    /// Pair-averaged q². Its plateau over the last windows is the usual
13    /// equilibration criterion where the energy identity does not apply (bimodal
14    /// couplings).
15    pub overlap2_avg: Vec<f64>,
16}
17
18/// The series a checkpoint averages: energy, link overlap, squared overlap.
19const N_SERIES: usize = 3;
20
21pub struct EquilDiagnosticAccum {
22    n_temps: usize,
23    checkpoints: Vec<usize>,
24    next_ckpt_idx: usize,
25    next_start_idx: usize,
26    count: usize,
27    sums: [Vec<f64>; N_SERIES],
28    /// Cumulative sums at the start of each opened window, indexed like `checkpoints`.
29    window_starts: Vec<[Vec<f64>; N_SERIES]>,
30    snapshots: Vec<EquilCheckpoint>,
31}
32
33impl EquilDiagnosticAccum {
34    pub fn new(n_temps: usize, n_sweeps: usize) -> Self {
35        let mut checkpoints = Vec::new();
36        let mut p = 128usize;
37        while p < n_sweeps {
38            checkpoints.push(p);
39            p *= 2;
40        }
41        if checkpoints.last() != Some(&n_sweeps) {
42            checkpoints.push(n_sweeps);
43        }
44
45        let mut accum = Self {
46            n_temps,
47            checkpoints,
48            next_ckpt_idx: 0,
49            next_start_idx: 0,
50            count: 0,
51            sums: std::array::from_fn(|_| vec![0.0; n_temps]),
52            window_starts: Vec::new(),
53            snapshots: Vec::new(),
54        };
55        accum.open_windows();
56        accum
57    }
58
59    fn window_start(checkpoint: usize) -> usize {
60        checkpoint / 2
61    }
62
63    fn open_windows(&mut self) {
64        while self
65            .checkpoints
66            .get(self.next_start_idx)
67            .is_some_and(|&c| Self::window_start(c) == self.count)
68        {
69            self.window_starts.push(self.sums.clone());
70            self.next_start_idx += 1;
71        }
72    }
73
74    /// Adds one sweep's replica-averaged energy, link overlap and squared overlap.
75    pub fn push(&mut self, energies: &[f32], link_overlaps: &[f32], overlaps2: &[f32]) {
76        self.count += 1;
77        for (sums, values) in self
78            .sums
79            .iter_mut()
80            .zip([energies, link_overlaps, overlaps2])
81        {
82            for (sum, &value) in sums.iter_mut().zip(&values[..self.n_temps]) {
83                *sum += value as f64;
84            }
85        }
86
87        while self.checkpoints.get(self.next_ckpt_idx) == Some(&self.count) {
88            let starts = &self.window_starts[self.next_ckpt_idx];
89            let width = (self.count - Self::window_start(self.count)) as f64;
90            let window_avg = |series: usize| -> Vec<f64> {
91                self.sums[series]
92                    .iter()
93                    .zip(&starts[series])
94                    .map(|(sum, start)| (sum - start) / width)
95                    .collect()
96            };
97            self.snapshots.push(EquilCheckpoint {
98                sweep: self.count,
99                energy_avg: window_avg(0),
100                link_overlap_avg: window_avg(1),
101                overlap2_avg: window_avg(2),
102            });
103            self.next_ckpt_idx += 1;
104        }
105        self.open_windows();
106    }
107
108    pub fn finish(self) -> Vec<EquilCheckpoint> {
109        self.snapshots
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use super::EquilDiagnosticAccum;
116
117    /// Energy equal to the 1-based sweep index; link overlap equal to its square.
118    fn run(n_sweeps: usize) -> Vec<(usize, f64, f64)> {
119        let mut accum = EquilDiagnosticAccum::new(1, n_sweeps);
120        for sweep in 1..=n_sweeps {
121            let value = sweep as f32;
122            accum.push(&[value], &[value * value], &[2.0 * value]);
123        }
124        accum
125            .finish()
126            .into_iter()
127            .map(|c| {
128                // The third series is twice the first, and so is its window mean.
129                assert_eq!(c.overlap2_avg[0], 2.0 * c.energy_avg[0]);
130                (c.sweep, c.energy_avg[0], c.link_overlap_avg[0])
131            })
132            .collect()
133    }
134
135    fn window_mean(lo: usize, hi: usize, f: impl Fn(f64) -> f64) -> f64 {
136        (lo..=hi).map(|s| f(s as f64)).sum::<f64>() / (hi - lo + 1) as f64
137    }
138
139    #[test]
140    fn checkpoints_average_the_last_half_window() {
141        let got = run(600);
142        let sweeps: Vec<usize> = got.iter().map(|c| c.0).collect();
143        assert_eq!(sweeps, vec![128, 256, 512, 600]);
144        for (sweep, energy, link) in got {
145            let lo = sweep / 2 + 1;
146            assert_eq!(energy, window_mean(lo, sweep, |s| s), "sweep {sweep}");
147            assert_eq!(link, window_mean(lo, sweep, |s| s * s), "sweep {sweep}");
148        }
149        // The window [65, 128] excludes early sweeps: mean 96.5, not the cumulative 64.5.
150        assert_eq!(run(128)[0].1, 96.5);
151    }
152
153    #[test]
154    fn short_and_overlapping_windows() {
155        assert_eq!(run(1), vec![(1, 1.0, 1.0)]);
156        assert_eq!(run(3), vec![(3, 2.5, 6.5)]);
157        // Windows (64, 128] and (64, 129] open at the same sweep.
158        let got = run(129);
159        assert_eq!(got[0], (128, 96.5, window_mean(65, 128, |s| s * s)));
160        assert_eq!(got[1], (129, 97.0, window_mean(65, 129, |s| s * s)));
161        assert!(run(0).is_empty());
162    }
163}