gam_terms/

grid_spline_2d.rs

//! Streaming scatter-add 2-D smoother: K×K tensor-product cubic B-splines
//! with the EXACT anisotropic biharmonic penalty and REML-selected λ.
//!
//! Basis. Each axis carries K equal-width cells over the data's bounding box
//! `[lo, hi]` with uniform extended knots `t_j = lo + (j−3)·h`, `h = (hi−lo)/K`,
//! giving `m = K+3` cubic B-splines per axis; the tensor product has
//! `p = (K+3)²` coefficients. A point in cell `i` activates exactly the four
//! splines `i..i+3` per axis, hence exactly 4×4 = 16 tensor basis entries per
//! data row.
//!
//! Streaming normal equations. ONE pass over the rows `(x1, x2, y_·, w)`
//! scatter-adds `X'WX` and `X'Wy_d` (any number of response dimensions share
//! the design, the penalty, and one REML λ — the multi-output "one surface
//! smoothness" contract of the ANOVA pair component): O(n·(16² + 16·D)) work,
//! no n×p design is ever materialized. Two tensor bases overlap only when both per-axis indices
//! differ by ≤ 3, so under the row-major coefficient index
//! `g = j1·(K+3) + j2` both `X'WX` and the penalty `S` are banded with
//! half-bandwidth `3(K+3)+3`; they are stored as upper bands — O(K³) numbers.
//!
//! Penalty. The FULL anisotropic biharmonic form for the diagonal metric
//! `A = diag(a1, a2)`,
//!   `J(f) = ∫∫ a1²·f_{x1x1}² + 2·a1·a2·f_{x1x2}² + a2²·f_{x2x2}²  dx1 dx2`,
//! INCLUDING the mixed `f_{x1x2}` term (the axis-wise P-spline difference
//! shortcut drops it), assembled per knot cell by 4-point Gauss–Legendre per
//! axis. Exactness degree arithmetic: on a knot cell every basis function is
//! a single cubic polynomial per axis, so each entry of `S` is a sum over
//! cells of integrands that factorize per axis as one of value·value
//! (degree 3+3 = 6), deriv·deriv (2+2 = 4) or 2nd-deriv·2nd-deriv (1+1 = 2);
//! every channel pairs a low-degree factor on one axis with at worst the
//! degree-6 value·value factor on the other. 4-point Gauss–Legendre is exact
//! through degree 2·4−1 = 7 ≥ 6, so the assembled `S` is the EXACT integral,
//! not a quadrature approximation.
//!
//! Solve and selection. A single reference factorization of `H₀=X'WX+S`
//! produces the affine generalized-eigenvalue pencil
//! `H(λ)=L[I+(λ-1)U diag(μ)U']L'`. The profiled score and its first two
//! analytic log-λ derivatives are then O(pD), and outward interval enclosures
//! isolate every stationary interval without a grid. The selected system is
//! factored once more for coefficients/posterior covariance. `p ≤ (32+3)² =
//! 1225`; K is capped at 32 to keep the dense reference factor and eigensystem
//! sizing contract honest. λ maximizes the
//! profiled-σ² restricted (REML) criterion
//!   `ℓ_R(λ) = −½[ log|X'WX+λS| − r·log λ + (n−3)·log σ̂²(λ) ] + const`,
//! where `r = p−3` is the penalty rank — the null space of `J` is
//! span{1, x1, x2} (the mixed term penalizes `x1·x2`, whose cross derivative
//! is 1 ≠ 0, so it is NOT in the null space), `σ̂²(λ) = (y'Wy − c'X'Wy)/(n−3)`
//! is the profiled scale, and the λ-free additive constants (`log|S|₊` on the
//! row space of S, `Σ log w`, 2π factors) are dropped: differences across λ
//! are exact REML criterion differences. The exact bounded-domain endpoints
//! (including the null-recovery end) compete with every certified stationary
//! point; no RNG or lattice is involved, so the same data imply the same fit.
//!
//! Prediction. `predict(x1, x2)` builds the 16-entry basis row; the mean is
//! its dot with `c` and the variance is the Bayesian posterior
//! `σ̂²·x'(X'WX+λS)⁻¹x` through the retained Cholesky factor. Outside the
//! bounding box the boundary cell's cubic polynomial extends naturally (the
//! cell index clamps, the local coordinate does not).

/// Dimension of the penalty null space: span{1, x1, x2}. The mixed
/// `2·a1·a2·f_{x1x2}²` term excludes `x1·x2` (its cross derivative is 1).
const PENALTY_NULLITY: usize = 3;

/// Cholesky pivot floor below which the penalized system is declared singular.
/// Dense-Cholesky sizing contract documented in the module header.
const MAX_CELLS_PER_AXIS: usize = 32;

/// 4-point Gauss–Legendre nodes and weights on [−1, 1]. Exact through degree
/// 2·4−1 = 7, which dominates the degree-6 worst per-axis factor of the
/// penalty integrands (see the module header for the degree arithmetic).
const GL4_NODES: [f64; 4] = [
    -0.861_136_311_594_052_6,
    -0.339_981_043_584_856_26,
    0.339_981_043_584_856_26,
    0.861_136_311_594_052_6,
];
const GL4_WEIGHTS: [f64; 4] = [
    0.347_854_845_137_453_85,
    0.652_145_154_862_546_2,
    0.652_145_154_862_546_2,
    0.347_854_845_137_453_85,
];

/// Cubic B-spline segment values at local coordinate `u` within a cell.
/// Entry `m` weights basis `cell + m`: m = 0 is the spline ENDING in this
/// cell (`(1−u)³/6`), m = 3 the one STARTING (`u³/6`). The four entries sum
/// to 1 (partition of unity) for u ∈ [0, 1].
#[inline]
fn bspline_value(u: f64) -> [f64; 4] {
    let v = 1.0 - u;
    [
        v * v * v / 6.0,
        (3.0 * u * u * u - 6.0 * u * u + 4.0) / 6.0,
        (-3.0 * u * u * u + 3.0 * u * u + 3.0 * u + 1.0) / 6.0,
        u * u * u / 6.0,
    ]
}

/// d/du of `bspline_value` (caller scales by 1/h for d/dx). Entries sum to 0.
#[inline]
fn bspline_d1(u: f64) -> [f64; 4] {
    let v = 1.0 - u;
    [
        -0.5 * v * v,
        0.5 * (3.0 * u * u - 4.0 * u),
        0.5 * (-3.0 * u * u + 2.0 * u + 1.0),
        0.5 * u * u,
    ]
}

/// d²/du² of `bspline_value` (caller scales by 1/h²). Piecewise LINEAR in u —
/// the degree-1 factor in the quadrature-exactness argument. Entries sum to 0.
#[inline]
fn bspline_d2(u: f64) -> [f64; 4] {
    [1.0 - u, 3.0 * u - 2.0, 1.0 - 3.0 * u, u]
}

/// One uniform B-spline axis over `[lo, lo + cells·h]`.
#[derive(Clone, Copy, Debug)]
struct Axis {
    lo: f64,
    h: f64,
    cells: usize,
}

impl Axis {
    /// Cell index and local coordinate. Inside the box `u ∈ [0, 1]`; outside,
    /// the cell clamps and `u` leaves [0, 1], extending the boundary cell's
    /// cubic polynomial (deterministic extrapolation, no special casing).
    #[inline]
    fn locate(&self, x: f64) -> (usize, f64) {
        let t = (x - self.lo) / self.h;
        let cell = (t.floor().max(0.0) as usize).min(self.cells - 1);
        (cell, t - cell as f64)
    }
}

/// The four active cubic B-spline values of one uniform axis `(lo, h, cells)`
/// at `x`: `(first basis index, values)`, where `values[i]` weights basis
/// `first + i` of the `cells + 3` axis splines. Outside `[lo, lo + cells·h]`
/// the boundary cell's cubic polynomial extends — the single convention
/// shared by fitting, prediction, and every consumer-rebuilt basis row.
pub fn axis_basis_at(lo: f64, h: f64, cells: usize, x: f64) -> (usize, [f64; 4]) {
    let (cell, u) = Axis { lo, h, cells }.locate(x);
    (cell, bspline_value(u))
}

/// The 16 active tensor-basis entries `(flat index, value)` at `(x1, x2)`.
/// Flat indices are strictly increasing across the returned arrays.
#[inline]
fn basis_row(axes: &[Axis; 2], m_axis: usize, x1: f64, x2: f64) -> ([usize; 16], [f64; 16]) {
    let (c1, u1) = axes[0].locate(x1);
    let (c2, u2) = axes[1].locate(x2);
    let b1 = bspline_value(u1);
    let b2 = bspline_value(u2);
    let mut idx = [0usize; 16];
    let mut val = [0f64; 16];
    for i in 0..4 {
        for j in 0..4 {
            idx[4 * i + j] = (c1 + i) * m_axis + (c2 + j);
            val[4 * i + j] = b1[i] * b2[j];
        }
    }
    (idx, val)
}

/// Dense lower-Cholesky in place (row-major `p×p`); returns the exact
/// `log det` (twice the log of the pivot products). The strict upper triangle
/// is zeroed so the buffer is exactly `L` afterwards.
pub fn cholesky_logdet(a: &mut [f64], p: usize) -> Result<f64, String> {
    let mut logdet = 0.0;
    for j in 0..p {
        let diag = a[j * p + j];
        let mut s = diag;
        let mut subtracted = 0.0_f64;
        for t in 0..j {
            let sq = a[j * p + t] * a[j * p + t];
            s -= sq;
            subtracted += sq;
        }
        // The pivot is `a_jj − Σ_t l_jt²`: one product and one subtraction per
        // term, so a pivot inside the rounding band of that accumulation is
        // not distinguishable from zero and the system is not positive definite
        // as computed (#2469). An absolute `1e-300` passed pure roundoff of an
        // O(1) row and refused an honest tiny pivot.
        let band = gam_linalg::roundoff::accumulation_growth(2 * j + 1) * (diag.abs() + subtracted);
        if !(s.is_finite() && s > band) {
            return Err(format!(
                "grid spline 2d: penalized system not positive definite at pivot {j} (value {s})"
            ));
        }
        let l = s.sqrt();
        a[j * p + j] = l;
        logdet += 2.0 * l.ln();
        for i in j + 1..p {
            let mut s2 = a[i * p + j];
            for t in 0..j {
                s2 -= a[i * p + t] * a[j * p + t];
            }
            a[i * p + j] = s2 / l;
        }
    }
    for i in 0..p {
        for j in i + 1..p {
            a[i * p + j] = 0.0;
        }
    }
    Ok(logdet)
}

/// Solve `L z = b` from a dense row-major lower-triangular factor.
fn lower_solve(l: &[f64], p: usize, b: &[f64]) -> Vec<f64> {
    let mut z = b.to_vec();
    for i in 0..p {
        let mut s = z[i];
        for t in 0..i {
            s -= l[i * p + t] * z[t];
        }
        z[i] = s / l[i * p + i];
    }
    z
}

/// Solve `L Lᵀ x = b` from the stored lower factor.
pub fn chol_solve(l: &[f64], p: usize, b: &[f64]) -> Vec<f64> {
    let mut z = lower_solve(l, p, b);
    for i in (0..p).rev() {
        let mut s = z[i];
        for t in i + 1..p {
            s -= l[t * p + i] * z[t];
        }
        z[i] = s / l[i * p + i];
    }
    z
}

/// Banded sufficient statistics of one streaming pass plus the exact penalty:
/// everything needed to evaluate the REML criterion and solve at any λ.
pub struct GridSpline2dDesign {
    axes: [Axis; 2],
    /// Basis count per axis, `K + 3`.
    /// Total coefficients, `(K + 3)²`.
    p: usize,
    /// Upper half-bandwidth `3·(K+3) + 3` of both banded matrices.
    band_half: usize,
    /// Upper band of `X'WX`: entry `(g, g+d)` at `g·(band_half+1) + d`.
    gram_band: Vec<f64>,
    /// Upper band of the exact anisotropic biharmonic penalty `S`.
    pen_band: Vec<f64>,
    /// `X'Wy_d`, one length-`p` vector per response dimension. The design
    /// (gram and penalty bands) is shared across dimensions; only these
    /// right-hand sides and the response cross-moments are per-dimension.
    rhs: Vec<Vec<f64>>,
    /// Response cross-moments `y_d'W y_e` (`D × D` row-major), for the
    /// profiled-σ² residual quadratics and the residual cross-covariance.
    cross_moments: Vec<f64>,
    n_obs: usize,
}

impl GridSpline2dDesign {
    /// Single-response entry: see [`Self::build_multi`].
    pub fn build(
        x1: &[f64],
        x2: &[f64],
        y: &[f64],
        w: &[f64],
        k: usize,
        metric: [f64; 2],
    ) -> Result<Self, String> {
        Self::build_multi(x1, x2, &[y], w, k, metric)
    }

    /// One streaming pass over the rows plus the exact per-cell quadrature
    /// assembly of the penalty. `k` is the number of cells per axis;
    /// `metric = [a1, a2]` is the diagonal anisotropy of the biharmonic form.
    /// `responses` holds one length-`n` response per dimension; the design,
    /// penalty, and the REML-shared λ are common to all dimensions (one
    /// surface smoothness), only the right-hand sides differ.
    pub fn build_multi(
        x1: &[f64],
        x2: &[f64],
        responses: &[&[f64]],
        w: &[f64],
        k: usize,
        metric: [f64; 2],
    ) -> Result<Self, String> {
        let n = x1.len();
        if responses.is_empty() {
            return Err("grid spline 2d: no response dimensions supplied".to_string());
        }
        if x2.len() != n || w.len() != n {
            return Err(format!(
                "grid spline 2d: length mismatch x1={n}, x2={}, w={}",
                x2.len(),
                w.len()
            ));
        }
        for (d, y) in responses.iter().enumerate() {
            if y.len() != n {
                return Err(format!(
                    "grid spline 2d: response dimension {d} has length {} != {n}",
                    y.len()
                ));
            }
        }
        if n <= PENALTY_NULLITY {
            return Err(format!(
                "grid spline 2d: needs more than {PENALTY_NULLITY} rows for the profiled REML \
                 degrees of freedom, got {n}"
            ));
        }
        if k == 0 || k > MAX_CELLS_PER_AXIS {
            return Err(format!(
                "grid spline 2d: k must be in 1..={MAX_CELLS_PER_AXIS} (dense Cholesky on \
                 (k+3)² coefficients — see module sizing contract), got {k}"
            ));
        }
        if !(metric[0].is_finite() && metric[0] > 0.0 && metric[1].is_finite() && metric[1] > 0.0) {
            return Err(format!(
                "grid spline 2d: metric diagonal must be finite and positive, got [{}, {}]",
                metric[0], metric[1]
            ));
        }
        for i in 0..n {
            if !(x1[i].is_finite() && x2[i].is_finite()) || !(w[i] > 0.0) || !w[i].is_finite() {
                return Err(format!(
                    "grid spline 2d: non-finite or non-positive input at row {i} \
                     (x1={}, x2={}, w={})",
                    x1[i], x2[i], w[i]
                ));
            }
            for (d, y) in responses.iter().enumerate() {
                if !y[i].is_finite() {
                    return Err(format!(
                        "grid spline 2d: non-finite response at row {i}, dimension {d} ({})",
                        y[i]
                    ));
                }
            }
        }
        let mut axes = [Axis {
            lo: 0.0,
            h: 1.0,
            cells: k,
        }; 2];
        for (axis, xs) in axes.iter_mut().zip([x1, x2]) {
            let mut lo = f64::INFINITY;
            let mut hi = f64::NEG_INFINITY;
            for &v in xs {
                lo = lo.min(v);
                hi = hi.max(v);
            }
            if !(hi > lo) {
                return Err(format!(
                    "grid spline 2d: degenerate axis bounding box [{lo}, {hi}]"
                ));
            }
            axis.lo = lo;
            axis.h = (hi - lo) / k as f64;
        }
        let m_axis = k + 3;
        let p = m_axis * m_axis;
        let band_half = 3 * m_axis + 3;
        let stride = band_half + 1;
        let n_dims = responses.len();
        let mut gram_band = vec![0.0_f64; p * stride];
        let mut rhs = vec![vec![0.0_f64; p]; n_dims];
        let mut cross_moments = vec![0.0_f64; n_dims * n_dims];

        // ── ONE streaming pass: scatter-add X'WX (upper band) and X'Wy_d ──
        // Each row touches exactly 16 basis entries with strictly increasing
        // flat indices, so the in-row pair loop (a ≤ b) lands directly in the
        // upper band: O(n·(16² + 16·D)) total work.
        for i in 0..n {
            let (idx, val) = basis_row(&axes, m_axis, x1[i], x2[i]);
            let wi = w[i];
            for (d, y) in responses.iter().enumerate() {
                let wy = wi * y[i];
                for e in 0..16 {
                    rhs[d][idx[e]] += wy * val[e];
                }
                for (e, ye) in responses.iter().enumerate().skip(d) {
                    cross_moments[d * n_dims + e] += wy * ye[i];
                }
            }
            for a in 0..16 {
                let base = idx[a] * stride - idx[a];
                let wa = wi * val[a];
                for b in a..16 {
                    gram_band[base + idx[b]] += wa * val[b];
                }
            }
        }
        for d in 0..n_dims {
            for e in 0..d {
                cross_moments[d * n_dims + e] = cross_moments[e * n_dims + d];
            }
        }

        // ── Exact penalty assembly: 4-pt Gauss–Legendre per axis per cell ──
        // Per-axis quadrature tables (cell-independent on a uniform grid):
        // values, d/dx (scaled 1/h), d²/dx² (scaled 1/h²) at each GL node.
        let mut tab = [[[[0.0_f64; 4]; 4]; 3]; 2]; // [axis][channel][node][basis offset]
        for ax in 0..2 {
            let h = axes[ax].h;
            for q in 0..4 {
                let u = 0.5 * (1.0 + GL4_NODES[q]);
                let v0 = bspline_value(u);
                let v1 = bspline_d1(u);
                let v2 = bspline_d2(u);
                for e in 0..4 {
                    tab[ax][0][q][e] = v0[e];
                    tab[ax][1][q][e] = v1[e] / h;
                    tab[ax][2][q][e] = v2[e] / (h * h);
                }
            }
        }
        // Channel scales: J = ∫ a1²·f11² + 2·a1·a2·f12² + a2²·f22².
        let s11 = metric[0] * metric[0];
        let s12 = 2.0 * metric[0] * metric[1];
        let s22 = metric[1] * metric[1];
        let cell_area_jac = 0.25 * axes[0].h * axes[1].h; // d(x1,x2)/d(ξ1,ξ2) on [−1,1]²
        let mut pen_band = vec![0.0_f64; p * stride];
        let mut r11 = [0.0_f64; 16];
        let mut r12 = [0.0_f64; 16];
        let mut r22 = [0.0_f64; 16];
        let mut idx = [0usize; 16];
        for c1 in 0..k {
            for c2 in 0..k {
                for i in 0..4 {
                    for j in 0..4 {
                        idx[4 * i + j] = (c1 + i) * m_axis + (c2 + j);
                    }
                }
                for q1 in 0..4 {
                    for q2 in 0..4 {
                        let wq = cell_area_jac * GL4_WEIGHTS[q1] * GL4_WEIGHTS[q2];
                        for i in 0..4 {
                            for j in 0..4 {
                                let e = 4 * i + j;
                                r11[e] = tab[0][2][q1][i] * tab[1][0][q2][j];
                                r12[e] = tab[0][1][q1][i] * tab[1][1][q2][j];
                                r22[e] = tab[0][0][q1][i] * tab[1][2][q2][j];
                            }
                        }
                        for a in 0..16 {
                            let base = idx[a] * stride - idx[a];
                            let (pa11, pa12, pa22) =
                                (wq * s11 * r11[a], wq * s12 * r12[a], wq * s22 * r22[a]);
                            for b in a..16 {
                                pen_band[base + idx[b]] +=
                                    pa11 * r11[b] + pa12 * r12[b] + pa22 * r22[b];
                            }
                        }
                    }
                }
            }
        }

        Ok(GridSpline2dDesign {
            axes,
            p,
            band_half,
            gram_band,
            pen_band,
            rhs,
            cross_moments,
            n_obs: n,
        })
    }

    /// Lower corner of the data bounding box per axis.
    pub fn lower_corner(&self) -> [f64; 2] {
        [self.axes[0].lo, self.axes[1].lo]
    }

    /// Knot-cell width per axis.
    pub fn cell_widths(&self) -> [f64; 2] {
        [self.axes[0].h, self.axes[1].h]
    }

    /// Exact penalty quadratic form `J(f) = c'Sc` of a coefficient vector —
    /// the assembled anisotropic biharmonic energy of the spline it encodes.
    pub fn penalty_value(&self, coeff: &[f64]) -> Result<f64, String> {
        if coeff.len() != self.p {
            return Err(format!(
                "grid spline 2d: coefficient length {} != {}",
                coeff.len(),
                self.p
            ));
        }
        let stride = self.band_half + 1;
        let mut j = 0.0;
        for g in 0..self.p {
            let dmax = self.band_half.min(self.p - 1 - g);
            j += self.pen_band[g * stride] * coeff[g] * coeff[g];
            for d in 1..=dmax {
                j += 2.0 * self.pen_band[g * stride + d] * coeff[g] * coeff[g + d];
            }
        }
        Ok(j)
    }

    /// `a'(X'WX)b` through the retained upper band (exact, O(p·bandwidth)).
    fn gram_quadratic(&self, a: &[f64], b: &[f64]) -> f64 {
        let stride = self.band_half + 1;
        let mut q = 0.0;
        for g in 0..self.p {
            let dmax = self.band_half.min(self.p - 1 - g);
            q += self.gram_band[g * stride] * a[g] * b[g];
            for d in 1..=dmax {
                q += self.gram_band[g * stride + d] * (a[g] * b[g + d] + a[g + d] * b[g]);
            }
        }
        q
    }

    /// Posterior summary of a fit FROM THIS DESIGN, in the exact algebra of
    /// the solved system (no approximation):
    /// - `unit_covariance = (X'WX + λS)⁻¹` (scale-free Bayesian posterior
    ///   covariance of the row-major coefficient vec, shared by dimensions);
    /// - `edf = tr[(X'WX + λS)⁻¹ X'WX]` (the smoother's effective degrees of
    ///   freedom at the fitted λ);
    /// - `residual_cross_cov[d,e] = r_d'W r_e / (n − edf)` assembled from the
    ///   streamed sufficient statistics
    ///   (`y_d'Wy_e − c_d'X'Wy_e − c_e'X'Wy_d + c_d'X'WX c_e`).
    pub fn posterior(&self, fit: &GridSpline2dFit) -> Result<GridSpline2dPosterior, String> {
        let p = self.p;
        let n_dims = self.rhs.len();
        if fit.coeffs.len() != n_dims || fit.coeffs.iter().any(|c| c.len() != p) {
            return Err(format!(
                "grid spline 2d: posterior asked for a fit with {} dimensions of length {}, \
                 design has {n_dims} of {p}",
                fit.coeffs.len(),
                fit.coeffs.first().map_or(0, Vec::len)
            ));
        }
        // H⁻¹ column by column through the retained factor (symmetric, O(p³)).
        let mut unit_covariance = vec![0.0_f64; p * p];
        let mut e_g = vec![0.0_f64; p];
        for g in 0..p {
            e_g[g] = 1.0;
            let col = chol_solve(&fit.chol, p, &e_g);
            e_g[g] = 0.0;
            for (r, &v) in col.iter().enumerate() {
                unit_covariance[r * p + g] = v;
            }
        }
        // edf = tr(H⁻¹ X'WX) via the gram band (diagonal once, off-band twice).
        let stride = self.band_half + 1;
        let mut edf = 0.0;
        for g in 0..p {
            let dmax = self.band_half.min(p - 1 - g);
            edf += self.gram_band[g * stride] * unit_covariance[g * p + g];
            for d in 1..=dmax {
                edf += 2.0 * self.gram_band[g * stride + d] * unit_covariance[g * p + g + d];
            }
        }
        let residual_df = self.n_obs as f64 - edf;
        if !(residual_df >= 1.0) {
            return Err(format!(
                "grid spline 2d: too few rows for a scale estimate \
                 (n = {}, edf = {edf:.2}; need n − edf ≥ 1)",
                self.n_obs
            ));
        }
        let mut residual_cross_cov = vec![0.0_f64; n_dims * n_dims];
        for d in 0..n_dims {
            for e in d..n_dims {
                let mut cd_rhse = 0.0;
                let mut ce_rhsd = 0.0;
                for g in 0..p {
                    cd_rhse += fit.coeffs[d][g] * self.rhs[e][g];
                    ce_rhsd += fit.coeffs[e][g] * self.rhs[d][g];
                }
                let quad = self.gram_quadratic(&fit.coeffs[d], &fit.coeffs[e]);
                let v =
                    (self.cross_moments[d * n_dims + e] - cd_rhse - ce_rhsd + quad) / residual_df;
                residual_cross_cov[d * n_dims + e] = v;
                residual_cross_cov[e * n_dims + d] = v;
            }
        }
        Ok(GridSpline2dPosterior {
            unit_covariance,
            edf,
            residual_df,
            residual_cross_cov,
        })
    }
}

/// Exact posterior summary of a [`GridSpline2dFit`] (see
/// [`GridSpline2dDesign::posterior`]): the bridge from the streaming engine
/// to covariance-consuming clients (the ANOVA pair-component carve).
pub struct GridSpline2dPosterior {
    /// `(X'WX + λS)⁻¹`, `p × p` row-major — scale-free posterior covariance
    /// of the row-major coefficient vec, shared by all response dimensions.
    pub unit_covariance: Vec<f64>,
    /// `tr[(X'WX + λS)⁻¹ X'WX]`.
    pub edf: f64,
    /// `n − edf`.
    pub residual_df: f64,
    /// `D × D` row-major residual cross-covariance at `n − edf`.
    pub residual_cross_cov: Vec<f64>,
}

/// Fitted penalized tensor-product smoother with its factored covariance.
pub struct GridSpline2dFit {
    /// Per response dimension: coefficients in row-major flat order
    /// `g = j1·(K+3) + j2`.
    pub coeffs: Vec<Vec<f64>>,
    /// Selected (or supplied) log smoothing parameter, shared by all
    /// response dimensions.
    pub log_lambda: f64,
    /// Per response dimension: profiled (or supplied) observation variance σ².
    pub sigma2: Vec<f64>,
    /// Pooled restricted log-likelihood at the optimum, up to λ- and
    /// data-independent additive constants (exact REML differences across λ).
    pub restricted_loglik: f64,
    /// Lower Cholesky factor of `X'WX + λS` — the factored posterior precision
    /// (unit-σ² scale) used for prediction variances, shared by all dimensions.
    chol: Vec<f64>,
    axes: [Axis; 2],
    m_axis: usize,
}

/// Serializable snapshot of a [`GridSpline2dFit`] (#1031 persistence
/// prerequisite). The grid is deliberately NOT a formula fast path — it is an
/// ANOVA pair component (#975 carve) — so there is no `FitResult` variant; this
/// state is what the carve's persistence payload serializes and what
/// `from_state` replays for an exact predict.
///
/// Predict needs the MEAN (`coeffs` + the 16-entry tensor basis row, which is a
/// pure function of `axes`/`m_axis`) and the VARIANCE
/// (`σ²·x'(X'WX+λS)⁻¹x` through the retained Cholesky factor `chol`). All of
/// that — and nothing about the training rows — lives on the fit already, so the
/// state is a verbatim snapshot: no design CSR, no re-factor on load.
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub struct GridSpline2dState {
    /// Per response dimension: row-major coefficients `g = j1·(K+3) + j2`.
    pub coeffs: Vec<Vec<f64>>,
    pub log_lambda: f64,
    /// Per response dimension: profiled (or supplied) observation variance σ².
    pub sigma2: Vec<f64>,
    pub restricted_loglik: f64,
    /// Lower Cholesky factor of `X'WX + λS` (unit-σ² scale), `p × p` row-major —
    /// the factored posterior precision the variance term solves against.
    pub chol: Vec<f64>,
    /// Per axis lower corner of the basis bounding box.
    pub axis_lo: [f64; 2],
    /// Per axis cell width `h = (hi − lo)/K`.
    pub axis_h: [f64; 2],
    /// Per axis cell count `K`.
    pub axis_cells: [u64; 2],
    /// Basis count per axis, `K + 3` (so `p = m_axis²`).
    pub m_axis: u64,
}

impl GridSpline2dFit {

    /// Posterior `(mean, variance)` of response dimension `dim` at an
    /// arbitrary point: the 16-entry basis row dotted with the coefficients,
    /// and `σ̂²_dim·x'(X'WX+λS)⁻¹x` through the retained Cholesky factor.
    /// Outside the bounding box the boundary cell's cubic polynomial extends.
    pub fn predict(&self, dim: usize, x1: f64, x2: f64) -> Result<(f64, f64), String> {
        if dim >= self.coeffs.len() {
            return Err(format!(
                "grid spline 2d: response dimension {dim} out of range (D = {})",
                self.coeffs.len()
            ));
        }
        if !(x1.is_finite() && x2.is_finite()) {
            return Err(format!(
                "grid spline 2d: non-finite prediction point ({x1}, {x2})"
            ));
        }
        let (idx, val) = basis_row(&self.axes, self.m_axis, x1, x2);
        let p = self.coeffs[dim].len();
        let mut mean = 0.0;
        let mut row = vec![0.0_f64; p];
        for e in 0..16 {
            mean += val[e] * self.coeffs[dim][idx[e]];
            row[idx[e]] += val[e];
        }
        let z = chol_solve(&self.chol, p, &row);
        let mut quad = 0.0;
        for g in 0..p {
            quad += row[g] * z[g];
        }
        Ok((mean, self.sigma2[dim] * quad))
    }
}