Skip to main content

gam_solve/inference/
residual_factor.rs

1//! #974 — the structured-residual covariance estimator and the single producer
2//! of [`MetricProvenance::WhitenedStructured`](gam_problem::MetricProvenance::WhitenedStructured).
3//!
4//! # What this estimates
5//!
6//! Given a residual matrix `R ∈ ℝ^{n×p}` (one `p`-dimensional reconstruction
7//! residual per row) and a smooth *activity coordinate* `z ∈ ℝ^n`, this fits the
8//! **structured residual-covariance model**
9//!
10//! ```text
11//!     Cov(r_n) = Σ_n = Λ · c(z_n) · Λᵀ + D ,
12//! ```
13//!
14//! where
15//!
16//! * `Λ ∈ ℝ^{p×r}` is a **low-rank interference factor** (the shared
17//!   off-isotropic subspace the residuals correlate along — e.g. a planted
18//!   interference subspace or a topology-race confound),
19//! * `D = diag(d) ≻ 0` is the **idiosyncratic diagonal** (per-channel
20//!   independent noise), and
21//! * `c(z) > 0` is the **smooth activity-scale law**: a strictly-positive scalar
22//!   that modulates the factor energy with the activity coordinate, recovered as
23//!   a binned-then-smoothed function of `z`.
24//!
25//! The fit is a deterministic, fixed-iteration **alternation** (no clock, no
26//! RNG; any tie is broken by index): it alternates
27//!
28//! 1. *(scale | Λ, D)* — re-estimate the per-row factor activity `c(z_n)` and
29//!    smooth it across `z`, holding the factor model fixed; and
30//! 2. *(Λ, D | scale)* — re-estimate the factor and diagonal from the
31//!    scale-deflated second-moment, holding the activity law fixed,
32//!
33//! a fixed small number of times. The **factor count `r`** is chosen by an
34//! evidence ladder: each candidate `r` is scored by its penalized Gaussian
35//! log-evidence and the best is kept.
36//!
37//! # What it produces
38//!
39//! `StructuredResidualModel::row_metric` materializes the **per-row precision
40//! factor** `U_n ∈ ℝ^{p×p}` with `U_n U_nᵀ = Σ_n^{-1}`, packaged as a
41//! [`RowMetric`](gam_problem::RowMetric) with
42//! [`MetricProvenance::WhitenedStructured`](gam_problem::MetricProvenance::WhitenedStructured).
43//! Whitening a residual `r_n` through it (`U_nᵀ r_n`) yields a vector whose
44//! squared Euclidean norm is `r_nᵀ Σ_n^{-1} r_n` — the Mahalanobis residual under
45//! the estimated noise model, which is exactly the likelihood-correct data-fit.
46//! The factor is built from `Σ_n^{-1}` computed in **Woodbury form** (an
47//! `r × r` solve, never a `p × p` inverse), so the estimator scales with the
48//! factor rank, not the dense output dimension.
49//!
50//! This is the first real producer of `WhitenedStructured`, and therefore the
51//! first metric whose `whitens_likelihood()` is `true`: see
52//! [`RowMetric::whitens_likelihood`](gam_problem::RowMetric::whitens_likelihood).
53
54use std::sync::Arc;
55
56use ndarray::{Array1, Array2, ArrayView1, ArrayView2};
57
58use faer::Side;
59use gam_linalg::faer_ndarray::{FaerCholesky, FaerEigh};
60use gam_problem::RowMetric;
61
62/// Alternation sweeps the structured-Gaussian fit may take before it is
63/// refused. The alternation is a block-coordinate ascent on the penalized
64/// log-evidence (with smoothed, floored scale updates, so not an exact
65/// ascent), and it stops on its own at the first sweep that fails to improve
66/// the evidence by more than the evidence's rounding band, keeping the best
67/// state it saw. This budget bounds an ascent that keeps improving, and its
68/// exhaustion is a refusal, not a result (#2469 — it used to run exactly 8
69/// sweeps and return whatever state that left).
70const ALTERNATION_MAX_SWEEPS: usize = 256;
71
72/// Number of bins the activity coordinate `z` is partitioned into for the smooth
73/// activity-scale `c(z)`. The per-bin factor activity is estimated then linearly
74/// interpolated across bin centers, giving a continuous piecewise-linear scale
75/// law. Chosen as a fixed structural constant (magic-by-default): enough bins to
76/// resolve a smooth monotone or unimodal scale trend without over-fitting the
77/// per-row noise.
78const ACTIVITY_SCALE_BINS: usize = 8;
79
80/// Relative floor on the idiosyncratic diagonal `D`, as a fraction of the mean
81/// residual variance. Keeps `Σ_n ≻ 0` and the Woodbury `r × r` capacitance
82/// invertible even when a channel is (near-)perfectly explained by the factor.
83const DIAGONAL_REL_FLOOR: f64 = 1e-6;
84
85/// Relative floor on the activity scale `c(z)`, as a fraction of its mean. Keeps
86/// `c(z) > 0` (a covariance scale) across the whole `z` range.
87const SCALE_REL_FLOOR: f64 = 1e-4;
88
89/// The fitted structured residual-covariance model: low-rank factor `Λ`,
90/// idiosyncratic diagonal `D`, and the smooth activity-scale `c(z)` evaluated at
91/// every row. Produces per-row precision factors and the
92/// [`MetricProvenance::WhitenedStructured`](gam_problem::MetricProvenance::WhitenedStructured)
93/// `RowMetric`.
94#[derive(Clone, Debug)]
95pub struct StructuredResidualModel {
96    /// Output dimensionality `p` (residual width).
97    p: usize,
98    /// Selected factor rank `r` (`0 ≤ r ≤ p`). `0` ⇒ pure-diagonal noise model.
99    factor_rank: usize,
100    /// Interference factor `Λ ∈ ℝ^{p×r}` (the shared off-diagonal subspace).
101    lambda: Array2<f64>,
102    /// Idiosyncratic diagonal `d ∈ ℝ^p` (`D = diag(d)`), floored `≻ 0`.
103    diagonal: Array1<f64>,
104    /// Per-row activity scale `c(z_n) > 0`, length `n`.
105    row_scale: Array1<f64>,
106    /// Penalized Gaussian log-evidence of the selected model (higher is better).
107    /// The value the evidence ladder maximized over the candidate ranks.
108    log_evidence: f64,
109}
110
111/// Estimator inputs: the residual matrix and the smooth activity coordinate.
112///
113/// `residuals` is `R ∈ ℝ^{n×p}`. `activity` is `z ∈ ℝ^n` — the coordinate the
114/// scale law `c(z)` is smooth in (e.g. an assignment-mass or activation-strength
115/// summary per row). When no genuine activity coordinate is available, passing a
116/// constant `z` recovers a homoscedastic factor model (`c(z) ≡ const`).
117pub struct ResidualFactorInput<'a> {
118    /// Residual matrix `R ∈ ℝ^{n×p}`.
119    pub residuals: ArrayView2<'a, f64>,
120    /// Activity coordinate `z ∈ ℝ^n` the scale law is smooth in.
121    pub activity: ArrayView1<'a, f64>,
122    /// Maximum factor rank the evidence ladder is allowed to consider. The
123    /// ladder scores `r = 0, 1, …, min(max_factor_rank, p−1)` and keeps the
124    /// penalized-evidence maximizer. `0` forces the pure-diagonal model.
125    pub max_factor_rank: usize,
126}
127
128/// A persistent, evidence-earning residual factor direction — a promotion
129/// candidate for the #2021 Λ nursery→promotion birth channel. Emitted by
130/// [`StructuredResidualModel::promotion_candidates`] when a column of this
131/// pass's `Λ` both (a) aligns with a column of the previous pass's `Λ` (it
132/// *persisted* across the outer alternation) and (b) explains residual energy
133/// above the idiosyncratic-noise floor (it *earns its complexity*). The driver
134/// accumulates persistence across passes (the nursery) and, once a direction
135/// survives long enough, promotes it to a new curved/linear atom seeded by
136/// [`Self::direction`].
137#[derive(Clone, Debug)]
138pub struct FactorPromotion {
139    /// Unit-norm factor direction in output space (`p`-vector): the L2-normalized
140    /// column of `Λ`. This is the decoder direction a promoted atom is born with.
141    pub direction: Array1<f64>,
142    /// Explained residual energy `‖Λ_:,j‖²` (pre-normalization squared column
143    /// norm) — the factor's contribution to `Σ = c·ΛΛᵀ + D`. Candidates are
144    /// returned in descending energy so the driver promotes the strongest first.
145    pub energy: f64,
146    /// `|cos|` alignment (∈ `[0, 1]`) between this direction and the best-matching
147    /// column of the previous pass's `Λ` — the persistence score gating promotion.
148    pub persistence_alignment: f64,
149    /// Index of the best-matching previous-pass `Λ` column (the nursery lineage
150    /// this candidate continues), so the driver can track a stable identity for a
151    /// direction across passes.
152    pub prev_column: usize,
153}
154
155impl StructuredResidualModel {
156    /// Fit the structured residual-covariance model by the deterministic
157    /// fixed-iteration alternation, selecting the factor rank by the evidence
158    /// ladder. Returns an error only on shape / non-finite-input violations; the
159    /// numerical path is total (every floor and solve is guarded).
160    pub fn fit(input: ResidualFactorInput<'_>) -> Result<Self, String> {
161        let r = input.residuals;
162        let z = input.activity;
163        let n = r.nrows();
164        let p = r.ncols();
165        if n == 0 || p == 0 {
166            return Err(format!(
167                "StructuredResidualModel::fit: residuals must be non-empty; got ({n}, {p})"
168            ));
169        }
170        if z.len() != n {
171            return Err(format!(
172                "StructuredResidualModel::fit: activity length {} != residual rows {n}",
173                z.len()
174            ));
175        }
176        if !r.iter().all(|v| v.is_finite()) {
177            return Err("StructuredResidualModel::fit: residuals must be finite".to_string());
178        }
179        if !z.iter().all(|v| v.is_finite()) {
180            return Err("StructuredResidualModel::fit: activity must be finite".to_string());
181        }
182
183        // Bin assignment for the activity-scale law: deterministic equal-width
184        // bins over the observed z-range. A degenerate (zero-width) range maps
185        // every row to bin 0, recovering a single homoscedastic scale.
186        let bins = ACTIVITY_SCALE_BINS.max(1);
187        let z_min = z.iter().copied().fold(f64::INFINITY, f64::min);
188        let z_max = z.iter().copied().fold(f64::NEG_INFINITY, f64::max);
189        let z_span = z_max - z_min;
190        let row_bin: Vec<usize> = (0..n)
191            .map(|i| {
192                if z_span <= 0.0 {
193                    0
194                } else {
195                    let frac = (z[i] - z_min) / z_span;
196                    let idx = (frac * bins as f64).floor() as isize;
197                    idx.clamp(0, bins as isize - 1) as usize
198                }
199            })
200            .collect();
201
202        let max_rank = input.max_factor_rank.min(p.saturating_sub(1));
203
204        // Evidence ladder over candidate factor ranks. Each candidate is fit by
205        // the full alternation and scored by its penalized Gaussian log-evidence;
206        // the maximizer is kept. Index order breaks any tie (lowest rank wins on
207        // an exact tie — Occam).
208        let mut best: Option<StructuredResidualModel> = None;
209        for rank in 0..=max_rank {
210            let model = Self::fit_fixed_rank(r, &row_bin, bins, rank)?;
211            let take = match &best {
212                None => true,
213                Some(b) => model.log_evidence > b.log_evidence,
214            };
215            if take {
216                best = Some(model);
217            }
218        }
219        best.ok_or_else(|| "StructuredResidualModel::fit: evidence ladder empty".to_string())
220    }
221
222    /// Fit the model at a fixed factor rank by the deterministic alternation.
223    fn fit_fixed_rank(
224        r: ArrayView2<'_, f64>,
225        row_bin: &[usize],
226        bins: usize,
227        rank: usize,
228    ) -> Result<Self, String> {
229        let n = r.nrows();
230        let p = r.ncols();
231
232        // Mean residual variance — the scale reference for the diagonal floor.
233        let mut total_var = 0.0_f64;
234        for i in 0..n {
235            for j in 0..p {
236                total_var += r[[i, j]] * r[[i, j]];
237            }
238        }
239        let mean_var = (total_var / (n as f64 * p as f64)).max(f64::MIN_POSITIVE);
240        let diag_floor = DIAGONAL_REL_FLOOR * mean_var;
241
242        // Initialize the per-row scale to 1 (homoscedastic start), the diagonal
243        // to the per-channel sample variance, and Λ to the leading eigenvectors
244        // of the (scale-1) second moment. The alternation refines all three.
245        let mut row_scale = Array1::<f64>::ones(n);
246        let mut bin_scale = Array1::<f64>::ones(bins);
247        // Raw (undeflated) per-channel second moment — the D estimator's data
248        // term. Constant across sweeps.
249        let raw_diag = column_variances(r);
250        let mut diagonal = raw_diag.mapv(|v| v.max(diag_floor));
251        let mut lambda = Array2::<f64>::zeros((p, rank));
252
253        // Best state seen, restored when a sweep stops improving the evidence.
254        let mut best: Option<(f64, Array2<f64>, Array1<f64>, Array1<f64>)> = None;
255        let mut converged = false;
256        for _sweep in 0..ALTERNATION_MAX_SWEEPS {
257            // (Λ, D | scale): scale-deflated second moment
258            //   S = (1/n) Σ_n (r_n r_nᵀ) / c(z_n).
259            // Under the model E[r_n r_nᵀ] = c_n ΛΛᵀ + D, so S ≈ ΛΛᵀ + D̄ with
260            // D̄ the scale-averaged diagonal; the leading eigenpairs of S − D
261            // give Λ, the residual diagonal gives D.
262            let s = scaled_second_moment(r, &row_scale);
263            let (evals, evecs) = symmetric_eig_ascending(&s)?;
264            // Leading `rank` eigenpairs (eigenvalues ascending ⇒ take the tail).
265            if rank > 0 {
266                for k in 0..rank {
267                    let col = p - 1 - k;
268                    // Factor energy above the idiosyncratic floor: the part of
269                    // the eigenvalue not explained by the mean diagonal.
270                    let mean_diag = diagonal.iter().copied().sum::<f64>() / p as f64;
271                    let energy = (evals[col] - mean_diag).max(0.0);
272                    let amp = energy.sqrt();
273                    for row in 0..p {
274                        lambda[[row, k]] = amp * evecs[[row, col]];
275                    }
276                }
277            }
278            // D update from the RAW (undeflated) moment, floored ≻ 0. The model
279            // is Σ_n = c_n·ΛΛᵀ + D with D NOT scale-multiplied, and c is mean-1
280            // normalized, so E[(1/n)Σ r_n r_nᵀ] = ΛΛᵀ + D exactly. The deflated
281            // moment `s` is the right object for the FACTOR block (its factor
282            // part is scale-free) but its diagonal carries D·mean(1/c) — a
283            // Jensen-inflated D (mean(1/c) > 1 for any non-constant law), which
284            // biased D upward by exactly mean(1/c̃) and let a spurious
285            // higher-rank candidate win the evidence ladder on a better D
286            // alone (the probe's rank-2 winner had a zero second column).
287            for j in 0..p {
288                let mut factor_var = 0.0_f64;
289                for k in 0..rank {
290                    factor_var += lambda[[j, k]] * lambda[[j, k]];
291                }
292                diagonal[j] = (raw_diag[j] - factor_var).max(diag_floor);
293            }
294
295            // (scale | Λ, D): per-row factor activity. With residual r_n, the
296            // factor-subspace energy is r_nᵀ P r_n where P projects onto
297            // range(Λ) in the D-whitened metric; the maximum-likelihood scalar
298            // multiplier on ΛΛᵀ that matches the row's factor-subspace energy is
299            //   c_n = (r̃_nᵀ B (BᵀB)^{-1} Bᵀ r̃_n) / tr(...)-normalizer.
300            // We use a stable closed-form proxy: the row's factor-coordinate
301            // energy ‖Λ⁺ r_n‖² normalized by the unit-scale expectation, then
302            // bin-smoothed across z. With rank 0 there is no factor ⇒ c ≡ 1.
303            if rank > 0 {
304                let mut bin_num = Array1::<f64>::zeros(bins);
305                let mut bin_den = Array1::<f64>::zeros(bins);
306                let coords = factor_coordinates(&lambda, &diagonal, r)?;
307                for i in 0..n {
308                    let mut energy = 0.0_f64;
309                    for k in 0..rank {
310                        energy += coords[[i, k]] * coords[[i, k]];
311                    }
312                    let b = row_bin[i];
313                    bin_num[b] += energy;
314                    bin_den[b] += rank as f64;
315                }
316                // Per-bin mean factor energy = activity scale. Empty bins inherit
317                // the global mean so the scale law stays defined everywhere.
318                let global = {
319                    let num: f64 = bin_num.iter().sum();
320                    let den: f64 = bin_den.iter().sum();
321                    if den > 0.0 { num / den } else { 1.0 }
322                };
323                for b in 0..bins {
324                    bin_scale[b] = if bin_den[b] > 0.0 {
325                        bin_num[b] / bin_den[b]
326                    } else {
327                        global
328                    };
329                }
330                // Smooth (3-point moving average over bins) for a continuous law,
331                // then floor ≻ 0.
332                let scale_floor = SCALE_REL_FLOOR * global.max(f64::MIN_POSITIVE);
333                let smoothed = moving_average_3(&bin_scale);
334                for b in 0..bins {
335                    bin_scale[b] = smoothed[b].max(scale_floor);
336                }
337                // Re-normalize so the mean scale is 1 (the factor amplitude lives
338                // in Λ; c(z) carries only the relative activity law). This keeps
339                // the (Λ, D) ↔ (scale) split identified.
340                //
341                // The mean MUST be taken over ROWS, not over bins. The identity
342                // that makes `raw_diag` an unbiased ΛΛᵀ + D estimator is
343                //   E[(1/n) Σ_n r_n r_nᵀ] = (1/n) Σ_i c(z_i) · ΛΛᵀ + D,
344                // which reduces to ΛΛᵀ + D iff the ROW mean of c is 1:
345                //   (1/n) Σ_i c(z_i) = Σ_b (n_b / n) · bin_scale[b] = 1,
346                // where n_b is the occupancy (row count) of bin b. Under uneven
347                // occupancy (the common case — z is data-driven) the bin-UNIFORM
348                // mean (1/bins) Σ_b bin_scale[b] ≠ this occupancy-weighted mean, so
349                // normalizing by it would leave raw_diag = ΛΛᵀ + D biased by
350                // exactly (occupancy mean / bin mean). Divide by the occupancy-
351                // weighted mean instead, so (1/n) Σ_i row_scale[i] is exactly 1.
352                // ORDERING: the positivity floor was applied above FIRST, so the
353                // floored per-bin values are the ones this normalization sees; the
354                // per-row assignment below therefore needs no second clamp (a
355                // re-clamp would use pre-normalization floor units and break the
356                // exact row-mean-1 invariant just established).
357                let mut bin_count = vec![0.0_f64; bins];
358                for &b in row_bin.iter() {
359                    bin_count[b] += 1.0;
360                }
361                let mean_scale =
362                    (0..bins).map(|b| bin_count[b] * bin_scale[b]).sum::<f64>() / n as f64;
363                if mean_scale > 0.0 {
364                    bin_scale.mapv_inplace(|v| v / mean_scale);
365                }
366                // Each bin_scale[b] is already ≥ scale_floor / mean_scale > 0.
367                for i in 0..n {
368                    row_scale[i] = bin_scale[row_bin[i]];
369                }
370            }
371            // Stop at the first sweep that does not improve the penalized
372            // log-evidence by more than the evidence's own rounding band: the
373            // band is the accumulation of every term the evidence sums, so an
374            // improvement inside it is not one the arithmetic can attest to,
375            // and a decrease is the smoothed, floored scale update overshooting.
376            // The best state is kept either way.
377            let (evidence, band) =
378                penalized_log_evidence_with_band(r, &lambda, &diagonal, &row_scale, rank);
379            match &best {
380                Some((best_evidence, _, _, _)) if evidence <= best_evidence + band => {
381                    converged = true;
382                    break;
383                }
384                _ => {
385                    best = Some((evidence, lambda.clone(), diagonal.clone(), row_scale.clone()));
386                }
387            }
388        }
389        let Some((_, best_lambda, best_diagonal, best_row_scale)) = best else {
390            return Err(format!(
391                "structured residual factor fit: no alternation sweep produced a state \
392                 (rank {rank}, p {p}, n {n})"
393            ));
394        };
395        if !converged {
396            return Err(format!(
397                "structured residual factor fit: the alternation was still improving the \
398                 penalized log-evidence beyond its rounding band after \
399                 {ALTERNATION_MAX_SWEEPS} sweeps (rank {rank}, p {p}, n {n})"
400            ));
401        }
402        let (lambda, diagonal, row_scale) = (best_lambda, best_diagonal, best_row_scale);
403
404        let log_evidence = penalized_log_evidence(r, &lambda, &diagonal, &row_scale, rank);
405        let mut model = Self {
406            p,
407            factor_rank: rank,
408            lambda,
409            diagonal,
410            row_scale,
411            log_evidence,
412        };
413        // Guard against any non-finite leak from a degenerate fit: fall back to a
414        // pure-diagonal model with the same evidence accounting.
415        if !model.is_finite() {
416            model.lambda = Array2::<f64>::zeros((p, rank));
417            model.row_scale = Array1::<f64>::ones(n);
418        }
419        Ok(model)
420    }
421
422    fn is_finite(&self) -> bool {
423        self.lambda.iter().all(|v| v.is_finite())
424            && self.diagonal.iter().all(|v| v.is_finite() && *v > 0.0)
425            && self.row_scale.iter().all(|v| v.is_finite() && *v > 0.0)
426            && self.log_evidence.is_finite()
427    }
428
429    /// Selected factor rank `r`.
430    pub fn factor_rank(&self) -> usize {
431        self.factor_rank
432    }
433
434    /// The fitted interference factor `Λ ∈ ℝ^{p×r}` (the shared off-isotropic
435    /// residual subspace). Consumed by the planted-subspace recovery test to
436    /// compare `range(Λ)` against the planted interference subspace.
437    pub fn factor(&self) -> ArrayView2<'_, f64> {
438        self.lambda.view()
439    }
440
441    /// The idiosyncratic diagonal `d ∈ ℝ^p` (`D = diag(d)`).
442    pub fn diagonal(&self) -> ArrayView1<'_, f64> {
443        self.diagonal.view()
444    }
445
446    /// The penalized Gaussian log-evidence the rank-selection ladder maximized.
447    pub fn log_evidence(&self) -> f64 {
448        self.log_evidence
449    }
450
451    /// #2021 Λ nursery→promotion: detect *persistent, evidence-earning* factor
452    /// directions relative to the previous outer-alternation pass's model.
453    ///
454    /// A column `j` of this model's `Λ` is a [`FactorPromotion`] candidate iff
455    /// both gates hold:
456    /// 1. **Earns its complexity** (evidence gate): its explained energy
457    ///    `‖Λ_:,j‖² ≥ energy_floor_mult · mean(diag(D))`. Every column is already
458    ///    inside the evidence-ladder-selected rank (so it cleared the BIC
459    ///    penalty globally); this per-direction floor additionally requires the
460    ///    factor to explain more than an average channel's idiosyncratic noise,
461    ///    so we never promote a direction that only barely survived rank
462    ///    selection.
463    /// 2. **Persists** (nursery gate): its `|cos|` alignment with the best-
464    ///    matching column of `prev`'s `Λ` is `≥ align_min` — the direction is the
465    ///    same subspace the previous pass already found, not a new
466    ///    pass-to-pass artifact.
467    ///
468    /// Returns candidates sorted by energy (descending). `prev = None` (the first
469    /// structured pass, damping toward `I`) yields no candidates — a direction
470    /// must survive at least one pass to enter the nursery. The driver holds the
471    /// cross-pass persistence count (promote after it clears the direction's
472    /// nursery dwell) and does the actual atom birth; this method is the pure,
473    /// per-pass detector.
474    ///
475    /// Errors on non-finite / out-of-range gates (`align_min ∈ [0,1]`,
476    /// `energy_floor_mult ≥ 0`) or a `prev` with a different output dim `p`.
477    pub fn promotion_candidates(
478        &self,
479        prev: Option<&StructuredResidualModel>,
480        align_min: f64,
481        energy_floor_mult: f64,
482    ) -> Result<Vec<FactorPromotion>, String> {
483        if !align_min.is_finite() || !(0.0..=1.0).contains(&align_min) {
484            return Err(format!(
485                "StructuredResidualModel::promotion_candidates: align_min must be finite in [0,1]; got {align_min}"
486            ));
487        }
488        if !energy_floor_mult.is_finite() || energy_floor_mult < 0.0 {
489            return Err(format!(
490                "StructuredResidualModel::promotion_candidates: energy_floor_mult must be finite and ≥ 0; got {energy_floor_mult}"
491            ));
492        }
493        let prev = match prev {
494            Some(pv) => pv,
495            None => return Ok(Vec::new()),
496        };
497        if prev.p != self.p {
498            return Err(format!(
499                "StructuredResidualModel::promotion_candidates: prev output dim {} != {}",
500                prev.p, self.p
501            ));
502        }
503        let r = self.factor_rank;
504        let prev_r = prev.factor_rank;
505        if r == 0 || prev_r == 0 {
506            return Ok(Vec::new());
507        }
508        // Idiosyncratic-noise floor: a promoted direction must explain more than
509        // an average channel's independent variance.
510        let mean_d = self.diagonal.iter().copied().sum::<f64>() / self.p as f64;
511        let energy_floor = energy_floor_mult * mean_d;
512
513        let mut out: Vec<FactorPromotion> = Vec::new();
514        for j in 0..r {
515            let col = self.lambda.column(j);
516            let energy: f64 = col.iter().map(|v| v * v).sum();
517            if energy <= 0.0 || energy < energy_floor {
518                continue;
519            }
520            let norm = energy.sqrt();
521            // Best |cos| against the previous pass's columns.
522            let mut best_align = 0.0_f64;
523            let mut best_k = 0usize;
524            for k in 0..prev_r {
525                let pcol = prev.lambda.column(k);
526                let pnorm: f64 = pcol.iter().map(|v| v * v).sum::<f64>().sqrt();
527                if pnorm <= 0.0 {
528                    continue;
529                }
530                let dot: f64 = col.iter().zip(pcol.iter()).map(|(a, b)| a * b).sum();
531                let cos_abs = (dot / (norm * pnorm)).abs();
532                if cos_abs > best_align {
533                    best_align = cos_abs;
534                    best_k = k;
535                }
536            }
537            if best_align >= align_min {
538                out.push(FactorPromotion {
539                    direction: col.mapv(|v| v / norm),
540                    energy,
541                    persistence_alignment: best_align,
542                    prev_column: best_k,
543                });
544            }
545        }
546        out.sort_by(|a, b| b.energy.total_cmp(&a.energy));
547        Ok(out)
548    }
549
550    /// Build the per-row precision factor stack `U_n ∈ ℝ^{p×p}` with
551    /// `U_n U_nᵀ = Σ_n^{-1}` and package it as a
552    /// [`MetricProvenance::WhitenedStructured`](gam_problem::MetricProvenance::WhitenedStructured)
553    /// `RowMetric`. This is the single
554    /// production site of `WhitenedStructured`.
555    ///
556    /// The precision is formed in **Woodbury form**:
557    /// ```text
558    ///   Σ_n^{-1} = D^{-1} − D^{-1} Λ ( c^{-1} I_r + Λᵀ D^{-1} Λ )^{-1} Λᵀ D^{-1},
559    /// ```
560    /// an `r × r` capacitance solve (never a `p × p` inverse). The factor `U_n`
561    /// is the lower-Cholesky of the assembled `Σ_n^{-1}` (`rank = p`), so
562    /// `whiten_residual_row` returns coordinates whose squared norm is the exact
563    /// Mahalanobis residual `r_nᵀ Σ_n^{-1} r_n`.
564    pub fn row_metric(&self, n_rows: usize) -> Result<RowMetric, String> {
565        if n_rows != self.row_scale.len() {
566            return Err(format!(
567                "StructuredResidualModel::row_metric: requested {n_rows} rows but model has {}",
568                self.row_scale.len()
569            ));
570        }
571        let p = self.p;
572        let r = self.factor_rank;
573        // Hoist every row-INDEPENDENT Woodbury part out of the per-row loop: the
574        // inverse diagonal D^{-1}, B = D^{-1}Λ, its transpose Bᵀ, and the Gram
575        // M0 = ΛᵀD^{-1}Λ. Only the c_n^{-1} I_r shift on the capacitance is
576        // per-row, so the per-row capacitance is M_n = M0 + c_n^{-1} I_r — a
577        // scalar-diagonal reweight of the SAME M0 (mirroring the Fix-B hoist in
578        // `penalized_log_evidence`). Building the n-row U_n stack now costs
579        // O(p·r² + n·(p·r + r³ + p³)) instead of rebuilding B and the Gram every
580        // row. The summation order per row is unchanged, so the assembled U_n is
581        // bit-for-bit identical to the per-row-rebuild it replaces.
582        let d_inv: Vec<f64> = (0..p).map(|i| 1.0 / self.diagonal[i]).collect();
583        let mut b = Array2::<f64>::zeros((p, r));
584        let mut bt = Array2::<f64>::zeros((r, p));
585        let mut m0 = Array2::<f64>::zeros((r, r));
586        if r > 0 {
587            for i in 0..p {
588                for k in 0..r {
589                    b[[i, k]] = d_inv[i] * self.lambda[[i, k]];
590                }
591            }
592            for a in 0..r {
593                for bk in 0..r {
594                    let mut acc = 0.0_f64;
595                    for i in 0..p {
596                        acc += self.lambda[[i, a]] * b[[i, bk]];
597                    }
598                    m0[[a, bk]] = acc;
599                }
600            }
601            for k in 0..r {
602                for i in 0..p {
603                    bt[[k, i]] = b[[i, k]];
604                }
605            }
606        }
607        // Row-major flat factor matrix: u[n, i*p + k] = U_n[i, k]. Each row's
608        // Woodbury assemble + p×p Cholesky is independent of every other row's
609        // (no cross-row reduction), so the stack parallelizes over rows with
610        // BIT-IDENTICAL output — every row runs the exact serial arithmetic and
611        // writes only its own p² chunk. This was the dominant serial wall of the
612        // #974 metric install (n_rows × O(p³) on one core while the inner fit
613        // parallelizes cleanly). Same engagement discipline as
614        // `scaled_second_moment`: only above a row threshold (serial avoids
615        // rayon overhead on small stacks) and only when not already inside a
616        // rayon worker (nested calls keep the outer region's cores). Error
617        // selection stays deterministic: the indexed collect preserves row
618        // order, and the first `Some` scanned in that order is the same
619        // lowest-row error the serial loop returned.
620        let mut u = Array2::<f64>::zeros((n_rows, p * p));
621        let first_error = {
622            use rayon::prelude::*;
623            const PARALLEL_ROW_MIN: usize = 64;
624            let build_row = |row: usize, urow: &mut [f64]| -> Option<String> {
625                let precision = match self.row_precision(&d_inv, &b, &bt, &m0, row) {
626                    Ok(m) => m,
627                    Err(err) => return Some(err),
628                };
629                let factor = match lower_cholesky_psd(&precision) {
630                    Ok(f) => f,
631                    Err(err) => return Some(err),
632                };
633                for i in 0..p {
634                    for k in 0..p {
635                        urow[i * p + k] = factor[[i, k]];
636                    }
637                }
638                None
639            };
640            let u_flat = u.as_slice_mut().ok_or_else(|| {
641                "StructuredResidualModel::row_metric: factor stack must be standard-layout"
642                    .to_string()
643            })?;
644            if p > 0 && n_rows >= PARALLEL_ROW_MIN && rayon::current_thread_index().is_none() {
645                u_flat
646                    .par_chunks_mut(p * p)
647                    .enumerate()
648                    .map(|(row, urow)| build_row(row, urow))
649                    .collect::<Vec<Option<String>>>()
650                    .into_iter()
651                    .flatten()
652                    .next()
653            } else if p > 0 {
654                u_flat
655                    .chunks_mut(p * p)
656                    .enumerate()
657                    .find_map(|(row, urow)| build_row(row, urow))
658            } else {
659                None
660            }
661        };
662        if let Some(err) = first_error {
663            return Err(err);
664        }
665        RowMetric::whitened_structured(Arc::new(u), p, p)
666    }
667
668    /// The model's isotropic (iid-MLE) dispersion: the per-coordinate average of
669    /// its own fitted total residual variance,
670    /// ```text
671    ///   φ̂ = (1/p) · mean_row tr(Σ̂(row)) = mean(c)·‖Λ‖²_F / p + mean(d).
672    /// ```
673    /// This is the honest single-scalar summary of the SAME second moment the
674    /// structured model fits — the scale an iid Gaussian residual model would
675    /// estimate from these residuals. Used as the first-pass damping anchor in
676    /// [`Self::row_metric_damped`] (#2243 cap #2): anchoring at `φ̂·I` instead of
677    /// the unit `I_p` means near-noiseless data is whitened by its MEASURED
678    /// noise, so the downstream unit-dispersion REML criterion prices the
679    /// smoothing penalty against the real dispersion rather than an assumed
680    /// unit one. Floored at `f64::MIN_POSITIVE` so the blend stays SPD.
681    pub fn isotropic_dispersion(&self) -> f64 {
682        let p = self.p.max(1) as f64;
683        let n = self.row_scale.len().max(1) as f64;
684        let mean_c = self.row_scale.iter().copied().sum::<f64>() / n;
685        let lambda_energy: f64 = self.lambda.iter().map(|v| v * v).sum();
686        let mean_d = self.diagonal.iter().copied().sum::<f64>() / p;
687        (mean_c * lambda_energy / p + mean_d).max(f64::MIN_POSITIVE)
688    }
689
690    /// Damped per-row metric for the #2021 driver: blend covariances in the
691    /// **covariance domain** (before the Woodbury→Cholesky) between this model's
692    /// estimate and a previous one,
693    /// ```text
694    ///   Σ_t(row) = (1 − γ) · Σ_prev(row) + γ · Σ̂_t(row),
695    /// ```
696    /// where `Σ̂_t(row) = c_t(z)·ΛΛᵀ + D` is this model's per-row covariance
697    /// (built from the hoisted-M0 / occupancy-weighted `c(z)` path), and
698    /// `Σ_prev(row)` is `prev`'s per-row covariance when `Some`, else the
699    /// MEASURED iid anchor `φ̂·I_p` ([`Self::isotropic_dispersion`], #2243 cap
700    /// #2: a unit `I_p` anchor silently assumed unit noise, which on
701    /// near-noiseless data pinned the whitened likelihood — and therefore the
702    /// REML smoothing balance — at a noise scale ~1/φ̂ too coarse, i.e. the
703    /// clean-data over-penalization).
704    ///
705    /// Endpoints (exact):
706    /// * `γ = 1.0` ⇒ this model's [`Self::row_metric`] exactly (Woodbury path,
707    ///   byte-identical);
708    /// * `γ = 0.0` ⇒ `prev`'s [`Self::row_metric`] when `Some` (byte-identical),
709    ///   else the measured-scale identity `(φ̂·I)^{-1}` factors.
710    ///
711    /// `γ` must be finite and in `[0, 1]`; when `prev` is `Some` it must share
712    /// this model's `p` and row count.
713    pub fn row_metric_damped(
714        &self,
715        n_rows: usize,
716        gamma: f64,
717        prev: Option<&StructuredResidualModel>,
718    ) -> Result<RowMetric, String> {
719        if n_rows != self.row_scale.len() {
720            return Err(format!(
721                "StructuredResidualModel::row_metric_damped: requested {n_rows} rows but model has {}",
722                self.row_scale.len()
723            ));
724        }
725        if !gamma.is_finite() || !(0.0..=1.0).contains(&gamma) {
726            return Err(format!(
727                "StructuredResidualModel::row_metric_damped: gamma must be finite in [0,1]; got {gamma}"
728            ));
729        }
730        if let Some(pv) = prev {
731            if pv.p != self.p {
732                return Err(format!(
733                    "StructuredResidualModel::row_metric_damped: prev output dim {} != {}",
734                    pv.p, self.p
735                ));
736            }
737            if pv.row_scale.len() != n_rows {
738                return Err(format!(
739                    "StructuredResidualModel::row_metric_damped: prev has {} rows but requested {n_rows}",
740                    pv.row_scale.len()
741                ));
742            }
743        }
744        // Exact endpoints — reuse the undamped producers so the result is
745        // byte-identical (γ=1 ⇒ this model; γ=0 ⇒ prev, or Euclidean identity).
746        if gamma == 1.0 {
747            return self.row_metric(n_rows);
748        }
749        if gamma == 0.0 {
750            if let Some(pv) = prev {
751                return pv.row_metric(n_rows);
752            }
753            // prev = None falls through to the general path: the blend is then
754            // exactly the measured-scale identity `φ̂·I` (#2243 cap #2), not the
755            // unit Euclidean identity.
756        }
757
758        let p = self.p;
759        // #2243 cap #2 — the first-pass (prev = None) damping anchor is the
760        // model's own measured isotropic dispersion, hoisted out of the row loop.
761        let iid_anchor = self.isotropic_dispersion();
762        // Row-INDEPENDENT outer products ΛΛᵀ (this model and, if present, prev):
763        // only the per-row activity scale c(z) multiplies them, so hoist the Gram
764        // out of the per-row loop (mirroring the row_metric / penalized_log_evidence
765        // hoist).
766        let self_gram = outer_product(&self.lambda);
767        let prev_gram = prev.map(|pv| outer_product(&pv.lambda));
768
769        // Per-row blend + invert + Cholesky, parallelized over rows exactly as
770        // in [`Self::row_metric`]: rows are independent (each writes only its
771        // own p² chunk of `u`, no cross-row reduction), so the parallel stack is
772        // bit-identical to the serial one, with the same engagement discipline
773        // (row threshold, never nested inside a rayon worker) and the same
774        // deterministic lowest-row error selection.
775        let mut u = Array2::<f64>::zeros((n_rows, p * p));
776        let first_error = {
777            use rayon::prelude::*;
778            const PARALLEL_ROW_MIN: usize = 64;
779            let build_row = |row: usize, urow: &mut [f64]| -> Option<String> {
780                self.damped_row_factor(
781                    row,
782                    gamma,
783                    prev,
784                    iid_anchor,
785                    &self_gram,
786                    prev_gram.as_ref(),
787                    urow,
788                )
789                .err()
790            };
791            let u_flat = u
792                .as_slice_mut()
793                .ok_or_else(|| "StructuredResidualModel::row_metric_damped: factor stack must be standard-layout".to_string())?;
794            if p > 0 && n_rows >= PARALLEL_ROW_MIN && rayon::current_thread_index().is_none() {
795                u_flat
796                    .par_chunks_mut(p * p)
797                    .enumerate()
798                    .map(|(row, urow)| build_row(row, urow))
799                    .collect::<Vec<Option<String>>>()
800                    .into_iter()
801                    .flatten()
802                    .next()
803            } else if p > 0 {
804                u_flat
805                    .chunks_mut(p * p)
806                    .enumerate()
807                    .find_map(|(row, urow)| build_row(row, urow))
808            } else {
809                None
810            }
811        };
812        if let Some(err) = first_error {
813            return Err(err);
814        }
815        RowMetric::whitened_structured(Arc::new(u), p, p)
816    }
817
818    /// One row of the damped-metric stack: assemble the convex-blend covariance
819    /// `Σ_t(row) = γ·(c·ΛΛᵀ + D) + (1−γ)·Σ_prev(row)`, symmetrize, invert, and
820    /// write the lower-Cholesky precision factor into `urow` (row-major `p×p`).
821    /// Factored out of [`Self::row_metric_damped`] so the serial and parallel
822    /// row drivers share one arithmetic body. `iid_anchor` is the measured
823    /// isotropic dispersion `φ̂` used as `Σ_prev = φ̂·I` when `prev` is `None`
824    /// (#2243 cap #2).
825    fn damped_row_factor(
826        &self,
827        row: usize,
828        gamma: f64,
829        prev: Option<&StructuredResidualModel>,
830        iid_anchor: f64,
831        self_gram: &Array2<f64>,
832        prev_gram: Option<&Array2<f64>>,
833        urow: &mut [f64],
834    ) -> Result<(), String> {
835        let p = self.p;
836        let c = self.row_scale[row].max(f64::MIN_POSITIVE);
837        // γ · Σ̂_t = γ·(c·ΛΛᵀ + D).
838        let mut sigma = Array2::<f64>::zeros((p, p));
839        for a in 0..p {
840            for b in 0..p {
841                sigma[[a, b]] = gamma * c * self_gram[[a, b]];
842            }
843            sigma[[a, a]] += gamma * self.diagonal[a];
844        }
845        // (1−γ) · Σ_prev  (prev's per-row Σ, or I_p when prev is None).
846        match (prev, prev_gram) {
847            (Some(pv), Some(pg)) => {
848                let cp = pv.row_scale[row].max(f64::MIN_POSITIVE);
849                for a in 0..p {
850                    for b in 0..p {
851                        sigma[[a, b]] += (1.0 - gamma) * cp * pg[[a, b]];
852                    }
853                    sigma[[a, a]] += (1.0 - gamma) * pv.diagonal[a];
854                }
855            }
856            (None, _) => {
857                // First-pass anchor: the MEASURED iid dispersion φ̂·I, not the
858                // unit identity (#2243 cap #2 — clean-data over-penalization).
859                for a in 0..p {
860                    sigma[[a, a]] += (1.0 - gamma) * iid_anchor;
861                }
862            }
863            (Some(_), None) => {
864                return Err(
865                    "previous residual model supplied without its precomputed gram".to_string()
866                );
867            }
868        }
869        // Symmetrize against round-off before inversion.
870        for a in 0..p {
871            for b in (a + 1)..p {
872                let avg = 0.5 * (sigma[[a, b]] + sigma[[b, a]]);
873                sigma[[a, b]] = avg;
874                sigma[[b, a]] = avg;
875            }
876        }
877        // Σ_t is a convex combination of SPD matrices (D ≻ 0 / I ≻ 0) ⇒ SPD.
878        // Precision = Σ_t^{-1} via a Cholesky solve against I_p, then the U_n
879        // factor is the lower-Cholesky of the precision (row_metric's U
880        // convention).
881        let precision = invert_spd(&sigma)?;
882        let factor = lower_cholesky_psd(&precision)?;
883        for i in 0..p {
884            for k in 0..p {
885                urow[i * p + k] = factor[[i, k]];
886            }
887        }
888        Ok(())
889    }
890
891    /// Per-row precision `Σ_n^{-1}` via the Woodbury identity (an `r × r` solve),
892    /// given the row-independent parts precomputed by [`Self::row_metric`]:
893    /// `d_inv = D^{-1}`, `b = D^{-1}Λ`, `bt = Bᵀ`, and the Gram `m0 = ΛᵀD^{-1}Λ`.
894    /// Only the per-row capacitance `M_n = m0 + c_n^{-1} I_r` and the back-solve
895    /// depend on the row.
896    fn row_precision(
897        &self,
898        d_inv: &[f64],
899        b: &Array2<f64>,
900        bt: &Array2<f64>,
901        m0: &Array2<f64>,
902        row: usize,
903    ) -> Result<Array2<f64>, String> {
904        let p = self.p;
905        let r = self.factor_rank;
906        // Start from D^{-1}.
907        let mut precision = Array2::<f64>::zeros((p, p));
908        for i in 0..p {
909            precision[[i, i]] = d_inv[i];
910        }
911        if r == 0 {
912            return Ok(precision);
913        }
914        let c = self.row_scale[row].max(f64::MIN_POSITIVE);
915        // Per-row capacitance M_n = M0 + c^{-1} I_r (copy the hoisted Gram, then
916        // add c^{-1} to the diagonal). M_n ≻ 0 since c^{-1} > 0 and M0 ⪰ 0.
917        let mut cap = m0.clone();
918        for a in 0..r {
919            cap[[a, a]] += 1.0 / c;
920        }
921        // Σ_n^{-1} = D^{-1} − B M_n^{-1} Bᵀ. Solve M_n X = Bᵀ for X = M_n^{-1} Bᵀ
922        // (r × p) via Cholesky.
923        let chol = cap
924            .cholesky(Side::Lower)
925            .map_err(|e| format!("StructuredResidualModel::row_precision capacitance: {e:?}"))?;
926        let x = chol.solve_mat(bt); // r × p
927        for i in 0..p {
928            for j in 0..p {
929                let mut acc = 0.0_f64;
930                for k in 0..r {
931                    acc += b[[i, k]] * x[[k, j]];
932                }
933                precision[[i, j]] -= acc;
934            }
935        }
936        // Symmetrize against round-off so the Cholesky downstream sees an exactly
937        // symmetric PSD matrix.
938        for i in 0..p {
939            for j in (i + 1)..p {
940                let avg = 0.5 * (precision[[i, j]] + precision[[j, i]]);
941                precision[[i, j]] = avg;
942                precision[[j, i]] = avg;
943            }
944        }
945        Ok(precision)
946    }
947}
948
949/// Outer product `Λ Λᵀ ∈ ℝ^{p×p}` of a factor matrix `Λ ∈ ℝ^{p×r}` — the
950/// row-independent factor covariance the per-row activity scale multiplies.
951/// Used by [`StructuredResidualModel::row_metric_damped`] to hoist the Gram out
952/// of its per-row covariance-blend loop.
953fn outer_product(lambda: &Array2<f64>) -> Array2<f64> {
954    let p = lambda.nrows();
955    let r = lambda.ncols();
956    let mut g = Array2::<f64>::zeros((p, p));
957    for a in 0..p {
958        for b in 0..p {
959            let mut acc = 0.0_f64;
960            for k in 0..r {
961                acc += lambda[[a, k]] * lambda[[b, k]];
962            }
963            g[[a, b]] = acc;
964        }
965    }
966    g
967}
968
969/// Inverse of a symmetric positive-definite matrix via a Cholesky solve against
970/// the identity, symmetrized against round-off. Used to form `Σ_t^{-1}` from a
971/// densely-blended covariance in [`StructuredResidualModel::row_metric_damped`]
972/// (the blended covariance is no longer low-rank-plus-diagonal, so Woodbury does
973/// not apply).
974fn invert_spd(a: &Array2<f64>) -> Result<Array2<f64>, String> {
975    let p = a.nrows();
976    let chol = a
977        .cholesky(Side::Lower)
978        .map_err(|e| format!("invert_spd: blended covariance not SPD: {e:?}"))?;
979    let mut inv = chol.solve_mat(&Array2::<f64>::eye(p));
980    for i in 0..p {
981        for j in (i + 1)..p {
982            let avg = 0.5 * (inv[[i, j]] + inv[[j, i]]);
983            inv[[i, j]] = avg;
984            inv[[j, i]] = avg;
985        }
986    }
987    Ok(inv)
988}
989
990/// Per-channel (column) sample second moment of the residual matrix.
991fn column_variances(r: ArrayView2<'_, f64>) -> Array1<f64> {
992    let n = r.nrows();
993    let p = r.ncols();
994    let mut v = Array1::<f64>::zeros(p);
995    for j in 0..p {
996        let mut acc = 0.0_f64;
997        for i in 0..n {
998            acc += r[[i, j]] * r[[i, j]];
999        }
1000        v[j] = acc / n as f64;
1001    }
1002    v
1003}
1004
1005/// Scale-deflated second moment `S = (1/n) Σ_n (r_n r_nᵀ) / c_n`.
1006/// Per-row-chunk contribution to the scaled second moment — the inner
1007/// `p×p` accumulation of one contiguous row block, summed in row order.
1008fn scaled_second_moment_chunk(
1009    r: ArrayView2<'_, f64>,
1010    row_scale: &Array1<f64>,
1011    lo: usize,
1012    hi: usize,
1013) -> Array2<f64> {
1014    let p = r.ncols();
1015    let mut s = Array2::<f64>::zeros((p, p));
1016    for i in lo..hi {
1017        let w = 1.0 / row_scale[i].max(f64::MIN_POSITIVE);
1018        for a in 0..p {
1019            let ra = r[[i, a]];
1020            for b in 0..p {
1021                s[[a, b]] += w * ra * r[[i, b]];
1022            }
1023        }
1024    }
1025    s
1026}
1027
1028/// `S = (1/n) Σ_n (r_n r_nᵀ) / c(z_n)` — the O(N·p²) scale-deflated second moment
1029/// that dominates each alternation sweep of the residual-factor fit.
1030///
1031/// Reduced over the deterministic length-only pairwise tree
1032/// [`par_deterministic_block_fold`]: the `p×p` partial of each `BASE_CHUNK`-row
1033/// base block is combined by the same `left_split` association regardless of
1034/// thread count OR nesting, so the result is a pure function of the ordered rows
1035/// (#2228 reduction doctrine — parallel and nested-serial evaluation are
1036/// bit-identical, not merely run-to-run reproducible). The tree self-serializes
1037/// below `BASE_CHUNK` rows (a base block is folded directly with no `rayon::join`),
1038/// so small inputs and nested calls stay on a single thread without a separate
1039/// branch that could associate the round-off differently.
1040fn scaled_second_moment(r: ArrayView2<'_, f64>, row_scale: &Array1<f64>) -> Array2<f64> {
1041    use gam_linalg::pairwise_reduce::par_deterministic_block_fold;
1042    let n = r.nrows();
1043    let p = r.ncols();
1044
1045    let mut s = par_deterministic_block_fold(
1046        n,
1047        |range: core::ops::Range<usize>| {
1048            scaled_second_moment_chunk(r, row_scale, range.start, range.end)
1049        },
1050        |mut acc: Array2<f64>, part: Array2<f64>| {
1051            acc += &part;
1052            acc
1053        },
1054    )
1055    .unwrap_or_else(|| Array2::<f64>::zeros((p, p)));
1056
1057    s.mapv_inplace(|v| v / n as f64);
1058    // Symmetrize against accumulation round-off.
1059    for a in 0..p {
1060        for b in (a + 1)..p {
1061            let avg = 0.5 * (s[[a, b]] + s[[b, a]]);
1062            s[[a, b]] = avg;
1063            s[[b, a]] = avg;
1064        }
1065    }
1066    s
1067}
1068
1069/// Factor coordinates `Λ⁺_D r_n` per row: the generalized-least-squares
1070/// projection of each residual onto `range(Λ)` in the `D^{-1}` metric, returned
1071/// as an `n × r` matrix. Solves the `r × r` normal equations
1072/// `(Λᵀ D^{-1} Λ) γ = Λᵀ D^{-1} r_n` per row (shared factorization).
1073fn factor_coordinates(
1074    lambda: &Array2<f64>,
1075    diagonal: &Array1<f64>,
1076    r: ArrayView2<'_, f64>,
1077) -> Result<Array2<f64>, String> {
1078    let p = lambda.nrows();
1079    let rank = lambda.ncols();
1080    let n = r.nrows();
1081    // GLS weights 1/D_ii, with zero-variance channels DROPPED (weight 0): a
1082    // channel whose residual is identically zero carries no factor information,
1083    // and its 1/0 = ∞ weight poisons the whole normal matrix into NaN — the
1084    // fully-explained-target abort (Cholesky NonPositivePivot) that killed
1085    // stagewise runs on targets the dictionary explains exactly. Dropping the
1086    // channel is the pseudo-inverse limit; with every channel degenerate the
1087    // ridged normal matrix stays PD and the coordinates are the least-norm 0.
1088    let d_inv: Vec<f64> = (0..p)
1089        .map(|i| {
1090            let d = diagonal[i];
1091            if !(d > 0.0 && d.is_finite()) {
1092                return 0.0;
1093            }
1094            // A subnormal-floored variance (the zero-residual case floors the
1095            // scale reference at f64::MIN_POSITIVE) passes `d > 0` but its
1096            // reciprocal OVERFLOWS to ∞ — the same NaN poisoning through the
1097            // second door. A non-finite weight is the same degenerate-channel
1098            // verdict: drop it.
1099            let w = d.recip();
1100            if w.is_finite() { w } else { 0.0 }
1101        })
1102        .collect();
1103    // Normal matrix ΛᵀD^{-1}Λ (+ tiny ridge for invertibility).
1104    let mut normal = Array2::<f64>::zeros((rank, rank));
1105    for a in 0..rank {
1106        for b in 0..rank {
1107            let mut acc = 0.0_f64;
1108            for i in 0..p {
1109                acc += lambda[[i, a]] * d_inv[i] * lambda[[i, b]];
1110            }
1111            normal[[a, b]] = acc;
1112        }
1113    }
1114    let trace = (0..rank).map(|k| normal[[k, k]]).sum::<f64>().max(1.0);
1115    let ridge = 1e-10 * trace / rank.max(1) as f64;
1116    for k in 0..rank {
1117        normal[[k, k]] += ridge;
1118    }
1119    let chol = normal
1120        .cholesky(Side::Lower)
1121        .map_err(|e| format!("factor_coordinates normal solve: {e:?}"))?;
1122    let mut coords = Array2::<f64>::zeros((n, rank));
1123    let mut rhs = Array1::<f64>::zeros(rank);
1124    for i in 0..n {
1125        for a in 0..rank {
1126            let mut acc = 0.0_f64;
1127            for j in 0..p {
1128                acc += lambda[[j, a]] * d_inv[j] * r[[i, j]];
1129            }
1130            rhs[a] = acc;
1131        }
1132        let gamma = chol.solvevec(&rhs);
1133        for a in 0..rank {
1134            coords[[i, a]] = gamma[a];
1135        }
1136    }
1137    Ok(coords)
1138}
1139
1140/// 3-point moving average over a bin vector (edge-clamped), giving the smooth
1141/// activity-scale law a continuous, low-curvature shape.
1142fn moving_average_3(v: &Array1<f64>) -> Array1<f64> {
1143    let m = v.len();
1144    let mut out = Array1::<f64>::zeros(m);
1145    for i in 0..m {
1146        let lo = i.saturating_sub(1);
1147        let hi = (i + 1).min(m - 1);
1148        let mut acc = 0.0_f64;
1149        let mut cnt = 0.0_f64;
1150        for j in lo..=hi {
1151            acc += v[j];
1152            cnt += 1.0;
1153        }
1154        out[i] = acc / cnt;
1155    }
1156    out
1157}
1158
1159/// Ascending-eigenvalue symmetric eigendecomposition (faer convention).
1160fn symmetric_eig_ascending(m: &Array2<f64>) -> Result<(Array1<f64>, Array2<f64>), String> {
1161    m.eigh(Side::Lower)
1162        .map_err(|e| format!("symmetric_eig: {e:?}"))
1163}
1164
1165/// Lower-triangular Cholesky factor `L` of a (numerically) PSD matrix `A` with
1166/// `L Lᵀ = A`, with a relative spectral floor so a marginally-indefinite
1167/// precision (round-off) still factors. Used to turn `Σ_n^{-1}` into the
1168/// `RowMetric` factor `U_n` (here `U_n = L`).
1169fn lower_cholesky_psd(a: &Array2<f64>) -> Result<Array2<f64>, String> {
1170    if let Ok(chol) = a.cholesky(Side::Lower) {
1171        return Ok(chol.lower_triangular());
1172    }
1173    // Eigen-repair: clamp eigenvalues to a small positive floor, rebuild the
1174    // REPAIRED matrix Q·diag(λ_clamped)·Qᵀ itself, and Cholesky that (always
1175    // succeeds, PD). The returned factor must satisfy L·Lᵀ = A_repaired —
1176    // rebuilding the symmetric square root here and factoring THAT would hand
1177    // callers a factor with L·Lᵀ = A^{1/2}, silently taking every whitened
1178    // quadratic form against the square root of the intended precision.
1179    let (evals, evecs) = symmetric_eig_ascending(a)?;
1180    let max_ev = evals.iter().copied().fold(0.0_f64, f64::max).max(1.0);
1181    let floor = 1e-10 * max_ev;
1182    let p = a.nrows();
1183    let mut repaired = Array2::<f64>::zeros((p, p));
1184    for i in 0..p {
1185        for j in 0..p {
1186            let mut acc = 0.0_f64;
1187            for k in 0..p {
1188                let ev = evals[k].max(floor);
1189                acc += evecs[[i, k]] * ev * evecs[[j, k]];
1190            }
1191            repaired[[i, j]] = acc;
1192        }
1193    }
1194    repaired
1195        .cholesky(Side::Lower)
1196        .map(|c| c.lower_triangular())
1197        .map_err(|e| format!("lower_cholesky_psd eigen-repair: {e:?}"))
1198}
1199
1200/// Penalized Gaussian log-evidence of the structured model at the fitted
1201/// parameters — the evidence ladder's rank-selection score.
1202///
1203/// The per-row log-density of `r_n ~ N(0, Σ_n)` is
1204/// `−½ ( log|Σ_n| + r_nᵀ Σ_n^{-1} r_n + p log 2π )`. We sum it across rows and
1205/// subtract a parameter-count penalty `½ k_params · log n` (a BIC-style Occam
1206/// term over the `p·r` factor entries + `p` diagonal entries + the bin scales),
1207/// so adding a spurious factor that does not improve the fit is rejected. Both
1208/// `log|Σ_n|` and the quadratic use the Woodbury / matrix-determinant lemma so no
1209/// dense `p × p` inverse or determinant is formed.
1210fn penalized_log_evidence(
1211    r: ArrayView2<'_, f64>,
1212    lambda: &Array2<f64>,
1213    diagonal: &Array1<f64>,
1214    row_scale: &Array1<f64>,
1215    rank: usize,
1216) -> f64 {
1217    penalized_log_evidence_with_band(r, lambda, diagonal, row_scale, rank).0
1218}
1219
1220/// The penalized log-evidence and its rounding band: Wilkinson's growth factor
1221/// for the number of floating-point operations the evidence accumulates, times
1222/// the sum of the magnitudes of every term it adds, so a difference between
1223/// two evidence values inside the band is not attested by the arithmetic
1224/// (#2469).
1225fn penalized_log_evidence_with_band(
1226    r: ArrayView2<'_, f64>,
1227    lambda: &Array2<f64>,
1228    diagonal: &Array1<f64>,
1229    row_scale: &Array1<f64>,
1230    rank: usize,
1231) -> (f64, f64) {
1232    let n = r.nrows();
1233    let p = r.ncols();
1234    let d_inv: Vec<f64> = (0..p).map(|i| 1.0 / diagonal[i]).collect();
1235    let log_det_d: f64 = diagonal.iter().map(|&d| d.ln()).sum();
1236    let two_pi_ln = (2.0 * std::f64::consts::PI).ln();
1237
1238    // Row-INDEPENDENT Gram M0 = ΛᵀD^{-1}Λ (r × r). This does not depend on the
1239    // row, so build it ONCE here rather than rebuilding it inside the per-row loop
1240    // (which was O(n·p·r²)). The per-row capacitance is only a scalar-diagonal
1241    // reweight of this SAME M0 — M_n = M0 + (1/c_n) I_r — so each row copies M0 and
1242    // adds 1/c_n to its diagonal (cheap, O(r)) before its own Cholesky. The
1243    // summation order over j (0..p) is preserved exactly and the diagonal add is
1244    // the identical `+= 1.0 / c` op, so the hoist is bit-for-bit identical to the
1245    // pre-hoist per-row rebuild (same log|Σ_n|, same quadratic, same evidence).
1246    let mut m0 = Array2::<f64>::zeros((rank, rank));
1247    if rank > 0 {
1248        for a in 0..rank {
1249            for b in 0..rank {
1250                let mut acc = 0.0_f64;
1251                for j in 0..p {
1252                    acc += lambda[[j, a]] * d_inv[j] * lambda[[j, b]];
1253                }
1254                m0[[a, b]] = acc;
1255            }
1256        }
1257    }
1258
1259    let mut log_lik = 0.0_f64;
1260    let mut magnitude = 0.0_f64;
1261    for i in 0..n {
1262        let c = row_scale[i].max(f64::MIN_POSITIVE);
1263        // Quadratic r_nᵀ Σ_n^{-1} r_n via Woodbury:
1264        //   r_nᵀ D^{-1} r_n − (Bᵀ r_n)ᵀ M^{-1} (Bᵀ r_n),
1265        // with B = D^{-1}Λ and M = c^{-1}I + ΛᵀD^{-1}Λ.
1266        let mut quad = 0.0_f64;
1267        for j in 0..p {
1268            quad += r[[i, j]] * d_inv[j] * r[[i, j]];
1269        }
1270        let mut log_det = log_det_d;
1271        if rank > 0 {
1272            // Per-row capacitance M_n = M0 + (1/c) I_r (copy the hoisted M0, then
1273            // add 1/c to the diagonal), and w = Bᵀ r_n = ΛᵀD^{-1} r_n.
1274            let mut m = m0.clone();
1275            for a in 0..rank {
1276                m[[a, a]] += 1.0 / c;
1277            }
1278            let mut w = Array1::<f64>::zeros(rank);
1279            for a in 0..rank {
1280                let mut wa = 0.0_f64;
1281                for j in 0..p {
1282                    wa += lambda[[j, a]] * d_inv[j] * r[[i, j]];
1283                }
1284                w[a] = wa;
1285            }
1286            // Cholesky M = R Rᵀ → log|M|, and solve M y = w.
1287            match m.cholesky(Side::Lower) {
1288                Ok(chol) => {
1289                    let y = chol.solvevec(&w);
1290                    let mut wy = 0.0_f64;
1291                    for a in 0..rank {
1292                        wy += w[a] * y[a];
1293                    }
1294                    quad -= wy;
1295                    // log|Σ_n| = log|D| + log|M| + r·log c   (matrix-determinant
1296                    // lemma; the c^{-1}I shift carries the +r·log c).
1297                    let diag = chol.diag();
1298                    let log_det_m: f64 = diag.iter().map(|&l| (l * l).ln()).sum();
1299                    log_det = log_det_d + log_det_m + rank as f64 * c.ln();
1300                }
1301                Err(_) => {
1302                    // Degenerate capacitance — fall back to the diagonal model's
1303                    // accounting for this row (no factor correction).
1304                    log_det = log_det_d;
1305                }
1306            }
1307        }
1308        log_lik += -0.5 * (log_det + quad + p as f64 * two_pi_ln);
1309        magnitude += 0.5 * (log_det.abs() + quad.abs() + p as f64 * two_pi_ln);
1310    }
1311
1312    let k_params = (p * rank + p + ACTIVITY_SCALE_BINS) as f64;
1313    let penalty = 0.5 * k_params * (n.max(2) as f64).ln();
1314    magnitude += penalty.abs();
1315    // Operations accumulated per row: the quadratic form (`p`), the factor
1316    // projections (`rank·p`), the capacitance solve (`rank²`) and the
1317    // `rank`-term contraction; plus the `p` log-determinant terms and the
1318    // `rank²·p` capacitance assembly once.
1319    let operations = n * (p * (1 + rank) + rank * (rank + 1)) + p * (1 + rank * rank);
1320    let band = gam_linalg::roundoff::accumulation_growth(operations) * magnitude;
1321    (log_lik - penalty, band)
1322}
1323
1324#[cfg(test)]
1325mod tests {
1326    use super::*;
1327    use ndarray::{Array1, Array2};
1328
1329    fn lcg_uniform(state: &mut u64) -> f64 {
1330        *state = state
1331            .wrapping_mul(6364136223846793005)
1332            .wrapping_add(1442695040888963407);
1333        ((*state >> 11) as f64) / ((1u64 << 53) as f64)
1334    }
1335
1336    fn lcg_normal(state: &mut u64) -> f64 {
1337        let u1 = lcg_uniform(state).max(1e-12);
1338        let u2 = lcg_uniform(state);
1339        (-2.0 * u1.ln()).sqrt() * (std::f64::consts::TAU * u2).cos()
1340    }
1341
1342    /// Per-rank evidence breakdown on the planted single-factor activity-law
1343    /// DGP (the `fitted_scale_recovers_planted_activity_law` plant). Pins the
1344    /// rank-selection decision itself: the ladder must prefer rank 1, and this
1345    /// test names the margin so an over-selection regression is diagnosable
1346    /// from the failure message alone.
1347    #[test]
1348    fn evidence_ladder_prefers_planted_rank_one() {
1349        let n = 5000usize;
1350        let p = 4usize;
1351        let lambda0 = ndarray::array![[1.5], [1.2], [-0.4], [0.3]];
1352        let sigma_eps = 0.2_f64;
1353        let slope = 1.3_f64;
1354        let mut seed = 0xD1B54A32D192ED03_u64;
1355        let mut residuals = Array2::<f64>::zeros((n, p));
1356        let mut activity = Array1::<f64>::zeros(n);
1357        for row in 0..n {
1358            let z = (row as f64) / (n as f64 - 1.0);
1359            activity[row] = z;
1360            let amp = (slope * z).exp().sqrt();
1361            let f = lcg_normal(&mut seed);
1362            for i in 0..p {
1363                residuals[[row, i]] = amp * lambda0[[i, 0]] * f + sigma_eps * lcg_normal(&mut seed);
1364            }
1365        }
1366        // Reproduce fit()'s bin assignment, then score each rank directly.
1367        let bins = ACTIVITY_SCALE_BINS.max(1);
1368        let row_bin: Vec<usize> = (0..n)
1369            .map(|i| {
1370                let frac = activity[i];
1371                (frac * bins as f64).floor().clamp(0.0, bins as f64 - 1.0) as usize
1372            })
1373            .collect();
1374        let mut report = String::new();
1375        let mut ev = Vec::new();
1376        for rank in 0..=2usize {
1377            let m = StructuredResidualModel::fit_fixed_rank(residuals.view(), &row_bin, bins, rank)
1378                .expect("fixed-rank fit");
1379            let k_params = (p * rank + p + ACTIVITY_SCALE_BINS) as f64;
1380            let log_lik = m.log_evidence() + 0.5 * k_params * (n as f64).ln();
1381            let col_norms: Vec<f64> = (0..rank)
1382                .map(|k| {
1383                    m.factor()
1384                        .column(k)
1385                        .iter()
1386                        .map(|v| v * v)
1387                        .sum::<f64>()
1388                        .sqrt()
1389                })
1390                .collect();
1391            report.push_str(&format!(
1392                "rank {rank}: evidence={:.3} loglik={:.3} penalty={:.3} col_norms={:?} diag={:?}\n",
1393                m.log_evidence(),
1394                log_lik,
1395                0.5 * k_params * (n as f64).ln(),
1396                col_norms,
1397                m.diagonal()
1398                    .iter()
1399                    .map(|v| (v * 1e4).round() / 1e4)
1400                    .collect::<Vec<_>>()
1401            ));
1402            ev.push(m.log_evidence());
1403        }
1404        assert!(
1405            ev[1] > ev[0] && ev[1] > ev[2],
1406            "evidence ladder must prefer the planted rank 1; breakdown:\n{report}"
1407        );
1408    }
1409
1410    /// Orthonormalize the columns of `m` (modified Gram–Schmidt), dropping
1411    /// numerically-null columns. Test-side helper for subspace comparisons.
1412    fn orthonormal_columns(m: ArrayView2<'_, f64>) -> Vec<Array1<f64>> {
1413        let mut basis: Vec<Array1<f64>> = Vec::new();
1414        for k in 0..m.ncols() {
1415            let mut v = m.column(k).to_owned();
1416            for q in &basis {
1417                let c = v.dot(q);
1418                v = &v - &(q * c);
1419            }
1420            let norm = v.dot(&v).sqrt();
1421            if norm > 1e-10 {
1422                basis.push(v / norm);
1423            }
1424        }
1425        basis
1426    }
1427
1428    /// Squared norm of the projection of unit vector `v` onto span(basis) —
1429    /// `cos²` of the principal angle between `v` and the subspace.
1430    fn projection_energy(v: &Array1<f64>, basis: &[Array1<f64>]) -> f64 {
1431        basis.iter().map(|q| v.dot(q).powi(2)).sum()
1432    }
1433
1434    /// #974 verification arm (a): the fitted factor must recover the PLANTED
1435    /// interference subspace. Two orthogonal planted directions with distinct
1436    /// strengths; the principal angles between each planted direction and
1437    /// range(Λ̂) must be small, and the evidence ladder must select rank 2.
1438    #[test]
1439    fn factor_recovers_planted_interference_subspace() {
1440        let n = 6000usize;
1441        let p = 6usize;
1442        // Two orthogonal planted unit directions.
1443        let raw1: Array1<f64> = ndarray::array![1.0, 1.0, 1.0, 1.0, 1.0, 1.0];
1444        let raw2: Array1<f64> = ndarray::array![1.0, -1.0, 1.0, -1.0, 1.0, -1.0];
1445        let v1 = &raw1 / raw1.dot(&raw1).sqrt();
1446        let v2 = &raw2 / raw2.dot(&raw2).sqrt();
1447        let (amp1, amp2) = (1.4_f64, 0.9_f64);
1448        let sigma_eps = 0.15_f64;
1449
1450        let mut seed = 0x9E3779B97F4A7C15_u64;
1451        let mut residuals = Array2::<f64>::zeros((n, p));
1452        let activity = Array1::<f64>::zeros(n); // constant ⇒ homoscedastic law
1453        for row in 0..n {
1454            let f1 = amp1 * lcg_normal(&mut seed);
1455            let f2 = amp2 * lcg_normal(&mut seed);
1456            for i in 0..p {
1457                residuals[[row, i]] = f1 * v1[i] + f2 * v2[i] + sigma_eps * lcg_normal(&mut seed);
1458            }
1459        }
1460
1461        let model = StructuredResidualModel::fit(ResidualFactorInput {
1462            residuals: residuals.view(),
1463            activity: activity.view(),
1464            max_factor_rank: 4,
1465        })
1466        .expect("fit");
1467
1468        assert_eq!(
1469            model.factor_rank(),
1470            2,
1471            "ladder must select the planted rank 2 (got {}, evidence {:.3})",
1472            model.factor_rank(),
1473            model.log_evidence()
1474        );
1475        let basis = orthonormal_columns(model.factor());
1476        assert_eq!(basis.len(), 2, "fitted factor must span 2 directions");
1477        let e1 = projection_energy(&v1, &basis);
1478        let e2 = projection_energy(&v2, &basis);
1479        // cos² of each principal angle ≥ 0.95 ⇒ angle ≤ ~13°.
1480        assert!(
1481            e1 > 0.95 && e2 > 0.95,
1482            "planted directions must lie in range(Λ̂): cos² = ({e1:.4}, {e2:.4})"
1483        );
1484    }
1485
1486    /// Reproduce `fit`'s equal-width bin assignment for a test activity vector.
1487    fn assign_bins(activity: &Array1<f64>, bins: usize) -> Vec<usize> {
1488        let n = activity.len();
1489        let z_min = activity.iter().copied().fold(f64::INFINITY, f64::min);
1490        let z_max = activity.iter().copied().fold(f64::NEG_INFINITY, f64::max);
1491        let span = z_max - z_min;
1492        (0..n)
1493            .map(|i| {
1494                if span <= 0.0 {
1495                    0
1496                } else {
1497                    let frac = (activity[i] - z_min) / span;
1498                    (frac * bins as f64).floor().clamp(0.0, bins as f64 - 1.0) as usize
1499                }
1500            })
1501            .collect()
1502    }
1503
1504    /// Naive, pre-hoist reference for `penalized_log_evidence`: rebuilds the
1505    /// row-independent Gram M0 = ΛᵀD⁻¹Λ INSIDE the per-row loop (the original
1506    /// formula). The production function hoists M0 out; the two must agree.
1507    fn naive_penalized_log_evidence(
1508        r: ArrayView2<'_, f64>,
1509        lambda: &Array2<f64>,
1510        diagonal: &Array1<f64>,
1511        row_scale: &Array1<f64>,
1512        rank: usize,
1513    ) -> f64 {
1514        let n = r.nrows();
1515        let p = r.ncols();
1516        let d_inv: Vec<f64> = (0..p).map(|i| 1.0 / diagonal[i]).collect();
1517        let log_det_d: f64 = diagonal.iter().map(|&d| d.ln()).sum();
1518        let two_pi_ln = (2.0 * std::f64::consts::PI).ln();
1519        let mut log_lik = 0.0_f64;
1520        for i in 0..n {
1521            let c = row_scale[i].max(f64::MIN_POSITIVE);
1522            let mut quad = 0.0_f64;
1523            for j in 0..p {
1524                quad += r[[i, j]] * d_inv[j] * r[[i, j]];
1525            }
1526            let mut log_det = log_det_d;
1527            if rank > 0 {
1528                let mut m = Array2::<f64>::zeros((rank, rank));
1529                let mut w = Array1::<f64>::zeros(rank);
1530                for a in 0..rank {
1531                    let mut wa = 0.0_f64;
1532                    for j in 0..p {
1533                        wa += lambda[[j, a]] * d_inv[j] * r[[i, j]];
1534                    }
1535                    w[a] = wa;
1536                    for b in 0..rank {
1537                        let mut acc = 0.0_f64;
1538                        for j in 0..p {
1539                            acc += lambda[[j, a]] * d_inv[j] * lambda[[j, b]];
1540                        }
1541                        m[[a, b]] = acc;
1542                    }
1543                    m[[a, a]] += 1.0 / c;
1544                }
1545                match m.cholesky(Side::Lower) {
1546                    Ok(chol) => {
1547                        let y = chol.solvevec(&w);
1548                        let mut wy = 0.0_f64;
1549                        for a in 0..rank {
1550                            wy += w[a] * y[a];
1551                        }
1552                        quad -= wy;
1553                        let diag = chol.diag();
1554                        let log_det_m: f64 = diag.iter().map(|&l| (l * l).ln()).sum();
1555                        log_det = log_det_d + log_det_m + rank as f64 * c.ln();
1556                    }
1557                    Err(_) => {
1558                        log_det = log_det_d;
1559                    }
1560                }
1561            }
1562            log_lik += -0.5 * (log_det + quad + p as f64 * two_pi_ln);
1563        }
1564        let k_params = (p * rank + p + ACTIVITY_SCALE_BINS) as f64;
1565        log_lik - 0.5 * k_params * (n.max(2) as f64).ln()
1566    }
1567
1568    /// Naive, per-row-rebuild reference for `factor_coordinates`: rebuilds and
1569    /// re-factors the (row-independent) normal matrix ΛᵀD⁻¹Λ for EVERY row.
1570    /// Mathematically identical to the shared-factorization production path.
1571    fn naive_factor_coordinates(
1572        lambda: &Array2<f64>,
1573        diagonal: &Array1<f64>,
1574        r: ArrayView2<'_, f64>,
1575    ) -> Array2<f64> {
1576        let p = lambda.nrows();
1577        let rank = lambda.ncols();
1578        let n = r.nrows();
1579        let d_inv: Vec<f64> = (0..p).map(|i| 1.0 / diagonal[i]).collect();
1580        let mut coords = Array2::<f64>::zeros((n, rank));
1581        for i in 0..n {
1582            let mut normal = Array2::<f64>::zeros((rank, rank));
1583            for a in 0..rank {
1584                for b in 0..rank {
1585                    let mut acc = 0.0_f64;
1586                    for j in 0..p {
1587                        acc += lambda[[j, a]] * d_inv[j] * lambda[[j, b]];
1588                    }
1589                    normal[[a, b]] = acc;
1590                }
1591            }
1592            let trace = (0..rank).map(|k| normal[[k, k]]).sum::<f64>().max(1.0);
1593            let ridge = 1e-10 * trace / rank.max(1) as f64;
1594            for k in 0..rank {
1595                normal[[k, k]] += ridge;
1596            }
1597            let chol = normal.cholesky(Side::Lower).expect("naive normal solve");
1598            let mut rhs = Array1::<f64>::zeros(rank);
1599            for a in 0..rank {
1600                let mut acc = 0.0_f64;
1601                for j in 0..p {
1602                    acc += lambda[[j, a]] * d_inv[j] * r[[i, j]];
1603                }
1604                rhs[a] = acc;
1605            }
1606            let gamma = chol.solvevec(&rhs);
1607            for a in 0..rank {
1608                coords[[i, a]] = gamma[a];
1609            }
1610        }
1611        coords
1612    }
1613
1614    /// FIX B equivalence: the hoisted `penalized_log_evidence` and the shared-
1615    /// factorization `factor_coordinates` must equal their naive per-row-rebuild
1616    /// references to ~1e-10 (in fact bit-for-bit — the hoist preserves op order).
1617    #[test]
1618    fn hoisted_gram_matches_naive_per_row_rebuild() {
1619        let n = 200usize;
1620        let p = 5usize;
1621        let rank = 2usize;
1622        let mut seed = 0x243F6A8885A308D3_u64;
1623        let mut lambda = Array2::<f64>::zeros((p, rank));
1624        for i in 0..p {
1625            for k in 0..rank {
1626                lambda[[i, k]] = lcg_normal(&mut seed);
1627            }
1628        }
1629        let mut diagonal = Array1::<f64>::zeros(p);
1630        for j in 0..p {
1631            diagonal[j] = 0.3 + lcg_uniform(&mut seed); // strictly positive
1632        }
1633        let mut row_scale = Array1::<f64>::zeros(n);
1634        for i in 0..n {
1635            row_scale[i] = 0.5 + 1.5 * lcg_uniform(&mut seed); // strictly positive
1636        }
1637        let mut residuals = Array2::<f64>::zeros((n, p));
1638        for i in 0..n {
1639            for j in 0..p {
1640                residuals[[i, j]] = lcg_normal(&mut seed);
1641            }
1642        }
1643
1644        let ev_hoisted =
1645            penalized_log_evidence(residuals.view(), &lambda, &diagonal, &row_scale, rank);
1646        let ev_naive =
1647            naive_penalized_log_evidence(residuals.view(), &lambda, &diagonal, &row_scale, rank);
1648        assert!(
1649            (ev_hoisted - ev_naive).abs() <= 1e-10 * (1.0 + ev_naive.abs()),
1650            "hoisted log-evidence must equal naive rebuild: {ev_hoisted} vs {ev_naive}"
1651        );
1652
1653        let coords_hoisted =
1654            factor_coordinates(&lambda, &diagonal, residuals.view()).expect("coords");
1655        let coords_naive = naive_factor_coordinates(&lambda, &diagonal, residuals.view());
1656        let mut max_abs = 0.0_f64;
1657        for i in 0..n {
1658            for a in 0..rank {
1659                max_abs = max_abs.max((coords_hoisted[[i, a]] - coords_naive[[i, a]]).abs());
1660            }
1661        }
1662        assert!(
1663            max_abs <= 1e-10,
1664            "hoisted factor coordinates must equal naive rebuild; max |Δ| = {max_abs:e}"
1665        );
1666    }
1667
1668    /// FIX A regression: on an uneven-bin synthetic with a KNOWN planted single
1669    /// factor, the low-rank reconstruction ΛΛᵀ + D built from the OCCUPANCY-
1670    /// weighted scale law reconstructs the empirical second moment
1671    /// (1/n) Σ_n r_n r_nᵀ strictly better (Frobenius) than the one built from
1672    /// the bin-UNIFORM scale law. Uses the module's own `scaled_second_moment` /
1673    /// eigen path so it exercises the real (Λ, D | scale) step.
1674    #[test]
1675    fn occupancy_scale_improves_second_moment_reconstruction() {
1676        let n = 4000usize;
1677        let p = 4usize;
1678        let lambda0 = ndarray::array![1.5, 1.2, -0.4, 0.3];
1679        let sigma_eps = 0.2_f64;
1680        let slope = 2.0_f64;
1681        let bins = ACTIVITY_SCALE_BINS.max(1);
1682        let mut seed = 0xCA62C1D6_u64 ^ 0x9B05688C_u64;
1683        let mut residuals = Array2::<f64>::zeros((n, p));
1684        let mut activity = Array1::<f64>::zeros(n);
1685        let mut c_true = Array1::<f64>::zeros(n);
1686        for row in 0..n {
1687            let u = (row as f64) / (n as f64 - 1.0);
1688            let z = u * u * u; // cubic warp ⇒ uneven bin occupancy
1689            activity[row] = z;
1690            let c = (slope * z).exp();
1691            c_true[row] = c;
1692            let amp = c.sqrt();
1693            let f = lcg_normal(&mut seed);
1694            for i in 0..p {
1695                residuals[[row, i]] = amp * lambda0[i] * f + sigma_eps * lcg_normal(&mut seed);
1696            }
1697        }
1698
1699        // Empirical (undeflated) second moment T = (1/n) Σ_n r_n r_nᵀ — the
1700        // object the model's ΛΛᵀ + D must reconstruct.
1701        let mut t = Array2::<f64>::zeros((p, p));
1702        for i in 0..n {
1703            for a in 0..p {
1704                for b in 0..p {
1705                    t[[a, b]] += residuals[[i, a]] * residuals[[i, b]];
1706                }
1707            }
1708        }
1709        t.mapv_inplace(|v| v / n as f64);
1710
1711        let raw_diag = column_variances(residuals.view());
1712        let mean_var = raw_diag.iter().sum::<f64>() / p as f64;
1713        let diag_floor = DIAGONAL_REL_FLOOR * mean_var.max(f64::MIN_POSITIVE);
1714
1715        // Per-bin raw scale law: mean of the true c(z) within each bin.
1716        let row_bin = assign_bins(&activity, bins);
1717        let mut bin_sum = vec![0.0_f64; bins];
1718        let mut bin_cnt = vec![0.0_f64; bins];
1719        for i in 0..n {
1720            bin_sum[row_bin[i]] += c_true[i];
1721            bin_cnt[row_bin[i]] += 1.0;
1722        }
1723        let bin_raw: Vec<f64> = (0..bins)
1724            .map(|b| {
1725                if bin_cnt[b] > 0.0 {
1726                    bin_sum[b] / bin_cnt[b]
1727                } else {
1728                    1.0
1729                }
1730            })
1731            .collect();
1732
1733        // Occupancy-weighted mean-1 (Fix A) vs bin-uniform mean-1 (old).
1734        let mean_occ = (0..bins).map(|b| bin_cnt[b] * bin_raw[b]).sum::<f64>() / n as f64;
1735        let occupied: Vec<usize> = (0..bins).filter(|&b| bin_cnt[b] > 0.0).collect();
1736        let mean_uni = occupied.iter().map(|&b| bin_raw[b]).sum::<f64>() / occupied.len() as f64;
1737        let row_scale_occ: Array1<f64> = (0..n).map(|i| bin_raw[row_bin[i]] / mean_occ).collect();
1738        let row_scale_uni: Array1<f64> = (0..n).map(|i| bin_raw[row_bin[i]] / mean_uni).collect();
1739
1740        // One (Λ, D | scale) extraction from the deflated moment, mirroring the
1741        // production first sweep, returning the reconstruction ΛΛᵀ + D.
1742        let extract_recon = |row_scale: &Array1<f64>| -> Array2<f64> {
1743            let s = scaled_second_moment(residuals.view(), row_scale);
1744            let (evals, evecs) = symmetric_eig_ascending(&s).expect("eig");
1745            let mean_diag = raw_diag.iter().map(|&v| v.max(diag_floor)).sum::<f64>() / p as f64;
1746            let col = p - 1;
1747            let amp = (evals[col] - mean_diag).max(0.0).sqrt();
1748            let mut lam = Array1::<f64>::zeros(p);
1749            for j in 0..p {
1750                lam[j] = amp * evecs[[j, col]];
1751            }
1752            let mut recon = Array2::<f64>::zeros((p, p));
1753            for a in 0..p {
1754                for b in 0..p {
1755                    recon[[a, b]] = lam[a] * lam[b];
1756                }
1757            }
1758            for j in 0..p {
1759                let d = (raw_diag[j] - lam[j] * lam[j]).max(diag_floor);
1760                recon[[j, j]] += d;
1761            }
1762            recon
1763        };
1764
1765        let frob = |m: &Array2<f64>| -> f64 {
1766            let mut acc = 0.0_f64;
1767            for a in 0..p {
1768                for b in 0..p {
1769                    let d = m[[a, b]] - t[[a, b]];
1770                    acc += d * d;
1771                }
1772            }
1773            acc.sqrt()
1774        };
1775
1776        let dist_occ = frob(&extract_recon(&row_scale_occ));
1777        let dist_uni = frob(&extract_recon(&row_scale_uni));
1778        assert!(
1779            dist_occ < dist_uni,
1780            "occupancy-weighted reconstruction must beat bin-uniform: \
1781             ‖·‖_F occ = {dist_occ:.6} vs uni = {dist_uni:.6}"
1782        );
1783    }
1784
1785    /// Fit a small structured model on a planted single-factor DGP — shared
1786    /// fixture builder for the producer / damped-metric integration tests.
1787    fn fit_small_model(seed0: u64, lambda0: &Array1<f64>) -> (usize, StructuredResidualModel) {
1788        let n = 300usize;
1789        let p = lambda0.len();
1790        let sigma_eps = 0.25_f64;
1791        let slope = 1.4_f64;
1792        let mut seed = seed0;
1793        let mut residuals = Array2::<f64>::zeros((n, p));
1794        let mut activity = Array1::<f64>::zeros(n);
1795        for row in 0..n {
1796            let u = (row as f64) / (n as f64 - 1.0);
1797            let z = u * u;
1798            activity[row] = z;
1799            let amp = (slope * z).exp().sqrt();
1800            let f = lcg_normal(&mut seed);
1801            for i in 0..p {
1802                residuals[[row, i]] = amp * lambda0[i] * f + sigma_eps * lcg_normal(&mut seed);
1803            }
1804        }
1805        let model = StructuredResidualModel::fit(ResidualFactorInput {
1806            residuals: residuals.view(),
1807            activity: activity.view(),
1808            max_factor_rank: 2,
1809        })
1810        .expect("fit");
1811        (n, model)
1812    }
1813
1814    /// WAVE-2 #2021 damped metric endpoint contracts: γ=1 ≡ row_metric (ignores
1815    /// prev), γ=0 ≡ prev.row_metric (or Euclidean identity), 0<γ<1 is SPD, and
1816    /// out-of-range / non-finite γ is rejected.
1817    #[test]
1818    fn row_metric_damped_endpoints() {
1819        let lambda_a = ndarray::array![1.5, 1.2, -0.4, 0.3];
1820        let lambda_b = ndarray::array![-0.6, 1.1, 0.9, -1.3];
1821        let (n, model) = fit_small_model(0x51ED270B_u64 ^ 0xF3A5C7D1_u64, &lambda_a);
1822        let (n_prev, prev) = fit_small_model(0x2545F491_u64 ^ 0x4F6CDD1D_u64, &lambda_b);
1823        assert_eq!(n, n_prev);
1824        let p = 4usize;
1825        let v: Array1<f64> = ndarray::array![0.7, -1.3, 0.4, 0.9];
1826
1827        // γ = 1 ⇒ byte-identical to this model's row_metric, regardless of prev.
1828        let base = model.row_metric(n).expect("row_metric");
1829        for prev_opt in [None, Some(&prev)] {
1830            let damped = model
1831                .row_metric_damped(n, 1.0, prev_opt)
1832                .expect("damped γ=1");
1833            for row in [0usize, n / 2, n - 1] {
1834                for i in 0..p {
1835                    for k in 0..p {
1836                        assert_eq!(
1837                            damped.factor_entry(row, i, k),
1838                            base.factor_entry(row, i, k),
1839                            "γ=1 must be byte-identical to row_metric at ({row},{i},{k})"
1840                        );
1841                    }
1842                }
1843            }
1844        }
1845
1846        // γ = 0, prev = None ⇒ the MEASURED-scale identity (#2243 cap #2):
1847        // Σ = φ̂·I with φ̂ the model's own isotropic dispersion, so
1848        // quad_form = ‖v‖² / φ̂ (a unit-I anchor would silently assume unit
1849        // noise and over-penalize clean data).
1850        let ident = model
1851            .row_metric_damped(n, 0.0, None)
1852            .expect("damped γ=0 None");
1853        let phi = model.isotropic_dispersion();
1854        assert!(phi.is_finite() && phi > 0.0, "φ̂ must be positive");
1855        let sumsq: f64 = v.iter().map(|x| x * x).sum();
1856        let expected = sumsq / phi;
1857        for row in [0usize, n / 2, n - 1] {
1858            let q = ident.quad_form(row, v.view());
1859            assert!(
1860                (q - expected).abs() <= 1e-9 * (1.0 + expected),
1861                "γ=0/None must be the measured-scale identity: quad_form {q} vs ‖v‖²/φ̂ {expected}"
1862            );
1863        }
1864
1865        // γ = 0, prev = Some ⇒ byte-identical to prev.row_metric.
1866        let prev_metric = prev.row_metric(n).expect("prev row_metric");
1867        let damped0 = model
1868            .row_metric_damped(n, 0.0, Some(&prev))
1869            .expect("damped γ=0 Some");
1870        for row in [0usize, n / 2, n - 1] {
1871            for i in 0..p {
1872                for k in 0..p {
1873                    assert_eq!(
1874                        damped0.factor_entry(row, i, k),
1875                        prev_metric.factor_entry(row, i, k),
1876                        "γ=0/Some must be byte-identical to prev.row_metric at ({row},{i},{k})"
1877                    );
1878                }
1879            }
1880        }
1881
1882        // 0 < γ < 1 ⇒ valid SPD metric.
1883        let mid = model
1884            .row_metric_damped(n, 0.5, Some(&prev))
1885            .expect("damped γ=0.5");
1886        for row in [0usize, n / 2, n - 1] {
1887            let q = mid.quad_form(row, v.view());
1888            assert!(
1889                q.is_finite() && q > 0.0,
1890                "γ=0.5 metric must be SPD; got {q}"
1891            );
1892        }
1893
1894        // Invalid γ rejected.
1895        assert!(model.row_metric_damped(n, 1.5, None).is_err());
1896        assert!(model.row_metric_damped(n, -0.1, None).is_err());
1897        assert!(model.row_metric_damped(n, f64::NAN, None).is_err());
1898    }
1899
1900    /// WAVE-2 #2021 Λ nursery→promotion: `promotion_candidates` must fire only
1901    /// for a factor that BOTH persists across passes (aligns with the previous
1902    /// model's Λ) AND clears the idiosyncratic-noise energy floor; a fresh
1903    /// orthogonal direction, an over-high energy floor, and `prev = None` all
1904    /// yield no candidates, and out-of-range gates are rejected.
1905    #[test]
1906    fn promotion_candidates_gates_on_persistence_and_energy() {
1907        let lambda_a = ndarray::array![1.5, 1.2, -0.4, 0.3];
1908        let lambda_b = ndarray::array![-0.6, 1.1, 0.9, -1.3];
1909        // Same planted direction across two passes ⇒ persistent.
1910        let (_, prev) = fit_small_model(0xA1B2C3D4_u64 ^ 0x0F0F0F0F_u64, &lambda_a);
1911        let (_, cur) = fit_small_model(0x5566778899AABBCC_u64, &lambda_a);
1912        // A different (well-separated) planted direction ⇒ NOT aligned with cur.
1913        let (_, other) = fit_small_model(0x1122334455667788_u64, &lambda_b);
1914
1915        assert!(prev.factor_rank() >= 1 && cur.factor_rank() >= 1 && other.factor_rank() >= 1);
1916
1917        // Persistent + energetic ⇒ at least one candidate, aligned with the
1918        // planted direction and above the noise floor.
1919        let cands = cur
1920            .promotion_candidates(Some(&prev), 0.9, 1.0)
1921            .expect("promotion_candidates");
1922        assert!(
1923            !cands.is_empty(),
1924            "a persistent, energetic factor must yield a promotion candidate"
1925        );
1926        let top = &cands[0];
1927        assert!(
1928            top.persistence_alignment >= 0.9,
1929            "top candidate must clear the alignment gate; got {}",
1930            top.persistence_alignment
1931        );
1932        // The promoted unit direction must align with the planted (unit) lambda_a.
1933        let la_norm = lambda_a.dot(&lambda_a).sqrt();
1934        let la_unit = lambda_a.mapv(|v| v / la_norm);
1935        let dir_cos = top.direction.dot(&la_unit).abs();
1936        assert!(
1937            dir_cos > 0.9,
1938            "promoted direction must recover the planted factor; |cos| = {dir_cos:.4}"
1939        );
1940        assert!(
1941            (top.direction.dot(&top.direction) - 1.0).abs() < 1e-10,
1942            "promoted direction must be unit-norm"
1943        );
1944        assert!(top.energy > 0.0);
1945
1946        // A fresh, well-separated direction does NOT persist ⇒ no candidate at 0.9.
1947        let cross = cur
1948            .promotion_candidates(Some(&other), 0.9, 1.0)
1949            .expect("promotion_candidates cross");
1950        assert!(
1951            cross.is_empty(),
1952            "a non-persistent (unaligned) factor must not be promoted; got {} candidate(s)",
1953            cross.len()
1954        );
1955
1956        // An over-high energy floor rejects even the persistent factor.
1957        let floored = cur
1958            .promotion_candidates(Some(&prev), 0.9, 1.0e6)
1959            .expect("promotion_candidates floored");
1960        assert!(
1961            floored.is_empty(),
1962            "energy floor must gate out factors below the noise-scaled threshold"
1963        );
1964
1965        // prev = None (first structured pass, damping toward I) ⇒ no candidates.
1966        assert!(cur.promotion_candidates(None, 0.9, 1.0).unwrap().is_empty());
1967
1968        // Invalid gates rejected.
1969        assert!(cur.promotion_candidates(Some(&prev), 1.5, 1.0).is_err());
1970        assert!(cur.promotion_candidates(Some(&prev), -0.1, 1.0).is_err());
1971        assert!(cur.promotion_candidates(Some(&prev), 0.9, -1.0).is_err());
1972        assert!(
1973            cur.promotion_candidates(Some(&prev), f64::NAN, 1.0)
1974                .is_err()
1975        );
1976    }
1977}