gam_terms/
kronecker.rs

1//! λ-invariant Kronecker tensor structure.
2//!
3//! Every numeric primitive this module needs — the faer/ndarray bridges, the
4//! symmetric sanitiser, the robust marginal eigensolve and the multi-index
5//! walk — is owned by [`crate::construction`], which also owns the per-iterate
6//! Kronecker engine that consumes this structure. This file used to carry a
7//! second copy of all of them, differing only in the error type it wrapped
8//! them in (#2470); the eigensolve failure is an estimation failure and is now
9//! reported as one on both routes.
10
11use crate::construction::{kronecker_marginal_eigensystems, kronecker_multi_index_advance};
12use gam_problem::EstimationError;
13use ndarray::{Array1, Array2};
14use std::sync::Arc;
15
16pub use gam_problem::penalty_matrix::kronecker_product;
17
18/// λ-invariant Kronecker tensor structure: everything in a tensor-product fit
19/// that depends ONLY on the marginal designs/penalties (which are fixed for the
20/// whole fit) and NOT on the smoothing parameters λ = exp(ρ).
21///
22/// The marginal eigendecomposition (`O(Σ q_k³)`), the reparameterized marginals
23/// `B_k · U_k`, and the balanced-penalty shrinkage scale `max_bal` are all
24/// functions of the fixed marginal data alone. Caching them once per fit lets
25/// every outer REML iterate (50+ per fit on the #1082 tensor cases) skip the
26/// repeated `eigh()` calls and `B_k U_k` GEMMs; only the cheap
27/// `kronecker_logdet_and_derivatives` λ-grid sweep is redone per iterate.
28#[derive(Clone, Debug)]
29pub struct KroneckerInvariantStructure {
30    /// Marginal eigenvalues from each marginal penalty eigendecomposition.
31    ///
32    /// `Arc`-shared so handing this structure to the per-iterate memoized
33    /// engine is an O(1) refcount bump, not a deep array copy.
34    pub marginal_eigenvalues: Arc<Vec<Array1<f64>>>,
35    /// Marginal eigenvector matrices U_k.
36    pub marginal_qs: Arc<Vec<Array2<f64>>>,
37    /// Reparameterized marginal designs: `B_k · U_k` for each marginal k.
38    pub reparameterized_marginals: Arc<Vec<Array2<f64>>>,
39    /// Max balanced-penalty eigenvalue scale `max_k-grid Σ_k μ_{k,j_k}/||S_k||_F`,
40    /// used to form the shrinkage ridge `floor * max_bal`. λ-independent.
41    pub max_balanced_eigenvalue: f64,
42}
43
44impl KroneckerInvariantStructure {
45    /// Compute the λ-invariant tensor structure once from the fixed marginal data.
46    pub fn compute(
47        marginal_designs: &[Array2<f64>],
48        marginal_penalties: &[Array2<f64>],
49        marginal_dims: &[usize],
50    ) -> Result<Self, EstimationError> {
51        let d = marginal_dims.len();
52        // Eigendecompose each marginal penalty once through the same robust path
53        // used by KroneckerPenaltySystem so every Kronecker caller sees the same
54        // eigensystem and pseudo-logdet surface.
55        let mut marginal_eigenvalues = Vec::with_capacity(d);
56        let mut marginal_qs = Vec::with_capacity(d);
57        for (evals, evecs) in kronecker_marginal_eigensystems(
58            marginal_penalties,
59            "kronecker_reparameterization_engine",
60        )? {
61            marginal_eigenvalues.push(evals);
62            marginal_qs.push(evecs);
63        }
64
65        // Reparameterized marginals: B_k · U_k.
66        let reparameterized_marginals: Vec<Array2<f64>> = marginal_designs
67            .iter()
68            .zip(marginal_qs.iter())
69            .map(|(b_k, u_k)| gam_linalg::faer_ndarray::fast_ab(b_k, u_k))
70            .collect();
71
72        // Max balanced eigenvalue: for Kronecker, the balanced penalty's max
73        // eigenvalue is the max over multi-indices of Σ_k (1/||S_k||_F) μ_{k,j_k}.
74        let mut max_balanced_eigenvalue = 0.0_f64;
75        let mut multi_idx = vec![0usize; d];
76        // A marginal with an all-zero penalty (‖S_k‖_F = 0) has no spectrum
77        // to balance and contributes nothing to the balanced eigenvalue; it is
78        // skipped rather than divided by a floor.
79        let frob_norms: Vec<f64> = marginal_penalties
80            .iter()
81            .map(|s| s.iter().map(|v| v * v).sum::<f64>().sqrt())
82            .collect();
83        loop {
84            let mut sigma = 0.0;
85            for k in 0..d {
86                if frob_norms[k] > 0.0 {
87                    sigma += marginal_eigenvalues[k][multi_idx[k]] / frob_norms[k];
88                }
89            }
90            max_balanced_eigenvalue = max_balanced_eigenvalue.max(sigma);
91
92            if kronecker_multi_index_advance(&mut multi_idx, marginal_dims) {
93                break;
94            }
95        }
96
97        Ok(Self {
98            marginal_eigenvalues: Arc::new(marginal_eigenvalues),
99            marginal_qs: Arc::new(marginal_qs),
100            reparameterized_marginals: Arc::new(reparameterized_marginals),
101            max_balanced_eigenvalue,
102        })
103    }
104}
gam_terms/kronecker.rs

gam_terms/
kronecker.rs
