Skip to main content

gam_terms/basis/
measure_jet_smooth.rs

1//! Measure-jet spline smooth: multiscale local-jet-residual energy of the
2//! empirical measure (center-quadratured current implementation).
3//!
4//! The term penalizes, at every quadrature point and every scale, the failure
5//! of `f` to be locally affine *in the measure*:
6//!
7//! ```text
8//!   Q = Σ_ℓ  w_ℓ · Σ_i  mass_i · q_i(ε_ℓ)^(1−2α) · R_{i,ℓ},
9//!   w_ℓ = log_step · ε_ℓ^(−η),   η = 2s + d(2−2α),
10//! ```
11//!
12//! where `R_{i,ℓ}` is the residual quadratic form of the exact weighted
13//! local affine projection at center `i` and scale `ε_ℓ`: kernel weights
14//! `w_j = mass_j · exp(−d_ij²/(2ε_ℓ²))`, kernel mass `q_i = Σ_j w_j`, and the
15//! fit `min_b ‖Cv − Φ̃b‖²_W` over weighted-centered values
16//! `Cv = v − (uᵀv)·1` (`u = w/q`) and weighted-centered scaled features
17//! `Φ̃` (rows `(c_j − c_i)/ε`, column means removed under `u`). Rank-deficient
18//! cells use the machine-precision pseudo-inverse of `Φ̃ᵀWΦ̃/q`, so ambient
19//! affine values are projected away exactly instead of paying a ridge toll.
20//!
21//! # Contracts (each is load-bearing; tests pin them)
22//!
23//! - **Exact constant annihilation.** The constant is removed by the weighted
24//!   mean projection `C`, never ridged: `Q·1 = 0` to machine precision at
25//!   every scale, so the penalty carries NO mass term and the fit has no
26//!   prior mean to revert to. This is the no-mean-reversion contract of the
27//!   measure-jet design; ridging the constant would silently reintroduce
28//!   mean reversion.
29//! - **Exact affine projection / rank adaptation.** The slope block uses the
30//!   rank-revealing pseudo-inverse of the dimensionless local Gram
31//!   `G = Φ̃ᵀWΦ̃/q`, not a Tikhonov ridge. On a 1-D filament in ambient
32//!   dimension d the resolved tangent slope is absorbed (not penalized);
33//!   unresolved directions have no variation after weighted centering and do
34//!   not create an affine toll. The retained rank is a numerical property of
35//!   the weighted cell, not a smoothing dial.
36//! - **Mellin band.** Scales form a geometric grid from the center-spacing
37//!   floor to the half-diameter; `w_ℓ = log_step · ε_ℓ^(−η)`, with
38//!   `η = 2s + d(2−2α)`, is the fixed-order quadrature weight used by this
39//!   implementation. It keeps the advertised continuous smoothness order
40//!   `s ∈ (0, 2)` from silently changing when `α` changes.
41//! - **Density normalization.** The outer quadrature weight
42//!   `mass_i · q_i^(1−2α)` realizes `dμ(x)/q_ε(x)^(2α−1)`. On a p-dimensional
43//!   stratum with sampling density `ρ`, `q_ε ~ Cρ ε^p` and the local residual
44//!   contributes an extra `ε^{p(2−2α)}` factor. The fixed-order scale weight
45//!   cancels that factor using the available dimension parameter; without that
46//!   correction, the symbol exponent would be `2s + 2p(α−1)`.
47//! - **Frozen-quadrature replay.** The penalty and extrapolation diagnostic
48//!   depend on the FIT data through center masses, the realized band, on-web
49//!   support anchors, and penalty normalization scales. The freeze step
50//!   persists all of them ([`MeasureJetFrozenQuadrature`]) so predict-time
51//!   rebuilds replay the exact fit-time penalty instead of recomputing it from
52//!   predict rows.
53//! - **Single assembly source.** Every quadratic form this module emits —
54//!   the energy, its (s, α) jets, the per-scale spectrum — is produced by
55//!   ONE workhorse ([`assemble_weighted_forms`]) that walks the local
56//!   residual blocks exactly once per request and differs only in the
57//!   scalar weights applied per block. Criterion value and criterion
58//!   derivatives cannot drift apart (the objective↔gradient desync class is
59//!   structurally excluded).
60//! - **single-scale/multiscale opt-in (#1039/#1116).** The per-scale spectrum
61//!   and the `(α, ln τ)` ψ dials are the multiscale-mode realization, engaged
62//!   ONLY when the spec opts in (`MeasureJetBasisSpec::multiscale = true`, the
63//!   DSL `mjs(…, multiscale=true)`); see [`measure_jet_multiscale_mode`]. There
64//!   is NO center-count auto-gate: at ANY center count the default is
65//!   single-scale — one jet-energy Primary at the auto order with no energy ψ
66//!   dials. The independent function-space null-component candidate requested
67//!   by `double_penalty` is present in either mode and has its own REML λ. The
68//!   flag is persisted on the spec, so freeze→replay re-enters the same mode
69//!   verbatim.
70//!
71//! # ψ-differentiability contract (what the ψ-channel stage consumes)
72//!
73//! Mirroring the constant-curvature κ-contract (#944): centers, masses, the
74//! band are deliberately hyperparameter-FIXED at build time; the representer
75//! range ℓ is the ONE opt-in design-moving dial (#1116). Consequences:
76//!
77//! - **Penalty-dial design drift is identically zero**: the (s, α, τ) dials
78//!   reweight only the jet-energy penalty, never the Gaussian representer
79//!   design (`∂X/∂{s,α,τ} ≡ 0`), so those channels are penalty-only
80//!   (`is_penalty_like` auto-derives true in the outer engine's
81//!   `DirectionalHyperParam`).
82//! - **The representer range ℓ is a design-and-pullback-moving dial** (matérn's
83//!   `log_kappa` analog, #1116): `X = K(data, centers; ℓ)·z` and the center
84//!   evaluation map `E = K(centers, centers; ℓ)·z` both depend on ℓ. The
85//!   center-value forms `Q` and `H₀` are ℓ-invariant, but their coefficient
86//!   pullbacks `EᵀQE` and `EᵀH₀E` are not; exact product-rule jets are shipped
87//!   alongside the design jets. When explicitly enabled, ℓ rebuilds the design
88//!   per outer trial; it does not change the frozen basis rank. FD-gated by
89//!   `psi_producer_matches_fd_length_scale`.
90//! - **Exact (s, α) penalty jets are shipped**:
91//!   [`measure_jet_energy_form_with_jets`] returns `∂Q/∂s`, `∂²Q/∂s²`,
92//!   `∂Q/∂α`, `∂²Q/∂α²`, `∂²Q/∂s∂α` in closed form — both dials enter only
93//!   through the per-block log-weights (`∂ln w/∂s = −2 ln ε`,
94//!   `∂ln w/∂α = −2 ln q`), so the jets are reweighted re-scatters of the
95//!   SAME residual blocks, FD-gated in this module's tests.
96//!   The retained τ coordinate is inert under the exact projection, so its
97//!   derivative slots are identically zero.
98//!
99//! # Cost shape (and the upgrade ladder above it)
100//!
101//! The outer sum is coarsened per scale to a deterministic ε/2-net (the
102//! outer Riemann sum needs resolution ε, not the center-spacing floor), so
103//! the band totals ~O(m²·d) instead of O(L·m³) — the current realization of
104//! the pyramid principle that each scale interacts at its own level. This is
105//! mass-lumped quadrature of the displayed outer integral; it is first-
106//! moment exact for the cell locations and carries the usual
107//! `O(diam²/ε²)` relative scale for smooth Gaussian-weighted functionals,
108//! not an estimand-preserving identity.
109//! The long-form home for the ladder and the substrate contracts is the
110//! frame notes (`docs/measure_jet_frame.md`); its §2 moment substrate is
111//! `measure_jet_moments.rs`, its §5 extrapolation pricing
112//! `measure_jet_predict.rs`.
113
114use ndarray::{Array1, Array2, ArrayView1, ArrayView2, Axis};
115use rayon::prelude::*;
116use serde::{Deserialize, Serialize};
117
118use faer::Side;
119
120use gam_linalg::faer_ndarray::{FaerEigh, default_rrqr_rank_alpha, rrqr_nullspace_basis};
121
122use super::{
123    AnisoBasisPsiDerivatives, AnisoPenaltyCrossProvider, BasisBuildResult, BasisError,
124    BasisMetadata, CenterStrategy, PenaltyCandidate, PenaltySource,
125    filter_active_penalty_candidates_with_ops, normalize_penalty,
126    normalize_penalty_cross_psi_derivative, normalize_penaltywith_psi_derivatives,
127    select_centers_by_strategy, trace_of_product,
128};
129
130/// Truncation radius of the Gaussian profile in units of the scale ε: weights
131/// beyond `3ε` are below `e^{-4.5} ≈ 1.1e-2` of the peak and are dropped from
132/// both the local fit and the `q^(1−2α)` outer weight. This is an absolute
133/// kernel-weight cutoff; using the same truncated q keeps the discrete
134/// functional self-consistent, but it is not a relative tail-error bound.
135pub(crate) const MEASURE_JET_PROFILE_CUTOFF: f64 = 3.0;
136
137/// Relative eigenvalue threshold for rank-revealing pseudo-inverses of local
138/// Gram matrices. Directions at the roundoff floor are treated as unresolved
139/// and excluded from the affine fit.
140pub(crate) const MEASURE_JET_PSEUDOINVERSE_RTOL: f64 = 64.0 * f64::EPSILON;
141
142/// Default continuous smoothness order `s` realized by the `0.0` auto
143/// sentinel. Sits mid-band in the admissible `(0, 2)` for the affine-jet
144/// (r = 2) energy: rough enough to stay pointwise-defined on filaments and
145/// sheets (`s > p/2` for intrinsic `p ≤ 2`), smooth enough to bridge gaps
146/// with attested trends.
147pub(crate) const MEASURE_JET_DEFAULT_ORDER_S: f64 = 1.5;
148
149/// Auto-band scale-count clamp: at least 3 octave-ish nodes so the energy is
150/// genuinely multiscale, at most 8 so degenerate spacing cannot explode the
151/// build.
152pub(crate) const MEASURE_JET_MIN_AUTO_SCALES: usize = 3;
153pub(crate) const MEASURE_JET_MAX_AUTO_SCALES: usize = 8;
154
155/// Representer-range multiple of the median nearest-center spacing used by
156/// the `0.0` auto sentinel.
157///
158/// Set to ×1: a Gaussian representer of range `ℓ = h` (the median
159/// nearest-center spacing) already overlaps its neighbors at
160/// `exp(−h²/(2ℓ²)) = exp(−1/2) ≈ 0.61`, so adjacent bumps blend smoothly while
161/// each center keeps a *distinct* response. The old ×2 made every column
162/// `exp(−1/8) ≈ 0.88` at its neighbor — the representers became nearly
163/// collinear, which (a) over-smoothed the fitted surface (the #1116/#1041
164/// accuracy deficit: measure-jet sat ~1.6× the matern/duchon truth-RMSE) and
165/// (b) drove the design Gram toward rank deficiency, so the inner PIRLS /
166/// outer REML conditioning degraded and the smoothing-parameter search cycled
167/// for hundreds of seconds (the #1116 timeout). One spacing-width kernel fixes
168/// both at the root without touching the energy penalty or the dials.
169pub(crate) const MEASURE_JET_AUTO_LENGTH_SCALE_FACTOR: f64 = 1.0;
170
171/// Memory budget (in f64 entries) above which the multi-form assembly stops
172/// parallelizing over scales: parallel scale partials cost
173/// `L · n_forms · m²` doubles; past this budget the scales run sequentially
174/// (same numbers — the per-scale loop and the ordered sum are deterministic
175/// either way).
176pub(crate) const MEASURE_JET_PARALLEL_FORM_BUDGET_DOUBLES: usize = 1 << 26;
177
178/// Realized-design identifiability policy for the measure-jet smooth.
179/// Mirrors [`super::ConstantCurvatureIdentifiability`] (#532): the fit-time
180/// section gets the parametric orthogonalization composed onto it by the global
181/// identifiability pipeline, and the composed transform is frozen so
182/// predict-time (and per-ψ-trial) rebuilds replay it verbatim.
183#[derive(Debug, Clone, Serialize, Deserialize, Default)]
184pub enum MeasureJetIdentifiability {
185    /// Fit-time default. With the single-scale affine head, the RBF center
186    /// values are mass-orthogonalized against the affine value space and the
187    /// head passes through exactly; without a head, the representer coefficient
188    /// sum-to-zero section is used. Global parametric residualization follows.
189    #[default]
190    CenterSumToZero,
191    /// Predict-time replay: the frozen composed transform captured at fit
192    /// time. `transform.nrows()` equals representer count plus affine-head width.
193    FrozenTransform { transform: Array2<f64> },
194}
195
196/// Fit-time quadrature of the empirical measure (center masses + realized
197/// scale band), frozen onto the spec so predict-time rebuilds replay the
198/// exact fit-time penalty. Recomputing either from predict rows would
199/// silently change the penalty the coefficients were estimated under.
200#[derive(Debug, Clone, Serialize, Deserialize)]
201pub struct MeasureJetFrozenQuadrature {
202    /// Per-center masses `m_i` (nearest-center fractions of the FIT rows).
203    pub masses: Array1<f64>,
204    /// Realized geometric scale band `ε_0 < … < ε_{L−1}`.
205    pub eps_band: Vec<f64>,
206    /// Per-scale on-web support anchor
207    /// `q̄_ℓ = (Σ_i m_i q_ℓ(c_i)) / (Σ_i m_i)`.
208    pub support_means: Vec<f64>,
209    /// Frobenius scales of the emitted per-level normalized penalties. Empty in
210    /// fused mode, where the band emits one primary penalty instead.
211    pub penalty_normalization_scales: Vec<f64>,
212    /// Frobenius scales of the raw per-level forms before the arbitrary Mellin
213    /// `log_step · ε_ℓ^(-2s0)` gauge is folded in.
214    pub raw_penalty_normalization_scales: Vec<f64>,
215    /// Frobenius scale of the single fused primary penalty. `None` in per-level
216    /// mode.
217    pub fused_penalty_normalization_scale: Option<f64>,
218    /// Ambient input-measurement-error scale `σ_coord` (issue #2225): the
219    /// perpendicular off-manifold residual spread of the fit-time empirical
220    /// measure, in the frozen centers' (standardized) coordinate frame. Frozen
221    /// so the predict-time errors-in-variables variance term
222    /// `Var_input = σ_coord²·‖∇f̂‖²` uses the same input-noise scale the fit
223    /// saw. `None` when it could not be estimated (no cell spanned a tangent),
224    /// leaving `Var_input` disabled. Defaults to `None` for models persisted
225    /// before the term existed.
226    #[serde(default)]
227    pub sigma_coord: Option<f64>,
228}
229
230/// Serde default for [`MeasureJetBasisSpec::learn_length_scale`]: freeze ℓ at
231/// the realized auto/user value unless a fit explicitly opts into the
232/// design-moving outer coordinate.
233fn measure_jet_learn_length_scale_default() -> bool {
234    false
235}
236
237/// Measure-jet smooth configuration (`mjs(x0, …, xd)`).
238///
239/// The feature columns are ambient coordinates of data concentrated near an
240/// unknown low-dimensional (possibly stratified) set; the term learns the
241/// geometry from the empirical measure itself — centers as quadrature nodes,
242/// masses as μ-weights, local jet residuals as the roughness carrier — with
243/// no graph, mesh, or neighbor-set inside the statistical object.
244#[derive(Debug, Clone, Serialize, Deserialize)]
245pub struct MeasureJetBasisSpec {
246    /// Center/knot selection strategy (deterministic; quadrature of μ).
247    pub center_strategy: CenterStrategy,
248    /// Continuous smoothness order `s ∈ (0, 2)`; `0.0` sentinel = auto
249    /// ([`MEASURE_JET_DEFAULT_ORDER_S`]).
250    pub order_s: f64,
251    /// Density-normalization exponent α (outer weight `q^{1−2α}`).
252    pub alpha: f64,
253    /// Historical τ coordinate retained for frozen specs and ψ layout. The
254    /// measure-jet energy itself uses the exact weighted affine projection and
255    /// is independent of τ; the τ ψ derivatives are therefore zero.
256    pub tau0: f64,
257    /// Number of scale nodes; `0` sentinel = auto dyadic band.
258    pub num_scales: usize,
259    /// Representer (Gaussian RBF) range ℓ; `0.0` sentinel = auto
260    /// (median nearest-center spacing × [`MEASURE_JET_AUTO_LENGTH_SCALE_FACTOR`]).
261    pub length_scale: f64,
262    /// Add a separate function-space affine/null-component penalty alongside
263    /// the jet-energy penalty. Its strength is independently REML-selected.
264    pub double_penalty: bool,
265    /// REML-learn the representer range ℓ as a design-moving outer dial
266    /// (opt-in), mirroring Matérn's `log_kappa`. The Gaussian kernel is
267    /// strictly PD for every ℓ > 0, so ℓ does NOT change the basis rank (always
268    /// `m` centers) — but it changes WHICH `m`-dim subspace the representers
269    /// span, i.e. the span alignment with the true surface. The stable default
270    /// freezes ℓ at the auto/user value; `true` enrolls the outer coordinate for
271    /// experiments that need REML-selected representer range.
272    #[serde(default = "measure_jet_learn_length_scale_default")]
273    pub learn_length_scale: bool,
274    /// Explicit opt-in for multiscale mode: the per-scale spectral penalty
275    /// split plus the `(α, ln τ)` outer ψ dials. `false` (default) keeps the
276    /// energy in single-scale mode at ANY center count. The separate
277    /// `double_penalty` null component is available in both modes. There is no
278    /// center-count auto-gate; the user opts in via
279    /// `mjs(…, multiscale=true)`. Persisted on the spec so freeze→replay enters
280    /// the same mode.
281    #[serde(default)]
282    pub multiscale: bool,
283    /// Realized-design identifiability policy (see type docs).
284    #[serde(default)]
285    pub identifiability: MeasureJetIdentifiability,
286    /// Fit-time quadrature replay (see type docs). `None` at fit time;
287    /// `Some` on the frozen predict/ψ-trial path.
288    #[serde(default)]
289    pub frozen_quadrature: Option<MeasureJetFrozenQuadrature>,
290}
291
292impl Default for MeasureJetBasisSpec {
293    fn default() -> Self {
294        Self {
295            center_strategy: CenterStrategy::FarthestPoint { num_centers: 50 },
296            order_s: 0.0,
297            // Density-WEIGHTED Hessian energy (the module-header default): the
298            // outer weight is q^{1−2α} = q^{−1} at α = 1. The density-free
299            // variant α = 3/2 gives q^{−2}, which on a low-intrinsic-dimension
300            // stratum (data on a 1-D/2-D manifold embedded in higher ambient d)
301            // makes the local kernel mass q tiny AND spatially varying along
302            // the manifold, so q^{−2} amplifies the penalty unevenly and
303            // over-smooths the high-frequency signal there (MEASURED #1116: on
304            // the 1-D-curve-in-3-D fixture α = 3/2 left mjs ~13× worse than
305            // matérn). α = 1's q^{−1} weighting is far gentler and is the
306            // header-derived default; an explicit `alpha=` still overrides for
307            // genuinely density-free use on a full-dimensional stratum.
308            alpha: 1.0,
309            tau0: 1e-3,
310            num_scales: 0,
311            length_scale: 0.0,
312            double_penalty: true,
313            learn_length_scale: false,
314            multiscale: false,
315            identifiability: MeasureJetIdentifiability::CenterSumToZero,
316            frozen_quadrature: None,
317        }
318    }
319}
320
321/// Realized geometric scale band: `eps` ascending, `log_step` the constant
322/// log-spacing `ln(eps[ℓ+1]/eps[ℓ])` used as the Mellin quadrature weight.
323pub struct MeasureJetBand {
324    pub eps: Vec<f64>,
325    pub log_step: f64,
326}
327
328/// The energy and its exact hyperparameter jets in the live dials. `s` and
329/// `α` enter only through per-block log-weights. The retained `ln τ` slots
330/// are zero because the local fit is the exact weighted affine projection
331/// and no longer depends on τ. All forms are scattered from the SAME local
332/// residual blocks, and the ψ-channel consumes them with zero design drift.
333pub struct MeasureJetEnergyJets {
334    pub q: Array2<f64>,
335    pub dq_ds: Array2<f64>,
336    pub d2q_ds2: Array2<f64>,
337    pub dq_dalpha: Array2<f64>,
338    pub d2q_dalpha2: Array2<f64>,
339    pub d2q_ds_dalpha: Array2<f64>,
340    pub dq_dlogtau: Array2<f64>,
341    pub d2q_dlogtau2: Array2<f64>,
342    pub d2q_ds_dlogtau: Array2<f64>,
343    pub d2q_dalpha_dlogtau: Array2<f64>,
344}
345
346/// Householder vector `u` for the uniform sum-to-zero constraint: the
347/// reflection `H = I − 2uuᵀ` maps `c̄ = 1/√m·1` onto `e₁`, so columns 2..m
348/// of `H` are an orthonormal basis of `1⊥` — the same model space as the
349/// generic RRQR nullspace basis, but with O(rows·m) STRUCTURED application
350/// (`X·z = (X − 2(Xu)uᵀ) minus column 1`) instead of the O(rows·m²)
351/// constraint GEMM that the scale-smoke gate identified as the dominant
352/// build cost.
353pub(crate) fn householder_sum_to_zero_u(m: usize) -> Array1<f64> {
354    let c = 1.0 / (m as f64).sqrt();
355    let mut u = Array1::<f64>::from_elem(m, c);
356    u[0] -= 1.0;
357    let norm = u.dot(&u).sqrt();
358    u.mapv_inplace(|v| v / norm);
359    u
360}
361
362/// Materialize the Householder sum-to-zero basis `z` (m × (m−1)) — columns
363/// 2..m of `H = I − 2uuᵀ` — for the frozen-replay metadata. O(m²), built
364/// once per fit.
365pub(crate) fn householder_sum_to_zero_z(u: &Array1<f64>) -> Array2<f64> {
366    let m = u.len();
367    let mut z = Array2::<f64>::zeros((m, m - 1));
368    for j in 0..(m - 1) {
369        for i in 0..m {
370            let h = if i == j + 1 { 1.0 } else { 0.0 } - 2.0 * u[i] * u[j + 1];
371            z[(i, j)] = h;
372        }
373    }
374    z
375}
376
377pub(crate) fn symmetric_pseudoinverse(
378    a: &Array2<f64>,
379    label: &str,
380) -> Result<Array2<f64>, BasisError> {
381    let n = a.nrows();
382    if a.ncols() != n {
383        crate::bail_dim_basis!(
384            "measure-jet pseudo-inverse `{label}` needs a square matrix, got {:?}",
385            a.dim()
386        );
387    }
388    let (evals, evecs) = a.eigh(Side::Lower).map_err(|e| {
389        BasisError::InvalidInput(format!(
390            "measure-jet pseudo-inverse `{label}` eigendecomposition failed: {e}"
391        ))
392    })?;
393    let lam_max = evals.iter().fold(0.0_f64, |acc, v| acc.max((*v).max(0.0)));
394    let rank_tol = MEASURE_JET_PSEUDOINVERSE_RTOL * (n.max(1) as f64) * lam_max;
395    let mut scaled = evecs.clone();
396    for (k, mut col) in scaled.axis_iter_mut(Axis(1)).enumerate() {
397        let lam = evals[k].max(0.0);
398        let inv = if lam > rank_tol { 1.0 / lam } else { 0.0 };
399        col.mapv_inplace(|v| v * inv);
400    }
401    Ok(scaled.dot(&evecs.t()))
402}
403
404/// Rank-adapted center values of the measure-jet energy's affine null space.
405///
406/// The first column is the constant. The remaining columns are the supported
407/// ambient-linear directions returned by [`measure_jet_affine_head_transform`].
408/// Using that transform makes the basis full-column-rank even when the centers
409/// lie on a lower-dimensional affine stratum of the ambient coordinates.
410fn measure_jet_affine_value_basis(
411    centers: ArrayView2<'_, f64>,
412    masses: ArrayView1<'_, f64>,
413) -> Array2<f64> {
414    let m = centers.nrows();
415    let head_transform = measure_jet_affine_head_transform(centers, masses);
416    let head_rank = head_transform.ncols();
417    let mut affine = Array2::<f64>::ones((m, head_rank + 1));
418    if head_rank > 0 {
419        affine
420            .slice_mut(ndarray::s![.., 1..])
421            .assign(&centers.dot(&head_transform));
422    }
423    affine
424}
425
426/// Mass-metric quadratic form selecting the affine/null component of center
427/// function values:
428///
429/// `H₀ = W A (Aᵀ W A)⁺ Aᵀ W`.
430///
431/// This is a function-space object: `vᵀH₀v` is the squared mass norm of the
432/// affine projection of the center values `v`. No coefficient metric enters.
433fn affine_function_nullspace_form(
434    centers: ArrayView2<'_, f64>,
435    masses: ArrayView1<'_, f64>,
436) -> Result<Array2<f64>, BasisError> {
437    let m = centers.nrows();
438    if masses.len() != m {
439        crate::bail_dim_basis!(
440            "measure-jet affine function-space form shape mismatch: centers {:?}, masses {}",
441            centers.dim(),
442            masses.len()
443        );
444    }
445    let affine = measure_jet_affine_value_basis(centers, masses);
446    let mut weighted_affine = affine.clone();
447    for (i, mut row) in weighted_affine.outer_iter_mut().enumerate() {
448        row.mapv_inplace(|v| v * masses[i]);
449    }
450    let affine_gram = affine.t().dot(&weighted_affine);
451    let affine_gram_pinv = symmetric_pseudoinverse(&affine_gram, "affine function-space Gram")?;
452    let form = weighted_affine
453        .dot(&affine_gram_pinv)
454        .dot(&weighted_affine.t());
455    Ok((&form + &form.t()) * 0.5)
456}
457
458/// Pull a center-value quadratic form back through an evaluation map.
459fn pullback_center_form(evaluation: &Array2<f64>, form: &Array2<f64>) -> Array2<f64> {
460    let pulled = evaluation.t().dot(form).dot(evaluation);
461    (&pulled + &pulled.t()) * 0.5
462}
463
464/// First and diagonal-second `u = ln ℓ` derivatives of `E(u)ᵀ H E(u)` for a
465/// `u`-invariant center-value form `H`.
466fn pullback_center_form_log_length_jets(
467    evaluation: &Array2<f64>,
468    evaluation_first: &Array2<f64>,
469    evaluation_second: &Array2<f64>,
470    form: &Array2<f64>,
471) -> (Array2<f64>, Array2<f64>) {
472    let h_e = form.dot(evaluation);
473    let h_e_first = form.dot(evaluation_first);
474    let h_e_second = form.dot(evaluation_second);
475    let first_raw = evaluation_first.t().dot(&h_e) + evaluation.t().dot(&h_e_first);
476    let second_raw = evaluation_second.t().dot(&h_e)
477        + evaluation.t().dot(&h_e_second)
478        + evaluation_first.t().dot(&h_e_first) * 2.0;
479    (
480        (&first_raw + &first_raw.t()) * 0.5,
481        (&second_raw + &second_raw.t()) * 0.5,
482    )
483}
484
485/// Mixed derivative `∂²(EᵀH(ψ)E)/(∂lnℓ ∂ψ)` when only `E` depends on `ℓ`.
486fn pullback_center_form_log_length_cross(
487    evaluation: &Array2<f64>,
488    evaluation_first: &Array2<f64>,
489    form_first: &Array2<f64>,
490) -> Array2<f64> {
491    let h_e = form_first.dot(evaluation);
492    let h_e_first = form_first.dot(evaluation_first);
493    let cross_raw = evaluation_first.t().dot(&h_e) + evaluation.t().dot(&h_e_first);
494    (&cross_raw + &cross_raw.t()) * 0.5
495}
496
497/// Function-space affine/null-component penalty in the current coefficient
498/// chart. Under a coefficient reparameterization `E -> E R`, this matrix
499/// transforms covariantly as `S₀ -> Rᵀ S₀ R`; the statistical functional is
500/// therefore independent of coefficient scaling.
501pub(crate) fn affine_function_nullspace_penalty(
502    evaluation: &Array2<f64>,
503    centers: ArrayView2<'_, f64>,
504    masses: ArrayView1<'_, f64>,
505) -> Result<Array2<f64>, BasisError> {
506    if evaluation.nrows() != centers.nrows() {
507        crate::bail_dim_basis!(
508            "measure-jet affine function-space penalty shape mismatch: evaluation {:?}, centers {:?}",
509            evaluation.dim(),
510            centers.dim()
511        );
512    }
513    let form = affine_function_nullspace_form(centers, masses)?;
514    Ok(pullback_center_form(evaluation, &form))
515}
516
517/// Pairwise squared distances `‖a_i − b_j‖²` via the GEMM identity
518/// `‖a − b‖² = ‖a‖² + ‖b‖² − 2·aᵀb`: one (n×d)·(d×m) matrix product carries
519/// every FMA at tile speed instead of n·m scalar distance loops — the
520/// machine-native form of this kernel, and the module's ONLY distance
521/// source (representer design, support curve, and the center-pair geometry:
522/// band floor, median spacing, ε/2-net, neighbor cutoffs). The cancellation
523/// error near-coincident points pay is O(ε_f64·‖x‖²) ABSOLUTE, harmless
524/// under a Gaussian profile (the kernel is flat at d ≈ 0); clamped at zero
525/// so roundoff cannot emit tiny negatives (the a = b diagonal therefore
526/// lands at roundoff scale, not an exact 0 — no caller pins it).
527pub(crate) fn pairwise_sq_dists(a: ArrayView2<'_, f64>, b: ArrayView2<'_, f64>) -> Array2<f64> {
528    let an: Vec<f64> = a.outer_iter().map(|r| r.dot(&r)).collect();
529    let bn: Vec<f64> = b.outer_iter().map(|r| r.dot(&r)).collect();
530    let mut g = a.dot(&b.t());
531    g.axis_iter_mut(Axis(0))
532        .into_par_iter()
533        .enumerate()
534        .for_each(|(i, mut row)| {
535            for (j, v) in row.iter_mut().enumerate() {
536                *v = (an[i] + bn[j] - 2.0 * *v).max(0.0);
537            }
538        });
539    g
540}
541
542/// Row-block size for streaming GEMM passes that must not materialize the
543/// full n×m distance matrix (nearest-node assignment): 64Ki rows × m ≤ a
544/// few hundred MB of transient per block, GEMM-speed throughout.
545pub(crate) const MEASURE_JET_ASSIGN_BLOCK_ROWS: usize = 65_536;
546
547pub(crate) fn validate_finite_points(
548    points: ArrayView2<'_, f64>,
549    what: &str,
550) -> Result<(), BasisError> {
551    for (i, row) in points.outer_iter().enumerate() {
552        if row.iter().any(|v| !v.is_finite()) {
553            crate::bail_invalid_basis!("measure-jet {what} row {i} has a non-finite coordinate");
554        }
555    }
556    Ok(())
557}
558
559/// Median nearest-OTHER-center distance — the resolution floor of the center
560/// quadrature, used for the band floor and the auto representer range.
561pub(crate) fn median_nearest_center_spacing(dist2: &Array2<f64>) -> Result<f64, BasisError> {
562    let m = dist2.nrows();
563    if m < 2 {
564        return Err(BasisError::InsufficientColumnsForConstraint { found: m });
565    }
566    let mut nearest: Vec<f64> = Vec::with_capacity(m);
567    for i in 0..m {
568        let mut best = f64::INFINITY;
569        for j in 0..m {
570            if j != i && dist2[(i, j)] < best {
571                best = dist2[(i, j)];
572            }
573        }
574        nearest.push(best.sqrt());
575    }
576    nearest.sort_by(|a, b| a.partial_cmp(b).expect("finite center spacings"));
577    let median = nearest[nearest.len() / 2];
578    if !(median.is_finite() && median > 0.0) {
579        crate::bail_invalid_basis!(
580            "measure-jet centers are degenerate (median nearest-center spacing = {median}); \
581             duplicate centers cannot carry a scale band"
582        );
583    }
584    Ok(median)
585}
586
587/// Build the realized geometric scale band from the center set: floor at the
588/// median nearest-center spacing (below it the quadrature resolves nothing),
589/// ceiling at half the bounding-box diagonal (a deterministic diameter-scale
590/// cap; local fits remain center-weighted and distinct there).
591/// `num_scales == 0` requests the auto count `clamp(⌈log2(ε_max/ε_min)⌉ + 1,
592/// 3, 8)`; a degenerate band (ceiling ≤ floor) collapses to the single floor
593/// scale with `log_step = ln 2`.
594pub fn measure_jet_band(
595    centers: ArrayView2<'_, f64>,
596    num_scales: usize,
597) -> Result<MeasureJetBand, BasisError> {
598    validate_finite_points(centers, "centers")?;
599    let dist2 = pairwise_sq_dists(centers, centers);
600    let eps_min = median_nearest_center_spacing(&dist2)?;
601    // Half the bounding-box diagonal: a cheap, deterministic diameter proxy.
602    let d = centers.ncols();
603    let mut diag2 = 0.0_f64;
604    for k in 0..d {
605        let col = centers.column(k);
606        let mut lo = f64::INFINITY;
607        let mut hi = f64::NEG_INFINITY;
608        for &v in col.iter() {
609            lo = lo.min(v);
610            hi = hi.max(v);
611        }
612        diag2 += (hi - lo) * (hi - lo);
613    }
614    let eps_max = 0.5 * diag2.sqrt();
615    if !(eps_max.is_finite() && eps_max > eps_min) {
616        return Ok(MeasureJetBand {
617            eps: vec![eps_min],
618            log_step: std::f64::consts::LN_2,
619        });
620    }
621    let auto = ((eps_max / eps_min).log2().ceil() as usize + 1)
622        .clamp(MEASURE_JET_MIN_AUTO_SCALES, MEASURE_JET_MAX_AUTO_SCALES);
623    let count = if num_scales == 0 { auto } else { num_scales };
624    if count == 1 {
625        return Ok(MeasureJetBand {
626            eps: vec![eps_min],
627            log_step: std::f64::consts::LN_2,
628        });
629    }
630    let ratio = (eps_max / eps_min).powf(1.0 / (count as f64 - 1.0));
631    let mut eps = Vec::with_capacity(count);
632    let mut e = eps_min;
633    for _ in 0..count {
634        eps.push(e);
635        e *= ratio;
636    }
637    Ok(MeasureJetBand {
638        eps,
639        log_step: ratio.ln(),
640    })
641}
642
643/// First-moment-exact quadrature of the empirical measure on the cell partition
644/// induced by the seed centers: nearest-center assignment (deterministic
645/// tie-break: lowest center index) yields per-cell masses, and each non-empty
646/// cell's quadrature node is its mass-weighted barycenter. Empty cells keep
647/// their seed coordinates with zero mass (the assembly skips them; their
648/// representer columns remain valid).
649pub fn measure_jet_quadrature_nodes(
650    data: ArrayView2<'_, f64>,
651    centers: ArrayView2<'_, f64>,
652) -> Result<(Array2<f64>, Array1<f64>), BasisError> {
653    if data.ncols() != centers.ncols() {
654        crate::bail_dim_basis!(
655            "measure-jet mass assignment dimension mismatch: data d={} centers d={}",
656            data.ncols(),
657            centers.ncols()
658        );
659    }
660    validate_finite_points(data, "data")?;
661    validate_finite_points(centers, "centers")?;
662    let n = data.nrows();
663    let m = centers.nrows();
664    let d = centers.ncols();
665    if n == 0 || m == 0 {
666        crate::bail_invalid_basis!("measure-jet mass assignment needs nonempty data and centers");
667    }
668    // Nearest-node assignment in streamed GEMM blocks: argmin_j ‖x−c_j‖² =
669    // argmin_j (‖c_j‖² − 2·xᵀc_j), so each block is one (rows×d)·(d×m)
670    // product plus a row-wise argmin — tile-speed FMAs, O(block·m) transient
671    // memory, deterministic ties to the lowest center index.
672    let cn: Vec<f64> = centers.outer_iter().map(|r| r.dot(&r)).collect();
673    let assignments: Vec<usize> = (0..n)
674        .step_by(MEASURE_JET_ASSIGN_BLOCK_ROWS)
675        .flat_map(|start| {
676            let end = (start + MEASURE_JET_ASSIGN_BLOCK_ROWS).min(n);
677            let g = data.slice(ndarray::s![start..end, ..]).dot(&centers.t());
678            let block: Vec<usize> = g
679                .axis_iter(Axis(0))
680                .into_par_iter()
681                .map(|row| {
682                    let mut best_j = 0usize;
683                    let mut best = f64::INFINITY;
684                    for (j, &gij) in row.iter().enumerate() {
685                        let s = cn[j] - 2.0 * gij;
686                        if s < best {
687                            best = s;
688                            best_j = j;
689                        }
690                    }
691                    best_j
692                })
693                .collect();
694            block
695        })
696        .collect();
697    let mut masses = Array1::<f64>::zeros(m);
698    let mut nodes = centers.to_owned();
699    let mut sums = Array2::<f64>::zeros((m, d));
700    let unit = 1.0 / n as f64;
701    for (i, &j) in assignments.iter().enumerate() {
702        masses[j] += unit;
703        for k in 0..d {
704            sums[(j, k)] += data[(i, k)];
705        }
706    }
707    // Cell barycenters: the first moment of μ on each cell. These are the
708    // realized nodes for first-moment-exact lumping.
709    let mut barycenter = sums;
710    for j in 0..m {
711        let count = masses[j] * n as f64;
712        if count > 0.0 {
713            for k in 0..d {
714                barycenter[(j, k)] /= count;
715                nodes[(j, k)] = barycenter[(j, k)];
716            }
717        }
718    }
719    Ok((nodes, masses))
720}
721
722/// Per-center masses of the empirical measure (the zeroth-moment half of
723/// [`measure_jet_quadrature_nodes`]; single assignment source).
724pub fn measure_jet_center_masses(
725    data: ArrayView2<'_, f64>,
726    centers: ArrayView2<'_, f64>,
727) -> Result<Array1<f64>, BasisError> {
728    measure_jet_quadrature_nodes(data, centers).map(|(_, masses)| masses)
729}
730
731/// THE single assembly source: walk every (scale, outer-net center) local
732/// residual block exactly once and scatter it into `n_forms` accumulators
733/// with caller-chosen scalar weights. The energy, its (s, α) jets, and the
734/// per-scale spectrum are all this routine with different weight closures,
735/// so a value/derivative desync is structurally impossible.
736///
737/// Per block the closure receives `(scale_idx, eps, q, base)` where `q` is
738/// the truncated kernel sum used by the local residual and `base`
739/// is the fully-assembled outer weight
740/// `log_step · ε^(−η) · net_mass_i · q^(1−2α)`, with
741/// `η = 2s + d(2−2α)` for the available dimension parameter, and writes, per requested
742/// form, one weight triple `[w_R, w_2, w_3]`. Only `w_R` is live:
743/// `R = CᵀWC − B·G⁺·Bᵀ/q`, with `G⁺` the rank-revealing pseudo-inverse.
744/// The extra slots are retained for the ψ layout and receive zero local
745/// channels because τ no longer changes the energy.
746///
747/// The outer sum over centers is coarsened per scale to a deterministic
748/// ε/2-net with nearest-member mass aggregation (the outer Riemann sum needs
749/// resolution ε, not the center-spacing floor), so each scale's cost sits at
750/// its own level and the band totals ~O(m²·d) instead of O(L·m³). The inner
751/// (local-fit) quadrature always uses the full center set, so the local
752/// residual identities (exact constant annihilation, PSD) are untouched.
753pub(crate) fn assemble_weighted_forms<F>(
754    centers: ArrayView2<'_, f64>,
755    masses: ArrayView1<'_, f64>,
756    band: &MeasureJetBand,
757    order_s: f64,
758    alpha: f64,
759    tau0: f64,
760    n_forms: usize,
761    channels: usize,
762    weights: &F,
763) -> Result<Vec<Array2<f64>>, BasisError>
764where
765    F: Fn(usize, f64, f64, f64, &mut [[f64; 3]]) + Sync,
766{
767    let m = centers.nrows();
768    let d = centers.ncols();
769    if n_forms == 0 || !(1..=3).contains(&channels) {
770        crate::bail_invalid_basis!(
771            "measure-jet assembly needs at least one output form and 1..=3 block channels"
772        );
773    }
774    if masses.len() != m {
775        crate::bail_dim_basis!(
776            "measure-jet energy mass/center mismatch: {} masses for {} centers",
777            masses.len(),
778            m
779        );
780    }
781    if band.eps.is_empty() || band.eps.iter().any(|e| !(e.is_finite() && *e > 0.0)) {
782        crate::bail_invalid_basis!("measure-jet energy needs a nonempty positive scale band");
783    }
784    if !(order_s.is_finite() && order_s > 0.0 && order_s < 2.0) {
785        crate::bail_invalid_basis!(
786            "measure-jet order s must lie in (0, 2) for the affine-jet energy; got {order_s}"
787        );
788    }
789    if !(alpha.is_finite() && tau0.is_finite() && tau0 >= 0.0) {
790        crate::bail_invalid_basis!(
791            "measure-jet energy needs finite alpha and finite tau0 >= 0; got alpha={alpha}, tau0={tau0}"
792        );
793    }
794    if masses.iter().any(|v| !(v.is_finite() && *v >= 0.0)) {
795        crate::bail_invalid_basis!("measure-jet energy needs finite nonnegative center masses");
796    }
797    let dist2 = pairwise_sq_dists(centers, centers);
798
799    // One block of `n_forms` m×m accumulators per scale. Each scale's center
800    // loop is sequential and the cross-scale sum below runs in band order,
801    // so the result is bit-deterministic whether or not the scales
802    // themselves run in parallel.
803    let assemble_scale = |scale_idx: usize, eps: f64| -> Result<Vec<Array2<f64>>, BasisError> {
804        let mut out: Vec<Array2<f64>> =
805            (0..n_forms).map(|_| Array2::<f64>::zeros((m, m))).collect();
806        let cutoff2 = (MEASURE_JET_PROFILE_CUTOFF * eps) * (MEASURE_JET_PROFILE_CUTOFF * eps);
807        let inv_two_eps2 = 1.0 / (2.0 * eps * eps);
808        let eta = 2.0 * order_s + (d as f64) * (2.0 - 2.0 * alpha);
809        let scale_weight = band.log_step * eps.powf(-eta);
810        // Outer-quadrature coarsening: greedy ε/2-net over the centers in
811        // fixed index order (deterministic), with every center's mass
812        // aggregated to its nearest net member (lowest-index tie break).
813        let net_radius2 = 0.25 * eps * eps;
814        let mut outer: Vec<usize> = Vec::new();
815        for i in 0..m {
816            if masses[i] <= 0.0 {
817                continue;
818            }
819            let covered = outer.iter().any(|&o| dist2[(i, o)] <= net_radius2);
820            if !covered {
821                outer.push(i);
822            }
823        }
824        let mut net_mass = vec![0.0_f64; m];
825        for i in 0..m {
826            if masses[i] <= 0.0 {
827                continue;
828            }
829            let mut best = f64::INFINITY;
830            let mut best_o = usize::MAX;
831            for &o in &outer {
832                if dist2[(i, o)] < best {
833                    best = dist2[(i, o)];
834                    best_o = o;
835                }
836            }
837            if best_o != usize::MAX {
838                net_mass[best_o] += masses[i];
839            }
840        }
841        let mut wbuf = vec![[0.0_f64; 3]; n_forms];
842        for &i in &outer {
843            // Local neighbor set (always includes i itself).
844            let mut idx: Vec<usize> = Vec::new();
845            for j in 0..m {
846                if dist2[(i, j)] <= cutoff2 {
847                    idx.push(j);
848                }
849            }
850            let ml = idx.len();
851            // Kernel weights and mass.
852            let mut w = Array1::<f64>::zeros(ml);
853            let mut q = 0.0_f64;
854            for (a, &j) in idx.iter().enumerate() {
855                let wj = masses[j] * (-dist2[(i, j)] * inv_two_eps2).exp();
856                w[a] = wj;
857                q += wj;
858            }
859            if !(q > 0.0) {
860                continue;
861            }
862            // Scaled local features Φ (ml × d) and weighted column means a.
863            let mut phi = Array2::<f64>::zeros((ml, d));
864            for (a, &j) in idx.iter().enumerate() {
865                for k in 0..d {
866                    phi[(a, k)] = (centers[(j, k)] - centers[(i, k)]) / eps;
867                }
868            }
869            let a_mean = phi.t().dot(&w) / q;
870            // B = WΦ − w·aᵀ and G = (ΦᵀWΦ)/q − a·aᵀ.
871            let mut wphi = phi.clone();
872            for (a, mut row) in wphi.outer_iter_mut().enumerate() {
873                row.mapv_inplace(|v| v * w[a]);
874            }
875            let mut b = wphi.clone();
876            for (a, mut row) in b.outer_iter_mut().enumerate() {
877                for k in 0..d {
878                    row[k] -= w[a] * a_mean[k];
879                }
880            }
881            let mut g = phi.t().dot(&wphi);
882            g.mapv_inplace(|v| v / q);
883            for r in 0..d {
884                for c in 0..d {
885                    g[(r, c)] -= a_mean[r] * a_mean[c];
886                }
887            }
888            let g_pinv = symmetric_pseudoinverse(&g, "local affine Gram")?;
889            let bm = b.dot(&g_pinv);
890            let base = scale_weight * net_mass[i] * q.powf(1.0 - 2.0 * alpha);
891            weights(scale_idx, eps, q, base, &mut wbuf);
892            // Scatter-add Σ_k wbuf[k]·R into each form. The τ channels are
893            // zero because the exact projection is τ-independent.
894            for (a, &ja) in idx.iter().enumerate() {
895                let bma = bm.row(a);
896                for (c, &jc) in idx.iter().enumerate() {
897                    let b_c = b.row(c);
898                    let mut val_r = -w[a] * w[c] / q - bma.dot(&b_c) / q;
899                    if a == c {
900                        val_r += w[a];
901                    }
902                    for (k, out_k) in out.iter_mut().enumerate() {
903                        let wk = wbuf[k];
904                        out_k[(ja, jc)] += wk[0] * val_r;
905                    }
906                }
907            }
908        }
909        Ok(out)
910    };
911
912    let n_scales = band.eps.len();
913    let parallel_ok = m
914        .saturating_mul(m)
915        .saturating_mul(n_scales)
916        .saturating_mul(n_forms)
917        <= MEASURE_JET_PARALLEL_FORM_BUDGET_DOUBLES;
918    let per_scale: Vec<Vec<Array2<f64>>> = if parallel_ok {
919        band.eps
920            .par_iter()
921            .enumerate()
922            .map(|(scale_idx, &eps)| assemble_scale(scale_idx, eps))
923            .collect::<Result<Vec<_>, BasisError>>()?
924    } else {
925        band.eps
926            .iter()
927            .enumerate()
928            .map(|(scale_idx, &eps)| assemble_scale(scale_idx, eps))
929            .collect::<Result<Vec<_>, BasisError>>()?
930    };
931
932    let mut totals: Vec<Array2<f64>> = (0..n_forms).map(|_| Array2::<f64>::zeros((m, m))).collect();
933    for scale_forms in per_scale {
934        for (total, part) in totals.iter_mut().zip(scale_forms) {
935            *total += &part;
936        }
937    }
938    // Numerical symmetrization (every analytic form here is symmetric).
939    Ok(totals.into_iter().map(|t| (&t + &t.t()) * 0.5).collect())
940}
941
942/// The multiscale jet-residual energy `Q` (m × m, symmetric PSD) on the
943/// center set. See the module docs for the formula and contracts; the local
944/// residual form is assembled through the closed-form identities
945///
946/// ```text
947///   CᵀWC          = W − w·wᵀ/q,
948///   B = CᵀWΦ̃     = WΦ − w·aᵀ          (a = Φᵀw/q),
949///   G = Φ̃ᵀWΦ̃/q  = (ΦᵀWΦ)/q − a·aᵀ,
950///   R_loc         = CᵀWC − B·G⁺·Bᵀ/q,
951/// ```
952///
953/// with `G⁺` realized through the symmetric eigendecomposition and a
954/// machine-precision rank cutoff. One walk of [`assemble_weighted_forms`]
955/// with the unit weight.
956pub fn measure_jet_energy_form(
957    centers: ArrayView2<'_, f64>,
958    masses: ArrayView1<'_, f64>,
959    band: &MeasureJetBand,
960    order_s: f64,
961    alpha: f64,
962    tau0: f64,
963) -> Result<Array2<f64>, BasisError> {
964    let mut forms = assemble_weighted_forms(
965        centers,
966        masses,
967        band,
968        order_s,
969        alpha,
970        tau0,
971        1,
972        1,
973        &|_, _, _, base, out: &mut [[f64; 3]]| out[0] = [base, 0.0, 0.0],
974    )?;
975    let q = forms.swap_remove(0);
976    // The energy `Q = Σ wᵢ Rᵢ` is a nonnegative combination of analytically
977    // PSD local residual forms, so it is PSD in exact arithmetic. The affine
978    // span is annihilated to machine zero, where roundoff in the per-block
979    // pseudo-inverse and the centering cancellation leaves the smallest
980    // eigenvalue at ±ε_mach·‖Q‖. Project onto the PSD cone (floor negative
981    // eigenvalues at 0) so `vᵀQv ≥ 0` holds exactly for every `v`, including
982    // the affine directions the energy must annihilate.
983    project_symmetric_psd(q, "measure-jet energy form")
984}
985
986/// Project a symmetric matrix onto the PSD cone by flooring its negative
987/// eigenvalues at 0. Only sub-machine-precision negative eigenvalues are
988/// expected here (the form is analytically PSD); a meaningfully negative
989/// eigenvalue would indicate an assembly bug, so it is floored but the
990/// reconstruction otherwise preserves the spectrum exactly.
991pub(crate) fn project_symmetric_psd(
992    a: Array2<f64>,
993    label: &str,
994) -> Result<Array2<f64>, BasisError> {
995    let n = a.nrows();
996    if n == 0 {
997        return Ok(a);
998    }
999    let (evals, evecs) = a.eigh(Side::Lower).map_err(|e| {
1000        BasisError::InvalidInput(format!(
1001            "measure-jet PSD projection `{label}` eigendecomposition failed: {e}"
1002        ))
1003    })?;
1004    if evals.iter().all(|&lam| lam >= 0.0) {
1005        return Ok(a);
1006    }
1007    let mut scaled = evecs.clone();
1008    for (k, mut col) in scaled.axis_iter_mut(Axis(1)).enumerate() {
1009        let lam = evals[k].max(0.0);
1010        col.mapv_inplace(|v| v * lam);
1011    }
1012    let psd = scaled.dot(&evecs.t());
1013    Ok((&psd + &psd.t()) * 0.5)
1014}
1015
1016/// The energy together with its exact first and second jets in the live
1017/// dials, plus zero slots for the retained `ψ_τ = ln τ` coordinate. With
1018/// `g_s = −2 ln ε`, `g_α = −2 ln q`:
1019///
1020/// ```text
1021///   ∂Q/∂s   = Σ g_s·w·R,        ∂²Q/∂s²   = Σ g_s²·w·R,
1022///   ∂Q/∂α   = Σ g_α·w·R,        ∂²Q/∂α²   = Σ g_α²·w·R,
1023///   ∂²Q/∂s∂α = Σ g_s·g_α·w·R,
1024///   ∂Q/∂ψ_τ = ∂²Q/∂ψ_τ² = ∂²Q/∂s∂ψ_τ = ∂²Q/∂α∂ψ_τ = 0.
1025/// ```
1026///
1027/// all scattered from the SAME local blocks as `Q` in one pass (no second
1028/// assembly that could drift). FD-gated in this module's tests. Requires
1029/// `tau0 > 0` only because the retained coordinate is `ln τ`.
1030pub fn measure_jet_energy_form_with_jets(
1031    centers: ArrayView2<'_, f64>,
1032    masses: ArrayView1<'_, f64>,
1033    band: &MeasureJetBand,
1034    order_s: f64,
1035    alpha: f64,
1036    tau0: f64,
1037) -> Result<MeasureJetEnergyJets, BasisError> {
1038    if !(tau0.is_finite() && tau0 > 0.0) {
1039        crate::bail_invalid_basis!(
1040            "measure-jet jets need tau0 > 0 because the retained τ coordinate is ln τ; got {tau0}"
1041        );
1042    }
1043    let mut forms = assemble_weighted_forms(
1044        centers,
1045        masses,
1046        band,
1047        order_s,
1048        alpha,
1049        tau0,
1050        10,
1051        3,
1052        &|_, eps: f64, q: f64, base: f64, out: &mut [[f64; 3]]| {
1053            let gs = -2.0 * eps.ln();
1054            let intrinsic_dim = centers.ncols() as f64;
1055            let ga = 2.0 * intrinsic_dim * eps.ln() - 2.0 * q.max(f64::MIN_POSITIVE).ln();
1056            out[0] = [base, 0.0, 0.0];
1057            out[1] = [gs * base, 0.0, 0.0];
1058            out[2] = [gs * gs * base, 0.0, 0.0];
1059            out[3] = [ga * base, 0.0, 0.0];
1060            out[4] = [ga * ga * base, 0.0, 0.0];
1061            out[5] = [gs * ga * base, 0.0, 0.0];
1062            out[6] = [0.0, 0.0, 0.0];
1063            out[7] = [0.0, 0.0, 0.0];
1064            out[8] = [0.0, 0.0, 0.0];
1065            out[9] = [0.0, 0.0, 0.0];
1066        },
1067    )?;
1068    let d2q_dalpha_dlogtau = forms.pop().expect("ten assembled forms");
1069    let d2q_ds_dlogtau = forms.pop().expect("ten assembled forms");
1070    let d2q_dlogtau2 = forms.pop().expect("ten assembled forms");
1071    let dq_dlogtau = forms.pop().expect("ten assembled forms");
1072    let d2q_ds_dalpha = forms.pop().expect("ten assembled forms");
1073    let d2q_dalpha2 = forms.pop().expect("ten assembled forms");
1074    let dq_dalpha = forms.pop().expect("ten assembled forms");
1075    let d2q_ds2 = forms.pop().expect("ten assembled forms");
1076    let dq_ds = forms.pop().expect("ten assembled forms");
1077    let q = forms.pop().expect("ten assembled forms");
1078    Ok(MeasureJetEnergyJets {
1079        q,
1080        dq_ds,
1081        d2q_ds2,
1082        dq_dalpha,
1083        d2q_dalpha2,
1084        d2q_ds_dalpha,
1085        dq_dlogtau,
1086        d2q_dlogtau2,
1087        d2q_ds_dlogtau,
1088        d2q_dalpha_dlogtau,
1089    })
1090}
1091
1092/// Per-scale energy decomposition of center values `v`: element ℓ is
1093/// `vᵀ Q_ℓ v`, the detail energy charged at scale `ε_ℓ`. Sums exactly to
1094/// `vᵀQv` (same blocks, one-hot weights) and doubles as the scale spectrum
1095/// diagnostic of the fitted intensity field — where along the band the
1096/// signal lives, and the analytic carrier of `∂/∂s` reweightings.
1097pub fn measure_jet_scale_spectrum(
1098    centers: ArrayView2<'_, f64>,
1099    masses: ArrayView1<'_, f64>,
1100    band: &MeasureJetBand,
1101    order_s: f64,
1102    alpha: f64,
1103    tau0: f64,
1104    values: ArrayView1<'_, f64>,
1105) -> Result<Vec<f64>, BasisError> {
1106    if values.len() != centers.nrows() {
1107        crate::bail_dim_basis!(
1108            "measure-jet scale spectrum needs one value per center: {} values for {} centers",
1109            values.len(),
1110            centers.nrows()
1111        );
1112    }
1113    let forms = measure_jet_energy_forms_per_scale(centers, masses, band, order_s, alpha, tau0)?;
1114    Ok(forms
1115        .iter()
1116        .map(|q_l| values.dot(&q_l.dot(&values)))
1117        .collect())
1118}
1119
1120/// The per-scale energy forms `Q_ℓ` (each m × m, symmetric PSD), with
1121/// `Σ_ℓ Q_ℓ = Q` exactly (same blocks, one-hot weights). These are the
1122/// spectral-split carriers: emitted as separate penalty candidates they let
1123/// the multi-penalty REML engine learn per-level amplitudes λ_ℓ directly —
1124/// scale adaptivity at ρ-speed with no rebuild and no new optimizer code.
1125pub fn measure_jet_energy_forms_per_scale(
1126    centers: ArrayView2<'_, f64>,
1127    masses: ArrayView1<'_, f64>,
1128    band: &MeasureJetBand,
1129    order_s: f64,
1130    alpha: f64,
1131    tau0: f64,
1132) -> Result<Vec<Array2<f64>>, BasisError> {
1133    let n_scales = band.eps.len();
1134    assemble_weighted_forms(
1135        centers,
1136        masses,
1137        band,
1138        order_s,
1139        alpha,
1140        tau0,
1141        n_scales,
1142        1,
1143        &|scale_idx, _, _, base, out: &mut [[f64; 3]]| {
1144            for (k, slot) in out.iter_mut().enumerate() {
1145                *slot = if k == scale_idx {
1146                    [base, 0.0, 0.0]
1147                } else {
1148                    [0.0, 0.0, 0.0]
1149                };
1150            }
1151        },
1152    )
1153}
1154
1155/// The support diagnostic `ε ↦ q_ε(x★)`: kernel mass of the (frozen) center
1156/// quadrature seen from each query point at every band scale (n_query × L).
1157/// A query ON the web sees its strand's mass already at fine scales; a query
1158/// OFF the web accumulates mass only once ε reaches its distance to the
1159/// support. This is the on-web-ness statistic shipped alongside predictions
1160/// — smooth, multiresolution, derived from the measure with no neighbor
1161/// sets.
1162pub fn measure_jet_support_curve(
1163    queries: ArrayView2<'_, f64>,
1164    centers: ArrayView2<'_, f64>,
1165    masses: ArrayView1<'_, f64>,
1166    eps_band: &[f64],
1167) -> Result<Array2<f64>, BasisError> {
1168    if queries.ncols() != centers.ncols() {
1169        crate::bail_dim_basis!(
1170            "measure-jet support curve dimension mismatch: queries d={} centers d={}",
1171            queries.ncols(),
1172            centers.ncols()
1173        );
1174    }
1175    if masses.len() != centers.nrows() {
1176        crate::bail_dim_basis!(
1177            "measure-jet support curve mass/center mismatch: {} masses for {} centers",
1178            masses.len(),
1179            centers.nrows()
1180        );
1181    }
1182    if eps_band.is_empty() || eps_band.iter().any(|e| !(e.is_finite() && *e > 0.0)) {
1183        crate::bail_invalid_basis!("measure-jet support curve needs a nonempty positive band");
1184    }
1185    validate_finite_points(queries, "queries")?;
1186    validate_finite_points(centers, "centers")?;
1187    let nq = queries.nrows();
1188    let nl = eps_band.len();
1189    // Distances once (GEMM), then every band scale reads the same d² row —
1190    // an L-fold saving over per-scale distance recomputation.
1191    let d2 = pairwise_sq_dists(queries, centers);
1192    let mut out = Array2::<f64>::zeros((nq, nl));
1193    out.axis_iter_mut(Axis(0))
1194        .into_par_iter()
1195        .enumerate()
1196        .for_each(|(qi, mut row)| {
1197            let d2_row = d2.row(qi);
1198            for (li, &eps) in eps_band.iter().enumerate() {
1199                let inv_two_eps2 = 1.0 / (2.0 * eps * eps);
1200                let mut acc = 0.0_f64;
1201                for (j, &dd) in d2_row.iter().enumerate() {
1202                    acc += masses[j] * (-dd * inv_two_eps2).exp();
1203                }
1204                row[li] = acc;
1205            }
1206        });
1207    Ok(out)
1208}
1209
1210pub(crate) fn measure_jet_support_means(
1211    centers: ArrayView2<'_, f64>,
1212    masses: ArrayView1<'_, f64>,
1213    eps_band: &[f64],
1214) -> Result<Vec<f64>, BasisError> {
1215    let total_mass = masses.sum();
1216    if !(total_mass.is_finite() && total_mass > 0.0) {
1217        crate::bail_invalid_basis!(
1218            "measure-jet support means need positive finite total mass; got {total_mass}"
1219        );
1220    }
1221    let support = measure_jet_support_curve(centers, centers, masses, eps_band)?;
1222    let mut means = vec![0.0_f64; eps_band.len()];
1223    for (i, row) in support.rows().into_iter().enumerate() {
1224        let mass = masses[i];
1225        for (mean, &q) in means.iter_mut().zip(row.iter()) {
1226            *mean += mass * q;
1227        }
1228    }
1229    for mean in &mut means {
1230        *mean /= total_mass;
1231        if !(*mean).is_finite() || *mean <= 0.0 {
1232            crate::bail_invalid_basis!(
1233                "measure-jet support mean must be positive and finite; got {mean}"
1234            );
1235        }
1236    }
1237    Ok(means)
1238}
1239
1240/// Gaussian representer features `exp(−‖x − c‖²/(2ℓ²))` (n × m).
1241pub fn measure_jet_design_matrix(
1242    data: ArrayView2<'_, f64>,
1243    centers: ArrayView2<'_, f64>,
1244    length_scale: f64,
1245) -> Result<Array2<f64>, BasisError> {
1246    if data.ncols() != centers.ncols() {
1247        crate::bail_dim_basis!(
1248            "measure-jet design dimension mismatch: data d={} centers d={}",
1249            data.ncols(),
1250            centers.ncols()
1251        );
1252    }
1253    if !(length_scale.is_finite() && length_scale > 0.0) {
1254        crate::bail_invalid_basis!(
1255            "measure-jet design needs a positive finite length_scale; got {length_scale}"
1256        );
1257    }
1258    validate_finite_points(data, "data")?;
1259    validate_finite_points(centers, "centers")?;
1260    let inv_two_l2 = 1.0 / (2.0 * length_scale * length_scale);
1261    // One GEMM for every distance, then the Gaussian applied in place — the
1262    // n×m allocation IS the output, no transient copy.
1263    let mut out = pairwise_sq_dists(data, centers);
1264    out.axis_iter_mut(Axis(0))
1265        .into_par_iter()
1266        .for_each(|mut row| {
1267            row.mapv_inplace(|d2| (-d2 * inv_two_l2).exp());
1268        });
1269    Ok(out)
1270}
1271
1272/// Exact first and diagonal-second derivatives of the Gaussian representer
1273/// design with respect to `u = ln ℓ`.
1274fn measure_jet_design_log_length_jets(
1275    data: ArrayView2<'_, f64>,
1276    centers: ArrayView2<'_, f64>,
1277    length_scale: f64,
1278) -> Result<(Array2<f64>, Array2<f64>), BasisError> {
1279    let kernel = measure_jet_design_matrix(data, centers, length_scale)?;
1280    let squared_distances = pairwise_sq_dists(data, centers);
1281    let inv_l2 = 1.0 / (length_scale * length_scale);
1282    let mut first = kernel.clone();
1283    let mut second = kernel;
1284    for ((first_value, second_value), &distance_squared) in first
1285        .iter_mut()
1286        .zip(second.iter_mut())
1287        .zip(squared_distances.iter())
1288    {
1289        let a = distance_squared * inv_l2;
1290        let kernel_value = *first_value;
1291        *first_value = kernel_value * a;
1292        *second_value = kernel_value * (a * a - 2.0 * a);
1293    }
1294    Ok((first, second))
1295}
1296
1297/// Rank-revealing ambient-linear head lift `T` (d × head_rank) for the
1298/// extrapolation null space (#1845).
1299///
1300/// The measure-jet energy annihilates ambient-affine functions EXACTLY (the
1301/// no-mass contract), so the affine functions are the penalty's null space —
1302/// the directions the fit is free to extend across a training gap. But the
1303/// Gaussian representer design cannot REPRESENT a global affine function off
1304/// its support: a finite sum of decaying bumps reverts to the parametric
1305/// backbone away from the centers, so in a gap the fit collapses toward the
1306/// training mean instead of carrying the flank-attested trend. Completing the
1307/// smoothing-spline structure, the builder appends this ambient-linear null
1308/// space to the design as an UNPENALIZED head (the `{x_1..x_d}` head the frame
1309/// notes §1 pin as the property the representer basis lacked).
1310///
1311/// The head is data-derived and magic-free. Ambient coordinates of data on a
1312/// low intrinsic-dimension stratum are rank-deficient as linear trends, so the
1313/// coordinate columns are orthonormalized on the centers and the
1314/// numerically-degenerate directions dropped. Working in the mean-CENTERED
1315/// coordinate columns (the mass-weighted mean is the intercept's, not the
1316/// head's) makes the rank test measure the genuine spread of the centers along
1317/// each direction rather than its offset; the relative floor
1318/// [`MEASURE_JET_PSEUDOINVERSE_RTOL`] is the module's own numerical rank
1319/// tolerance (the same one the local Gram pseudo-inverses use). The returned
1320/// `T` satisfies `head(points) = points · T` (the mean-centering only informs
1321/// the keep/drop decision; the constant component of `points · T` is removed
1322/// downstream by the global parametric orthogonalization). `T` is a
1323/// deterministic function of the frozen centers + masses, so the frozen replay
1324/// path reconstructs the identical head with no persisted state. Public so the
1325/// predict-side errors-in-variables gradient (#2225) can reconstruct the head
1326/// lift `T` from the frozen centers + masses to differentiate the head block.
1327pub fn measure_jet_affine_head_transform(
1328    centers: ArrayView2<'_, f64>,
1329    masses: ArrayView1<'_, f64>,
1330) -> Array2<f64> {
1331    let m = centers.nrows();
1332    let d = centers.ncols();
1333    let total_mass = masses.sum();
1334    // Mass inner product on center values.
1335    let mdot = |u: &Array1<f64>, v: &Array1<f64>| -> f64 {
1336        let mut acc = 0.0;
1337        for i in 0..m {
1338            acc += masses[i] * u[i] * v[i];
1339        }
1340        acc
1341    };
1342    // Mean-centered coordinate columns: the mass-weighted mean is removed so the
1343    // residual mass-norm is the genuine spread of the centers along a direction,
1344    // not dominated by the coordinate's offset (which the intercept owns).
1345    let cols: Vec<Array1<f64>> = (0..d)
1346        .map(|k| {
1347            let col = centers.column(k).to_owned();
1348            let mean = if total_mass > 0.0 {
1349                mdot(&col, &Array1::ones(m)) / total_mass
1350            } else {
1351                0.0
1352            };
1353            col.mapv(|x| x - mean)
1354        })
1355        .collect();
1356    // Relative numerical rank floor from the centered coordinate-column scale.
1357    let max_norm = cols
1358        .iter()
1359        .fold(0.0_f64, |acc, c| acc.max(mdot(c, c).sqrt()));
1360    let drop_below =
1361        (MEASURE_JET_PSEUDOINVERSE_RTOL * (d.max(1) as f64) * max_norm).max(f64::MIN_POSITIVE);
1362    // Mass-weighted modified Gram–Schmidt on the centered columns; `t`
1363    // accumulates the lift in the ORIGINAL coordinate basis, so every kept head
1364    // column is `points · t_r` (up to the intercept-owned constant).
1365    let mut q_cols: Vec<Array1<f64>> = Vec::new();
1366    let mut t_cols: Vec<Array1<f64>> = Vec::new();
1367    for k in 0..d {
1368        let mut v = cols[k].clone();
1369        let mut t = Array1::<f64>::zeros(d);
1370        t[k] = 1.0;
1371        for (q, tq) in q_cols.iter().zip(t_cols.iter()) {
1372            let proj = mdot(q, &v);
1373            v.scaled_add(-proj, q);
1374            t.scaled_add(-proj, tq);
1375        }
1376        let norm = mdot(&v, &v).sqrt();
1377        if norm > drop_below {
1378            v.mapv_inplace(|x| x / norm);
1379            t.mapv_inplace(|x| x / norm);
1380            q_cols.push(v);
1381            t_cols.push(t);
1382        }
1383    }
1384    let head_rank = t_cols.len();
1385    let mut t_mat = Array2::<f64>::zeros((d, head_rank));
1386    for (r, t) in t_cols.into_iter().enumerate() {
1387        t_mat.column_mut(r).assign(&t);
1388    }
1389    t_mat
1390}
1391
1392/// Resolve the realized representer range ℓ. An explicit positive
1393/// `spec_length_scale` is used verbatim; the `0.0` sentinel auto-initializes
1394/// from the median nearest-center spacing (one spacing width: neighbors
1395/// overlap at exp(−1/2) ≈ 0.61, smooth blend without collinearity).
1396pub fn realized_measure_jet_length_scale(
1397    centers: ArrayView2<'_, f64>,
1398    spec_length_scale: f64,
1399) -> Result<f64, BasisError> {
1400    if spec_length_scale.is_finite() && spec_length_scale > 0.0 {
1401        return Ok(spec_length_scale);
1402    }
1403    if spec_length_scale != 0.0 {
1404        crate::bail_invalid_basis!(
1405            "measure-jet length_scale must be positive (or 0.0 for auto); got {spec_length_scale}"
1406        );
1407    }
1408    let dist2 = pairwise_sq_dists(centers, centers);
1409    let spacing = median_nearest_center_spacing(&dist2)?;
1410    Ok(MEASURE_JET_AUTO_LENGTH_SCALE_FACTOR * spacing)
1411}
1412
1413/// The realized, ψ-FIXED geometry shared by the basis builder and the
1414/// ψ-derivative producer — ONE realization source, so the penalty the fit
1415/// uses and the penalty the ψ-channel differentiates can never drift apart
1416/// (the #901 desync class, excluded structurally).
1417pub(crate) struct RealizedMeasureJetGeometry {
1418    pub(crate) centers: Array2<f64>,
1419    pub(crate) masses: Array1<f64>,
1420    pub(crate) eps_band: Vec<f64>,
1421    pub(crate) log_step: f64,
1422    pub(crate) length_scale: f64,
1423    /// Assembly order for the energy weights: the realized default in
1424    /// per-level mode (absorbed per candidate by normalization), the
1425    /// explicit value in fused mode.
1426    pub(crate) order_s_eval: f64,
1427    /// Spectral-split mode marker (`order_s == 0.0` sentinel).
1428    pub(crate) per_level: bool,
1429    pub(crate) z: Array2<f64>,
1430    pub(crate) coefficient_gauge: gam_problem::Gauge,
1431    pub(crate) kz: Array2<f64>,
1432    /// Ambient-linear head lift `T` (d × head_rank): the extrapolation
1433    /// null-space basis appended to the representer design (#1845). The head
1434    /// columns evaluate as `points · T`; empty (`d × 0`) when the geometry
1435    /// resolves no supported linear direction. Deterministic in the frozen
1436    /// centers + masses, so predict-time replay rebuilds it verbatim.
1437    pub(crate) head_transform: Array2<f64>,
1438}
1439
1440pub(crate) fn realize_measure_jet_geometry(
1441    data: ArrayView2<'_, f64>,
1442    spec: &MeasureJetBasisSpec,
1443) -> Result<RealizedMeasureJetGeometry, BasisError> {
1444    if data.ncols() == 0 {
1445        crate::bail_invalid_basis!("measure-jet smooth needs at least one feature column");
1446    }
1447    validate_finite_points(data, "data")?;
1448    let seed_centers = select_centers_by_strategy(data, &spec.center_strategy)?;
1449    let m = seed_centers.nrows();
1450    if m < 3 {
1451        return Err(BasisError::InsufficientColumnsForConstraint { found: m });
1452    }
1453    let order_s = if spec.order_s == 0.0 {
1454        MEASURE_JET_DEFAULT_ORDER_S
1455    } else {
1456        spec.order_s
1457    };
1458    // Quadrature realization. Fit path: the realized nodes are the cell
1459    // BARYCENTERS of the seed partition (first-moment-exact lumping of μ —
1460    // see `measure_jet_quadrature_nodes`), so the metadata's `centers` are
1461    // already the realized nodes and the frozen path (predict / ψ-trial,
1462    // `CenterStrategy::UserProvided`) replays them verbatim with the frozen
1463    // masses, band, support anchors, and normalization scales.
1464    let (centers, masses, eps_band, log_step) = match &spec.frozen_quadrature {
1465        Some(frozen) => {
1466            if frozen.masses.len() != m {
1467                crate::bail_dim_basis!(
1468                    "frozen measure-jet quadrature mismatch: {} masses for {} centers",
1469                    frozen.masses.len(),
1470                    m
1471                );
1472            }
1473            if frozen.eps_band.is_empty() {
1474                crate::bail_invalid_basis!("frozen measure-jet quadrature has an empty band");
1475            }
1476            let log_step = if frozen.eps_band.len() >= 2 {
1477                (frozen.eps_band[1] / frozen.eps_band[0]).ln()
1478            } else {
1479                std::f64::consts::LN_2
1480            };
1481            (
1482                seed_centers,
1483                frozen.masses.clone(),
1484                frozen.eps_band.clone(),
1485                log_step,
1486            )
1487        }
1488        None => {
1489            let (nodes, masses) = measure_jet_quadrature_nodes(data, seed_centers.view())?;
1490            let band = measure_jet_band(nodes.view(), spec.num_scales)?;
1491            (nodes, masses, band.eps, band.log_step)
1492        }
1493    };
1494    let length_scale = realized_measure_jet_length_scale(centers.view(), spec.length_scale)?;
1495    // Ambient-linear extrapolation head (#1845): the raw center space becomes
1496    // `[ m Gaussian representers | head_rank ambient-linear columns ]`. The head
1497    // carries the penalty's affine null space explicitly so the fit no longer
1498    // reverts to the parametric backbone (the training mean) across an
1499    // unsupported gap.
1500    // The extrapolation head is the single-scale (fused) gap-bridge path. In
1501    // multiscale mode the per-scale spectral penalties carry their own
1502    // structure and the design stays the pure representer basis (the per-level
1503    // replay + width contracts pin `m − 1` columns), so the head is added only
1504    // when the term is single-scale.
1505    let head_transform = if spec.multiscale {
1506        Array2::<f64>::zeros((centers.ncols(), 0))
1507    } else {
1508        measure_jet_affine_head_transform(centers.view(), masses.view())
1509    };
1510    let head_rank = head_transform.ncols();
1511    let m_aug = m + head_rank;
1512    let k_cc = measure_jet_design_matrix(centers.view(), centers.view(), length_scale)?;
1513    let head_cc = centers.dot(&head_transform);
1514    // Realized-design constraint transform. In single-scale mode the explicit
1515    // affine head and Gaussian representers can otherwise carry the same affine
1516    // CENTER values in two different ways. That is a genuine gauge redundancy,
1517    // not a reason to ridge either coefficient block. At fit time remove it
1518    // exactly by restricting the RBF center values to the mass-orthogonal
1519    // complement of the supported affine space:
1520    //
1521    //   C = A^T W K_cc,       Z_rbf = null(C).
1522    //
1523    // The head then passes through as an identity block. The frozen composed
1524    // `z · z_parametric` is replayed verbatim at prediction/ψ trials (#532), so
1525    // the rank-revealed section never changes after fit-time realization. In
1526    // multiscale mode there is no explicit head, hence no affine duplication;
1527    // retain the existing representer sum-to-zero section there.
1528    let (z, coefficient_gauge) = match &spec.identifiability {
1529        MeasureJetIdentifiability::FrozenTransform { transform } => {
1530            if transform.nrows() != m_aug {
1531                crate::bail_dim_basis!(
1532                    "frozen measure-jet identifiability transform mismatch: {} representers + {} head columns but transform has {} rows",
1533                    m,
1534                    head_rank,
1535                    transform.nrows()
1536                );
1537            }
1538            (
1539                transform.clone(),
1540                gam_problem::Gauge::from_block_transforms(&[transform.clone()]),
1541            )
1542        }
1543        MeasureJetIdentifiability::CenterSumToZero => {
1544            let z_rbf = if head_rank > 0 {
1545                let affine = measure_jet_affine_value_basis(centers.view(), masses.view());
1546                let mut weighted_affine = affine.clone();
1547                for (i, mut row) in weighted_affine.outer_iter_mut().enumerate() {
1548                    row.mapv_inplace(|v| v * masses[i]);
1549                }
1550                // `rrqr_nullspace_basis(B)` returns null(B^T). Here
1551                // `B = K_cc^T W A = C^T`, hence the returned columns span
1552                // null(C), exactly the required RBF coefficient section.
1553                let constraint_cross = k_cc.t().dot(&weighted_affine);
1554                rrqr_nullspace_basis(&constraint_cross, default_rrqr_rank_alpha())
1555                    .map_err(BasisError::LinalgError)?
1556                    .0
1557            } else {
1558                let u = householder_sum_to_zero_u(m);
1559                householder_sum_to_zero_z(&u)
1560            };
1561            let rbf_rank = z_rbf.ncols();
1562            let mut z_block = Array2::<f64>::zeros((m_aug, rbf_rank + head_rank));
1563            z_block
1564                .slice_mut(ndarray::s![..m, ..rbf_rank])
1565                .assign(&z_rbf);
1566            for r in 0..head_rank {
1567                z_block[(m + r, rbf_rank + r)] = 1.0;
1568            }
1569            (
1570                z_block.clone(),
1571                gam_problem::Gauge::from_block_transforms(&[z_block]),
1572            )
1573        }
1574    };
1575    // Augmented raw center matrix `[K(centers, centers) | centers · T]`, so the
1576    // restricted `kz` maps constrained coefficients to center nodal values for
1577    // BOTH the representers and the head; the energy annihilates the head block
1578    // (affine) to machine precision, so it stays the unpenalized null space.
1579    let mut k_aug = Array2::<f64>::zeros((m, m_aug));
1580    k_aug.slice_mut(ndarray::s![.., ..m]).assign(&k_cc);
1581    if head_rank > 0 {
1582        k_aug.slice_mut(ndarray::s![.., m..]).assign(&head_cc);
1583    }
1584    let kz = coefficient_gauge.restrict_design(&k_aug);
1585    Ok(RealizedMeasureJetGeometry {
1586        centers,
1587        masses,
1588        eps_band,
1589        log_step,
1590        length_scale,
1591        order_s_eval: order_s,
1592        // Multiscale (per-scale spectral) energy is an EXPLICIT opt-in (#1116):
1593        // one Primary energy at any center count unless the spec asks for the
1594        // scale split. The independent null-component candidate is orthogonal
1595        // to this mode decision. No center-count auto-gate.
1596        per_level: spec.multiscale,
1597        z,
1598        coefficient_gauge,
1599        kz,
1600        head_transform,
1601    })
1602}
1603
1604/// Estimate the ambient input-measurement-error scale `σ_coord` — the
1605/// perpendicular off-manifold residual spread of the empirical measure — for
1606/// the errors-in-variables predictive-variance term `Var_input = ∇f̂ᵀΣ_x∇f̂`,
1607/// `Σ_x = σ_coord²·I` (issue #2225).
1608///
1609/// The measure-jet models data concentrated near an unknown low-intrinsic-
1610/// dimension set sampled with isotropic ambient coordinate noise. In a
1611/// neighborhood the set is locally affine, so the noise lives in the ambient
1612/// directions ORTHOGONAL to the local tangent — exactly the smallest principal
1613/// directions of the local data covariance. This is the standard local-PCA
1614/// noise floor: for each center's nearest-assignment cell with enough points to
1615/// span a tangent (`≥ d + 1`, the linear-algebra rank requirement — not a tuned
1616/// knob), the smallest eigenvalue of the cell-local covariance estimates the
1617/// perpendicular variance `σ_coord²`; averaging over cells (weighted by the
1618/// cell count) pools the estimate. No response values, no smoothing dial, and
1619/// no magic constant enter — it is a pure function of the ambient point cloud
1620/// and the frozen centers, in the centers' (standardized) coordinate frame.
1621///
1622/// Returns `None` when no cell can span a tangent (e.g. `d`-dimensional data
1623/// with fewer than `d + 1` points per cell, or a full-dimensional stratum with
1624/// no separable perpendicular direction) — the caller then leaves `Var_input`
1625/// disabled rather than invent a scale.
1626pub fn measure_jet_input_noise_scale(
1627    data: ArrayView2<'_, f64>,
1628    centers: ArrayView2<'_, f64>,
1629) -> Result<Option<f64>, BasisError> {
1630    let d = data.ncols();
1631    let m = centers.nrows();
1632    if d == 0 || m == 0 || data.nrows() == 0 {
1633        return Ok(None);
1634    }
1635    if centers.ncols() != d {
1636        crate::bail_dim_basis!(
1637            "measure-jet input-noise estimate: data d={d} disagrees with centers d={}",
1638            centers.ncols()
1639        );
1640    }
1641    validate_finite_points(data, "data")?;
1642    validate_finite_points(centers, "centers")?;
1643    // Nearest-center assignment (the same rule that lumps the quadrature
1644    // masses): the squared-distance Gram, argmin per row.
1645    let sq = pairwise_sq_dists(data, centers);
1646    let mut members: Vec<Vec<usize>> = vec![Vec::new(); m];
1647    for (j, row) in sq.axis_iter(Axis(0)).enumerate() {
1648        let mut best = 0usize;
1649        let mut best_d = f64::INFINITY;
1650        for (i, &dij) in row.iter().enumerate() {
1651            if dij < best_d {
1652                best_d = dij;
1653                best = i;
1654            }
1655        }
1656        members[best].push(j);
1657    }
1658    let mut weighted_sum = 0.0_f64;
1659    let mut weight = 0.0_f64;
1660    for cell in &members {
1661        let n_i = cell.len();
1662        // A cell needs at least d + 1 points to define a full-rank local
1663        // covariance; otherwise its smallest eigenvalue is a spurious zero.
1664        if n_i < d + 1 {
1665            continue;
1666        }
1667        // Cell-local mean and covariance in ambient coordinates.
1668        let mut mean = Array1::<f64>::zeros(d);
1669        for &j in cell {
1670            mean += &data.row(j);
1671        }
1672        mean /= n_i as f64;
1673        let mut cov = Array2::<f64>::zeros((d, d));
1674        for &j in cell {
1675            let mut centered = data.row(j).to_owned();
1676            centered -= &mean;
1677            for a in 0..d {
1678                for b in 0..d {
1679                    cov[(a, b)] += centered[a] * centered[b];
1680                }
1681            }
1682        }
1683        cov /= n_i as f64;
1684        // Symmetrize against accumulation asymmetry, then read the smallest
1685        // eigenvalue = the perpendicular (noise) principal variance.
1686        let cov_sym = (&cov + &cov.t()) * 0.5;
1687        let (evals, _) = cov_sym.eigh(Side::Lower).map_err(|e| {
1688            BasisError::InvalidInput(format!(
1689                "measure-jet input-noise estimate: local covariance eigendecomposition failed: {e}"
1690            ))
1691        })?;
1692        let smallest = evals
1693            .iter()
1694            .copied()
1695            .fold(f64::INFINITY, |acc, v| acc.min(v))
1696            .max(0.0);
1697        if smallest.is_finite() {
1698            weighted_sum += n_i as f64 * smallest;
1699            weight += n_i as f64;
1700        }
1701    }
1702    if weight <= 0.0 {
1703        return Ok(None);
1704    }
1705    let sigma2 = weighted_sum / weight;
1706    if !(sigma2.is_finite() && sigma2 > 0.0) {
1707        return Ok(None);
1708    }
1709    Ok(Some(sigma2.sqrt()))
1710}
1711
1712/// Whether a measure-jet spec runs in multiscale mode (per-scale spectral
1713/// energies + `(α, ln τ)` ψ dials). The separate `double_penalty`
1714/// affine/null-component candidate is available in both modes. This is the
1715/// single source of truth shared by the builder and outer enrollment predicates,
1716/// so the energy layout and ψ dimension cannot disagree. Multiscale is an
1717/// explicit opt-in (`spec.multiscale`); there is no center-count auto-gate
1718/// (#1116).
1719pub fn measure_jet_multiscale_mode(spec: &MeasureJetBasisSpec) -> bool {
1720    spec.multiscale
1721}
1722
1723/// Build the measure-jet smooth: Gaussian representer design `K(data,
1724/// centers)·z`, multiscale jet-residual penalty (one candidate per scale in
1725/// spectral mode, one Primary in pinned-order mode), an optional separate
1726/// function-space null-component candidate, and the replayable
1727/// [`BasisMetadata::MeasureJet`]. The geometry comes from the
1728/// empirical measure (centers + masses + band) through the shared
1729/// realization helper — the same source the ψ-derivative producer uses.
1730pub fn build_measure_jet_basis(
1731    data: ArrayView2<'_, f64>,
1732    spec: &MeasureJetBasisSpec,
1733) -> Result<BasisBuildResult, BasisError> {
1734    let RealizedMeasureJetGeometry {
1735        centers,
1736        masses,
1737        eps_band,
1738        log_step,
1739        length_scale,
1740        order_s_eval: order_s,
1741        per_level,
1742        z,
1743        coefficient_gauge,
1744        kz,
1745        head_transform,
1746    } = realize_measure_jet_geometry(data, spec)?;
1747    let band = MeasureJetBand {
1748        eps: eps_band.clone(),
1749        log_step,
1750    };
1751    let m = centers.nrows();
1752    let head_rank = head_transform.ncols();
1753    let m_aug = m + head_rank;
1754    // Augmented raw design `[K(data, centers) | data · T]` (#1845): the head
1755    // columns are the ambient-linear extrapolation basis. The gauge restricts
1756    // BOTH blocks together, so the frozen composed transform replays the head
1757    // verbatim at predict time.
1758    let kernel_design = measure_jet_design_matrix(data, centers.view(), length_scale)?;
1759    let mut raw_design = Array2::<f64>::zeros((data.nrows(), m_aug));
1760    raw_design
1761        .slice_mut(ndarray::s![.., ..m])
1762        .assign(&kernel_design);
1763    if head_rank > 0 {
1764        let head_design = data.dot(&head_transform);
1765        raw_design
1766            .slice_mut(ndarray::s![.., m..])
1767            .assign(&head_design);
1768    }
1769    let constrained_design = coefficient_gauge.restrict_design(&raw_design);
1770    let design = gam_linalg::matrix::DesignMatrix::Dense(
1771        gam_linalg::matrix::DenseDesignMatrix::from(constrained_design),
1772    );
1773    let support_means = measure_jet_support_means(centers.view(), masses.view(), &eps_band)?;
1774    // Spectral/geometric split. With the auto order sentinel (order_s == 0.0)
1775    // the term emits one candidate PER scale: the multi-penalty REML engine
1776    // then learns the level amplitudes λ_ℓ directly — scale adaptivity at
1777    // ρ-speed, dead scales REML-deselected (the Duchon-ARD pattern) — and the
1778    // fitted order is read off the spectrum (ŝ = −½ · slope of ln λ̂_ℓ on
1779    // ln ε_ℓ) instead of being optimized. An explicit s > 0 pins the Mellin
1780    // weights and fuses the band into one candidate. The Mellin prefactor
1781    // ε^(−η)·log_step inside each per-scale form is absorbed by the
1782    // per-candidate Frobenius normalization, so REML owns the amplitudes
1783    // outright. The sentinel itself is persisted in the metadata as the mode
1784    // marker: a replay MUST re-enter the same mode or the penalty count
1785    // desyncs (the gam#860 trap class).
1786    let mut candidates = Vec::new();
1787    let mut penalty_normalization_scales = Vec::new();
1788    let mut raw_penalty_normalization_scales = Vec::new();
1789    let mut fused_penalty_normalization_scale = None;
1790    if per_level {
1791        let forms = measure_jet_energy_forms_per_scale(
1792            centers.view(),
1793            masses.view(),
1794            &band,
1795            order_s,
1796            spec.alpha,
1797            spec.tau0,
1798        )?;
1799        for (level, q_l) in forms.into_iter().enumerate() {
1800            let s_l = kz.t().dot(&q_l).dot(&kz);
1801            let (s_norm, c_l) = normalize_penalty(&((&s_l + &s_l.t()) * 0.5));
1802            let intrinsic_dim = centers.ncols() as f64;
1803            let eta = 2.0 * order_s + intrinsic_dim * (2.0 - 2.0 * spec.alpha);
1804            let scale_weight = log_step * eps_band[level].powf(-eta);
1805            penalty_normalization_scales.push(c_l);
1806            raw_penalty_normalization_scales.push(c_l / scale_weight);
1807            candidates.push(PenaltyCandidate {
1808                matrix: s_norm,
1809                nullspace_dim_hint: 0,
1810                source: PenaltySource::Other(format!("measure_jet_scale_{level}")),
1811                normalization_scale: c_l,
1812                kronecker_factors: None,
1813                op: None,
1814            });
1815        }
1816    } else {
1817        let q_form = measure_jet_energy_form(
1818            centers.view(),
1819            masses.view(),
1820            &band,
1821            order_s,
1822            spec.alpha,
1823            spec.tau0,
1824        )?;
1825        // The Primary is exactly the jet-energy functional pulled back through
1826        // the center evaluation map. It is independent of `double_penalty`:
1827        // statistical selection is a distinct REML component below, never a
1828        // fixed coefficient toll fused into this estimand.
1829        let penalty = pullback_center_form(&kz, &q_form);
1830        let (penalty_norm, c_primary) = normalize_penalty(&penalty);
1831        fused_penalty_normalization_scale = Some(c_primary);
1832        candidates.push(PenaltyCandidate {
1833            matrix: penalty_norm,
1834            nullspace_dim_hint: 0,
1835            source: PenaltySource::Primary,
1836            normalization_scale: c_primary,
1837            kronecker_factors: None,
1838            op: None,
1839        });
1840    }
1841    // Explicit null recovery is a genuine statistical component: penalize the
1842    // affine/null FUNCTION projection under the empirical-measure mass metric,
1843    // and let REML select its strength independently in both modes. This is the
1844    // standard double-penalty decomposition (roughness + null component); no
1845    // coefficient identity and no hard-coded mixture changes the Primary.
1846    if spec.double_penalty {
1847        let null_penalty = affine_function_nullspace_penalty(&kz, centers.view(), masses.view())?;
1848        let (null_penalty_norm, c_null) = normalize_penalty(&null_penalty);
1849        candidates.push(PenaltyCandidate {
1850            matrix: null_penalty_norm,
1851            nullspace_dim_hint: 0,
1852            source: PenaltySource::DoublePenaltyNullspace,
1853            normalization_scale: c_null,
1854            kronecker_factors: None,
1855            op: None,
1856        });
1857    }
1858    let (penalties, nullspace_dims, penaltyinfo, null_eigenvectors, ops) =
1859        filter_active_penalty_candidates_with_ops(candidates)?;
1860    // #2225: compute the errors-in-variables input-noise scale while `centers`
1861    // is still owned; it is moved into the metadata `centers` field below.
1862    let sigma_coord = measure_jet_input_noise_scale(data, centers.view())?;
1863    Ok(BasisBuildResult {
1864        design,
1865        penalties,
1866        nullspace_dims,
1867        penaltyinfo,
1868        metadata: BasisMetadata::MeasureJet {
1869            centers,
1870            input_scales: None,
1871            length_scale,
1872            eps_band,
1873            // The SPEC's order field, sentinel included: 0.0 marks per-level
1874            // (spectral) mode and must replay as per-level — persisting the
1875            // realized default here would silently flip the rebuild into
1876            // fused mode and desync the penalty count.
1877            order_s: spec.order_s,
1878            alpha: spec.alpha,
1879            tau0: spec.tau0,
1880            masses,
1881            support_means,
1882            penalty_normalization_scales,
1883            raw_penalty_normalization_scales,
1884            fused_penalty_normalization_scale,
1885            constraint_transform: Some(z),
1886            // Perpendicular off-manifold residual scale of the fit rows in the
1887            // centers' frame — the errors-in-variables input-noise scale (#2225).
1888            sigma_coord,
1889        },
1890        kronecker_factored: None,
1891        ops,
1892        null_eigenvectors,
1893        joint_null_rotation: None,
1894    })
1895}
1896
1897/// Exact ψ-jets of the REALIZED measure-jet penalty candidates, adapted to
1898/// the anisotropic group-ψ carrier the spatial optimizer consumes.
1899///
1900/// Coordinates (the layout contract for the registration arm):
1901/// - per-level (spectral) mode: `[ln ℓ?, α, ln τ]` — order is absorbed by the
1902///   REML-learned scale amplitudes; `ln τ` is retained as an inert coordinate;
1903/// - single-scale mode: `[ln ℓ?]`, because its energy dials are fixed.
1904///
1905/// Only `ln ℓ` moves the design. It also moves every coefficient-space penalty
1906/// pullback through the center evaluation map `E(ℓ)`; `(α, ln τ)` move only the
1907/// per-scale center-value forms. Exact diagonal and mixed product-rule jets are
1908/// emitted before Frobenius normalization.
1909/// Penalty derivatives are routed through the SAME constrained Frobenius
1910/// normalization as the fit-time candidates
1911/// (`normalize_penaltywith_psi_derivatives` + the cross rule), so criterion
1912/// value and criterion derivative share one normalization — the #901 lesson
1913/// made structural. The function-space null candidate has nonzero `ln ℓ` jets
1914/// and zero `(α, ln τ)` jets. The per-candidate layout follows the builder's
1915/// ORIGINAL order (scale candidates or Primary, then null component); consumers
1916/// align to the FITTED penalty list via
1917/// `PenaltyInfo.original_index` when the active-candidate filter dropped
1918/// any.
1919pub fn build_measure_jet_basis_psi_derivatives(
1920    data: ArrayView2<'_, f64>,
1921    spec: &MeasureJetBasisSpec,
1922) -> Result<AnisoBasisPsiDerivatives, BasisError> {
1923    if !(spec.tau0.is_finite() && spec.tau0 > 0.0) {
1924        crate::bail_invalid_basis!(
1925            "measure-jet ψ derivatives need tau0 > 0 because the retained τ coordinate is ln τ; got {}",
1926            spec.tau0
1927        );
1928    }
1929    let geom = realize_measure_jet_geometry(data, spec)?;
1930    let band = MeasureJetBand {
1931        eps: geom.eps_band.clone(),
1932        log_step: geom.log_step,
1933    };
1934    let n = data.nrows();
1935    let p = geom.kz.ncols();
1936    let m = geom.centers.nrows();
1937    let m_aug = m + geom.head_transform.ncols();
1938
1939    struct LengthScaleJets {
1940        evaluation_first: Array2<f64>,
1941        evaluation_second: Array2<f64>,
1942        design_first: Array2<f64>,
1943        design_second: Array2<f64>,
1944    }
1945
1946    // The Gaussian representer range moves both the FIT design and the center
1947    // evaluation map `E = [K_cc | A_head] Z`. The affine head is ℓ-invariant,
1948    // so its raw derivative columns are exactly zero before applying the frozen
1949    // Gauge section. Keeping `Z` frozen is the replay contract: rank/gauge
1950    // realization happens once at fit time, then every ψ trial differentiates
1951    // the same coefficient chart.
1952    let length_scale_jets = if spec.learn_length_scale {
1953        let (dk_data, d2k_data) =
1954            measure_jet_design_log_length_jets(data, geom.centers.view(), geom.length_scale)?;
1955        let mut dk_data_aug = Array2::<f64>::zeros((n, m_aug));
1956        let mut d2k_data_aug = Array2::<f64>::zeros((n, m_aug));
1957        dk_data_aug.slice_mut(ndarray::s![.., ..m]).assign(&dk_data);
1958        d2k_data_aug
1959            .slice_mut(ndarray::s![.., ..m])
1960            .assign(&d2k_data);
1961
1962        let (dk_centers, d2k_centers) = measure_jet_design_log_length_jets(
1963            geom.centers.view(),
1964            geom.centers.view(),
1965            geom.length_scale,
1966        )?;
1967        let mut dk_centers_aug = Array2::<f64>::zeros((m, m_aug));
1968        let mut d2k_centers_aug = Array2::<f64>::zeros((m, m_aug));
1969        dk_centers_aug
1970            .slice_mut(ndarray::s![.., ..m])
1971            .assign(&dk_centers);
1972        d2k_centers_aug
1973            .slice_mut(ndarray::s![.., ..m])
1974            .assign(&d2k_centers);
1975
1976        Some(LengthScaleJets {
1977            evaluation_first: geom.coefficient_gauge.restrict_design(&dk_centers_aug),
1978            evaluation_second: geom.coefficient_gauge.restrict_design(&d2k_centers_aug),
1979            design_first: geom.coefficient_gauge.restrict_design(&dk_data_aug),
1980            design_second: geom.coefficient_gauge.restrict_design(&d2k_data_aug),
1981        })
1982    } else {
1983        None
1984    };
1985
1986    let coord_offset = usize::from(length_scale_jets.is_some());
1987    let n_coords = coord_offset + if geom.per_level { 2 } else { 0 };
1988    let pairs: Vec<(usize, usize)> = (0..n_coords)
1989        .flat_map(|a| ((a + 1)..n_coords).map(move |b| (a, b)))
1990        .collect();
1991    let zero_p = || Array2::<f64>::zeros((p, p));
1992
1993    struct RawPenaltyJets {
1994        value: Array2<f64>,
1995        first: Vec<Array2<f64>>,
1996        second_diag: Vec<Array2<f64>>,
1997        cross: Vec<Array2<f64>>,
1998    }
1999
2000    let sandwich = |form: &Array2<f64>| pullback_center_form(&geom.kz, form);
2001    let length_diag = |form: &Array2<f64>| {
2002        let jets = length_scale_jets
2003            .as_ref()
2004            .expect("length-scale form jets require an enrolled length coordinate");
2005        pullback_center_form_log_length_jets(
2006            &geom.kz,
2007            &jets.evaluation_first,
2008            &jets.evaluation_second,
2009            form,
2010        )
2011    };
2012    let length_cross = |form_first: &Array2<f64>| {
2013        let jets = length_scale_jets
2014            .as_ref()
2015            .expect("length-scale cross jets require an enrolled length coordinate");
2016        pullback_center_form_log_length_cross(&geom.kz, &jets.evaluation_first, form_first)
2017    };
2018
2019    // Raw (pre-normalization) value + exact jet stacks per ORIGINAL candidate.
2020    // Coordinate order is `[lnℓ?, α, lnτ]` in multiscale mode and `[lnℓ?]`
2021    // in single-scale mode. Candidate order exactly mirrors the value builder:
2022    // scale candidates or Primary first, then the optional null-component
2023    // candidate. Active filtering aligns through `PenaltyInfo::original_index`.
2024    let mut raw: Vec<RawPenaltyJets> = if geom.per_level {
2025        let l_count = band.eps.len();
2026        // Six forms per scale: value, ∂α, ∂α², and zero τ slots — same
2027        // blocks, one walk (single-source rule).
2028        let forms = assemble_weighted_forms(
2029            geom.centers.view(),
2030            geom.masses.view(),
2031            &band,
2032            geom.order_s_eval,
2033            spec.alpha,
2034            spec.tau0,
2035            6 * l_count,
2036            3,
2037            &|scale_idx, eps: f64, q: f64, base: f64, out: &mut [[f64; 3]]| {
2038                for slot in out.iter_mut() {
2039                    *slot = [0.0, 0.0, 0.0];
2040                }
2041                let intrinsic_dim = geom.centers.ncols() as f64;
2042                let ga = 2.0 * intrinsic_dim * eps.ln() - 2.0 * q.max(f64::MIN_POSITIVE).ln();
2043                let k0 = 6 * scale_idx;
2044                out[k0] = [base, 0.0, 0.0];
2045                out[k0 + 1] = [ga * base, 0.0, 0.0];
2046                out[k0 + 2] = [ga * ga * base, 0.0, 0.0];
2047                out[k0 + 3] = [0.0, 0.0, 0.0];
2048                out[k0 + 4] = [0.0, 0.0, 0.0];
2049                out[k0 + 5] = [0.0, 0.0, 0.0];
2050            },
2051        )?;
2052        let alpha_coord = coord_offset;
2053        let tau_coord = coord_offset + 1;
2054        let mut raw = Vec::with_capacity(l_count + usize::from(spec.double_penalty));
2055        for level in 0..l_count {
2056            let chunk = &forms[6 * level..6 * level + 6];
2057            let mut first: Vec<Array2<f64>> = (0..n_coords).map(|_| zero_p()).collect();
2058            let mut second_diag: Vec<Array2<f64>> = (0..n_coords).map(|_| zero_p()).collect();
2059            first[alpha_coord] = sandwich(&chunk[1]);
2060            first[tau_coord] = sandwich(&chunk[3]);
2061            second_diag[alpha_coord] = sandwich(&chunk[2]);
2062            second_diag[tau_coord] = sandwich(&chunk[4]);
2063            if coord_offset == 1 {
2064                let (ell_first, ell_second) = length_diag(&chunk[0]);
2065                first[0] = ell_first;
2066                second_diag[0] = ell_second;
2067            }
2068            let mut cross: Vec<Array2<f64>> = (0..pairs.len()).map(|_| zero_p()).collect();
2069            for (pair_idx, &(a, b)) in pairs.iter().enumerate() {
2070                cross[pair_idx] = if coord_offset == 1 && a == 0 && b == alpha_coord {
2071                    length_cross(&chunk[1])
2072                } else if coord_offset == 1 && a == 0 && b == tau_coord {
2073                    length_cross(&chunk[3])
2074                } else if a == alpha_coord && b == tau_coord {
2075                    sandwich(&chunk[5])
2076                } else {
2077                    zero_p()
2078                };
2079            }
2080            raw.push(RawPenaltyJets {
2081                value: sandwich(&chunk[0]),
2082                first,
2083                second_diag,
2084                cross,
2085            });
2086        }
2087        raw
2088    } else {
2089        // Single-scale mode enrolls no `(s, α, lnτ)` penalty dials. It still
2090        // emits the pure Primary and, when requested, a separate REML null
2091        // component; an opt-in `lnℓ` coordinate differentiates both pullbacks.
2092        let q_form = measure_jet_energy_form(
2093            geom.centers.view(),
2094            geom.masses.view(),
2095            &band,
2096            geom.order_s_eval,
2097            spec.alpha,
2098            spec.tau0,
2099        )?;
2100        let mut first: Vec<Array2<f64>> = (0..n_coords).map(|_| zero_p()).collect();
2101        let mut second_diag: Vec<Array2<f64>> = (0..n_coords).map(|_| zero_p()).collect();
2102        if coord_offset == 1 {
2103            let (ell_first, ell_second) = length_diag(&q_form);
2104            first[0] = ell_first;
2105            second_diag[0] = ell_second;
2106        }
2107        vec![RawPenaltyJets {
2108            value: sandwich(&q_form),
2109            first,
2110            second_diag,
2111            cross: Vec::new(),
2112        }]
2113    };
2114
2115    if spec.double_penalty {
2116        let null_form = affine_function_nullspace_form(geom.centers.view(), geom.masses.view())?;
2117        let mut first: Vec<Array2<f64>> = (0..n_coords).map(|_| zero_p()).collect();
2118        let mut second_diag: Vec<Array2<f64>> = (0..n_coords).map(|_| zero_p()).collect();
2119        if coord_offset == 1 {
2120            let (ell_first, ell_second) = length_diag(&null_form);
2121            first[0] = ell_first;
2122            second_diag[0] = ell_second;
2123        }
2124        raw.push(RawPenaltyJets {
2125            value: sandwich(&null_form),
2126            first,
2127            second_diag,
2128            // H₀ is independent of α and τ; its only moving object is E(ℓ),
2129            // so every mixed coordinate derivative is zero.
2130            cross: (0..pairs.len()).map(|_| zero_p()).collect(),
2131        });
2132    }
2133
2134    let n_cands = raw.len();
2135    let mut penalties_first: Vec<Vec<Array2<f64>>> =
2136        (0..n_coords).map(|_| Vec::with_capacity(n_cands)).collect();
2137    let mut penalties_second_diag: Vec<Vec<Array2<f64>>> =
2138        (0..n_coords).map(|_| Vec::with_capacity(n_cands)).collect();
2139    // Cross matrices per pair per candidate, precomputed eagerly (the
2140    // candidate count is the band length, not the data size) and served
2141    // through the on-demand provider.
2142    let mut crosses: Vec<Vec<Array2<f64>>> = (0..pairs.len()).map(|_| Vec::new()).collect();
2143    for candidate in &raw {
2144        let s_raw = &candidate.value;
2145        // ONE Frobenius scale per candidate, fixed up front from `s_raw`
2146        // alone: c anchors the value and every derivative of this candidate.
2147        // `normalize_penaltywith_psi_derivatives` recomputes the identical c
2148        // per coordinate (same trace_of_product + sqrt on the same `s_raw`),
2149        // and its degenerate convention is mirrored here: ‖S‖_F ≤ 1e-12 (or
2150        // non-finite) reports scale 1.0 — the value passes through unscaled,
2151        // and the cross helper receives that same 1.0, never a collapsed
2152        // near-zero scale.
2153        let fro = trace_of_product(s_raw, s_raw).sqrt();
2154        let c = if fro.is_finite() && fro > 1e-12 {
2155            fro
2156        } else {
2157            1.0
2158        };
2159        for coord in 0..n_coords {
2160            let (_, s_first, s_second, _) = normalize_penaltywith_psi_derivatives(
2161                s_raw,
2162                &candidate.first[coord],
2163                &candidate.second_diag[coord],
2164            );
2165            penalties_first[coord].push(s_first);
2166            penalties_second_diag[coord].push(s_second);
2167        }
2168        for (pair_idx, &(a, b)) in pairs.iter().enumerate() {
2169            let cross_raw_mat = normalize_penalty_cross_psi_derivative(
2170                s_raw,
2171                &candidate.first[a],
2172                &candidate.first[b],
2173                &candidate.cross[pair_idx],
2174                c,
2175            );
2176            crosses[pair_idx].push(cross_raw_mat);
2177        }
2178    }
2179
2180    let pair_index: Vec<((usize, usize), Vec<Array2<f64>>)> =
2181        pairs.iter().copied().zip(crosses.into_iter()).collect();
2182    let provider = AnisoPenaltyCrossProvider::new(move |a, b| {
2183        pair_index
2184            .iter()
2185            .find(|((pa, pb), _)| (*pa, *pb) == (a, b) || (*pa, *pb) == (b, a))
2186            .map(|(_, mats)| mats.clone())
2187            .ok_or_else(|| {
2188                BasisError::InvalidInput(format!(
2189                    "measure-jet ψ cross derivative requested for unknown pair ({a}, {b})"
2190                ))
2191            })
2192    });
2193    let mut design_first: Vec<Array2<f64>> = (0..n_coords)
2194        .map(|_| Array2::<f64>::zeros((n, p)))
2195        .collect();
2196    let mut design_second_diag: Vec<Array2<f64>> = (0..n_coords)
2197        .map(|_| Array2::<f64>::zeros((n, p)))
2198        .collect();
2199    if let Some(jets) = &length_scale_jets {
2200        design_first[0] = jets.design_first.clone();
2201        design_second_diag[0] = jets.design_second.clone();
2202    }
2203    Ok(AnisoBasisPsiDerivatives {
2204        design_first,
2205        design_second_diag,
2206        design_second_cross: Vec::new(),
2207        design_second_cross_pairs: Vec::new(),
2208        penalties_first,
2209        penalties_second_diag,
2210        penalties_cross_pairs: pairs,
2211        penalties_cross_provider: Some(provider),
2212        implicit_operator: None,
2213    })
2214}
2215
2216#[cfg(test)]
2217mod tests {
2218    use super::*;
2219
2220    /// Deterministic Box–Muller standard normal from a 64-bit LCG state — a
2221    /// self-contained noise generator (no external RNG dependency).
2222    fn lcg_normal(state: &mut u64) -> f64 {
2223        let mut next = || {
2224            *state = state
2225                .wrapping_mul(6364136223846793005)
2226                .wrapping_add(1442695040888963407);
2227            // Top 53 bits → uniform (0, 1).
2228            (((*state >> 11) as f64) + 0.5) / (1u64 << 53) as f64
2229        };
2230        let u1 = next();
2231        let u2 = next();
2232        (-2.0 * u1.ln()).sqrt() * (std::f64::consts::TAU * u2).cos()
2233    }
2234
2235    /// The perpendicular off-manifold residual estimator recovers a KNOWN
2236    /// ambient noise scale on a 1-D manifold (a line) embedded in 2-D: points
2237    /// sampled along the tangent with isotropic-perpendicular Gaussian noise of
2238    /// scale σ, centers spaced along the line. The local-PCA smallest-eigenvalue
2239    /// floor must return ≈ σ (#2225).
2240    #[test]
2241    pub(crate) fn input_noise_scale_recovers_known_perpendicular_sigma() {
2242        // Line direction (unit) and its perpendicular in 2-D.
2243        let tang = [1.0 / 5f64.sqrt(), 2.0 / 5f64.sqrt()];
2244        let perp = [2.0 / 5f64.sqrt(), -1.0 / 5f64.sqrt()];
2245        let sigma = 0.05_f64;
2246        let n = 600usize;
2247        let mut state = 0x1234_5678_9abc_def0u64;
2248        let mut data = Array2::<f64>::zeros((n, 2));
2249        for j in 0..n {
2250            // Tangential coordinate marches deterministically over [0, 3].
2251            let t = 3.0 * (j as f64) / (n as f64 - 1.0);
2252            let noise = sigma * lcg_normal(&mut state);
2253            for a in 0..2 {
2254                data[(j, a)] = t * tang[a] + noise * perp[a];
2255            }
2256        }
2257        // Centers along the line (on the noiseless manifold): plenty of points
2258        // per cell to span the tangent.
2259        let n_centers = 8usize;
2260        let mut centers = Array2::<f64>::zeros((n_centers, 2));
2261        for i in 0..n_centers {
2262            let t = 3.0 * (i as f64 + 0.5) / (n_centers as f64);
2263            for a in 0..2 {
2264                centers[(i, a)] = t * tang[a];
2265            }
2266        }
2267        let est = measure_jet_input_noise_scale(data.view(), centers.view())
2268            .expect("estimate ok")
2269            .expect("noise scale present");
2270        // Sample smallest-eigenvalue floor is mildly downward-biased; require it
2271        // within 40% of the truth (central estimate, not a tuned tolerance).
2272        assert!(
2273            (est - sigma).abs() <= 0.4 * sigma,
2274            "estimated σ_coord {est} far from true {sigma}"
2275        );
2276    }
2277
2278    /// Too few points per cell (cannot span a d-dim tangent) ⇒ no estimate,
2279    /// so the caller leaves Var_input disabled rather than invent a scale.
2280    #[test]
2281    pub(crate) fn input_noise_scale_none_when_cells_too_small() {
2282        let data = array![[0.0, 0.0], [1.0, 2.0], [2.0, 4.0]];
2283        let centers = array![[0.0, 0.0], [1.0, 2.0], [2.0, 4.0]];
2284        // Each point is its own nearest center (1 point per cell < d + 1 = 3).
2285        assert!(
2286            measure_jet_input_noise_scale(data.view(), centers.view())
2287                .expect("estimate ok")
2288                .is_none()
2289        );
2290    }
2291
2292    pub(crate) fn two_cluster_centers() -> (ndarray::Array2<f64>, ndarray::Array1<f64>) {
2293        let centers = array![
2294            [0.00, 0.00],
2295            [0.31, 0.05],
2296            [0.58, -0.07],
2297            [0.93, 0.11],
2298            [1.22, 0.02],
2299            [1.49, -0.04],
2300            [3.10, 2.00],
2301            [3.42, 2.13],
2302            [3.71, 1.91],
2303            [4.05, 2.07],
2304            [4.33, 1.96],
2305            [4.61, 2.12],
2306        ];
2307        let m = centers.nrows();
2308        let masses = ndarray::Array1::<f64>::from_elem(m, 1.0 / m as f64);
2309        (centers, masses)
2310    }
2311    use ndarray::array;
2312
2313    pub(crate) fn band_for(centers: &Array2<f64>) -> MeasureJetBand {
2314        measure_jet_band(centers.view(), 0).expect("band")
2315    }
2316
2317    /// The no-mass contract: constants must be annihilated to machine
2318    /// precision at every scale (the constant is projected, never ridged).
2319    #[test]
2320    pub(crate) fn energy_form_annihilates_constants_exactly() {
2321        let (centers, masses) = two_cluster_centers();
2322        let band = band_for(&centers);
2323        let q = measure_jet_energy_form(centers.view(), masses.view(), &band, 1.5, 1.0, 1e-3)
2324            .expect("energy form");
2325        let m = q.nrows();
2326        let ones = Array1::<f64>::ones(m);
2327        let qv = q.dot(&ones);
2328        let scale = q.iter().fold(0.0_f64, |acc, v| acc.max(v.abs()));
2329        assert!(scale > 0.0, "energy form is identically zero");
2330        for (i, v) in qv.iter().enumerate() {
2331            assert!(
2332                v.abs() <= 1e-12 * scale,
2333                "Q·1 leak at row {i}: {v:.3e} vs scale {scale:.3e}"
2334            );
2335        }
2336        let vqv = ones.dot(&qv);
2337        assert!(
2338            vqv.abs() <= 1e-12 * scale,
2339            "constant carries energy: 1ᵀQ1 = {vqv:.3e}"
2340        );
2341    }
2342
2343    /// The default local projection annihilates ambient affine functions
2344    /// exactly; τ is retained for ψ layout but no longer adds an affine toll.
2345    #[test]
2346    pub(crate) fn energy_form_annihilates_affine_at_default_tau() {
2347        let (centers, masses) = two_cluster_centers();
2348        let band = band_for(&centers);
2349        let m = centers.nrows();
2350        // Affine values v = 0.7 + 1.3·x − 0.4·y, and a rough ±1 checkerboard.
2351        let mut affine = Array1::<f64>::zeros(m);
2352        let mut rough = Array1::<f64>::zeros(m);
2353        for i in 0..m {
2354            affine[i] = 0.7 + 1.3 * centers[(i, 0)] - 0.4 * centers[(i, 1)];
2355            rough[i] = if i % 2 == 0 { 1.0 } else { -1.0 };
2356        }
2357        let q = measure_jet_energy_form(centers.view(), masses.view(), &band, 1.5, 1.0, 1e-3)
2358            .expect("energy form");
2359        let e_affine = affine.dot(&q.dot(&affine));
2360        let e_rough = rough.dot(&q.dot(&rough));
2361        assert!(e_rough > 0.0, "rough vector must pay energy");
2362        assert!(
2363            e_affine.abs() <= 1e-12 * e_rough,
2364            "default affine energy {e_affine:.3e} vs rough {e_rough:.3e}"
2365        );
2366    }
2367
2368    /// PSD: the energy is a sum of weighted least-squares residuals.
2369    #[test]
2370    pub(crate) fn energy_form_is_psd() {
2371        let (centers, masses) = two_cluster_centers();
2372        let band = band_for(&centers);
2373        let q = measure_jet_energy_form(centers.view(), masses.view(), &band, 1.5, 1.0, 1e-3)
2374            .expect("energy form");
2375        let m = q.nrows();
2376        for trial in 0..5usize {
2377            let v = Array1::<f64>::from_shape_fn(m, |i| {
2378                ((i * 7 + trial * 13) % 11) as f64 / 11.0 - 0.5
2379            });
2380            let e = v.dot(&q.dot(&v));
2381            assert!(e >= -1e-10, "vᵀQv = {e:.3e} < 0 on trial {trial}");
2382        }
2383    }
2384
2385    /// A 1-D filament embedded in 2-D: high-frequency center values along the
2386    /// strand pay strictly more energy than a slow trend.
2387    #[test]
2388    pub(crate) fn rough_vector_pays_more_than_smooth() {
2389        let m = 24usize;
2390        let centers = Array2::<f64>::from_shape_fn((m, 2), |(i, k)| {
2391            let t = i as f64 / (m as f64 - 1.0);
2392            if k == 0 {
2393                t * 4.0
2394            } else {
2395                0.3 * (t * 4.0).sin()
2396            }
2397        });
2398        let masses = Array1::<f64>::from_elem(m, 1.0 / m as f64);
2399        let band = band_for(&centers);
2400        let q = measure_jet_energy_form(centers.view(), masses.view(), &band, 1.5, 1.0, 1e-3)
2401            .expect("energy form");
2402        let slow = Array1::<f64>::from_shape_fn(m, |i| (i as f64 / (m as f64 - 1.0)).powi(2));
2403        let fast = Array1::<f64>::from_shape_fn(m, |i| if i % 2 == 0 { 0.5 } else { -0.5 });
2404        let e_slow = slow.dot(&q.dot(&slow));
2405        let e_fast = fast.dot(&q.dot(&fast));
2406        assert!(
2407            e_fast > 10.0 * e_slow,
2408            "alternating values must pay >> a slow trend: fast {e_fast:.3e} vs slow {e_slow:.3e}"
2409        );
2410    }
2411
2412    /// The exact (s, α) jets and zero τ slots must match central finite
2413    /// differences of the energy — the FD gate the ψ-channel stage will
2414    /// inherit (the discipline whose absence is exactly the
2415    /// objective↔gradient desync bug class).
2416    #[test]
2417    pub(crate) fn energy_jets_match_finite_differences() {
2418        let (centers, masses) = two_cluster_centers();
2419        let band = band_for(&centers);
2420        let (s0, a0, tau) = (1.3, 0.8, 1e-3);
2421        let jets =
2422            measure_jet_energy_form_with_jets(centers.view(), masses.view(), &band, s0, a0, tau)
2423                .expect("jets");
2424        let q_at = |s: f64, a: f64| {
2425            measure_jet_energy_form(centers.view(), masses.view(), &band, s, a, tau)
2426                .expect("energy form")
2427        };
2428        // Base form must equal the plain assembly bit-for-bit (same walk).
2429        let q_plain = q_at(s0, a0);
2430        for (a, b) in jets.q.iter().zip(q_plain.iter()) {
2431            assert!(
2432                (a - b).abs() <= 1e-14 * (1.0 + b.abs()),
2433                "Q drift {a} vs {b}"
2434            );
2435        }
2436        let lt0 = tau.ln();
2437        let q_at_lt = |lt: f64| {
2438            measure_jet_energy_form(centers.view(), masses.view(), &band, s0, a0, lt.exp())
2439                .expect("energy form")
2440        };
2441        // FD step calibrated for the SECOND differences: their roundoff
2442        // floor is ~4·ε_f64·scale/h² (assembly noise amplified by 1/h²), so
2443        // h = 1e-4 ≈ ε^(1/4) balances it against the O(h²) truncation —
2444        // both land ≥3 orders below the unchanged 5e-5·scale gate. h = 1e-5
2445        // sits ON the roundoff floor and fails spuriously.
2446        let h = 1e-4;
2447        let checks: [(&str, &Array2<f64>, Array2<f64>); 9] = [
2448            ("dq_ds", &jets.dq_ds, {
2449                let (p, m_) = (q_at(s0 + h, a0), q_at(s0 - h, a0));
2450                (&p - &m_) / (2.0 * h)
2451            }),
2452            ("d2q_ds2", &jets.d2q_ds2, {
2453                let (p, c, m_) = (q_at(s0 + h, a0), q_at(s0, a0), q_at(s0 - h, a0));
2454                (&(&p + &m_) - &(&c * 2.0)) / (h * h)
2455            }),
2456            ("dq_dalpha", &jets.dq_dalpha, {
2457                let (p, m_) = (q_at(s0, a0 + h), q_at(s0, a0 - h));
2458                (&p - &m_) / (2.0 * h)
2459            }),
2460            ("d2q_dalpha2", &jets.d2q_dalpha2, {
2461                let (p, c, m_) = (q_at(s0, a0 + h), q_at(s0, a0), q_at(s0, a0 - h));
2462                (&(&p + &m_) - &(&c * 2.0)) / (h * h)
2463            }),
2464            ("d2q_ds_dalpha", &jets.d2q_ds_dalpha, {
2465                let pp = q_at(s0 + h, a0 + h);
2466                let pm = q_at(s0 + h, a0 - h);
2467                let mp = q_at(s0 - h, a0 + h);
2468                let mm = q_at(s0 - h, a0 - h);
2469                (&(&pp - &pm) - &(&mp - &mm)) / (4.0 * h * h)
2470            }),
2471            ("dq_dlogtau", &jets.dq_dlogtau, {
2472                let (p, m_) = (q_at_lt(lt0 + h), q_at_lt(lt0 - h));
2473                (&p - &m_) / (2.0 * h)
2474            }),
2475            ("d2q_dlogtau2", &jets.d2q_dlogtau2, {
2476                let (p, c, m_) = (q_at_lt(lt0 + h), q_at_lt(lt0), q_at_lt(lt0 - h));
2477                (&(&p + &m_) - &(&c * 2.0)) / (h * h)
2478            }),
2479            ("d2q_ds_dlogtau", &jets.d2q_ds_dlogtau, {
2480                let f = |s: f64, lt: f64| {
2481                    measure_jet_energy_form(centers.view(), masses.view(), &band, s, a0, lt.exp())
2482                        .expect("energy form")
2483                };
2484                let pp = f(s0 + h, lt0 + h);
2485                let pm = f(s0 + h, lt0 - h);
2486                let mp = f(s0 - h, lt0 + h);
2487                let mm = f(s0 - h, lt0 - h);
2488                (&(&pp - &pm) - &(&mp - &mm)) / (4.0 * h * h)
2489            }),
2490            ("d2q_dalpha_dlogtau", &jets.d2q_dalpha_dlogtau, {
2491                let f = |a: f64, lt: f64| {
2492                    measure_jet_energy_form(centers.view(), masses.view(), &band, s0, a, lt.exp())
2493                        .expect("energy form")
2494                };
2495                let pp = f(a0 + h, lt0 + h);
2496                let pm = f(a0 + h, lt0 - h);
2497                let mp = f(a0 - h, lt0 + h);
2498                let mm = f(a0 - h, lt0 - h);
2499                (&(&pp - &pm) - &(&mp - &mm)) / (4.0 * h * h)
2500            }),
2501        ];
2502        for (name, analytic, fd) in checks.iter() {
2503            let scale = fd.iter().fold(1e-30_f64, |acc, v| acc.max(v.abs()));
2504            for (a, b) in analytic.iter().zip(fd.iter()) {
2505                assert!(
2506                    (a - b).abs() <= 5e-5 * scale,
2507                    "{name} jet mismatch: analytic {a:.6e} vs FD {b:.6e} (scale {scale:.3e})"
2508                );
2509            }
2510        }
2511    }
2512
2513    /// The per-scale spectrum must sum exactly to the total energy (same
2514    /// blocks, one-hot weights) and concentrate rough content at fine
2515    /// scales.
2516    #[test]
2517    pub(crate) fn scale_spectrum_sums_to_total_and_localizes_roughness() {
2518        let m = 24usize;
2519        let centers = Array2::<f64>::from_shape_fn((m, 2), |(i, k)| {
2520            let t = i as f64 / (m as f64 - 1.0);
2521            if k == 0 { t * 4.0 } else { 0.0 }
2522        });
2523        let masses = Array1::<f64>::from_elem(m, 1.0 / m as f64);
2524        let band = band_for(&centers);
2525        let q = measure_jet_energy_form(centers.view(), masses.view(), &band, 1.5, 1.0, 1e-3)
2526            .expect("energy form");
2527        let fast = Array1::<f64>::from_shape_fn(m, |i| if i % 2 == 0 { 0.5 } else { -0.5 });
2528        let spec = measure_jet_scale_spectrum(
2529            centers.view(),
2530            masses.view(),
2531            &band,
2532            1.5,
2533            1.0,
2534            1e-3,
2535            fast.view(),
2536        )
2537        .expect("spectrum");
2538        assert_eq!(spec.len(), band.eps.len());
2539        let total = fast.dot(&q.dot(&fast));
2540        let sum: f64 = spec.iter().sum();
2541        assert!(
2542            (sum - total).abs() <= 1e-10 * total.abs().max(1e-30),
2543            "spectrum must sum to vᵀQv: {sum:.6e} vs {total:.6e}"
2544        );
2545        // Alternating-sign content lives at the finest scale of the band.
2546        let finest = spec[0];
2547        let coarsest = *spec.last().expect("nonempty spectrum");
2548        assert!(
2549            finest > coarsest,
2550            "alternating values must charge fine scales hardest: fine {finest:.3e} vs coarse {coarsest:.3e}"
2551        );
2552    }
2553
2554    /// The support curve separates on-web from off-web queries at fine
2555    /// scales and grows monotonically in ε for any query.
2556    #[test]
2557    pub(crate) fn support_curve_separates_on_web_from_off_web() {
2558        let m = 24usize;
2559        let centers = Array2::<f64>::from_shape_fn((m, 2), |(i, k)| {
2560            let t = i as f64 / (m as f64 - 1.0);
2561            if k == 0 { t * 4.0 } else { 0.0 }
2562        });
2563        let masses = Array1::<f64>::from_elem(m, 1.0 / m as f64);
2564        let band = band_for(&centers);
2565        let queries = array![[2.0, 0.0], [2.0, 1.5]];
2566        let curves =
2567            measure_jet_support_curve(queries.view(), centers.view(), masses.view(), &band.eps)
2568                .expect("support curve");
2569        // On-web sees strictly more mass than off-web at the finest scale.
2570        assert!(
2571            curves[(0, 0)] > 10.0 * curves[(1, 0)],
2572            "fine-scale support must separate web from void: on {:.3e} vs off {:.3e}",
2573            curves[(0, 0)],
2574            curves[(1, 0)]
2575        );
2576        // Kernel mass is monotone in ε for every query.
2577        for qi in 0..2 {
2578            for li in 1..band.eps.len() {
2579                assert!(
2580                    curves[(qi, li)] >= curves[(qi, li - 1)] - 1e-15,
2581                    "support curve must be monotone in scale (query {qi}, level {li})"
2582                );
2583            }
2584        }
2585    }
2586
2587    /// The default is single-scale mode at ANY center count: one Primary
2588    /// jet-energy candidate plus the independently REML-selected affine/null
2589    /// component requested by the default `double_penalty`. Multiscale (the
2590    /// per-scale spectral split + ψ dials) is an EXPLICIT opt-in
2591    /// (`spec.multiscale`, the DSL `mjs(…, multiscale=true)`) — there is no
2592    /// center-count auto-gate (#1116). `measure_jet_multiscale_mode` is the
2593    /// single source for this decision.
2594    #[test]
2595    pub(crate) fn default_stays_single_scale_until_multiscale_opt_in() {
2596        let n = 200usize;
2597        let data = Array2::<f64>::from_shape_fn((n, 2), |(i, k)| {
2598            let t = i as f64 / (n as f64 - 1.0);
2599            if k == 0 {
2600                t * 3.0
2601            } else {
2602                0.4 * (t * 3.0).sin()
2603            }
2604        });
2605        // Default (multiscale = false) stays single-scale even at a LARGE center
2606        // count that, under the deleted auto-gate, would have flipped to
2607        // multiscale: one pure Primary plus one function-space null component.
2608        let single = MeasureJetBasisSpec {
2609            center_strategy: CenterStrategy::FarthestPoint { num_centers: 80 },
2610            ..MeasureJetBasisSpec::default()
2611        };
2612        assert!(
2613            !measure_jet_multiscale_mode(&single),
2614            "default must resolve to single-scale at any center count"
2615        );
2616        let built_single =
2617            build_measure_jet_basis(data.view(), &single).expect("single-scale build");
2618        assert_eq!(
2619            built_single.penalties.len(),
2620            2,
2621            "single-scale double-penalty mode emits Primary + affine/null component"
2622        );
2623        assert!(matches!(
2624            built_single.penaltyinfo[0].source,
2625            PenaltySource::Primary
2626        ));
2627        assert!(matches!(
2628            built_single.penaltyinfo[1].source,
2629            PenaltySource::DoublePenaltyNullspace
2630        ));
2631        // The explicit opt-in flips to multiscale at the SAME center count: the
2632        // per-scale spectral split (several candidates) plus the same explicit
2633        // null-component candidate, strictly more candidates than single-scale.
2634        let multi = MeasureJetBasisSpec {
2635            center_strategy: CenterStrategy::FarthestPoint { num_centers: 80 },
2636            multiscale: true,
2637            ..MeasureJetBasisSpec::default()
2638        };
2639        assert!(
2640            measure_jet_multiscale_mode(&multi),
2641            "multiscale=true must resolve to multiscale mode"
2642        );
2643        let built_multi = build_measure_jet_basis(data.view(), &multi).expect("multiscale build");
2644        assert!(
2645            built_multi.penalties.len() > built_single.penalties.len(),
2646            "multiscale mode emits the per-scale spectral split plus null selection, got {} (vs single-scale {})",
2647            built_multi.penalties.len(),
2648            built_single.penalties.len()
2649        );
2650    }
2651
2652    /// An explicit order pins the Mellin weights and fuses the band into a
2653    /// single Primary candidate. Disabling explicit null recovery leaves exactly
2654    /// that candidate; enabling it must never alter the Primary itself.
2655    #[test]
2656    pub(crate) fn fused_mode_without_double_penalty_emits_single_primary_candidate() {
2657        let n = 40usize;
2658        let data = Array2::<f64>::from_shape_fn((n, 2), |(i, k)| {
2659            let t = i as f64 / (n as f64 - 1.0);
2660            if k == 0 {
2661                t * 3.0
2662            } else {
2663                0.4 * (t * 3.0).sin()
2664            }
2665        });
2666        let spec = MeasureJetBasisSpec {
2667            center_strategy: CenterStrategy::FarthestPoint { num_centers: 14 },
2668            order_s: 1.3,
2669            double_penalty: false,
2670            ..MeasureJetBasisSpec::default()
2671        };
2672        let built = build_measure_jet_basis(data.view(), &spec).expect("fused build");
2673        assert_eq!(
2674            built.penalties.len(),
2675            1,
2676            "single-scale mode without null recovery emits exactly one Primary"
2677        );
2678        assert!(matches!(
2679            built.penaltyinfo[0].source,
2680            PenaltySource::Primary
2681        ));
2682        let BasisMetadata::MeasureJet { order_s, .. } = &built.metadata else {
2683            panic!("measure-jet build must return MeasureJet metadata");
2684        };
2685        assert_eq!(*order_s, 1.3, "explicit order must persist verbatim");
2686    }
2687
2688    /// The single-scale affine head is a gauge-fixed decomposition, not a
2689    /// coefficient ridge: RBF center values are exactly mass-orthogonal to the
2690    /// supported affine space and replacing those directions with the head keeps
2691    /// the total reduced width at `m - 1`.
2692    #[test]
2693    pub(crate) fn single_scale_affine_head_gauge_annihilates_center_cross() {
2694        let n = 90usize;
2695        let data = Array2::<f64>::from_shape_fn((n, 2), |(i, k)| {
2696            let t = i as f64 / (n as f64 - 1.0);
2697            if k == 0 {
2698                3.0 * t
2699            } else {
2700                (2.0 * std::f64::consts::PI * t).sin() + 0.2 * t
2701            }
2702        });
2703        let spec = MeasureJetBasisSpec {
2704            center_strategy: CenterStrategy::FarthestPoint { num_centers: 18 },
2705            double_penalty: false,
2706            multiscale: false,
2707            ..MeasureJetBasisSpec::default()
2708        };
2709        let geom = realize_measure_jet_geometry(data.view(), &spec).expect("realized geometry");
2710        let m = geom.centers.nrows();
2711        let head_rank = geom.head_transform.ncols();
2712        assert!(head_rank > 0, "fixture must realize an affine head");
2713        assert_eq!(
2714            geom.z.ncols(),
2715            m - 1,
2716            "affine gauge replaces duplicated RBF directions without widening the smooth"
2717        );
2718        let rbf_rank = m - (head_rank + 1);
2719        let z_rbf = geom.z.slice(ndarray::s![..m, ..rbf_rank]).to_owned();
2720        let k_cc =
2721            measure_jet_design_matrix(geom.centers.view(), geom.centers.view(), geom.length_scale)
2722                .expect("center kernel");
2723        let affine = measure_jet_affine_value_basis(geom.centers.view(), geom.masses.view());
2724        assert_eq!(affine.ncols(), head_rank + 1);
2725        let mut weighted_affine = affine.clone();
2726        for (i, mut row) in weighted_affine.outer_iter_mut().enumerate() {
2727            row.mapv_inplace(|v| v * geom.masses[i]);
2728        }
2729        let constraint_cross = k_cc.t().dot(&weighted_affine);
2730        let residual = constraint_cross.t().dot(&z_rbf);
2731        let scale = constraint_cross
2732            .iter()
2733            .fold(1.0_f64, |acc, value| acc.max(value.abs()));
2734        assert!(
2735            residual.iter().all(|value| value.abs() <= 1e-10 * scale),
2736            "A^T W Kcc Z_rbf must vanish; max residual {:.3e}",
2737            residual
2738                .iter()
2739                .fold(0.0_f64, |acc, value| acc.max(value.abs()))
2740        );
2741    }
2742
2743    /// A function-space penalty must transform covariantly with its evaluation
2744    /// map. This directly excludes any hidden Euclidean coefficient projector.
2745    #[test]
2746    pub(crate) fn affine_null_penalty_is_covariant_under_coefficient_reparameterization() {
2747        let centers = array![
2748            [-1.0, 0.2],
2749            [-0.4, -0.3],
2750            [0.1, 0.5],
2751            [0.7, -0.2],
2752            [1.2, 0.4],
2753            [1.8, -0.1],
2754        ];
2755        let masses = array![0.08, 0.12, 0.18, 0.22, 0.17, 0.23];
2756        let evaluation = Array2::<f64>::from_shape_fn((centers.nrows(), 3), |(i, j)| {
2757            ((i + 2 * j + 1) as f64).sin() + 0.15 * (i * (j + 1)) as f64
2758        });
2759        let reparameterization = array![[1.7, 0.2, -0.1], [0.0, 0.6, 0.3], [0.0, 0.0, 1.3]];
2760        let base = affine_function_nullspace_penalty(&evaluation, centers.view(), masses.view())
2761            .expect("base function-space penalty");
2762        let transformed_evaluation = evaluation.dot(&reparameterization);
2763        let transformed = affine_function_nullspace_penalty(
2764            &transformed_evaluation,
2765            centers.view(),
2766            masses.view(),
2767        )
2768        .expect("reparameterized function-space penalty");
2769        let expected = reparameterization.t().dot(&base).dot(&reparameterization);
2770        let scale = expected
2771            .iter()
2772            .fold(1.0_f64, |acc, value| acc.max(value.abs()));
2773        assert!(
2774            transformed
2775                .iter()
2776                .zip(expected.iter())
2777                .all(|(actual, want)| (actual - want).abs() <= 1e-11 * scale),
2778            "S(E R) must equal R^T S(E) R"
2779        );
2780    }
2781
2782    /// `double_penalty` adds a distinct evidence-selected component and cannot
2783    /// mutate the jet-energy estimand carried by Primary.
2784    #[test]
2785    pub(crate) fn double_penalty_leaves_primary_matrix_unchanged() {
2786        let n = 64usize;
2787        let data = Array2::<f64>::from_shape_fn((n, 2), |(i, k)| {
2788            let t = i as f64 / (n as f64 - 1.0);
2789            if k == 0 { 2.5 * t } else { (4.0 * t).cos() }
2790        });
2791        let base = MeasureJetBasisSpec {
2792            center_strategy: CenterStrategy::FarthestPoint { num_centers: 16 },
2793            order_s: 1.25,
2794            double_penalty: false,
2795            ..MeasureJetBasisSpec::default()
2796        };
2797        let without = build_measure_jet_basis(data.view(), &base).expect("primary-only build");
2798        let with = build_measure_jet_basis(
2799            data.view(),
2800            &MeasureJetBasisSpec {
2801                double_penalty: true,
2802                ..base.clone()
2803            },
2804        )
2805        .expect("double-penalty build");
2806        assert_eq!(without.penalties.len(), 1);
2807        assert_eq!(with.penalties.len(), 2);
2808        assert!(matches!(
2809            without.penaltyinfo[0].source,
2810            PenaltySource::Primary
2811        ));
2812        assert!(matches!(with.penaltyinfo[0].source, PenaltySource::Primary));
2813        assert!(matches!(
2814            with.penaltyinfo[1].source,
2815            PenaltySource::DoublePenaltyNullspace
2816        ));
2817        assert!(
2818            without.penalties[0]
2819                .iter()
2820                .zip(with.penalties[0].iter())
2821                .all(|(a, b)| (a - b).abs() <= 1e-13),
2822            "turning on null recovery must not modify Primary"
2823        );
2824    }
2825
2826    /// The Householder basis must be orthonormal with sum-to-zero columns.
2827    #[test]
2828    pub(crate) fn householder_sum_to_zero_basis_is_orthonormal() {
2829        let m = 9usize;
2830        let u = householder_sum_to_zero_u(m);
2831        let z = householder_sum_to_zero_z(&u);
2832        for j in 0..(m - 1) {
2833            let col_j = z.column(j);
2834            assert!(col_j.sum().abs() <= 1e-12, "column {j} must sum to zero");
2835            for j2 in j..(m - 1) {
2836                let dot = col_j.dot(&z.column(j2));
2837                let want = if j == j2 { 1.0 } else { 0.0 };
2838                assert!(
2839                    (dot - want).abs() <= 1e-12,
2840                    "orthonormality failure at ({j}, {j2}): {dot}"
2841                );
2842            }
2843        }
2844    }
2845
2846    /// Frozen-geometry fixture shared by the ψ-producer FD gates: build
2847    /// once, pin everything (nodes, masses, band, transform, realized ℓ),
2848    /// and return the pinned spec so dial-perturbed rebuilds move ONLY the
2849    /// dials — the per-trial contract the optimizer relies on.
2850    pub(crate) fn frozen_spec_fixture(
2851        order_s: f64,
2852        multiscale: bool,
2853    ) -> (Array2<f64>, MeasureJetBasisSpec) {
2854        // Multiscale (per-scale + ψ) mode is the explicit opt-in (#1116); the
2855        // per-level fixture passes `multiscale = true`, the fused fixture
2856        // `false`. A large center count is kept so the multiscale spectrum is
2857        // identifiable when opted in.
2858        let n = 140usize;
2859        let data = Array2::<f64>::from_shape_fn((n, 2), |(i, k)| {
2860            let t = i as f64 / (n as f64 - 1.0);
2861            if k == 0 {
2862                t * 3.0
2863            } else {
2864                0.5 * (t * 3.0).cos() + if i % 9 == 0 { 0.8 } else { 0.0 }
2865            }
2866        });
2867        let spec = MeasureJetBasisSpec {
2868            center_strategy: CenterStrategy::FarthestPoint { num_centers: 70 },
2869            order_s,
2870            multiscale,
2871            // These fixtures gate the PENALTY-dial derivatives; freeze ℓ so the
2872            // coordinate layout is exactly the penalty dials (the design-moving
2873            // ℓ dial has its own FD gate, `psi_producer_matches_fd_length_scale`).
2874            learn_length_scale: false,
2875            ..MeasureJetBasisSpec::default()
2876        };
2877        let first = build_measure_jet_basis(data.view(), &spec).expect("fixture build");
2878        let BasisMetadata::MeasureJet {
2879            centers,
2880            length_scale,
2881            eps_band,
2882            masses,
2883            support_means,
2884            penalty_normalization_scales,
2885            raw_penalty_normalization_scales,
2886            fused_penalty_normalization_scale,
2887            constraint_transform,
2888            ..
2889        } = &first.metadata
2890        else {
2891            panic!("measure-jet build must return MeasureJet metadata");
2892        };
2893        let frozen = MeasureJetBasisSpec {
2894            center_strategy: CenterStrategy::UserProvided(centers.clone()),
2895            order_s,
2896            alpha: spec.alpha,
2897            tau0: spec.tau0,
2898            num_scales: eps_band.len(),
2899            length_scale: *length_scale,
2900            double_penalty: spec.double_penalty,
2901            learn_length_scale: false,
2902            multiscale,
2903            identifiability: MeasureJetIdentifiability::FrozenTransform {
2904                transform: constraint_transform.clone().expect("fit-time z"),
2905            },
2906            frozen_quadrature: Some(MeasureJetFrozenQuadrature {
2907                masses: masses.clone(),
2908                eps_band: eps_band.clone(),
2909                support_means: support_means.clone(),
2910                penalty_normalization_scales: penalty_normalization_scales.clone(),
2911                raw_penalty_normalization_scales: raw_penalty_normalization_scales.clone(),
2912                fused_penalty_normalization_scale: *fused_penalty_normalization_scale,
2913                sigma_coord: None,
2914            }),
2915        };
2916        (data, frozen)
2917    }
2918
2919    /// ψ-producer vs central finite differences of the NORMALIZED fit-time
2920    /// candidates under frozen geometry — per-level mode (coords α, lnτ).
2921    /// This is the end-to-end gate #901 never had: the derivative is checked
2922    /// against the exact object the optimizer consumes.
2923    #[test]
2924    pub(crate) fn psi_producer_matches_fd_per_level_mode() {
2925        let (data, frozen) = frozen_spec_fixture(0.0, true);
2926        let derivs =
2927            build_measure_jet_basis_psi_derivatives(data.view(), &frozen).expect("psi derivatives");
2928        let l_count = frozen
2929            .frozen_quadrature
2930            .as_ref()
2931            .expect("frozen quadrature")
2932            .eps_band
2933            .len();
2934        assert_eq!(
2935            derivs.penalties_first.len(),
2936            2,
2937            "per-level coords are (α, lnτ)"
2938        );
2939        assert_eq!(derivs.penalties_first[0].len(), l_count + 1);
2940        assert_eq!(derivs.penalties_cross_pairs, vec![(0, 1)]);
2941        let pen_at = |alpha: f64, tau0: f64| {
2942            let trial = MeasureJetBasisSpec {
2943                alpha,
2944                tau0,
2945                ..frozen.clone()
2946            };
2947            build_measure_jet_basis(data.view(), &trial)
2948                .expect("trial build")
2949                .penalties
2950        };
2951        // Second-difference-optimal step (see the jets FD test): the 4-point
2952        // cross stencil shares the ~ε·scale/h² roundoff floor.
2953        let h = 1e-4;
2954        let (a0, t0) = (frozen.alpha, frozen.tau0);
2955        let ap = pen_at(a0 + h, t0);
2956        let am = pen_at(a0 - h, t0);
2957        let tp = pen_at(a0, t0 * h.exp());
2958        let tm = pen_at(a0, t0 * (-h).exp());
2959        assert_eq!(
2960            ap.len(),
2961            l_count + 1,
2962            "fixture must keep every scale active"
2963        );
2964        for level in 0..l_count {
2965            let fd_alpha = (&ap[level] - &am[level]) / (2.0 * h);
2966            let fd_tau = (&tp[level] - &tm[level]) / (2.0 * h);
2967            for (name, analytic, fd) in [
2968                ("alpha", &derivs.penalties_first[0][level], fd_alpha),
2969                ("ln_tau", &derivs.penalties_first[1][level], fd_tau),
2970            ] {
2971                let scale = fd.iter().fold(1e-30_f64, |acc, v| acc.max(v.abs()));
2972                for (x, y) in analytic.iter().zip(fd.iter()) {
2973                    assert!(
2974                        (x - y).abs() <= 5e-5 * scale,
2975                        "{name} jet of scale-candidate {level}: analytic {x:.6e} vs FD {y:.6e}"
2976                    );
2977                }
2978            }
2979        }
2980        // The function-space null candidate is independent of α and τ.
2981        for coord in 0..2 {
2982            assert!(
2983                derivs.penalties_first[coord][l_count]
2984                    .iter()
2985                    .all(|v| *v == 0.0),
2986                "null-component candidate must have zero (α, lnτ) drift"
2987            );
2988        }
2989        // Cross derivative through the provider, against a 4-point FD.
2990        let provider = derivs
2991            .penalties_cross_provider
2992            .as_ref()
2993            .expect("cross provider");
2994        let cross = provider.evaluate(0, 1).expect("cross pair (α, lnτ)");
2995        let pp = pen_at(a0 + h, t0 * h.exp());
2996        let pm = pen_at(a0 + h, t0 * (-h).exp());
2997        let mp = pen_at(a0 - h, t0 * h.exp());
2998        let mm = pen_at(a0 - h, t0 * (-h).exp());
2999        for level in 0..l_count {
3000            let fd = (&(&pp[level] - &pm[level]) - &(&mp[level] - &mm[level])) / (4.0 * h * h);
3001            let scale = fd.iter().fold(1e-30_f64, |acc, v| acc.max(v.abs()));
3002            for (x, y) in cross[level].iter().zip(fd.iter()) {
3003                assert!(
3004                    (x - y).abs() <= 5e-4 * scale,
3005                    "cross (α, lnτ) jet of scale-candidate {level}: analytic {x:.6e} vs FD {y:.6e}"
3006                );
3007            }
3008        }
3009    }
3010
3011    /// Design-moving ℓ dial (#1116): the producer's design jets and every
3012    /// normalized penalty candidate's jets must match central differences of the
3013    /// REBUILT objects under frozen geometry. Although the center-value forms
3014    /// `Q` and `H₀` are ℓ-invariant, their coefficient pullbacks `E(ℓ)ᵀQ E(ℓ)`
3015    /// and `E(ℓ)ᵀH₀E(ℓ)` are not.
3016    #[test]
3017    pub(crate) fn psi_producer_matches_fd_length_scale() {
3018        // Single-scale with opt-in ℓ learning; frozen geometry so only ℓ moves
3019        // across the FD trials.
3020        let (data, mut frozen) = frozen_spec_fixture(0.0, false);
3021        frozen.learn_length_scale = true;
3022        let derivs =
3023            build_measure_jet_basis_psi_derivatives(data.view(), &frozen).expect("psi derivatives");
3024        // ℓ is the only coordinate in single-scale + learn_length_scale.
3025        assert_eq!(
3026            derivs.design_first.len(),
3027            1,
3028            "single-scale + learn_length_scale enrolls exactly the ℓ coordinate"
3029        );
3030        assert_eq!(
3031            derivs.penalties_first[0].len(),
3032            2,
3033            "single-scale double penalty carries Primary + affine/null component"
3034        );
3035        // Rebuild design and normalized penalties at ℓ·e^{±h}; the explicit
3036        // positive length_scale is honored verbatim while the frozen transform
3037        // keeps the coefficient chart fixed.
3038        let ell0 = frozen.length_scale;
3039        let build_at = |ell: f64| {
3040            let trial = MeasureJetBasisSpec {
3041                length_scale: ell,
3042                ..frozen.clone()
3043            };
3044            build_measure_jet_basis(data.view(), &trial).expect("trial build")
3045        };
3046        let h: f64 = 1e-4;
3047        let plus = build_at(ell0 * h.exp());
3048        let minus = build_at(ell0 * (-h).exp());
3049        let at = build_at(ell0);
3050        assert_eq!(
3051            plus.penalties.len(),
3052            2,
3053            "fixture must keep both candidates active"
3054        );
3055        assert_eq!(
3056            minus.penalties.len(),
3057            2,
3058            "fixture must keep both candidates active"
3059        );
3060        assert_eq!(
3061            at.penalties.len(),
3062            2,
3063            "fixture must keep both candidates active"
3064        );
3065
3066        let x_plus = plus.design.to_dense();
3067        let x_minus = minus.design.to_dense();
3068        let x_0 = at.design.to_dense();
3069        let fd_first = (&x_plus - &x_minus) / (2.0 * h);
3070        let fd_second = (&x_plus - &(&x_0 * 2.0) + &x_minus) / (h * h);
3071        let scale1 = fd_first.iter().fold(1e-30_f64, |acc, v| acc.max(v.abs()));
3072        for (x, y) in derivs.design_first[0].iter().zip(fd_first.iter()) {
3073            assert!(
3074                (x - y).abs() <= 5e-5 * scale1,
3075                "∂X/∂lnℓ: analytic {x:.6e} vs FD {y:.6e}"
3076            );
3077        }
3078        let scale2 = fd_second.iter().fold(1e-30_f64, |acc, v| acc.max(v.abs()));
3079        for (x, y) in derivs.design_second_diag[0].iter().zip(fd_second.iter()) {
3080            assert!(
3081                (x - y).abs() <= 1e-3 * scale2,
3082                "∂²X/∂lnℓ²: analytic {x:.6e} vs FD {y:.6e}"
3083            );
3084        }
3085
3086        for candidate in 0..2 {
3087            let fd_penalty_first =
3088                (&plus.penalties[candidate] - &minus.penalties[candidate]) / (2.0 * h);
3089            let fd_penalty_second = (&plus.penalties[candidate]
3090                - &(&at.penalties[candidate] * 2.0)
3091                + &minus.penalties[candidate])
3092                / (h * h);
3093            let first_scale = fd_penalty_first
3094                .iter()
3095                .fold(1e-12_f64, |acc, value| acc.max(value.abs()));
3096            let second_scale = fd_penalty_second
3097                .iter()
3098                .fold(1e-10_f64, |acc, value| acc.max(value.abs()));
3099            for (analytic, finite_difference) in derivs.penalties_first[0][candidate]
3100                .iter()
3101                .zip(fd_penalty_first.iter())
3102            {
3103                assert!(
3104                    (analytic - finite_difference).abs() <= 1e-4 * first_scale,
3105                    "candidate {candidate} ∂S~/∂lnℓ: analytic {analytic:.6e} vs FD {finite_difference:.6e}"
3106                );
3107            }
3108            for (analytic, finite_difference) in derivs.penalties_second_diag[0][candidate]
3109                .iter()
3110                .zip(fd_penalty_second.iter())
3111            {
3112                assert!(
3113                    (analytic - finite_difference).abs() <= 5e-3 * second_scale,
3114                    "candidate {candidate} ∂²S~/∂lnℓ²: analytic {analytic:.6e} vs FD {finite_difference:.6e}"
3115                );
3116            }
3117        }
3118    }
3119
3120    /// Quadrature nodes must be the mass-weighted cell barycenters
3121    /// (first-moment-exact lumping), with empty cells keeping their seed
3122    /// coordinates at zero mass.
3123    #[test]
3124    pub(crate) fn quadrature_nodes_are_cell_barycenters() {
3125        // Two tight groups around (0,0) and (10,10); a third seed far away
3126        // captures nothing.
3127        let data = array![
3128            [0.0, 0.2],
3129            [0.4, -0.2],
3130            [0.2, 0.0],
3131            [9.8, 10.1],
3132            [10.2, 9.9],
3133        ];
3134        let seeds = array![[0.1, 0.1], [10.0, 10.0], [-50.0, -50.0]];
3135        let (nodes, masses) =
3136            measure_jet_quadrature_nodes(data.view(), seeds.view()).expect("quadrature nodes");
3137        assert!((masses.sum() - 1.0).abs() <= 1e-15, "masses must sum to 1");
3138        assert!((masses[0] - 0.6).abs() <= 1e-15);
3139        assert!((masses[1] - 0.4).abs() <= 1e-15);
3140        assert_eq!(masses[2], 0.0);
3141        // Cell 0 barycenter = (0.2, 0.0).
3142        assert_eq!(nodes[(0, 0)], 0.2);
3143        assert_eq!(nodes[(0, 1)], 0.0);
3144        // Cell 1 barycenter = (10.0, 10.0), which is not a sampled row.
3145        assert_eq!(nodes[(1, 0)], 10.0);
3146        assert_eq!(nodes[(1, 1)], 10.0);
3147        // Empty cell keeps its seed coordinates.
3148        assert_eq!(nodes[(2, 0)], -50.0);
3149        assert_eq!(nodes[(2, 1)], -50.0);
3150    }
3151
3152    /// Freeze→replay: rebuilding from the first build's frozen transform and
3153    /// frozen quadrature must reproduce design and penalty bit-for-bit (the
3154    /// predict-path contract).
3155    #[test]
3156    pub(crate) fn build_replay_roundtrip_reproduces_design_and_penalty() {
3157        // A bent filament with a side cluster; multiscale opt-in so this
3158        // exercises the per-scale (spectral) replay path (#1116).
3159        let n = 140usize;
3160        let data = Array2::<f64>::from_shape_fn((n, 2), |(i, k)| {
3161            let t = i as f64 / (n as f64 - 1.0);
3162            if k == 0 {
3163                t * 3.0
3164            } else {
3165                0.5 * (t * 3.0).cos() + if i % 9 == 0 { 0.8 } else { 0.0 }
3166            }
3167        });
3168        let spec = MeasureJetBasisSpec {
3169            center_strategy: CenterStrategy::FarthestPoint { num_centers: 70 },
3170            multiscale: true,
3171            ..MeasureJetBasisSpec::default()
3172        };
3173        let first = build_measure_jet_basis(data.view(), &spec).expect("first build");
3174        let BasisMetadata::MeasureJet {
3175            centers,
3176            length_scale,
3177            eps_band,
3178            order_s,
3179            alpha,
3180            tau0,
3181            masses,
3182            support_means,
3183            penalty_normalization_scales,
3184            raw_penalty_normalization_scales,
3185            fused_penalty_normalization_scale,
3186            constraint_transform,
3187            ..
3188        } = &first.metadata
3189        else {
3190            panic!("measure-jet build must return MeasureJet metadata");
3191        };
3192        let replay_spec = MeasureJetBasisSpec {
3193            center_strategy: CenterStrategy::UserProvided(centers.clone()),
3194            order_s: *order_s,
3195            alpha: *alpha,
3196            tau0: *tau0,
3197            num_scales: eps_band.len(),
3198            length_scale: *length_scale,
3199            double_penalty: spec.double_penalty,
3200            learn_length_scale: spec.learn_length_scale,
3201            multiscale: spec.multiscale,
3202            identifiability: MeasureJetIdentifiability::FrozenTransform {
3203                transform: constraint_transform.clone().expect("fit-time z"),
3204            },
3205            frozen_quadrature: Some(MeasureJetFrozenQuadrature {
3206                masses: masses.clone(),
3207                eps_band: eps_band.clone(),
3208                support_means: support_means.clone(),
3209                penalty_normalization_scales: penalty_normalization_scales.clone(),
3210                raw_penalty_normalization_scales: raw_penalty_normalization_scales.clone(),
3211                fused_penalty_normalization_scale: *fused_penalty_normalization_scale,
3212                sigma_coord: None,
3213            }),
3214        };
3215        // Per-level mode: one candidate per band scale plus the function-space
3216        // null component, and the count must survive replay bit-for-bit.
3217        assert_eq!(
3218            first.penalties.len(),
3219            eps_band.len() + 1,
3220            "per-level mode must emit one candidate per scale + null component"
3221        );
3222        let second = build_measure_jet_basis(data.view(), &replay_spec).expect("replay build");
3223        let x1 = first.design.to_dense();
3224        let x2 = second.design.to_dense();
3225        assert_eq!(x1.shape(), x2.shape());
3226        for (a, b) in x1.iter().zip(x2.iter()) {
3227            assert!((a - b).abs() <= 1e-12, "design replay drift: {a} vs {b}");
3228        }
3229        assert_eq!(first.penalties.len(), second.penalties.len());
3230        for (p1, p2) in first.penalties.iter().zip(second.penalties.iter()) {
3231            for (a, b) in p1.iter().zip(p2.iter()) {
3232                assert!((a - b).abs() <= 1e-12, "penalty replay drift: {a} vs {b}");
3233            }
3234        }
3235    }
3236}