gam_solve/
measure_jet_gram_cache.rs

1//! Sufficient-statistic caches for #1033 mechanism (a), the measure-jet
2//! fixed-design case.
3//!
4//! This module is for single-scale-mode measure jets where `dX/dpsi == 0`: the
5//! design matrix `X` is theta-invariant across the lambda/rho outer loop, while
6//! the penalty and, for GLM PIRLS, the scalar working-weight diagonal `W` may
7//! change. It is distinct from `GaussianFixedCache`, which covers only the
8//! Gaussian+identity lane with constant `W`, and from `PsiGramTensor`, which
9//! covers design-moving psi via Chebyshev expansions, #1033 mechanism (b).
10//!
11//! Invariant: n-row work from the measure-jet basis builder happens once per fit
12//! at construction. Gaussian constant-`W` accessors are O(p^3) or cheaper and do
13//! not re-touch the n design rows. The GLM changing-`W` lane keeps the fixed
14//! rows cached and performs only the irreducible weighted contractions needed
15//! when PIRLS weights move.
16
17use gam_linalg::faer_ndarray::{fast_xt_diag_x, fast_xt_diag_y};
18use ndarray::{Array1, Array2, ArrayView1, ArrayView2};
19
20/// Gaussian / constant-`W` sufficient statistics for a fixed design.
21///
22/// This stores `X'WX`, `X'W(y - offset)`, and `(y - offset)'W(y - offset)` so
23/// per-lambda assembly and RSS/evidence terms are n-free. It generalizes the
24/// constant-design idea beyond the existing Gaussian+identity-only cache while
25/// keeping the same fixed-`W` requirement for this lane.
26pub struct FixedDesignGramCache {
27    xtwx: Array2<f64>,
28    xtwy: Array1<f64>,
29    ywy: f64,
30    n: usize,
31    p: usize,
32}
33
34impl FixedDesignGramCache {
35    /// Build fixed-design Gaussian sufficient statistics.
36    ///
37    /// The right-hand side is routed through `fast_xt_diag_y`, the same weighted
38    /// contraction primitive used by the runtime recompute path.
39    pub fn build(
40        x: ArrayView2<'_, f64>,
41        y: ArrayView1<'_, f64>,
42        offset: Option<ArrayView1<'_, f64>>,
43        weights: Option<ArrayView1<'_, f64>>,
44    ) -> Result<Self, String> {
45        let n = x.nrows();
46        let p = x.ncols();
47        if y.len() != n {
48            return Err(format!(
49                "y length {} must match design row count {}",
50                y.len(),
51                n
52            ));
53        }
54        if let Some(offset_values) = offset {
55            if offset_values.len() != n {
56                return Err(format!(
57                    "offset length {} must match design row count {}",
58                    offset_values.len(),
59                    n
60                ));
61            }
62        }
63        if let Some(weight_values) = weights {
64            if weight_values.len() != n {
65                return Err(format!(
66                    "weights length {} must match design row count {}",
67                    weight_values.len(),
68                    n
69                ));
70            }
71            validate_nonnegative_finite_weights(weight_values)?;
72        }
73        validate_finite_vector("y", y)?;
74        if let Some(offset_values) = offset {
75            validate_finite_vector("offset", offset_values)?;
76        }
77        validate_finite_matrix("x", x)?;
78
79        let r = match offset {
80            Some(offset_values) => &y.to_owned() - &offset_values.to_owned(),
81            None => y.to_owned(),
82        };
83        let w = match weights {
84            Some(weight_values) => weight_values.to_owned(),
85            None => Array1::ones(n),
86        };
87        let x_owned = x.to_owned();
88        let xtwx = fast_xt_diag_x(&x_owned, &w);
89        let r2 = r.view().insert_axis(ndarray::Axis(1));
90        let xtwy_mat = fast_xt_diag_y(&x_owned, &w, &r2);
91        let xtwy = xtwy_mat.column(0).to_owned();
92        let ywy = weighted_sum_squares(w.view(), r.view());
93
94        Ok(Self {
95            xtwx,
96            xtwy,
97            ywy,
98            n,
99            p,
100        })
101    }
102
103    pub fn n(&self) -> usize {
104        self.n
105    }
106
107    pub fn p(&self) -> usize {
108        self.p
109    }
110
111    pub fn xtwx(&self) -> ArrayView2<'_, f64> {
112        self.xtwx.view()
113    }
114
115    pub fn xtwy(&self) -> ArrayView1<'_, f64> {
116        self.xtwy.view()
117    }
118
119    pub fn ywy(&self) -> f64 {
120        self.ywy
121    }
122
123}
124
125/// Cached fixed design rows for GLM / changing-`W` PIRLS trials.
126///
127/// This cache owns the theta-invariant `X` rows once. Each trial recomputes
128/// `X'WX` and `X'Wz` because the scalar working weights and working response
129/// genuinely move during PIRLS. The Gaussian constant-Gram trick does not apply
130/// when `W` changes; the saved work is the expensive measure-jet basis/design
131/// construction, not the unavoidable weighted contraction over fixed rows.
132pub struct FixedDesignRowCache {
133    x: Array2<f64>,
134    n: usize,
135    p: usize,
136}
137
138impl FixedDesignRowCache {
139    /// Cache a finite, non-empty fixed design.
140    pub fn build(x: ArrayView2<'_, f64>) -> Result<Self, String> {
141        if x.nrows() == 0 || x.ncols() == 0 {
142            return Err(format!(
143                "design must be non-empty, got shape {}x{}",
144                x.nrows(),
145                x.ncols()
146            ));
147        }
148        validate_finite_matrix("x", x)?;
149        let n = x.nrows();
150        let p = x.ncols();
151        Ok(Self {
152            x: x.to_owned(),
153            n,
154            p,
155        })
156    }
157
158    pub fn n(&self) -> usize {
159        self.n
160    }
161
162    pub fn p(&self) -> usize {
163        self.p
164    }
165
166    pub fn design(&self) -> ArrayView2<'_, f64> {
167        self.x.view()
168    }
169
170    /// Recompute `X' diag(weights) X` over cached rows.
171    ///
172    /// This remains O(n p^2), the irreducible weighted contraction when `W`
173    /// changes. It avoids rebuilding the n-row measure-jet design.
174    pub fn xtwx(&self, weights: ArrayView1<'_, f64>) -> Result<Array2<f64>, String> {
175        self.validate_changing_weights(weights)?;
176        Ok(fast_xt_diag_x(&self.x, &weights))
177    }
178
179    fn validate_changing_weights(&self, weights: ArrayView1<'_, f64>) -> Result<(), String> {
180        if weights.len() != self.n {
181            return Err(format!(
182                "weights length {} must match design row count {}",
183                weights.len(),
184                self.n
185            ));
186        }
187        validate_finite_vector("weights", weights)
188    }
189}
190
191fn validate_finite_matrix(name: &str, matrix: ArrayView2<'_, f64>) -> Result<(), String> {
192    for ((row, col), value) in matrix.indexed_iter() {
193        if !(*value).is_finite() {
194            return Err(format!("{name}[{row},{col}] must be finite"));
195        }
196    }
197    Ok(())
198}
199
200fn validate_finite_vector(name: &str, vector: ArrayView1<'_, f64>) -> Result<(), String> {
201    for (index, value) in vector.iter().enumerate() {
202        if !(*value).is_finite() {
203            return Err(format!("{name}[{index}] must be finite"));
204        }
205    }
206    Ok(())
207}
208
209fn validate_nonnegative_finite_weights(weights: ArrayView1<'_, f64>) -> Result<(), String> {
210    for (index, weight) in weights.iter().enumerate() {
211        if !(*weight).is_finite() {
212            return Err(format!("weights[{index}] must be finite"));
213        }
214        if *weight < 0.0 {
215            return Err(format!("weights[{index}] must be non-negative"));
216        }
217    }
218    Ok(())
219}
220
221fn weighted_sum_squares(weights: ArrayView1<'_, f64>, values: ArrayView1<'_, f64>) -> f64 {
222    weights
223        .iter()
224        .zip(values.iter())
225        .map(|(weight, value)| *weight * *value * *value)
226        .sum()
227}
228
229#[cfg(test)]
230mod tests {
231    use super::{FixedDesignGramCache, FixedDesignRowCache};
232    use approx::assert_abs_diff_eq;
233    use gam_linalg::faer_ndarray::fast_xt_diag_x;
234    use ndarray::{Array1, Array2};
235
236    fn deterministic_design(n: usize, p: usize) -> Array2<f64> {
237        Array2::from_shape_fn((n, p), |(i, j)| {
238            let row = i as f64 + 1.0;
239            let col = j as f64 + 1.0;
240            ((row * 0.17 + col * 0.31).sin()) + row * col * 0.002
241        })
242    }
243
244    fn deterministic_response(n: usize) -> Array1<f64> {
245        Array1::from_shape_fn(n, |i| {
246            let row = i as f64 + 1.0;
247            (row * 0.23).cos() + row * 0.015
248        })
249    }
250
251    fn deterministic_offset(n: usize) -> Array1<f64> {
252        Array1::from_shape_fn(n, |i| {
253            let row = i as f64 + 1.0;
254            0.2 * (row * 0.11).sin() - 0.01 * row
255        })
256    }
257
258    fn deterministic_weights(n: usize, scale: f64) -> Array1<f64> {
259        Array1::from_shape_fn(n, |i| {
260            let row = i as f64 + 1.0;
261            0.4 + scale * (1.0 + (row * 0.19).sin())
262        })
263    }
264
265    fn naive_xtx(x: &Array2<f64>) -> Array2<f64> {
266        let n = x.nrows();
267        let p = x.ncols();
268        let mut out = Array2::zeros((p, p));
269        for row in 0..n {
270            for a in 0..p {
271                for b in 0..p {
272                    out[[a, b]] += x[[row, a]] * x[[row, b]];
273                }
274            }
275        }
276        out
277    }
278
279    fn naive_xtwy(x: &Array2<f64>, weights: &Array1<f64>, r: &Array1<f64>) -> Array1<f64> {
280        let n = x.nrows();
281        let p = x.ncols();
282        let mut out = Array1::zeros(p);
283        for row in 0..n {
284            for col in 0..p {
285                out[col] += x[[row, col]] * weights[row] * r[row];
286            }
287        }
288        out
289    }
290
291    fn naive_ywy(weights: &Array1<f64>, r: &Array1<f64>) -> f64 {
292        let mut sum = 0.0;
293        for row in 0..weights.len() {
294            sum += weights[row] * r[row] * r[row];
295        }
296        sum
297    }
298
299    fn assert_matrix_close(actual: ndarray::ArrayView2<'_, f64>, expected: &Array2<f64>, eps: f64) {
300        assert_eq!(actual.nrows(), expected.nrows());
301        assert_eq!(actual.ncols(), expected.ncols());
302        for row in 0..expected.nrows() {
303            for col in 0..expected.ncols() {
304                assert_abs_diff_eq!(actual[[row, col]], expected[[row, col]], epsilon = eps);
305            }
306        }
307    }
308
309    fn assert_vector_close(actual: ndarray::ArrayView1<'_, f64>, expected: &Array1<f64>, eps: f64) {
310        assert_eq!(actual.len(), expected.len());
311        for index in 0..expected.len() {
312            assert_abs_diff_eq!(actual[index], expected[index], epsilon = eps);
313        }
314    }
315
316    #[test]
317    fn gaussian_xtwx_matches_naive() {
318        let n = 40;
319        let p = 4;
320        let x = deterministic_design(n, p);
321        let y = deterministic_response(n);
322        let cache = FixedDesignGramCache::build(x.view(), y.view(), None, None).unwrap();
323        let naive = naive_xtx(&x);
324        assert_matrix_close(cache.xtwx(), &naive, 1.0e-9);
325    }
326
327    #[test]
328    fn gaussian_xtwy_and_ywy_match_naive() {
329        let n = 40;
330        let p = 4;
331        let x = deterministic_design(n, p);
332        let y = deterministic_response(n);
333        let offset = deterministic_offset(n);
334        let weights = deterministic_weights(n, 0.35);
335        let r = &y - &offset;
336        let cache = FixedDesignGramCache::build(
337            x.view(),
338            y.view(),
339            Some(offset.view()),
340            Some(weights.view()),
341        )
342        .unwrap();
343        let expected_xtwy = naive_xtwy(&x, &weights, &r);
344        let expected_ywy = naive_ywy(&weights, &r);
345        assert_vector_close(cache.xtwy(), &expected_xtwy, 1.0e-9);
346        assert_abs_diff_eq!(cache.ywy(), expected_ywy, epsilon = 1.0e-9);
347    }
348
349    #[test]
350    fn row_cache_xtwx_matches_fresh_build_across_weights() {
351        let n = 40;
352        let p = 4;
353        let x = deterministic_design(n, p);
354        let cache = FixedDesignRowCache::build(x.view()).unwrap();
355        let weight_sets = [
356            deterministic_weights(n, 0.12),
357            deterministic_weights(n, 0.27),
358            deterministic_weights(n, 0.41),
359        ];
360        for weights in weight_sets.iter() {
361            let cached = cache.xtwx(weights.view()).unwrap();
362            let fresh = fast_xt_diag_x(&x, weights);
363            assert_matrix_close(cached.view(), &fresh, 1.0e-12);
364        }
365    }
366
367    #[test]
368    fn build_rejects_shape_mismatch() {
369        let n = 40;
370        let p = 4;
371        let x = deterministic_design(n, p);
372        let mismatched_y = deterministic_response(n - 1);
373        assert!(FixedDesignGramCache::build(x.view(), mismatched_y.view(), None, None).is_err());
374
375        let y = deterministic_response(n);
376        let mut weights = deterministic_weights(n, 0.2);
377        weights[3] = f64::NAN;
378        assert!(
379            FixedDesignGramCache::build(x.view(), y.view(), None, Some(weights.view())).is_err()
380        );
381    }
382}
gam_solve/measure_jet_gram_cache.rs

gam_solve/
measure_jet_gram_cache.rs
