Skip to main content

spectrafit_models/
gaussian_nd.rs

1//! N-dimensional axis-aligned Gaussian kernel.
2//!
3//! A single parametric kernel for any dimensionality `D`, carrying `d` as a
4//! field (set at construction from the node's explicit `n_dims`). Generalizes
5//! the validated 3-D spike over `0..D` axes. Axis-aligned (no rotation).
6
7use std::borrow::Cow;
8
9use crate::Model;
10
11/// Axis-aligned `D`-dimensional Gaussian peak (`n_dims() == d`).
12///
13/// `f(x) = A · exp( −Σ_{i=0}^{D-1} (x_i − center_i)² / (2·σ_i²) )`
14///
15/// Parameters (in order): `[amplitude, center_0 … center_{D-1}, sigma_0 … sigma_{D-1}]`
16/// — length `1 + 2D`. Indexed (not `x/y/z`) so the naming scales to arbitrary N.
17pub struct GaussianND {
18    /// Number of coordinate dimensions.
19    d: usize,
20}
21
22impl GaussianND {
23    /// Construct a `D`-dimensional Gaussian. `d` comes from the node's explicit
24    /// `n_dims` field (the compiler passes it in).
25    pub fn new(d: usize) -> Self {
26        GaussianND { d }
27    }
28}
29
30impl Model for GaussianND {
31    fn eval(&self, x: &[f64], params: &[f64]) -> f64 {
32        let d = self.d;
33        let a = params[0];
34        let mut z = 0.0_f64;
35        for i in 0..d {
36            let dx = x[i] - params[1 + i];
37            let s = params[1 + d + i];
38            z -= (dx * dx) / (2.0 * s * s);
39        }
40        a * z.exp()
41    }
42
43    fn jacobian(&self, x: &[f64], params: &[f64]) -> Vec<f64> {
44        let mut out = vec![0.0_f64; 1 + 2 * self.d];
45        self.jacobian_into(x, params, &mut out);
46        out
47    }
48
49    fn jacobian_into(&self, x: &[f64], params: &[f64], out: &mut [f64]) {
50        let d = self.d;
51        let a = params[0];
52        // Shared exponential factor g = exp(Σ …).
53        let mut z = 0.0_f64;
54        for i in 0..d {
55            let dx = x[i] - params[1 + i];
56            let s = params[1 + d + i];
57            z -= (dx * dx) / (2.0 * s * s);
58        }
59        let g = z.exp();
60
61        out[0] = g; // ∂/∂amplitude
62        for i in 0..d {
63            let dx = x[i] - params[1 + i];
64            let s = params[1 + d + i];
65            let s2 = s * s;
66            out[1 + i] = a * g * dx / s2; // ∂/∂center_i
67            out[1 + d + i] = a * g * dx * dx / (s2 * s); // ∂/∂sigma_i
68        }
69    }
70
71    fn param_names(&self) -> Vec<Cow<'static, str>> {
72        let mut names: Vec<Cow<'static, str>> = Vec::with_capacity(1 + 2 * self.d);
73        names.push(Cow::Borrowed("amplitude"));
74        for i in 0..self.d {
75            names.push(Cow::Owned(format!("center_{i}")));
76        }
77        for i in 0..self.d {
78            names.push(Cow::Owned(format!("sigma_{i}")));
79        }
80        names
81    }
82
83    fn n_dims(&self) -> usize {
84        self.d
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91    use approx::assert_relative_eq;
92
93    #[test]
94    fn param_names_indexed_and_sized() {
95        let m = GaussianND::new(3);
96        let names: Vec<String> = m.param_names().iter().map(|c| c.to_string()).collect();
97        assert_eq!(
98            names,
99            vec![
100                "amplitude",
101                "center_0",
102                "center_1",
103                "center_2",
104                "sigma_0",
105                "sigma_1",
106                "sigma_2"
107            ]
108        );
109        assert_eq!(m.n_dims(), 3);
110    }
111
112    #[test]
113    fn eval_at_center_equals_amplitude() {
114        let m = GaussianND::new(2);
115        // params: A, c0, c1, s0, s1
116        let v = m.eval(&[0.5, -1.0], &[3.0, 0.5, -1.0, 1.0, 2.0]);
117        assert_relative_eq!(v, 3.0, epsilon = 1e-12);
118    }
119
120    #[test]
121    fn jacobian_into_matches_finite_difference_5d() {
122        // Upgraded from a single-point forward difference (h=1e-6,
123        // epsilon=1e-4) to a central difference (O(h^2)) swept over three
124        // parameter regimes and several x-points relative to each regime's
125        // own centre.
126        let m = GaussianND::new(5);
127        // A, c0..c4, s0..s4
128        let param_sets = [
129            [2.0, 0.5, -1.0, 1.0, -0.5, 0.3, 1.0, 1.5, 0.9, 1.2, 0.8], // nominal
130            [1e-3, 0.0, 0.0, 0.0, 0.0, 0.0, 0.05, 0.05, 0.05, 0.05, 0.05], // small amplitude/width
131            [5.0, -2.0, 3.0, -1.5, 2.5, -0.5, 2.0, 3.0, 1.5, 2.5, 1.0], // offset centres, wide
132        ];
133        let offsets: [[f64; 5]; 3] = [
134            [0.0, 0.0, 0.0, 0.0, 0.0],    // at centre
135            [0.7, -0.3, 1.1, -0.8, 0.2],  // off-centre (original point)
136            [-2.5, 2.0, -1.5, 2.5, -2.0], // far off-centre / near a few sigma
137        ];
138        for params in param_sets {
139            let centers: [f64; 5] = std::array::from_fn(|i| params[1 + i]);
140            for off in offsets {
141                let x: [f64; 5] = std::array::from_fn(|i| centers[i] + off[i]);
142                let mut analytic = vec![0.0_f64; params.len()];
143                m.jacobian_into(&x, &params, &mut analytic);
144                for i in 0..params.len() {
145                    let h = 1e-6 * params[i].abs().max(1.0);
146                    let (mut a, mut b) = (params, params);
147                    a[i] += h;
148                    b[i] -= h;
149                    let fd = (m.eval(&x, &a) - m.eval(&x, &b)) / (2.0 * h);
150                    assert_relative_eq!(analytic[i], fd, epsilon = 1e-7, max_relative = 1e-6);
151                }
152            }
153        }
154    }
155}