Skip to main content

gam_problem/
row_metric.rs

1//! `RowMetric` — the single provenance-carrying per-row inner product shared by
2//! the SAE-manifold **likelihood** (residual whitening) and the **gauge**
3//! (isometry pullback weight).
4//!
5//! # Why this exists
6//!
7//! The SAE-manifold machine historically carried *two* independent inner
8//! products:
9//!
10//! * the **likelihood** measured reconstruction residuals isotropically — a
11//!   single scalar dispersion `φ̂ = RSS / residual-dof`, the data-fit loop
12//!   summing the bare `½ rᵀr`; there was no per-row metric at all; and
13//! * the **gauge** carried its own per-row metric in
14//!   `IsometryPenalty.weight: WeightField` — a low-rank `W_n = U_n U_nᵀ`
15//!   pullback `g_n = J_nᵀ W_n J_n`, settable independently of anything the
16//!   likelihood saw.
17//!
18//! Nothing structurally forced "the metric the likelihood whitens by" to equal
19//! "the metric the gauge pulls back through". That is exactly the
20//! objective↔gradient-desync bug class wearing geometry clothing: a
21//! likelihood-metric ≠ gauge-metric state was *representable*.
22//!
23//! `RowMetric` collapses the two into one object. The likelihood whitens
24//! through it; the gauge `WeightField` is *constructed from* it. A
25//! divergent-metric state is therefore unrepresentable — there is only one
26//! per-row factor stack `U_n`, with one [`MetricProvenance`] tag.
27//!
28//! # Magic-by-default selector
29//!
30//! There is no flag. The provenance is chosen by whether per-row Fisher factors
31//! exist:
32//!
33//! * no factors supplied ⇒ [`MetricProvenance::Euclidean`]; `W_n = I_p`;
34//!   whitening is the identity, so `φ̂` and the data-fit loop are
35//!   **bit-for-bit** the prior isotropic path; and
36//! * per-row Fisher factors supplied ⇒ [`MetricProvenance::OutputFisher`]; the
37//!   residual is whitened by `U_nᵀ` and the gauge pulls back through the same
38//!   `U_n`.
39//!
40//! # Validation
41//!
42//! Every metric block is constructed **through**
43//! [`crate::normalize_fisher_rao_blocks`], which
44//! broadcasts and eigenvalue-validates PSD-ness. `RowMetric` does not
45//! reimplement that validation; it materializes `W_n = U_n U_nᵀ` (which is PSD
46//! by construction) and runs it through the shared normalizer as the
47//! single point of truth for "is this a valid precision metric".
48//!
49//! Any rank floor used to make a block invertible for an internal solve is
50//! **solver-only** (mirroring `RidgePolicy::solver_only`, #747): it never enters
51//! the residual the objective sums, so `δ` cannot bias the criterion.
52//!
53//! # Rung 1 — the behavioral metric *in the reconstruction loss* (nats currency)
54//!
55//! [`MetricProvenance::OutputFisher`] installs the output-Fisher inner product
56//! as a **gauge** metric only: it whitens *nothing* (`whitens_likelihood()` is
57//! `false`), by deliberate #980 contract, so reconstruction stays the isotropic
58//! `½‖r‖²`. That answers "what coordinate is canonical", not "what does a
59//! reconstruction error *cost*".
60//!
61//! [`MetricProvenance::BehavioralFisher`] is the opposite deliberate choice:
62//! the **same** low-rank output-Fisher factors, but installed as the
63//! reconstruction *likelihood weight*. Plain MSE prices a reconstruction error
64//! `e = x − x̂` by its Euclidean size; the model, however, reads the activation
65//! only through the rest of the network, so the behavioral cost of `e` is the
66//! KL between the clean and corrupted next-token distributions,
67//! `KL ≈ ½ eᵀ G(x) e` with `G = JᵀFJ` the network-Jacobian pullback of the
68//! output Fisher `F` (units: **nats**). Minimizing `(x−x̂)ᵀ G (x−x̂)` instead of
69//! `‖x−x̂‖²` is **generalized least squares**: for a *fixed* per-row `G` it is
70//! still a linear Gaussian model in the coefficients, so the entire
71//! REML/evidence/EDF/certificate stack survives verbatim — this is why the
72//! metric rides the identical `whitens_likelihood()` plumbing the
73//! [`MetricProvenance::WhitenedStructured`] noise model uses, and why the G=I
74//! limit reproduces the plain-MSE fit bit-for-bit (see the module tests).
75//!
76//! This is the principled form of Braun's end-to-end **KL + MSE** objective.
77//! Anchoring to the activation keeps it *reconstruction* (it does not collapse
78//! to "match the logits by any means" — the decoder still has to reproduce `x`),
79//! while pricing the residual in nats through `G`. The payoff is automatic
80//! selection for *mattering*: `G`'s null directions — activation structure the
81//! rest of the network cannot read — are penalized nothing, because
82//! `eᵀ G e = 0` there. MSE in a behaviorally-inert direction goes free, which is
83//! the correct behavior, not a bug: nothing downstream changes, so nothing
84//! should be paid.
85//!
86//! **The d×d `G` is never materialized.** `G` is sketched by `s` random probes,
87//! `vᵢ = Jᵀ F^{1/2} uᵢ` (`uᵢ` iid, `s ≈ 4…16`), computed by `s` backward passes
88//! per token at *harvest* time (the model-interaction boundary) and stored as
89//! the columns of the per-row factor `U_n = [v₁ … v_s] ∈ ℝ^{p×s}`. Then
90//! `G ≈ Σᵢ vᵢ vᵢᵀ = U_n U_nᵀ` and the criterion-facing
91//! `eᵀ G e ≈ Σᵢ (vᵢᵀ e)² = ‖U_nᵀ e‖²` is exactly what
92//! [`RowMetric::quad_form`] / [`RowMetric::whiten_residual_row`] already
93//! compute — zero train-time model cost, `O(p·s)` per row. See
94//! [`RowMetric::behavioral_fisher`] and the probe-packing helper
95//! [`pack_probe_factors`].
96
97use ndarray::{Array2, Array3, ArrayView1};
98use std::sync::Arc;
99
100use crate::normalize_fisher_rao_blocks;
101
102/// Per-observation behavioral-metric field `W_n ∈ ℝ^{p × p}`, stored in
103/// **low-rank factored form** `W_n = U_n U_n^T` with `U_n ∈ ℝ^{p × r_n}`.
104///
105/// The canonical coordinate is the one where one unit of motion in `t` is one
106/// unit of behavioral change in the output space, so the `W_n` weighting is
107/// load-bearing: the pullback metric is `g_n = J_n^T W_n J_n`. Storing as
108/// `U_n` lets every contraction in this module run in
109/// `(J^T U_n)(U_n^T J)` order, which is `O(p · r · d + r · d²)` per row — we
110/// **never** materialize the `p × p` `W_n`, which is essential when `p`
111/// (number of observation channels) is large but rank is small (e.g. one or
112/// two behavioral dimensions per latent observation).
113///
114/// `Identity` is the gauge-fix default and corresponds to `U_n = I_p` so the
115/// pullback reduces to the standard `J_n^T J_n`. `Factored` stores the
116/// per-row `U_n` blocks contiguously: every row's factor is `p × rank`, and
117/// rows may share the same rank (uniform-rank case) or vary if the field is
118/// data-driven. For the uniform-rank case the storage is
119/// `(n_obs, p * rank)` row-major.
120#[derive(Clone)]
121pub enum WeightField {
122    /// `W_n = I_p` for every `n`. Reduces to the bare pullback `J^T J`.
123    Identity,
124    /// Per-row low-rank factor `U_n ∈ ℝ^{p × rank}`. Storage layout: a
125    /// `(n_obs, p * rank)` row-major matrix where row `n` packs `U_n` in
126    /// column-major-within-row order `U_n[i, k] = u[n, i * rank + k]`.
127    Factored {
128        u: Arc<Array2<f64>>,
129        rank: usize,
130        p_out: usize,
131    },
132}
133
134impl std::fmt::Debug for WeightField {
135    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
136        match self {
137            WeightField::Identity => f.write_str("Identity"),
138            WeightField::Factored { u, rank, p_out } => f
139                .debug_struct("Factored")
140                .field("shape", &format_args!("{}×{}", u.nrows(), u.ncols()))
141                .field("rank", rank)
142                .field("p_out", p_out)
143                .finish(),
144        }
145    }
146}
147
148impl WeightField {
149    /// Apply `U_n^T J_n` for a specific row, given both the row's `J_n` flat
150    /// `(p * d)` slice and the row's `U_n` flat `(p * rank)` slice. Returns
151    /// the `(rank × d)` matrix and its row count.
152    pub fn project_jac_row_with_u(
153        u_row: &[f64],
154        jac_row: &[f64],
155        p: usize,
156        rank: usize,
157        d: usize,
158    ) -> Array2<f64> {
159        // M[k, a] = Σ_i U[i, k] · J[i, a].
160        let mut m = Array2::<f64>::zeros((rank, d));
161        for k in 0..rank {
162            for a in 0..d {
163                let mut s = 0.0;
164                for i in 0..p {
165                    s += u_row[i * rank + k] * jac_row[i * d + a];
166                }
167                m[[k, a]] = s;
168            }
169        }
170        m
171    }
172}
173
174/// Where the per-row metric came from — the provenance that makes
175/// "likelihood-metric ≠ gauge-metric" diagnosable instead of silent.
176///
177/// Object 4 (the gauge object) reads this to certify which inner product the
178/// fit actually used; #974 fills [`MetricProvenance::WhitenedStructured`] with a
179/// factor-analytic residual-covariance whitening.
180#[derive(Clone, Copy, PartialEq, Eq, Debug)]
181pub enum MetricProvenance {
182    /// `M_n = I_p` for every row. The likelihood is isotropic and the gauge
183    /// pullback reduces to the bare `J_nᵀ J_n`. This is the default and is
184    /// bit-for-bit the historical isotropic-`φ̂` path.
185    Euclidean,
186    /// `M_n = U_n U_nᵀ (+ solver-only δI)` from supplied per-row output-Fisher
187    /// factors `U_n ∈ ℝ^{p × rank}`. The canonical "one unit of latent motion ↦
188    /// one unit of behavioral change" metric: residuals are whitened in the
189    /// output-Fisher inner product and the gauge pulls back through the same
190    /// factors. The `rank` is carried in the provenance so a consumer (Object 4)
191    /// can certify the factor rank that produced the inner product.
192    OutputFisher { rank: usize },
193    /// `M_n = U_n U_nᵀ` from per-row output-Fisher factors that aggregate the
194    /// **downstream** influence of position `n` over future positions through
195    /// the KV path, rather than the same-position logits of
196    /// [`MetricProvenance::OutputFisher`] (#980, mechanism 2).
197    ///
198    /// The same-position pullback `∂logits_t/∂x_t` can be ≈ 0 for a feature
199    /// whose entire causal effect lands many tokens later (information carried
200    /// forward through attention); a gauge built on it is blind to exactly that
201    /// content. This provenance is the forward-looking alternative: each row's
202    /// factor `U_n` is the top-`rank` factorization of the aggregated output
203    /// Fisher `Σ_{t ≥ n} (∂logits_t/∂x_n)ᵀ F_t (∂logits_t/∂x_n)` over future
204    /// positions the residual stream at `n` reaches. It is provenance-generic:
205    /// it whitens nothing (`Self::whitens_likelihood` is `false`, like
206    /// [`MetricProvenance::OutputFisher`]) and drives the gauge / lens /
207    /// enrichment unchanged (`Self::is_output_fisher_like`). The lens/gauge
208    /// machinery consumes it identically; only the *scientific* reading
209    /// changes — dormant-feature detection becomes forward-looking (a feature
210    /// driving far-future tokens now registers behavioral coupling that the
211    /// same-position metric reported as ≈ 0).
212    OutputFisherDownstream { rank: usize },
213    /// **Rung 1** — the output-Fisher metric installed as the reconstruction
214    /// **likelihood weight** (generalized least squares in nats), not merely as
215    /// a gauge. `M_n = U_n U_nᵀ ≈ G_n = J_nᵀ F_n J_n` is the `s`-probe sketch of
216    /// the pulled-back output Fisher, with `U_n = [v₁ … v_s]`,
217    /// `vᵢ = J_nᵀ F_n^{1/2} uᵢ`, and `probes = s` the number of random probes
218    /// (the factor rank).
219    ///
220    /// This is the *only* [`RowMetric::is_output_fisher_like`]-adjacent
221    /// provenance for which [`RowMetric::whitens_likelihood`] is `true`: the
222    /// data-fit sums `½ eᵀ G_n e = ½ ‖U_nᵀ e‖²` (nats) instead of `½‖e‖²`. It is
223    /// distinct from [`Self::OutputFisher`] precisely because the choice to let
224    /// the metric enter the *loss* (rather than only the gauge) is deliberate and
225    /// must not be silently inherited by the #980 gauge / two-tier-harvest
226    /// contract — that contract relies on [`Self::OutputFisher`] whitening
227    /// nothing. Because `G_n` is a *fixed* per-row metric, the whitened problem
228    /// is again linear-Gaussian in the coefficients, so REML/evidence/EDF are
229    /// unchanged (the GLS-preserves-REML property, verified in the module tests
230    /// against the `G=I` plain-MSE limit).
231    BehavioralFisher { probes: usize },
232    /// Structured-residual whitening: `M_n = Σ_n^{-1}` from the **estimated**
233    /// factor-analytic residual covariance `Σ_n = Λ c(z_n) Λᵀ + D` (#974), with
234    /// `factor_rank` the selected factor count. Produced by
235    /// Structured-residual producers materialize this provenance when they fit
236    /// a residual-covariance whitening model;
237    /// the only provenance for which
238    /// [`whitens_likelihood`](RowMetric::whitens_likelihood) is `true`. It
239    /// carries the same low-rank factor layout as
240    /// [`MetricProvenance::OutputFisher`].
241    WhitenedStructured { factor_rank: usize },
242}
243
244/// Scientific status of a factored output-Fisher approximation.
245///
246/// This is independent of both factor rank and a scalar trace diagnostic. A
247/// zero estimated tail trace does not prove an exact factorization, and a
248/// positive tail estimate does not prove the retained operator is below the
249/// true Fisher in Loewner order (#2249).
250#[derive(Clone, Copy, PartialEq, Eq, Debug)]
251pub enum FisherFactorKind {
252    /// The supplied factor exactly represents the complete local Fisher.
253    ExactFull,
254    /// The producer certified `0 <= U U^T <= F` as an operator inequality.
255    CertifiedPsdLowerBound,
256    /// Randomized, stochastic, truncated, or otherwise uncertified factor.
257    UncertifiedApproximation,
258}
259
260impl FisherFactorKind {
261    /// Stable artifact/FFI tag. Factor status is part of the scientific data,
262    /// not an inference a consumer may recreate from rank or trace metadata.
263    pub const fn tag(self) -> &'static str {
264        match self {
265            Self::ExactFull => "exact_full",
266            Self::CertifiedPsdLowerBound => "certified_psd_lower_bound",
267            Self::UncertifiedApproximation => "uncertified_approximation",
268        }
269    }
270
271    /// Parse the required public factor-status tag (#2249).
272    pub fn from_tag(tag: &str) -> Result<Self, String> {
273        match tag {
274            "exact_full" => Ok(Self::ExactFull),
275            "certified_psd_lower_bound" => Ok(Self::CertifiedPsdLowerBound),
276            "uncertified_approximation" => Ok(Self::UncertifiedApproximation),
277            other => Err(format!(
278                "fisher_factor_kind must be 'exact_full', 'certified_psd_lower_bound', or \
279                 'uncertified_approximation'; got {other:?}"
280            )),
281        }
282    }
283}
284
285/// The single per-row metric object. Holds one low-rank factor stack `U_n` (or
286/// none, for Euclidean) plus the validated PSD blocks, tagged with its
287/// [`MetricProvenance`].
288///
289/// `p` is the output dimensionality (residual / Jacobian-column dimension); the
290/// per-row factor `U_n ∈ ℝ^{p × rank}` so `W_n = U_n U_nᵀ ∈ ℝ^{p × p}` without
291/// ever being materialized as `p × p` in any hot path.
292#[derive(Clone, Debug)]
293pub struct RowMetric {
294    provenance: MetricProvenance,
295    n_rows: usize,
296    p: usize,
297    rank: usize,
298    /// `(n_rows, p * rank)` row-major: `U_n[i, k] = u[n, i * rank + k]`. `None`
299    /// for [`MetricProvenance::Euclidean`] (the identity factor is implicit).
300    factors: Option<Arc<Array2<f64>>>,
301    /// **Solver-only** Tikhonov floor `δ` added as `δ I_p` to make a
302    /// rank-deficient `U_n U_nᵀ` invertible for an *internal solve only*.
303    ///
304    /// Invariant (mirrors `RidgePolicy::solver_only`, #747): `δ` **never** enters
305    /// any quantity that feeds the evidence criterion. The criterion-facing
306    /// quad-form / whitening / fisher-mass methods all use the *un-floored*
307    /// `U_n U_nᵀ`; only [`Self::solve_floor`]-tagged solver helpers see `δ`. A
308    /// nonzero floor therefore cannot bias the objective the optimizer reports.
309    solver_delta: f64,
310    /// Per-row traces `tr(M_n)` of the criterion-facing (un-floored) metric.
311    ///
312    /// This is the only dense-block reduction any consumer reads (the #980
313    /// Fisher-mass row measure); the `(n_rows, p, p)` block stack itself is
314    /// validated **streamingly** at construction through
315    /// [`normalize_fisher_rao_blocks`] one row at a time and then dropped.
316    /// Retaining it was `n·p²·8` bytes — 13 GiB at `(n=2000, p=896)` and an
317    /// OOM at LLM-scale `p` — for a record nothing ever re-read. The solver
318    /// `δ` is deliberately *not* baked in here, so this is the
319    /// criterion-facing trace.
320    traces: ndarray::Array1<f64>,
321    /// Explicit output-Fisher factor status. `None` for Euclidean and structured
322    /// residual metrics, which do not claim to approximate an output Fisher.
323    fisher_factor_kind: Option<FisherFactorKind>,
324    /// Optional per-row non-negative Fisher tail-trace diagnostic. It never
325    /// changes `M_n = U_n U_n^T`, the criterion, solver, or factor status: only
326    /// [`FisherFactorKind`] can distinguish exact, certified-lower-bound, and
327    /// uncertified operators (#2249/#2263).
328    truncation_mass_residual: Option<Arc<ndarray::Array1<f64>>>,
329}
330
331impl RowMetric {
332    /// Euclidean metric: `W_n = I_p` for all `n`. Whitening is the identity, so
333    /// the likelihood residual path is bit-for-bit the prior isotropic `φ̂`.
334    ///
335    /// Constructed directly: the identity stack is PSD axiomatically, so
336    /// routing it through the dense normalizer would materialize and
337    /// spectrum-check `n` identity blocks (`n·p²` memory, `n·p³` flops) to
338    /// validate a tautology. `tr(I_p) = p` per row.
339    pub fn euclidean(n_rows: usize, p: usize) -> Result<Self, String> {
340        Ok(Self {
341            provenance: MetricProvenance::Euclidean,
342            n_rows,
343            p,
344            rank: p,
345            factors: None,
346            solver_delta: 0.0,
347            traces: ndarray::Array1::<f64>::from_elem(n_rows, p as f64),
348            fisher_factor_kind: None,
349            truncation_mass_residual: None,
350        })
351    }
352
353    /// Output-Fisher metric: per-row low-rank factors `U_n ∈ ℝ^{p × rank}`
354    /// supplied as a `(n_rows, p * rank)` row-major matrix (`U_n[i, k] =
355    /// u[n, i * rank + k]`). The induced `M_n = U_n U_nᵀ` is PSD by
356    /// construction; it is validated through [`normalize_fisher_rao_blocks`] so
357    /// the validation path is shared. No solver floor (`δ = 0`).
358    pub fn output_fisher(u: Arc<Array2<f64>>, p: usize, rank: usize) -> Result<Self, String> {
359        Self::from_factors(MetricProvenance::OutputFisher { rank }, u, p, rank, 0.0)
360    }
361
362    /// Downstream-influence output-Fisher metric: per-row factors `U_n ∈
363    /// ℝ^{p × rank}` whose `M_n = U_n U_nᵀ` is the aggregated output Fisher of
364    /// position `n` over the **future** positions it reaches through the KV path
365    /// ([`MetricProvenance::OutputFisherDownstream`], #980 mechanism 2). The
366    /// factor layout is identical to [`Self::output_fisher`]; only the
367    /// provenance tag (and hence the scientific reading) differs. Whitens
368    /// nothing, drives the gauge / lens / enrichment exactly as the
369    /// same-position metric does — the consuming machinery is provenance-generic
370    /// (see [`Self::is_output_fisher_like`]).
371    pub fn output_fisher_downstream(
372        u: Arc<Array2<f64>>,
373        p: usize,
374        rank: usize,
375    ) -> Result<Self, String> {
376        Self::from_factors(
377            MetricProvenance::OutputFisherDownstream { rank },
378            u,
379            p,
380            rank,
381            0.0,
382        )
383    }
384
385    /// **Rung 1** — the output-Fisher metric as a reconstruction *likelihood
386    /// weight* (GLS in nats): per-row `s`-probe factors `U_n ∈ ℝ^{p × probes}`
387    /// supplied as a `(n_rows, p * probes)` row-major matrix
388    /// (`U_n[i, k] = u[n, i * probes + k]`), so that column `k` is the probe
389    /// vector `v_k = J_nᵀ F_n^{1/2} u_k` and `M_n = U_n U_nᵀ ≈ G_n`. Unlike
390    /// [`Self::output_fisher`], the resulting metric returns
391    /// `whitens_likelihood() == true`: the data-fit prices reconstruction error
392    /// as `½ eᵀ G_n e`. Validated through [`normalize_fisher_rao_blocks`] like
393    /// every factored metric; no solver floor (`δ = 0`).
394    ///
395    /// See [`pack_probe_factors`] to build `u` from a natural `(n, p, s)` probe
396    /// stack emitted at harvest time.
397    pub fn behavioral_fisher(u: Arc<Array2<f64>>, p: usize, probes: usize) -> Result<Self, String> {
398        Self::from_factors(
399            MetricProvenance::BehavioralFisher { probes },
400            u,
401            p,
402            probes,
403            0.0,
404        )
405    }
406
407    /// Like [`Self::output_fisher`] but with a **solver-only** Tikhonov floor
408    /// `δ ≥ 0`. The floor is recorded for solver helpers only; every
409    /// criterion-facing method (`quad_form`, `whiten_residual`, `fisher_mass`)
410    /// ignores it (#747 discipline), so the evidence criterion is `δ`-free.
411    pub fn output_fisher_with_solver_floor(
412        u: Arc<Array2<f64>>,
413        p: usize,
414        rank: usize,
415        solver_delta: f64,
416    ) -> Result<Self, String> {
417        if !(solver_delta.is_finite() && solver_delta >= 0.0) {
418            return Err(format!(
419                "RowMetric::output_fisher_with_solver_floor: solver_delta must be finite and \
420                 non-negative; got {solver_delta}"
421            ));
422        }
423        Self::from_factors(
424            MetricProvenance::OutputFisher { rank },
425            u,
426            p,
427            rank,
428            solver_delta,
429        )
430    }
431
432    /// Structured-residual whitening from supplied per-row precision factors.
433    ///
434    /// `u` carries the per-row factor stack `U_n ∈ ℝ^{p × rank}` (row-major flat)
435    /// with `U_n U_nᵀ = M_n = Σ_n^{-1}` — the precision of the **estimated**
436    /// residual-covariance noise model. This is the low-level constructor; #974
437    /// producers that *fit* `Σ_n` (a low-rank factor + diagonal + smooth
438    /// activity-scale) assemble these factors and call through here. Because the
439    /// provenance is
440    /// [`MetricProvenance::WhitenedStructured`], [`Self::whitens_likelihood`] is
441    /// `true`: a metric built this way is the first that whitens the likelihood.
442    pub fn whitened_structured(u: Arc<Array2<f64>>, p: usize, rank: usize) -> Result<Self, String> {
443        Self::from_factors(
444            MetricProvenance::WhitenedStructured { factor_rank: rank },
445            u,
446            p,
447            rank,
448            0.0,
449        )
450    }
451
452    fn from_factors(
453        provenance: MetricProvenance,
454        u: Arc<Array2<f64>>,
455        p: usize,
456        rank: usize,
457        solver_delta: f64,
458    ) -> Result<Self, String> {
459        let n_rows = u.nrows();
460        if u.ncols() != p * rank {
461            return Err(format!(
462                "RowMetric::from_factors: factor matrix has {} cols; expected p*rank = {}*{} = {}",
463                u.ncols(),
464                p,
465                rank,
466                p * rank
467            ));
468        }
469        if !u.iter().all(|v| v.is_finite()) {
470            return Err("RowMetric::from_factors: factors must be finite".to_string());
471        }
472        // Materialize W_n = U_n U_nᵀ one row at a time (PSD by construction),
473        // validate each through the single shared normalizer rather than
474        // reimplementing the PSD check, record its trace, and drop the block.
475        // Streaming keeps construction O(p²) memory; the former whole-stack
476        // materialization retained `n·p²` doubles nothing ever re-read.
477        let mut traces = ndarray::Array1::<f64>::zeros(n_rows);
478        let mut full = Array3::<f64>::zeros((1, p, p));
479        for row in 0..n_rows {
480            for i in 0..p {
481                for j in 0..p {
482                    let mut acc = 0.0;
483                    for k in 0..rank {
484                        acc += u[[row, i * rank + k]] * u[[row, j * rank + k]];
485                    }
486                    full[[0, i, j]] = acc;
487                }
488            }
489            normalize_fisher_rao_blocks(full.view().into_dyn(), 1, p)
490                .map_err(|e| format!("RowMetric::from_factors: row {row}: {e}"))?;
491            let mut tr = 0.0_f64;
492            for i in 0..p {
493                tr += full[[0, i, i]];
494            }
495            traces[row] = tr;
496        }
497        Ok(Self {
498            provenance,
499            n_rows,
500            p,
501            rank,
502            factors: Some(u),
503            solver_delta,
504            traces,
505            fisher_factor_kind: match provenance {
506                MetricProvenance::OutputFisher { .. }
507                | MetricProvenance::OutputFisherDownstream { .. }
508                | MetricProvenance::BehavioralFisher { .. } => {
509                    Some(FisherFactorKind::UncertifiedApproximation)
510                }
511                MetricProvenance::Euclidean | MetricProvenance::WhitenedStructured { .. } => None,
512            },
513            truncation_mass_residual: None,
514        })
515    }
516
517    /// Attach an explicit mathematical certificate to a factored output-Fisher
518    /// metric. Constructors deliberately default to `UncertifiedApproximation`:
519    /// exactness or Loewner-order dominance must be asserted by the producer,
520    /// never inferred from rank or residual trace.
521    pub fn with_fisher_factor_kind(mut self, kind: FisherFactorKind) -> Result<Self, String> {
522        if self.fisher_factor_kind.is_none() {
523            return Err(
524                "RowMetric::with_fisher_factor_kind requires an output-Fisher metric".to_string(),
525            );
526        }
527        match kind {
528            FisherFactorKind::ExactFull if self.truncation_mass_residual.is_some() => {
529                return Err(
530                    "RowMetric::with_fisher_factor_kind ExactFull forbids an omitted-trace record"
531                        .to_string(),
532                );
533            }
534            FisherFactorKind::CertifiedPsdLowerBound if self.truncation_mass_residual.is_none() => {
535                return Err(
536                    "RowMetric::with_fisher_factor_kind CertifiedPsdLowerBound requires an exact omitted-trace record"
537                        .to_string(),
538                );
539            }
540            // The remaining (kind, record) pairs are exactly the consistent
541            // ones: an ExactFull factor with no omitted-trace record, a
542            // CertifiedPsdLowerBound with one, and UncertifiedApproximation,
543            // which asserts nothing about the omitted mass either way.
544            FisherFactorKind::ExactFull
545            | FisherFactorKind::CertifiedPsdLowerBound
546            | FisherFactorKind::UncertifiedApproximation => {}
547        }
548        self.fisher_factor_kind = Some(kind);
549        Ok(self)
550    }
551
552    /// Attach the harvested per-row omitted Fisher trace to this factored
553    /// metric. This is an audit channel, not a metric modification.
554    pub fn with_truncation_mass_residual(
555        mut self,
556        residual: Arc<ndarray::Array1<f64>>,
557    ) -> Result<Self, String> {
558        if self.factors.is_none() {
559            return Err(
560                "RowMetric::with_truncation_mass_residual requires a factored metric".to_string(),
561            );
562        }
563        if residual.len() != self.n_rows {
564            return Err(format!(
565                "RowMetric::with_truncation_mass_residual requires {} rows; got {}",
566                self.n_rows,
567                residual.len()
568            ));
569        }
570        for (row, &value) in residual.iter().enumerate() {
571            if !(value.is_finite() && value >= 0.0) {
572                return Err(format!(
573                    "RowMetric::with_truncation_mass_residual row {row} must be finite and non-negative; got {value}"
574                ));
575            }
576        }
577        self.truncation_mass_residual = Some(residual);
578        Ok(self)
579    }
580
581    /// Restrict the metric to the rows `rows` (an index subset or permutation),
582    /// preserving provenance, `p`, `rank`, and the solver floor. The
583    /// outer-criterion row subsample uses this to whiten the subsampled fit
584    /// through the SAME per-row metric the full-`N` fit uses, so the ρ search
585    /// ranks the delivered criterion (e.g. a #974 structured-whitening fit is not
586    /// silently searched unwhitened). Each gathered row's factor block is copied
587    /// verbatim, so the induced `M_n = U_n U_nᵀ` is bit-identical to the full
588    /// metric's on every selected row.
589    pub fn gather_rows(&self, rows: &[usize]) -> Result<Self, String> {
590        for (pos, &r) in rows.iter().enumerate() {
591            if r >= self.n_rows {
592                return Err(format!(
593                    "RowMetric::gather_rows: row index {r} at position {pos} is out of bounds \
594                     (n_rows = {})",
595                    self.n_rows
596                ));
597            }
598        }
599        match self.factors.as_ref() {
600            // Euclidean carries an implicit identity factor per row, so the subset
601            // is just a smaller identity stack — no factor storage to gather.
602            None => Self::euclidean(rows.len(), self.p),
603            Some(factors) => {
604                let cols = self.p * self.rank;
605                let mut sub = Array2::<f64>::zeros((rows.len(), cols));
606                for (pos, &r) in rows.iter().enumerate() {
607                    sub.row_mut(pos).assign(&factors.row(r));
608                }
609                // Re-runs the shared PSD normalizer on the subset (a subset of
610                // valid rows stays valid) and preserves the exact provenance and
611                // solver floor.
612                let mut metric = Self::from_factors(
613                    self.provenance,
614                    Arc::new(sub),
615                    self.p,
616                    self.rank,
617                    self.solver_delta,
618                )?;
619                metric.fisher_factor_kind = self.fisher_factor_kind;
620                match self.truncation_mass_residual.as_ref() {
621                    None => Ok(metric),
622                    Some(residual) => {
623                        let gathered =
624                            ndarray::Array1::from_iter(rows.iter().map(|&row| residual[row]));
625                        metric.with_truncation_mass_residual(Arc::new(gathered))
626                    }
627                }
628            }
629        }
630    }
631
632    /// The provenance tag (consumed by Object 4 to certify the inner product).
633    pub fn provenance(&self) -> MetricProvenance {
634        self.provenance
635    }
636
637    /// Explicit output-Fisher factor status, never inferred from diagnostics.
638    pub fn fisher_factor_kind(&self) -> Option<FisherFactorKind> {
639        self.fisher_factor_kind
640    }
641
642    /// Whether this metric is allowed to **whiten the likelihood** (i.e. replace
643    /// the isotropic reconstruction data-fit `½ rᵀr` with the whitened
644    /// `½ rᵀ M_n r`).
645    ///
646    /// This is TRUE for two provenances, for two distinct reasons:
647    ///
648    /// * [`MetricProvenance::WhitenedStructured`] — a genuinely *estimated noise
649    ///   model* (a factor-analytic residual covariance, #974), for which
650    ///   whitening the likelihood is the statistically correct thing to do; and
651    /// * [`MetricProvenance::BehavioralFisher`] — the **Rung 1** deliberate
652    ///   choice to price reconstruction error in nats: the output-Fisher metric
653    ///   `G_n` installed *as the loss weight* (`½ eᵀ G_n e`), a generalized
654    ///   least-squares reconstruction. Because `G_n` is a fixed per-row metric
655    ///   the problem stays linear-Gaussian, so REML/evidence/EDF are preserved.
656    ///
657    /// It is FALSE for [`MetricProvenance::Euclidean`] (nothing to whiten by) and
658    /// for the *gauge-only* [`MetricProvenance::OutputFisher`] /
659    /// [`MetricProvenance::OutputFisherDownstream`]: there the output-Fisher
660    /// inner product is an **output-geometry gauge**, and whitening the
661    /// likelihood by it *implicitly* (without the caller electing GLS) would
662    /// silently replace the reconstruction loss with a Fisher pullback — the #980
663    /// failure mode, and the reason the two-tier harvest can withhold factors
664    /// from a row without changing its loss. `BehavioralFisher` is the *explicit*
665    /// election of that same arithmetic as the intended objective.
666    pub fn whitens_likelihood(&self) -> bool {
667        matches!(
668            self.provenance,
669            MetricProvenance::WhitenedStructured { .. } | MetricProvenance::BehavioralFisher { .. }
670        )
671    }
672
673    /// Whether this metric **drives the gauge** — i.e. the isometry-penalty
674    /// pullback weight is taken from it rather than the identity.
675    ///
676    /// TRUE for any non-[`MetricProvenance::Euclidean`] provenance: both
677    /// [`MetricProvenance::OutputFisher`] and
678    /// [`MetricProvenance::WhitenedStructured`] supply a non-identity per-row
679    /// inner product the gauge pulls back through. Euclidean reduces the gauge
680    /// pullback to the bare `J_nᵀ J_n`, so it does not drive the gauge.
681    pub fn drives_gauge(&self) -> bool {
682        !matches!(self.provenance, MetricProvenance::Euclidean)
683    }
684
685    /// Whether this metric is an **output-Fisher gauge** — either the
686    /// same-position [`MetricProvenance::OutputFisher`] or the downstream
687    /// [`MetricProvenance::OutputFisherDownstream`] (#980). The two share every
688    /// consumer behavior (Sym(F) separation under the gauge, two-lens coupling,
689    /// steering geometry, enrichment); they differ only in the *scientific*
690    /// reading of what behavioral coupling means (same-position vs
691    /// forward-looking). Consumers that gate on "is this an output-Fisher
692    /// pullback" should use this predicate rather than matching one variant, so
693    /// the downstream metric rides the identical path.
694    pub fn is_output_fisher_like(&self) -> bool {
695        matches!(
696            self.provenance,
697            MetricProvenance::OutputFisher { .. } | MetricProvenance::OutputFisherDownstream { .. }
698        )
699    }
700
701    /// Number of rows the metric is defined over.
702    pub fn n_rows(&self) -> usize {
703        self.n_rows
704    }
705
706    /// Output dimensionality `p` (residual / Jacobian-column dimension).
707    pub fn p_out(&self) -> usize {
708        self.p
709    }
710
711    /// The factor rank: the dimension of the whitened residual
712    /// [`Self::whiten_residual_row`] returns (and the column count of the per-row
713    /// factor `U_n ∈ ℝ^{p × rank}`). For [`MetricProvenance::Euclidean`] this is
714    /// `p` (the implicit identity factor), so a consumer that sizes a whitened
715    /// buffer by `metric_rank()` gets the right length in every provenance.
716    pub fn metric_rank(&self) -> usize {
717        self.rank
718    }
719
720    /// Per-row traces `tr(M_n)` of the criterion-facing (un-floored) metric —
721    /// the Fisher-mass reduction the #980 row measure consumes. The dense
722    /// `(n_rows, p, p)` stack is validated streamingly at construction and
723    /// never retained; consumers wanting an explicit `W_n` rebuild it from
724    /// [`Self::metric_rank`]-sized factors.
725    pub fn row_traces(&self) -> ndarray::ArrayView1<'_, f64> {
726        self.traces.view()
727    }
728
729    /// Omitted Fisher trace for `row`, when the harvest supplied one.
730    pub fn truncation_mass_residual(&self, row: usize) -> Option<f64> {
731        self.truncation_mass_residual
732            .as_ref()
733            .map(|residual| residual[row])
734    }
735
736    /// Fraction of total audited Fisher trace omitted at `row`. `None` means the
737    /// factor stack carried no truncation audit. A zero-total metric has zero
738    /// omitted fraction when its reported residual is also zero.
739    pub fn truncation_mass_residual_fraction(&self, row: usize) -> Option<f64> {
740        self.truncation_mass_residual(row).map(|residual| {
741            let total = self.traces[row] + residual;
742            if total > 0.0 { residual / total } else { 0.0 }
743        })
744    }
745
746    /// Whiten a single `p`-dimensional residual row `r` into the coordinates
747    /// whose squared Euclidean norm equals `rᵀ W_n r`.
748    ///
749    /// * Euclidean: returns `r` unchanged (`‖r‖² = rᵀ I r`), so the likelihood
750    ///   reproduces the isotropic `½ rᵀr` data-fit bit-for-bit.
751    /// * Factored: returns `U_nᵀ r ∈ ℝ^{rank}`, with
752    ///   `‖U_nᵀ r‖² = rᵀ U_n U_nᵀ r = rᵀ W_n r`.
753    ///
754    /// This is the load-bearing identity that lets the data-fit loop sum
755    /// `0.5 * Σ whitened²` and recover exactly `rᵀ W_n r` whatever the
756    /// provenance.
757    pub fn whiten_residual_row(&self, row: usize, r: ArrayView1<'_, f64>) -> Vec<f64> {
758        match &self.factors {
759            None => r.iter().copied().collect(),
760            Some(u) => {
761                let mut out = vec![0.0_f64; self.rank];
762                for k in 0..self.rank {
763                    let mut acc = 0.0;
764                    for i in 0..self.p {
765                        acc += u[[row, i * self.rank + k]] * r[i];
766                    }
767                    out[k] = acc;
768                }
769                out
770            }
771        }
772    }
773
774    /// The factor entry `U_n[i, k]` for one row (`i ∈ [0, p)`, `k ∈ [0, rank)`).
775    /// For [`MetricProvenance::Euclidean`] the implicit factor is `I_p`, so this
776    /// returns `1.0` when `i == k` and `0.0` otherwise — letting a consumer that
777    /// whitens a Jacobian via `factor_entry` produce the identity whitening
778    /// without a provenance branch. Reads the **un-floored** factors (criterion
779    /// face, #747).
780    #[inline]
781    pub fn factor_entry(&self, row: usize, i: usize, k: usize) -> f64 {
782        match &self.factors {
783            None => {
784                if i == k {
785                    1.0
786                } else {
787                    0.0
788                }
789            }
790            Some(u) => u[[row, i * self.rank + k]],
791        }
792    }
793
794    /// Apply the full per-row metric `M_n x = U_n (U_nᵀ x) ∈ ℝ^p` for one
795    /// `p`-vector `x`, formed factored (`rank` flops in, `p` flops out) — never
796    /// materializing `M_n` as `p × p`. Euclidean returns `x` unchanged
797    /// (`M_n = I_p`). This is the p-space metric-applied vector the SAE β-tier
798    /// data-fit gradient contracts (β lives in p-output space, so its gradient
799    /// needs `M_n r_n`, not the rank-space whitened residual `U_nᵀ r_n`). Uses the
800    /// **un-floored** factors (criterion face, `δ`-free, #747 invariant).
801    pub fn apply_metric_row(&self, row: usize, x: ArrayView1<'_, f64>) -> Vec<f64> {
802        match &self.factors {
803            None => x.iter().copied().collect(),
804            Some(u) => {
805                // w = U_nᵀ x ∈ ℝ^{rank}.
806                let mut w = vec![0.0_f64; self.rank];
807                for k in 0..self.rank {
808                    let mut acc = 0.0;
809                    for i in 0..self.p {
810                        acc += u[[row, i * self.rank + k]] * x[i];
811                    }
812                    w[k] = acc;
813                }
814                // out = U_n w ∈ ℝ^p.
815                let mut out = vec![0.0_f64; self.p];
816                for i in 0..self.p {
817                    let mut acc = 0.0;
818                    for k in 0..self.rank {
819                        acc += u[[row, i * self.rank + k]] * w[k];
820                    }
821                    out[i] = acc;
822                }
823                out
824            }
825        }
826    }
827
828    /// Pullback metric `g_n = J_nᵀ W_n J_n` for one row, formed as
829    /// `(J_nᵀ U_n)(U_nᵀ J_n)` — never materializing the `p × p` `W_n`.
830    ///
831    /// `j_row` is the row's Jacobian `J_n ∈ ℝ^{p × d}` flattened row-major
832    /// (`J_n[i, a] = j_row[i * d + a]`). Returns the `d × d` `g_n`.
833    pub fn pullback(&self, row: usize, j_row: &[f64], d: usize) -> Array2<f64> {
834        match &self.factors {
835            None => {
836                // W_n = I_p ⇒ g_n = J_nᵀ J_n.
837                let mut g = Array2::<f64>::zeros((d, d));
838                for a in 0..d {
839                    for b in a..d {
840                        let mut acc = 0.0;
841                        for i in 0..self.p {
842                            acc += j_row[i * d + a] * j_row[i * d + b];
843                        }
844                        g[[a, b]] = acc;
845                        g[[b, a]] = acc;
846                    }
847                }
848                g
849            }
850            Some(u) => {
851                // M_n = U_nᵀ J_n ∈ ℝ^{rank × d}; g_n = M_nᵀ M_n.
852                let mut m = Array2::<f64>::zeros((self.rank, d));
853                for k in 0..self.rank {
854                    for a in 0..d {
855                        let mut acc = 0.0;
856                        for i in 0..self.p {
857                            acc += u[[row, i * self.rank + k]] * j_row[i * d + a];
858                        }
859                        m[[k, a]] = acc;
860                    }
861                }
862                let mut g = Array2::<f64>::zeros((d, d));
863                for a in 0..d {
864                    for b in a..d {
865                        let mut acc = 0.0;
866                        for k in 0..self.rank {
867                            acc += m[[k, a]] * m[[k, b]];
868                        }
869                        g[[a, b]] = acc;
870                        g[[b, a]] = acc;
871                    }
872                }
873                g
874            }
875        }
876    }
877
878    /// Quadratic form `r_nᵀ M_n r_n` for one row's residual `r_n ∈ ℝ^p`, formed
879    /// **factored** as `‖U_nᵀ r_n‖²` — never materializing the `p × p` `M_n`.
880    ///
881    /// This is the criterion-facing squared residual the likelihood sums; it uses
882    /// the **un-floored** `U_n U_nᵀ`, so the solver `δ` does not enter it
883    /// (#747 invariant). Euclidean provenance returns the bit-identical `‖r_n‖²`.
884    #[inline]
885    pub fn quad_form(&self, row: usize, r: ArrayView1<'_, f64>) -> f64 {
886        match &self.factors {
887            None => r.iter().map(|&v| v * v).sum(),
888            Some(_) => self
889                .whiten_residual_row(row, r)
890                .iter()
891                .map(|&w| w * w)
892                .sum(),
893        }
894    }
895
896    /// Whiten a per-row Jacobian `J_n ∈ ℝ^{p × d}` (row-major flat,
897    /// `J_n[i, a] = j_row[i * d + a]`) into `M_n = U_nᵀ J_n ∈ ℝ^{rank × d}` so
898    /// that `M_nᵀ M_n = J_nᵀ (U_n U_nᵀ) J_n = J_nᵀ W_n J_n` is the pullback
899    /// **without** any `p × p` intermediate. Euclidean returns `J_n` reshaped to
900    /// `(p, d)` (the identity whitening). Solver `δ` is not applied (criterion
901    /// face).
902    pub fn whiten_jacobian(&self, row: usize, j_row: &[f64], d: usize) -> Array2<f64> {
903        match &self.factors {
904            None => {
905                let mut out = Array2::<f64>::zeros((self.p, d));
906                for i in 0..self.p {
907                    for a in 0..d {
908                        out[[i, a]] = j_row[i * d + a];
909                    }
910                }
911                out
912            }
913            Some(u) => {
914                let mut m = Array2::<f64>::zeros((self.rank, d));
915                for k in 0..self.rank {
916                    for a in 0..d {
917                        let mut acc = 0.0;
918                        for i in 0..self.p {
919                            acc += u[[row, i * self.rank + k]] * j_row[i * d + a];
920                        }
921                        m[[k, a]] = acc;
922                    }
923                }
924                m
925            }
926        }
927    }
928
929    /// Fisher mass of a per-row output vector `x_n ∈ ℝ^p`: the scalar
930    /// `x_nᵀ M_n x_n` (alias of [`Self::quad_form`] read as an information mass
931    /// rather than a residual square). Factored, never `p × p`, `δ`-free.
932    #[inline]
933    pub fn fisher_mass(&self, row: usize, x: ArrayView1<'_, f64>) -> f64 {
934        self.quad_form(row, x)
935    }
936
937    /// The **solver-only** Tikhonov floor `δ` (#747). Returned for internal
938    /// solver helpers that need `U_n U_nᵀ + δ I` to be invertible; by contract
939    /// no caller may fold this into a criterion-facing quantity. Always `0` for
940    /// Euclidean and for factored metrics built without an explicit floor.
941    pub fn solver_floor(&self) -> f64 {
942        self.solver_delta
943    }
944
945    /// The gauge view of this metric: the
946    /// [`crate::WeightField`] the isometry penalty pulls back through.
947    ///
948    /// This is the **single** way an `IsometryPenalty` acquires a non-identity
949    /// gauge metric — the independent `WeightField` setter has been removed — so
950    /// the gauge metric is, by construction, the same object the likelihood
951    /// whitens with.
952    pub fn to_weight_field(&self) -> crate::WeightField {
953        use crate::WeightField;
954        match &self.factors {
955            None => WeightField::Identity,
956            Some(u) => WeightField::Factored {
957                u: Arc::clone(u),
958                rank: self.rank,
959                p_out: self.p,
960            },
961        }
962    }
963}
964
965/// Pack a harvest-emitted probe stack into the row-major factor layout
966/// [`RowMetric::behavioral_fisher`] expects.
967///
968/// The harvest boundary (the model-interaction side) emits, per token, `s`
969/// probe vectors `vₖ = J_nᵀ F_n^{1/2} uₖ ∈ ℝ^p` — the natural shape is
970/// `probes[n, i, k] = (vₖ)ᵢ`, an `(n_rows, p, probes)` stack. This assembles the
971/// `(n_rows, p · probes)` row-major matrix `u[n, i·probes + k] = probes[n, i, k]`
972/// that the constructor consumes so that column `k` of the per-row factor `U_n`
973/// is exactly probe `vₖ` and `M_n = U_n U_nᵀ = Σₖ vₖ vₖᵀ ≈ G_n`.
974///
975/// This is a pure repack of the standard C-order flattening; it exists so the
976/// harvest → metric seam is a single named, validated Rust surface rather than
977/// an ad-hoc reshape at each call site. Errors on non-finite entries so the
978/// failure is caught here rather than deep in [`normalize_fisher_rao_blocks`].
979pub fn pack_probe_factors(probes: ndarray::ArrayView3<'_, f64>) -> Result<Array2<f64>, String> {
980    let (n_rows, p, s) = probes.dim();
981    if s == 0 {
982        return Err("pack_probe_factors: need at least one probe (s == 0)".to_string());
983    }
984    if !probes.iter().all(|v| v.is_finite()) {
985        return Err("pack_probe_factors: probe entries must be finite".to_string());
986    }
987    let mut u = Array2::<f64>::zeros((n_rows, p * s));
988    for n in 0..n_rows {
989        for i in 0..p {
990            for k in 0..s {
991                u[[n, i * s + k]] = probes[[n, i, k]];
992            }
993        }
994    }
995    Ok(u)
996}
997
998#[cfg(test)]
999mod tests {
1000    use super::*;
1001    use ndarray::array;
1002
1003    // ── RowMetric::euclidean ──────────────────────────────────────────────────
1004
1005    #[test]
1006    fn euclidean_metric_has_correct_dimensions() {
1007        let m = RowMetric::euclidean(5, 3).unwrap();
1008        assert_eq!(m.n_rows(), 5);
1009        assert_eq!(m.p_out(), 3);
1010        assert_eq!(m.metric_rank(), 3);
1011    }
1012
1013    #[test]
1014    fn euclidean_metric_traces_equal_p() {
1015        let p = 4_usize;
1016        let m = RowMetric::euclidean(3, p).unwrap();
1017        for tr in m.row_traces().iter() {
1018            assert!((*tr - p as f64).abs() < 1e-14, "trace {tr} != p={p}");
1019        }
1020    }
1021
1022    #[test]
1023    fn euclidean_provenance_is_euclidean() {
1024        let m = RowMetric::euclidean(1, 2).unwrap();
1025        assert_eq!(m.provenance(), MetricProvenance::Euclidean);
1026    }
1027
1028    #[test]
1029    fn euclidean_does_not_whiten_likelihood() {
1030        let m = RowMetric::euclidean(1, 2).unwrap();
1031        assert!(!m.whitens_likelihood());
1032    }
1033
1034    #[test]
1035    fn euclidean_does_not_drive_gauge() {
1036        let m = RowMetric::euclidean(1, 2).unwrap();
1037        assert!(!m.drives_gauge());
1038    }
1039
1040    #[test]
1041    fn euclidean_is_not_output_fisher_like() {
1042        let m = RowMetric::euclidean(1, 2).unwrap();
1043        assert!(!m.is_output_fisher_like());
1044    }
1045
1046    #[test]
1047    fn euclidean_solver_floor_is_zero() {
1048        let m = RowMetric::euclidean(1, 2).unwrap();
1049        assert_eq!(m.solver_floor(), 0.0);
1050    }
1051
1052    #[test]
1053    fn euclidean_to_weight_field_is_identity() {
1054        let m = RowMetric::euclidean(1, 2).unwrap();
1055        assert!(matches!(m.to_weight_field(), WeightField::Identity));
1056    }
1057
1058    #[test]
1059    fn euclidean_whiten_residual_is_passthrough() {
1060        let m = RowMetric::euclidean(1, 3).unwrap();
1061        let r = array![1.0_f64, 2.0, 3.0];
1062        let w = m.whiten_residual_row(0, r.view());
1063        assert_eq!(w, vec![1.0, 2.0, 3.0]);
1064    }
1065
1066    #[test]
1067    fn euclidean_factor_entry_is_identity() {
1068        let m = RowMetric::euclidean(1, 3).unwrap();
1069        assert_eq!(m.factor_entry(0, 0, 0), 1.0);
1070        assert_eq!(m.factor_entry(0, 1, 1), 1.0);
1071        assert_eq!(m.factor_entry(0, 2, 2), 1.0);
1072        assert_eq!(m.factor_entry(0, 0, 1), 0.0);
1073        assert_eq!(m.factor_entry(0, 1, 0), 0.0);
1074    }
1075
1076    #[test]
1077    fn euclidean_quad_form_is_squared_norm() {
1078        let m = RowMetric::euclidean(1, 3).unwrap();
1079        let r = array![1.0_f64, 2.0, 2.0];
1080        assert!((m.quad_form(0, r.view()) - 9.0).abs() < 1e-14);
1081    }
1082
1083    // ── MetricProvenance predicates ───────────────────────────────────────────
1084
1085    #[test]
1086    fn output_fisher_drives_gauge_but_not_likelihood() {
1087        let u = Arc::new(array![[1.0_f64]]);
1088        let m = RowMetric::output_fisher(u, 1, 1).unwrap();
1089        assert!(m.drives_gauge());
1090        assert!(!m.whitens_likelihood());
1091        assert!(m.is_output_fisher_like());
1092    }
1093
1094    #[test]
1095    fn whitened_structured_whitens_likelihood_and_drives_gauge() {
1096        let u = Arc::new(array![[1.0_f64]]);
1097        let m = RowMetric::whitened_structured(u, 1, 1).unwrap();
1098        assert!(m.whitens_likelihood());
1099        assert!(m.drives_gauge());
1100        assert!(!m.is_output_fisher_like());
1101    }
1102
1103    #[test]
1104    fn behavioral_fisher_whitens_likelihood_and_drives_gauge() {
1105        // The Rung-1 deliberate GLS metric: unlike the gauge-only OutputFisher,
1106        // it whitens the reconstruction likelihood.
1107        let u = Arc::new(array![[1.0_f64, 0.5]]); // p=1, probes=2
1108        let m = RowMetric::behavioral_fisher(u, 1, 2).unwrap();
1109        assert!(m.whitens_likelihood());
1110        assert!(m.drives_gauge());
1111        assert_eq!(
1112            m.provenance(),
1113            MetricProvenance::BehavioralFisher { probes: 2 }
1114        );
1115        assert_eq!(m.metric_rank(), 2);
1116    }
1117
1118    #[test]
1119    fn behavioral_fisher_quad_form_is_probe_sum() {
1120        // p=2, s=2 probes v1=(1,0), v2=(0,2) → G = diag(1,4);
1121        // e=(3,1) → eᵀGe = 9·1 + 1·4 = 13 = Σ (vᵢᵀe)² = 3² + 2² = 13.
1122        // Column-major-within-row layout U[i,k]=u[i*probes+k]:
1123        //   U[0,0]=1 U[0,1]=0  U[1,0]=0 U[1,1]=2
1124        let u = Arc::new(array![[1.0_f64, 0.0, 0.0, 2.0]]);
1125        let m = RowMetric::behavioral_fisher(u, 2, 2).unwrap();
1126        let e = array![3.0_f64, 1.0];
1127        assert!((m.quad_form(0, e.view()) - 13.0).abs() < 1e-12);
1128    }
1129
1130    #[test]
1131    fn behavioral_fisher_g_identity_reproduces_euclidean_quad_form() {
1132        // GLS with G=I must reduce to plain MSE. Identity probes (s=p, U=I_p)
1133        // ⇒ M_n = I ⇒ quad_form == ‖e‖², matching Euclidean bit-for-bit, and
1134        // metric_rank == p so the whitened residual-dof accounting is unchanged.
1135        let p = 3;
1136        let mut u = Array2::<f64>::zeros((1, p * p));
1137        for i in 0..p {
1138            u[[0, i * p + i]] = 1.0;
1139        }
1140        let bf = RowMetric::behavioral_fisher(Arc::new(u), p, p).unwrap();
1141        let euc = RowMetric::euclidean(1, p).unwrap();
1142        let e = array![1.5_f64, -2.0, 0.25];
1143        assert_eq!(bf.metric_rank(), euc.metric_rank());
1144        assert!((bf.quad_form(0, e.view()) - euc.quad_form(0, e.view())).abs() < 1e-14);
1145        // and whitened residual is the residual itself (identity whitening)
1146        assert_eq!(bf.whiten_residual_row(0, e.view()), vec![1.5, -2.0, 0.25]);
1147    }
1148
1149    #[test]
1150    fn pack_probe_factors_matches_manual_layout() {
1151        use ndarray::Array3;
1152        // n=1, p=2, s=2: probes[0,i,k] = v_k[i]; v0=(1,3), v1=(2,4)
1153        let mut probes = Array3::<f64>::zeros((1, 2, 2));
1154        probes[[0, 0, 0]] = 1.0; // v0[0]
1155        probes[[0, 1, 0]] = 3.0; // v0[1]
1156        probes[[0, 0, 1]] = 2.0; // v1[0]
1157        probes[[0, 1, 1]] = 4.0; // v1[1]
1158        let u = pack_probe_factors(probes.view()).unwrap();
1159        // Layout U[i,k] = u[i*s + k]: [v0[0],v1[0], v0[1],v1[1]] = [1,2,3,4]
1160        assert_eq!(u.as_slice().unwrap(), &[1.0, 2.0, 3.0, 4.0]);
1161        // Round-trips into a valid metric whose G = v0 v0ᵀ + v1 v1ᵀ.
1162        let m = RowMetric::behavioral_fisher(Arc::new(u), 2, 2).unwrap();
1163        // e=(1,0): eᵀGe = v0[0]²+v1[0]² = 1+4 = 5.
1164        let e = array![1.0_f64, 0.0];
1165        assert!((m.quad_form(0, e.view()) - 5.0).abs() < 1e-12);
1166    }
1167
1168    #[test]
1169    fn pack_probe_factors_rejects_zero_probes() {
1170        use ndarray::Array3;
1171        let probes = Array3::<f64>::zeros((2, 3, 0));
1172        assert!(pack_probe_factors(probes.view()).is_err());
1173    }
1174
1175    #[test]
1176    fn output_fisher_downstream_is_output_fisher_like() {
1177        let u = Arc::new(array![[1.0_f64]]);
1178        let m = RowMetric::output_fisher_downstream(u, 1, 1).unwrap();
1179        assert!(m.is_output_fisher_like());
1180        assert!(m.drives_gauge());
1181    }
1182
1183    #[test]
1184    fn fisher_factor_status_is_never_inferred_from_zero_residual_2249() {
1185        let factors = Arc::new(Array2::from_elem((1, 1), 2.0));
1186        let metric = RowMetric::output_fisher(factors, 1, 1)
1187            .unwrap()
1188            .with_truncation_mass_residual(Arc::new(array![0.0]))
1189            .unwrap();
1190        assert_eq!(
1191            metric.fisher_factor_kind(),
1192            Some(FisherFactorKind::UncertifiedApproximation)
1193        );
1194        let certified = metric
1195            .clone()
1196            .with_fisher_factor_kind(FisherFactorKind::CertifiedPsdLowerBound)
1197            .unwrap();
1198        assert_eq!(
1199            certified.fisher_factor_kind(),
1200            Some(FisherFactorKind::CertifiedPsdLowerBound)
1201        );
1202        assert!(
1203            metric
1204                .with_fisher_factor_kind(FisherFactorKind::ExactFull)
1205                .is_err(),
1206            "an omitted-trace record is incompatible with an exact-full claim"
1207        );
1208    }
1209
1210    // ── WeightField::project_jac_row_with_u ──────────────────────────────────
1211
1212    #[test]
1213    fn project_jac_with_identity_returns_jac() {
1214        // p=2, rank=2, d=2; U=I_2, J=[[1,2],[3,4]] → M = U^T J = J
1215        let u_row = [1.0_f64, 0.0, 0.0, 1.0]; // U[i,k]=u[i*rank+k], I_2
1216        let j_row = [1.0_f64, 2.0, 3.0, 4.0]; // J[i,a]=j[i*d+a]
1217        let m = WeightField::project_jac_row_with_u(&u_row, &j_row, 2, 2, 2);
1218        assert!((m[[0, 0]] - 1.0).abs() < 1e-14);
1219        assert!((m[[0, 1]] - 2.0).abs() < 1e-14);
1220        assert!((m[[1, 0]] - 3.0).abs() < 1e-14);
1221        assert!((m[[1, 1]] - 4.0).abs() < 1e-14);
1222    }
1223
1224    #[test]
1225    fn project_jac_with_zeros_returns_zero_matrix() {
1226        let u_row = [0.0_f64, 0.0];
1227        let j_row = [1.0_f64, 2.0];
1228        let m = WeightField::project_jac_row_with_u(&u_row, &j_row, 2, 1, 1);
1229        assert_eq!(m[[0, 0]], 0.0);
1230    }
1231}