flow_pacmap/
adam.rs

1//! Adam optimizer step for the 2-D embedding.
2//!
3//! All state operates on `Vec<[f32; 2]>` — no matrix types, no large allocations.
4//! Rayon `par_chunks_mut` parallelises the per-point update.
5
6use rayon::prelude::*;
7
8/// Persistent Adam state buffers. Allocated once, live for the full run.
9pub struct AdamState {
10    pub m: Vec<[f32; 2]>, // first moment
11    pub v: Vec<[f32; 2]>, // second moment
12}
13
14impl AdamState {
15    pub fn new(n: usize) -> Self {
16        Self {
17            m: vec![[0.0; 2]; n],
18            v: vec![[0.0; 2]; n],
19        }
20    }
21}
22
23/// Apply one Adam step: update `embedding` in place.
24///
25/// `t` is the 1-indexed iteration number (used for bias correction).
26pub fn adam_step(
27    embedding: &mut [[f32; 2]],
28    grad: &[[f32; 2]],
29    state: &mut AdamState,
30    t: usize,
31    lr: f32,
32) {
33    let beta1 = 0.9_f32;
34    let beta2 = 0.999_f32;
35    let eps = 1e-7_f32;
36    let t = t as f32;
37
38    // Bias-correction scalars (computed once per iteration)
39    let bc1 = 1.0 - beta1.powf(t);
40    let bc2 = 1.0 - beta2.powf(t);
41    let lr_t = lr * bc2.sqrt() / bc1;
42
43    // Parallel per-point update
44    embedding
45        .par_chunks_mut(4096)
46        .zip(grad.par_chunks(4096))
47        .zip(state.m.par_chunks_mut(4096))
48        .zip(state.v.par_chunks_mut(4096))
49        .for_each(|(((y_chunk, g_chunk), m_chunk), v_chunk)| {
50            for (((y, g), m), v) in y_chunk
51                .iter_mut()
52                .zip(g_chunk)
53                .zip(m_chunk.iter_mut())
54                .zip(v_chunk.iter_mut())
55            {
56                for dim in 0..2 {
57                    m[dim] = beta1 * m[dim] + (1.0 - beta1) * g[dim];
58                    v[dim] = beta2 * v[dim] + (1.0 - beta2) * g[dim] * g[dim];
59                    y[dim] -= lr_t * m[dim] / (v[dim].sqrt() + eps);
60                }
61            }
62        });
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68
69    #[test]
70    fn adam_moves_in_gradient_direction() {
71        let n = 4;
72        let mut emb = vec![[0.0_f32; 2]; n];
73        let grad = vec![[1.0_f32, -1.0]; n]; // constant gradient
74        let mut state = AdamState::new(n);
75
76        adam_step(&mut emb, &grad, &mut state, 1, 1.0);
77
78        // After one step, each point should have moved in the -gradient direction
79        for y in &emb {
80            assert!(y[0] < 0.0, "should move in -x direction");
81            assert!(y[1] > 0.0, "should move in +y direction");
82        }
83    }
84}
flow_pacmap/adam.rs

flow_pacmap/
adam.rs
