|
| 1 | +//! `signature_pde_sweep` — the Chen-Lyons signature kernel via the Goursat |
| 2 | +//! PDE, on ndarray's canonical SIMD substrate. |
| 3 | +//! |
| 4 | +//! This is the real, general-dimension successor to `jc`'s |
| 5 | +//! `goursat_substrate_probe` example (lance-graph |
| 6 | +//! `TD-PILLAR11-SCIENTIFIC-LOOPS-BYPASS-NDARRAY-SIMD-1` / |
| 7 | +//! `.claude/knowledge/vertical-simd-consumer-contract.md` W1.5 item #6). The |
| 8 | +//! probe fixed path dimension = 2 as a throwaway falsifier; this module lifts |
| 9 | +//! the same wavefront to arbitrary dimension so it is a drop-in replacement |
| 10 | +//! for `sigker::signature_kernel_pde(x: &[Vec<f64>], y: &[Vec<f64>]) -> f64`. |
| 11 | +//! |
| 12 | +//! ## The recurrence |
| 13 | +//! |
| 14 | +//! For paths `x` (length `n`) and `y` (length `m`) in `R^dim`, the depth-∞ |
| 15 | +//! signature kernel solves |
| 16 | +//! |
| 17 | +//! ```text |
| 18 | +//! K[i+1][j+1] = K[i+1][j] + K[i][j+1] - K[i][j] + c_ij * K[i][j] |
| 19 | +//! c_ij = <x[i+1]-x[i], y[j+1]-y[j]> (Euclidean inner product) |
| 20 | +//! K[i][0] = K[0][j] = 1 (boundary) |
| 21 | +//! ``` |
| 22 | +//! |
| 23 | +//! in `O(n*m*dim)` flops, with no signature materialization (Hambly-Lyons |
| 24 | +//! 2010; see `jc::hambly_lyons` for the uniqueness certificate this kernel |
| 25 | +//! exists to serve). |
| 26 | +//! |
| 27 | +//! ## Why a wavefront, not row-major |
| 28 | +//! |
| 29 | +//! `K[i+1][j+1]` depends on three earlier cells, none of which lie on the |
| 30 | +//! same row — but every cell on the anti-diagonal `i+j = d` depends only on |
| 31 | +//! cells on diagonals `d-1` and `d-2`. So the solve sweeps diagonals with |
| 32 | +//! three rolling row-indexed buffers (`prev2`, `prev1`, `cur`) instead of a |
| 33 | +//! full `n*m` grid, and — because the interior of one diagonal has no |
| 34 | +//! cross-cell dependency — each diagonal's interior is computed |
| 35 | +//! [`LANES`]-wide via [`F64x8`]. |
| 36 | +//! |
| 37 | +//! `y`'s increments are stored **reversed per dimension** (`dyr`): the |
| 38 | +//! diagonal walk needs `dy[j-1]` for `j` decreasing as row `i` increases, |
| 39 | +//! and reversing turns that backward walk into a forward, contiguous read — |
| 40 | +//! no gather primitive is needed. This is an architectural property of the |
| 41 | +//! recurrence, not a hand-tuned trick: verified bit-for-bit against the |
| 42 | +//! dim=2 probe before this generalization (`E-OCR`-style falsifier, see the |
| 43 | +//! module tests below). |
| 44 | +//! |
| 45 | +//! ## Numerics |
| 46 | +//! |
| 47 | +//! The three-FMA body (`t = 1*left + up`, `u = -1*diag + t`, |
| 48 | +//! `new = c*diag + u`) fuses only the last multiply-add; `±1.0` multipliers |
| 49 | +//! round exactly like a plain add/subtract, so this differs from a |
| 50 | +//! non-fused scalar evaluation only in the rounding of `c*diag`, at the |
| 51 | +//! `f64` ULP level — matching the probe's predeclared A1<->A2 tolerance. |
| 52 | +
|
| 53 | +use crate::simd::F64x8; |
| 54 | + |
| 55 | +/// SIMD lane width for the interior-diagonal sweep (matches [`F64x8`]). |
| 56 | +const LANES: usize = 8; |
| 57 | + |
| 58 | +/// Per-dimension increments, one contiguous `Vec<f64>` per coordinate axis: |
| 59 | +/// `out[a][i] = path[i+1][a] - path[i][a]`. Storing per-axis (rather than |
| 60 | +/// interleaved) is what lets the SIMD sweep read each axis as a plain |
| 61 | +/// contiguous slice. |
| 62 | +fn increments_soa(path: &[Vec<f64>], dim: usize) -> Vec<Vec<f64>> { |
| 63 | + let mut out = vec![Vec::with_capacity(path.len().saturating_sub(1)); dim]; |
| 64 | + for w in path.windows(2) { |
| 65 | + for (axis, lane) in out.iter_mut().enumerate() { |
| 66 | + lane.push(w[1][axis] - w[0][axis]); |
| 67 | + } |
| 68 | + } |
| 69 | + out |
| 70 | +} |
| 71 | + |
| 72 | +/// Signature kernel `<S(x), S(y)>` via the depth-infinity Goursat PDE. |
| 73 | +/// |
| 74 | +/// Drop-in for `sigker::signature_kernel_pde` — identical signature, any |
| 75 | +/// path dimension (`x[0].len()`), any (possibly unequal) path lengths. |
| 76 | +/// |
| 77 | +/// # Panics |
| 78 | +/// |
| 79 | +/// Panics (debug only) if `x` and `y` disagree on coordinate dimension. |
| 80 | +/// Panics (always) if either path is empty — a path needs at least one |
| 81 | +/// point. |
| 82 | +/// |
| 83 | +/// # Examples |
| 84 | +/// |
| 85 | +/// ``` |
| 86 | +/// use ndarray::hpc::signature_pde::signature_pde_sweep; |
| 87 | +/// // A single-point path has no increments; the kernel is the empty-word 1. |
| 88 | +/// let k = signature_pde_sweep(&[vec![0.0, 0.0]], &[vec![1.0, 1.0]]); |
| 89 | +/// assert!((k - 1.0).abs() < 1e-12); |
| 90 | +/// ``` |
| 91 | +pub fn signature_pde_sweep(x: &[Vec<f64>], y: &[Vec<f64>]) -> f64 { |
| 92 | + let (n, m) = (x.len(), y.len()); |
| 93 | + assert!(n >= 1 && m >= 1, "signature_pde_sweep: paths must have at least one point"); |
| 94 | + let dim = x[0].len(); |
| 95 | + debug_assert_eq!(dim, y[0].len(), "signature_pde_sweep: x and y must share coordinate dimension"); |
| 96 | + |
| 97 | + let dx = increments_soa(x, dim); |
| 98 | + let mut dyr = increments_soa(y, dim); |
| 99 | + for lane in &mut dyr { |
| 100 | + lane.reverse(); |
| 101 | + } |
| 102 | + |
| 103 | + let mut prev2 = vec![1.0f64; n]; |
| 104 | + let mut prev1 = vec![1.0f64; n]; |
| 105 | + let mut cur = vec![1.0f64; n]; |
| 106 | + let (one, neg_one, zero) = (F64x8::splat(1.0), F64x8::splat(-1.0), F64x8::splat(0.0)); |
| 107 | + let mut lane_out = [0.0f64; LANES]; |
| 108 | + |
| 109 | + for d in 2..(n + m - 1) { |
| 110 | + // Boundary of this diagonal: k[0][d] and k[d][0] are always 1. |
| 111 | + if d < m { |
| 112 | + cur[0] = 1.0; |
| 113 | + } |
| 114 | + if d < n { |
| 115 | + cur[d] = 1.0; |
| 116 | + } |
| 117 | + // Interior rows: i >= 1, j = d - i >= 1, i <= n-1, j <= m-1. |
| 118 | + let lo = 1usize.max(d.saturating_sub(m - 1)); |
| 119 | + let hi = (d - 1).min(n - 1); |
| 120 | + if lo > hi { |
| 121 | + std::mem::swap(&mut prev2, &mut prev1); |
| 122 | + std::mem::swap(&mut prev1, &mut cur); |
| 123 | + continue; |
| 124 | + } |
| 125 | + // dyr index for row i is (m-1-d)+i: transiently negative in isize |
| 126 | + // before adding i, always in-range once i is in the interior band. |
| 127 | + let base = (m as isize) - 1 - (d as isize); |
| 128 | + let mut i = lo; |
| 129 | + while i + LANES <= hi + 1 { |
| 130 | + let left = F64x8::from_slice(&prev1[i..i + LANES]); |
| 131 | + let up = F64x8::from_slice(&prev1[i - 1..i - 1 + LANES]); |
| 132 | + let diag = F64x8::from_slice(&prev2[i - 1..i - 1 + LANES]); |
| 133 | + let r = (base + i as isize) as usize; |
| 134 | + debug_assert!(r + LANES <= m - 1, "signature_pde_sweep: dyr SIMD window out of range"); |
| 135 | + let mut c = zero; |
| 136 | + for a in 0..dim { |
| 137 | + let av = F64x8::from_slice(&dx[a][i - 1..i - 1 + LANES]); |
| 138 | + let bv = F64x8::from_slice(&dyr[a][r..r + LANES]); |
| 139 | + c = av.mul_add(bv, c); |
| 140 | + } |
| 141 | + let t = one.mul_add(left, up); |
| 142 | + let u = neg_one.mul_add(diag, t); |
| 143 | + c.mul_add(diag, u).copy_to_slice(&mut lane_out); |
| 144 | + cur[i..i + LANES].copy_from_slice(&lane_out); |
| 145 | + i += LANES; |
| 146 | + } |
| 147 | + // Scalar tail: same three-FMA arithmetic, so A2 stays internally uniform. |
| 148 | + while i <= hi { |
| 149 | + let r = (base + i as isize) as usize; |
| 150 | + debug_assert!(r < m - 1, "signature_pde_sweep: dyr scalar index out of range"); |
| 151 | + let mut c = 0.0f64; |
| 152 | + for a in 0..dim { |
| 153 | + c = dx[a][i - 1].mul_add(dyr[a][r], c); |
| 154 | + } |
| 155 | + let t = 1.0f64.mul_add(prev1[i], prev1[i - 1]); |
| 156 | + let u = (-1.0f64).mul_add(prev2[i - 1], t); |
| 157 | + cur[i] = c.mul_add(prev2[i - 1], u); |
| 158 | + i += 1; |
| 159 | + } |
| 160 | + std::mem::swap(&mut prev2, &mut prev1); |
| 161 | + std::mem::swap(&mut prev1, &mut cur); |
| 162 | + } |
| 163 | + prev1[n - 1] |
| 164 | +} |
| 165 | + |
| 166 | +#[cfg(test)] |
| 167 | +mod tests { |
| 168 | + use super::*; |
| 169 | + |
| 170 | + /// Test-only scalar reference: the shipped recurrence, row-major, no |
| 171 | + /// SIMD, arbitrary dimension. Deliberately NOT shared code with |
| 172 | + /// `signature_pde_sweep` — this is the independent oracle the parity |
| 173 | + /// tests check against, mirroring `sigker::signature_kernel_pde`'s own |
| 174 | + /// row-major evaluation without depending on that crate (ndarray must |
| 175 | + /// not depend on sigker). |
| 176 | + fn goursat_reference(x: &[Vec<f64>], y: &[Vec<f64>]) -> f64 { |
| 177 | + let (n, m) = (x.len(), y.len()); |
| 178 | + let dim = x[0].len(); |
| 179 | + let mut k = vec![1.0f64; n * m]; |
| 180 | + for i in 0..n.saturating_sub(1) { |
| 181 | + for j in 0..m.saturating_sub(1) { |
| 182 | + let c: f64 = (0..dim) |
| 183 | + .map(|a| (x[i + 1][a] - x[i][a]) * (y[j + 1][a] - y[j][a])) |
| 184 | + .sum(); |
| 185 | + let (left, up, diag) = (k[(i + 1) * m + j], k[i * m + j + 1], k[i * m + j]); |
| 186 | + k[(i + 1) * m + j + 1] = left + up - diag + c * diag; |
| 187 | + } |
| 188 | + } |
| 189 | + k[n * m - 1] |
| 190 | + } |
| 191 | + |
| 192 | + fn wiggly_path(n: usize, dim: usize, seed: f64) -> Vec<Vec<f64>> { |
| 193 | + (0..=n) |
| 194 | + .map(|i| { |
| 195 | + let t = i as f64 / n.max(1) as f64; |
| 196 | + (0..dim) |
| 197 | + .map(|a| { |
| 198 | + let phase = seed + a as f64 * 1.7; |
| 199 | + t * (a as f64 + 1.0) + 0.05 * (37.0 * t + phase).cos() |
| 200 | + }) |
| 201 | + .collect() |
| 202 | + }) |
| 203 | + .collect() |
| 204 | + } |
| 205 | + |
| 206 | + fn assert_matches_reference(x: &[Vec<f64>], y: &[Vec<f64>]) { |
| 207 | + let expected = goursat_reference(x, y); |
| 208 | + let actual = signature_pde_sweep(x, y); |
| 209 | + let tol = 1e-9 * expected.abs().max(1.0); |
| 210 | + assert!( |
| 211 | + (expected - actual).abs() <= tol, |
| 212 | + "signature_pde_sweep mismatch: reference={expected:e} actual={actual:e} \ |
| 213 | + (n={}, m={}, dim={})", |
| 214 | + x.len(), |
| 215 | + y.len(), |
| 216 | + x[0].len() |
| 217 | + ); |
| 218 | + } |
| 219 | + |
| 220 | + #[test] |
| 221 | + fn parity_across_dimensions() { |
| 222 | + for &dim in &[1usize, 2, 3, 5] { |
| 223 | + let x = wiggly_path(37, dim, 0.3); |
| 224 | + let y = wiggly_path(37, dim, 1.1); |
| 225 | + assert_matches_reference(&x, &y); |
| 226 | + } |
| 227 | + } |
| 228 | + |
| 229 | + #[test] |
| 230 | + fn parity_rectangular_grids() { |
| 231 | + // n != m, both well past one SIMD lane, neither a multiple of LANES. |
| 232 | + let x = wiggly_path(50, 3, 0.2); |
| 233 | + let y = wiggly_path(23, 3, 0.9); |
| 234 | + assert_matches_reference(&x, &y); |
| 235 | + let x = wiggly_path(23, 3, 0.2); |
| 236 | + let y = wiggly_path(50, 3, 0.9); |
| 237 | + assert_matches_reference(&x, &y); |
| 238 | + } |
| 239 | + |
| 240 | + #[test] |
| 241 | + fn parity_lengths_not_multiple_of_lanes() { |
| 242 | + for &n in &[2usize, 3, 9, 15, 17, 33, 65] { |
| 243 | + let x = wiggly_path(n, 2, 0.4); |
| 244 | + let y = wiggly_path(n + 3, 2, 0.6); |
| 245 | + assert_matches_reference(&x, &y); |
| 246 | + } |
| 247 | + } |
| 248 | + |
| 249 | + #[test] |
| 250 | + fn degenerate_single_point_path_is_one() { |
| 251 | + let x = vec![vec![0.3, -1.2, 4.0]]; |
| 252 | + let y = wiggly_path(10, 3, 0.5); |
| 253 | + assert!((signature_pde_sweep(&x, &y) - 1.0).abs() < 1e-12); |
| 254 | + assert!((signature_pde_sweep(&y, &x) - 1.0).abs() < 1e-12); |
| 255 | + let both_single = vec![vec![1.0, 2.0]]; |
| 256 | + assert!((signature_pde_sweep(&both_single, &both_single) - 1.0).abs() < 1e-12); |
| 257 | + } |
| 258 | + |
| 259 | + #[test] |
| 260 | + fn zero_increment_path_keeps_kernel_at_one() { |
| 261 | + // A constant path has zero increments everywhere; c_ij = 0 for all |
| 262 | + // (i, j), and by induction on the boundary K[i][0] = K[0][j] = 1, |
| 263 | + // the whole grid stays 1. This is a real invariant of the |
| 264 | + // recurrence (verified by hand, not asserted from the code under |
| 265 | + // test), so it discriminates a broken sweep from a correct one — |
| 266 | + // not a vacuous "the function returns *something*" check. |
| 267 | + let x: Vec<Vec<f64>> = (0..20).map(|_| vec![7.0, -3.0]).collect(); |
| 268 | + let y: Vec<Vec<f64>> = (0..15).map(|_| vec![7.0, -3.0]).collect(); |
| 269 | + let k = signature_pde_sweep(&x, &y); |
| 270 | + assert!((k - 1.0).abs() < 1e-12, "expected K == 1 everywhere, got {k}"); |
| 271 | + } |
| 272 | +} |
0 commit comments