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 += ∂
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}