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 fused jet-energy penalty at the auto order with no ψ
66//!   dials, the same one-λ outer footprint as Duchon/Matérn — so a fit pays
67//!   Duchon-class cost unless the user explicitly asks for the per-scale
68//!   spectrum. The flag is persisted on the spec, so freeze→replay re-enters
69//!   the same mode 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-moving dial** (matérn's `log_kappa`
83//!   analog, #1116): `X = K(data, centers; ℓ)·z` depends on ℓ, so
84//!   `∂X/∂lnℓ = (K ⊙ r²/ℓ²)·z ≠ 0`, while the jet-energy penalty does NOT
85//!   depend on ℓ (`∂S/∂lnℓ ≡ 0`). When explicitly enabled, ℓ rebuilds the
86//!   design per outer trial; it does not change the basis RANK (the Gaussian
87//!   kernel is PD ∀ℓ > 0 → rank ≡ m centers) but it does change which m-dim
88//!   subspace the representers span. 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;
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/// Single-scale fused-mode weight of the affine-preserving nullspace ridge,
172/// as a fraction of the primary energy penalty's Frobenius scale (#1116). The
173/// ridge's role is identifiability — flooring the fused energy's near-nullspace
174/// — not smoothing selection, so it is fused at a small FIXED fraction rather
175/// than carrying its own REML λ (which would double the single-scale outer
176/// search). 1e-2 lifts the unpenalized directions ~two orders below the
177/// dominant penalized modes: a definite identifiability floor that leaves the
178/// genuine spectrum (and the single primary λ) in control of the fit.
179pub(crate) const MEASURE_JET_FUSED_RIDGE_FRACTION: f64 = 1e-2;
180
181/// Memory budget (in f64 entries) above which the multi-form assembly stops
182/// parallelizing over scales: parallel scale partials cost
183/// `L · n_forms · m²` doubles; past this budget the scales run sequentially
184/// (same numbers — the per-scale loop and the ordered sum are deterministic
185/// either way).
186pub(crate) const MEASURE_JET_PARALLEL_FORM_BUDGET_DOUBLES: usize = 1 << 26;
187
188/// Realized-design identifiability policy for the measure-jet smooth.
189/// Mirrors [`super::ConstantCurvatureIdentifiability`] (#532): the fit-time
190/// center-space sum-to-zero `z` gets the parametric orthogonalization
191/// composed onto it by the global identifiability pipeline, and the composed
192/// transform is frozen so predict-time (and per-ψ-trial) rebuilds replay it
193/// verbatim instead of recomputing `z` from the centers.
194#[derive(Debug, Clone, Serialize, Deserialize, Default)]
195pub enum MeasureJetIdentifiability {
196    /// Fit-time default: uniform-weight coefficient sum-to-zero over the
197    /// centers (`Σ_j α_j = 0`), then global parametric residualization.
198    #[default]
199    CenterSumToZero,
200    /// Predict-time replay: the frozen composed transform captured at fit
201    /// time. `transform.nrows()` equals the number of centers.
202    FrozenTransform { transform: Array2<f64> },
203}
204
205/// Fit-time quadrature of the empirical measure (center masses + realized
206/// scale band), frozen onto the spec so predict-time rebuilds replay the
207/// exact fit-time penalty. Recomputing either from predict rows would
208/// silently change the penalty the coefficients were estimated under.
209#[derive(Debug, Clone, Serialize, Deserialize)]
210pub struct MeasureJetFrozenQuadrature {
211    /// Per-center masses `m_i` (nearest-center fractions of the FIT rows).
212    pub masses: Array1<f64>,
213    /// Realized geometric scale band `ε_0 < … < ε_{L−1}`.
214    pub eps_band: Vec<f64>,
215    /// Per-scale on-web support anchor
216    /// `q̄_ℓ = (Σ_i m_i q_ℓ(c_i)) / (Σ_i m_i)`.
217    pub support_means: Vec<f64>,
218    /// Frobenius scales of the emitted per-level normalized penalties. Empty in
219    /// fused mode, where the band emits one primary penalty instead.
220    pub penalty_normalization_scales: Vec<f64>,
221    /// Frobenius scales of the raw per-level forms before the arbitrary Mellin
222    /// `log_step · ε_ℓ^(-2s0)` gauge is folded in.
223    pub raw_penalty_normalization_scales: Vec<f64>,
224    /// Frobenius scale of the single fused primary penalty. `None` in per-level
225    /// mode.
226    pub fused_penalty_normalization_scale: Option<f64>,
227}
228
229/// Serde default for [`MeasureJetBasisSpec::learn_length_scale`]: freeze ℓ at
230/// the realized auto/user value unless a fit explicitly opts into the
231/// design-moving outer coordinate.
232fn measure_jet_learn_length_scale_default() -> bool {
233    false
234}
235
236/// Measure-jet smooth configuration (`mjs(x0, …, xd)`).
237///
238/// The feature columns are ambient coordinates of data concentrated near an
239/// unknown low-dimensional (possibly stratified) set; the term learns the
240/// geometry from the empirical measure itself — centers as quadrature nodes,
241/// masses as μ-weights, local jet residuals as the roughness carrier — with
242/// no graph, mesh, or neighbor-set inside the statistical object.
243#[derive(Debug, Clone, Serialize, Deserialize)]
244pub struct MeasureJetBasisSpec {
245    /// Center/knot selection strategy (deterministic; quadrature of μ).
246    pub center_strategy: CenterStrategy,
247    /// Continuous smoothness order `s ∈ (0, 2)`; `0.0` sentinel = auto
248    /// ([`MEASURE_JET_DEFAULT_ORDER_S`]).
249    pub order_s: f64,
250    /// Density-normalization exponent α (outer weight `q^{1−2α}`).
251    pub alpha: f64,
252    /// Historical τ coordinate retained for frozen specs and ψ layout. The
253    /// measure-jet energy itself uses the exact weighted affine projection and
254    /// is independent of τ; the τ ψ derivatives are therefore zero.
255    pub tau0: f64,
256    /// Number of scale nodes; `0` sentinel = auto dyadic band.
257    pub num_scales: usize,
258    /// Representer (Gaussian RBF) range ℓ; `0.0` sentinel = auto
259    /// (median nearest-center spacing × [`MEASURE_JET_AUTO_LENGTH_SCALE_FACTOR`]).
260    pub length_scale: f64,
261    /// Add an affine-preserving shrinkage penalty alongside the jet-energy
262    /// penalty.
263    pub double_penalty: bool,
264    /// REML-learn the representer range ℓ as a design-moving outer dial
265    /// (opt-in), mirroring Matérn's `log_kappa`. The Gaussian kernel is
266    /// strictly PD for every ℓ > 0, so ℓ does NOT change the basis rank (always
267    /// `m` centers) — but it changes WHICH `m`-dim subspace the representers
268    /// span, i.e. the span alignment with the true surface. The stable default
269    /// freezes ℓ at the auto/user value; `true` enrolls the outer coordinate for
270    /// experiments that need REML-selected representer range.
271    #[serde(default = "measure_jet_learn_length_scale_default")]
272    pub learn_length_scale: bool,
273    /// Explicit opt-in for multiscale mode: the per-scale spectral penalty
274    /// split plus the `(α, ln τ)` outer ψ dials and the affine-preserving
275    /// ridge. `false` (default) keeps the term in single-scale mode at ANY
276    /// center count — one fused jet-energy penalty, the one-λ Duchon/Matérn
277    /// footprint. There is no center-count auto-gate; the user opts in via
278    /// `mjs(…, multiscale=true)`. Persisted on the spec so freeze→replay
279    /// re-enters the same mode.
280    #[serde(default)]
281    pub multiscale: bool,
282    /// Realized-design identifiability policy (see type docs).
283    #[serde(default)]
284    pub identifiability: MeasureJetIdentifiability,
285    /// Fit-time quadrature replay (see type docs). `None` at fit time;
286    /// `Some` on the frozen predict/ψ-trial path.
287    #[serde(default)]
288    pub frozen_quadrature: Option<MeasureJetFrozenQuadrature>,
289}
290
291impl Default for MeasureJetBasisSpec {
292    fn default() -> Self {
293        Self {
294            center_strategy: CenterStrategy::FarthestPoint { num_centers: 50 },
295            order_s: 0.0,
296            // Density-WEIGHTED Hessian energy (the module-header default): the
297            // outer weight is q^{1−2α} = q^{−1} at α = 1. The density-free
298            // variant α = 3/2 gives q^{−2}, which on a low-intrinsic-dimension
299            // stratum (data on a 1-D/2-D manifold embedded in higher ambient d)
300            // makes the local kernel mass q tiny AND spatially varying along
301            // the manifold, so q^{−2} amplifies the penalty unevenly and
302            // over-smooths the high-frequency signal there (MEASURED #1116: on
303            // the 1-D-curve-in-3-D fixture α = 3/2 left mjs ~13× worse than
304            // matérn). α = 1's q^{−1} weighting is far gentler and is the
305            // header-derived default; an explicit `alpha=` still overrides for
306            // genuinely density-free use on a full-dimensional stratum.
307            alpha: 1.0,
308            tau0: 1e-3,
309            num_scales: 0,
310            length_scale: 0.0,
311            double_penalty: true,
312            learn_length_scale: false,
313            multiscale: false,
314            identifiability: MeasureJetIdentifiability::CenterSumToZero,
315            frozen_quadrature: None,
316        }
317    }
318}
319
320/// Realized geometric scale band: `eps` ascending, `log_step` the constant
321/// log-spacing `ln(eps[ℓ+1]/eps[ℓ])` used as the Mellin quadrature weight.
322pub struct MeasureJetBand {
323    pub eps: Vec<f64>,
324    pub log_step: f64,
325}
326
327/// The energy and its exact hyperparameter jets in the live dials. `s` and
328/// `α` enter only through per-block log-weights. The retained `ln τ` slots
329/// are zero because the local fit is the exact weighted affine projection
330/// and no longer depends on τ. All forms are scattered from the SAME local
331/// residual blocks, and the ψ-channel consumes them with zero design drift.
332pub struct MeasureJetEnergyJets {
333    pub q: Array2<f64>,
334    pub dq_ds: Array2<f64>,
335    pub d2q_ds2: Array2<f64>,
336    pub dq_dalpha: Array2<f64>,
337    pub d2q_dalpha2: Array2<f64>,
338    pub d2q_ds_dalpha: Array2<f64>,
339    pub dq_dlogtau: Array2<f64>,
340    pub d2q_dlogtau2: Array2<f64>,
341    pub d2q_ds_dlogtau: Array2<f64>,
342    pub d2q_dalpha_dlogtau: Array2<f64>,
343}
344
345/// Householder vector `u` for the uniform sum-to-zero constraint: the
346/// reflection `H = I − 2uuᵀ` maps `c̄ = 1/√m·1` onto `e₁`, so columns 2..m
347/// of `H` are an orthonormal basis of `1⊥` — the same model space as the
348/// generic RRQR nullspace basis, but with O(rows·m) STRUCTURED application
349/// (`X·z = (X − 2(Xu)uᵀ) minus column 1`) instead of the O(rows·m²)
350/// constraint GEMM that the scale-smoke gate identified as the dominant
351/// build cost.
352pub(crate) fn householder_sum_to_zero_u(m: usize) -> Array1<f64> {
353    let c = 1.0 / (m as f64).sqrt();
354    let mut u = Array1::<f64>::from_elem(m, c);
355    u[0] -= 1.0;
356    let norm = u.dot(&u).sqrt();
357    u.mapv_inplace(|v| v / norm);
358    u
359}
360
361/// Materialize the Householder sum-to-zero basis `z` (m × (m−1)) — columns
362/// 2..m of `H = I − 2uuᵀ` — for the frozen-replay metadata. O(m²), built
363/// once per fit.
364pub(crate) fn householder_sum_to_zero_z(u: &Array1<f64>) -> Array2<f64> {
365    let m = u.len();
366    let mut z = Array2::<f64>::zeros((m, m - 1));
367    for j in 0..(m - 1) {
368        for i in 0..m {
369            let h = if i == j + 1 { 1.0 } else { 0.0 } - 2.0 * u[i] * u[j + 1];
370            z[(i, j)] = h;
371        }
372    }
373    z
374}
375
376pub(crate) fn symmetric_pseudoinverse(
377    a: &Array2<f64>,
378    label: &str,
379) -> Result<Array2<f64>, BasisError> {
380    let n = a.nrows();
381    if a.ncols() != n {
382        crate::bail_dim_basis!(
383            "measure-jet pseudo-inverse `{label}` needs a square matrix, got {:?}",
384            a.dim()
385        );
386    }
387    let (evals, evecs) = a.eigh(Side::Lower).map_err(|e| {
388        BasisError::InvalidInput(format!(
389            "measure-jet pseudo-inverse `{label}` eigendecomposition failed: {e}"
390        ))
391    })?;
392    let lam_max = evals.iter().fold(0.0_f64, |acc, v| acc.max((*v).max(0.0)));
393    let rank_tol = MEASURE_JET_PSEUDOINVERSE_RTOL * (n.max(1) as f64) * lam_max;
394    let mut scaled = evecs.clone();
395    for (k, mut col) in scaled.axis_iter_mut(Axis(1)).enumerate() {
396        let lam = evals[k].max(0.0);
397        let inv = if lam > rank_tol { 1.0 / lam } else { 0.0 };
398        col.mapv_inplace(|v| v * inv);
399    }
400    Ok(scaled.dot(&evecs.t()))
401}
402
403/// Affine-preserving shrinkage ridge `I − P_pseudoaffine` in the constrained
404/// Gaussian-coefficient space, where `P_pseudoaffine` projects onto the
405/// coefficient directions `b̂` that — pushed through the representer `K·z` — best
406/// reproduce the affine center functions `[1, c₁…c_d]` under the mass-weighted
407/// normal equations. The double penalty thus shrinks the wiggle space while
408/// preserving the (pseudo-)affine component, the measure-jet analogue of leaving
409/// a smooth's null space lightly damped.
410pub(crate) fn affine_preserving_coefficient_ridge(
411    kz: &Array2<f64>,
412    centers: ArrayView2<'_, f64>,
413    masses: ArrayView1<'_, f64>,
414) -> Result<Array2<f64>, BasisError> {
415    let m = centers.nrows();
416    let d = centers.ncols();
417    let p = kz.ncols();
418    if kz.nrows() != m || masses.len() != m {
419        crate::bail_dim_basis!(
420            "measure-jet affine-preserving ridge shape mismatch: kz {:?}, centers {:?}, masses {}",
421            kz.dim(),
422            centers.dim(),
423            masses.len()
424        );
425    }
426    if p == 0 {
427        return Ok(Array2::<f64>::zeros((0, 0)));
428    }
429    let mut weighted_kz = kz.clone();
430    for (i, mut row) in weighted_kz.outer_iter_mut().enumerate() {
431        row.mapv_inplace(|v| v * masses[i]);
432    }
433    let normal = kz.t().dot(&weighted_kz);
434    let normal_pinv = symmetric_pseudoinverse(&normal, "affine ridge normal")?;
435    let mut affine = Array2::<f64>::ones((m, d + 1));
436    for i in 0..m {
437        for k in 0..d {
438            affine[(i, k + 1)] = centers[(i, k)];
439        }
440    }
441    let mut weighted_affine = affine.clone();
442    for (i, mut row) in weighted_affine.outer_iter_mut().enumerate() {
443        row.mapv_inplace(|v| v * masses[i]);
444    }
445    let rhs = kz.t().dot(&weighted_affine);
446    let beta = normal_pinv.dot(&rhs);
447    let beta_gram = beta.t().dot(&beta);
448    let (evals, evecs) = beta_gram.eigh(Side::Lower).map_err(|e| {
449        BasisError::InvalidInput(format!(
450            "measure-jet affine ridge subspace eigendecomposition failed: {e}"
451        ))
452    })?;
453    let lam_max = evals.iter().fold(0.0_f64, |acc, v| acc.max((*v).max(0.0)));
454    let rank_tol = MEASURE_JET_PSEUDOINVERSE_RTOL * ((d + 1).max(1) as f64) * lam_max;
455    let mut ridge = Array2::<f64>::eye(p);
456    for k in 0..(d + 1) {
457        let lam = evals[k].max(0.0);
458        if lam <= rank_tol {
459            continue;
460        }
461        let dir = beta.dot(&evecs.column(k).to_owned()) / lam.sqrt();
462        for r in 0..p {
463            for c in 0..p {
464                ridge[(r, c)] -= dir[r] * dir[c];
465            }
466        }
467    }
468    Ok((&ridge + &ridge.t()) * 0.5)
469}
470
471/// Pairwise squared distances `‖a_i − b_j‖²` via the GEMM identity
472/// `‖a − b‖² = ‖a‖² + ‖b‖² − 2·aᵀb`: one (n×d)·(d×m) matrix product carries
473/// every FMA at tile speed instead of n·m scalar distance loops — the
474/// machine-native form of this kernel, and the module's ONLY distance
475/// source (representer design, support curve, and the center-pair geometry:
476/// band floor, median spacing, ε/2-net, neighbor cutoffs). The cancellation
477/// error near-coincident points pay is O(ε_f64·‖x‖²) ABSOLUTE, harmless
478/// under a Gaussian profile (the kernel is flat at d ≈ 0); clamped at zero
479/// so roundoff cannot emit tiny negatives (the a = b diagonal therefore
480/// lands at roundoff scale, not an exact 0 — no caller pins it).
481pub(crate) fn pairwise_sq_dists(a: ArrayView2<'_, f64>, b: ArrayView2<'_, f64>) -> Array2<f64> {
482    let an: Vec<f64> = a.outer_iter().map(|r| r.dot(&r)).collect();
483    let bn: Vec<f64> = b.outer_iter().map(|r| r.dot(&r)).collect();
484    let mut g = a.dot(&b.t());
485    g.axis_iter_mut(Axis(0))
486        .into_par_iter()
487        .enumerate()
488        .for_each(|(i, mut row)| {
489            for (j, v) in row.iter_mut().enumerate() {
490                *v = (an[i] + bn[j] - 2.0 * *v).max(0.0);
491            }
492        });
493    g
494}
495
496/// Row-block size for streaming GEMM passes that must not materialize the
497/// full n×m distance matrix (nearest-node assignment): 64Ki rows × m ≤ a
498/// few hundred MB of transient per block, GEMM-speed throughout.
499pub(crate) const MEASURE_JET_ASSIGN_BLOCK_ROWS: usize = 65_536;
500
501pub(crate) fn validate_finite_points(
502    points: ArrayView2<'_, f64>,
503    what: &str,
504) -> Result<(), BasisError> {
505    for (i, row) in points.outer_iter().enumerate() {
506        if row.iter().any(|v| !v.is_finite()) {
507            crate::bail_invalid_basis!("measure-jet {what} row {i} has a non-finite coordinate");
508        }
509    }
510    Ok(())
511}
512
513/// Median nearest-OTHER-center distance — the resolution floor of the center
514/// quadrature, used for the band floor and the auto representer range.
515pub(crate) fn median_nearest_center_spacing(dist2: &Array2<f64>) -> Result<f64, BasisError> {
516    let m = dist2.nrows();
517    if m < 2 {
518        return Err(BasisError::InsufficientColumnsForConstraint { found: m });
519    }
520    let mut nearest: Vec<f64> = Vec::with_capacity(m);
521    for i in 0..m {
522        let mut best = f64::INFINITY;
523        for j in 0..m {
524            if j != i && dist2[(i, j)] < best {
525                best = dist2[(i, j)];
526            }
527        }
528        nearest.push(best.sqrt());
529    }
530    nearest.sort_by(|a, b| a.partial_cmp(b).expect("finite center spacings"));
531    let median = nearest[nearest.len() / 2];
532    if !(median.is_finite() && median > 0.0) {
533        crate::bail_invalid_basis!(
534            "measure-jet centers are degenerate (median nearest-center spacing = {median}); \
535             duplicate centers cannot carry a scale band"
536        );
537    }
538    Ok(median)
539}
540
541/// Build the realized geometric scale band from the center set: floor at the
542/// median nearest-center spacing (below it the quadrature resolves nothing),
543/// ceiling at half the bounding-box diagonal (a deterministic diameter-scale
544/// cap; local fits remain center-weighted and distinct there).
545/// `num_scales == 0` requests the auto count `clamp(⌈log2(ε_max/ε_min)⌉ + 1,
546/// 3, 8)`; a degenerate band (ceiling ≤ floor) collapses to the single floor
547/// scale with `log_step = ln 2`.
548pub fn measure_jet_band(
549    centers: ArrayView2<'_, f64>,
550    num_scales: usize,
551) -> Result<MeasureJetBand, BasisError> {
552    validate_finite_points(centers, "centers")?;
553    let dist2 = pairwise_sq_dists(centers, centers);
554    let eps_min = median_nearest_center_spacing(&dist2)?;
555    // Half the bounding-box diagonal: a cheap, deterministic diameter proxy.
556    let d = centers.ncols();
557    let mut diag2 = 0.0_f64;
558    for k in 0..d {
559        let col = centers.column(k);
560        let mut lo = f64::INFINITY;
561        let mut hi = f64::NEG_INFINITY;
562        for &v in col.iter() {
563            lo = lo.min(v);
564            hi = hi.max(v);
565        }
566        diag2 += (hi - lo) * (hi - lo);
567    }
568    let eps_max = 0.5 * diag2.sqrt();
569    if !(eps_max.is_finite() && eps_max > eps_min) {
570        return Ok(MeasureJetBand {
571            eps: vec![eps_min],
572            log_step: std::f64::consts::LN_2,
573        });
574    }
575    let auto = ((eps_max / eps_min).log2().ceil() as usize + 1)
576        .clamp(MEASURE_JET_MIN_AUTO_SCALES, MEASURE_JET_MAX_AUTO_SCALES);
577    let count = if num_scales == 0 { auto } else { num_scales };
578    if count == 1 {
579        return Ok(MeasureJetBand {
580            eps: vec![eps_min],
581            log_step: std::f64::consts::LN_2,
582        });
583    }
584    let ratio = (eps_max / eps_min).powf(1.0 / (count as f64 - 1.0));
585    let mut eps = Vec::with_capacity(count);
586    let mut e = eps_min;
587    for _ in 0..count {
588        eps.push(e);
589        e *= ratio;
590    }
591    Ok(MeasureJetBand {
592        eps,
593        log_step: ratio.ln(),
594    })
595}
596
597/// First-moment-exact quadrature of the empirical measure on the cell partition
598/// induced by the seed centers: nearest-center assignment (deterministic
599/// tie-break: lowest center index) yields per-cell masses, and each non-empty
600/// cell's quadrature node is its mass-weighted barycenter. Empty cells keep
601/// their seed coordinates with zero mass (the assembly skips them; their
602/// representer columns remain valid).
603pub fn measure_jet_quadrature_nodes(
604    data: ArrayView2<'_, f64>,
605    centers: ArrayView2<'_, f64>,
606) -> Result<(Array2<f64>, Array1<f64>), BasisError> {
607    if data.ncols() != centers.ncols() {
608        crate::bail_dim_basis!(
609            "measure-jet mass assignment dimension mismatch: data d={} centers d={}",
610            data.ncols(),
611            centers.ncols()
612        );
613    }
614    validate_finite_points(data, "data")?;
615    validate_finite_points(centers, "centers")?;
616    let n = data.nrows();
617    let m = centers.nrows();
618    let d = centers.ncols();
619    if n == 0 || m == 0 {
620        crate::bail_invalid_basis!("measure-jet mass assignment needs nonempty data and centers");
621    }
622    // Nearest-node assignment in streamed GEMM blocks: argmin_j ‖x−c_j‖² =
623    // argmin_j (‖c_j‖² − 2·xᵀc_j), so each block is one (rows×d)·(d×m)
624    // product plus a row-wise argmin — tile-speed FMAs, O(block·m) transient
625    // memory, deterministic ties to the lowest center index.
626    let cn: Vec<f64> = centers.outer_iter().map(|r| r.dot(&r)).collect();
627    let assignments: Vec<usize> = (0..n)
628        .step_by(MEASURE_JET_ASSIGN_BLOCK_ROWS)
629        .flat_map(|start| {
630            let end = (start + MEASURE_JET_ASSIGN_BLOCK_ROWS).min(n);
631            let g = data.slice(ndarray::s![start..end, ..]).dot(&centers.t());
632            let block: Vec<usize> = g
633                .axis_iter(Axis(0))
634                .into_par_iter()
635                .map(|row| {
636                    let mut best_j = 0usize;
637                    let mut best = f64::INFINITY;
638                    for (j, &gij) in row.iter().enumerate() {
639                        let s = cn[j] - 2.0 * gij;
640                        if s < best {
641                            best = s;
642                            best_j = j;
643                        }
644                    }
645                    best_j
646                })
647                .collect();
648            block
649        })
650        .collect();
651    let mut masses = Array1::<f64>::zeros(m);
652    let mut nodes = centers.to_owned();
653    let mut sums = Array2::<f64>::zeros((m, d));
654    let unit = 1.0 / n as f64;
655    for (i, &j) in assignments.iter().enumerate() {
656        masses[j] += unit;
657        for k in 0..d {
658            sums[(j, k)] += data[(i, k)];
659        }
660    }
661    // Cell barycenters: the first moment of μ on each cell. These are the
662    // realized nodes for first-moment-exact lumping.
663    let mut barycenter = sums;
664    for j in 0..m {
665        let count = masses[j] * n as f64;
666        if count > 0.0 {
667            for k in 0..d {
668                barycenter[(j, k)] /= count;
669                nodes[(j, k)] = barycenter[(j, k)];
670            }
671        }
672    }
673    Ok((nodes, masses))
674}
675
676/// Per-center masses of the empirical measure (the zeroth-moment half of
677/// [`measure_jet_quadrature_nodes`]; single assignment source).
678pub fn measure_jet_center_masses(
679    data: ArrayView2<'_, f64>,
680    centers: ArrayView2<'_, f64>,
681) -> Result<Array1<f64>, BasisError> {
682    measure_jet_quadrature_nodes(data, centers).map(|(_, masses)| masses)
683}
684
685/// THE single assembly source: walk every (scale, outer-net center) local
686/// residual block exactly once and scatter it into `n_forms` accumulators
687/// with caller-chosen scalar weights. The energy, its (s, α) jets, and the
688/// per-scale spectrum are all this routine with different weight closures,
689/// so a value/derivative desync is structurally impossible.
690///
691/// Per block the closure receives `(scale_idx, eps, q, base)` where `q` is
692/// the truncated kernel sum used by the local residual and `base`
693/// is the fully-assembled outer weight
694/// `log_step · ε^(−η) · net_mass_i · q^(1−2α)`, with
695/// `η = 2s + d(2−2α)` for the available dimension parameter, and writes, per requested
696/// form, one weight triple `[w_R, w_2, w_3]`. Only `w_R` is live:
697/// `R = CᵀWC − B·G⁺·Bᵀ/q`, with `G⁺` the rank-revealing pseudo-inverse.
698/// The extra slots are retained for the ψ layout and receive zero local
699/// channels because τ no longer changes the energy.
700///
701/// The outer sum over centers is coarsened per scale to a deterministic
702/// ε/2-net with nearest-member mass aggregation (the outer Riemann sum needs
703/// resolution ε, not the center-spacing floor), so each scale's cost sits at
704/// its own level and the band totals ~O(m²·d) instead of O(L·m³). The inner
705/// (local-fit) quadrature always uses the full center set, so the local
706/// residual identities (exact constant annihilation, PSD) are untouched.
707pub(crate) fn assemble_weighted_forms<F>(
708    centers: ArrayView2<'_, f64>,
709    masses: ArrayView1<'_, f64>,
710    band: &MeasureJetBand,
711    order_s: f64,
712    alpha: f64,
713    tau0: f64,
714    n_forms: usize,
715    channels: usize,
716    weights: &F,
717) -> Result<Vec<Array2<f64>>, BasisError>
718where
719    F: Fn(usize, f64, f64, f64, &mut [[f64; 3]]) + Sync,
720{
721    let m = centers.nrows();
722    let d = centers.ncols();
723    if n_forms == 0 || !(1..=3).contains(&channels) {
724        crate::bail_invalid_basis!(
725            "measure-jet assembly needs at least one output form and 1..=3 block channels"
726        );
727    }
728    if masses.len() != m {
729        crate::bail_dim_basis!(
730            "measure-jet energy mass/center mismatch: {} masses for {} centers",
731            masses.len(),
732            m
733        );
734    }
735    if band.eps.is_empty() || band.eps.iter().any(|e| !(e.is_finite() && *e > 0.0)) {
736        crate::bail_invalid_basis!("measure-jet energy needs a nonempty positive scale band");
737    }
738    if !(order_s.is_finite() && order_s > 0.0 && order_s < 2.0) {
739        crate::bail_invalid_basis!(
740            "measure-jet order s must lie in (0, 2) for the affine-jet energy; got {order_s}"
741        );
742    }
743    if !(alpha.is_finite() && tau0.is_finite() && tau0 >= 0.0) {
744        crate::bail_invalid_basis!(
745            "measure-jet energy needs finite alpha and finite tau0 >= 0; got alpha={alpha}, tau0={tau0}"
746        );
747    }
748    if masses.iter().any(|v| !(v.is_finite() && *v >= 0.0)) {
749        crate::bail_invalid_basis!("measure-jet energy needs finite nonnegative center masses");
750    }
751    let dist2 = pairwise_sq_dists(centers, centers);
752
753    // One block of `n_forms` m×m accumulators per scale. Each scale's center
754    // loop is sequential and the cross-scale sum below runs in band order,
755    // so the result is bit-deterministic whether or not the scales
756    // themselves run in parallel.
757    let assemble_scale = |scale_idx: usize, eps: f64| -> Result<Vec<Array2<f64>>, BasisError> {
758        let mut out: Vec<Array2<f64>> =
759            (0..n_forms).map(|_| Array2::<f64>::zeros((m, m))).collect();
760        let cutoff2 = (MEASURE_JET_PROFILE_CUTOFF * eps) * (MEASURE_JET_PROFILE_CUTOFF * eps);
761        let inv_two_eps2 = 1.0 / (2.0 * eps * eps);
762        let eta = 2.0 * order_s + (d as f64) * (2.0 - 2.0 * alpha);
763        let scale_weight = band.log_step * eps.powf(-eta);
764        // Outer-quadrature coarsening: greedy ε/2-net over the centers in
765        // fixed index order (deterministic), with every center's mass
766        // aggregated to its nearest net member (lowest-index tie break).
767        let net_radius2 = 0.25 * eps * eps;
768        let mut outer: Vec<usize> = Vec::new();
769        for i in 0..m {
770            if masses[i] <= 0.0 {
771                continue;
772            }
773            let covered = outer.iter().any(|&o| dist2[(i, o)] <= net_radius2);
774            if !covered {
775                outer.push(i);
776            }
777        }
778        let mut net_mass = vec![0.0_f64; m];
779        for i in 0..m {
780            if masses[i] <= 0.0 {
781                continue;
782            }
783            let mut best = f64::INFINITY;
784            let mut best_o = usize::MAX;
785            for &o in &outer {
786                if dist2[(i, o)] < best {
787                    best = dist2[(i, o)];
788                    best_o = o;
789                }
790            }
791            if best_o != usize::MAX {
792                net_mass[best_o] += masses[i];
793            }
794        }
795        let mut wbuf = vec![[0.0_f64; 3]; n_forms];
796        for &i in &outer {
797            // Local neighbor set (always includes i itself).
798            let mut idx: Vec<usize> = Vec::new();
799            for j in 0..m {
800                if dist2[(i, j)] <= cutoff2 {
801                    idx.push(j);
802                }
803            }
804            let ml = idx.len();
805            // Kernel weights and mass.
806            let mut w = Array1::<f64>::zeros(ml);
807            let mut q = 0.0_f64;
808            for (a, &j) in idx.iter().enumerate() {
809                let wj = masses[j] * (-dist2[(i, j)] * inv_two_eps2).exp();
810                w[a] = wj;
811                q += wj;
812            }
813            if !(q > 0.0) {
814                continue;
815            }
816            // Scaled local features Φ (ml × d) and weighted column means a.
817            let mut phi = Array2::<f64>::zeros((ml, d));
818            for (a, &j) in idx.iter().enumerate() {
819                for k in 0..d {
820                    phi[(a, k)] = (centers[(j, k)] - centers[(i, k)]) / eps;
821                }
822            }
823            let a_mean = phi.t().dot(&w) / q;
824            // B = WΦ − w·aᵀ and G = (ΦᵀWΦ)/q − a·aᵀ.
825            let mut wphi = phi.clone();
826            for (a, mut row) in wphi.outer_iter_mut().enumerate() {
827                row.mapv_inplace(|v| v * w[a]);
828            }
829            let mut b = wphi.clone();
830            for (a, mut row) in b.outer_iter_mut().enumerate() {
831                for k in 0..d {
832                    row[k] -= w[a] * a_mean[k];
833                }
834            }
835            let mut g = phi.t().dot(&wphi);
836            g.mapv_inplace(|v| v / q);
837            for r in 0..d {
838                for c in 0..d {
839                    g[(r, c)] -= a_mean[r] * a_mean[c];
840                }
841            }
842            let g_pinv = symmetric_pseudoinverse(&g, "local affine Gram")?;
843            let bm = b.dot(&g_pinv);
844            let base = scale_weight * net_mass[i] * q.powf(1.0 - 2.0 * alpha);
845            weights(scale_idx, eps, q, base, &mut wbuf);
846            // Scatter-add Σ_k wbuf[k]·R into each form. The τ channels are
847            // zero because the exact projection is τ-independent.
848            for (a, &ja) in idx.iter().enumerate() {
849                let bma = bm.row(a);
850                for (c, &jc) in idx.iter().enumerate() {
851                    let b_c = b.row(c);
852                    let mut val_r = -w[a] * w[c] / q - bma.dot(&b_c) / q;
853                    if a == c {
854                        val_r += w[a];
855                    }
856                    for (k, out_k) in out.iter_mut().enumerate() {
857                        let wk = wbuf[k];
858                        out_k[(ja, jc)] += wk[0] * val_r;
859                    }
860                }
861            }
862        }
863        Ok(out)
864    };
865
866    let n_scales = band.eps.len();
867    let parallel_ok = m
868        .saturating_mul(m)
869        .saturating_mul(n_scales)
870        .saturating_mul(n_forms)
871        <= MEASURE_JET_PARALLEL_FORM_BUDGET_DOUBLES;
872    let per_scale: Vec<Vec<Array2<f64>>> = if parallel_ok {
873        band.eps
874            .par_iter()
875            .enumerate()
876            .map(|(scale_idx, &eps)| assemble_scale(scale_idx, eps))
877            .collect::<Result<Vec<_>, BasisError>>()?
878    } else {
879        band.eps
880            .iter()
881            .enumerate()
882            .map(|(scale_idx, &eps)| assemble_scale(scale_idx, eps))
883            .collect::<Result<Vec<_>, BasisError>>()?
884    };
885
886    let mut totals: Vec<Array2<f64>> = (0..n_forms).map(|_| Array2::<f64>::zeros((m, m))).collect();
887    for scale_forms in per_scale {
888        for (total, part) in totals.iter_mut().zip(scale_forms) {
889            *total += &part;
890        }
891    }
892    // Numerical symmetrization (every analytic form here is symmetric).
893    Ok(totals.into_iter().map(|t| (&t + &t.t()) * 0.5).collect())
894}
895
896/// The multiscale jet-residual energy `Q` (m × m, symmetric PSD) on the
897/// center set. See the module docs for the formula and contracts; the local
898/// residual form is assembled through the closed-form identities
899///
900/// ```text
901///   CᵀWC          = W − w·wᵀ/q,
902///   B = CᵀWΦ̃     = WΦ − w·aᵀ          (a = Φᵀw/q),
903///   G = Φ̃ᵀWΦ̃/q  = (ΦᵀWΦ)/q − a·aᵀ,
904///   R_loc         = CᵀWC − B·G⁺·Bᵀ/q,
905/// ```
906///
907/// with `G⁺` realized through the symmetric eigendecomposition and a
908/// machine-precision rank cutoff. One walk of [`assemble_weighted_forms`]
909/// with the unit weight.
910pub fn measure_jet_energy_form(
911    centers: ArrayView2<'_, f64>,
912    masses: ArrayView1<'_, f64>,
913    band: &MeasureJetBand,
914    order_s: f64,
915    alpha: f64,
916    tau0: f64,
917) -> Result<Array2<f64>, BasisError> {
918    let mut forms = assemble_weighted_forms(
919        centers,
920        masses,
921        band,
922        order_s,
923        alpha,
924        tau0,
925        1,
926        1,
927        &|_, _, _, base, out: &mut [[f64; 3]]| out[0] = [base, 0.0, 0.0],
928    )?;
929    let q = forms.swap_remove(0);
930    // The energy `Q = Σ wᵢ Rᵢ` is a nonnegative combination of analytically
931    // PSD local residual forms, so it is PSD in exact arithmetic. The affine
932    // span is annihilated to machine zero, where roundoff in the per-block
933    // pseudo-inverse and the centering cancellation leaves the smallest
934    // eigenvalue at ±ε_mach·‖Q‖. Project onto the PSD cone (floor negative
935    // eigenvalues at 0) so `vᵀQv ≥ 0` holds exactly for every `v`, including
936    // the affine directions the energy must annihilate.
937    project_symmetric_psd(q, "measure-jet energy form")
938}
939
940/// Project a symmetric matrix onto the PSD cone by flooring its negative
941/// eigenvalues at 0. Only sub-machine-precision negative eigenvalues are
942/// expected here (the form is analytically PSD); a meaningfully negative
943/// eigenvalue would indicate an assembly bug, so it is floored but the
944/// reconstruction otherwise preserves the spectrum exactly.
945pub(crate) fn project_symmetric_psd(
946    a: Array2<f64>,
947    label: &str,
948) -> Result<Array2<f64>, BasisError> {
949    let n = a.nrows();
950    if n == 0 {
951        return Ok(a);
952    }
953    let (evals, evecs) = a.eigh(Side::Lower).map_err(|e| {
954        BasisError::InvalidInput(format!(
955            "measure-jet PSD projection `{label}` eigendecomposition failed: {e}"
956        ))
957    })?;
958    if evals.iter().all(|&lam| lam >= 0.0) {
959        return Ok(a);
960    }
961    let mut scaled = evecs.clone();
962    for (k, mut col) in scaled.axis_iter_mut(Axis(1)).enumerate() {
963        let lam = evals[k].max(0.0);
964        col.mapv_inplace(|v| v * lam);
965    }
966    let psd = scaled.dot(&evecs.t());
967    Ok((&psd + &psd.t()) * 0.5)
968}
969
970/// The energy together with its exact first and second jets in the live
971/// dials, plus zero slots for the retained `ψ_τ = ln τ` coordinate. With
972/// `g_s = −2 ln ε`, `g_α = −2 ln q`:
973///
974/// ```text
975///   ∂Q/∂s   = Σ g_s·w·R,        ∂²Q/∂s²   = Σ g_s²·w·R,
976///   ∂Q/∂α   = Σ g_α·w·R,        ∂²Q/∂α²   = Σ g_α²·w·R,
977///   ∂²Q/∂s∂α = Σ g_s·g_α·w·R,
978///   ∂Q/∂ψ_τ = ∂²Q/∂ψ_τ² = ∂²Q/∂s∂ψ_τ = ∂²Q/∂α∂ψ_τ = 0.
979/// ```
980///
981/// all scattered from the SAME local blocks as `Q` in one pass (no second
982/// assembly that could drift). FD-gated in this module's tests. Requires
983/// `tau0 > 0` only because the retained coordinate is `ln τ`.
984pub fn measure_jet_energy_form_with_jets(
985    centers: ArrayView2<'_, f64>,
986    masses: ArrayView1<'_, f64>,
987    band: &MeasureJetBand,
988    order_s: f64,
989    alpha: f64,
990    tau0: f64,
991) -> Result<MeasureJetEnergyJets, BasisError> {
992    if !(tau0.is_finite() && tau0 > 0.0) {
993        crate::bail_invalid_basis!(
994            "measure-jet jets need tau0 > 0 because the retained τ coordinate is ln τ; got {tau0}"
995        );
996    }
997    let mut forms = assemble_weighted_forms(
998        centers,
999        masses,
1000        band,
1001        order_s,
1002        alpha,
1003        tau0,
1004        10,
1005        3,
1006        &|_, eps: f64, q: f64, base: f64, out: &mut [[f64; 3]]| {
1007            let gs = -2.0 * eps.ln();
1008            let intrinsic_dim = centers.ncols() as f64;
1009            let ga = 2.0 * intrinsic_dim * eps.ln() - 2.0 * q.max(f64::MIN_POSITIVE).ln();
1010            out[0] = [base, 0.0, 0.0];
1011            out[1] = [gs * base, 0.0, 0.0];
1012            out[2] = [gs * gs * base, 0.0, 0.0];
1013            out[3] = [ga * base, 0.0, 0.0];
1014            out[4] = [ga * ga * base, 0.0, 0.0];
1015            out[5] = [gs * ga * base, 0.0, 0.0];
1016            out[6] = [0.0, 0.0, 0.0];
1017            out[7] = [0.0, 0.0, 0.0];
1018            out[8] = [0.0, 0.0, 0.0];
1019            out[9] = [0.0, 0.0, 0.0];
1020        },
1021    )?;
1022    let d2q_dalpha_dlogtau = forms.pop().expect("ten assembled forms");
1023    let d2q_ds_dlogtau = forms.pop().expect("ten assembled forms");
1024    let d2q_dlogtau2 = forms.pop().expect("ten assembled forms");
1025    let dq_dlogtau = forms.pop().expect("ten assembled forms");
1026    let d2q_ds_dalpha = forms.pop().expect("ten assembled forms");
1027    let d2q_dalpha2 = forms.pop().expect("ten assembled forms");
1028    let dq_dalpha = forms.pop().expect("ten assembled forms");
1029    let d2q_ds2 = forms.pop().expect("ten assembled forms");
1030    let dq_ds = forms.pop().expect("ten assembled forms");
1031    let q = forms.pop().expect("ten assembled forms");
1032    Ok(MeasureJetEnergyJets {
1033        q,
1034        dq_ds,
1035        d2q_ds2,
1036        dq_dalpha,
1037        d2q_dalpha2,
1038        d2q_ds_dalpha,
1039        dq_dlogtau,
1040        d2q_dlogtau2,
1041        d2q_ds_dlogtau,
1042        d2q_dalpha_dlogtau,
1043    })
1044}
1045
1046/// Per-scale energy decomposition of center values `v`: element ℓ is
1047/// `vᵀ Q_ℓ v`, the detail energy charged at scale `ε_ℓ`. Sums exactly to
1048/// `vᵀQv` (same blocks, one-hot weights) and doubles as the scale spectrum
1049/// diagnostic of the fitted intensity field — where along the band the
1050/// signal lives, and the analytic carrier of `∂/∂s` reweightings.
1051pub fn measure_jet_scale_spectrum(
1052    centers: ArrayView2<'_, f64>,
1053    masses: ArrayView1<'_, f64>,
1054    band: &MeasureJetBand,
1055    order_s: f64,
1056    alpha: f64,
1057    tau0: f64,
1058    values: ArrayView1<'_, f64>,
1059) -> Result<Vec<f64>, BasisError> {
1060    if values.len() != centers.nrows() {
1061        crate::bail_dim_basis!(
1062            "measure-jet scale spectrum needs one value per center: {} values for {} centers",
1063            values.len(),
1064            centers.nrows()
1065        );
1066    }
1067    let forms = measure_jet_energy_forms_per_scale(centers, masses, band, order_s, alpha, tau0)?;
1068    Ok(forms
1069        .iter()
1070        .map(|q_l| values.dot(&q_l.dot(&values)))
1071        .collect())
1072}
1073
1074/// The per-scale energy forms `Q_ℓ` (each m × m, symmetric PSD), with
1075/// `Σ_ℓ Q_ℓ = Q` exactly (same blocks, one-hot weights). These are the
1076/// spectral-split carriers: emitted as separate penalty candidates they let
1077/// the multi-penalty REML engine learn per-level amplitudes λ_ℓ directly —
1078/// scale adaptivity at ρ-speed with no rebuild and no new optimizer code.
1079pub fn measure_jet_energy_forms_per_scale(
1080    centers: ArrayView2<'_, f64>,
1081    masses: ArrayView1<'_, f64>,
1082    band: &MeasureJetBand,
1083    order_s: f64,
1084    alpha: f64,
1085    tau0: f64,
1086) -> Result<Vec<Array2<f64>>, BasisError> {
1087    let n_scales = band.eps.len();
1088    assemble_weighted_forms(
1089        centers,
1090        masses,
1091        band,
1092        order_s,
1093        alpha,
1094        tau0,
1095        n_scales,
1096        1,
1097        &|scale_idx, _, _, base, out: &mut [[f64; 3]]| {
1098            for (k, slot) in out.iter_mut().enumerate() {
1099                *slot = if k == scale_idx {
1100                    [base, 0.0, 0.0]
1101                } else {
1102                    [0.0, 0.0, 0.0]
1103                };
1104            }
1105        },
1106    )
1107}
1108
1109/// The support diagnostic `ε ↦ q_ε(x★)`: kernel mass of the (frozen) center
1110/// quadrature seen from each query point at every band scale (n_query × L).
1111/// A query ON the web sees its strand's mass already at fine scales; a query
1112/// OFF the web accumulates mass only once ε reaches its distance to the
1113/// support. This is the on-web-ness statistic shipped alongside predictions
1114/// — smooth, multiresolution, derived from the measure with no neighbor
1115/// sets.
1116pub fn measure_jet_support_curve(
1117    queries: ArrayView2<'_, f64>,
1118    centers: ArrayView2<'_, f64>,
1119    masses: ArrayView1<'_, f64>,
1120    eps_band: &[f64],
1121) -> Result<Array2<f64>, BasisError> {
1122    if queries.ncols() != centers.ncols() {
1123        crate::bail_dim_basis!(
1124            "measure-jet support curve dimension mismatch: queries d={} centers d={}",
1125            queries.ncols(),
1126            centers.ncols()
1127        );
1128    }
1129    if masses.len() != centers.nrows() {
1130        crate::bail_dim_basis!(
1131            "measure-jet support curve mass/center mismatch: {} masses for {} centers",
1132            masses.len(),
1133            centers.nrows()
1134        );
1135    }
1136    if eps_band.is_empty() || eps_band.iter().any(|e| !(e.is_finite() && *e > 0.0)) {
1137        crate::bail_invalid_basis!("measure-jet support curve needs a nonempty positive band");
1138    }
1139    validate_finite_points(queries, "queries")?;
1140    validate_finite_points(centers, "centers")?;
1141    let nq = queries.nrows();
1142    let nl = eps_band.len();
1143    // Distances once (GEMM), then every band scale reads the same d² row —
1144    // an L-fold saving over per-scale distance recomputation.
1145    let d2 = pairwise_sq_dists(queries, centers);
1146    let mut out = Array2::<f64>::zeros((nq, nl));
1147    out.axis_iter_mut(Axis(0))
1148        .into_par_iter()
1149        .enumerate()
1150        .for_each(|(qi, mut row)| {
1151            let d2_row = d2.row(qi);
1152            for (li, &eps) in eps_band.iter().enumerate() {
1153                let inv_two_eps2 = 1.0 / (2.0 * eps * eps);
1154                let mut acc = 0.0_f64;
1155                for (j, &dd) in d2_row.iter().enumerate() {
1156                    acc += masses[j] * (-dd * inv_two_eps2).exp();
1157                }
1158                row[li] = acc;
1159            }
1160        });
1161    Ok(out)
1162}
1163
1164pub(crate) fn measure_jet_support_means(
1165    centers: ArrayView2<'_, f64>,
1166    masses: ArrayView1<'_, f64>,
1167    eps_band: &[f64],
1168) -> Result<Vec<f64>, BasisError> {
1169    let total_mass = masses.sum();
1170    if !(total_mass.is_finite() && total_mass > 0.0) {
1171        crate::bail_invalid_basis!(
1172            "measure-jet support means need positive finite total mass; got {total_mass}"
1173        );
1174    }
1175    let support = measure_jet_support_curve(centers, centers, masses, eps_band)?;
1176    let mut means = vec![0.0_f64; eps_band.len()];
1177    for (i, row) in support.rows().into_iter().enumerate() {
1178        let mass = masses[i];
1179        for (mean, &q) in means.iter_mut().zip(row.iter()) {
1180            *mean += mass * q;
1181        }
1182    }
1183    for mean in &mut means {
1184        *mean /= total_mass;
1185        if !(*mean).is_finite() || *mean <= 0.0 {
1186            crate::bail_invalid_basis!(
1187                "measure-jet support mean must be positive and finite; got {mean}"
1188            );
1189        }
1190    }
1191    Ok(means)
1192}
1193
1194/// Gaussian representer features `exp(−‖x − c‖²/(2ℓ²))` (n × m).
1195pub fn measure_jet_design_matrix(
1196    data: ArrayView2<'_, f64>,
1197    centers: ArrayView2<'_, f64>,
1198    length_scale: f64,
1199) -> Result<Array2<f64>, BasisError> {
1200    if data.ncols() != centers.ncols() {
1201        crate::bail_dim_basis!(
1202            "measure-jet design dimension mismatch: data d={} centers d={}",
1203            data.ncols(),
1204            centers.ncols()
1205        );
1206    }
1207    if !(length_scale.is_finite() && length_scale > 0.0) {
1208        crate::bail_invalid_basis!(
1209            "measure-jet design needs a positive finite length_scale; got {length_scale}"
1210        );
1211    }
1212    validate_finite_points(data, "data")?;
1213    validate_finite_points(centers, "centers")?;
1214    let inv_two_l2 = 1.0 / (2.0 * length_scale * length_scale);
1215    // One GEMM for every distance, then the Gaussian applied in place — the
1216    // n×m allocation IS the output, no transient copy.
1217    let mut out = pairwise_sq_dists(data, centers);
1218    out.axis_iter_mut(Axis(0))
1219        .into_par_iter()
1220        .for_each(|mut row| {
1221            row.mapv_inplace(|d2| (-d2 * inv_two_l2).exp());
1222        });
1223    Ok(out)
1224}
1225
1226/// Rank-revealing ambient-linear head lift `T` (d × head_rank) for the
1227/// extrapolation null space (#1845).
1228///
1229/// The measure-jet energy annihilates ambient-affine functions EXACTLY (the
1230/// no-mass contract), so the affine functions are the penalty's null space —
1231/// the directions the fit is free to extend across a training gap. But the
1232/// Gaussian representer design cannot REPRESENT a global affine function off
1233/// its support: a finite sum of decaying bumps reverts to the parametric
1234/// backbone away from the centers, so in a gap the fit collapses toward the
1235/// training mean instead of carrying the flank-attested trend. Completing the
1236/// smoothing-spline structure, the builder appends this ambient-linear null
1237/// space to the design as an UNPENALIZED head (the `{x_1..x_d}` head the frame
1238/// notes §1 pin as the property the representer basis lacked).
1239///
1240/// The head is data-derived and magic-free. Ambient coordinates of data on a
1241/// low intrinsic-dimension stratum are rank-deficient as linear trends, so the
1242/// coordinate columns are orthonormalized on the centers and the
1243/// numerically-degenerate directions dropped. Working in the mean-CENTERED
1244/// coordinate columns (the mass-weighted mean is the intercept's, not the
1245/// head's) makes the rank test measure the genuine spread of the centers along
1246/// each direction rather than its offset; the relative floor
1247/// [`MEASURE_JET_PSEUDOINVERSE_RTOL`] is the module's own numerical rank
1248/// tolerance (the same one the local Gram pseudo-inverses use). The returned
1249/// `T` satisfies `head(points) = points · T` (the mean-centering only informs
1250/// the keep/drop decision; the constant component of `points · T` is removed
1251/// downstream by the global parametric orthogonalization). `T` is a
1252/// deterministic function of the frozen centers + masses, so the frozen replay
1253/// path reconstructs the identical head with no persisted state.
1254fn measure_jet_affine_head_transform(
1255    centers: ArrayView2<'_, f64>,
1256    masses: ArrayView1<'_, f64>,
1257) -> Array2<f64> {
1258    let m = centers.nrows();
1259    let d = centers.ncols();
1260    let total_mass = masses.sum();
1261    // Mass inner product on center values.
1262    let mdot = |u: &Array1<f64>, v: &Array1<f64>| -> f64 {
1263        let mut acc = 0.0;
1264        for i in 0..m {
1265            acc += masses[i] * u[i] * v[i];
1266        }
1267        acc
1268    };
1269    // Mean-centered coordinate columns: the mass-weighted mean is removed so the
1270    // residual mass-norm is the genuine spread of the centers along a direction,
1271    // not dominated by the coordinate's offset (which the intercept owns).
1272    let cols: Vec<Array1<f64>> = (0..d)
1273        .map(|k| {
1274            let col = centers.column(k).to_owned();
1275            let mean = if total_mass > 0.0 {
1276                mdot(&col, &Array1::ones(m)) / total_mass
1277            } else {
1278                0.0
1279            };
1280            col.mapv(|x| x - mean)
1281        })
1282        .collect();
1283    // Relative numerical rank floor from the centered coordinate-column scale.
1284    let max_norm = cols.iter().fold(0.0_f64, |acc, c| acc.max(mdot(c, c).sqrt()));
1285    let drop_below =
1286        (MEASURE_JET_PSEUDOINVERSE_RTOL * (d.max(1) as f64) * max_norm).max(f64::MIN_POSITIVE);
1287    // Mass-weighted modified Gram–Schmidt on the centered columns; `t`
1288    // accumulates the lift in the ORIGINAL coordinate basis, so every kept head
1289    // column is `points · t_r` (up to the intercept-owned constant).
1290    let mut q_cols: Vec<Array1<f64>> = Vec::new();
1291    let mut t_cols: Vec<Array1<f64>> = Vec::new();
1292    for k in 0..d {
1293        let mut v = cols[k].clone();
1294        let mut t = Array1::<f64>::zeros(d);
1295        t[k] = 1.0;
1296        for (q, tq) in q_cols.iter().zip(t_cols.iter()) {
1297            let proj = mdot(q, &v);
1298            v.scaled_add(-proj, q);
1299            t.scaled_add(-proj, tq);
1300        }
1301        let norm = mdot(&v, &v).sqrt();
1302        if norm > drop_below {
1303            v.mapv_inplace(|x| x / norm);
1304            t.mapv_inplace(|x| x / norm);
1305            q_cols.push(v);
1306            t_cols.push(t);
1307        }
1308    }
1309    let head_rank = t_cols.len();
1310    let mut t_mat = Array2::<f64>::zeros((d, head_rank));
1311    for (r, t) in t_cols.into_iter().enumerate() {
1312        t_mat.column_mut(r).assign(&t);
1313    }
1314    t_mat
1315}
1316
1317/// Resolve the realized representer range ℓ. An explicit positive
1318/// `spec_length_scale` is used verbatim; the `0.0` sentinel auto-initializes
1319/// from the median nearest-center spacing (one spacing width: neighbors
1320/// overlap at exp(−1/2) ≈ 0.61, smooth blend without collinearity).
1321pub fn realized_measure_jet_length_scale(
1322    centers: ArrayView2<'_, f64>,
1323    spec_length_scale: f64,
1324) -> Result<f64, BasisError> {
1325    if spec_length_scale.is_finite() && spec_length_scale > 0.0 {
1326        return Ok(spec_length_scale);
1327    }
1328    if spec_length_scale != 0.0 {
1329        crate::bail_invalid_basis!(
1330            "measure-jet length_scale must be positive (or 0.0 for auto); got {spec_length_scale}"
1331        );
1332    }
1333    let dist2 = pairwise_sq_dists(centers, centers);
1334    let spacing = median_nearest_center_spacing(&dist2)?;
1335    Ok(MEASURE_JET_AUTO_LENGTH_SCALE_FACTOR * spacing)
1336}
1337
1338/// The realized, ψ-FIXED geometry shared by the basis builder and the
1339/// ψ-derivative producer — ONE realization source, so the penalty the fit
1340/// uses and the penalty the ψ-channel differentiates can never drift apart
1341/// (the #901 desync class, excluded structurally).
1342pub(crate) struct RealizedMeasureJetGeometry {
1343    pub(crate) centers: Array2<f64>,
1344    pub(crate) masses: Array1<f64>,
1345    pub(crate) eps_band: Vec<f64>,
1346    pub(crate) log_step: f64,
1347    pub(crate) length_scale: f64,
1348    /// Assembly order for the energy weights: the realized default in
1349    /// per-level mode (absorbed per candidate by normalization), the
1350    /// explicit value in fused mode.
1351    pub(crate) order_s_eval: f64,
1352    /// Spectral-split mode marker (`order_s == 0.0` sentinel).
1353    pub(crate) per_level: bool,
1354    pub(crate) z: Array2<f64>,
1355    pub(crate) coefficient_gauge: gam_problem::Gauge,
1356    pub(crate) kz: Array2<f64>,
1357    /// Ambient-linear head lift `T` (d × head_rank): the extrapolation
1358    /// null-space basis appended to the representer design (#1845). The head
1359    /// columns evaluate as `points · T`; empty (`d × 0`) when the geometry
1360    /// resolves no supported linear direction. Deterministic in the frozen
1361    /// centers + masses, so predict-time replay rebuilds it verbatim.
1362    pub(crate) head_transform: Array2<f64>,
1363}
1364
1365pub(crate) fn realize_measure_jet_geometry(
1366    data: ArrayView2<'_, f64>,
1367    spec: &MeasureJetBasisSpec,
1368) -> Result<RealizedMeasureJetGeometry, BasisError> {
1369    if data.ncols() == 0 {
1370        crate::bail_invalid_basis!("measure-jet smooth needs at least one feature column");
1371    }
1372    validate_finite_points(data, "data")?;
1373    let seed_centers = select_centers_by_strategy(data, &spec.center_strategy)?;
1374    let m = seed_centers.nrows();
1375    if m < 3 {
1376        return Err(BasisError::InsufficientColumnsForConstraint { found: m });
1377    }
1378    let order_s = if spec.order_s == 0.0 {
1379        MEASURE_JET_DEFAULT_ORDER_S
1380    } else {
1381        spec.order_s
1382    };
1383    // Quadrature realization. Fit path: the realized nodes are the cell
1384    // BARYCENTERS of the seed partition (first-moment-exact lumping of μ —
1385    // see `measure_jet_quadrature_nodes`), so the metadata's `centers` are
1386    // already the realized nodes and the frozen path (predict / ψ-trial,
1387    // `CenterStrategy::UserProvided`) replays them verbatim with the frozen
1388    // masses, band, support anchors, and normalization scales.
1389    let (centers, masses, eps_band, log_step) = match &spec.frozen_quadrature {
1390        Some(frozen) => {
1391            if frozen.masses.len() != m {
1392                crate::bail_dim_basis!(
1393                    "frozen measure-jet quadrature mismatch: {} masses for {} centers",
1394                    frozen.masses.len(),
1395                    m
1396                );
1397            }
1398            if frozen.eps_band.is_empty() {
1399                crate::bail_invalid_basis!("frozen measure-jet quadrature has an empty band");
1400            }
1401            let log_step = if frozen.eps_band.len() >= 2 {
1402                (frozen.eps_band[1] / frozen.eps_band[0]).ln()
1403            } else {
1404                std::f64::consts::LN_2
1405            };
1406            (
1407                seed_centers,
1408                frozen.masses.clone(),
1409                frozen.eps_band.clone(),
1410                log_step,
1411            )
1412        }
1413        None => {
1414            let (nodes, masses) = measure_jet_quadrature_nodes(data, seed_centers.view())?;
1415            let band = measure_jet_band(nodes.view(), spec.num_scales)?;
1416            (nodes, masses, band.eps, band.log_step)
1417        }
1418    };
1419    let length_scale = realized_measure_jet_length_scale(centers.view(), spec.length_scale)?;
1420    // Ambient-linear extrapolation head (#1845): the raw center space becomes
1421    // `[ m Gaussian representers | head_rank ambient-linear columns ]`. The head
1422    // carries the penalty's affine null space explicitly so the fit no longer
1423    // reverts to the parametric backbone (the training mean) across an
1424    // unsupported gap.
1425    // The extrapolation head is the single-scale (fused) gap-bridge path. In
1426    // multiscale mode the per-scale spectral penalties carry their own
1427    // structure and the design stays the pure representer basis (the per-level
1428    // replay + width contracts pin `m − 1` columns), so the head is added only
1429    // when the term is single-scale.
1430    let head_transform = if spec.multiscale {
1431        Array2::<f64>::zeros((centers.ncols(), 0))
1432    } else {
1433        measure_jet_affine_head_transform(centers.view(), masses.view())
1434    };
1435    let head_rank = head_transform.ncols();
1436    let m_aug = m + head_rank;
1437    // Realized-design constraint transform: block sum-to-zero on the
1438    // representer coefficients composed with an identity pass-through on the
1439    // head at fit time; the frozen composed `z · z_parametric` at predict time
1440    // (#532 pattern — see MeasureJetIdentifiability). The head rides the SAME
1441    // global parametric orthogonalization, so its constant component cannot
1442    // collide with the intercept.
1443    let (z, coefficient_gauge) = match &spec.identifiability {
1444        MeasureJetIdentifiability::FrozenTransform { transform } => {
1445            if transform.nrows() != m_aug {
1446                crate::bail_dim_basis!(
1447                    "frozen measure-jet identifiability transform mismatch: {} representers + {} head columns but transform has {} rows",
1448                    m,
1449                    head_rank,
1450                    transform.nrows()
1451                );
1452            }
1453            (
1454                transform.clone(),
1455                gam_problem::Gauge::from_block_transforms(&[transform.clone()]),
1456            )
1457        }
1458        MeasureJetIdentifiability::CenterSumToZero => {
1459            // Householder sum-to-zero on the representers, block-composed with
1460            // the identity on the head columns. The active projection is owned
1461            // by Gauge; the dense z is persisted for frozen replay metadata.
1462            let u = householder_sum_to_zero_u(m);
1463            let z_rbf = householder_sum_to_zero_z(&u);
1464            let mut z_block = Array2::<f64>::zeros((m_aug, (m - 1) + head_rank));
1465            z_block.slice_mut(ndarray::s![..m, ..m - 1]).assign(&z_rbf);
1466            for r in 0..head_rank {
1467                z_block[(m + r, (m - 1) + r)] = 1.0;
1468            }
1469            (
1470                z_block.clone(),
1471                gam_problem::Gauge::from_block_transforms(&[z_block]),
1472            )
1473        }
1474    };
1475    // Augmented raw center matrix `[K(centers, centers) | centers · T]`, so the
1476    // restricted `kz` maps constrained coefficients to center nodal values for
1477    // BOTH the representers and the head; the energy annihilates the head block
1478    // (affine) to machine precision, so it stays the unpenalized null space.
1479    let k_cc = measure_jet_design_matrix(centers.view(), centers.view(), length_scale)?;
1480    let mut k_aug = Array2::<f64>::zeros((m, m_aug));
1481    k_aug.slice_mut(ndarray::s![.., ..m]).assign(&k_cc);
1482    if head_rank > 0 {
1483        let head_cc = centers.dot(&head_transform);
1484        k_aug.slice_mut(ndarray::s![.., m..]).assign(&head_cc);
1485    }
1486    let kz = coefficient_gauge.restrict_design(&k_aug);
1487    Ok(RealizedMeasureJetGeometry {
1488        centers,
1489        masses,
1490        eps_band,
1491        log_step,
1492        length_scale,
1493        order_s_eval: order_s,
1494        // multiscale (per-scale spectral) mode is an EXPLICIT opt-in
1495        // (#1116): single-scale (one fused penalty, the Duchon/Matérn outer
1496        // footprint) at any center count unless the spec asks for it. No
1497        // center-count auto-gate.
1498        per_level: spec.multiscale,
1499        z,
1500        coefficient_gauge,
1501        kz,
1502        head_transform,
1503    })
1504}
1505
1506/// Whether a measure-jet spec runs in multiscale mode (per-scale spectral
1507/// penalties + `(α, ln τ)` ψ dials + the affine-preserving ridge). The single
1508/// source of truth for the mode decision, shared by the builder and the
1509/// outer-engine enrollment predicates so the penalty count and the ψ-dimension
1510/// cannot disagree. Multiscale is an explicit opt-in (`spec.multiscale`); there
1511/// is no center-count auto-gate (#1116).
1512pub fn measure_jet_multiscale_mode(spec: &MeasureJetBasisSpec) -> bool {
1513    spec.multiscale
1514}
1515
1516/// Build the measure-jet smooth: Gaussian representer design `K(data,
1517/// centers)·z`, multiscale jet-residual penalty (one candidate per scale in
1518/// spectral mode, one fused candidate in pinned-order mode), and the
1519/// replayable [`BasisMetadata::MeasureJet`]. The geometry comes from the
1520/// empirical measure (centers + masses + band) through the shared
1521/// realization helper — the same source the ψ-derivative producer uses.
1522pub fn build_measure_jet_basis(
1523    data: ArrayView2<'_, f64>,
1524    spec: &MeasureJetBasisSpec,
1525) -> Result<BasisBuildResult, BasisError> {
1526    let RealizedMeasureJetGeometry {
1527        centers,
1528        masses,
1529        eps_band,
1530        log_step,
1531        length_scale,
1532        order_s_eval: order_s,
1533        per_level,
1534        z,
1535        coefficient_gauge,
1536        kz,
1537        head_transform,
1538    } = realize_measure_jet_geometry(data, spec)?;
1539    let band = MeasureJetBand {
1540        eps: eps_band.clone(),
1541        log_step,
1542    };
1543    let m = centers.nrows();
1544    let head_rank = head_transform.ncols();
1545    let m_aug = m + head_rank;
1546    // Augmented raw design `[K(data, centers) | data · T]` (#1845): the head
1547    // columns are the ambient-linear extrapolation basis. The gauge restricts
1548    // BOTH blocks together, so the frozen composed transform replays the head
1549    // verbatim at predict time.
1550    let kernel_design = measure_jet_design_matrix(data, centers.view(), length_scale)?;
1551    let mut raw_design = Array2::<f64>::zeros((data.nrows(), m_aug));
1552    raw_design
1553        .slice_mut(ndarray::s![.., ..m])
1554        .assign(&kernel_design);
1555    if head_rank > 0 {
1556        let head_design = data.dot(&head_transform);
1557        raw_design
1558            .slice_mut(ndarray::s![.., m..])
1559            .assign(&head_design);
1560    }
1561    let constrained_design = coefficient_gauge.restrict_design(&raw_design);
1562    let design = gam_linalg::matrix::DesignMatrix::Dense(
1563        gam_linalg::matrix::DenseDesignMatrix::from(constrained_design),
1564    );
1565    let support_means = measure_jet_support_means(centers.view(), masses.view(), &eps_band)?;
1566    // Spectral/geometric split. With the auto order sentinel (order_s == 0.0)
1567    // the term emits one candidate PER scale: the multi-penalty REML engine
1568    // then learns the level amplitudes λ_ℓ directly — scale adaptivity at
1569    // ρ-speed, dead scales REML-deselected (the Duchon-ARD pattern) — and the
1570    // fitted order is read off the spectrum (ŝ = −½ · slope of ln λ̂_ℓ on
1571    // ln ε_ℓ) instead of being optimized. An explicit s > 0 pins the Mellin
1572    // weights and fuses the band into one candidate. The Mellin prefactor
1573    // ε^(−η)·log_step inside each per-scale form is absorbed by the
1574    // per-candidate Frobenius normalization, so REML owns the amplitudes
1575    // outright. The sentinel itself is persisted in the metadata as the mode
1576    // marker: a replay MUST re-enter the same mode or the penalty count
1577    // desyncs (the gam#860 trap class).
1578    let mut candidates = Vec::new();
1579    let mut penalty_normalization_scales = Vec::new();
1580    let mut raw_penalty_normalization_scales = Vec::new();
1581    let mut fused_penalty_normalization_scale = None;
1582    if per_level {
1583        let forms = measure_jet_energy_forms_per_scale(
1584            centers.view(),
1585            masses.view(),
1586            &band,
1587            order_s,
1588            spec.alpha,
1589            spec.tau0,
1590        )?;
1591        for (level, q_l) in forms.into_iter().enumerate() {
1592            let s_l = kz.t().dot(&q_l).dot(&kz);
1593            let (s_norm, c_l) = normalize_penalty(&((&s_l + &s_l.t()) * 0.5));
1594            let intrinsic_dim = centers.ncols() as f64;
1595            let eta = 2.0 * order_s + intrinsic_dim * (2.0 - 2.0 * spec.alpha);
1596            let scale_weight = log_step * eps_band[level].powf(-eta);
1597            penalty_normalization_scales.push(c_l);
1598            raw_penalty_normalization_scales.push(c_l / scale_weight);
1599            candidates.push(PenaltyCandidate {
1600                matrix: s_norm,
1601                nullspace_dim_hint: 0,
1602                source: PenaltySource::Other(format!("measure_jet_scale_{level}")),
1603                normalization_scale: c_l,
1604                kronecker_factors: None,
1605                op: None,
1606            });
1607        }
1608    } else {
1609        let q_form = measure_jet_energy_form(
1610            centers.view(),
1611            masses.view(),
1612            &band,
1613            order_s,
1614            spec.alpha,
1615            spec.tau0,
1616        )?;
1617        let mut penalty = kz.t().dot(&q_form).dot(&kz);
1618        penalty = (&penalty + &penalty.t()) * 0.5;
1619        // Single-scale nullspace stabilization, FUSED (not a second REML λ).
1620        //
1621        // The fused jet-residual energy penalizes only the LOCAL jet residual,
1622        // so it has a large near-nullspace of
1623        // globally-smooth-but-not-locally-affine modes; on low-effective-rank
1624        // data those weakly-penalized directions are unidentified and the inner
1625        // solve / outer REML cycle (MEASURED #1116: with NO stabilizer the
1626        // bms-accuracy fit was 3s→59s and less accurate). The
1627        // affine-preserving ridge `I − P_affine` stabilizes exactly those
1628        // directions. But its job is IDENTIFIABILITY, not smoothing selection —
1629        // it does not need its own learned λ. As a separate candidate it
1630        // doubled the single-scale outer-REML search into 2-D (the #1116
1631        // perf-parity 2.8× vs Duchon's 1-λ fit). Instead fold it into the
1632        // primary at a small FIXED fraction of the primary's Frobenius scale:
1633        // the near-nullspace gets a definite floor (κ ≈ 1/ε_ridge bounded), the
1634        // smoothing strength stays owned by the single primary λ, and the outer
1635        // footprint is one λ — Duchon/Matérn parity. Multiscale mode keeps the
1636        // ridge as a genuine learned candidate (its per-scale spectrum IS the
1637        // estimand).
1638        if spec.double_penalty {
1639            // With the ambient-linear head carrying the affine null space
1640            // explicitly (#1845), the representer's own affine-reproducing
1641            // directions become redundant with the (unpenalized) head: leaving
1642            // them unpenalized would split the flank trend arbitrarily between a
1643            // decaying representer combination and the head, and only the head
1644            // survives across the gap. A plain identity floor on the representer
1645            // coefficients (`blockdiag(I_m, 0_head)` restricted through the same
1646            // gauge) breaks that tie deterministically toward the free head —
1647            // any affine mass placed on the representer now pays a strictly
1648            // positive toll, so the penalized-least-squares optimum drives it
1649            // into the head — while still flooring the energy's smooth
1650            // near-nullspace exactly as the affine-preserving ridge did. Without
1651            // a head the geometry keeps the original affine-preserving ridge.
1652            let ridge = if head_rank > 0 {
1653                let mut r_raw = Array2::<f64>::zeros((m_aug, m_aug));
1654                for i in 0..m {
1655                    r_raw[(i, i)] = 1.0;
1656                }
1657                coefficient_gauge.restrict_penalty(&r_raw)
1658            } else {
1659                affine_preserving_coefficient_ridge(&kz, centers.view(), masses.view())?
1660            };
1661            let primary_fro = trace_of_product(&penalty, &penalty).sqrt();
1662            let ridge_fro = trace_of_product(&ridge, &ridge).sqrt();
1663            if primary_fro.is_finite()
1664                && primary_fro > 0.0
1665                && ridge_fro.is_finite()
1666                && ridge_fro > 0.0
1667            {
1668                let w = MEASURE_JET_FUSED_RIDGE_FRACTION * primary_fro / ridge_fro;
1669                penalty = &penalty + &(&ridge * w);
1670            }
1671        }
1672        let (penalty_norm, c_primary) = normalize_penalty(&penalty);
1673        fused_penalty_normalization_scale = Some(c_primary);
1674        candidates.push(PenaltyCandidate {
1675            matrix: penalty_norm,
1676            nullspace_dim_hint: 0,
1677            source: PenaltySource::Primary,
1678            normalization_scale: c_primary,
1679            kronecker_factors: None,
1680            op: None,
1681        });
1682    }
1683    // Multiscale mode: the affine-preserving ridge is a genuine learned
1684    // candidate (the per-scale split leaves its own spectrum to estimate). In
1685    // single-scale mode it is fused into the primary above (a fixed
1686    // identifiability floor, not a second λ) — see that block.
1687    if spec.double_penalty && per_level {
1688        let ridge = affine_preserving_coefficient_ridge(&kz, centers.view(), masses.view())?;
1689        let (ridge_norm, c_ridge) = normalize_penalty(&ridge);
1690        candidates.push(PenaltyCandidate {
1691            matrix: ridge_norm,
1692            nullspace_dim_hint: 0,
1693            source: PenaltySource::DoublePenaltyNullspace,
1694            normalization_scale: c_ridge,
1695            kronecker_factors: None,
1696            op: None,
1697        });
1698    }
1699    let (penalties, nullspace_dims, penaltyinfo, null_eigenvectors, ops) =
1700        filter_active_penalty_candidates_with_ops(candidates)?;
1701    Ok(BasisBuildResult {
1702        design,
1703        penalties,
1704        nullspace_dims,
1705        penaltyinfo,
1706        metadata: BasisMetadata::MeasureJet {
1707            centers,
1708            input_scales: None,
1709            length_scale,
1710            eps_band,
1711            // The SPEC's order field, sentinel included: 0.0 marks per-level
1712            // (spectral) mode and must replay as per-level — persisting the
1713            // realized default here would silently flip the rebuild into
1714            // fused mode and desync the penalty count.
1715            order_s: spec.order_s,
1716            alpha: spec.alpha,
1717            tau0: spec.tau0,
1718            masses,
1719            support_means,
1720            penalty_normalization_scales,
1721            raw_penalty_normalization_scales,
1722            fused_penalty_normalization_scale,
1723            constraint_transform: Some(z),
1724        },
1725        kronecker_factored: None,
1726        ops,
1727        null_eigenvectors,
1728        joint_null_rotation: None,
1729    })
1730}
1731
1732/// Exact ψ-jets of the REALIZED measure-jet penalty candidates, adapted to
1733/// the anisotropic group-ψ carrier the spatial optimizer consumes.
1734///
1735/// Coordinates (the layout contract for the registration arm):
1736/// - per-level (spectral) mode, `order_s == 0.0`: `[α, ln τ]` — the order is
1737///   absorbed by the REML-learned per-scale amplitudes and is NOT a ψ
1738///   coordinate; `ln τ` is retained as an inert coordinate with zero
1739///   derivatives because the exact affine projection is τ-independent;
1740/// - fused (pinned-order) mode: `[s, α, ln τ]`.
1741///
1742/// Design drift is identically zero for every coordinate (the Gaussian
1743/// representer is ψ-fixed), so `design_first`/`design_second_diag` are
1744/// correctly-shaped zero matrices and there are no design cross terms.
1745/// Penalty derivatives are routed through the SAME constrained Frobenius
1746/// normalization as the fit-time candidates
1747/// (`normalize_penaltywith_psi_derivatives` + the cross rule), so criterion
1748/// value and criterion derivative share one normalization — the #901 lesson
1749/// made structural. The affine-preserving ridge candidate (when
1750/// `double_penalty` is on) carries identically-zero derivatives. The
1751/// per-candidate layout follows the builder's ORIGINAL candidate order
1752/// (scale candidates then ridge / primary then ridge); consumers align to
1753/// the FITTED penalty list via
1754/// `PenaltyInfo.original_index` when the active-candidate filter dropped
1755/// any.
1756pub fn build_measure_jet_basis_psi_derivatives(
1757    data: ArrayView2<'_, f64>,
1758    spec: &MeasureJetBasisSpec,
1759) -> Result<AnisoBasisPsiDerivatives, BasisError> {
1760    if !(spec.tau0.is_finite() && spec.tau0 > 0.0) {
1761        crate::bail_invalid_basis!(
1762            "measure-jet ψ derivatives need tau0 > 0 because the retained τ coordinate is ln τ; got {}",
1763            spec.tau0
1764        );
1765    }
1766    let geom = realize_measure_jet_geometry(data, spec)?;
1767    let band = MeasureJetBand {
1768        eps: geom.eps_band.clone(),
1769        log_step: geom.log_step,
1770    };
1771    let n = data.nrows();
1772    let p = geom.kz.ncols(); // design width: constrained Gaussian-coefficient space
1773    let kz = &geom.kz;
1774    let sandwich = |j: &Array2<f64>| {
1775        let s = kz.t().dot(j).dot(kz);
1776        (&s + &s.t()) * 0.5
1777    };
1778    // Raw (pre-normalization) value + jet stacks per ORIGINAL candidate, in
1779    // coordinate order. `raw[cand] = (S, firsts, seconds, crosses)` with
1780    // crosses keyed by coordinate pairs in `pairs` order.
1781    let (n_coords, pairs, raw): (
1782        usize,
1783        Vec<(usize, usize)>,
1784        Vec<(
1785            Array2<f64>,
1786            Vec<Array2<f64>>,
1787            Vec<Array2<f64>>,
1788            Vec<Array2<f64>>,
1789        )>,
1790    ) = if geom.per_level {
1791        let l_count = band.eps.len();
1792        // Six forms per scale: value, ∂α, ∂α², and zero τ slots — same
1793        // blocks, one walk (single-source rule).
1794        let forms = assemble_weighted_forms(
1795            geom.centers.view(),
1796            geom.masses.view(),
1797            &band,
1798            geom.order_s_eval,
1799            spec.alpha,
1800            spec.tau0,
1801            6 * l_count,
1802            3,
1803            &|scale_idx, eps: f64, q: f64, base: f64, out: &mut [[f64; 3]]| {
1804                for slot in out.iter_mut() {
1805                    *slot = [0.0, 0.0, 0.0];
1806                }
1807                let intrinsic_dim = geom.centers.ncols() as f64;
1808                let ga = 2.0 * intrinsic_dim * eps.ln() - 2.0 * q.max(f64::MIN_POSITIVE).ln();
1809                let k0 = 6 * scale_idx;
1810                out[k0] = [base, 0.0, 0.0];
1811                out[k0 + 1] = [ga * base, 0.0, 0.0];
1812                out[k0 + 2] = [ga * ga * base, 0.0, 0.0];
1813                out[k0 + 3] = [0.0, 0.0, 0.0];
1814                out[k0 + 4] = [0.0, 0.0, 0.0];
1815                out[k0 + 5] = [0.0, 0.0, 0.0];
1816            },
1817        )?;
1818        let mut raw = Vec::with_capacity(l_count);
1819        for level in 0..l_count {
1820            let chunk = &forms[6 * level..6 * level + 6];
1821            raw.push((
1822                sandwich(&chunk[0]),
1823                vec![sandwich(&chunk[1]), sandwich(&chunk[3])],
1824                vec![sandwich(&chunk[2]), sandwich(&chunk[4])],
1825                vec![sandwich(&chunk[5])],
1826            ));
1827        }
1828        (2usize, vec![(0usize, 1usize)], raw)
1829    } else {
1830        // Single-scale mode (the default): the penalty dials (s, α, lnτ) are NOT
1831        // enrolled — the fused penalty is fixed and only the design-moving ℓ
1832        // coordinate (added below) varies. So there are ZERO penalty-derivative
1833        // COORDINATES (n_coords = 0), but the term still has ONE fitted penalty
1834        // candidate (the fused primary), so we emit one `raw` entry with EMPTY
1835        // per-coordinate derivative vecs. That keeps the candidate COUNT
1836        // (n_active = 1) aligned with the build's penalty list while contributing
1837        // no penalty ψ-derivative (the penalty is ℓ-independent). MUST match
1838        // `measure_jet_penalty_psi_dim` (= 0 coords) and the build's 1 candidate.
1839        let q_value = sandwich(&measure_jet_energy_form(
1840            geom.centers.view(),
1841            geom.masses.view(),
1842            &band,
1843            geom.order_s_eval,
1844            spec.alpha,
1845            spec.tau0,
1846        )?);
1847        let raw = vec![(q_value, Vec::new(), Vec::new(), Vec::new())];
1848        (0usize, Vec::new(), raw)
1849    };
1850    // Design-moving ℓ coordinate (#1116). The Gaussian representer design
1851    // `X = K(data, centers; ℓ)·z` is the ONLY ℓ-dependent object (the jet-energy
1852    // penalty uses the ε-band, not ℓ; the constraint transform `z` acts in
1853    // center-coefficient space). With `u = lnℓ` and `a = r²/ℓ²`:
1854    //   ∂K/∂u = K ⊙ a,     ∂²K/∂u² = K ⊙ (a² − 2a),
1855    // and `∂X/∂u = (∂K/∂u)·z`, `∂²X/∂u² = (∂²K/∂u²)·z`. Cross terms with the
1856    // (s, α, lnτ) penalty dials vanish: those carry zero design drift and the
1857    // penalty carries zero ℓ drift. FD-gated by `psi_producer_matches_fd_length_scale`.
1858    let length_scale_design = if spec.learn_length_scale {
1859        let ell = geom.length_scale;
1860        let k = measure_jet_design_matrix(data, geom.centers.view(), ell)?;
1861        let r2 = pairwise_sq_dists(data, geom.centers.view());
1862        let inv_l2 = 1.0 / (ell * ell);
1863        let mut dk = k.clone();
1864        let mut d2k = k.clone();
1865        for ((dk_v, d2k_v), &r2_v) in dk.iter_mut().zip(d2k.iter_mut()).zip(r2.iter()) {
1866            let a = r2_v * inv_l2;
1867            let kij = *dk_v;
1868            *dk_v = kij * a;
1869            *d2k_v = kij * (a * a - 2.0 * a);
1870        }
1871        // The ambient-linear head (#1845) is ℓ-independent, so its derivative
1872        // columns are exactly zero: pad the representer derivatives to the
1873        // augmented raw width before restricting through the SAME Gauge-owned
1874        // coefficient section the design uses.
1875        let m = geom.centers.nrows();
1876        let m_aug = m + geom.head_transform.ncols();
1877        let mut dk_aug = Array2::<f64>::zeros((data.nrows(), m_aug));
1878        let mut d2k_aug = Array2::<f64>::zeros((data.nrows(), m_aug));
1879        dk_aug.slice_mut(ndarray::s![.., ..m]).assign(&dk);
1880        d2k_aug.slice_mut(ndarray::s![.., ..m]).assign(&d2k);
1881        let dx_du = geom.coefficient_gauge.restrict_design(&dk_aug);
1882        let d2x_du2 = geom.coefficient_gauge.restrict_design(&d2k_aug);
1883        Some((dx_du, d2x_du2))
1884    } else {
1885        None
1886    };
1887    let n_active = raw.len();
1888    // Mirror the build-path gate (`build_measure_jet_basis`): the
1889    // affine-preserving ridge is a SEPARATE candidate only in multiscale mode.
1890    // In single-scale mode it is FUSED into the primary value above (#1116), so
1891    // no separate candidate — the ψ-derivative penalty count must match the
1892    // build path, count desync = the gam#860 trap class.
1893    let ridge = spec.double_penalty && geom.per_level;
1894    let n_cands = n_active + usize::from(ridge);
1895    let zero_p = || Array2::<f64>::zeros((p, p));
1896    let mut penalties_first: Vec<Vec<Array2<f64>>> =
1897        (0..n_coords).map(|_| Vec::with_capacity(n_cands)).collect();
1898    let mut penalties_second_diag: Vec<Vec<Array2<f64>>> =
1899        (0..n_coords).map(|_| Vec::with_capacity(n_cands)).collect();
1900    // Cross matrices per pair per candidate, precomputed eagerly (the
1901    // candidate count is the band length, not the data size) and served
1902    // through the on-demand provider.
1903    let mut crosses: Vec<Vec<Array2<f64>>> = (0..pairs.len()).map(|_| Vec::new()).collect();
1904    for (s_raw, firsts, seconds, cross_raw) in &raw {
1905        // ONE Frobenius scale per candidate, fixed up front from `s_raw`
1906        // alone: c anchors the value and every derivative of this candidate.
1907        // `normalize_penaltywith_psi_derivatives` recomputes the identical c
1908        // per coordinate (same trace_of_product + sqrt on the same `s_raw`),
1909        // and its degenerate convention is mirrored here: ‖S‖_F ≤ 1e-12 (or
1910        // non-finite) reports scale 1.0 — the value passes through unscaled,
1911        // and the cross helper receives that same 1.0, never a collapsed
1912        // near-zero scale.
1913        let fro = trace_of_product(s_raw, s_raw).sqrt();
1914        let c = if fro.is_finite() && fro > 1e-12 {
1915            fro
1916        } else {
1917            1.0
1918        };
1919        for coord in 0..n_coords {
1920            let (_, s_first, s_second, _) =
1921                normalize_penaltywith_psi_derivatives(s_raw, &firsts[coord], &seconds[coord]);
1922            penalties_first[coord].push(s_first);
1923            penalties_second_diag[coord].push(s_second);
1924        }
1925        for (pair_idx, &(a, b)) in pairs.iter().enumerate() {
1926            let cross_raw_mat = normalize_penalty_cross_psi_derivative(
1927                s_raw,
1928                &firsts[a],
1929                &firsts[b],
1930                &cross_raw[pair_idx],
1931                c,
1932            );
1933            crosses[pair_idx].push(cross_raw_mat);
1934        }
1935    }
1936    if ridge {
1937        for coord in 0..n_coords {
1938            penalties_first[coord].push(zero_p());
1939            penalties_second_diag[coord].push(zero_p());
1940        }
1941        for pair_crosses in crosses.iter_mut() {
1942            pair_crosses.push(zero_p());
1943        }
1944    }
1945    // Prepend the design-moving ℓ coordinate at index 0 when enrolled. ℓ
1946    // carries the nonzero `design_first/second`, ZERO penalty derivatives (the
1947    // penalty is ℓ-independent), and no cross terms with the penalty dials
1948    // (those carry zero design drift). Every penalty-dial coordinate and cross
1949    // pair therefore shifts up by one.
1950    let coord_offset = usize::from(length_scale_design.is_some());
1951    if coord_offset == 1 {
1952        penalties_first.insert(0, (0..n_cands).map(|_| zero_p()).collect());
1953        penalties_second_diag.insert(0, (0..n_cands).map(|_| zero_p()).collect());
1954    }
1955    let n_coords_total = n_coords + coord_offset;
1956    // Penalty-dial cross pairs, shifted past the ℓ coordinate.
1957    let mut all_pairs: Vec<(usize, usize)> = pairs
1958        .iter()
1959        .map(|&(a, b)| (a + coord_offset, b + coord_offset))
1960        .collect();
1961    let mut all_crosses: Vec<Vec<Array2<f64>>> = crosses;
1962    // ℓ (coord 0) × every penalty coordinate: the penalty is ℓ-independent and
1963    // the penalty dials carry zero design drift, so each cross block is zero —
1964    // but the pair MUST be registered (the consumer queries every coordinate
1965    // pair `(a, b)`), else the provider errors on an unregistered ℓ pair.
1966    if coord_offset == 1 {
1967        for c in 1..n_coords_total {
1968            all_pairs.push((0, c));
1969            all_crosses.push((0..n_cands).map(|_| zero_p()).collect());
1970        }
1971    }
1972    let pair_index: Vec<((usize, usize), Vec<Array2<f64>>)> = all_pairs
1973        .iter()
1974        .copied()
1975        .zip(all_crosses.into_iter())
1976        .collect();
1977    let shifted_pairs = all_pairs;
1978    let provider = AnisoPenaltyCrossProvider::new(move |a, b| {
1979        pair_index
1980            .iter()
1981            .find(|((pa, pb), _)| (*pa, *pb) == (a, b) || (*pa, *pb) == (b, a))
1982            .map(|(_, mats)| mats.clone())
1983            .ok_or_else(|| {
1984                BasisError::InvalidInput(format!(
1985                    "measure-jet ψ cross derivative requested for unknown pair ({a}, {b})"
1986                ))
1987            })
1988    });
1989    let mut design_first: Vec<Array2<f64>> = (0..n_coords_total)
1990        .map(|_| Array2::<f64>::zeros((n, p)))
1991        .collect();
1992    let mut design_second_diag: Vec<Array2<f64>> = (0..n_coords_total)
1993        .map(|_| Array2::<f64>::zeros((n, p)))
1994        .collect();
1995    if let Some((dx_du, d2x_du2)) = length_scale_design {
1996        design_first[0] = dx_du;
1997        design_second_diag[0] = d2x_du2;
1998    }
1999    Ok(AnisoBasisPsiDerivatives {
2000        design_first,
2001        design_second_diag,
2002        design_second_cross: Vec::new(),
2003        design_second_cross_pairs: Vec::new(),
2004        penalties_first,
2005        penalties_second_diag,
2006        penalties_cross_pairs: shifted_pairs,
2007        penalties_cross_provider: Some(provider),
2008        implicit_operator: None,
2009    })
2010}
2011
2012#[cfg(test)]
2013mod tests {
2014    use super::*;
2015    pub(crate) fn two_cluster_centers() -> (ndarray::Array2<f64>, ndarray::Array1<f64>) {
2016        let centers = array![
2017            [0.00, 0.00],
2018            [0.31, 0.05],
2019            [0.58, -0.07],
2020            [0.93, 0.11],
2021            [1.22, 0.02],
2022            [1.49, -0.04],
2023            [3.10, 2.00],
2024            [3.42, 2.13],
2025            [3.71, 1.91],
2026            [4.05, 2.07],
2027            [4.33, 1.96],
2028            [4.61, 2.12],
2029        ];
2030        let m = centers.nrows();
2031        let masses = ndarray::Array1::<f64>::from_elem(m, 1.0 / m as f64);
2032        (centers, masses)
2033    }
2034    use ndarray::array;
2035
2036    pub(crate) fn band_for(centers: &Array2<f64>) -> MeasureJetBand {
2037        measure_jet_band(centers.view(), 0).expect("band")
2038    }
2039
2040    /// The no-mass contract: constants must be annihilated to machine
2041    /// precision at every scale (the constant is projected, never ridged).
2042    #[test]
2043    pub(crate) fn energy_form_annihilates_constants_exactly() {
2044        let (centers, masses) = two_cluster_centers();
2045        let band = band_for(&centers);
2046        let q = measure_jet_energy_form(centers.view(), masses.view(), &band, 1.5, 1.0, 1e-3)
2047            .expect("energy form");
2048        let m = q.nrows();
2049        let ones = Array1::<f64>::ones(m);
2050        let qv = q.dot(&ones);
2051        let scale = q.iter().fold(0.0_f64, |acc, v| acc.max(v.abs()));
2052        assert!(scale > 0.0, "energy form is identically zero");
2053        for (i, v) in qv.iter().enumerate() {
2054            assert!(
2055                v.abs() <= 1e-12 * scale,
2056                "Q·1 leak at row {i}: {v:.3e} vs scale {scale:.3e}"
2057            );
2058        }
2059        let vqv = ones.dot(&qv);
2060        assert!(
2061            vqv.abs() <= 1e-12 * scale,
2062            "constant carries energy: 1ᵀQ1 = {vqv:.3e}"
2063        );
2064    }
2065
2066    /// The default local projection annihilates ambient affine functions
2067    /// exactly; τ is retained for ψ layout but no longer adds an affine toll.
2068    #[test]
2069    pub(crate) fn energy_form_annihilates_affine_at_default_tau() {
2070        let (centers, masses) = two_cluster_centers();
2071        let band = band_for(&centers);
2072        let m = centers.nrows();
2073        // Affine values v = 0.7 + 1.3·x − 0.4·y, and a rough ±1 checkerboard.
2074        let mut affine = Array1::<f64>::zeros(m);
2075        let mut rough = Array1::<f64>::zeros(m);
2076        for i in 0..m {
2077            affine[i] = 0.7 + 1.3 * centers[(i, 0)] - 0.4 * centers[(i, 1)];
2078            rough[i] = if i % 2 == 0 { 1.0 } else { -1.0 };
2079        }
2080        let q = measure_jet_energy_form(centers.view(), masses.view(), &band, 1.5, 1.0, 1e-3)
2081            .expect("energy form");
2082        let e_affine = affine.dot(&q.dot(&affine));
2083        let e_rough = rough.dot(&q.dot(&rough));
2084        assert!(e_rough > 0.0, "rough vector must pay energy");
2085        assert!(
2086            e_affine.abs() <= 1e-12 * e_rough,
2087            "default affine energy {e_affine:.3e} vs rough {e_rough:.3e}"
2088        );
2089    }
2090
2091    /// PSD: the energy is a sum of weighted least-squares residuals.
2092    #[test]
2093    pub(crate) fn energy_form_is_psd() {
2094        let (centers, masses) = two_cluster_centers();
2095        let band = band_for(&centers);
2096        let q = measure_jet_energy_form(centers.view(), masses.view(), &band, 1.5, 1.0, 1e-3)
2097            .expect("energy form");
2098        let m = q.nrows();
2099        for trial in 0..5usize {
2100            let v = Array1::<f64>::from_shape_fn(m, |i| {
2101                ((i * 7 + trial * 13) % 11) as f64 / 11.0 - 0.5
2102            });
2103            let e = v.dot(&q.dot(&v));
2104            assert!(e >= -1e-10, "vᵀQv = {e:.3e} < 0 on trial {trial}");
2105        }
2106    }
2107
2108    /// A 1-D filament embedded in 2-D: high-frequency center values along the
2109    /// strand pay strictly more energy than a slow trend.
2110    #[test]
2111    pub(crate) fn rough_vector_pays_more_than_smooth() {
2112        let m = 24usize;
2113        let centers = Array2::<f64>::from_shape_fn((m, 2), |(i, k)| {
2114            let t = i as f64 / (m as f64 - 1.0);
2115            if k == 0 {
2116                t * 4.0
2117            } else {
2118                0.3 * (t * 4.0).sin()
2119            }
2120        });
2121        let masses = Array1::<f64>::from_elem(m, 1.0 / m as f64);
2122        let band = band_for(&centers);
2123        let q = measure_jet_energy_form(centers.view(), masses.view(), &band, 1.5, 1.0, 1e-3)
2124            .expect("energy form");
2125        let slow = Array1::<f64>::from_shape_fn(m, |i| (i as f64 / (m as f64 - 1.0)).powi(2));
2126        let fast = Array1::<f64>::from_shape_fn(m, |i| if i % 2 == 0 { 0.5 } else { -0.5 });
2127        let e_slow = slow.dot(&q.dot(&slow));
2128        let e_fast = fast.dot(&q.dot(&fast));
2129        assert!(
2130            e_fast > 10.0 * e_slow,
2131            "alternating values must pay >> a slow trend: fast {e_fast:.3e} vs slow {e_slow:.3e}"
2132        );
2133    }
2134
2135    /// The exact (s, α) jets and zero τ slots must match central finite
2136    /// differences of the energy — the FD gate the ψ-channel stage will
2137    /// inherit (the discipline whose absence is exactly the
2138    /// objective↔gradient desync bug class).
2139    #[test]
2140    pub(crate) fn energy_jets_match_finite_differences() {
2141        let (centers, masses) = two_cluster_centers();
2142        let band = band_for(&centers);
2143        let (s0, a0, tau) = (1.3, 0.8, 1e-3);
2144        let jets =
2145            measure_jet_energy_form_with_jets(centers.view(), masses.view(), &band, s0, a0, tau)
2146                .expect("jets");
2147        let q_at = |s: f64, a: f64| {
2148            measure_jet_energy_form(centers.view(), masses.view(), &band, s, a, tau)
2149                .expect("energy form")
2150        };
2151        // Base form must equal the plain assembly bit-for-bit (same walk).
2152        let q_plain = q_at(s0, a0);
2153        for (a, b) in jets.q.iter().zip(q_plain.iter()) {
2154            assert!(
2155                (a - b).abs() <= 1e-14 * (1.0 + b.abs()),
2156                "Q drift {a} vs {b}"
2157            );
2158        }
2159        let lt0 = tau.ln();
2160        let q_at_lt = |lt: f64| {
2161            measure_jet_energy_form(centers.view(), masses.view(), &band, s0, a0, lt.exp())
2162                .expect("energy form")
2163        };
2164        // FD step calibrated for the SECOND differences: their roundoff
2165        // floor is ~4·ε_f64·scale/h² (assembly noise amplified by 1/h²), so
2166        // h = 1e-4 ≈ ε^(1/4) balances it against the O(h²) truncation —
2167        // both land ≥3 orders below the unchanged 5e-5·scale gate. h = 1e-5
2168        // sits ON the roundoff floor and fails spuriously.
2169        let h = 1e-4;
2170        let checks: [(&str, &Array2<f64>, Array2<f64>); 9] = [
2171            ("dq_ds", &jets.dq_ds, {
2172                let (p, m_) = (q_at(s0 + h, a0), q_at(s0 - h, a0));
2173                (&p - &m_) / (2.0 * h)
2174            }),
2175            ("d2q_ds2", &jets.d2q_ds2, {
2176                let (p, c, m_) = (q_at(s0 + h, a0), q_at(s0, a0), q_at(s0 - h, a0));
2177                (&(&p + &m_) - &(&c * 2.0)) / (h * h)
2178            }),
2179            ("dq_dalpha", &jets.dq_dalpha, {
2180                let (p, m_) = (q_at(s0, a0 + h), q_at(s0, a0 - h));
2181                (&p - &m_) / (2.0 * h)
2182            }),
2183            ("d2q_dalpha2", &jets.d2q_dalpha2, {
2184                let (p, c, m_) = (q_at(s0, a0 + h), q_at(s0, a0), q_at(s0, a0 - h));
2185                (&(&p + &m_) - &(&c * 2.0)) / (h * h)
2186            }),
2187            ("d2q_ds_dalpha", &jets.d2q_ds_dalpha, {
2188                let pp = q_at(s0 + h, a0 + h);
2189                let pm = q_at(s0 + h, a0 - h);
2190                let mp = q_at(s0 - h, a0 + h);
2191                let mm = q_at(s0 - h, a0 - h);
2192                (&(&pp - &pm) - &(&mp - &mm)) / (4.0 * h * h)
2193            }),
2194            ("dq_dlogtau", &jets.dq_dlogtau, {
2195                let (p, m_) = (q_at_lt(lt0 + h), q_at_lt(lt0 - h));
2196                (&p - &m_) / (2.0 * h)
2197            }),
2198            ("d2q_dlogtau2", &jets.d2q_dlogtau2, {
2199                let (p, c, m_) = (q_at_lt(lt0 + h), q_at_lt(lt0), q_at_lt(lt0 - h));
2200                (&(&p + &m_) - &(&c * 2.0)) / (h * h)
2201            }),
2202            ("d2q_ds_dlogtau", &jets.d2q_ds_dlogtau, {
2203                let f = |s: f64, lt: f64| {
2204                    measure_jet_energy_form(centers.view(), masses.view(), &band, s, a0, lt.exp())
2205                        .expect("energy form")
2206                };
2207                let pp = f(s0 + h, lt0 + h);
2208                let pm = f(s0 + h, lt0 - h);
2209                let mp = f(s0 - h, lt0 + h);
2210                let mm = f(s0 - h, lt0 - h);
2211                (&(&pp - &pm) - &(&mp - &mm)) / (4.0 * h * h)
2212            }),
2213            ("d2q_dalpha_dlogtau", &jets.d2q_dalpha_dlogtau, {
2214                let f = |a: f64, lt: f64| {
2215                    measure_jet_energy_form(centers.view(), masses.view(), &band, s0, a, lt.exp())
2216                        .expect("energy form")
2217                };
2218                let pp = f(a0 + h, lt0 + h);
2219                let pm = f(a0 + h, lt0 - h);
2220                let mp = f(a0 - h, lt0 + h);
2221                let mm = f(a0 - h, lt0 - h);
2222                (&(&pp - &pm) - &(&mp - &mm)) / (4.0 * h * h)
2223            }),
2224        ];
2225        for (name, analytic, fd) in checks.iter() {
2226            let scale = fd.iter().fold(1e-30_f64, |acc, v| acc.max(v.abs()));
2227            for (a, b) in analytic.iter().zip(fd.iter()) {
2228                assert!(
2229                    (a - b).abs() <= 5e-5 * scale,
2230                    "{name} jet mismatch: analytic {a:.6e} vs FD {b:.6e} (scale {scale:.3e})"
2231                );
2232            }
2233        }
2234    }
2235
2236    /// The per-scale spectrum must sum exactly to the total energy (same
2237    /// blocks, one-hot weights) and concentrate rough content at fine
2238    /// scales.
2239    #[test]
2240    pub(crate) fn scale_spectrum_sums_to_total_and_localizes_roughness() {
2241        let m = 24usize;
2242        let centers = Array2::<f64>::from_shape_fn((m, 2), |(i, k)| {
2243            let t = i as f64 / (m as f64 - 1.0);
2244            if k == 0 { t * 4.0 } else { 0.0 }
2245        });
2246        let masses = Array1::<f64>::from_elem(m, 1.0 / m as f64);
2247        let band = band_for(&centers);
2248        let q = measure_jet_energy_form(centers.view(), masses.view(), &band, 1.5, 1.0, 1e-3)
2249            .expect("energy form");
2250        let fast = Array1::<f64>::from_shape_fn(m, |i| if i % 2 == 0 { 0.5 } else { -0.5 });
2251        let spec = measure_jet_scale_spectrum(
2252            centers.view(),
2253            masses.view(),
2254            &band,
2255            1.5,
2256            1.0,
2257            1e-3,
2258            fast.view(),
2259        )
2260        .expect("spectrum");
2261        assert_eq!(spec.len(), band.eps.len());
2262        let total = fast.dot(&q.dot(&fast));
2263        let sum: f64 = spec.iter().sum();
2264        assert!(
2265            (sum - total).abs() <= 1e-10 * total.abs().max(1e-30),
2266            "spectrum must sum to vᵀQv: {sum:.6e} vs {total:.6e}"
2267        );
2268        // Alternating-sign content lives at the finest scale of the band.
2269        let finest = spec[0];
2270        let coarsest = *spec.last().expect("nonempty spectrum");
2271        assert!(
2272            finest > coarsest,
2273            "alternating values must charge fine scales hardest: fine {finest:.3e} vs coarse {coarsest:.3e}"
2274        );
2275    }
2276
2277    /// The support curve separates on-web from off-web queries at fine
2278    /// scales and grows monotonically in ε for any query.
2279    #[test]
2280    pub(crate) fn support_curve_separates_on_web_from_off_web() {
2281        let m = 24usize;
2282        let centers = Array2::<f64>::from_shape_fn((m, 2), |(i, k)| {
2283            let t = i as f64 / (m as f64 - 1.0);
2284            if k == 0 { t * 4.0 } else { 0.0 }
2285        });
2286        let masses = Array1::<f64>::from_elem(m, 1.0 / m as f64);
2287        let band = band_for(&centers);
2288        let queries = array![[2.0, 0.0], [2.0, 1.5]];
2289        let curves =
2290            measure_jet_support_curve(queries.view(), centers.view(), masses.view(), &band.eps)
2291                .expect("support curve");
2292        // On-web sees strictly more mass than off-web at the finest scale.
2293        assert!(
2294            curves[(0, 0)] > 10.0 * curves[(1, 0)],
2295            "fine-scale support must separate web from void: on {:.3e} vs off {:.3e}",
2296            curves[(0, 0)],
2297            curves[(1, 0)]
2298        );
2299        // Kernel mass is monotone in ε for every query.
2300        for qi in 0..2 {
2301            for li in 1..band.eps.len() {
2302                assert!(
2303                    curves[(qi, li)] >= curves[(qi, li - 1)] - 1e-15,
2304                    "support curve must be monotone in scale (query {qi}, level {li})"
2305                );
2306            }
2307        }
2308    }
2309
2310    /// The default is single-scale mode at ANY center count: ONE fused penalty,
2311    /// the same one-λ outer footprint as Duchon/Matérn. Multiscale (the
2312    /// per-scale spectral split + ψ dials + the affine-preserving ridge) is an
2313    /// EXPLICIT opt-in (`spec.multiscale`, the DSL `mjs(…, multiscale=true)`) —
2314    /// there is no center-count auto-gate (#1116). `measure_jet_multiscale_mode`
2315    /// is the single source for this decision.
2316    #[test]
2317    pub(crate) fn default_stays_single_scale_until_multiscale_opt_in() {
2318        let n = 200usize;
2319        let data = Array2::<f64>::from_shape_fn((n, 2), |(i, k)| {
2320            let t = i as f64 / (n as f64 - 1.0);
2321            if k == 0 {
2322                t * 3.0
2323            } else {
2324                0.4 * (t * 3.0).sin()
2325            }
2326        });
2327        // Default (multiscale = false) stays single-scale even at a LARGE center
2328        // count that, under the deleted auto-gate, would have flipped to
2329        // multiscale: exactly ONE candidate — the fused jet-energy penalty with
2330        // the affine-preserving nullspace ridge folded in at a fixed fraction
2331        // (#1116), the one-λ Duchon/Matérn footprint.
2332        let single = MeasureJetBasisSpec {
2333            center_strategy: CenterStrategy::FarthestPoint { num_centers: 80 },
2334            ..MeasureJetBasisSpec::default()
2335        };
2336        assert!(
2337            !measure_jet_multiscale_mode(&single),
2338            "default must resolve to single-scale at any center count"
2339        );
2340        let built_single =
2341            build_measure_jet_basis(data.view(), &single).expect("single-scale build");
2342        assert_eq!(
2343            built_single.penalties.len(),
2344            1,
2345            "single-scale mode emits one fused penalty (ridge folded in, not a 2nd λ)"
2346        );
2347        // The explicit opt-in flips to multiscale at the SAME center count: the
2348        // per-scale spectral split (several candidates) + the ridge, strictly
2349        // more candidates than the single fused penalty.
2350        let multi = MeasureJetBasisSpec {
2351            center_strategy: CenterStrategy::FarthestPoint { num_centers: 80 },
2352            multiscale: true,
2353            ..MeasureJetBasisSpec::default()
2354        };
2355        assert!(
2356            measure_jet_multiscale_mode(&multi),
2357            "multiscale=true must resolve to multiscale mode"
2358        );
2359        let built_multi = build_measure_jet_basis(data.view(), &multi).expect("multiscale build");
2360        assert!(
2361            built_multi.penalties.len() > built_single.penalties.len(),
2362            "multiscale mode emits the per-scale spectral split plus the ridge, got {} (vs single-scale {})",
2363            built_multi.penalties.len(),
2364            built_single.penalties.len()
2365        );
2366    }
2367
2368    /// An explicit order pins the Mellin weights and fuses the band into a
2369    /// single Primary candidate. With multiscale off (the default) the mode is
2370    /// single-scale: one fused Primary penalty with the affine-preserving
2371    /// nullspace ridge folded in (#1116) = exactly one candidate.
2372    #[test]
2373    pub(crate) fn fused_mode_emits_single_primary_candidate() {
2374        let n = 40usize;
2375        let data = Array2::<f64>::from_shape_fn((n, 2), |(i, k)| {
2376            let t = i as f64 / (n as f64 - 1.0);
2377            if k == 0 {
2378                t * 3.0
2379            } else {
2380                0.4 * (t * 3.0).sin()
2381            }
2382        });
2383        let spec = MeasureJetBasisSpec {
2384            center_strategy: CenterStrategy::FarthestPoint { num_centers: 14 },
2385            order_s: 1.3,
2386            ..MeasureJetBasisSpec::default()
2387        };
2388        let built = build_measure_jet_basis(data.view(), &spec).expect("fused build");
2389        assert_eq!(
2390            built.penalties.len(),
2391            1,
2392            "fused single-scale mode emits exactly one Primary candidate (ridge folded in)"
2393        );
2394        let BasisMetadata::MeasureJet { order_s, .. } = &built.metadata else {
2395            panic!("measure-jet build must return MeasureJet metadata");
2396        };
2397        assert_eq!(*order_s, 1.3, "explicit order must persist verbatim");
2398    }
2399
2400    /// The Householder basis must be orthonormal with sum-to-zero columns.
2401    #[test]
2402    pub(crate) fn householder_sum_to_zero_basis_is_orthonormal() {
2403        let m = 9usize;
2404        let u = householder_sum_to_zero_u(m);
2405        let z = householder_sum_to_zero_z(&u);
2406        for j in 0..(m - 1) {
2407            let col_j = z.column(j);
2408            assert!(col_j.sum().abs() <= 1e-12, "column {j} must sum to zero");
2409            for j2 in j..(m - 1) {
2410                let dot = col_j.dot(&z.column(j2));
2411                let want = if j == j2 { 1.0 } else { 0.0 };
2412                assert!(
2413                    (dot - want).abs() <= 1e-12,
2414                    "orthonormality failure at ({j}, {j2}): {dot}"
2415                );
2416            }
2417        }
2418    }
2419
2420    /// Frozen-geometry fixture shared by the ψ-producer FD gates: build
2421    /// once, pin everything (nodes, masses, band, transform, realized ℓ),
2422    /// and return the pinned spec so dial-perturbed rebuilds move ONLY the
2423    /// dials — the per-trial contract the optimizer relies on.
2424    pub(crate) fn frozen_spec_fixture(
2425        order_s: f64,
2426        multiscale: bool,
2427    ) -> (Array2<f64>, MeasureJetBasisSpec) {
2428        // Multiscale (per-scale + ψ) mode is the explicit opt-in (#1116); the
2429        // per-level fixture passes `multiscale = true`, the fused fixture
2430        // `false`. A large center count is kept so the multiscale spectrum is
2431        // identifiable when opted in.
2432        let n = 140usize;
2433        let data = Array2::<f64>::from_shape_fn((n, 2), |(i, k)| {
2434            let t = i as f64 / (n as f64 - 1.0);
2435            if k == 0 {
2436                t * 3.0
2437            } else {
2438                0.5 * (t * 3.0).cos() + if i % 9 == 0 { 0.8 } else { 0.0 }
2439            }
2440        });
2441        let spec = MeasureJetBasisSpec {
2442            center_strategy: CenterStrategy::FarthestPoint { num_centers: 70 },
2443            order_s,
2444            multiscale,
2445            // These fixtures gate the PENALTY-dial derivatives; freeze ℓ so the
2446            // coordinate layout is exactly the penalty dials (the design-moving
2447            // ℓ dial has its own FD gate, `psi_producer_matches_fd_length_scale`).
2448            learn_length_scale: false,
2449            ..MeasureJetBasisSpec::default()
2450        };
2451        let first = build_measure_jet_basis(data.view(), &spec).expect("fixture build");
2452        let BasisMetadata::MeasureJet {
2453            centers,
2454            length_scale,
2455            eps_band,
2456            masses,
2457            support_means,
2458            penalty_normalization_scales,
2459            raw_penalty_normalization_scales,
2460            fused_penalty_normalization_scale,
2461            constraint_transform,
2462            ..
2463        } = &first.metadata
2464        else {
2465            panic!("measure-jet build must return MeasureJet metadata");
2466        };
2467        let frozen = MeasureJetBasisSpec {
2468            center_strategy: CenterStrategy::UserProvided(centers.clone()),
2469            order_s,
2470            alpha: spec.alpha,
2471            tau0: spec.tau0,
2472            num_scales: eps_band.len(),
2473            length_scale: *length_scale,
2474            double_penalty: spec.double_penalty,
2475            learn_length_scale: false,
2476            multiscale,
2477            identifiability: MeasureJetIdentifiability::FrozenTransform {
2478                transform: constraint_transform.clone().expect("fit-time z"),
2479            },
2480            frozen_quadrature: Some(MeasureJetFrozenQuadrature {
2481                masses: masses.clone(),
2482                eps_band: eps_band.clone(),
2483                support_means: support_means.clone(),
2484                penalty_normalization_scales: penalty_normalization_scales.clone(),
2485                raw_penalty_normalization_scales: raw_penalty_normalization_scales.clone(),
2486                fused_penalty_normalization_scale: *fused_penalty_normalization_scale,
2487            }),
2488        };
2489        (data, frozen)
2490    }
2491
2492    /// ψ-producer vs central finite differences of the NORMALIZED fit-time
2493    /// candidates under frozen geometry — per-level mode (coords α, lnτ).
2494    /// This is the end-to-end gate #901 never had: the derivative is checked
2495    /// against the exact object the optimizer consumes.
2496    #[test]
2497    pub(crate) fn psi_producer_matches_fd_per_level_mode() {
2498        let (data, frozen) = frozen_spec_fixture(0.0, true);
2499        let derivs =
2500            build_measure_jet_basis_psi_derivatives(data.view(), &frozen).expect("psi derivatives");
2501        let l_count = frozen
2502            .frozen_quadrature
2503            .as_ref()
2504            .expect("frozen quadrature")
2505            .eps_band
2506            .len();
2507        assert_eq!(
2508            derivs.penalties_first.len(),
2509            2,
2510            "per-level coords are (α, lnτ)"
2511        );
2512        assert_eq!(derivs.penalties_first[0].len(), l_count + 1);
2513        assert_eq!(derivs.penalties_cross_pairs, vec![(0, 1)]);
2514        let pen_at = |alpha: f64, tau0: f64| {
2515            let trial = MeasureJetBasisSpec {
2516                alpha,
2517                tau0,
2518                ..frozen.clone()
2519            };
2520            build_measure_jet_basis(data.view(), &trial)
2521                .expect("trial build")
2522                .penalties
2523        };
2524        // Second-difference-optimal step (see the jets FD test): the 4-point
2525        // cross stencil shares the ~ε·scale/h² roundoff floor.
2526        let h = 1e-4;
2527        let (a0, t0) = (frozen.alpha, frozen.tau0);
2528        let ap = pen_at(a0 + h, t0);
2529        let am = pen_at(a0 - h, t0);
2530        let tp = pen_at(a0, t0 * h.exp());
2531        let tm = pen_at(a0, t0 * (-h).exp());
2532        assert_eq!(
2533            ap.len(),
2534            l_count + 1,
2535            "fixture must keep every scale active"
2536        );
2537        for level in 0..l_count {
2538            let fd_alpha = (&ap[level] - &am[level]) / (2.0 * h);
2539            let fd_tau = (&tp[level] - &tm[level]) / (2.0 * h);
2540            for (name, analytic, fd) in [
2541                ("alpha", &derivs.penalties_first[0][level], fd_alpha),
2542                ("ln_tau", &derivs.penalties_first[1][level], fd_tau),
2543            ] {
2544                let scale = fd.iter().fold(1e-30_f64, |acc, v| acc.max(v.abs()));
2545                for (x, y) in analytic.iter().zip(fd.iter()) {
2546                    assert!(
2547                        (x - y).abs() <= 5e-5 * scale,
2548                        "{name} jet of scale-candidate {level}: analytic {x:.6e} vs FD {y:.6e}"
2549                    );
2550                }
2551            }
2552        }
2553        // The ridge candidate carries identically-zero derivatives.
2554        for coord in 0..2 {
2555            assert!(
2556                derivs.penalties_first[coord][l_count]
2557                    .iter()
2558                    .all(|v| *v == 0.0),
2559                "ridge candidate must have zero ψ drift"
2560            );
2561        }
2562        // Cross derivative through the provider, against a 4-point FD.
2563        let provider = derivs
2564            .penalties_cross_provider
2565            .as_ref()
2566            .expect("cross provider");
2567        let cross = provider.evaluate(0, 1).expect("cross pair (α, lnτ)");
2568        let pp = pen_at(a0 + h, t0 * h.exp());
2569        let pm = pen_at(a0 + h, t0 * (-h).exp());
2570        let mp = pen_at(a0 - h, t0 * h.exp());
2571        let mm = pen_at(a0 - h, t0 * (-h).exp());
2572        for level in 0..l_count {
2573            let fd = (&(&pp[level] - &pm[level]) - &(&mp[level] - &mm[level])) / (4.0 * h * h);
2574            let scale = fd.iter().fold(1e-30_f64, |acc, v| acc.max(v.abs()));
2575            for (x, y) in cross[level].iter().zip(fd.iter()) {
2576                assert!(
2577                    (x - y).abs() <= 5e-4 * scale,
2578                    "cross (α, lnτ) jet of scale-candidate {level}: analytic {x:.6e} vs FD {y:.6e}"
2579                );
2580            }
2581        }
2582    }
2583
2584    /// Design-moving ℓ dial (#1116): the producer's `design_first[0]` /
2585    /// `design_second_diag[0]` must match central finite differences of the
2586    /// REBUILT design `X(ℓ)` under frozen geometry, and the ℓ coordinate must
2587    /// carry IDENTICALLY ZERO penalty derivatives (the jet-energy penalty is
2588    /// ℓ-independent). This is the FD gate matérn's log_kappa has and mjs lacked.
2589    #[test]
2590    pub(crate) fn psi_producer_matches_fd_length_scale() {
2591        // Single-scale with opt-in ℓ learning; frozen geometry so only ℓ moves
2592        // across the FD trials.
2593        let (data, mut frozen) = frozen_spec_fixture(0.0, false);
2594        frozen.learn_length_scale = true;
2595        let derivs =
2596            build_measure_jet_basis_psi_derivatives(data.view(), &frozen).expect("psi derivatives");
2597        // ℓ is the only coordinate in single-scale + learn_length_scale.
2598        assert_eq!(
2599            derivs.design_first.len(),
2600            1,
2601            "single-scale + learn_length_scale enrolls exactly the ℓ coordinate"
2602        );
2603        // The ℓ coordinate carries one (all-zero) penalty candidate: the penalty
2604        // is ℓ-independent, so its ψ-derivative is identically zero.
2605        assert_eq!(
2606            derivs.penalties_first[0].len(),
2607            1,
2608            "one fitted penalty candidate"
2609        );
2610        assert!(
2611            derivs.penalties_first[0][0].iter().all(|v| *v == 0.0)
2612                && derivs.penalties_second_diag[0][0].iter().all(|v| *v == 0.0),
2613            "the jet-energy penalty must not move with ℓ"
2614        );
2615        // Rebuild the DESIGN at ℓ·e^{±h} (perturb via spec.length_scale; the
2616        // realized resolver uses an explicit positive length_scale verbatim).
2617        let ell0 = frozen.length_scale;
2618        let design_at = |ell: f64| {
2619            let trial = MeasureJetBasisSpec {
2620                length_scale: ell,
2621                ..frozen.clone()
2622            };
2623            build_measure_jet_basis(data.view(), &trial)
2624                .expect("trial build")
2625                .design
2626                .to_dense()
2627        };
2628        let h: f64 = 1e-4;
2629        let x_plus = design_at(ell0 * h.exp());
2630        let x_minus = design_at(ell0 * (-h).exp());
2631        let x_0 = design_at(ell0);
2632        let fd_first = (&x_plus - &x_minus) / (2.0 * h);
2633        let fd_second = (&x_plus - &(&x_0 * 2.0) + &x_minus) / (h * h);
2634        let scale1 = fd_first.iter().fold(1e-30_f64, |acc, v| acc.max(v.abs()));
2635        for (x, y) in derivs.design_first[0].iter().zip(fd_first.iter()) {
2636            assert!(
2637                (x - y).abs() <= 5e-5 * scale1,
2638                "∂X/∂lnℓ: analytic {x:.6e} vs FD {y:.6e}"
2639            );
2640        }
2641        let scale2 = fd_second.iter().fold(1e-30_f64, |acc, v| acc.max(v.abs()));
2642        for (x, y) in derivs.design_second_diag[0].iter().zip(fd_second.iter()) {
2643            assert!(
2644                (x - y).abs() <= 1e-3 * scale2,
2645                "∂²X/∂lnℓ²: analytic {x:.6e} vs FD {y:.6e}"
2646            );
2647        }
2648    }
2649
2650    /// Quadrature nodes must be the mass-weighted cell barycenters
2651    /// (first-moment-exact lumping), with empty cells keeping their seed
2652    /// coordinates at zero mass.
2653    #[test]
2654    pub(crate) fn quadrature_nodes_are_cell_barycenters() {
2655        // Two tight groups around (0,0) and (10,10); a third seed far away
2656        // captures nothing.
2657        let data = array![
2658            [0.0, 0.2],
2659            [0.4, -0.2],
2660            [0.2, 0.0],
2661            [9.8, 10.1],
2662            [10.2, 9.9],
2663        ];
2664        let seeds = array![[0.1, 0.1], [10.0, 10.0], [-50.0, -50.0]];
2665        let (nodes, masses) =
2666            measure_jet_quadrature_nodes(data.view(), seeds.view()).expect("quadrature nodes");
2667        assert!((masses.sum() - 1.0).abs() <= 1e-15, "masses must sum to 1");
2668        assert!((masses[0] - 0.6).abs() <= 1e-15);
2669        assert!((masses[1] - 0.4).abs() <= 1e-15);
2670        assert_eq!(masses[2], 0.0);
2671        // Cell 0 barycenter = (0.2, 0.0).
2672        assert_eq!(nodes[(0, 0)], 0.2);
2673        assert_eq!(nodes[(0, 1)], 0.0);
2674        // Cell 1 barycenter = (10.0, 10.0), which is not a sampled row.
2675        assert_eq!(nodes[(1, 0)], 10.0);
2676        assert_eq!(nodes[(1, 1)], 10.0);
2677        // Empty cell keeps its seed coordinates.
2678        assert_eq!(nodes[(2, 0)], -50.0);
2679        assert_eq!(nodes[(2, 1)], -50.0);
2680    }
2681
2682    /// Freeze→replay: rebuilding from the first build's frozen transform and
2683    /// frozen quadrature must reproduce design and penalty bit-for-bit (the
2684    /// predict-path contract).
2685    #[test]
2686    pub(crate) fn build_replay_roundtrip_reproduces_design_and_penalty() {
2687        // A bent filament with a side cluster; multiscale opt-in so this
2688        // exercises the per-scale (spectral) replay path (#1116).
2689        let n = 140usize;
2690        let data = Array2::<f64>::from_shape_fn((n, 2), |(i, k)| {
2691            let t = i as f64 / (n as f64 - 1.0);
2692            if k == 0 {
2693                t * 3.0
2694            } else {
2695                0.5 * (t * 3.0).cos() + if i % 9 == 0 { 0.8 } else { 0.0 }
2696            }
2697        });
2698        let spec = MeasureJetBasisSpec {
2699            center_strategy: CenterStrategy::FarthestPoint { num_centers: 70 },
2700            multiscale: true,
2701            ..MeasureJetBasisSpec::default()
2702        };
2703        let first = build_measure_jet_basis(data.view(), &spec).expect("first build");
2704        let BasisMetadata::MeasureJet {
2705            centers,
2706            length_scale,
2707            eps_band,
2708            order_s,
2709            alpha,
2710            tau0,
2711            masses,
2712            support_means,
2713            penalty_normalization_scales,
2714            raw_penalty_normalization_scales,
2715            fused_penalty_normalization_scale,
2716            constraint_transform,
2717            ..
2718        } = &first.metadata
2719        else {
2720            panic!("measure-jet build must return MeasureJet metadata");
2721        };
2722        let replay_spec = MeasureJetBasisSpec {
2723            center_strategy: CenterStrategy::UserProvided(centers.clone()),
2724            order_s: *order_s,
2725            alpha: *alpha,
2726            tau0: *tau0,
2727            num_scales: eps_band.len(),
2728            length_scale: *length_scale,
2729            double_penalty: spec.double_penalty,
2730            learn_length_scale: spec.learn_length_scale,
2731            multiscale: spec.multiscale,
2732            identifiability: MeasureJetIdentifiability::FrozenTransform {
2733                transform: constraint_transform.clone().expect("fit-time z"),
2734            },
2735            frozen_quadrature: Some(MeasureJetFrozenQuadrature {
2736                masses: masses.clone(),
2737                eps_band: eps_band.clone(),
2738                support_means: support_means.clone(),
2739                penalty_normalization_scales: penalty_normalization_scales.clone(),
2740                raw_penalty_normalization_scales: raw_penalty_normalization_scales.clone(),
2741                fused_penalty_normalization_scale: *fused_penalty_normalization_scale,
2742            }),
2743        };
2744        // Per-level (auto-sentinel) mode: one candidate per band scale plus
2745        // the ridge, and the count must survive freeze→replay bit-for-bit.
2746        assert_eq!(
2747            first.penalties.len(),
2748            eps_band.len() + 1,
2749            "per-level mode must emit one candidate per scale + ridge"
2750        );
2751        let second = build_measure_jet_basis(data.view(), &replay_spec).expect("replay build");
2752        let x1 = first.design.to_dense();
2753        let x2 = second.design.to_dense();
2754        assert_eq!(x1.shape(), x2.shape());
2755        for (a, b) in x1.iter().zip(x2.iter()) {
2756            assert!((a - b).abs() <= 1e-12, "design replay drift: {a} vs {b}");
2757        }
2758        assert_eq!(first.penalties.len(), second.penalties.len());
2759        for (p1, p2) in first.penalties.iter().zip(second.penalties.iter()) {
2760            for (a, b) in p1.iter().zip(p2.iter()) {
2761                assert!((a - b).abs() <= 1e-12, "penalty replay drift: {a} vs {b}");
2762            }
2763        }
2764    }
2765}