Skip to main content

gam_terms/smooth/
term_specs.rs

1use coefficient_transforms::{
2    convex_derivative_control_transform_matrix, cumulative_exp, cumulative_sum_transform_matrix,
3    second_cumulative_exp,
4};
5
6pub use error::SmoothError;
7
8use input_standardization::estimate_isotropic_scale;
9
10use shape_constraints::{
11    bspline_first_derivative_control_spans, shape_lower_bounds_local, shape_order_and_sign,
12    shape_supports_basis, shape_uses_box_reparameterization,
13};
14
15pub fn describe_thin_plate_center_request(strategy: &CenterStrategy) -> String {
16    match strategy {
17        CenterStrategy::Auto(inner) => describe_thin_plate_center_request(inner),
18        CenterStrategy::DuchonSpectral { knots, basis } => format!(
19            "{} with Duchon spectral rank {}",
20            describe_thin_plate_center_request(knots),
21            basis.rank()
22        ),
23        CenterStrategy::UserProvided(centers) => format!("{} centers", centers.nrows()),
24        CenterStrategy::EqualMass { num_centers }
25        | CenterStrategy::EqualMassCovarRepresentative { num_centers }
26        | CenterStrategy::FarthestPoint { num_centers }
27        | CenterStrategy::KMeans { num_centers, .. } => format!("{num_centers} centers"),
28        CenterStrategy::UniformGrid { points_per_dim } => {
29            format!("uniform grid with {points_per_dim} points per dimension")
30        }
31    }
32}
33
34pub fn rewrite_thin_plate_knots_error(
35    err: BasisError,
36    termname: &str,
37    feature_count: usize,
38    spec: &ThinPlateBasisSpec,
39) -> BasisError {
40    match err {
41        // Polynomial-nullspace shortfall reported directly by the kernel
42        // builder ("thin-plate spline requires at least N centers to span ...").
43        BasisError::InvalidInput(msg)
44            if msg.contains("thin-plate spline requires at least")
45                && (msg.contains("centers to span") || msg.contains("knots to span")) =>
46        {
47            let min_centers = crate::basis::thin_plate_polynomial_basis_dimension(feature_count);
48            let requested = describe_thin_plate_center_request(&spec.center_strategy);
49            BasisError::InvalidInput(format!(
50                "joint TPS term '{termname}' over {feature_count} covariates with {requested} is invalid; minimum centers is {min_centers}"
51            ))
52        }
53        // Insufficient-rows shortfall raised by `select_thin_plate_knots` when
54        // the requested center count exceeds the available row count. Rewrite
55        // it in term language so the diagnostic points at the smooth term and
56        // the polynomial-nullspace minimum the user needs to satisfy.
57        BasisError::InvalidInput(msg)
58            if msg.starts_with("requested ") && msg.contains(" knots but only ") =>
59        {
60            let min_centers = crate::basis::thin_plate_polynomial_basis_dimension(feature_count);
61            let requested = describe_thin_plate_center_request(&spec.center_strategy);
62            BasisError::InvalidInput(format!(
63                "joint TPS term '{termname}' over {feature_count} covariates with {requested} is invalid; minimum centers is {min_centers}"
64            ))
65        }
66        other => other,
67    }
68}
69
70#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
71pub enum ShapeConstraint {
72    None,
73    MonotoneIncreasing,
74    MonotoneDecreasing,
75    Convex,
76    Concave,
77}
78
79/// Parse a shape-constraint string into a [`ShapeConstraint`].
80///
81/// This is the single source of truth shared by the formula DSL
82/// (`s(x, shape=...)`) and the `smooths={...}` override path
83/// (`Smooth.shape_constraint`). The accepted spellings cover the canonical
84/// Python `ShapeConstraintLiteral` strings exactly
85/// (`"none"` / `"monotone_increasing"` / `"monotone_decreasing"` /
86/// `"convex"` / `"concave"`) plus a few common aliases. Hyphens and case are
87/// normalized, so `"Monotone-Increasing"` and `"mono_inc"` both resolve to
88/// [`ShapeConstraint::MonotoneIncreasing`].
89pub fn parse_shape_constraint(raw: &str) -> Result<ShapeConstraint, String> {
90    let normalized = raw.trim().to_ascii_lowercase().replace('-', "_");
91    match normalized.as_str() {
92        "" | "none" => Ok(ShapeConstraint::None),
93        "monotone_increasing" | "monotonic_increasing" | "increasing" | "mono_inc" | "mpi" => {
94            Ok(ShapeConstraint::MonotoneIncreasing)
95        }
96        "monotone_decreasing" | "monotonic_decreasing" | "decreasing" | "mono_dec" | "mpd" => {
97            Ok(ShapeConstraint::MonotoneDecreasing)
98        }
99        "convex" | "cvx" => Ok(ShapeConstraint::Convex),
100        "concave" | "ccv" => Ok(ShapeConstraint::Concave),
101        other => Err(format!(
102            "unknown shape constraint {other:?}; expected one of \
103             \"none\", \"monotone_increasing\", \"monotone_decreasing\", \
104             \"convex\", \"concave\""
105        )),
106    }
107}
108
109impl ShapeConstraint {
110    /// Canonical formula-DSL spelling, i.e. the text emitted into
111    /// `s(x, shape=...)`. Round-trips through [`parse_shape_constraint`].
112    pub fn dsl_str(&self) -> &'static str {
113        match self {
114            ShapeConstraint::None => "none",
115            ShapeConstraint::MonotoneIncreasing => "monotone_increasing",
116            ShapeConstraint::MonotoneDecreasing => "monotone_decreasing",
117            ShapeConstraint::Convex => "convex",
118            ShapeConstraint::Concave => "concave",
119        }
120    }
121}
122
123/// Smooth-term head keywords recognised by the formula DSL. A `shape=` option
124/// may be attached to any term whose head is one of these.
125pub const SMOOTH_HEAD_KEYWORDS: [&str; 11] = [
126    "s",
127    "smooth",
128    "te",
129    "tensor",
130    "thinplate",
131    "tps",
132    "duchon",
133    "matern",
134    "sphere",
135    "bs",
136    "bspline",
137];
138
139/// Rewrite smooth-term calls in `formula` so each named smooth carries a
140/// `shape=<kind>` option understood by the formula DSL.
141///
142/// `constraints` pairs the smooth-term text as it appears in the formula
143/// (e.g. `"s(x)"` or `"s(x, type=duchon, centers=8)"`) with a shape-constraint
144/// spelling accepted by [`parse_shape_constraint`]; comparison is exact after
145/// whitespace removal. A `"none"` constraint is a no-op. Referencing a term not
146/// present in the formula is an error.
147///
148/// This is the single source of truth for the `gamfit.fit(..., constraints=…)`
149/// rewrite — the Python wrapper only marshals the mapping across the FFI and
150/// holds no formula-parsing or alias-normalization logic of its own.
151pub fn apply_shape_constraints_to_formula(
152    formula: &str,
153    constraints: &[(String, String)],
154) -> Result<String, String> {
155    use std::collections::{BTreeMap, BTreeSet};
156
157    if constraints.is_empty() {
158        return Ok(formula.to_string());
159    }
160    let strip_ws = |s: &str| -> String { s.chars().filter(|c| !c.is_whitespace()).collect() };
161
162    // Whitespace-stripped term text -> canonical shape spelling.
163    let mut wanted: BTreeMap<String, &'static str> = BTreeMap::new();
164    // Whitespace-stripped term text -> original key (for error labels).
165    let mut originals: BTreeMap<String, String> = BTreeMap::new();
166    for (key, kind_raw) in constraints {
167        let kind = parse_shape_constraint(kind_raw)?;
168        let nk = strip_ws(key);
169        originals.entry(nk.clone()).or_insert_with(|| key.clone());
170        if kind != ShapeConstraint::None {
171            wanted.insert(nk, kind.dsl_str());
172        }
173    }
174    if wanted.is_empty() {
175        return Ok(formula.to_string());
176    }
177
178    let chars: Vec<char> = formula.chars().collect();
179    let n = chars.len();
180    let is_ident = |c: char| c.is_ascii_alphanumeric() || c == '_';
181
182    let mut out = String::with_capacity(formula.len() + 32);
183    let mut matched: BTreeSet<String> = BTreeSet::new();
184    let mut i = 0usize;
185    while i < n {
186        // Locate the next smooth-term head (`<keyword> \s* (`) at or after `i`,
187        // respecting word boundaries so `abs(` never matches the `s(` head.
188        let mut head: Option<(usize, usize)> = None; // (head_start, paren_index)
189        let mut p = i;
190        while p < n {
191            let boundary = p == 0 || !is_ident(chars[p - 1]);
192            if boundary {
193                for kw in SMOOTH_HEAD_KEYWORDS.iter() {
194                    let klen = kw.chars().count();
195                    if p + klen > n || chars[p..p + klen].iter().collect::<String>() != **kw {
196                        continue;
197                    }
198                    let mut q = p + klen;
199                    while q < n && chars[q].is_whitespace() {
200                        q += 1;
201                    }
202                    if q < n && chars[q] == '(' {
203                        head = Some((p, q));
204                        break;
205                    }
206                }
207            }
208            if head.is_some() {
209                break;
210            }
211            p += 1;
212        }
213        let (head_start, paren_open) = match head {
214            Some(h) => h,
215            None => {
216                out.extend(chars[i..].iter());
217                break;
218            }
219        };
220        out.extend(chars[i..head_start].iter());
221
222        // Find the matching close paren, honoring nesting and string literals.
223        let body_start = paren_open + 1;
224        let mut depth = 1i32;
225        let mut j = body_start;
226        let mut in_str: Option<char> = None;
227        let mut closed = false;
228        while j < n {
229            let ch = chars[j];
230            if let Some(quote) = in_str {
231                if ch == quote {
232                    in_str = None;
233                }
234            } else if ch == '\'' || ch == '"' {
235                in_str = Some(ch);
236            } else if ch == '(' {
237                depth += 1;
238            } else if ch == ')' {
239                depth -= 1;
240                if depth == 0 {
241                    closed = true;
242                    break;
243                }
244            }
245            j += 1;
246        }
247
248        if !closed {
249            // Unbalanced — emit the remainder verbatim; the DSL parser will
250            // produce the canonical error.
251            out.extend(chars[head_start..].iter());
252            break;
253        }
254
255        let term_text: String = chars[head_start..=j].iter().collect();
256
257        let key_norm = strip_ws(&term_text);
258
259        match wanted.get(&key_norm) {
260            None => out.extend(chars[head_start..=j].iter()),
261            Some(kind) => {
262                let head_paren: String = chars[head_start..body_start].iter().collect();
263                let inside: String = chars[body_start..j].iter().collect();
264                let inside = inside.trim();
265                if inside.is_empty() {
266                    out.push_str(&format!("{head_paren}shape={kind})"));
267                } else {
268                    out.push_str(&format!("{head_paren}{inside}, shape={kind})"));
269                }
270                matched.insert(key_norm);
271            }
272        }
273
274        i = j + 1;
275    }
276
277    let mut missing: Vec<String> = wanted
278        .keys()
279        .filter(|k| !matched.contains(*k))
280        .map(|k| originals.get(k).cloned().unwrap_or_else(|| k.clone()))
281        .collect();
282
283    if !missing.is_empty() {
284        missing.sort();
285        return Err(format!(
286            "shape constraints referenced smooth term(s) not found in formula: {}",
287            missing.join(", ")
288        ));
289    }
290
291    Ok(out)
292}
293
294#[derive(Debug, Clone, Serialize, Deserialize)]
295pub enum BySmoothKind {
296    Numeric,
297    Level { level_bits: u64 },
298}
299
300#[derive(Debug, Clone, Serialize, Deserialize)]
301#[serde(deny_unknown_fields)]
302pub enum SmoothBasisSpec {
303    /// Row-gated wrapper used for mgcv-style ``by=`` smooths.
304    ///
305    /// ``ByNumeric`` multiplies the inner smooth by a numeric column.
306    /// ``ByLevel`` keeps the inner smooth active only for rows whose encoded
307    /// categorical value has the stored bit pattern.  Unordered factor-by
308    /// smooths are represented as one independent ``ByLevel`` term per level.
309    ///
310    /// `kind` preserves the compact structural discriminator, while `by`
311    /// carries the full row-gating spec used to build the local design.
312    ByVariable {
313        inner: Box<SmoothBasisSpec>,
314        by_col: usize,
315        kind: BySmoothKind,
316        by: ByVariableSpec,
317    },
318    /// Sum-to-zero factor smooth (`bs="sz"`): with L levels, estimate L-1
319    /// deviation coefficient blocks and use the final level as the negative
320    /// sum of the others, enforcing coefficient-wise zero sums across levels.
321    FactorSumToZero {
322        inner: Box<SmoothBasisSpec>,
323        by_col: usize,
324        levels: Vec<u64>,
325        /// Global-orthogonality column map `Z` captured at fit time when this
326        /// term overlapped an owner smooth (`s(x) + s(g, x, bs=sz)`, #978):
327        /// the hierarchical-ownership pass residualized this term's realized
328        /// design as `X ← X·Z`, shrinking its coefficient block. `Z` depends
329        /// on the *training-row* owner designs, so prediction cannot rederive
330        /// it — it must be persisted and replayed
331        /// (`apply_global_smooth_identifiability` consumes it verbatim).
332        /// Chart convention: `Z` lives in the post-restack, post-joint-null-Q
333        /// coordinates — the raw `sz` rebuild reapplies `Q` deterministically
334        /// (#700), then `Z` applies on top. `None` for non-overlapping terms.
335        #[serde(default)]
336        frozen_global_orthogonality: Option<Array2<f64>>,
337    },
338    BSpline1D {
339        feature_col: usize,
340        spec: BSplineBasisSpec,
341    },
342    /// A smooth modulated by a `by=` variable. Numeric `by` scales one inner
343    /// smooth; factor `by` replicates the inner smooth by level.
344    BySmooth {
345        smooth: Box<SmoothBasisSpec>,
346        by_kind: ByVarKind,
347    },
348    /// Factor-smooth interaction families (`bs="fs"`, `bs="sz"`) and
349    /// random slopes (`bs="re"`).
350    FactorSmooth { spec: FactorSmoothSpec },
351    ThinPlate {
352        feature_cols: Vec<usize>,
353        spec: ThinPlateBasisSpec,
354        /// Uniform coordinate scale estimated on a fresh build and persisted
355        /// for exact frozen replay.
356        input_scale: Option<crate::IsotropicScale>,
357    },
358    Sphere {
359        feature_cols: Vec<usize>,
360        spec: SphericalSplineBasisSpec,
361    },
362    /// Constant-curvature (`M_κ`) geodesic-kernel smooth over κ-stereographic
363    /// chart coordinates (#944): one construction interpolating
364    /// S^d → ℝ^d → H^d through the spec's fixed κ. The Wahba S² smooth is the
365    /// structural template; the geometry comes from
366    /// `geometry::constant_curvature::ConstantCurvature`.
367    ConstantCurvature {
368        feature_cols: Vec<usize>,
369        spec: ConstantCurvatureBasisSpec,
370    },
371    Matern {
372        feature_cols: Vec<usize>,
373        spec: MaternBasisSpec,
374        input_scale: Option<crate::IsotropicScale>,
375    },
376    /// Measure-jet spline smooth: multiscale local-jet-residual energy of the
377    /// empirical measure (centers as μ-quadrature, masses as μ-weights — no
378    /// graph, mesh, or neighbor set inside the statistical object). The
379    /// feature columns are ambient coordinates of data concentrated near an
380    /// unknown low-dimensional, possibly stratified set.
381    MeasureJet {
382        feature_cols: Vec<usize>,
383        spec: MeasureJetBasisSpec,
384        input_scale: Option<crate::IsotropicScale>,
385    },
386    Duchon {
387        feature_cols: Vec<usize>,
388        spec: DuchonBasisSpec,
389        input_scale: Option<crate::IsotropicScale>,
390    },
391    Pca {
392        feature_cols: Vec<usize>,
393        basis_matrix: Array2<f64>,
394        centered: bool,
395        #[serde(default = "default_pca_smooth_penalty")]
396        smooth_penalty: f64,
397        #[serde(default)]
398        center_mean: Option<Array1<f64>>,
399        #[serde(default)]
400        pca_basis_path: Option<PathBuf>,
401        #[serde(default = "default_pca_chunk_size")]
402        chunk_size: usize,
403    },
404    /// Tensor-product smooth built from 1D B-spline marginals.
405    ///
406    /// This is the `te()`-style construction used when axes have different units/scales
407    /// (for example, space x time) and isotropic radial kernels are not appropriate.
408    TensorBSpline {
409        feature_cols: Vec<usize>,
410        spec: TensorBSplineSpec,
411    },
412}
413
414impl SmoothBasisSpec {
415    /// Conservative lower bound on the number of sample rows needed for this
416    /// smooth basis to have a well-posed REML fit.
417    ///
418    /// Each basis kind answers the question for itself, so the workflow does
419    /// not have to know how many columns a B-spline, tensor product, PCA
420    /// projection, or spatial kernel emits. The contract is a *lower bound*:
421    /// returning too small a number is permitted (the inner solver will catch
422    /// any genuine n-vs-rank failure that slips past); returning too large a
423    /// number is a regression because it rejects legitimate fits.
424    ///
425    /// Rationale: B-spline / tensor / PCA bases have a closed-form column
426    /// count, so we use the exact dimension. Radial bases (TPS, Matern,
427    /// Duchon, Sphere) and factor smooths choose their column count from the
428    /// data (`heuristic_centers`, `unique_count`); we fall back to a small
429    /// constant floor because a fit on fewer than five rows cannot stabilise
430    /// any radial smooth regardless of the configured kernel scale.
431    pub fn min_sample_rows(&self) -> usize {
432        // Floor used for data-driven bases whose column count is not known
433        // from the spec alone. Five rows is the minimum at which the inner
434        // pivot/QR + REML smoothing-parameter search has any chance of being
435        // well-posed for a non-parametric smooth.
436        const RADIAL_FLOOR: usize = 5;
437
438        match self {
439            Self::ByVariable { inner, .. } => inner.min_sample_rows(),
440            Self::FactorSumToZero { inner, levels, .. } => {
441                // L-1 independent deviation blocks each carrying the inner
442                // basis dimension. Skip the levels-multiplier if it doesn't
443                // bring more rows; we want the *lower bound* not the rank.
444                let inner_min = inner.min_sample_rows();
445                let lvls = levels.len().saturating_sub(1).max(1);
446                inner_min.saturating_mul(lvls)
447            }
448            Self::BSpline1D { spec, .. } => bspline_basis_min_rows(spec),
449            Self::BySmooth { smooth, .. } => smooth.min_sample_rows(),
450            Self::FactorSmooth { spec } => {
451                // Replicates the marginal once per level; without a known
452                // level count we conservatively require at least the marginal
453                // basis dimension.
454                bspline_basis_min_rows(&spec.marginal)
455            }
456            Self::ThinPlate { .. }
457            | Self::Sphere { .. }
458            | Self::ConstantCurvature { .. }
459            | Self::Matern { .. }
460            | Self::MeasureJet { .. }
461            | Self::Duchon { .. } => RADIAL_FLOOR,
462            Self::Pca { basis_matrix, .. } => basis_matrix.ncols().max(1),
463            Self::TensorBSpline { spec, .. } => {
464                // A `te(...)` smooth is *penalized*: each margin carries a
465                // difference (wiggliness) penalty and the tensor inherits a
466                // Kronecker-sum penalty `S = Σ_i I ⊗ … ⊗ S_i ⊗ … ⊗ I`. The raw
467                // column count is the *product* of the per-marginal column
468                // counts, but that product is the lower bound for an
469                // *unpenalized* tensor regression — it is the number of rows you
470                // would need to identify every interaction column with no
471                // regularization. The penalty regularizes all of those
472                // interaction directions; only the combined penalty *null space*
473                // (the tensor product of the per-margin polynomial trends, a
474                // handful of columns) must be identified by the data, and the
475                // smoothing-parameter search shrinks the rest. The effective
476                // degrees of freedom of the fitted `te()` are therefore a small
477                // fraction of the column product, which is exactly why mgcv
478                // fits a default `te(x, y)` on a couple hundred rows.
479                //
480                // The honest *penalized* lower bound is the **sum** of the
481                // per-marginal column counts, not their product: a row floor of
482                // `Σ_i k_i` still guarantees enough data to identify each
483                // margin's additive main-effect (the largest sub-block the
484                // penalty cannot shrink to zero), while no longer conflating
485                // unpenalized column-count identifiability with penalized
486                // well-posedness. This accepts moderate-`n` penalized tensors
487                // (e.g. a 20×20 default basis on n=200) yet still rejects a
488                // genuinely undersized fit where `n < Σ_i k_i` and even the
489                // additive part is rank-deficient.
490                //
491                // Binary / low-cardinality margins (#724): gam will accept a
492                // `te(x, badh)` whose `badh ∈ {0, 1}` margin nominally requests
493                // more basis columns than `badh` has unique values, where mgcv
494                // refuses the unpenalized term as ill-posed ("badh has
495                // insufficient unique values to support k knots"). This is
496                // correct-by-design, *not* a degenerate fit: the marginal
497                // wiggliness penalty on the `badh` axis has a null space that is
498                // exactly its identifiable trend (the two cell means of a binary
499                // covariate), and the Kronecker-sum penalty shrinks every tensor
500                // column outside that null space toward zero. The resulting fit
501                // is the well-posed "per-level `x` smooth + binary main effect"
502                // that mgcv reaches only after manually collapsing the basis —
503                // gam reaches it automatically because the penalty, not the raw
504                // column count, sets the effective rank. A genuinely
505                // rank-deficient design (penalty null space wider than the data
506                // can support) is still caught downstream by the inner pivoted
507                // factorization, which owns the exact n-vs-rank decision; this
508                // pre-fit gate only refuses the grossly-undersized formula.
509                let mut total: usize = 0;
510                for marginal in &spec.marginalspecs {
511                    let m = bspline_basis_min_rows(marginal);
512                    total = total.saturating_add(m.max(1));
513                }
514                total.max(RADIAL_FLOOR)
515            }
516        }
517    }
518
519    /// Stable structural discriminant for warm-start cache keying (#869).
520    ///
521    /// Two smooths that produce different bases / penalty structures must map
522    /// to different strings here so they cannot collide on the persistent
523    /// warm-start `cache_key` (which is otherwise blind to topology: it hashes
524    /// only the raw input column count, so e.g. `sphere` vs `torus` vs
525    /// `euclidean` candidates fit on the *same* data would otherwise share one
526    /// key and cross-contaminate each other's β/ρ seed). The string is the
527    /// topology identity, not the fitted coefficients, so same-topology refits
528    /// (the screen→full-refit cascade) still hit the same key and reuse work.
529    pub fn structural_kind(&self) -> &'static str {
530        match self {
531            Self::ByVariable { .. } => "by_variable",
532            Self::FactorSumToZero { .. } => "factor_sum_to_zero",
533            Self::BSpline1D { .. } => "bspline_1d",
534            Self::BySmooth { .. } => "by_smooth",
535            Self::FactorSmooth { .. } => "factor_smooth",
536            Self::ThinPlate { .. } => "thin_plate",
537            Self::Sphere { .. } => "sphere",
538            Self::ConstantCurvature { .. } => "constant_curvature",
539            Self::Matern { .. } => "matern",
540            Self::MeasureJet { .. } => "measurejet",
541            Self::Duchon { .. } => "duchon",
542            Self::Pca { .. } => "pca",
543            Self::TensorBSpline { .. } => "tensor_bspline",
544        }
545    }
546
547    /// True for a tensor-product smooth that is only *marginally* centered
548    /// (`ti(...)`, [`TensorBSplineIdentifiability::MarginalSumToZero`]): its
549    /// per-margin sum-to-zero reparameterization `(B_xZ_x)⊗(B_zZ_z)` has ALREADY
550    /// removed each axis's main effect analytically (mgcv-identical), so its
551    /// main-effect removal is complete and it must take NO additional
552    /// owner-residualization block. Residualizing it a second time against the
553    /// realized main-effect designs is a grid-fragile no-op on an exact tensor
554    /// grid but eats genuine pure-interaction curvature off-grid (#1470).
555    pub fn is_marginally_centered_tensor(&self) -> bool {
556        matches!(
557            self,
558            Self::TensorBSpline { spec, .. }
559                if matches!(spec.identifiability, TensorBSplineIdentifiability::MarginalSumToZero)
560        )
561    }
562
563    /// A sum-to-zero factor smooth (`bs="sz"`) has ALREADY removed the
564    /// cross-group main effect analytically, in coefficient space, via its
565    /// `Σ_g d_g(x) ≡ 0` reparameterization (`L-1` deviation blocks with the
566    /// reference level the negative sum of the others) — exactly mgcv's `sz`
567    /// construction, which is self-identifiable against an overlapping `s(x)`
568    /// with no further constraint. Residualizing it a SECOND time against the
569    /// realized B-spline span of the explicit `s(x)` smooth is redundant in
570    /// exact arithmetic (the common-to-all-groups component is zero by
571    /// construction) and actively HARMFUL on finite data: each deviation block
572    /// is a B-spline in `x` whose realized columns share `s(x)`'s span, so the
573    /// joint residualization collapses the full `L·k`-column deviation design to
574    /// `L·k − rank(s(x))` columns and eats the within-group curvature `s(x)`
575    /// cannot represent. REML then rails the deviation smoothing parameter and
576    /// the factor smooth under-recovers (#1605). This is the exact analogue of
577    /// the marginally-centered tensor (`ti`) exemption (#1470), so such a term
578    /// takes NO owner-residualization block.
579    pub fn is_sum_to_zero_factor_smooth(&self) -> bool {
580        matches!(
581            self,
582            Self::FactorSumToZero { .. }
583                | Self::FactorSmooth {
584                    spec: FactorSmoothSpec {
585                        flavour: FactorSmoothFlavour::Sz,
586                        ..
587                    }
588                }
589        )
590    }
591
592    /// Feature columns this basis consumes, used alongside `structural_kind`
593    /// to disambiguate two same-kind smooths on different axes. Wrapper
594    /// variants delegate to their inner basis.
595    pub fn structural_feature_cols(&self) -> Vec<usize> {
596        match self {
597            Self::ByVariable { inner, .. } | Self::FactorSumToZero { inner, .. } => {
598                inner.structural_feature_cols()
599            }
600            Self::BySmooth { smooth, .. } => smooth.structural_feature_cols(),
601            Self::FactorSmooth { .. } => Vec::new(),
602            Self::BSpline1D { feature_col, .. } => vec![*feature_col],
603            Self::ThinPlate { feature_cols, .. }
604            | Self::Sphere { feature_cols, .. }
605            | Self::ConstantCurvature { feature_cols, .. }
606            | Self::Matern { feature_cols, .. }
607            | Self::MeasureJet { feature_cols, .. }
608            | Self::Duchon { feature_cols, .. }
609            | Self::Pca { feature_cols, .. }
610            | Self::TensorBSpline { feature_cols, .. } => feature_cols.clone(),
611        }
612    }
613}
614
615/// Lower bound on the number of sample rows a 1D B-spline smooth needs for a
616/// well-posed *penalized* REML fit. Used as the per-smooth row floor in
617/// [`SmoothBasisSpec::min_sample_rows`].
618///
619/// For a *singly*-penalized smooth the floor is the full column count: the
620/// wiggliness penalty leaves the order-`m` polynomial trend unpenalized, and
621/// gam's original gate conservatively required enough rows for the whole basis.
622/// That conservative floor is kept here unchanged.
623///
624/// A *double*-penalized smooth (mgcv `select=TRUE`) is different: it adds a
625/// second penalty on the wiggliness penalty's null space, so even the
626/// polynomial trend is shrinkable toward zero and *nothing* in the basis
627/// requires unpenalized identification by the data — exactly the reasoning the
628/// `TensorBSpline` arm of [`SmoothBasisSpec::min_sample_rows`] already applies
629/// to a penalized tensor. Its honest floor is therefore a small stabilization
630/// constant, not the column count. This is what lets mgcv (and now gam) fit
631/// several `select=TRUE` smooths on a dataset whose row count is below the
632/// summed basis width (e.g. the n≈30 `wine_gamair` fold, 5 `ps` smooths,
633/// p≈51): the penalties, not the data, set the effective rank. The bounded
634/// outer REML loop still terminates, and the genuine n-vs-rank decision is
635/// owned downstream by the inner pivoted factorization. Without this, gam
636/// rejected the fit outright (or, before the gate existed, the outer REML loop
637/// wandered the flat overparameterized surface until the benchmark wall budget
638/// killed it — #1089).
639pub fn bspline_basis_min_rows(spec: &crate::basis::BSplineBasisSpec) -> usize {
640    use crate::basis::BSplineKnotSpec;
641    let columns = match &spec.knotspec {
642        BSplineKnotSpec::Generate {
643            num_internal_knots, ..
644        } => *num_internal_knots + spec.degree + 1,
645        BSplineKnotSpec::Automatic {
646            num_internal_knots: Some(k),
647            ..
648        } => *k + spec.degree + 1,
649        BSplineKnotSpec::Automatic {
650            num_internal_knots: None,
651            ..
652        } => {
653            // Knot count is data-derived (`default_internal_knot_count_for_data`).
654            // A minimal cubic basis is `degree + 2` columns; below that the
655            // basis cannot represent a non-parametric smooth.
656            spec.degree + 2
657        }
658        BSplineKnotSpec::Provided(knots) => knots.len().saturating_sub(spec.degree + 1).max(1),
659        // cr basis dimension equals the knot count (no degree offset).
660        BSplineKnotSpec::NaturalCubicRegression { knots } => knots.len(),
661        BSplineKnotSpec::PeriodicUniform { num_basis, .. } => *num_basis,
662    };
663    let columns = columns.max(spec.degree + 2);
664
665    if spec.double_penalty {
666        // Fully shrinkable basis: only a small stabilization floor must be
667        // identified by the data, capped by the actual column count.
668        const DOUBLE_PENALTY_FLOOR: usize = 2;
669        DOUBLE_PENALTY_FLOOR.min(columns).max(1)
670    } else {
671        columns
672    }
673}
674
675#[derive(Debug, Clone, Serialize, Deserialize)]
676pub enum ByVariableSpec {
677    Numeric,
678    Level { value_bits: u64, label: String },
679}
680
681#[derive(Debug, Clone, Serialize, Deserialize)]
682pub enum ByVarKind {
683    Numeric {
684        feature_col: usize,
685    },
686    Factor {
687        feature_col: usize,
688        ordered: bool,
689        frozen_levels: Option<Vec<u64>>,
690    },
691}
692
693#[derive(Debug, Clone, Serialize, Deserialize)]
694pub struct FactorSmoothSpec {
695    pub continuous_cols: Vec<usize>,
696    pub group_col: usize,
697    pub marginal: BSplineBasisSpec,
698    pub flavour: FactorSmoothFlavour,
699    pub group_frozen_levels: Option<Vec<u64>>,
700    /// Fit-time global-orthogonality chart `Z` for this term (`s(x) + fs(x, g)`
701    /// overlap residualization, #978), in the post-joint-null-`Q` coordinates
702    /// (the raw rebuild recomputes any `Q` itself; `fs` penalties are
703    /// typically full-rank so `Q` is absent). Training-row dependent, hence
704    /// persisted; replayed verbatim by `apply_global_smooth_identifiability`.
705    #[serde(default)]
706    pub frozen_global_orthogonality: Option<Array2<f64>>,
707}
708
709#[derive(Debug, Clone, Serialize, Deserialize)]
710pub enum FactorSmoothFlavour {
711    Fs { m_null_penalty_orders: Vec<usize> },
712    Sz,
713    Re,
714}
715
716#[derive(Debug, Clone, Serialize, Deserialize)]
717pub struct TensorBSplineSpec {
718    pub marginalspecs: Vec<BSplineBasisSpec>,
719    #[serde(default)]
720    pub periods: Vec<Option<f64>>,
721    #[serde(default = "default_tensor_double_penalty")]
722    pub double_penalty: bool,
723    #[serde(default)]
724    pub identifiability: TensorBSplineIdentifiability,
725    #[serde(default)]
726    pub penalty_decomposition: TensorBSplinePenaltyDecomposition,
727}
728
729pub const fn default_tensor_double_penalty() -> bool {
730    true
731}
732
733impl Default for TensorBSplineSpec {
734    fn default() -> Self {
735        Self {
736            marginalspecs: Vec::new(),
737            periods: Vec::new(),
738            double_penalty: default_tensor_double_penalty(),
739            identifiability: TensorBSplineIdentifiability::default(),
740            penalty_decomposition: TensorBSplinePenaltyDecomposition::default(),
741        }
742    }
743}
744
745#[derive(Debug, Default, Clone, Serialize, Deserialize)]
746pub enum TensorBSplineIdentifiability {
747    None,
748    #[default]
749    SumToZero,
750    /// mgcv `ti(...)` semantics: a *tensor interaction* smooth that excludes the
751    /// marginal main effects. A sum-to-zero constraint is applied to **each
752    /// marginal basis independently** before forming the tensor product, so the
753    /// resulting column space contains no function of a single variable alone —
754    /// only the pure interaction survives. The realized identifiability
755    /// transform is the Kronecker product `Z = Z₀ ⊗ Z₁ ⊗ … ⊗ Z_{d-1}` of the
756    /// per-margin sum-to-zero null-space bases, which is exactly the
757    /// reparameterization that turns the full-tensor design into the tensor
758    /// product of the centered margins.
759    MarginalSumToZero,
760    FrozenTransform {
761        transform: Array2<f64>,
762    },
763}
764
765#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
766pub enum TensorBSplinePenaltyDecomposition {
767    /// mgcv `te(...)`: one overlapping Kronecker-product penalty per margin,
768    /// `S_j` embedded against identities in the other tensor factors.
769    #[default]
770    MarginalKroneckerSum,
771    /// mgcv `t2(...)`: split every marginal coefficient space into penalized
772    /// range and penalty-null subspaces, then emit one disjoint tensor-subspace
773    /// penalty for every non-empty penalized/null combination.
774    Separable,
775}
776
777#[derive(Debug, Clone, Serialize, Deserialize)]
778pub struct SmoothTermSpec {
779    pub name: String,
780    pub basis: SmoothBasisSpec,
781    pub shape: ShapeConstraint,
782    /// Joint-null absorption rotation captured at fit time. `Some(Q)` means
783    /// the fitted coefficient vector lives in `γ`-coordinates with
784    /// `β_raw = Q · γ`; prediction must rotate the raw-basis design via
785    /// `X_new = X_new_raw · Q` to match. `None` means either the smooth had
786    /// no joint null space (penalty already full-rank) or rotation was
787    /// suppressed (smooth carries shape constraints whose cone geometry
788    /// would not survive an arbitrary orthogonal rotation). Persisted so
789    /// `save → load → predict` is bit-equivalent to in-memory prediction.
790    #[serde(default)]
791    pub joint_null_rotation: Option<crate::basis::JointNullRotation>,
792    /// The span-preserving parametric orthogonalization this term carries
793    /// (#2747), when it takes that arm. Persisted so `save → load → predict`
794    /// rebuilds the design the fitted coefficients belong to.
795    #[serde(default)]
796    pub frozen_parametric_residualization: Option<ParametricResidualizationChart>,
797}
798
799/// The predict-time replay record of a span-preserving parametric
800/// orthogonalization (#2747): this term's realized block is `X·T − C·R`.
801///
802/// `T` is not here — it is the ordinary coefficient transform, already absorbed
803/// into the basis metadata, so a rebuilt `design_local` arrives with it applied.
804/// What cannot be absorbed is `R`, because it multiplies the CONSTRAINT block
805/// rather than the basis.
806///
807/// `R` is TRAINING-ROW data and is never re-derived, for the same reason
808/// `frozen_global_orthogonality` is not (#978): the fit already decided this
809/// term's residualization, and rederiving it from prediction rows would
810/// silently evaluate a different model. `C` itself IS rebuilt at the new rows —
811/// that is what `C` is — from a recipe that is deterministic given the spec: the
812/// intercept and owned linear axes come from
813/// `build_parametric_constraint_block_for_term`, and the owner smooths are named
814/// here rather than re-derived, because which owners bound is decided by a
815/// cross-residual on the FIT rows.
816#[derive(Debug, Clone, Serialize, Deserialize)]
817pub struct ParametricResidualizationChart {
818    /// Indices into `TermCollectionSpec::smooth_terms` of the owner smooths
819    /// whose realized designs joined the constraint block, in the order
820    /// `build_constraint_block` stacked them.
821    pub owner_terms: Vec<usize>,
822    /// Whether the parametric block (intercept + owned linear axes) led the
823    /// constraint block. Recorded rather than recomputed so a spec that changes
824    /// shape cannot silently re-order the columns `correction`'s rows index.
825    pub has_parametric_block: bool,
826    /// `R`, `q × k`, stated against the RAW constraint columns.
827    pub correction: Array2<f64>,
828}
829
830/// Which construction `apply_global_smooth_identifiability` chose to make one
831/// term's realized block orthogonal to its constraint block.
832///
833/// Frozen with the block rather than re-decided per rebuild: the containment
834/// test that picks between them is a function of the realized design, hence of
835/// any basis parameter an outer search is moving, and a fit that re-decided it
836/// per trial would step the model DIMENSION mid-search (`Delete` costs a
837/// coefficient direction, `Residualize` costs none).
838#[derive(Debug, Clone, Copy, PartialEq, Eq)]
839pub enum SmoothCollectionGaugeArm {
840    /// `X·Z`, `Z` spanning `null((XᵀC)ᵀ)` — free only under containment, and
841    /// then identical to the pre-`76a520c45` path.
842    Delete,
843    /// `X·T − C·R`, span-preserving — always licensed.
844    Residualize,
845}
846
847/// The COLLECTION's gauge for one smooth term: the realized constraint block it
848/// was made orthogonal to, and which construction did it.
849///
850/// # Why this is freezable, and why a term-local rebuild needs it
851///
852/// `C = [intercept | owned linear axes | owner smooths' realized blocks]` is a
853/// function of the data and of OTHER terms — never of this term's own basis
854/// parameters. So `C` is invariant along an outer search over this term's `ψ`
855/// (a length scale, a curvature, a measure-jet dial), while the two objects
856/// derived FROM it are not: `T` and `R` are both functions of the realized
857/// design, so both move with `ψ`.
858///
859/// That asymmetry is the whole content of gam#2747's second half. The spatial
860/// outer search rebuilds one term LOCALLY per trial and splices the result into
861/// the collection design; a term-local build cannot see `C`, so before this
862/// existed the splice replaced the design and its chart while leaving the
863/// collection's `R` behind, and the fit shipped `X(ψ̂)·Z − C·R(ψ₀)` — a block
864/// orthogonal to nothing, with `‖XᵀC‖/(‖X‖‖C‖)` measured at `4.15e-1` against
865/// the `1e-8` bar the same step asserts whenever it applies a transform.
866///
867/// Carrying the ψ-INDEPENDENT half forward and re-deriving the ψ-dependent half
868/// is therefore not an optimization — it is the only formulation under which
869/// the criterion an outer search minimizes and the model the fit ships are the
870/// same object.
871#[derive(Debug, Clone)]
872pub struct SmoothCollectionGauge {
873    /// Which construction the collection chose.
874    pub arm: SmoothCollectionGaugeArm,
875    /// `C`, `n × q`, exactly as `build_constraint_block` stacked it.
876    pub constraint_block: Array2<f64>,
877    /// Indices into the collection's smooth terms of the owner smooths whose
878    /// realized designs joined `C`, in stack order. Recorded for the same
879    /// reason [`ParametricResidualizationChart::owner_terms`] is.
880    pub owner_terms: Vec<usize>,
881    /// Whether the parametric block led `C`.
882    pub has_parametric_block: bool,
883    /// The TERM-LOCAL identifiability chart the gauged block was derived ON —
884    /// `z_local`, before this gauge composed its own `T` on top of it (gam#2760).
885    ///
886    /// The term's `BasisMetadata` records the COMPOSITION `z_local · T`, which is
887    /// what a predict-time replay wants: predict does not move `ψ`, so replaying
888    /// the composed chart reproduces the fitted block exactly. A caller that DOES
889    /// move `ψ` re-derives `T` here and must therefore rebuild in `z_local`, not
890    /// in the composition — otherwise the collection's orthogonalization is
891    /// applied twice, once in the stale `ψ₀` chart carried by the spec and once
892    /// freshly by this gauge.
893    ///
894    /// Measured before this existed, on a one-Duchon-term collection with
895    /// `C = [1]` and the `Delete` arm: the replay spec carried
896    /// `FrozenTransform(12, 11)`, the doubly-charted rebuild had orthogonality
897    /// residual `1.5e-12` at the fit's own `ℓ` and `9.0e-1` one octave away, so
898    /// the gauge resolved a direction and deleted a column at every `ψ`. The
899    /// term reached the splice one column short and every κ fixture on a Duchon
900    /// term refused in 0.2 s. On the `Residualize` arm the second application is
901    /// idempotent, which is why the same defect was invisible on Matérn.
902    ///
903    /// `None` means the term-local build applied no chart of its own — the usual
904    /// case for the radial families, whose `OrthogonalToParametric` policy defers
905    /// entirely to this gauge.
906    ///
907    /// Like `C` and the arm, this is ψ-INDEPENDENT: it is a center-space
908    /// constraint (`1ᵀα = 0`, a linear-orthogonality frame) or a frozen replay
909    /// chart, never a function of the realized design.
910    pub local_identifiability_transform: Option<Array2<f64>>,
911    /// The width of the TERM-LOCAL block this gauge was derived on, before the
912    /// arm ran (gam#2760).
913    ///
914    /// The collection's realized width is this minus whatever the arm removed,
915    /// and both halves are needed to read a rebuild that comes out narrow. A
916    /// rebuild whose LOCAL width differs from this is not the same basis and is
917    /// a defect; a rebuild that matches here and still comes out narrow after the
918    /// arm has hit a ψ at which the realized design loses rank in this gauge's
919    /// chart — a statement about the trial point, not about the rebuild.
920    pub local_columns: usize,
921}
922
923#[derive(Debug, Clone)]
924pub struct SmoothTerm {
925    pub name: String,
926    pub coeff_range: Range<usize>,
927    pub shape: ShapeConstraint,
928    /// Active local penalty identities. Numerical and semantic channels are
929    /// inseparable, including after a preceding candidate is dropped.
930    pub active_penalties: Vec<ActivePenalty>,
931    pub dropped_penalties: Vec<DroppedPenaltyInfo>,
932    pub metadata: BasisMetadata,
933    /// Optional term-local lower bounds for constrained coefficients.
934    /// `-inf` means unconstrained.
935    pub lower_bounds_local: Option<Array1<f64>>,
936    /// Optional term-local inequality constraints in local coefficient coordinates.
937    /// `A_local * beta_local >= b_local`.
938    pub linear_constraints_local: Option<LinearInequalityConstraints>,
939    /// Optional factored tensor-product representation preserved for operator-backed
940    /// assembly in the main design builder.
941    pub kronecker_factored: Option<KroneckerFactoredBasis>,
942    /// Joint-null absorption rotation. `Some(Q)` records the orthonormal
943    /// `(p_local × p_local)` matrix that was applied to this term's design
944    /// and per-block penalties at construction time:
945    /// `term_design ← X_raw · Q`, `active_penalties[k].matrix ← Qᵀ · S_raw · Q`.
946    /// The smooth's coefficient vector therefore lives in the rotated
947    /// (`γ`) coordinate system, with `β_raw = Q · γ` recovering the raw
948    /// pre-rotation parameterization. `None` means either no joint null
949    /// space (penalty already full-rank) or rotation was suppressed —
950    /// suppression fires when the smooth carries shape constraints
951    /// (lower bounds or local linear inequalities) that would lose their
952    /// cone geometry under a general orthogonal rotation.
953    ///
954    /// Prediction-side replay: callers building a new-data design `X_new_raw`
955    /// from the *raw* basis must call [`SmoothTerm::apply_rotation_to_predict`]
956    /// (or equivalent) to obtain `X_new = X_new_raw · Q` matching this
957    /// term's coefficient system.
958    ///
959    /// Persistence replay: `freeze_term_collection_from_design` copies this
960    /// rotation into `SmoothTermSpec`, which is serialized with fitted-model
961    /// payloads and reused by the predict-time basis builder. Saved models
962    /// therefore replay the same `X_new_raw · Q` transform as in-memory
963    /// prediction.
964    pub joint_null_rotation: Option<crate::basis::JointNullRotation>,
965    /// Global-orthogonality transform that `apply_global_smooth_identifiability`
966    /// applied to this term's design but could NOT embed into `metadata`
967    /// (factor-smooth kinds: `sz` metadata is per-marginal, `fs` metadata has
968    /// no transform slot — #978). `freeze_term_collection_from_design` copies
969    /// it onto the term's basis spec (`frozen_global_orthogonality`) so the
970    /// predict-side rebuild replays it instead of emitting the unresidualized
971    /// (wider) design that the fitted coefficients no longer match.
972    /// Chart convention is per kind: post-`Q` `Z` for `sz` (the raw rebuild
973    /// reapplies `Q` itself, #700), full `Q·Z` chart for `fs`.
974    pub unabsorbed_global_orthogonality: Option<Array2<f64>>,
975    /// The span-preserving parametric orthogonalization applied to this term
976    /// (#2747): its realized block is `X·T − C·R`.
977    /// `freeze_term_collection_from_design` copies this onto
978    /// `SmoothTermSpec::frozen_parametric_residualization` so a predict-time
979    /// rebuild subtracts the same correction instead of emitting a design the
980    /// fitted coefficients do not match.
981    pub parametric_residualization: Option<ParametricResidualizationChart>,
982    /// The ψ-INDEPENDENT half of the global step, exported so a TERM-LOCAL
983    /// rebuild can put its result back into this collection's gauge instead of
984    /// silently leaving it out (#2747). `None` when the collection applied no
985    /// global transform to this term, and therefore also on every replay — a
986    /// replay is already in the gauge by construction.
987    ///
988    /// This is deliberately NOT frozen onto `SmoothTermSpec`: it is a fit-time
989    /// object for fit-time rebuilds. Predict-time replay is served by
990    /// [`ParametricResidualizationChart`] / `frozen_global_orthogonality`,
991    /// which carry the ψ-DEPENDENT half the fit settled on.
992    pub collection_gauge: Option<SmoothCollectionGauge>,
993}
994
995impl SmoothTerm {
996    /// Apply the joint-null absorption rotation to a raw new-data design
997    /// matrix, returning `X_new_raw · Q` when this term was rotated at
998    /// fit time, or `X_new_raw` unchanged when no rotation was applied.
999    ///
1000    /// Callers in the prediction path: after building the smooth's basis
1001    /// at new data via the *raw* basis builder (the same builder used at
1002    /// fit time, applied to `x_new` instead of the training rows), call
1003    /// this method on the resulting matrix before forming `X · β`. The
1004    /// fitted `β` lives in `γ`-coordinates if Q was applied; multiplying
1005    /// the un-rotated `X_new_raw` by `β` would give a wrong η.
1006    ///
1007    /// Returns an error if the raw design's column count does not match
1008    /// the rotation's `p_local`. The width invariant must hold: the raw
1009    /// basis builder MUST emit the same `p_local` columns that the
1010    /// fit-time builder did, and the rotation is `(p_local × p_local)`.
1011    pub fn apply_rotation_to_predict(
1012        &self,
1013        x_new_raw: Array2<f64>,
1014    ) -> Result<Array2<f64>, BasisError> {
1015        let Some(rot) = self.joint_null_rotation.as_ref() else {
1016            return Ok(x_new_raw);
1017        };
1018        let p_local = rot.rotation.nrows();
1019        if x_new_raw.ncols() != p_local {
1020            crate::bail_dim_basis!(
1021                "joint-null rotation replay for term '{}': raw design has {} columns, \
1022                 rotation expects {} (the raw basis builder must emit the same column \
1023                 count as at fit time)",
1024                self.name,
1025                x_new_raw.ncols(),
1026                p_local,
1027            );
1028        }
1029        Ok(gam_linalg::faer_ndarray::fast_ab(&x_new_raw, &rot.rotation))
1030    }
1031
1032    /// Dimension of the **joint** null space of this term's active penalties:
1033    /// the coefficient directions penalized by *no* penalty. The smooth-component
1034    /// Wald test ([`crate::inference::smooth_test::wood_smooth_test`]) treats this
1035    /// many leading coefficients as genuine unpenalized fixed effects and tests
1036    /// them at full rank; the remainder is the penalized sub-block tested with a
1037    /// rank-`≈EDF` truncated pseudo-inverse.
1038    ///
1039    /// Because every penalty block `S_k` is positive semi-definite,
1040    /// `vᵀ(Σ_k S_k)v = Σ_k vᵀ S_k v = 0` iff `S_k v = 0` for *every* `k`; the
1041    /// joint null space is therefore exactly `null(Σ_k S_k)`, of dimension
1042    /// `p_local − rank(Σ_k S_k)`. This is the **intersection** of the per-penalty
1043    /// null spaces, not their sum.
1044    ///
1045    /// Summing the per-penalty `nullspace_dims` instead (the historical defect
1046    /// behind #1360) *unions* the null spaces and badly over-counts: a
1047    /// double-penalty smooth carries a bending penalty (null space = its
1048    /// polynomial part) plus a complementary null-space ridge (which penalizes
1049    /// exactly that polynomial part), so the two null spaces are disjoint and the
1050    /// joint null space is empty — yet the per-penalty dims sum to nearly
1051    /// `p_local`. Feeding that inflated count to the Wald test makes it test
1052    /// almost the whole shrunk block at full rank, manufacturing overwhelming
1053    /// "significance" for a term the fit drove to ~0 EDF.
1054    pub fn wald_unpenalized_dim(&self) -> usize {
1055        joint_unpenalized_dim(self.coeff_range.len(), &self.active_penalties)
1056    }
1057}
1058
1059/// Numeric core of [`SmoothTerm::wald_unpenalized_dim`]: the dimension of the
1060/// joint null space `∩_k null(S_k) = null(Σ_k S_k)` of a term's local penalty
1061/// blocks, with a conservative fallback when a penalty is not materialized as a
1062/// full `p_local × p_local` matrix (e.g. a Kronecker tensor factor).
1063pub fn joint_unpenalized_dim(p_local: usize, active_penalties: &[ActivePenalty]) -> usize {
1064    use gam_linalg::faer_ndarray::FaerEigh;
1065    if p_local == 0 {
1066        return 0;
1067    }
1068    if active_penalties.is_empty() {
1069        // No penalty ⇒ a wholly unpenalized (fixed-effect) block.
1070        return p_local;
1071    }
1072    // Sum the penalties that are materialized as full `p_local × p_local`
1073    // blocks (the common smooth case). The covariance block the Wald test
1074    // slices lives in this same coefficient basis (post joint-null rotation),
1075    // so the rank is computed in the right metric.
1076    let mut s_total = Array2::<f64>::zeros((p_local, p_local));
1077    let mut materialized = 0usize;
1078    for penalty in active_penalties {
1079        let s = &penalty.matrix;
1080        if s.nrows() == p_local && s.ncols() == p_local {
1081            s_total += s;
1082            materialized += 1;
1083        }
1084    }
1085    if materialized == active_penalties.len() {
1086        let symmetric = {
1087            let transpose = s_total.t().to_owned();
1088            (&s_total + &transpose) * 0.5
1089        };
1090        if let Ok((evals, _)) = symmetric.eigh(faer::Side::Lower) {
1091            let max_abs = evals.iter().fold(0.0_f64, |acc, &v| acc.max(v.abs()));
1092            if max_abs == 0.0 {
1093                // All penalties identically zero ⇒ unpenalized block.
1094                return p_local;
1095            }
1096            let tol = max_abs * (p_local as f64) * 1e-12;
1097            let rank = evals.iter().filter(|&&v| v > tol).count();
1098            return p_local.saturating_sub(rank);
1099        }
1100    }
1101    // Conservative fallback when a penalty is not a materialized full block
1102    // (e.g. a Kronecker tensor factor): with ≥2 active penalties the joint
1103    // null space is almost always empty (the only over-rejecting direction);
1104    // with a single penalty it is exactly that penalty's own null space.
1105    if active_penalties.len() >= 2 {
1106        0
1107    } else {
1108        active_penalties
1109            .iter()
1110            .map(|penalty| penalty.nullity)
1111            .min()
1112            .unwrap_or(0)
1113            .min(p_local)
1114    }
1115}
1116
1117#[derive(Debug, Clone, Serialize, Deserialize)]
1118pub struct PenaltyBlockInfo {
1119    pub global_index: usize,
1120    pub termname: Option<String>,
1121    pub penalty: ActivePenaltyInfo,
1122}
1123
1124#[derive(Debug, Clone, Serialize, Deserialize)]
1125pub struct DroppedPenaltyBlockInfo {
1126    pub termname: Option<String>,
1127    pub penalty: DroppedPenaltyInfo,
1128}
1129
1130#[derive(Debug, Clone)]
1131pub struct SmoothDesign {
1132    pub term_designs: Vec<DesignMatrix>,
1133    /// Per-term block-local penalties.  Each `col_range` is relative to the
1134    /// smooth block (i.e. indexing into the concatenation of `term_designs`).
1135    pub penalties: Vec<BlockwisePenalty>,
1136    pub nullspace_dims: Vec<usize>,
1137    pub penaltyinfo: Vec<PenaltyBlockInfo>,
1138    pub dropped_penaltyinfo: Vec<DroppedPenaltyBlockInfo>,
1139    pub terms: Vec<SmoothTerm>,
1140    /// Optional smooth-block lower bounds in smooth coefficient coordinates.
1141    /// Length equals `total_smooth_cols()` when present.
1142    pub coefficient_lower_bounds: Option<Array1<f64>>,
1143    /// Optional smooth-block inequality constraints:
1144    /// `A_smooth * beta_smooth >= b`.
1145    pub linear_constraints: Option<LinearInequalityConstraints>,
1146}
1147
1148impl SmoothDesign {
1149    pub fn total_smooth_cols(&self) -> usize {
1150        self.term_designs.iter().map(DesignMatrix::ncols).sum()
1151    }
1152    pub fn nrows(&self) -> usize {
1153        self.term_designs.first().map_or(0, DesignMatrix::nrows)
1154    }
1155}
1156
1157#[derive(Debug, Clone)]
1158pub struct RawSmoothDesign {
1159    pub term_designs: Vec<DesignMatrix>,
1160    /// Sum of every fixed affine term contribution on the realized rows.
1161    pub affine_offset: Array1<f64>,
1162    /// Per-term block-local penalties.  Each `col_range` is relative to the
1163    /// smooth block (i.e. indexing into the concatenation of `term_designs`).
1164    pub penalties: Vec<BlockwisePenalty>,
1165    pub nullspace_dims: Vec<usize>,
1166    pub penaltyinfo: Vec<PenaltyBlockInfo>,
1167    pub dropped_penaltyinfo: Vec<DroppedPenaltyBlockInfo>,
1168    pub terms: Vec<SmoothTerm>,
1169    pub coefficient_lower_bounds: Option<Array1<f64>>,
1170    pub linear_constraints: Option<LinearInequalityConstraints>,
1171}
1172
1173impl RawSmoothDesign {
1174    pub fn total_smooth_cols(&self) -> usize {
1175        self.term_designs.iter().map(DesignMatrix::ncols).sum()
1176    }
1177    pub fn nrows(&self) -> usize {
1178        self.term_designs.first().map_or(0, DesignMatrix::nrows)
1179    }
1180}
1181
1182#[derive(Debug, Default, Clone, Serialize, Deserialize)]
1183pub enum BoundedCoefficientPriorSpec {
1184    #[default]
1185    None,
1186    Uniform,
1187    Beta {
1188        a: f64,
1189        b: f64,
1190    },
1191}
1192
1193#[derive(Debug, Clone, Serialize, Deserialize, Default)]
1194pub enum LinearCoefficientGeometry {
1195    #[default]
1196    Unconstrained,
1197    Bounded {
1198        min: f64,
1199        max: f64,
1200        #[serde(default)]
1201        prior: BoundedCoefficientPriorSpec,
1202    },
1203}
1204
1205#[derive(Debug, Clone, Serialize, Deserialize)]
1206pub struct LinearTermSpec {
1207    pub name: String,
1208    /// Primary feature column index. For Wilkinson-Rogers `:` interaction
1209    /// terms (`a:b[:c...]`) this is the first column in `feature_cols`; the
1210    /// realized design column is the elementwise product across every entry
1211    /// of `feature_cols`. Plain (non-interaction) linear terms set
1212    /// `feature_cols == vec![feature_col]`.
1213    pub feature_col: usize,
1214    /// Full list of columns whose elementwise product yields this term's
1215    /// design column. `len() >= 1`; `len() == 1` is a plain linear effect.
1216    #[serde(default)]
1217    pub feature_cols: Vec<usize>,
1218    /// Categorical-level gates for a factor-aware `:` interaction.
1219    ///
1220    /// Each `(col, level_bits)` multiplies the realized design column by the
1221    /// indicator `1[canonical_level_bits(data[row, col]) == level_bits]` (the
1222    /// canonical key collapses `-0.0`/`+0.0` and NaN payloads so numerically
1223    /// equal codes name one level; see `gam_data::canonical_level_bits`). This
1224    /// is how a
1225    /// `factor:x` (or `factor:factor`) interaction is expanded: `build_termspec`
1226    /// emits one `LinearTermSpec` per surviving cell of the categorical
1227    /// operand(s) (treatment-coded, first level dropped per factor), each
1228    /// carrying the numeric operands in `feature_cols` and the cell's level
1229    /// gate(s) here. Empty for a plain numeric `:` interaction or main effect,
1230    /// in which case the realized column is exactly the numeric product.
1231    #[serde(default)]
1232    pub categorical_levels: Vec<(usize, u64)>,
1233    /// Optional zero-centered shrinkage ridge with a REML-selected `λ`.
1234    /// Parametric effects are unpenalized/MLE by default;
1235    /// `linear(x, double_penalty=true)` opts into shrinkage.
1236    #[serde(default = "default_linear_term_double_penalty")]
1237    pub double_penalty: bool,
1238    #[serde(default)]
1239    pub coefficient_geometry: LinearCoefficientGeometry,
1240    #[serde(default)]
1241    pub coefficient_min: Option<f64>,
1242    #[serde(default)]
1243    pub coefficient_max: Option<f64>,
1244    /// The term's empirical function mass (`n⁻¹bᵀb` on the TRAINING design
1245    /// column), captured once at fit time by
1246    /// [`crate::smooth::freeze_term_collection_from_design`] and baked into the
1247    /// frozen/saved spec. `None` while a spec is still being resolved during
1248    /// fitting (before freezing) or for a `double_penalty=false` term, which
1249    /// carries no ridge and never needs a mass.
1250    ///
1251    /// A rebuild of this term's design column at PREDICT time (a held-out
1252    /// grid, a small group/class-anchor set, a single test row, …) must reuse
1253    /// this persisted value rather than recomputing it from the evaluation
1254    /// rows: a covariate that varies fine across the training set can easily
1255    /// be constant across a tiny evaluation subset by chance, and recomputing
1256    /// there would misfire the fit-time "identically zero" identifiability
1257    /// guard on a perfectly good term (#1561 REF_ERROR/METRIC_OFF triage).
1258    #[serde(default)]
1259    pub frozen_function_mass: Option<f64>,
1260}
1261
1262impl LinearTermSpec {
1263    /// Return the effective list of feature columns. Backfills from
1264    /// `feature_col` for legacy specs that predate the multi-column field.
1265    pub fn effective_feature_cols(&self) -> Vec<usize> {
1266        if self.feature_cols.is_empty() {
1267            vec![self.feature_col]
1268        } else {
1269            self.feature_cols.clone()
1270        }
1271    }
1272
1273    /// True when this term is a Wilkinson-Rogers `:` interaction (multi-col).
1274    pub fn is_interaction(&self) -> bool {
1275        self.feature_cols.len() > 1 || !self.categorical_levels.is_empty()
1276    }
1277
1278    /// Realize this linear term's `(n,)` design column from `data`.
1279    ///
1280    /// The column is the elementwise product of every numeric feature column
1281    /// (`effective_feature_cols`) gated by the categorical-level indicators in
1282    /// `categorical_levels`: each `(col, level_bits)` multiplies the running
1283    /// column by `1[canonical_level_bits(data[row, col]) == level_bits]` (signed
1284    /// zero / NaN canonicalized so numerically equal codes match). A plain numeric
1285    /// term (no `categorical_levels`) reduces to the bare product, matching the
1286    /// historical behaviour. A pure categorical interaction (empty
1287    /// `feature_cols`, non-empty `categorical_levels`) reduces to the cell
1288    /// indicator. Bounds are validated here; the returned column has length
1289    /// `data.nrows()`.
1290    pub fn realized_design_column(&self, data: ArrayView2<'_, f64>) -> Result<Array1<f64>, String> {
1291        let n = data.nrows();
1292        let p = data.ncols();
1293        let bounds = |col: usize| -> Result<(), String> {
1294            if col >= p {
1295                Err(format!(
1296                    "linear term '{}' feature column {} out of bounds for {} columns",
1297                    self.name, col, p
1298                ))
1299            } else {
1300                Ok(())
1301            }
1302        };
1303
1304        // Numeric operands. When `categorical_levels` is set we treat
1305        // `feature_cols` as the (possibly empty) numeric operand list and start
1306        // from a column of ones; otherwise we preserve the legacy backfill from
1307        // `feature_col` so a plain term with no `feature_cols` still resolves.
1308        let mut column = if self.categorical_levels.is_empty() {
1309            let cols = self.effective_feature_cols();
1310            for &c in &cols {
1311                bounds(c)?;
1312            }
1313            let mut acc = data.column(cols[0]).to_owned();
1314            for &c in cols.iter().skip(1) {
1315                acc *= &data.column(c);
1316            }
1317            acc
1318        } else {
1319            let mut acc = Array1::<f64>::ones(n);
1320            for &c in &self.feature_cols {
1321                bounds(c)?;
1322                acc *= &data.column(c);
1323            }
1324            acc
1325        };
1326
1327        for &(col, level_bits) in &self.categorical_levels {
1328            bounds(col)?;
1329            // Canonicalize the stored key once (loop-invariant) so the gate is
1330            // robust to level sets interned before signed-zero canonicalization
1331            // landed, not just to canonical data rows (#2146).
1332            let level_bits = gam_data::canonical_level_bits(f64::from_bits(level_bits));
1333            let gate = data.column(col);
1334            for (out, &v) in column.iter_mut().zip(gate.iter()) {
1335                if gam_data::canonical_level_bits(v) != level_bits {
1336                    *out = 0.0;
1337                }
1338            }
1339        }
1340
1341        Ok(column)
1342    }
1343}
1344
1345pub const fn default_linear_term_double_penalty() -> bool {
1346    false
1347}
1348
1349pub const fn default_pca_smooth_penalty() -> f64 {
1350    1.0
1351}
1352
1353pub const fn default_pca_chunk_size() -> usize {
1354    4096
1355}
1356
1357/// Random-effects term specification.
1358///
1359/// The selected feature column is interpreted as a categorical grouping variable.
1360/// The term contributes a one-hot dummy block with an identity penalty on group
1361/// coefficients, equivalent to i.i.d. Gaussian random effects.
1362#[derive(Debug, Clone, Serialize, Deserialize)]
1363pub struct RandomEffectTermSpec {
1364    pub name: String,
1365    pub feature_col: usize,
1366    /// If true, drop the lexicographically first group level to use treatment coding.
1367    /// If false, keep all levels (full one-hot block, still identifiable under ridge).
1368    pub drop_first_level: bool,
1369    /// If true, add a ridge penalty and estimate this block as a random effect.
1370    /// If false, leave the one-hot/treatment-coded block unpenalized so it is a
1371    /// fixed categorical main effect.  The default preserves older saved models.
1372    #[serde(default = "default_random_effect_penalized")]
1373    pub penalized: bool,
1374    /// Optional fixed kept-level set (sorted by f64 bit pattern) captured at fit time.
1375    /// When present, prediction uses exactly these columns to avoid design drift.
1376    #[serde(default)]
1377    pub frozen_levels: Option<Vec<u64>>,
1378    /// Whether an *unseen* level of this grouping column is tolerated at predict
1379    /// time (encoded as an out-of-vocabulary code and shrunk toward the
1380    /// population mean) instead of raising a schema mismatch.
1381    ///
1382    /// Only a genuine random effect — `group(g)`/`re(g)`/`s(g, bs="re")` — is
1383    /// lenient: the held-out-group policy is a deliberate contract. A FIXED
1384    /// categorical factor — a bare `+ g` OR an explicit `factor(g)` — although
1385    /// materialized as a penalized one-hot block, must raise on an
1386    /// out-of-vocabulary level at predict rather than being silently mapped to
1387    /// the factor's centering point (#2102/#2137). `factor(g)` originally shared
1388    /// the `group()`/`re()` parse arm and so wrongly inherited the lenient policy
1389    /// (#2137). For a string factor the typed schema encode rejects the unseen
1390    /// level upstream; for a numeric-coded `factor(year)` the reject is enforced
1391    /// by `build_random_effect_block`, which owns the frozen vocabulary. The
1392    /// `true` default preserves the pre-#2102 (uniformly lenient) behavior for
1393    /// models serialized before this field existed.
1394    #[serde(default = "default_random_effect_lenient_unseen")]
1395    pub lenient_unseen: bool,
1396}
1397
1398pub fn default_random_effect_penalized() -> bool {
1399    true
1400}
1401
1402pub fn default_random_effect_lenient_unseen() -> bool {
1403    true
1404}
1405
1406pub fn validate_measure_jet_positive_vec_len(
1407    label: &str,
1408    term_name: &str,
1409    field: &str,
1410    values: &[f64],
1411    expected: usize,
1412) -> Result<(), String> {
1413    if values.len() != expected {
1414        return Err(SmoothError::invalid_config(format!(
1415            "{label} term '{term_name}' frozen MeasureJet {field} has length {}, expected {expected}",
1416            values.len()
1417        ))
1418        .into());
1419    }
1420    if values
1421        .iter()
1422        .any(|value| !(value.is_finite() && *value > 0.0))
1423    {
1424        return Err(SmoothError::invalid_config(format!(
1425            "{label} term '{term_name}' frozen MeasureJet {field} values must be positive and finite"
1426        ))
1427        .into());
1428    }
1429    Ok(())
1430}
1431
1432#[derive(Debug, Clone, Serialize, Deserialize)]
1433pub struct TermCollectionSpec {
1434    pub linear_terms: Vec<LinearTermSpec>,
1435    pub random_effect_terms: Vec<RandomEffectTermSpec>,
1436    pub smooth_terms: Vec<SmoothTermSpec>,
1437}
1438
1439pub fn validate_smooth_basis_frozen(
1440    basis: &SmoothBasisSpec,
1441    label: &str,
1442    term_name: &str,
1443) -> Result<(), String> {
1444    if let Err(error) = basis.validate_scale_configuration() {
1445        return Err(SmoothError::invalid_config(format!(
1446            "{label} term '{term_name}' has an invalid scale contract: {error}"
1447        ))
1448        .into());
1449    }
1450    match basis {
1451        SmoothBasisSpec::ByVariable { inner, .. }
1452        | SmoothBasisSpec::FactorSumToZero { inner, .. } => {
1453            validate_smooth_basis_frozen(inner, label, term_name)
1454        }
1455        SmoothBasisSpec::BSpline1D { spec, .. } => {
1456            if !matches!(
1457                spec.knotspec,
1458                BSplineKnotSpec::Provided(_)
1459                    | BSplineKnotSpec::PeriodicUniform { .. }
1460                    | BSplineKnotSpec::NaturalCubicRegression { .. }
1461            ) {
1462                return Err(format!(
1463                    "{label} term '{term_name}' is not frozen: BSpline knotspec must be Provided, PeriodicUniform, or NaturalCubicRegression"
1464                ));
1465            }
1466            Ok(())
1467        }
1468        SmoothBasisSpec::ThinPlate { spec, .. } => {
1469            if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
1470                return Err(format!(
1471                    "{label} term '{term_name}' is not frozen: ThinPlate centers must be UserProvided"
1472                ));
1473            }
1474            if matches!(
1475                spec.identifiability,
1476                SpatialIdentifiability::OrthogonalToParametric
1477            ) {
1478                return Err(format!(
1479                    "{label} term '{term_name}' is not frozen: ThinPlate identifiability must be FrozenTransform or None"
1480                ));
1481            }
1482            Ok(())
1483        }
1484        _ => Ok(()),
1485    }
1486}
1487
1488impl TermCollectionSpec {
1489    /// Write this collection's topology identity into a warm-start cache
1490    /// fingerprint (#869).
1491    ///
1492    /// The persistent warm-start `cache_key` hashes only family + raw input
1493    /// dimensions, so two fits on the same data that differ *only* in their
1494    /// smooth topology (the `s(..., type=AUTO)` candidate enumeration: sphere
1495    /// vs torus vs euclidean vs duchon) collide on one key and seed each other
1496    /// with geometrically incompatible β/ρ. Folding the per-term structural
1497    /// kind + feature columns + linear/random-effect counts into the shape hash
1498    /// gives each candidate its own key, so the screen→full-refit reuse of one
1499    /// candidate is preserved while cross-candidate contamination is removed.
1500    /// Only the structural identity is hashed (not fitted coefficients or
1501    /// frozen knot values), so a refit of the *same* topology still hits.
1502    pub fn write_structural_shape_hash(&self, h: &mut gam_runtime::warm_start::Fingerprinter) {
1503        h.write_str("term-collection");
1504        h.write_usize(self.linear_terms.len());
1505        for linear in &self.linear_terms {
1506            h.write_str(&linear.name);
1507        }
1508        h.write_usize(self.random_effect_terms.len());
1509        h.write_usize(self.smooth_terms.len());
1510        for smooth in &self.smooth_terms {
1511            h.write_str(&smooth.name);
1512            h.write_str(smooth.basis.structural_kind());
1513            for col in smooth.basis.structural_feature_cols() {
1514                h.write_usize(col);
1515            }
1516        }
1517    }
1518
1519    /// Validate that a term collection spec represents a fully frozen model
1520    /// (i.e. all knots/centers are pre-computed, identifiability transforms are
1521    /// baked in, and random-effect levels are fixed).
1522    pub fn validate_frozen(&self, label: &str) -> Result<(), String> {
1523        for linear in &self.linear_terms {
1524            if let (Some(min), Some(max)) = (linear.coefficient_min, linear.coefficient_max)
1525                && (!min.is_finite() || !max.is_finite() || min > max)
1526            {
1527                return Err(SmoothError::invalid_config(format!(
1528                    "{label} linear term '{}' has invalid coefficient constraint [{min}, {max}]",
1529                    linear.name
1530                ))
1531                .into());
1532            }
1533            if let Some(min) = linear.coefficient_min
1534                && !min.is_finite()
1535            {
1536                return Err(SmoothError::invalid_config(format!(
1537                    "{label} linear term '{}' has non-finite coefficient minimum {min}",
1538                    linear.name
1539                ))
1540                .into());
1541            }
1542            if let Some(max) = linear.coefficient_max
1543                && !max.is_finite()
1544            {
1545                return Err(SmoothError::invalid_config(format!(
1546                    "{label} linear term '{}' has non-finite coefficient maximum {max}",
1547                    linear.name
1548                ))
1549                .into());
1550            }
1551            if let LinearCoefficientGeometry::Bounded { min, max, prior } =
1552                &linear.coefficient_geometry
1553            {
1554                if !min.is_finite() || !max.is_finite() || min >= max {
1555                    return Err(SmoothError::invalid_config(format!(
1556                        "{label} bounded term '{}' has invalid bounds [{min}, {max}]",
1557                        linear.name
1558                    ))
1559                    .into());
1560                }
1561                match prior {
1562                    BoundedCoefficientPriorSpec::None | BoundedCoefficientPriorSpec::Uniform => {}
1563                    BoundedCoefficientPriorSpec::Beta { a, b } => {
1564                        if !a.is_finite() || !b.is_finite() || *a < 1.0 || *b < 1.0 {
1565                            return Err(SmoothError::invalid_config(format!(
1566                                "{label} bounded term '{}' has invalid Beta prior ({a}, {b})",
1567                                linear.name
1568                            ))
1569                            .into());
1570                        }
1571                    }
1572                }
1573            }
1574        }
1575        for st in &self.smooth_terms {
1576            if let Err(error) = st.basis.validate_scale_configuration() {
1577                return Err(SmoothError::invalid_config(format!(
1578                    "{label} term '{}' has an invalid scale contract: {error}",
1579                    st.name
1580                ))
1581                .into());
1582            }
1583            match &st.basis {
1584                SmoothBasisSpec::ByVariable { inner, .. } => {
1585                    validate_smooth_basis_frozen(inner, label, &st.name)?;
1586                    let nested = SmoothTermSpec {
1587            frozen_parametric_residualization: None,
1588                        name: st.name.clone(),
1589                        basis: (**inner).clone(),
1590                        shape: st.shape,
1591                        joint_null_rotation: None,
1592                    };
1593                    TermCollectionSpec {
1594                        linear_terms: Vec::new(),
1595                        random_effect_terms: Vec::new(),
1596                        smooth_terms: vec![nested],
1597                    }
1598                    .validate_frozen(label)?;
1599                }
1600                SmoothBasisSpec::FactorSumToZero { inner, levels, .. } => {
1601                    if levels.len() < 2 {
1602                        return Err(format!(
1603                            "{label} term '{}' has invalid frozen sz levels",
1604                            st.name
1605                        ));
1606                    }
1607                    validate_smooth_basis_frozen(inner, label, &st.name)?;
1608                }
1609                SmoothBasisSpec::BSpline1D { spec, .. } => {
1610                    if !matches!(
1611                        spec.knotspec,
1612                        BSplineKnotSpec::Provided(_)
1613                            | BSplineKnotSpec::PeriodicUniform { .. }
1614                            | BSplineKnotSpec::NaturalCubicRegression { .. }
1615                    ) {
1616                        return Err(SmoothError::invalid_config(format!(
1617                            "{label} term '{}' is not frozen: BSpline knotspec must be Provided, PeriodicUniform, or NaturalCubicRegression",
1618                            st.name
1619                        ))
1620                        .into());
1621                    }
1622                }
1623                SmoothBasisSpec::ThinPlate {
1624                    spec, input_scale, ..
1625                } => {
1626                    if input_scale.is_none() {
1627                        return Err(SmoothError::invalid_config(format!(
1628                            "{label} term '{}' is not frozen: ThinPlate input_scale is missing",
1629                            st.name
1630                        ))
1631                        .into());
1632                    }
1633                    if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
1634                        return Err(SmoothError::invalid_config(format!(
1635                            "{label} term '{}' is not frozen: ThinPlate centers must be UserProvided",
1636                            st.name
1637                        ))
1638                        .into());
1639                    }
1640                    if matches!(
1641                        spec.identifiability,
1642                        SpatialIdentifiability::OrthogonalToParametric
1643                    ) {
1644                        return Err(SmoothError::invalid_config(format!(
1645                            "{label} term '{}' is not frozen: ThinPlate identifiability must be FrozenTransform or None",
1646                            st.name
1647                        ))
1648                        .into());
1649                    }
1650                }
1651                SmoothBasisSpec::Sphere { spec, .. } => {
1652                    if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
1653                        return Err(SmoothError::invalid_config(format!(
1654                            "{label} term '{}' is not frozen: Sphere centers must be UserProvided",
1655                            st.name
1656                        ))
1657                        .into());
1658                    }
1659                    if matches!(spec.method, crate::basis::SphereMethod::Harmonic)
1660                        && spec.max_degree.is_none_or(|d| d == 0)
1661                    {
1662                        return Err(format!(
1663                            "{label} term '{}' is not frozen: sphere max_degree must be positive",
1664                            st.name
1665                        ));
1666                    }
1667                }
1668                SmoothBasisSpec::ConstantCurvature { spec, .. } => {
1669                    if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
1670                        return Err(SmoothError::invalid_config(format!(
1671                            "{label} term '{}' is not frozen: ConstantCurvature centers must be UserProvided",
1672                            st.name
1673                        ))
1674                        .into());
1675                    }
1676                    if !(spec.length_scale.is_finite() && spec.length_scale > 0.0) {
1677                        return Err(SmoothError::invalid_config(format!(
1678                            "{label} term '{}' is not frozen: ConstantCurvature length_scale must be the realized positive value",
1679                            st.name
1680                        ))
1681                        .into());
1682                    }
1683                }
1684                SmoothBasisSpec::MeasureJet {
1685                    spec, input_scale, ..
1686                } => {
1687                    if input_scale.is_none() {
1688                        return Err(SmoothError::invalid_config(format!(
1689                            "{label} term '{}' is not frozen: MeasureJet input_scale is missing",
1690                            st.name
1691                        ))
1692                        .into());
1693                    }
1694                    let centers = match &spec.center_strategy {
1695                        CenterStrategy::UserProvided(centers) => centers,
1696                        _ => {
1697                            return Err(SmoothError::invalid_config(format!(
1698                                "{label} term '{}' is not frozen: MeasureJet centers must be UserProvided",
1699                                st.name
1700                            ))
1701                            .into());
1702                        }
1703                    };
1704                    if centers.nrows() == 0 {
1705                        return Err(SmoothError::invalid_config(format!(
1706                            "{label} term '{}' is not frozen: MeasureJet centers are empty",
1707                            st.name
1708                        ))
1709                        .into());
1710                    }
1711                    if !(spec.length_scale.is_finite() && spec.length_scale > 0.0) {
1712                        return Err(SmoothError::invalid_config(format!(
1713                            "{label} term '{}' is not frozen: MeasureJet length_scale must be the realized positive value",
1714                            st.name
1715                        ))
1716                        .into());
1717                    }
1718                    // Exact replay needs the fit-data penalty quadrature and
1719                    // normalization payload (`BasisMetadata::MeasureJet`).
1720                    let frozen = spec.frozen_quadrature.as_ref().ok_or_else(|| {
1721                        SmoothError::invalid_config(format!(
1722                            "{label} term '{}' is not frozen: MeasureJet frozen_quadrature payload is missing",
1723                            st.name
1724                        ))
1725                    })?;
1726                    if frozen.masses.len() != centers.nrows() {
1727                        return Err(SmoothError::invalid_config(format!(
1728                            "{label} term '{}' frozen MeasureJet has {} masses for {} centers",
1729                            st.name,
1730                            frozen.masses.len(),
1731                            centers.nrows()
1732                        ))
1733                        .into());
1734                    }
1735                    let total_mass = frozen.masses.sum();
1736                    if frozen
1737                        .masses
1738                        .iter()
1739                        .any(|mass| !(mass.is_finite() && *mass >= 0.0))
1740                        || !(total_mass.is_finite() && total_mass > 0.0)
1741                    {
1742                        return Err(SmoothError::invalid_config(format!(
1743                            "{label} term '{}' frozen MeasureJet masses must be finite, nonnegative, and have positive total mass",
1744                            st.name
1745                        ))
1746                        .into());
1747                    }
1748                    let n_levels = frozen.eps_band.len();
1749                    if n_levels == 0
1750                        || frozen
1751                            .eps_band
1752                            .iter()
1753                            .any(|eps| !(eps.is_finite() && *eps > 0.0))
1754                    {
1755                        return Err(SmoothError::invalid_config(format!(
1756                            "{label} term '{}' frozen MeasureJet eps_band must be nonempty, finite, and positive",
1757                            st.name
1758                        ))
1759                        .into());
1760                    }
1761                    for (idx, pair) in frozen.eps_band.windows(2).enumerate() {
1762                        if pair[1] <= pair[0] {
1763                            return Err(SmoothError::invalid_config(format!(
1764                                "{label} term '{}' frozen MeasureJet eps_band is not strictly ascending at {idx}: {} then {}",
1765                                st.name,
1766                                pair[0],
1767                                pair[1]
1768                            ))
1769                            .into());
1770                        }
1771                    }
1772                    validate_measure_jet_positive_vec_len(
1773                        label,
1774                        &st.name,
1775                        "support_means",
1776                        &frozen.support_means,
1777                        n_levels,
1778                    )?;
1779                    // Mode predicate MUST match the builder's
1780                    // (`measure_jet_multiscale_mode`): per-level/multiscale is the
1781                    // explicit `spec.multiscale` opt-in (#1116). In single-scale
1782                    // mode the builder emits a single FUSED penalty (empty
1783                    // per-level scales + `fused_penalty_normalization_scale:
1784                    // Some`); only the multiscale opt-in carries `n_levels`
1785                    // per-level scales.
1786                    let per_level = crate::basis::measure_jet_multiscale_mode(spec);
1787                    if per_level {
1788                        validate_measure_jet_positive_vec_len(
1789                            label,
1790                            &st.name,
1791                            "penalty_normalization_scales",
1792                            &frozen.penalty_normalization_scales,
1793                            n_levels,
1794                        )?;
1795                        validate_measure_jet_positive_vec_len(
1796                            label,
1797                            &st.name,
1798                            "raw_penalty_normalization_scales",
1799                            &frozen.raw_penalty_normalization_scales,
1800                            n_levels,
1801                        )?;
1802                        if frozen.fused_penalty_normalization_scale.is_some() {
1803                            return Err(SmoothError::invalid_config(format!(
1804                                "{label} term '{}' per-level MeasureJet must not carry a fused penalty normalization scale",
1805                                st.name
1806                            ))
1807                            .into());
1808                        }
1809                    } else {
1810                        if !frozen.penalty_normalization_scales.is_empty()
1811                            || !frozen.raw_penalty_normalization_scales.is_empty()
1812                        {
1813                            return Err(SmoothError::invalid_config(format!(
1814                                "{label} term '{}' fused MeasureJet must not carry per-level penalty normalization scales",
1815                                st.name
1816                            ))
1817                            .into());
1818                        }
1819                        match frozen.fused_penalty_normalization_scale {
1820                            Some(scale) if scale.is_finite() && scale > 0.0 => {}
1821                            Some(scale) => {
1822                                return Err(SmoothError::invalid_config(format!(
1823                                    "{label} term '{}' fused MeasureJet penalty normalization scale must be positive and finite, got {scale}",
1824                                    st.name
1825                                ))
1826                                .into());
1827                            }
1828                            None => {
1829                                return Err(SmoothError::invalid_config(format!(
1830                                    "{label} term '{}' fused MeasureJet is missing its penalty normalization scale",
1831                                    st.name
1832                                ))
1833                                .into());
1834                            }
1835                        }
1836                    }
1837                }
1838                SmoothBasisSpec::Matern {
1839                    spec, input_scale, ..
1840                } => {
1841                    if input_scale.is_none() {
1842                        return Err(SmoothError::invalid_config(format!(
1843                            "{label} term '{}' is not frozen: Matern input_scale is missing",
1844                            st.name
1845                        ))
1846                        .into());
1847                    }
1848                    if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
1849                        return Err(SmoothError::invalid_config(format!(
1850                            "{label} term '{}' is not frozen: Matern centers must be UserProvided",
1851                            st.name
1852                        ))
1853                        .into());
1854                    }
1855                    if spec
1856                        .length_scale
1857                        .resolved()
1858                        .is_none_or(|value| !value.is_finite() || value <= 0.0)
1859                    {
1860                        return Err(SmoothError::invalid_config(format!(
1861                            "{label} term '{}' is not frozen: Matern length_scale must be resolved, finite, and positive",
1862                            st.name
1863                        ))
1864                        .into());
1865                    }
1866                }
1867                SmoothBasisSpec::Duchon {
1868                    spec, input_scale, ..
1869                } => {
1870                    if input_scale.is_none() {
1871                        return Err(SmoothError::invalid_config(format!(
1872                            "{label} term '{}' is not frozen: Duchon input_scale is missing",
1873                            st.name
1874                        ))
1875                        .into());
1876                    }
1877                    if !crate::basis::duchon_center_strategy_is_frozen(&spec.center_strategy) {
1878                        return Err(SmoothError::invalid_config(format!(
1879                            "{label} term '{}' is not frozen: Duchon knots and spectral basis must be resolved",
1880                            st.name
1881                        ))
1882                        .into());
1883                    }
1884                    if matches!(
1885                        spec.identifiability,
1886                        SpatialIdentifiability::OrthogonalToParametric
1887                    ) {
1888                        return Err(SmoothError::invalid_config(format!(
1889                            "{label} term '{}' is not frozen: Duchon identifiability must be FrozenTransform or None",
1890                            st.name
1891                        ))
1892                        .into());
1893                    }
1894                }
1895                SmoothBasisSpec::Pca {
1896                    centered,
1897                    center_mean,
1898                    pca_basis_path,
1899                    ..
1900                } => {
1901                    if *centered && center_mean.is_none() && pca_basis_path.is_none() {
1902                        return Err(SmoothError::invalid_config(format!(
1903                            "{label} term '{}' is not frozen: centered Pca missing center_mean",
1904                            st.name
1905                        ))
1906                        .into());
1907                    }
1908                }
1909                SmoothBasisSpec::BySmooth { smooth, by_kind } => {
1910                    if let SmoothBasisSpec::BySmooth { .. } = smooth.as_ref() {
1911                        return Err(format!("{label} term '{}' has nested by-smooths", st.name));
1912                    }
1913                    match by_kind {
1914                        ByVarKind::Numeric { .. } => {}
1915                        ByVarKind::Factor { frozen_levels, .. } if frozen_levels.is_none() => {
1916                            return Err(format!(
1917                                "{label} term '{}' is not frozen: by-factor levels missing",
1918                                st.name
1919                            ));
1920                        }
1921                        ByVarKind::Factor { .. } => {}
1922                    }
1923                    let nested = TermCollectionSpec {
1924                        linear_terms: vec![],
1925                        random_effect_terms: vec![],
1926                        smooth_terms: vec![SmoothTermSpec {
1927            frozen_parametric_residualization: None,
1928                            name: st.name.clone(),
1929                            basis: (**smooth).clone(),
1930                            shape: st.shape,
1931                            joint_null_rotation: None,
1932                        }],
1933                    };
1934                    nested.validate_frozen(label)?;
1935                }
1936                SmoothBasisSpec::FactorSmooth { spec } => {
1937                    if spec.group_frozen_levels.is_none() {
1938                        return Err(format!(
1939                            "{label} term '{}' is not frozen: factor-smooth levels missing",
1940                            st.name
1941                        ));
1942                    }
1943                    if !matches!(
1944                        spec.marginal.knotspec,
1945                        BSplineKnotSpec::Provided(_)
1946                            | BSplineKnotSpec::PeriodicUniform { .. }
1947                            // mgcv's `bs="sz"` default marginal is a cubic
1948                            // regression spline (#1074), and the freeze step
1949                            // restores it as a `NaturalCubicRegression` knotspec
1950                            // carrying its `k` value-knots (spatial_optimization.rs
1951                            // `marginal_is_cr` branch) — the SAME treatment the
1952                            // tensor margin already gets in the arm below. Without
1953                            // this variant a frozen `sz` factor smooth fails its own
1954                            // predict-time freeze check ("factor-smooth marginal
1955                            // knots missing") even though its knots are fully
1956                            // materialized; the validation simply was not updated
1957                            // when the cr marginal landed.
1958                            | BSplineKnotSpec::NaturalCubicRegression { .. }
1959                    ) {
1960                        return Err(format!(
1961                            "{label} term '{}' is not frozen: factor-smooth marginal knots missing",
1962                            st.name
1963                        ));
1964                    }
1965                }
1966                SmoothBasisSpec::TensorBSpline { spec, .. } => {
1967                    for (dim, marginal) in spec.marginalspecs.iter().enumerate() {
1968                        if !matches!(
1969                            marginal.knotspec,
1970                            BSplineKnotSpec::Provided(_)
1971                                | BSplineKnotSpec::PeriodicUniform { .. }
1972                                | BSplineKnotSpec::NaturalCubicRegression { .. }
1973                        ) {
1974                            return Err(SmoothError::invalid_config(format!(
1975                                "{label} term '{}' dim {} is not frozen: tensor marginal knotspec must be Provided, PeriodicUniform, or NaturalCubicRegression",
1976                                st.name, dim
1977                            ))
1978                            .into());
1979                        }
1980                    }
1981                    if matches!(
1982                        spec.identifiability,
1983                        TensorBSplineIdentifiability::SumToZero
1984                            | TensorBSplineIdentifiability::MarginalSumToZero
1985                    ) {
1986                        return Err(SmoothError::invalid_config(format!(
1987                            "{label} term '{}' is not frozen: tensor identifiability must be FrozenTransform or None",
1988                            st.name
1989                        ))
1990                        .into());
1991                    }
1992                }
1993            }
1994        }
1995
1996        for rt in &self.random_effect_terms {
1997            if rt.frozen_levels.is_none() {
1998                return Err(SmoothError::invalid_config(format!(
1999                    "{label} random-effect term '{}' is not frozen: missing frozen_levels",
2000                    rt.name
2001                ))
2002                .into());
2003            }
2004        }
2005
2006        Ok(())
2007    }
2008
2009    /// Re-resolve every stored feature-column index through `remap`, returning a
2010    /// spec that addresses a different column layout.
2011    ///
2012    /// A frozen `TermCollectionSpec` stores feature columns as *absolute indices
2013    /// into the training table*. To replay it on a fresh dataset whose columns
2014    /// sit at different positions — the common case at prediction time, where
2015    /// the response column is unknown and may be absent entirely — every index
2016    /// must be re-resolved against the new layout. `remap` receives each
2017    /// training-table index and returns its position in the runtime table;
2018    /// callers typically implement it as "look the name up in the training
2019    /// headers, then resolve that name against the prediction dataset".
2020    ///
2021    /// This is the single authority on *which* fields carry a column index
2022    /// across every basis variant (linear, random-effect, the `by=` column of
2023    /// `ByVariable`/`FactorSumToZero`/`BySmooth`, the continuous and group
2024    /// columns of a `FactorSmooth`, and the multi-axis `feature_cols` of every
2025    /// spatial/tensor basis), so a predict-time realignment cannot silently miss
2026    /// one and dereference a stale training index.
2027    pub fn remap_feature_columns<E, F>(&self, mut remap: F) -> Result<TermCollectionSpec, E>
2028    where
2029        F: FnMut(usize) -> Result<usize, E>,
2030    {
2031        let mut out = self.clone();
2032        for lt in &mut out.linear_terms {
2033            lt.feature_col = remap(lt.feature_col)?;
2034            // Also remap the full interaction-factor list. The design builder
2035            // (`build_term_collection_design_inner`) materializes the column from
2036            // `effective_feature_cols()` — which returns `feature_cols` whenever
2037            // it is non-empty (i.e. essentially always, including a plain linear
2038            // term where `feature_cols == [feature_col]`). Remapping only the
2039            // singular `feature_col` left these at their saved *training* indices
2040            // at predict time, so a parametric `Surv(...) ~ x` (and any `:`
2041            // interaction) bailed with "feature column N out of bounds" once the
2042            // response/time columns shift the runtime layout (issue #898).
2043            for fc in lt.feature_cols.iter_mut() {
2044                *fc = remap(*fc)?;
2045            }
2046            // A factor-aware `:` interaction also gates on categorical columns;
2047            // those indices live in the same training-time layout and must be
2048            // realigned to the runtime table alongside the numeric operands, or
2049            // the predict-time level indicator would dereference a stale column.
2050            for (col, _bits) in lt.categorical_levels.iter_mut() {
2051                *col = remap(*col)?;
2052            }
2053        }
2054        for rt in &mut out.random_effect_terms {
2055            rt.feature_col = remap(rt.feature_col)?;
2056        }
2057        for st in &mut out.smooth_terms {
2058            remap_smooth_basis_feature_columns(&mut st.basis, &mut remap)?;
2059        }
2060        Ok(out)
2061    }
2062}
2063
2064/// Walk a `SmoothBasisSpec` tree, re-resolving every column index through
2065/// `remap`. Shared by all predict-time column realignment (see
2066/// [`TermCollectionSpec::remap_feature_columns`]); kept exhaustive so a newly
2067/// added index-bearing variant fails to compile until it is handled here.
2068pub fn remap_smooth_basis_feature_columns<E, F>(
2069    basis: &mut SmoothBasisSpec,
2070    remap: &mut F,
2071) -> Result<(), E>
2072where
2073    F: FnMut(usize) -> Result<usize, E>,
2074{
2075    match basis {
2076        SmoothBasisSpec::ByVariable { inner, by_col, .. }
2077        | SmoothBasisSpec::FactorSumToZero { inner, by_col, .. } => {
2078            *by_col = remap(*by_col)?;
2079            remap_smooth_basis_feature_columns(inner, remap)?;
2080        }
2081        SmoothBasisSpec::BSpline1D { feature_col, .. } => {
2082            *feature_col = remap(*feature_col)?;
2083        }
2084        SmoothBasisSpec::BySmooth { smooth, by_kind } => {
2085            let by_feature_col = match by_kind {
2086                ByVarKind::Numeric { feature_col } | ByVarKind::Factor { feature_col, .. } => {
2087                    feature_col
2088                }
2089            };
2090            *by_feature_col = remap(*by_feature_col)?;
2091            remap_smooth_basis_feature_columns(smooth, remap)?;
2092        }
2093        SmoothBasisSpec::FactorSmooth { spec } => {
2094            for fc in spec.continuous_cols.iter_mut() {
2095                *fc = remap(*fc)?;
2096            }
2097            spec.group_col = remap(spec.group_col)?;
2098        }
2099        SmoothBasisSpec::ThinPlate { feature_cols, .. }
2100        | SmoothBasisSpec::Sphere { feature_cols, .. }
2101        | SmoothBasisSpec::ConstantCurvature { feature_cols, .. }
2102        | SmoothBasisSpec::Matern { feature_cols, .. }
2103        | SmoothBasisSpec::MeasureJet { feature_cols, .. }
2104        | SmoothBasisSpec::Duchon { feature_cols, .. }
2105        | SmoothBasisSpec::Pca { feature_cols, .. }
2106        | SmoothBasisSpec::TensorBSpline { feature_cols, .. } => {
2107            for fc in feature_cols.iter_mut() {
2108                *fc = remap(*fc)?;
2109            }
2110        }
2111    }
2112    Ok(())
2113}
2114
2115#[derive(Debug, Clone)]
2116pub enum PenaltyStructureHint {
2117    Ridge(f64),
2118    Kronecker(Vec<Array2<f64>>),
2119}
2120
2121/// A penalty matrix stored at its natural block size together with the
2122/// column range it occupies in the global coefficient vector.
2123///
2124/// Instead of embedding every penalty into a full `p_total × p_total` dense
2125/// matrix filled with zeros, we keep the compact local matrix and reconstruct
2126/// the global view only when a downstream consumer explicitly requires it.
2127#[derive(Clone)]
2128pub struct BlockwisePenalty {
2129    /// Column range in the global coefficient vector that this penalty covers.
2130    pub col_range: Range<usize>,
2131    /// The local penalty matrix — dimensions `block_p × block_p` where
2132    /// `block_p = col_range.len()`.
2133    pub local: Array2<f64>,
2134    /// Optional nonzero centering vector for this coefficient block.
2135    pub prior_mean: gam_problem::CoefficientPriorMean,
2136    /// Optional structural hint so downstream spectral/logdet code can stay
2137    /// block-local or factorized without reverse-engineering the matrix.
2138    pub structure_hint: Option<PenaltyStructureHint>,
2139    /// Optional operator-form handle bit-equivalent to `local`. Populated when
2140    /// the originating closed-form factory emitted an op-form penalty so exact
2141    /// operator algebra can use matvec instead of materializing the dense
2142    /// `block_p × block_p` Gram. `None` for ordinary dense penalties.
2143    pub op: Option<std::sync::Arc<dyn crate::analytic_penalties::PenaltyOp>>,
2144}
2145
2146impl std::fmt::Debug for BlockwisePenalty {
2147    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2148        f.debug_struct("BlockwisePenalty")
2149            .field("col_range", &self.col_range)
2150            .field(
2151                "local",
2152                &format_args!("{}×{}", self.local.nrows(), self.local.ncols()),
2153            )
2154            .field("prior_mean", &self.prior_mean)
2155            .field("structure_hint", &self.structure_hint)
2156            .field("op", &self.op.as_ref().map(|o| o.dim()))
2157            .finish()
2158    }
2159}
2160
2161impl BlockwisePenalty {
2162    /// Create a new blockwise penalty.
2163    pub fn new(col_range: Range<usize>, local: Array2<f64>) -> Self {
2164        assert_eq!(col_range.len(), local.nrows());
2165        assert_eq!(col_range.len(), local.ncols());
2166        Self {
2167            col_range,
2168            local,
2169            prior_mean: gam_problem::CoefficientPriorMean::Zero,
2170            structure_hint: None,
2171            op: None,
2172        }
2173    }
2174
2175    pub fn with_prior_mean(mut self, prior_mean: gam_problem::CoefficientPriorMean) -> Self {
2176        self.prior_mean = prior_mean;
2177        self
2178    }
2179
2180    /// Attach an op-form penalty handle bit-equivalent to `local`.
2181    pub fn with_op(
2182        mut self,
2183        op: Option<std::sync::Arc<dyn crate::analytic_penalties::PenaltyOp>>,
2184    ) -> Self {
2185        self.op = op;
2186        self
2187    }
2188
2189    pub fn ridge(col_range: Range<usize>, scale: f64) -> Self {
2190        let block_size = col_range.len();
2191        let mut local = Array2::<f64>::zeros((block_size, block_size));
2192        for i in 0..block_size {
2193            local[[i, i]] = scale;
2194        }
2195        Self {
2196            col_range,
2197            local,
2198            prior_mean: gam_problem::CoefficientPriorMean::Zero,
2199            structure_hint: Some(PenaltyStructureHint::Ridge(scale)),
2200            op: None,
2201        }
2202    }
2203
2204    pub fn kronecker(
2205        col_range: Range<usize>,
2206        local: Array2<f64>,
2207        factors: Vec<Array2<f64>>,
2208    ) -> Self {
2209        assert_eq!(col_range.len(), local.nrows());
2210        assert_eq!(col_range.len(), local.ncols());
2211        Self {
2212            col_range,
2213            local,
2214            prior_mean: gam_problem::CoefficientPriorMean::Zero,
2215            structure_hint: Some(PenaltyStructureHint::Kronecker(factors)),
2216            op: None,
2217        }
2218    }
2219
2220    /// Expand this blockwise penalty into a full `p_total × p_total` dense
2221    /// matrix (mostly zeros). Use sparingly — the whole point of blockwise
2222    /// storage is to avoid this allocation.
2223    pub fn to_global(&self, p_total: usize) -> Array2<f64> {
2224        let mut g = Array2::<f64>::zeros((p_total, p_total));
2225        let r = &self.col_range;
2226        assert!(
2227            r.end <= p_total && self.local.nrows() == r.len() && self.local.ncols() == r.len(),
2228            "BlockwisePenalty::to_global shape invariant violated: \
2229             col_range={}..{}, local={}x{}, p_total={}",
2230            r.start,
2231            r.end,
2232            self.local.nrows(),
2233            self.local.ncols(),
2234            p_total,
2235        );
2236        g.slice_mut(s![r.start..r.end, r.start..r.end])
2237            .assign(&self.local);
2238        g
2239    }
2240
2241    /// Convert into a blockwise [`gam_problem::PenaltyMatrix`] without
2242    /// expanding to full dimensions.
2243    pub fn to_penalty_matrix(&self, total_dim: usize) -> gam_problem::PenaltyMatrix {
2244        gam_problem::PenaltyMatrix::Blockwise {
2245            local: self.local.clone(),
2246            col_range: self.col_range.clone(),
2247            total_dim,
2248        }
2249    }
2250
2251    /// The block size of this penalty.
2252    #[inline]
2253    pub fn block_size(&self) -> usize {
2254        self.col_range.len()
2255    }
2256}
2257
2258/// Compute `Σ_k λ_k S_k` directly from blockwise penalties, accumulating
2259/// into a pre-allocated `p_total × p_total` output without ever materializing
2260/// individual global matrices.
2261pub fn weighted_blockwise_penalty_sum(
2262    penalties: &[BlockwisePenalty],
2263    lambdas: &[f64],
2264    p_total: usize,
2265) -> Array2<f64> {
2266    assert_eq!(penalties.len(), lambdas.len());
2267    // Smoothing parameters λ_k must be non-negative and finite. A negative
2268    // λ would flip the sign of the corresponding block S_k, turning the
2269    // total penalty matrix indefinite and silently corrupting every
2270    // downstream Cholesky / PIRLS / REML / pseudo-logdet computation that
2271    // assumes S_λ ⪰ 0. Catch this at the boundary rather than after it
2272    // has propagated.
2273    for (idx, &lam) in lambdas.iter().enumerate() {
2274        assert!(
2275            lam.is_finite() && lam >= 0.0,
2276            "weighted_blockwise_penalty_sum: lambdas[{idx}] = {lam} is invalid (must be finite and non-negative; negative smoothing parameters violate S_λ ⪰ 0)",
2277        );
2278    }
2279    // Block column ranges must also fit inside the declared total parameter
2280    // dimension; an out-of-bounds slice would otherwise panic from ndarray
2281    // with a far less informative message.
2282    for (idx, bp) in penalties.iter().enumerate() {
2283        let r = &bp.col_range;
2284        assert!(
2285            r.end <= p_total,
2286            "weighted_blockwise_penalty_sum: penalties[{idx}] col_range {:?} exceeds p_total = {p_total}",
2287            r,
2288        );
2289    }
2290    let mut out = Array2::<f64>::zeros((p_total, p_total));
2291    for (bp, &lam) in penalties.iter().zip(lambdas.iter()) {
2292        let r = &bp.col_range;
2293        let mut slice = out.slice_mut(s![r.start..r.end, r.start..r.end]);
2294        slice.scaled_add(lam, &bp.local);
2295    }
2296    out
2297}
2298
2299// ---------------------------------------------------------------------------
2300// KroneckerPenaltySystem — factored tensor-product penalty representation
2301// ---------------------------------------------------------------------------
2302
2303/// Factored representation of tensor-product penalties with precomputed
2304/// marginal eigensystems for O(∏q_j) logdet and penalty operations.
2305#[derive(Debug, Clone)]
2306pub struct KroneckerPenaltySystem {
2307    /// Marginal penalty matrices: `marginal_penalties[k]` is `(q_k, q_k)`.
2308    pub marginal_penalties: Vec<Array2<f64>>,
2309    /// Precomputed eigensystems: `(eigenvalues, eigenvectors)` per marginal.
2310    pub marginal_eigensystems: Vec<(Array1<f64>, Array2<f64>)>,
2311    /// Marginal basis dimensions.
2312    pub marginal_dims: Vec<usize>,
2313    /// Whether a global ridge (double) penalty is present.
2314    pub has_double_penalty: bool,
2315}
2316
2317impl KroneckerPenaltySystem {
2318    pub fn new(
2319        marginal_penalties: Vec<Array2<f64>>,
2320        marginal_dims: Vec<usize>,
2321        has_double_penalty: bool,
2322    ) -> Result<Self, BasisError> {
2323        if marginal_penalties.len() != marginal_dims.len() {
2324            crate::bail_dim_basis!(
2325                "KroneckerPenaltySystem: {} penalties vs {} dims",
2326                marginal_penalties.len(),
2327                marginal_dims.len()
2328            );
2329        }
2330        let eigensystems =
2331            kronecker_marginal_eigensystems(&marginal_penalties, "KroneckerPenaltySystem")
2332                .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
2333        Ok(Self {
2334            marginal_penalties,
2335            marginal_eigensystems: eigensystems,
2336            marginal_dims,
2337            has_double_penalty,
2338        })
2339    }
2340
2341    pub fn p_total(&self) -> usize {
2342        self.marginal_dims.iter().copied().product()
2343    }
2344
2345    pub fn ndim(&self) -> usize {
2346        self.marginal_dims.len()
2347    }
2348
2349    pub fn num_penalties(&self) -> usize {
2350        self.marginal_dims.len() + if self.has_double_penalty { 1 } else { 0 }
2351    }
2352
2353    /// Compute `log|S|₊` and its first/second derivatives w.r.t. `ρ_k = log(λ_k)`.
2354    ///
2355    /// Iterates over the ∏q_j multi-index grid. Cost: O(d · ∏q_j), no O(p²) storage.
2356    pub fn logdet_and_derivatives(
2357        &self,
2358        lambdas: &[f64],
2359        objective_ridge: f64,
2360    ) -> (f64, Array1<f64>, Array2<f64>) {
2361        let n_pen = self.num_penalties();
2362        assert_eq!(lambdas.len(), n_pen, "lambda count mismatch");
2363        let marginal_evals: Vec<_> = self
2364            .marginal_eigensystems
2365            .iter()
2366            .map(|(evals, _)| evals.view())
2367            .collect();
2368        kronecker_logdet_and_derivatives(
2369            &marginal_evals,
2370            &self.marginal_dims,
2371            lambdas,
2372            self.has_double_penalty,
2373            objective_ridge,
2374        )
2375    }
2376
2377    pub fn logdet_rank_and_derivatives(
2378        &self,
2379        lambdas: &[f64],
2380        objective_ridge: f64,
2381    ) -> (f64, usize, Array1<f64>, Array2<f64>) {
2382        let n_pen = self.num_penalties();
2383        assert_eq!(lambdas.len(), n_pen, "lambda count mismatch");
2384        let d = self.marginal_dims.len();
2385        let mut logdet = 0.0;
2386        let mut rank = 0usize;
2387        let mut grad = Array1::<f64>::zeros(n_pen);
2388        let mut hess = Array2::<f64>::zeros((n_pen, n_pen));
2389        // Positivity floor for a penalized eigenvalue `σ`: below this the mode
2390        // is treated as an unpenalized (null-space) direction and excluded from
2391        // both the rank count and the pseudo-log-determinant.
2392        const EIGENVALUE_POSITIVITY_FLOOR: f64 = 1e-12;
2393        // Floor on the *structural* eigenvalue sum (λ-independent) used to
2394        // classify the joint null space.
2395        const STRUCTURAL_ZERO_FLOOR: f64 = 1e-12;
2396        let mut multi_idx = vec![0usize; d];
2397        loop {
2398            let mut sigma = 0.0;
2399            let mut structural_sigma = 0.0;
2400            for k in 0..d {
2401                let marginal_eigenvalue = self.marginal_eigensystems[k].0[multi_idx[k]];
2402                structural_sigma += marginal_eigenvalue;
2403                sigma += lambdas[k] * marginal_eigenvalue;
2404            }
2405            let joint_null = structural_sigma <= STRUCTURAL_ZERO_FLOOR;
2406            if self.has_double_penalty && joint_null {
2407                sigma += lambdas[d];
2408            }
2409            if structural_sigma > STRUCTURAL_ZERO_FLOOR {
2410                sigma += objective_ridge;
2411            }
2412
2413            if sigma > EIGENVALUE_POSITIVITY_FLOOR {
2414                rank += 1;
2415                logdet += sigma.ln();
2416                let inv_sigma = 1.0 / sigma;
2417                let inv_sigma2 = inv_sigma * inv_sigma;
2418                for k in 0..n_pen {
2419                    let ck = if k < d {
2420                        lambdas[k] * self.marginal_eigensystems[k].0[multi_idx[k]]
2421                    } else if joint_null {
2422                        lambdas[d]
2423                    } else {
2424                        0.0
2425                    };
2426                    grad[k] += ck * inv_sigma;
2427                    hess[[k, k]] += ck * inv_sigma - ck * ck * inv_sigma2;
2428                    for l in (k + 1)..n_pen {
2429                        let cl = if l < d {
2430                            lambdas[l] * self.marginal_eigensystems[l].0[multi_idx[l]]
2431                        } else if joint_null {
2432                            lambdas[d]
2433                        } else {
2434                            0.0
2435                        };
2436                        let off = -ck * cl * inv_sigma2;
2437                        hess[[k, l]] += off;
2438                        hess[[l, k]] += off;
2439                    }
2440                }
2441            }
2442
2443            let mut carry = true;
2444            for dim in (0..d).rev() {
2445                if carry {
2446                    multi_idx[dim] += 1;
2447                    if multi_idx[dim] < self.marginal_dims[dim] {
2448                        carry = false;
2449                    } else {
2450                        multi_idx[dim] = 0;
2451                    }
2452                }
2453            }
2454            if carry {
2455                break;
2456            }
2457        }
2458        (logdet, rank, grad, hess)
2459    }
2460}
2461
2462#[cfg(test)]
2463mod joint_unpenalized_dim_tests {
2464    use super::{ActivePenalty, ActivePenaltyInfo, PenaltySource, joint_unpenalized_dim};
2465    use ndarray::{Array2, array};
2466
2467    fn active_penalty(
2468        matrix: Array2<f64>,
2469        effective_rank: usize,
2470        nullity: usize,
2471        original_index: usize,
2472        source: PenaltySource,
2473    ) -> ActivePenalty {
2474        ActivePenalty {
2475            matrix,
2476            nullity,
2477            null_eigenvectors: None,
2478            op: None,
2479            info: ActivePenaltyInfo {
2480                source,
2481                original_index,
2482                effective_rank,
2483                normalization_scale: 1.0,
2484                kronecker_factors: None,
2485                structural_null_frame: None,
2486            },
2487        }
2488    }
2489
2490    #[test]
2491    fn no_penalty_is_fully_unpenalized() {
2492        assert_eq!(joint_unpenalized_dim(4, &[]), 4);
2493    }
2494
2495    #[test]
2496    fn single_penalty_returns_its_own_null_space() {
2497        // A 3×3 penalty that penalizes only the last coordinate ⇒ 2-dim null
2498        // space (the first two coordinates are unpenalized).
2499        let s = array![[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 5.0]];
2500        let penalties = [active_penalty(s, 1, 2, 0, PenaltySource::Primary)];
2501        assert_eq!(joint_unpenalized_dim(3, &penalties), 2);
2502    }
2503
2504    #[test]
2505    fn complementary_double_penalty_has_empty_joint_null_space() {
2506        // The #1360 case in miniature: a "bending" penalty that leaves the
2507        // first coordinate (its 2-dim... here 1-dim) null, plus a
2508        // complementary "null-space ridge" that penalizes exactly that
2509        // coordinate. Per-penalty null dims are {1, 2} and sum to 3 (≈ p),
2510        // but the INTERSECTION is empty: every coordinate is penalized by
2511        // someone, so the joint unpenalized dim is 0.
2512        let bending = array![[0.0, 0.0, 0.0], [0.0, 4.0, 0.0], [0.0, 0.0, 4.0]];
2513        let ridge = array![[2.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]];
2514        let penalties = [
2515            active_penalty(bending, 2, 1, 0, PenaltySource::Primary),
2516            active_penalty(ridge, 1, 2, 1, PenaltySource::DoublePenaltyNullspace),
2517        ];
2518        assert_eq!(joint_unpenalized_dim(3, &penalties), 0);
2519    }
2520
2521    #[test]
2522    fn partial_overlap_keeps_shared_null_direction() {
2523        // Two penalties that BOTH leave coordinate 0 unpenalized ⇒ the shared
2524        // null direction survives the intersection (joint unpenalized dim 1),
2525        // even though naively summing the per-penalty dims would give 4.
2526        let a = array![[0.0, 0.0, 0.0], [0.0, 3.0, 0.0], [0.0, 0.0, 0.0]];
2527        let b = array![[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 3.0]];
2528        let penalties = [
2529            active_penalty(a, 1, 2, 0, PenaltySource::Primary),
2530            active_penalty(b, 1, 2, 1, PenaltySource::OperatorStiffness),
2531        ];
2532        assert_eq!(joint_unpenalized_dim(3, &penalties), 1);
2533    }
2534
2535    #[test]
2536    fn non_materialized_penalty_falls_back_conservatively() {
2537        // A penalty whose stored block is not p_local × p_local (e.g. a
2538        // Kronecker tensor factor). With ≥2 penalties the conservative joint
2539        // dim is 0 (never over-rejecting).
2540        let full: Array2<f64> = array![[0.0, 0.0], [0.0, 1.0]];
2541        let factor: Array2<f64> = array![[1.0]]; // wrong shape for p_local=2
2542        let mixed_penalties = [
2543            active_penalty(full, 1, 1, 0, PenaltySource::Primary),
2544            active_penalty(
2545                factor.clone(),
2546                2,
2547                0,
2548                1,
2549                PenaltySource::TensorMarginal { dim: 0 },
2550            ),
2551        ];
2552        assert_eq!(joint_unpenalized_dim(2, &mixed_penalties), 0);
2553        // With a single non-materialized penalty, fall back to its own null dim.
2554        let factor_penalties = [active_penalty(
2555            factor,
2556            2,
2557            2,
2558            0,
2559            PenaltySource::TensorMarginal { dim: 0 },
2560        )];
2561        assert_eq!(joint_unpenalized_dim(4, &factor_penalties), 2);
2562    }
2563}
2564
2565#[cfg(test)]
2566mod kronecker_penalty_system_tests {
2567    use super::KroneckerPenaltySystem;
2568    use ndarray::array;
2569
2570    #[test]
2571    fn double_penalty_rank_derivatives_use_only_joint_null_space() {
2572        let penalties = vec![
2573            array![[0.0, 0.0], [0.0, 2.0]],
2574            array![[0.0, 0.0], [0.0, 3.0]],
2575        ];
2576        let system = KroneckerPenaltySystem::new(penalties, vec![2usize, 2usize], true).unwrap();
2577        let lambdas = vec![5.0, 7.0, 11.0];
2578
2579        let (logdet, rank, grad, hess) = system.logdet_rank_and_derivatives(&lambdas, 0.0);
2580
2581        let expected_diag = [11.0_f64, 21.0, 10.0, 31.0];
2582        let expected_logdet: f64 = expected_diag.iter().map(|v| v.ln()).sum();
2583        assert_eq!(rank, 4);
2584        assert!((logdet - expected_logdet).abs() <= 1e-12);
2585        assert!(
2586            (grad[2] - 1.0).abs() <= 1e-12,
2587            "double-penalty rank derivative must count only the joint null mode, got {}",
2588            grad[2]
2589        );
2590        assert!(hess[[2, 2]].abs() <= 1e-12);
2591    }
2592}
2593
2594#[derive(Clone, Debug)]
2595pub struct TermCollectionDesign {
2596    /// The full design matrix.
2597    ///
2598    /// Prefer a true sparse matrix when every block is sparse-compatible.
2599    /// If the collection already contains intrinsically sparse blocks, preserve
2600    /// that storage and let PIRLS decide later whether the penalized system is
2601    /// sparse-native eligible. Purely dense materialized blocks still fall back
2602    /// to the lazy block operator when sparse storage would just re-encode a
2603    /// dense matrix.
2604    pub design: DesignMatrix,
2605    /// Known row-wise affine contribution to the linear predictor.
2606    ///
2607    /// The realized predictor is `affine_offset + design * beta`. This channel
2608    /// is deliberately separate from `design`: folding it into an estimated
2609    /// intercept would make inhomogeneous term constraints coefficient-
2610    /// dependent and would corrupt all linear-operator/Hessian identities.
2611    pub affine_offset: Array1<f64>,
2612    pub penalties: Vec<BlockwisePenalty>,
2613    pub nullspace_dims: Vec<usize>,
2614    pub penaltyinfo: Vec<PenaltyBlockInfo>,
2615    pub dropped_penaltyinfo: Vec<DroppedPenaltyBlockInfo>,
2616    /// Optional global coefficient lower bounds for constrained fitting.
2617    /// Length equals `design.ncols()` when present. Unconstrained entries are `-inf`.
2618    pub coefficient_lower_bounds: Option<Array1<f64>>,
2619    /// Optional global inequality constraints:
2620    /// `A * beta >= b`.
2621    pub linear_constraints: Option<LinearInequalityConstraints>,
2622    pub intercept_range: Range<usize>,
2623    pub linear_ranges: Vec<(String, Range<usize>)>,
2624    /// Per-linear-term empirical function mass used for this build's ridge
2625    /// penalty (parallel to `spec.linear_terms`; `Some` only for
2626    /// `double_penalty=true` terms). This is the TRAINING-time value the
2627    /// first (fit-time) build actually computed from `data` — never a value
2628    /// recomputed from a prediction/rebuild call's (possibly tiny or
2629    /// zero-variance) rows. `freeze_term_collection_from_design` copies this
2630    /// into each term's `frozen_function_mass` so a later rebuild reuses it
2631    /// instead of recomputing.
2632    pub linear_function_masses: Vec<Option<f64>>,
2633    pub random_effect_ranges: Vec<(String, Range<usize>)>,
2634    pub random_effect_levels: Vec<(String, Vec<u64>)>,
2635    pub smooth: SmoothDesign,
2636}
2637
2638impl TermCollectionDesign {
2639    /// Add this collection's fixed affine channel to a caller-owned likelihood
2640    /// offset, validating the universal row/finite-value contract at the seam
2641    /// where the two offset sources become one.
2642    pub fn compose_offset(
2643        &self,
2644        base: ArrayView1<'_, f64>,
2645        context: &str,
2646    ) -> Result<Array1<f64>, BasisError> {
2647        let n = self.design.nrows();
2648        if self.affine_offset.len() != n || base.len() != n {
2649            crate::bail_dim_basis!(
2650                "{context}: design rows={n}, affine offset rows={}, base offset rows={}",
2651                self.affine_offset.len(),
2652                base.len()
2653            );
2654        }
2655        if self.affine_offset.iter().any(|value| !value.is_finite())
2656            || base.iter().any(|value| !value.is_finite())
2657        {
2658            crate::bail_invalid_basis!("{context}: offsets must be finite");
2659        }
2660        Ok(base.to_owned() + &self.affine_offset)
2661    }
2662
2663    /// Evaluate `affine_offset + design * beta` with a checked coefficient
2664    /// width. Prediction/reporting code should use this instead of applying the
2665    /// linear operator alone whenever it has no separate user offset channel.
2666    pub fn apply(&self, beta: ArrayView1<'_, f64>) -> Result<Array1<f64>, BasisError> {
2667        if beta.len() != self.design.ncols() {
2668            crate::bail_dim_basis!(
2669                "term-collection predictor coefficient length {} does not match design width {}",
2670                beta.len(),
2671                self.design.ncols()
2672            );
2673        }
2674        if beta.iter().any(|value| !value.is_finite()) {
2675            crate::bail_invalid_basis!("term-collection predictor coefficients must be finite");
2676        }
2677        if self.affine_offset.len() != self.design.nrows() {
2678            crate::bail_dim_basis!(
2679                "term-collection affine offset has {} rows but design has {}",
2680                self.affine_offset.len(),
2681                self.design.nrows()
2682            );
2683        }
2684        if self.affine_offset.iter().any(|value| !value.is_finite()) {
2685            crate::bail_invalid_basis!("term-collection affine offset must be finite");
2686        }
2687        Ok(self.design.apply(&beta.to_owned()) + &self.affine_offset)
2688    }
2689
2690    /// Number of global penalty blocks that precede the smooth-term penalty
2691    /// blocks in the flat smoothing-parameter / EDF-trace layout.
2692    ///
2693    /// The prefix is not the same as `random_effect_ranges.len()`: unpenalized
2694    /// (or empty) random-effect ranges contribute coefficient columns but no
2695    /// penalty block. The authoritative layout is the recorded `penaltyinfo`
2696    /// sequence assembled alongside `penalties`.
2697    pub fn leading_penalty_blocks_before_smooth(&self) -> usize {
2698        self.penaltyinfo
2699            .iter()
2700            .take_while(|info| {
2701                matches!(
2702                    &info.penalty.source,
2703                    crate::basis::PenaltySource::Other(source)
2704                        if source == "LinearTermRidge"
2705                            || source.starts_with("RandomEffectRidge(")
2706                )
2707            })
2708            .count()
2709    }
2710
2711    /// Global flat-penalty range owned by one realized smooth term.
2712    ///
2713    /// This is the only supported translation from a smooth-local penalty
2714    /// index to the model-global smoothing-parameter layout. The prefix comes
2715    /// from the metadata recorded alongside the matrices that were actually
2716    /// emitted; it must not be reconstructed from term specs or coefficient
2717    /// ranges, because unpenalized effects have columns but no penalty block.
2718    pub fn smooth_term_penalty_range(
2719        &self,
2720        term_idx: usize,
2721    ) -> Result<Option<Range<usize>>, String> {
2722        let Some(term) = self.smooth.terms.get(term_idx) else {
2723            return Ok(None);
2724        };
2725        if term.active_penalties.is_empty() {
2726            return Ok(None);
2727        }
2728
2729        let leading = self.leading_penalty_blocks_before_smooth();
2730        let smooth_count = self
2731            .smooth
2732            .terms
2733            .iter()
2734            .map(|smooth| smooth.active_penalties.len())
2735            .sum::<usize>();
2736        let expected = leading
2737            .checked_add(smooth_count)
2738            .ok_or_else(|| "term-collection penalty count overflow".to_string())?;
2739        if expected != self.penalties.len() || self.penaltyinfo.len() != self.penalties.len() {
2740            return Err(format!(
2741                "term-collection penalty layout is inconsistent: {leading} leading blocks + \
2742                 {smooth_count} smooth blocks = {expected}, but there are {} penalties and {} \
2743                 metadata records",
2744                self.penalties.len(),
2745                self.penaltyinfo.len()
2746            ));
2747        }
2748
2749        let local_offset = self
2750            .smooth
2751            .terms
2752            .iter()
2753            .take(term_idx)
2754            .map(|smooth| smooth.active_penalties.len())
2755            .sum::<usize>();
2756        let start = leading
2757            .checked_add(local_offset)
2758            .ok_or_else(|| "smooth penalty offset overflow".to_string())?;
2759        let end = start
2760            .checked_add(term.active_penalties.len())
2761            .ok_or_else(|| "smooth penalty range overflow".to_string())?;
2762        Ok(Some(start..end))
2763    }
2764
2765    /// Convert blockwise penalties to `PenaltyMatrix::Blockwise` without
2766    /// expanding to `p_total × p_total`. This is the preferred path for
2767    /// family modules that accept `Vec<PenaltyMatrix>`.
2768    pub fn penalties_as_penalty_matrix(&self) -> Vec<gam_problem::PenaltyMatrix> {
2769        let p = self.design.ncols();
2770        self.penalties
2771            .iter()
2772            .map(|bp| bp.to_penalty_matrix(p))
2773            .collect()
2774    }
2775
2776    /// Number of penalty blocks.
2777    #[inline]
2778    pub fn num_penalties(&self) -> usize {
2779        self.penalties.len()
2780    }
2781
2782    /// Resolve coefficient groups against this design's global coefficient
2783    /// layout and append their penalties after the existing term penalties.
2784    pub fn realize_coefficient_groups(
2785        &self,
2786        groups: &[CoefficientGroupSpec],
2787        base_prior: &gam_spec::RhoPrior,
2788    ) -> Result<RealizedCoefficientGroups, BasisError> {
2789        realize_coefficient_groups(self, groups, base_prior)
2790    }
2791
2792    /// Extract a `KroneckerPenaltySystem` when the model's *only* smooth term is
2793    /// a single Kronecker-factored tensor.
2794    ///
2795    /// This is a deliberate single-tensor fast path, not a partial feature: any
2796    /// other shape — zero Kronecker terms, several of them, or a tensor mixed
2797    /// with non-tensor smooth terms — is served correctly by the standard
2798    /// block-separable assembly, so this returns `None` and the caller falls
2799    /// back to it. The two former conditions (`len != 1` and "a non-Kronecker
2800    /// smooth term exists") are jointly equivalent to "the sole smooth term is
2801    /// Kronecker", which the slice pattern below expresses directly in one pass.
2802    pub fn kronecker_penalty_system(&self) -> Option<KroneckerPenaltySystem> {
2803        let [only_term] = self.smooth.terms.as_slice() else {
2804            return None;
2805        };
2806        let kron = only_term.kronecker_factored.as_ref()?;
2807        // A genuine tensor product needs at least two margins, and the marginal
2808        // design / penalty / dim collections must agree in length. A degenerate
2809        // (single-margin) or internally inconsistent factored basis cannot feed
2810        // the Kronecker fast path, so fall back to the standard assembly rather
2811        // than construct a malformed `KroneckerPenaltySystem` from it.
2812        if kron.marginal_dims.len() < 2
2813            || kron.marginal_penalties.len() != kron.marginal_dims.len()
2814            || kron.marginal_designs.len() != kron.marginal_dims.len()
2815        {
2816            return None;
2817        }
2818        KroneckerPenaltySystem::new(
2819            kron.marginal_penalties.clone(),
2820            kron.marginal_dims.clone(),
2821            kron.has_double_penalty,
2822        )
2823        .ok()
2824    }
2825}
2826
2827// `FittedTermCollection`, `SpatialLengthScaleOptimizationTiming`, and
2828// `FittedTermCollectionWithSpec` were relocated with the GAM fit-orchestration
2829// drivers to `gam-models` (`crate::fit_orchestration::drivers`) — they hold a
2830// `gam_solve::UnifiedFitResult` and are consumed only by those drivers (#1521).
2831
2832#[derive(Clone)]
2833pub struct StandardLatentCoordConfig {
2834    pub values: std::sync::Arc<crate::latent::LatentCoordValues>,
2835    pub term_index: gam_problem::types::SmoothTermIdx,
2836    pub feature_cols: Vec<usize>,
2837    pub manifold: crate::latent::LatentManifold,
2838    pub manifold_auto: bool,
2839    pub retraction_registry: gam_problem::LatentRetractionRegistry,
2840    pub analytic_penalties: Option<std::sync::Arc<crate::AnalyticPenaltyRegistry>>,
2841}
2842
2843#[derive(Clone, Debug, Serialize, Deserialize)]
2844pub struct AdaptiveSpatialMap {
2845    pub termname: String,
2846    pub feature_cols: Vec<usize>,
2847    pub collocation_points: Array2<f64>,
2848    pub inv_magweight: Array1<f64>,
2849    pub invgradweight: Array1<f64>,
2850    pub inv_lapweight: Array1<f64>,
2851}
2852
2853#[derive(Clone, Debug, Serialize, Deserialize)]
2854pub struct AdaptiveRegularizationDiagnostics {
2855    pub epsilon_0: f64,
2856    pub epsilon_g: f64,
2857    pub epsilon_c: f64,
2858    pub epsilon_outer_iterations: usize,
2859    pub mm_iterations: usize,
2860    pub converged: bool,
2861    pub maps: Vec<AdaptiveSpatialMap>,
2862}
2863
2864#[derive(Debug, Clone)]
2865pub struct LinearColumnConditioning {
2866    col_idx: usize,
2867    mean: f64,
2868    scale: f64,
2869}
2870
2871#[derive(Debug, Clone, Default)]
2872pub struct LinearFitConditioning {
2873    pub intercept_idx: usize,
2874    pub columns: Vec<LinearColumnConditioning>,
2875}
2876
2877#[derive(Clone)]
2878pub struct SpatialPsiDerivative {
2879    // These are derivatives with respect to psi = log(kappa), not log(length_scale).
2880    pub penalty_index: usize,
2881    pub penalty_indices: Vec<usize>,
2882    pub global_range: Range<usize>,
2883    pub total_p: usize,
2884    pub x_psi_local: Array2<f64>,
2885    pub s_psi_components_local: Vec<Array2<f64>>,
2886    pub x_psi_psi_local: Array2<f64>,
2887    pub s_psi_psi_components_local: Vec<Array2<f64>>,
2888    pub aniso_group_id: Option<usize>,
2889    /// Pre-computed cross-derivative design matrices for other axes
2890    /// in the same aniso group: Vec of (axis_offset_in_group, matrix).
2891    pub aniso_cross_designs: Option<Vec<(usize, Array2<f64>)>>,
2892    /// On-demand cross-penalty second derivatives ∂²S_m/∂ψ_a∂ψ_b for axes in
2893    /// the same anisotropy group. The input is the other axis offset in the
2894    /// group, and the output is one local penalty matrix per active penalty.
2895    pub aniso_cross_penalty_provider: Option<
2896        std::sync::Arc<
2897            dyn Fn(usize) -> Result<Vec<Array2<f64>>, EstimationError> + Send + Sync + 'static,
2898        >,
2899    >,
2900    /// Optional implicit design-derivative operator (shared across all axes
2901    /// in the same aniso group). When present, `x_psi_local` and
2902    /// `x_psi_psi_local` may be zero-sized, and design-derivative matvecs
2903    /// should go through this operator using `implicit_axis` as the axis index.
2904    pub implicit_operator: Option<std::sync::Arc<crate::basis::ImplicitDesignPsiDerivative>>,
2905    /// Which axis in the implicit operator this entry corresponds to.
2906    pub implicit_axis: usize,
2907}
2908
2909#[derive(Debug, Clone)]
2910pub struct SpatialLogKappaCoords {
2911    /// Flattened ψ values. For isotropic terms, one entry per term.
2912    /// For anisotropic terms, d entries per term (one ψ_a per axis).
2913    pub values: Array1<f64>,
2914    /// Dimensionality of each term: 1 for isotropic, d for anisotropic.
2915    pub dims_per_term: Vec<usize>,
2916}
2917
2918/// Which end of the ψ bound the shared `aniso_bounds_from_data` helper is
2919/// computing. The lower end consumes the `.0` element of
2920/// `spatial_term_psi_bounds`; the upper end consumes `.1`.
2921#[derive(Clone, Copy)]
2922pub enum AnisoBoundEnd {
2923    Lower,
2924    Upper,
2925}
2926
2927impl SpatialLogKappaCoords {
2928    /// Construct from an explicit dims layout plus values.
2929    pub fn new_with_dims(values: Array1<f64>, dims_per_term: Vec<usize>) -> Self {
2930        assert_eq!(
2931            values.len(),
2932            dims_per_term.iter().sum::<usize>(),
2933            "SpatialLogKappaCoords: values length {} != sum of dims_per_term {}",
2934            values.len(),
2935            dims_per_term.iter().sum::<usize>(),
2936        );
2937        Self {
2938            values,
2939            dims_per_term,
2940        }
2941    }
2942
2943    /// Isotropic initialization.
2944    pub fn from_length_scales(
2945        spec: &TermCollectionSpec,
2946        term_indices: &[usize],
2947        options: &SpatialLengthScaleOptimizationOptions,
2948    ) -> Self {
2949        let mut out = Array1::<f64>::zeros(term_indices.len());
2950        for (slot, &term_idx) in term_indices.iter().enumerate() {
2951            // Constant-curvature: the single ψ slot is the raw signed κ, seeded
2952            // from the spec (default κ = 0). The −ln(length_scale) convention is
2953            // log-κ semantics and must not touch the raw-κ coordinate; the κ
2954            // window projection happens later via `clamp_to_bounds`. Mirrors the
2955            // aniso constructor's κ branch.
2956            if let Some(cc) = constant_curvature_term_spec(spec, term_idx) {
2957                out[slot] = cc.kappa;
2958                continue;
2959            }
2960            // ψ = −ln(length_scale) through the SAME expression the aniso
2961            // constructor and the upstream spec projection use (#2726) — see
2962            // `spatial_length_scale_window::spatial_term_seed_psi`.
2963            out[slot] = spatial_term_seed_psi(spec, term_idx, options);
2964        }
2965        Self {
2966            values: out,
2967            dims_per_term: vec![1; term_indices.len()],
2968        }
2969    }
2970
2971    /// Anisotropic-aware initialization.
2972    ///
2973    /// The input frame is uniformly standardized by `IsotropicScale`; it never
2974    /// manufactures an axis preference. Genuine axis contrasts come only from
2975    /// the term's explicit, centered `aniso_log_scales` state.
2976    ///
2977    /// For each term, checks whether it has `aniso_log_scales` set on its basis spec.
2978    /// - If isotropic (no aniso_log_scales, or 1-D): 1 entry = −ln(length_scale).
2979    /// - If anisotropic with a scalar length scale: d entries, one ψ_a per axis.
2980    ///   Initialized as ψ_a = −ln(length_scale) + η_a  where η_a are the existing
2981    ///   aniso_log_scales (which sum to zero). Multi-dimensional terms without
2982    ///   explicit anisotropy stay scalar here so the seed dimensionality matches
2983    ///   `spatial_dims_per_term`.
2984    pub fn from_length_scales_aniso(
2985        spec: &TermCollectionSpec,
2986        term_indices: &[usize],
2987        options: &SpatialLengthScaleOptimizationOptions,
2988    ) -> Self {
2989        let mut vals = Vec::new();
2990        let mut dims = Vec::new();
2991        for &term_idx in term_indices {
2992            // Measure-jet: dial coordinates seeded directly from the term's
2993            // realized (α, τ[, s]); the −ln(length_scale) convention below is
2994            // κ-semantics and never applies to dials.
2995            if let Some(mj) = measure_jet_term_spec(spec, term_idx) {
2996                let seed = measure_jet_psi_seed(mj);
2997                dims.push(seed.len());
2998                vals.extend(seed);
2999                continue;
3000            }
3001            // Constant-curvature: one signed κ slot seeded from the spec's κ
3002            // (clamped feasible). The −ln(length_scale) convention below is
3003            // log-κ semantics and must not touch the raw-κ coordinate. Bounds
3004            // are unavailable here (no data view), so this is the raw spec κ;
3005            // `reseed_from_data` / `clamp_to_bounds` later project it feasible.
3006            if let Some(cc) = constant_curvature_term_spec(spec, term_idx) {
3007                vals.push(cc.kappa);
3008                dims.push(1);
3009                continue;
3010            }
3011            // Global scale ψ̄ = −ln(length_scale), through the SAME expression
3012            // the isotropic constructor and the upstream spec projection use
3013            // (#2726).
3014            let psi_bar = spatial_term_seed_psi(spec, term_idx, options);
3015
3016            if spatial_term_uses_per_axis_psi(spec, term_idx) {
3017                // Per-axis anisotropy is enrolled in the joint outer vector:
3018                // ψ_a = ψ̄ + η_a, one slot per axis. The hyper_dirs builder
3019                // produces matching per-axis derivatives in
3020                // `try_build_spatial_term_log_kappa_aniso_derivativeinfos`.
3021                let d = get_spatial_feature_dim(spec, term_idx).unwrap_or(1);
3022                let eta_raw = get_spatial_aniso_log_scales(spec, term_idx)
3023                    .expect("predicate guarantees aniso_log_scales is Some");
3024                let eta = center_aniso_log_scales(&eta_raw);
3025                for &eta_a in &eta {
3026                    vals.push(psi_bar + eta_a);
3027                }
3028                dims.push(d);
3029            } else {
3030                // Isotropic enrollment — either a 1-D term, a multi-D term
3031                // without explicit anisotropy, or a basis (e.g. Duchon) whose
3032                // η is a fixed geometry parameter rather than a REML hyper
3033                // axis. Exactly one ψ̄ slot, matching the single
3034                // `SpatialPsiDerivative` produced by
3035                // `try_build_spatial_term_log_kappa_derivativeinfo`.
3036                vals.push(psi_bar);
3037                dims.push(1);
3038            }
3039        }
3040        Self {
3041            values: Array1::from_vec(vals),
3042            dims_per_term: dims,
3043        }
3044    }
3045
3046    /// Isotropic lower bounds for the κ SEARCH BOX.
3047    ///
3048    /// Each entry gets the ψ_lo edge of [`spatial_term_psi_search_box`] — the
3049    /// data-geometry window widened to contain the term's own incumbent length
3050    /// scale, so the point the search starts at and is graded against is inside
3051    /// the set being searched (#2454).
3052    pub fn lower_bounds_from_data(
3053        data: ArrayView2<'_, f64>,
3054        spec: &TermCollectionSpec,
3055        term_indices: &[usize],
3056        options: &SpatialLengthScaleOptimizationOptions,
3057    ) -> Result<Self, BasisError> {
3058        let mut values = Array1::<f64>::zeros(term_indices.len());
3059        for (slot, &term_idx) in term_indices.iter().enumerate() {
3060            values[slot] = spatial_term_psi_search_box(data, spec, term_idx, options)?.0;
3061        }
3062        Ok(Self {
3063            values,
3064            dims_per_term: vec![1; term_indices.len()],
3065        })
3066    }
3067
3068    /// Isotropic upper bounds for the κ SEARCH BOX — the ψ_hi edge of
3069    /// [`spatial_term_psi_search_box`]; see [`Self::lower_bounds_from_data`].
3070    pub fn upper_bounds_from_data(
3071        data: ArrayView2<'_, f64>,
3072        spec: &TermCollectionSpec,
3073        term_indices: &[usize],
3074        options: &SpatialLengthScaleOptimizationOptions,
3075    ) -> Result<Self, BasisError> {
3076        let mut values = Array1::<f64>::zeros(term_indices.len());
3077        for (slot, &term_idx) in term_indices.iter().enumerate() {
3078            values[slot] = spatial_term_psi_search_box(data, spec, term_idx, options)?.1;
3079        }
3080        Ok(Self {
3081            values,
3082            dims_per_term: vec![1; term_indices.len()],
3083        })
3084    }
3085
3086    /// Anisotropic-aware lower bounds derived from per-term data geometry.
3087    /// For hybrid anisotropic terms the scalar ψ_lo bound applies to the
3088    /// mean `ψ̄`, not directly to every raw axis coordinate `ψ_a = ψ̄ + η_a`.
3089    /// Shift each axis by the current centered `η_a` so projecting/clamping
3090    /// the seed moves only the global scale direction and does not silently
3091    /// shrink anisotropy that is already consistent with the current
3092    /// `length_scale`.
3093    ///
3094    pub fn lower_bounds_aniso_from_data(
3095        data: ArrayView2<'_, f64>,
3096        spec: &TermCollectionSpec,
3097        term_indices: &[usize],
3098        dims_per_term: &[usize],
3099        options: &SpatialLengthScaleOptimizationOptions,
3100    ) -> Result<Self, BasisError> {
3101        Self::aniso_bounds_from_data(
3102            data,
3103            spec,
3104            term_indices,
3105            dims_per_term,
3106            options,
3107            AnisoBoundEnd::Lower,
3108        )
3109    }
3110
3111    /// Anisotropic-aware upper bounds derived from per-term data geometry.
3112    /// See `lower_bounds_aniso_from_data` for the hybrid-aniso offsetting and
3113    /// pure-Duchon dispatch rationale.
3114    pub fn upper_bounds_aniso_from_data(
3115        data: ArrayView2<'_, f64>,
3116        spec: &TermCollectionSpec,
3117        term_indices: &[usize],
3118        dims_per_term: &[usize],
3119        options: &SpatialLengthScaleOptimizationOptions,
3120    ) -> Result<Self, BasisError> {
3121        Self::aniso_bounds_from_data(
3122            data,
3123            spec,
3124            term_indices,
3125            dims_per_term,
3126            options,
3127            AnisoBoundEnd::Upper,
3128        )
3129    }
3130
3131    /// Shared implementation for the lower/upper anisotropic bounds. The bound
3132    /// end selects one element of the typed `(lo, hi)` data-geometry result;
3133    /// the per-term cursor walk and anisotropy-offset handling are identical.
3134    fn aniso_bounds_from_data(
3135        data: ArrayView2<'_, f64>,
3136        spec: &TermCollectionSpec,
3137        term_indices: &[usize],
3138        dims_per_term: &[usize],
3139        options: &SpatialLengthScaleOptimizationOptions,
3140        end: AnisoBoundEnd,
3141    ) -> Result<Self, BasisError> {
3142        assert_eq!(term_indices.len(), dims_per_term.len());
3143        let total: usize = dims_per_term.iter().sum();
3144        let mut values = Array1::<f64>::zeros(total);
3145        let mut cursor = 0;
3146        for (slot, &term_idx) in term_indices.iter().enumerate() {
3147            let d = dims_per_term[slot];
3148            // Measure-jet: per-coordinate dial boxes, never the κ-window
3149            // geometry (`-ln(length_scale)` is not this term's chart, and that
3150            // window would reject legitimate dial values outright). The
3151            // design-moving `ln ℓ` dial still gets a DATA-derived window — its
3152            // own node-spacing floor and node-diameter ceiling — because it is a
3153            // length in the chart the basis is realized in (gam#2750).
3154            if measure_jet_term_spec(spec, term_idx).is_some() {
3155                let term = spec
3156                    .smooth_terms
3157                    .get(term_idx)
3158                    .expect("measure_jet_term_spec resolved this index");
3159                let bounds = measure_jet_psi_bound_values(
3160                    data,
3161                    &term.basis,
3162                    matches!(end, AnisoBoundEnd::Upper),
3163                )?;
3164                for (offset, bound) in bounds.into_iter().enumerate() {
3165                    if offset < d {
3166                        values[cursor + offset] = bound;
3167                    }
3168                }
3169                cursor += d;
3170                continue;
3171            }
3172            // Constant-curvature: the single signed-κ box from the data chart
3173            // window (symmetric about κ = 0), never a κ = log-scale window.
3174            if constant_curvature_term_spec(spec, term_idx).is_some() {
3175                let (lo, hi) = constant_curvature_kappa_bounds(data, spec, term_idx);
3176                if d >= 1 {
3177                    values[cursor] = match end {
3178                        AnisoBoundEnd::Lower => lo,
3179                        AnisoBoundEnd::Upper => hi,
3180                    };
3181                }
3182                cursor += d;
3183                continue;
3184            }
3185            let psi_bound = {
3186                // The SEARCH box, not the bare geometry window: the same
3187                // incumbent-containment argument applies per axis, because the
3188                // axis offsets `η_a` are added to a common scalar ψ̄ bound
3189                // (#2454).
3190                let (lo, hi) = spatial_term_psi_search_box(data, spec, term_idx, options)?;
3191                match end {
3192                    AnisoBoundEnd::Lower => lo,
3193                    AnisoBoundEnd::Upper => hi,
3194                }
3195            };
3196            let axis_offsets = if d <= 1 {
3197                vec![0.0; d]
3198            } else {
3199                get_spatial_aniso_log_scales(spec, term_idx)
3200                    .filter(|eta| eta.len() == d)
3201                    .map(|eta| center_aniso_log_scales(&eta))
3202                    .unwrap_or_else(|| vec![0.0; d])
3203            };
3204            for offset in 0..d {
3205                values[cursor + offset] = psi_bound + axis_offsets[offset];
3206            }
3207            cursor += d;
3208        }
3209        Ok(Self {
3210            values,
3211            dims_per_term: dims_per_term.to_vec(),
3212        })
3213    }
3214
3215    /// Rewrite any ψ entries whose originating term lacks an explicit
3216    /// `length_scale` so they sit at the midpoint of the per-term data-derived
3217    /// ψ window. Used so the outer optimizer starts inside the physically
3218    /// meaningful region instead of at an arbitrary `options.max_length_scale`
3219    /// derived seed. For terms with an explicit length_scale, the user's
3220    /// choice is respected. Anisotropy offsets η_a (those stored by
3221    /// `from_length_scales_aniso`) are preserved: we re-center around the new
3222    /// ψ̄, keeping Ση_a = 0.
3223    pub fn reseed_from_data(
3224        mut self,
3225        data: ArrayView2<'_, f64>,
3226        spec: &TermCollectionSpec,
3227        term_indices: &[usize],
3228        options: &SpatialLengthScaleOptimizationOptions,
3229    ) -> Result<Self, BasisError> {
3230        assert_eq!(term_indices.len(), self.dims_per_term.len());
3231        let mut cursor = 0;
3232        for (slot, &term_idx) in term_indices.iter().enumerate() {
3233            let d = self.dims_per_term[slot];
3234            // Measure-jet dials are seeded from the realized spec and must
3235            // not be recentered into a κ data window.
3236            if measure_jet_term_spec(spec, term_idx).is_some() {
3237                cursor += d;
3238                continue;
3239            }
3240            // Constant-curvature κ is seeded from the spec (the user's curvature
3241            // hint, default κ = 0); `clamp_to_bounds` projects it feasible. It
3242            // is not a log-scale, so the log-κ recenter below never applies.
3243            if constant_curvature_term_spec(spec, term_idx).is_some() {
3244                cursor += d;
3245                continue;
3246            }
3247            let Some(psi_bar_new) = spatial_term_psi_seed(data, spec, term_idx, options)? else {
3248                cursor += d;
3249                continue;
3250            };
3251            if d == 0 {
3252                continue;
3253            }
3254            let current: Vec<f64> = self.values.slice(s![cursor..cursor + d]).to_vec();
3255            let psi_bar_old = current.iter().sum::<f64>() / d as f64;
3256            for (offset, &old_value) in current.iter().enumerate() {
3257                self.values[cursor + offset] = psi_bar_new + (old_value - psi_bar_old);
3258            }
3259            cursor += d;
3260        }
3261        Ok(self)
3262    }
3263
3264    /// Project ψ values into `[lower, upper]` element-wise. Used after
3265    /// `from_length_scales*` + `reseed_from_data` when a user-supplied
3266    /// `spec.length_scale` falls outside the data-derived ψ window set by
3267    /// `{lower,upper}_bounds*_from_data`. BFGS requires theta0 ∈ [lower,
3268    /// upper]; projecting is the unique closest feasible seed. The user's
3269    /// length_scale was always a hint for the outer optimizer (the optimizer
3270    /// is authoritative for κ), not a hard constraint — so clipping preserves
3271    /// their intent as far as the geometry allows. Emits `log::info!` when
3272    /// any coordinate moves, so the outside-window case is diagnostically
3273    /// visible (not silent).
3274    pub fn clamp_to_bounds(
3275        mut self,
3276        lower: &SpatialLogKappaCoords,
3277        upper: &SpatialLogKappaCoords,
3278    ) -> Self {
3279        assert_eq!(self.values.len(), lower.values.len());
3280        assert_eq!(self.values.len(), upper.values.len());
3281        let mut n_projected = 0usize;
3282        let mut worst_delta = 0.0_f64;
3283        for idx in 0..self.values.len() {
3284            let lo = lower.values[idx];
3285            let hi = upper.values[idx];
3286            if !(lo.is_finite() && hi.is_finite()) {
3287                continue;
3288            }
3289            let v = self.values[idx];
3290            if v < lo {
3291                worst_delta = worst_delta.max(lo - v);
3292                self.values[idx] = lo;
3293                n_projected += 1;
3294            } else if v > hi {
3295                worst_delta = worst_delta.max(v - hi);
3296                self.values[idx] = hi;
3297                n_projected += 1;
3298            }
3299        }
3300        if n_projected > 0 {
3301            log::info!(
3302                "[spatial-kappa] projected {n_projected}/{} ψ seed coords into data-derived bounds \
3303                 (worst excess={worst_delta:.3} log units); user length_scale falls outside \
3304                 [{KERNEL_RANGE_MIN_DIAMETER_FRACTION}/r_max, {KERNEL_RANGE_MAX_SPACING_MULTIPLE}/r_min] geometry window",
3305                self.values.len()
3306            );
3307        }
3308        self
3309    }
3310
3311    /// Reconstruct from theta tail with known dimensionality layout.
3312    pub fn from_theta_tail_with_dims(
3313        theta: &Array1<f64>,
3314        start: usize,
3315        dims_per_term: Vec<usize>,
3316    ) -> Self {
3317        let total: usize = dims_per_term.iter().sum();
3318        Self {
3319            values: theta.slice(s![start..start + total]).to_owned(),
3320            dims_per_term,
3321        }
3322    }
3323
3324    /// Total number of ψ values in the flat array (= sum of dims_per_term).
3325    pub fn len(&self) -> usize {
3326        self.values.len()
3327    }
3328
3329    /// Dimensionality layout: how many ψ values each term contributes.
3330    pub fn dims_per_term(&self) -> &[usize] {
3331        &self.dims_per_term
3332    }
3333
3334    /// Get the offset into the flat array for logical term i.
3335    fn term_offset(&self, term_idx: usize) -> usize {
3336        self.dims_per_term[..term_idx].iter().sum()
3337    }
3338
3339    /// Get the slice of ψ values for logical term i.
3340    pub fn term_slice(&self, term_idx: usize) -> &[f64] {
3341        let offset = self.term_offset(term_idx);
3342        let d = self.dims_per_term[term_idx];
3343        &self
3344            .values
3345            .as_slice()
3346            .expect("psi values are an owned contiguous Array1")[offset..offset + d]
3347    }
3348
3349    pub fn as_array(&self) -> &Array1<f64> {
3350        &self.values
3351    }
3352
3353    /// #1464: overwrite the single ψ value of a scalar (1-D) logical term by its
3354    /// position `slot` in this coords vector (the same ordering as the
3355    /// `term_indices` slice the constructors were built from). Used to install
3356    /// the continuously selected curvature-likelihood optimum into a
3357    /// constant-curvature term's raw-κ slot before the nuisance joint solve.
3358    /// No-op (returns `false`) when the slot is not scalar.
3359    pub fn set_scalar_slot(&mut self, slot: usize, value: f64) -> bool {
3360        if slot >= self.dims_per_term.len() || self.dims_per_term[slot] != 1 {
3361            return false;
3362        }
3363        let offset = self.term_offset(slot);
3364        self.values[offset] = value;
3365        true
3366    }
3367
3368    /// Split at a logical-term boundary. `mid` is the number of terms in the
3369    /// first half (not a flat-array index).
3370    pub fn split_at(&self, mid: usize) -> (Self, Self) {
3371        let flat_mid: usize = self.dims_per_term[..mid].iter().sum();
3372        (
3373            Self {
3374                values: self.values.slice(s![0..flat_mid]).to_owned(),
3375                dims_per_term: self.dims_per_term[..mid].to_vec(),
3376            },
3377            Self {
3378                values: self.values.slice(s![flat_mid..]).to_owned(),
3379                dims_per_term: self.dims_per_term[mid..].to_vec(),
3380            },
3381        )
3382    }
3383
3384    /// Apply optimized ψ values back to the spec.
3385    ///
3386    /// For isotropic terms (dims=1): sets scalar length_scale = exp(−ψ).
3387    /// For anisotropic terms (dims=d): hybrid/isotropic families set
3388    /// length_scale = exp(−ψ̄) with centered η_a = ψ_a − ψ̄, while pure Duchon
3389    /// writes only centered η_a and leaves length_scale = None.
3390    pub fn apply_tospec(
3391        &self,
3392        spec: &TermCollectionSpec,
3393        term_indices: &[usize],
3394    ) -> Result<TermCollectionSpec, EstimationError> {
3395        if term_indices.len() != self.dims_per_term.len() {
3396            crate::bail_invalid_estim!(
3397                "SpatialLogKappaCoords::apply_tospec: term count mismatch: \
3398                 term_indices={} dims_per_term={}",
3399                term_indices.len(),
3400                self.dims_per_term.len()
3401            );
3402        }
3403        let mut updated = spec.clone();
3404        for (slot, &term_idx) in term_indices.iter().enumerate() {
3405            let psi = self.term_slice(slot);
3406            let d = self.dims_per_term[slot];
3407            // Measure-jet: write the dial coordinates straight back; the
3408            // κ-translation below would misread them as log-scales.
3409            if measure_jet_term_spec(&updated, term_idx).is_some() {
3410                set_measure_jet_psi_dials(&mut updated, term_idx, psi)?;
3411                continue;
3412            }
3413            // Constant-curvature: write the optimized signed κ straight back;
3414            // the −exp(ψ) length-scale translation below is log-κ semantics and
3415            // would misread the raw curvature.
3416            if constant_curvature_term_spec(&updated, term_idx).is_some() {
3417                set_constant_curvature_kappa(&mut updated, term_idx, psi)?;
3418                continue;
3419            }
3420            let (next_length_scale, next_aniso) = spatial_term_psi_to_length_scale_and_aniso(psi);
3421            if (d == 1 || next_length_scale.is_some())
3422                && let Some(length_scale) = next_length_scale
3423            {
3424                set_spatial_length_scale(&mut updated, term_idx, length_scale)?;
3425            }
3426            if let Some(eta) = next_aniso {
3427                set_spatial_aniso_log_scales(&mut updated, term_idx, eta)?;
3428            }
3429        }
3430        Ok(updated)
3431    }
3432}
3433
3434pub fn center_aniso_log_scales(eta: &[f64]) -> Vec<f64> {
3435    if eta.len() <= 1 {
3436        return eta.to_vec();
3437    }
3438    let mean = eta.iter().sum::<f64>() / eta.len() as f64;
3439    eta.iter()
3440        .map(|&v| {
3441            let centered = v - mean;
3442            if centered.abs() <= 1e-15 {
3443                0.0
3444            } else {
3445                centered
3446            }
3447        })
3448        .collect()
3449}
3450
3451/// Whether a spatial term contributes per-axis ψ entries to the outer joint
3452/// hyperparameter vector.
3453pub fn spatial_term_uses_per_axis_psi(resolvedspec: &TermCollectionSpec, term_idx: usize) -> bool {
3454    if let Some(mj) = measure_jet_term_spec(resolvedspec, term_idx) {
3455        return measure_jet_enrolls_psi(mj);
3456    }
3457    let Some(d) = get_spatial_feature_dim(resolvedspec, term_idx) else {
3458        return false;
3459    };
3460    if d <= 1 {
3461        return false;
3462    }
3463    let Some(eta) = get_spatial_aniso_log_scales(resolvedspec, term_idx) else {
3464        return false;
3465    };
3466    if eta.len() != d {
3467        return false;
3468    }
3469    // gam#2735 — a hybrid Duchon's per-axis η IS a REML coordinate wherever its
3470    // per-axis ψ derivative surface is complete.
3471    //
3472    // It used to be excluded outright, on the reading that "η is a FIXED,
3473    // geometry-derived basis parameter … standardize the geometry, then learn
3474    // the smoothness". That reading is right about the SEED and wrong about the
3475    // estimand: `initial_aniso_contrasts` reads the per-axis spread of the KNOT
3476    // CLOUD, which is a property of where the inputs are and carries no
3477    // information about which axis the RESPONSE varies along. On a design whose
3478    // inputs are isotropic and whose signal is not — `large_scale_reml_stress`
3479    // draws `X ~ N(0, I)` and puts its entire non-linear content on one axis —
3480    // the seeded contrasts are sampling noise, and freezing them there costs
3481    // 1795 nats of criterion and 3.6x of held-out reconstruction error against
3482    // the η the criterion itself prefers. The contrasts are identifiable (they
3483    // change the kernel's shape, not merely its scale), so REML can and should
3484    // estimate them, exactly as it already does for the anisotropic Matérn.
3485    //
3486    // The capability predicate — not this site — decides which specs qualify;
3487    // anything it declines stays on its single isotropic ψ axis, bit-identically
3488    // to before.
3489    let Some(term) = resolvedspec.smooth_terms.get(term_idx) else {
3490        return false;
3491    };
3492    match &term.basis {
3493        SmoothBasisSpec::Duchon { spec, .. } => {
3494            // A joint null rotation `Q` has to be applied to every ψ-derivative
3495            // block — the isotropic arm does it explicitly in
3496            // `try_build_spatial_term_log_kappa_derivative`. The per-axis
3497            // consumer does not, so a rotated Duchon term stays isotropic
3498            // rather than shipping an unrotated per-axis derivative against a
3499            // rotated design. (The anisotropic Matérn has the same gap; it is
3500            // pre-existing and not touched here.)
3501            term.joint_null_rotation.is_none()
3502                && crate::basis::duchon_spec_supports_axis_psi(spec, d)
3503        }
3504        _ => true,
3505    }
3506}
3507
3508pub fn set_spatial_length_scale(
3509    spec: &mut TermCollectionSpec,
3510    term_idx: usize,
3511    length_scale: f64,
3512) -> Result<(), EstimationError> {
3513    let Some(term) = spec.smooth_terms.get_mut(term_idx) else {
3514        crate::bail_invalid_estim!("spatial length-scale term index {term_idx} out of range");
3515    };
3516    match &mut term.basis {
3517        SmoothBasisSpec::ThinPlate { spec, .. } => {
3518            spec.length_scale = length_scale;
3519            Ok(())
3520        }
3521        SmoothBasisSpec::Matern { spec, .. } => {
3522            spec.length_scale.set_resolved(length_scale);
3523            Ok(())
3524        }
3525        SmoothBasisSpec::Duchon { spec, .. } => {
3526            spec.length_scale = Some(length_scale);
3527            Ok(())
3528        }
3529        _ => Err(EstimationError::InvalidInput(format!(
3530            "term '{}' does not expose a spatial length scale",
3531            term.name
3532        ))),
3533    }
3534}
3535
3536pub fn get_spatial_length_scale(spec: &TermCollectionSpec, term_idx: usize) -> Option<f64> {
3537    spec.smooth_terms
3538        .get(term_idx)
3539        .and_then(|term| match &term.basis {
3540            SmoothBasisSpec::ThinPlate { spec, .. } => Some(spec.length_scale),
3541            SmoothBasisSpec::Matern { spec, .. } => spec.length_scale.resolved(),
3542            SmoothBasisSpec::Duchon { spec, .. } => spec.length_scale,
3543            _ => None,
3544        })
3545}
3546
3547pub fn spatial_term_supports_hyper_optimization(
3548    spec: &TermCollectionSpec,
3549    term_idx: usize,
3550) -> bool {
3551    // Ordinary penalized thin-plate regression splines do not have an
3552    // identifiable kernel scale once REML is already learning the smoothing
3553    // penalty. Treat the resolved length scale as fixed geometry; enrolling a
3554    // scalar TPS kappa axis creates the flat ρ/κ valleys reported in #718,
3555    // #721, #731, and #732.
3556    if let Some(term) = spec.smooth_terms.get(term_idx)
3557        && let SmoothBasisSpec::ThinPlate { .. } = &term.basis
3558    {
3559        return false;
3560    }
3561
3562    // Duchon anisotropy η is SEEDED from geometry and ESTIMATED by REML
3563    // (gam#2735). `auto_seed_aniso_contrasts` still supplies the starting
3564    // contrasts from the knot-cloud spread on every Duchon basis build — that is
3565    // a good seed, because it standardizes a genuinely elongated input cloud
3566    // before the search begins. What it cannot do is finish the job: the knot
3567    // cloud says where the inputs ARE, not which axis the RESPONSE varies along,
3568    // so on an isotropic design it seeds noise and a frozen η is then simply
3569    // wrong. `spatial_term_uses_per_axis_psi` enrolls the contrasts as outer ψ
3570    // coordinates wherever `duchon_spec_supports_axis_psi` certifies the
3571    // per-axis derivative surface; a pure Duchon (no κ) is one of the
3572    // configurations it declines, so that path is unchanged. Only an explicit
3573    // kernel length scale κ (the Matérn / hybrid path) is optimized here.
3574    //
3575    // ISOTROPIC Matérn: the *default* `matern(x1, x2)` is isotropic
3576    // (`scale_dims=false` → `aniso_log_scales = None`). It contributes exactly
3577    // ONE κ optimization axis — its scalar log-κ. The shared GAMLSS /
3578    // location-scale exact-joint ψ engine and the spatial-κ joint outer solver
3579    // both require an isotropic Matérn block to expose this single isotropic κ
3580    // axis (#822/#851); without it the per-block ψ-derivative lists are empty
3581    // and the joint-ψ hooks degenerate to `None`. The isotropic κ is the lone
3582    // kernel hyper axis here, mirroring the per-axis ψ ARD that the anisotropic
3583    // path exposes (just collapsed to one dimension).
3584    //
3585    // ANISOTROPIC Matérn (`scale_dims=true` → `aniso_log_scales = Some`) keeps
3586    // its per-axis kernel-η ARD: the d-dimensional ψ search is the *point* of
3587    // the anisotropic request ("Matérn keeps its kernel-η ARD").
3588    //
3589    // Either way a Matérn term always enrolls a κ/ψ axis (1 isotropic, or d
3590    // anisotropic), so `spatial_dims_per_term` reports the correct count.
3591    if let Some(term) = spec.smooth_terms.get(term_idx)
3592        && let SmoothBasisSpec::Matern { .. } = &term.basis
3593    {
3594        return true;
3595    }
3596
3597    // Measure-jet geometry dials are outer ψ coordinates; enrollment is
3598    // owned by `measure_jet_enrolls_psi`.
3599    if let Some(mj) = measure_jet_term_spec(spec, term_idx) {
3600        return measure_jet_enrolls_psi(mj);
3601    }
3602
3603    // Constant-curvature smooths always enroll (#944 stage 3): κ̂ is the headline
3604    // estimand, so unlike a fixed-ℓ kernel the geometry is fitted by default and
3605    // not gated on a user-supplied scale. The term carries TWO outer
3606    // coordinates — the raw signed κ (interior κ = 0) and the log range
3607    // η = ln ℓ, which must be estimated because it is confounded with κ
3608    // (gam#2747) — but both are owned by the term-local curvature profile
3609    // (`constant_curvature_kappa_profile_optimum`), which certifies them BEFORE
3610    // the joint spatial solve and is why `fit_term_collectionwith_spatial_
3611    // length_scale_optimization` filters constant-curvature terms out of its
3612    // `spatial_terms` list. Enrolling here keeps the term in the eligibility
3613    // scan that reaches that profile.
3614    if constant_curvature_term_spec(spec, term_idx).is_some() {
3615        return true;
3616    }
3617
3618    get_spatial_length_scale(spec, term_idx).is_some()
3619}
3620
3621/// The constant-curvature smooth's spec, when `term_idx` is one. Single
3622/// accessor for every κ-ψ dispatch below, mirroring `measure_jet_term_spec`.
3623pub fn constant_curvature_term_spec(
3624    spec: &TermCollectionSpec,
3625    term_idx: usize,
3626) -> Option<&crate::basis::ConstantCurvatureBasisSpec> {
3627    spec.smooth_terms
3628        .get(term_idx)
3629        .and_then(|term| match &term.basis {
3630            SmoothBasisSpec::ConstantCurvature { spec, .. } => Some(spec),
3631            _ => None,
3632        })
3633}
3634
3635/// The fraction of the way from κ = 0 to the nearest singular κ that the outer
3636/// search is allowed to consume. Each branch has its own singularity and its own
3637/// gauge, and this one fraction is the retreat applied to both (#2687):
3638///
3639/// * **κ < 0 — the chart edge, a PER-POINT gauge.** `λ(p) = 1 + κ‖p‖²` is the
3640///   conformal factor's denominator and vanishes at `κ = −1/‖p‖²`: the point
3641///   reaches the boundary of the Poincaré ball. The retreat is `λ ≥ 1 − F`, so
3642///   the gauge may lose at most this fraction of its flat (κ = 0) value of 1.
3643/// * **κ > 0 — the antipodal fold, a PER-PAIR gauge.** The per-point gauge is
3644///   vacuous there (`‖p‖² ≥ 0`), and it is not what the kernel evaluates anyway.
3645///   Every distance goes through `w = (−x) ⊕_κ c`, whose Möbius denominator
3646///   `D = 1 + 2κ⟨x,c⟩ + κ²‖x‖²‖c‖²` is `(1 − κ‖x‖‖c‖)²` for an anti-aligned pair
3647///   and **vanishes at `κ = +1/(‖x‖‖c‖)`** — the two points are exactly
3648///   antipodal and `w` passes through infinity. Past the fold the chart is not
3649///   merely inaccurate but folded: the pair's scale-free geodesic separation is
3650///   exactly invariant under `κ ↦ 1/(κ‖x‖²‖c‖²)`, so every κ beyond it
3651///   duplicates one before it. `√D` is the pair's gauge — literally
3652///   `|1 ∓ κ‖x‖‖c‖|`, the same shape as `λ` — and the retreat is the same
3653///   `√D ≥ 1 − F`, i.e. `κ‖x‖‖c‖ ≤ F`.
3654///
3655/// Both gauges hit their wall at `|κ| = 1/R²` when `‖x‖ = ‖c‖ = R`, which is why
3656/// the window is symmetric on a cloud whose centers reach the data radius; it is
3657/// not one branch's constraint mirrored onto the other. Gated by
3658/// `spherical_branch_folds_at_kappa_r2_one_so_the_kappa_window_is_symmetric_2687`
3659/// in `gam-geometry`, which pins the fold, the refusal, and the involution.
3660///
3661/// ## What `0.5` buys, measured (#2687)
3662///
3663/// This is a MODELLING retreat, not a numerical one, and the measurement that
3664/// separates the two is in
3665/// `gam_geometry::manifolds::constant_curvature_antipodal_resolution_tests`.
3666/// Differencing the shipped Möbius route against the cancellation-free closed
3667/// form `d = (4/√κ)·arctan(√κ R)` gives
3668///
3669/// ```text
3670///   rel_err(∂d/∂κ)   ≈ ε / D
3671///   rel_err(∂²d/∂κ²) ≈ ε / D^{3/2}
3672/// ```
3673///
3674/// At `F = 0.5` the worst evaluated pair has `D ≥ (1 − F)² = 0.25`, where the
3675/// κ-Hessian the outer route consumes (`Derivative::Analytic`, exact `d²V/dκ²`)
3676/// carries ~14 of 16 digits. Half the mantissa — the bar a Newton Hessian
3677/// actually needs — is reached only at `D = ε^{1/3} ≈ 6.1e-6`, i.e. `κR² ≈
3678/// 0.9975`. **The arithmetic permits a box 203× closer to the fold in `1 − κR²`
3679/// than this fraction goes.** So moving `F` needs an argument about the
3680/// ESTIMATOR, not about the arithmetic, and #2687 carries the measurement that
3681/// argues against widening it: on the fixture whose κ̂ is railed here, the
3682/// profiled criterion is monotone across the entire interval, so a wider box
3683/// only moves the rail — to κ̂ = 2.78 against a planted 1.5, further from the
3684/// truth than the shipped box gives.
3685///
3686/// κ = 0 (flat) is the centre of the window, an interior point of the
3687/// `S^d ← ℝ^d → H^d` family — exactly the reachability the raw-κ (not log-κ)
3688/// coordinate exists to preserve.
3689pub const CONSTANT_CURVATURE_KAPPA_CHART_FRACTION: f64 = 0.5;
3690
3691/// Floor on the data's squared chart radius used to scale the κ window, so a
3692/// degenerate (near-origin) point cloud still yields a finite, usable bracket
3693/// rather than an unbounded one.
3694pub const CONSTANT_CURVATURE_MIN_CHART_RADIUS2: f64 = 1e-8;
3695
3696/// `(κ_min, κ_max)` outer-optimization window for a constant-curvature term,
3697/// derived over the configuration the basis actually EVALUATES.
3698///
3699/// Let `R = max‖p‖` over every point the basis touches — the term's feature
3700/// columns AND the centers. Then
3701///
3702/// ```text
3703///   (κ_min, κ_max) = ( −F/R² , +F/R² )
3704/// ```
3705///
3706/// which is the pre-#2716 formula with `R` taken over the right set. `R_c` comes
3707/// from [`constant_curvature_center_chart_radius2`](crate::basis::constant_curvature_center_chart_radius2),
3708/// which bounds `max‖c‖²` per strategy WITHOUT materializing the centers.
3709///
3710/// ## Why `R` and not the two walls separately
3711///
3712/// The two walls really are different objects — the κ<0 wall is the per-POINT
3713/// chart gauge `1 + κ‖p‖²`, the κ>0 wall is the per-PAIR antipodal fold at
3714/// `κ‖p‖‖q‖ = 1` — and taking each over its own set would give
3715/// `κ_max = F/(R_x·R_c)`, wider than `F/R²` whenever the centers sit inside the
3716/// data hull. That is mathematically the tighter statement and it is the wrong
3717/// bound to ship, for a reason that is not visible from the geometry:
3718///
3719/// **the box has to be freeze-invariant.** `freeze_term_collection_from_design`
3720/// rewrites a fitted term's `CenterStrategy` as `UserProvided(realized centers)`,
3721/// so the SAME configuration is described data-driven before the fit and
3722/// user-provided after it. A `κ_max` that reads the realized center radius moves
3723/// between those two descriptions — measured: on the #944 coverage fixture the
3724/// fit railed at `1.412031543260163` (its data-driven `F/R_x²`) and inference
3725/// then computed `1.4127975943783915` from the frozen centers, so κ̂ was no
3726/// longer at its own bound, was classified interior, and the whole fit was
3727/// refused as a non-stationary point estimate. A bound that changes when the
3728/// same geometry is re-described is not a bound on the geometry.
3729///
3730/// `max(R_x, R_c)` is the smallest radius that survives the re-description, and
3731/// it is conservative against the true evaluated-pair maximum: the RKHS penalty
3732/// Gram evaluates `K_κ(centers, centers)` as well as `K_κ(data, centers)`, so the
3733/// evaluated products run up to `R_c·max(R_x, R_c) ≤ R²`.
3734///
3735/// ## What moves, and what does not
3736///
3737/// Every data-driven strategy selects data rows verbatim or convex combinations
3738/// of them, so `R_c ≤ R_x`, `R = R_x`, and the box is **bit-identical** to its
3739/// pre-#2716 value. Two strategies move, both in the direction that was wrong:
3740///
3741/// * `UserProvided` — verbatim, any radius. The old upper end passed the fold
3742///   once `R_c ≥ 2·R_x`, admitting two κ that produce an identical scale-free
3743///   geometry for the extreme pair (the doubly-covered box #2716 measured), and
3744///   the old lower end let the search reach a κ at which a center is outside the
3745///   chart, where `validate_chart_points` refuses — turning a box excursion into
3746///   a hard basis-build error mid-optimization rather than a rail.
3747/// * `UniformGrid` — the Cartesian product of per-axis linspaces over the data's
3748///   BOUNDING BOX, so a corner center sits at up to `√d·R_x`, crossing the same
3749///   threshold at `d ≥ 4` with no user input at all.
3750///
3751/// See [`CONSTANT_CURVATURE_KAPPA_CHART_FRACTION`] for the two gauges and for
3752/// what its `0.5` buys, measured.
3753pub fn constant_curvature_kappa_bounds(
3754    data: ArrayView2<'_, f64>,
3755    spec: &TermCollectionSpec,
3756    term_idx: usize,
3757) -> (f64, f64) {
3758    let (feature_cols, cc) = match spec.smooth_terms.get(term_idx).map(|t| &t.basis) {
3759        Some(SmoothBasisSpec::ConstantCurvature {
3760            feature_cols, spec, ..
3761        }) => (feature_cols, spec),
3762        _ => return (-1.0, 1.0),
3763    };
3764    let data_r2 = crate::basis::constant_curvature_data_chart_radius2(data, feature_cols);
3765    let center_r2 = crate::basis::constant_curvature_center_chart_radius2(
3766        data,
3767        feature_cols,
3768        &cc.center_strategy,
3769    );
3770    let max_r2 = data_r2
3771        .max(center_r2)
3772        .max(CONSTANT_CURVATURE_MIN_CHART_RADIUS2);
3773    let half = CONSTANT_CURVATURE_KAPPA_CHART_FRACTION / max_r2;
3774    (-half, half)
3775}
3776
3777/// Write the optimized κ back into a constant-curvature term spec. Returns
3778/// `true` when κ moved. Centers, ℓ, and the constraint transform `z` are
3779/// κ-FIXED by the basis κ-contract, so only `kappa` changes.
3780pub fn set_constant_curvature_kappa(
3781    spec: &mut TermCollectionSpec,
3782    term_idx: usize,
3783    psi: &[f64],
3784) -> Result<bool, EstimationError> {
3785    let Some(term) = spec.smooth_terms.get_mut(term_idx) else {
3786        crate::bail_invalid_estim!(
3787            "constant-curvature κ write-back: term index {term_idx} out of range"
3788        );
3789    };
3790    set_single_term_constant_curvature_kappa(term, psi)
3791}
3792
3793/// Single-term κ write-back: the shared validate+apply core, also used directly
3794/// on the cached per-trial build spec in the incremental realizer (whose caller
3795/// has already change-checked at the collection level and rebuilds regardless
3796/// of the moved flag). Mirrors [`set_single_term_measure_jet_psi_dials`].
3797pub fn set_single_term_constant_curvature_kappa(
3798    term: &mut SmoothTermSpec,
3799    psi: &[f64],
3800) -> Result<bool, EstimationError> {
3801    if psi.len() != 1 {
3802        crate::bail_invalid_estim!(
3803            "constant-curvature κ write-back expects exactly one value, got {}",
3804            psi.len()
3805        );
3806    }
3807    let next_kappa = psi[0];
3808    if !next_kappa.is_finite() {
3809        crate::bail_invalid_estim!(
3810            "constant-curvature κ write-back produced a non-finite κ = {next_kappa}"
3811        );
3812    }
3813    let SmoothBasisSpec::ConstantCurvature { spec: cc, .. } = &mut term.basis else {
3814        crate::bail_invalid_estim!(
3815            "constant-curvature κ write-back targeted a non-constant-curvature term"
3816        );
3817    };
3818    if cc.kappa != next_kappa {
3819        cc.kappa = next_kappa;
3820        Ok(true)
3821    } else {
3822        Ok(false)
3823    }
3824}
3825
3826/// Returns `true` when a spatial term has NO outer optimization axes — i.e.
3827/// the user provided an explicit `length_scale` and the term does not enroll
3828/// REML-side per-axis ψ contrasts, so both the scalar κ and any fixed geometry
3829/// anisotropy are anchored.
3830///
3831/// This is the per-term predicate that distinguishes "fixed kernel scale"
3832/// from "optimize the kernel scale" within the family entry points that
3833/// want to honor an explicit user-supplied scale (e.g. Bernoulli
3834/// marginal-slope, where the joint-spatial outer solver otherwise spends
3835/// ~80 iters stalled on the user's chosen ρ at high gradient).
3836pub fn spatial_term_has_locked_kappa(spec: &TermCollectionSpec, term_idx: usize) -> bool {
3837    let explicitly_fixed = spec
3838        .smooth_terms
3839        .get(term_idx)
3840        .is_some_and(|term| match &term.basis {
3841            SmoothBasisSpec::Matern { spec, .. } => spec.length_scale.is_fixed(),
3842            SmoothBasisSpec::ThinPlate { .. } => true,
3843            SmoothBasisSpec::Duchon { spec, .. } => spec.length_scale.is_some(),
3844            _ => false,
3845        });
3846    explicitly_fixed && !spatial_term_uses_per_axis_psi(spec, term_idx)
3847}
3848
3849/// Returns `true` when every spatial term in `spec` has a locked kernel scale
3850/// (explicit `length_scale=X` without anisotropy) and therefore contributes no
3851/// outer ψ/κ optimization axis. Empty term collections also return `true` —
3852/// there are no kappas to optimize.
3853///
3854/// Used by family entry points that want to honor a user-supplied scalar length
3855/// scale exactly: when all spatial terms are locked the n-block joint-spatial
3856/// outer solver has nothing to optimize, and routing through it merely spends
3857/// ~80 outer iters chasing a stalled ARC at the user's chosen ρ. Skipping
3858/// straight to the rho-only path avoids that waste and respects the user's
3859/// explicit kernel-scale input.
3860pub fn all_spatial_terms_kappa_fixed(spec: &TermCollectionSpec) -> bool {
3861    spec.smooth_terms.iter().enumerate().all(|(idx, _)| {
3862        !spatial_term_supports_hyper_optimization(spec, idx)
3863            || spatial_term_has_locked_kappa(spec, idx)
3864    })
3865}
3866
3867pub fn spatial_identifiability_policy(
3868    termspec: &SmoothTermSpec,
3869) -> Option<&SpatialIdentifiability> {
3870    match &termspec.basis {
3871        SmoothBasisSpec::ThinPlate { spec, .. } => Some(&spec.identifiability),
3872        SmoothBasisSpec::Duchon { spec, .. } => Some(&spec.identifiability),
3873        _ => None,
3874    }
3875}
3876
3877/// Standard deviation of the wide, weakly-informative symmetric `Normal` prior
3878/// placed on a relaxable double-penalty smooth's `DoublePenaltyNullspace`
3879/// selection coordinate in every design regime.
3880pub const NULLSPACE_DEGENERACY_RHO_SD: f64 = 15.0;
3881
3882
3883/// Per-term data-derived ψ = log κ bounds.
3884///
3885/// Uses the same safe operating range documented in
3886/// [`crate::basis::build_matern_basis`] / [`crate::basis::build_duchon_basis`]:
3887///   κ ∈ [2 / r_max, 1e2 / r_min]
3888/// where (r_min, r_max) are pairwise-distance extrema of the term's resolved
3889/// centers (post-fit) or the standardized feature data columns (pre-fit).
3890/// Lower edge of the data-derived kernel-range window, as a fraction of the
3891/// maximum pairwise distance `r_max`: length scales below `2/r_max` resolve
3892/// structure finer than the closest center pair, so the kernel range floor is
3893/// set at twice the maximum spacing.
3894pub const KERNEL_RANGE_MIN_DIAMETER_FRACTION: f64 = 2.0;
3895
3896/// Upper edge of the data-derived kernel-range window, as a multiple of the
3897/// minimum pairwise distance `r_min`: beyond `100/r_min` the radial columns go
3898/// nearly collinear with the polynomial nullspace, so the kernel range is
3899/// capped here to keep the basis geometry well-conditioned.
3900pub const KERNEL_RANGE_MAX_SPACING_MULTIPLE: f64 = 1e2;
3901
3902fn spatial_term_stored_input_scale(term: &SmoothTermSpec) -> Option<crate::IsotropicScale> {
3903    match &term.basis {
3904        SmoothBasisSpec::ThinPlate { input_scale, .. }
3905        | SmoothBasisSpec::Matern { input_scale, .. }
3906        | SmoothBasisSpec::Duchon { input_scale, .. } => *input_scale,
3907        _ => None,
3908    }
3909}
3910
3911fn spatial_term_realized_input_scale(
3912    data: ArrayView2<'_, f64>,
3913    term: &SmoothTermSpec,
3914) -> Result<crate::IsotropicScale, BasisError> {
3915    let (feature_cols, stored) = match &term.basis {
3916        SmoothBasisSpec::ThinPlate {
3917            feature_cols,
3918            input_scale,
3919            ..
3920        }
3921        | SmoothBasisSpec::Matern {
3922            feature_cols,
3923            input_scale,
3924            ..
3925        }
3926        | SmoothBasisSpec::Duchon {
3927            feature_cols,
3928            input_scale,
3929            ..
3930        } => (feature_cols, input_scale),
3931        _ => {
3932            return Err(BasisError::InvalidInput(format!(
3933                "term '{}' does not have an isotropic Euclidean input frame",
3934                term.name
3935            )));
3936        }
3937    };
3938    if let Some(scale) = stored {
3939        return Ok(*scale);
3940    }
3941    let x = select_columns(data, feature_cols)?;
3942    estimate_isotropic_scale(x.view())
3943}
3944
3945/// Returns ψ-space bounds (ψ_lo = ln(κ_lo), ψ_hi = ln(κ_hi)).
3946///
3947/// The returned window is intersected with the options window so user-set
3948/// `min_length_scale` / `max_length_scale` remain hard limits. Degenerate
3949/// geometry or an empty intersection is a typed error: changing to a generic
3950/// options window would silently optimize in a different coordinate chart.
3951pub fn spatial_term_psi_bounds(
3952    data: ArrayView2<'_, f64>,
3953    spec: &TermCollectionSpec,
3954    term_idx: usize,
3955    options: &SpatialLengthScaleOptimizationOptions,
3956) -> Result<(f64, f64), BasisError> {
3957    let options_window = (
3958        -options.max_length_scale.ln(),
3959        -options.min_length_scale.ln(),
3960    );
3961    // Constant-curvature: the ψ coordinate is the raw signed κ, so its window is
3962    // the chart-feasible κ bracket, NOT a log-ℓ window. Mirrors the aniso bounds
3963    // path's `constant_curvature_kappa_bounds` branch so the isotropic
3964    // (non-aniso) seed clamp projects κ into the right interval.
3965    if constant_curvature_term_spec(spec, term_idx).is_some() {
3966        return Ok(constant_curvature_kappa_bounds(data, spec, term_idx));
3967    }
3968    let term = spec.smooth_terms.get(term_idx).ok_or_else(|| {
3969        BasisError::InvalidInput(format!(
3970            "spatial term index {term_idx} is out of bounds for {} smooth terms",
3971            spec.smooth_terms.len()
3972        ))
3973    })?;
3974    // Prefer resolved centers (post-fit) since they live in the same standardized
3975    // space the kernel actually sees. Centers are capped at `default_num_centers`
3976    // (<=2000), so exact pairwise bounds are cheap (<4M ops). If centers are
3977    // not yet UserProvided, fall back to the standardized feature data columns
3978    // with the capped-sample path (O(K²·d), K=1024) — the sample is
3979    // conservative for κ bounds (see `pairwise_distance_bounds_sampled`
3980    // docs): it never excludes a feasible κ the exact method would include.
3981    //
3982    // Under anisotropy the kernel metric is y-space (y_a = exp(η_a) x_a),
3983    // so r_min/r_max must be y-space distances. This matters only when the
3984    // spec already carries calibrated η_a at setup time (e.g., warm-start
3985    // or refit paths); for fresh optimization η_a starts at 0 and y = x.
3986    let aniso = get_spatial_aniso_log_scales(spec, term_idx);
3987    let stored_input_scale = spatial_term_stored_input_scale(term);
3988    let input_scale = spatial_term_realized_input_scale(data, term)?;
3989    let r_bounds = match spatial_term_center_strategy(term) {
3990        Some(CenterStrategy::UserProvided(centers)) if centers.nrows() >= 2 => {
3991            let mut centers_in_frame = centers.clone();
3992            if stored_input_scale.is_none() {
3993                input_scale.standardize(&mut centers_in_frame);
3994            }
3995            let bounds = match aniso.as_deref() {
3996                Some(eta) if eta.len() == centers_in_frame.ncols() => {
3997                    let y = points_in_aniso_y_space(centers_in_frame.view(), eta);
3998                    pairwise_distance_bounds(y.view())
3999                }
4000                _ => pairwise_distance_bounds(centers_in_frame.view()),
4001            };
4002            bounds
4003        }
4004        _ => {
4005            let x = standardized_spatial_term_data(data, term)?;
4006            match aniso.as_deref() {
4007                Some(eta) if eta.len() == x.ncols() => {
4008                    let y = points_in_aniso_y_space(x.view(), eta);
4009                    pairwise_distance_bounds_sampled(y.view())
4010                }
4011                _ => pairwise_distance_bounds_sampled(x.view()),
4012            }
4013        }
4014    };
4015    let (r_min, r_max) = r_bounds.ok_or_else(|| {
4016        BasisError::InvalidInput(format!(
4017            "term '{}' has no positive finite pairwise-distance range",
4018            term.name
4019        ))
4020    })?;
4021    // Length scales substantially larger than the data diameter make radial
4022    // TPS/Matern columns nearly collinear with their polynomial nullspace.
4023    // The nullspace already carries constant/linear low-frequency structure,
4024    // so cap the kernel range at the diameter scale instead of letting the
4025    // optimizer enter a numerically degenerate basis geometry.
4026    // `r_min`/`r_max` are measured in the standardized kernel frame, where
4027    // ℓ_eff = ℓ_original / σ_geom. The optimizer/spec ψ coordinate is
4028    // ψ_original = log(1/ℓ_original), hence
4029    //
4030    //   κ_original = κ_eff / σ_geom
4031    //              = κ_eff * compensate_length_scale(1, scales).
4032    //
4033    // Convert exactly once here before intersecting the data window with the
4034    // original-coordinate user options. Previously these standardized κ bounds
4035    // were written directly into the spec; the basis builder then divided ℓ by
4036    // σ_geom again, making the realized endpoint too long by 1/σ_geom.
4037    let inverse_sigma = input_scale.reciprocal();
4038    let psi_chart_offset = inverse_sigma.ln();
4039    let psi_lo_data = (KERNEL_RANGE_MIN_DIAMETER_FRACTION / r_max).ln() + psi_chart_offset;
4040    let psi_hi_data = (KERNEL_RANGE_MAX_SPACING_MULTIPLE / r_min).ln() + psi_chart_offset;
4041    // #1074: the Matérn-specific length-scale ceiling that used to live here was
4042    // deleted. It was masking, not fixing, the real defect: a hard upper bound on
4043    // the kernel range that pinned the κ-optimizer short rather than letting the
4044    // optimizer find the REML optimum. Matérn now shares the same generic geometry
4045    // window as Duchon / TPS (`KERNEL_RANGE_MIN_DIAMETER_FRACTION / r_max` floor,
4046    // `KERNEL_RANGE_MAX_SPACING_MULTIPLE / r_min` ceiling); the #1357 fully-flat
4047    // collapse corner is guarded by the EDF-collapse guard in
4048    // `spatial_optimization.rs`, which acts on the realized fit, not on a clamp.
4049    // Intersect with the options window so min/max_length_scale remain hard caps.
4050    let psi_lo = psi_lo_data.max(options_window.0);
4051    let psi_hi = psi_hi_data.min(options_window.1);
4052    if psi_lo >= psi_hi {
4053        return Err(BasisError::InvalidInput(format!(
4054            "term '{}' has an empty spatial ψ window after intersecting data bounds [{psi_lo_data}, {psi_hi_data}] with configured bounds [{}, {}]",
4055            term.name, options_window.0, options_window.1
4056        )));
4057    }
4058    Ok((psi_lo, psi_hi))
4059}
4060
4061/// The ψ box the κ optimizer SEARCHES, as opposed to the data-geometry window
4062/// [`spatial_term_psi_bounds`] describes.
4063///
4064/// # Why these are two objects and not one (#2454)
4065///
4066/// [`spatial_term_psi_bounds`] answers a question about the DATA: over what
4067/// kernel ranges does this point cloud keep a well-conditioned radial basis? It
4068/// is a pure function of the geometry, equivariant under rotation and dilation
4069/// of the inputs, and it knows nothing about which length scale anyone has
4070/// actually used.
4071///
4072/// A search box has to answer a second question the geometry cannot: is the
4073/// point the search STARTS at, and whose fit its answer will be GRADED against,
4074/// inside the set being searched? The term's own resolved `length_scale` is an
4075/// incumbent, not a hypothesis — a design was built and a fit converged at it,
4076/// which is direct evidence that the geometry is admissible there, and evidence
4077/// outranks a heuristic that was trying to predict admissibility.
4078///
4079/// Excluding the incumbent is not the conservative choice; it silently poses a
4080/// different problem. Every consumer seeds at the incumbent,
4081/// [`SpatialLogKappaCoords::clamp_to_bounds`] projects that seed onto the
4082/// window's edge, and the joint optimizer's answer is then compared against a
4083/// baseline fit at the point the projection discarded. Measured on the fixture
4084/// #2454 was opened against — `length_scale = 12` against a geometry window
4085/// `[e^-5.719, e^+0.0118]` — ψ is projected from −2.4849 onto the bound
4086/// −0.011839795303285494, the optimizer starts THERE, its gradient points into
4087/// the excluded region at every iteration, and it terminates on that bound with
4088/// a criterion 14.7 nats WORSE than the baseline (−53.75 against −68.41). The
4089/// route then reports the monotonicity failure as a solver failure when it is a
4090/// feasible-set failure: `min` over a set that excludes the incumbent is under
4091/// no obligation to beat the incumbent.
4092///
4093/// So the box is the geometry window WIDENED — never narrowed — to contain the
4094/// incumbent. `min_length_scale` / `max_length_scale` stay hard caps: those are
4095/// the caller's own explicit constraint, and an incumbent outside them is a
4096/// contradiction the caller stated, not one this function invented.
4097pub fn spatial_term_psi_search_box(
4098    data: ArrayView2<'_, f64>,
4099    spec: &TermCollectionSpec,
4100    term_idx: usize,
4101    options: &SpatialLengthScaleOptimizationOptions,
4102) -> Result<(f64, f64), BasisError> {
4103    let (mut psi_lo, mut psi_hi) = spatial_term_psi_bounds(data, spec, term_idx, options)?;
4104    // Constant-curvature terms carry a signed-κ chart, not a log-ℓ chart, so
4105    // `-ln(length_scale)` is not their coordinate and the geometry bracket is
4106    // already the feasible set. Leave that box exactly as it was.
4107    if constant_curvature_term_spec(spec, term_idx).is_some() {
4108        return Ok((psi_lo, psi_hi));
4109    }
4110    let options_window = (
4111        -options.max_length_scale.ln(),
4112        -options.min_length_scale.ln(),
4113    );
4114    if let Some(length_scale) = get_spatial_length_scale(spec, term_idx)
4115        && length_scale.is_finite()
4116        && length_scale > 0.0
4117    {
4118        let psi_incumbent = -length_scale.ln();
4119        if psi_incumbent.is_finite() {
4120            psi_lo = psi_lo.min(psi_incumbent.max(options_window.0));
4121            psi_hi = psi_hi.max(psi_incumbent.min(options_window.1));
4122        }
4123    }
4124    Ok((psi_lo, psi_hi))
4125}
4126
4127#[cfg(test)]
4128mod spatial_psi_bound_coordinate_tests {
4129    use super::*;
4130    use crate::basis::{MaternIdentifiability, MaternNu};
4131    use ndarray::array;
4132
4133    fn frozen_matern_bounds(theta: f64, dilation: f64) -> (f64, f64) {
4134        let source = array![
4135            [-1.7, -0.4],
4136            [-1.1, 0.8],
4137            [-0.2, -1.3],
4138            [0.5, 1.6],
4139            [1.4, -0.7],
4140            [2.1, 0.5],
4141        ];
4142        let (cos_theta, sin_theta) = (theta.cos(), theta.sin());
4143        let mut data = Array2::<f64>::zeros(source.raw_dim());
4144        for row in 0..source.nrows() {
4145            let x = source[[row, 0]];
4146            let y = source[[row, 1]];
4147            data[[row, 0]] = dilation * (cos_theta * x - sin_theta * y);
4148            data[[row, 1]] = dilation * (sin_theta * x + cos_theta * y);
4149        }
4150        let input_scale = estimate_isotropic_scale(data.view()).expect("isotropic input scale");
4151        let mut centers = data.clone();
4152        input_scale.standardize(&mut centers);
4153        let spec = TermCollectionSpec {
4154            linear_terms: Vec::new(),
4155            random_effect_terms: Vec::new(),
4156            smooth_terms: vec![SmoothTermSpec {
4157            frozen_parametric_residualization: None,
4158                name: "matern".to_string(),
4159                basis: SmoothBasisSpec::Matern {
4160                    feature_cols: vec![0, 1],
4161                    spec: MaternBasisSpec {
4162                        periodic: None,
4163                        center_strategy: CenterStrategy::UserProvided(centers),
4164                        length_scale: crate::basis::MaternLengthScale::fixed(1.0),
4165                        nu: MaternNu::FiveHalves,
4166                        include_intercept: false,
4167                        double_penalty: true,
4168                        identifiability: MaternIdentifiability::CenterSumToZero,
4169                        aniso_log_scales: None,
4170                    },
4171                    input_scale: Some(input_scale),
4172                },
4173                shape: ShapeConstraint::None,
4174                joint_null_rotation: None,
4175            }],
4176        };
4177        spatial_term_psi_bounds(
4178            data.view(),
4179            &spec,
4180            0,
4181            &SpatialLengthScaleOptimizationOptions::default(),
4182        )
4183        .expect("finite spatial ψ bounds")
4184    }
4185
4186    fn assert_close(left: f64, right: f64) {
4187        assert!(
4188            (left - right).abs() <= 1e-12,
4189            "coordinate-equivalent bounds differ: left={left:.16e}, right={right:.16e}"
4190        );
4191    }
4192
4193    /// The search box a κ optimizer is handed must contain the length scale it
4194    /// is seeded at and graded against (#2454).
4195    ///
4196    /// Stated as containment rather than as a numeric window, because the point
4197    /// is not where the edge lands — it is that `clamp_to_bounds` has nothing to
4198    /// do. A window that excludes the incumbent makes `min` over the box free to
4199    /// return something strictly worse than the incumbent, which is exactly the
4200    /// "optimizing κ made the fit worse" refusal the monotone fixtures reported
4201    /// as a solver failure.
4202    ///
4203    /// Both directions are pinned: an incumbent far OUTSIDE the geometry window
4204    /// must be inside the search box, and an incumbent inside it must not move
4205    /// the box at all (widened, never narrowed, and never gratuitously).
4206    #[test]
4207    fn psi_search_box_contains_the_incumbent_length_scale_2454() {
4208        let source = array![
4209            [-1.7, -0.4],
4210            [-1.1, 0.8],
4211            [-0.2, -1.3],
4212            [0.5, 1.6],
4213            [1.4, -0.7],
4214            [2.1, 0.5],
4215        ];
4216        let options = SpatialLengthScaleOptimizationOptions::default();
4217        let box_for = |length_scale: f64| -> ((f64, f64), (f64, f64)) {
4218            let input_scale =
4219                estimate_isotropic_scale(source.view()).expect("isotropic input scale");
4220            let mut centers = source.clone();
4221            input_scale.standardize(&mut centers);
4222            let spec = TermCollectionSpec {
4223                linear_terms: Vec::new(),
4224                random_effect_terms: Vec::new(),
4225                smooth_terms: vec![SmoothTermSpec {
4226            frozen_parametric_residualization: None,
4227                    name: "matern".to_string(),
4228                    basis: SmoothBasisSpec::Matern {
4229                        feature_cols: vec![0, 1],
4230                        spec: MaternBasisSpec {
4231                            periodic: None,
4232                            center_strategy: CenterStrategy::UserProvided(centers),
4233                            length_scale: crate::basis::MaternLengthScale::fixed(length_scale),
4234                            nu: MaternNu::FiveHalves,
4235                            include_intercept: false,
4236                            double_penalty: true,
4237                            identifiability: MaternIdentifiability::CenterSumToZero,
4238                            aniso_log_scales: None,
4239                        },
4240                        input_scale: Some(input_scale),
4241                    },
4242                    shape: ShapeConstraint::None,
4243                    joint_null_rotation: None,
4244                }],
4245            };
4246            let geometry = spatial_term_psi_bounds(source.view(), &spec, 0, &options)
4247                .expect("finite geometry window");
4248            let search = spatial_term_psi_search_box(source.view(), &spec, 0, &options)
4249                .expect("finite search box");
4250            (geometry, search)
4251        };
4252
4253        // An incumbent far past the long-range edge of the geometry window —
4254        // #2454's fixture shape, where `length_scale = 12` sits about six data
4255        // diameters out.
4256        let far = 1.0e3_f64;
4257        let (geometry, search) = box_for(far);
4258        let psi_far = -far.ln();
4259        assert!(
4260            psi_far < geometry.0,
4261            "fixture must place the incumbent OUTSIDE the geometry window, got \
4262             psi={psi_far} against [{}, {}]",
4263            geometry.0,
4264            geometry.1
4265        );
4266        assert!(
4267            search.0 <= psi_far && psi_far <= search.1,
4268            "the search box [{}, {}] must contain the incumbent psi={psi_far}; a seed \
4269             the box excludes is projected onto its edge and the optimum is then taken \
4270             over a set that does not contain the point it is graded against (#2454)",
4271            search.0,
4272            search.1
4273        );
4274        assert!(
4275            search.1 == geometry.1 && search.0 <= geometry.0,
4276            "the search box must be the geometry window WIDENED, never narrowed: \
4277             geometry=[{}, {}] search=[{}, {}]",
4278            geometry.0,
4279            geometry.1,
4280            search.0,
4281            search.1
4282        );
4283
4284        // An incumbent already inside the window must leave the box untouched.
4285        let (geometry_mid, search_mid) = box_for((-0.5 * (geometry.0 + geometry.1)).exp());
4286        assert!(
4287            search_mid == geometry_mid,
4288            "an incumbent inside the geometry window must not move the search box: \
4289             geometry=[{}, {}] search=[{}, {}]",
4290            geometry_mid.0,
4291            geometry_mid.1,
4292            search_mid.0,
4293            search_mid.1
4294        );
4295    }
4296
4297    #[test]
4298    fn standardized_center_bounds_return_to_original_units_under_rotation_and_scaling() {
4299        let base = frozen_matern_bounds(0.0, 1.0);
4300        let rotated = frozen_matern_bounds(0.61, 1.0);
4301        assert_close(rotated.0, base.0);
4302        assert_close(rotated.1, base.1);
4303
4304        let dilation = 4.0_f64;
4305        let rotated_scaled = frozen_matern_bounds(0.61, dilation);
4306        let expected_shift = dilation.ln();
4307        assert_close(rotated_scaled.0, base.0 - expected_shift);
4308        assert_close(rotated_scaled.1, base.1 - expected_shift);
4309    }
4310}
4311
4312/// Data-derived ψ seed for a spatial term when the user has not set an
4313/// explicit length_scale on its basis spec. Uses the geometric mean of the
4314/// data-informed kappa range (i.e., the midpoint of the ψ window).
4315pub fn spatial_term_psi_seed(
4316    data: ArrayView2<'_, f64>,
4317    spec: &TermCollectionSpec,
4318    term_idx: usize,
4319    options: &SpatialLengthScaleOptimizationOptions,
4320) -> Result<Option<f64>, BasisError> {
4321    if get_spatial_length_scale(spec, term_idx).is_some() {
4322        return Ok(None); // user/spec-provided length_scale wins
4323    }
4324    let (psi_lo, psi_hi) = spatial_term_psi_bounds(data, spec, term_idx, options)?;
4325    Ok(Some(0.5 * (psi_lo + psi_hi)))
4326}
4327
4328pub fn spatial_term_psi_to_length_scale_and_aniso(psi: &[f64]) -> (Option<f64>, Option<Vec<f64>>) {
4329    if psi.len() <= 1 {
4330        (Some((-psi.first().copied().unwrap_or(0.0)).exp()), None)
4331    } else {
4332        let psi_bar = psi.iter().sum::<f64>() / psi.len() as f64;
4333        (
4334            Some((-psi_bar).exp()),
4335            Some(psi.iter().map(|&value| value - psi_bar).collect()),
4336        )
4337    }
4338}
4339
4340/// Get the `aniso_log_scales` from a spatial term, if present.
4341pub fn get_spatial_aniso_log_scales(
4342    spec: &TermCollectionSpec,
4343    term_idx: usize,
4344) -> Option<Vec<f64>> {
4345    spec.smooth_terms
4346        .get(term_idx)
4347        .and_then(|term| match &term.basis {
4348            SmoothBasisSpec::Matern { spec, .. } => spec.aniso_log_scales.clone(),
4349            SmoothBasisSpec::Duchon { spec, .. } => spec.aniso_log_scales.clone(),
4350            _ => None,
4351        })
4352}
4353
4354/// Per-axis response-structure score for anisotropy seeding.
4355///
4356/// For each spatial axis `a`, sort the response `y` by the axis coordinate
4357/// `x_a` and measure the total squared successive variation of the sorted
4358/// response, `tv_a = Σ_i (y_{σ(i+1)} − y_{σ(i)})²` where `σ` orders rows by
4359/// `x_a`. An axis that carries real (possibly nonlinear) signal makes `y` vary
4360/// SMOOTHLY when the rows are walked in that axis's order, so `tv_a` is SMALL;
4361/// a pure-nuisance axis leaves `y` looking unordered, so `tv_a` is LARGE.
4362///
4363/// This deliberately does NOT use a linear correlation `corr(x_a, y)`: for an
4364/// odd, symmetric signal such as `sin(2·x1)` over a symmetric domain the linear
4365/// correlation is ~0 on the *signal* axis, which would misdirect the seed. The
4366/// total-variation-of-sorted-response score captures nonlinear association.
4367///
4368/// Returns `score_a = −½·ln(tv_a + ε)` (larger ⇒ more signal on axis `a`),
4369/// centered to sum to zero, or `None` when the data is degenerate (too few
4370/// rows, non-finite, or all axes equally (un)structured). The caller adds a
4371/// BOUNDED multiple of this to the geometry seed — it is a conservative nudge,
4372/// never a hard override.
4373pub fn response_aware_axis_contrasts(
4374    x: ndarray::ArrayView2<'_, f64>,
4375    y: ndarray::ArrayView1<'_, f64>,
4376) -> Option<Vec<f64>> {
4377    let n = x.nrows();
4378    let d = x.ncols();
4379    if d <= 1 || n < 4 || y.len() != n {
4380        return None;
4381    }
4382    if x.iter().any(|v| !v.is_finite()) || y.iter().any(|v| !v.is_finite()) {
4383        return None;
4384    }
4385    let mut scores = Vec::with_capacity(d);
4386    for a in 0..d {
4387        let mut order: Vec<usize> = (0..n).collect();
4388        let col = x.column(a);
4389        order.sort_by(|&i, &j| {
4390            col[i]
4391                .partial_cmp(&col[j])
4392                .unwrap_or(std::cmp::Ordering::Equal)
4393        });
4394        let mut tv = 0.0_f64;
4395        for w in order.windows(2) {
4396            let diff = y[w[1]] - y[w[0]];
4397            tv += diff * diff;
4398        }
4399        // ε guards against ln(0) on a perfectly flat / constant response.
4400        scores.push(-0.5 * (tv + 1e-12).ln());
4401    }
4402    if scores.iter().any(|v| !v.is_finite()) {
4403        return None;
4404    }
4405    let mean = scores.iter().sum::<f64>() / d as f64;
4406    let centered: Vec<f64> = scores.iter().map(|&s| s - mean).collect();
4407    // If every axis is equally structured the centered scores are ~0 and the
4408    // nudge is a no-op — return None so the geometry seed is used unchanged.
4409    if centered.iter().all(|&v| v.abs() < 1e-9) {
4410        return None;
4411    }
4412    Some(centered)
4413}
4414
4415/// Conservative, response-aware anisotropy seed nudge applied before the κ outer
4416/// loop. For each anisotropic spatial term it adds a BOUNDED multiple of the
4417/// per-axis response-structure contrast (`response_aware_axis_contrasts`) on top
4418/// of the existing geometry seed, so the optimizer starts in the correct basin
4419/// instead of at a response-blind near-symmetric point (the #1376 under-recovery
4420/// where a signal axis and a nuisance axis with equal coordinate spread seed to
4421/// ~\[0,0\]). The nudge is clamped to keep this a perturbation, never a hard
4422/// override, so shared aniso Matérn/Duchon fits cannot be destabilized by it.
4423pub fn apply_response_aware_anisotropy_seed(
4424    data: ArrayView2<'_, f64>,
4425    y: ndarray::ArrayView1<'_, f64>,
4426    spec: &mut TermCollectionSpec,
4427    spatial_terms: &[usize],
4428) {
4429    // Bound on the per-axis contrast nudge (in η units). One LN_2 ≈ 0.69 halves
4430    // the effective per-axis length scale; capping at LN_2 keeps the seed within
4431    // one optimizer log-step of the geometry seed while still breaking the
4432    // symmetric-seed trap.
4433    const MAX_NUDGE: f64 = std::f64::consts::LN_2;
4434    for &term_idx in spatial_terms {
4435        let Some(current_eta) = get_spatial_aniso_log_scales(spec, term_idx) else {
4436            continue;
4437        };
4438        let d = current_eta.len();
4439        if d <= 1 {
4440            continue;
4441        }
4442        let Some(term) = spec.smooth_terms.get(term_idx) else {
4443            continue;
4444        };
4445        let feature_cols = term.basis.structural_feature_cols();
4446        if feature_cols.len() != d {
4447            continue;
4448        }
4449        let Ok(x) = select_columns(data, &feature_cols) else {
4450            continue;
4451        };
4452        let Some(contrast) = response_aware_axis_contrasts(x.view(), y) else {
4453            continue;
4454        };
4455        let nudged: Vec<f64> = current_eta
4456            .iter()
4457            .zip(contrast.iter())
4458            .map(|(&eta_a, &c_a)| eta_a + c_a.clamp(-MAX_NUDGE, MAX_NUDGE))
4459            .collect();
4460        // `set_spatial_aniso_log_scales` re-centers to Σ η = 0. A term that does
4461        // not support aniso scales is silently skipped (the seed is optional).
4462        if let Err(err) = set_spatial_aniso_log_scales(spec, term_idx, nudged) {
4463            log::debug!(
4464                "[spatial-kappa] response-aware anisotropy seed skipped for term {term_idx}: {err}"
4465            );
4466        }
4467    }
4468}
4469
4470/// Get the number of feature columns (spatial dimensionality) for a spatial term.
4471pub fn get_spatial_feature_dim(spec: &TermCollectionSpec, term_idx: usize) -> Option<usize> {
4472    spec.smooth_terms
4473        .get(term_idx)
4474        .and_then(|term| match &term.basis {
4475            SmoothBasisSpec::ThinPlate { feature_cols, .. } => Some(feature_cols.len()),
4476            SmoothBasisSpec::Matern { feature_cols, .. } => Some(feature_cols.len()),
4477            SmoothBasisSpec::Duchon { feature_cols, .. } => Some(feature_cols.len()),
4478            _ => None,
4479        })
4480}
4481
4482/// Log the learned per-axis spatial anisotropy for all spatial terms that
4483/// have `aniso_log_scales` set after optimization.
4484///
4485/// For scalar-scale families this reports eta, effective per-axis length
4486/// scales, and per-axis kappa values. For pure Duchon it reports the centered
4487/// eta contrasts only.
4488pub fn log_spatial_aniso_scales(spec: &TermCollectionSpec) {
4489    for (term_idx, term) in spec.smooth_terms.iter().enumerate() {
4490        let (aniso, length_scale) = match &term.basis {
4491            SmoothBasisSpec::Matern { spec, .. } => {
4492                (spec.aniso_log_scales.as_ref(), spec.length_scale.resolved())
4493            }
4494            SmoothBasisSpec::Duchon { spec, .. } => {
4495                (spec.aniso_log_scales.as_ref(), spec.length_scale)
4496            }
4497            _ => (None, None),
4498        };
4499        let Some(eta) = aniso else { continue };
4500        if eta.is_empty() {
4501            continue;
4502        }
4503        let mut lines = match length_scale {
4504            Some(ls) => format!(
4505                "[spatial-kappa] term {} (\"{}\"): anisotropic length scales optimized (global length_scale={:.4})",
4506                term_idx, term.name, ls
4507            ),
4508            None => format!(
4509                "[spatial-kappa] term {} (\"{}\"): pure Duchon shape anisotropy optimized",
4510                term_idx, term.name
4511            ),
4512        };
4513        for (a, &eta_a) in eta.iter().enumerate() {
4514            if let Some(ls) = length_scale {
4515                let length_a = ls * (-eta_a).exp();
4516                let kappa_a = (1.0 / ls) * eta_a.exp();
4517                lines.push_str(&format!(
4518                    "\n  axis {}: eta={:+.4}, length={:.4}, kappa={:.4}",
4519                    a, eta_a, length_a, kappa_a
4520                ));
4521            } else {
4522                lines.push_str(&format!("\n  axis {}: eta={:+.4}", a, eta_a));
4523            }
4524        }
4525        log::info!("{}", lines);
4526    }
4527}
4528
4529/// Set `aniso_log_scales` on a spatial term's basis spec.
4530pub fn set_spatial_aniso_log_scales(
4531    spec: &mut TermCollectionSpec,
4532    term_idx: usize,
4533    eta: Vec<f64>,
4534) -> Result<(), EstimationError> {
4535    let eta = center_aniso_log_scales(&eta);
4536    let Some(term) = spec.smooth_terms.get_mut(term_idx) else {
4537        crate::bail_invalid_estim!("spatial aniso_log_scales term index {term_idx} out of range");
4538    };
4539    match &mut term.basis {
4540        SmoothBasisSpec::Matern { spec, .. } => {
4541            spec.aniso_log_scales = Some(eta);
4542            Ok(())
4543        }
4544        SmoothBasisSpec::Duchon { spec, .. } => {
4545            spec.aniso_log_scales = Some(eta);
4546            Ok(())
4547        }
4548        _ => Err(EstimationError::InvalidInput(format!(
4549            "term '{}' does not support aniso_log_scales",
4550            term.name
4551        ))),
4552    }
4553}
4554
4555/// Sync knot-cloud-derived anisotropy contrasts from basis metadata back into
4556/// the mutable spec so the optimizer starts from the correct eta values.
4557///
4558/// Call this after building the smooth design but before initializing the
4559/// optimizer's psi coordinates. For each spatial term whose metadata contains
4560/// computed `aniso_log_scales`, this writes them into the spec.
4561pub fn sync_aniso_contrasts_from_metadata(spec: &mut TermCollectionSpec, design: &SmoothDesign) {
4562    for (term_idx, term) in design.terms.iter().enumerate() {
4563        let meta_aniso = match &term.metadata {
4564            BasisMetadata::Matern {
4565                aniso_log_scales, ..
4566            } => aniso_log_scales.clone(),
4567            BasisMetadata::Duchon {
4568                aniso_log_scales, ..
4569            } => aniso_log_scales.clone(),
4570            _ => None,
4571        };
4572        if let Some(eta) = meta_aniso
4573            && eta.len() > 1
4574        {
4575            if let Err(err) = set_spatial_aniso_log_scales(spec, term_idx, eta) {
4576                log::debug!(
4577                    "term {term_idx}: anisotropic log-scale sync skipped, keeping the existing scales: {err}"
4578                );
4579            }
4580        }
4581    }
4582}
4583
4584#[derive(Debug, Clone)]
4585pub struct SpatialLengthScaleOptimizationOptions {
4586    /// Enable outer-loop optimization over spatial κ (= 1 / length_scale)
4587    /// for supported radial-kernel smooths.
4588    /// This applies to ThinPlate, Matérn, and Duchon terms.
4589    pub enabled: bool,
4590    /// Maximum number of outer iterations in the exact joint [rho, psi] solve.
4591    pub max_outer_iter: usize,
4592    /// Relative improvement threshold for terminating the outer solve.
4593    pub rel_tol: f64,
4594    /// Initial log(length_scale) perturbation used for seed construction.
4595    pub log_step: f64,
4596    /// Minimum allowed length_scale during κ search.
4597    pub min_length_scale: f64,
4598    /// Maximum allowed length_scale during κ search.
4599    pub max_length_scale: f64,
4600    /// Automatic geometry-initializer threshold for large-scale spatial fits.
4601    ///
4602    /// When n exceeds twice this value, the fitter uses a spatially stratified
4603    /// subsample only to seed κ/anisotropy geometry: centers are resolved,
4604    /// axis contrasts are initialized from center/data spread, and one or two
4605    /// cheap ψ reseeding updates are applied. It never runs PIRLS, REML, ARC,
4606    /// BFGS, or any recursive optimizer on the pilot.
4607    ///
4608    /// The final coefficients, smoothing parameters, and spatial geometry are
4609    /// always optimized on the full dataset.
4610    ///
4611    /// Set to 0 to skip the pilot geometry initializer.
4612    pub pilot_subsample_threshold: usize,
4613}
4614
4615impl Default for SpatialLengthScaleOptimizationOptions {
4616    fn default() -> Self {
4617        Self {
4618            enabled: true,
4619            max_outer_iter: 80,
4620            rel_tol: 1e-4,
4621            log_step: std::f64::consts::LN_2,
4622            min_length_scale: 1e-3,
4623            max_length_scale: 1e3,
4624            pilot_subsample_threshold: 10_000,
4625        }
4626    }
4627}
4628
4629impl SpatialLengthScaleOptimizationOptions {
4630    /// Validate the struct's invariants. Callers that construct these options
4631    /// from external input (CLI, config, Python API) should call this before
4632    /// passing the options into the fitter. Returns `Err` with a descriptive
4633    /// message when an invariant is violated; the fitter then panics or
4634    /// returns `EstimationError` at its own boundary.
4635    ///
4636    /// Invariants:
4637    ///   * `min_length_scale > 0`, finite
4638    ///   * `max_length_scale > 0`, finite
4639    ///   * `min_length_scale < max_length_scale`
4640    ///   * `rel_tol > 0`, finite
4641    ///   * `log_step > 0`, finite
4642    ///
4643    /// These invariants are what the downstream κ-bound and ψ-window code
4644    /// assumes (`-log(max_ls)` must be finite, `(min,max)` must not be
4645    /// inverted, etc.). Without validation, invalid options produce silent
4646    /// NaN-propagation inside the outer optimizer.
4647    pub fn validate(&self) -> Result<(), String> {
4648        if !self.min_length_scale.is_finite() || self.min_length_scale <= 0.0 {
4649            return Err(SmoothError::invalid_config(format!(
4650                "SpatialLengthScaleOptimizationOptions::min_length_scale must be > 0 and finite, got {}",
4651                self.min_length_scale
4652            ))
4653            .into());
4654        }
4655        if !self.max_length_scale.is_finite() || self.max_length_scale <= 0.0 {
4656            return Err(SmoothError::invalid_config(format!(
4657                "SpatialLengthScaleOptimizationOptions::max_length_scale must be > 0 and finite, got {}",
4658                self.max_length_scale
4659            ))
4660            .into());
4661        }
4662        if self.min_length_scale >= self.max_length_scale {
4663            return Err(SmoothError::invalid_config(format!(
4664                "SpatialLengthScaleOptimizationOptions requires min_length_scale < max_length_scale, got min={} max={}",
4665                self.min_length_scale, self.max_length_scale
4666            ))
4667            .into());
4668        }
4669        if !self.rel_tol.is_finite() || self.rel_tol <= 0.0 {
4670            return Err(SmoothError::invalid_config(format!(
4671                "SpatialLengthScaleOptimizationOptions::rel_tol must be > 0 and finite, got {}",
4672                self.rel_tol
4673            ))
4674            .into());
4675        }
4676        if !self.log_step.is_finite() || self.log_step <= 0.0 {
4677            return Err(SmoothError::invalid_config(format!(
4678                "SpatialLengthScaleOptimizationOptions::log_step must be > 0 and finite, got {}",
4679                self.log_step
4680            ))
4681            .into());
4682        }
4683        Ok(())
4684    }
4685}
4686
4687#[derive(Debug, Clone)]
4688pub struct RandomEffectBlock {
4689    pub name: String,
4690    /// O(n) group-label vector: group_ids\[i\] = column index in [0, num_groups).
4691    /// `None` if the observation's level is not in the kept set.
4692    pub group_ids: Vec<Option<usize>>,
4693    pub num_groups: usize,
4694    pub kept_levels: Vec<u64>,
4695}
4696
4697pub const BLOCK_SPARSE_ZERO_EPS: f64 = 1e-12;
4698
4699pub const BLOCK_SPARSE_MAX_DENSITY: f64 = 0.20;
4700
4701pub fn blocks_have_intrinsic_sparse_structure(blocks: &[DesignBlock]) -> bool {
4702    blocks
4703        .iter()
4704        .any(|block| matches!(block, DesignBlock::Sparse(_) | DesignBlock::RandomEffect(_)))
4705}
4706
4707pub fn sparse_compatible_block_nnz(block: &DesignBlock) -> Option<usize> {
4708    match block {
4709        DesignBlock::Intercept(n) => Some(*n),
4710        DesignBlock::RandomEffect(op) => {
4711            Some(op.group_ids.iter().filter(|gid| gid.is_some()).count())
4712        }
4713        DesignBlock::Sparse(sparse) => Some(sparse.val().len()),
4714        DesignBlock::Dense(dense) => dense.as_dense_ref().map(|matrix| {
4715            matrix
4716                .iter()
4717                .filter(|&&value| value.abs() > BLOCK_SPARSE_ZERO_EPS)
4718                .count()
4719        }),
4720    }
4721}
4722
4723pub fn try_build_sparse_design_from_blocks(
4724    blocks: &[DesignBlock],
4725) -> Result<Option<DesignMatrix>, BasisError> {
4726    if blocks.is_empty() {
4727        return Ok(None);
4728    }
4729    let nrows = blocks[0].nrows();
4730    let ncols: usize = blocks.iter().map(DesignBlock::ncols).sum();
4731    if nrows == 0 || ncols == 0 || ncols <= 32 {
4732        return Ok(None);
4733    }
4734
4735    let preserve_sparse_storage = blocks_have_intrinsic_sparse_structure(blocks);
4736    let sparse_nnz_limit = if preserve_sparse_storage {
4737        usize::MAX
4738    } else {
4739        let total_cells = nrows.saturating_mul(ncols);
4740        ((total_cells as f64) * BLOCK_SPARSE_MAX_DENSITY).floor() as usize
4741    };
4742    let mut nnz = 0usize;
4743    for block in blocks {
4744        let block_nnz = if let Some(block_nnz) = sparse_compatible_block_nnz(block) {
4745            block_nnz
4746        } else {
4747            return Ok(None);
4748        };
4749        nnz = nnz.saturating_add(block_nnz);
4750        if nnz > sparse_nnz_limit {
4751            return Ok(None);
4752        }
4753    }
4754
4755    let mut triplets = Vec::<Triplet<usize, usize, f64>>::with_capacity(nnz);
4756    let mut col_offset = 0usize;
4757    for block in blocks {
4758        match block {
4759            DesignBlock::Intercept(n) => {
4760                for row in 0..*n {
4761                    triplets.push(Triplet::new(row, col_offset, 1.0));
4762                }
4763            }
4764            DesignBlock::RandomEffect(op) => {
4765                for (row, group_id) in op.group_ids.iter().enumerate() {
4766                    if let Some(group) = group_id {
4767                        triplets.push(Triplet::new(row, col_offset + group, 1.0));
4768                    }
4769                }
4770            }
4771            DesignBlock::Sparse(sparse) => {
4772                let (symbolic, values) = sparse.parts();
4773                let col_ptr = symbolic.col_ptr();
4774                let row_idx = symbolic.row_idx();
4775                for col in 0..sparse.ncols() {
4776                    for idx in col_ptr[col]..col_ptr[col + 1] {
4777                        let value = values[idx];
4778                        if value.abs() > BLOCK_SPARSE_ZERO_EPS {
4779                            triplets.push(Triplet::new(row_idx[idx], col_offset + col, value));
4780                        }
4781                    }
4782                }
4783            }
4784            DesignBlock::Dense(dense) => {
4785                let matrix = dense.as_dense_ref().ok_or_else(|| {
4786                    BasisError::InvalidInput(
4787                        "sparse-compatible block assembly requires materialized dense blocks"
4788                            .to_string(),
4789                    )
4790                })?;
4791                for row in 0..matrix.nrows() {
4792                    for col in 0..matrix.ncols() {
4793                        let value = matrix[[row, col]];
4794                        if value.abs() > BLOCK_SPARSE_ZERO_EPS {
4795                            triplets.push(Triplet::new(row, col_offset + col, value));
4796                        }
4797                    }
4798                }
4799            }
4800        }
4801        col_offset += block.ncols();
4802    }
4803
4804    let sparse = SparseColMat::try_new_from_triplets(nrows, ncols, &triplets).map_err(|_| {
4805        BasisError::SparseCreation("failed to assemble sparse term-collection design".to_string())
4806    })?;
4807    Ok(Some(DesignMatrix::Sparse(
4808        gam_linalg::matrix::SparseDesignMatrix::new(sparse),
4809    )))
4810}
4811
4812pub fn assemble_term_collection_design_matrix(
4813    blocks: Vec<DesignBlock>,
4814) -> Result<DesignMatrix, BasisError> {
4815    if let Some(sparse) = try_build_sparse_design_from_blocks(&blocks)? {
4816        return Ok(sparse);
4817    }
4818    let block_op = BlockDesignOperator::new(blocks).map_err(|e| {
4819        BasisError::InvalidInput(format!("failed to build block design operator: {e}"))
4820    })?;
4821    Ok(DesignMatrix::Dense(
4822        gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(block_op)),
4823    ))
4824}
4825
4826pub fn select_columns(
4827    data: ArrayView2<'_, f64>,
4828    cols: &[usize],
4829) -> Result<Array2<f64>, BasisError> {
4830    let n = data.nrows();
4831    let p = data.ncols();
4832    for &c in cols {
4833        if c >= p {
4834            crate::bail_dim_basis!("feature column {c} is out of bounds for data with {p} columns");
4835        }
4836    }
4837    let mut out = Array2::<f64>::zeros((n, cols.len()));
4838    for (j, &c) in cols.iter().enumerate() {
4839        out.column_mut(j).assign(&data.column(c));
4840    }
4841    Ok(out)
4842}
4843
4844pub fn nonfinite_value_label(value: f64) -> &'static str {
4845    if value.is_nan() {
4846        "NaN"
4847    } else if value.is_sign_positive() {
4848        "+Inf"
4849    } else {
4850        "-Inf"
4851    }
4852}
4853
4854pub fn validate_term_feature_column_finite(
4855    data: ArrayView2<'_, f64>,
4856    term_kind: &str,
4857    term_name: &str,
4858    feature_col: usize,
4859) -> Result<(), BasisError> {
4860    let p = data.ncols();
4861    if feature_col >= p {
4862        crate::bail_dim_basis!(
4863            "{term_kind} term '{term_name}' feature column {feature_col} out of bounds for {p} columns"
4864        );
4865    }
4866    for (row, &value) in data.column(feature_col).iter().enumerate() {
4867        if !value.is_finite() {
4868            crate::bail_invalid_basis!(
4869                "{term_kind} term '{term_name}' feature column {feature_col} row {row} contains non-finite value {}",
4870                nonfinite_value_label(value)
4871            );
4872        }
4873    }
4874    Ok(())
4875}
4876
4877pub fn validate_smooth_terms_finite_inputs(
4878    data: ArrayView2<'_, f64>,
4879    terms: &[SmoothTermSpec],
4880) -> Result<(), BasisError> {
4881    for term in terms {
4882        for feature_col in smooth_term_feature_cols(term) {
4883            validate_term_feature_column_finite(data, "smooth", &term.name, feature_col)?;
4884        }
4885    }
4886    Ok(())
4887}
4888
4889pub fn validate_term_collection_finite_inputs(
4890    data: ArrayView2<'_, f64>,
4891    spec: &TermCollectionSpec,
4892) -> Result<(), BasisError> {
4893    for term in &spec.linear_terms {
4894        validate_term_feature_column_finite(data, "linear", &term.name, term.feature_col)?;
4895    }
4896    for term in &spec.random_effect_terms {
4897        validate_term_feature_column_finite(data, "random-effect", &term.name, term.feature_col)?;
4898    }
4899    validate_smooth_terms_finite_inputs(data, &spec.smooth_terms)
4900}
4901
4902#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
4903pub struct JointSpatialCenterGroupKey {
4904    feature_cols: Vec<usize>,
4905    strategy_kind: CenterStrategyKind,
4906    strategy_aux: usize,
4907    requested_num_centers: usize,
4908    input_scale_bits: Option<u64>,
4909}
4910
4911pub fn spatial_term_min_center_count(term: &SmoothTermSpec) -> usize {
4912    match &term.basis {
4913        SmoothBasisSpec::ThinPlate { feature_cols, .. } => feature_cols.len() + 1,
4914        SmoothBasisSpec::Duchon {
4915            feature_cols, spec, ..
4916        } => match spec.nullspace_order {
4917            crate::basis::DuchonNullspaceOrder::Zero => 1,
4918            crate::basis::DuchonNullspaceOrder::Linear => feature_cols.len() + 1,
4919            crate::basis::DuchonNullspaceOrder::Degree(degree) => {
4920                crate::basis::duchon_nullspace_dimension(feature_cols.len(), degree)
4921            }
4922        },
4923        SmoothBasisSpec::Matern { .. } => 1,
4924        _ => 1,
4925    }
4926}
4927
4928pub fn spatial_term_group_key(term: &SmoothTermSpec) -> Option<JointSpatialCenterGroupKey> {
4929    let (feature_cols, strategy, input_scale) = match &term.basis {
4930        SmoothBasisSpec::ThinPlate {
4931            feature_cols,
4932            spec,
4933            input_scale,
4934        } => (feature_cols, &spec.center_strategy, *input_scale),
4935        SmoothBasisSpec::Matern {
4936            feature_cols,
4937            spec,
4938            input_scale,
4939        } => (feature_cols, &spec.center_strategy, *input_scale),
4940        SmoothBasisSpec::Duchon {
4941            feature_cols,
4942            spec,
4943            input_scale,
4944        } => (feature_cols, &spec.center_strategy, *input_scale),
4945        _ => return None,
4946    };
4947    let strategy_kind = center_strategy_kind(strategy);
4948    let strategy_aux = match strategy {
4949        CenterStrategy::Auto(inner) => match inner.as_ref() {
4950            CenterStrategy::KMeans { max_iter, .. } => *max_iter,
4951            CenterStrategy::UniformGrid { points_per_dim } => *points_per_dim,
4952            _ => 0,
4953        },
4954        CenterStrategy::KMeans { max_iter, .. } => *max_iter,
4955        CenterStrategy::UniformGrid { points_per_dim } => *points_per_dim,
4956        _ => 0,
4957    };
4958    Some(JointSpatialCenterGroupKey {
4959        feature_cols: feature_cols.clone(),
4960        strategy_kind,
4961        strategy_aux,
4962        requested_num_centers: strategy.planned_num_centers(feature_cols.len()),
4963        input_scale_bits: input_scale.map(crate::IsotropicScale::to_bits),
4964    })
4965}
4966
4967pub fn spatial_term_center_strategy(term: &SmoothTermSpec) -> Option<&CenterStrategy> {
4968    match &term.basis {
4969        SmoothBasisSpec::ThinPlate { spec, .. } => Some(&spec.center_strategy),
4970        SmoothBasisSpec::Matern { spec, .. } => Some(&spec.center_strategy),
4971        SmoothBasisSpec::Duchon { spec, .. } => Some(&spec.center_strategy),
4972        _ => None,
4973    }
4974}
4975
4976pub fn set_spatial_term_centers(
4977    term: &mut SmoothTermSpec,
4978    centers: Array2<f64>,
4979) -> Result<(), BasisError> {
4980    match &mut term.basis {
4981        SmoothBasisSpec::ThinPlate { spec, .. } => {
4982            spec.center_strategy = CenterStrategy::UserProvided(centers);
4983            Ok(())
4984        }
4985        SmoothBasisSpec::Matern { spec, .. } => {
4986            spec.center_strategy = CenterStrategy::UserProvided(centers);
4987            Ok(())
4988        }
4989        SmoothBasisSpec::Duchon { spec, .. } => {
4990            spec.center_strategy = CenterStrategy::UserProvided(centers);
4991            Ok(())
4992        }
4993        _ => Err(BasisError::InvalidInput(format!(
4994            "term '{}' does not support spatial center planning",
4995            term.name
4996        ))),
4997    }
4998}
4999
5000pub fn standardized_spatial_term_data(
5001    data: ArrayView2<'_, f64>,
5002    term: &SmoothTermSpec,
5003) -> Result<Array2<f64>, BasisError> {
5004    let (feature_cols, input_scale) = match &term.basis {
5005        SmoothBasisSpec::ThinPlate {
5006            feature_cols,
5007            input_scale,
5008            ..
5009        }
5010        | SmoothBasisSpec::Matern {
5011            feature_cols,
5012            input_scale,
5013            ..
5014        }
5015        | SmoothBasisSpec::Duchon {
5016            feature_cols,
5017            input_scale,
5018            ..
5019        } => (feature_cols, *input_scale),
5020        _ => {
5021            crate::bail_invalid_basis!("term '{}' is not a spatial smooth", term.name);
5022        }
5023    };
5024    let mut x = select_columns(data, feature_cols)?;
5025    input_scale
5026        .map_or_else(|| estimate_isotropic_scale(x.view()), Ok)?
5027        .standardize(&mut x);
5028    Ok(x)
5029}
5030
5031pub fn plan_joint_spatial_centers_for_term_blocks(
5032    data: ArrayView2<'_, f64>,
5033    term_blocks: &[Vec<SmoothTermSpec>],
5034) -> Result<Vec<Vec<SmoothTermSpec>>, BasisError> {
5035    let mut planned_blocks = term_blocks.to_vec();
5036    let n = data.nrows();
5037    let mut groups: BTreeMap<JointSpatialCenterGroupKey, Vec<(usize, usize)>> = BTreeMap::new();
5038
5039    for (block_idx, terms) in planned_blocks.iter().enumerate() {
5040        for (term_idx, term) in terms.iter().enumerate() {
5041            let Some(strategy) = spatial_term_center_strategy(term) else {
5042                continue;
5043            };
5044            if !center_strategy_is_auto(strategy) {
5045                continue;
5046            }
5047            let Some(group_key) = spatial_term_group_key(term) else {
5048                continue;
5049            };
5050            if !matches!(
5051                group_key.strategy_kind,
5052                CenterStrategyKind::EqualMass
5053                    | CenterStrategyKind::EqualMassCovarRepresentative
5054                    | CenterStrategyKind::FarthestPoint
5055                    | CenterStrategyKind::KMeans
5056                    | CenterStrategyKind::UniformGrid
5057            ) {
5058                continue;
5059            }
5060            groups
5061                .entry(group_key)
5062                .or_default()
5063                .push((block_idx, term_idx));
5064        }
5065    }
5066
5067    for (group_key, members) in groups {
5068        if members.len() < 2 {
5069            continue;
5070        }
5071        let min_required = members
5072            .iter()
5073            .map(|&(block_idx, term_idx)| {
5074                spatial_term_min_center_count(&planned_blocks[block_idx][term_idx])
5075            })
5076            .max()
5077            .unwrap_or(1);
5078        let joint_centers = group_key
5079            .requested_num_centers
5080            .max(min_required)
5081            .min(n.max(1));
5082        let (first_block_idx, first_term_idx) = members[0];
5083        let prototype = &planned_blocks[first_block_idx][first_term_idx];
5084        let standardized = standardized_spatial_term_data(data, prototype)?;
5085        let strategy = spatial_term_center_strategy(prototype).ok_or_else(|| {
5086            BasisError::InvalidInput(format!(
5087                "term '{}' lost its spatial center strategy during joint planning",
5088                prototype.name
5089            ))
5090        })?;
5091        let joint_strategy = center_strategy_with_num_centers(
5092            strategy,
5093            joint_centers,
5094            group_key.feature_cols.len(),
5095        )?;
5096        let shared_centers = select_centers_by_strategy(standardized.view(), &joint_strategy)?;
5097        log::info!(
5098            "sharing {} spatial centers across {} smooth terms over columns {:?} (requested {} centers)",
5099            shared_centers.nrows(),
5100            members.len(),
5101            group_key.feature_cols,
5102            group_key.requested_num_centers,
5103        );
5104        for (block_idx, term_idx) in members {
5105            set_spatial_term_centers(
5106                &mut planned_blocks[block_idx][term_idx],
5107                shared_centers.clone(),
5108            )?;
5109        }
5110    }
5111
5112    // Resolve typed Matérn Auto scales and the legacy thin-plate numeric
5113    // auto marker to a data-driven initialization here so REML starts in a
5114    // regime where it can escape. Matérn retains Auto provenance after this
5115    // numeric seed is installed, so it cannot later masquerade as user-fixed.
5116    for block in planned_blocks.iter_mut() {
5117        for term in block.iter_mut() {
5118            auto_init_length_scale_in_place(data, term);
5119        }
5120    }
5121
5122    Ok(planned_blocks)
5123}
5124
5125/// Tiny positive floor for the auto length scale, guarding against a zero
5126/// kernel range when every feature column is (near-)constant.
5127const AUTO_LENGTH_SCALE_FLOOR: f64 = 1e-6;
5128
5129/// Widest per-axis range of the selected feature columns. Returns `None` when
5130/// every selected column is constant / non-finite (no usable spatial scale).
5131fn feature_columns_max_range(data: ArrayView2<'_, f64>, feature_cols: &[usize]) -> Option<f64> {
5132    let mut max_range = 0.0_f64;
5133    for &c in feature_cols {
5134        if c >= data.ncols() {
5135            continue;
5136        }
5137        let col = data.column(c);
5138        let mut lo = f64::INFINITY;
5139        let mut hi = f64::NEG_INFINITY;
5140        for &v in col.iter() {
5141            if v.is_finite() {
5142                if v < lo {
5143                    lo = v;
5144                }
5145                if v > hi {
5146                    hi = v;
5147                }
5148            }
5149        }
5150        if hi > lo {
5151            let r = hi - lo;
5152            if r > max_range {
5153                max_range = r;
5154            }
5155        }
5156    }
5157    if max_range.is_finite() && max_range > 0.0 {
5158        Some(max_range)
5159    } else {
5160        None
5161    }
5162}
5163
5164/// Rotation-invariant analogue of [`feature_columns_max_range`], calibrated to
5165/// the span of the cloud's longest direction.
5166///
5167/// For a uniform interval of width `L`, the leading covariance eigenvalue is
5168/// `L²/12`, so `sqrt(12·λ_max)` recovers `L`. The same identity holds for the
5169/// longest side of an axis-aligned uniform box, while `λ_max` is invariant
5170/// under every orthogonal change of coordinates. This preserves the scale of
5171/// the former widest-axis seed without making it frame-dependent (gam#2252).
5172/// Sorting the complete points lexicographically makes each frame stable under
5173/// a pure row permutation (gam#1378).
5174fn feature_columns_rotation_invariant_range(
5175    data: ArrayView2<'_, f64>,
5176    feature_cols: &[usize],
5177) -> Option<f64> {
5178    let cols: Vec<usize> = feature_cols
5179        .iter()
5180        .copied()
5181        .filter(|&c| c < data.ncols())
5182        .collect();
5183    if cols.is_empty() {
5184        return None;
5185    }
5186    let mut points: Vec<Vec<f64>> = data
5187        .rows()
5188        .into_iter()
5189        .filter_map(|row| {
5190            let point: Vec<f64> = cols.iter().map(|&column| row[column]).collect();
5191            point.iter().all(|value| value.is_finite()).then_some(point)
5192        })
5193        .collect();
5194    if points.is_empty() {
5195        return None;
5196    }
5197    points.sort_by(|left, right| {
5198        left.iter()
5199            .zip(right)
5200            .find_map(|(a, b)| {
5201                let ordering = a.total_cmp(b);
5202                ordering.is_ne().then_some(ordering)
5203            })
5204            .unwrap_or(std::cmp::Ordering::Equal)
5205    });
5206
5207    let dimensions = cols.len();
5208    let count = points.len() as f64;
5209    let mut centroid = vec![0.0_f64; dimensions];
5210    for point in &points {
5211        for (coordinate, value) in centroid.iter_mut().zip(point) {
5212            *coordinate += *value;
5213        }
5214    }
5215    for coordinate in &mut centroid {
5216        *coordinate /= count;
5217    }
5218
5219    let mut covariance = Array2::<f64>::zeros((dimensions, dimensions));
5220    for point in &points {
5221        for row in 0..dimensions {
5222            let centered_row = point[row] - centroid[row];
5223            for column in 0..=row {
5224                covariance[[row, column]] += centered_row * (point[column] - centroid[column]);
5225            }
5226        }
5227    }
5228    for row in 0..dimensions {
5229        for column in 0..=row {
5230            let value = covariance[[row, column]] / count;
5231            covariance[[row, column]] = value;
5232            covariance[[column, row]] = value;
5233        }
5234    }
5235
5236    use gam_linalg::faer_ndarray::FaerEigh;
5237    let (eigenvalues, _) = covariance
5238        .eigh(faer::Side::Lower)
5239        .expect("finite covariance must have a symmetric eigendecomposition");
5240    let leading_variance = eigenvalues[eigenvalues.len() - 1];
5241    let extent = (12.0 * leading_variance).sqrt();
5242    if extent.is_finite() && extent > 0.0 {
5243        Some(extent)
5244    } else {
5245        None
5246    }
5247}
5248
5249/// Compute a data-driven initial length scale from the per-axis range of the
5250/// feature columns. The heuristic `max_range / sqrt(n)` puts the kernel on
5251/// the wiggly side of REML's basin so the optimizer can grow it back if the
5252/// signal is smooth, but is small enough that high-frequency truths remain
5253/// reachable for smoother kernels (ν ≥ 5/2). Clamped to a tiny positive
5254/// floor so degenerate constant-input columns can't produce 0.
5255pub fn auto_initial_length_scale(data: ArrayView2<'_, f64>, feature_cols: &[usize]) -> f64 {
5256    let n = data.nrows();
5257    if n == 0 || feature_cols.is_empty() {
5258        return 1.0;
5259    }
5260    let Some(max_range) = feature_columns_max_range(data, feature_cols) else {
5261        return 1.0;
5262    };
5263    let init = max_range / (n as f64).sqrt();
5264    init.max(AUTO_LENGTH_SCALE_FLOOR).min(max_range)
5265}
5266
5267/// Density-adaptive auto length scale for a kernel basis with `num_centers`
5268/// requested centers (#1731).
5269///
5270/// The plain [`auto_initial_length_scale`] seed `max_range / sqrt(n)` is the
5271/// fill distance of the *n data points*; it is independent of the requested
5272/// center count `k`. For a radial kernel at a FIXED length scale, packing more
5273/// centers into the same cloud makes neighbouring basis functions overlap and
5274/// go numerically collinear, so the realized basis saturates in rank (the
5275/// `matern_rank_reduce_centers` cap) and a richer `k` becomes a no-op — or even
5276/// shrinks the basis. The kernel stays well-conditioned only while the length
5277/// scale tracks the *center* spacing, not the data spacing.
5278///
5279/// We seed the length scale at the fill distance of `max(n, k)` points,
5280/// `max_range / sqrt(max(n, k))`. When `n ≥ k` (the usual case) this is exactly
5281/// the existing `max_range / sqrt(n)` seed, so every current result and small-`k`
5282/// basis size is preserved bit-for-bit (in every covariate dimension). When
5283/// `k > n` (a dense center request on a small cloud, the regime where an
5284/// `n`-sized seed sits above the center spacing and over-smooths the centers
5285/// into collinearity) the seed shrinks with `k` to the center spacing, keeping
5286/// the requested centers numerically independent. This is the Matérn analogue of
5287/// the Duchon-promotion "length_scale from center spacing" rule
5288/// (`hybrid_duchon_promotion_length_scale`).
5289pub fn auto_initial_length_scale_for_centers(
5290    data: ArrayView2<'_, f64>,
5291    feature_cols: &[usize],
5292    num_centers: usize,
5293) -> f64 {
5294    let n = data.nrows();
5295    if n == 0 || feature_cols.is_empty() {
5296        return 1.0;
5297    }
5298    // #2252: rotation-invariant extent for the Matérn seed so the enrolled κ/range
5299    // solve — which is basin-/seed-sensitive (see the matern geometry-stall path)
5300    // — starts from a frame-independent point and lands in the SAME basin in every
5301    // rotated frame, making the isotropic Matérn fit rotation-invariant. The
5302    // per-axis span (`feature_columns_max_range`) is a projection of the cloud and
5303    // is rotation-variant; the covariance spectral extent `sqrt(12·λ_max)` is
5304    // invariant under any orthogonal map and retains the former span calibration.
5305    // Duchon/thin-plate seeds are computed by separate helpers and are unchanged,
5306    // so those (seed-robust) bases stay bit-identical — this fix is scoped to the
5307    // seed-sensitive Matérn path.
5308    let Some(max_range) = feature_columns_rotation_invariant_range(data, feature_cols) else {
5309        return 1.0;
5310    };
5311    // Resolution density: at least the data points, but no coarser than the
5312    // center spacing once more centers than data are requested. Using the same
5313    // `sqrt` fill-distance law as `auto_initial_length_scale` keeps the seed
5314    // bit-identical whenever `n ≥ num_centers` (every dimension), and only
5315    // shrinks it — never grows it — when `num_centers > n`.
5316    let resolution_points = n.max(num_centers).max(1) as f64;
5317    let spacing = max_range / resolution_points.sqrt();
5318    spacing.max(AUTO_LENGTH_SCALE_FLOOR).min(max_range)
5319}
5320
5321/// Rotation-invariant center-resolution range for the companion Matérn basin.
5322///
5323/// A reduced-rank Matérn basis has two distinct geometric resolutions: the
5324/// observation fill distance used by the short/rich cold seed, and the coarser
5325/// fill distance of its `k` retained centers. The latter is the canonical
5326/// response-free representative of the overlapping, long-range basin:
5327/// `sqrt(12 * lambda_max(cov(x))) / sqrt(k)`. It uses the same covariance
5328/// extent as [`auto_initial_length_scale_for_centers`], so rigid rotations and
5329/// row permutations leave it unchanged.
5330pub fn matern_low_rank_center_resolution_length_scale(
5331    data: ArrayView2<'_, f64>,
5332    feature_cols: &[usize],
5333    num_centers: usize,
5334) -> Option<f64> {
5335    if data.nrows() == 0 || feature_cols.is_empty() || num_centers == 0 {
5336        return None;
5337    }
5338    let extent = feature_columns_rotation_invariant_range(data, feature_cols)?;
5339    let length_scale = extent / (num_centers as f64).sqrt();
5340    Some(length_scale.max(AUTO_LENGTH_SCALE_FLOOR).min(extent))
5341}
5342
5343/// Low-rank radial-basis length-scale seed tied to the requested center spacing.
5344///
5345/// Thin-plate regression splines with `k << n` represent the surface through a
5346/// compact set of centers; seeding the kernel at the observation fill distance
5347/// (`max_range / sqrt(n)`) makes the center Gram nearly diagonal and turns the
5348/// bending penalty into an ill-scaled ridge on the radial coefficients. REML then
5349/// sees a weakly identified smoothing surface and can settle on under-recovered
5350/// spatial fits. Seed at the center fill distance instead, so neighbouring
5351/// centers interact at O(1) scale before REML tunes the smoothing parameter.
5352pub fn auto_initial_length_scale_for_low_rank_centers(
5353    data: ArrayView2<'_, f64>,
5354    feature_cols: &[usize],
5355    num_centers: usize,
5356) -> f64 {
5357    if data.nrows() == 0 || feature_cols.is_empty() {
5358        return 1.0;
5359    }
5360    let Some(max_range) = feature_columns_max_range(data, feature_cols) else {
5361        return 1.0;
5362    };
5363    let resolution_points = num_centers.max(1) as f64;
5364    let spacing = max_range / resolution_points.sqrt();
5365    spacing.max(AUTO_LENGTH_SCALE_FLOOR).min(max_range)
5366}
5367
5368/// Requested center count encoded by a [`CenterStrategy`], if it carries an
5369/// explicit count (used to make the Matérn auto length scale density-adaptive).
5370fn center_strategy_requested_count(strategy: &CenterStrategy) -> Option<usize> {
5371    match strategy {
5372        CenterStrategy::Auto(inner) => center_strategy_requested_count(inner),
5373        CenterStrategy::DuchonSpectral { knots, .. } => center_strategy_requested_count(knots),
5374        CenterStrategy::UserProvided(centers) => Some(centers.nrows()),
5375        CenterStrategy::EqualMass { num_centers }
5376        | CenterStrategy::EqualMassCovarRepresentative { num_centers }
5377        | CenterStrategy::FarthestPoint { num_centers }
5378        | CenterStrategy::KMeans { num_centers, .. } => Some(*num_centers),
5379        CenterStrategy::UniformGrid { .. } => None,
5380    }
5381}
5382
5383/// Walk a term and resolve an omitted Matérn length scale, or a thin-plate
5384/// smooth still carrying its numeric auto marker, with
5385/// [`auto_initial_length_scale`]. Matérn's typed Auto provenance survives.
5386pub fn auto_init_length_scale_in_place(data: ArrayView2<'_, f64>, term: &mut SmoothTermSpec) {
5387    auto_init_length_scale_in_basis(data, &mut term.basis);
5388}
5389
5390/// Resolve the typed Matérn Auto length scale (and thin-plate's numeric auto
5391/// marker) with a data-derived value for any reachable kernel — including the
5392/// inner kernel of a `by=`/factor-smooth wrapper.
5393///
5394/// `by=<factor>` and the sum-to-zero factor smooth wrap a spatial kernel inside
5395/// `SmoothBasisSpec::ByVariable` / `SmoothBasisSpec::FactorSumToZero` /
5396/// `SmoothBasisSpec::BySmooth`, so the wrapper variant is what the planner sees.
5397/// Without recursing into the wrapped basis the inner Matérn remains unresolved
5398/// (and ThinPlate keeps its numeric marker), so no valid kernel scale exists at
5399/// fit or predict time. Recurse so the inner kernel is initialized identically
5400/// to a top-level one.
5401pub fn auto_init_length_scale_in_basis(data: ArrayView2<'_, f64>, basis: &mut SmoothBasisSpec) {
5402    match basis {
5403        SmoothBasisSpec::Matern {
5404            feature_cols, spec, ..
5405        } => {
5406            if spec.length_scale.resolved().is_none() {
5407                // Density-adaptive seed (#1731): when the requested center count
5408                // is known, scale the auto length scale with the *center*
5409                // spacing so a richer `k` stays numerically full-rank instead of
5410                // saturating against `matern_rank_reduce_centers`. For `n ≥ k`
5411                // (the usual case) this is identical to the plain `max_range /
5412                // sqrt(n)` seed in 2-D, so small-`k` results are unchanged. The
5413                // unconstrained / non-explicit `UniformGrid` strategy falls back
5414                // to the plain seed.
5415                let resolved = match center_strategy_requested_count(&spec.center_strategy) {
5416                    Some(k) => auto_initial_length_scale_for_centers(data, feature_cols, k),
5417                    None => auto_initial_length_scale(data, feature_cols),
5418                };
5419                spec.length_scale.resolve_auto_once(resolved);
5420            }
5421        }
5422        SmoothBasisSpec::ThinPlate {
5423            feature_cols, spec, ..
5424        } => {
5425            if spec.length_scale == 0.0 {
5426                spec.length_scale = match center_strategy_requested_count(&spec.center_strategy) {
5427                    Some(k) => {
5428                        auto_initial_length_scale_for_low_rank_centers(data, feature_cols, k)
5429                    }
5430                    None => auto_initial_length_scale(data, feature_cols),
5431                };
5432            }
5433        }
5434        SmoothBasisSpec::ByVariable { inner, .. }
5435        | SmoothBasisSpec::FactorSumToZero { inner, .. } => {
5436            auto_init_length_scale_in_basis(data, inner);
5437        }
5438        SmoothBasisSpec::BySmooth { smooth, .. } => {
5439            auto_init_length_scale_in_basis(data, smooth);
5440        }
5441        // Enumerated rather than wildcarded so a new basis family has to state
5442        // whether it carries an auto-seeded length scale. None of these do:
5443        // B-spline marginals (`FactorSmooth`, `TensorBSpline`) and `Pca` are
5444        // knot/column constructions with no kernel bandwidth, and the
5445        // `Sphere` / `ConstantCurvature` / `MeasureJet` / `Duchon` specs carry
5446        // no `length_scale` field to resolve.
5447        SmoothBasisSpec::BSpline1D { .. }
5448        | SmoothBasisSpec::FactorSmooth { .. }
5449        | SmoothBasisSpec::Sphere { .. }
5450        | SmoothBasisSpec::ConstantCurvature { .. }
5451        | SmoothBasisSpec::MeasureJet { .. }
5452        | SmoothBasisSpec::Duchon { .. }
5453        | SmoothBasisSpec::Pca { .. }
5454        | SmoothBasisSpec::TensorBSpline { .. } => {}
5455    }
5456}
5457
5458impl LinearFitConditioning {
5459    pub fn from_columns(design: &TermCollectionDesign, selected_cols: &[usize]) -> Self {
5460        const SCALE_EPS: f64 = 1e-12;
5461        let n = design.design.nrows();
5462        let p = design.design.ncols();
5463        let mut columns = Vec::with_capacity(selected_cols.len());
5464        if n == 0 || selected_cols.is_empty() {
5465            return Self {
5466                intercept_idx: design.intercept_range.start,
5467                columns,
5468            };
5469        }
5470        let chunk_rows = gam_linalg::utils::row_chunk_for_byte_budget(n, p);
5471        // Two-pass mean/variance so operator-backed designs don't need to
5472        // materialize the full dense matrix. Pass 1 accumulates per-column
5473        // sums; pass 2 accumulates the sum of squared deviations from the
5474        // pass-1 mean. This matches the original `Σ (x − mean)² / n` formula
5475        // without the catastrophic cancellation of `E[X²] − E[X]²`.
5476        let mut sums = vec![0.0_f64; selected_cols.len()];
5477        for start in (0..n).step_by(chunk_rows) {
5478            let end = (start + chunk_rows).min(n);
5479            let chunk = design
5480                .design
5481                .try_row_chunk(start..end)
5482                .expect("LinearFitConditioning::from_columns row chunk failed");
5483            for (k, &col_idx) in selected_cols.iter().enumerate() {
5484                let column = chunk.column(col_idx);
5485                for &v in column.iter() {
5486                    sums[k] += v;
5487                }
5488            }
5489        }
5490        let inv_n = 1.0_f64 / n as f64;
5491        let means: Vec<f64> = sums.iter().map(|&s| s * inv_n).collect();
5492        let mut sq_devs = vec![0.0_f64; selected_cols.len()];
5493        for start in (0..n).step_by(chunk_rows) {
5494            let end = (start + chunk_rows).min(n);
5495            let chunk = design
5496                .design
5497                .try_row_chunk(start..end)
5498                .expect("LinearFitConditioning::from_columns row chunk failed");
5499            for (k, &col_idx) in selected_cols.iter().enumerate() {
5500                let mean_k = means[k];
5501                let column = chunk.column(col_idx);
5502                for &v in column.iter() {
5503                    let d = v - mean_k;
5504                    sq_devs[k] += d * d;
5505                }
5506            }
5507        }
5508        for (k, &col_idx) in selected_cols.iter().enumerate() {
5509            let mean = means[k];
5510            let var = sq_devs[k] * inv_n;
5511            let (mean, scale) = if var.is_finite() && var > SCALE_EPS * SCALE_EPS {
5512                (mean, var.sqrt())
5513            } else {
5514                // Leave nearly-constant columns untouched; centering them would collapse
5515                // the design column to ~0 and change the model rather than just condition it.
5516                (0.0, 1.0)
5517            };
5518            columns.push(LinearColumnConditioning {
5519                col_idx,
5520                mean,
5521                scale,
5522            });
5523        }
5524        Self {
5525            intercept_idx: design.intercept_range.start,
5526            columns,
5527        }
5528    }
5529
5530    pub fn apply_to_design(&self, design: &Array2<f64>) -> Array2<f64> {
5531        let mut out = design.clone();
5532        for col in &self.columns {
5533            {
5534                let mut dst = out.column_mut(col.col_idx);
5535                dst -= col.mean;
5536            }
5537            if col.scale != 1.0 {
5538                out.column_mut(col.col_idx).mapv_inplace(|v| v / col.scale);
5539            }
5540        }
5541        out
5542    }
5543
5544    fn transform_matrix_columnswith_a(&self, mat: &Array2<f64>) -> Array2<f64> {
5545        let mut out = mat.clone();
5546        let intercept = self.intercept_idx;
5547        for col in &self.columns {
5548            let intercept_col = out.column(intercept).to_owned();
5549            let mut target = out.column_mut(col.col_idx);
5550            target -= &(intercept_col * col.mean);
5551            if col.scale != 1.0 {
5552                target.mapv_inplace(|v| v / col.scale);
5553            }
5554        }
5555        out
5556    }
5557
5558    fn transform_matrixrowswith_a_transpose(&self, mat: &Array2<f64>) -> Array2<f64> {
5559        let mut out = mat.clone();
5560        let intercept = self.intercept_idx;
5561        for col in &self.columns {
5562            let interceptrow = out.row(intercept).to_owned();
5563            let mut target = out.row_mut(col.col_idx);
5564            target -= &(interceptrow * col.mean);
5565            if col.scale != 1.0 {
5566                target.mapv_inplace(|v| v / col.scale);
5567            }
5568        }
5569        out
5570    }
5571
5572    /// Left-multiply `mat_internal` by `M⁻ᵀ` where `M⁻¹[intercept, j] = mean_j`
5573    /// and `M⁻¹[j, j] = scale_j` for each conditioned column. Used together
5574    /// with [`Self::right_multiply_by_m_inv`] to back-transform an internal
5575    /// penalized Hessian to the original coefficient basis.
5576    fn left_multiply_by_m_inv_transpose(&self, mat_internal: &Array2<f64>) -> Array2<f64> {
5577        let mut out = mat_internal.clone();
5578        let intercept = self.intercept_idx;
5579        let interceptrow_snapshot = mat_internal.row(intercept).to_owned();
5580        for col in &self.columns {
5581            if col.scale != 1.0 {
5582                out.row_mut(col.col_idx).mapv_inplace(|v| v * col.scale);
5583            }
5584            if col.mean != 0.0 {
5585                let mut target = out.row_mut(col.col_idx);
5586                target += &(&interceptrow_snapshot * col.mean);
5587            }
5588        }
5589        out
5590    }
5591
5592    /// Right-multiply `mat_internal` by `M⁻¹`. Mirror of
5593    /// [`Self::left_multiply_by_m_inv_transpose`] on columns.
5594    fn right_multiply_by_m_inv(&self, mat_internal: &Array2<f64>) -> Array2<f64> {
5595        let mut out = mat_internal.clone();
5596        let intercept = self.intercept_idx;
5597        let intercept_col_snapshot = mat_internal.column(intercept).to_owned();
5598        for col in &self.columns {
5599            if col.scale != 1.0 {
5600                out.column_mut(col.col_idx).mapv_inplace(|v| v * col.scale);
5601            }
5602            if col.mean != 0.0 {
5603                let mut target = out.column_mut(col.col_idx);
5604                target += &(&intercept_col_snapshot * col.mean);
5605            }
5606        }
5607        out
5608    }
5609
5610    /// Transform blockwise penalties through the conditioning.
5611    ///
5612    /// For block-local penalties whose `col_range` does not overlap with any
5613    /// conditioning column, the transform is identity (the conditioning only
5614    /// affects unpenalized linear columns). In that common case the penalty
5615    /// passes through unchanged, avoiding O(p²) materialization entirely.
5616    pub fn transform_blockwise_penalties_to_internal(
5617        &self,
5618        penalties: &[BlockwisePenalty],
5619        p: usize,
5620    ) -> Vec<crate::penalty_spec::PenaltySpec> {
5621        let conditioning_cols: std::collections::HashSet<usize> =
5622            self.columns.iter().map(|c| c.col_idx).collect();
5623        penalties
5624            .iter()
5625            .map(|bp| {
5626                let overlaps =
5627                    (bp.col_range.start..bp.col_range.end).any(|j| conditioning_cols.contains(&j));
5628                if overlaps {
5629                    // Rare: penalty block overlaps conditioning columns.
5630                    // Fall back to dense transform.
5631                    let global = bp.to_global(p);
5632                    let right = self.transform_matrix_columnswith_a(&global);
5633                    let transformed = self.transform_matrixrowswith_a_transpose(&right);
5634                    crate::penalty_spec::PenaltySpec::Dense(transformed)
5635                } else {
5636                    // Common: smooth penalty block doesn't touch linear columns.
5637                    // The conditioning is identity on this block.
5638                    crate::penalty_spec::PenaltySpec::from_blockwise(bp.clone())
5639                }
5640            })
5641            .collect()
5642    }
5643
5644    pub fn backtransform_beta(&self, beta_internal: &Array1<f64>) -> Array1<f64> {
5645        let mut beta = beta_internal.clone();
5646        let intercept = self.intercept_idx;
5647        for col in &self.columns {
5648            beta[intercept] -= beta_internal[col.col_idx] * col.mean / col.scale;
5649            beta[col.col_idx] = beta_internal[col.col_idx] / col.scale;
5650        }
5651        beta
5652    }
5653
5654    /// `H_orig = M⁻ᵀ · H_int · M⁻¹`, derived from
5655    /// `L_int(β_int) = L_orig(M · β_int)` via the chain rule.
5656    pub fn transform_penalized_hessian_to_original(&self, h_internal: &Array2<f64>) -> Array2<f64> {
5657        let right = self.right_multiply_by_m_inv(h_internal);
5658        self.left_multiply_by_m_inv_transpose(&right)
5659    }
5660
5661    pub fn internal_bounds_for(&self, col_idx: usize, min: f64, max: f64) -> (f64, f64) {
5662        if let Some(col) = self.columns.iter().find(|c| c.col_idx == col_idx) {
5663            (min * col.scale, max * col.scale)
5664        } else {
5665            (min, max)
5666        }
5667    }
5668}
5669
5670pub fn freeze_raw_spatial_metadata(metadata: BasisMetadata, raw_cols: usize) -> BasisMetadata {
5671    match metadata {
5672        BasisMetadata::ThinPlate {
5673            centers,
5674            length_scale,
5675            periodic,
5676            identifiability_transform: None,
5677            input_scale,
5678            radial_reparam,
5679        } => BasisMetadata::ThinPlate {
5680            centers,
5681            length_scale,
5682            periodic,
5683            identifiability_transform: Some(Array2::eye(raw_cols)),
5684            input_scale,
5685            radial_reparam,
5686        },
5687        BasisMetadata::Duchon {
5688            centers,
5689            length_scale,
5690            periodic,
5691            power,
5692            nullspace_order,
5693            identifiability_transform: None,
5694            input_scale,
5695            aniso_log_scales,
5696            operator_collocation_points,
5697            radial_reparam,
5698            spectral_basis,
5699        } => BasisMetadata::Duchon {
5700            centers,
5701            length_scale,
5702            periodic,
5703            power,
5704            nullspace_order,
5705            identifiability_transform: Some(Array2::eye(raw_cols)),
5706            input_scale,
5707            aniso_log_scales,
5708            operator_collocation_points,
5709            radial_reparam,
5710            spectral_basis,
5711        },
5712        other => other,
5713    }
5714}
5715
5716pub fn matern_operator_penalty_triplet_from_metadata(
5717    metadata: &BasisMetadata,
5718) -> Result<crate::basis::FilteredPenalties, BasisError> {
5719    let BasisMetadata::Matern {
5720        centers,
5721        length_scale,
5722        periodic,
5723        nu,
5724        include_intercept,
5725        identifiability_transform,
5726        aniso_log_scales,
5727        input_scale,
5728        ..
5729    } = metadata
5730    else {
5731        crate::bail_invalid_basis!("Matérn operator penalties require Matérn metadata");
5732    };
5733    // The metadata records `length_scale` in *original* (un-standardized) data
5734    // coordinates, while `centers` live in the *standardized* coordinate frame
5735    // (uniform division by `input_scale`). The realized design built the
5736    // kernel against those standardized centers using the compensated
5737    // effective length scale `length_scale / input_scale`. The collocation operators
5738    // here are evaluated on the same standardized centers, so they must use the
5739    // SAME effective length scale — otherwise the penalty regularizes a
5740    // different RKHS range than the design lives in, leaving rough coefficient
5741    // directions effectively unpenalized. That mismatch is benign in 1-D
5742    // (no standardization) but produces a catastrophic out-of-sample blow-up in
5743    // every dimension where the input scale differs from one (#706).
5744    // Since #2636 the two frames are distinct types, so this conversion is the
5745    // only way to obtain the standardized scalar the callee requires.
5746    let penalty_length_scale = input_scale
5747        .to_standardized_units(*length_scale)
5748        .standardized_value();
5749    matern_operator_penalty_triplet_at_length_scale(
5750        centers.view(),
5751        periodic.as_deref(),
5752        identifiability_transform.as_ref(),
5753        *nu,
5754        *include_intercept,
5755        aniso_log_scales.as_deref(),
5756        penalty_length_scale,
5757    )
5758}
5759
5760/// Build the canonical Matérn operator-penalty triplet (mass / tension /
5761/// stiffness) at an explicit **effective** length scale — i.e. the
5762/// isotropic-scale-compensated, standardized-frame scale the design's kernel was built
5763/// against (NOT the original-coordinate `length_scale` stored in metadata).
5764///
5765/// This is the SINGLE source of truth for the Matérn penalty topology. Two
5766/// callers route through it and must therefore stay byte-for-byte consistent:
5767///   * the cold/slow design rebuild (`matern_operator_penalty_triplet_from_metadata`,
5768///     compensating the frozen metadata `length_scale`), and
5769///   * the n-free κ-optimizer re-key (`FrozenTermCollectionIncrementalRealizer::
5770///     canonical_penalties_at_psi`, compensating the trial `ψ → exp(-ψ)` scale).
5771///
5772/// Sharing the body makes the penalty BLOCK COUNT and the per-block numerics
5773/// one deterministic function of `(geometry, ν, η, ℓ_eff)`. The active-operator
5774/// gate is `m = ν + d/2`, which is independent of ℓ, so the block count is
5775/// **ψ-stable by construction**: the re-key can never produce a different number
5776/// of blocks than the frozen design (the desync that #1270 hard-errored on).
5777pub fn matern_operator_penalty_triplet_at_length_scale(
5778    centers: ArrayView2<'_, f64>,
5779    periodic: Option<&[Option<f64>]>,
5780    identifiability_transform: Option<&Array2<f64>>,
5781    nu: crate::basis::MaternNu,
5782    include_intercept: bool,
5783    aniso_log_scales: Option<&[f64]>,
5784    effective_length_scale: f64,
5785) -> Result<crate::basis::FilteredPenalties, BasisError> {
5786    let penalty_centers = crate::basis::expand_periodic_centers(&centers.to_owned(), periodic)?;
5787    let ops = build_matern_collocation_operator_matrices(
5788        penalty_centers.view(),
5789        None,
5790        effective_length_scale,
5791        nu,
5792        include_intercept,
5793        identifiability_transform.map(|z| z.view()),
5794        aniso_log_scales,
5795    )?;
5796    // Gate operator dials on the Matérn-ν RKHS Sobolev order m = ν + d/2.
5797    // Derivative energies through j=m belong to H^m inclusively, so the 1-D
5798    // ν=3/2 kernel (m=2) carries stiffness as well as mass+tension. The sole
5799    // exception is ν=1/2: its center cusp makes collocated D1/D2 undefined and
5800    // it therefore retains mass only (#707). The matching topology gate lives
5801    // at `DuchonOperatorPenaltySpec::matern_for_smoothness`.
5802    const ORDER_EPS: f64 = 1e-9;
5803    let d = penalty_centers.ncols();
5804    let m = nu.half_integer_value() + 0.5 * d as f64;
5805    let mut candidates = Vec::with_capacity(3);
5806    for (raw, source, min_order) in [
5807        (ops.d0.t().dot(&ops.d0), PenaltySource::OperatorMass, 0.0),
5808        (ops.d1.t().dot(&ops.d1), PenaltySource::OperatorTension, 1.0),
5809        (
5810            ops.d2.t().dot(&ops.d2),
5811            PenaltySource::OperatorStiffness,
5812            2.0,
5813        ),
5814    ] {
5815        let nondifferentiable_ou = matches!(nu, crate::basis::MaternNu::Half);
5816        if min_order > 0.0 && (nondifferentiable_ou || m + ORDER_EPS < min_order) {
5817            continue;
5818        }
5819        let sym = (&raw + &raw.t()) * 0.5;
5820        let (matrix, normalization_scale) = normalize_penalty_in_constrained_space(&sym);
5821        candidates.push(PenaltyCandidate {
5822            matrix: ConstructiveQuadratic::try_from_dense_psd(matrix, "Matérn operator penalty")?,
5823            source,
5824            normalization_scale,
5825            kronecker_factors: None,
5826            op: None,
5827        });
5828    }
5829    filter_penalty_candidates(candidates)
5830}
5831
5832pub fn normalize_penalty_in_constrained_space(matrix: &Array2<f64>) -> (Array2<f64>, f64) {
5833    // Constrained-space normalization:
5834    //   c = ||S_con||_F,  S_tilde = S_con / c.
5835    // This is the only normalization coherent with a REML objective that is
5836    // evaluated entirely in constrained coordinates.
5837    let matrix = (matrix + &matrix.t().to_owned()) * 0.5;
5838    // Clamp noise-floor negative eigenvalues so β'Sβ is non-negative as a contract, not just in exact arithmetic.
5839    let matrix = crate::basis::project_penalty_to_psd_cone(&matrix);
5840    let c = matrix.iter().map(|v| v * v).sum::<f64>().sqrt();
5841    if c.is_finite() && c > 0.0 {
5842        (matrix.mapv(|v| v / c), c)
5843    } else {
5844        (matrix, 1.0)
5845    }
5846}
5847
5848pub fn tensor_product_design_from_sparse_marginals(
5849    marginal_sparse: &[&SparseColMat<usize, f64>],
5850) -> Result<SparseColMat<usize, f64>, BasisError> {
5851    if marginal_sparse.is_empty() {
5852        crate::bail_invalid_basis!("TensorBSpline requires at least one marginal basis");
5853    }
5854    let n = marginal_sparse[0].nrows();
5855    for (i, m) in marginal_sparse.iter().enumerate().skip(1) {
5856        if m.nrows() != n {
5857            crate::bail_dim_basis!(
5858                "tensor sparse marginal row mismatch at dim {i}: expected {n}, got {}",
5859                m.nrows()
5860            );
5861        }
5862    }
5863    let dims: Vec<usize> = marginal_sparse.iter().map(|m| m.ncols()).collect();
5864    let total_cols = dims.iter().try_fold(1usize, |acc, &q| {
5865        acc.checked_mul(q)
5866            .ok_or_else(|| BasisError::DimensionMismatch("tensor basis too large".to_string()))
5867    })?;
5868    let mut strides = vec![1usize; dims.len()];
5869    for d in (0..dims.len().saturating_sub(1)).rev() {
5870        strides[d] = strides[d + 1]
5871            .checked_mul(dims[d + 1])
5872            .ok_or_else(|| BasisError::DimensionMismatch("tensor basis too large".to_string()))?;
5873    }
5874
5875    use faer::sparse::SparseRowMat;
5876    let csrs: Vec<SparseRowMat<usize, f64>> = marginal_sparse
5877        .iter()
5878        .enumerate()
5879        .map(|(d, m)| {
5880            m.as_ref().to_row_major().map_err(|e| {
5881                BasisError::SparseCreation(format!(
5882                    "tensor sparse marginal {d} CSR conversion failed: {e:?}"
5883                ))
5884            })
5885        })
5886        .collect::<Result<Vec<_>, _>>()?;
5887    let row_ptrs: Vec<&[usize]> = csrs.iter().map(|c| c.symbolic().row_ptr()).collect();
5888    let col_idxs: Vec<&[usize]> = csrs.iter().map(|c| c.symbolic().col_idx()).collect();
5889    let vals: Vec<&[f64]> = csrs.iter().map(|c| c.val()).collect();
5890
5891    use rayon::prelude::*;
5892    const CHUNK: usize = 1024;
5893    let num_chunks = n.div_ceil(CHUNK);
5894    let per_chunk: Vec<Vec<Triplet<usize, usize, f64>>> = (0..num_chunks)
5895        .into_par_iter()
5896        .map(|chunk_idx| {
5897            let row_start = chunk_idx * CHUNK;
5898            let row_end = (row_start + CHUNK).min(n);
5899            let mut chunk_triplets = Vec::<Triplet<usize, usize, f64>>::new();
5900            let mut cur_cols = Vec::<usize>::with_capacity(64);
5901            let mut cur_vals = Vec::<f64>::with_capacity(64);
5902            let mut next_cols = Vec::<usize>::with_capacity(64);
5903            let mut next_vals = Vec::<f64>::with_capacity(64);
5904            for i in row_start..row_end {
5905                cur_cols.clear();
5906                cur_vals.clear();
5907                cur_cols.push(0);
5908                cur_vals.push(1.0);
5909                let mut row_is_zero = false;
5910                for d in 0..dims.len() {
5911                    let row_start_d = row_ptrs[d][i];
5912                    let row_end_d = row_ptrs[d][i + 1];
5913                    if row_start_d == row_end_d {
5914                        row_is_zero = true;
5915                        break;
5916                    }
5917                    let stride = strides[d];
5918                    next_cols.clear();
5919                    next_vals.clear();
5920                    next_cols.reserve(cur_cols.len() * (row_end_d - row_start_d));
5921                    next_vals.reserve(cur_vals.len() * (row_end_d - row_start_d));
5922                    for (&prev_col, &prev_val) in cur_cols.iter().zip(cur_vals.iter()) {
5923                        for ptr in row_start_d..row_end_d {
5924                            let cj = col_idxs[d][ptr];
5925                            let vj = vals[d][ptr];
5926                            next_cols.push(prev_col + cj * stride);
5927                            next_vals.push(prev_val * vj);
5928                        }
5929                    }
5930                    std::mem::swap(&mut cur_cols, &mut next_cols);
5931                    std::mem::swap(&mut cur_vals, &mut next_vals);
5932                }
5933                if row_is_zero {
5934                    continue;
5935                }
5936                for (&col, &val) in cur_cols.iter().zip(cur_vals.iter()) {
5937                    chunk_triplets.push(Triplet::new(i, col, val));
5938                }
5939            }
5940            chunk_triplets
5941        })
5942        .collect();
5943    let total_nnz: usize = per_chunk.iter().map(Vec::len).sum();
5944    let mut triplets = Vec::<Triplet<usize, usize, f64>>::with_capacity(total_nnz);
5945    for chunk in per_chunk {
5946        triplets.extend(chunk);
5947    }
5948    SparseColMat::try_new_from_triplets(n, total_cols, &triplets).map_err(|e| {
5949        BasisError::SparseCreation(format!(
5950            "failed to assemble sparse tensor product design: {e:?}"
5951        ))
5952    })
5953}
5954
5955pub fn dense_local_margin_to_sparse(
5956    dense: &Array2<f64>,
5957) -> Result<SparseColMat<usize, f64>, BasisError> {
5958    let expected_row_nnz = dense.ncols().min(4);
5959    let mut triplets =
5960        Vec::<Triplet<usize, usize, f64>>::with_capacity(dense.nrows() * expected_row_nnz);
5961    for ((row, col), &value) in dense.indexed_iter() {
5962        if value != 0.0 {
5963            triplets.push(Triplet::new(row, col, value));
5964        }
5965    }
5966    SparseColMat::try_new_from_triplets(dense.nrows(), dense.ncols(), &triplets).map_err(|e| {
5967        BasisError::SparseCreation(format!(
5968            "failed to convert tensor marginal design to sparse form: {e:?}"
5969        ))
5970    })
5971}
5972
5973pub struct TensorMarginRangeNullProjectors {
5974    range: Array2<f64>,
5975    null: Array2<f64>,
5976}
5977
5978pub fn projector_from_columns(columns: &Array2<f64>, indices: &[usize]) -> Array2<f64> {
5979    if indices.is_empty() {
5980        return Array2::<f64>::zeros((columns.nrows(), columns.nrows()));
5981    }
5982    let basis = columns.select(Axis(1), indices);
5983    basis.dot(&basis.t())
5984}
5985
5986pub fn tensor_margin_range_null_projectors(
5987    normalized_marginal_penalties: &[(Array2<f64>, f64)],
5988) -> Result<Vec<TensorMarginRangeNullProjectors>, BasisError> {
5989    normalized_marginal_penalties
5990        .iter()
5991        .enumerate()
5992        .map(|(dim, (penalty, _))| {
5993            let analysis = crate::basis::analyze_penalty_block(penalty)?;
5994            if analysis.rank == 0 {
5995                crate::bail_invalid_basis!(
5996                    "t2 separable tensor penalty margin {dim} has rank-zero penalty; \
5997                     cannot split penalized and null subspaces"
5998                );
5999            }
6000            let mut range_idx = Vec::<usize>::new();
6001            let mut null_idx = Vec::<usize>::new();
6002            for (idx, &ev) in analysis.eigenvalues.iter().enumerate() {
6003                if ev > analysis.rank_tol {
6004                    range_idx.push(idx);
6005                } else {
6006                    null_idx.push(idx);
6007                }
6008            }
6009            Ok(TensorMarginRangeNullProjectors {
6010                range: projector_from_columns(&analysis.eigenvectors, &range_idx),
6011                null: projector_from_columns(&analysis.eigenvectors, &null_idx),
6012            })
6013        })
6014        .collect()
6015}
6016
6017pub fn build_tensor_bspline_basis(
6018    data: ArrayView2<'_, f64>,
6019    feature_cols: &[usize],
6020    spec: &TensorBSplineSpec,
6021) -> Result<BasisBuildResult, BasisError> {
6022    if feature_cols.is_empty() {
6023        crate::bail_invalid_basis!("TensorBSpline requires at least one feature column");
6024    }
6025    if feature_cols.len() != spec.marginalspecs.len() {
6026        crate::bail_dim_basis!(
6027            "TensorBSpline feature/spec mismatch: feature_cols={}, marginalspecs={}",
6028            feature_cols.len(),
6029            spec.marginalspecs.len()
6030        );
6031    }
6032    if let Some((margin, _)) = spec
6033        .marginalspecs
6034        .iter()
6035        .enumerate()
6036        .find(|(_, marginal)| marginal.boundary_conditions.has_nonzero_anchor())
6037    {
6038        crate::bail_invalid_basis!(
6039            "TensorBSpline margin {margin} has a non-zero endpoint anchor. An inhomogeneous \
6040             marginal constraint cannot be represented by the tensor's homogeneous coefficient \
6041             chart plus one scalar row offset; use a separate anchored 1-D smooth or an explicit \
6042             model offset"
6043        );
6044    }
6045    if !spec.periods.is_empty() && spec.periods.len() != feature_cols.len() {
6046        crate::bail_dim_basis!(
6047            "TensorBSpline periods length {} does not match feature count {}",
6048            spec.periods.len(),
6049            feature_cols.len()
6050        );
6051    }
6052    let p = data.ncols();
6053    for &c in feature_cols {
6054        if c >= p {
6055            crate::bail_dim_basis!(
6056                "tensor feature column {c} is out of bounds for data with {p} columns"
6057            );
6058        }
6059    }
6060
6061    let mut marginal_knots = Vec::<Array1<f64>>::with_capacity(feature_cols.len());
6062    // Per-margin cr flag (#1074): `true` when the margin is a natural cubic
6063    // regression spline, so the tensor freeze rebuilds the cr knotspec.
6064    let mut marginal_is_cr_flags = Vec::<bool>::with_capacity(feature_cols.len());
6065    let mut marginal_degrees = Vec::<usize>::with_capacity(feature_cols.len());
6066    let mut marginalnum_basis = Vec::<usize>::with_capacity(feature_cols.len());
6067    let mut marginal_penalties = Vec::<Array2<f64>>::with_capacity(feature_cols.len());
6068    let mut marginal_function_grams = Vec::<Array2<f64>>::with_capacity(feature_cols.len());
6069    let mut marginal_designs = Vec::<Array2<f64>>::with_capacity(feature_cols.len());
6070    // Per-margin effective period: either user-set via `spec.periods` or
6071    // implied by a `PeriodicUniform` marginal knotspec (which the 1D B-spline
6072    // builder realizes as a cyclic B-spline basis).
6073    // Captured here so freeze→reload round-trips both routes back to a
6074    // `PeriodicUniform` marginal knotspec; otherwise a `PeriodicUniform`
6075    // margin specified without `spec.periods` would freeze as a plain
6076    // `Provided(knots)` open spline and lose its wrap-around at predict time.
6077    let mut marginal_effective_periods = Vec::<Option<f64>>::with_capacity(feature_cols.len());
6078    // Per-marginal sparse representation, populated when the 1D builder returned
6079    // a `DesignMatrix::Sparse`. Used to assemble the Khatri-Rao tensor product
6080    // sparsely (only ∏(degree+1) nonzeros per row) instead of densifying to
6081    // shape (n, ∏ q_j) up front. Periodic B-spline margins are local-support
6082    // bases too; when the 1D builder returns them densely, we convert that
6083    // marginal back to sparse form so cylinder/torus tensor products keep the
6084    // same scale behavior as open tensor products.
6085    let mut marginal_sparse =
6086        Vec::<Option<SparseColMat<usize, f64>>>::with_capacity(feature_cols.len());
6087
6088    // Reuse the robust 1D builder to ensure the same knot validation and
6089    // marginal difference-penalty construction as standalone smooth terms.
6090    for (dim, (&col, marginalspec)) in feature_cols
6091        .iter()
6092        .zip(spec.marginalspecs.iter())
6093        .enumerate()
6094    {
6095        // Tensor basis uses raw marginal knot-product columns. Applying 1D
6096        // identifiability constraints here would change marginal penalty sizes
6097        // without changing the tensor design construction, causing dimension
6098        // mismatch. Keep marginal builders unconstrained at this stage.
6099        let mut marginal_unconstrained = marginalspec.clone();
6100        marginal_unconstrained.identifiability = BSplineIdentifiability::None;
6101        let built = build_bspline_basis_1d(data.column(col), &marginal_unconstrained)?;
6102        // A cr (`NaturalCubicRegression`) margin emits `CubicRegression1D`
6103        // metadata whose `knots` are the k value-knots; a B-spline margin emits
6104        // `BSpline1D` with the clamped knot vector. Capture either so the
6105        // tensor freeze can rebuild the exact same marginal knotspec (#1074).
6106        let (knots, marginal_is_cr, effective_degree, function_gram) = match built.metadata {
6107            BasisMetadata::BSpline1D {
6108                knots,
6109                periodic,
6110                degree,
6111                ..
6112            } => {
6113                let effective_degree = degree.unwrap_or(marginal_unconstrained.degree);
6114                let gram = if spec.double_penalty {
6115                    Some(match periodic {
6116                        Some((start, period, num_basis)) => {
6117                            crate::basis::periodic_bspline_function_gram(
6118                                start,
6119                                start + period,
6120                                effective_degree,
6121                                num_basis,
6122                            )?
6123                        }
6124                        None => crate::basis::bspline_function_gram(&knots, effective_degree)?,
6125                    })
6126                } else {
6127                    None
6128                };
6129                (knots, false, effective_degree, gram)
6130            }
6131            BasisMetadata::CubicRegression1D { knots, .. } => {
6132                let gram = spec
6133                    .double_penalty
6134                    .then(|| crate::basis::cubic_regression_function_gram(&knots))
6135                    .transpose()?;
6136                (knots, true, marginalspec.degree, gram)
6137            }
6138            _ => {
6139                crate::bail_invalid_basis!(
6140                    "internal TensorBSpline error at dim {dim}: expected BSpline1D or CubicRegression1D metadata"
6141                );
6142            }
6143        };
6144        let metadata_knots = match marginalspec.knotspec {
6145            BSplineKnotSpec::PeriodicUniform {
6146                data_range,
6147                num_basis,
6148            } => Array1::linspace(data_range.0, data_range.1, num_basis),
6149            _ => knots,
6150        };
6151        if let Some(function_gram) = function_gram {
6152            if function_gram.dim() != (built.design.ncols(), built.design.ncols()) {
6153                crate::bail_dim_basis!(
6154                    "internal TensorBSpline error at dim {dim}: function Gram is {:?}, basis has {} columns",
6155                    function_gram.dim(),
6156                    built.design.ncols()
6157                );
6158            }
6159            marginal_function_grams.push(function_gram);
6160        }
6161        marginal_knots.push(metadata_knots);
6162        marginal_is_cr_flags.push(marginal_is_cr);
6163        marginal_degrees.push(effective_degree);
6164        marginalnum_basis.push(built.design.ncols());
6165        // Capture the sparse representation of this marginal (when the
6166        // 1D builder produced one) before densifying for the dense
6167        // marginal cache used by `tensor_product_design_from_marginals`
6168        // and `TensorProductDesignOperator`.
6169        let dense_marginal = built.design.to_dense();
6170        let sparse_view: Option<SparseColMat<usize, f64>> = match built.design.as_sparse() {
6171            Some(sd) => {
6172                let inner: &SparseColMat<usize, f64> = sd;
6173                Some(inner.clone())
6174            }
6175            None => match marginalspec.knotspec {
6176                BSplineKnotSpec::PeriodicUniform { .. } => {
6177                    Some(dense_local_margin_to_sparse(&dense_marginal)?)
6178                }
6179                _ => None,
6180            },
6181        };
6182        marginal_sparse.push(sparse_view);
6183        marginal_designs.push(dense_marginal);
6184        marginal_penalties.push(
6185            built
6186                .active_penalties
6187                .first()
6188                .ok_or_else(|| {
6189                    BasisError::InvalidInput(format!(
6190                        "internal TensorBSpline error at dim {dim}: missing marginal penalty"
6191                    ))
6192                })?
6193                .matrix
6194                .clone(),
6195        );
6196        built.active_penalties.first().ok_or_else(|| {
6197            BasisError::InvalidInput(format!(
6198                "internal TensorBSpline error at dim {dim}: missing marginal nullspace dim"
6199            ))
6200        })?;
6201        // A `PeriodicUniform` marginal knotspec implies the margin is
6202        // wrap-around: the 1D builder already realized it as a periodic
6203        // basis, so the tensor product inherits that periodicity. Record
6204        // the period derived from the knotspec's data range so freeze
6205        // restores `PeriodicUniform` on the marginal — otherwise the
6206        // round-trip downgrades it to `Provided(knots)` (an open spline)
6207        // and predict-time wraps disappear.
6208        let implied_period = match marginalspec.knotspec {
6209            BSplineKnotSpec::PeriodicUniform { data_range, .. } => {
6210                Some(data_range.1 - data_range.0)
6211            }
6212            _ => spec.periods.get(dim).and_then(|p| *p),
6213        };
6214        marginal_effective_periods.push(implied_period);
6215    }
6216
6217    let total_cols: usize = marginalnum_basis.iter().product();
6218    let mut dense_design = (!matches!(spec.identifiability, TensorBSplineIdentifiability::None))
6219        .then(|| tensor_product_design_from_marginals(&marginal_designs))
6220        .transpose()?;
6221    let mut candidates = Vec::<PenaltyCandidate>::with_capacity(
6222        match spec.penalty_decomposition {
6223            TensorBSplinePenaltyDecomposition::MarginalKroneckerSum => marginal_penalties.len(),
6224            TensorBSplinePenaltyDecomposition::Separable => marginal_penalties.len() * 2,
6225        } + if spec.double_penalty { 1 } else { 0 },
6226    );
6227
6228    // Tensor-product smoothing parameters are one-per-margin.  Therefore the
6229    // physical penalty attached to a margin must be normalized in that margin's
6230    // own working coordinates before it is embedded in the full tensor product.
6231    // Normalizing only the already-Kroneckered matrix would fold arbitrary
6232    // dimension-dependent identity factors into the margin's lambda and would
6233    // make anisotropic REML/LAML smoothing depend on the other margins' basis
6234    // sizes rather than on the marginal roughness operator itself.
6235    let normalized_marginal_penalties: Vec<(Array2<f64>, f64)> = marginal_penalties
6236        .iter()
6237        .map(normalize_penalty_in_constrained_space)
6238        .collect();
6239    let tensor_function_gram = if spec.double_penalty {
6240        if marginal_function_grams.len() != marginalnum_basis.len() {
6241            crate::bail_dim_basis!(
6242                "TensorBSpline double penalty requires one function Gram per margin; got {} for {} margins",
6243                marginal_function_grams.len(),
6244                marginalnum_basis.len()
6245            );
6246        }
6247        let mut gram = Array2::<f64>::eye(1);
6248        for marginal_gram in &marginal_function_grams {
6249            gram = kronecker_product(&gram, marginal_gram);
6250        }
6251        Some(gram)
6252    } else {
6253        None
6254    };
6255    // A single PSD sum has exactly the joint null space shared by every
6256    // marginal roughness block. It is used only to define the global
6257    // null-component penalty; the ordinary tensor candidates below retain
6258    // their one-coordinate-per-margin decomposition.
6259    let joint_wiggliness = if spec.double_penalty {
6260        let mut sum = Array2::<f64>::zeros((total_cols, total_cols));
6261        for dim in 0..normalized_marginal_penalties.len() {
6262            let mut embedded = Array2::<f64>::eye(1);
6263            for (margin, &width) in marginalnum_basis.iter().enumerate() {
6264                let factor = if margin == dim {
6265                    normalized_marginal_penalties[margin].0.clone()
6266                } else {
6267                    Array2::<f64>::eye(width)
6268                };
6269                embedded = kronecker_product(&embedded, &factor);
6270            }
6271            sum += &embedded;
6272        }
6273        Some(sum)
6274    } else {
6275        None
6276    };
6277    let mut kronecker_marginal_penalties =
6278        Vec::<Array2<f64>>::with_capacity(normalized_marginal_penalties.len());
6279
6280    match spec.penalty_decomposition {
6281        TensorBSplinePenaltyDecomposition::MarginalKroneckerSum => {
6282            // Accumulate the Kronecker-sum of the per-margin penalties,
6283            // `Σ_dim S_dim`, whose null space is exactly the *joint* null space
6284            // of all marginal penalties — the tensor of marginal polynomial
6285            // null spaces. The tensor double penalty (below) shrinks only this
6286            // joint null, never the already-penalized interaction range.
6287            for dim in 0..normalized_marginal_penalties.len() {
6288                let mut s_dim = Array2::<f64>::eye(1);
6289                let mut factors = Vec::<Array2<f64>>::with_capacity(marginalnum_basis.len());
6290                for (j, &qj) in marginalnum_basis.iter().enumerate() {
6291                    let factor = if j == dim {
6292                        normalized_marginal_penalties[j].0.clone()
6293                    } else {
6294                        Array2::<f64>::eye(qj)
6295                    };
6296                    factors.push(factor.clone());
6297                    s_dim = kronecker_product(&s_dim, &factor);
6298                }
6299                if dim == kronecker_marginal_penalties.len() {
6300                    kronecker_marginal_penalties.push(normalized_marginal_penalties[dim].0.clone());
6301                }
6302                candidates.push(PenaltyCandidate {
6303                    matrix: ConstructiveQuadratic::try_from_dense_psd(
6304                        s_dim,
6305                        "tensor marginal penalty",
6306                    )?,
6307                    source: PenaltySource::TensorMarginal { dim },
6308                    normalization_scale: normalized_marginal_penalties[dim].1,
6309                    kronecker_factors: Some(factors),
6310                    op: None,
6311                });
6312            }
6313
6314            if let (Some(primary), Some(gram)) =
6315                (joint_wiggliness.as_ref(), tensor_function_gram.as_ref())
6316                && let Some(shrink) =
6317                    crate::basis::function_space_nullspace_shrinkage(primary, gram)?
6318            {
6319                let (matrix, normalization_scale) = normalize_penalty_in_constrained_space(&shrink);
6320                candidates.push(PenaltyCandidate {
6321                    matrix: ConstructiveQuadratic::try_from_dense_psd(
6322                        matrix,
6323                        "tensor global null-function ridge",
6324                    )?,
6325                    source: PenaltySource::TensorGlobalRidge,
6326                    normalization_scale,
6327                    kronecker_factors: None,
6328                    op: None,
6329                });
6330            }
6331        }
6332        TensorBSplinePenaltyDecomposition::Separable => {
6333            let projectors = tensor_margin_range_null_projectors(&normalized_marginal_penalties)?;
6334            let n_masks = 1usize.checked_shl(projectors.len() as u32).ok_or_else(|| {
6335                BasisError::InvalidInput(format!(
6336                    "t2 separable tensor penalty supports at most {} margins, got {}",
6337                    usize::BITS - 1,
6338                    projectors.len()
6339                ))
6340            })?;
6341            for mask in 1..n_masks {
6342                let mut matrix = Array2::<f64>::eye(1);
6343                let mut factors = Vec::<Array2<f64>>::with_capacity(projectors.len());
6344                let mut penalized_margins = Vec::<usize>::new();
6345                for (dim, projector) in projectors.iter().enumerate() {
6346                    let use_range = ((mask >> dim) & 1) == 1;
6347                    let factor = if use_range {
6348                        penalized_margins.push(dim);
6349                        projector.range.clone()
6350                    } else {
6351                        projector.null.clone()
6352                    };
6353                    matrix = kronecker_product(&matrix, &factor);
6354                    factors.push(factor);
6355                }
6356                let (matrix, normalization_scale) = normalize_penalty_in_constrained_space(&matrix);
6357                candidates.push(PenaltyCandidate {
6358                    matrix: ConstructiveQuadratic::try_from_dense_psd(
6359                        matrix,
6360                        "tensor separable penalty",
6361                    )?,
6362                    source: PenaltySource::TensorSeparable { penalized_margins },
6363                    normalization_scale,
6364                    kronecker_factors: Some(factors),
6365                    op: None,
6366                });
6367            }
6368
6369            if let (Some(primary), Some(gram)) =
6370                (joint_wiggliness.as_ref(), tensor_function_gram.as_ref())
6371                && let Some(matrix) =
6372                    crate::basis::function_space_nullspace_shrinkage(primary, gram)?
6373            {
6374                let (matrix, normalization_scale) = normalize_penalty_in_constrained_space(&matrix);
6375                candidates.push(PenaltyCandidate {
6376                    matrix: ConstructiveQuadratic::try_from_dense_psd(
6377                        matrix,
6378                        "separable tensor global null-function ridge",
6379                    )?,
6380                    source: PenaltySource::TensorGlobalRidge,
6381                    normalization_scale,
6382                    kronecker_factors: None,
6383                    op: None,
6384                });
6385            }
6386        }
6387    }
6388
6389    let z_opt = match &spec.identifiability {
6390        TensorBSplineIdentifiability::None => None,
6391        TensorBSplineIdentifiability::SumToZero => {
6392            if total_cols < 2 {
6393                crate::bail_invalid_basis!(
6394                    "TensorBSpline requires at least 2 basis coefficients to enforce sum-to-zero identifiability"
6395                );
6396            }
6397            let dense_design_ref = dense_design.as_ref().ok_or_else(|| {
6398                BasisError::InvalidInput(
6399                    "tensor sum-to-zero identifiability requires a realized basis".to_string(),
6400                )
6401            })?;
6402            let (_, z) = apply_sum_to_zero_constraint(dense_design_ref.view(), None)?;
6403            let gauge = gam_problem::Gauge::sum_to_zero(z);
6404            Some(gauge.block_transform(0))
6405        }
6406        TensorBSplineIdentifiability::MarginalSumToZero => {
6407            // `ti(...)`: drop the marginal main effects by centering every
6408            // margin independently, then form the tensor product of the
6409            // centered margins. Concretely, each margin `j` is reparameterized
6410            // by its own sum-to-zero null basis `Z_j` (so the constant — i.e.
6411            // the marginal intercept — is removed from that axis), and the
6412            // combined reparameterization is the Kronecker product
6413            // `Z = Z₀ ⊗ Z₁ ⊗ … ⊗ Z_{d-1}`. Applying `Z` to the full-tensor
6414            // design `B = B₀ ⊗ … ⊗ B_{d-1}` yields `B Z = (B₀ Z₀) ⊗ … ⊗
6415            // (B_{d-1} Z_{d-1})`, the tensor product of the centered margins,
6416            // which by construction contains no pure main effect.
6417            if marginal_designs.len() < 2 {
6418                crate::bail_invalid_basis!(
6419                    "tensor interaction (ti) identifiability requires at least 2 margins"
6420                );
6421            }
6422            let mut z = Array2::<f64>::eye(1);
6423            for (dim, marginal) in marginal_designs.iter().enumerate() {
6424                if marginal.ncols() < 2 {
6425                    crate::bail_invalid_basis!(
6426                        "tensor interaction (ti) margin {dim} has fewer than 2 basis functions; \
6427                         cannot remove its marginal main effect"
6428                    );
6429                }
6430                let (_, z_dim) = apply_sum_to_zero_constraint(marginal.view(), None)?;
6431                let gauge_dim = gam_problem::Gauge::sum_to_zero(z_dim);
6432                let z_dim = gauge_dim.block_transform(0);
6433                z = kronecker_product(&z, &z_dim);
6434            }
6435            Some(z)
6436        }
6437        TensorBSplineIdentifiability::FrozenTransform { transform } => {
6438            if transform.nrows() != total_cols {
6439                crate::bail_dim_basis!(
6440                    "frozen tensor identifiability transform mismatch: design has {} columns but transform has {} rows",
6441                    total_cols,
6442                    transform.nrows()
6443                );
6444            }
6445            Some(transform.clone())
6446        }
6447    };
6448
6449    if let Some(z) = z_opt.as_ref() {
6450        let gauge = gam_problem::Gauge::from_block_transforms(&[z.clone()]);
6451        let dense = dense_design.as_mut().ok_or_else(|| {
6452            BasisError::InvalidInput(
6453                "tensor identifiability transform requires a realized basis".to_string(),
6454            )
6455        })?;
6456        let restricted_design = gauge.restrict_design(dense);
6457        *dense = restricted_design;
6458        candidates = candidates
6459            .into_iter()
6460            .map(|candidate| -> Result<PenaltyCandidate, BasisError> {
6461                let restricted = candidate
6462                    .matrix
6463                    .restricted(&gauge, "tensor identifiability restriction")?;
6464                // Re-normalize in the *actual* coefficient chart used by the
6465                // fit.  The tensor sum-to-zero transform is not norm-preserving
6466                // for each overlapping marginal penalty, so carrying the raw
6467                // marginal Frobenius scale into the restricted space changes the
6468                // relative amount of smoothing seen by the LAML/REML optimizer.
6469                // Keep the physical scale in metadata and give the optimizer
6470                // unit-scale constrained penalties for every tensor margin.
6471                let (_, c_new) = normalize_penalty_in_constrained_space(restricted.dense());
6472                let matrix = restricted.scaled(
6473                    1.0 / c_new,
6474                    "normalized tensor penalty after identifiability",
6475                )?;
6476                Ok(PenaltyCandidate {
6477                    matrix,
6478                    source: candidate.source,
6479                    normalization_scale: candidate.normalization_scale * c_new,
6480                    // Z^T S Z is no longer a Kronecker product of the original
6481                    // marginal factors, so the Kronecker fast path in construction.rs
6482                    // must not be taken. Clearing kronecker_factors forces the generic
6483                    // block-local eigendecomposition path, which operates on the
6484                    // transformed matrix and is correct.
6485                    kronecker_factors: None,
6486                    op: candidate.op.clone(),
6487                })
6488            })
6489            .collect::<Result<Vec<_>, _>>()?;
6490
6491        if candidates
6492            .iter()
6493            .any(|candidate| matches!(candidate.source, PenaltySource::TensorGlobalRidge))
6494        {
6495            let width = candidates
6496                .first()
6497                .ok_or_else(|| {
6498                    BasisError::InvalidInput(
6499                        "TensorBSpline global ridge has no penalty candidates".to_string(),
6500                    )
6501                })?
6502                .matrix
6503                .nrows();
6504            let physical_primary_terms = candidates
6505                .iter()
6506                .filter(|candidate| !matches!(candidate.source, PenaltySource::TensorGlobalRidge))
6507                .map(|candidate| {
6508                    candidate.matrix.scaled(
6509                        candidate.normalization_scale,
6510                        "physical tensor primary penalty",
6511                    )
6512                })
6513                .collect::<Result<Vec<_>, _>>()?;
6514            let joint_primary = ConstructiveQuadratic::sum(
6515                &physical_primary_terms,
6516                "joint tensor primary penalty",
6517            )?;
6518            for candidate in &mut candidates {
6519                if !matches!(candidate.source, PenaltySource::TensorGlobalRidge) {
6520                    continue;
6521                }
6522                let physical_ridge = candidate
6523                    .matrix
6524                    .scaled(candidate.normalization_scale, "physical tensor null ridge")?;
6525                match crate::basis::rebuild_metric_consistent_ridge(
6526                    &joint_primary,
6527                    &physical_ridge,
6528                )? {
6529                    Some(rebuilt) => {
6530                        let (_, scale) = normalize_penalty_in_constrained_space(rebuilt.dense());
6531                        candidate.matrix =
6532                            rebuilt.scaled(1.0 / scale, "normalized rebuilt tensor null ridge")?;
6533                        candidate.normalization_scale = scale;
6534                    }
6535                    None => {
6536                        candidate.matrix = ConstructiveQuadratic::zero(width);
6537                        candidate.normalization_scale = 1.0;
6538                    }
6539                }
6540                candidate.kronecker_factors = None;
6541                candidate.op = None;
6542            }
6543        }
6544    }
6545
6546    let filtered = filter_penalty_candidates(candidates)?;
6547    let identifiability_is_none =
6548        matches!(spec.identifiability, TensorBSplineIdentifiability::None);
6549    // All marginals expose a sparse representation iff each `marginal_sparse`
6550    // slot is `Some(...)`. Currently this is true when every marginal is a
6551    // free-boundary, non-periodic 1D B-spline returned as
6552    // `DesignMatrix::Sparse` from `build_bspline_basis_1d`. Periodic B-splines
6553    // and other dense-only marginals leave a `None` and trigger the fall-back
6554    // path. Identifiability transforms (`SumToZero`, `FrozenTransform`) make
6555    // the tensor design dense in general, so we also gate on that.
6556    let all_marginals_sparse = marginal_sparse.iter().all(Option::is_some);
6557    let design = if let Some(dense_design) = dense_design {
6558        DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(dense_design))
6559    } else if identifiability_is_none && all_marginals_sparse {
6560        // Sparse Khatri-Rao path: assemble the (n, ∏ q_j) tensor product
6561        // directly as a SparseColMat, preserving the ∏(degree_j+1) nonzero
6562        // structure per row instead of densifying to ∏ q_j columns. This is
6563        // mathematically identical to `tensor_product_design_from_marginals`
6564        // applied to the corresponding dense marginals.
6565        let sparse_marginals: Vec<&SparseColMat<usize, f64>> = marginal_sparse
6566            .iter()
6567            .map(|m| m.as_ref().expect("all_marginals_sparse just verified"))
6568            .collect();
6569        let sparse_design = tensor_product_design_from_sparse_marginals(&sparse_marginals)?;
6570        DesignMatrix::Sparse(gam_linalg::matrix::SparseDesignMatrix::new(sparse_design))
6571    } else {
6572        let marginals: Vec<Arc<Array2<f64>>> = marginal_designs
6573            .iter()
6574            .map(|m| Arc::new(m.clone()))
6575            .collect();
6576        let op = TensorProductDesignOperator::new(marginals).map_err(|e| {
6577            BasisError::InvalidInput(format!("TensorProductDesignOperator build failed: {e}"))
6578        })?;
6579        DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(op)))
6580    };
6581
6582    Ok(BasisBuildResult {
6583        design,
6584        affine_offset: None,
6585        active_penalties: filtered.active,
6586        dropped_penalties: filtered.dropped,
6587        joint_null_rotation: None,
6588        metadata: BasisMetadata::TensorBSpline {
6589            feature_cols: feature_cols.to_vec(),
6590            knots: marginal_knots,
6591            degrees: marginal_degrees,
6592            // Prefer the per-margin effective period derived in the loop —
6593            // it captures both the explicit `spec.periods` route and the
6594            // implied period from a `PeriodicUniform` marginal knotspec.
6595            // Falling back to `spec.periods` when populated keeps any
6596            // user-supplied explicit period authoritative even if the
6597            // marginal knotspec carried no periodicity hint.
6598            periods: marginal_effective_periods,
6599            is_cr: marginal_is_cr_flags,
6600            identifiability_transform: z_opt,
6601        },
6602        // The current Kronecker runtime diagonalizes only the marginal
6603        // roughness operators and represents its optional joint-null block as
6604        // a Euclidean selector. A function-space ridge generally does not
6605        // commute with those marginals, so advertising it as factored would
6606        // make PIRLS and REML solve a different objective. Keep the exact
6607        // canonical matrices whenever null recovery is active.
6608        kronecker_factored: if !spec.double_penalty
6609            && matches!(spec.identifiability, TensorBSplineIdentifiability::None)
6610            && matches!(
6611                spec.penalty_decomposition,
6612                TensorBSplinePenaltyDecomposition::MarginalKroneckerSum
6613            ) {
6614            Some(KroneckerFactoredBasis::new(
6615                marginal_designs,
6616                kronecker_marginal_penalties,
6617                marginalnum_basis.clone(),
6618                spec.double_penalty,
6619            ))
6620        } else {
6621            None
6622        },
6623    })
6624}
6625
6626#[cfg(test)]
6627mod tensor_function_space_runtime_tests {
6628    use super::*;
6629    use crate::basis::{
6630        BSplineBoundaryConditions, BSplineEndpointBoundaryCondition, OneDimensionalBoundary,
6631    };
6632    use ndarray::array;
6633
6634    fn marginal() -> BSplineBasisSpec {
6635        BSplineBasisSpec {
6636            degree: 2,
6637            penalty_order: 1,
6638            knotspec: BSplineKnotSpec::Generate {
6639                data_range: (0.0, 1.0),
6640                num_internal_knots: 2,
6641            },
6642            double_penalty: false,
6643            identifiability: BSplineIdentifiability::None,
6644            boundary: OneDimensionalBoundary::Open,
6645            boundary_conditions: BSplineBoundaryConditions::default(),
6646        }
6647    }
6648
6649    #[test]
6650    fn function_space_tensor_ridge_uses_exact_canonical_runtime() {
6651        let data = array![
6652            [0.00, 0.13],
6653            [0.15, 0.82],
6654            [0.29, 0.37],
6655            [0.43, 0.95],
6656            [0.58, 0.21],
6657            [0.71, 0.66],
6658            [0.86, 0.48],
6659            [1.00, 0.04]
6660        ];
6661        let mut spec = TensorBSplineSpec {
6662            marginalspecs: vec![marginal(), marginal()],
6663            periods: Vec::new(),
6664            double_penalty: true,
6665            identifiability: TensorBSplineIdentifiability::None,
6666            penalty_decomposition: TensorBSplinePenaltyDecomposition::MarginalKroneckerSum,
6667        };
6668        let built = build_tensor_bspline_basis(data.view(), &[0, 1], &spec)
6669            .expect("double-penalty tensor basis");
6670        assert!(
6671            built
6672                .active_penalties
6673                .iter()
6674                .any(|penalty| { matches!(penalty.info.source, PenaltySource::TensorGlobalRidge) })
6675        );
6676        assert!(
6677            built.kronecker_factored.is_none(),
6678            "the legacy factored runtime cannot represent a function-metric global ridge"
6679        );
6680
6681        spec.double_penalty = false;
6682        let singly_penalized = build_tensor_bspline_basis(data.view(), &[0, 1], &spec)
6683            .expect("single-penalty tensor basis");
6684        assert!(
6685            singly_penalized.kronecker_factored.is_some(),
6686            "the exact marginal-only fast path must remain available"
6687        );
6688    }
6689
6690    #[test]
6691    fn tensor_nonzero_anchor_is_rejected_before_its_affine_lift_can_be_dropped() {
6692        let data = array![[0.0, 0.0], [0.25, 0.75], [0.75, 0.25], [1.0, 1.0]];
6693        let mut anchored = marginal();
6694        anchored.boundary_conditions.left =
6695            BSplineEndpointBoundaryCondition::Anchored { value: 1.25 };
6696        let spec = TensorBSplineSpec {
6697            marginalspecs: vec![anchored, marginal()],
6698            periods: Vec::new(),
6699            double_penalty: false,
6700            identifiability: TensorBSplineIdentifiability::None,
6701            penalty_decomposition: TensorBSplinePenaltyDecomposition::MarginalKroneckerSum,
6702        };
6703
6704        let error = build_tensor_bspline_basis(data.view(), &[0, 1], &spec)
6705            .expect_err("a tensor margin cannot silently discard an inhomogeneous lift");
6706        let message = error.to_string();
6707        assert!(message.contains("TensorBSpline margin 0"));
6708        assert!(message.contains("non-zero endpoint anchor"));
6709        assert!(message.contains("explicit model offset"));
6710    }
6711}
6712
6713pub fn tensor_product_design_from_marginals(
6714    marginal_designs: &[Array2<f64>],
6715) -> Result<Array2<f64>, BasisError> {
6716    if marginal_designs.is_empty() {
6717        crate::bail_invalid_basis!("TensorBSpline requires at least one marginal basis");
6718    }
6719    let n = marginal_designs[0].nrows();
6720    for (i, b) in marginal_designs.iter().enumerate().skip(1) {
6721        if b.nrows() != n {
6722            crate::bail_dim_basis!(
6723                "tensor marginal row mismatch at dim {i}: expected {n}, got {}",
6724                b.nrows()
6725            );
6726        }
6727    }
6728    let total_cols = marginal_designs.iter().try_fold(1usize, |acc, b| {
6729        acc.checked_mul(b.ncols())
6730            .ok_or_else(|| BasisError::DimensionMismatch("tensor basis too large".to_string()))
6731    })?;
6732    // Tensor-product Khatri-Rao: design[i, j] = Π_d marginal_d[i, j_d]
6733    // where j is the multi-index (j_1, ..., j_D) flattened. Independent
6734    // across rows; parallelize row chunks and fill the pre-allocated
6735    // contiguous Array2 in place (no Vec-flatten-collect intermediate,
6736    // which doubled the peak memory at large-scale N).
6737    use ndarray::parallel::prelude::*;
6738    use rayon::iter::{IntoParallelIterator, ParallelIterator};
6739    let mut design = Array2::<f64>::zeros((n, total_cols));
6740    design
6741        .axis_chunks_iter_mut(ndarray::Axis(0), 1024)
6742        .into_par_iter()
6743        .enumerate()
6744        .for_each(|(chunk_idx, mut block)| {
6745            let row_offset = chunk_idx * 1024;
6746            // Scratch buffers reused across rows in this chunk.
6747            let mut cur = Vec::<f64>::with_capacity(total_cols);
6748            let mut next = Vec::<f64>::with_capacity(total_cols);
6749            for (local_i, mut out_row) in block.outer_iter_mut().enumerate() {
6750                let i = row_offset + local_i;
6751                cur.clear();
6752                cur.push(1.0);
6753                for b in marginal_designs {
6754                    let q = b.ncols();
6755                    next.clear();
6756                    next.resize(cur.len() * q, 0.0);
6757                    // Hoist the row view out of the inner `col` loop so the
6758                    // q reads per `a_idx` reuse a single contiguous slice
6759                    // instead of recomputing `b[[i, col]]` strides per cell.
6760                    let b_row = b.row(i);
6761                    let b_slice = b_row
6762                        .as_slice()
6763                        .expect("Array2 row from outer_iter is contiguous");
6764                    for (a_idx, &aval) in cur.iter().enumerate() {
6765                        let off = a_idx * q;
6766                        let dst = &mut next[off..off + q];
6767                        for col in 0..q {
6768                            dst[col] = aval * b_slice[col];
6769                        }
6770                    }
6771                    std::mem::swap(&mut cur, &mut next);
6772                }
6773                // `out_row` is a row of the contiguous C-major `design`
6774                // Array2, so it is backed by a contiguous slice. Use a
6775                // bulk slice copy instead of an element-by-element write
6776                // loop.
6777                let out_slice = out_row
6778                    .as_slice_mut()
6779                    .expect("design row is contiguous in C-major Array2");
6780                out_slice.copy_from_slice(&cur);
6781            }
6782        });
6783    Ok(design)
6784}
6785
6786/// Render a numeric factor level for an error message: an integer-valued code
6787/// (`1999.0`) prints as `1999`, so an unseen-level message names the level the
6788/// user actually wrote rather than a spurious `.0`.
6789fn fmt_level_value(v: f64) -> String {
6790    if v.is_finite() && v.fract() == 0.0 && v.abs() < 1e15 {
6791        format!("{}", v as i64)
6792    } else {
6793        format!("{v}")
6794    }
6795}
6796
6797pub fn build_random_effect_block(
6798    data: ArrayView2<'_, f64>,
6799    spec: &RandomEffectTermSpec,
6800) -> Result<RandomEffectBlock, BasisError> {
6801    let n = data.nrows();
6802    let p = data.ncols();
6803    if spec.feature_col >= p {
6804        crate::bail_dim_basis!(
6805            "random-effect term '{}' feature column {} out of bounds for {} columns",
6806            spec.name,
6807            spec.feature_col,
6808            p
6809        );
6810    }
6811
6812    let col = data.column(spec.feature_col);
6813    if col.iter().any(|v| !v.is_finite()) {
6814        crate::bail_invalid_basis!(
6815            "random-effect term '{}' contains non-finite group values",
6816            spec.name
6817        );
6818    }
6819
6820    let kept_levels: Vec<u64> = if let Some(levels) = spec.frozen_levels.as_ref() {
6821        if levels.is_empty() {
6822            crate::bail_invalid_basis!(
6823                "random-effect term '{}' has empty frozen_levels",
6824                spec.name
6825            );
6826        }
6827        // Canonicalize a possibly-legacy frozen set: a `-0.0` group interned
6828        // before signed-zero canonicalization landed would otherwise never match
6829        // a canonicalized data row. Idempotent on already-canonical sets (#2145).
6830        levels
6831            .iter()
6832            .map(|&b| gam_data::canonical_level_bits(f64::from_bits(b)))
6833            .collect()
6834    } else {
6835        let mut seen = BTreeSet::<u64>::new();
6836        let mut levels = Vec::<u64>::new();
6837        for &v in col {
6838            let bits = gam_data::canonical_level_bits(v);
6839            if seen.insert(bits) {
6840                levels.push(bits);
6841            }
6842        }
6843        if levels.is_empty() {
6844            crate::bail_invalid_basis!("random-effect term '{}' has no observed levels", spec.name);
6845        }
6846        let start_idx = if spec.drop_first_level && levels.len() > 1 {
6847            1usize
6848        } else {
6849            0usize
6850        };
6851        levels[start_idx..].to_vec()
6852    };
6853
6854    if kept_levels.is_empty() {
6855        crate::bail_invalid_basis!(
6856            "random-effect term '{}' drops all levels; keep at least one level",
6857            spec.name
6858        );
6859    }
6860
6861    let q = kept_levels.len();
6862    let mut level_to_col = BTreeMap::<u64, usize>::new();
6863    for (idx, &bits) in kept_levels.iter().enumerate() {
6864        if level_to_col.insert(bits, idx).is_some() {
6865            crate::bail_invalid_basis!(
6866                "random-effect term '{}' has duplicate frozen level bits {bits}",
6867                spec.name
6868            );
6869        }
6870    }
6871    // A FIXED categorical factor (`factor(g)` or a bare `+ g`; `lenient_unseen
6872    // == false`) must reject an out-of-vocabulary level rather than silently
6873    // encode it as an all-zero dummy row that collapses onto the factor's
6874    // centering point (#2137/#2102). For a *string* factor the typed schema
6875    // encode rejects the unseen level before we get here; a *numeric-coded*
6876    // `factor(year)` column, however, reaches Rust as a plain numeric column
6877    // with no categorical schema, so the operator that owns the frozen level
6878    // vocabulary is the enforcement point that closes the same gap. Only when
6879    // the full one-hot block is kept (`!drop_first_level`) does an absent level
6880    // unambiguously mean "unseen" — with treatment coding the dropped baseline
6881    // is a legitimate absent column, so we do not gate that path. `frozen_levels`
6882    // presence marks the predict/frozen context; at fit the vocabulary is
6883    // derived from this very data, so no row is unseen.
6884    let strict_unseen =
6885        !spec.lenient_unseen && !spec.drop_first_level && spec.frozen_levels.is_some();
6886    let mut group_ids = Vec::with_capacity(n);
6887    for (row, &v) in col.iter().enumerate() {
6888        let bits = gam_data::canonical_level_bits(v);
6889        let group_id = level_to_col.get(&bits).copied();
6890        if strict_unseen && group_id.is_none() {
6891            crate::bail_invalid_basis!(
6892                "unseen level '{}' in fixed factor column '{}' at row {}; the factor's levels \
6893                 were fixed at fit time and an out-of-vocabulary level cannot be predicted \
6894                 (use group({}) for a random effect that tolerates held-out levels)",
6895                fmt_level_value(v),
6896                spec.name,
6897                row,
6898                spec.name
6899            );
6900        }
6901        group_ids.push(group_id);
6902    }
6903
6904    Ok(RandomEffectBlock {
6905        name: spec.name.clone(),
6906        group_ids,
6907        num_groups: q,
6908        kept_levels,
6909    })
6910}
6911
6912#[cfg(test)]
6913mod random_effect_signed_zero_tests {
6914    use super::{RandomEffectTermSpec, build_random_effect_block};
6915    use ndarray::array;
6916
6917    fn spec() -> RandomEffectTermSpec {
6918        RandomEffectTermSpec {
6919            name: "g".to_string(),
6920            feature_col: 0,
6921            drop_first_level: false,
6922            penalized: true,
6923            frozen_levels: None,
6924            lenient_unseen: true,
6925        }
6926    }
6927
6928    #[test]
6929    fn signed_zero_rows_share_one_group() {
6930        // A column mixing +0.0 and -0.0 for the physically same group must
6931        // intern as ONE level, and every row (either spelling) must resolve to
6932        // that single group column — the #2145 fit-side regression.
6933        let data = array![[-0.0_f64], [0.0], [1.0], [-0.0], [1.0]];
6934        let block = build_random_effect_block(data.view(), &spec()).unwrap();
6935        assert_eq!(
6936            block.num_groups, 2,
6937            "0.0/-0.0 must not split into two groups"
6938        );
6939        // Rows 0,1,3 are the same group; rows 2,4 the other.
6940        assert_eq!(block.group_ids[0], block.group_ids[1]);
6941        assert_eq!(block.group_ids[0], block.group_ids[3]);
6942        assert_eq!(block.group_ids[2], block.group_ids[4]);
6943        assert_ne!(block.group_ids[0], block.group_ids[2]);
6944    }
6945
6946    #[test]
6947    fn frozen_positive_zero_matches_negative_zero_row() {
6948        // A model frozen on +0.0 must resolve a -0.0 prediction row to the same
6949        // column — the #2145 predict-side regression that dropped the effect.
6950        let mut s = spec();
6951        s.frozen_levels = Some(vec![0.0_f64.to_bits(), 1.0_f64.to_bits()]);
6952        let data = array![[-0.0_f64], [1.0]];
6953        let block = build_random_effect_block(data.view(), &s).unwrap();
6954        assert_eq!(
6955            block.group_ids[0],
6956            Some(0),
6957            "-0.0 must match the +0.0 column"
6958        );
6959        assert_eq!(block.group_ids[1], Some(1));
6960    }
6961
6962    #[test]
6963    fn frozen_negative_zero_matches_positive_zero_row() {
6964        // The symmetric direction: a legacy model interned on -0.0 (pre-fix)
6965        // must still resolve a +0.0 prediction row after canonicalization.
6966        let mut s = spec();
6967        s.frozen_levels = Some(vec![(-0.0_f64).to_bits(), 1.0_f64.to_bits()]);
6968        let data = array![[0.0_f64], [1.0]];
6969        let block = build_random_effect_block(data.view(), &s).unwrap();
6970        assert_eq!(
6971            block.group_ids[0],
6972            Some(0),
6973            "+0.0 must match the -0.0 column"
6974        );
6975    }
6976
6977    // ---- #2137: fixed factor (`factor(g)`) strict-unseen enforcement --------
6978
6979    fn fixed_factor_spec() -> RandomEffectTermSpec {
6980        // A numeric-coded `factor(year)`: full one-hot (`drop_first_level=false`),
6981        // FIXED (`lenient_unseen=false`), vocabulary pinned at fit.
6982        let mut s = spec();
6983        s.name = "year".to_string();
6984        s.lenient_unseen = false;
6985        s
6986    }
6987
6988    #[test]
6989    fn fixed_factor_rejects_unseen_numeric_level_at_predict() {
6990        // The numeric-coded `factor(year)` gap (#2137): the column reaches the
6991        // operator as plain numbers (no categorical schema to pre-filter it), so
6992        // the operator that owns the frozen vocabulary must reject an unseen
6993        // code rather than encode an all-zero (centering-point) row.
6994        let mut s = fixed_factor_spec();
6995        s.frozen_levels = Some(vec![2000.0_f64.to_bits(), 2001.0_f64.to_bits()]);
6996        let data = array![[2000.0_f64], [1999.0]];
6997        let err = build_random_effect_block(data.view(), &s)
6998            .expect_err("an unseen fixed-factor level must be rejected");
6999        let msg = format!("{err}");
7000        assert!(
7001            msg.contains("unseen level"),
7002            "message must name the defect: {msg}"
7003        );
7004        assert!(
7005            msg.contains("1999"),
7006            "message must name the integer level (not 1999.0): {msg}"
7007        );
7008        assert!(msg.contains("year"), "message must name the column: {msg}");
7009    }
7010
7011    #[test]
7012    fn fixed_factor_accepts_seen_numeric_levels_at_predict() {
7013        // Control: every seen level still resolves; strictness rejects only the
7014        // genuinely out-of-vocabulary code.
7015        let mut s = fixed_factor_spec();
7016        s.frozen_levels = Some(vec![2000.0_f64.to_bits(), 2001.0_f64.to_bits()]);
7017        let data = array![[2001.0_f64], [2000.0]];
7018        let block = build_random_effect_block(data.view(), &s).unwrap();
7019        assert_eq!(block.group_ids[0], Some(1));
7020        assert_eq!(block.group_ids[1], Some(0));
7021    }
7022
7023    #[test]
7024    fn fixed_factor_at_fit_time_derives_vocabulary_and_never_false_rejects() {
7025        // At FIT (`frozen_levels=None`) the vocabulary is derived from this very
7026        // data, so no row is unseen — the strict guard must not fire even though
7027        // the factor is strict.
7028        let mut s = fixed_factor_spec();
7029        s.frozen_levels = None;
7030        let data = array![[2000.0_f64], [2001.0], [2002.0], [2000.0]];
7031        let block = build_random_effect_block(data.view(), &s)
7032            .expect("fit-time build must not reject its own levels");
7033        assert_eq!(block.num_groups, 3);
7034    }
7035
7036    #[test]
7037    fn random_effect_still_tolerates_unseen_numeric_level() {
7038        // Non-regression: a lenient random effect (`group`/`re`/`s(bs="re")`)
7039        // encodes an unseen level as an all-zero (population-mean) row, NOT a
7040        // rejection — the held-out-group contract (#2102) is unchanged.
7041        let mut s = spec(); // lenient_unseen = true
7042        s.frozen_levels = Some(vec![2000.0_f64.to_bits(), 2001.0_f64.to_bits()]);
7043        let data = array![[2000.0_f64], [1999.0]];
7044        let block = build_random_effect_block(data.view(), &s)
7045            .expect("a random effect tolerates unseen levels");
7046        assert_eq!(block.group_ids[0], Some(0));
7047        assert_eq!(
7048            block.group_ids[1], None,
7049            "unseen level → population mean, not a reject"
7050        );
7051    }
7052}
7053
7054impl SmoothDesign {
7055    /// Map an unconstrained term coefficient vector to its constrained shape space.
7056    /// This is useful for nonlinear fits that optimize unconstrained parameters.
7057    pub fn map_term_coefficients(
7058        unconstrained: &Array1<f64>,
7059        shape: ShapeConstraint,
7060    ) -> Result<Array1<f64>, BasisError> {
7061        if unconstrained.is_empty() {
7062            crate::bail_invalid_basis!("unconstrained coefficient vector cannot be empty");
7063        }
7064        let mapped = match shape {
7065            ShapeConstraint::None => unconstrained.clone(),
7066            ShapeConstraint::MonotoneIncreasing => cumulative_exp(unconstrained, 1.0),
7067            ShapeConstraint::MonotoneDecreasing => cumulative_exp(unconstrained, -1.0),
7068            ShapeConstraint::Convex => second_cumulative_exp(unconstrained, 1.0),
7069            ShapeConstraint::Concave => second_cumulative_exp(unconstrained, -1.0),
7070        };
7071        Ok(mapped)
7072    }
7073}
7074
7075pub struct LocalSmoothTermBuild {
7076    pub dim: usize,
7077    pub design: DesignMatrix,
7078    /// Fixed row-wise term contribution for an affine basis chart.
7079    pub affine_offset: Option<Array1<f64>>,
7080    pub active_penalties: Vec<ActivePenalty>,
7081    /// Joint-null absorption rotation for this smooth. `Some(rotation)`
7082    /// records `Q = [U_range | U_null]` spanning `null(Σ_k penalties[k])`,
7083    /// the joint null across all active penalty blocks on this smooth.
7084    /// `None` means the joint penalty is full-rank (joint nullity = 0) or
7085    /// there are no penalties. Stage-2 commit A: plumbing only — populated
7086    /// by commit B, applied by commit D.
7087    pub joint_null_rotation: Option<crate::basis::JointNullRotation>,
7088    pub dropped_penalties: Vec<DroppedPenaltyInfo>,
7089    pub metadata: BasisMetadata,
7090    pub linear_constraints: Option<LinearInequalityConstraints>,
7091    pub box_reparam: bool,
7092    pub kronecker_factored: Option<KroneckerFactoredBasis>,
7093}
7094
7095#[derive(Clone)]
7096pub struct PcaScoresMemmapDesignOperator {
7097    mmap: Arc<memmap2::Mmap>,
7098    data_offset: usize,
7099    nrows: usize,
7100    ncols: usize,
7101    chunk_size: usize,
7102}
7103
7104impl PcaScoresMemmapDesignOperator {
7105    fn open(path: PathBuf, chunk_size: usize) -> Result<Self, BasisError> {
7106        let file = File::open(&path).map_err(|err| {
7107            BasisError::InvalidInput(format!(
7108                "failed to open lazy Pca .npy scores '{}': {err}",
7109                path.display()
7110            ))
7111        })?;
7112        // The .npy scores file is read-only training-cache data; this
7113        // module never mutates it. The error path below converts mmap
7114        // failure to a typed `BasisError::InvalidInput`.
7115        // SAFETY: `memmap2::Mmap::map` requires no concurrent writers; the
7116        // contract is held by this module's read-only access pattern.
7117        let mmap = unsafe {
7118            memmap2::Mmap::map(&file).map_err(|err| {
7119                BasisError::InvalidInput(format!(
7120                    "failed to memmap lazy Pca .npy scores '{}': {err}",
7121                    path.display()
7122                ))
7123            })?
7124        };
7125        let (data_offset, nrows, ncols) = parse_f64_2d_npy_header(&mmap, &path)?;
7126        let expected = data_offset
7127            .checked_add(nrows.saturating_mul(ncols).saturating_mul(8))
7128            .ok_or_else(|| {
7129                BasisError::InvalidInput(format!(
7130                    "lazy Pca .npy scores '{}' shape is too large",
7131                    path.display()
7132                ))
7133            })?;
7134        if mmap.len() < expected {
7135            crate::bail_invalid_basis!(
7136                "lazy Pca .npy scores '{}' is truncated: header expects {} bytes, file has {}",
7137                path.display(),
7138                expected,
7139                mmap.len()
7140            );
7141        }
7142        Ok(Self {
7143            mmap: Arc::new(mmap),
7144            data_offset,
7145            nrows,
7146            ncols,
7147            chunk_size: chunk_size.max(1),
7148        })
7149    }
7150
7151    fn value(&self, row: usize, col: usize) -> f64 {
7152        let offset = self.data_offset + (row * self.ncols + col) * 8;
7153        let mut bytes = [0_u8; 8];
7154        bytes.copy_from_slice(&self.mmap[offset..offset + 8]);
7155        f64::from_le_bytes(bytes)
7156    }
7157
7158    fn chunk_rows(&self) -> usize {
7159        self.chunk_size.min(self.nrows.max(1))
7160    }
7161}
7162
7163impl LinearOperator for PcaScoresMemmapDesignOperator {
7164    fn nrows(&self) -> usize {
7165        self.nrows
7166    }
7167
7168    fn ncols(&self) -> usize {
7169        self.ncols
7170    }
7171
7172    fn apply(&self, vector: &Array1<f64>) -> Array1<f64> {
7173        assert_eq!(
7174            vector.len(),
7175            self.ncols,
7176            "lazy Pca apply vector length mismatch"
7177        );
7178        let mut out = Array1::<f64>::zeros(self.nrows);
7179        for start in (0..self.nrows).step_by(self.chunk_rows()) {
7180            let end = (start + self.chunk_rows()).min(self.nrows);
7181            for row in start..end {
7182                let mut acc = 0.0;
7183                for col in 0..self.ncols {
7184                    acc += self.value(row, col) * vector[col];
7185                }
7186                out[row] = acc;
7187            }
7188        }
7189        out
7190    }
7191
7192    fn apply_transpose(&self, vector: &Array1<f64>) -> Array1<f64> {
7193        assert_eq!(
7194            vector.len(),
7195            self.nrows,
7196            "lazy Pca apply_transpose vector length mismatch"
7197        );
7198        let mut out = Array1::<f64>::zeros(self.ncols);
7199        for start in (0..self.nrows).step_by(self.chunk_rows()) {
7200            let end = (start + self.chunk_rows()).min(self.nrows);
7201            for row in start..end {
7202                let scale = vector[row];
7203                if scale == 0.0 {
7204                    continue;
7205                }
7206                for col in 0..self.ncols {
7207                    out[col] += scale * self.value(row, col);
7208                }
7209            }
7210        }
7211        out
7212    }
7213
7214    fn diag_xtw_x(&self, weights: &Array1<f64>) -> Result<Array2<f64>, String> {
7215        if weights.len() != self.nrows {
7216            return Err(format!(
7217                "lazy Pca diag_xtw_x weight length mismatch: weights={}, nrows={}",
7218                weights.len(),
7219                self.nrows
7220            ));
7221        }
7222        FiniteSignedWeightsView::try_from_array(weights)
7223            .map_err(|reason| format!("lazy Pca diag_xtw_x: {reason}"))?;
7224        let mut gram = Array2::<f64>::zeros((self.ncols, self.ncols));
7225        for start in (0..self.nrows).step_by(self.chunk_rows()) {
7226            let end = (start + self.chunk_rows()).min(self.nrows);
7227            for row in start..end {
7228                let w = weights[row];
7229                if w == 0.0 {
7230                    continue;
7231                }
7232                for a in 0..self.ncols {
7233                    let xa = self.value(row, a);
7234                    if xa == 0.0 {
7235                        continue;
7236                    }
7237                    for b in a..self.ncols {
7238                        gram[[a, b]] += w * xa * self.value(row, b);
7239                    }
7240                }
7241            }
7242        }
7243        for a in 0..self.ncols {
7244            for b in 0..a {
7245                gram[[a, b]] = gram[[b, a]];
7246            }
7247        }
7248        Ok(gram)
7249    }
7250
7251    fn apply_weighted_normal(
7252        &self,
7253        weights: FiniteSignedWeightsView<'_>,
7254        vector: &Array1<f64>,
7255        penalty: Option<&Array2<f64>>,
7256        ridge: f64,
7257    ) -> Array1<f64> {
7258        assert_eq!(
7259            weights.len(),
7260            self.nrows,
7261            "lazy Pca weighted-normal weight mismatch"
7262        );
7263        assert_eq!(
7264            vector.len(),
7265            self.ncols,
7266            "lazy Pca weighted-normal vector mismatch"
7267        );
7268        let weights = weights.view();
7269        let mut out = Array1::<f64>::zeros(self.ncols);
7270        for start in (0..self.nrows).step_by(self.chunk_rows()) {
7271            let end = (start + self.chunk_rows()).min(self.nrows);
7272            for row in start..end {
7273                let w = weights[row];
7274                if w == 0.0 {
7275                    continue;
7276                }
7277                let mut row_dot = 0.0;
7278                for col in 0..self.ncols {
7279                    row_dot += self.value(row, col) * vector[col];
7280                }
7281                if row_dot == 0.0 {
7282                    continue;
7283                }
7284                let scaled = w * row_dot;
7285                for col in 0..self.ncols {
7286                    out[col] += scaled * self.value(row, col);
7287                }
7288            }
7289        }
7290        if let Some(pen) = penalty {
7291            out += &pen.dot(vector);
7292        }
7293        if ridge > 0.0 {
7294            out += &vector.mapv(|x| ridge * x);
7295        }
7296        out
7297    }
7298}
7299
7300impl DenseDesignOperator for PcaScoresMemmapDesignOperator {
7301    fn compute_xtwy(&self, weights: &Array1<f64>, y: &Array1<f64>) -> Result<Array1<f64>, String> {
7302        if weights.len() != self.nrows || y.len() != self.nrows {
7303            return Err(format!(
7304                "lazy Pca compute_xtwy dimension mismatch: weights={}, y={}, nrows={}",
7305                weights.len(),
7306                y.len(),
7307                self.nrows
7308            ));
7309        }
7310        FiniteSignedWeightsView::try_from_array(weights)
7311            .map_err(|reason| format!("lazy Pca compute_xtwy: {reason}"))?;
7312        let mut out = Array1::<f64>::zeros(self.ncols);
7313        for start in (0..self.nrows).step_by(self.chunk_rows()) {
7314            let end = (start + self.chunk_rows()).min(self.nrows);
7315            for row in start..end {
7316                let scale = weights[row] * y[row];
7317                if scale == 0.0 {
7318                    continue;
7319                }
7320                for col in 0..self.ncols {
7321                    out[col] += scale * self.value(row, col);
7322                }
7323            }
7324        }
7325        Ok(out)
7326    }
7327
7328    fn row_chunk_into(
7329        &self,
7330        rows: Range<usize>,
7331        mut out: ArrayViewMut2<'_, f64>,
7332    ) -> Result<(), MatrixMaterializationError> {
7333        if rows.end > self.nrows || rows.start > rows.end {
7334            return Err(MatrixMaterializationError::MissingRowChunk {
7335                context: "lazy Pca row range out of bounds",
7336            });
7337        }
7338        if out.nrows() != rows.end - rows.start || out.ncols() != self.ncols {
7339            return Err(MatrixMaterializationError::MissingRowChunk {
7340                context: "lazy Pca row_chunk_into shape mismatch",
7341            });
7342        }
7343        for (local, row) in (rows.start..rows.end).enumerate() {
7344            for col in 0..self.ncols {
7345                out[[local, col]] = self.value(row, col);
7346            }
7347        }
7348        Ok(())
7349    }
7350
7351    fn to_dense(&self) -> Array2<f64> {
7352        let mut out = Array2::<f64>::zeros((self.nrows, self.ncols));
7353        self.row_chunk_into(0..self.nrows, out.view_mut())
7354            .expect("lazy Pca full materialization failed");
7355        out
7356    }
7357}
7358
7359pub fn parse_f64_2d_npy_header(
7360    bytes: &[u8],
7361    path: &PathBuf,
7362) -> Result<(usize, usize, usize), BasisError> {
7363    let mut reader = std::io::Cursor::new(bytes);
7364    let header = npyz::NpyHeader::from_reader(&mut reader).map_err(|err| {
7365        BasisError::InvalidInput(format!(
7366            "lazy Pca scores '{}' has an invalid .npy header: {err}",
7367            path.display()
7368        ))
7369    })?;
7370    let is_little_endian_f64 = matches!(
7371        header.dtype(),
7372        npyz::DType::Plain(ref dtype)
7373            if dtype.type_char() == npyz::TypeChar::Float
7374                && dtype.size_field() == 8
7375                && dtype.endianness() == npyz::Endianness::Little
7376    );
7377    if !is_little_endian_f64 {
7378        crate::bail_invalid_basis!(
7379            "lazy Pca scores '{}' must be scalar little-endian float64 .npy, got {}",
7380            path.display(),
7381            header.dtype().descr()
7382        );
7383    }
7384    if header.order() != npyz::Order::C {
7385        crate::bail_invalid_basis!(
7386            "lazy Pca scores '{}' must be C-contiguous, not Fortran-ordered",
7387            path.display()
7388        );
7389    }
7390    if header.shape().len() != 2 {
7391        crate::bail_invalid_basis!(
7392            "lazy Pca scores '{}' must have shape (N, K), got {:?}",
7393            path.display(),
7394            header.shape()
7395        );
7396    }
7397    let nrows = usize::try_from(header.shape()[0]).map_err(|_| {
7398        BasisError::InvalidInput(format!(
7399            "lazy Pca scores '{}' row count {} exceeds this platform's address space",
7400            path.display(),
7401            header.shape()[0]
7402        ))
7403    })?;
7404    let ncols = usize::try_from(header.shape()[1]).map_err(|_| {
7405        BasisError::InvalidInput(format!(
7406            "lazy Pca scores '{}' column count {} exceeds this platform's address space",
7407            path.display(),
7408            header.shape()[1]
7409        ))
7410    })?;
7411    let data_offset = usize::try_from(reader.position()).map_err(|_| {
7412        BasisError::InvalidInput(format!(
7413            "lazy Pca scores '{}' header offset exceeds this platform's address space",
7414            path.display()
7415        ))
7416    })?;
7417    Ok((data_offset, nrows, ncols))
7418}
7419
7420pub fn pca_center_mean(x: ArrayView2<'_, f64>) -> Result<Array1<f64>, BasisError> {
7421    if x.nrows() == 0 {
7422        crate::bail_invalid_basis!("Pca basis requires at least one row to compute center mean");
7423    }
7424    let mut mean = Array1::<f64>::zeros(x.ncols());
7425    for row in x.rows() {
7426        mean += &row;
7427    }
7428    mean.mapv_inplace(|v| v / x.nrows() as f64);
7429    Ok(mean)
7430}
7431
7432/// Build the empirical final-function mass penalty from the raw score Gram.
7433///
7434/// For the realized PCA score design `Z`, the quadratic form is
7435///
7436/// `beta^T S beta = smooth_penalty * mean_i((Z beta)_i^2)`.
7437///
7438/// Thus `smooth_penalty` chooses the reference-measure scale only; the existing
7439/// REML smoothing coordinate multiplying this penalty learns the shrinkage
7440/// strength.  In particular, this is not an identity ridge on whichever
7441/// coefficient chart happened to encode the score columns.
7442fn pca_function_mass_penalty(
7443    mut raw_score_gram: Array2<f64>,
7444    n_rows: usize,
7445    smooth_penalty: f64,
7446) -> Result<Array2<f64>, BasisError> {
7447    let k = raw_score_gram.ncols();
7448    if raw_score_gram.nrows() != k {
7449        crate::bail_dim_basis!(
7450            "Pca score Gram must be square, got {}x{}",
7451            raw_score_gram.nrows(),
7452            k
7453        );
7454    }
7455    if n_rows == 0 {
7456        crate::bail_invalid_basis!("Pca basis requires at least one score row");
7457    }
7458    if k == 0 {
7459        crate::bail_invalid_basis!("Pca basis requires at least one score column");
7460    }
7461    if k > n_rows {
7462        crate::bail_invalid_basis!(
7463            "Pca score design is rank deficient: {} score columns cannot have full column rank with only {} rows; remove redundant components",
7464            k,
7465            n_rows
7466        );
7467    }
7468    if raw_score_gram.iter().any(|value| !value.is_finite()) {
7469        crate::bail_invalid_basis!("Pca score design produced a non-finite function Gram");
7470    }
7471
7472    // Use the same design-rank convention as the global identifiability audit.
7473    // `rrqr_from_gram_with_permutation` recovers the column-pivoted QR verdict
7474    // from Z^T Z while retaining the tall design's row-count-aware tolerance.
7475    let rrqr = gam_linalg::faer_ndarray::rrqr_from_gram_with_permutation(
7476        &raw_score_gram,
7477        n_rows,
7478        gam_linalg::faer_ndarray::default_rrqr_rank_alpha(),
7479    )
7480    .map_err(BasisError::LinalgError)?;
7481    if rrqr.rank != k {
7482        let redundant_columns = &rrqr.column_permutation[rrqr.rank..];
7483        crate::bail_invalid_basis!(
7484            "Pca score design is rank deficient under canonical RRQR: rank {} < {} (tolerance {:.6e}); redundant score columns {:?}; remove zero or dependent components instead of stabilizing them with a coefficient ridge",
7485            rrqr.rank,
7486            k,
7487            rrqr.rank_tol,
7488            redundant_columns
7489        );
7490    }
7491
7492    raw_score_gram.mapv_inplace(|value| value * smooth_penalty / n_rows as f64);
7493    Ok(raw_score_gram)
7494}
7495
7496pub fn build_pca_smooth_basis(
7497    data: ArrayView2<'_, f64>,
7498    feature_cols: &[usize],
7499    basis_matrix: &Array2<f64>,
7500    centered: bool,
7501    smooth_penalty: f64,
7502    center_mean: Option<&Array1<f64>>,
7503    pca_basis_path: Option<&PathBuf>,
7504    chunk_size: usize,
7505) -> Result<BasisBuildResult, BasisError> {
7506    if !smooth_penalty.is_finite() || smooth_penalty < 0.0 {
7507        crate::bail_invalid_basis!(
7508            "Pca smooth_penalty must be finite and non-negative, got {}",
7509            smooth_penalty
7510        );
7511    }
7512    if data.nrows() == 0 {
7513        crate::bail_invalid_basis!("Pca basis requires at least one data row");
7514    }
7515
7516    if let Some(path) = pca_basis_path {
7517        let op = PcaScoresMemmapDesignOperator::open(path.clone(), chunk_size)?;
7518        if op.nrows != data.nrows() {
7519            crate::bail_dim_basis!(
7520                "lazy Pca scores row mismatch: .npy has {}, data has {}",
7521                op.nrows,
7522                data.nrows()
7523            );
7524        }
7525        // The out-of-core scores are already the realized final-function
7526        // design. Stream Z^T Z without materializing its n-by-k rows.
7527        let raw_score_gram = op
7528            .diag_xtw_x(&Array1::<f64>::ones(op.nrows))
7529            .map_err(|err| {
7530                BasisError::InvalidInput(format!(
7531                    "lazy Pca function-mass Gram construction failed: {err}"
7532                ))
7533            })?;
7534        let penalty = pca_function_mass_penalty(raw_score_gram, op.nrows, smooth_penalty)?;
7535        let filtered = filter_penalty_candidates(vec![PenaltyCandidate {
7536            matrix: ConstructiveQuadratic::try_from_dense_psd(
7537                penalty,
7538                "lazy PCA function-mass penalty",
7539            )?,
7540            source: PenaltySource::OperatorMass,
7541            normalization_scale: 1.0,
7542            kronecker_factors: None,
7543            op: None,
7544        }])?;
7545        return Ok(BasisBuildResult {
7546            design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(op))),
7547            affine_offset: None,
7548            active_penalties: filtered.active,
7549            dropped_penalties: filtered.dropped,
7550            joint_null_rotation: None,
7551            metadata: BasisMetadata::Pca {
7552                feature_cols: feature_cols.to_vec(),
7553                basis_matrix: basis_matrix.clone(),
7554                centered,
7555                smooth_penalty,
7556                center_mean: center_mean.cloned(),
7557                pca_basis_path: Some(path.clone()),
7558                chunk_size: chunk_size.max(1),
7559            },
7560            kronecker_factored: None,
7561        });
7562    }
7563    if basis_matrix.nrows() != feature_cols.len() {
7564        crate::bail_dim_basis!(
7565            "Pca basis row mismatch: basis rows={}, feature columns={}",
7566            basis_matrix.nrows(),
7567            feature_cols.len()
7568        );
7569    }
7570    let mut x = select_columns(data, feature_cols)?;
7571    let mean = if centered {
7572        match center_mean {
7573            Some(mean) => mean.clone(),
7574            None => pca_center_mean(x.view())?,
7575        }
7576    } else {
7577        Array1::<f64>::zeros(feature_cols.len())
7578    };
7579    if centered {
7580        for mut row in x.rows_mut() {
7581            row -= &mean;
7582        }
7583    }
7584    let design = fast_ab(&x, basis_matrix);
7585    let raw_score_gram = gam_linalg::faer_ndarray::fast_ata(&design);
7586    let penalty = pca_function_mass_penalty(raw_score_gram, design.nrows(), smooth_penalty)?;
7587    let filtered = filter_penalty_candidates(vec![PenaltyCandidate {
7588        matrix: ConstructiveQuadratic::try_from_dense_psd(penalty, "PCA function-mass penalty")?,
7589        source: PenaltySource::OperatorMass,
7590        normalization_scale: 1.0,
7591        kronecker_factors: None,
7592        op: None,
7593    }])?;
7594    Ok(BasisBuildResult {
7595        design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(design)),
7596        affine_offset: None,
7597        active_penalties: filtered.active,
7598        dropped_penalties: filtered.dropped,
7599        joint_null_rotation: None,
7600        metadata: BasisMetadata::Pca {
7601            feature_cols: feature_cols.to_vec(),
7602            basis_matrix: basis_matrix.clone(),
7603            centered,
7604            smooth_penalty,
7605            center_mean: centered.then_some(mean),
7606            pca_basis_path: None,
7607            chunk_size: chunk_size.max(1),
7608        },
7609        kronecker_factored: None,
7610    })
7611}
7612
7613#[cfg(test)]
7614mod pca_function_mass_tests {
7615    use super::{PenaltySource, build_pca_smooth_basis, parse_f64_2d_npy_header};
7616    use ndarray::{Array1, Array2, array};
7617    use std::io::Write;
7618    use std::path::PathBuf;
7619
7620    fn quadratic_form(matrix: &Array2<f64>, coefficients: &Array1<f64>) -> f64 {
7621        coefficients.dot(&matrix.dot(coefficients))
7622    }
7623
7624    fn assert_close(left: f64, right: f64) {
7625        let scale = left.abs().max(right.abs()).max(1.0);
7626        assert!(
7627            (left - right).abs() <= 1e-11 * scale,
7628            "values differ: left={left:.16e}, right={right:.16e}"
7629        );
7630    }
7631
7632    fn write_f64_npy(scores: &Array2<f64>) -> PathBuf {
7633        let path = std::env::temp_dir().join(format!(
7634            "gam_terms_pca_function_mass_{}.npy",
7635            std::process::id()
7636        ));
7637        let mut header = format!(
7638            "{{'descr': '<f8', 'fortran_order': False, 'shape': ({}, {}), }}",
7639            scores.nrows(),
7640            scores.ncols()
7641        );
7642        while (10 + header.len() + 1) % 16 != 0 {
7643            header.push(' ');
7644        }
7645        header.push('\n');
7646        let header_len = u16::try_from(header.len()).expect("test .npy header fits u16");
7647
7648        let mut file = std::fs::File::create(&path).expect("create test .npy");
7649        file.write_all(b"\x93NUMPY").expect("write .npy magic");
7650        file.write_all(&[1, 0]).expect("write .npy version");
7651        file.write_all(&header_len.to_le_bytes())
7652            .expect("write .npy header length");
7653        file.write_all(header.as_bytes())
7654            .expect("write .npy header");
7655        for &value in scores {
7656            file.write_all(&value.to_le_bytes())
7657                .expect("write .npy score");
7658        }
7659        path
7660    }
7661
7662    fn npy_v1_bytes(mut header: String) -> Vec<u8> {
7663        while (10 + header.len() + 1) % 16 != 0 {
7664            header.push(' ');
7665        }
7666        header.push('\n');
7667        let header_len = u16::try_from(header.len()).expect("test header fits v1");
7668        let mut bytes = b"\x93NUMPY".to_vec();
7669        bytes.extend_from_slice(&[1, 0]);
7670        bytes.extend_from_slice(&header_len.to_le_bytes());
7671        bytes.extend_from_slice(header.as_bytes());
7672        bytes
7673    }
7674
7675    #[test]
7676    fn npy_header_parser_uses_exact_ast_fields_2293() {
7677        let path = PathBuf::from("scores.npy");
7678        let bytes = npy_v1_bytes(
7679            "{'shape':(3, 2), 'note':'True', 'descr':'<f8', 'fortran_order':False,}".to_string(),
7680        );
7681        let (offset, rows, cols) =
7682            parse_f64_2d_npy_header(&bytes, &path).expect("valid reordered header");
7683        assert_eq!((rows, cols), (3, 2));
7684        assert_eq!(offset, bytes.len());
7685
7686        for header in [
7687            "{'descr':'<f8','fortran_order':True,'shape':(3,2),}",
7688            "{'descr':'>f8','fortran_order':False,'shape':(3,2),}",
7689            "{'descr':'<f8','shape':(3,2),}",
7690            "{'descr':'<f8','fortran_order':'False','shape':(3,2),}",
7691            "{'descr':'<f8','fortran_order':False,'shape':(6,),}",
7692        ] {
7693            let invalid = npy_v1_bytes(header.to_string());
7694            assert!(
7695                parse_f64_2d_npy_header(&invalid, &path).is_err(),
7696                "{header}"
7697            );
7698        }
7699    }
7700
7701    #[test]
7702    fn pca_penalty_quadratic_equals_empirical_fitted_function_norm() {
7703        let data = array![[1.0, 2.0], [-1.0, 0.5], [2.0, -0.5], [0.25, -1.5]];
7704        let basis = array![[1.0, 0.5], [-0.25, 2.0]];
7705        let smooth_penalty = 2.5;
7706        let built = build_pca_smooth_basis(
7707            data.view(),
7708            &[0, 1],
7709            &basis,
7710            false,
7711            smooth_penalty,
7712            None,
7713            None,
7714            2,
7715        )
7716        .expect("full-rank PCA basis");
7717        let coefficients = array![0.7, -1.2];
7718        let design = built.design.to_dense();
7719        let fitted = design.dot(&coefficients);
7720        let expected = smooth_penalty * fitted.dot(&fitted) / fitted.len() as f64;
7721        let actual = quadratic_form(&built.active_penalties[0].matrix, &coefficients);
7722
7723        assert_close(actual, expected);
7724        assert_eq!(built.active_penalties[0].nullity, 0);
7725        assert_eq!(
7726            built.active_penalties[0].info.source,
7727            PenaltySource::OperatorMass
7728        );
7729    }
7730
7731    #[test]
7732    fn pca_function_mass_is_invariant_to_nonorthogonal_score_reparameterization() {
7733        let scores = array![[1.0, 2.0], [-1.0, 0.5], [2.0, -0.5], [0.25, -1.5]];
7734        let identity = Array2::<f64>::eye(2);
7735        // An invertible scale-plus-shear, deliberately not orthogonal.
7736        let transform = array![[2.0, 0.5], [0.0, 0.25]];
7737        let base_coefficients = array![0.8, -1.1];
7738        // transform * transformed_coefficients == base_coefficients.
7739        let transformed_coefficients = array![1.5, -4.4];
7740        let smooth_penalty = 1.7;
7741
7742        let base = build_pca_smooth_basis(
7743            scores.view(),
7744            &[0, 1],
7745            &identity,
7746            false,
7747            smooth_penalty,
7748            None,
7749            None,
7750            2,
7751        )
7752        .expect("base PCA chart");
7753        let transformed = build_pca_smooth_basis(
7754            scores.view(),
7755            &[0, 1],
7756            &transform,
7757            false,
7758            smooth_penalty,
7759            None,
7760            None,
7761            2,
7762        )
7763        .expect("reparameterized PCA chart");
7764
7765        let fitted_base = base.design.to_dense().dot(&base_coefficients);
7766        let fitted_transformed = transformed.design.to_dense().dot(&transformed_coefficients);
7767        for (&left, &right) in fitted_base.iter().zip(fitted_transformed.iter()) {
7768            assert_close(left, right);
7769        }
7770        assert_close(
7771            quadratic_form(&base.active_penalties[0].matrix, &base_coefficients),
7772            quadratic_form(
7773                &transformed.active_penalties[0].matrix,
7774                &transformed_coefficients,
7775            ),
7776        );
7777    }
7778
7779    #[test]
7780    fn rank_deficient_pca_score_design_is_rejected() {
7781        let scores = array![[1.0, 0.0], [2.0, 0.0], [3.0, 0.0], [4.0, 0.0]];
7782        let result = build_pca_smooth_basis(
7783            scores.view(),
7784            &[0, 1],
7785            &Array2::<f64>::eye(2),
7786            false,
7787            1.0,
7788            None,
7789            None,
7790            2,
7791        );
7792        let err = result.err().expect("zero score column must be rejected");
7793        let message = err.to_string();
7794        assert!(
7795            message.contains("rank deficient"),
7796            "unexpected error: {message}"
7797        );
7798        assert!(
7799            message.contains("rank 1 < 2"),
7800            "missing RRQR evidence: {message}"
7801        );
7802    }
7803
7804    #[test]
7805    fn lazy_and_dense_pca_function_mass_penalties_match() {
7806        let scores = array![[1.0, 2.0], [-1.0, 0.5], [2.0, -0.5], [0.25, -1.5]];
7807        let smooth_penalty = 2.25;
7808        let path = write_f64_npy(&scores);
7809        let dense = build_pca_smooth_basis(
7810            scores.view(),
7811            &[0, 1],
7812            &Array2::<f64>::eye(2),
7813            false,
7814            smooth_penalty,
7815            None,
7816            None,
7817            2,
7818        )
7819        .expect("dense PCA basis");
7820        let lazy_data = Array2::<f64>::zeros((scores.nrows(), 0));
7821        let lazy = build_pca_smooth_basis(
7822            lazy_data.view(),
7823            &[],
7824            &Array2::<f64>::zeros((0, scores.ncols())),
7825            false,
7826            smooth_penalty,
7827            None,
7828            Some(&path),
7829            2,
7830        )
7831        .expect("lazy PCA basis");
7832        std::fs::remove_file(&path).expect("remove test .npy");
7833
7834        for (&left, &right) in dense.active_penalties[0]
7835            .matrix
7836            .iter()
7837            .zip(lazy.active_penalties[0].matrix.iter())
7838        {
7839            assert_close(left, right);
7840        }
7841        for (&left, &right) in dense
7842            .design
7843            .to_dense()
7844            .iter()
7845            .zip(lazy.design.to_dense().iter())
7846        {
7847            assert_close(left, right);
7848        }
7849    }
7850}
7851
7852/// A factor-level `by=` wrapper owns the model-space centering of its inner
7853/// smooth: it gates the raw/structurally-constrained basis to the level rows
7854/// and then centers that gated block exactly once against the level indicator
7855/// (`build_parametric_constraint_block_for_term` in `design_construction`).
7856/// Leaving the inner B-spline's default pooled weighted-sum-to-zero active here
7857/// would impose two generically-independent constraints — the pooled column
7858/// moment `m = Σ_h m_h` and the per-level moment `m_g` — so a raw `k`-column
7859/// basis collapses to `k-2` columns per level instead of `k-1`, deleting one
7860/// genuine nonconstant spline direction *before REML runs* (#1427). The group
7861/// main effect carries only the constant, so it cannot restore that direction.
7862///
7863/// Only the *default model-space* centering is deferred. Explicit structural or
7864/// frozen transforms (`RemoveLinearTrend`, `OrthogonalToDesignColumns`,
7865/// `FrozenTransform`, `None`) are user/structural choices and are preserved
7866/// verbatim.
7867pub fn defer_inner_model_centering_to_factor_level_wrapper(basis: &mut SmoothBasisSpec) {
7868    if let SmoothBasisSpec::BSpline1D { spec, .. } = basis
7869        && matches!(
7870            spec.identifiability,
7871            BSplineIdentifiability::WeightedSumToZero { .. }
7872        )
7873    {
7874        spec.identifiability = BSplineIdentifiability::None;
7875    }
7876}
7877
7878pub fn apply_by_variable_to_local_build(
7879    mut built: LocalSmoothTermBuild,
7880    data: ArrayView2<'_, f64>,
7881    by_col: usize,
7882    by: &ByVariableSpec,
7883    term_name: &str,
7884) -> Result<LocalSmoothTermBuild, BasisError> {
7885    if by_col >= data.ncols() {
7886        crate::bail_dim_basis!(
7887            "by-variable smooth term '{term_name}' references column {by_col}, but data has {} columns",
7888            data.ncols()
7889        );
7890    }
7891    let weights = match by {
7892        ByVariableSpec::Numeric => data.column(by_col).to_owned(),
7893        ByVariableSpec::Level { value_bits, .. } => {
7894            let value_bits = gam_data::canonical_level_bits(f64::from_bits(*value_bits));
7895            data.column(by_col).mapv(|value| {
7896                if gam_data::canonical_level_bits(value) == value_bits {
7897                    1.0
7898                } else {
7899                    0.0
7900                }
7901            })
7902        }
7903    };
7904    if weights.iter().any(|value| !value.is_finite()) {
7905        crate::bail_invalid_basis!(
7906            "by-variable smooth term '{term_name}' has non-finite by-column values"
7907        );
7908    }
7909
7910    let mut dense = built
7911        .design
7912        .try_to_dense_by_chunks("by-variable smooth row gating")
7913        .map_err(BasisError::InvalidInput)?;
7914    for (mut row, &weight) in dense.rows_mut().into_iter().zip(weights.iter()) {
7915        row.mapv_inplace(|value| value * weight);
7916    }
7917    if let Some(offset) = built.affine_offset.as_mut() {
7918        if offset.len() != weights.len() {
7919            crate::bail_dim_basis!(
7920                "by-variable smooth term '{term_name}' affine offset has {} rows but the by-variable has {}",
7921                offset.len(),
7922                weights.len()
7923            );
7924        }
7925        *offset *= &weights;
7926    }
7927    built.design = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(dense));
7928    built.kronecker_factored = None;
7929    Ok(built)
7930}
7931
7932/// Build the local smooth term for a `BySmooth` spec, which unifies numeric-by
7933/// and factor-by modulation into a single `SmoothTermSpec`.
7934///
7935/// For a **numeric** by-variable the inner smooth is built once and every row
7936/// is multiplied by the by-column value (identical to `ByVariable::Numeric`).
7937///
7938/// For a **factor** by-variable the inner smooth is built once and gated per
7939/// level into side-by-side column blocks, producing a `n × (L * p)` design
7940/// matrix.  The penalties are block-diagonalised (one copy of the inner penalty
7941/// per level) exactly as `build_factor_smooth` does for `bs="fs"/"sz"`.
7942pub fn build_by_smooth_local(
7943    data: ArrayView2<'_, f64>,
7944    term: &SmoothTermSpec,
7945    smooth: &SmoothBasisSpec,
7946    by_kind: &ByVarKind,
7947    workspace: &mut crate::basis::BasisWorkspace,
7948) -> Result<LocalSmoothTermBuild, BasisError> {
7949    let inner_term = SmoothTermSpec {
7950            frozen_parametric_residualization: None,
7951        name: term.name.clone(),
7952        basis: (*smooth).clone(),
7953        shape: term.shape,
7954        joint_null_rotation: None,
7955    };
7956    let inner = build_single_local_smooth_term(data, &inner_term, workspace)?;
7957
7958    match by_kind {
7959        ByVarKind::Numeric { feature_col } => {
7960            let inner_meta = inner.metadata.clone();
7961            let mut built = apply_by_variable_to_local_build(
7962                inner,
7963                data,
7964                *feature_col,
7965                &ByVariableSpec::Numeric,
7966                &term.name,
7967            )?;
7968            built.metadata = BasisMetadata::BySmooth {
7969                inner: Box::new(inner_meta),
7970                by_col: *feature_col,
7971                levels: None,
7972                ordered: false,
7973            };
7974            Ok(built)
7975        }
7976        ByVarKind::Factor {
7977            feature_col,
7978            frozen_levels,
7979            ordered,
7980        } => {
7981            // Collect factor levels: prefer the frozen set (replay path), else
7982            // scan the data column (first-fit path).
7983            let level_bits: Vec<u64> = if let Some(fl) = frozen_levels {
7984                fl.iter()
7985                    .map(|&b| gam_data::canonical_level_bits(f64::from_bits(b)))
7986                    .collect()
7987            } else {
7988                let col = data.column(*feature_col);
7989                let mut seen = BTreeSet::<u64>::new();
7990                for &v in col.iter() {
7991                    if v.is_finite() {
7992                        seen.insert(gam_data::canonical_level_bits(v));
7993                    }
7994                }
7995                seen.into_iter().collect()
7996            };
7997            let n_levels = level_bits.len();
7998            if n_levels == 0 {
7999                crate::bail_invalid_basis!(
8000                    "by-factor smooth term '{}': factor column {} has no observed levels",
8001                    term.name,
8002                    feature_col
8003                );
8004            }
8005            let p = inner.dim;
8006            let q = n_levels * p;
8007            let n = data.nrows();
8008
8009            let inner_dense = inner
8010                .design
8011                .try_to_dense_by_chunks("by-factor smooth design gating")
8012                .map_err(BasisError::InvalidInput)?;
8013
8014            // Gate each level into its own p-wide column block.
8015            let mut combined = Array2::<f64>::zeros((n, q));
8016            for (lvl_idx, &bits) in level_bits.iter().enumerate() {
8017                let col_start = lvl_idx * p;
8018                for row in 0..n {
8019                    if gam_data::canonical_level_bits(data[[row, *feature_col]]) == bits {
8020                        combined
8021                            .slice_mut(s![row, col_start..col_start + p])
8022                            .assign(&inner_dense.row(row));
8023                    }
8024                }
8025            }
8026
8027            // Build per-level INDEPENDENT penalties (#1427): one copy of each
8028            // inner penalty per level, but each confined to that single level's
8029            // diagonal block, so every (level, inner-penalty) pair is its OWN
8030            // smoothing-parameter coordinate. `s(x, by=g)` selects the per-group
8031            // curve wiggliness independently — the design is block-diagonal and
8032            // block-separable, so a correct REML must reproduce gamfit's own
8033            // independent per-group fits. Tiling a single inner penalty across
8034            // every level (as the `bs="fs"` shared-λ random-effect construction
8035            // does) collapses all groups onto ONE λ, which cannot match uneven
8036            // per-level smoothness and degrades as data grows (under-recovery up
8037            // to ~16× at n=2000). Emit `n_levels * n_penalties` blocks instead.
8038            let inner_meta = inner.metadata.clone();
8039            let n_penalties = inner.active_penalties.len();
8040            let n_blocks = n_penalties.saturating_mul(n_levels);
8041            let mut candidates = Vec::<PenaltyCandidate>::with_capacity(n_blocks);
8042            for base_penalty in &inner.active_penalties {
8043                for lvl in 0..n_levels {
8044                    let off = lvl * p;
8045                    let mut s_big = Array2::<f64>::zeros((q, q));
8046                    s_big
8047                        .slice_mut(s![off..off + p, off..off + p])
8048                        .assign(&base_penalty.matrix);
8049                    let (s_big, scale) = normalize_penalty_in_constrained_space(&s_big);
8050                    candidates.push(PenaltyCandidate {
8051                        matrix: ConstructiveQuadratic::try_from_dense_psd(
8052                            s_big,
8053                            "factor-smooth replicated penalty",
8054                        )?,
8055                        source: base_penalty.info.source.clone(),
8056                        normalization_scale: base_penalty.info.normalization_scale * scale,
8057                        kronecker_factors: None,
8058                        op: None,
8059                    });
8060                }
8061            }
8062
8063            // Re-analyze the completed q×q blocks in their actual coefficient
8064            // space. Copying the p×p marginal nullity understated every block's
8065            // null space by `(n_levels-1)·p`, while leaving its null basis and
8066            // joint-null rotation absent. The canonical filter authors matrix,
8067            // rank, nullity, null basis, and metadata together.
8068            let filtered = crate::basis::filter_penalty_candidates(candidates)?;
8069            let joint_null_rotation = crate::basis::compute_joint_null_rotation(&filtered.active)?;
8070            let mut dropped_penalties = inner.dropped_penalties;
8071            dropped_penalties.extend(filtered.dropped);
8072
8073            Ok(LocalSmoothTermBuild {
8074                dim: q,
8075                design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(combined)),
8076                // Exactly one factor block is active on every row. Each block
8077                // represents the same anchored marginal, so its fixed lift is
8078                // the inner lift once, not once per level.
8079                affine_offset: inner.affine_offset,
8080                active_penalties: filtered.active,
8081                joint_null_rotation,
8082                dropped_penalties,
8083                metadata: BasisMetadata::BySmooth {
8084                    inner: Box::new(inner_meta),
8085                    by_col: *feature_col,
8086                    levels: Some(level_bits),
8087                    ordered: *ordered,
8088                },
8089                linear_constraints: None,
8090                box_reparam: false,
8091                kronecker_factored: None,
8092            })
8093        }
8094    }
8095}
8096
8097pub fn ensure_by_variable_specs_match(
8098    kind: &BySmoothKind,
8099    by: &ByVariableSpec,
8100    term_name: &str,
8101) -> Result<(), BasisError> {
8102    match (kind, by) {
8103        (BySmoothKind::Numeric, ByVariableSpec::Numeric) => Ok(()),
8104        (BySmoothKind::Level { level_bits }, ByVariableSpec::Level { value_bits, .. })
8105            if level_bits == value_bits =>
8106        {
8107            Ok(())
8108        }
8109        _ => Err(BasisError::InvalidInput(format!(
8110            "by-variable smooth term '{term_name}' has inconsistent by-variable specifications"
8111        ))),
8112    }
8113}
8114
8115/// Choose a deterministic orthonormal basis for a subspace from its projector.
8116///
8117/// Eigenvectors belonging to a repeated eigenvalue are defined only up to an
8118/// arbitrary orthogonal rotation.  That freedom is harmless when consumers use
8119/// the whole projector `ZZ^T`, but it changes model semantics when each column
8120/// receives its own smoothing parameter.  We remove the eigensolver gauge by
8121/// repeatedly projecting coefficient coordinate axes into the subspace and
8122/// selecting the largest residual (with stable lowest-index tie breaking).
8123/// The result depends only on the subspace projector and the declared
8124/// coefficient chart, never on the orientation returned by an eigensolver.
8125fn canonical_nullspace_directions(z: &Array2<f64>) -> Result<Array2<f64>, BasisError> {
8126    let (coefficient_dim, nullity) = z.dim();
8127    if nullity == 0 {
8128        return Ok(Array2::zeros((coefficient_dim, 0)));
8129    }
8130    if coefficient_dim < nullity || z.iter().any(|value| !value.is_finite()) {
8131        crate::bail_invalid_basis!(
8132            "null-space basis must be finite with rows >= columns, got {}x{}",
8133            coefficient_dim,
8134            nullity
8135        );
8136    }
8137
8138    let tolerance = 128.0 * f64::EPSILON * coefficient_dim.max(1) as f64;
8139    let mut canonical = Array2::<f64>::zeros((coefficient_dim, nullity));
8140    for accepted in 0..nullity {
8141        let mut best_coordinate = usize::MAX;
8142        let mut best_norm = 0.0_f64;
8143        let mut best = Array1::<f64>::zeros(coefficient_dim);
8144
8145        for coordinate in 0..coefficient_dim {
8146            // `P e_j = Z (Z^T e_j)` without materializing the full projector.
8147            let mut candidate = Array1::<f64>::zeros(coefficient_dim);
8148            for row in 0..coefficient_dim {
8149                candidate[row] = (0..nullity)
8150                    .map(|axis| z[[row, axis]] * z[[coordinate, axis]])
8151                    .sum();
8152            }
8153            // Two-pass modified Gram--Schmidt keeps the selected directions
8154            // orthogonal even when successive projected coordinates are close.
8155            for _ in 0..2 {
8156                for axis in 0..accepted {
8157                    let direction = canonical.column(axis);
8158                    let projection = direction.dot(&candidate);
8159                    candidate.scaled_add(-projection, &direction);
8160                }
8161            }
8162            let norm = candidate.dot(&candidate).sqrt();
8163            let tie_band = tolerance * best_norm.max(1.0);
8164            if best_coordinate == usize::MAX || norm > best_norm + tie_band {
8165                best_coordinate = coordinate;
8166                best_norm = norm;
8167                best = candidate;
8168            }
8169        }
8170
8171        if best_coordinate == usize::MAX || best_norm <= tolerance {
8172            crate::bail_invalid_basis!(
8173                "null-space projector exposed only {} of {} independent directions",
8174                accepted,
8175                nullity
8176            );
8177        }
8178        best.mapv_inplace(|value| value / best_norm);
8179        // Fix the remaining sign gauge for reproducible metadata/debug output.
8180        let sign_anchor = best
8181            .iter()
8182            .enumerate()
8183            .max_by(|(left_index, left), (right_index, right)| {
8184                left.abs()
8185                    .partial_cmp(&right.abs())
8186                    .unwrap_or(std::cmp::Ordering::Equal)
8187                    .then_with(|| right_index.cmp(left_index))
8188            })
8189            .map(|(_, value)| *value)
8190            .unwrap_or(1.0);
8191        if sign_anchor < 0.0 {
8192            best.mapv_inplace(|value| -value);
8193        }
8194        canonical.column_mut(accepted).assign(&best);
8195    }
8196    Ok(canonical)
8197}
8198
8199#[cfg(test)]
8200mod canonical_nullspace_direction_tests {
8201    use super::*;
8202    use ndarray::array;
8203
8204    #[test]
8205    fn per_axis_null_penalties_are_invariant_to_eigensolver_gauge_2315() {
8206        let inv_sqrt_two = 0.5_f64.sqrt();
8207        let z = array![
8208            [inv_sqrt_two, 0.0],
8209            [inv_sqrt_two, 0.0],
8210            [0.0, 1.0],
8211            [0.0, 0.0]
8212        ];
8213        let rotation = array![[0.6, -0.8], [0.8, 0.6]];
8214        let rotated = z.dot(&rotation);
8215        let reference = canonical_nullspace_directions(&z).expect("canonical null basis");
8216        let actual =
8217            canonical_nullspace_directions(&rotated).expect("rotated canonical null basis");
8218        for axis in 0..reference.ncols() {
8219            let reference_penalty = reference
8220                .column(axis)
8221                .to_owned()
8222                .insert_axis(Axis(1))
8223                .dot(&reference.column(axis).insert_axis(Axis(0)));
8224            let actual_penalty = actual
8225                .column(axis)
8226                .to_owned()
8227                .insert_axis(Axis(1))
8228                .dot(&actual.column(axis).insert_axis(Axis(0)));
8229            let max_error = reference_penalty
8230                .iter()
8231                .zip(actual_penalty.iter())
8232                .map(|(left, right)| (left - right).abs())
8233                .fold(0.0_f64, f64::max);
8234            assert!(
8235                max_error <= 256.0 * f64::EPSILON,
8236                "axis {axis} changed by {max_error:e}"
8237            );
8238        }
8239    }
8240}
8241
8242/// Build a factor-smooth interaction basis (`bs="fs"`/`"sz"`/`"re"`).
8243///
8244/// A factor smooth replicates a shared marginal smooth in the continuous
8245/// covariate(s) once per level of a grouping factor, coupling all level blocks
8246/// through a *single* set of smoothing parameters (one per marginal penalty).
8247/// This is mgcv's `smooth.construct.fs.smooth.spec` realization and the
8248/// random-effect interpretation of a smooth: the per-level deviations are an
8249/// exchangeable family whose joint wiggliness/shrinkage is governed by the
8250/// shared λ, so the construction scales to many levels with a fixed parameter
8251/// count.
8252///
8253/// Flavours:
8254/// * `Fs` — full random factor-smooth. The marginal carries its wiggliness
8255///   penalty *and* a null-space ridge (double penalty), so the replicated
8256///   design is a proper full-rank random effect: each level's curve is shrunk
8257///   toward zero (intercept + linear trend included), recovering the mgcv
8258///   `bs="fs"` penalty structure `I_L ⊗ S_j` for every marginal penalty `S_j`.
8259/// * `Sz` — sum-to-zero factor smooth. Delegates to the existing
8260///   [`SmoothBasisSpec::FactorSumToZero`] construction (`L-1` deviation blocks,
8261///   coefficient-wise zero sum across levels).
8262/// * `Re` — pure random effect / random slope (`bs="re"`). A degree-1 marginal
8263///   gives the per-level `[1, x]` span; the penalty is the identity over each
8264///   level block (iid Gaussian coefficients), matching mgcv's `bs="re"` ridge.
8265///
8266/// The grouping levels are resolved once at fit time (sorted unique bit
8267/// patterns of the factor column) and frozen into the returned metadata so the
8268/// predict-time rebuild evaluates every row against its own level's block.
8269pub fn build_factor_smooth(
8270    data: ArrayView2<'_, f64>,
8271    spec: &FactorSmoothSpec,
8272    term_name: &str,
8273    workspace: &mut crate::basis::BasisWorkspace,
8274) -> Result<LocalSmoothTermBuild, BasisError> {
8275    if spec.continuous_cols.len() != 1 {
8276        crate::bail_invalid_basis!(
8277            "factor smooth term '{}' currently supports exactly one continuous covariate; found {}",
8278            term_name,
8279            spec.continuous_cols.len()
8280        );
8281    }
8282    let feature_col = spec.continuous_cols[0];
8283    let group_col = spec.group_col;
8284    if feature_col >= data.ncols() || group_col >= data.ncols() {
8285        crate::bail_dim_basis!(
8286            "factor smooth term '{}' references columns ({}, {}) out of bounds for {} columns",
8287            term_name,
8288            feature_col,
8289            group_col,
8290            data.ncols()
8291        );
8292    }
8293
8294    // `Sz` is exactly the existing sum-to-zero factor smooth: reuse it verbatim
8295    // so there is a single source of truth for the zero-sum construction.
8296    if matches!(spec.flavour, FactorSmoothFlavour::Sz) {
8297        let levels = resolve_factor_smooth_levels(data, group_col, spec, term_name)?;
8298        let inner = SmoothBasisSpec::BSpline1D {
8299            feature_col,
8300            spec: factor_smooth_marginal_for_replay(&spec.marginal),
8301        };
8302        let sz_term = SmoothTermSpec {
8303            frozen_parametric_residualization: None,
8304            name: term_name.to_string(),
8305            basis: SmoothBasisSpec::FactorSumToZero {
8306                inner: Box::new(inner),
8307                by_col: group_col,
8308                levels: levels.clone(),
8309                frozen_global_orthogonality: None,
8310            },
8311            shape: ShapeConstraint::None,
8312            joint_null_rotation: None,
8313        };
8314        let mut built = build_single_local_smooth_term(data, &sz_term, workspace)?;
8315        // The delegated `FactorSumToZero` build returns the BARE inner B-spline
8316        // metadata (`BasisMetadata::BSpline1D`), but the term that owns this
8317        // build carries a `SmoothBasisSpec::FactorSmooth { Sz }` spec. Two
8318        // things break if we hand that mismatched pair downstream:
8319        //   1. `freeze_smooth_basis_from_metadata` matches on (spec, metadata)
8320        //      and has no `(FactorSmooth, BSpline1D)` arm, so any refit / spatial
8321        //      re-optimization that freezes the basis aborts with a "smooth
8322        //      metadata/spec type mismatch" error.
8323        //   2. The bare B-spline metadata carries no grouping levels, so a
8324        //      predict-time rebuild cannot replay the SAME replicated design.
8325        // Re-wrap the marginal geometry as `FactorSmooth` metadata exactly as
8326        // the Fs/Re path below does, giving all three factor-smooth flavours a
8327        // single, freeze-consistent metadata shape that also pins the levels.
8328        // Since #1605 the sz marginal is ALWAYS the penalized B-spline the `fs`
8329        // sibling uses (a natural cubic regression marginal hard-enforces f''=0
8330        // at the boundary and cannot represent curved deviations — a consistency
8331        // failure). The `CubicRegression1D` arm below is therefore unreachable on
8332        // a freshly-built sz spec; it is retained only as defense / backward
8333        // compatibility for a frozen spec that still carries a cr marginal, so
8334        // the predict-time freeze restores whatever marginal class it finds.
8335        let (knots, degree, periodic, marginal_is_cr) = match &built.metadata {
8336            BasisMetadata::BSpline1D {
8337                knots,
8338                periodic,
8339                degree,
8340                ..
8341            } => (
8342                knots.clone(),
8343                degree.unwrap_or(spec.marginal.degree),
8344                *periodic,
8345                false,
8346            ),
8347            BasisMetadata::CubicRegression1D { knots, .. } => {
8348                (knots.clone(), spec.marginal.degree, None, true)
8349            }
8350            other => {
8351                crate::bail_invalid_basis!(
8352                    "sz factor smooth term '{}' produced an unexpected marginal metadata variant {:?}",
8353                    term_name,
8354                    other
8355                );
8356            }
8357        };
8358        built.metadata = BasisMetadata::FactorSmooth {
8359            continuous_cols: spec.continuous_cols.clone(),
8360            group_col,
8361            knots,
8362            degree,
8363            periodic,
8364            group_levels: levels,
8365            flavour: "sz".to_string(),
8366            marginal_is_cr,
8367        };
8368        return Ok(built);
8369    }
8370
8371    let levels = resolve_factor_smooth_levels(data, group_col, spec, term_name)?;
8372    let n_levels = levels.len();
8373    if n_levels < 2 {
8374        crate::bail_invalid_basis!(
8375            "factor smooth term '{}' requires at least two grouping levels; found {}",
8376            term_name,
8377            n_levels
8378        );
8379    }
8380
8381    // `Fs` (order ≥ 1, the default) is the random-effect flavour: it penalizes
8382    // each null-space dimension of the marginal wiggliness penalty separately
8383    // below (mgcv's `bs="fs"` construction). That replaces the marginal's single
8384    // *combined* double penalty, so disable the latter here to avoid penalizing
8385    // the null space twice (once combined, once per dimension). The explicit
8386    // `m=0` opt-out keeps the legacy combined double penalty and adds no
8387    // per-dimension penalties.
8388    let use_per_dim_null = matches!(
8389        &spec.flavour,
8390        FactorSmoothFlavour::Fs { m_null_penalty_orders }
8391            if m_null_penalty_orders.iter().copied().max().unwrap_or(0) >= 1
8392    );
8393
8394    // Build the shared marginal design + penalties from the 1-D B-spline.
8395    // `Re` forces a degree-1 marginal (linear span) and replaces the marginal
8396    // wiggliness with an identity ridge below; `Fs` keeps the user's marginal
8397    // (cubic by default) and, under the per-dimension null path, gets its null
8398    // space penalized one dimension at a time after replication.
8399    let mut marginal_spec = factor_smooth_marginal_for_replay(&spec.marginal);
8400    if use_per_dim_null {
8401        marginal_spec.double_penalty = false;
8402    }
8403    let inner_term = SmoothTermSpec {
8404            frozen_parametric_residualization: None,
8405        name: format!("{term_name}::marginal"),
8406        basis: SmoothBasisSpec::BSpline1D {
8407            feature_col,
8408            spec: marginal_spec,
8409        },
8410        shape: ShapeConstraint::None,
8411        joint_null_rotation: None,
8412    };
8413    let inner = build_single_local_smooth_term(data, &inner_term, workspace)?;
8414    let mut base = inner
8415        .design
8416        .try_to_dense_by_chunks("factor smooth marginal")
8417        .map_err(BasisError::InvalidInput)?;
8418    if matches!(spec.flavour, FactorSmoothFlavour::Re) {
8419        // `bs="re"` is a parametric random intercept+slope, not a B-spline
8420        // smooth evaluated through clamped knot support.  A degree-1 B-spline
8421        // with no internal knots spans the training rows, but outside the
8422        // boundary knots its basis is not the model matrix for `(1 + x | g)`;
8423        // held-out extrapolation then loses the random slope contribution.
8424        // Build the random-effect marginal directly as `[1, x - c]`, centered
8425        // at the frozen marginal domain, so fit-time and replay-time rows use
8426        // the same well-conditioned parametric columns on and off the training
8427        // interval.
8428        let center = match &inner.metadata {
8429            BasisMetadata::BSpline1D { knots, .. } if !knots.is_empty() => {
8430                0.5 * (knots[0] + knots[knots.len() - 1])
8431            }
8432            _ => 0.0,
8433        };
8434        let mut linear = Array2::<f64>::ones((data.nrows(), 2));
8435        linear
8436            .column_mut(1)
8437            .assign(&data.column(feature_col).mapv(|x| x - center));
8438        base = linear;
8439    }
8440    let n = base.nrows();
8441    let p = base.ncols();
8442    let q = p * n_levels;
8443
8444    // Block-diagonal replicated design: row i contributes its marginal row to
8445    // the column block owned by its grouping level, zeros elsewhere.
8446    let mut dense = Array2::<f64>::zeros((n, q));
8447    for i in 0..n {
8448        let bits = gam_data::canonical_level_bits(data[[i, group_col]]);
8449        let Some(level_idx) = levels.iter().position(|b| *b == bits) else {
8450            // Held-out-group contract (#2365): `bs="re"` is a genuine random
8451            // effect, so in the frozen (predict/replay) context a row whose
8452            // group is outside the training vocabulary carries NO fitted
8453            // deviation — its row stays all-zero across every group block and
8454            // the prediction is the population component, mirroring
8455            // `build_random_effect_block`'s lenient `group_id = None` rows.
8456            // `fs`/`sz` estimate a per-level deviation FUNCTION (an unseen
8457            // level has no zero-deviation fallback that means "population"),
8458            // and the fit path derives `levels` from this very data, so both
8459            // stay strict.
8460            if matches!(spec.flavour, FactorSmoothFlavour::Re) && spec.group_frozen_levels.is_some()
8461            {
8462                continue;
8463            }
8464            return Err(BasisError::InvalidInput(format!(
8465                "factor smooth term '{term_name}' saw an unseen grouping level at row {}",
8466                i + 1
8467            )));
8468        };
8469        let start = level_idx * p;
8470        dense
8471            .slice_mut(s![i, start..start + p])
8472            .assign(&base.row(i));
8473    }
8474
8475    // Penalties: replicate each marginal penalty into a block-diagonal
8476    // `I_L ⊗ S_j` so every level shares the same smoothing parameter λ_j (one
8477    // λ per marginal penalty), the defining feature of a factor smooth. For
8478    // `Re` the marginal penalty is replaced by one ridge per parametric
8479    // coordinate so intercept and slope variances can be learned separately.
8480    let marginal_penalties: Vec<(Array2<f64>, PenaltySource, f64)> =
8481        if matches!(spec.flavour, FactorSmoothFlavour::Re) {
8482            (0..p)
8483                .map(|j| {
8484                    let mut matrix = Array2::<f64>::zeros((p, p));
8485                    matrix[[j, j]] = 1.0;
8486                    (matrix, PenaltySource::Primary, 1.0)
8487                })
8488                .collect()
8489        } else {
8490            inner
8491                .active_penalties
8492                .iter()
8493                .map(|penalty| {
8494                    (
8495                        penalty.matrix.clone(),
8496                        penalty.info.source.clone(),
8497                        penalty.info.normalization_scale,
8498                    )
8499                })
8500                .collect()
8501        };
8502
8503    let mut candidates = Vec::<PenaltyCandidate>::with_capacity(marginal_penalties.len());
8504    for (s_inner, source, base_scale) in marginal_penalties {
8505        let mut s_big = Array2::<f64>::zeros((q, q));
8506        for level in 0..n_levels {
8507            let start = level * p;
8508            s_big
8509                .slice_mut(s![start..start + p, start..start + p])
8510                .assign(&s_inner);
8511        }
8512        let (s_big, factor_smooth_scale) = normalize_penalty_in_constrained_space(&s_big);
8513        candidates.push(PenaltyCandidate {
8514            matrix: ConstructiveQuadratic::try_from_dense_psd(
8515                s_big,
8516                "factor-smooth shared penalty",
8517            )?,
8518            source,
8519            normalization_scale: base_scale * factor_smooth_scale,
8520            kronecker_factors: None,
8521            op: None,
8522        });
8523    }
8524
8525    // `Fs` is the random-effect flavour of a smooth: the per-group curve is an
8526    // exchangeable Gaussian *function*, so EVERY coefficient — including the
8527    // {const, linear} null space of the marginal wiggliness penalty — must be
8528    // shrinkable toward zero under its own shared variance. The wiggliness
8529    // penalty `S_wiggle` shapes curvature but leaves the per-group intercept and
8530    // slope (its null space) completely UNPENALIZED. With the null space free,
8531    // each group fits its own intercept and slope with NO partial pooling, so
8532    // the held-out per-subject forecast inherits the full no-pooling variance
8533    // and curves away from the true per-group line (gam#712 real arm, gam#713;
8534    // gam#903 sleepstudy forecast ran ~74% over the lme4 BLUP bar).
8535    //
8536    // mgcv's `bs="fs"` fixes this by penalizing each null-space dimension
8537    // SEPARATELY (`smooth.construct.fs.smooth.spec` adds one rank-1 penalty per
8538    // null coordinate), each replicated block-diagonally across levels under a
8539    // single shared smoothing parameter — so REML fits a distinct
8540    // random-intercept variance and random-slope variance, the partial pooling
8541    // that makes the forecast track lme4's correlated random-effect BLUP. A
8542    // single *combined* null penalty (one λ for intercept+slope together) cannot
8543    // express the typically very different intercept and slope variances, which
8544    // is the residual forecast gap. We mirror mgcv exactly: for each orthonormal
8545    // canonical null direction `z_k` of the marginal wiggliness penalty, add
8546    // `I_L ⊗ (z_k z_kᵀ)` as its own penalty. The marginal's combined double
8547    // penalty was disabled above, so the null space is penalized once, per
8548    // dimension. With linear data REML drives the curvature λ up and degrades
8549    // `fs` to a linear random slope (edf → ≈2/group); with genuine curvature the
8550    // wiggliness λ stays small and the wiggle survives (data-adaptive, not a
8551    // cap). Gated by `m_null_penalty_orders`: order ≥ 1 (default) enables the
8552    // per-dimension null penalties; `m=0` keeps the legacy combined double
8553    // penalty and adds nothing here.
8554    if use_per_dim_null
8555        && let Some(Some(z)) = inner
8556            .active_penalties
8557            .first()
8558            .map(|penalty| &penalty.null_eigenvectors)
8559        && z.nrows() == p
8560    {
8561        let z = canonical_nullspace_directions(z)?;
8562        for k in 0..z.ncols() {
8563            // Rank-1 marginal penalty `z_k z_kᵀ`, replicated block-diagonally
8564            // across levels into `I_L ⊗ (z_k z_kᵀ)`. Its own λ is one shared
8565            // variance for this null component (intercept or slope) across all
8566            // groups — the random-effect structure of mgcv `fs`.
8567            let zk = z.column(k);
8568            let mut p_k = Array2::<f64>::zeros((p, p));
8569            for a in 0..p {
8570                for b in 0..p {
8571                    p_k[[a, b]] = zk[a] * zk[b];
8572                }
8573            }
8574            let mut s_null = Array2::<f64>::zeros((q, q));
8575            for level in 0..n_levels {
8576                let start = level * p;
8577                s_null
8578                    .slice_mut(s![start..start + p, start..start + p])
8579                    .assign(&p_k);
8580            }
8581            let (s_null, null_scale) = normalize_penalty_in_constrained_space(&s_null);
8582            candidates.push(PenaltyCandidate {
8583                matrix: ConstructiveQuadratic::try_from_dense_psd(
8584                    s_null,
8585                    "factor-smooth null-function penalty",
8586                )?,
8587                source: PenaltySource::Primary,
8588                normalization_scale: null_scale,
8589                kronecker_factors: None,
8590                op: None,
8591            });
8592        }
8593    }
8594    let filtered = crate::basis::filter_penalty_candidates(candidates)?;
8595    let joint_null_rotation = crate::basis::compute_joint_null_rotation(&filtered.active)?;
8596    let mut dropped_penalties = inner.dropped_penalties;
8597    dropped_penalties.extend(filtered.dropped);
8598
8599    // Metadata: carry the marginal knot geometry + frozen levels so prediction
8600    // reconstructs an identical replicated design.
8601    let (knots, degree, periodic) = match &inner.metadata {
8602        BasisMetadata::BSpline1D {
8603            knots,
8604            periodic,
8605            degree,
8606            ..
8607        } => (
8608            knots.clone(),
8609            degree.unwrap_or(spec.marginal.degree),
8610            *periodic,
8611        ),
8612        other => {
8613            crate::bail_invalid_basis!(
8614                "factor smooth term '{}' produced an unexpected marginal metadata variant {:?}",
8615                term_name,
8616                other
8617            );
8618        }
8619    };
8620    let flavour_tag = match &spec.flavour {
8621        FactorSmoothFlavour::Fs { .. } => "fs",
8622        FactorSmoothFlavour::Sz => "sz",
8623        FactorSmoothFlavour::Re => "re",
8624    }
8625    .to_string();
8626    let metadata = BasisMetadata::FactorSmooth {
8627        continuous_cols: spec.continuous_cols.clone(),
8628        group_col,
8629        knots,
8630        degree,
8631        periodic,
8632        group_levels: levels,
8633        flavour: flavour_tag,
8634        // fs/re marginals are always B-spline; the cr marginal is sz-only and
8635        // handled on the dedicated Sz path above.
8636        marginal_is_cr: false,
8637    };
8638
8639    Ok(LocalSmoothTermBuild {
8640        dim: q,
8641        design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(dense)),
8642        // Exactly one group block is active on every row, so a common
8643        // marginal anchor contributes its lift once for that row.
8644        affine_offset: inner.affine_offset,
8645        active_penalties: filtered.active,
8646        joint_null_rotation,
8647        dropped_penalties,
8648        metadata,
8649        linear_constraints: None,
8650        box_reparam: false,
8651        kronecker_factored: None,
8652    })
8653}
8654
8655/// Resolve the grouping levels for a factor smooth: replay the frozen level
8656/// list when present (predict path), otherwise discover the sorted unique bit
8657/// patterns of the factor column (fit path).
8658pub fn resolve_factor_smooth_levels(
8659    data: ArrayView2<'_, f64>,
8660    group_col: usize,
8661    spec: &FactorSmoothSpec,
8662    term_name: &str,
8663) -> Result<Vec<u64>, BasisError> {
8664    if let Some(frozen) = &spec.group_frozen_levels {
8665        if frozen.is_empty() {
8666            crate::bail_invalid_basis!(
8667                "factor smooth term '{}' has an empty frozen level list",
8668                term_name
8669            );
8670        }
8671        return Ok(frozen
8672            .iter()
8673            .map(|&b| gam_data::canonical_level_bits(f64::from_bits(b)))
8674            .collect());
8675    }
8676    let mut bits: Vec<u64> = data
8677        .column(group_col)
8678        .iter()
8679        .map(|v| gam_data::canonical_level_bits(*v))
8680        .collect();
8681    bits.sort_by(|a, b| {
8682        f64::from_bits(*a)
8683            .partial_cmp(&f64::from_bits(*b))
8684            .unwrap_or(std::cmp::Ordering::Equal)
8685    });
8686    bits.dedup();
8687    Ok(bits)
8688}
8689
8690/// Marginal B-spline spec for a factor-smooth block. The marginal always builds
8691/// without an identifiability constraint (the per-level replication, not a
8692/// sum-to-zero side constraint, provides identifiability against the parametric
8693/// block). At predict time the marginal's knot geometry has already been pinned
8694/// into `marginal.knotspec` by the metadata replay, so the spec is used
8695/// verbatim aside from clearing the identifiability transform.
8696pub fn factor_smooth_marginal_for_replay(marginal: &BSplineBasisSpec) -> BSplineBasisSpec {
8697    let mut m = marginal.clone();
8698    m.identifiability = BSplineIdentifiability::None;
8699    m
8700}
8701
8702pub fn build_single_local_smooth_term(
8703    data: ArrayView2<'_, f64>,
8704    term: &SmoothTermSpec,
8705    workspace: &mut crate::basis::BasisWorkspace,
8706) -> Result<LocalSmoothTermBuild, BasisError> {
8707    term.basis.validate_scale_configuration()?;
8708    if term.shape != ShapeConstraint::None && !shape_supports_basis(term) {
8709        crate::bail_invalid_basis!(
8710            "ShapeConstraint::{:?} is unsupported for term '{}'",
8711            term.shape,
8712            term.name
8713        );
8714    }
8715    if let SmoothBasisSpec::ByVariable {
8716        inner,
8717        by_col,
8718        kind,
8719        by,
8720    } = &term.basis
8721    {
8722        ensure_by_variable_specs_match(kind, by, &term.name)?;
8723        let mut inner_basis = (**inner).clone();
8724        // Factor-level `by=` owns model-space centering (it centers the gated
8725        // block against the level indicator downstream). Defer the inner
8726        // basis's default pooled centering so the level block is not
8727        // double-centered down to `k-2` columns (#1427). Numeric-by smooths are
8728        // untouched: they are not row-gated to a level and keep ordinary
8729        // intercept centering.
8730        if matches!(by, ByVariableSpec::Level { .. }) {
8731            defer_inner_model_centering_to_factor_level_wrapper(&mut inner_basis);
8732        }
8733        let inner_term = SmoothTermSpec {
8734            frozen_parametric_residualization: None,
8735            name: term.name.clone(),
8736            basis: inner_basis,
8737            shape: term.shape,
8738            joint_null_rotation: None,
8739        };
8740        let built = build_single_local_smooth_term(data, &inner_term, workspace)?;
8741        return apply_by_variable_to_local_build(built, data, *by_col, by, &term.name);
8742    }
8743
8744    // BySmooth: a `by=` smooth that unifies numeric or factor modulation into a
8745    // single term.  Lower it here so the downstream match does not need an arm.
8746    if let SmoothBasisSpec::BySmooth { smooth, by_kind } = &term.basis {
8747        return build_by_smooth_local(data, term, smooth, by_kind, workspace);
8748    }
8749
8750    let mut built: BasisBuildResult = match &term.basis {
8751        SmoothBasisSpec::FactorSumToZero {
8752            inner,
8753            by_col,
8754            levels,
8755            ..
8756        } => {
8757            if *by_col >= data.ncols() {
8758                crate::bail_dim_basis!(
8759                    "term '{}' by column {} out of bounds for {} columns",
8760                    term.name,
8761                    by_col,
8762                    data.ncols()
8763                );
8764            }
8765            if levels.len() < 2 {
8766                crate::bail_invalid_basis!(
8767                    "sum-to-zero factor smooth term '{}' requires at least two levels",
8768                    term.name
8769                );
8770            }
8771            if term.shape != ShapeConstraint::None {
8772                crate::bail_invalid_basis!(
8773                    "ShapeConstraint::{:?} is unsupported for sum-to-zero factor smooth term '{}'",
8774                    term.shape,
8775                    term.name
8776                );
8777            }
8778            let inner_term = SmoothTermSpec {
8779            frozen_parametric_residualization: None,
8780                name: format!("{}::inner", term.name),
8781                basis: (**inner).clone(),
8782                shape: ShapeConstraint::None,
8783                joint_null_rotation: None,
8784            };
8785            let mut inner_built = build_single_local_smooth_term(data, &inner_term, workspace)?;
8786            if inner_built.affine_offset.is_some() {
8787                crate::bail_invalid_basis!(
8788                    "sum-to-zero factor smooth term '{}' cannot contain a non-zero endpoint anchor: a shared fixed affine lift would violate the per-covariate zero-sum deviation identity",
8789                    term.name
8790                );
8791            }
8792            // Capture the marginal penalty's null directions BEFORE the penalty
8793            // vector is rebuilt below; the sum-to-zero null-space ridge replicates
8794            // these `z_k` into the contrast space (mgcv `bs="fs"` double-penalty).
8795            let inner_null_eigenvectors = inner_built
8796                .active_penalties
8797                .first()
8798                .and_then(|penalty| penalty.null_eigenvectors.clone());
8799            let base = inner_built
8800                .design
8801                .try_to_dense_by_chunks("sum-to-zero factor smooth")
8802                .map_err(BasisError::InvalidInput)?;
8803            let n = base.nrows();
8804            let p = base.ncols();
8805            let l_minus_one = levels.len() - 1;
8806            // Canonicalize the stored level keys once so signed-zero / NaN codes
8807            // match regardless of how the level set was interned (#2145/#2146).
8808            let canon_levels: Vec<u64> = levels
8809                .iter()
8810                .map(|&b| gam_data::canonical_level_bits(f64::from_bits(b)))
8811                .collect();
8812            let mut dense = Array2::<f64>::zeros((n, p * l_minus_one));
8813            for i in 0..n {
8814                let bits = gam_data::canonical_level_bits(data[[i, *by_col]]);
8815                let level_idx = canon_levels
8816                    .iter()
8817                    .position(|b| *b == bits)
8818                    .ok_or_else(|| {
8819                        BasisError::InvalidInput(format!(
8820                            "sum-to-zero factor smooth term '{}' saw an unseen level at row {}",
8821                            term.name,
8822                            i + 1
8823                        ))
8824                    })?;
8825                if level_idx < l_minus_one {
8826                    let start = level_idx * p;
8827                    dense
8828                        .slice_mut(s![i, start..start + p])
8829                        .assign(&base.row(i));
8830                } else {
8831                    for level in 0..l_minus_one {
8832                        let start = level * p;
8833                        dense
8834                            .slice_mut(s![i, start..start + p])
8835                            .assign(&base.row(i).mapv(|v| -v));
8836                    }
8837                }
8838            }
8839            let mut candidates = Vec::<PenaltyCandidate>::with_capacity(
8840                inner_built.active_penalties.len() * levels.len(),
8841            );
8842            // Replicate each marginal penalty into the sum-to-zero contrast
8843            // space. With `L-1` free deviation blocks and the reference level
8844            // `d_L = -Σ_{k<L} d_k`, the marginal penalty summed over ALL `L`
8845            // levels, `Σ_{k=1}^{L} d_kᵀ S d_k`, expands to the `(I + 11ᵀ) ⊗ S`
8846            // contrast form (factor 2 on the diagonal blocks, 1 off-diagonal).
8847            //
8848            // PER-GROUP SMOOTHING PARAMETERS (#1074). mgcv's `bs="sz"` does NOT
8849            // pool that sum under one λ: `smooth.construct.sz` emits ONE penalty
8850            // matrix per factor level (here 6 separate `S`s, each with its own
8851            // smoothing parameter), so REML can shrink a low-amplitude group's
8852            // deviation curve hard while leaving a high-amplitude group nearly
8853            // unpenalized. A single shared wiggliness λ (the old construction)
8854            // forces every group to the SAME curvature budget, so a group whose
8855            // true curve is flat drags curvature into the noise of the busy
8856            // groups and vice-versa — systematic truth-recovery loss even when
8857            // the pooled total edf matches mgcv's (the observed `sz` 1.23× gap).
8858            //
8859            // We mirror mgcv exactly by splitting the per-marginal penalty
8860            // `Σ_{k=1}^{L} d_kᵀ S d_k` back into its `L` independent
8861            // rank-controlled summands BEFORE mapping to the contrast space, each
8862            // carrying its own λ:
8863            //   * level k < L (free block):  `d_kᵀ S d_k` → block-diagonal
8864            //     `(e_k e_kᵀ) ⊗ S`  (only the (k,k) block is `S`).
8865            //   * level L (reference):       `d_Lᵀ S d_L = (Σ_{j<L} d_j)ᵀ S (·)`
8866            //     → the fully-coupled `(11ᵀ) ⊗ S` block.
8867            // Summed at equal λ these `L` blocks recover the old `(I + 11ᵀ) ⊗ S`
8868            // exactly (`Σ_k e_k e_kᵀ = I`), so this is a strict generalization:
8869            // the pooled fit is still reachable, REML only GAINS the freedom to
8870            // spend curvature per group. The zero-sum reparameterization (hence
8871            // the `sz` vs `fs` identifiability) is untouched.
8872            //
8873            // `which_level ∈ 0..=l_minus_one`: `< l_minus_one` selects the single
8874            // free deviation block; `== l_minus_one` selects the reference-level
8875            // coupling block.
8876            let stz_per_group_penalty =
8877                |s_inner: &Array2<f64>, which_level: usize| -> Array2<f64> {
8878                    let mut s_big = Array2::<f64>::zeros((p * l_minus_one, p * l_minus_one));
8879                    if which_level < l_minus_one {
8880                        // (e_k e_kᵀ) ⊗ S: a single diagonal block.
8881                        let k = which_level;
8882                        let mut block = s_big.slice_mut(s![k * p..(k + 1) * p, k * p..(k + 1) * p]);
8883                        block.assign(s_inner);
8884                    } else {
8885                        // (11ᵀ) ⊗ S: every block (diagonal and off-diagonal) is S.
8886                        for a in 0..l_minus_one {
8887                            for b in 0..l_minus_one {
8888                                let mut block =
8889                                    s_big.slice_mut(s![a * p..(a + 1) * p, b * p..(b + 1) * p]);
8890                                block.assign(s_inner);
8891                            }
8892                        }
8893                    }
8894                    s_big
8895                };
8896            for base_penalty in &inner_built.active_penalties {
8897                // Emit `L` independent per-level blocks for this marginal penalty.
8898                for which_level in 0..=l_minus_one {
8899                    let raw = stz_per_group_penalty(&base_penalty.matrix, which_level);
8900                    let (s_big, group_scale) = normalize_penalty_in_constrained_space(&raw);
8901                    candidates.push(PenaltyCandidate {
8902                        matrix: ConstructiveQuadratic::try_from_dense_psd(
8903                            s_big,
8904                            "grouped factor-smooth penalty",
8905                        )?,
8906                        source: base_penalty.info.source.clone(),
8907                        normalization_scale: base_penalty.info.normalization_scale * group_scale,
8908                        kronecker_factors: None,
8909                        op: None,
8910                    });
8911                }
8912            }
8913
8914            // Null-space ridge, mirroring the `bs="fs"` double-penalty
8915            // construction (#1605, same defect class as #700/#712/#713). The
8916            // marginal wiggliness penalty `S` shapes curvature but leaves the
8917            // {const, linear} null space of each deviation curve COMPLETELY
8918            // unpenalized. With that null space free, the single combined
8919            // wiggliness smoothing parameter cannot separate the per-group
8920            // intercept/slope variance from the curvature variance, so REML
8921            // parks the wiggliness `λ` high — over-smoothing (under-fitting) the
8922            // deviation blocks even when the truth lives in their span (the `sz`
8923            // recovery gap vs the `fs` superset). mgcv's `bs="fs"` fixes the
8924            // analogous gap by penalizing each null-space dimension SEPARATELY
8925            // under its own shared variance; we mirror that here while keeping
8926            // the zero-sum reparameterization, so the constraint (and the
8927            // identifiability of `sz` vs `fs`) is preserved. For each orthonormal
8928            // canonical null direction `z_k` of the marginal penalty, add the
8929            // rank-1 marginal penalty `z_k z_kᵀ` mapped into the SAME `(I + 11ᵀ)`
8930            // sum-to-zero contrast space, each carrying its own `λ`.
8931            if let Some(z) = inner_null_eigenvectors.as_ref()
8932                && z.nrows() == p
8933            {
8934                let z = canonical_nullspace_directions(z)?;
8935                for k in 0..z.ncols() {
8936                    let zk = z.column(k);
8937                    let mut p_k = Array2::<f64>::zeros((p, p));
8938                    for a in 0..p {
8939                        for b in 0..p {
8940                            p_k[[a, b]] = zk[a] * zk[b];
8941                        }
8942                    }
8943                    // Null ridges stay POOLED (the `(I + 11ᵀ) ⊗ z_k z_kᵀ` form):
8944                    // they govern the per-group intercept/slope shrinkage, which
8945                    // mgcv pools under one variance even for `sz`; only the
8946                    // curvature (wiggliness) penalty is split per group above.
8947                    let stz_pooled_null = {
8948                        let mut s_big = Array2::<f64>::zeros((p * l_minus_one, p * l_minus_one));
8949                        for a in 0..l_minus_one {
8950                            for b in 0..l_minus_one {
8951                                let factor = if a == b { 2.0 } else { 1.0 };
8952                                let mut block =
8953                                    s_big.slice_mut(s![a * p..(a + 1) * p, b * p..(b + 1) * p]);
8954                                block.assign(&p_k.mapv(|v| v * factor));
8955                            }
8956                        }
8957                        s_big
8958                    };
8959                    let (s_null, null_scale) =
8960                        normalize_penalty_in_constrained_space(&stz_pooled_null);
8961                    candidates.push(PenaltyCandidate {
8962                        matrix: ConstructiveQuadratic::try_from_dense_psd(
8963                            s_null,
8964                            "grouped factor-smooth null penalty",
8965                        )?,
8966                        source: PenaltySource::DoublePenaltyNullspace,
8967                        normalization_scale: null_scale,
8968                        kronecker_factors: None,
8969                        op: None,
8970                    });
8971                }
8972            }
8973            let filtered = crate::basis::filter_penalty_candidates(candidates)?;
8974            let mut dropped_penalties = std::mem::take(&mut inner_built.dropped_penalties);
8975            dropped_penalties.extend(filtered.dropped);
8976            inner_built.dim = p * l_minus_one;
8977            inner_built.design =
8978                DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(dense));
8979            inner_built.active_penalties = filtered.active;
8980            inner_built.dropped_penalties = dropped_penalties;
8981            inner_built.joint_null_rotation =
8982                crate::basis::compute_joint_null_rotation(&inner_built.active_penalties)?;
8983            inner_built.kronecker_factored = None;
8984            return Ok(inner_built);
8985        }
8986        SmoothBasisSpec::BSpline1D { feature_col, spec } => {
8987            if *feature_col >= data.ncols() {
8988                crate::bail_dim_basis!(
8989                    "term '{}' feature column {} out of bounds for {} columns",
8990                    term.name,
8991                    feature_col,
8992                    data.ncols()
8993                );
8994            }
8995            let mut spec_local = spec.clone();
8996            if term.shape != ShapeConstraint::None {
8997                // Shape-constrained B-splines are anchored by construction.
8998                // Sum-to-zero side constraints conflict with monotonic/convex cones.
8999                spec_local.identifiability = BSplineIdentifiability::None;
9000            }
9001            // Endpoint boundary conditions are structural for B-splines: the
9002            // basis builder bakes their homogeneous nullspace transform into
9003            // the design, penalties, and stored raw-basis transform.
9004            build_bspline_basis_1d(data.column(*feature_col), &spec_local)?
9005        }
9006        SmoothBasisSpec::ThinPlate {
9007            feature_cols,
9008            spec,
9009            input_scale,
9010        } => {
9011            if term.shape != ShapeConstraint::None {
9012                if feature_cols.len() != 1 {
9013                    crate::bail_invalid_basis!(
9014                        "ShapeConstraint::{:?} for term '{}' on ThinPlate basis requires exactly 1 feature axis; found {}",
9015                        term.shape,
9016                        term.name,
9017                        feature_cols.len()
9018                    );
9019                }
9020            }
9021            let mut spec_local = spec.clone();
9022            let frame = term.basis.scale_contract().normalize_euclidean_frame(
9023                select_columns(data, feature_cols)?,
9024                *input_scale,
9025                Some(spec.length_scale),
9026                &mut spec_local.center_strategy,
9027            )?;
9028            let x = frame.coordinates;
9029            let realized_input_scale = frame.input_scale;
9030            let length_scale_eff = frame
9031                .length_scale
9032                .expect("ThinPlate declares a required length-scale coordinate");
9033            spec_local.length_scale = length_scale_eff.standardized_value();
9034            if matches!(
9035                spec_local.identifiability,
9036                SpatialIdentifiability::OrthogonalToParametric
9037            ) {
9038                spec_local.identifiability = SpatialIdentifiability::None;
9039            }
9040            let mut result = build_thin_plate_basis(x.view(), &spec_local).map_err(|err| {
9041                rewrite_thin_plate_knots_error(err, &term.name, feature_cols.len(), spec)
9042            })?;
9043            // Inject the input scale into metadata; also restore the user's
9044            // original length_scale (not the σ_geom-compensated one) so a
9045            // metadata-driven rebuild that re-applies compensation does not
9046            // double-divide. The build may auto-promote to Duchon when
9047            // canonical TPS is infeasible (k < polynomial-nullspace size);
9048            // in that case patch the Duchon metadata variant so predict-time
9049            // round-trips through the same standardized data path.
9050            match &mut result.metadata {
9051                BasisMetadata::ThinPlate {
9052                    input_scale: metadata_scale,
9053                    length_scale,
9054                    ..
9055                } => {
9056                    *metadata_scale = realized_input_scale;
9057                    *length_scale = crate::OriginalUnits::new(spec.length_scale);
9058                }
9059                BasisMetadata::Duchon {
9060                    input_scale: metadata_scale,
9061                    length_scale,
9062                    ..
9063                } => {
9064                    // Auto-promotion (canonical TPS infeasible at this (d, k)).
9065                    // Since #1091 the promotion does NOT forward the incoming
9066                    // σ_geom-compensated `spec_local.length_scale` to the Duchon
9067                    // builder — it DISCARDS it and substitutes the geometric mean
9068                    // of the center pairwise distances (`promotion_length_scale`,
9069                    // the natural radial-kernel scale where κ·r ≈ O(1)). So the
9070                    // realized kernel bandwidth recorded in this metadata bears no
9071                    // fixed relation to the user's `spec.length_scale`; clobbering
9072                    // it to the user-facing value (the pre-#1091 behavior) makes
9073                    // freeze→replay re-derive `compensate(spec.length_scale, σ) ≠
9074                    // promotion_length_scale`, evaluating the kernel at the wrong
9075                    // bandwidth and corrupting the replayed design (#1091 broke the
9076                    // e7ff5ed83 freeze contract for the auto-promoted path).
9077                    //
9078                    // The freeze→replay round trip rebuilds through the Duchon arm,
9079                    // which re-applies σ_geom compensation:
9080                    //   replay_eff = compensate(metadata.length_scale, σ)
9081                    //              = metadata.length_scale / σ_geom.
9082                    // For replay_eff to reproduce the realized `promotion_length_scale`
9083                    // we must store the UN-compensated value `promotion_length_scale
9084                    // · σ_geom`. `compensate(1.0, σ) = 1/σ_geom`, so divide the
9085                    // realized scale by it to multiply back through σ_geom. With no
9086                    // standardization. Restore original units before freezing.
9087                    if let Some(promoted) = *length_scale {
9088                        // The promotion builder tagged its output against its
9089                        // own `input_scale: ONE`, but the coordinates it saw
9090                        // were already standardized by `realized_input_scale`.
9091                        // So the value is a STANDARDIZED length wearing the
9092                        // builder's local tag; re-base it onto the scale this
9093                        // metadata is about to carry. Same arithmetic as
9094                        // before (multiply through by σ_geom), now stated as
9095                        // the frame conversion it always was.
9096                        let promoted = crate::StandardizedUnits::new(promoted.original_value());
9097                        *length_scale = Some(realized_input_scale.to_original_units(promoted));
9098                    }
9099                    *metadata_scale = realized_input_scale;
9100                }
9101                // `build_thin_plate_basis` emits exactly one of the two arms
9102                // above (TPS, or its Duchon auto-promotion). Any other metadata
9103                // means the standardized-coordinate patch above silently did not
9104                // run, which would leave the frozen input scale unset and break
9105                // freeze→replay; refuse instead of shipping the unpatched design.
9106                _ => {
9107                    crate::bail_invalid_basis!(
9108                        "term '{}' thin-plate build produced metadata that is neither ThinPlate \
9109                         nor its Duchon auto-promotion, so the realized input scale cannot be frozen",
9110                        term.name
9111                    );
9112                }
9113            }
9114            result
9115        }
9116        SmoothBasisSpec::Sphere { feature_cols, spec } => {
9117            if term.shape != ShapeConstraint::None {
9118                crate::bail_invalid_basis!(
9119                    "ShapeConstraint::{:?} for term '{}' is not supported on spherical splines",
9120                    term.shape,
9121                    term.name
9122                );
9123            }
9124            let x = select_columns(data, feature_cols)?;
9125            build_spherical_spline_basis(x.view(), spec)?
9126        }
9127        SmoothBasisSpec::ConstantCurvature { feature_cols, spec } => {
9128            if term.shape != ShapeConstraint::None {
9129                crate::bail_invalid_basis!(
9130                    "ShapeConstraint::{:?} for term '{}' is not supported on constant-curvature smooths",
9131                    term.shape,
9132                    term.name
9133                );
9134            }
9135            // Chart coordinates are consumed verbatim: NO auto-standardization.
9136            // Rescaling axes would change the chart gauge `1 + κ‖x‖²` and
9137            // silently redefine which curvature κ refers to (the same point
9138            // cloud at a different chart scale has a different κ̂); the user's
9139            // coordinates ARE the geometry here, exactly as for the sphere
9140            // smooth's (lat, lon).
9141            let x = select_columns(data, feature_cols)?;
9142            build_constant_curvature_basis(x.view(), spec)?
9143        }
9144        SmoothBasisSpec::MeasureJet {
9145            feature_cols,
9146            spec,
9147            input_scale,
9148        } => {
9149            if term.shape != ShapeConstraint::None {
9150                crate::bail_invalid_basis!(
9151                    "ShapeConstraint::{:?} for term '{}' is not supported on measure-jet smooths",
9152                    term.shape,
9153                    term.name
9154                );
9155            }
9156            // The typed scale contract owns the intentionally asymmetric
9157            // fresh/replay rule: fresh explicit ranges are in original units,
9158            // while a frozen MeasureJet range is already in its realized frame.
9159            let mut spec_local = spec.clone();
9160            let frame = term.basis.scale_contract().normalize_euclidean_frame(
9161                select_columns(data, feature_cols)?,
9162                *input_scale,
9163                Some(spec.length_scale),
9164                &mut spec_local.center_strategy,
9165            )?;
9166            let x = frame.coordinates;
9167            let realized_input_scale = frame.input_scale;
9168            let length_scale_eff = frame
9169                .length_scale
9170                .expect("MeasureJet declares a required length-scale coordinate");
9171            spec_local.length_scale = length_scale_eff.standardized_value();
9172            let mut result = build_measure_jet_basis(x.view(), &spec_local)?;
9173            if let BasisMetadata::MeasureJet {
9174                input_scale: metadata_scale,
9175                ..
9176            } = &mut result.metadata
9177            {
9178                *metadata_scale = realized_input_scale;
9179            }
9180            result
9181        }
9182        SmoothBasisSpec::Matern {
9183            feature_cols,
9184            spec,
9185            input_scale,
9186        } => {
9187            if term.shape != ShapeConstraint::None {
9188                if feature_cols.len() != 1 {
9189                    crate::bail_invalid_basis!(
9190                        "ShapeConstraint::{:?} for term '{}' on Matern basis requires exactly 1 feature axis; found {}",
9191                        term.shape,
9192                        term.name,
9193                        feature_cols.len()
9194                    );
9195                }
9196            }
9197            let original_length_scale = spec.length_scale.resolved().ok_or_else(|| {
9198                BasisError::InvalidInput(format!(
9199                    "term '{}' reached Matérn construction before its Auto length scale was resolved",
9200                    term.name
9201                ))
9202            })?;
9203            let mut spec_local = spec.clone();
9204            let frame = term.basis.scale_contract().normalize_euclidean_frame(
9205                select_columns(data, feature_cols)?,
9206                *input_scale,
9207                Some(original_length_scale),
9208                &mut spec_local.center_strategy,
9209            )?;
9210            let x = frame.coordinates;
9211            let realized_input_scale = frame.input_scale;
9212            let length_scale_eff = frame
9213                .length_scale
9214                .expect("Matérn declares a required length-scale coordinate");
9215            spec_local
9216                .length_scale
9217                .set_resolved(length_scale_eff.standardized_value());
9218            let mut result = build_matern_basiswithworkspace(x.view(), &spec_local, workspace)?;
9219            if let BasisMetadata::Matern {
9220                input_scale: metadata_scale,
9221                length_scale,
9222                ..
9223            } = &mut result.metadata
9224            {
9225                *metadata_scale = realized_input_scale;
9226                *length_scale = crate::OriginalUnits::new(original_length_scale);
9227            }
9228            result
9229        }
9230        SmoothBasisSpec::Duchon {
9231            feature_cols,
9232            spec,
9233            input_scale,
9234        } => {
9235            if term.shape != ShapeConstraint::None {
9236                if feature_cols.len() != 1 {
9237                    crate::bail_invalid_basis!(
9238                        "ShapeConstraint::{:?} for term '{}' on Duchon basis requires exactly 1 feature axis; found {}",
9239                        term.shape,
9240                        term.name,
9241                        feature_cols.len()
9242                    );
9243                }
9244            }
9245            let mut spec_local = spec.clone();
9246            let frame = term.basis.scale_contract().normalize_euclidean_frame(
9247                select_columns(data, feature_cols)?,
9248                *input_scale,
9249                spec.length_scale,
9250                &mut spec_local.center_strategy,
9251            )?;
9252            let x = frame.coordinates;
9253            let realized_input_scale = frame.input_scale;
9254            let length_scale_eff = frame.length_scale;
9255            spec_local.length_scale =
9256                length_scale_eff.map(crate::StandardizedUnits::standardized_value);
9257            // The Duchon input axis is standardized in place above (`x → x/σ`,
9258            // scale-only, no centering). A 1-D cyclic boundary `[start, end)`
9259            // declared in ORIGINAL covariate units must move into that same
9260            // standardized frame, or the periodic wrap in
9261            // `build_periodic_duchon_basis_1d` (which the cyclic-boundary
9262            // dispatch in `build_duchon_basis_uncached` normalizes onto) folds
9263            // the standardized coordinate against an original-unit period: the
9264            // seam never closes and the basis silently degrades to
9265            // non-periodic (#1074: `duchon(x, periodic=true)` predictions
9266            // diverged across the wrap, f(0) ≠ f(2π)). Rescale by the same
9267            // 1/σ applied to the data so training and predict share one
9268            // periodic geometry.
9269            if let crate::basis::OneDimensionalBoundary::Cyclic { start, end } =
9270                spec_local.boundary.clone()
9271            {
9272                spec_local.boundary = crate::basis::OneDimensionalBoundary::Cyclic {
9273                    start: realized_input_scale
9274                        .to_standardized_units(crate::OriginalUnits::new(start))
9275                        .standardized_value(),
9276                    end: realized_input_scale
9277                        .to_standardized_units(crate::OriginalUnits::new(end))
9278                        .standardized_value(),
9279                };
9280            }
9281            // The SAME original-units-vs-standardized-frame reasoning applies
9282            // to `spec.periodic` (the per-axis period vector the position API
9283            // and mixed-periodicity tensor paths use): each declared period is
9284            // in original covariate units and must be divided by that axis's
9285            // uniform input scale, or the wrap folds standardized coordinates against an
9286            // original-unit period (the #1074 seam failure, previously fixed
9287            // only for the 1-D `boundary` spelling above).
9288            if let Some(periods) = spec_local.periodic.as_mut() {
9289                for axis_period in periods {
9290                    if let Some(period) = axis_period.as_mut() {
9291                        *period = realized_input_scale
9292                            .to_standardized_units(crate::OriginalUnits::new(*period))
9293                            .standardized_value();
9294                    }
9295                }
9296            }
9297            if matches!(
9298                spec_local.identifiability,
9299                SpatialIdentifiability::OrthogonalToParametric
9300            ) {
9301                spec_local.identifiability = SpatialIdentifiability::None;
9302            }
9303            let mut result = build_duchon_basiswithworkspace(x.view(), &spec_local, workspace)?;
9304            if let BasisMetadata::Duchon {
9305                input_scale: metadata_scale,
9306                length_scale,
9307                periodic,
9308                ..
9309            } = &mut result.metadata
9310            {
9311                *metadata_scale = realized_input_scale;
9312                *length_scale = spec.length_scale.map(crate::OriginalUnits::new);
9313                // Same convention as `length_scale`: metadata (and hence the
9314                // frozen replay spec design_freezing copies it into) always
9315                // stores the period in ORIGINAL covariate units, and the
9316                // standardization rescale above recomputes the standardized
9317                // period fresh from `input_scale` on EVERY build — fresh fit
9318                // and frozen replay alike — so the compensation stays
9319                // idempotent with no fit-vs-replay branch. Leaving the
9320                // builder-resolved (standardized-frame) period here would
9321                // double-divide on replay. Invariant this relies on: every
9322                // producer that sets a Cyclic `boundary` also sets
9323                // `spec.periodic` from the same original-units source (the
9324                // formula DSL does; see `parse_periodic_axes_option` /
9325                // `parse_cyclic_boundary` in term_builder.rs), so the pristine
9326                // `spec.periodic` is a valid original-units record for the
9327                // boundary spelling too.
9328                if spec.periodic.is_some() || spec.boundary.period().is_some() {
9329                    *periodic = spec
9330                        .periodic
9331                        .clone()
9332                        .or_else(|| spec.boundary.period().map(|(_, _, p)| vec![Some(p)]));
9333                }
9334            }
9335            result
9336        }
9337        SmoothBasisSpec::Pca {
9338            feature_cols,
9339            basis_matrix,
9340            centered,
9341            smooth_penalty,
9342            center_mean,
9343            pca_basis_path,
9344            chunk_size,
9345        } => {
9346            if term.shape != ShapeConstraint::None {
9347                crate::bail_invalid_basis!(
9348                    "ShapeConstraint::{:?} for term '{}' is not supported on Pca basis",
9349                    term.shape,
9350                    term.name
9351                );
9352            }
9353            build_pca_smooth_basis(
9354                data,
9355                feature_cols,
9356                basis_matrix,
9357                *centered,
9358                *smooth_penalty,
9359                center_mean.as_ref(),
9360                pca_basis_path.as_ref(),
9361                *chunk_size,
9362            )?
9363        }
9364        SmoothBasisSpec::TensorBSpline { feature_cols, spec } => {
9365            build_tensor_bspline_basis(data, feature_cols, spec)?
9366        }
9367        SmoothBasisSpec::ByVariable { .. } => {
9368            crate::bail_invalid_basis!(
9369                "internal: ByVariable smooths must return before inner basis dispatch"
9370            );
9371        }
9372        SmoothBasisSpec::BySmooth { .. } => {
9373            crate::bail_invalid_basis!("internal: BySmooth smooths must be lowered to ByVariable before inner basis dispatch"
9374                    .to_string(),);
9375        }
9376        SmoothBasisSpec::FactorSmooth { spec } => {
9377            if term.shape != ShapeConstraint::None {
9378                crate::bail_invalid_basis!(
9379                    "ShapeConstraint::{:?} is unsupported for factor smooth term '{}'",
9380                    term.shape,
9381                    term.name
9382                );
9383            }
9384            return build_factor_smooth(data, spec, &term.name, workspace);
9385        }
9386    };
9387
9388    // The Matérn design ALWAYS uses the operator-collocation {mass, tension,
9389    // stiffness} penalty triplet, overriding whatever penalty
9390    // `build_matern_basis_seeded` produced for the `double_penalty` flag.
9391    //
9392    // #1074 investigated swapping this for the genuine RKHS kernel penalty
9393    // `β' K_CC β` (mgcv `bs="gp"` / fields kriging) on the theory that the
9394    // operator triplet under-smooths the rougher half-integer kernels. MSI
9395    // truth-recovery measurement REFUTED that: the kernel penalty did NOT
9396    // improve ν=3/2 recovery (`matern(x,nu=1.5)` RMSE-vs-truth stayed 0.0554)
9397    // and it REGRESSED the high-frequency-init guard — `matern(x,nu≥5/2)` on
9398    // sin(2π·8·x) collapsed (span 0.53, RMSE 0.70) because the single RKHS
9399    // norm over-smooths a high-frequency truth where the Sobolev-order operator
9400    // dials do not. The operator triplet is therefore retained as the Matérn
9401    // penalty, and the κ-optimizer re-key / ψ-derivative paths route through the
9402    // same triplet builder so the block count stays ψ-stable (#1270).
9403    if let SmoothBasisSpec::Matern { .. } = &term.basis {
9404        let filtered = matern_operator_penalty_triplet_from_metadata(&built.metadata)?;
9405        built.active_penalties = filtered.active;
9406        built.dropped_penalties = filtered.dropped;
9407    }
9408
9409    if built.affine_offset.is_some() && term.shape != ShapeConstraint::None {
9410        crate::bail_invalid_basis!(
9411            "non-zero endpoint anchors cannot be combined with ShapeConstraint::{:?} on term '{}': the coefficient cone constrains only the homogeneous spline and would not certify the final affine function",
9412            term.shape,
9413            term.name
9414        );
9415    }
9416    let p_local = built.design.ncols();
9417    let affine_offset = built.affine_offset;
9418    let mut metadata = built.metadata.clone();
9419    // Extract factored Kronecker representation before consuming fields.
9420    // Invalidate it if shape transforms will be applied (they break structure).
9421    let kron_factored = if term.shape == ShapeConstraint::None {
9422        built.kronecker_factored
9423    } else {
9424        None
9425    };
9426    let mut design_t = built.design;
9427    let mut penalties_t = built.active_penalties;
9428    let mut dropped_penalties_t = built.dropped_penalties;
9429    if matches!(
9430        spatial_identifiability_policy(term),
9431        Some(SpatialIdentifiability::OrthogonalToParametric)
9432    ) {
9433        metadata = freeze_raw_spatial_metadata(metadata, design_t.ncols());
9434    }
9435
9436    let use_box_reparam =
9437        term.shape != ShapeConstraint::None && shape_uses_box_reparameterization(&term.basis);
9438    if let Some((order, sign)) = shape_order_and_sign(term.shape)
9439        && use_box_reparam
9440    {
9441        // Order 1 (monotone): the plain first-difference cone θ_{i+1}−θ_i ≥ 0 is
9442        // the control-polygon monotonicity criterion, which is independent of
9443        // Greville-abscissa spacing (it only fixes the *sign* of consecutive
9444        // control-point gaps), so the integer-difference transform is exact.
9445        //
9446        // Order 2 (convex/concave): the plain second-difference cone is only
9447        // correct for evenly spaced Greville abscissae. gam's B-splines are
9448        // clamped (and may use quantile knots), so the abscissae are not
9449        // uniform and the geometrically-correct cone is the second *divided*
9450        // difference. Build the knot-span-scaled transform so γ_{≥2} ≥ 0
9451        // certifies convexity of the function, not of the raw coefficient
9452        // index. Periodic splines are rejected by the exact-support gate: their
9453        // cyclic coefficient chart cannot use this open divided-difference cone.
9454        let t = if order == 2 {
9455            let (knots, degree) = match &metadata {
9456                BasisMetadata::BSpline1D {
9457                    knots,
9458                    degree: Some(degree),
9459                    periodic,
9460                    ..
9461                } if periodic.is_none() => (knots, *degree),
9462                _ => {
9463                    crate::bail_invalid_basis!(
9464                        "shape-constrained convex/concave term '{}' requires realized open B-spline knot and degree metadata",
9465                        term.name
9466                    );
9467                }
9468            };
9469            let spans = bspline_first_derivative_control_spans(knots.view(), degree)?;
9470            if spans.len() + 1 != p_local {
9471                crate::bail_invalid_basis!(
9472                    "shape-constraint derivative-control span count {} does not match basis dim {} for term '{}'",
9473                    spans.len(),
9474                    p_local,
9475                    term.name
9476                );
9477            }
9478            convex_derivative_control_transform_matrix(&spans, sign)?
9479        } else {
9480            cumulative_sum_transform_matrix(p_local, order, sign)
9481        };
9482        // Coefficient-side transform: wrap the design in an operator that
9483        // applies T on the coefficient side, preserving sparsity/operator
9484        // structure of the inner design.
9485        let inner_dense = match design_t {
9486            DesignMatrix::Dense(d) => d,
9487            DesignMatrix::Sparse(sp) => gam_linalg::matrix::DenseDesignMatrix::from(
9488                sp.try_to_dense_arc("shape-constrained coefficient transform")
9489                    .map_err(BasisError::InvalidInput)?,
9490            ),
9491        };
9492        let coeff_op =
9493            gam_linalg::matrix::CoefficientTransformOperator::new(inner_dense, t.clone()).map_err(
9494                |e| BasisError::InvalidInput(format!("CoefficientTransformOperator: {e}")),
9495            )?;
9496        design_t = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(
9497            coeff_op,
9498        )));
9499        // `β = Tγ` is an invertible change of coefficient chart. Every
9500        // physical quadratic functional, including the function-space
9501        // null-component penalty, therefore transforms by the same congruence
9502        // `S_γ = Tᵀ S_β T`. Rebuilding `ZZᵀ` in the γ Euclidean metric
9503        // would change the represented functional under this harmless chart
9504        // change and violate SPEC 5.
9505        for penalty in &mut penalties_t {
9506            let tt_s = fast_atb(&t, &penalty.matrix);
9507            penalty.matrix = fast_ab(&tt_s, &t);
9508            penalty.op = None;
9509            penalty.info.kronecker_factors = None;
9510            // A declared structural null frame does NOT survive this chart.
9511            // `null(Tᵀ S T) = T⁻¹ null(S)`, and `T` is a cumulative-sum /
9512            // derivative-control transform — invertible but not orthogonal, so
9513            // the image of an orthonormal frame is not orthonormal and the
9514            // declaration's own contract cannot carry it. Withdraw it here
9515            // rather than ship a frame that no longer spans the null space:
9516            // consumers measure when nothing is declared, which is the honest
9517            // fallback, whereas a stale frame is a wrong theorem.
9518            penalty.info.structural_null_frame = None;
9519        }
9520    }
9521    let penalty_candidates = penalties_t
9522        .into_iter()
9523        .map(|penalty| -> Result<PenaltyCandidate, BasisError> {
9524            let ActivePenalty {
9525                matrix,
9526                op: op_in,
9527                info,
9528                ..
9529            } = penalty;
9530            let (matrix, c_new) = normalize_penalty_in_constrained_space(&matrix);
9531            let normalization_scale = info.normalization_scale * c_new;
9532            let op_scale = 1.0 / c_new;
9533            let kronecker_scale = 1.0 / c_new;
9534            // Frobenius rescale: wrap inner op in `ScaledPenaltyOp(1/c_new)`
9535            // so `op.as_dense() == matrix` post-normalization.
9536            let scaled_op = if op_scale > 0.0 && op_scale.is_finite() {
9537                op_in.map(|op| {
9538                    std::sync::Arc::new(crate::analytic_penalties::ScaledPenaltyOp::new(
9539                        op, op_scale,
9540                    ))
9541                        as std::sync::Arc<dyn crate::analytic_penalties::PenaltyOp>
9542                })
9543            } else {
9544                None
9545            };
9546            let kronecker_factors = info.kronecker_factors.map(|mut factors| {
9547                if let Some(first) = factors.first_mut() {
9548                    first.mapv_inplace(|v| v * kronecker_scale);
9549                }
9550                factors
9551            });
9552            // Re-attach the structural null frame the basis factory declared
9553            // (#2445, extended in #2761). This is the SECOND chokepoint where
9554            // `try_from_dense_psd` sees only a dense matrix — the first, in
9555            // `term_design`, already re-attaches — and a declaration dropped
9556            // here never reaches the double-penalty rebuild, which then decides
9557            // the ridge's existence by a rank test on the shipped matrix. For a
9558            // basis whose Primary loses numerical rank as an outer coordinate
9559            // moves (measure-jet's representer range: rank 47 of 48 at the auto
9560            // seed, 14 of 48 at 8x it), that turns the penalty TOPOLOGY into a
9561            // function of ψ and aborts the incremental realizer mid-search.
9562            // A positive Frobenius rescale does not move a null space, so the
9563            // frame transports verbatim through this chart.
9564            let structural_null_frame = info.structural_null_frame;
9565            let matrix = ConstructiveQuadratic::try_from_dense_psd(
9566                matrix,
9567                "shape-constrained transformed penalty",
9568            )?;
9569            let matrix = match structural_null_frame {
9570                Some(frame) => matrix.with_structural_null_frame(
9571                    frame,
9572                    "renormalized smooth penalty structural frame",
9573                )?,
9574                None => matrix,
9575            };
9576            Ok(PenaltyCandidate {
9577                matrix,
9578                source: info.source,
9579                normalization_scale,
9580                kronecker_factors,
9581                op: scaled_op,
9582            })
9583        })
9584        .collect::<Result<Vec<_>, _>>()?;
9585    let filtered = crate::basis::filter_penalty_candidates(penalty_candidates)?;
9586    dropped_penalties_t.extend(filtered.dropped);
9587    // Joint-null absorption rotation. Fresh fit specs compute Q from the final
9588    // per-smooth penalty set (after all in-smooth reparameterizations have
9589    // already been applied). Frozen specs already carry the complete realized
9590    // coefficient chart in their `FrozenTransform`; recomputing Q there would
9591    // rotate an already-frozen chart a second time and desynchronize value
9592    // rebuilds from derivative operators.
9593    //
9594    // Kronecker-factored smooths (tensor B-splines under `TensorBSplineIdentifiability::None`)
9595    // carry their joint penalty as `Σ_d S_d` with `S_d = I ⊗ … ⊗ S_d^{1D} ⊗ … ⊗ I`.
9596    // The joint null space is the tensor of marginal nulls and is handled directly
9597    // by the REML runtime's `kronecker_penalty_system` path (see
9598    // `runtime.rs:8334-8344`). Applying a dense (p × p) Q here would densify
9599    // `X_raw = mx ⊗ my` into `X_raw · Q`, destroying the Kronecker product
9600    // structure that the runtime relies on for fast log-det/derivative
9601    // assembly — and the rotation block at the wrapper site also unconditionally
9602    // wipes `kronecker_factored`, leaving the runtime to fall back to the
9603    // dense per-block log-det. Skip the rotation for Kronecker-factored terms
9604    // so the factored representation survives end-to-end.
9605    let joint_null_rotation = match term.joint_null_rotation.clone() {
9606        Some(persisted) => Some(persisted),
9607        None if smooth_has_frozen_identifiability(term) => None,
9608        None if kron_factored.is_some() => None,
9609        None => crate::basis::compute_joint_null_rotation(&filtered.active)?,
9610    };
9611
9612    Ok(LocalSmoothTermBuild {
9613        dim: p_local,
9614        design: design_t,
9615        affine_offset,
9616        active_penalties: filtered.active,
9617        joint_null_rotation,
9618        dropped_penalties: dropped_penalties_t,
9619        metadata,
9620        linear_constraints: None,
9621        box_reparam: use_box_reparam,
9622        kronecker_factored: kron_factored,
9623    })
9624}
9625
9626pub fn build_smooth_design(
9627    data: ArrayView2<'_, f64>,
9628    terms: &[SmoothTermSpec],
9629) -> Result<RawSmoothDesign, BasisError> {
9630    let mut ws = crate::basis::BasisWorkspace::new();
9631    build_smooth_design_withworkspace(data, terms, &mut ws)
9632}
9633
9634/// Like `build_smooth_design`, but honors the caller workspace policy while
9635/// building each planned smooth term with an independent per-term workspace.
9636///
9637/// Independent workspaces avoid shared mutable distance-cache state during the
9638/// parallel term build; the final design, penalties, and metadata are assembled
9639/// in the original smooth-term order.
9640pub fn build_smooth_design_withworkspace(
9641    data: ArrayView2<'_, f64>,
9642    terms: &[SmoothTermSpec],
9643    workspace: &mut crate::basis::BasisWorkspace,
9644) -> Result<RawSmoothDesign, BasisError> {
9645    validate_smooth_terms_finite_inputs(data, terms)?;
9646    build_smooth_design_withworkspace_unvalidated(data, terms, workspace)
9647}
9648
9649pub fn build_smooth_design_withworkspace_unvalidated(
9650    data: ArrayView2<'_, f64>,
9651    terms: &[SmoothTermSpec],
9652    workspace: &mut crate::basis::BasisWorkspace,
9653) -> Result<RawSmoothDesign, BasisError> {
9654    let mut planned_blocks = plan_joint_spatial_centers_for_term_blocks(data, &[terms.to_vec()])?;
9655    let planned_terms = planned_blocks.pop().ok_or_else(|| {
9656        BasisError::InvalidInput(
9657            "joint spatial center planner returned no smooth blocks".to_string(),
9658        )
9659    })?;
9660    let policy = workspace.policy().clone();
9661    let local_builds: Vec<LocalSmoothTermBuild> = {
9662        use rayon::iter::{IntoParallelIterator, ParallelIterator};
9663        planned_terms
9664            .into_par_iter()
9665            .map(|term| {
9666                let mut term_workspace = crate::basis::BasisWorkspace::with_policy(policy.clone());
9667                build_single_local_smooth_term(data, &term, &mut term_workspace)
9668            })
9669            .collect::<Result<Vec<_>, _>>()?
9670    };
9671
9672    let total_p: usize = local_builds.iter().map(|built| built.dim).sum();
9673
9674    let mut local_designs: Vec<DesignMatrix> = Vec::with_capacity(local_builds.len());
9675    let mut affine_offset = Array1::<f64>::zeros(data.nrows());
9676    let mut terms_out = Vec::<SmoothTerm>::with_capacity(terms.len());
9677    let mut penalties_global = Vec::<BlockwisePenalty>::new();
9678    let mut nullspace_dims_global = Vec::<usize>::new();
9679    let mut penaltyinfo_global = Vec::<PenaltyBlockInfo>::new();
9680    let mut dropped_penaltyinfo_global = Vec::<DroppedPenaltyBlockInfo>::new();
9681    let mut coefficient_lower_bounds = Array1::<f64>::from_elem(total_p, f64::NEG_INFINITY);
9682    let mut any_bounds = false;
9683    // Each linear-constraint row only touches the current term's column slice.
9684    // Track `(col_start, col_end, local_row_values)` and assemble the final
9685    // dense `Array2` in one pass, avoiding per-row `Array1::zeros(total_p)`
9686    // allocation plus a row-by-row copy at the end.
9687    let mut linear_constraintsrows: Vec<(usize, usize, Array1<f64>)> = Vec::new();
9688    let mut linear_constraints_b: Vec<f64> = Vec::new();
9689
9690    let mut col_start = 0usize;
9691    for (term, mut built) in terms.iter().zip(local_builds.into_iter()) {
9692        let p_local = built.dim;
9693        let col_end = col_start + p_local;
9694        let lb_local = if built.box_reparam {
9695            shape_lower_bounds_local(term.shape, p_local)
9696        } else {
9697            None
9698        };
9699
9700        // Stage-2 joint-null absorption rotation. Fired *before* the
9701        // penalty / design / global aggregation loops below so that every
9702        // subsequent reference to `built.active_penalties` and `built.design`
9703        // sees the post-rotation values.
9704        //
9705        // The math: when the smooth's joint penalty `Σ_k S_k` has a
9706        // non-trivial null space, eigh selects `Q = [U_range | U_null]`
9707        // with null columns at the tail. Setting `β_raw = Q · γ` and
9708        // applying:
9709        //     design        ← X · Q
9710        //     penalties[k]  ← Qᵀ · S_k · Q   (block-diag, zero null tail)
9711        // yields a model whose fitted γ is invariant to the rotation
9712        // (since likelihood depends only on `X · β_raw = X · Q · γ`), but
9713        // whose penalty is full-rank on the range columns. The large-scale
9714        // failing case (cert refusal in the joint-Newton inner solve)
9715        // resolves because `H_pen = H_loglik + S` becomes full rank on
9716        // the smooth's range columns.
9717        //
9718        // Rotation is suppressed when the smooth carries coordinate-wise
9719        // shape constraints (`lb_local` or `built.linear_constraints`):
9720        // those encode a cone in the original coordinate system and a
9721        // general orthogonal rotation breaks the cone geometry. Smooths
9722        // with shape constraints typically have full-rank joint penalty
9723        // (their structural shape comes from the cone, not from null
9724        // directions in the penalty), so suppression is rarely a loss.
9725        //
9726        // `applied_rotation` carries the Q that was applied (or `None`
9727        // if no rotation fired). It is persisted onto `SmoothTerm` below
9728        // so prediction-side `X_new_raw · Q` replay can reproduce the
9729        // exact rotation. Persistence through the saved-model artifact
9730        // is a follow-up — see the doc on `SmoothTerm.joint_null_rotation`.
9731        let applied_rotation: Option<crate::basis::JointNullRotation> = match (
9732            built.joint_null_rotation.take(),
9733            lb_local.is_some(),
9734            built.linear_constraints.is_some(),
9735        ) {
9736            (Some(rot), false, false) => {
9737                let q = &rot.rotation;
9738                built.design =
9739                    apply_smooth_transform_to_design(built.design.clone(), q, &term.name)?;
9740                for penalty in &mut built.active_penalties {
9741                    let qt_s = gam_linalg::faer_ndarray::fast_atb(q, &penalty.matrix);
9742                    penalty.matrix = gam_linalg::faer_ndarray::fast_ab(&qt_s, q);
9743                    penalty.null_eigenvectors = penalty
9744                        .null_eigenvectors
9745                        .as_ref()
9746                        .map(|basis| gam_linalg::faer_ndarray::fast_atb(q, basis));
9747                    // `Q` is orthogonal, so `null(Qᵀ S Q) = Qᵀ null(S)` and the
9748                    // image of an orthonormal frame is orthonormal — a declared
9749                    // structural null frame transports EXACTLY here (unlike the
9750                    // non-orthogonal shape-constraint chart in
9751                    // `build_single_local_smooth_term`, which has to withdraw
9752                    // it). Rotating it alongside `null_eigenvectors`
9753                    // is what keeps the declaration a statement about the same
9754                    // subspace after the rotation instead of a stale frame in
9755                    // the pre-rotation coordinates (#2761).
9756                    penalty.info.structural_null_frame = penalty
9757                        .info
9758                        .structural_null_frame
9759                        .as_ref()
9760                        .map(|frame| gam_linalg::faer_ndarray::fast_atb(q, frame));
9761                    penalty.op = None;
9762                    penalty.info.kronecker_factors = None;
9763                }
9764                built.kronecker_factored = None;
9765                Some(rot)
9766            }
9767            (Some(_), _, _) => None,
9768            (None, _, _) => None,
9769        };
9770
9771        for active_penalty in &built.active_penalties {
9772            let global_index = penalties_global.len();
9773            penalties_global.push(
9774                BlockwisePenalty::new(col_start..col_end, active_penalty.matrix.clone())
9775                    .with_op(active_penalty.op.clone()),
9776            );
9777            nullspace_dims_global.push(active_penalty.nullity);
9778            penaltyinfo_global.push(PenaltyBlockInfo {
9779                global_index,
9780                termname: Some(term.name.clone()),
9781                penalty: active_penalty.info.clone(),
9782            });
9783        }
9784        for info in &built.dropped_penalties {
9785            dropped_penaltyinfo_global.push(DroppedPenaltyBlockInfo {
9786                termname: Some(term.name.clone()),
9787                penalty: info.clone(),
9788            });
9789        }
9790
9791        if let Some(lin_local) = &built.linear_constraints {
9792            for r in 0..lin_local.a.nrows() {
9793                linear_constraintsrows.push((col_start, col_end, lin_local.a.row(r).to_owned()));
9794                linear_constraints_b.push(lin_local.b[r]);
9795            }
9796        }
9797        if let Some(lb_local) = &lb_local {
9798            coefficient_lower_bounds
9799                .slice_mut(s![col_start..col_end])
9800                .assign(lb_local);
9801            any_bounds = true;
9802        }
9803
9804        if let Some(term_offset) = built.affine_offset.as_ref() {
9805            if term_offset.len() != data.nrows() {
9806                crate::bail_dim_basis!(
9807                    "smooth term '{}' affine offset has {} rows but the realized data has {}",
9808                    term.name,
9809                    term_offset.len(),
9810                    data.nrows()
9811                );
9812            }
9813            affine_offset += term_offset;
9814        }
9815
9816        // Move the per-term design out of `built` rather than cloning it.
9817        local_designs.push(built.design);
9818
9819        terms_out.push(SmoothTerm {
9820            parametric_residualization: None,
9821            name: term.name.clone(),
9822            coeff_range: col_start..col_end,
9823            shape: term.shape,
9824            active_penalties: built.active_penalties,
9825            dropped_penalties: built.dropped_penalties,
9826            metadata: built.metadata,
9827            lower_bounds_local: lb_local,
9828            linear_constraints_local: built.linear_constraints,
9829            kronecker_factored: built.kronecker_factored.take(),
9830            joint_null_rotation: applied_rotation,
9831            unabsorbed_global_orthogonality: None,
9832            // The RAW build precedes the global step, so it decides no gauge;
9833            // `apply_global_smooth_identifiability` fills this in (#2747).
9834            collection_gauge: None,
9835        });
9836
9837        col_start = col_end;
9838    }
9839
9840    assert_eq!(
9841        penalties_global.len(),
9842        nullspace_dims_global.len(),
9843        "global smooth penalty/nullspace bookkeeping diverged"
9844    );
9845    assert_eq!(
9846        penalties_global.len(),
9847        penaltyinfo_global.len(),
9848        "global smooth penalty metadata bookkeeping diverged"
9849    );
9850
9851    Ok(RawSmoothDesign {
9852        term_designs: local_designs,
9853        affine_offset,
9854        penalties: penalties_global,
9855        nullspace_dims: nullspace_dims_global,
9856        penaltyinfo: penaltyinfo_global,
9857        dropped_penaltyinfo: dropped_penaltyinfo_global,
9858        terms: terms_out,
9859        coefficient_lower_bounds: if any_bounds {
9860            Some(coefficient_lower_bounds)
9861        } else {
9862            None
9863        },
9864        linear_constraints: if linear_constraintsrows.is_empty() {
9865            None
9866        } else {
9867            let mut a = Array2::<f64>::zeros((linear_constraintsrows.len(), total_p));
9868            for (i, (cs, ce, values)) in linear_constraintsrows.iter().enumerate() {
9869                a.row_mut(i).slice_mut(s![*cs..*ce]).assign(values);
9870            }
9871            Some(LinearInequalityConstraints {
9872                a,
9873                b: Array1::from_vec(linear_constraints_b),
9874            })
9875        },
9876    })
9877}
9878
9879#[cfg(test)]
9880mod factor_smooth_heldout_group_tests {
9881    use super::*;
9882    use crate::basis::BasisWorkspace;
9883    use ndarray::{Array1, array};
9884
9885    fn pinned_marginal() -> BSplineBasisSpec {
9886        BSplineBasisSpec {
9887            degree: 3,
9888            penalty_order: 2,
9889            knotspec: BSplineKnotSpec::Provided(Array1::from(vec![
9890                0.0, 0.0, 0.0, 0.0, 0.25, 0.6, 1.0, 1.0, 1.0, 1.0,
9891            ])),
9892            double_penalty: false,
9893            identifiability: BSplineIdentifiability::None,
9894            boundary: crate::basis::OneDimensionalBoundary::Open,
9895            boundary_conditions: crate::basis::BSplineBoundaryConditions::default(),
9896        }
9897    }
9898
9899    fn factor_smooth_term(
9900        flavour: FactorSmoothFlavour,
9901        frozen: Option<Vec<u64>>,
9902    ) -> SmoothTermSpec {
9903        SmoothTermSpec {
9904            frozen_parametric_residualization: None,
9905            name: "fs_heldout".to_string(),
9906            basis: SmoothBasisSpec::FactorSmooth {
9907                spec: FactorSmoothSpec {
9908                    continuous_cols: vec![0],
9909                    group_col: 1,
9910                    marginal: pinned_marginal(),
9911                    flavour,
9912                    group_frozen_levels: frozen,
9913                    frozen_global_orthogonality: None,
9914                },
9915            },
9916            shape: ShapeConstraint::None,
9917            joint_null_rotation: None,
9918        }
9919    }
9920
9921    const FROZEN_01: [f64; 2] = [0.0, 1.0];
9922
9923    fn frozen_bits() -> Vec<u64> {
9924        FROZEN_01.iter().map(|v| v.to_bits()).collect()
9925    }
9926
9927    /// #2365: in the frozen (predict/replay) context, a `bs="re"` row whose
9928    /// group is outside the training vocabulary must build with an all-zero
9929    /// row — zero fitted deviation, population prediction — instead of
9930    /// erroring before the random-effect operator can apply its held-out-group
9931    /// contract.
9932    #[test]
9933    fn re_heldout_group_row_is_zero_deviation() {
9934        let data = array![[0.1, 0.0], [0.5, 1.0], [0.9, 7.0]];
9935        let term = factor_smooth_term(FactorSmoothFlavour::Re, Some(frozen_bits()));
9936        let mut workspace = BasisWorkspace::default();
9937        let build = build_single_local_smooth_term(data.view(), &term, &mut workspace)
9938            .expect("a held-out group must not fail the bs=\"re\" design build");
9939        let dense = build
9940            .design
9941            .try_to_dense_by_chunks("heldout test")
9942            .expect("dense");
9943        assert!(
9944            dense.row(2).iter().all(|&v| v == 0.0),
9945            "unseen-group row must carry zero deviation across every group block, got {:?}",
9946            dense.row(2)
9947        );
9948        assert!(
9949            dense.row(0).iter().any(|&v| v != 0.0) && dense.row(1).iter().any(|&v| v != 0.0),
9950            "in-vocabulary rows must still populate their group blocks"
9951        );
9952    }
9953
9954    /// The `fs` flavour estimates a per-level deviation FUNCTION — an unseen
9955    /// level has no zero-deviation population fallback — so the frozen-context
9956    /// build must stay strict (#2102/#2137 must not regress through #2365).
9957    #[test]
9958    fn fs_heldout_group_stays_strict() {
9959        let data = array![[0.1, 0.0], [0.5, 1.0], [0.9, 7.0]];
9960        let term = factor_smooth_term(
9961            FactorSmoothFlavour::Fs {
9962                m_null_penalty_orders: vec![1],
9963            },
9964            Some(frozen_bits()),
9965        );
9966        let mut workspace = BasisWorkspace::default();
9967        let err = match build_single_local_smooth_term(data.view(), &term, &mut workspace) {
9968            Ok(_) => panic!("fs must reject an unseen grouping level"),
9969            Err(err) => err,
9970        };
9971        assert!(
9972            err.to_string().contains("unseen grouping level"),
9973            "fs unseen-level refusal must name the defect, got: {err}"
9974        );
9975    }
9976}
9977
9978#[cfg(test)]
9979mod linear_term_contract_tests {
9980    use super::LinearTermSpec;
9981
9982    #[test]
9983    fn missing_linear_double_penalty_deserializes_to_unpenalized_mle() {
9984        let term: LinearTermSpec = serde_json::from_str(r#"{"name":"x","feature_col":0}"#)
9985            .expect("minimal saved linear term");
9986        assert!(
9987            !term.double_penalty,
9988            "descriptor and formula defaults must both preserve parametric MLE semantics"
9989        );
9990    }
9991}