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    /// Structured-residual whitening from supplied per-row precision factors.
408    ///
409    /// `u` carries the per-row factor stack `U_n ∈ ℝ^{p × rank}` (row-major flat)
410    /// with `U_n U_nᵀ = M_n = Σ_n^{-1}` — the precision of the **estimated**
411    /// residual-covariance noise model. This is the low-level constructor; #974
412    /// producers that *fit* `Σ_n` (a low-rank factor + diagonal + smooth
413    /// activity-scale) assemble these factors and call through here. Because the
414    /// provenance is
415    /// [`MetricProvenance::WhitenedStructured`], [`Self::whitens_likelihood`] is
416    /// `true`: a metric built this way is the first that whitens the likelihood.
417    pub fn whitened_structured(u: Arc<Array2<f64>>, p: usize, rank: usize) -> Result<Self, String> {
418        Self::from_factors(
419            MetricProvenance::WhitenedStructured { factor_rank: rank },
420            u,
421            p,
422            rank,
423            0.0,
424        )
425    }
426
427    fn from_factors(
428        provenance: MetricProvenance,
429        u: Arc<Array2<f64>>,
430        p: usize,
431        rank: usize,
432        solver_delta: f64,
433    ) -> Result<Self, String> {
434        let n_rows = u.nrows();
435        if u.ncols() != p * rank {
436            return Err(format!(
437                "RowMetric::from_factors: factor matrix has {} cols; expected p*rank = {}*{} = {}",
438                u.ncols(),
439                p,
440                rank,
441                p * rank
442            ));
443        }
444        if !u.iter().all(|v| v.is_finite()) {
445            return Err("RowMetric::from_factors: factors must be finite".to_string());
446        }
447        // Materialize W_n = U_n U_nᵀ one row at a time (PSD by construction),
448        // validate each through the single shared normalizer rather than
449        // reimplementing the PSD check, record its trace, and drop the block.
450        // Streaming keeps construction O(p²) memory; the former whole-stack
451        // materialization retained `n·p²` doubles nothing ever re-read.
452        let mut traces = ndarray::Array1::<f64>::zeros(n_rows);
453        let mut full = Array3::<f64>::zeros((1, p, p));
454        for row in 0..n_rows {
455            for i in 0..p {
456                for j in 0..p {
457                    let mut acc = 0.0;
458                    for k in 0..rank {
459                        acc += u[[row, i * rank + k]] * u[[row, j * rank + k]];
460                    }
461                    full[[0, i, j]] = acc;
462                }
463            }
464            normalize_fisher_rao_blocks(full.view().into_dyn(), 1, p)
465                .map_err(|e| format!("RowMetric::from_factors: row {row}: {e}"))?;
466            let mut tr = 0.0_f64;
467            for i in 0..p {
468                tr += full[[0, i, i]];
469            }
470            traces[row] = tr;
471        }
472        Ok(Self {
473            provenance,
474            n_rows,
475            p,
476            rank,
477            factors: Some(u),
478            solver_delta,
479            traces,
480            fisher_factor_kind: match provenance {
481                MetricProvenance::OutputFisher { .. }
482                | MetricProvenance::OutputFisherDownstream { .. }
483                | MetricProvenance::BehavioralFisher { .. } => {
484                    Some(FisherFactorKind::UncertifiedApproximation)
485                }
486                MetricProvenance::Euclidean | MetricProvenance::WhitenedStructured { .. } => None,
487            },
488            truncation_mass_residual: None,
489        })
490    }
491
492    /// Attach an explicit mathematical certificate to a factored output-Fisher
493    /// metric. Constructors deliberately default to `UncertifiedApproximation`:
494    /// exactness or Loewner-order dominance must be asserted by the producer,
495    /// never inferred from rank or residual trace.
496    pub fn with_fisher_factor_kind(mut self, kind: FisherFactorKind) -> Result<Self, String> {
497        if self.fisher_factor_kind.is_none() {
498            return Err(
499                "RowMetric::with_fisher_factor_kind requires an output-Fisher metric".to_string(),
500            );
501        }
502        match kind {
503            FisherFactorKind::ExactFull if self.truncation_mass_residual.is_some() => {
504                return Err(
505                    "RowMetric::with_fisher_factor_kind ExactFull forbids an omitted-trace record"
506                        .to_string(),
507                );
508            }
509            FisherFactorKind::CertifiedPsdLowerBound if self.truncation_mass_residual.is_none() => {
510                return Err(
511                    "RowMetric::with_fisher_factor_kind CertifiedPsdLowerBound requires an exact omitted-trace record"
512                        .to_string(),
513                );
514            }
515            // The remaining (kind, record) pairs are exactly the consistent
516            // ones: an ExactFull factor with no omitted-trace record, a
517            // CertifiedPsdLowerBound with one, and UncertifiedApproximation,
518            // which asserts nothing about the omitted mass either way.
519            FisherFactorKind::ExactFull
520            | FisherFactorKind::CertifiedPsdLowerBound
521            | FisherFactorKind::UncertifiedApproximation => {}
522        }
523        self.fisher_factor_kind = Some(kind);
524        Ok(self)
525    }
526
527    /// Attach the harvested per-row omitted Fisher trace to this factored
528    /// metric. This is an audit channel, not a metric modification.
529    pub fn with_truncation_mass_residual(
530        mut self,
531        residual: Arc<ndarray::Array1<f64>>,
532    ) -> Result<Self, String> {
533        if self.factors.is_none() {
534            return Err(
535                "RowMetric::with_truncation_mass_residual requires a factored metric".to_string(),
536            );
537        }
538        if residual.len() != self.n_rows {
539            return Err(format!(
540                "RowMetric::with_truncation_mass_residual requires {} rows; got {}",
541                self.n_rows,
542                residual.len()
543            ));
544        }
545        for (row, &value) in residual.iter().enumerate() {
546            if !(value.is_finite() && value >= 0.0) {
547                return Err(format!(
548                    "RowMetric::with_truncation_mass_residual row {row} must be finite and non-negative; got {value}"
549                ));
550            }
551        }
552        self.truncation_mass_residual = Some(residual);
553        Ok(self)
554    }
555
556    /// Restrict the metric to the rows `rows` (an index subset or permutation),
557    /// preserving provenance, `p`, `rank`, and the solver floor. The
558    /// outer-criterion row subsample uses this to whiten the subsampled fit
559    /// through the SAME per-row metric the full-`N` fit uses, so the ρ search
560    /// ranks the delivered criterion (e.g. a #974 structured-whitening fit is not
561    /// silently searched unwhitened). Each gathered row's factor block is copied
562    /// verbatim, so the induced `M_n = U_n U_nᵀ` is bit-identical to the full
563    /// metric's on every selected row.
564    pub fn gather_rows(&self, rows: &[usize]) -> Result<Self, String> {
565        for (pos, &r) in rows.iter().enumerate() {
566            if r >= self.n_rows {
567                return Err(format!(
568                    "RowMetric::gather_rows: row index {r} at position {pos} is out of bounds \
569                     (n_rows = {})",
570                    self.n_rows
571                ));
572            }
573        }
574        match self.factors.as_ref() {
575            // Euclidean carries an implicit identity factor per row, so the subset
576            // is just a smaller identity stack — no factor storage to gather.
577            None => Self::euclidean(rows.len(), self.p),
578            Some(factors) => {
579                let cols = self.p * self.rank;
580                let mut sub = Array2::<f64>::zeros((rows.len(), cols));
581                for (pos, &r) in rows.iter().enumerate() {
582                    sub.row_mut(pos).assign(&factors.row(r));
583                }
584                // Re-runs the shared PSD normalizer on the subset (a subset of
585                // valid rows stays valid) and preserves the exact provenance and
586                // solver floor.
587                let mut metric = Self::from_factors(
588                    self.provenance,
589                    Arc::new(sub),
590                    self.p,
591                    self.rank,
592                    self.solver_delta,
593                )?;
594                metric.fisher_factor_kind = self.fisher_factor_kind;
595                match self.truncation_mass_residual.as_ref() {
596                    None => Ok(metric),
597                    Some(residual) => {
598                        let gathered =
599                            ndarray::Array1::from_iter(rows.iter().map(|&row| residual[row]));
600                        metric.with_truncation_mass_residual(Arc::new(gathered))
601                    }
602                }
603            }
604        }
605    }
606
607    /// The provenance tag (consumed by Object 4 to certify the inner product).
608    pub fn provenance(&self) -> MetricProvenance {
609        self.provenance
610    }
611
612    /// Explicit output-Fisher factor status, never inferred from diagnostics.
613    pub fn fisher_factor_kind(&self) -> Option<FisherFactorKind> {
614        self.fisher_factor_kind
615    }
616
617    /// Whether this metric is allowed to **whiten the likelihood** (i.e. replace
618    /// the isotropic reconstruction data-fit `½ rᵀr` with the whitened
619    /// `½ rᵀ M_n r`).
620    ///
621    /// This is TRUE for two provenances, for two distinct reasons:
622    ///
623    /// * [`MetricProvenance::WhitenedStructured`] — a genuinely *estimated noise
624    ///   model* (a factor-analytic residual covariance, #974), for which
625    ///   whitening the likelihood is the statistically correct thing to do; and
626    /// * [`MetricProvenance::BehavioralFisher`] — the **Rung 1** deliberate
627    ///   choice to price reconstruction error in nats: the output-Fisher metric
628    ///   `G_n` installed *as the loss weight* (`½ eᵀ G_n e`), a generalized
629    ///   least-squares reconstruction. Because `G_n` is a fixed per-row metric
630    ///   the problem stays linear-Gaussian, so REML/evidence/EDF are preserved.
631    ///
632    /// It is FALSE for [`MetricProvenance::Euclidean`] (nothing to whiten by) and
633    /// for the *gauge-only* [`MetricProvenance::OutputFisher`] /
634    /// [`MetricProvenance::OutputFisherDownstream`]: there the output-Fisher
635    /// inner product is an **output-geometry gauge**, and whitening the
636    /// likelihood by it *implicitly* (without the caller electing GLS) would
637    /// silently replace the reconstruction loss with a Fisher pullback — the #980
638    /// failure mode, and the reason the two-tier harvest can withhold factors
639    /// from a row without changing its loss. `BehavioralFisher` is the *explicit*
640    /// election of that same arithmetic as the intended objective.
641    pub fn whitens_likelihood(&self) -> bool {
642        matches!(
643            self.provenance,
644            MetricProvenance::WhitenedStructured { .. } | MetricProvenance::BehavioralFisher { .. }
645        )
646    }
647
648    /// Whether this metric **drives the gauge** — i.e. the isometry-penalty
649    /// pullback weight is taken from it rather than the identity.
650    ///
651    /// TRUE for any non-[`MetricProvenance::Euclidean`] provenance: both
652    /// [`MetricProvenance::OutputFisher`] and
653    /// [`MetricProvenance::WhitenedStructured`] supply a non-identity per-row
654    /// inner product the gauge pulls back through. Euclidean reduces the gauge
655    /// pullback to the bare `J_nᵀ J_n`, so it does not drive the gauge.
656    pub fn drives_gauge(&self) -> bool {
657        !matches!(self.provenance, MetricProvenance::Euclidean)
658    }
659
660    /// Number of rows the metric is defined over.
661    pub fn n_rows(&self) -> usize {
662        self.n_rows
663    }
664
665    /// Output dimensionality `p` (residual / Jacobian-column dimension).
666    pub fn p_out(&self) -> usize {
667        self.p
668    }
669
670    /// The factor rank: the dimension of the whitened residual
671    /// [`Self::whiten_residual_row`] returns (and the column count of the per-row
672    /// factor `U_n ∈ ℝ^{p × rank}`). For [`MetricProvenance::Euclidean`] this is
673    /// `p` (the implicit identity factor), so a consumer that sizes a whitened
674    /// buffer by `metric_rank()` gets the right length in every provenance.
675    pub fn metric_rank(&self) -> usize {
676        self.rank
677    }
678
679    /// Per-row traces `tr(M_n)` of the criterion-facing (un-floored) metric —
680    /// the Fisher-mass reduction the #980 row measure consumes. The dense
681    /// `(n_rows, p, p)` stack is validated streamingly at construction and
682    /// never retained; consumers wanting an explicit `W_n` rebuild it from
683    /// [`Self::metric_rank`]-sized factors.
684    pub fn row_traces(&self) -> ndarray::ArrayView1<'_, f64> {
685        self.traces.view()
686    }
687
688    /// Omitted Fisher trace for `row`, when the harvest supplied one.
689    pub fn truncation_mass_residual(&self, row: usize) -> Option<f64> {
690        self.truncation_mass_residual
691            .as_ref()
692            .map(|residual| residual[row])
693    }
694
695    /// Fraction of total audited Fisher trace omitted at `row`. `None` means the
696    /// factor stack carried no truncation audit. A zero-total metric has zero
697    /// omitted fraction when its reported residual is also zero.
698    pub fn truncation_mass_residual_fraction(&self, row: usize) -> Option<f64> {
699        self.truncation_mass_residual(row).map(|residual| {
700            let total = self.traces[row] + residual;
701            if total > 0.0 { residual / total } else { 0.0 }
702        })
703    }
704
705    /// Whiten a single `p`-dimensional residual row `r` into the coordinates
706    /// whose squared Euclidean norm equals `rᵀ W_n r`.
707    ///
708    /// * Euclidean: returns `r` unchanged (`‖r‖² = rᵀ I r`), so the likelihood
709    ///   reproduces the isotropic `½ rᵀr` data-fit bit-for-bit.
710    /// * Factored: returns `U_nᵀ r ∈ ℝ^{rank}`, with
711    ///   `‖U_nᵀ r‖² = rᵀ U_n U_nᵀ r = rᵀ W_n r`.
712    ///
713    /// This is the load-bearing identity that lets the data-fit loop sum
714    /// `0.5 * Σ whitened²` and recover exactly `rᵀ W_n r` whatever the
715    /// provenance.
716    pub fn whiten_residual_row(&self, row: usize, r: ArrayView1<'_, f64>) -> Vec<f64> {
717        match &self.factors {
718            None => r.iter().copied().collect(),
719            Some(u) => {
720                let mut out = vec![0.0_f64; self.rank];
721                for k in 0..self.rank {
722                    let mut acc = 0.0;
723                    for i in 0..self.p {
724                        acc += u[[row, i * self.rank + k]] * r[i];
725                    }
726                    out[k] = acc;
727                }
728                out
729            }
730        }
731    }
732
733    /// The factor entry `U_n[i, k]` for one row (`i ∈ [0, p)`, `k ∈ [0, rank)`).
734    /// For [`MetricProvenance::Euclidean`] the implicit factor is `I_p`, so this
735    /// returns `1.0` when `i == k` and `0.0` otherwise — letting a consumer that
736    /// whitens a Jacobian via `factor_entry` produce the identity whitening
737    /// without a provenance branch. Reads the **un-floored** factors (criterion
738    /// face, #747).
739    #[inline]
740    pub fn factor_entry(&self, row: usize, i: usize, k: usize) -> f64 {
741        match &self.factors {
742            None => {
743                if i == k {
744                    1.0
745                } else {
746                    0.0
747                }
748            }
749            Some(u) => u[[row, i * self.rank + k]],
750        }
751    }
752
753    /// Apply the full per-row metric `M_n x = U_n (U_nᵀ x) ∈ ℝ^p` for one
754    /// `p`-vector `x`, formed factored (`rank` flops in, `p` flops out) — never
755    /// materializing `M_n` as `p × p`. Euclidean returns `x` unchanged
756    /// (`M_n = I_p`). This is the p-space metric-applied vector the SAE β-tier
757    /// data-fit gradient contracts (β lives in p-output space, so its gradient
758    /// needs `M_n r_n`, not the rank-space whitened residual `U_nᵀ r_n`). Uses the
759    /// **un-floored** factors (criterion face, `δ`-free, #747 invariant).
760    pub fn apply_metric_row(&self, row: usize, x: ArrayView1<'_, f64>) -> Vec<f64> {
761        match &self.factors {
762            None => x.iter().copied().collect(),
763            Some(u) => {
764                // w = U_nᵀ x ∈ ℝ^{rank}.
765                let mut w = vec![0.0_f64; self.rank];
766                for k in 0..self.rank {
767                    let mut acc = 0.0;
768                    for i in 0..self.p {
769                        acc += u[[row, i * self.rank + k]] * x[i];
770                    }
771                    w[k] = acc;
772                }
773                // out = U_n w ∈ ℝ^p.
774                let mut out = vec![0.0_f64; self.p];
775                for i in 0..self.p {
776                    let mut acc = 0.0;
777                    for k in 0..self.rank {
778                        acc += u[[row, i * self.rank + k]] * w[k];
779                    }
780                    out[i] = acc;
781                }
782                out
783            }
784        }
785    }
786
787    /// Pullback metric `g_n = J_nᵀ W_n J_n` for one row, formed as
788    /// `(J_nᵀ U_n)(U_nᵀ J_n)` — never materializing the `p × p` `W_n`.
789    ///
790    /// `j_row` is the row's Jacobian `J_n ∈ ℝ^{p × d}` flattened row-major
791    /// (`J_n[i, a] = j_row[i * d + a]`). Returns the `d × d` `g_n`.
792    pub fn pullback(&self, row: usize, j_row: &[f64], d: usize) -> Array2<f64> {
793        match &self.factors {
794            None => {
795                // W_n = I_p ⇒ g_n = J_nᵀ J_n.
796                let mut g = Array2::<f64>::zeros((d, d));
797                for a in 0..d {
798                    for b in a..d {
799                        let mut acc = 0.0;
800                        for i in 0..self.p {
801                            acc += j_row[i * d + a] * j_row[i * d + b];
802                        }
803                        g[[a, b]] = acc;
804                        g[[b, a]] = acc;
805                    }
806                }
807                g
808            }
809            Some(u) => {
810                // M_n = U_nᵀ J_n ∈ ℝ^{rank × d}; g_n = M_nᵀ M_n.
811                let mut m = Array2::<f64>::zeros((self.rank, d));
812                for k in 0..self.rank {
813                    for a in 0..d {
814                        let mut acc = 0.0;
815                        for i in 0..self.p {
816                            acc += u[[row, i * self.rank + k]] * j_row[i * d + a];
817                        }
818                        m[[k, a]] = acc;
819                    }
820                }
821                let mut g = Array2::<f64>::zeros((d, d));
822                for a in 0..d {
823                    for b in a..d {
824                        let mut acc = 0.0;
825                        for k in 0..self.rank {
826                            acc += m[[k, a]] * m[[k, b]];
827                        }
828                        g[[a, b]] = acc;
829                        g[[b, a]] = acc;
830                    }
831                }
832                g
833            }
834        }
835    }
836
837    /// Quadratic form `r_nᵀ M_n r_n` for one row's residual `r_n ∈ ℝ^p`, formed
838    /// **factored** as `‖U_nᵀ r_n‖²` — never materializing the `p × p` `M_n`.
839    ///
840    /// This is the criterion-facing squared residual the likelihood sums; it uses
841    /// the **un-floored** `U_n U_nᵀ`, so the solver `δ` does not enter it
842    /// (#747 invariant). Euclidean provenance returns the bit-identical `‖r_n‖²`.
843    #[inline]
844    pub fn quad_form(&self, row: usize, r: ArrayView1<'_, f64>) -> f64 {
845        match &self.factors {
846            None => r.iter().map(|&v| v * v).sum(),
847            Some(_) => self
848                .whiten_residual_row(row, r)
849                .iter()
850                .map(|&w| w * w)
851                .sum(),
852        }
853    }
854
855    /// Fisher mass of a per-row output vector `x_n ∈ ℝ^p`: the scalar
856    /// `x_nᵀ M_n x_n` (alias of [`Self::quad_form`] read as an information mass
857    /// rather than a residual square). Factored, never `p × p`, `δ`-free.
858    #[inline]
859    pub fn fisher_mass(&self, row: usize, x: ArrayView1<'_, f64>) -> f64 {
860        self.quad_form(row, x)
861    }
862
863    /// The gauge view of this metric: the
864    /// [`crate::WeightField`] the isometry penalty pulls back through.
865    ///
866    /// This is the **single** way an `IsometryPenalty` acquires a non-identity
867    /// gauge metric — the independent `WeightField` setter has been removed — so
868    /// the gauge metric is, by construction, the same object the likelihood
869    /// whitens with.
870    pub fn to_weight_field(&self) -> crate::WeightField {
871        use crate::WeightField;
872        match &self.factors {
873            None => WeightField::Identity,
874            Some(u) => WeightField::Factored {
875                u: Arc::clone(u),
876                rank: self.rank,
877                p_out: self.p,
878            },
879        }
880    }
881}
882
883/// Pack a harvest-emitted probe stack into the row-major factor layout
884/// [`RowMetric::behavioral_fisher`] expects.
885///
886/// The harvest boundary (the model-interaction side) emits, per token, `s`
887/// probe vectors `vₖ = J_nᵀ F_n^{1/2} uₖ ∈ ℝ^p` — the natural shape is
888/// `probes[n, i, k] = (vₖ)ᵢ`, an `(n_rows, p, probes)` stack. This assembles the
889/// `(n_rows, p · probes)` row-major matrix `u[n, i·probes + k] = probes[n, i, k]`
890/// that the constructor consumes so that column `k` of the per-row factor `U_n`
891/// is exactly probe `vₖ` and `M_n = U_n U_nᵀ = Σₖ vₖ vₖᵀ ≈ G_n`.
892///
893/// This is a pure repack of the standard C-order flattening; it exists so the
894/// harvest → metric seam is a single named, validated Rust surface rather than
895/// an ad-hoc reshape at each call site. Errors on non-finite entries so the
896/// failure is caught here rather than deep in [`normalize_fisher_rao_blocks`].
897pub fn pack_probe_factors(probes: ndarray::ArrayView3<'_, f64>) -> Result<Array2<f64>, String> {
898    let (n_rows, p, s) = probes.dim();
899    if s == 0 {
900        return Err("pack_probe_factors: need at least one probe (s == 0)".to_string());
901    }
902    if !probes.iter().all(|v| v.is_finite()) {
903        return Err("pack_probe_factors: probe entries must be finite".to_string());
904    }
905    let mut u = Array2::<f64>::zeros((n_rows, p * s));
906    for n in 0..n_rows {
907        for i in 0..p {
908            for k in 0..s {
909                u[[n, i * s + k]] = probes[[n, i, k]];
910            }
911        }
912    }
913    Ok(u)
914}
915
916#[cfg(test)]
917mod tests {
918    use super::*;
919    use ndarray::array;
920
921    // ── RowMetric::euclidean ──────────────────────────────────────────────────
922
923    #[test]
924    fn euclidean_metric_has_correct_dimensions() {
925        let m = RowMetric::euclidean(5, 3).unwrap();
926        assert_eq!(m.n_rows(), 5);
927        assert_eq!(m.p_out(), 3);
928        assert_eq!(m.metric_rank(), 3);
929    }
930
931    #[test]
932    fn euclidean_metric_traces_equal_p() {
933        let p = 4_usize;
934        let m = RowMetric::euclidean(3, p).unwrap();
935        for tr in m.row_traces().iter() {
936            assert!((*tr - p as f64).abs() < 1e-14, "trace {tr} != p={p}");
937        }
938    }
939
940    #[test]
941    fn euclidean_provenance_is_euclidean() {
942        let m = RowMetric::euclidean(1, 2).unwrap();
943        assert_eq!(m.provenance(), MetricProvenance::Euclidean);
944    }
945
946    #[test]
947    fn euclidean_does_not_whiten_likelihood() {
948        let m = RowMetric::euclidean(1, 2).unwrap();
949        assert!(!m.whitens_likelihood());
950    }
951
952    #[test]
953    fn euclidean_does_not_drive_gauge() {
954        let m = RowMetric::euclidean(1, 2).unwrap();
955        assert!(!m.drives_gauge());
956    }
957
958    #[test]
959    fn euclidean_to_weight_field_is_identity() {
960        let m = RowMetric::euclidean(1, 2).unwrap();
961        assert!(matches!(m.to_weight_field(), WeightField::Identity));
962    }
963
964    #[test]
965    fn euclidean_whiten_residual_is_passthrough() {
966        let m = RowMetric::euclidean(1, 3).unwrap();
967        let r = array![1.0_f64, 2.0, 3.0];
968        let w = m.whiten_residual_row(0, r.view());
969        assert_eq!(w, vec![1.0, 2.0, 3.0]);
970    }
971
972    #[test]
973    fn euclidean_factor_entry_is_identity() {
974        let m = RowMetric::euclidean(1, 3).unwrap();
975        assert_eq!(m.factor_entry(0, 0, 0), 1.0);
976        assert_eq!(m.factor_entry(0, 1, 1), 1.0);
977        assert_eq!(m.factor_entry(0, 2, 2), 1.0);
978        assert_eq!(m.factor_entry(0, 0, 1), 0.0);
979        assert_eq!(m.factor_entry(0, 1, 0), 0.0);
980    }
981
982    #[test]
983    fn euclidean_quad_form_is_squared_norm() {
984        let m = RowMetric::euclidean(1, 3).unwrap();
985        let r = array![1.0_f64, 2.0, 2.0];
986        assert!((m.quad_form(0, r.view()) - 9.0).abs() < 1e-14);
987    }
988
989    // ── MetricProvenance predicates ───────────────────────────────────────────
990
991    #[test]
992    fn behavioral_fisher_whitens_likelihood_and_drives_gauge() {
993        // The Rung-1 deliberate GLS metric: unlike the gauge-only OutputFisher,
994        // it whitens the reconstruction likelihood.
995        let u = Arc::new(array![[1.0_f64, 0.5]]); // p=1, probes=2
996        let m = RowMetric::behavioral_fisher(u, 1, 2).unwrap();
997        assert!(m.whitens_likelihood());
998        assert!(m.drives_gauge());
999        assert_eq!(
1000            m.provenance(),
1001            MetricProvenance::BehavioralFisher { probes: 2 }
1002        );
1003        assert_eq!(m.metric_rank(), 2);
1004    }
1005
1006    #[test]
1007    fn behavioral_fisher_quad_form_is_probe_sum() {
1008        // p=2, s=2 probes v1=(1,0), v2=(0,2) → G = diag(1,4);
1009        // e=(3,1) → eᵀGe = 9·1 + 1·4 = 13 = Σ (vᵢᵀe)² = 3² + 2² = 13.
1010        // Column-major-within-row layout U[i,k]=u[i*probes+k]:
1011        //   U[0,0]=1 U[0,1]=0  U[1,0]=0 U[1,1]=2
1012        let u = Arc::new(array![[1.0_f64, 0.0, 0.0, 2.0]]);
1013        let m = RowMetric::behavioral_fisher(u, 2, 2).unwrap();
1014        let e = array![3.0_f64, 1.0];
1015        assert!((m.quad_form(0, e.view()) - 13.0).abs() < 1e-12);
1016    }
1017
1018    #[test]
1019    fn behavioral_fisher_g_identity_reproduces_euclidean_quad_form() {
1020        // GLS with G=I must reduce to plain MSE. Identity probes (s=p, U=I_p)
1021        // ⇒ M_n = I ⇒ quad_form == ‖e‖², matching Euclidean bit-for-bit, and
1022        // metric_rank == p so the whitened residual-dof accounting is unchanged.
1023        let p = 3;
1024        let mut u = Array2::<f64>::zeros((1, p * p));
1025        for i in 0..p {
1026            u[[0, i * p + i]] = 1.0;
1027        }
1028        let bf = RowMetric::behavioral_fisher(Arc::new(u), p, p).unwrap();
1029        let euc = RowMetric::euclidean(1, p).unwrap();
1030        let e = array![1.5_f64, -2.0, 0.25];
1031        assert_eq!(bf.metric_rank(), euc.metric_rank());
1032        assert!((bf.quad_form(0, e.view()) - euc.quad_form(0, e.view())).abs() < 1e-14);
1033        // and whitened residual is the residual itself (identity whitening)
1034        assert_eq!(bf.whiten_residual_row(0, e.view()), vec![1.5, -2.0, 0.25]);
1035    }
1036
1037    #[test]
1038    fn pack_probe_factors_matches_manual_layout() {
1039        use ndarray::Array3;
1040        // n=1, p=2, s=2: probes[0,i,k] = v_k[i]; v0=(1,3), v1=(2,4)
1041        let mut probes = Array3::<f64>::zeros((1, 2, 2));
1042        probes[[0, 0, 0]] = 1.0; // v0[0]
1043        probes[[0, 1, 0]] = 3.0; // v0[1]
1044        probes[[0, 0, 1]] = 2.0; // v1[0]
1045        probes[[0, 1, 1]] = 4.0; // v1[1]
1046        let u = pack_probe_factors(probes.view()).unwrap();
1047        // Layout U[i,k] = u[i*s + k]: [v0[0],v1[0], v0[1],v1[1]] = [1,2,3,4]
1048        assert_eq!(u.as_slice().unwrap(), &[1.0, 2.0, 3.0, 4.0]);
1049        // Round-trips into a valid metric whose G = v0 v0ᵀ + v1 v1ᵀ.
1050        let m = RowMetric::behavioral_fisher(Arc::new(u), 2, 2).unwrap();
1051        // e=(1,0): eᵀGe = v0[0]²+v1[0]² = 1+4 = 5.
1052        let e = array![1.0_f64, 0.0];
1053        assert!((m.quad_form(0, e.view()) - 5.0).abs() < 1e-12);
1054    }
1055
1056    #[test]
1057    fn pack_probe_factors_rejects_zero_probes() {
1058        use ndarray::Array3;
1059        let probes = Array3::<f64>::zeros((2, 3, 0));
1060        assert!(pack_probe_factors(probes.view()).is_err());
1061    }
1062
1063    #[test]
1064    fn fisher_factor_status_is_never_inferred_from_zero_residual_2249() {
1065        let factors = Arc::new(Array2::from_elem((1, 1), 2.0));
1066        let metric = RowMetric::output_fisher(factors, 1, 1)
1067            .unwrap()
1068            .with_truncation_mass_residual(Arc::new(array![0.0]))
1069            .unwrap();
1070        assert_eq!(
1071            metric.fisher_factor_kind(),
1072            Some(FisherFactorKind::UncertifiedApproximation)
1073        );
1074        let certified = metric
1075            .clone()
1076            .with_fisher_factor_kind(FisherFactorKind::CertifiedPsdLowerBound)
1077            .unwrap();
1078        assert_eq!(
1079            certified.fisher_factor_kind(),
1080            Some(FisherFactorKind::CertifiedPsdLowerBound)
1081        );
1082        assert!(
1083            metric
1084                .with_fisher_factor_kind(FisherFactorKind::ExactFull)
1085                .is_err(),
1086            "an omitted-trace record is incompatible with an exact-full claim"
1087        );
1088    }
1089
1090    // ── WeightField::project_jac_row_with_u ──────────────────────────────────
1091
1092    #[test]
1093    fn project_jac_with_identity_returns_jac() {
1094        // p=2, rank=2, d=2; U=I_2, J=[[1,2],[3,4]] → M = U^T J = J
1095        let u_row = [1.0_f64, 0.0, 0.0, 1.0]; // U[i,k]=u[i*rank+k], I_2
1096        let j_row = [1.0_f64, 2.0, 3.0, 4.0]; // J[i,a]=j[i*d+a]
1097        let m = WeightField::project_jac_row_with_u(&u_row, &j_row, 2, 2, 2);
1098        assert!((m[[0, 0]] - 1.0).abs() < 1e-14);
1099        assert!((m[[0, 1]] - 2.0).abs() < 1e-14);
1100        assert!((m[[1, 0]] - 3.0).abs() < 1e-14);
1101        assert!((m[[1, 1]] - 4.0).abs() < 1e-14);
1102    }
1103
1104    #[test]
1105    fn project_jac_with_zeros_returns_zero_matrix() {
1106        let u_row = [0.0_f64, 0.0];
1107        let j_row = [1.0_f64, 2.0];
1108        let m = WeightField::project_jac_row_with_u(&u_row, &j_row, 2, 1, 1);
1109        assert_eq!(m[[0, 0]], 0.0);
1110    }
1111}