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::UserProvided(centers) => format!("{} centers", centers.nrows()),
19        CenterStrategy::EqualMass { num_centers }
20        | CenterStrategy::EqualMassCovarRepresentative { num_centers }
21        | CenterStrategy::FarthestPoint { num_centers }
22        | CenterStrategy::KMeans { num_centers, .. } => format!("{num_centers} centers"),
23        CenterStrategy::UniformGrid { points_per_dim } => {
24            format!("uniform grid with {points_per_dim} points per dimension")
25        }
26    }
27}
28
29pub fn rewrite_thin_plate_knots_error(
30    err: BasisError,
31    termname: &str,
32    feature_count: usize,
33    spec: &ThinPlateBasisSpec,
34) -> BasisError {
35    match err {
36        // Polynomial-nullspace shortfall reported directly by the kernel
37        // builder ("thin-plate spline requires at least N centers to span ...").
38        BasisError::InvalidInput(msg)
39            if msg.contains("thin-plate spline requires at least")
40                && (msg.contains("centers to span") || msg.contains("knots to span")) =>
41        {
42            let min_centers = crate::basis::thin_plate_polynomial_basis_dimension(feature_count);
43            let requested = describe_thin_plate_center_request(&spec.center_strategy);
44            BasisError::InvalidInput(format!(
45                "joint TPS term '{termname}' over {feature_count} covariates with {requested} is invalid; minimum centers is {min_centers}"
46            ))
47        }
48        // Insufficient-rows shortfall raised by `select_thin_plate_knots` when
49        // the requested center count exceeds the available row count. Rewrite
50        // it in term language so the diagnostic points at the smooth term and
51        // the polynomial-nullspace minimum the user needs to satisfy.
52        BasisError::InvalidInput(msg)
53            if msg.starts_with("requested ") && msg.contains(" knots but only ") =>
54        {
55            let min_centers = crate::basis::thin_plate_polynomial_basis_dimension(feature_count);
56            let requested = describe_thin_plate_center_request(&spec.center_strategy);
57            BasisError::InvalidInput(format!(
58                "joint TPS term '{termname}' over {feature_count} covariates with {requested} is invalid; minimum centers is {min_centers}"
59            ))
60        }
61        other => other,
62    }
63}
64
65#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
66pub enum ShapeConstraint {
67    None,
68    MonotoneIncreasing,
69    MonotoneDecreasing,
70    Convex,
71    Concave,
72}
73
74/// Parse a shape-constraint string into a [`ShapeConstraint`].
75///
76/// This is the single source of truth shared by the formula DSL
77/// (`s(x, shape=...)`) and the `smooths={...}` override path
78/// (`Smooth.shape_constraint`). The accepted spellings cover the canonical
79/// Python `ShapeConstraintLiteral` strings exactly
80/// (`"none"` / `"monotone_increasing"` / `"monotone_decreasing"` /
81/// `"convex"` / `"concave"`) plus a few common aliases. Hyphens and case are
82/// normalized, so `"Monotone-Increasing"` and `"mono_inc"` both resolve to
83/// [`ShapeConstraint::MonotoneIncreasing`].
84pub fn parse_shape_constraint(raw: &str) -> Result<ShapeConstraint, String> {
85    let normalized = raw.trim().to_ascii_lowercase().replace('-', "_");
86    match normalized.as_str() {
87        "" | "none" => Ok(ShapeConstraint::None),
88        "monotone_increasing" | "monotonic_increasing" | "increasing" | "mono_inc" | "mpi" => {
89            Ok(ShapeConstraint::MonotoneIncreasing)
90        }
91        "monotone_decreasing" | "monotonic_decreasing" | "decreasing" | "mono_dec" | "mpd" => {
92            Ok(ShapeConstraint::MonotoneDecreasing)
93        }
94        "convex" | "cvx" => Ok(ShapeConstraint::Convex),
95        "concave" | "ccv" => Ok(ShapeConstraint::Concave),
96        other => Err(format!(
97            "unknown shape constraint {other:?}; expected one of \
98             \"none\", \"monotone_increasing\", \"monotone_decreasing\", \
99             \"convex\", \"concave\""
100        )),
101    }
102}
103
104impl ShapeConstraint {
105    /// Canonical formula-DSL spelling, i.e. the text emitted into
106    /// `s(x, shape=...)`. Round-trips through [`parse_shape_constraint`].
107    pub fn dsl_str(&self) -> &'static str {
108        match self {
109            ShapeConstraint::None => "none",
110            ShapeConstraint::MonotoneIncreasing => "monotone_increasing",
111            ShapeConstraint::MonotoneDecreasing => "monotone_decreasing",
112            ShapeConstraint::Convex => "convex",
113            ShapeConstraint::Concave => "concave",
114        }
115    }
116}
117
118/// Smooth-term head keywords recognised by the formula DSL. A `shape=` option
119/// may be attached to any term whose head is one of these.
120pub const SMOOTH_HEAD_KEYWORDS: [&str; 11] = [
121    "s",
122    "smooth",
123    "te",
124    "tensor",
125    "thinplate",
126    "tps",
127    "duchon",
128    "matern",
129    "sphere",
130    "bs",
131    "bspline",
132];
133
134/// Rewrite smooth-term calls in `formula` so each named smooth carries a
135/// `shape=<kind>` option understood by the formula DSL.
136///
137/// `constraints` pairs the smooth-term text as it appears in the formula
138/// (e.g. `"s(x)"` or `"s(x, type=duchon, centers=8)"`) with a shape-constraint
139/// spelling accepted by [`parse_shape_constraint`]; comparison is exact after
140/// whitespace removal. A `"none"` constraint is a no-op. Referencing a term not
141/// present in the formula is an error.
142///
143/// This is the single source of truth for the `gamfit.fit(..., constraints=…)`
144/// rewrite — the Python wrapper only marshals the mapping across the FFI and
145/// holds no formula-parsing or alias-normalization logic of its own.
146pub fn apply_shape_constraints_to_formula(
147    formula: &str,
148    constraints: &[(String, String)],
149) -> Result<String, String> {
150    use std::collections::{BTreeMap, BTreeSet};
151
152    if constraints.is_empty() {
153        return Ok(formula.to_string());
154    }
155    let strip_ws = |s: &str| -> String { s.chars().filter(|c| !c.is_whitespace()).collect() };
156
157    // Whitespace-stripped term text -> canonical shape spelling.
158    let mut wanted: BTreeMap<String, &'static str> = BTreeMap::new();
159    // Whitespace-stripped term text -> original key (for error labels).
160    let mut originals: BTreeMap<String, String> = BTreeMap::new();
161    for (key, kind_raw) in constraints {
162        let kind = parse_shape_constraint(kind_raw)?;
163        let nk = strip_ws(key);
164        originals.entry(nk.clone()).or_insert_with(|| key.clone());
165        if kind != ShapeConstraint::None {
166            wanted.insert(nk, kind.dsl_str());
167        }
168    }
169    if wanted.is_empty() {
170        return Ok(formula.to_string());
171    }
172
173    let chars: Vec<char> = formula.chars().collect();
174    let n = chars.len();
175    let is_ident = |c: char| c.is_ascii_alphanumeric() || c == '_';
176
177    let mut out = String::with_capacity(formula.len() + 32);
178    let mut matched: BTreeSet<String> = BTreeSet::new();
179    let mut i = 0usize;
180    while i < n {
181        // Locate the next smooth-term head (`<keyword> \s* (`) at or after `i`,
182        // respecting word boundaries so `abs(` never matches the `s(` head.
183        let mut head: Option<(usize, usize)> = None; // (head_start, paren_index)
184        let mut p = i;
185        while p < n {
186            let boundary = p == 0 || !is_ident(chars[p - 1]);
187            if boundary {
188                for kw in SMOOTH_HEAD_KEYWORDS.iter() {
189                    let klen = kw.chars().count();
190                    if p + klen > n || chars[p..p + klen].iter().collect::<String>() != **kw {
191                        continue;
192                    }
193                    let mut q = p + klen;
194                    while q < n && chars[q].is_whitespace() {
195                        q += 1;
196                    }
197                    if q < n && chars[q] == '(' {
198                        head = Some((p, q));
199                        break;
200                    }
201                }
202            }
203            if head.is_some() {
204                break;
205            }
206            p += 1;
207        }
208        let (head_start, paren_open) = match head {
209            Some(h) => h,
210            None => {
211                out.extend(chars[i..].iter());
212                break;
213            }
214        };
215        out.extend(chars[i..head_start].iter());
216
217        // Find the matching close paren, honoring nesting and string literals.
218        let body_start = paren_open + 1;
219        let mut depth = 1i32;
220        let mut j = body_start;
221        let mut in_str: Option<char> = None;
222        let mut closed = false;
223        while j < n {
224            let ch = chars[j];
225            if let Some(quote) = in_str {
226                if ch == quote {
227                    in_str = None;
228                }
229            } else if ch == '\'' || ch == '"' {
230                in_str = Some(ch);
231            } else if ch == '(' {
232                depth += 1;
233            } else if ch == ')' {
234                depth -= 1;
235                if depth == 0 {
236                    closed = true;
237                    break;
238                }
239            }
240            j += 1;
241        }
242
243        if !closed {
244            // Unbalanced — emit the remainder verbatim; the DSL parser will
245            // produce the canonical error.
246            out.extend(chars[head_start..].iter());
247            break;
248        }
249
250        let term_text: String = chars[head_start..=j].iter().collect();
251
252        let key_norm = strip_ws(&term_text);
253
254        match wanted.get(&key_norm) {
255            None => out.extend(chars[head_start..=j].iter()),
256            Some(kind) => {
257                let head_paren: String = chars[head_start..body_start].iter().collect();
258                let inside: String = chars[body_start..j].iter().collect();
259                let inside = inside.trim();
260                if inside.is_empty() {
261                    out.push_str(&format!("{head_paren}shape={kind})"));
262                } else {
263                    out.push_str(&format!("{head_paren}{inside}, shape={kind})"));
264                }
265                matched.insert(key_norm);
266            }
267        }
268
269        i = j + 1;
270    }
271
272    let mut missing: Vec<String> = wanted
273        .keys()
274        .filter(|k| !matched.contains(*k))
275        .map(|k| originals.get(k).cloned().unwrap_or_else(|| k.clone()))
276        .collect();
277
278    if !missing.is_empty() {
279        missing.sort();
280        return Err(format!(
281            "shape constraints referenced smooth term(s) not found in formula: {}",
282            missing.join(", ")
283        ));
284    }
285
286    Ok(out)
287}
288
289#[derive(Debug, Clone, Serialize, Deserialize)]
290pub enum BySmoothKind {
291    Numeric,
292    Level { level_bits: u64 },
293}
294
295#[derive(Debug, Clone, Serialize, Deserialize)]
296#[serde(deny_unknown_fields)]
297pub enum SmoothBasisSpec {
298    /// Row-gated wrapper used for mgcv-style ``by=`` smooths.
299    ///
300    /// ``ByNumeric`` multiplies the inner smooth by a numeric column.
301    /// ``ByLevel`` keeps the inner smooth active only for rows whose encoded
302    /// categorical value has the stored bit pattern.  Unordered factor-by
303    /// smooths are represented as one independent ``ByLevel`` term per level.
304    ///
305    /// `kind` preserves the compact structural discriminator, while `by`
306    /// carries the full row-gating spec used to build the local design.
307    ByVariable {
308        inner: Box<SmoothBasisSpec>,
309        by_col: usize,
310        kind: BySmoothKind,
311        by: ByVariableSpec,
312    },
313    /// Sum-to-zero factor smooth (`bs="sz"`): with L levels, estimate L-1
314    /// deviation coefficient blocks and use the final level as the negative
315    /// sum of the others, enforcing coefficient-wise zero sums across levels.
316    FactorSumToZero {
317        inner: Box<SmoothBasisSpec>,
318        by_col: usize,
319        levels: Vec<u64>,
320        /// Global-orthogonality column map `Z` captured at fit time when this
321        /// term overlapped an owner smooth (`s(x) + s(g, x, bs=sz)`, #978):
322        /// the hierarchical-ownership pass residualized this term's realized
323        /// design as `X ← X·Z`, shrinking its coefficient block. `Z` depends
324        /// on the *training-row* owner designs, so prediction cannot rederive
325        /// it — it must be persisted and replayed
326        /// (`apply_global_smooth_identifiability` consumes it verbatim).
327        /// Chart convention: `Z` lives in the post-restack, post-joint-null-Q
328        /// coordinates — the raw `sz` rebuild reapplies `Q` deterministically
329        /// (#700), then `Z` applies on top. `None` for non-overlapping terms.
330        #[serde(default)]
331        frozen_global_orthogonality: Option<Array2<f64>>,
332    },
333    BSpline1D {
334        feature_col: usize,
335        spec: BSplineBasisSpec,
336    },
337    /// A smooth modulated by a `by=` variable. Numeric `by` scales one inner
338    /// smooth; factor `by` replicates the inner smooth by level.
339    BySmooth {
340        smooth: Box<SmoothBasisSpec>,
341        by_kind: ByVarKind,
342    },
343    /// Factor-smooth interaction families (`bs="fs"`, `bs="sz"`) and
344    /// random slopes (`bs="re"`).
345    FactorSmooth { spec: FactorSmoothSpec },
346    ThinPlate {
347        feature_cols: Vec<usize>,
348        spec: ThinPlateBasisSpec,
349        /// Uniform coordinate scale estimated on a fresh build and persisted
350        /// for exact frozen replay.
351        input_scale: Option<crate::IsotropicScale>,
352    },
353    Sphere {
354        feature_cols: Vec<usize>,
355        spec: SphericalSplineBasisSpec,
356    },
357    /// Constant-curvature (`M_κ`) geodesic-kernel smooth over κ-stereographic
358    /// chart coordinates (#944): one construction interpolating
359    /// S^d → ℝ^d → H^d through the spec's fixed κ. The Wahba S² smooth is the
360    /// structural template; the geometry comes from
361    /// `geometry::constant_curvature::ConstantCurvature`.
362    ConstantCurvature {
363        feature_cols: Vec<usize>,
364        spec: ConstantCurvatureBasisSpec,
365    },
366    Matern {
367        feature_cols: Vec<usize>,
368        spec: MaternBasisSpec,
369        input_scale: Option<crate::IsotropicScale>,
370    },
371    /// Measure-jet spline smooth: multiscale local-jet-residual energy of the
372    /// empirical measure (centers as μ-quadrature, masses as μ-weights — no
373    /// graph, mesh, or neighbor set inside the statistical object). The
374    /// feature columns are ambient coordinates of data concentrated near an
375    /// unknown low-dimensional, possibly stratified set.
376    MeasureJet {
377        feature_cols: Vec<usize>,
378        spec: MeasureJetBasisSpec,
379        input_scale: Option<crate::IsotropicScale>,
380    },
381    Duchon {
382        feature_cols: Vec<usize>,
383        spec: DuchonBasisSpec,
384        input_scale: Option<crate::IsotropicScale>,
385    },
386    Pca {
387        feature_cols: Vec<usize>,
388        basis_matrix: Array2<f64>,
389        centered: bool,
390        #[serde(default = "default_pca_smooth_penalty")]
391        smooth_penalty: f64,
392        #[serde(default)]
393        center_mean: Option<Array1<f64>>,
394        #[serde(default)]
395        pca_basis_path: Option<PathBuf>,
396        #[serde(default = "default_pca_chunk_size")]
397        chunk_size: usize,
398    },
399    /// Tensor-product smooth built from 1D B-spline marginals.
400    ///
401    /// This is the `te()`-style construction used when axes have different units/scales
402    /// (for example, space x time) and isotropic radial kernels are not appropriate.
403    TensorBSpline {
404        feature_cols: Vec<usize>,
405        spec: TensorBSplineSpec,
406    },
407}
408
409impl SmoothBasisSpec {
410    /// Conservative lower bound on the number of sample rows needed for this
411    /// smooth basis to have a well-posed REML fit.
412    ///
413    /// Each basis kind answers the question for itself, so the workflow does
414    /// not have to know how many columns a B-spline, tensor product, PCA
415    /// projection, or spatial kernel emits. The contract is a *lower bound*:
416    /// returning too small a number is permitted (the inner solver will catch
417    /// any genuine n-vs-rank failure that slips past); returning too large a
418    /// number is a regression because it rejects legitimate fits.
419    ///
420    /// Rationale: B-spline / tensor / PCA bases have a closed-form column
421    /// count, so we use the exact dimension. Radial bases (TPS, Matern,
422    /// Duchon, Sphere) and factor smooths choose their column count from the
423    /// data (`heuristic_centers`, `unique_count`); we fall back to a small
424    /// constant floor because a fit on fewer than five rows cannot stabilise
425    /// any radial smooth regardless of the configured kernel scale.
426    pub fn min_sample_rows(&self) -> usize {
427        // Floor used for data-driven bases whose column count is not known
428        // from the spec alone. Five rows is the minimum at which the inner
429        // pivot/QR + REML smoothing-parameter search has any chance of being
430        // well-posed for a non-parametric smooth.
431        const RADIAL_FLOOR: usize = 5;
432
433        match self {
434            Self::ByVariable { inner, .. } => inner.min_sample_rows(),
435            Self::FactorSumToZero { inner, levels, .. } => {
436                // L-1 independent deviation blocks each carrying the inner
437                // basis dimension. Skip the levels-multiplier if it doesn't
438                // bring more rows; we want the *lower bound* not the rank.
439                let inner_min = inner.min_sample_rows();
440                let lvls = levels.len().saturating_sub(1).max(1);
441                inner_min.saturating_mul(lvls)
442            }
443            Self::BSpline1D { spec, .. } => bspline_basis_min_rows(spec),
444            Self::BySmooth { smooth, .. } => smooth.min_sample_rows(),
445            Self::FactorSmooth { spec } => {
446                // Replicates the marginal once per level; without a known
447                // level count we conservatively require at least the marginal
448                // basis dimension.
449                bspline_basis_min_rows(&spec.marginal)
450            }
451            Self::ThinPlate { .. }
452            | Self::Sphere { .. }
453            | Self::ConstantCurvature { .. }
454            | Self::Matern { .. }
455            | Self::MeasureJet { .. }
456            | Self::Duchon { .. } => RADIAL_FLOOR,
457            Self::Pca { basis_matrix, .. } => basis_matrix.ncols().max(1),
458            Self::TensorBSpline { spec, .. } => {
459                // A `te(...)` smooth is *penalized*: each margin carries a
460                // difference (wiggliness) penalty and the tensor inherits a
461                // Kronecker-sum penalty `S = Σ_i I ⊗ … ⊗ S_i ⊗ … ⊗ I`. The raw
462                // column count is the *product* of the per-marginal column
463                // counts, but that product is the lower bound for an
464                // *unpenalized* tensor regression — it is the number of rows you
465                // would need to identify every interaction column with no
466                // regularization. The penalty regularizes all of those
467                // interaction directions; only the combined penalty *null space*
468                // (the tensor product of the per-margin polynomial trends, a
469                // handful of columns) must be identified by the data, and the
470                // smoothing-parameter search shrinks the rest. The effective
471                // degrees of freedom of the fitted `te()` are therefore a small
472                // fraction of the column product, which is exactly why mgcv
473                // fits a default `te(x, y)` on a couple hundred rows.
474                //
475                // The honest *penalized* lower bound is the **sum** of the
476                // per-marginal column counts, not their product: a row floor of
477                // `Σ_i k_i` still guarantees enough data to identify each
478                // margin's additive main-effect (the largest sub-block the
479                // penalty cannot shrink to zero), while no longer conflating
480                // unpenalized column-count identifiability with penalized
481                // well-posedness. This accepts moderate-`n` penalized tensors
482                // (e.g. a 20×20 default basis on n=200) yet still rejects a
483                // genuinely undersized fit where `n < Σ_i k_i` and even the
484                // additive part is rank-deficient.
485                //
486                // Binary / low-cardinality margins (#724): gam will accept a
487                // `te(x, badh)` whose `badh ∈ {0, 1}` margin nominally requests
488                // more basis columns than `badh` has unique values, where mgcv
489                // refuses the unpenalized term as ill-posed ("badh has
490                // insufficient unique values to support k knots"). This is
491                // correct-by-design, *not* a degenerate fit: the marginal
492                // wiggliness penalty on the `badh` axis has a null space that is
493                // exactly its identifiable trend (the two cell means of a binary
494                // covariate), and the Kronecker-sum penalty shrinks every tensor
495                // column outside that null space toward zero. The resulting fit
496                // is the well-posed "per-level `x` smooth + binary main effect"
497                // that mgcv reaches only after manually collapsing the basis —
498                // gam reaches it automatically because the penalty, not the raw
499                // column count, sets the effective rank. A genuinely
500                // rank-deficient design (penalty null space wider than the data
501                // can support) is still caught downstream by the inner pivoted
502                // factorization, which owns the exact n-vs-rank decision; this
503                // pre-fit gate only refuses the grossly-undersized formula.
504                let mut total: usize = 0;
505                for marginal in &spec.marginalspecs {
506                    let m = bspline_basis_min_rows(marginal);
507                    total = total.saturating_add(m.max(1));
508                }
509                total.max(RADIAL_FLOOR)
510            }
511        }
512    }
513
514    /// Stable structural discriminant for warm-start cache keying (#869).
515    ///
516    /// Two smooths that produce different bases / penalty structures must map
517    /// to different strings here so they cannot collide on the persistent
518    /// warm-start `cache_key` (which is otherwise blind to topology: it hashes
519    /// only the raw input column count, so e.g. `sphere` vs `torus` vs
520    /// `euclidean` candidates fit on the *same* data would otherwise share one
521    /// key and cross-contaminate each other's β/ρ seed). The string is the
522    /// topology identity, not the fitted coefficients, so same-topology refits
523    /// (the screen→full-refit cascade) still hit the same key and reuse work.
524    pub fn structural_kind(&self) -> &'static str {
525        match self {
526            Self::ByVariable { .. } => "by_variable",
527            Self::FactorSumToZero { .. } => "factor_sum_to_zero",
528            Self::BSpline1D { .. } => "bspline_1d",
529            Self::BySmooth { .. } => "by_smooth",
530            Self::FactorSmooth { .. } => "factor_smooth",
531            Self::ThinPlate { .. } => "thin_plate",
532            Self::Sphere { .. } => "sphere",
533            Self::ConstantCurvature { .. } => "constant_curvature",
534            Self::Matern { .. } => "matern",
535            Self::MeasureJet { .. } => "measurejet",
536            Self::Duchon { .. } => "duchon",
537            Self::Pca { .. } => "pca",
538            Self::TensorBSpline { .. } => "tensor_bspline",
539        }
540    }
541
542    /// True for a tensor-product smooth that is only *marginally* centered
543    /// (`ti(...)`, [`TensorBSplineIdentifiability::MarginalSumToZero`]): its
544    /// per-margin sum-to-zero reparameterization `(B_xZ_x)⊗(B_zZ_z)` has ALREADY
545    /// removed each axis's main effect analytically (mgcv-identical), so its
546    /// main-effect removal is complete and it must take NO additional
547    /// owner-residualization block. Residualizing it a second time against the
548    /// realized main-effect designs is a grid-fragile no-op on an exact tensor
549    /// grid but eats genuine pure-interaction curvature off-grid (#1470).
550    pub fn is_marginally_centered_tensor(&self) -> bool {
551        matches!(
552            self,
553            Self::TensorBSpline { spec, .. }
554                if matches!(spec.identifiability, TensorBSplineIdentifiability::MarginalSumToZero)
555        )
556    }
557
558    /// A sum-to-zero factor smooth (`bs="sz"`) has ALREADY removed the
559    /// cross-group main effect analytically, in coefficient space, via its
560    /// `Σ_g d_g(x) ≡ 0` reparameterization (`L-1` deviation blocks with the
561    /// reference level the negative sum of the others) — exactly mgcv's `sz`
562    /// construction, which is self-identifiable against an overlapping `s(x)`
563    /// with no further constraint. Residualizing it a SECOND time against the
564    /// realized B-spline span of the explicit `s(x)` smooth is redundant in
565    /// exact arithmetic (the common-to-all-groups component is zero by
566    /// construction) and actively HARMFUL on finite data: each deviation block
567    /// is a B-spline in `x` whose realized columns share `s(x)`'s span, so the
568    /// joint residualization collapses the full `L·k`-column deviation design to
569    /// `L·k − rank(s(x))` columns and eats the within-group curvature `s(x)`
570    /// cannot represent. REML then rails the deviation smoothing parameter and
571    /// the factor smooth under-recovers (#1605). This is the exact analogue of
572    /// the marginally-centered tensor (`ti`) exemption (#1470), so such a term
573    /// takes NO owner-residualization block.
574    pub fn is_sum_to_zero_factor_smooth(&self) -> bool {
575        matches!(
576            self,
577            Self::FactorSumToZero { .. }
578                | Self::FactorSmooth {
579                    spec: FactorSmoothSpec {
580                        flavour: FactorSmoothFlavour::Sz,
581                        ..
582                    }
583                }
584        )
585    }
586
587    /// Feature columns this basis consumes, used alongside [`structural_kind`]
588    /// to disambiguate two same-kind smooths on different axes. Wrapper
589    /// variants delegate to their inner basis.
590    pub fn structural_feature_cols(&self) -> Vec<usize> {
591        match self {
592            Self::ByVariable { inner, .. } | Self::FactorSumToZero { inner, .. } => {
593                inner.structural_feature_cols()
594            }
595            Self::BySmooth { smooth, .. } => smooth.structural_feature_cols(),
596            Self::FactorSmooth { .. } => Vec::new(),
597            Self::BSpline1D { feature_col, .. } => vec![*feature_col],
598            Self::ThinPlate { feature_cols, .. }
599            | Self::Sphere { feature_cols, .. }
600            | Self::ConstantCurvature { feature_cols, .. }
601            | Self::Matern { feature_cols, .. }
602            | Self::MeasureJet { feature_cols, .. }
603            | Self::Duchon { feature_cols, .. }
604            | Self::Pca { feature_cols, .. }
605            | Self::TensorBSpline { feature_cols, .. } => feature_cols.clone(),
606        }
607    }
608}
609
610/// Lower bound on the number of sample rows a 1D B-spline smooth needs for a
611/// well-posed *penalized* REML fit. Used as the per-smooth row floor in
612/// [`SmoothBasisSpec::min_sample_rows`].
613///
614/// For a *singly*-penalized smooth the floor is the full column count: the
615/// wiggliness penalty leaves the order-`m` polynomial trend unpenalized, and
616/// gam's original gate conservatively required enough rows for the whole basis.
617/// That conservative floor is kept here unchanged.
618///
619/// A *double*-penalized smooth (mgcv `select=TRUE`) is different: it adds a
620/// second penalty on the wiggliness penalty's null space, so even the
621/// polynomial trend is shrinkable toward zero and *nothing* in the basis
622/// requires unpenalized identification by the data — exactly the reasoning the
623/// `TensorBSpline` arm of [`SmoothBasisSpec::min_sample_rows`] already applies
624/// to a penalized tensor. Its honest floor is therefore a small stabilization
625/// constant, not the column count. This is what lets mgcv (and now gam) fit
626/// several `select=TRUE` smooths on a dataset whose row count is below the
627/// summed basis width (e.g. the n≈30 `wine_gamair` fold, 5 `ps` smooths,
628/// p≈51): the penalties, not the data, set the effective rank. The bounded
629/// outer REML loop still terminates, and the genuine n-vs-rank decision is
630/// owned downstream by the inner pivoted factorization. Without this, gam
631/// rejected the fit outright (or, before the gate existed, the outer REML loop
632/// wandered the flat overparameterized surface until the benchmark wall budget
633/// killed it — #1089).
634pub fn bspline_basis_min_rows(spec: &crate::basis::BSplineBasisSpec) -> usize {
635    use crate::basis::BSplineKnotSpec;
636    let columns = match &spec.knotspec {
637        BSplineKnotSpec::Generate {
638            num_internal_knots, ..
639        } => *num_internal_knots + spec.degree + 1,
640        BSplineKnotSpec::Automatic {
641            num_internal_knots: Some(k),
642            ..
643        } => *k + spec.degree + 1,
644        BSplineKnotSpec::Automatic {
645            num_internal_knots: None,
646            ..
647        } => {
648            // Knot count is data-derived (`default_internal_knot_count_for_data`).
649            // A minimal cubic basis is `degree + 2` columns; below that the
650            // basis cannot represent a non-parametric smooth.
651            spec.degree + 2
652        }
653        BSplineKnotSpec::Provided(knots) => knots.len().saturating_sub(spec.degree + 1).max(1),
654        // cr basis dimension equals the knot count (no degree offset).
655        BSplineKnotSpec::NaturalCubicRegression { knots } => knots.len(),
656        BSplineKnotSpec::PeriodicUniform { num_basis, .. } => *num_basis,
657    };
658    let columns = columns.max(spec.degree + 2);
659
660    if spec.double_penalty {
661        // Fully shrinkable basis: only a small stabilization floor must be
662        // identified by the data, capped by the actual column count.
663        const DOUBLE_PENALTY_FLOOR: usize = 2;
664        DOUBLE_PENALTY_FLOOR.min(columns).max(1)
665    } else {
666        columns
667    }
668}
669
670#[derive(Debug, Clone, Serialize, Deserialize)]
671pub enum ByVariableSpec {
672    Numeric,
673    Level { value_bits: u64, label: String },
674}
675
676#[derive(Debug, Clone, Serialize, Deserialize)]
677pub enum ByVarKind {
678    Numeric {
679        feature_col: usize,
680    },
681    Factor {
682        feature_col: usize,
683        ordered: bool,
684        frozen_levels: Option<Vec<u64>>,
685    },
686}
687
688#[derive(Debug, Clone, Serialize, Deserialize)]
689pub struct FactorSmoothSpec {
690    pub continuous_cols: Vec<usize>,
691    pub group_col: usize,
692    pub marginal: BSplineBasisSpec,
693    pub flavour: FactorSmoothFlavour,
694    pub group_frozen_levels: Option<Vec<u64>>,
695    /// Fit-time global-orthogonality chart `Z` for this term (`s(x) + fs(x, g)`
696    /// overlap residualization, #978), in the post-joint-null-`Q` coordinates
697    /// (the raw rebuild recomputes any `Q` itself; `fs` penalties are
698    /// typically full-rank so `Q` is absent). Training-row dependent, hence
699    /// persisted; replayed verbatim by `apply_global_smooth_identifiability`.
700    #[serde(default)]
701    pub frozen_global_orthogonality: Option<Array2<f64>>,
702}
703
704#[derive(Debug, Clone, Serialize, Deserialize)]
705pub enum FactorSmoothFlavour {
706    Fs { m_null_penalty_orders: Vec<usize> },
707    Sz,
708    Re,
709}
710
711#[derive(Debug, Clone, Serialize, Deserialize)]
712pub struct TensorBSplineSpec {
713    pub marginalspecs: Vec<BSplineBasisSpec>,
714    #[serde(default)]
715    pub periods: Vec<Option<f64>>,
716    #[serde(default = "default_tensor_double_penalty")]
717    pub double_penalty: bool,
718    #[serde(default)]
719    pub identifiability: TensorBSplineIdentifiability,
720    #[serde(default)]
721    pub penalty_decomposition: TensorBSplinePenaltyDecomposition,
722}
723
724pub const fn default_tensor_double_penalty() -> bool {
725    true
726}
727
728impl Default for TensorBSplineSpec {
729    fn default() -> Self {
730        Self {
731            marginalspecs: Vec::new(),
732            periods: Vec::new(),
733            double_penalty: default_tensor_double_penalty(),
734            identifiability: TensorBSplineIdentifiability::default(),
735            penalty_decomposition: TensorBSplinePenaltyDecomposition::default(),
736        }
737    }
738}
739
740#[derive(Debug, Default, Clone, Serialize, Deserialize)]
741pub enum TensorBSplineIdentifiability {
742    None,
743    #[default]
744    SumToZero,
745    /// mgcv `ti(...)` semantics: a *tensor interaction* smooth that excludes the
746    /// marginal main effects. A sum-to-zero constraint is applied to **each
747    /// marginal basis independently** before forming the tensor product, so the
748    /// resulting column space contains no function of a single variable alone —
749    /// only the pure interaction survives. The realized identifiability
750    /// transform is the Kronecker product `Z = Z₀ ⊗ Z₁ ⊗ … ⊗ Z_{d-1}` of the
751    /// per-margin sum-to-zero null-space bases, which is exactly the
752    /// reparameterization that turns the full-tensor design into the tensor
753    /// product of the centered margins.
754    MarginalSumToZero,
755    FrozenTransform {
756        transform: Array2<f64>,
757    },
758}
759
760#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
761pub enum TensorBSplinePenaltyDecomposition {
762    /// mgcv `te(...)`: one overlapping Kronecker-product penalty per margin,
763    /// `S_j` embedded against identities in the other tensor factors.
764    #[default]
765    MarginalKroneckerSum,
766    /// mgcv `t2(...)`: split every marginal coefficient space into penalized
767    /// range and penalty-null subspaces, then emit one disjoint tensor-subspace
768    /// penalty for every non-empty penalized/null combination.
769    Separable,
770}
771
772#[derive(Debug, Clone, Serialize, Deserialize)]
773pub struct SmoothTermSpec {
774    pub name: String,
775    pub basis: SmoothBasisSpec,
776    pub shape: ShapeConstraint,
777    /// Joint-null absorption rotation captured at fit time. `Some(Q)` means
778    /// the fitted coefficient vector lives in `γ`-coordinates with
779    /// `β_raw = Q · γ`; prediction must rotate the raw-basis design via
780    /// `X_new = X_new_raw · Q` to match. `None` means either the smooth had
781    /// no joint null space (penalty already full-rank) or rotation was
782    /// suppressed (smooth carries shape constraints whose cone geometry
783    /// would not survive an arbitrary orthogonal rotation). Persisted so
784    /// `save → load → predict` is bit-equivalent to in-memory prediction.
785    #[serde(default)]
786    pub joint_null_rotation: Option<crate::basis::JointNullRotation>,
787}
788
789#[derive(Debug, Clone)]
790pub struct SmoothTerm {
791    pub name: String,
792    pub coeff_range: Range<usize>,
793    pub shape: ShapeConstraint,
794    /// Active local penalty identities. Numerical and semantic channels are
795    /// inseparable, including after a preceding candidate is dropped.
796    pub active_penalties: Vec<ActivePenalty>,
797    pub dropped_penalties: Vec<DroppedPenaltyInfo>,
798    pub metadata: BasisMetadata,
799    /// Optional term-local lower bounds for constrained coefficients.
800    /// `-inf` means unconstrained.
801    pub lower_bounds_local: Option<Array1<f64>>,
802    /// Optional term-local inequality constraints in local coefficient coordinates.
803    /// `A_local * beta_local >= b_local`.
804    pub linear_constraints_local: Option<LinearInequalityConstraints>,
805    /// Optional factored tensor-product representation preserved for operator-backed
806    /// assembly in the main design builder.
807    pub kronecker_factored: Option<KroneckerFactoredBasis>,
808    /// Joint-null absorption rotation. `Some(Q)` records the orthonormal
809    /// `(p_local × p_local)` matrix that was applied to this term's design
810    /// and per-block penalties at construction time:
811    /// `term_design ← X_raw · Q`, `active_penalties[k].matrix ← Qᵀ · S_raw · Q`.
812    /// The smooth's coefficient vector therefore lives in the rotated
813    /// (`γ`) coordinate system, with `β_raw = Q · γ` recovering the raw
814    /// pre-rotation parameterization. `None` means either no joint null
815    /// space (penalty already full-rank) or rotation was suppressed —
816    /// suppression fires when the smooth carries shape constraints
817    /// (lower bounds or local linear inequalities) that would lose their
818    /// cone geometry under a general orthogonal rotation.
819    ///
820    /// Prediction-side replay: callers building a new-data design `X_new_raw`
821    /// from the *raw* basis must call [`SmoothTerm::apply_rotation_to_predict`]
822    /// (or equivalent) to obtain `X_new = X_new_raw · Q` matching this
823    /// term's coefficient system.
824    ///
825    /// Persistence replay: `freeze_term_collection_from_design` copies this
826    /// rotation into `SmoothTermSpec`, which is serialized with fitted-model
827    /// payloads and reused by the predict-time basis builder. Saved models
828    /// therefore replay the same `X_new_raw · Q` transform as in-memory
829    /// prediction.
830    pub joint_null_rotation: Option<crate::basis::JointNullRotation>,
831    /// Global-orthogonality transform that `apply_global_smooth_identifiability`
832    /// applied to this term's design but could NOT embed into `metadata`
833    /// (factor-smooth kinds: `sz` metadata is per-marginal, `fs` metadata has
834    /// no transform slot — #978). `freeze_term_collection_from_design` copies
835    /// it onto the term's basis spec (`frozen_global_orthogonality`) so the
836    /// predict-side rebuild replays it instead of emitting the unresidualized
837    /// (wider) design that the fitted coefficients no longer match.
838    /// Chart convention is per kind: post-`Q` `Z` for `sz` (the raw rebuild
839    /// reapplies `Q` itself, #700), full `Q·Z` chart for `fs`.
840    pub unabsorbed_global_orthogonality: Option<Array2<f64>>,
841}
842
843impl SmoothTerm {
844    /// Apply the joint-null absorption rotation to a raw new-data design
845    /// matrix, returning `X_new_raw · Q` when this term was rotated at
846    /// fit time, or `X_new_raw` unchanged when no rotation was applied.
847    ///
848    /// Callers in the prediction path: after building the smooth's basis
849    /// at new data via the *raw* basis builder (the same builder used at
850    /// fit time, applied to `x_new` instead of the training rows), call
851    /// this method on the resulting matrix before forming `X · β`. The
852    /// fitted `β` lives in `γ`-coordinates if Q was applied; multiplying
853    /// the un-rotated `X_new_raw` by `β` would give a wrong η.
854    ///
855    /// Returns an error if the raw design's column count does not match
856    /// the rotation's `p_local`. The width invariant must hold: the raw
857    /// basis builder MUST emit the same `p_local` columns that the
858    /// fit-time builder did, and the rotation is `(p_local × p_local)`.
859    pub fn apply_rotation_to_predict(
860        &self,
861        x_new_raw: Array2<f64>,
862    ) -> Result<Array2<f64>, BasisError> {
863        let Some(rot) = self.joint_null_rotation.as_ref() else {
864            return Ok(x_new_raw);
865        };
866        let p_local = rot.rotation.nrows();
867        if x_new_raw.ncols() != p_local {
868            crate::bail_dim_basis!(
869                "joint-null rotation replay for term '{}': raw design has {} columns, \
870                 rotation expects {} (the raw basis builder must emit the same column \
871                 count as at fit time)",
872                self.name,
873                x_new_raw.ncols(),
874                p_local,
875            );
876        }
877        Ok(gam_linalg::faer_ndarray::fast_ab(&x_new_raw, &rot.rotation))
878    }
879
880    /// Dimension of the **joint** null space of this term's active penalties:
881    /// the coefficient directions penalized by *no* penalty. The smooth-component
882    /// Wald test ([`crate::inference::smooth_test::wood_smooth_test`]) treats this
883    /// many leading coefficients as genuine unpenalized fixed effects and tests
884    /// them at full rank; the remainder is the penalized sub-block tested with a
885    /// rank-`≈EDF` truncated pseudo-inverse.
886    ///
887    /// Because every penalty block `S_k` is positive semi-definite,
888    /// `vᵀ(Σ_k S_k)v = Σ_k vᵀ S_k v = 0` iff `S_k v = 0` for *every* `k`; the
889    /// joint null space is therefore exactly `null(Σ_k S_k)`, of dimension
890    /// `p_local − rank(Σ_k S_k)`. This is the **intersection** of the per-penalty
891    /// null spaces, not their sum.
892    ///
893    /// Summing the per-penalty `nullspace_dims` instead (the historical defect
894    /// behind #1360) *unions* the null spaces and badly over-counts: a
895    /// double-penalty smooth carries a bending penalty (null space = its
896    /// polynomial part) plus a complementary null-space ridge (which penalizes
897    /// exactly that polynomial part), so the two null spaces are disjoint and the
898    /// joint null space is empty — yet the per-penalty dims sum to nearly
899    /// `p_local`. Feeding that inflated count to the Wald test makes it test
900    /// almost the whole shrunk block at full rank, manufacturing overwhelming
901    /// "significance" for a term the fit drove to ~0 EDF.
902    pub fn wald_unpenalized_dim(&self) -> usize {
903        joint_unpenalized_dim(self.coeff_range.len(), &self.active_penalties)
904    }
905}
906
907/// Numeric core of [`SmoothTerm::wald_unpenalized_dim`]: the dimension of the
908/// joint null space `∩_k null(S_k) = null(Σ_k S_k)` of a term's local penalty
909/// blocks, with a conservative fallback when a penalty is not materialized as a
910/// full `p_local × p_local` matrix (e.g. a Kronecker tensor factor).
911pub fn joint_unpenalized_dim(p_local: usize, active_penalties: &[ActivePenalty]) -> usize {
912    use gam_linalg::faer_ndarray::FaerEigh;
913    if p_local == 0 {
914        return 0;
915    }
916    if active_penalties.is_empty() {
917        // No penalty ⇒ a wholly unpenalized (fixed-effect) block.
918        return p_local;
919    }
920    // Sum the penalties that are materialized as full `p_local × p_local`
921    // blocks (the common smooth case). The covariance block the Wald test
922    // slices lives in this same coefficient basis (post joint-null rotation),
923    // so the rank is computed in the right metric.
924    let mut s_total = Array2::<f64>::zeros((p_local, p_local));
925    let mut materialized = 0usize;
926    for penalty in active_penalties {
927        let s = &penalty.matrix;
928        if s.nrows() == p_local && s.ncols() == p_local {
929            s_total += s;
930            materialized += 1;
931        }
932    }
933    if materialized == active_penalties.len() {
934        let symmetric = {
935            let transpose = s_total.t().to_owned();
936            (&s_total + &transpose) * 0.5
937        };
938        if let Ok((evals, _)) = symmetric.eigh(faer::Side::Lower) {
939            let max_abs = evals.iter().fold(0.0_f64, |acc, &v| acc.max(v.abs()));
940            if max_abs == 0.0 {
941                // All penalties identically zero ⇒ unpenalized block.
942                return p_local;
943            }
944            let tol = max_abs * (p_local as f64) * 1e-12;
945            let rank = evals.iter().filter(|&&v| v > tol).count();
946            return p_local.saturating_sub(rank);
947        }
948    }
949    // Conservative fallback when a penalty is not a materialized full block
950    // (e.g. a Kronecker tensor factor): with ≥2 active penalties the joint
951    // null space is almost always empty (the only over-rejecting direction);
952    // with a single penalty it is exactly that penalty's own null space.
953    if active_penalties.len() >= 2 {
954        0
955    } else {
956        active_penalties
957            .iter()
958            .map(|penalty| penalty.nullity)
959            .min()
960            .unwrap_or(0)
961            .min(p_local)
962    }
963}
964
965#[derive(Debug, Clone, Serialize, Deserialize)]
966pub struct PenaltyBlockInfo {
967    pub global_index: usize,
968    pub termname: Option<String>,
969    pub penalty: ActivePenaltyInfo,
970}
971
972#[derive(Debug, Clone, Serialize, Deserialize)]
973pub struct DroppedPenaltyBlockInfo {
974    pub termname: Option<String>,
975    pub penalty: DroppedPenaltyInfo,
976}
977
978#[derive(Debug, Clone)]
979pub struct SmoothDesign {
980    pub term_designs: Vec<DesignMatrix>,
981    /// Per-term block-local penalties.  Each `col_range` is relative to the
982    /// smooth block (i.e. indexing into the concatenation of `term_designs`).
983    pub penalties: Vec<BlockwisePenalty>,
984    pub nullspace_dims: Vec<usize>,
985    pub penaltyinfo: Vec<PenaltyBlockInfo>,
986    pub dropped_penaltyinfo: Vec<DroppedPenaltyBlockInfo>,
987    pub terms: Vec<SmoothTerm>,
988    /// Optional smooth-block lower bounds in smooth coefficient coordinates.
989    /// Length equals `total_smooth_cols()` when present.
990    pub coefficient_lower_bounds: Option<Array1<f64>>,
991    /// Optional smooth-block inequality constraints:
992    /// `A_smooth * beta_smooth >= b`.
993    pub linear_constraints: Option<LinearInequalityConstraints>,
994}
995
996impl SmoothDesign {
997    pub fn total_smooth_cols(&self) -> usize {
998        self.term_designs.iter().map(DesignMatrix::ncols).sum()
999    }
1000    pub fn nrows(&self) -> usize {
1001        self.term_designs.first().map_or(0, DesignMatrix::nrows)
1002    }
1003}
1004
1005#[derive(Debug, Clone)]
1006pub struct RawSmoothDesign {
1007    pub term_designs: Vec<DesignMatrix>,
1008    /// Sum of every fixed affine term contribution on the realized rows.
1009    pub affine_offset: Array1<f64>,
1010    /// Per-term block-local penalties.  Each `col_range` is relative to the
1011    /// smooth block (i.e. indexing into the concatenation of `term_designs`).
1012    pub penalties: Vec<BlockwisePenalty>,
1013    pub nullspace_dims: Vec<usize>,
1014    pub penaltyinfo: Vec<PenaltyBlockInfo>,
1015    pub dropped_penaltyinfo: Vec<DroppedPenaltyBlockInfo>,
1016    pub terms: Vec<SmoothTerm>,
1017    pub coefficient_lower_bounds: Option<Array1<f64>>,
1018    pub linear_constraints: Option<LinearInequalityConstraints>,
1019}
1020
1021impl RawSmoothDesign {
1022    pub fn total_smooth_cols(&self) -> usize {
1023        self.term_designs.iter().map(DesignMatrix::ncols).sum()
1024    }
1025    pub fn nrows(&self) -> usize {
1026        self.term_designs.first().map_or(0, DesignMatrix::nrows)
1027    }
1028}
1029
1030#[derive(Debug, Default, Clone, Serialize, Deserialize)]
1031pub enum BoundedCoefficientPriorSpec {
1032    #[default]
1033    None,
1034    Uniform,
1035    Beta {
1036        a: f64,
1037        b: f64,
1038    },
1039}
1040
1041#[derive(Debug, Clone, Serialize, Deserialize, Default)]
1042pub enum LinearCoefficientGeometry {
1043    #[default]
1044    Unconstrained,
1045    Bounded {
1046        min: f64,
1047        max: f64,
1048        #[serde(default)]
1049        prior: BoundedCoefficientPriorSpec,
1050    },
1051}
1052
1053#[derive(Debug, Clone, Serialize, Deserialize)]
1054pub struct LinearTermSpec {
1055    pub name: String,
1056    /// Primary feature column index. For Wilkinson-Rogers `:` interaction
1057    /// terms (`a:b[:c...]`) this is the first column in `feature_cols`; the
1058    /// realized design column is the elementwise product across every entry
1059    /// of `feature_cols`. Plain (non-interaction) linear terms set
1060    /// `feature_cols == vec![feature_col]`.
1061    pub feature_col: usize,
1062    /// Full list of columns whose elementwise product yields this term's
1063    /// design column. `len() >= 1`; `len() == 1` is a plain linear effect.
1064    #[serde(default)]
1065    pub feature_cols: Vec<usize>,
1066    /// Categorical-level gates for a factor-aware `:` interaction.
1067    ///
1068    /// Each `(col, level_bits)` multiplies the realized design column by the
1069    /// indicator `1[canonical_level_bits(data[row, col]) == level_bits]` (the
1070    /// canonical key collapses `-0.0`/`+0.0` and NaN payloads so numerically
1071    /// equal codes name one level; see `gam_data::canonical_level_bits`). This
1072    /// is how a
1073    /// `factor:x` (or `factor:factor`) interaction is expanded: `build_termspec`
1074    /// emits one `LinearTermSpec` per surviving cell of the categorical
1075    /// operand(s) (treatment-coded, first level dropped per factor), each
1076    /// carrying the numeric operands in `feature_cols` and the cell's level
1077    /// gate(s) here. Empty for a plain numeric `:` interaction or main effect,
1078    /// in which case the realized column is exactly the numeric product.
1079    #[serde(default)]
1080    pub categorical_levels: Vec<(usize, u64)>,
1081    /// Zero-centered shrinkage ridge with a REML-selected `λ`. It is enabled
1082    /// by default so an unsupported non-intercept effect can be recovered as
1083    /// zero; `linear(x, double_penalty=false)` requests an explicit
1084    /// unpenalized/MLE effect.
1085    #[serde(default = "default_linear_term_double_penalty")]
1086    pub double_penalty: bool,
1087    #[serde(default)]
1088    pub coefficient_geometry: LinearCoefficientGeometry,
1089    #[serde(default)]
1090    pub coefficient_min: Option<f64>,
1091    #[serde(default)]
1092    pub coefficient_max: Option<f64>,
1093    /// The term's empirical function mass (`n⁻¹bᵀb` on the TRAINING design
1094    /// column), captured once at fit time by
1095    /// [`crate::smooth::freeze_term_collection_from_design`] and baked into the
1096    /// frozen/saved spec. `None` while a spec is still being resolved during
1097    /// fitting (before freezing) or for a `double_penalty=false` term, which
1098    /// carries no ridge and never needs a mass.
1099    ///
1100    /// A rebuild of this term's design column at PREDICT time (a held-out
1101    /// grid, a small group/class-anchor set, a single test row, …) must reuse
1102    /// this persisted value rather than recomputing it from the evaluation
1103    /// rows: a covariate that varies fine across the training set can easily
1104    /// be constant across a tiny evaluation subset by chance, and recomputing
1105    /// there would misfire the fit-time "identically zero" identifiability
1106    /// guard on a perfectly good term (#1561 REF_ERROR/METRIC_OFF triage).
1107    #[serde(default)]
1108    pub frozen_function_mass: Option<f64>,
1109}
1110
1111impl LinearTermSpec {
1112    /// Return the effective list of feature columns. Backfills from
1113    /// `feature_col` for legacy specs that predate the multi-column field.
1114    pub fn effective_feature_cols(&self) -> Vec<usize> {
1115        if self.feature_cols.is_empty() {
1116            vec![self.feature_col]
1117        } else {
1118            self.feature_cols.clone()
1119        }
1120    }
1121
1122    /// True when this term is a Wilkinson-Rogers `:` interaction (multi-col).
1123    pub fn is_interaction(&self) -> bool {
1124        self.feature_cols.len() > 1 || !self.categorical_levels.is_empty()
1125    }
1126
1127    /// Realize this linear term's `(n,)` design column from `data`.
1128    ///
1129    /// The column is the elementwise product of every numeric feature column
1130    /// (`effective_feature_cols`) gated by the categorical-level indicators in
1131    /// `categorical_levels`: each `(col, level_bits)` multiplies the running
1132    /// column by `1[canonical_level_bits(data[row, col]) == level_bits]` (signed
1133    /// zero / NaN canonicalized so numerically equal codes match). A plain numeric
1134    /// term (no `categorical_levels`) reduces to the bare product, matching the
1135    /// historical behaviour. A pure categorical interaction (empty
1136    /// `feature_cols`, non-empty `categorical_levels`) reduces to the cell
1137    /// indicator. Bounds are validated here; the returned column has length
1138    /// `data.nrows()`.
1139    pub fn realized_design_column(&self, data: ArrayView2<'_, f64>) -> Result<Array1<f64>, String> {
1140        let n = data.nrows();
1141        let p = data.ncols();
1142        let bounds = |col: usize| -> Result<(), String> {
1143            if col >= p {
1144                Err(format!(
1145                    "linear term '{}' feature column {} out of bounds for {} columns",
1146                    self.name, col, p
1147                ))
1148            } else {
1149                Ok(())
1150            }
1151        };
1152
1153        // Numeric operands. When `categorical_levels` is set we treat
1154        // `feature_cols` as the (possibly empty) numeric operand list and start
1155        // from a column of ones; otherwise we preserve the legacy backfill from
1156        // `feature_col` so a plain term with no `feature_cols` still resolves.
1157        let mut column = if self.categorical_levels.is_empty() {
1158            let cols = self.effective_feature_cols();
1159            for &c in &cols {
1160                bounds(c)?;
1161            }
1162            let mut acc = data.column(cols[0]).to_owned();
1163            for &c in cols.iter().skip(1) {
1164                acc *= &data.column(c);
1165            }
1166            acc
1167        } else {
1168            let mut acc = Array1::<f64>::ones(n);
1169            for &c in &self.feature_cols {
1170                bounds(c)?;
1171                acc *= &data.column(c);
1172            }
1173            acc
1174        };
1175
1176        for &(col, level_bits) in &self.categorical_levels {
1177            bounds(col)?;
1178            // Canonicalize the stored key once (loop-invariant) so the gate is
1179            // robust to level sets interned before signed-zero canonicalization
1180            // landed, not just to canonical data rows (#2146).
1181            let level_bits = gam_data::canonical_level_bits(f64::from_bits(level_bits));
1182            let gate = data.column(col);
1183            for (out, &v) in column.iter_mut().zip(gate.iter()) {
1184                if gam_data::canonical_level_bits(v) != level_bits {
1185                    *out = 0.0;
1186                }
1187            }
1188        }
1189
1190        Ok(column)
1191    }
1192}
1193
1194pub const fn default_linear_term_double_penalty() -> bool {
1195    true
1196}
1197
1198pub const fn default_pca_smooth_penalty() -> f64 {
1199    1.0
1200}
1201
1202pub const fn default_pca_chunk_size() -> usize {
1203    4096
1204}
1205
1206/// Random-effects term specification.
1207///
1208/// The selected feature column is interpreted as a categorical grouping variable.
1209/// The term contributes a one-hot dummy block with an identity penalty on group
1210/// coefficients, equivalent to i.i.d. Gaussian random effects.
1211#[derive(Debug, Clone, Serialize, Deserialize)]
1212pub struct RandomEffectTermSpec {
1213    pub name: String,
1214    pub feature_col: usize,
1215    /// If true, drop the lexicographically first group level to use treatment coding.
1216    /// If false, keep all levels (full one-hot block, still identifiable under ridge).
1217    pub drop_first_level: bool,
1218    /// If true, add a ridge penalty and estimate this block as a random effect.
1219    /// If false, leave the one-hot/treatment-coded block unpenalized so it is a
1220    /// fixed categorical main effect.  The default preserves older saved models.
1221    #[serde(default = "default_random_effect_penalized")]
1222    pub penalized: bool,
1223    /// Optional fixed kept-level set (sorted by f64 bit pattern) captured at fit time.
1224    /// When present, prediction uses exactly these columns to avoid design drift.
1225    #[serde(default)]
1226    pub frozen_levels: Option<Vec<u64>>,
1227    /// Whether an *unseen* level of this grouping column is tolerated at predict
1228    /// time (encoded as an out-of-vocabulary code and shrunk toward the
1229    /// population mean) instead of raising a schema mismatch.
1230    ///
1231    /// Only a genuine random effect — `group(g)`/`re(g)`/`s(g, bs="re")` — is
1232    /// lenient: the held-out-group policy is a deliberate contract. A FIXED
1233    /// categorical factor — a bare `+ g` OR an explicit `factor(g)` — although
1234    /// materialized as a penalized one-hot block, must raise on an
1235    /// out-of-vocabulary level at predict rather than being silently mapped to
1236    /// the factor's centering point (#2102/#2137). `factor(g)` originally shared
1237    /// the `group()`/`re()` parse arm and so wrongly inherited the lenient policy
1238    /// (#2137). For a string factor the typed schema encode rejects the unseen
1239    /// level upstream; for a numeric-coded `factor(year)` the reject is enforced
1240    /// by `build_random_effect_block`, which owns the frozen vocabulary. The
1241    /// `true` default preserves the pre-#2102 (uniformly lenient) behavior for
1242    /// models serialized before this field existed.
1243    #[serde(default = "default_random_effect_lenient_unseen")]
1244    pub lenient_unseen: bool,
1245}
1246
1247pub fn default_random_effect_penalized() -> bool {
1248    true
1249}
1250
1251pub fn default_random_effect_lenient_unseen() -> bool {
1252    true
1253}
1254
1255pub fn validate_measure_jet_positive_vec_len(
1256    label: &str,
1257    term_name: &str,
1258    field: &str,
1259    values: &[f64],
1260    expected: usize,
1261) -> Result<(), String> {
1262    if values.len() != expected {
1263        return Err(SmoothError::invalid_config(format!(
1264            "{label} term '{term_name}' frozen MeasureJet {field} has length {}, expected {expected}",
1265            values.len()
1266        ))
1267        .into());
1268    }
1269    if values
1270        .iter()
1271        .any(|value| !(value.is_finite() && *value > 0.0))
1272    {
1273        return Err(SmoothError::invalid_config(format!(
1274            "{label} term '{term_name}' frozen MeasureJet {field} values must be positive and finite"
1275        ))
1276        .into());
1277    }
1278    Ok(())
1279}
1280
1281#[derive(Debug, Clone, Serialize, Deserialize)]
1282pub struct TermCollectionSpec {
1283    pub linear_terms: Vec<LinearTermSpec>,
1284    pub random_effect_terms: Vec<RandomEffectTermSpec>,
1285    pub smooth_terms: Vec<SmoothTermSpec>,
1286}
1287
1288pub fn validate_smooth_basis_frozen(
1289    basis: &SmoothBasisSpec,
1290    label: &str,
1291    term_name: &str,
1292) -> Result<(), String> {
1293    if let Err(error) = basis.validate_scale_configuration() {
1294        return Err(SmoothError::invalid_config(format!(
1295            "{label} term '{term_name}' has an invalid scale contract: {error}"
1296        ))
1297        .into());
1298    }
1299    match basis {
1300        SmoothBasisSpec::ByVariable { inner, .. }
1301        | SmoothBasisSpec::FactorSumToZero { inner, .. } => {
1302            validate_smooth_basis_frozen(inner, label, term_name)
1303        }
1304        SmoothBasisSpec::BSpline1D { spec, .. } => {
1305            if !matches!(
1306                spec.knotspec,
1307                BSplineKnotSpec::Provided(_)
1308                    | BSplineKnotSpec::PeriodicUniform { .. }
1309                    | BSplineKnotSpec::NaturalCubicRegression { .. }
1310            ) {
1311                return Err(format!(
1312                    "{label} term '{term_name}' is not frozen: BSpline knotspec must be Provided, PeriodicUniform, or NaturalCubicRegression"
1313                ));
1314            }
1315            Ok(())
1316        }
1317        SmoothBasisSpec::ThinPlate { spec, .. } => {
1318            if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
1319                return Err(format!(
1320                    "{label} term '{term_name}' is not frozen: ThinPlate centers must be UserProvided"
1321                ));
1322            }
1323            if matches!(
1324                spec.identifiability,
1325                SpatialIdentifiability::OrthogonalToParametric
1326            ) {
1327                return Err(format!(
1328                    "{label} term '{term_name}' is not frozen: ThinPlate identifiability must be FrozenTransform or None"
1329                ));
1330            }
1331            Ok(())
1332        }
1333        _ => Ok(()),
1334    }
1335}
1336
1337impl TermCollectionSpec {
1338    /// Write this collection's topology identity into a warm-start cache
1339    /// fingerprint (#869).
1340    ///
1341    /// The persistent warm-start `cache_key` hashes only family + raw input
1342    /// dimensions, so two fits on the same data that differ *only* in their
1343    /// smooth topology (the `s(..., type=AUTO)` candidate enumeration: sphere
1344    /// vs torus vs euclidean vs duchon) collide on one key and seed each other
1345    /// with geometrically incompatible β/ρ. Folding the per-term structural
1346    /// kind + feature columns + linear/random-effect counts into the shape hash
1347    /// gives each candidate its own key, so the screen→full-refit reuse of one
1348    /// candidate is preserved while cross-candidate contamination is removed.
1349    /// Only the structural identity is hashed (not fitted coefficients or
1350    /// frozen knot values), so a refit of the *same* topology still hits.
1351    pub fn write_structural_shape_hash(&self, h: &mut gam_runtime::warm_start::Fingerprinter) {
1352        h.write_str("term-collection");
1353        h.write_usize(self.linear_terms.len());
1354        for linear in &self.linear_terms {
1355            h.write_str(&linear.name);
1356        }
1357        h.write_usize(self.random_effect_terms.len());
1358        h.write_usize(self.smooth_terms.len());
1359        for smooth in &self.smooth_terms {
1360            h.write_str(&smooth.name);
1361            h.write_str(smooth.basis.structural_kind());
1362            for col in smooth.basis.structural_feature_cols() {
1363                h.write_usize(col);
1364            }
1365        }
1366    }
1367
1368    /// Validate that a term collection spec represents a fully frozen model
1369    /// (i.e. all knots/centers are pre-computed, identifiability transforms are
1370    /// baked in, and random-effect levels are fixed).
1371    pub fn validate_frozen(&self, label: &str) -> Result<(), String> {
1372        for linear in &self.linear_terms {
1373            if let (Some(min), Some(max)) = (linear.coefficient_min, linear.coefficient_max)
1374                && (!min.is_finite() || !max.is_finite() || min > max)
1375            {
1376                return Err(SmoothError::invalid_config(format!(
1377                    "{label} linear term '{}' has invalid coefficient constraint [{min}, {max}]",
1378                    linear.name
1379                ))
1380                .into());
1381            }
1382            if let Some(min) = linear.coefficient_min
1383                && !min.is_finite()
1384            {
1385                return Err(SmoothError::invalid_config(format!(
1386                    "{label} linear term '{}' has non-finite coefficient minimum {min}",
1387                    linear.name
1388                ))
1389                .into());
1390            }
1391            if let Some(max) = linear.coefficient_max
1392                && !max.is_finite()
1393            {
1394                return Err(SmoothError::invalid_config(format!(
1395                    "{label} linear term '{}' has non-finite coefficient maximum {max}",
1396                    linear.name
1397                ))
1398                .into());
1399            }
1400            if let LinearCoefficientGeometry::Bounded { min, max, prior } =
1401                &linear.coefficient_geometry
1402            {
1403                if !min.is_finite() || !max.is_finite() || min >= max {
1404                    return Err(SmoothError::invalid_config(format!(
1405                        "{label} bounded term '{}' has invalid bounds [{min}, {max}]",
1406                        linear.name
1407                    ))
1408                    .into());
1409                }
1410                match prior {
1411                    BoundedCoefficientPriorSpec::None | BoundedCoefficientPriorSpec::Uniform => {}
1412                    BoundedCoefficientPriorSpec::Beta { a, b } => {
1413                        if !a.is_finite() || !b.is_finite() || *a < 1.0 || *b < 1.0 {
1414                            return Err(SmoothError::invalid_config(format!(
1415                                "{label} bounded term '{}' has invalid Beta prior ({a}, {b})",
1416                                linear.name
1417                            ))
1418                            .into());
1419                        }
1420                    }
1421                }
1422            }
1423        }
1424        for st in &self.smooth_terms {
1425            if let Err(error) = st.basis.validate_scale_configuration() {
1426                return Err(SmoothError::invalid_config(format!(
1427                    "{label} term '{}' has an invalid scale contract: {error}",
1428                    st.name
1429                ))
1430                .into());
1431            }
1432            match &st.basis {
1433                SmoothBasisSpec::ByVariable { inner, .. } => {
1434                    validate_smooth_basis_frozen(inner, label, &st.name)?;
1435                    let nested = SmoothTermSpec {
1436                        name: st.name.clone(),
1437                        basis: (**inner).clone(),
1438                        shape: st.shape,
1439                        joint_null_rotation: None,
1440                    };
1441                    TermCollectionSpec {
1442                        linear_terms: Vec::new(),
1443                        random_effect_terms: Vec::new(),
1444                        smooth_terms: vec![nested],
1445                    }
1446                    .validate_frozen(label)?;
1447                }
1448                SmoothBasisSpec::FactorSumToZero { inner, levels, .. } => {
1449                    if levels.len() < 2 {
1450                        return Err(format!(
1451                            "{label} term '{}' has invalid frozen sz levels",
1452                            st.name
1453                        ));
1454                    }
1455                    validate_smooth_basis_frozen(inner, label, &st.name)?;
1456                }
1457                SmoothBasisSpec::BSpline1D { spec, .. } => {
1458                    if !matches!(
1459                        spec.knotspec,
1460                        BSplineKnotSpec::Provided(_)
1461                            | BSplineKnotSpec::PeriodicUniform { .. }
1462                            | BSplineKnotSpec::NaturalCubicRegression { .. }
1463                    ) {
1464                        return Err(SmoothError::invalid_config(format!(
1465                            "{label} term '{}' is not frozen: BSpline knotspec must be Provided, PeriodicUniform, or NaturalCubicRegression",
1466                            st.name
1467                        ))
1468                        .into());
1469                    }
1470                }
1471                SmoothBasisSpec::ThinPlate {
1472                    spec, input_scale, ..
1473                } => {
1474                    if input_scale.is_none() {
1475                        return Err(SmoothError::invalid_config(format!(
1476                            "{label} term '{}' is not frozen: ThinPlate input_scale is missing",
1477                            st.name
1478                        ))
1479                        .into());
1480                    }
1481                    if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
1482                        return Err(SmoothError::invalid_config(format!(
1483                            "{label} term '{}' is not frozen: ThinPlate centers must be UserProvided",
1484                            st.name
1485                        ))
1486                        .into());
1487                    }
1488                    if matches!(
1489                        spec.identifiability,
1490                        SpatialIdentifiability::OrthogonalToParametric
1491                    ) {
1492                        return Err(SmoothError::invalid_config(format!(
1493                            "{label} term '{}' is not frozen: ThinPlate identifiability must be FrozenTransform or None",
1494                            st.name
1495                        ))
1496                        .into());
1497                    }
1498                }
1499                SmoothBasisSpec::Sphere { spec, .. } => {
1500                    if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
1501                        return Err(SmoothError::invalid_config(format!(
1502                            "{label} term '{}' is not frozen: Sphere centers must be UserProvided",
1503                            st.name
1504                        ))
1505                        .into());
1506                    }
1507                    if matches!(spec.method, crate::basis::SphereMethod::Harmonic)
1508                        && spec.max_degree.is_none_or(|d| d == 0)
1509                    {
1510                        return Err(format!(
1511                            "{label} term '{}' is not frozen: sphere max_degree must be positive",
1512                            st.name
1513                        ));
1514                    }
1515                }
1516                SmoothBasisSpec::ConstantCurvature { spec, .. } => {
1517                    if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
1518                        return Err(SmoothError::invalid_config(format!(
1519                            "{label} term '{}' is not frozen: ConstantCurvature centers must be UserProvided",
1520                            st.name
1521                        ))
1522                        .into());
1523                    }
1524                    if !(spec.length_scale.is_finite() && spec.length_scale > 0.0) {
1525                        return Err(SmoothError::invalid_config(format!(
1526                            "{label} term '{}' is not frozen: ConstantCurvature length_scale must be the realized positive value",
1527                            st.name
1528                        ))
1529                        .into());
1530                    }
1531                }
1532                SmoothBasisSpec::MeasureJet {
1533                    spec, input_scale, ..
1534                } => {
1535                    if input_scale.is_none() {
1536                        return Err(SmoothError::invalid_config(format!(
1537                            "{label} term '{}' is not frozen: MeasureJet input_scale is missing",
1538                            st.name
1539                        ))
1540                        .into());
1541                    }
1542                    let centers = match &spec.center_strategy {
1543                        CenterStrategy::UserProvided(centers) => centers,
1544                        _ => {
1545                            return Err(SmoothError::invalid_config(format!(
1546                                "{label} term '{}' is not frozen: MeasureJet centers must be UserProvided",
1547                                st.name
1548                            ))
1549                            .into());
1550                        }
1551                    };
1552                    if centers.nrows() == 0 {
1553                        return Err(SmoothError::invalid_config(format!(
1554                            "{label} term '{}' is not frozen: MeasureJet centers are empty",
1555                            st.name
1556                        ))
1557                        .into());
1558                    }
1559                    if !(spec.length_scale.is_finite() && spec.length_scale > 0.0) {
1560                        return Err(SmoothError::invalid_config(format!(
1561                            "{label} term '{}' is not frozen: MeasureJet length_scale must be the realized positive value",
1562                            st.name
1563                        ))
1564                        .into());
1565                    }
1566                    // Exact replay needs the fit-data penalty quadrature and
1567                    // normalization payload (`BasisMetadata::MeasureJet`).
1568                    let frozen = spec.frozen_quadrature.as_ref().ok_or_else(|| {
1569                        SmoothError::invalid_config(format!(
1570                            "{label} term '{}' is not frozen: MeasureJet frozen_quadrature payload is missing",
1571                            st.name
1572                        ))
1573                    })?;
1574                    if frozen.masses.len() != centers.nrows() {
1575                        return Err(SmoothError::invalid_config(format!(
1576                            "{label} term '{}' frozen MeasureJet has {} masses for {} centers",
1577                            st.name,
1578                            frozen.masses.len(),
1579                            centers.nrows()
1580                        ))
1581                        .into());
1582                    }
1583                    let total_mass = frozen.masses.sum();
1584                    if frozen
1585                        .masses
1586                        .iter()
1587                        .any(|mass| !(mass.is_finite() && *mass >= 0.0))
1588                        || !(total_mass.is_finite() && total_mass > 0.0)
1589                    {
1590                        return Err(SmoothError::invalid_config(format!(
1591                            "{label} term '{}' frozen MeasureJet masses must be finite, nonnegative, and have positive total mass",
1592                            st.name
1593                        ))
1594                        .into());
1595                    }
1596                    let n_levels = frozen.eps_band.len();
1597                    if n_levels == 0
1598                        || frozen
1599                            .eps_band
1600                            .iter()
1601                            .any(|eps| !(eps.is_finite() && *eps > 0.0))
1602                    {
1603                        return Err(SmoothError::invalid_config(format!(
1604                            "{label} term '{}' frozen MeasureJet eps_band must be nonempty, finite, and positive",
1605                            st.name
1606                        ))
1607                        .into());
1608                    }
1609                    for (idx, pair) in frozen.eps_band.windows(2).enumerate() {
1610                        if pair[1] <= pair[0] {
1611                            return Err(SmoothError::invalid_config(format!(
1612                                "{label} term '{}' frozen MeasureJet eps_band is not strictly ascending at {idx}: {} then {}",
1613                                st.name,
1614                                pair[0],
1615                                pair[1]
1616                            ))
1617                            .into());
1618                        }
1619                    }
1620                    validate_measure_jet_positive_vec_len(
1621                        label,
1622                        &st.name,
1623                        "support_means",
1624                        &frozen.support_means,
1625                        n_levels,
1626                    )?;
1627                    // Mode predicate MUST match the builder's
1628                    // (`measure_jet_multiscale_mode`): per-level/multiscale is the
1629                    // explicit `spec.multiscale` opt-in (#1116). In single-scale
1630                    // mode the builder emits a single FUSED penalty (empty
1631                    // per-level scales + `fused_penalty_normalization_scale:
1632                    // Some`); only the multiscale opt-in carries `n_levels`
1633                    // per-level scales.
1634                    let per_level = crate::basis::measure_jet_multiscale_mode(spec);
1635                    if per_level {
1636                        validate_measure_jet_positive_vec_len(
1637                            label,
1638                            &st.name,
1639                            "penalty_normalization_scales",
1640                            &frozen.penalty_normalization_scales,
1641                            n_levels,
1642                        )?;
1643                        validate_measure_jet_positive_vec_len(
1644                            label,
1645                            &st.name,
1646                            "raw_penalty_normalization_scales",
1647                            &frozen.raw_penalty_normalization_scales,
1648                            n_levels,
1649                        )?;
1650                        if frozen.fused_penalty_normalization_scale.is_some() {
1651                            return Err(SmoothError::invalid_config(format!(
1652                                "{label} term '{}' per-level MeasureJet must not carry a fused penalty normalization scale",
1653                                st.name
1654                            ))
1655                            .into());
1656                        }
1657                    } else {
1658                        if !frozen.penalty_normalization_scales.is_empty()
1659                            || !frozen.raw_penalty_normalization_scales.is_empty()
1660                        {
1661                            return Err(SmoothError::invalid_config(format!(
1662                                "{label} term '{}' fused MeasureJet must not carry per-level penalty normalization scales",
1663                                st.name
1664                            ))
1665                            .into());
1666                        }
1667                        match frozen.fused_penalty_normalization_scale {
1668                            Some(scale) if scale.is_finite() && scale > 0.0 => {}
1669                            Some(scale) => {
1670                                return Err(SmoothError::invalid_config(format!(
1671                                    "{label} term '{}' fused MeasureJet penalty normalization scale must be positive and finite, got {scale}",
1672                                    st.name
1673                                ))
1674                                .into());
1675                            }
1676                            None => {
1677                                return Err(SmoothError::invalid_config(format!(
1678                                    "{label} term '{}' fused MeasureJet is missing its penalty normalization scale",
1679                                    st.name
1680                                ))
1681                                .into());
1682                            }
1683                        }
1684                    }
1685                }
1686                SmoothBasisSpec::Matern {
1687                    spec, input_scale, ..
1688                } => {
1689                    if input_scale.is_none() {
1690                        return Err(SmoothError::invalid_config(format!(
1691                            "{label} term '{}' is not frozen: Matern input_scale is missing",
1692                            st.name
1693                        ))
1694                        .into());
1695                    }
1696                    if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
1697                        return Err(SmoothError::invalid_config(format!(
1698                            "{label} term '{}' is not frozen: Matern centers must be UserProvided",
1699                            st.name
1700                        ))
1701                            .into());
1702                    }
1703                    if spec
1704                        .length_scale
1705                        .resolved()
1706                        .is_none_or(|value| !value.is_finite() || value <= 0.0)
1707                    {
1708                        return Err(SmoothError::invalid_config(format!(
1709                            "{label} term '{}' is not frozen: Matern length_scale must be resolved, finite, and positive",
1710                            st.name
1711                        ))
1712                        .into());
1713                    }
1714                }
1715                SmoothBasisSpec::Duchon {
1716                    spec, input_scale, ..
1717                } => {
1718                    if input_scale.is_none() {
1719                        return Err(SmoothError::invalid_config(format!(
1720                            "{label} term '{}' is not frozen: Duchon input_scale is missing",
1721                            st.name
1722                        ))
1723                        .into());
1724                    }
1725                    if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
1726                        return Err(SmoothError::invalid_config(format!(
1727                            "{label} term '{}' is not frozen: Duchon centers must be UserProvided",
1728                            st.name
1729                        ))
1730                        .into());
1731                    }
1732                    if matches!(
1733                        spec.identifiability,
1734                        SpatialIdentifiability::OrthogonalToParametric
1735                    ) {
1736                        return Err(SmoothError::invalid_config(format!(
1737                            "{label} term '{}' is not frozen: Duchon identifiability must be FrozenTransform or None",
1738                            st.name
1739                        ))
1740                        .into());
1741                    }
1742                }
1743                SmoothBasisSpec::Pca {
1744                    centered,
1745                    center_mean,
1746                    pca_basis_path,
1747                    ..
1748                } => {
1749                    if *centered && center_mean.is_none() && pca_basis_path.is_none() {
1750                        return Err(SmoothError::invalid_config(format!(
1751                            "{label} term '{}' is not frozen: centered Pca missing center_mean",
1752                            st.name
1753                        ))
1754                        .into());
1755                    }
1756                }
1757                SmoothBasisSpec::BySmooth { smooth, by_kind } => {
1758                    if let SmoothBasisSpec::BySmooth { .. } = smooth.as_ref() {
1759                        return Err(format!("{label} term '{}' has nested by-smooths", st.name));
1760                    }
1761                    match by_kind {
1762                        ByVarKind::Numeric { .. } => {}
1763                        ByVarKind::Factor { frozen_levels, .. } if frozen_levels.is_none() => {
1764                            return Err(format!(
1765                                "{label} term '{}' is not frozen: by-factor levels missing",
1766                                st.name
1767                            ));
1768                        }
1769                        ByVarKind::Factor { .. } => {}
1770                    }
1771                    let nested = TermCollectionSpec {
1772                        linear_terms: vec![],
1773                        random_effect_terms: vec![],
1774                        smooth_terms: vec![SmoothTermSpec {
1775                            name: st.name.clone(),
1776                            basis: (**smooth).clone(),
1777                            shape: st.shape,
1778                            joint_null_rotation: None,
1779                        }],
1780                    };
1781                    nested.validate_frozen(label)?;
1782                }
1783                SmoothBasisSpec::FactorSmooth { spec } => {
1784                    if spec.group_frozen_levels.is_none() {
1785                        return Err(format!(
1786                            "{label} term '{}' is not frozen: factor-smooth levels missing",
1787                            st.name
1788                        ));
1789                    }
1790                    if !matches!(
1791                        spec.marginal.knotspec,
1792                        BSplineKnotSpec::Provided(_)
1793                            | BSplineKnotSpec::PeriodicUniform { .. }
1794                            // mgcv's `bs="sz"` default marginal is a cubic
1795                            // regression spline (#1074), and the freeze step
1796                            // restores it as a `NaturalCubicRegression` knotspec
1797                            // carrying its `k` value-knots (spatial_optimization.rs
1798                            // `marginal_is_cr` branch) — the SAME treatment the
1799                            // tensor margin already gets in the arm below. Without
1800                            // this variant a frozen `sz` factor smooth fails its own
1801                            // predict-time freeze check ("factor-smooth marginal
1802                            // knots missing") even though its knots are fully
1803                            // materialized; the validation simply was not updated
1804                            // when the cr marginal landed.
1805                            | BSplineKnotSpec::NaturalCubicRegression { .. }
1806                    ) {
1807                        return Err(format!(
1808                            "{label} term '{}' is not frozen: factor-smooth marginal knots missing",
1809                            st.name
1810                        ));
1811                    }
1812                }
1813                SmoothBasisSpec::TensorBSpline { spec, .. } => {
1814                    for (dim, marginal) in spec.marginalspecs.iter().enumerate() {
1815                        if !matches!(
1816                            marginal.knotspec,
1817                            BSplineKnotSpec::Provided(_)
1818                                | BSplineKnotSpec::PeriodicUniform { .. }
1819                                | BSplineKnotSpec::NaturalCubicRegression { .. }
1820                        ) {
1821                            return Err(SmoothError::invalid_config(format!(
1822                                "{label} term '{}' dim {} is not frozen: tensor marginal knotspec must be Provided, PeriodicUniform, or NaturalCubicRegression",
1823                                st.name, dim
1824                            ))
1825                            .into());
1826                        }
1827                    }
1828                    if matches!(
1829                        spec.identifiability,
1830                        TensorBSplineIdentifiability::SumToZero
1831                            | TensorBSplineIdentifiability::MarginalSumToZero
1832                    ) {
1833                        return Err(SmoothError::invalid_config(format!(
1834                            "{label} term '{}' is not frozen: tensor identifiability must be FrozenTransform or None",
1835                            st.name
1836                        ))
1837                        .into());
1838                    }
1839                }
1840            }
1841        }
1842
1843        for rt in &self.random_effect_terms {
1844            if rt.frozen_levels.is_none() {
1845                return Err(SmoothError::invalid_config(format!(
1846                    "{label} random-effect term '{}' is not frozen: missing frozen_levels",
1847                    rt.name
1848                ))
1849                .into());
1850            }
1851        }
1852
1853        Ok(())
1854    }
1855
1856    /// Re-resolve every stored feature-column index through `remap`, returning a
1857    /// spec that addresses a different column layout.
1858    ///
1859    /// A frozen `TermCollectionSpec` stores feature columns as *absolute indices
1860    /// into the training table*. To replay it on a fresh dataset whose columns
1861    /// sit at different positions — the common case at prediction time, where
1862    /// the response column is unknown and may be absent entirely — every index
1863    /// must be re-resolved against the new layout. `remap` receives each
1864    /// training-table index and returns its position in the runtime table;
1865    /// callers typically implement it as "look the name up in the training
1866    /// headers, then resolve that name against the prediction dataset".
1867    ///
1868    /// This is the single authority on *which* fields carry a column index
1869    /// across every basis variant (linear, random-effect, the `by=` column of
1870    /// `ByVariable`/`FactorSumToZero`/`BySmooth`, the continuous and group
1871    /// columns of a `FactorSmooth`, and the multi-axis `feature_cols` of every
1872    /// spatial/tensor basis), so a predict-time realignment cannot silently miss
1873    /// one and dereference a stale training index.
1874    pub fn remap_feature_columns<E, F>(&self, mut remap: F) -> Result<TermCollectionSpec, E>
1875    where
1876        F: FnMut(usize) -> Result<usize, E>,
1877    {
1878        let mut out = self.clone();
1879        for lt in &mut out.linear_terms {
1880            lt.feature_col = remap(lt.feature_col)?;
1881            // Also remap the full interaction-factor list. The design builder
1882            // (`build_term_collection_design_inner`) materializes the column from
1883            // `effective_feature_cols()` — which returns `feature_cols` whenever
1884            // it is non-empty (i.e. essentially always, including a plain linear
1885            // term where `feature_cols == [feature_col]`). Remapping only the
1886            // singular `feature_col` left these at their saved *training* indices
1887            // at predict time, so a parametric `Surv(...) ~ x` (and any `:`
1888            // interaction) bailed with "feature column N out of bounds" once the
1889            // response/time columns shift the runtime layout (issue #898).
1890            for fc in lt.feature_cols.iter_mut() {
1891                *fc = remap(*fc)?;
1892            }
1893            // A factor-aware `:` interaction also gates on categorical columns;
1894            // those indices live in the same training-time layout and must be
1895            // realigned to the runtime table alongside the numeric operands, or
1896            // the predict-time level indicator would dereference a stale column.
1897            for (col, _bits) in lt.categorical_levels.iter_mut() {
1898                *col = remap(*col)?;
1899            }
1900        }
1901        for rt in &mut out.random_effect_terms {
1902            rt.feature_col = remap(rt.feature_col)?;
1903        }
1904        for st in &mut out.smooth_terms {
1905            remap_smooth_basis_feature_columns(&mut st.basis, &mut remap)?;
1906        }
1907        Ok(out)
1908    }
1909}
1910
1911/// Walk a `SmoothBasisSpec` tree, re-resolving every column index through
1912/// `remap`. Shared by all predict-time column realignment (see
1913/// [`TermCollectionSpec::remap_feature_columns`]); kept exhaustive so a newly
1914/// added index-bearing variant fails to compile until it is handled here.
1915pub fn remap_smooth_basis_feature_columns<E, F>(
1916    basis: &mut SmoothBasisSpec,
1917    remap: &mut F,
1918) -> Result<(), E>
1919where
1920    F: FnMut(usize) -> Result<usize, E>,
1921{
1922    match basis {
1923        SmoothBasisSpec::ByVariable { inner, by_col, .. }
1924        | SmoothBasisSpec::FactorSumToZero { inner, by_col, .. } => {
1925            *by_col = remap(*by_col)?;
1926            remap_smooth_basis_feature_columns(inner, remap)?;
1927        }
1928        SmoothBasisSpec::BSpline1D { feature_col, .. } => {
1929            *feature_col = remap(*feature_col)?;
1930        }
1931        SmoothBasisSpec::BySmooth { smooth, by_kind } => {
1932            let by_feature_col = match by_kind {
1933                ByVarKind::Numeric { feature_col } | ByVarKind::Factor { feature_col, .. } => {
1934                    feature_col
1935                }
1936            };
1937            *by_feature_col = remap(*by_feature_col)?;
1938            remap_smooth_basis_feature_columns(smooth, remap)?;
1939        }
1940        SmoothBasisSpec::FactorSmooth { spec } => {
1941            for fc in spec.continuous_cols.iter_mut() {
1942                *fc = remap(*fc)?;
1943            }
1944            spec.group_col = remap(spec.group_col)?;
1945        }
1946        SmoothBasisSpec::ThinPlate { feature_cols, .. }
1947        | SmoothBasisSpec::Sphere { feature_cols, .. }
1948        | SmoothBasisSpec::ConstantCurvature { feature_cols, .. }
1949        | SmoothBasisSpec::Matern { feature_cols, .. }
1950        | SmoothBasisSpec::MeasureJet { feature_cols, .. }
1951        | SmoothBasisSpec::Duchon { feature_cols, .. }
1952        | SmoothBasisSpec::Pca { feature_cols, .. }
1953        | SmoothBasisSpec::TensorBSpline { feature_cols, .. } => {
1954            for fc in feature_cols.iter_mut() {
1955                *fc = remap(*fc)?;
1956            }
1957        }
1958    }
1959    Ok(())
1960}
1961
1962#[derive(Debug, Clone)]
1963pub enum PenaltyStructureHint {
1964    Ridge(f64),
1965    Kronecker(Vec<Array2<f64>>),
1966}
1967
1968/// A penalty matrix stored at its natural block size together with the
1969/// column range it occupies in the global coefficient vector.
1970///
1971/// Instead of embedding every penalty into a full `p_total × p_total` dense
1972/// matrix filled with zeros, we keep the compact local matrix and reconstruct
1973/// the global view only when a downstream consumer explicitly requires it.
1974#[derive(Clone)]
1975pub struct BlockwisePenalty {
1976    /// Column range in the global coefficient vector that this penalty covers.
1977    pub col_range: Range<usize>,
1978    /// The local penalty matrix — dimensions `block_p × block_p` where
1979    /// `block_p = col_range.len()`.
1980    pub local: Array2<f64>,
1981    /// Optional nonzero centering vector for this coefficient block.
1982    pub prior_mean: gam_problem::CoefficientPriorMean,
1983    /// Optional structural hint so downstream spectral/logdet code can stay
1984    /// block-local or factorized without reverse-engineering the matrix.
1985    pub structure_hint: Option<PenaltyStructureHint>,
1986    /// Optional operator-form handle bit-equivalent to `local`. Populated when
1987    /// the originating closed-form factory emitted an op-form penalty so exact
1988    /// operator algebra can use matvec instead of materializing the dense
1989    /// `block_p × block_p` Gram. `None` for ordinary dense penalties.
1990    pub op: Option<std::sync::Arc<dyn crate::analytic_penalties::PenaltyOp>>,
1991}
1992
1993impl std::fmt::Debug for BlockwisePenalty {
1994    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1995        f.debug_struct("BlockwisePenalty")
1996            .field("col_range", &self.col_range)
1997            .field(
1998                "local",
1999                &format_args!("{}×{}", self.local.nrows(), self.local.ncols()),
2000            )
2001            .field("prior_mean", &self.prior_mean)
2002            .field("structure_hint", &self.structure_hint)
2003            .field("op", &self.op.as_ref().map(|o| o.dim()))
2004            .finish()
2005    }
2006}
2007
2008impl BlockwisePenalty {
2009    /// Create a new blockwise penalty.
2010    pub fn new(col_range: Range<usize>, local: Array2<f64>) -> Self {
2011        assert_eq!(col_range.len(), local.nrows());
2012        assert_eq!(col_range.len(), local.ncols());
2013        Self {
2014            col_range,
2015            local,
2016            prior_mean: gam_problem::CoefficientPriorMean::Zero,
2017            structure_hint: None,
2018            op: None,
2019        }
2020    }
2021
2022    pub fn with_prior_mean(mut self, prior_mean: gam_problem::CoefficientPriorMean) -> Self {
2023        self.prior_mean = prior_mean;
2024        self
2025    }
2026
2027    /// Attach an op-form penalty handle bit-equivalent to `local`.
2028    pub fn with_op(
2029        mut self,
2030        op: Option<std::sync::Arc<dyn crate::analytic_penalties::PenaltyOp>>,
2031    ) -> Self {
2032        self.op = op;
2033        self
2034    }
2035
2036    pub fn ridge(col_range: Range<usize>, scale: f64) -> Self {
2037        let block_size = col_range.len();
2038        let mut local = Array2::<f64>::zeros((block_size, block_size));
2039        for i in 0..block_size {
2040            local[[i, i]] = scale;
2041        }
2042        Self {
2043            col_range,
2044            local,
2045            prior_mean: gam_problem::CoefficientPriorMean::Zero,
2046            structure_hint: Some(PenaltyStructureHint::Ridge(scale)),
2047            op: None,
2048        }
2049    }
2050
2051    pub fn kronecker(
2052        col_range: Range<usize>,
2053        local: Array2<f64>,
2054        factors: Vec<Array2<f64>>,
2055    ) -> Self {
2056        assert_eq!(col_range.len(), local.nrows());
2057        assert_eq!(col_range.len(), local.ncols());
2058        Self {
2059            col_range,
2060            local,
2061            prior_mean: gam_problem::CoefficientPriorMean::Zero,
2062            structure_hint: Some(PenaltyStructureHint::Kronecker(factors)),
2063            op: None,
2064        }
2065    }
2066
2067    /// Expand this blockwise penalty into a full `p_total × p_total` dense
2068    /// matrix (mostly zeros). Use sparingly — the whole point of blockwise
2069    /// storage is to avoid this allocation.
2070    pub fn to_global(&self, p_total: usize) -> Array2<f64> {
2071        let mut g = Array2::<f64>::zeros((p_total, p_total));
2072        let r = &self.col_range;
2073        assert!(
2074            r.end <= p_total && self.local.nrows() == r.len() && self.local.ncols() == r.len(),
2075            "BlockwisePenalty::to_global shape invariant violated: \
2076             col_range={}..{}, local={}x{}, p_total={}",
2077            r.start,
2078            r.end,
2079            self.local.nrows(),
2080            self.local.ncols(),
2081            p_total,
2082        );
2083        g.slice_mut(s![r.start..r.end, r.start..r.end])
2084            .assign(&self.local);
2085        g
2086    }
2087
2088    /// Convert into a blockwise [`gam_problem::PenaltyMatrix`] without
2089    /// expanding to full dimensions.
2090    pub fn to_penalty_matrix(&self, total_dim: usize) -> gam_problem::PenaltyMatrix {
2091        gam_problem::PenaltyMatrix::Blockwise {
2092            local: self.local.clone(),
2093            col_range: self.col_range.clone(),
2094            total_dim,
2095        }
2096    }
2097
2098    /// The block size of this penalty.
2099    #[inline]
2100    pub fn block_size(&self) -> usize {
2101        self.col_range.len()
2102    }
2103}
2104
2105/// Compute `Σ_k λ_k S_k` directly from blockwise penalties, accumulating
2106/// into a pre-allocated `p_total × p_total` output without ever materializing
2107/// individual global matrices.
2108pub fn weighted_blockwise_penalty_sum(
2109    penalties: &[BlockwisePenalty],
2110    lambdas: &[f64],
2111    p_total: usize,
2112) -> Array2<f64> {
2113    assert_eq!(penalties.len(), lambdas.len());
2114    // Smoothing parameters λ_k must be non-negative and finite. A negative
2115    // λ would flip the sign of the corresponding block S_k, turning the
2116    // total penalty matrix indefinite and silently corrupting every
2117    // downstream Cholesky / PIRLS / REML / pseudo-logdet computation that
2118    // assumes S_λ ⪰ 0. Catch this at the boundary rather than after it
2119    // has propagated.
2120    for (idx, &lam) in lambdas.iter().enumerate() {
2121        assert!(
2122            lam.is_finite() && lam >= 0.0,
2123            "weighted_blockwise_penalty_sum: lambdas[{idx}] = {lam} is invalid (must be finite and non-negative; negative smoothing parameters violate S_λ ⪰ 0)",
2124        );
2125    }
2126    // Block column ranges must also fit inside the declared total parameter
2127    // dimension; an out-of-bounds slice would otherwise panic from ndarray
2128    // with a far less informative message.
2129    for (idx, bp) in penalties.iter().enumerate() {
2130        let r = &bp.col_range;
2131        assert!(
2132            r.end <= p_total,
2133            "weighted_blockwise_penalty_sum: penalties[{idx}] col_range {:?} exceeds p_total = {p_total}",
2134            r,
2135        );
2136    }
2137    let mut out = Array2::<f64>::zeros((p_total, p_total));
2138    for (bp, &lam) in penalties.iter().zip(lambdas.iter()) {
2139        let r = &bp.col_range;
2140        let mut slice = out.slice_mut(s![r.start..r.end, r.start..r.end]);
2141        slice.scaled_add(lam, &bp.local);
2142    }
2143    out
2144}
2145
2146// ---------------------------------------------------------------------------
2147// KroneckerPenaltySystem — factored tensor-product penalty representation
2148// ---------------------------------------------------------------------------
2149
2150/// Factored representation of tensor-product penalties with precomputed
2151/// marginal eigensystems for O(∏q_j) logdet and penalty operations.
2152#[derive(Debug, Clone)]
2153pub struct KroneckerPenaltySystem {
2154    /// Marginal penalty matrices: `marginal_penalties[k]` is `(q_k, q_k)`.
2155    pub marginal_penalties: Vec<Array2<f64>>,
2156    /// Precomputed eigensystems: `(eigenvalues, eigenvectors)` per marginal.
2157    pub marginal_eigensystems: Vec<(Array1<f64>, Array2<f64>)>,
2158    /// Marginal basis dimensions.
2159    pub marginal_dims: Vec<usize>,
2160    /// Whether a global ridge (double) penalty is present.
2161    pub has_double_penalty: bool,
2162}
2163
2164impl KroneckerPenaltySystem {
2165    pub fn new(
2166        marginal_penalties: Vec<Array2<f64>>,
2167        marginal_dims: Vec<usize>,
2168        has_double_penalty: bool,
2169    ) -> Result<Self, BasisError> {
2170        if marginal_penalties.len() != marginal_dims.len() {
2171            crate::bail_dim_basis!(
2172                "KroneckerPenaltySystem: {} penalties vs {} dims",
2173                marginal_penalties.len(),
2174                marginal_dims.len()
2175            );
2176        }
2177        let eigensystems =
2178            kronecker_marginal_eigensystems(&marginal_penalties, "KroneckerPenaltySystem")
2179                .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
2180        Ok(Self {
2181            marginal_penalties,
2182            marginal_eigensystems: eigensystems,
2183            marginal_dims,
2184            has_double_penalty,
2185        })
2186    }
2187
2188    pub fn p_total(&self) -> usize {
2189        self.marginal_dims.iter().copied().product()
2190    }
2191
2192    pub fn ndim(&self) -> usize {
2193        self.marginal_dims.len()
2194    }
2195
2196    pub fn num_penalties(&self) -> usize {
2197        self.marginal_dims.len() + if self.has_double_penalty { 1 } else { 0 }
2198    }
2199
2200    /// Compute `log|S|₊` and its first/second derivatives w.r.t. `ρ_k = log(λ_k)`.
2201    ///
2202    /// Iterates over the ∏q_j multi-index grid. Cost: O(d · ∏q_j), no O(p²) storage.
2203    pub fn logdet_and_derivatives(
2204        &self,
2205        lambdas: &[f64],
2206        ridge: f64,
2207    ) -> (f64, Array1<f64>, Array2<f64>) {
2208        let n_pen = self.num_penalties();
2209        assert_eq!(lambdas.len(), n_pen, "lambda count mismatch");
2210        let marginal_evals: Vec<_> = self
2211            .marginal_eigensystems
2212            .iter()
2213            .map(|(evals, _)| evals.view())
2214            .collect();
2215        kronecker_logdet_and_derivatives(
2216            &marginal_evals,
2217            &self.marginal_dims,
2218            lambdas,
2219            self.has_double_penalty,
2220            ridge,
2221        )
2222    }
2223
2224    pub fn logdet_rank_and_derivatives(
2225        &self,
2226        lambdas: &[f64],
2227        ridge: f64,
2228    ) -> (f64, usize, Array1<f64>, Array2<f64>) {
2229        let n_pen = self.num_penalties();
2230        assert_eq!(lambdas.len(), n_pen, "lambda count mismatch");
2231        let d = self.marginal_dims.len();
2232        let mut logdet = 0.0;
2233        let mut rank = 0usize;
2234        let mut grad = Array1::<f64>::zeros(n_pen);
2235        let mut hess = Array2::<f64>::zeros((n_pen, n_pen));
2236        // Positivity floor for a penalized eigenvalue `σ`: below this the mode
2237        // is treated as an unpenalized (null-space) direction and excluded from
2238        // both the rank count and the pseudo-log-determinant.
2239        const EIGENVALUE_POSITIVITY_FLOOR: f64 = 1e-12;
2240        // Floor on the *structural* eigenvalue sum (λ-independent) used to decide
2241        // whether a mode lives in the penalty range space and so should receive
2242        // the stabilizing ridge; a structurally-null mode gets no ridge.
2243        const STRUCTURAL_ZERO_FLOOR: f64 = 1e-12;
2244        let mut multi_idx = vec![0usize; d];
2245        loop {
2246            let mut sigma = 0.0;
2247            let mut structural_sigma = 0.0;
2248            for k in 0..d {
2249                let marginal_eigenvalue = self.marginal_eigensystems[k].0[multi_idx[k]];
2250                structural_sigma += marginal_eigenvalue;
2251                sigma += lambdas[k] * marginal_eigenvalue;
2252            }
2253            let joint_null = structural_sigma <= STRUCTURAL_ZERO_FLOOR;
2254            if self.has_double_penalty && joint_null {
2255                sigma += lambdas[d];
2256            }
2257            if structural_sigma > STRUCTURAL_ZERO_FLOOR {
2258                sigma += ridge;
2259            }
2260
2261            if sigma > EIGENVALUE_POSITIVITY_FLOOR {
2262                rank += 1;
2263                logdet += sigma.ln();
2264                let inv_sigma = 1.0 / sigma;
2265                let inv_sigma2 = inv_sigma * inv_sigma;
2266                for k in 0..n_pen {
2267                    let ck = if k < d {
2268                        lambdas[k] * self.marginal_eigensystems[k].0[multi_idx[k]]
2269                    } else if joint_null {
2270                        lambdas[d]
2271                    } else {
2272                        0.0
2273                    };
2274                    grad[k] += ck * inv_sigma;
2275                    hess[[k, k]] += ck * inv_sigma - ck * ck * inv_sigma2;
2276                    for l in (k + 1)..n_pen {
2277                        let cl = if l < d {
2278                            lambdas[l] * self.marginal_eigensystems[l].0[multi_idx[l]]
2279                        } else if joint_null {
2280                            lambdas[d]
2281                        } else {
2282                            0.0
2283                        };
2284                        let off = -ck * cl * inv_sigma2;
2285                        hess[[k, l]] += off;
2286                        hess[[l, k]] += off;
2287                    }
2288                }
2289            }
2290
2291            let mut carry = true;
2292            for dim in (0..d).rev() {
2293                if carry {
2294                    multi_idx[dim] += 1;
2295                    if multi_idx[dim] < self.marginal_dims[dim] {
2296                        carry = false;
2297                    } else {
2298                        multi_idx[dim] = 0;
2299                    }
2300                }
2301            }
2302            if carry {
2303                break;
2304            }
2305        }
2306        (logdet, rank, grad, hess)
2307    }
2308}
2309
2310#[cfg(test)]
2311mod joint_unpenalized_dim_tests {
2312    use super::{ActivePenalty, ActivePenaltyInfo, PenaltySource, joint_unpenalized_dim};
2313    use ndarray::{Array2, array};
2314
2315    fn active_penalty(
2316        matrix: Array2<f64>,
2317        effective_rank: usize,
2318        nullity: usize,
2319        original_index: usize,
2320        source: PenaltySource,
2321    ) -> ActivePenalty {
2322        ActivePenalty {
2323            matrix,
2324            nullity,
2325            null_eigenvectors: None,
2326            op: None,
2327            info: ActivePenaltyInfo {
2328                source,
2329                original_index,
2330                effective_rank,
2331                normalization_scale: 1.0,
2332                kronecker_factors: None,
2333            },
2334        }
2335    }
2336
2337    #[test]
2338    fn no_penalty_is_fully_unpenalized() {
2339        assert_eq!(joint_unpenalized_dim(4, &[]), 4);
2340    }
2341
2342    #[test]
2343    fn single_penalty_returns_its_own_null_space() {
2344        // A 3×3 penalty that penalizes only the last coordinate ⇒ 2-dim null
2345        // space (the first two coordinates are unpenalized).
2346        let s = array![[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 5.0]];
2347        let penalties = [active_penalty(s, 1, 2, 0, PenaltySource::Primary)];
2348        assert_eq!(joint_unpenalized_dim(3, &penalties), 2);
2349    }
2350
2351    #[test]
2352    fn complementary_double_penalty_has_empty_joint_null_space() {
2353        // The #1360 case in miniature: a "bending" penalty that leaves the
2354        // first coordinate (its 2-dim... here 1-dim) null, plus a
2355        // complementary "null-space ridge" that penalizes exactly that
2356        // coordinate. Per-penalty null dims are {1, 2} and sum to 3 (≈ p),
2357        // but the INTERSECTION is empty: every coordinate is penalized by
2358        // someone, so the joint unpenalized dim is 0.
2359        let bending = array![[0.0, 0.0, 0.0], [0.0, 4.0, 0.0], [0.0, 0.0, 4.0]];
2360        let ridge = array![[2.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]];
2361        let penalties = [
2362            active_penalty(bending, 2, 1, 0, PenaltySource::Primary),
2363            active_penalty(ridge, 1, 2, 1, PenaltySource::DoublePenaltyNullspace),
2364        ];
2365        assert_eq!(joint_unpenalized_dim(3, &penalties), 0);
2366    }
2367
2368    #[test]
2369    fn partial_overlap_keeps_shared_null_direction() {
2370        // Two penalties that BOTH leave coordinate 0 unpenalized ⇒ the shared
2371        // null direction survives the intersection (joint unpenalized dim 1),
2372        // even though naively summing the per-penalty dims would give 4.
2373        let a = array![[0.0, 0.0, 0.0], [0.0, 3.0, 0.0], [0.0, 0.0, 0.0]];
2374        let b = array![[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 3.0]];
2375        let penalties = [
2376            active_penalty(a, 1, 2, 0, PenaltySource::Primary),
2377            active_penalty(b, 1, 2, 1, PenaltySource::OperatorStiffness),
2378        ];
2379        assert_eq!(joint_unpenalized_dim(3, &penalties), 1);
2380    }
2381
2382    #[test]
2383    fn non_materialized_penalty_falls_back_conservatively() {
2384        // A penalty whose stored block is not p_local × p_local (e.g. a
2385        // Kronecker tensor factor). With ≥2 penalties the conservative joint
2386        // dim is 0 (never over-rejecting).
2387        let full: Array2<f64> = array![[0.0, 0.0], [0.0, 1.0]];
2388        let factor: Array2<f64> = array![[1.0]]; // wrong shape for p_local=2
2389        let mixed_penalties = [
2390            active_penalty(full, 1, 1, 0, PenaltySource::Primary),
2391            active_penalty(
2392                factor.clone(),
2393                2,
2394                0,
2395                1,
2396                PenaltySource::TensorMarginal { dim: 0 },
2397            ),
2398        ];
2399        assert_eq!(joint_unpenalized_dim(2, &mixed_penalties), 0);
2400        // With a single non-materialized penalty, fall back to its own null dim.
2401        let factor_penalties = [active_penalty(
2402            factor,
2403            2,
2404            2,
2405            0,
2406            PenaltySource::TensorMarginal { dim: 0 },
2407        )];
2408        assert_eq!(joint_unpenalized_dim(4, &factor_penalties), 2);
2409    }
2410}
2411
2412#[cfg(test)]
2413mod kronecker_penalty_system_tests {
2414    use super::KroneckerPenaltySystem;
2415    use ndarray::array;
2416
2417    #[test]
2418    fn double_penalty_rank_derivatives_use_only_joint_null_space() {
2419        let penalties = vec![
2420            array![[0.0, 0.0], [0.0, 2.0]],
2421            array![[0.0, 0.0], [0.0, 3.0]],
2422        ];
2423        let system = KroneckerPenaltySystem::new(penalties, vec![2usize, 2usize], true).unwrap();
2424        let lambdas = vec![5.0, 7.0, 11.0];
2425
2426        let (logdet, rank, grad, hess) = system.logdet_rank_and_derivatives(&lambdas, 0.0);
2427
2428        let expected_diag = [11.0_f64, 21.0, 10.0, 31.0];
2429        let expected_logdet: f64 = expected_diag.iter().map(|v| v.ln()).sum();
2430        assert_eq!(rank, 4);
2431        assert!((logdet - expected_logdet).abs() <= 1e-12);
2432        assert!(
2433            (grad[2] - 1.0).abs() <= 1e-12,
2434            "double-penalty rank derivative must count only the joint null mode, got {}",
2435            grad[2]
2436        );
2437        assert!(hess[[2, 2]].abs() <= 1e-12);
2438    }
2439}
2440
2441#[derive(Clone, Debug)]
2442pub struct TermCollectionDesign {
2443    /// The full design matrix.
2444    ///
2445    /// Prefer a true sparse matrix when every block is sparse-compatible.
2446    /// If the collection already contains intrinsically sparse blocks, preserve
2447    /// that storage and let PIRLS decide later whether the penalized system is
2448    /// sparse-native eligible. Purely dense materialized blocks still fall back
2449    /// to the lazy block operator when sparse storage would just re-encode a
2450    /// dense matrix.
2451    pub design: DesignMatrix,
2452    /// Known row-wise affine contribution to the linear predictor.
2453    ///
2454    /// The realized predictor is `affine_offset + design * beta`. This channel
2455    /// is deliberately separate from `design`: folding it into an estimated
2456    /// intercept would make inhomogeneous term constraints coefficient-
2457    /// dependent and would corrupt all linear-operator/Hessian identities.
2458    pub affine_offset: Array1<f64>,
2459    pub penalties: Vec<BlockwisePenalty>,
2460    pub nullspace_dims: Vec<usize>,
2461    pub penaltyinfo: Vec<PenaltyBlockInfo>,
2462    pub dropped_penaltyinfo: Vec<DroppedPenaltyBlockInfo>,
2463    /// Optional global coefficient lower bounds for constrained fitting.
2464    /// Length equals `design.ncols()` when present. Unconstrained entries are `-inf`.
2465    pub coefficient_lower_bounds: Option<Array1<f64>>,
2466    /// Optional global inequality constraints:
2467    /// `A * beta >= b`.
2468    pub linear_constraints: Option<LinearInequalityConstraints>,
2469    pub intercept_range: Range<usize>,
2470    pub linear_ranges: Vec<(String, Range<usize>)>,
2471    /// Per-linear-term empirical function mass used for this build's ridge
2472    /// penalty (parallel to `spec.linear_terms`; `Some` only for
2473    /// `double_penalty=true` terms). This is the TRAINING-time value the
2474    /// first (fit-time) build actually computed from `data` — never a value
2475    /// recomputed from a prediction/rebuild call's (possibly tiny or
2476    /// zero-variance) rows. `freeze_term_collection_from_design` copies this
2477    /// into each term's `frozen_function_mass` so a later rebuild reuses it
2478    /// instead of recomputing.
2479    pub linear_function_masses: Vec<Option<f64>>,
2480    pub random_effect_ranges: Vec<(String, Range<usize>)>,
2481    pub random_effect_levels: Vec<(String, Vec<u64>)>,
2482    pub smooth: SmoothDesign,
2483}
2484
2485impl TermCollectionDesign {
2486    /// Add this collection's fixed affine channel to a caller-owned likelihood
2487    /// offset, validating the universal row/finite-value contract at the seam
2488    /// where the two offset sources become one.
2489    pub fn compose_offset(
2490        &self,
2491        base: ArrayView1<'_, f64>,
2492        context: &str,
2493    ) -> Result<Array1<f64>, BasisError> {
2494        let n = self.design.nrows();
2495        if self.affine_offset.len() != n || base.len() != n {
2496            crate::bail_dim_basis!(
2497                "{context}: design rows={n}, affine offset rows={}, base offset rows={}",
2498                self.affine_offset.len(),
2499                base.len()
2500            );
2501        }
2502        if self.affine_offset.iter().any(|value| !value.is_finite())
2503            || base.iter().any(|value| !value.is_finite())
2504        {
2505            crate::bail_invalid_basis!("{context}: offsets must be finite");
2506        }
2507        Ok(base.to_owned() + &self.affine_offset)
2508    }
2509
2510    /// Evaluate `affine_offset + design * beta` with a checked coefficient
2511    /// width. Prediction/reporting code should use this instead of applying the
2512    /// linear operator alone whenever it has no separate user offset channel.
2513    pub fn apply(&self, beta: ArrayView1<'_, f64>) -> Result<Array1<f64>, BasisError> {
2514        if beta.len() != self.design.ncols() {
2515            crate::bail_dim_basis!(
2516                "term-collection predictor coefficient length {} does not match design width {}",
2517                beta.len(),
2518                self.design.ncols()
2519            );
2520        }
2521        if beta.iter().any(|value| !value.is_finite()) {
2522            crate::bail_invalid_basis!("term-collection predictor coefficients must be finite");
2523        }
2524        if self.affine_offset.len() != self.design.nrows() {
2525            crate::bail_dim_basis!(
2526                "term-collection affine offset has {} rows but design has {}",
2527                self.affine_offset.len(),
2528                self.design.nrows()
2529            );
2530        }
2531        if self.affine_offset.iter().any(|value| !value.is_finite()) {
2532            crate::bail_invalid_basis!("term-collection affine offset must be finite");
2533        }
2534        Ok(self.design.apply(&beta.to_owned()) + &self.affine_offset)
2535    }
2536
2537    /// Number of global penalty blocks that precede the smooth-term penalty
2538    /// blocks in the flat smoothing-parameter / EDF-trace layout.
2539    ///
2540    /// The prefix is not the same as `random_effect_ranges.len()`: unpenalized
2541    /// (or empty) random-effect ranges contribute coefficient columns but no
2542    /// penalty block. The authoritative layout is the recorded `penaltyinfo`
2543    /// sequence assembled alongside `penalties`.
2544    pub fn leading_penalty_blocks_before_smooth(&self) -> usize {
2545        self.penaltyinfo
2546            .iter()
2547            .take_while(|info| {
2548                matches!(
2549                    &info.penalty.source,
2550                    crate::basis::PenaltySource::Other(source)
2551                        if source == "LinearTermRidge"
2552                            || source.starts_with("RandomEffectRidge(")
2553                )
2554            })
2555            .count()
2556    }
2557
2558    /// Global flat-penalty range owned by one realized smooth term.
2559    ///
2560    /// This is the only supported translation from a smooth-local penalty
2561    /// index to the model-global smoothing-parameter layout. The prefix comes
2562    /// from the metadata recorded alongside the matrices that were actually
2563    /// emitted; it must not be reconstructed from term specs or coefficient
2564    /// ranges, because unpenalized effects have columns but no penalty block.
2565    pub fn smooth_term_penalty_range(
2566        &self,
2567        term_idx: usize,
2568    ) -> Result<Option<Range<usize>>, String> {
2569        let Some(term) = self.smooth.terms.get(term_idx) else {
2570            return Ok(None);
2571        };
2572        if term.active_penalties.is_empty() {
2573            return Ok(None);
2574        }
2575
2576        let leading = self.leading_penalty_blocks_before_smooth();
2577        let smooth_count = self
2578            .smooth
2579            .terms
2580            .iter()
2581            .map(|smooth| smooth.active_penalties.len())
2582            .sum::<usize>();
2583        let expected = leading
2584            .checked_add(smooth_count)
2585            .ok_or_else(|| "term-collection penalty count overflow".to_string())?;
2586        if expected != self.penalties.len() || self.penaltyinfo.len() != self.penalties.len() {
2587            return Err(format!(
2588                "term-collection penalty layout is inconsistent: {leading} leading blocks + \
2589                 {smooth_count} smooth blocks = {expected}, but there are {} penalties and {} \
2590                 metadata records",
2591                self.penalties.len(),
2592                self.penaltyinfo.len()
2593            ));
2594        }
2595
2596        let local_offset = self
2597            .smooth
2598            .terms
2599            .iter()
2600            .take(term_idx)
2601            .map(|smooth| smooth.active_penalties.len())
2602            .sum::<usize>();
2603        let start = leading
2604            .checked_add(local_offset)
2605            .ok_or_else(|| "smooth penalty offset overflow".to_string())?;
2606        let end = start
2607            .checked_add(term.active_penalties.len())
2608            .ok_or_else(|| "smooth penalty range overflow".to_string())?;
2609        Ok(Some(start..end))
2610    }
2611
2612    /// Convert blockwise penalties to `PenaltyMatrix::Blockwise` without
2613    /// expanding to `p_total × p_total`. This is the preferred path for
2614    /// family modules that accept `Vec<PenaltyMatrix>`.
2615    pub fn penalties_as_penalty_matrix(&self) -> Vec<gam_problem::PenaltyMatrix> {
2616        let p = self.design.ncols();
2617        self.penalties
2618            .iter()
2619            .map(|bp| bp.to_penalty_matrix(p))
2620            .collect()
2621    }
2622
2623    /// Number of penalty blocks.
2624    #[inline]
2625    pub fn num_penalties(&self) -> usize {
2626        self.penalties.len()
2627    }
2628
2629    /// Resolve coefficient groups against this design's global coefficient
2630    /// layout and append their penalties after the existing term penalties.
2631    pub fn realize_coefficient_groups(
2632        &self,
2633        groups: &[CoefficientGroupSpec],
2634        base_prior: &gam_spec::RhoPrior,
2635    ) -> Result<RealizedCoefficientGroups, BasisError> {
2636        realize_coefficient_groups(self, groups, base_prior)
2637    }
2638
2639    /// Extract a `KroneckerPenaltySystem` when the model's *only* smooth term is
2640    /// a single Kronecker-factored tensor.
2641    ///
2642    /// This is a deliberate single-tensor fast path, not a partial feature: any
2643    /// other shape — zero Kronecker terms, several of them, or a tensor mixed
2644    /// with non-tensor smooth terms — is served correctly by the standard
2645    /// block-separable assembly, so this returns `None` and the caller falls
2646    /// back to it. The two former conditions (`len != 1` and "a non-Kronecker
2647    /// smooth term exists") are jointly equivalent to "the sole smooth term is
2648    /// Kronecker", which the slice pattern below expresses directly in one pass.
2649    pub fn kronecker_penalty_system(&self) -> Option<KroneckerPenaltySystem> {
2650        let [only_term] = self.smooth.terms.as_slice() else {
2651            return None;
2652        };
2653        let kron = only_term.kronecker_factored.as_ref()?;
2654        // A genuine tensor product needs at least two margins, and the marginal
2655        // design / penalty / dim collections must agree in length. A degenerate
2656        // (single-margin) or internally inconsistent factored basis cannot feed
2657        // the Kronecker fast path, so fall back to the standard assembly rather
2658        // than construct a malformed `KroneckerPenaltySystem` from it.
2659        if kron.marginal_dims.len() < 2
2660            || kron.marginal_penalties.len() != kron.marginal_dims.len()
2661            || kron.marginal_designs.len() != kron.marginal_dims.len()
2662        {
2663            return None;
2664        }
2665        KroneckerPenaltySystem::new(
2666            kron.marginal_penalties.clone(),
2667            kron.marginal_dims.clone(),
2668            kron.has_double_penalty,
2669        )
2670        .ok()
2671    }
2672}
2673
2674// `FittedTermCollection`, `SpatialLengthScaleOptimizationTiming`, and
2675// `FittedTermCollectionWithSpec` were relocated with the GAM fit-orchestration
2676// drivers to `gam-models` (`crate::fit_orchestration::drivers`) — they hold a
2677// `gam_solve::UnifiedFitResult` and are consumed only by those drivers (#1521).
2678
2679#[derive(Clone)]
2680pub struct StandardLatentCoordConfig {
2681    pub values: std::sync::Arc<crate::latent::LatentCoordValues>,
2682    pub term_index: gam_problem::types::SmoothTermIdx,
2683    pub feature_cols: Vec<usize>,
2684    pub manifold: crate::latent::LatentManifold,
2685    pub manifold_auto: bool,
2686    pub retraction_registry: gam_problem::LatentRetractionRegistry,
2687    pub analytic_penalties: Option<std::sync::Arc<crate::AnalyticPenaltyRegistry>>,
2688}
2689
2690#[derive(Clone, Debug, Serialize, Deserialize)]
2691pub struct AdaptiveSpatialMap {
2692    pub termname: String,
2693    pub feature_cols: Vec<usize>,
2694    pub collocation_points: Array2<f64>,
2695    pub inv_magweight: Array1<f64>,
2696    pub invgradweight: Array1<f64>,
2697    pub inv_lapweight: Array1<f64>,
2698}
2699
2700#[derive(Clone, Debug, Serialize, Deserialize)]
2701pub struct AdaptiveRegularizationDiagnostics {
2702    pub epsilon_0: f64,
2703    pub epsilon_g: f64,
2704    pub epsilon_c: f64,
2705    pub epsilon_outer_iterations: usize,
2706    pub mm_iterations: usize,
2707    pub converged: bool,
2708    pub maps: Vec<AdaptiveSpatialMap>,
2709}
2710
2711#[derive(Debug, Clone)]
2712pub struct LinearColumnConditioning {
2713    col_idx: usize,
2714    mean: f64,
2715    scale: f64,
2716}
2717
2718#[derive(Debug, Clone, Default)]
2719pub struct LinearFitConditioning {
2720    pub intercept_idx: usize,
2721    pub columns: Vec<LinearColumnConditioning>,
2722}
2723
2724#[derive(Clone)]
2725pub struct SpatialPsiDerivative {
2726    // These are derivatives with respect to psi = log(kappa), not log(length_scale).
2727    pub penalty_index: usize,
2728    pub penalty_indices: Vec<usize>,
2729    pub global_range: Range<usize>,
2730    pub total_p: usize,
2731    pub x_psi_local: Array2<f64>,
2732    pub s_psi_components_local: Vec<Array2<f64>>,
2733    pub x_psi_psi_local: Array2<f64>,
2734    pub s_psi_psi_components_local: Vec<Array2<f64>>,
2735    pub aniso_group_id: Option<usize>,
2736    /// Pre-computed cross-derivative design matrices for other axes
2737    /// in the same aniso group: Vec of (axis_offset_in_group, matrix).
2738    pub aniso_cross_designs: Option<Vec<(usize, Array2<f64>)>>,
2739    /// On-demand cross-penalty second derivatives ∂²S_m/∂ψ_a∂ψ_b for axes in
2740    /// the same anisotropy group. The input is the other axis offset in the
2741    /// group, and the output is one local penalty matrix per active penalty.
2742    pub aniso_cross_penalty_provider: Option<
2743        std::sync::Arc<
2744            dyn Fn(usize) -> Result<Vec<Array2<f64>>, EstimationError> + Send + Sync + 'static,
2745        >,
2746    >,
2747    /// Optional implicit design-derivative operator (shared across all axes
2748    /// in the same aniso group). When present, `x_psi_local` and
2749    /// `x_psi_psi_local` may be zero-sized, and design-derivative matvecs
2750    /// should go through this operator using `implicit_axis` as the axis index.
2751    pub implicit_operator: Option<std::sync::Arc<crate::basis::ImplicitDesignPsiDerivative>>,
2752    /// Which axis in the implicit operator this entry corresponds to.
2753    pub implicit_axis: usize,
2754}
2755
2756#[derive(Debug, Clone)]
2757pub struct SpatialLogKappaCoords {
2758    /// Flattened ψ values. For isotropic terms, one entry per term.
2759    /// For anisotropic terms, d entries per term (one ψ_a per axis).
2760    pub values: Array1<f64>,
2761    /// Dimensionality of each term: 1 for isotropic, d for anisotropic.
2762    pub dims_per_term: Vec<usize>,
2763}
2764
2765/// Which end of the ψ bound the shared `aniso_bounds_from_data` helper is
2766/// computing. The lower end consumes the `.0` element of
2767/// `spatial_term_psi_bounds`; the upper end consumes `.1`.
2768#[derive(Clone, Copy)]
2769pub enum AnisoBoundEnd {
2770    Lower,
2771    Upper,
2772}
2773
2774impl SpatialLogKappaCoords {
2775    /// Construct from an explicit dims layout plus values.
2776    pub fn new_with_dims(values: Array1<f64>, dims_per_term: Vec<usize>) -> Self {
2777        assert_eq!(
2778            values.len(),
2779            dims_per_term.iter().sum::<usize>(),
2780            "SpatialLogKappaCoords: values length {} != sum of dims_per_term {}",
2781            values.len(),
2782            dims_per_term.iter().sum::<usize>(),
2783        );
2784        Self {
2785            values,
2786            dims_per_term,
2787        }
2788    }
2789
2790    /// Isotropic initialization.
2791    pub fn from_length_scales(
2792        spec: &TermCollectionSpec,
2793        term_indices: &[usize],
2794        options: &SpatialLengthScaleOptimizationOptions,
2795    ) -> Self {
2796        let mut out = Array1::<f64>::zeros(term_indices.len());
2797        for (slot, &term_idx) in term_indices.iter().enumerate() {
2798            // Constant-curvature: the single ψ slot is the raw signed κ, seeded
2799            // from the spec (default κ = 0). The −ln(length_scale) convention is
2800            // log-κ semantics and must not touch the raw-κ coordinate; the κ
2801            // window projection happens later via `clamp_to_bounds`. Mirrors the
2802            // aniso constructor's κ branch.
2803            if let Some(cc) = constant_curvature_term_spec(spec, term_idx) {
2804                out[slot] = cc.kappa;
2805                continue;
2806            }
2807            let length_scale = get_spatial_length_scale(spec, term_idx)
2808                .unwrap_or(options.min_length_scale)
2809                .clamp(options.min_length_scale, options.max_length_scale);
2810            out[slot] = -length_scale.ln();
2811        }
2812        Self {
2813            values: out,
2814            dims_per_term: vec![1; term_indices.len()],
2815        }
2816    }
2817
2818    /// Anisotropic-aware initialization.
2819    ///
2820    /// The input frame is uniformly standardized by `IsotropicScale`; it never
2821    /// manufactures an axis preference. Genuine axis contrasts come only from
2822    /// the term's explicit, centered `aniso_log_scales` state.
2823    ///
2824    /// For each term, checks whether it has `aniso_log_scales` set on its basis spec.
2825    /// - If isotropic (no aniso_log_scales, or 1-D): 1 entry = −ln(length_scale).
2826    /// - If anisotropic with a scalar length scale: d entries, one ψ_a per axis.
2827    ///   Initialized as ψ_a = −ln(length_scale) + η_a  where η_a are the existing
2828    ///   aniso_log_scales (which sum to zero). Multi-dimensional terms without
2829    ///   explicit anisotropy stay scalar here so the seed dimensionality matches
2830    ///   `spatial_dims_per_term`.
2831    pub fn from_length_scales_aniso(
2832        spec: &TermCollectionSpec,
2833        term_indices: &[usize],
2834        options: &SpatialLengthScaleOptimizationOptions,
2835    ) -> Self {
2836        let mut vals = Vec::new();
2837        let mut dims = Vec::new();
2838        for &term_idx in term_indices {
2839            // Measure-jet: dial coordinates seeded directly from the term's
2840            // realized (α, τ[, s]); the −ln(length_scale) convention below is
2841            // κ-semantics and never applies to dials.
2842            if let Some(mj) = measure_jet_term_spec(spec, term_idx) {
2843                let seed = measure_jet_psi_seed(mj);
2844                dims.push(seed.len());
2845                vals.extend(seed);
2846                continue;
2847            }
2848            // Constant-curvature: one signed κ slot seeded from the spec's κ
2849            // (clamped feasible). The −ln(length_scale) convention below is
2850            // log-κ semantics and must not touch the raw-κ coordinate. Bounds
2851            // are unavailable here (no data view), so this is the raw spec κ;
2852            // `reseed_from_data` / `clamp_to_bounds` later project it feasible.
2853            if let Some(cc) = constant_curvature_term_spec(spec, term_idx) {
2854                vals.push(cc.kappa);
2855                dims.push(1);
2856                continue;
2857            }
2858            let length_scale = get_spatial_length_scale(spec, term_idx)
2859                .unwrap_or(options.min_length_scale)
2860                .clamp(options.min_length_scale, options.max_length_scale);
2861            let psi_bar = -length_scale.ln(); // global scale = −ln(length_scale)
2862
2863            if spatial_term_uses_per_axis_psi(spec, term_idx) {
2864                // Per-axis anisotropy is enrolled in the joint outer vector:
2865                // ψ_a = ψ̄ + η_a, one slot per axis. The hyper_dirs builder
2866                // produces matching per-axis derivatives in
2867                // `try_build_spatial_term_log_kappa_aniso_derivativeinfos`.
2868                let d = get_spatial_feature_dim(spec, term_idx).unwrap_or(1);
2869                let eta_raw = get_spatial_aniso_log_scales(spec, term_idx)
2870                    .expect("predicate guarantees aniso_log_scales is Some");
2871                let eta = center_aniso_log_scales(&eta_raw);
2872                for &eta_a in &eta {
2873                    vals.push(psi_bar + eta_a);
2874                }
2875                dims.push(d);
2876            } else {
2877                // Isotropic enrollment — either a 1-D term, a multi-D term
2878                // without explicit anisotropy, or a basis (e.g. Duchon) whose
2879                // η is a fixed geometry parameter rather than a REML hyper
2880                // axis. Exactly one ψ̄ slot, matching the single
2881                // `SpatialPsiDerivative` produced by
2882                // `try_build_spatial_term_log_kappa_derivativeinfo`.
2883                vals.push(psi_bar);
2884                dims.push(1);
2885            }
2886        }
2887        Self {
2888            values: Array1::from_vec(vals),
2889            dims_per_term: dims,
2890        }
2891    }
2892
2893    /// Isotropic lower bounds derived from per-term data geometry.
2894    /// Each entry gets the ψ_lo bound returned by `spatial_term_psi_bounds`
2895    /// for the corresponding term, intersected with the options window.
2896    pub fn lower_bounds_from_data(
2897        data: ArrayView2<'_, f64>,
2898        spec: &TermCollectionSpec,
2899        term_indices: &[usize],
2900        options: &SpatialLengthScaleOptimizationOptions,
2901    ) -> Result<Self, BasisError> {
2902        let mut values = Array1::<f64>::zeros(term_indices.len());
2903        for (slot, &term_idx) in term_indices.iter().enumerate() {
2904            values[slot] = spatial_term_psi_bounds(data, spec, term_idx, options)?.0;
2905        }
2906        Ok(Self {
2907            values,
2908            dims_per_term: vec![1; term_indices.len()],
2909        })
2910    }
2911
2912    /// Isotropic upper bounds derived from per-term data geometry.
2913    pub fn upper_bounds_from_data(
2914        data: ArrayView2<'_, f64>,
2915        spec: &TermCollectionSpec,
2916        term_indices: &[usize],
2917        options: &SpatialLengthScaleOptimizationOptions,
2918    ) -> Result<Self, BasisError> {
2919        let mut values = Array1::<f64>::zeros(term_indices.len());
2920        for (slot, &term_idx) in term_indices.iter().enumerate() {
2921            values[slot] = spatial_term_psi_bounds(data, spec, term_idx, options)?.1;
2922        }
2923        Ok(Self {
2924            values,
2925            dims_per_term: vec![1; term_indices.len()],
2926        })
2927    }
2928
2929    /// Anisotropic-aware lower bounds derived from per-term data geometry.
2930    /// For hybrid anisotropic terms the scalar ψ_lo bound applies to the
2931    /// mean `ψ̄`, not directly to every raw axis coordinate `ψ_a = ψ̄ + η_a`.
2932    /// Shift each axis by the current centered `η_a` so projecting/clamping
2933    /// the seed moves only the global scale direction and does not silently
2934    /// shrink anisotropy that is already consistent with the current
2935    /// `length_scale`.
2936    ///
2937    pub fn lower_bounds_aniso_from_data(
2938        data: ArrayView2<'_, f64>,
2939        spec: &TermCollectionSpec,
2940        term_indices: &[usize],
2941        dims_per_term: &[usize],
2942        options: &SpatialLengthScaleOptimizationOptions,
2943    ) -> Result<Self, BasisError> {
2944        Self::aniso_bounds_from_data(
2945            data,
2946            spec,
2947            term_indices,
2948            dims_per_term,
2949            options,
2950            AnisoBoundEnd::Lower,
2951        )
2952    }
2953
2954    /// Anisotropic-aware upper bounds derived from per-term data geometry.
2955    /// See `lower_bounds_aniso_from_data` for the hybrid-aniso offsetting and
2956    /// pure-Duchon dispatch rationale.
2957    pub fn upper_bounds_aniso_from_data(
2958        data: ArrayView2<'_, f64>,
2959        spec: &TermCollectionSpec,
2960        term_indices: &[usize],
2961        dims_per_term: &[usize],
2962        options: &SpatialLengthScaleOptimizationOptions,
2963    ) -> Result<Self, BasisError> {
2964        Self::aniso_bounds_from_data(
2965            data,
2966            spec,
2967            term_indices,
2968            dims_per_term,
2969            options,
2970            AnisoBoundEnd::Upper,
2971        )
2972    }
2973
2974    /// Shared implementation for the lower/upper anisotropic bounds. The bound
2975    /// end selects one element of the typed `(lo, hi)` data-geometry result;
2976    /// the per-term cursor walk and anisotropy-offset handling are identical.
2977    fn aniso_bounds_from_data(
2978        data: ArrayView2<'_, f64>,
2979        spec: &TermCollectionSpec,
2980        term_indices: &[usize],
2981        dims_per_term: &[usize],
2982        options: &SpatialLengthScaleOptimizationOptions,
2983        end: AnisoBoundEnd,
2984    ) -> Result<Self, BasisError> {
2985        assert_eq!(term_indices.len(), dims_per_term.len());
2986        let total: usize = dims_per_term.iter().sum();
2987        let mut values = Array1::<f64>::zeros(total);
2988        let mut cursor = 0;
2989        for (slot, &term_idx) in term_indices.iter().enumerate() {
2990            let d = dims_per_term[slot];
2991            // Measure-jet: per-coordinate dial boxes, never κ-window geometry
2992            // (which would reject legitimate dial values outright).
2993            if let Some(mj) = measure_jet_term_spec(spec, term_idx) {
2994                let bounds = measure_jet_psi_bound_values(mj, matches!(end, AnisoBoundEnd::Upper));
2995                for (offset, bound) in bounds.into_iter().enumerate() {
2996                    if offset < d {
2997                        values[cursor + offset] = bound;
2998                    }
2999                }
3000                cursor += d;
3001                continue;
3002            }
3003            // Constant-curvature: the single signed-κ box from the data chart
3004            // window (symmetric about κ = 0), never a κ = log-scale window.
3005            if constant_curvature_term_spec(spec, term_idx).is_some() {
3006                let (lo, hi) = constant_curvature_kappa_bounds(data, spec, term_idx);
3007                if d >= 1 {
3008                    values[cursor] = match end {
3009                        AnisoBoundEnd::Lower => lo,
3010                        AnisoBoundEnd::Upper => hi,
3011                    };
3012                }
3013                cursor += d;
3014                continue;
3015            }
3016            let psi_bound = {
3017                let (lo, hi) = spatial_term_psi_bounds(data, spec, term_idx, options)?;
3018                match end {
3019                    AnisoBoundEnd::Lower => lo,
3020                    AnisoBoundEnd::Upper => hi,
3021                }
3022            };
3023            let axis_offsets = if d <= 1 {
3024                vec![0.0; d]
3025            } else {
3026                get_spatial_aniso_log_scales(spec, term_idx)
3027                    .filter(|eta| eta.len() == d)
3028                    .map(|eta| center_aniso_log_scales(&eta))
3029                    .unwrap_or_else(|| vec![0.0; d])
3030            };
3031            for offset in 0..d {
3032                values[cursor + offset] = psi_bound + axis_offsets[offset];
3033            }
3034            cursor += d;
3035        }
3036        Ok(Self {
3037            values,
3038            dims_per_term: dims_per_term.to_vec(),
3039        })
3040    }
3041
3042    /// Rewrite any ψ entries whose originating term lacks an explicit
3043    /// `length_scale` so they sit at the midpoint of the per-term data-derived
3044    /// ψ window. Used so the outer optimizer starts inside the physically
3045    /// meaningful region instead of at an arbitrary `options.max_length_scale`
3046    /// derived seed. For terms with an explicit length_scale, the user's
3047    /// choice is respected. Anisotropy offsets η_a (those stored by
3048    /// `from_length_scales_aniso`) are preserved: we re-center around the new
3049    /// ψ̄, keeping Ση_a = 0.
3050    pub fn reseed_from_data(
3051        mut self,
3052        data: ArrayView2<'_, f64>,
3053        spec: &TermCollectionSpec,
3054        term_indices: &[usize],
3055        options: &SpatialLengthScaleOptimizationOptions,
3056    ) -> Result<Self, BasisError> {
3057        assert_eq!(term_indices.len(), self.dims_per_term.len());
3058        let mut cursor = 0;
3059        for (slot, &term_idx) in term_indices.iter().enumerate() {
3060            let d = self.dims_per_term[slot];
3061            // Measure-jet dials are seeded from the realized spec and must
3062            // not be recentered into a κ data window.
3063            if measure_jet_term_spec(spec, term_idx).is_some() {
3064                cursor += d;
3065                continue;
3066            }
3067            // Constant-curvature κ is seeded from the spec (the user's curvature
3068            // hint, default κ = 0); `clamp_to_bounds` projects it feasible. It
3069            // is not a log-scale, so the log-κ recenter below never applies.
3070            if constant_curvature_term_spec(spec, term_idx).is_some() {
3071                cursor += d;
3072                continue;
3073            }
3074            let Some(psi_bar_new) = spatial_term_psi_seed(data, spec, term_idx, options)? else {
3075                cursor += d;
3076                continue;
3077            };
3078            if d == 0 {
3079                continue;
3080            }
3081            let current: Vec<f64> = self.values.slice(s![cursor..cursor + d]).to_vec();
3082            let psi_bar_old = current.iter().sum::<f64>() / d as f64;
3083            for (offset, &old_value) in current.iter().enumerate() {
3084                self.values[cursor + offset] = psi_bar_new + (old_value - psi_bar_old);
3085            }
3086            cursor += d;
3087        }
3088        Ok(self)
3089    }
3090
3091    /// Project ψ values into `[lower, upper]` element-wise. Used after
3092    /// `from_length_scales*` + `reseed_from_data` when a user-supplied
3093    /// `spec.length_scale` falls outside the data-derived ψ window set by
3094    /// `{lower,upper}_bounds*_from_data`. BFGS requires theta0 ∈ [lower,
3095    /// upper]; projecting is the unique closest feasible seed. The user's
3096    /// length_scale was always a hint for the outer optimizer (the optimizer
3097    /// is authoritative for κ), not a hard constraint — so clipping preserves
3098    /// their intent as far as the geometry allows. Emits `log::info!` when
3099    /// any coordinate moves, so the outside-window case is diagnostically
3100    /// visible (not silent).
3101    pub fn clamp_to_bounds(
3102        mut self,
3103        lower: &SpatialLogKappaCoords,
3104        upper: &SpatialLogKappaCoords,
3105    ) -> Self {
3106        assert_eq!(self.values.len(), lower.values.len());
3107        assert_eq!(self.values.len(), upper.values.len());
3108        let mut n_projected = 0usize;
3109        let mut worst_delta = 0.0_f64;
3110        for idx in 0..self.values.len() {
3111            let lo = lower.values[idx];
3112            let hi = upper.values[idx];
3113            if !(lo.is_finite() && hi.is_finite()) {
3114                continue;
3115            }
3116            let v = self.values[idx];
3117            if v < lo {
3118                worst_delta = worst_delta.max(lo - v);
3119                self.values[idx] = lo;
3120                n_projected += 1;
3121            } else if v > hi {
3122                worst_delta = worst_delta.max(v - hi);
3123                self.values[idx] = hi;
3124                n_projected += 1;
3125            }
3126        }
3127        if n_projected > 0 {
3128            log::info!(
3129                "[spatial-kappa] projected {n_projected}/{} ψ seed coords into data-derived bounds \
3130                 (worst excess={worst_delta:.3} log units); user length_scale falls outside \
3131                 [{KERNEL_RANGE_MIN_DIAMETER_FRACTION}/r_max, {KERNEL_RANGE_MAX_SPACING_MULTIPLE}/r_min] geometry window",
3132                self.values.len()
3133            );
3134        }
3135        self
3136    }
3137
3138    /// Reconstruct from theta tail with known dimensionality layout.
3139    pub fn from_theta_tail_with_dims(
3140        theta: &Array1<f64>,
3141        start: usize,
3142        dims_per_term: Vec<usize>,
3143    ) -> Self {
3144        let total: usize = dims_per_term.iter().sum();
3145        Self {
3146            values: theta.slice(s![start..start + total]).to_owned(),
3147            dims_per_term,
3148        }
3149    }
3150
3151    /// Total number of ψ values in the flat array (= sum of dims_per_term).
3152    pub fn len(&self) -> usize {
3153        self.values.len()
3154    }
3155
3156    /// Dimensionality layout: how many ψ values each term contributes.
3157    pub fn dims_per_term(&self) -> &[usize] {
3158        &self.dims_per_term
3159    }
3160
3161    /// Get the offset into the flat array for logical term i.
3162    fn term_offset(&self, term_idx: usize) -> usize {
3163        self.dims_per_term[..term_idx].iter().sum()
3164    }
3165
3166    /// Get the slice of ψ values for logical term i.
3167    pub fn term_slice(&self, term_idx: usize) -> &[f64] {
3168        let offset = self.term_offset(term_idx);
3169        let d = self.dims_per_term[term_idx];
3170        &self.values.as_slice().unwrap()[offset..offset + d]
3171    }
3172
3173    pub fn as_array(&self) -> &Array1<f64> {
3174        &self.values
3175    }
3176
3177    /// #1464: overwrite the single ψ value of a scalar (1-D) logical term by its
3178    /// position `slot` in this coords vector (the same ordering as the
3179    /// `term_indices` slice the constructors were built from). Used to inject the
3180    /// fixed-κ sign-basin seed into a constant-curvature term's raw-κ slot before
3181    /// the joint solve. No-op (returns `false`) when the slot is not scalar.
3182    pub fn set_scalar_slot(&mut self, slot: usize, value: f64) -> bool {
3183        if slot >= self.dims_per_term.len() || self.dims_per_term[slot] != 1 {
3184            return false;
3185        }
3186        let offset = self.term_offset(slot);
3187        self.values[offset] = value;
3188        true
3189    }
3190
3191    /// Split at a logical-term boundary. `mid` is the number of terms in the
3192    /// first half (not a flat-array index).
3193    pub fn split_at(&self, mid: usize) -> (Self, Self) {
3194        let flat_mid: usize = self.dims_per_term[..mid].iter().sum();
3195        (
3196            Self {
3197                values: self.values.slice(s![0..flat_mid]).to_owned(),
3198                dims_per_term: self.dims_per_term[..mid].to_vec(),
3199            },
3200            Self {
3201                values: self.values.slice(s![flat_mid..]).to_owned(),
3202                dims_per_term: self.dims_per_term[mid..].to_vec(),
3203            },
3204        )
3205    }
3206
3207    /// Apply optimized ψ values back to the spec.
3208    ///
3209    /// For isotropic terms (dims=1): sets scalar length_scale = exp(−ψ).
3210    /// For anisotropic terms (dims=d): hybrid/isotropic families set
3211    /// length_scale = exp(−ψ̄) with centered η_a = ψ_a − ψ̄, while pure Duchon
3212    /// writes only centered η_a and leaves length_scale = None.
3213    pub fn apply_tospec(
3214        &self,
3215        spec: &TermCollectionSpec,
3216        term_indices: &[usize],
3217    ) -> Result<TermCollectionSpec, EstimationError> {
3218        if term_indices.len() != self.dims_per_term.len() {
3219            crate::bail_invalid_estim!(
3220                "SpatialLogKappaCoords::apply_tospec: term count mismatch: \
3221                 term_indices={} dims_per_term={}",
3222                term_indices.len(),
3223                self.dims_per_term.len()
3224            );
3225        }
3226        let mut updated = spec.clone();
3227        for (slot, &term_idx) in term_indices.iter().enumerate() {
3228            let psi = self.term_slice(slot);
3229            let d = self.dims_per_term[slot];
3230            // Measure-jet: write the dial coordinates straight back; the
3231            // κ-translation below would misread them as log-scales.
3232            if measure_jet_term_spec(&updated, term_idx).is_some() {
3233                set_measure_jet_psi_dials(&mut updated, term_idx, psi)?;
3234                continue;
3235            }
3236            // Constant-curvature: write the optimized signed κ straight back;
3237            // the −exp(ψ) length-scale translation below is log-κ semantics and
3238            // would misread the raw curvature.
3239            if constant_curvature_term_spec(&updated, term_idx).is_some() {
3240                set_constant_curvature_kappa(&mut updated, term_idx, psi)?;
3241                continue;
3242            }
3243            let (next_length_scale, next_aniso) = spatial_term_psi_to_length_scale_and_aniso(psi);
3244            if (d == 1 || next_length_scale.is_some())
3245                && let Some(length_scale) = next_length_scale
3246            {
3247                set_spatial_length_scale(&mut updated, term_idx, length_scale)?;
3248            }
3249            if let Some(eta) = next_aniso {
3250                set_spatial_aniso_log_scales(&mut updated, term_idx, eta)?;
3251            }
3252        }
3253        Ok(updated)
3254    }
3255}
3256
3257pub fn center_aniso_log_scales(eta: &[f64]) -> Vec<f64> {
3258    if eta.len() <= 1 {
3259        return eta.to_vec();
3260    }
3261    let mean = eta.iter().sum::<f64>() / eta.len() as f64;
3262    eta.iter()
3263        .map(|&v| {
3264            let centered = v - mean;
3265            if centered.abs() <= 1e-15 {
3266                0.0
3267            } else {
3268                centered
3269            }
3270        })
3271        .collect()
3272}
3273
3274/// Whether a spatial term contributes per-axis ψ entries to the outer joint
3275/// hyperparameter vector.
3276pub fn spatial_term_uses_per_axis_psi(resolvedspec: &TermCollectionSpec, term_idx: usize) -> bool {
3277    if let Some(mj) = measure_jet_term_spec(resolvedspec, term_idx) {
3278        return measure_jet_enrolls_psi(mj);
3279    }
3280    let Some(d) = get_spatial_feature_dim(resolvedspec, term_idx) else {
3281        return false;
3282    };
3283    if d <= 1 {
3284        return false;
3285    }
3286    let Some(eta) = get_spatial_aniso_log_scales(resolvedspec, term_idx) else {
3287        return false;
3288    };
3289    if eta.len() != d {
3290        return false;
3291    }
3292    !matches!(
3293        resolvedspec
3294            .smooth_terms
3295            .get(term_idx)
3296            .map(|term| &term.basis),
3297        Some(SmoothBasisSpec::Duchon { .. })
3298    )
3299}
3300
3301pub fn set_spatial_length_scale(
3302    spec: &mut TermCollectionSpec,
3303    term_idx: usize,
3304    length_scale: f64,
3305) -> Result<(), EstimationError> {
3306    let Some(term) = spec.smooth_terms.get_mut(term_idx) else {
3307        crate::bail_invalid_estim!("spatial length-scale term index {term_idx} out of range");
3308    };
3309    match &mut term.basis {
3310        SmoothBasisSpec::ThinPlate { spec, .. } => {
3311            spec.length_scale = length_scale;
3312            Ok(())
3313        }
3314        SmoothBasisSpec::Matern { spec, .. } => {
3315            spec.length_scale.set_resolved(length_scale);
3316            Ok(())
3317        }
3318        SmoothBasisSpec::Duchon { spec, .. } => {
3319            spec.length_scale = Some(length_scale);
3320            Ok(())
3321        }
3322        _ => Err(EstimationError::InvalidInput(format!(
3323            "term '{}' does not expose a spatial length scale",
3324            term.name
3325        ))),
3326    }
3327}
3328
3329pub fn get_spatial_length_scale(spec: &TermCollectionSpec, term_idx: usize) -> Option<f64> {
3330    spec.smooth_terms
3331        .get(term_idx)
3332        .and_then(|term| match &term.basis {
3333            SmoothBasisSpec::ThinPlate { spec, .. } => Some(spec.length_scale),
3334            SmoothBasisSpec::Matern { spec, .. } => spec.length_scale.resolved(),
3335            SmoothBasisSpec::Duchon { spec, .. } => spec.length_scale,
3336            _ => None,
3337        })
3338}
3339
3340pub fn spatial_term_supports_hyper_optimization(
3341    spec: &TermCollectionSpec,
3342    term_idx: usize,
3343) -> bool {
3344    // Ordinary penalized thin-plate regression splines do not have an
3345    // identifiable kernel scale once REML is already learning the smoothing
3346    // penalty. Treat the resolved length scale as fixed geometry; enrolling a
3347    // scalar TPS kappa axis creates the flat ρ/κ valleys reported in #718,
3348    // #721, #731, and #732.
3349    if let Some(term) = spec.smooth_terms.get(term_idx)
3350        && let SmoothBasisSpec::ThinPlate { .. } = &term.basis
3351    {
3352        return false;
3353    }
3354
3355    // Duchon anisotropy η is a FIXED, geometry-derived basis parameter, NOT a
3356    // REML hyper axis: the metric is estimated once from the knot-cloud spread
3357    // (`auto_seed_aniso_contrasts`, applied on every Duchon basis build) and the
3358    // Hilbert-scale λ's carry all learned smoothness. So a pure Duchon (no κ)
3359    // contributes no outer optimization axis even when `scale_dims` is on —
3360    // "standardize the geometry, then learn the smoothness." Only an explicit
3361    // kernel length scale κ (the Matérn / hybrid path) is optimized here.
3362    //
3363    // ISOTROPIC Matérn: the *default* `matern(x1, x2)` is isotropic
3364    // (`scale_dims=false` → `aniso_log_scales = None`). It contributes exactly
3365    // ONE κ optimization axis — its scalar log-κ. The shared GAMLSS /
3366    // location-scale exact-joint ψ engine and the spatial-κ joint outer solver
3367    // both require an isotropic Matérn block to expose this single isotropic κ
3368    // axis (#822/#851); without it the per-block ψ-derivative lists are empty
3369    // and the joint-ψ hooks degenerate to `None`. The isotropic κ is the lone
3370    // kernel hyper axis here, mirroring the per-axis ψ ARD that the anisotropic
3371    // path exposes (just collapsed to one dimension).
3372    //
3373    // ANISOTROPIC Matérn (`scale_dims=true` → `aniso_log_scales = Some`) keeps
3374    // its per-axis kernel-η ARD: the d-dimensional ψ search is the *point* of
3375    // the anisotropic request ("Matérn keeps its kernel-η ARD").
3376    //
3377    // Either way a Matérn term always enrolls a κ/ψ axis (1 isotropic, or d
3378    // anisotropic), so `spatial_dims_per_term` reports the correct count.
3379    if let Some(term) = spec.smooth_terms.get(term_idx)
3380        && let SmoothBasisSpec::Matern { .. } = &term.basis
3381    {
3382        return true;
3383    }
3384
3385    // Measure-jet geometry dials are outer ψ coordinates; enrollment is
3386    // owned by `measure_jet_enrolls_psi`.
3387    if let Some(mj) = measure_jet_term_spec(spec, term_idx) {
3388        return measure_jet_enrolls_psi(mj);
3389    }
3390
3391    // Constant-curvature smooths always enroll their single signed curvature κ
3392    // as an outer ψ-coordinate (#944 stage 3): κ̂ is the headline estimand, so
3393    // unlike a fixed-ℓ kernel it is fitted by default, not gated on a
3394    // user-supplied scale. The coordinate is raw κ (interior κ = 0), and its
3395    // exact design/penalty κ-derivatives come from
3396    // `build_constant_curvature_basis_kappa_derivatives`.
3397    if constant_curvature_term_spec(spec, term_idx).is_some() {
3398        return true;
3399    }
3400
3401    get_spatial_length_scale(spec, term_idx).is_some()
3402}
3403
3404/// The measure-jet term's spec, when `term_idx` is a measure-jet smooth.
3405/// Single accessor for every dial-plumbing dispatch below.
3406pub fn measure_jet_term_spec(
3407    spec: &TermCollectionSpec,
3408    term_idx: usize,
3409) -> Option<&crate::basis::MeasureJetBasisSpec> {
3410    spec.smooth_terms
3411        .get(term_idx)
3412        .and_then(|term| match &term.basis {
3413            SmoothBasisSpec::MeasureJet { spec, .. } => Some(spec),
3414            _ => None,
3415        })
3416}
3417
3418/// Single source for measure-jet outer-ψ enrollment: the lnτ dial is
3419/// undefined in the τ = 0 pseudo-inverse oracle mode (see
3420/// `build_measure_jet_basis_psi_derivatives`), so only a positive ridge
3421/// enrolls the dial group. `spatial_term_supports_hyper_optimization` and
3422/// `spatial_term_uses_per_axis_psi` both defer here so the θ-layout
3423/// sources cannot disagree.
3424pub fn measure_jet_enrolls_psi(mj: &crate::basis::MeasureJetBasisSpec) -> bool {
3425    // Two independent enrollment sources (#1116), both explicit:
3426    //   * the design-moving representer length-scale ℓ (`learn_length_scale`),
3427    //     available in every mode when the spec opts in;
3428    //   * the multiscale penalty dials (s, α, lnτ): the per-scale spectral
3429    //     split's (α, lnτ) ride the explicit `multiscale` opt-in, and the lnτ
3430    //     channel additionally needs a positive ridge (τ = 0 is the
3431    //     pseudo-inverse oracle mode where lnτ is undefined).
3432    // A term enrolls if EITHER source is active.
3433    measure_jet_learns_length_scale(mj)
3434        || (mj.tau0 > 0.0 && crate::basis::measure_jet_multiscale_mode(mj))
3435}
3436
3437/// Whether the design-moving ℓ dial is enrolled for this term. ℓ is fixed by
3438/// default and learnable in every mode only when `learn_length_scale = true`.
3439pub fn measure_jet_learns_length_scale(mj: &crate::basis::MeasureJetBasisSpec) -> bool {
3440    mj.learn_length_scale
3441}
3442
3443pub fn freeze_measure_jet_length_scale_learning(spec: &mut TermCollectionSpec) -> usize {
3444    let mut frozen = 0;
3445    for term in spec.smooth_terms.iter_mut() {
3446        if let SmoothBasisSpec::MeasureJet { spec: mj, .. } = &mut term.basis
3447            && mj.learn_length_scale
3448        {
3449            mj.learn_length_scale = false;
3450            frozen += 1;
3451        }
3452    }
3453    frozen
3454}
3455
3456/// Measure-jet ψ dial boxes. The dials are NOT log-kernel-scales, so the
3457/// κ-window machinery never applies: `α` spans density-weighted (0) through
3458/// past-Coifman–Lafon (>1) normalization, and `lnτ` covers the ridge from
3459/// numerically-exact-projection to heavy noise-floor damping. (The energy
3460/// order `s` is the pinned explicit value or absorbed by the REML-learned
3461/// per-scale amplitudes — see `measure_jet_penalty_psi_dim` — so it carries no
3462/// dial box.)
3463pub const MEASURE_JET_PSI_ALPHA_BOUNDS: (f64, f64) = (-1.0, 3.0);
3464
3465pub const MEASURE_JET_PSI_LN_TAU_BOUNDS: (f64, f64) = (-18.420680743952367, 4.605170185988092);
3466
3467/// Log-ℓ box for the design-moving representer length-scale dial (#1116). An
3468/// ABSOLUTE window in the data coordinate scale (ln of ℓ ∈ [1e-3, 1e2]) used
3469/// only when the spec explicitly enrolls the learned representer range. Absolute
3470/// (not seed-relative) so the bound producer needs no data view, matching the
3471/// other dial boxes. `ln(1e-3) = -6.9077…`, `ln(1e2) = 4.6051…`.
3472pub const MEASURE_JET_PSI_LN_LENGTH_SCALE_BOUNDS: (f64, f64) =
3473    (-6.907755278982137, 4.605170185988092);
3474
3475/// Number of multiscale PENALTY dials (excluding the design-moving ℓ):
3476/// multiscale (per-scale spectral) mode carries (α, lnτ) = 2 — the order is
3477/// either the pinned explicit `s` or absorbed by the REML-learned per-scale
3478/// amplitudes, so it is NOT a dial; single-scale (the default) carries none.
3479/// MUST agree with the penalty-coordinate layout of
3480/// `build_measure_jet_basis_psi_derivatives` (its `per_level` branch always
3481/// emits exactly the (α, lnτ) coordinate pair).
3482pub fn measure_jet_penalty_psi_dim(mj: &crate::basis::MeasureJetBasisSpec) -> usize {
3483    if crate::basis::measure_jet_multiscale_mode(mj) {
3484        2
3485    } else {
3486        0
3487    }
3488}
3489
3490/// ψ dimension of a measure-jet term. The design-moving ℓ dial (when enrolled)
3491/// is coordinate 0; the multiscale penalty dials follow. MUST agree with the
3492/// coordinate layout of `build_measure_jet_basis_psi_derivatives` (ℓ first).
3493pub fn measure_jet_psi_dim(mj: &crate::basis::MeasureJetBasisSpec) -> usize {
3494    usize::from(measure_jet_learns_length_scale(mj)) + measure_jet_penalty_psi_dim(mj)
3495}
3496
3497/// Seed ψ from the term's realized dials, in producer coordinate order: ℓ first
3498/// (when enrolled), then the multiscale penalty dials. The ℓ seed is the
3499/// realized representer range `ln(length_scale)` (the resolved spec carries the
3500/// concrete auto value after the design build/freeze).
3501pub fn measure_jet_psi_seed(mj: &crate::basis::MeasureJetBasisSpec) -> Vec<f64> {
3502    let mut seed = Vec::with_capacity(measure_jet_psi_dim(mj));
3503    if measure_jet_learns_length_scale(mj) {
3504        // length_scale > 0 after resolution; the 0.0 sentinel (pre-resolution)
3505        // falls back to the centre of the log-ℓ box so the optimizer still
3506        // starts feasible and the first data-aware reseed corrects it.
3507        let ell = if mj.length_scale > 0.0 {
3508            mj.length_scale
3509        } else {
3510            1.0
3511        };
3512        seed.push(ell.ln());
3513    }
3514    if measure_jet_penalty_psi_dim(mj) > 0 {
3515        // Multiscale penalty dials, producer order: (α, lnτ).
3516        let ln_tau = mj.tau0.max(f64::MIN_POSITIVE).ln();
3517        seed.extend_from_slice(&[mj.alpha, ln_tau]);
3518    }
3519    seed
3520}
3521
3522/// One end of the per-coordinate dial boxes, in producer coordinate order
3523/// (ℓ first when enrolled, then the multiscale penalty dials).
3524pub fn measure_jet_psi_bound_values(
3525    mj: &crate::basis::MeasureJetBasisSpec,
3526    upper: bool,
3527) -> Vec<f64> {
3528    let pick = |b: (f64, f64)| if upper { b.1 } else { b.0 };
3529    let mut bounds = Vec::with_capacity(measure_jet_psi_dim(mj));
3530    if measure_jet_learns_length_scale(mj) {
3531        bounds.push(pick(MEASURE_JET_PSI_LN_LENGTH_SCALE_BOUNDS));
3532    }
3533    if measure_jet_penalty_psi_dim(mj) > 0 {
3534        // Multiscale penalty dials, producer order: (α, lnτ).
3535        bounds.push(pick(MEASURE_JET_PSI_ALPHA_BOUNDS));
3536        bounds.push(pick(MEASURE_JET_PSI_LN_TAU_BOUNDS));
3537    }
3538    bounds
3539}
3540
3541/// Write optimized ψ dials back into a measure-jet spec. Returns `true` when
3542/// any dial actually moved. The geometry (centers, masses, band, ℓ, z) is
3543/// ψ-FIXED by contract — only the dials change, so frozen-quadrature
3544/// rebuilds reproduce the identical penalty layout at the new dials.
3545pub fn apply_measure_jet_psi(
3546    mj: &mut crate::basis::MeasureJetBasisSpec,
3547    psi: &[f64],
3548) -> Result<bool, EstimationError> {
3549    if psi.len() != measure_jet_psi_dim(mj) {
3550        crate::bail_invalid_estim!(
3551            "measure-jet ψ write-back dimension mismatch: got {} values for a {}-dial term",
3552            psi.len(),
3553            measure_jet_psi_dim(mj)
3554        );
3555    }
3556    let mut changed = false;
3557    // Coordinate 0 (when enrolled) is the design-moving ln(ℓ); the multiscale
3558    // penalty dials follow. Same order as `measure_jet_psi_seed` and the
3559    // producer (`build_measure_jet_basis_psi_derivatives`).
3560    let mut cursor = 0usize;
3561    if measure_jet_learns_length_scale(mj) {
3562        let next_ell = psi[cursor].exp();
3563        cursor += 1;
3564        if !(next_ell.is_finite() && next_ell > 0.0) {
3565            crate::bail_invalid_estim!(
3566                "measure-jet ψ write-back produced a non-finite/non-positive length_scale (ℓ={next_ell})"
3567            );
3568        }
3569        if next_ell != mj.length_scale {
3570            mj.length_scale = next_ell;
3571            changed = true;
3572        }
3573    }
3574    if measure_jet_penalty_psi_dim(mj) > 0 {
3575        // Multiscale penalty dials, producer order: (α, lnτ). The order `s` is
3576        // not a dial (pinned explicit or absorbed by the per-scale amplitudes).
3577        let next_alpha = psi[cursor];
3578        let next_tau = psi[cursor + 1].exp();
3579        if !(next_alpha.is_finite() && next_tau.is_finite() && next_tau > 0.0) {
3580            crate::bail_invalid_estim!(
3581                "measure-jet ψ write-back produced non-finite dials (alpha={next_alpha}, tau={next_tau})"
3582            );
3583        }
3584        if next_alpha != mj.alpha {
3585            mj.alpha = next_alpha;
3586            changed = true;
3587        }
3588        if next_tau != mj.tau0 {
3589            mj.tau0 = next_tau;
3590            changed = true;
3591        }
3592    }
3593    Ok(changed)
3594}
3595
3596/// Collection-level measure-jet dial write-back (the `apply_tospec` /
3597/// realizer-side entry). Returns whether anything moved.
3598pub fn set_measure_jet_psi_dials(
3599    spec: &mut TermCollectionSpec,
3600    term_idx: usize,
3601    psi: &[f64],
3602) -> Result<bool, EstimationError> {
3603    let Some(term) = spec.smooth_terms.get_mut(term_idx) else {
3604        crate::bail_invalid_estim!("measure-jet ψ write-back: term index {term_idx} out of range");
3605    };
3606    set_single_term_measure_jet_psi_dials(term, psi)
3607}
3608
3609/// Single-term dial write-back: the shared match+apply core, also used
3610/// directly on the cached per-trial build spec (whose caller has already
3611/// change-checked at the collection level and rebuilds regardless of the
3612/// moved flag).
3613pub fn set_single_term_measure_jet_psi_dials(
3614    term: &mut SmoothTermSpec,
3615    psi: &[f64],
3616) -> Result<bool, EstimationError> {
3617    let SmoothBasisSpec::MeasureJet { spec: mj, .. } = &mut term.basis else {
3618        crate::bail_invalid_estim!("measure-jet ψ write-back targeted a non-measure-jet term");
3619    };
3620    apply_measure_jet_psi(mj, psi)
3621}
3622
3623/// The constant-curvature smooth's spec, when `term_idx` is one. Single
3624/// accessor for every κ-ψ dispatch below, mirroring `measure_jet_term_spec`.
3625pub fn constant_curvature_term_spec(
3626    spec: &TermCollectionSpec,
3627    term_idx: usize,
3628) -> Option<&crate::basis::ConstantCurvatureBasisSpec> {
3629    spec.smooth_terms
3630        .get(term_idx)
3631        .and_then(|term| match &term.basis {
3632            SmoothBasisSpec::ConstantCurvature { spec, .. } => Some(spec),
3633            _ => None,
3634        })
3635}
3636
3637/// Hard positive cap on |κ| relative to the data's inverse squared chart
3638/// radius. The κ-stereographic chart is valid for `1 + κ‖x‖² > 0`; at
3639/// `|κ| = 1/R²` (R² = max squared chart radius) the gauge `1 + κ‖x‖²` reaches
3640/// the chart edge for the farthest data point, so the optimizer is boxed to a
3641/// safe fraction of that scale on both sides. κ = 0 (flat) is the centre of
3642/// the window, an interior point of the `S^d ← ℝ^d → H^d` family — exactly the
3643/// reachability the raw-κ (not log-κ) coordinate exists to preserve.
3644pub const CONSTANT_CURVATURE_KAPPA_CHART_FRACTION: f64 = 0.5;
3645
3646/// Floor on the data's squared chart radius used to scale the κ window, so a
3647/// degenerate (near-origin) point cloud still yields a finite, usable bracket
3648/// rather than an unbounded one.
3649pub const CONSTANT_CURVATURE_MIN_CHART_RADIUS2: f64 = 1e-8;
3650
3651/// `(κ_min, κ_max)` outer-optimization window for a constant-curvature term,
3652/// derived from the data's maximum squared chart radius `R²` so the κ-jets
3653/// never leave the κ-stereographic chart. Symmetric about κ = 0:
3654/// `±CONSTANT_CURVATURE_KAPPA_CHART_FRACTION / R²`.
3655pub fn constant_curvature_kappa_bounds(
3656    data: ArrayView2<'_, f64>,
3657    spec: &TermCollectionSpec,
3658    term_idx: usize,
3659) -> (f64, f64) {
3660    let feature_cols = match spec.smooth_terms.get(term_idx).map(|t| &t.basis) {
3661        Some(SmoothBasisSpec::ConstantCurvature { feature_cols, .. }) => feature_cols,
3662        _ => return (-1.0, 1.0),
3663    };
3664    let mut max_r2 = CONSTANT_CURVATURE_MIN_CHART_RADIUS2;
3665    for row in data.outer_iter() {
3666        let mut r2 = 0.0_f64;
3667        for &c in feature_cols.iter() {
3668            if let Some(&v) = row.get(c)
3669                && v.is_finite()
3670            {
3671                r2 += v * v;
3672            }
3673        }
3674        if r2 > max_r2 {
3675            max_r2 = r2;
3676        }
3677    }
3678    let half = CONSTANT_CURVATURE_KAPPA_CHART_FRACTION / max_r2;
3679    (-half, half)
3680}
3681
3682/// Write the optimized κ back into a constant-curvature term spec. Returns
3683/// `true` when κ moved. Centers, ℓ, and the constraint transform `z` are
3684/// κ-FIXED by the basis κ-contract, so only `kappa` changes.
3685pub fn set_constant_curvature_kappa(
3686    spec: &mut TermCollectionSpec,
3687    term_idx: usize,
3688    psi: &[f64],
3689) -> Result<bool, EstimationError> {
3690    let Some(term) = spec.smooth_terms.get_mut(term_idx) else {
3691        crate::bail_invalid_estim!(
3692            "constant-curvature κ write-back: term index {term_idx} out of range"
3693        );
3694    };
3695    set_single_term_constant_curvature_kappa(term, psi)
3696}
3697
3698/// Single-term κ write-back: the shared validate+apply core, also used directly
3699/// on the cached per-trial build spec in the incremental realizer (whose caller
3700/// has already change-checked at the collection level and rebuilds regardless
3701/// of the moved flag). Mirrors [`set_single_term_measure_jet_psi_dials`].
3702pub fn set_single_term_constant_curvature_kappa(
3703    term: &mut SmoothTermSpec,
3704    psi: &[f64],
3705) -> Result<bool, EstimationError> {
3706    if psi.len() != 1 {
3707        crate::bail_invalid_estim!(
3708            "constant-curvature κ write-back expects exactly one value, got {}",
3709            psi.len()
3710        );
3711    }
3712    let next_kappa = psi[0];
3713    if !next_kappa.is_finite() {
3714        crate::bail_invalid_estim!(
3715            "constant-curvature κ write-back produced a non-finite κ = {next_kappa}"
3716        );
3717    }
3718    let SmoothBasisSpec::ConstantCurvature { spec: cc, .. } = &mut term.basis else {
3719        crate::bail_invalid_estim!(
3720            "constant-curvature κ write-back targeted a non-constant-curvature term"
3721        );
3722    };
3723    if cc.kappa != next_kappa {
3724        cc.kappa = next_kappa;
3725        Ok(true)
3726    } else {
3727        Ok(false)
3728    }
3729}
3730
3731/// Returns `true` when a spatial term has NO outer optimization axes — i.e.
3732/// the user provided an explicit `length_scale` and the term does not enroll
3733/// REML-side per-axis ψ contrasts, so both the scalar κ and any fixed geometry
3734/// anisotropy are anchored.
3735///
3736/// This is the per-term predicate that distinguishes "fixed kernel scale"
3737/// from "optimize the kernel scale" within the family entry points that
3738/// want to honor an explicit user-supplied scale (e.g. Bernoulli
3739/// marginal-slope, where the joint-spatial outer solver otherwise spends
3740/// ~80 iters stalled on the user's chosen ρ at high gradient).
3741pub fn spatial_term_has_locked_kappa(spec: &TermCollectionSpec, term_idx: usize) -> bool {
3742    let explicitly_fixed = spec
3743        .smooth_terms
3744        .get(term_idx)
3745        .is_some_and(|term| match &term.basis {
3746            SmoothBasisSpec::Matern { spec, .. } => spec.length_scale.is_fixed(),
3747            SmoothBasisSpec::ThinPlate { .. } => true,
3748            SmoothBasisSpec::Duchon { spec, .. } => spec.length_scale.is_some(),
3749            _ => false,
3750        });
3751    explicitly_fixed && !spatial_term_uses_per_axis_psi(spec, term_idx)
3752}
3753
3754pub fn all_spatial_terms_kappa_fixed(spec: &TermCollectionSpec) -> bool {
3755    spec.smooth_terms.iter().enumerate().all(|(idx, _)| {
3756        !spatial_term_supports_hyper_optimization(spec, idx)
3757            || spatial_term_has_locked_kappa(spec, idx)
3758    })
3759}
3760
3761pub fn spatial_identifiability_policy(
3762    termspec: &SmoothTermSpec,
3763) -> Option<&SpatialIdentifiability> {
3764    match &termspec.basis {
3765        SmoothBasisSpec::ThinPlate { spec, .. } => Some(&spec.identifiability),
3766        SmoothBasisSpec::Duchon { spec, .. } => Some(&spec.identifiability),
3767        _ => None,
3768    }
3769}
3770
3771/// Standard deviation of the wide, weakly-informative symmetric `Normal` prior
3772/// placed on a relaxable double-penalty smooth's `DoublePenaltyNullspace`
3773/// selection coordinate when the fit is well-determined.
3774pub const NULLSPACE_WELLDET_DEGENERACY_RHO_SD: f64 = 15.0;
3775
3776/// True iff `prior` is the well-determined double-penalty null-space
3777/// degeneracy prior placed on a `DoublePenaltyNullspace` selection coordinate.
3778pub fn is_nullspace_degeneracy_prior(prior: &gam_spec::RhoPrior) -> bool {
3779    matches!(
3780        prior,
3781        gam_spec::RhoPrior::Normal { mean, sd }
3782            if *mean == 0.0 && *sd == NULLSPACE_WELLDET_DEGENERACY_RHO_SD
3783    )
3784}
3785
3786/// Per-term data-derived ψ = log κ bounds.
3787///
3788/// Uses the same safe operating range documented in
3789/// [`crate::basis::build_matern_basis`] / [`crate::basis::build_duchon_basis`]:
3790///   κ ∈ [2 / r_max, 1e2 / r_min]
3791/// where (r_min, r_max) are pairwise-distance extrema of the term's resolved
3792/// centers (post-fit) or the standardized feature data columns (pre-fit).
3793/// Lower edge of the data-derived kernel-range window, as a fraction of the
3794/// maximum pairwise distance `r_max`: length scales below `2/r_max` resolve
3795/// structure finer than the closest center pair, so the kernel range floor is
3796/// set at twice the maximum spacing.
3797pub const KERNEL_RANGE_MIN_DIAMETER_FRACTION: f64 = 2.0;
3798
3799/// Upper edge of the data-derived kernel-range window, as a multiple of the
3800/// minimum pairwise distance `r_min`: beyond `100/r_min` the radial columns go
3801/// nearly collinear with the polynomial nullspace, so the kernel range is
3802/// capped here to keep the basis geometry well-conditioned.
3803pub const KERNEL_RANGE_MAX_SPACING_MULTIPLE: f64 = 1e2;
3804
3805fn spatial_term_stored_input_scale(term: &SmoothTermSpec) -> Option<crate::IsotropicScale> {
3806    match &term.basis {
3807        SmoothBasisSpec::ThinPlate { input_scale, .. }
3808        | SmoothBasisSpec::Matern { input_scale, .. }
3809        | SmoothBasisSpec::Duchon { input_scale, .. } => *input_scale,
3810        _ => None,
3811    }
3812}
3813
3814fn spatial_term_realized_input_scale(
3815    data: ArrayView2<'_, f64>,
3816    term: &SmoothTermSpec,
3817) -> Result<crate::IsotropicScale, BasisError> {
3818    let (feature_cols, stored) = match &term.basis {
3819        SmoothBasisSpec::ThinPlate {
3820            feature_cols,
3821            input_scale,
3822            ..
3823        }
3824        | SmoothBasisSpec::Matern {
3825            feature_cols,
3826            input_scale,
3827            ..
3828        }
3829        | SmoothBasisSpec::Duchon {
3830            feature_cols,
3831            input_scale,
3832            ..
3833        } => (feature_cols, input_scale),
3834        _ => {
3835            return Err(BasisError::InvalidInput(format!(
3836                "term '{}' does not have an isotropic Euclidean input frame",
3837                term.name
3838            )));
3839        }
3840    };
3841    if let Some(scale) = stored {
3842        return Ok(*scale);
3843    }
3844    let x = select_columns(data, feature_cols)?;
3845    estimate_isotropic_scale(x.view())
3846}
3847
3848/// Returns ψ-space bounds (ψ_lo = ln(κ_lo), ψ_hi = ln(κ_hi)).
3849///
3850/// The returned window is intersected with the options window so user-set
3851/// `min_length_scale` / `max_length_scale` remain hard limits. Degenerate
3852/// geometry or an empty intersection is a typed error: changing to a generic
3853/// options window would silently optimize in a different coordinate chart.
3854pub fn spatial_term_psi_bounds(
3855    data: ArrayView2<'_, f64>,
3856    spec: &TermCollectionSpec,
3857    term_idx: usize,
3858    options: &SpatialLengthScaleOptimizationOptions,
3859) -> Result<(f64, f64), BasisError> {
3860    let options_window = (
3861        -options.max_length_scale.ln(),
3862        -options.min_length_scale.ln(),
3863    );
3864    // Constant-curvature: the ψ coordinate is the raw signed κ, so its window is
3865    // the chart-feasible κ bracket, NOT a log-ℓ window. Mirrors the aniso bounds
3866    // path's `constant_curvature_kappa_bounds` branch so the isotropic
3867    // (non-aniso) seed clamp projects κ into the right interval.
3868    if constant_curvature_term_spec(spec, term_idx).is_some() {
3869        return Ok(constant_curvature_kappa_bounds(data, spec, term_idx));
3870    }
3871    let term = spec.smooth_terms.get(term_idx).ok_or_else(|| {
3872        BasisError::InvalidInput(format!(
3873            "spatial term index {term_idx} is out of bounds for {} smooth terms",
3874            spec.smooth_terms.len()
3875        ))
3876    })?;
3877    // Prefer resolved centers (post-fit) since they live in the same standardized
3878    // space the kernel actually sees. Centers are capped at `default_num_centers`
3879    // (<=2000), so exact pairwise bounds are cheap (<4M ops). If centers are
3880    // not yet UserProvided, fall back to the standardized feature data columns
3881    // with the capped-sample path (O(K²·d), K=1024) — the sample is
3882    // conservative for κ bounds (see `pairwise_distance_bounds_sampled`
3883    // docs): it never excludes a feasible κ the exact method would include.
3884    //
3885    // Under anisotropy the kernel metric is y-space (y_a = exp(η_a) x_a),
3886    // so r_min/r_max must be y-space distances. This matters only when the
3887    // spec already carries calibrated η_a at setup time (e.g., warm-start
3888    // or refit paths); for fresh optimization η_a starts at 0 and y = x.
3889    let aniso = get_spatial_aniso_log_scales(spec, term_idx);
3890    let stored_input_scale = spatial_term_stored_input_scale(term);
3891    let input_scale = spatial_term_realized_input_scale(data, term)?;
3892    let r_bounds = match spatial_term_center_strategy(term) {
3893        Some(CenterStrategy::UserProvided(centers)) if centers.nrows() >= 2 => {
3894            let mut centers_in_frame = centers.clone();
3895            if stored_input_scale.is_none() {
3896                input_scale.standardize(&mut centers_in_frame);
3897            }
3898            let bounds = match aniso.as_deref() {
3899                Some(eta) if eta.len() == centers_in_frame.ncols() => {
3900                    let y = points_in_aniso_y_space(centers_in_frame.view(), eta);
3901                    pairwise_distance_bounds(y.view())
3902                }
3903                _ => pairwise_distance_bounds(centers_in_frame.view()),
3904            };
3905            bounds
3906        }
3907        _ => {
3908            let x = standardized_spatial_term_data(data, term)?;
3909            match aniso.as_deref() {
3910                Some(eta) if eta.len() == x.ncols() => {
3911                    let y = points_in_aniso_y_space(x.view(), eta);
3912                    pairwise_distance_bounds_sampled(y.view())
3913                }
3914                _ => pairwise_distance_bounds_sampled(x.view()),
3915            }
3916        }
3917    };
3918    let (r_min, r_max) = r_bounds.ok_or_else(|| {
3919        BasisError::InvalidInput(format!(
3920            "term '{}' has no positive finite pairwise-distance range",
3921            term.name
3922        ))
3923    })?;
3924    // Length scales substantially larger than the data diameter make radial
3925    // TPS/Matern columns nearly collinear with their polynomial nullspace.
3926    // The nullspace already carries constant/linear low-frequency structure,
3927    // so cap the kernel range at the diameter scale instead of letting the
3928    // optimizer enter a numerically degenerate basis geometry.
3929    // `r_min`/`r_max` are measured in the standardized kernel frame, where
3930    // ℓ_eff = ℓ_original / σ_geom. The optimizer/spec ψ coordinate is
3931    // ψ_original = log(1/ℓ_original), hence
3932    //
3933    //   κ_original = κ_eff / σ_geom
3934    //              = κ_eff * compensate_length_scale(1, scales).
3935    //
3936    // Convert exactly once here before intersecting the data window with the
3937    // original-coordinate user options. Previously these standardized κ bounds
3938    // were written directly into the spec; the basis builder then divided ℓ by
3939    // σ_geom again, making the realized endpoint too long by 1/σ_geom.
3940    let inverse_sigma = input_scale.reciprocal();
3941    let psi_chart_offset = inverse_sigma.ln();
3942    let psi_lo_data = (KERNEL_RANGE_MIN_DIAMETER_FRACTION / r_max).ln() + psi_chart_offset;
3943    let psi_hi_data = (KERNEL_RANGE_MAX_SPACING_MULTIPLE / r_min).ln() + psi_chart_offset;
3944    // #1074: the Matérn-specific length-scale ceiling that used to live here was
3945    // deleted. It was masking, not fixing, the real defect: a hard upper bound on
3946    // the kernel range that pinned the κ-optimizer short rather than letting the
3947    // optimizer find the REML optimum. Matérn now shares the same generic geometry
3948    // window as Duchon / TPS (`KERNEL_RANGE_MIN_DIAMETER_FRACTION / r_max` floor,
3949    // `KERNEL_RANGE_MAX_SPACING_MULTIPLE / r_min` ceiling); the #1357 fully-flat
3950    // collapse corner is guarded by the EDF-collapse guard in
3951    // `spatial_optimization.rs`, which acts on the realized fit, not on a clamp.
3952    // Intersect with the options window so min/max_length_scale remain hard caps.
3953    let psi_lo = psi_lo_data.max(options_window.0);
3954    let psi_hi = psi_hi_data.min(options_window.1);
3955    if psi_lo >= psi_hi {
3956        return Err(BasisError::InvalidInput(format!(
3957            "term '{}' has an empty spatial ψ window after intersecting data bounds [{psi_lo_data}, {psi_hi_data}] with configured bounds [{}, {}]",
3958            term.name, options_window.0, options_window.1
3959        )));
3960    }
3961    Ok((psi_lo, psi_hi))
3962}
3963
3964#[cfg(test)]
3965mod spatial_psi_bound_coordinate_tests {
3966    use super::*;
3967    use crate::basis::{MaternIdentifiability, MaternNu};
3968    use ndarray::array;
3969
3970    fn frozen_matern_bounds(theta: f64, dilation: f64) -> (f64, f64) {
3971        let source = array![
3972            [-1.7, -0.4],
3973            [-1.1, 0.8],
3974            [-0.2, -1.3],
3975            [0.5, 1.6],
3976            [1.4, -0.7],
3977            [2.1, 0.5],
3978        ];
3979        let (cos_theta, sin_theta) = (theta.cos(), theta.sin());
3980        let mut data = Array2::<f64>::zeros(source.raw_dim());
3981        for row in 0..source.nrows() {
3982            let x = source[[row, 0]];
3983            let y = source[[row, 1]];
3984            data[[row, 0]] = dilation * (cos_theta * x - sin_theta * y);
3985            data[[row, 1]] = dilation * (sin_theta * x + cos_theta * y);
3986        }
3987        let input_scale = estimate_isotropic_scale(data.view()).expect("isotropic input scale");
3988        let mut centers = data.clone();
3989        input_scale.standardize(&mut centers);
3990        let spec = TermCollectionSpec {
3991            linear_terms: Vec::new(),
3992            random_effect_terms: Vec::new(),
3993            smooth_terms: vec![SmoothTermSpec {
3994                name: "matern".to_string(),
3995                basis: SmoothBasisSpec::Matern {
3996                    feature_cols: vec![0, 1],
3997                    spec: MaternBasisSpec {
3998                        periodic: None,
3999                        center_strategy: CenterStrategy::UserProvided(centers),
4000                        length_scale: crate::basis::MaternLengthScale::fixed(1.0),
4001                        nu: MaternNu::FiveHalves,
4002                        include_intercept: false,
4003                        double_penalty: true,
4004                        identifiability: MaternIdentifiability::CenterSumToZero,
4005                        aniso_log_scales: None,
4006                    },
4007                    input_scale: Some(input_scale),
4008                },
4009                shape: ShapeConstraint::None,
4010                joint_null_rotation: None,
4011            }],
4012        };
4013        spatial_term_psi_bounds(
4014            data.view(),
4015            &spec,
4016            0,
4017            &SpatialLengthScaleOptimizationOptions::default(),
4018        )
4019        .expect("finite spatial ψ bounds")
4020    }
4021
4022    fn assert_close(left: f64, right: f64) {
4023        assert!(
4024            (left - right).abs() <= 1e-12,
4025            "coordinate-equivalent bounds differ: left={left:.16e}, right={right:.16e}"
4026        );
4027    }
4028
4029    #[test]
4030    fn standardized_center_bounds_return_to_original_units_under_rotation_and_scaling() {
4031        let base = frozen_matern_bounds(0.0, 1.0);
4032        let rotated = frozen_matern_bounds(0.61, 1.0);
4033        assert_close(rotated.0, base.0);
4034        assert_close(rotated.1, base.1);
4035
4036        let dilation = 4.0_f64;
4037        let rotated_scaled = frozen_matern_bounds(0.61, dilation);
4038        let expected_shift = dilation.ln();
4039        assert_close(rotated_scaled.0, base.0 - expected_shift);
4040        assert_close(rotated_scaled.1, base.1 - expected_shift);
4041    }
4042}
4043
4044/// Data-derived ψ seed for a spatial term when the user has not set an
4045/// explicit length_scale on its basis spec. Uses the geometric mean of the
4046/// data-informed kappa range (i.e., the midpoint of the ψ window).
4047pub fn spatial_term_psi_seed(
4048    data: ArrayView2<'_, f64>,
4049    spec: &TermCollectionSpec,
4050    term_idx: usize,
4051    options: &SpatialLengthScaleOptimizationOptions,
4052) -> Result<Option<f64>, BasisError> {
4053    if get_spatial_length_scale(spec, term_idx).is_some() {
4054        return Ok(None); // user/spec-provided length_scale wins
4055    }
4056    let (psi_lo, psi_hi) = spatial_term_psi_bounds(data, spec, term_idx, options)?;
4057    Ok(Some(0.5 * (psi_lo + psi_hi)))
4058}
4059
4060pub fn spatial_term_psi_to_length_scale_and_aniso(psi: &[f64]) -> (Option<f64>, Option<Vec<f64>>) {
4061    if psi.len() <= 1 {
4062        (Some((-psi.first().copied().unwrap_or(0.0)).exp()), None)
4063    } else {
4064        let psi_bar = psi.iter().sum::<f64>() / psi.len() as f64;
4065        (
4066            Some((-psi_bar).exp()),
4067            Some(psi.iter().map(|&value| value - psi_bar).collect()),
4068        )
4069    }
4070}
4071
4072/// Get the `aniso_log_scales` from a spatial term, if present.
4073pub fn get_spatial_aniso_log_scales(
4074    spec: &TermCollectionSpec,
4075    term_idx: usize,
4076) -> Option<Vec<f64>> {
4077    spec.smooth_terms
4078        .get(term_idx)
4079        .and_then(|term| match &term.basis {
4080            SmoothBasisSpec::Matern { spec, .. } => spec.aniso_log_scales.clone(),
4081            SmoothBasisSpec::Duchon { spec, .. } => spec.aniso_log_scales.clone(),
4082            _ => None,
4083        })
4084}
4085
4086/// Per-axis response-structure score for anisotropy seeding.
4087///
4088/// For each spatial axis `a`, sort the response `y` by the axis coordinate
4089/// `x_a` and measure the total squared successive variation of the sorted
4090/// response, `tv_a = Σ_i (y_{σ(i+1)} − y_{σ(i)})²` where `σ` orders rows by
4091/// `x_a`. An axis that carries real (possibly nonlinear) signal makes `y` vary
4092/// SMOOTHLY when the rows are walked in that axis's order, so `tv_a` is SMALL;
4093/// a pure-nuisance axis leaves `y` looking unordered, so `tv_a` is LARGE.
4094///
4095/// This deliberately does NOT use a linear correlation `corr(x_a, y)`: for an
4096/// odd, symmetric signal such as `sin(2·x1)` over a symmetric domain the linear
4097/// correlation is ~0 on the *signal* axis, which would misdirect the seed. The
4098/// total-variation-of-sorted-response score captures nonlinear association.
4099///
4100/// Returns `score_a = −½·ln(tv_a + ε)` (larger ⇒ more signal on axis `a`),
4101/// centered to sum to zero, or `None` when the data is degenerate (too few
4102/// rows, non-finite, or all axes equally (un)structured). The caller adds a
4103/// BOUNDED multiple of this to the geometry seed — it is a conservative nudge,
4104/// never a hard override.
4105pub fn response_aware_axis_contrasts(
4106    x: ndarray::ArrayView2<'_, f64>,
4107    y: ndarray::ArrayView1<'_, f64>,
4108) -> Option<Vec<f64>> {
4109    let n = x.nrows();
4110    let d = x.ncols();
4111    if d <= 1 || n < 4 || y.len() != n {
4112        return None;
4113    }
4114    if x.iter().any(|v| !v.is_finite()) || y.iter().any(|v| !v.is_finite()) {
4115        return None;
4116    }
4117    let mut scores = Vec::with_capacity(d);
4118    for a in 0..d {
4119        let mut order: Vec<usize> = (0..n).collect();
4120        let col = x.column(a);
4121        order.sort_by(|&i, &j| {
4122            col[i]
4123                .partial_cmp(&col[j])
4124                .unwrap_or(std::cmp::Ordering::Equal)
4125        });
4126        let mut tv = 0.0_f64;
4127        for w in order.windows(2) {
4128            let diff = y[w[1]] - y[w[0]];
4129            tv += diff * diff;
4130        }
4131        // ε guards against ln(0) on a perfectly flat / constant response.
4132        scores.push(-0.5 * (tv + 1e-12).ln());
4133    }
4134    if scores.iter().any(|v| !v.is_finite()) {
4135        return None;
4136    }
4137    let mean = scores.iter().sum::<f64>() / d as f64;
4138    let centered: Vec<f64> = scores.iter().map(|&s| s - mean).collect();
4139    // If every axis is equally structured the centered scores are ~0 and the
4140    // nudge is a no-op — return None so the geometry seed is used unchanged.
4141    if centered.iter().all(|&v| v.abs() < 1e-9) {
4142        return None;
4143    }
4144    Some(centered)
4145}
4146
4147/// Conservative, response-aware anisotropy seed nudge applied before the κ outer
4148/// loop. For each anisotropic spatial term it adds a BOUNDED multiple of the
4149/// per-axis response-structure contrast (`response_aware_axis_contrasts`) on top
4150/// of the existing geometry seed, so the optimizer starts in the correct basin
4151/// instead of at a response-blind near-symmetric point (the #1376 under-recovery
4152/// where a signal axis and a nuisance axis with equal coordinate spread seed to
4153/// ~[0,0]). The nudge is clamped to keep this a perturbation, never a hard
4154/// override, so shared aniso Matérn/Duchon fits cannot be destabilized by it.
4155pub fn apply_response_aware_anisotropy_seed(
4156    data: ArrayView2<'_, f64>,
4157    y: ndarray::ArrayView1<'_, f64>,
4158    spec: &mut TermCollectionSpec,
4159    spatial_terms: &[usize],
4160) {
4161    // Bound on the per-axis contrast nudge (in η units). One LN_2 ≈ 0.69 halves
4162    // the effective per-axis length scale; capping at LN_2 keeps the seed within
4163    // one optimizer log-step of the geometry seed while still breaking the
4164    // symmetric-seed trap.
4165    const MAX_NUDGE: f64 = std::f64::consts::LN_2;
4166    for &term_idx in spatial_terms {
4167        let Some(current_eta) = get_spatial_aniso_log_scales(spec, term_idx) else {
4168            continue;
4169        };
4170        let d = current_eta.len();
4171        if d <= 1 {
4172            continue;
4173        }
4174        let Some(term) = spec.smooth_terms.get(term_idx) else {
4175            continue;
4176        };
4177        let feature_cols = term.basis.structural_feature_cols();
4178        if feature_cols.len() != d {
4179            continue;
4180        }
4181        let Ok(x) = select_columns(data, &feature_cols) else {
4182            continue;
4183        };
4184        let Some(contrast) = response_aware_axis_contrasts(x.view(), y) else {
4185            continue;
4186        };
4187        let nudged: Vec<f64> = current_eta
4188            .iter()
4189            .zip(contrast.iter())
4190            .map(|(&eta_a, &c_a)| eta_a + c_a.clamp(-MAX_NUDGE, MAX_NUDGE))
4191            .collect();
4192        // `set_spatial_aniso_log_scales` re-centers to Σ η = 0. A term that does
4193        // not support aniso scales is silently skipped (the seed is optional).
4194        if let Err(err) = set_spatial_aniso_log_scales(spec, term_idx, nudged) {
4195            log::debug!(
4196                "[spatial-kappa] response-aware anisotropy seed skipped for term {term_idx}: {err}"
4197            );
4198        }
4199    }
4200}
4201
4202/// Get the number of feature columns (spatial dimensionality) for a spatial term.
4203pub fn get_spatial_feature_dim(spec: &TermCollectionSpec, term_idx: usize) -> Option<usize> {
4204    spec.smooth_terms
4205        .get(term_idx)
4206        .and_then(|term| match &term.basis {
4207            SmoothBasisSpec::ThinPlate { feature_cols, .. } => Some(feature_cols.len()),
4208            SmoothBasisSpec::Matern { feature_cols, .. } => Some(feature_cols.len()),
4209            SmoothBasisSpec::Duchon { feature_cols, .. } => Some(feature_cols.len()),
4210            _ => None,
4211        })
4212}
4213
4214/// Log the learned per-axis spatial anisotropy for all spatial terms that
4215/// have `aniso_log_scales` set after optimization.
4216///
4217/// For scalar-scale families this reports eta, effective per-axis length
4218/// scales, and per-axis kappa values. For pure Duchon it reports the centered
4219/// eta contrasts only.
4220pub fn log_spatial_aniso_scales(spec: &TermCollectionSpec) {
4221    for (term_idx, term) in spec.smooth_terms.iter().enumerate() {
4222        let (aniso, length_scale) = match &term.basis {
4223            SmoothBasisSpec::Matern { spec, .. } => {
4224                (spec.aniso_log_scales.as_ref(), spec.length_scale.resolved())
4225            }
4226            SmoothBasisSpec::Duchon { spec, .. } => {
4227                (spec.aniso_log_scales.as_ref(), spec.length_scale)
4228            }
4229            _ => (None, None),
4230        };
4231        let Some(eta) = aniso else { continue };
4232        if eta.is_empty() {
4233            continue;
4234        }
4235        let mut lines = match length_scale {
4236            Some(ls) => format!(
4237                "[spatial-kappa] term {} (\"{}\"): anisotropic length scales optimized (global length_scale={:.4})",
4238                term_idx, term.name, ls
4239            ),
4240            None => format!(
4241                "[spatial-kappa] term {} (\"{}\"): pure Duchon shape anisotropy optimized",
4242                term_idx, term.name
4243            ),
4244        };
4245        for (a, &eta_a) in eta.iter().enumerate() {
4246            if let Some(ls) = length_scale {
4247                let length_a = ls * (-eta_a).exp();
4248                let kappa_a = (1.0 / ls) * eta_a.exp();
4249                lines.push_str(&format!(
4250                    "\n  axis {}: eta={:+.4}, length={:.4}, kappa={:.4}",
4251                    a, eta_a, length_a, kappa_a
4252                ));
4253            } else {
4254                lines.push_str(&format!("\n  axis {}: eta={:+.4}", a, eta_a));
4255            }
4256        }
4257        log::info!("{}", lines);
4258    }
4259}
4260
4261/// Set `aniso_log_scales` on a spatial term's basis spec.
4262pub fn set_spatial_aniso_log_scales(
4263    spec: &mut TermCollectionSpec,
4264    term_idx: usize,
4265    eta: Vec<f64>,
4266) -> Result<(), EstimationError> {
4267    let eta = center_aniso_log_scales(&eta);
4268    let Some(term) = spec.smooth_terms.get_mut(term_idx) else {
4269        crate::bail_invalid_estim!("spatial aniso_log_scales term index {term_idx} out of range");
4270    };
4271    match &mut term.basis {
4272        SmoothBasisSpec::Matern { spec, .. } => {
4273            spec.aniso_log_scales = Some(eta);
4274            Ok(())
4275        }
4276        SmoothBasisSpec::Duchon { spec, .. } => {
4277            spec.aniso_log_scales = Some(eta);
4278            Ok(())
4279        }
4280        _ => Err(EstimationError::InvalidInput(format!(
4281            "term '{}' does not support aniso_log_scales",
4282            term.name
4283        ))),
4284    }
4285}
4286
4287/// Sync knot-cloud-derived anisotropy contrasts from basis metadata back into
4288/// the mutable spec so the optimizer starts from the correct eta values.
4289///
4290/// Call this after building the smooth design but before initializing the
4291/// optimizer's psi coordinates. For each spatial term whose metadata contains
4292/// computed `aniso_log_scales`, this writes them into the spec.
4293pub fn sync_aniso_contrasts_from_metadata(spec: &mut TermCollectionSpec, design: &SmoothDesign) {
4294    for (term_idx, term) in design.terms.iter().enumerate() {
4295        let meta_aniso = match &term.metadata {
4296            BasisMetadata::Matern {
4297                aniso_log_scales, ..
4298            } => aniso_log_scales.clone(),
4299            BasisMetadata::Duchon {
4300                aniso_log_scales, ..
4301            } => aniso_log_scales.clone(),
4302            _ => None,
4303        };
4304        if let Some(eta) = meta_aniso
4305            && eta.len() > 1
4306        {
4307            set_spatial_aniso_log_scales(spec, term_idx, eta).ok();
4308        }
4309    }
4310}
4311
4312#[derive(Debug, Clone)]
4313pub struct SpatialLengthScaleOptimizationOptions {
4314    /// Enable outer-loop optimization over spatial κ (= 1 / length_scale)
4315    /// for supported radial-kernel smooths.
4316    /// This applies to ThinPlate, Matérn, and Duchon terms.
4317    pub enabled: bool,
4318    /// Maximum number of outer iterations in the exact joint [rho, psi] solve.
4319    pub max_outer_iter: usize,
4320    /// Relative improvement threshold for terminating the outer solve.
4321    pub rel_tol: f64,
4322    /// Initial log(length_scale) perturbation used for seed construction.
4323    pub log_step: f64,
4324    /// Minimum allowed length_scale during κ search.
4325    pub min_length_scale: f64,
4326    /// Maximum allowed length_scale during κ search.
4327    pub max_length_scale: f64,
4328    /// Automatic geometry-initializer threshold for large-scale spatial fits.
4329    ///
4330    /// When n exceeds twice this value, the fitter uses a spatially stratified
4331    /// subsample only to seed κ/anisotropy geometry: centers are resolved,
4332    /// axis contrasts are initialized from center/data spread, and one or two
4333    /// cheap ψ reseeding updates are applied. It never runs PIRLS, REML, ARC,
4334    /// BFGS, or any recursive optimizer on the pilot.
4335    ///
4336    /// The final coefficients, smoothing parameters, and spatial geometry are
4337    /// always optimized on the full dataset.
4338    ///
4339    /// Set to 0 to skip the pilot geometry initializer.
4340    pub pilot_subsample_threshold: usize,
4341}
4342
4343impl Default for SpatialLengthScaleOptimizationOptions {
4344    fn default() -> Self {
4345        Self {
4346            enabled: true,
4347            max_outer_iter: 80,
4348            rel_tol: 1e-4,
4349            log_step: std::f64::consts::LN_2,
4350            min_length_scale: 1e-3,
4351            max_length_scale: 1e3,
4352            pilot_subsample_threshold: 10_000,
4353        }
4354    }
4355}
4356
4357impl SpatialLengthScaleOptimizationOptions {
4358    /// Validate the struct's invariants. Callers that construct these options
4359    /// from external input (CLI, config, Python API) should call this before
4360    /// passing the options into the fitter. Returns `Err` with a descriptive
4361    /// message when an invariant is violated; the fitter then panics or
4362    /// returns `EstimationError` at its own boundary.
4363    ///
4364    /// Invariants:
4365    ///   * `min_length_scale > 0`, finite
4366    ///   * `max_length_scale > 0`, finite
4367    ///   * `min_length_scale < max_length_scale`
4368    ///   * `rel_tol > 0`, finite
4369    ///   * `log_step > 0`, finite
4370    ///
4371    /// These invariants are what the downstream κ-bound and ψ-window code
4372    /// assumes (`-log(max_ls)` must be finite, `(min,max)` must not be
4373    /// inverted, etc.). Without validation, invalid options produce silent
4374    /// NaN-propagation inside the outer optimizer.
4375    pub fn validate(&self) -> Result<(), String> {
4376        if !self.min_length_scale.is_finite() || self.min_length_scale <= 0.0 {
4377            return Err(SmoothError::invalid_config(format!(
4378                "SpatialLengthScaleOptimizationOptions::min_length_scale must be > 0 and finite, got {}",
4379                self.min_length_scale
4380            ))
4381            .into());
4382        }
4383        if !self.max_length_scale.is_finite() || self.max_length_scale <= 0.0 {
4384            return Err(SmoothError::invalid_config(format!(
4385                "SpatialLengthScaleOptimizationOptions::max_length_scale must be > 0 and finite, got {}",
4386                self.max_length_scale
4387            ))
4388            .into());
4389        }
4390        if self.min_length_scale >= self.max_length_scale {
4391            return Err(SmoothError::invalid_config(format!(
4392                "SpatialLengthScaleOptimizationOptions requires min_length_scale < max_length_scale, got min={} max={}",
4393                self.min_length_scale, self.max_length_scale
4394            ))
4395            .into());
4396        }
4397        if !self.rel_tol.is_finite() || self.rel_tol <= 0.0 {
4398            return Err(SmoothError::invalid_config(format!(
4399                "SpatialLengthScaleOptimizationOptions::rel_tol must be > 0 and finite, got {}",
4400                self.rel_tol
4401            ))
4402            .into());
4403        }
4404        if !self.log_step.is_finite() || self.log_step <= 0.0 {
4405            return Err(SmoothError::invalid_config(format!(
4406                "SpatialLengthScaleOptimizationOptions::log_step must be > 0 and finite, got {}",
4407                self.log_step
4408            ))
4409            .into());
4410        }
4411        Ok(())
4412    }
4413}
4414
4415#[derive(Debug, Clone)]
4416pub struct RandomEffectBlock {
4417    pub name: String,
4418    /// O(n) group-label vector: group_ids[i] = column index in [0, num_groups).
4419    /// `None` if the observation's level is not in the kept set.
4420    pub group_ids: Vec<Option<usize>>,
4421    pub num_groups: usize,
4422    pub kept_levels: Vec<u64>,
4423}
4424
4425pub const BLOCK_SPARSE_ZERO_EPS: f64 = 1e-12;
4426
4427pub const BLOCK_SPARSE_MAX_DENSITY: f64 = 0.20;
4428
4429pub fn blocks_have_intrinsic_sparse_structure(blocks: &[DesignBlock]) -> bool {
4430    blocks
4431        .iter()
4432        .any(|block| matches!(block, DesignBlock::Sparse(_) | DesignBlock::RandomEffect(_)))
4433}
4434
4435pub fn sparse_compatible_block_nnz(block: &DesignBlock) -> Option<usize> {
4436    match block {
4437        DesignBlock::Intercept(n) => Some(*n),
4438        DesignBlock::RandomEffect(op) => {
4439            Some(op.group_ids.iter().filter(|gid| gid.is_some()).count())
4440        }
4441        DesignBlock::Sparse(sparse) => Some(sparse.val().len()),
4442        DesignBlock::Dense(dense) => dense.as_dense_ref().map(|matrix| {
4443            matrix
4444                .iter()
4445                .filter(|&&value| value.abs() > BLOCK_SPARSE_ZERO_EPS)
4446                .count()
4447        }),
4448    }
4449}
4450
4451pub fn try_build_sparse_design_from_blocks(
4452    blocks: &[DesignBlock],
4453) -> Result<Option<DesignMatrix>, BasisError> {
4454    if blocks.is_empty() {
4455        return Ok(None);
4456    }
4457    let nrows = blocks[0].nrows();
4458    let ncols: usize = blocks.iter().map(DesignBlock::ncols).sum();
4459    if nrows == 0 || ncols == 0 || ncols <= 32 {
4460        return Ok(None);
4461    }
4462
4463    let preserve_sparse_storage = blocks_have_intrinsic_sparse_structure(blocks);
4464    let sparse_nnz_limit = if preserve_sparse_storage {
4465        usize::MAX
4466    } else {
4467        let total_cells = nrows.saturating_mul(ncols);
4468        ((total_cells as f64) * BLOCK_SPARSE_MAX_DENSITY).floor() as usize
4469    };
4470    let mut nnz = 0usize;
4471    for block in blocks {
4472        let block_nnz = if let Some(block_nnz) = sparse_compatible_block_nnz(block) {
4473            block_nnz
4474        } else {
4475            return Ok(None);
4476        };
4477        nnz = nnz.saturating_add(block_nnz);
4478        if nnz > sparse_nnz_limit {
4479            return Ok(None);
4480        }
4481    }
4482
4483    let mut triplets = Vec::<Triplet<usize, usize, f64>>::with_capacity(nnz);
4484    let mut col_offset = 0usize;
4485    for block in blocks {
4486        match block {
4487            DesignBlock::Intercept(n) => {
4488                for row in 0..*n {
4489                    triplets.push(Triplet::new(row, col_offset, 1.0));
4490                }
4491            }
4492            DesignBlock::RandomEffect(op) => {
4493                for (row, group_id) in op.group_ids.iter().enumerate() {
4494                    if let Some(group) = group_id {
4495                        triplets.push(Triplet::new(row, col_offset + group, 1.0));
4496                    }
4497                }
4498            }
4499            DesignBlock::Sparse(sparse) => {
4500                let (symbolic, values) = sparse.parts();
4501                let col_ptr = symbolic.col_ptr();
4502                let row_idx = symbolic.row_idx();
4503                for col in 0..sparse.ncols() {
4504                    for idx in col_ptr[col]..col_ptr[col + 1] {
4505                        let value = values[idx];
4506                        if value.abs() > BLOCK_SPARSE_ZERO_EPS {
4507                            triplets.push(Triplet::new(row_idx[idx], col_offset + col, value));
4508                        }
4509                    }
4510                }
4511            }
4512            DesignBlock::Dense(dense) => {
4513                let matrix = dense.as_dense_ref().ok_or_else(|| {
4514                    BasisError::InvalidInput(
4515                        "sparse-compatible block assembly requires materialized dense blocks"
4516                            .to_string(),
4517                    )
4518                })?;
4519                for row in 0..matrix.nrows() {
4520                    for col in 0..matrix.ncols() {
4521                        let value = matrix[[row, col]];
4522                        if value.abs() > BLOCK_SPARSE_ZERO_EPS {
4523                            triplets.push(Triplet::new(row, col_offset + col, value));
4524                        }
4525                    }
4526                }
4527            }
4528        }
4529        col_offset += block.ncols();
4530    }
4531
4532    let sparse = SparseColMat::try_new_from_triplets(nrows, ncols, &triplets).map_err(|_| {
4533        BasisError::SparseCreation("failed to assemble sparse term-collection design".to_string())
4534    })?;
4535    Ok(Some(DesignMatrix::Sparse(
4536        gam_linalg::matrix::SparseDesignMatrix::new(sparse),
4537    )))
4538}
4539
4540pub fn assemble_term_collection_design_matrix(
4541    blocks: Vec<DesignBlock>,
4542) -> Result<DesignMatrix, BasisError> {
4543    if let Some(sparse) = try_build_sparse_design_from_blocks(&blocks)? {
4544        return Ok(sparse);
4545    }
4546    let block_op = BlockDesignOperator::new(blocks).map_err(|e| {
4547        BasisError::InvalidInput(format!("failed to build block design operator: {e}"))
4548    })?;
4549    Ok(DesignMatrix::Dense(
4550        gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(block_op)),
4551    ))
4552}
4553
4554pub fn select_columns(
4555    data: ArrayView2<'_, f64>,
4556    cols: &[usize],
4557) -> Result<Array2<f64>, BasisError> {
4558    let n = data.nrows();
4559    let p = data.ncols();
4560    for &c in cols {
4561        if c >= p {
4562            crate::bail_dim_basis!("feature column {c} is out of bounds for data with {p} columns");
4563        }
4564    }
4565    let mut out = Array2::<f64>::zeros((n, cols.len()));
4566    for (j, &c) in cols.iter().enumerate() {
4567        out.column_mut(j).assign(&data.column(c));
4568    }
4569    Ok(out)
4570}
4571
4572pub fn nonfinite_value_label(value: f64) -> &'static str {
4573    if value.is_nan() {
4574        "NaN"
4575    } else if value.is_sign_positive() {
4576        "+Inf"
4577    } else {
4578        "-Inf"
4579    }
4580}
4581
4582pub fn validate_term_feature_column_finite(
4583    data: ArrayView2<'_, f64>,
4584    term_kind: &str,
4585    term_name: &str,
4586    feature_col: usize,
4587) -> Result<(), BasisError> {
4588    let p = data.ncols();
4589    if feature_col >= p {
4590        crate::bail_dim_basis!(
4591            "{term_kind} term '{term_name}' feature column {feature_col} out of bounds for {p} columns"
4592        );
4593    }
4594    for (row, &value) in data.column(feature_col).iter().enumerate() {
4595        if !value.is_finite() {
4596            crate::bail_invalid_basis!(
4597                "{term_kind} term '{term_name}' feature column {feature_col} row {row} contains non-finite value {}",
4598                nonfinite_value_label(value)
4599            );
4600        }
4601    }
4602    Ok(())
4603}
4604
4605pub fn validate_smooth_terms_finite_inputs(
4606    data: ArrayView2<'_, f64>,
4607    terms: &[SmoothTermSpec],
4608) -> Result<(), BasisError> {
4609    for term in terms {
4610        for feature_col in smooth_term_feature_cols(term) {
4611            validate_term_feature_column_finite(data, "smooth", &term.name, feature_col)?;
4612        }
4613    }
4614    Ok(())
4615}
4616
4617pub fn validate_term_collection_finite_inputs(
4618    data: ArrayView2<'_, f64>,
4619    spec: &TermCollectionSpec,
4620) -> Result<(), BasisError> {
4621    for term in &spec.linear_terms {
4622        validate_term_feature_column_finite(data, "linear", &term.name, term.feature_col)?;
4623    }
4624    for term in &spec.random_effect_terms {
4625        validate_term_feature_column_finite(data, "random-effect", &term.name, term.feature_col)?;
4626    }
4627    validate_smooth_terms_finite_inputs(data, &spec.smooth_terms)
4628}
4629
4630#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
4631pub struct JointSpatialCenterGroupKey {
4632    feature_cols: Vec<usize>,
4633    strategy_kind: CenterStrategyKind,
4634    strategy_aux: usize,
4635    requested_num_centers: usize,
4636    input_scale_bits: Option<u64>,
4637}
4638
4639pub fn spatial_term_min_center_count(term: &SmoothTermSpec) -> usize {
4640    match &term.basis {
4641        SmoothBasisSpec::ThinPlate { feature_cols, .. } => feature_cols.len() + 1,
4642        SmoothBasisSpec::Duchon {
4643            feature_cols, spec, ..
4644        } => match spec.nullspace_order {
4645            crate::basis::DuchonNullspaceOrder::Zero => 1,
4646            crate::basis::DuchonNullspaceOrder::Linear => feature_cols.len() + 1,
4647            crate::basis::DuchonNullspaceOrder::Degree(degree) => {
4648                crate::basis::duchon_nullspace_dimension(feature_cols.len(), degree)
4649            }
4650        },
4651        SmoothBasisSpec::Matern { .. } => 1,
4652        _ => 1,
4653    }
4654}
4655
4656pub fn spatial_term_group_key(term: &SmoothTermSpec) -> Option<JointSpatialCenterGroupKey> {
4657    let (feature_cols, strategy, input_scale) = match &term.basis {
4658        SmoothBasisSpec::ThinPlate {
4659            feature_cols,
4660            spec,
4661            input_scale,
4662        } => (feature_cols, &spec.center_strategy, *input_scale),
4663        SmoothBasisSpec::Matern {
4664            feature_cols,
4665            spec,
4666            input_scale,
4667        } => (feature_cols, &spec.center_strategy, *input_scale),
4668        SmoothBasisSpec::Duchon {
4669            feature_cols,
4670            spec,
4671            input_scale,
4672        } => (feature_cols, &spec.center_strategy, *input_scale),
4673        _ => return None,
4674    };
4675    let strategy_kind = center_strategy_kind(strategy);
4676    let strategy_aux = match strategy {
4677        CenterStrategy::Auto(inner) => match inner.as_ref() {
4678            CenterStrategy::KMeans { max_iter, .. } => *max_iter,
4679            CenterStrategy::UniformGrid { points_per_dim } => *points_per_dim,
4680            _ => 0,
4681        },
4682        CenterStrategy::KMeans { max_iter, .. } => *max_iter,
4683        CenterStrategy::UniformGrid { points_per_dim } => *points_per_dim,
4684        _ => 0,
4685    };
4686    Some(JointSpatialCenterGroupKey {
4687        feature_cols: feature_cols.clone(),
4688        strategy_kind,
4689        strategy_aux,
4690        requested_num_centers: strategy.planned_num_centers(feature_cols.len()),
4691        input_scale_bits: input_scale.map(crate::IsotropicScale::to_bits),
4692    })
4693}
4694
4695pub fn spatial_term_center_strategy(term: &SmoothTermSpec) -> Option<&CenterStrategy> {
4696    match &term.basis {
4697        SmoothBasisSpec::ThinPlate { spec, .. } => Some(&spec.center_strategy),
4698        SmoothBasisSpec::Matern { spec, .. } => Some(&spec.center_strategy),
4699        SmoothBasisSpec::Duchon { spec, .. } => Some(&spec.center_strategy),
4700        _ => None,
4701    }
4702}
4703
4704pub fn set_spatial_term_centers(
4705    term: &mut SmoothTermSpec,
4706    centers: Array2<f64>,
4707) -> Result<(), BasisError> {
4708    match &mut term.basis {
4709        SmoothBasisSpec::ThinPlate { spec, .. } => {
4710            spec.center_strategy = CenterStrategy::UserProvided(centers);
4711            Ok(())
4712        }
4713        SmoothBasisSpec::Matern { spec, .. } => {
4714            spec.center_strategy = CenterStrategy::UserProvided(centers);
4715            Ok(())
4716        }
4717        SmoothBasisSpec::Duchon { spec, .. } => {
4718            spec.center_strategy = CenterStrategy::UserProvided(centers);
4719            Ok(())
4720        }
4721        _ => Err(BasisError::InvalidInput(format!(
4722            "term '{}' does not support spatial center planning",
4723            term.name
4724        ))),
4725    }
4726}
4727
4728pub fn standardized_spatial_term_data(
4729    data: ArrayView2<'_, f64>,
4730    term: &SmoothTermSpec,
4731) -> Result<Array2<f64>, BasisError> {
4732    let (feature_cols, input_scale) = match &term.basis {
4733        SmoothBasisSpec::ThinPlate {
4734            feature_cols,
4735            input_scale,
4736            ..
4737        }
4738        | SmoothBasisSpec::Matern {
4739            feature_cols,
4740            input_scale,
4741            ..
4742        }
4743        | SmoothBasisSpec::Duchon {
4744            feature_cols,
4745            input_scale,
4746            ..
4747        } => (feature_cols, *input_scale),
4748        _ => {
4749            crate::bail_invalid_basis!("term '{}' is not a spatial smooth", term.name);
4750        }
4751    };
4752    let mut x = select_columns(data, feature_cols)?;
4753    input_scale
4754        .map_or_else(|| estimate_isotropic_scale(x.view()), Ok)?
4755        .standardize(&mut x);
4756    Ok(x)
4757}
4758
4759pub fn plan_joint_spatial_centers_for_term_blocks(
4760    data: ArrayView2<'_, f64>,
4761    term_blocks: &[Vec<SmoothTermSpec>],
4762) -> Result<Vec<Vec<SmoothTermSpec>>, BasisError> {
4763    let mut planned_blocks = term_blocks.to_vec();
4764    let n = data.nrows();
4765    let mut groups: BTreeMap<JointSpatialCenterGroupKey, Vec<(usize, usize)>> = BTreeMap::new();
4766
4767    for (block_idx, terms) in planned_blocks.iter().enumerate() {
4768        for (term_idx, term) in terms.iter().enumerate() {
4769            let Some(strategy) = spatial_term_center_strategy(term) else {
4770                continue;
4771            };
4772            if !center_strategy_is_auto(strategy) {
4773                continue;
4774            }
4775            let Some(group_key) = spatial_term_group_key(term) else {
4776                continue;
4777            };
4778            if !matches!(
4779                group_key.strategy_kind,
4780                CenterStrategyKind::EqualMass
4781                    | CenterStrategyKind::EqualMassCovarRepresentative
4782                    | CenterStrategyKind::FarthestPoint
4783                    | CenterStrategyKind::KMeans
4784                    | CenterStrategyKind::UniformGrid
4785            ) {
4786                continue;
4787            }
4788            groups
4789                .entry(group_key)
4790                .or_default()
4791                .push((block_idx, term_idx));
4792        }
4793    }
4794
4795    for (group_key, members) in groups {
4796        if members.len() < 2 {
4797            continue;
4798        }
4799        let min_required = members
4800            .iter()
4801            .map(|&(block_idx, term_idx)| {
4802                spatial_term_min_center_count(&planned_blocks[block_idx][term_idx])
4803            })
4804            .max()
4805            .unwrap_or(1);
4806        let joint_centers = group_key
4807            .requested_num_centers
4808            .max(min_required)
4809            .min(n.max(1));
4810        let (first_block_idx, first_term_idx) = members[0];
4811        let prototype = &planned_blocks[first_block_idx][first_term_idx];
4812        let standardized = standardized_spatial_term_data(data, prototype)?;
4813        let strategy = spatial_term_center_strategy(prototype).ok_or_else(|| {
4814            BasisError::InvalidInput(format!(
4815                "term '{}' lost its spatial center strategy during joint planning",
4816                prototype.name
4817            ))
4818        })?;
4819        let joint_strategy = center_strategy_with_num_centers(
4820            strategy,
4821            joint_centers,
4822            group_key.feature_cols.len(),
4823        )?;
4824        let shared_centers = select_centers_by_strategy(standardized.view(), &joint_strategy)?;
4825        log::info!(
4826            "sharing {} spatial centers across {} smooth terms over columns {:?} (requested {} centers)",
4827            shared_centers.nrows(),
4828            members.len(),
4829            group_key.feature_cols,
4830            group_key.requested_num_centers,
4831        );
4832        for (block_idx, term_idx) in members {
4833            set_spatial_term_centers(
4834                &mut planned_blocks[block_idx][term_idx],
4835                shared_centers.clone(),
4836            )?;
4837        }
4838    }
4839
4840    // Resolve typed Matérn Auto scales and the legacy thin-plate numeric
4841    // auto marker to a data-driven initialization here so REML starts in a
4842    // regime where it can escape. Matérn retains Auto provenance after this
4843    // numeric seed is installed, so it cannot later masquerade as user-fixed.
4844    for block in planned_blocks.iter_mut() {
4845        for term in block.iter_mut() {
4846            auto_init_length_scale_in_place(data, term);
4847        }
4848    }
4849
4850    Ok(planned_blocks)
4851}
4852
4853/// Tiny positive floor for the auto length scale, guarding against a zero
4854/// kernel range when every feature column is (near-)constant.
4855const AUTO_LENGTH_SCALE_FLOOR: f64 = 1e-6;
4856
4857/// Widest per-axis range of the selected feature columns. Returns `None` when
4858/// every selected column is constant / non-finite (no usable spatial scale).
4859fn feature_columns_max_range(data: ArrayView2<'_, f64>, feature_cols: &[usize]) -> Option<f64> {
4860    let mut max_range = 0.0_f64;
4861    for &c in feature_cols {
4862        if c >= data.ncols() {
4863            continue;
4864        }
4865        let col = data.column(c);
4866        let mut lo = f64::INFINITY;
4867        let mut hi = f64::NEG_INFINITY;
4868        for &v in col.iter() {
4869            if v.is_finite() {
4870                if v < lo {
4871                    lo = v;
4872                }
4873                if v > hi {
4874                    hi = v;
4875                }
4876            }
4877        }
4878        if hi > lo {
4879            let r = hi - lo;
4880            if r > max_range {
4881                max_range = r;
4882            }
4883        }
4884    }
4885    if max_range.is_finite() && max_range > 0.0 {
4886        Some(max_range)
4887    } else {
4888        None
4889    }
4890}
4891
4892/// Rotation-invariant analogue of [`feature_columns_max_range`], calibrated to
4893/// the span of the cloud's longest direction.
4894///
4895/// For a uniform interval of width `L`, the leading covariance eigenvalue is
4896/// `L²/12`, so `sqrt(12·λ_max)` recovers `L`. The same identity holds for the
4897/// longest side of an axis-aligned uniform box, while `λ_max` is invariant
4898/// under every orthogonal change of coordinates. This preserves the scale of
4899/// the former widest-axis seed without making it frame-dependent (gam#2252).
4900/// Sorting the complete points lexicographically makes each frame stable under
4901/// a pure row permutation (gam#1378).
4902fn feature_columns_rotation_invariant_range(
4903    data: ArrayView2<'_, f64>,
4904    feature_cols: &[usize],
4905) -> Option<f64> {
4906    let cols: Vec<usize> = feature_cols
4907        .iter()
4908        .copied()
4909        .filter(|&c| c < data.ncols())
4910        .collect();
4911    if cols.is_empty() {
4912        return None;
4913    }
4914    let mut points: Vec<Vec<f64>> = data
4915        .rows()
4916        .into_iter()
4917        .filter_map(|row| {
4918            let point: Vec<f64> = cols.iter().map(|&column| row[column]).collect();
4919            point.iter().all(|value| value.is_finite()).then_some(point)
4920        })
4921        .collect();
4922    if points.is_empty() {
4923        return None;
4924    }
4925    points.sort_by(|left, right| {
4926        left.iter()
4927            .zip(right)
4928            .find_map(|(a, b)| {
4929                let ordering = a.total_cmp(b);
4930                ordering.is_ne().then_some(ordering)
4931            })
4932            .unwrap_or(std::cmp::Ordering::Equal)
4933    });
4934
4935    let dimensions = cols.len();
4936    let count = points.len() as f64;
4937    let mut centroid = vec![0.0_f64; dimensions];
4938    for point in &points {
4939        for (coordinate, value) in centroid.iter_mut().zip(point) {
4940            *coordinate += *value;
4941        }
4942    }
4943    for coordinate in &mut centroid {
4944        *coordinate /= count;
4945    }
4946
4947    let mut covariance = Array2::<f64>::zeros((dimensions, dimensions));
4948    for point in &points {
4949        for row in 0..dimensions {
4950            let centered_row = point[row] - centroid[row];
4951            for column in 0..=row {
4952                covariance[[row, column]] += centered_row * (point[column] - centroid[column]);
4953            }
4954        }
4955    }
4956    for row in 0..dimensions {
4957        for column in 0..=row {
4958            let value = covariance[[row, column]] / count;
4959            covariance[[row, column]] = value;
4960            covariance[[column, row]] = value;
4961        }
4962    }
4963
4964    use gam_linalg::faer_ndarray::FaerEigh;
4965    let (eigenvalues, _) = covariance
4966        .eigh(faer::Side::Lower)
4967        .expect("finite covariance must have a symmetric eigendecomposition");
4968    let leading_variance = eigenvalues[eigenvalues.len() - 1];
4969    let extent = (12.0 * leading_variance).sqrt();
4970    if extent.is_finite() && extent > 0.0 {
4971        Some(extent)
4972    } else {
4973        None
4974    }
4975}
4976
4977/// Compute a data-driven initial length scale from the per-axis range of the
4978/// feature columns. The heuristic `max_range / sqrt(n)` puts the kernel on
4979/// the wiggly side of REML's basin so the optimizer can grow it back if the
4980/// signal is smooth, but is small enough that high-frequency truths remain
4981/// reachable for smoother kernels (ν ≥ 5/2). Clamped to a tiny positive
4982/// floor so degenerate constant-input columns can't produce 0.
4983pub fn auto_initial_length_scale(data: ArrayView2<'_, f64>, feature_cols: &[usize]) -> f64 {
4984    let n = data.nrows();
4985    if n == 0 || feature_cols.is_empty() {
4986        return 1.0;
4987    }
4988    let Some(max_range) = feature_columns_max_range(data, feature_cols) else {
4989        return 1.0;
4990    };
4991    let init = max_range / (n as f64).sqrt();
4992    init.max(AUTO_LENGTH_SCALE_FLOOR).min(max_range)
4993}
4994
4995/// Density-adaptive auto length scale for a kernel basis with `num_centers`
4996/// requested centers (#1731).
4997///
4998/// The plain [`auto_initial_length_scale`] seed `max_range / sqrt(n)` is the
4999/// fill distance of the *n data points*; it is independent of the requested
5000/// center count `k`. For a radial kernel at a FIXED length scale, packing more
5001/// centers into the same cloud makes neighbouring basis functions overlap and
5002/// go numerically collinear, so the realized basis saturates in rank (the
5003/// `matern_rank_reduce_centers` cap) and a richer `k` becomes a no-op — or even
5004/// shrinks the basis. The kernel stays well-conditioned only while the length
5005/// scale tracks the *center* spacing, not the data spacing.
5006///
5007/// We seed the length scale at the fill distance of `max(n, k)` points,
5008/// `max_range / sqrt(max(n, k))`. When `n ≥ k` (the usual case) this is exactly
5009/// the existing `max_range / sqrt(n)` seed, so every current result and small-`k`
5010/// basis size is preserved bit-for-bit (in every covariate dimension). When
5011/// `k > n` (a dense center request on a small cloud, the regime where an
5012/// `n`-sized seed sits above the center spacing and over-smooths the centers
5013/// into collinearity) the seed shrinks with `k` to the center spacing, keeping
5014/// the requested centers numerically independent. This is the Matérn analogue of
5015/// the Duchon-promotion "length_scale from center spacing" rule
5016/// (`hybrid_duchon_promotion_length_scale`).
5017pub fn auto_initial_length_scale_for_centers(
5018    data: ArrayView2<'_, f64>,
5019    feature_cols: &[usize],
5020    num_centers: usize,
5021) -> f64 {
5022    let n = data.nrows();
5023    if n == 0 || feature_cols.is_empty() {
5024        return 1.0;
5025    }
5026    // #2252: rotation-invariant extent for the Matérn seed so the enrolled κ/range
5027    // solve — which is basin-/seed-sensitive (see the matern geometry-stall path)
5028    // — starts from a frame-independent point and lands in the SAME basin in every
5029    // rotated frame, making the isotropic Matérn fit rotation-invariant. The
5030    // per-axis span (`feature_columns_max_range`) is a projection of the cloud and
5031    // is rotation-variant; the covariance spectral extent `sqrt(12·λ_max)` is
5032    // invariant under any orthogonal map and retains the former span calibration.
5033    // Duchon/thin-plate seeds are computed by separate helpers and are unchanged,
5034    // so those (seed-robust) bases stay bit-identical — this fix is scoped to the
5035    // seed-sensitive Matérn path.
5036    let Some(max_range) = feature_columns_rotation_invariant_range(data, feature_cols) else {
5037        return 1.0;
5038    };
5039    // Resolution density: at least the data points, but no coarser than the
5040    // center spacing once more centers than data are requested. Using the same
5041    // `sqrt` fill-distance law as `auto_initial_length_scale` keeps the seed
5042    // bit-identical whenever `n ≥ num_centers` (every dimension), and only
5043    // shrinks it — never grows it — when `num_centers > n`.
5044    let resolution_points = n.max(num_centers).max(1) as f64;
5045    let spacing = max_range / resolution_points.sqrt();
5046    spacing.max(AUTO_LENGTH_SCALE_FLOOR).min(max_range)
5047}
5048
5049/// Rotation-invariant center-resolution range for the companion Matérn basin.
5050///
5051/// A reduced-rank Matérn basis has two distinct geometric resolutions: the
5052/// observation fill distance used by the short/rich cold seed, and the coarser
5053/// fill distance of its `k` retained centers. The latter is the canonical
5054/// response-free representative of the overlapping, long-range basin:
5055/// `sqrt(12 * lambda_max(cov(x))) / sqrt(k)`. It uses the same covariance
5056/// extent as [`auto_initial_length_scale_for_centers`], so rigid rotations and
5057/// row permutations leave it unchanged.
5058pub fn matern_low_rank_center_resolution_length_scale(
5059    data: ArrayView2<'_, f64>,
5060    feature_cols: &[usize],
5061    num_centers: usize,
5062) -> Option<f64> {
5063    if data.nrows() == 0 || feature_cols.is_empty() || num_centers == 0 {
5064        return None;
5065    }
5066    let extent = feature_columns_rotation_invariant_range(data, feature_cols)?;
5067    let length_scale = extent / (num_centers as f64).sqrt();
5068    Some(length_scale.max(AUTO_LENGTH_SCALE_FLOOR).min(extent))
5069}
5070
5071/// Low-rank radial-basis length-scale seed tied to the requested center spacing.
5072///
5073/// Thin-plate regression splines with `k << n` represent the surface through a
5074/// compact set of centers; seeding the kernel at the observation fill distance
5075/// (`max_range / sqrt(n)`) makes the center Gram nearly diagonal and turns the
5076/// bending penalty into an ill-scaled ridge on the radial coefficients. REML then
5077/// sees a weakly identified smoothing surface and can settle on under-recovered
5078/// spatial fits. Seed at the center fill distance instead, so neighbouring
5079/// centers interact at O(1) scale before REML tunes the smoothing parameter.
5080pub fn auto_initial_length_scale_for_low_rank_centers(
5081    data: ArrayView2<'_, f64>,
5082    feature_cols: &[usize],
5083    num_centers: usize,
5084) -> f64 {
5085    if data.nrows() == 0 || feature_cols.is_empty() {
5086        return 1.0;
5087    }
5088    let Some(max_range) = feature_columns_max_range(data, feature_cols) else {
5089        return 1.0;
5090    };
5091    let resolution_points = num_centers.max(1) as f64;
5092    let spacing = max_range / resolution_points.sqrt();
5093    spacing.max(AUTO_LENGTH_SCALE_FLOOR).min(max_range)
5094}
5095
5096/// Requested center count encoded by a [`CenterStrategy`], if it carries an
5097/// explicit count (used to make the Matérn auto length scale density-adaptive).
5098fn center_strategy_requested_count(strategy: &CenterStrategy) -> Option<usize> {
5099    match strategy {
5100        CenterStrategy::Auto(inner) => center_strategy_requested_count(inner),
5101        CenterStrategy::UserProvided(centers) => Some(centers.nrows()),
5102        CenterStrategy::EqualMass { num_centers }
5103        | CenterStrategy::EqualMassCovarRepresentative { num_centers }
5104        | CenterStrategy::FarthestPoint { num_centers }
5105        | CenterStrategy::KMeans { num_centers, .. } => Some(*num_centers),
5106        CenterStrategy::UniformGrid { .. } => None,
5107    }
5108}
5109
5110/// Walk a term and resolve an omitted Matérn length scale, or a thin-plate
5111/// smooth still carrying its numeric auto marker, with
5112/// [`auto_initial_length_scale`]. Matérn's typed Auto provenance survives.
5113pub fn auto_init_length_scale_in_place(data: ArrayView2<'_, f64>, term: &mut SmoothTermSpec) {
5114    auto_init_length_scale_in_basis(data, &mut term.basis);
5115}
5116
5117/// Resolve the typed Matérn Auto length scale (and thin-plate's numeric auto
5118/// marker) with a data-derived value for any reachable kernel — including the
5119/// inner kernel of a `by=`/factor-smooth wrapper.
5120///
5121/// `by=<factor>` and the sum-to-zero factor smooth wrap a spatial kernel inside
5122/// `SmoothBasisSpec::ByVariable` / `SmoothBasisSpec::FactorSumToZero` /
5123/// `SmoothBasisSpec::BySmooth`, so the wrapper variant is what the planner sees.
5124/// Without recursing into the wrapped basis the inner Matérn remains unresolved
5125/// (and ThinPlate keeps its numeric marker), so no valid kernel scale exists at
5126/// fit or predict time. Recurse so the inner kernel is initialized identically
5127/// to a top-level one.
5128pub fn auto_init_length_scale_in_basis(data: ArrayView2<'_, f64>, basis: &mut SmoothBasisSpec) {
5129    match basis {
5130        SmoothBasisSpec::Matern {
5131            feature_cols, spec, ..
5132        } => {
5133            if spec.length_scale.resolved().is_none() {
5134                // Density-adaptive seed (#1731): when the requested center count
5135                // is known, scale the auto length scale with the *center*
5136                // spacing so a richer `k` stays numerically full-rank instead of
5137                // saturating against `matern_rank_reduce_centers`. For `n ≥ k`
5138                // (the usual case) this is identical to the plain `max_range /
5139                // sqrt(n)` seed in 2-D, so small-`k` results are unchanged. The
5140                // unconstrained / non-explicit `UniformGrid` strategy falls back
5141                // to the plain seed.
5142                let resolved = match center_strategy_requested_count(&spec.center_strategy) {
5143                    Some(k) => auto_initial_length_scale_for_centers(data, feature_cols, k),
5144                    None => auto_initial_length_scale(data, feature_cols),
5145                };
5146                spec.length_scale.resolve_auto_once(resolved);
5147            }
5148        }
5149        SmoothBasisSpec::ThinPlate {
5150            feature_cols, spec, ..
5151        } => {
5152            if spec.length_scale == 0.0 {
5153                spec.length_scale = match center_strategy_requested_count(&spec.center_strategy) {
5154                    Some(k) => {
5155                        auto_initial_length_scale_for_low_rank_centers(data, feature_cols, k)
5156                    }
5157                    None => auto_initial_length_scale(data, feature_cols),
5158                };
5159            }
5160        }
5161        SmoothBasisSpec::ByVariable { inner, .. }
5162        | SmoothBasisSpec::FactorSumToZero { inner, .. } => {
5163            auto_init_length_scale_in_basis(data, inner);
5164        }
5165        SmoothBasisSpec::BySmooth { smooth, .. } => {
5166            auto_init_length_scale_in_basis(data, smooth);
5167        }
5168        _ => {}
5169    }
5170}
5171
5172impl LinearFitConditioning {
5173    pub fn from_columns(design: &TermCollectionDesign, selected_cols: &[usize]) -> Self {
5174        const SCALE_EPS: f64 = 1e-12;
5175        let n = design.design.nrows();
5176        let p = design.design.ncols();
5177        let mut columns = Vec::with_capacity(selected_cols.len());
5178        if n == 0 || selected_cols.is_empty() {
5179            return Self {
5180                intercept_idx: design.intercept_range.start,
5181                columns,
5182            };
5183        }
5184        let chunk_rows = gam_linalg::utils::row_chunk_for_byte_budget(n, p);
5185        // Two-pass mean/variance so operator-backed designs don't need to
5186        // materialize the full dense matrix. Pass 1 accumulates per-column
5187        // sums; pass 2 accumulates the sum of squared deviations from the
5188        // pass-1 mean. This matches the original `Σ (x − mean)² / n` formula
5189        // without the catastrophic cancellation of `E[X²] − E[X]²`.
5190        let mut sums = vec![0.0_f64; selected_cols.len()];
5191        for start in (0..n).step_by(chunk_rows) {
5192            let end = (start + chunk_rows).min(n);
5193            let chunk = design
5194                .design
5195                .try_row_chunk(start..end)
5196                .expect("LinearFitConditioning::from_columns row chunk failed");
5197            for (k, &col_idx) in selected_cols.iter().enumerate() {
5198                let column = chunk.column(col_idx);
5199                for &v in column.iter() {
5200                    sums[k] += v;
5201                }
5202            }
5203        }
5204        let inv_n = 1.0_f64 / n as f64;
5205        let means: Vec<f64> = sums.iter().map(|&s| s * inv_n).collect();
5206        let mut sq_devs = vec![0.0_f64; selected_cols.len()];
5207        for start in (0..n).step_by(chunk_rows) {
5208            let end = (start + chunk_rows).min(n);
5209            let chunk = design
5210                .design
5211                .try_row_chunk(start..end)
5212                .expect("LinearFitConditioning::from_columns row chunk failed");
5213            for (k, &col_idx) in selected_cols.iter().enumerate() {
5214                let mean_k = means[k];
5215                let column = chunk.column(col_idx);
5216                for &v in column.iter() {
5217                    let d = v - mean_k;
5218                    sq_devs[k] += d * d;
5219                }
5220            }
5221        }
5222        for (k, &col_idx) in selected_cols.iter().enumerate() {
5223            let mean = means[k];
5224            let var = sq_devs[k] * inv_n;
5225            let (mean, scale) = if var.is_finite() && var > SCALE_EPS * SCALE_EPS {
5226                (mean, var.sqrt())
5227            } else {
5228                // Leave nearly-constant columns untouched; centering them would collapse
5229                // the design column to ~0 and change the model rather than just condition it.
5230                (0.0, 1.0)
5231            };
5232            columns.push(LinearColumnConditioning {
5233                col_idx,
5234                mean,
5235                scale,
5236            });
5237        }
5238        Self {
5239            intercept_idx: design.intercept_range.start,
5240            columns,
5241        }
5242    }
5243
5244    pub fn apply_to_design(&self, design: &Array2<f64>) -> Array2<f64> {
5245        let mut out = design.clone();
5246        for col in &self.columns {
5247            {
5248                let mut dst = out.column_mut(col.col_idx);
5249                dst -= col.mean;
5250            }
5251            if col.scale != 1.0 {
5252                out.column_mut(col.col_idx).mapv_inplace(|v| v / col.scale);
5253            }
5254        }
5255        out
5256    }
5257
5258    fn transform_matrix_columnswith_a(&self, mat: &Array2<f64>) -> Array2<f64> {
5259        let mut out = mat.clone();
5260        let intercept = self.intercept_idx;
5261        for col in &self.columns {
5262            let intercept_col = out.column(intercept).to_owned();
5263            let mut target = out.column_mut(col.col_idx);
5264            target -= &(intercept_col * col.mean);
5265            if col.scale != 1.0 {
5266                target.mapv_inplace(|v| v / col.scale);
5267            }
5268        }
5269        out
5270    }
5271
5272    fn transform_matrixrowswith_a_transpose(&self, mat: &Array2<f64>) -> Array2<f64> {
5273        let mut out = mat.clone();
5274        let intercept = self.intercept_idx;
5275        for col in &self.columns {
5276            let interceptrow = out.row(intercept).to_owned();
5277            let mut target = out.row_mut(col.col_idx);
5278            target -= &(interceptrow * col.mean);
5279            if col.scale != 1.0 {
5280                target.mapv_inplace(|v| v / col.scale);
5281            }
5282        }
5283        out
5284    }
5285
5286    /// Left-multiply `mat_internal` by `M⁻ᵀ` where `M⁻¹[intercept, j] = mean_j`
5287    /// and `M⁻¹[j, j] = scale_j` for each conditioned column. Used together
5288    /// with [`Self::right_multiply_by_m_inv`] to back-transform an internal
5289    /// penalized Hessian to the original coefficient basis.
5290    fn left_multiply_by_m_inv_transpose(&self, mat_internal: &Array2<f64>) -> Array2<f64> {
5291        let mut out = mat_internal.clone();
5292        let intercept = self.intercept_idx;
5293        let interceptrow_snapshot = mat_internal.row(intercept).to_owned();
5294        for col in &self.columns {
5295            if col.scale != 1.0 {
5296                out.row_mut(col.col_idx).mapv_inplace(|v| v * col.scale);
5297            }
5298            if col.mean != 0.0 {
5299                let mut target = out.row_mut(col.col_idx);
5300                target += &(&interceptrow_snapshot * col.mean);
5301            }
5302        }
5303        out
5304    }
5305
5306    /// Right-multiply `mat_internal` by `M⁻¹`. Mirror of
5307    /// [`Self::left_multiply_by_m_inv_transpose`] on columns.
5308    fn right_multiply_by_m_inv(&self, mat_internal: &Array2<f64>) -> Array2<f64> {
5309        let mut out = mat_internal.clone();
5310        let intercept = self.intercept_idx;
5311        let intercept_col_snapshot = mat_internal.column(intercept).to_owned();
5312        for col in &self.columns {
5313            if col.scale != 1.0 {
5314                out.column_mut(col.col_idx).mapv_inplace(|v| v * col.scale);
5315            }
5316            if col.mean != 0.0 {
5317                let mut target = out.column_mut(col.col_idx);
5318                target += &(&intercept_col_snapshot * col.mean);
5319            }
5320        }
5321        out
5322    }
5323
5324    /// Transform blockwise penalties through the conditioning.
5325    ///
5326    /// For block-local penalties whose `col_range` does not overlap with any
5327    /// conditioning column, the transform is identity (the conditioning only
5328    /// affects unpenalized linear columns). In that common case the penalty
5329    /// passes through unchanged, avoiding O(p²) materialization entirely.
5330    pub fn transform_blockwise_penalties_to_internal(
5331        &self,
5332        penalties: &[BlockwisePenalty],
5333        p: usize,
5334    ) -> Vec<crate::penalty_spec::PenaltySpec> {
5335        let conditioning_cols: std::collections::HashSet<usize> =
5336            self.columns.iter().map(|c| c.col_idx).collect();
5337        penalties
5338            .iter()
5339            .map(|bp| {
5340                let overlaps =
5341                    (bp.col_range.start..bp.col_range.end).any(|j| conditioning_cols.contains(&j));
5342                if overlaps {
5343                    // Rare: penalty block overlaps conditioning columns.
5344                    // Fall back to dense transform.
5345                    let global = bp.to_global(p);
5346                    let right = self.transform_matrix_columnswith_a(&global);
5347                    let transformed = self.transform_matrixrowswith_a_transpose(&right);
5348                    crate::penalty_spec::PenaltySpec::Dense(transformed)
5349                } else {
5350                    // Common: smooth penalty block doesn't touch linear columns.
5351                    // The conditioning is identity on this block.
5352                    crate::penalty_spec::PenaltySpec::from_blockwise(bp.clone())
5353                }
5354            })
5355            .collect()
5356    }
5357
5358    pub fn backtransform_beta(&self, beta_internal: &Array1<f64>) -> Array1<f64> {
5359        let mut beta = beta_internal.clone();
5360        let intercept = self.intercept_idx;
5361        for col in &self.columns {
5362            beta[intercept] -= beta_internal[col.col_idx] * col.mean / col.scale;
5363            beta[col.col_idx] = beta_internal[col.col_idx] / col.scale;
5364        }
5365        beta
5366    }
5367
5368    /// `H_orig = M⁻ᵀ · H_int · M⁻¹`, derived from
5369    /// `L_int(β_int) = L_orig(M · β_int)` via the chain rule.
5370    pub fn transform_penalized_hessian_to_original(&self, h_internal: &Array2<f64>) -> Array2<f64> {
5371        let right = self.right_multiply_by_m_inv(h_internal);
5372        self.left_multiply_by_m_inv_transpose(&right)
5373    }
5374
5375    pub fn internal_bounds_for(&self, col_idx: usize, min: f64, max: f64) -> (f64, f64) {
5376        if let Some(col) = self.columns.iter().find(|c| c.col_idx == col_idx) {
5377            (min * col.scale, max * col.scale)
5378        } else {
5379            (min, max)
5380        }
5381    }
5382}
5383
5384pub fn freeze_raw_spatial_metadata(metadata: BasisMetadata, raw_cols: usize) -> BasisMetadata {
5385    match metadata {
5386        BasisMetadata::ThinPlate {
5387            centers,
5388            length_scale,
5389            periodic,
5390            identifiability_transform: None,
5391            input_scale,
5392            radial_reparam,
5393        } => BasisMetadata::ThinPlate {
5394            centers,
5395            length_scale,
5396            periodic,
5397            identifiability_transform: Some(Array2::eye(raw_cols)),
5398            input_scale,
5399            radial_reparam,
5400        },
5401        BasisMetadata::Duchon {
5402            centers,
5403            length_scale,
5404            periodic,
5405            power,
5406            nullspace_order,
5407            identifiability_transform: None,
5408            input_scale,
5409            aniso_log_scales,
5410            operator_collocation_points,
5411            radial_reparam,
5412        } => BasisMetadata::Duchon {
5413            centers,
5414            length_scale,
5415            periodic,
5416            power,
5417            nullspace_order,
5418            identifiability_transform: Some(Array2::eye(raw_cols)),
5419            input_scale,
5420            aniso_log_scales,
5421            operator_collocation_points,
5422            radial_reparam,
5423        },
5424        other => other,
5425    }
5426}
5427
5428pub fn matern_operator_penalty_triplet_from_metadata(
5429    metadata: &BasisMetadata,
5430) -> Result<crate::basis::FilteredPenalties, BasisError> {
5431    let BasisMetadata::Matern {
5432        centers,
5433        length_scale,
5434        periodic,
5435        nu,
5436        include_intercept,
5437        identifiability_transform,
5438        aniso_log_scales,
5439        input_scale,
5440        ..
5441    } = metadata
5442    else {
5443        crate::bail_invalid_basis!("Matérn operator penalties require Matérn metadata");
5444    };
5445    // The metadata records `length_scale` in *original* (un-standardized) data
5446    // coordinates, while `centers` live in the *standardized* coordinate frame
5447    // (uniform division by `input_scale`). The realized design built the
5448    // kernel against those standardized centers using the compensated
5449    // effective length scale `length_scale / input_scale`. The collocation operators
5450    // here are evaluated on the same standardized centers, so they must use the
5451    // SAME effective length scale — otherwise the penalty regularizes a
5452    // different RKHS range than the design lives in, leaving rough coefficient
5453    // directions effectively unpenalized. That mismatch is benign in 1-D
5454    // (no standardization) but produces a catastrophic out-of-sample blow-up in
5455    // every dimension where the input scale differs from one (#706).
5456    let penalty_length_scale = input_scale.to_standardized_units(*length_scale);
5457    matern_operator_penalty_triplet_at_length_scale(
5458        centers.view(),
5459        periodic.as_deref(),
5460        identifiability_transform.as_ref(),
5461        *nu,
5462        *include_intercept,
5463        aniso_log_scales.as_deref(),
5464        penalty_length_scale,
5465    )
5466}
5467
5468/// Build the canonical Matérn operator-penalty triplet (mass / tension /
5469/// stiffness) at an explicit **effective** length scale — i.e. the
5470/// isotropic-scale-compensated, standardized-frame scale the design's kernel was built
5471/// against (NOT the original-coordinate `length_scale` stored in metadata).
5472///
5473/// This is the SINGLE source of truth for the Matérn penalty topology. Two
5474/// callers route through it and must therefore stay byte-for-byte consistent:
5475///   * the cold/slow design rebuild (`matern_operator_penalty_triplet_from_metadata`,
5476///     compensating the frozen metadata `length_scale`), and
5477///   * the n-free κ-optimizer re-key (`FrozenTermCollectionIncrementalRealizer::
5478///     canonical_penalties_at_psi`, compensating the trial `ψ → exp(-ψ)` scale).
5479///
5480/// Sharing the body makes the penalty BLOCK COUNT and the per-block numerics
5481/// one deterministic function of `(geometry, ν, η, ℓ_eff)`. The active-operator
5482/// gate is `m = ν + d/2`, which is independent of ℓ, so the block count is
5483/// **ψ-stable by construction**: the re-key can never produce a different number
5484/// of blocks than the frozen design (the desync that #1270 hard-errored on).
5485pub fn matern_operator_penalty_triplet_at_length_scale(
5486    centers: ArrayView2<'_, f64>,
5487    periodic: Option<&[Option<f64>]>,
5488    identifiability_transform: Option<&Array2<f64>>,
5489    nu: crate::basis::MaternNu,
5490    include_intercept: bool,
5491    aniso_log_scales: Option<&[f64]>,
5492    effective_length_scale: f64,
5493) -> Result<crate::basis::FilteredPenalties, BasisError> {
5494    let penalty_centers = crate::basis::expand_periodic_centers(&centers.to_owned(), periodic)?;
5495    let ops = build_matern_collocation_operator_matrices(
5496        penalty_centers.view(),
5497        None,
5498        effective_length_scale,
5499        nu,
5500        include_intercept,
5501        identifiability_transform.map(|z| z.view()),
5502        aniso_log_scales,
5503    )?;
5504    // Gate operator dials on the Matérn-ν RKHS Sobolev order m = ν + d/2.
5505    // Derivative energies through j=m belong to H^m inclusively, so the 1-D
5506    // ν=3/2 kernel (m=2) carries stiffness as well as mass+tension. The sole
5507    // exception is ν=1/2: its center cusp makes collocated D1/D2 undefined and
5508    // it therefore retains mass only (#707). The matching topology gate lives
5509    // at `DuchonOperatorPenaltySpec::matern_for_smoothness`.
5510    const ORDER_EPS: f64 = 1e-9;
5511    let d = penalty_centers.ncols();
5512    let m = nu.half_integer_value() + 0.5 * d as f64;
5513    let mut candidates = Vec::with_capacity(3);
5514    for (raw, source, min_order) in [
5515        (ops.d0.t().dot(&ops.d0), PenaltySource::OperatorMass, 0.0),
5516        (ops.d1.t().dot(&ops.d1), PenaltySource::OperatorTension, 1.0),
5517        (
5518            ops.d2.t().dot(&ops.d2),
5519            PenaltySource::OperatorStiffness,
5520            2.0,
5521        ),
5522    ] {
5523        let nondifferentiable_ou = matches!(nu, crate::basis::MaternNu::Half);
5524        if min_order > 0.0 && (nondifferentiable_ou || m + ORDER_EPS < min_order) {
5525            continue;
5526        }
5527        let sym = (&raw + &raw.t()) * 0.5;
5528        let (matrix, normalization_scale) = normalize_penalty_in_constrained_space(&sym);
5529        candidates.push(PenaltyCandidate {
5530            matrix: ConstructiveQuadratic::try_from_dense_psd(
5531                matrix,
5532                "Matérn operator penalty",
5533            )?,
5534            source,
5535            normalization_scale,
5536            kronecker_factors: None,
5537            op: None,
5538        });
5539    }
5540    filter_penalty_candidates(candidates)
5541}
5542
5543pub fn normalize_penalty_in_constrained_space(matrix: &Array2<f64>) -> (Array2<f64>, f64) {
5544    // Constrained-space normalization:
5545    //   c = ||S_con||_F,  S_tilde = S_con / c.
5546    // This is the only normalization coherent with a REML objective that is
5547    // evaluated entirely in constrained coordinates.
5548    let matrix = (matrix + &matrix.t().to_owned()) * 0.5;
5549    // Clamp noise-floor negative eigenvalues so β'Sβ is non-negative as a contract, not just in exact arithmetic.
5550    let matrix = crate::basis::project_penalty_to_psd_cone(&matrix);
5551    let c = matrix.iter().map(|v| v * v).sum::<f64>().sqrt();
5552    if c.is_finite() && c > 0.0 {
5553        (matrix.mapv(|v| v / c), c)
5554    } else {
5555        (matrix, 1.0)
5556    }
5557}
5558
5559pub fn tensor_product_design_from_sparse_marginals(
5560    marginal_sparse: &[&SparseColMat<usize, f64>],
5561) -> Result<SparseColMat<usize, f64>, BasisError> {
5562    if marginal_sparse.is_empty() {
5563        crate::bail_invalid_basis!("TensorBSpline requires at least one marginal basis");
5564    }
5565    let n = marginal_sparse[0].nrows();
5566    for (i, m) in marginal_sparse.iter().enumerate().skip(1) {
5567        if m.nrows() != n {
5568            crate::bail_dim_basis!(
5569                "tensor sparse marginal row mismatch at dim {i}: expected {n}, got {}",
5570                m.nrows()
5571            );
5572        }
5573    }
5574    let dims: Vec<usize> = marginal_sparse.iter().map(|m| m.ncols()).collect();
5575    let total_cols = dims.iter().try_fold(1usize, |acc, &q| {
5576        acc.checked_mul(q)
5577            .ok_or_else(|| BasisError::DimensionMismatch("tensor basis too large".to_string()))
5578    })?;
5579    let mut strides = vec![1usize; dims.len()];
5580    for d in (0..dims.len().saturating_sub(1)).rev() {
5581        strides[d] = strides[d + 1]
5582            .checked_mul(dims[d + 1])
5583            .ok_or_else(|| BasisError::DimensionMismatch("tensor basis too large".to_string()))?;
5584    }
5585
5586    use faer::sparse::SparseRowMat;
5587    let csrs: Vec<SparseRowMat<usize, f64>> = marginal_sparse
5588        .iter()
5589        .enumerate()
5590        .map(|(d, m)| {
5591            m.as_ref().to_row_major().map_err(|e| {
5592                BasisError::SparseCreation(format!(
5593                    "tensor sparse marginal {d} CSR conversion failed: {e:?}"
5594                ))
5595            })
5596        })
5597        .collect::<Result<Vec<_>, _>>()?;
5598    let row_ptrs: Vec<&[usize]> = csrs.iter().map(|c| c.symbolic().row_ptr()).collect();
5599    let col_idxs: Vec<&[usize]> = csrs.iter().map(|c| c.symbolic().col_idx()).collect();
5600    let vals: Vec<&[f64]> = csrs.iter().map(|c| c.val()).collect();
5601
5602    use rayon::prelude::*;
5603    const CHUNK: usize = 1024;
5604    let num_chunks = n.div_ceil(CHUNK);
5605    let per_chunk: Vec<Vec<Triplet<usize, usize, f64>>> = (0..num_chunks)
5606        .into_par_iter()
5607        .map(|chunk_idx| {
5608            let row_start = chunk_idx * CHUNK;
5609            let row_end = (row_start + CHUNK).min(n);
5610            let mut chunk_triplets = Vec::<Triplet<usize, usize, f64>>::new();
5611            let mut cur_cols = Vec::<usize>::with_capacity(64);
5612            let mut cur_vals = Vec::<f64>::with_capacity(64);
5613            let mut next_cols = Vec::<usize>::with_capacity(64);
5614            let mut next_vals = Vec::<f64>::with_capacity(64);
5615            for i in row_start..row_end {
5616                cur_cols.clear();
5617                cur_vals.clear();
5618                cur_cols.push(0);
5619                cur_vals.push(1.0);
5620                let mut row_is_zero = false;
5621                for d in 0..dims.len() {
5622                    let row_start_d = row_ptrs[d][i];
5623                    let row_end_d = row_ptrs[d][i + 1];
5624                    if row_start_d == row_end_d {
5625                        row_is_zero = true;
5626                        break;
5627                    }
5628                    let stride = strides[d];
5629                    next_cols.clear();
5630                    next_vals.clear();
5631                    next_cols.reserve(cur_cols.len() * (row_end_d - row_start_d));
5632                    next_vals.reserve(cur_vals.len() * (row_end_d - row_start_d));
5633                    for (&prev_col, &prev_val) in cur_cols.iter().zip(cur_vals.iter()) {
5634                        for ptr in row_start_d..row_end_d {
5635                            let cj = col_idxs[d][ptr];
5636                            let vj = vals[d][ptr];
5637                            next_cols.push(prev_col + cj * stride);
5638                            next_vals.push(prev_val * vj);
5639                        }
5640                    }
5641                    std::mem::swap(&mut cur_cols, &mut next_cols);
5642                    std::mem::swap(&mut cur_vals, &mut next_vals);
5643                }
5644                if row_is_zero {
5645                    continue;
5646                }
5647                for (&col, &val) in cur_cols.iter().zip(cur_vals.iter()) {
5648                    chunk_triplets.push(Triplet::new(i, col, val));
5649                }
5650            }
5651            chunk_triplets
5652        })
5653        .collect();
5654    let total_nnz: usize = per_chunk.iter().map(Vec::len).sum();
5655    let mut triplets = Vec::<Triplet<usize, usize, f64>>::with_capacity(total_nnz);
5656    for chunk in per_chunk {
5657        triplets.extend(chunk);
5658    }
5659    SparseColMat::try_new_from_triplets(n, total_cols, &triplets).map_err(|e| {
5660        BasisError::SparseCreation(format!(
5661            "failed to assemble sparse tensor product design: {e:?}"
5662        ))
5663    })
5664}
5665
5666pub fn dense_local_margin_to_sparse(
5667    dense: &Array2<f64>,
5668) -> Result<SparseColMat<usize, f64>, BasisError> {
5669    let expected_row_nnz = dense.ncols().min(4);
5670    let mut triplets =
5671        Vec::<Triplet<usize, usize, f64>>::with_capacity(dense.nrows() * expected_row_nnz);
5672    for ((row, col), &value) in dense.indexed_iter() {
5673        if value != 0.0 {
5674            triplets.push(Triplet::new(row, col, value));
5675        }
5676    }
5677    SparseColMat::try_new_from_triplets(dense.nrows(), dense.ncols(), &triplets).map_err(|e| {
5678        BasisError::SparseCreation(format!(
5679            "failed to convert tensor marginal design to sparse form: {e:?}"
5680        ))
5681    })
5682}
5683
5684pub struct TensorMarginRangeNullProjectors {
5685    range: Array2<f64>,
5686    null: Array2<f64>,
5687}
5688
5689pub fn projector_from_columns(columns: &Array2<f64>, indices: &[usize]) -> Array2<f64> {
5690    if indices.is_empty() {
5691        return Array2::<f64>::zeros((columns.nrows(), columns.nrows()));
5692    }
5693    let basis = columns.select(Axis(1), indices);
5694    basis.dot(&basis.t())
5695}
5696
5697pub fn tensor_margin_range_null_projectors(
5698    normalized_marginal_penalties: &[(Array2<f64>, f64)],
5699) -> Result<Vec<TensorMarginRangeNullProjectors>, BasisError> {
5700    normalized_marginal_penalties
5701        .iter()
5702        .enumerate()
5703        .map(|(dim, (penalty, _))| {
5704            let analysis = crate::basis::analyze_penalty_block(penalty)?;
5705            if analysis.rank == 0 {
5706                crate::bail_invalid_basis!(
5707                    "t2 separable tensor penalty margin {dim} has rank-zero penalty; \
5708                     cannot split penalized and null subspaces"
5709                );
5710            }
5711            let mut range_idx = Vec::<usize>::new();
5712            let mut null_idx = Vec::<usize>::new();
5713            for (idx, &ev) in analysis.eigenvalues.iter().enumerate() {
5714                if ev > analysis.tol {
5715                    range_idx.push(idx);
5716                } else {
5717                    null_idx.push(idx);
5718                }
5719            }
5720            Ok(TensorMarginRangeNullProjectors {
5721                range: projector_from_columns(&analysis.eigenvectors, &range_idx),
5722                null: projector_from_columns(&analysis.eigenvectors, &null_idx),
5723            })
5724        })
5725        .collect()
5726}
5727
5728pub fn build_tensor_bspline_basis(
5729    data: ArrayView2<'_, f64>,
5730    feature_cols: &[usize],
5731    spec: &TensorBSplineSpec,
5732) -> Result<BasisBuildResult, BasisError> {
5733    if feature_cols.is_empty() {
5734        crate::bail_invalid_basis!("TensorBSpline requires at least one feature column");
5735    }
5736    if feature_cols.len() != spec.marginalspecs.len() {
5737        crate::bail_dim_basis!(
5738            "TensorBSpline feature/spec mismatch: feature_cols={}, marginalspecs={}",
5739            feature_cols.len(),
5740            spec.marginalspecs.len()
5741        );
5742    }
5743    if let Some((margin, _)) = spec
5744        .marginalspecs
5745        .iter()
5746        .enumerate()
5747        .find(|(_, marginal)| marginal.boundary_conditions.has_nonzero_anchor())
5748    {
5749        crate::bail_invalid_basis!(
5750            "TensorBSpline margin {margin} has a non-zero endpoint anchor. An inhomogeneous \
5751             marginal constraint cannot be represented by the tensor's homogeneous coefficient \
5752             chart plus one scalar row offset; use a separate anchored 1-D smooth or an explicit \
5753             model offset"
5754        );
5755    }
5756    if !spec.periods.is_empty() && spec.periods.len() != feature_cols.len() {
5757        crate::bail_dim_basis!(
5758            "TensorBSpline periods length {} does not match feature count {}",
5759            spec.periods.len(),
5760            feature_cols.len()
5761        );
5762    }
5763    let p = data.ncols();
5764    for &c in feature_cols {
5765        if c >= p {
5766            crate::bail_dim_basis!(
5767                "tensor feature column {c} is out of bounds for data with {p} columns"
5768            );
5769        }
5770    }
5771
5772    let mut marginal_knots = Vec::<Array1<f64>>::with_capacity(feature_cols.len());
5773    // Per-margin cr flag (#1074): `true` when the margin is a natural cubic
5774    // regression spline, so the tensor freeze rebuilds the cr knotspec.
5775    let mut marginal_is_cr_flags = Vec::<bool>::with_capacity(feature_cols.len());
5776    let mut marginal_degrees = Vec::<usize>::with_capacity(feature_cols.len());
5777    let mut marginalnum_basis = Vec::<usize>::with_capacity(feature_cols.len());
5778    let mut marginal_penalties = Vec::<Array2<f64>>::with_capacity(feature_cols.len());
5779    let mut marginal_function_grams = Vec::<Array2<f64>>::with_capacity(feature_cols.len());
5780    let mut marginal_designs = Vec::<Array2<f64>>::with_capacity(feature_cols.len());
5781    // Per-margin effective period: either user-set via `spec.periods` or
5782    // implied by a `PeriodicUniform` marginal knotspec (which the 1D B-spline
5783    // builder realizes as a cyclic B-spline basis).
5784    // Captured here so freeze→reload round-trips both routes back to a
5785    // `PeriodicUniform` marginal knotspec; otherwise a `PeriodicUniform`
5786    // margin specified without `spec.periods` would freeze as a plain
5787    // `Provided(knots)` open spline and lose its wrap-around at predict time.
5788    let mut marginal_effective_periods = Vec::<Option<f64>>::with_capacity(feature_cols.len());
5789    // Per-marginal sparse representation, populated when the 1D builder returned
5790    // a `DesignMatrix::Sparse`. Used to assemble the Khatri-Rao tensor product
5791    // sparsely (only ∏(degree+1) nonzeros per row) instead of densifying to
5792    // shape (n, ∏ q_j) up front. Periodic B-spline margins are local-support
5793    // bases too; when the 1D builder returns them densely, we convert that
5794    // marginal back to sparse form so cylinder/torus tensor products keep the
5795    // same scale behavior as open tensor products.
5796    let mut marginal_sparse =
5797        Vec::<Option<SparseColMat<usize, f64>>>::with_capacity(feature_cols.len());
5798
5799    // Reuse the robust 1D builder to ensure the same knot validation and
5800    // marginal difference-penalty construction as standalone smooth terms.
5801    for (dim, (&col, marginalspec)) in feature_cols
5802        .iter()
5803        .zip(spec.marginalspecs.iter())
5804        .enumerate()
5805    {
5806        // Tensor basis uses raw marginal knot-product columns. Applying 1D
5807        // identifiability constraints here would change marginal penalty sizes
5808        // without changing the tensor design construction, causing dimension
5809        // mismatch. Keep marginal builders unconstrained at this stage.
5810        let mut marginal_unconstrained = marginalspec.clone();
5811        marginal_unconstrained.identifiability = BSplineIdentifiability::None;
5812        let built = build_bspline_basis_1d(data.column(col), &marginal_unconstrained)?;
5813        // A cr (`NaturalCubicRegression`) margin emits `CubicRegression1D`
5814        // metadata whose `knots` are the k value-knots; a B-spline margin emits
5815        // `BSpline1D` with the clamped knot vector. Capture either so the
5816        // tensor freeze can rebuild the exact same marginal knotspec (#1074).
5817        let (knots, marginal_is_cr, effective_degree, function_gram) = match built.metadata {
5818            BasisMetadata::BSpline1D {
5819                knots,
5820                periodic,
5821                degree,
5822                ..
5823            } => {
5824                let effective_degree = degree.unwrap_or(marginal_unconstrained.degree);
5825                let gram = if spec.double_penalty {
5826                    Some(match periodic {
5827                        Some((start, period, num_basis)) => {
5828                            crate::basis::periodic_bspline_function_gram(
5829                                start,
5830                                start + period,
5831                                effective_degree,
5832                                num_basis,
5833                            )?
5834                        }
5835                        None => crate::basis::bspline_function_gram(&knots, effective_degree)?,
5836                    })
5837                } else {
5838                    None
5839                };
5840                (knots, false, effective_degree, gram)
5841            }
5842            BasisMetadata::CubicRegression1D { knots, .. } => {
5843                let gram = spec
5844                    .double_penalty
5845                    .then(|| crate::basis::cubic_regression_function_gram(&knots))
5846                    .transpose()?;
5847                (knots, true, marginalspec.degree, gram)
5848            }
5849            _ => {
5850                crate::bail_invalid_basis!(
5851                    "internal TensorBSpline error at dim {dim}: expected BSpline1D or CubicRegression1D metadata"
5852                );
5853            }
5854        };
5855        let metadata_knots = match marginalspec.knotspec {
5856            BSplineKnotSpec::PeriodicUniform {
5857                data_range,
5858                num_basis,
5859            } => Array1::linspace(data_range.0, data_range.1, num_basis),
5860            _ => knots,
5861        };
5862        if let Some(function_gram) = function_gram {
5863            if function_gram.dim() != (built.design.ncols(), built.design.ncols()) {
5864                crate::bail_dim_basis!(
5865                    "internal TensorBSpline error at dim {dim}: function Gram is {:?}, basis has {} columns",
5866                    function_gram.dim(),
5867                    built.design.ncols()
5868                );
5869            }
5870            marginal_function_grams.push(function_gram);
5871        }
5872        marginal_knots.push(metadata_knots);
5873        marginal_is_cr_flags.push(marginal_is_cr);
5874        marginal_degrees.push(effective_degree);
5875        marginalnum_basis.push(built.design.ncols());
5876        // Capture the sparse representation of this marginal (when the
5877        // 1D builder produced one) before densifying for the dense
5878        // marginal cache used by `tensor_product_design_from_marginals`
5879        // and `TensorProductDesignOperator`.
5880        let dense_marginal = built.design.to_dense();
5881        let sparse_view: Option<SparseColMat<usize, f64>> = match built.design.as_sparse() {
5882            Some(sd) => {
5883                let inner: &SparseColMat<usize, f64> = sd;
5884                Some(inner.clone())
5885            }
5886            None => match marginalspec.knotspec {
5887                BSplineKnotSpec::PeriodicUniform { .. } => {
5888                    Some(dense_local_margin_to_sparse(&dense_marginal)?)
5889                }
5890                _ => None,
5891            },
5892        };
5893        marginal_sparse.push(sparse_view);
5894        marginal_designs.push(dense_marginal);
5895        marginal_penalties.push(
5896            built
5897                .active_penalties
5898                .first()
5899                .ok_or_else(|| {
5900                    BasisError::InvalidInput(format!(
5901                        "internal TensorBSpline error at dim {dim}: missing marginal penalty"
5902                    ))
5903                })?
5904                .matrix
5905                .clone(),
5906        );
5907        built.active_penalties.first().ok_or_else(|| {
5908            BasisError::InvalidInput(format!(
5909                "internal TensorBSpline error at dim {dim}: missing marginal nullspace dim"
5910            ))
5911        })?;
5912        // A `PeriodicUniform` marginal knotspec implies the margin is
5913        // wrap-around: the 1D builder already realized it as a periodic
5914        // basis, so the tensor product inherits that periodicity. Record
5915        // the period derived from the knotspec's data range so freeze
5916        // restores `PeriodicUniform` on the marginal — otherwise the
5917        // round-trip downgrades it to `Provided(knots)` (an open spline)
5918        // and predict-time wraps disappear.
5919        let implied_period = match marginalspec.knotspec {
5920            BSplineKnotSpec::PeriodicUniform { data_range, .. } => {
5921                Some(data_range.1 - data_range.0)
5922            }
5923            _ => spec.periods.get(dim).and_then(|p| *p),
5924        };
5925        marginal_effective_periods.push(implied_period);
5926    }
5927
5928    let total_cols: usize = marginalnum_basis.iter().product();
5929    let mut dense_design = (!matches!(spec.identifiability, TensorBSplineIdentifiability::None))
5930        .then(|| tensor_product_design_from_marginals(&marginal_designs))
5931        .transpose()?;
5932    let mut candidates = Vec::<PenaltyCandidate>::with_capacity(
5933        match spec.penalty_decomposition {
5934            TensorBSplinePenaltyDecomposition::MarginalKroneckerSum => marginal_penalties.len(),
5935            TensorBSplinePenaltyDecomposition::Separable => marginal_penalties.len() * 2,
5936        } + if spec.double_penalty { 1 } else { 0 },
5937    );
5938
5939    // Tensor-product smoothing parameters are one-per-margin.  Therefore the
5940    // physical penalty attached to a margin must be normalized in that margin's
5941    // own working coordinates before it is embedded in the full tensor product.
5942    // Normalizing only the already-Kroneckered matrix would fold arbitrary
5943    // dimension-dependent identity factors into the margin's lambda and would
5944    // make anisotropic REML/LAML smoothing depend on the other margins' basis
5945    // sizes rather than on the marginal roughness operator itself.
5946    let normalized_marginal_penalties: Vec<(Array2<f64>, f64)> = marginal_penalties
5947        .iter()
5948        .map(normalize_penalty_in_constrained_space)
5949        .collect();
5950    let tensor_function_gram = if spec.double_penalty {
5951        if marginal_function_grams.len() != marginalnum_basis.len() {
5952            crate::bail_dim_basis!(
5953                "TensorBSpline double penalty requires one function Gram per margin; got {} for {} margins",
5954                marginal_function_grams.len(),
5955                marginalnum_basis.len()
5956            );
5957        }
5958        let mut gram = Array2::<f64>::eye(1);
5959        for marginal_gram in &marginal_function_grams {
5960            gram = kronecker_product(&gram, marginal_gram);
5961        }
5962        Some(gram)
5963    } else {
5964        None
5965    };
5966    // A single PSD sum has exactly the joint null space shared by every
5967    // marginal roughness block. It is used only to define the global
5968    // null-component penalty; the ordinary tensor candidates below retain
5969    // their one-coordinate-per-margin decomposition.
5970    let joint_wiggliness = if spec.double_penalty {
5971        let mut sum = Array2::<f64>::zeros((total_cols, total_cols));
5972        for dim in 0..normalized_marginal_penalties.len() {
5973            let mut embedded = Array2::<f64>::eye(1);
5974            for (margin, &width) in marginalnum_basis.iter().enumerate() {
5975                let factor = if margin == dim {
5976                    normalized_marginal_penalties[margin].0.clone()
5977                } else {
5978                    Array2::<f64>::eye(width)
5979                };
5980                embedded = kronecker_product(&embedded, &factor);
5981            }
5982            sum += &embedded;
5983        }
5984        Some(sum)
5985    } else {
5986        None
5987    };
5988    let mut kronecker_marginal_penalties =
5989        Vec::<Array2<f64>>::with_capacity(normalized_marginal_penalties.len());
5990
5991    match spec.penalty_decomposition {
5992        TensorBSplinePenaltyDecomposition::MarginalKroneckerSum => {
5993            // Accumulate the Kronecker-sum of the per-margin penalties,
5994            // `Σ_dim S_dim`, whose null space is exactly the *joint* null space
5995            // of all marginal penalties — the tensor of marginal polynomial
5996            // null spaces. The tensor double penalty (below) shrinks only this
5997            // joint null, never the already-penalized interaction range.
5998            for dim in 0..normalized_marginal_penalties.len() {
5999                let mut s_dim = Array2::<f64>::eye(1);
6000                let mut factors = Vec::<Array2<f64>>::with_capacity(marginalnum_basis.len());
6001                for (j, &qj) in marginalnum_basis.iter().enumerate() {
6002                    let factor = if j == dim {
6003                        normalized_marginal_penalties[j].0.clone()
6004                    } else {
6005                        Array2::<f64>::eye(qj)
6006                    };
6007                    factors.push(factor.clone());
6008                    s_dim = kronecker_product(&s_dim, &factor);
6009                }
6010                if dim == kronecker_marginal_penalties.len() {
6011                    kronecker_marginal_penalties.push(normalized_marginal_penalties[dim].0.clone());
6012                }
6013                candidates.push(PenaltyCandidate {
6014                    matrix: ConstructiveQuadratic::try_from_dense_psd(
6015                        s_dim,
6016                        "tensor marginal penalty",
6017                    )?,
6018                    source: PenaltySource::TensorMarginal { dim },
6019                    normalization_scale: normalized_marginal_penalties[dim].1,
6020                    kronecker_factors: Some(factors),
6021                    op: None,
6022                });
6023            }
6024
6025            if let (Some(primary), Some(gram)) =
6026                (joint_wiggliness.as_ref(), tensor_function_gram.as_ref())
6027                && let Some(shrink) =
6028                    crate::basis::function_space_nullspace_shrinkage(primary, gram)?
6029            {
6030                let (matrix, normalization_scale) = normalize_penalty_in_constrained_space(&shrink);
6031                candidates.push(PenaltyCandidate {
6032                    matrix: ConstructiveQuadratic::try_from_dense_psd(
6033                        matrix,
6034                        "tensor global null-function ridge",
6035                    )?,
6036                    source: PenaltySource::TensorGlobalRidge,
6037                    normalization_scale,
6038                    kronecker_factors: None,
6039                    op: None,
6040                });
6041            }
6042        }
6043        TensorBSplinePenaltyDecomposition::Separable => {
6044            let projectors = tensor_margin_range_null_projectors(&normalized_marginal_penalties)?;
6045            let n_masks = 1usize.checked_shl(projectors.len() as u32).ok_or_else(|| {
6046                BasisError::InvalidInput(format!(
6047                    "t2 separable tensor penalty supports at most {} margins, got {}",
6048                    usize::BITS - 1,
6049                    projectors.len()
6050                ))
6051            })?;
6052            for mask in 1..n_masks {
6053                let mut matrix = Array2::<f64>::eye(1);
6054                let mut factors = Vec::<Array2<f64>>::with_capacity(projectors.len());
6055                let mut penalized_margins = Vec::<usize>::new();
6056                for (dim, projector) in projectors.iter().enumerate() {
6057                    let use_range = ((mask >> dim) & 1) == 1;
6058                    let factor = if use_range {
6059                        penalized_margins.push(dim);
6060                        projector.range.clone()
6061                    } else {
6062                        projector.null.clone()
6063                    };
6064                    matrix = kronecker_product(&matrix, &factor);
6065                    factors.push(factor);
6066                }
6067                let (matrix, normalization_scale) = normalize_penalty_in_constrained_space(&matrix);
6068                candidates.push(PenaltyCandidate {
6069                    matrix: ConstructiveQuadratic::try_from_dense_psd(
6070                        matrix,
6071                        "tensor separable penalty",
6072                    )?,
6073                    source: PenaltySource::TensorSeparable { penalized_margins },
6074                    normalization_scale,
6075                    kronecker_factors: Some(factors),
6076                    op: None,
6077                });
6078            }
6079
6080            if let (Some(primary), Some(gram)) =
6081                (joint_wiggliness.as_ref(), tensor_function_gram.as_ref())
6082                && let Some(matrix) =
6083                    crate::basis::function_space_nullspace_shrinkage(primary, gram)?
6084            {
6085                let (matrix, normalization_scale) = normalize_penalty_in_constrained_space(&matrix);
6086                candidates.push(PenaltyCandidate {
6087                    matrix: ConstructiveQuadratic::try_from_dense_psd(
6088                        matrix,
6089                        "separable tensor global null-function ridge",
6090                    )?,
6091                    source: PenaltySource::TensorGlobalRidge,
6092                    normalization_scale,
6093                    kronecker_factors: None,
6094                    op: None,
6095                });
6096            }
6097        }
6098    }
6099
6100    let z_opt = match &spec.identifiability {
6101        TensorBSplineIdentifiability::None => None,
6102        TensorBSplineIdentifiability::SumToZero => {
6103            if total_cols < 2 {
6104                crate::bail_invalid_basis!(
6105                    "TensorBSpline requires at least 2 basis coefficients to enforce sum-to-zero identifiability"
6106                );
6107            }
6108            let dense_design_ref = dense_design.as_ref().ok_or_else(|| {
6109                BasisError::InvalidInput(
6110                    "tensor sum-to-zero identifiability requires a realized basis".to_string(),
6111                )
6112            })?;
6113            let (_, z) = apply_sum_to_zero_constraint(dense_design_ref.view(), None)?;
6114            let gauge = gam_problem::Gauge::sum_to_zero(z);
6115            Some(gauge.block_transform(0))
6116        }
6117        TensorBSplineIdentifiability::MarginalSumToZero => {
6118            // `ti(...)`: drop the marginal main effects by centering every
6119            // margin independently, then form the tensor product of the
6120            // centered margins. Concretely, each margin `j` is reparameterized
6121            // by its own sum-to-zero null basis `Z_j` (so the constant — i.e.
6122            // the marginal intercept — is removed from that axis), and the
6123            // combined reparameterization is the Kronecker product
6124            // `Z = Z₀ ⊗ Z₁ ⊗ … ⊗ Z_{d-1}`. Applying `Z` to the full-tensor
6125            // design `B = B₀ ⊗ … ⊗ B_{d-1}` yields `B Z = (B₀ Z₀) ⊗ … ⊗
6126            // (B_{d-1} Z_{d-1})`, the tensor product of the centered margins,
6127            // which by construction contains no pure main effect.
6128            if marginal_designs.len() < 2 {
6129                crate::bail_invalid_basis!(
6130                    "tensor interaction (ti) identifiability requires at least 2 margins"
6131                );
6132            }
6133            let mut z = Array2::<f64>::eye(1);
6134            for (dim, marginal) in marginal_designs.iter().enumerate() {
6135                if marginal.ncols() < 2 {
6136                    crate::bail_invalid_basis!(
6137                        "tensor interaction (ti) margin {dim} has fewer than 2 basis functions; \
6138                         cannot remove its marginal main effect"
6139                    );
6140                }
6141                let (_, z_dim) = apply_sum_to_zero_constraint(marginal.view(), None)?;
6142                let gauge_dim = gam_problem::Gauge::sum_to_zero(z_dim);
6143                let z_dim = gauge_dim.block_transform(0);
6144                z = kronecker_product(&z, &z_dim);
6145            }
6146            Some(z)
6147        }
6148        TensorBSplineIdentifiability::FrozenTransform { transform } => {
6149            if transform.nrows() != total_cols {
6150                crate::bail_dim_basis!(
6151                    "frozen tensor identifiability transform mismatch: design has {} columns but transform has {} rows",
6152                    total_cols,
6153                    transform.nrows()
6154                );
6155            }
6156            Some(transform.clone())
6157        }
6158    };
6159
6160    if let Some(z) = z_opt.as_ref() {
6161        let gauge = gam_problem::Gauge::from_block_transforms(&[z.clone()]);
6162        let dense = dense_design.as_mut().ok_or_else(|| {
6163            BasisError::InvalidInput(
6164                "tensor identifiability transform requires a realized basis".to_string(),
6165            )
6166        })?;
6167        let restricted_design = gauge.restrict_design(dense);
6168        *dense = restricted_design;
6169        candidates = candidates
6170            .into_iter()
6171            .map(|candidate| -> Result<PenaltyCandidate, BasisError> {
6172                let restricted = candidate
6173                    .matrix
6174                    .restricted(&gauge, "tensor identifiability restriction")?;
6175                // Re-normalize in the *actual* coefficient chart used by the
6176                // fit.  The tensor sum-to-zero transform is not norm-preserving
6177                // for each overlapping marginal penalty, so carrying the raw
6178                // marginal Frobenius scale into the restricted space changes the
6179                // relative amount of smoothing seen by the LAML/REML optimizer.
6180                // Keep the physical scale in metadata and give the optimizer
6181                // unit-scale constrained penalties for every tensor margin.
6182                let (_, c_new) = normalize_penalty_in_constrained_space(restricted.dense());
6183                let matrix = restricted.scaled(
6184                    1.0 / c_new,
6185                    "normalized tensor penalty after identifiability",
6186                )?;
6187                Ok(PenaltyCandidate {
6188                    matrix,
6189                    source: candidate.source,
6190                    normalization_scale: candidate.normalization_scale * c_new,
6191                    // Z^T S Z is no longer a Kronecker product of the original
6192                    // marginal factors, so the Kronecker fast path in construction.rs
6193                    // must not be taken. Clearing kronecker_factors forces the generic
6194                    // block-local eigendecomposition path, which operates on the
6195                    // transformed matrix and is correct.
6196                    kronecker_factors: None,
6197                    op: candidate.op.clone(),
6198                })
6199            })
6200            .collect::<Result<Vec<_>, _>>()?;
6201
6202        if candidates
6203            .iter()
6204            .any(|candidate| matches!(candidate.source, PenaltySource::TensorGlobalRidge))
6205        {
6206            let width = candidates
6207                .first()
6208                .ok_or_else(|| {
6209                    BasisError::InvalidInput(
6210                        "TensorBSpline global ridge has no penalty candidates".to_string(),
6211                    )
6212                })?
6213                .matrix
6214                .nrows();
6215            let physical_primary_terms = candidates
6216                .iter()
6217                .filter(|candidate| {
6218                    !matches!(candidate.source, PenaltySource::TensorGlobalRidge)
6219                })
6220                .map(|candidate| {
6221                    candidate.matrix.scaled(
6222                        candidate.normalization_scale,
6223                        "physical tensor primary penalty",
6224                    )
6225                })
6226                .collect::<Result<Vec<_>, _>>()?;
6227            let joint_primary = ConstructiveQuadratic::sum(
6228                &physical_primary_terms,
6229                "joint tensor primary penalty",
6230            )?;
6231            for candidate in &mut candidates {
6232                if !matches!(candidate.source, PenaltySource::TensorGlobalRidge) {
6233                    continue;
6234                }
6235                let physical_ridge = candidate.matrix.scaled(
6236                    candidate.normalization_scale,
6237                    "physical tensor null ridge",
6238                )?;
6239                match crate::basis::rebuild_metric_consistent_ridge(
6240                    &joint_primary,
6241                    &physical_ridge,
6242                )? {
6243                    Some(rebuilt) => {
6244                        let (_, scale) =
6245                            normalize_penalty_in_constrained_space(rebuilt.dense());
6246                        candidate.matrix = rebuilt.scaled(
6247                            1.0 / scale,
6248                            "normalized rebuilt tensor null ridge",
6249                        )?;
6250                        candidate.normalization_scale = scale;
6251                    }
6252                    None => {
6253                        candidate.matrix = ConstructiveQuadratic::zero(width);
6254                        candidate.normalization_scale = 1.0;
6255                    }
6256                }
6257                candidate.kronecker_factors = None;
6258                candidate.op = None;
6259            }
6260        }
6261    }
6262
6263    let filtered = filter_penalty_candidates(candidates)?;
6264    let identifiability_is_none =
6265        matches!(spec.identifiability, TensorBSplineIdentifiability::None);
6266    // All marginals expose a sparse representation iff each `marginal_sparse`
6267    // slot is `Some(...)`. Currently this is true when every marginal is a
6268    // free-boundary, non-periodic 1D B-spline returned as
6269    // `DesignMatrix::Sparse` from `build_bspline_basis_1d`. Periodic B-splines
6270    // and other dense-only marginals leave a `None` and trigger the fall-back
6271    // path. Identifiability transforms (`SumToZero`, `FrozenTransform`) make
6272    // the tensor design dense in general, so we also gate on that.
6273    let all_marginals_sparse = marginal_sparse.iter().all(Option::is_some);
6274    let design = if let Some(dense_design) = dense_design {
6275        DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(dense_design))
6276    } else if identifiability_is_none && all_marginals_sparse {
6277        // Sparse Khatri-Rao path: assemble the (n, ∏ q_j) tensor product
6278        // directly as a SparseColMat, preserving the ∏(degree_j+1) nonzero
6279        // structure per row instead of densifying to ∏ q_j columns. This is
6280        // mathematically identical to `tensor_product_design_from_marginals`
6281        // applied to the corresponding dense marginals.
6282        let sparse_marginals: Vec<&SparseColMat<usize, f64>> = marginal_sparse
6283            .iter()
6284            .map(|m| m.as_ref().expect("all_marginals_sparse just verified"))
6285            .collect();
6286        let sparse_design = tensor_product_design_from_sparse_marginals(&sparse_marginals)?;
6287        DesignMatrix::Sparse(gam_linalg::matrix::SparseDesignMatrix::new(sparse_design))
6288    } else {
6289        let marginals: Vec<Arc<Array2<f64>>> = marginal_designs
6290            .iter()
6291            .map(|m| Arc::new(m.clone()))
6292            .collect();
6293        let op = TensorProductDesignOperator::new(marginals).map_err(|e| {
6294            BasisError::InvalidInput(format!("TensorProductDesignOperator build failed: {e}"))
6295        })?;
6296        DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(op)))
6297    };
6298
6299    Ok(BasisBuildResult {
6300        design,
6301        affine_offset: None,
6302        active_penalties: filtered.active,
6303        dropped_penalties: filtered.dropped,
6304        joint_null_rotation: None,
6305        metadata: BasisMetadata::TensorBSpline {
6306            feature_cols: feature_cols.to_vec(),
6307            knots: marginal_knots,
6308            degrees: marginal_degrees,
6309            // Prefer the per-margin effective period derived in the loop —
6310            // it captures both the explicit `spec.periods` route and the
6311            // implied period from a `PeriodicUniform` marginal knotspec.
6312            // Falling back to `spec.periods` when populated keeps any
6313            // user-supplied explicit period authoritative even if the
6314            // marginal knotspec carried no periodicity hint.
6315            periods: marginal_effective_periods,
6316            is_cr: marginal_is_cr_flags,
6317            identifiability_transform: z_opt,
6318        },
6319        // The current Kronecker runtime diagonalizes only the marginal
6320        // roughness operators and represents its optional joint-null block as
6321        // a Euclidean selector. A function-space ridge generally does not
6322        // commute with those marginals, so advertising it as factored would
6323        // make PIRLS and REML solve a different objective. Keep the exact
6324        // canonical matrices whenever null recovery is active.
6325        kronecker_factored: if !spec.double_penalty
6326            && matches!(spec.identifiability, TensorBSplineIdentifiability::None)
6327            && matches!(
6328                spec.penalty_decomposition,
6329                TensorBSplinePenaltyDecomposition::MarginalKroneckerSum
6330            ) {
6331            Some(KroneckerFactoredBasis::new(
6332                marginal_designs,
6333                kronecker_marginal_penalties,
6334                marginalnum_basis.clone(),
6335                spec.double_penalty,
6336            ))
6337        } else {
6338            None
6339        },
6340    })
6341}
6342
6343#[cfg(test)]
6344mod tensor_function_space_runtime_tests {
6345    use super::*;
6346    use crate::basis::{
6347        BSplineBoundaryConditions, BSplineEndpointBoundaryCondition, OneDimensionalBoundary,
6348    };
6349    use ndarray::array;
6350
6351    fn marginal() -> BSplineBasisSpec {
6352        BSplineBasisSpec {
6353            degree: 2,
6354            penalty_order: 1,
6355            knotspec: BSplineKnotSpec::Generate {
6356                data_range: (0.0, 1.0),
6357                num_internal_knots: 2,
6358            },
6359            double_penalty: false,
6360            identifiability: BSplineIdentifiability::None,
6361            boundary: OneDimensionalBoundary::Open,
6362            boundary_conditions: BSplineBoundaryConditions::default(),
6363        }
6364    }
6365
6366    #[test]
6367    fn function_space_tensor_ridge_uses_exact_canonical_runtime() {
6368        let data = array![
6369            [0.00, 0.13],
6370            [0.15, 0.82],
6371            [0.29, 0.37],
6372            [0.43, 0.95],
6373            [0.58, 0.21],
6374            [0.71, 0.66],
6375            [0.86, 0.48],
6376            [1.00, 0.04]
6377        ];
6378        let mut spec = TensorBSplineSpec {
6379            marginalspecs: vec![marginal(), marginal()],
6380            periods: Vec::new(),
6381            double_penalty: true,
6382            identifiability: TensorBSplineIdentifiability::None,
6383            penalty_decomposition: TensorBSplinePenaltyDecomposition::MarginalKroneckerSum,
6384        };
6385        let built = build_tensor_bspline_basis(data.view(), &[0, 1], &spec)
6386            .expect("double-penalty tensor basis");
6387        assert!(
6388            built
6389                .active_penalties
6390                .iter()
6391                .any(|penalty| { matches!(penalty.info.source, PenaltySource::TensorGlobalRidge) })
6392        );
6393        assert!(
6394            built.kronecker_factored.is_none(),
6395            "the legacy factored runtime cannot represent a function-metric global ridge"
6396        );
6397
6398        spec.double_penalty = false;
6399        let singly_penalized = build_tensor_bspline_basis(data.view(), &[0, 1], &spec)
6400            .expect("single-penalty tensor basis");
6401        assert!(
6402            singly_penalized.kronecker_factored.is_some(),
6403            "the exact marginal-only fast path must remain available"
6404        );
6405    }
6406
6407    #[test]
6408    fn tensor_nonzero_anchor_is_rejected_before_its_affine_lift_can_be_dropped() {
6409        let data = array![[0.0, 0.0], [0.25, 0.75], [0.75, 0.25], [1.0, 1.0]];
6410        let mut anchored = marginal();
6411        anchored.boundary_conditions.left =
6412            BSplineEndpointBoundaryCondition::Anchored { value: 1.25 };
6413        let spec = TensorBSplineSpec {
6414            marginalspecs: vec![anchored, marginal()],
6415            periods: Vec::new(),
6416            double_penalty: false,
6417            identifiability: TensorBSplineIdentifiability::None,
6418            penalty_decomposition: TensorBSplinePenaltyDecomposition::MarginalKroneckerSum,
6419        };
6420
6421        let error = build_tensor_bspline_basis(data.view(), &[0, 1], &spec)
6422            .expect_err("a tensor margin cannot silently discard an inhomogeneous lift");
6423        let message = error.to_string();
6424        assert!(message.contains("TensorBSpline margin 0"));
6425        assert!(message.contains("non-zero endpoint anchor"));
6426        assert!(message.contains("explicit model offset"));
6427    }
6428
6429}
6430
6431pub fn tensor_product_design_from_marginals(
6432    marginal_designs: &[Array2<f64>],
6433) -> Result<Array2<f64>, BasisError> {
6434    if marginal_designs.is_empty() {
6435        crate::bail_invalid_basis!("TensorBSpline requires at least one marginal basis");
6436    }
6437    let n = marginal_designs[0].nrows();
6438    for (i, b) in marginal_designs.iter().enumerate().skip(1) {
6439        if b.nrows() != n {
6440            crate::bail_dim_basis!(
6441                "tensor marginal row mismatch at dim {i}: expected {n}, got {}",
6442                b.nrows()
6443            );
6444        }
6445    }
6446    let total_cols = marginal_designs.iter().try_fold(1usize, |acc, b| {
6447        acc.checked_mul(b.ncols())
6448            .ok_or_else(|| BasisError::DimensionMismatch("tensor basis too large".to_string()))
6449    })?;
6450    // Tensor-product Khatri-Rao: design[i, j] = Π_d marginal_d[i, j_d]
6451    // where j is the multi-index (j_1, ..., j_D) flattened. Independent
6452    // across rows; parallelize row chunks and fill the pre-allocated
6453    // contiguous Array2 in place (no Vec-flatten-collect intermediate,
6454    // which doubled the peak memory at large-scale N).
6455    use ndarray::parallel::prelude::*;
6456    use rayon::iter::{IntoParallelIterator, ParallelIterator};
6457    let mut design = Array2::<f64>::zeros((n, total_cols));
6458    design
6459        .axis_chunks_iter_mut(ndarray::Axis(0), 1024)
6460        .into_par_iter()
6461        .enumerate()
6462        .for_each(|(chunk_idx, mut block)| {
6463            let row_offset = chunk_idx * 1024;
6464            // Scratch buffers reused across rows in this chunk.
6465            let mut cur = Vec::<f64>::with_capacity(total_cols);
6466            let mut next = Vec::<f64>::with_capacity(total_cols);
6467            for (local_i, mut out_row) in block.outer_iter_mut().enumerate() {
6468                let i = row_offset + local_i;
6469                cur.clear();
6470                cur.push(1.0);
6471                for b in marginal_designs {
6472                    let q = b.ncols();
6473                    next.clear();
6474                    next.resize(cur.len() * q, 0.0);
6475                    // Hoist the row view out of the inner `col` loop so the
6476                    // q reads per `a_idx` reuse a single contiguous slice
6477                    // instead of recomputing `b[[i, col]]` strides per cell.
6478                    let b_row = b.row(i);
6479                    let b_slice = b_row
6480                        .as_slice()
6481                        .expect("Array2 row from outer_iter is contiguous");
6482                    for (a_idx, &aval) in cur.iter().enumerate() {
6483                        let off = a_idx * q;
6484                        let dst = &mut next[off..off + q];
6485                        for col in 0..q {
6486                            dst[col] = aval * b_slice[col];
6487                        }
6488                    }
6489                    std::mem::swap(&mut cur, &mut next);
6490                }
6491                // `out_row` is a row of the contiguous C-major `design`
6492                // Array2, so it is backed by a contiguous slice. Use a
6493                // bulk slice copy instead of an element-by-element write
6494                // loop.
6495                let out_slice = out_row
6496                    .as_slice_mut()
6497                    .expect("design row is contiguous in C-major Array2");
6498                out_slice.copy_from_slice(&cur);
6499            }
6500        });
6501    Ok(design)
6502}
6503
6504/// Render a numeric factor level for an error message: an integer-valued code
6505/// (`1999.0`) prints as `1999`, so an unseen-level message names the level the
6506/// user actually wrote rather than a spurious `.0`.
6507fn fmt_level_value(v: f64) -> String {
6508    if v.is_finite() && v.fract() == 0.0 && v.abs() < 1e15 {
6509        format!("{}", v as i64)
6510    } else {
6511        format!("{v}")
6512    }
6513}
6514
6515pub fn build_random_effect_block(
6516    data: ArrayView2<'_, f64>,
6517    spec: &RandomEffectTermSpec,
6518) -> Result<RandomEffectBlock, BasisError> {
6519    let n = data.nrows();
6520    let p = data.ncols();
6521    if spec.feature_col >= p {
6522        crate::bail_dim_basis!(
6523            "random-effect term '{}' feature column {} out of bounds for {} columns",
6524            spec.name,
6525            spec.feature_col,
6526            p
6527        );
6528    }
6529
6530    let col = data.column(spec.feature_col);
6531    if col.iter().any(|v| !v.is_finite()) {
6532        crate::bail_invalid_basis!(
6533            "random-effect term '{}' contains non-finite group values",
6534            spec.name
6535        );
6536    }
6537
6538    let kept_levels: Vec<u64> = if let Some(levels) = spec.frozen_levels.as_ref() {
6539        if levels.is_empty() {
6540            crate::bail_invalid_basis!(
6541                "random-effect term '{}' has empty frozen_levels",
6542                spec.name
6543            );
6544        }
6545        // Canonicalize a possibly-legacy frozen set: a `-0.0` group interned
6546        // before signed-zero canonicalization landed would otherwise never match
6547        // a canonicalized data row. Idempotent on already-canonical sets (#2145).
6548        levels
6549            .iter()
6550            .map(|&b| gam_data::canonical_level_bits(f64::from_bits(b)))
6551            .collect()
6552    } else {
6553        let mut seen = BTreeSet::<u64>::new();
6554        let mut levels = Vec::<u64>::new();
6555        for &v in col {
6556            let bits = gam_data::canonical_level_bits(v);
6557            if seen.insert(bits) {
6558                levels.push(bits);
6559            }
6560        }
6561        if levels.is_empty() {
6562            crate::bail_invalid_basis!("random-effect term '{}' has no observed levels", spec.name);
6563        }
6564        let start_idx = if spec.drop_first_level && levels.len() > 1 {
6565            1usize
6566        } else {
6567            0usize
6568        };
6569        levels[start_idx..].to_vec()
6570    };
6571
6572    if kept_levels.is_empty() {
6573        crate::bail_invalid_basis!(
6574            "random-effect term '{}' drops all levels; keep at least one level",
6575            spec.name
6576        );
6577    }
6578
6579    let q = kept_levels.len();
6580    let mut level_to_col = BTreeMap::<u64, usize>::new();
6581    for (idx, &bits) in kept_levels.iter().enumerate() {
6582        if level_to_col.insert(bits, idx).is_some() {
6583            crate::bail_invalid_basis!(
6584                "random-effect term '{}' has duplicate frozen level bits {bits}",
6585                spec.name
6586            );
6587        }
6588    }
6589    // A FIXED categorical factor (`factor(g)` or a bare `+ g`; `lenient_unseen
6590    // == false`) must reject an out-of-vocabulary level rather than silently
6591    // encode it as an all-zero dummy row that collapses onto the factor's
6592    // centering point (#2137/#2102). For a *string* factor the typed schema
6593    // encode rejects the unseen level before we get here; a *numeric-coded*
6594    // `factor(year)` column, however, reaches Rust as a plain numeric column
6595    // with no categorical schema, so the operator that owns the frozen level
6596    // vocabulary is the enforcement point that closes the same gap. Only when
6597    // the full one-hot block is kept (`!drop_first_level`) does an absent level
6598    // unambiguously mean "unseen" — with treatment coding the dropped baseline
6599    // is a legitimate absent column, so we do not gate that path. `frozen_levels`
6600    // presence marks the predict/frozen context; at fit the vocabulary is
6601    // derived from this very data, so no row is unseen.
6602    let strict_unseen =
6603        !spec.lenient_unseen && !spec.drop_first_level && spec.frozen_levels.is_some();
6604    let mut group_ids = Vec::with_capacity(n);
6605    for (row, &v) in col.iter().enumerate() {
6606        let bits = gam_data::canonical_level_bits(v);
6607        let group_id = level_to_col.get(&bits).copied();
6608        if strict_unseen && group_id.is_none() {
6609            crate::bail_invalid_basis!(
6610                "unseen level '{}' in fixed factor column '{}' at row {}; the factor's levels \
6611                 were fixed at fit time and an out-of-vocabulary level cannot be predicted \
6612                 (use group({}) for a random effect that tolerates held-out levels)",
6613                fmt_level_value(v),
6614                spec.name,
6615                row,
6616                spec.name
6617            );
6618        }
6619        group_ids.push(group_id);
6620    }
6621
6622    Ok(RandomEffectBlock {
6623        name: spec.name.clone(),
6624        group_ids,
6625        num_groups: q,
6626        kept_levels,
6627    })
6628}
6629
6630#[cfg(test)]
6631mod random_effect_signed_zero_tests {
6632    use super::{RandomEffectTermSpec, build_random_effect_block};
6633    use ndarray::array;
6634
6635    fn spec() -> RandomEffectTermSpec {
6636        RandomEffectTermSpec {
6637            name: "g".to_string(),
6638            feature_col: 0,
6639            drop_first_level: false,
6640            penalized: true,
6641            frozen_levels: None,
6642            lenient_unseen: true,
6643        }
6644    }
6645
6646    #[test]
6647    fn signed_zero_rows_share_one_group() {
6648        // A column mixing +0.0 and -0.0 for the physically same group must
6649        // intern as ONE level, and every row (either spelling) must resolve to
6650        // that single group column — the #2145 fit-side regression.
6651        let data = array![[-0.0_f64], [0.0], [1.0], [-0.0], [1.0]];
6652        let block = build_random_effect_block(data.view(), &spec()).unwrap();
6653        assert_eq!(
6654            block.num_groups, 2,
6655            "0.0/-0.0 must not split into two groups"
6656        );
6657        // Rows 0,1,3 are the same group; rows 2,4 the other.
6658        assert_eq!(block.group_ids[0], block.group_ids[1]);
6659        assert_eq!(block.group_ids[0], block.group_ids[3]);
6660        assert_eq!(block.group_ids[2], block.group_ids[4]);
6661        assert_ne!(block.group_ids[0], block.group_ids[2]);
6662    }
6663
6664    #[test]
6665    fn frozen_positive_zero_matches_negative_zero_row() {
6666        // A model frozen on +0.0 must resolve a -0.0 prediction row to the same
6667        // column — the #2145 predict-side regression that dropped the effect.
6668        let mut s = spec();
6669        s.frozen_levels = Some(vec![0.0_f64.to_bits(), 1.0_f64.to_bits()]);
6670        let data = array![[-0.0_f64], [1.0]];
6671        let block = build_random_effect_block(data.view(), &s).unwrap();
6672        assert_eq!(
6673            block.group_ids[0],
6674            Some(0),
6675            "-0.0 must match the +0.0 column"
6676        );
6677        assert_eq!(block.group_ids[1], Some(1));
6678    }
6679
6680    #[test]
6681    fn frozen_negative_zero_matches_positive_zero_row() {
6682        // The symmetric direction: a legacy model interned on -0.0 (pre-fix)
6683        // must still resolve a +0.0 prediction row after canonicalization.
6684        let mut s = spec();
6685        s.frozen_levels = Some(vec![(-0.0_f64).to_bits(), 1.0_f64.to_bits()]);
6686        let data = array![[0.0_f64], [1.0]];
6687        let block = build_random_effect_block(data.view(), &s).unwrap();
6688        assert_eq!(
6689            block.group_ids[0],
6690            Some(0),
6691            "+0.0 must match the -0.0 column"
6692        );
6693    }
6694
6695    // ---- #2137: fixed factor (`factor(g)`) strict-unseen enforcement --------
6696
6697    fn fixed_factor_spec() -> RandomEffectTermSpec {
6698        // A numeric-coded `factor(year)`: full one-hot (`drop_first_level=false`),
6699        // FIXED (`lenient_unseen=false`), vocabulary pinned at fit.
6700        let mut s = spec();
6701        s.name = "year".to_string();
6702        s.lenient_unseen = false;
6703        s
6704    }
6705
6706    #[test]
6707    fn fixed_factor_rejects_unseen_numeric_level_at_predict() {
6708        // The numeric-coded `factor(year)` gap (#2137): the column reaches the
6709        // operator as plain numbers (no categorical schema to pre-filter it), so
6710        // the operator that owns the frozen vocabulary must reject an unseen
6711        // code rather than encode an all-zero (centering-point) row.
6712        let mut s = fixed_factor_spec();
6713        s.frozen_levels = Some(vec![2000.0_f64.to_bits(), 2001.0_f64.to_bits()]);
6714        let data = array![[2000.0_f64], [1999.0]];
6715        let err = build_random_effect_block(data.view(), &s)
6716            .expect_err("an unseen fixed-factor level must be rejected");
6717        let msg = format!("{err}");
6718        assert!(
6719            msg.contains("unseen level"),
6720            "message must name the defect: {msg}"
6721        );
6722        assert!(
6723            msg.contains("1999"),
6724            "message must name the integer level (not 1999.0): {msg}"
6725        );
6726        assert!(msg.contains("year"), "message must name the column: {msg}");
6727    }
6728
6729    #[test]
6730    fn fixed_factor_accepts_seen_numeric_levels_at_predict() {
6731        // Control: every seen level still resolves; strictness rejects only the
6732        // genuinely out-of-vocabulary code.
6733        let mut s = fixed_factor_spec();
6734        s.frozen_levels = Some(vec![2000.0_f64.to_bits(), 2001.0_f64.to_bits()]);
6735        let data = array![[2001.0_f64], [2000.0]];
6736        let block = build_random_effect_block(data.view(), &s).unwrap();
6737        assert_eq!(block.group_ids[0], Some(1));
6738        assert_eq!(block.group_ids[1], Some(0));
6739    }
6740
6741    #[test]
6742    fn fixed_factor_at_fit_time_derives_vocabulary_and_never_false_rejects() {
6743        // At FIT (`frozen_levels=None`) the vocabulary is derived from this very
6744        // data, so no row is unseen — the strict guard must not fire even though
6745        // the factor is strict.
6746        let mut s = fixed_factor_spec();
6747        s.frozen_levels = None;
6748        let data = array![[2000.0_f64], [2001.0], [2002.0], [2000.0]];
6749        let block = build_random_effect_block(data.view(), &s)
6750            .expect("fit-time build must not reject its own levels");
6751        assert_eq!(block.num_groups, 3);
6752    }
6753
6754    #[test]
6755    fn random_effect_still_tolerates_unseen_numeric_level() {
6756        // Non-regression: a lenient random effect (`group`/`re`/`s(bs="re")`)
6757        // encodes an unseen level as an all-zero (population-mean) row, NOT a
6758        // rejection — the held-out-group contract (#2102) is unchanged.
6759        let mut s = spec(); // lenient_unseen = true
6760        s.frozen_levels = Some(vec![2000.0_f64.to_bits(), 2001.0_f64.to_bits()]);
6761        let data = array![[2000.0_f64], [1999.0]];
6762        let block = build_random_effect_block(data.view(), &s)
6763            .expect("a random effect tolerates unseen levels");
6764        assert_eq!(block.group_ids[0], Some(0));
6765        assert_eq!(
6766            block.group_ids[1], None,
6767            "unseen level → population mean, not a reject"
6768        );
6769    }
6770}
6771
6772impl SmoothDesign {
6773    /// Map an unconstrained term coefficient vector to its constrained shape space.
6774    /// This is useful for nonlinear fits that optimize unconstrained parameters.
6775    pub fn map_term_coefficients(
6776        unconstrained: &Array1<f64>,
6777        shape: ShapeConstraint,
6778    ) -> Result<Array1<f64>, BasisError> {
6779        if unconstrained.is_empty() {
6780            crate::bail_invalid_basis!("unconstrained coefficient vector cannot be empty");
6781        }
6782        let mapped = match shape {
6783            ShapeConstraint::None => unconstrained.clone(),
6784            ShapeConstraint::MonotoneIncreasing => cumulative_exp(unconstrained, 1.0),
6785            ShapeConstraint::MonotoneDecreasing => cumulative_exp(unconstrained, -1.0),
6786            ShapeConstraint::Convex => second_cumulative_exp(unconstrained, 1.0),
6787            ShapeConstraint::Concave => second_cumulative_exp(unconstrained, -1.0),
6788        };
6789        Ok(mapped)
6790    }
6791}
6792
6793pub struct LocalSmoothTermBuild {
6794    pub dim: usize,
6795    pub design: DesignMatrix,
6796    /// Fixed row-wise term contribution for an affine basis chart.
6797    pub affine_offset: Option<Array1<f64>>,
6798    pub active_penalties: Vec<ActivePenalty>,
6799    /// Joint-null absorption rotation for this smooth. `Some(rotation)`
6800    /// records `Q = [U_range | U_null]` spanning `null(Σ_k penalties[k])`,
6801    /// the joint null across all active penalty blocks on this smooth.
6802    /// `None` means the joint penalty is full-rank (joint nullity = 0) or
6803    /// there are no penalties. Stage-2 commit A: plumbing only — populated
6804    /// by commit B, applied by commit D.
6805    pub joint_null_rotation: Option<crate::basis::JointNullRotation>,
6806    pub dropped_penalties: Vec<DroppedPenaltyInfo>,
6807    pub metadata: BasisMetadata,
6808    pub linear_constraints: Option<LinearInequalityConstraints>,
6809    pub box_reparam: bool,
6810    pub kronecker_factored: Option<KroneckerFactoredBasis>,
6811}
6812
6813#[derive(Clone)]
6814pub struct PcaScoresMemmapDesignOperator {
6815    mmap: Arc<memmap2::Mmap>,
6816    data_offset: usize,
6817    nrows: usize,
6818    ncols: usize,
6819    chunk_size: usize,
6820}
6821
6822impl PcaScoresMemmapDesignOperator {
6823    fn open(path: PathBuf, chunk_size: usize) -> Result<Self, BasisError> {
6824        let file = File::open(&path).map_err(|err| {
6825            BasisError::InvalidInput(format!(
6826                "failed to open lazy Pca .npy scores '{}': {err}",
6827                path.display()
6828            ))
6829        })?;
6830        // The .npy scores file is read-only training-cache data; this
6831        // module never mutates it. The error path below converts mmap
6832        // failure to a typed `BasisError::InvalidInput`.
6833        // SAFETY: `memmap2::Mmap::map` requires no concurrent writers; the
6834        // contract is held by this module's read-only access pattern.
6835        let mmap = unsafe {
6836            memmap2::Mmap::map(&file).map_err(|err| {
6837                BasisError::InvalidInput(format!(
6838                    "failed to memmap lazy Pca .npy scores '{}': {err}",
6839                    path.display()
6840                ))
6841            })?
6842        };
6843        let (data_offset, nrows, ncols) = parse_f64_2d_npy_header(&mmap, &path)?;
6844        let expected = data_offset
6845            .checked_add(nrows.saturating_mul(ncols).saturating_mul(8))
6846            .ok_or_else(|| {
6847                BasisError::InvalidInput(format!(
6848                    "lazy Pca .npy scores '{}' shape is too large",
6849                    path.display()
6850                ))
6851            })?;
6852        if mmap.len() < expected {
6853            crate::bail_invalid_basis!(
6854                "lazy Pca .npy scores '{}' is truncated: header expects {} bytes, file has {}",
6855                path.display(),
6856                expected,
6857                mmap.len()
6858            );
6859        }
6860        Ok(Self {
6861            mmap: Arc::new(mmap),
6862            data_offset,
6863            nrows,
6864            ncols,
6865            chunk_size: chunk_size.max(1),
6866        })
6867    }
6868
6869    fn value(&self, row: usize, col: usize) -> f64 {
6870        let offset = self.data_offset + (row * self.ncols + col) * 8;
6871        let mut bytes = [0_u8; 8];
6872        bytes.copy_from_slice(&self.mmap[offset..offset + 8]);
6873        f64::from_le_bytes(bytes)
6874    }
6875
6876    fn chunk_rows(&self) -> usize {
6877        self.chunk_size.min(self.nrows.max(1))
6878    }
6879}
6880
6881impl LinearOperator for PcaScoresMemmapDesignOperator {
6882    fn nrows(&self) -> usize {
6883        self.nrows
6884    }
6885
6886    fn ncols(&self) -> usize {
6887        self.ncols
6888    }
6889
6890    fn apply(&self, vector: &Array1<f64>) -> Array1<f64> {
6891        assert_eq!(
6892            vector.len(),
6893            self.ncols,
6894            "lazy Pca apply vector length mismatch"
6895        );
6896        let mut out = Array1::<f64>::zeros(self.nrows);
6897        for start in (0..self.nrows).step_by(self.chunk_rows()) {
6898            let end = (start + self.chunk_rows()).min(self.nrows);
6899            for row in start..end {
6900                let mut acc = 0.0;
6901                for col in 0..self.ncols {
6902                    acc += self.value(row, col) * vector[col];
6903                }
6904                out[row] = acc;
6905            }
6906        }
6907        out
6908    }
6909
6910    fn apply_transpose(&self, vector: &Array1<f64>) -> Array1<f64> {
6911        assert_eq!(
6912            vector.len(),
6913            self.nrows,
6914            "lazy Pca apply_transpose vector length mismatch"
6915        );
6916        let mut out = Array1::<f64>::zeros(self.ncols);
6917        for start in (0..self.nrows).step_by(self.chunk_rows()) {
6918            let end = (start + self.chunk_rows()).min(self.nrows);
6919            for row in start..end {
6920                let scale = vector[row];
6921                if scale == 0.0 {
6922                    continue;
6923                }
6924                for col in 0..self.ncols {
6925                    out[col] += scale * self.value(row, col);
6926                }
6927            }
6928        }
6929        out
6930    }
6931
6932    fn diag_xtw_x(&self, weights: &Array1<f64>) -> Result<Array2<f64>, String> {
6933        if weights.len() != self.nrows {
6934            return Err(format!(
6935                "lazy Pca diag_xtw_x weight length mismatch: weights={}, nrows={}",
6936                weights.len(),
6937                self.nrows
6938            ));
6939        }
6940        FiniteSignedWeightsView::try_from_array(weights)
6941            .map_err(|reason| format!("lazy Pca diag_xtw_x: {reason}"))?;
6942        let mut gram = Array2::<f64>::zeros((self.ncols, self.ncols));
6943        for start in (0..self.nrows).step_by(self.chunk_rows()) {
6944            let end = (start + self.chunk_rows()).min(self.nrows);
6945            for row in start..end {
6946                let w = weights[row];
6947                if w == 0.0 {
6948                    continue;
6949                }
6950                for a in 0..self.ncols {
6951                    let xa = self.value(row, a);
6952                    if xa == 0.0 {
6953                        continue;
6954                    }
6955                    for b in a..self.ncols {
6956                        gram[[a, b]] += w * xa * self.value(row, b);
6957                    }
6958                }
6959            }
6960        }
6961        for a in 0..self.ncols {
6962            for b in 0..a {
6963                gram[[a, b]] = gram[[b, a]];
6964            }
6965        }
6966        Ok(gram)
6967    }
6968
6969    fn apply_weighted_normal(
6970        &self,
6971        weights: FiniteSignedWeightsView<'_>,
6972        vector: &Array1<f64>,
6973        penalty: Option<&Array2<f64>>,
6974        ridge: f64,
6975    ) -> Array1<f64> {
6976        assert_eq!(
6977            weights.len(),
6978            self.nrows,
6979            "lazy Pca weighted-normal weight mismatch"
6980        );
6981        assert_eq!(
6982            vector.len(),
6983            self.ncols,
6984            "lazy Pca weighted-normal vector mismatch"
6985        );
6986        let weights = weights.view();
6987        let mut out = Array1::<f64>::zeros(self.ncols);
6988        for start in (0..self.nrows).step_by(self.chunk_rows()) {
6989            let end = (start + self.chunk_rows()).min(self.nrows);
6990            for row in start..end {
6991                let w = weights[row];
6992                if w == 0.0 {
6993                    continue;
6994                }
6995                let mut row_dot = 0.0;
6996                for col in 0..self.ncols {
6997                    row_dot += self.value(row, col) * vector[col];
6998                }
6999                if row_dot == 0.0 {
7000                    continue;
7001                }
7002                let scaled = w * row_dot;
7003                for col in 0..self.ncols {
7004                    out[col] += scaled * self.value(row, col);
7005                }
7006            }
7007        }
7008        if let Some(pen) = penalty {
7009            out += &pen.dot(vector);
7010        }
7011        if ridge > 0.0 {
7012            out += &vector.mapv(|x| ridge * x);
7013        }
7014        out
7015    }
7016}
7017
7018impl DenseDesignOperator for PcaScoresMemmapDesignOperator {
7019    fn compute_xtwy(&self, weights: &Array1<f64>, y: &Array1<f64>) -> Result<Array1<f64>, String> {
7020        if weights.len() != self.nrows || y.len() != self.nrows {
7021            return Err(format!(
7022                "lazy Pca compute_xtwy dimension mismatch: weights={}, y={}, nrows={}",
7023                weights.len(),
7024                y.len(),
7025                self.nrows
7026            ));
7027        }
7028        FiniteSignedWeightsView::try_from_array(weights)
7029            .map_err(|reason| format!("lazy Pca compute_xtwy: {reason}"))?;
7030        let mut out = Array1::<f64>::zeros(self.ncols);
7031        for start in (0..self.nrows).step_by(self.chunk_rows()) {
7032            let end = (start + self.chunk_rows()).min(self.nrows);
7033            for row in start..end {
7034                let scale = weights[row] * y[row];
7035                if scale == 0.0 {
7036                    continue;
7037                }
7038                for col in 0..self.ncols {
7039                    out[col] += scale * self.value(row, col);
7040                }
7041            }
7042        }
7043        Ok(out)
7044    }
7045
7046    fn row_chunk_into(
7047        &self,
7048        rows: Range<usize>,
7049        mut out: ArrayViewMut2<'_, f64>,
7050    ) -> Result<(), MatrixMaterializationError> {
7051        if rows.end > self.nrows || rows.start > rows.end {
7052            return Err(MatrixMaterializationError::MissingRowChunk {
7053                context: "lazy Pca row range out of bounds",
7054            });
7055        }
7056        if out.nrows() != rows.end - rows.start || out.ncols() != self.ncols {
7057            return Err(MatrixMaterializationError::MissingRowChunk {
7058                context: "lazy Pca row_chunk_into shape mismatch",
7059            });
7060        }
7061        for (local, row) in (rows.start..rows.end).enumerate() {
7062            for col in 0..self.ncols {
7063                out[[local, col]] = self.value(row, col);
7064            }
7065        }
7066        Ok(())
7067    }
7068
7069    fn to_dense(&self) -> Array2<f64> {
7070        let mut out = Array2::<f64>::zeros((self.nrows, self.ncols));
7071        self.row_chunk_into(0..self.nrows, out.view_mut())
7072            .expect("lazy Pca full materialization failed");
7073        out
7074    }
7075}
7076
7077pub fn parse_f64_2d_npy_header(
7078    bytes: &[u8],
7079    path: &PathBuf,
7080) -> Result<(usize, usize, usize), BasisError> {
7081    let mut reader = std::io::Cursor::new(bytes);
7082    let header = npyz::NpyHeader::from_reader(&mut reader).map_err(|err| {
7083        BasisError::InvalidInput(format!(
7084            "lazy Pca scores '{}' has an invalid .npy header: {err}",
7085            path.display()
7086        ))
7087    })?;
7088    let is_little_endian_f64 = matches!(
7089        header.dtype(),
7090        npyz::DType::Plain(ref dtype)
7091            if dtype.type_char() == npyz::TypeChar::Float
7092                && dtype.size_field() == 8
7093                && dtype.endianness() == npyz::Endianness::Little
7094    );
7095    if !is_little_endian_f64 {
7096        crate::bail_invalid_basis!(
7097            "lazy Pca scores '{}' must be scalar little-endian float64 .npy, got {}",
7098            path.display(),
7099            header.dtype().descr()
7100        );
7101    }
7102    if header.order() != npyz::Order::C {
7103        crate::bail_invalid_basis!(
7104            "lazy Pca scores '{}' must be C-contiguous, not Fortran-ordered",
7105            path.display()
7106        );
7107    }
7108    if header.shape().len() != 2 {
7109        crate::bail_invalid_basis!(
7110            "lazy Pca scores '{}' must have shape (N, K), got {:?}",
7111            path.display(),
7112            header.shape()
7113        );
7114    }
7115    let nrows = usize::try_from(header.shape()[0]).map_err(|_| {
7116        BasisError::InvalidInput(format!(
7117            "lazy Pca scores '{}' row count {} exceeds this platform's address space",
7118            path.display(),
7119            header.shape()[0]
7120        ))
7121    })?;
7122    let ncols = usize::try_from(header.shape()[1]).map_err(|_| {
7123        BasisError::InvalidInput(format!(
7124            "lazy Pca scores '{}' column count {} exceeds this platform's address space",
7125            path.display(),
7126            header.shape()[1]
7127        ))
7128    })?;
7129    let data_offset = usize::try_from(reader.position()).map_err(|_| {
7130        BasisError::InvalidInput(format!(
7131            "lazy Pca scores '{}' header offset exceeds this platform's address space",
7132            path.display()
7133        ))
7134    })?;
7135    Ok((data_offset, nrows, ncols))
7136}
7137
7138pub fn pca_center_mean(x: ArrayView2<'_, f64>) -> Result<Array1<f64>, BasisError> {
7139    if x.nrows() == 0 {
7140        crate::bail_invalid_basis!("Pca basis requires at least one row to compute center mean");
7141    }
7142    let mut mean = Array1::<f64>::zeros(x.ncols());
7143    for row in x.rows() {
7144        mean += &row;
7145    }
7146    mean.mapv_inplace(|v| v / x.nrows() as f64);
7147    Ok(mean)
7148}
7149
7150/// Build the empirical final-function mass penalty from the raw score Gram.
7151///
7152/// For the realized PCA score design `Z`, the quadratic form is
7153///
7154/// `beta^T S beta = smooth_penalty * mean_i((Z beta)_i^2)`.
7155///
7156/// Thus `smooth_penalty` chooses the reference-measure scale only; the existing
7157/// REML smoothing coordinate multiplying this penalty learns the shrinkage
7158/// strength.  In particular, this is not an identity ridge on whichever
7159/// coefficient chart happened to encode the score columns.
7160fn pca_function_mass_penalty(
7161    mut raw_score_gram: Array2<f64>,
7162    n_rows: usize,
7163    smooth_penalty: f64,
7164) -> Result<Array2<f64>, BasisError> {
7165    let k = raw_score_gram.ncols();
7166    if raw_score_gram.nrows() != k {
7167        crate::bail_dim_basis!(
7168            "Pca score Gram must be square, got {}x{}",
7169            raw_score_gram.nrows(),
7170            k
7171        );
7172    }
7173    if n_rows == 0 {
7174        crate::bail_invalid_basis!("Pca basis requires at least one score row");
7175    }
7176    if k == 0 {
7177        crate::bail_invalid_basis!("Pca basis requires at least one score column");
7178    }
7179    if k > n_rows {
7180        crate::bail_invalid_basis!(
7181            "Pca score design is rank deficient: {} score columns cannot have full column rank with only {} rows; remove redundant components",
7182            k,
7183            n_rows
7184        );
7185    }
7186    if raw_score_gram.iter().any(|value| !value.is_finite()) {
7187        crate::bail_invalid_basis!("Pca score design produced a non-finite function Gram");
7188    }
7189
7190    // Use the same design-rank convention as the global identifiability audit.
7191    // `rrqr_from_gram_with_permutation` recovers the column-pivoted QR verdict
7192    // from Z^T Z while retaining the tall design's row-count-aware tolerance.
7193    let rrqr = gam_linalg::faer_ndarray::rrqr_from_gram_with_permutation(
7194        &raw_score_gram,
7195        n_rows,
7196        gam_linalg::faer_ndarray::default_rrqr_rank_alpha(),
7197    )
7198    .map_err(BasisError::LinalgError)?;
7199    if rrqr.rank != k {
7200        let redundant_columns = &rrqr.column_permutation[rrqr.rank..];
7201        crate::bail_invalid_basis!(
7202            "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",
7203            rrqr.rank,
7204            k,
7205            rrqr.rank_tol,
7206            redundant_columns
7207        );
7208    }
7209
7210    raw_score_gram.mapv_inplace(|value| value * smooth_penalty / n_rows as f64);
7211    Ok(raw_score_gram)
7212}
7213
7214pub fn build_pca_smooth_basis(
7215    data: ArrayView2<'_, f64>,
7216    feature_cols: &[usize],
7217    basis_matrix: &Array2<f64>,
7218    centered: bool,
7219    smooth_penalty: f64,
7220    center_mean: Option<&Array1<f64>>,
7221    pca_basis_path: Option<&PathBuf>,
7222    chunk_size: usize,
7223) -> Result<BasisBuildResult, BasisError> {
7224    if !smooth_penalty.is_finite() || smooth_penalty < 0.0 {
7225        crate::bail_invalid_basis!(
7226            "Pca smooth_penalty must be finite and non-negative, got {}",
7227            smooth_penalty
7228        );
7229    }
7230    if data.nrows() == 0 {
7231        crate::bail_invalid_basis!("Pca basis requires at least one data row");
7232    }
7233
7234    if let Some(path) = pca_basis_path {
7235        let op = PcaScoresMemmapDesignOperator::open(path.clone(), chunk_size)?;
7236        if op.nrows != data.nrows() {
7237            crate::bail_dim_basis!(
7238                "lazy Pca scores row mismatch: .npy has {}, data has {}",
7239                op.nrows,
7240                data.nrows()
7241            );
7242        }
7243        // The out-of-core scores are already the realized final-function
7244        // design. Stream Z^T Z without materializing its n-by-k rows.
7245        let raw_score_gram = op
7246            .diag_xtw_x(&Array1::<f64>::ones(op.nrows))
7247            .map_err(|err| {
7248                BasisError::InvalidInput(format!(
7249                    "lazy Pca function-mass Gram construction failed: {err}"
7250                ))
7251            })?;
7252        let penalty = pca_function_mass_penalty(raw_score_gram, op.nrows, smooth_penalty)?;
7253        let filtered = filter_penalty_candidates(vec![PenaltyCandidate {
7254            matrix: ConstructiveQuadratic::try_from_dense_psd(
7255                penalty,
7256                "lazy PCA function-mass penalty",
7257            )?,
7258            source: PenaltySource::OperatorMass,
7259            normalization_scale: 1.0,
7260            kronecker_factors: None,
7261            op: None,
7262        }])?;
7263        return Ok(BasisBuildResult {
7264            design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(op))),
7265            affine_offset: None,
7266            active_penalties: filtered.active,
7267            dropped_penalties: filtered.dropped,
7268            joint_null_rotation: None,
7269            metadata: BasisMetadata::Pca {
7270                feature_cols: feature_cols.to_vec(),
7271                basis_matrix: basis_matrix.clone(),
7272                centered,
7273                smooth_penalty,
7274                center_mean: center_mean.cloned(),
7275                pca_basis_path: Some(path.clone()),
7276                chunk_size: chunk_size.max(1),
7277            },
7278            kronecker_factored: None,
7279        });
7280    }
7281    if basis_matrix.nrows() != feature_cols.len() {
7282        crate::bail_dim_basis!(
7283            "Pca basis row mismatch: basis rows={}, feature columns={}",
7284            basis_matrix.nrows(),
7285            feature_cols.len()
7286        );
7287    }
7288    let mut x = select_columns(data, feature_cols)?;
7289    let mean = if centered {
7290        match center_mean {
7291            Some(mean) => mean.clone(),
7292            None => pca_center_mean(x.view())?,
7293        }
7294    } else {
7295        Array1::<f64>::zeros(feature_cols.len())
7296    };
7297    if centered {
7298        for mut row in x.rows_mut() {
7299            row -= &mean;
7300        }
7301    }
7302    let design = fast_ab(&x, basis_matrix);
7303    let raw_score_gram = gam_linalg::faer_ndarray::fast_ata(&design);
7304    let penalty = pca_function_mass_penalty(raw_score_gram, design.nrows(), smooth_penalty)?;
7305    let filtered = filter_penalty_candidates(vec![PenaltyCandidate {
7306        matrix: ConstructiveQuadratic::try_from_dense_psd(
7307            penalty,
7308            "PCA function-mass penalty",
7309        )?,
7310        source: PenaltySource::OperatorMass,
7311        normalization_scale: 1.0,
7312        kronecker_factors: None,
7313        op: None,
7314    }])?;
7315    Ok(BasisBuildResult {
7316        design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(design)),
7317        affine_offset: None,
7318        active_penalties: filtered.active,
7319        dropped_penalties: filtered.dropped,
7320        joint_null_rotation: None,
7321        metadata: BasisMetadata::Pca {
7322            feature_cols: feature_cols.to_vec(),
7323            basis_matrix: basis_matrix.clone(),
7324            centered,
7325            smooth_penalty,
7326            center_mean: centered.then_some(mean),
7327            pca_basis_path: None,
7328            chunk_size: chunk_size.max(1),
7329        },
7330        kronecker_factored: None,
7331    })
7332}
7333
7334#[cfg(test)]
7335mod pca_function_mass_tests {
7336    use super::{PenaltySource, build_pca_smooth_basis, parse_f64_2d_npy_header};
7337    use ndarray::{Array1, Array2, array};
7338    use std::io::Write;
7339    use std::path::PathBuf;
7340
7341    fn quadratic_form(matrix: &Array2<f64>, coefficients: &Array1<f64>) -> f64 {
7342        coefficients.dot(&matrix.dot(coefficients))
7343    }
7344
7345    fn assert_close(left: f64, right: f64) {
7346        let scale = left.abs().max(right.abs()).max(1.0);
7347        assert!(
7348            (left - right).abs() <= 1e-11 * scale,
7349            "values differ: left={left:.16e}, right={right:.16e}"
7350        );
7351    }
7352
7353    fn write_f64_npy(scores: &Array2<f64>) -> PathBuf {
7354        let path = std::env::temp_dir().join(format!(
7355            "gam_terms_pca_function_mass_{}.npy",
7356            std::process::id()
7357        ));
7358        let mut header = format!(
7359            "{{'descr': '<f8', 'fortran_order': False, 'shape': ({}, {}), }}",
7360            scores.nrows(),
7361            scores.ncols()
7362        );
7363        while (10 + header.len() + 1) % 16 != 0 {
7364            header.push(' ');
7365        }
7366        header.push('\n');
7367        let header_len = u16::try_from(header.len()).expect("test .npy header fits u16");
7368
7369        let mut file = std::fs::File::create(&path).expect("create test .npy");
7370        file.write_all(b"\x93NUMPY").expect("write .npy magic");
7371        file.write_all(&[1, 0]).expect("write .npy version");
7372        file.write_all(&header_len.to_le_bytes())
7373            .expect("write .npy header length");
7374        file.write_all(header.as_bytes())
7375            .expect("write .npy header");
7376        for &value in scores {
7377            file.write_all(&value.to_le_bytes())
7378                .expect("write .npy score");
7379        }
7380        path
7381    }
7382
7383    fn npy_v1_bytes(mut header: String) -> Vec<u8> {
7384        while (10 + header.len() + 1) % 16 != 0 {
7385            header.push(' ');
7386        }
7387        header.push('\n');
7388        let header_len = u16::try_from(header.len()).expect("test header fits v1");
7389        let mut bytes = b"\x93NUMPY".to_vec();
7390        bytes.extend_from_slice(&[1, 0]);
7391        bytes.extend_from_slice(&header_len.to_le_bytes());
7392        bytes.extend_from_slice(header.as_bytes());
7393        bytes
7394    }
7395
7396    #[test]
7397    fn npy_header_parser_uses_exact_ast_fields_2293() {
7398        let path = PathBuf::from("scores.npy");
7399        let bytes = npy_v1_bytes(
7400            "{'shape':(3, 2), 'note':'True', 'descr':'<f8', 'fortran_order':False,}".to_string(),
7401        );
7402        let (offset, rows, cols) =
7403            parse_f64_2d_npy_header(&bytes, &path).expect("valid reordered header");
7404        assert_eq!((rows, cols), (3, 2));
7405        assert_eq!(offset, bytes.len());
7406
7407        for header in [
7408            "{'descr':'<f8','fortran_order':True,'shape':(3,2),}",
7409            "{'descr':'>f8','fortran_order':False,'shape':(3,2),}",
7410            "{'descr':'<f8','shape':(3,2),}",
7411            "{'descr':'<f8','fortran_order':'False','shape':(3,2),}",
7412            "{'descr':'<f8','fortran_order':False,'shape':(6,),}",
7413        ] {
7414            let invalid = npy_v1_bytes(header.to_string());
7415            assert!(
7416                parse_f64_2d_npy_header(&invalid, &path).is_err(),
7417                "{header}"
7418            );
7419        }
7420    }
7421
7422    #[test]
7423    fn pca_penalty_quadratic_equals_empirical_fitted_function_norm() {
7424        let data = array![[1.0, 2.0], [-1.0, 0.5], [2.0, -0.5], [0.25, -1.5]];
7425        let basis = array![[1.0, 0.5], [-0.25, 2.0]];
7426        let smooth_penalty = 2.5;
7427        let built = build_pca_smooth_basis(
7428            data.view(),
7429            &[0, 1],
7430            &basis,
7431            false,
7432            smooth_penalty,
7433            None,
7434            None,
7435            2,
7436        )
7437        .expect("full-rank PCA basis");
7438        let coefficients = array![0.7, -1.2];
7439        let design = built.design.to_dense();
7440        let fitted = design.dot(&coefficients);
7441        let expected = smooth_penalty * fitted.dot(&fitted) / fitted.len() as f64;
7442        let actual = quadratic_form(&built.active_penalties[0].matrix, &coefficients);
7443
7444        assert_close(actual, expected);
7445        assert_eq!(built.active_penalties[0].nullity, 0);
7446        assert_eq!(
7447            built.active_penalties[0].info.source,
7448            PenaltySource::OperatorMass
7449        );
7450    }
7451
7452    #[test]
7453    fn pca_function_mass_is_invariant_to_nonorthogonal_score_reparameterization() {
7454        let scores = array![[1.0, 2.0], [-1.0, 0.5], [2.0, -0.5], [0.25, -1.5]];
7455        let identity = Array2::<f64>::eye(2);
7456        // An invertible scale-plus-shear, deliberately not orthogonal.
7457        let transform = array![[2.0, 0.5], [0.0, 0.25]];
7458        let base_coefficients = array![0.8, -1.1];
7459        // transform * transformed_coefficients == base_coefficients.
7460        let transformed_coefficients = array![1.5, -4.4];
7461        let smooth_penalty = 1.7;
7462
7463        let base = build_pca_smooth_basis(
7464            scores.view(),
7465            &[0, 1],
7466            &identity,
7467            false,
7468            smooth_penalty,
7469            None,
7470            None,
7471            2,
7472        )
7473        .expect("base PCA chart");
7474        let transformed = build_pca_smooth_basis(
7475            scores.view(),
7476            &[0, 1],
7477            &transform,
7478            false,
7479            smooth_penalty,
7480            None,
7481            None,
7482            2,
7483        )
7484        .expect("reparameterized PCA chart");
7485
7486        let fitted_base = base.design.to_dense().dot(&base_coefficients);
7487        let fitted_transformed = transformed.design.to_dense().dot(&transformed_coefficients);
7488        for (&left, &right) in fitted_base.iter().zip(fitted_transformed.iter()) {
7489            assert_close(left, right);
7490        }
7491        assert_close(
7492            quadratic_form(&base.active_penalties[0].matrix, &base_coefficients),
7493            quadratic_form(
7494                &transformed.active_penalties[0].matrix,
7495                &transformed_coefficients,
7496            ),
7497        );
7498    }
7499
7500    #[test]
7501    fn rank_deficient_pca_score_design_is_rejected() {
7502        let scores = array![[1.0, 0.0], [2.0, 0.0], [3.0, 0.0], [4.0, 0.0]];
7503        let result = build_pca_smooth_basis(
7504            scores.view(),
7505            &[0, 1],
7506            &Array2::<f64>::eye(2),
7507            false,
7508            1.0,
7509            None,
7510            None,
7511            2,
7512        );
7513        let err = result.err().expect("zero score column must be rejected");
7514        let message = err.to_string();
7515        assert!(
7516            message.contains("rank deficient"),
7517            "unexpected error: {message}"
7518        );
7519        assert!(
7520            message.contains("rank 1 < 2"),
7521            "missing RRQR evidence: {message}"
7522        );
7523    }
7524
7525    #[test]
7526    fn lazy_and_dense_pca_function_mass_penalties_match() {
7527        let scores = array![[1.0, 2.0], [-1.0, 0.5], [2.0, -0.5], [0.25, -1.5]];
7528        let smooth_penalty = 2.25;
7529        let path = write_f64_npy(&scores);
7530        let dense = build_pca_smooth_basis(
7531            scores.view(),
7532            &[0, 1],
7533            &Array2::<f64>::eye(2),
7534            false,
7535            smooth_penalty,
7536            None,
7537            None,
7538            2,
7539        )
7540        .expect("dense PCA basis");
7541        let lazy_data = Array2::<f64>::zeros((scores.nrows(), 0));
7542        let lazy = build_pca_smooth_basis(
7543            lazy_data.view(),
7544            &[],
7545            &Array2::<f64>::zeros((0, scores.ncols())),
7546            false,
7547            smooth_penalty,
7548            None,
7549            Some(&path),
7550            2,
7551        )
7552        .expect("lazy PCA basis");
7553        std::fs::remove_file(&path).expect("remove test .npy");
7554
7555        for (&left, &right) in dense.active_penalties[0]
7556            .matrix
7557            .iter()
7558            .zip(lazy.active_penalties[0].matrix.iter())
7559        {
7560            assert_close(left, right);
7561        }
7562        for (&left, &right) in dense
7563            .design
7564            .to_dense()
7565            .iter()
7566            .zip(lazy.design.to_dense().iter())
7567        {
7568            assert_close(left, right);
7569        }
7570    }
7571}
7572
7573/// A factor-level `by=` wrapper owns the model-space centering of its inner
7574/// smooth: it gates the raw/structurally-constrained basis to the level rows
7575/// and then centers that gated block exactly once against the level indicator
7576/// (`build_parametric_constraint_block_for_term` in `design_construction`).
7577/// Leaving the inner B-spline's default pooled weighted-sum-to-zero active here
7578/// would impose two generically-independent constraints — the pooled column
7579/// moment `m = Σ_h m_h` and the per-level moment `m_g` — so a raw `k`-column
7580/// basis collapses to `k-2` columns per level instead of `k-1`, deleting one
7581/// genuine nonconstant spline direction *before REML runs* (#1427). The group
7582/// main effect carries only the constant, so it cannot restore that direction.
7583///
7584/// Only the *default model-space* centering is deferred. Explicit structural or
7585/// frozen transforms (`RemoveLinearTrend`, `OrthogonalToDesignColumns`,
7586/// `FrozenTransform`, `None`) are user/structural choices and are preserved
7587/// verbatim.
7588pub fn defer_inner_model_centering_to_factor_level_wrapper(basis: &mut SmoothBasisSpec) {
7589    if let SmoothBasisSpec::BSpline1D { spec, .. } = basis
7590        && matches!(
7591            spec.identifiability,
7592            BSplineIdentifiability::WeightedSumToZero { .. }
7593        )
7594    {
7595        spec.identifiability = BSplineIdentifiability::None;
7596    }
7597}
7598
7599pub fn apply_by_variable_to_local_build(
7600    mut built: LocalSmoothTermBuild,
7601    data: ArrayView2<'_, f64>,
7602    by_col: usize,
7603    by: &ByVariableSpec,
7604    term_name: &str,
7605) -> Result<LocalSmoothTermBuild, BasisError> {
7606    if by_col >= data.ncols() {
7607        crate::bail_dim_basis!(
7608            "by-variable smooth term '{term_name}' references column {by_col}, but data has {} columns",
7609            data.ncols()
7610        );
7611    }
7612    let weights = match by {
7613        ByVariableSpec::Numeric => data.column(by_col).to_owned(),
7614        ByVariableSpec::Level { value_bits, .. } => {
7615            let value_bits = gam_data::canonical_level_bits(f64::from_bits(*value_bits));
7616            data.column(by_col).mapv(|value| {
7617                if gam_data::canonical_level_bits(value) == value_bits {
7618                    1.0
7619                } else {
7620                    0.0
7621                }
7622            })
7623        }
7624    };
7625    if weights.iter().any(|value| !value.is_finite()) {
7626        crate::bail_invalid_basis!(
7627            "by-variable smooth term '{term_name}' has non-finite by-column values"
7628        );
7629    }
7630
7631    let mut dense = built
7632        .design
7633        .try_to_dense_by_chunks("by-variable smooth row gating")
7634        .map_err(BasisError::InvalidInput)?;
7635    for (mut row, &weight) in dense.rows_mut().into_iter().zip(weights.iter()) {
7636        row.mapv_inplace(|value| value * weight);
7637    }
7638    if let Some(offset) = built.affine_offset.as_mut() {
7639        if offset.len() != weights.len() {
7640            crate::bail_dim_basis!(
7641                "by-variable smooth term '{term_name}' affine offset has {} rows but the by-variable has {}",
7642                offset.len(),
7643                weights.len()
7644            );
7645        }
7646        *offset *= &weights;
7647    }
7648    built.design = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(dense));
7649    built.kronecker_factored = None;
7650    Ok(built)
7651}
7652
7653/// Build the local smooth term for a `BySmooth` spec, which unifies numeric-by
7654/// and factor-by modulation into a single `SmoothTermSpec`.
7655///
7656/// For a **numeric** by-variable the inner smooth is built once and every row
7657/// is multiplied by the by-column value (identical to `ByVariable::Numeric`).
7658///
7659/// For a **factor** by-variable the inner smooth is built once and gated per
7660/// level into side-by-side column blocks, producing a `n × (L * p)` design
7661/// matrix.  The penalties are block-diagonalised (one copy of the inner penalty
7662/// per level) exactly as `build_factor_smooth` does for `bs="fs"/"sz"`.
7663pub fn build_by_smooth_local(
7664    data: ArrayView2<'_, f64>,
7665    term: &SmoothTermSpec,
7666    smooth: &SmoothBasisSpec,
7667    by_kind: &ByVarKind,
7668    workspace: &mut crate::basis::BasisWorkspace,
7669) -> Result<LocalSmoothTermBuild, BasisError> {
7670    let inner_term = SmoothTermSpec {
7671        name: term.name.clone(),
7672        basis: (*smooth).clone(),
7673        shape: term.shape,
7674        joint_null_rotation: None,
7675    };
7676    let inner = build_single_local_smooth_term(data, &inner_term, workspace)?;
7677
7678    match by_kind {
7679        ByVarKind::Numeric { feature_col } => {
7680            let inner_meta = inner.metadata.clone();
7681            let mut built = apply_by_variable_to_local_build(
7682                inner,
7683                data,
7684                *feature_col,
7685                &ByVariableSpec::Numeric,
7686                &term.name,
7687            )?;
7688            built.metadata = BasisMetadata::BySmooth {
7689                inner: Box::new(inner_meta),
7690                by_col: *feature_col,
7691                levels: None,
7692                ordered: false,
7693            };
7694            Ok(built)
7695        }
7696        ByVarKind::Factor {
7697            feature_col,
7698            frozen_levels,
7699            ordered,
7700        } => {
7701            // Collect factor levels: prefer the frozen set (replay path), else
7702            // scan the data column (first-fit path).
7703            let level_bits: Vec<u64> = if let Some(fl) = frozen_levels {
7704                fl.iter()
7705                    .map(|&b| gam_data::canonical_level_bits(f64::from_bits(b)))
7706                    .collect()
7707            } else {
7708                let col = data.column(*feature_col);
7709                let mut seen = BTreeSet::<u64>::new();
7710                for &v in col.iter() {
7711                    if v.is_finite() {
7712                        seen.insert(gam_data::canonical_level_bits(v));
7713                    }
7714                }
7715                seen.into_iter().collect()
7716            };
7717            let n_levels = level_bits.len();
7718            if n_levels == 0 {
7719                crate::bail_invalid_basis!(
7720                    "by-factor smooth term '{}': factor column {} has no observed levels",
7721                    term.name,
7722                    feature_col
7723                );
7724            }
7725            let p = inner.dim;
7726            let q = n_levels * p;
7727            let n = data.nrows();
7728
7729            let inner_dense = inner
7730                .design
7731                .try_to_dense_by_chunks("by-factor smooth design gating")
7732                .map_err(BasisError::InvalidInput)?;
7733
7734            // Gate each level into its own p-wide column block.
7735            let mut combined = Array2::<f64>::zeros((n, q));
7736            for (lvl_idx, &bits) in level_bits.iter().enumerate() {
7737                let col_start = lvl_idx * p;
7738                for row in 0..n {
7739                    if gam_data::canonical_level_bits(data[[row, *feature_col]]) == bits {
7740                        combined
7741                            .slice_mut(s![row, col_start..col_start + p])
7742                            .assign(&inner_dense.row(row));
7743                    }
7744                }
7745            }
7746
7747            // Build per-level INDEPENDENT penalties (#1427): one copy of each
7748            // inner penalty per level, but each confined to that single level's
7749            // diagonal block, so every (level, inner-penalty) pair is its OWN
7750            // smoothing-parameter coordinate. `s(x, by=g)` selects the per-group
7751            // curve wiggliness independently — the design is block-diagonal and
7752            // block-separable, so a correct REML must reproduce gamfit's own
7753            // independent per-group fits. Tiling a single inner penalty across
7754            // every level (as the `bs="fs"` shared-λ random-effect construction
7755            // does) collapses all groups onto ONE λ, which cannot match uneven
7756            // per-level smoothness and degrades as data grows (under-recovery up
7757            // to ~16× at n=2000). Emit `n_levels * n_penalties` blocks instead.
7758            let inner_meta = inner.metadata.clone();
7759            let n_penalties = inner.active_penalties.len();
7760            let n_blocks = n_penalties.saturating_mul(n_levels);
7761            let mut candidates = Vec::<PenaltyCandidate>::with_capacity(n_blocks);
7762            for base_penalty in &inner.active_penalties {
7763                for lvl in 0..n_levels {
7764                    let off = lvl * p;
7765                    let mut s_big = Array2::<f64>::zeros((q, q));
7766                    s_big
7767                        .slice_mut(s![off..off + p, off..off + p])
7768                        .assign(&base_penalty.matrix);
7769                    let (s_big, scale) = normalize_penalty_in_constrained_space(&s_big);
7770                    candidates.push(PenaltyCandidate {
7771                        matrix: ConstructiveQuadratic::try_from_dense_psd(
7772                            s_big,
7773                            "factor-smooth replicated penalty",
7774                        )?,
7775                        source: base_penalty.info.source.clone(),
7776                        normalization_scale: base_penalty.info.normalization_scale * scale,
7777                        kronecker_factors: None,
7778                        op: None,
7779                    });
7780                }
7781            }
7782
7783            // Re-analyze the completed q×q blocks in their actual coefficient
7784            // space. Copying the p×p marginal nullity understated every block's
7785            // null space by `(n_levels-1)·p`, while leaving its null basis and
7786            // joint-null rotation absent. The canonical filter authors matrix,
7787            // rank, nullity, null basis, and metadata together.
7788            let filtered = crate::basis::filter_penalty_candidates(candidates)?;
7789            let joint_null_rotation = crate::basis::compute_joint_null_rotation(&filtered.active)?;
7790            let mut dropped_penalties = inner.dropped_penalties;
7791            dropped_penalties.extend(filtered.dropped);
7792
7793            Ok(LocalSmoothTermBuild {
7794                dim: q,
7795                design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(combined)),
7796                // Exactly one factor block is active on every row. Each block
7797                // represents the same anchored marginal, so its fixed lift is
7798                // the inner lift once, not once per level.
7799                affine_offset: inner.affine_offset,
7800                active_penalties: filtered.active,
7801                joint_null_rotation,
7802                dropped_penalties,
7803                metadata: BasisMetadata::BySmooth {
7804                    inner: Box::new(inner_meta),
7805                    by_col: *feature_col,
7806                    levels: Some(level_bits),
7807                    ordered: *ordered,
7808                },
7809                linear_constraints: None,
7810                box_reparam: false,
7811                kronecker_factored: None,
7812            })
7813        }
7814    }
7815}
7816
7817pub fn ensure_by_variable_specs_match(
7818    kind: &BySmoothKind,
7819    by: &ByVariableSpec,
7820    term_name: &str,
7821) -> Result<(), BasisError> {
7822    match (kind, by) {
7823        (BySmoothKind::Numeric, ByVariableSpec::Numeric) => Ok(()),
7824        (BySmoothKind::Level { level_bits }, ByVariableSpec::Level { value_bits, .. })
7825            if level_bits == value_bits =>
7826        {
7827            Ok(())
7828        }
7829        _ => Err(BasisError::InvalidInput(format!(
7830            "by-variable smooth term '{term_name}' has inconsistent by-variable specifications"
7831        ))),
7832    }
7833}
7834
7835/// Choose a deterministic orthonormal basis for a subspace from its projector.
7836///
7837/// Eigenvectors belonging to a repeated eigenvalue are defined only up to an
7838/// arbitrary orthogonal rotation.  That freedom is harmless when consumers use
7839/// the whole projector `ZZ^T`, but it changes model semantics when each column
7840/// receives its own smoothing parameter.  We remove the eigensolver gauge by
7841/// repeatedly projecting coefficient coordinate axes into the subspace and
7842/// selecting the largest residual (with stable lowest-index tie breaking).
7843/// The result depends only on the subspace projector and the declared
7844/// coefficient chart, never on the orientation returned by an eigensolver.
7845fn canonical_nullspace_directions(z: &Array2<f64>) -> Result<Array2<f64>, BasisError> {
7846    let (coefficient_dim, nullity) = z.dim();
7847    if nullity == 0 {
7848        return Ok(Array2::zeros((coefficient_dim, 0)));
7849    }
7850    if coefficient_dim < nullity || z.iter().any(|value| !value.is_finite()) {
7851        crate::bail_invalid_basis!(
7852            "null-space basis must be finite with rows >= columns, got {}x{}",
7853            coefficient_dim,
7854            nullity
7855        );
7856    }
7857
7858    let tolerance = 128.0 * f64::EPSILON * coefficient_dim.max(1) as f64;
7859    let mut canonical = Array2::<f64>::zeros((coefficient_dim, nullity));
7860    for accepted in 0..nullity {
7861        let mut best_coordinate = usize::MAX;
7862        let mut best_norm = 0.0_f64;
7863        let mut best = Array1::<f64>::zeros(coefficient_dim);
7864
7865        for coordinate in 0..coefficient_dim {
7866            // `P e_j = Z (Z^T e_j)` without materializing the full projector.
7867            let mut candidate = Array1::<f64>::zeros(coefficient_dim);
7868            for row in 0..coefficient_dim {
7869                candidate[row] = (0..nullity)
7870                    .map(|axis| z[[row, axis]] * z[[coordinate, axis]])
7871                    .sum();
7872            }
7873            // Two-pass modified Gram--Schmidt keeps the selected directions
7874            // orthogonal even when successive projected coordinates are close.
7875            for _ in 0..2 {
7876                for axis in 0..accepted {
7877                    let direction = canonical.column(axis);
7878                    let projection = direction.dot(&candidate);
7879                    candidate.scaled_add(-projection, &direction);
7880                }
7881            }
7882            let norm = candidate.dot(&candidate).sqrt();
7883            let tie_band = tolerance * best_norm.max(1.0);
7884            if best_coordinate == usize::MAX || norm > best_norm + tie_band {
7885                best_coordinate = coordinate;
7886                best_norm = norm;
7887                best = candidate;
7888            }
7889        }
7890
7891        if best_coordinate == usize::MAX || best_norm <= tolerance {
7892            crate::bail_invalid_basis!(
7893                "null-space projector exposed only {} of {} independent directions",
7894                accepted,
7895                nullity
7896            );
7897        }
7898        best.mapv_inplace(|value| value / best_norm);
7899        // Fix the remaining sign gauge for reproducible metadata/debug output.
7900        let sign_anchor = best
7901            .iter()
7902            .enumerate()
7903            .max_by(|(left_index, left), (right_index, right)| {
7904                left.abs()
7905                    .partial_cmp(&right.abs())
7906                    .unwrap_or(std::cmp::Ordering::Equal)
7907                    .then_with(|| right_index.cmp(left_index))
7908            })
7909            .map(|(_, value)| *value)
7910            .unwrap_or(1.0);
7911        if sign_anchor < 0.0 {
7912            best.mapv_inplace(|value| -value);
7913        }
7914        canonical.column_mut(accepted).assign(&best);
7915    }
7916    Ok(canonical)
7917}
7918
7919#[cfg(test)]
7920mod canonical_nullspace_direction_tests {
7921    use super::*;
7922    use ndarray::array;
7923
7924    #[test]
7925    fn per_axis_null_penalties_are_invariant_to_eigensolver_gauge_2315() {
7926        let inv_sqrt_two = 0.5_f64.sqrt();
7927        let z = array![
7928            [inv_sqrt_two, 0.0],
7929            [inv_sqrt_two, 0.0],
7930            [0.0, 1.0],
7931            [0.0, 0.0]
7932        ];
7933        let rotation = array![[0.6, -0.8], [0.8, 0.6]];
7934        let rotated = z.dot(&rotation);
7935        let reference = canonical_nullspace_directions(&z).expect("canonical null basis");
7936        let actual =
7937            canonical_nullspace_directions(&rotated).expect("rotated canonical null basis");
7938        for axis in 0..reference.ncols() {
7939            let reference_penalty = reference
7940                .column(axis)
7941                .to_owned()
7942                .insert_axis(Axis(1))
7943                .dot(&reference.column(axis).insert_axis(Axis(0)));
7944            let actual_penalty = actual
7945                .column(axis)
7946                .to_owned()
7947                .insert_axis(Axis(1))
7948                .dot(&actual.column(axis).insert_axis(Axis(0)));
7949            let max_error = reference_penalty
7950                .iter()
7951                .zip(actual_penalty.iter())
7952                .map(|(left, right)| (left - right).abs())
7953                .fold(0.0_f64, f64::max);
7954            assert!(
7955                max_error <= 256.0 * f64::EPSILON,
7956                "axis {axis} changed by {max_error:e}"
7957            );
7958        }
7959    }
7960}
7961
7962/// Build a factor-smooth interaction basis (`bs="fs"`/`"sz"`/`"re"`).
7963///
7964/// A factor smooth replicates a shared marginal smooth in the continuous
7965/// covariate(s) once per level of a grouping factor, coupling all level blocks
7966/// through a *single* set of smoothing parameters (one per marginal penalty).
7967/// This is mgcv's `smooth.construct.fs.smooth.spec` realization and the
7968/// random-effect interpretation of a smooth: the per-level deviations are an
7969/// exchangeable family whose joint wiggliness/shrinkage is governed by the
7970/// shared λ, so the construction scales to many levels with a fixed parameter
7971/// count.
7972///
7973/// Flavours:
7974/// * `Fs` — full random factor-smooth. The marginal carries its wiggliness
7975///   penalty *and* a null-space ridge (double penalty), so the replicated
7976///   design is a proper full-rank random effect: each level's curve is shrunk
7977///   toward zero (intercept + linear trend included), recovering the mgcv
7978///   `bs="fs"` penalty structure `I_L ⊗ S_j` for every marginal penalty `S_j`.
7979/// * `Sz` — sum-to-zero factor smooth. Delegates to the existing
7980///   [`SmoothBasisSpec::FactorSumToZero`] construction (`L-1` deviation blocks,
7981///   coefficient-wise zero sum across levels).
7982/// * `Re` — pure random effect / random slope (`bs="re"`). A degree-1 marginal
7983///   gives the per-level `[1, x]` span; the penalty is the identity over each
7984///   level block (iid Gaussian coefficients), matching mgcv's `bs="re"` ridge.
7985///
7986/// The grouping levels are resolved once at fit time (sorted unique bit
7987/// patterns of the factor column) and frozen into the returned metadata so the
7988/// predict-time rebuild evaluates every row against its own level's block.
7989pub fn build_factor_smooth(
7990    data: ArrayView2<'_, f64>,
7991    spec: &FactorSmoothSpec,
7992    term_name: &str,
7993    workspace: &mut crate::basis::BasisWorkspace,
7994) -> Result<LocalSmoothTermBuild, BasisError> {
7995    if spec.continuous_cols.len() != 1 {
7996        crate::bail_invalid_basis!(
7997            "factor smooth term '{}' currently supports exactly one continuous covariate; found {}",
7998            term_name,
7999            spec.continuous_cols.len()
8000        );
8001    }
8002    let feature_col = spec.continuous_cols[0];
8003    let group_col = spec.group_col;
8004    if feature_col >= data.ncols() || group_col >= data.ncols() {
8005        crate::bail_dim_basis!(
8006            "factor smooth term '{}' references columns ({}, {}) out of bounds for {} columns",
8007            term_name,
8008            feature_col,
8009            group_col,
8010            data.ncols()
8011        );
8012    }
8013
8014    // `Sz` is exactly the existing sum-to-zero factor smooth: reuse it verbatim
8015    // so there is a single source of truth for the zero-sum construction.
8016    if matches!(spec.flavour, FactorSmoothFlavour::Sz) {
8017        let levels = resolve_factor_smooth_levels(data, group_col, spec, term_name)?;
8018        let inner = SmoothBasisSpec::BSpline1D {
8019            feature_col,
8020            spec: factor_smooth_marginal_for_replay(&spec.marginal),
8021        };
8022        let sz_term = SmoothTermSpec {
8023            name: term_name.to_string(),
8024            basis: SmoothBasisSpec::FactorSumToZero {
8025                inner: Box::new(inner),
8026                by_col: group_col,
8027                levels: levels.clone(),
8028                frozen_global_orthogonality: None,
8029            },
8030            shape: ShapeConstraint::None,
8031            joint_null_rotation: None,
8032        };
8033        let mut built = build_single_local_smooth_term(data, &sz_term, workspace)?;
8034        // The delegated `FactorSumToZero` build returns the BARE inner B-spline
8035        // metadata (`BasisMetadata::BSpline1D`), but the term that owns this
8036        // build carries a `SmoothBasisSpec::FactorSmooth { Sz }` spec. Two
8037        // things break if we hand that mismatched pair downstream:
8038        //   1. `freeze_smooth_basis_from_metadata` matches on (spec, metadata)
8039        //      and has no `(FactorSmooth, BSpline1D)` arm, so any refit / spatial
8040        //      re-optimization that freezes the basis aborts with a "smooth
8041        //      metadata/spec type mismatch" error.
8042        //   2. The bare B-spline metadata carries no grouping levels, so a
8043        //      predict-time rebuild cannot replay the SAME replicated design.
8044        // Re-wrap the marginal geometry as `FactorSmooth` metadata exactly as
8045        // the Fs/Re path below does, giving all three factor-smooth flavours a
8046        // single, freeze-consistent metadata shape that also pins the levels.
8047        // Since #1605 the sz marginal is ALWAYS the penalized B-spline the `fs`
8048        // sibling uses (a natural cubic regression marginal hard-enforces f''=0
8049        // at the boundary and cannot represent curved deviations — a consistency
8050        // failure). The `CubicRegression1D` arm below is therefore unreachable on
8051        // a freshly-built sz spec; it is retained only as defense / backward
8052        // compatibility for a frozen spec that still carries a cr marginal, so
8053        // the predict-time freeze restores whatever marginal class it finds.
8054        let (knots, degree, periodic, marginal_is_cr) = match &built.metadata {
8055            BasisMetadata::BSpline1D {
8056                knots,
8057                periodic,
8058                degree,
8059                ..
8060            } => (
8061                knots.clone(),
8062                degree.unwrap_or(spec.marginal.degree),
8063                *periodic,
8064                false,
8065            ),
8066            BasisMetadata::CubicRegression1D { knots, .. } => {
8067                (knots.clone(), spec.marginal.degree, None, true)
8068            }
8069            other => {
8070                crate::bail_invalid_basis!(
8071                    "sz factor smooth term '{}' produced an unexpected marginal metadata variant {:?}",
8072                    term_name,
8073                    other
8074                );
8075            }
8076        };
8077        built.metadata = BasisMetadata::FactorSmooth {
8078            continuous_cols: spec.continuous_cols.clone(),
8079            group_col,
8080            knots,
8081            degree,
8082            periodic,
8083            group_levels: levels,
8084            flavour: "sz".to_string(),
8085            marginal_is_cr,
8086        };
8087        return Ok(built);
8088    }
8089
8090    let levels = resolve_factor_smooth_levels(data, group_col, spec, term_name)?;
8091    let n_levels = levels.len();
8092    if n_levels < 2 {
8093        crate::bail_invalid_basis!(
8094            "factor smooth term '{}' requires at least two grouping levels; found {}",
8095            term_name,
8096            n_levels
8097        );
8098    }
8099
8100    // `Fs` (order ≥ 1, the default) is the random-effect flavour: it penalizes
8101    // each null-space dimension of the marginal wiggliness penalty separately
8102    // below (mgcv's `bs="fs"` construction). That replaces the marginal's single
8103    // *combined* double penalty, so disable the latter here to avoid penalizing
8104    // the null space twice (once combined, once per dimension). The explicit
8105    // `m=0` opt-out keeps the legacy combined double penalty and adds no
8106    // per-dimension penalties.
8107    let use_per_dim_null = matches!(
8108        &spec.flavour,
8109        FactorSmoothFlavour::Fs { m_null_penalty_orders }
8110            if m_null_penalty_orders.iter().copied().max().unwrap_or(0) >= 1
8111    );
8112
8113    // Build the shared marginal design + penalties from the 1-D B-spline.
8114    // `Re` forces a degree-1 marginal (linear span) and replaces the marginal
8115    // wiggliness with an identity ridge below; `Fs` keeps the user's marginal
8116    // (cubic by default) and, under the per-dimension null path, gets its null
8117    // space penalized one dimension at a time after replication.
8118    let mut marginal_spec = factor_smooth_marginal_for_replay(&spec.marginal);
8119    if use_per_dim_null {
8120        marginal_spec.double_penalty = false;
8121    }
8122    let inner_term = SmoothTermSpec {
8123        name: format!("{term_name}::marginal"),
8124        basis: SmoothBasisSpec::BSpline1D {
8125            feature_col,
8126            spec: marginal_spec,
8127        },
8128        shape: ShapeConstraint::None,
8129        joint_null_rotation: None,
8130    };
8131    let inner = build_single_local_smooth_term(data, &inner_term, workspace)?;
8132    let mut base = inner
8133        .design
8134        .try_to_dense_by_chunks("factor smooth marginal")
8135        .map_err(BasisError::InvalidInput)?;
8136    if matches!(spec.flavour, FactorSmoothFlavour::Re) {
8137        // `bs="re"` is a parametric random intercept+slope, not a B-spline
8138        // smooth evaluated through clamped knot support.  A degree-1 B-spline
8139        // with no internal knots spans the training rows, but outside the
8140        // boundary knots its basis is not the model matrix for `(1 + x | g)`;
8141        // held-out extrapolation then loses the random slope contribution.
8142        // Build the random-effect marginal directly as `[1, x - c]`, centered
8143        // at the frozen marginal domain, so fit-time and replay-time rows use
8144        // the same well-conditioned parametric columns on and off the training
8145        // interval.
8146        let center = match &inner.metadata {
8147            BasisMetadata::BSpline1D { knots, .. } if !knots.is_empty() => {
8148                0.5 * (knots[0] + knots[knots.len() - 1])
8149            }
8150            _ => 0.0,
8151        };
8152        let mut linear = Array2::<f64>::ones((data.nrows(), 2));
8153        linear
8154            .column_mut(1)
8155            .assign(&data.column(feature_col).mapv(|x| x - center));
8156        base = linear;
8157    }
8158    let n = base.nrows();
8159    let p = base.ncols();
8160    let q = p * n_levels;
8161
8162    // Block-diagonal replicated design: row i contributes its marginal row to
8163    // the column block owned by its grouping level, zeros elsewhere.
8164    let mut dense = Array2::<f64>::zeros((n, q));
8165    for i in 0..n {
8166        let bits = gam_data::canonical_level_bits(data[[i, group_col]]);
8167        let Some(level_idx) = levels.iter().position(|b| *b == bits) else {
8168            // Held-out-group contract (#2365): `bs="re"` is a genuine random
8169            // effect, so in the frozen (predict/replay) context a row whose
8170            // group is outside the training vocabulary carries NO fitted
8171            // deviation — its row stays all-zero across every group block and
8172            // the prediction is the population component, mirroring
8173            // `build_random_effect_block`'s lenient `group_id = None` rows.
8174            // `fs`/`sz` estimate a per-level deviation FUNCTION (an unseen
8175            // level has no zero-deviation fallback that means "population"),
8176            // and the fit path derives `levels` from this very data, so both
8177            // stay strict.
8178            if matches!(spec.flavour, FactorSmoothFlavour::Re)
8179                && spec.group_frozen_levels.is_some()
8180            {
8181                continue;
8182            }
8183            return Err(BasisError::InvalidInput(format!(
8184                "factor smooth term '{term_name}' saw an unseen grouping level at row {}",
8185                i + 1
8186            )));
8187        };
8188        let start = level_idx * p;
8189        dense
8190            .slice_mut(s![i, start..start + p])
8191            .assign(&base.row(i));
8192    }
8193
8194    // Penalties: replicate each marginal penalty into a block-diagonal
8195    // `I_L ⊗ S_j` so every level shares the same smoothing parameter λ_j (one
8196    // λ per marginal penalty), the defining feature of a factor smooth. For
8197    // `Re` the marginal penalty is replaced by one ridge per parametric
8198    // coordinate so intercept and slope variances can be learned separately.
8199    let marginal_penalties: Vec<(Array2<f64>, PenaltySource, f64)> =
8200        if matches!(spec.flavour, FactorSmoothFlavour::Re) {
8201            (0..p)
8202                .map(|j| {
8203                    let mut matrix = Array2::<f64>::zeros((p, p));
8204                    matrix[[j, j]] = 1.0;
8205                    (matrix, PenaltySource::Primary, 1.0)
8206                })
8207                .collect()
8208        } else {
8209            inner
8210                .active_penalties
8211                .iter()
8212                .map(|penalty| {
8213                    (
8214                        penalty.matrix.clone(),
8215                        penalty.info.source.clone(),
8216                        penalty.info.normalization_scale,
8217                    )
8218                })
8219                .collect()
8220        };
8221
8222    let mut candidates = Vec::<PenaltyCandidate>::with_capacity(marginal_penalties.len());
8223    for (s_inner, source, base_scale) in marginal_penalties {
8224        let mut s_big = Array2::<f64>::zeros((q, q));
8225        for level in 0..n_levels {
8226            let start = level * p;
8227            s_big
8228                .slice_mut(s![start..start + p, start..start + p])
8229                .assign(&s_inner);
8230        }
8231        let (s_big, factor_smooth_scale) = normalize_penalty_in_constrained_space(&s_big);
8232        candidates.push(PenaltyCandidate {
8233            matrix: ConstructiveQuadratic::try_from_dense_psd(
8234                s_big,
8235                "factor-smooth shared penalty",
8236            )?,
8237            source,
8238            normalization_scale: base_scale * factor_smooth_scale,
8239            kronecker_factors: None,
8240            op: None,
8241        });
8242    }
8243
8244    // `Fs` is the random-effect flavour of a smooth: the per-group curve is an
8245    // exchangeable Gaussian *function*, so EVERY coefficient — including the
8246    // {const, linear} null space of the marginal wiggliness penalty — must be
8247    // shrinkable toward zero under its own shared variance. The wiggliness
8248    // penalty `S_wiggle` shapes curvature but leaves the per-group intercept and
8249    // slope (its null space) completely UNPENALIZED. With the null space free,
8250    // each group fits its own intercept and slope with NO partial pooling, so
8251    // the held-out per-subject forecast inherits the full no-pooling variance
8252    // and curves away from the true per-group line (gam#712 real arm, gam#713;
8253    // gam#903 sleepstudy forecast ran ~74% over the lme4 BLUP bar).
8254    //
8255    // mgcv's `bs="fs"` fixes this by penalizing each null-space dimension
8256    // SEPARATELY (`smooth.construct.fs.smooth.spec` adds one rank-1 penalty per
8257    // null coordinate), each replicated block-diagonally across levels under a
8258    // single shared smoothing parameter — so REML fits a distinct
8259    // random-intercept variance and random-slope variance, the partial pooling
8260    // that makes the forecast track lme4's correlated random-effect BLUP. A
8261    // single *combined* null penalty (one λ for intercept+slope together) cannot
8262    // express the typically very different intercept and slope variances, which
8263    // is the residual forecast gap. We mirror mgcv exactly: for each orthonormal
8264    // canonical null direction `z_k` of the marginal wiggliness penalty, add
8265    // `I_L ⊗ (z_k z_kᵀ)` as its own penalty. The marginal's combined double
8266    // penalty was disabled above, so the null space is penalized once, per
8267    // dimension. With linear data REML drives the curvature λ up and degrades
8268    // `fs` to a linear random slope (edf → ≈2/group); with genuine curvature the
8269    // wiggliness λ stays small and the wiggle survives (data-adaptive, not a
8270    // cap). Gated by `m_null_penalty_orders`: order ≥ 1 (default) enables the
8271    // per-dimension null penalties; `m=0` keeps the legacy combined double
8272    // penalty and adds nothing here.
8273    if use_per_dim_null
8274        && let Some(Some(z)) = inner
8275            .active_penalties
8276            .first()
8277            .map(|penalty| &penalty.null_eigenvectors)
8278        && z.nrows() == p
8279    {
8280        let z = canonical_nullspace_directions(z)?;
8281        for k in 0..z.ncols() {
8282            // Rank-1 marginal penalty `z_k z_kᵀ`, replicated block-diagonally
8283            // across levels into `I_L ⊗ (z_k z_kᵀ)`. Its own λ is one shared
8284            // variance for this null component (intercept or slope) across all
8285            // groups — the random-effect structure of mgcv `fs`.
8286            let zk = z.column(k);
8287            let mut p_k = Array2::<f64>::zeros((p, p));
8288            for a in 0..p {
8289                for b in 0..p {
8290                    p_k[[a, b]] = zk[a] * zk[b];
8291                }
8292            }
8293            let mut s_null = Array2::<f64>::zeros((q, q));
8294            for level in 0..n_levels {
8295                let start = level * p;
8296                s_null
8297                    .slice_mut(s![start..start + p, start..start + p])
8298                    .assign(&p_k);
8299            }
8300            let (s_null, null_scale) = normalize_penalty_in_constrained_space(&s_null);
8301            candidates.push(PenaltyCandidate {
8302                matrix: ConstructiveQuadratic::try_from_dense_psd(
8303                    s_null,
8304                    "factor-smooth null-function penalty",
8305                )?,
8306                source: PenaltySource::Primary,
8307                normalization_scale: null_scale,
8308                kronecker_factors: None,
8309                op: None,
8310            });
8311        }
8312    }
8313    let filtered = crate::basis::filter_penalty_candidates(candidates)?;
8314    let joint_null_rotation = crate::basis::compute_joint_null_rotation(&filtered.active)?;
8315    let mut dropped_penalties = inner.dropped_penalties;
8316    dropped_penalties.extend(filtered.dropped);
8317
8318    // Metadata: carry the marginal knot geometry + frozen levels so prediction
8319    // reconstructs an identical replicated design.
8320    let (knots, degree, periodic) = match &inner.metadata {
8321        BasisMetadata::BSpline1D {
8322            knots,
8323            periodic,
8324            degree,
8325            ..
8326        } => (
8327            knots.clone(),
8328            degree.unwrap_or(spec.marginal.degree),
8329            *periodic,
8330        ),
8331        other => {
8332            crate::bail_invalid_basis!(
8333                "factor smooth term '{}' produced an unexpected marginal metadata variant {:?}",
8334                term_name,
8335                other
8336            );
8337        }
8338    };
8339    let flavour_tag = match &spec.flavour {
8340        FactorSmoothFlavour::Fs { .. } => "fs",
8341        FactorSmoothFlavour::Sz => "sz",
8342        FactorSmoothFlavour::Re => "re",
8343    }
8344    .to_string();
8345    let metadata = BasisMetadata::FactorSmooth {
8346        continuous_cols: spec.continuous_cols.clone(),
8347        group_col,
8348        knots,
8349        degree,
8350        periodic,
8351        group_levels: levels,
8352        flavour: flavour_tag,
8353        // fs/re marginals are always B-spline; the cr marginal is sz-only and
8354        // handled on the dedicated Sz path above.
8355        marginal_is_cr: false,
8356    };
8357
8358    Ok(LocalSmoothTermBuild {
8359        dim: q,
8360        design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(dense)),
8361        // Exactly one group block is active on every row, so a common
8362        // marginal anchor contributes its lift once for that row.
8363        affine_offset: inner.affine_offset,
8364        active_penalties: filtered.active,
8365        joint_null_rotation,
8366        dropped_penalties,
8367        metadata,
8368        linear_constraints: None,
8369        box_reparam: false,
8370        kronecker_factored: None,
8371    })
8372}
8373
8374/// Resolve the grouping levels for a factor smooth: replay the frozen level
8375/// list when present (predict path), otherwise discover the sorted unique bit
8376/// patterns of the factor column (fit path).
8377pub fn resolve_factor_smooth_levels(
8378    data: ArrayView2<'_, f64>,
8379    group_col: usize,
8380    spec: &FactorSmoothSpec,
8381    term_name: &str,
8382) -> Result<Vec<u64>, BasisError> {
8383    if let Some(frozen) = &spec.group_frozen_levels {
8384        if frozen.is_empty() {
8385            crate::bail_invalid_basis!(
8386                "factor smooth term '{}' has an empty frozen level list",
8387                term_name
8388            );
8389        }
8390        return Ok(frozen
8391            .iter()
8392            .map(|&b| gam_data::canonical_level_bits(f64::from_bits(b)))
8393            .collect());
8394    }
8395    let mut bits: Vec<u64> = data
8396        .column(group_col)
8397        .iter()
8398        .map(|v| gam_data::canonical_level_bits(*v))
8399        .collect();
8400    bits.sort_by(|a, b| {
8401        f64::from_bits(*a)
8402            .partial_cmp(&f64::from_bits(*b))
8403            .unwrap_or(std::cmp::Ordering::Equal)
8404    });
8405    bits.dedup();
8406    Ok(bits)
8407}
8408
8409/// Marginal B-spline spec for a factor-smooth block. The marginal always builds
8410/// without an identifiability constraint (the per-level replication, not a
8411/// sum-to-zero side constraint, provides identifiability against the parametric
8412/// block). At predict time the marginal's knot geometry has already been pinned
8413/// into `marginal.knotspec` by the metadata replay, so the spec is used
8414/// verbatim aside from clearing the identifiability transform.
8415pub fn factor_smooth_marginal_for_replay(marginal: &BSplineBasisSpec) -> BSplineBasisSpec {
8416    let mut m = marginal.clone();
8417    m.identifiability = BSplineIdentifiability::None;
8418    m
8419}
8420
8421pub fn build_single_local_smooth_term(
8422    data: ArrayView2<'_, f64>,
8423    term: &SmoothTermSpec,
8424    workspace: &mut crate::basis::BasisWorkspace,
8425) -> Result<LocalSmoothTermBuild, BasisError> {
8426    term.basis.validate_scale_configuration()?;
8427    if term.shape != ShapeConstraint::None && !shape_supports_basis(term) {
8428        crate::bail_invalid_basis!(
8429            "ShapeConstraint::{:?} is unsupported for term '{}'",
8430            term.shape,
8431            term.name
8432        );
8433    }
8434    if let SmoothBasisSpec::ByVariable {
8435        inner,
8436        by_col,
8437        kind,
8438        by,
8439    } = &term.basis
8440    {
8441        ensure_by_variable_specs_match(kind, by, &term.name)?;
8442        let mut inner_basis = (**inner).clone();
8443        // Factor-level `by=` owns model-space centering (it centers the gated
8444        // block against the level indicator downstream). Defer the inner
8445        // basis's default pooled centering so the level block is not
8446        // double-centered down to `k-2` columns (#1427). Numeric-by smooths are
8447        // untouched: they are not row-gated to a level and keep ordinary
8448        // intercept centering.
8449        if matches!(by, ByVariableSpec::Level { .. }) {
8450            defer_inner_model_centering_to_factor_level_wrapper(&mut inner_basis);
8451        }
8452        let inner_term = SmoothTermSpec {
8453            name: term.name.clone(),
8454            basis: inner_basis,
8455            shape: term.shape,
8456            joint_null_rotation: None,
8457        };
8458        let built = build_single_local_smooth_term(data, &inner_term, workspace)?;
8459        return apply_by_variable_to_local_build(built, data, *by_col, by, &term.name);
8460    }
8461
8462    // BySmooth: a `by=` smooth that unifies numeric or factor modulation into a
8463    // single term.  Lower it here so the downstream match does not need an arm.
8464    if let SmoothBasisSpec::BySmooth { smooth, by_kind } = &term.basis {
8465        return build_by_smooth_local(data, term, smooth, by_kind, workspace);
8466    }
8467
8468    let mut built: BasisBuildResult = match &term.basis {
8469        SmoothBasisSpec::FactorSumToZero {
8470            inner,
8471            by_col,
8472            levels,
8473            ..
8474        } => {
8475            if *by_col >= data.ncols() {
8476                crate::bail_dim_basis!(
8477                    "term '{}' by column {} out of bounds for {} columns",
8478                    term.name,
8479                    by_col,
8480                    data.ncols()
8481                );
8482            }
8483            if levels.len() < 2 {
8484                crate::bail_invalid_basis!(
8485                    "sum-to-zero factor smooth term '{}' requires at least two levels",
8486                    term.name
8487                );
8488            }
8489            if term.shape != ShapeConstraint::None {
8490                crate::bail_invalid_basis!(
8491                    "ShapeConstraint::{:?} is unsupported for sum-to-zero factor smooth term '{}'",
8492                    term.shape,
8493                    term.name
8494                );
8495            }
8496            let inner_term = SmoothTermSpec {
8497                name: format!("{}::inner", term.name),
8498                basis: (**inner).clone(),
8499                shape: ShapeConstraint::None,
8500                joint_null_rotation: None,
8501            };
8502            let mut inner_built = build_single_local_smooth_term(data, &inner_term, workspace)?;
8503            if inner_built.affine_offset.is_some() {
8504                crate::bail_invalid_basis!(
8505                    "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",
8506                    term.name
8507                );
8508            }
8509            // Capture the marginal penalty's null directions BEFORE the penalty
8510            // vector is rebuilt below; the sum-to-zero null-space ridge replicates
8511            // these `z_k` into the contrast space (mgcv `bs="fs"` double-penalty).
8512            let inner_null_eigenvectors = inner_built
8513                .active_penalties
8514                .first()
8515                .and_then(|penalty| penalty.null_eigenvectors.clone());
8516            let base = inner_built
8517                .design
8518                .try_to_dense_by_chunks("sum-to-zero factor smooth")
8519                .map_err(BasisError::InvalidInput)?;
8520            let n = base.nrows();
8521            let p = base.ncols();
8522            let l_minus_one = levels.len() - 1;
8523            // Canonicalize the stored level keys once so signed-zero / NaN codes
8524            // match regardless of how the level set was interned (#2145/#2146).
8525            let canon_levels: Vec<u64> = levels
8526                .iter()
8527                .map(|&b| gam_data::canonical_level_bits(f64::from_bits(b)))
8528                .collect();
8529            let mut dense = Array2::<f64>::zeros((n, p * l_minus_one));
8530            for i in 0..n {
8531                let bits = gam_data::canonical_level_bits(data[[i, *by_col]]);
8532                let level_idx = canon_levels
8533                    .iter()
8534                    .position(|b| *b == bits)
8535                    .ok_or_else(|| {
8536                        BasisError::InvalidInput(format!(
8537                            "sum-to-zero factor smooth term '{}' saw an unseen level at row {}",
8538                            term.name,
8539                            i + 1
8540                        ))
8541                    })?;
8542                if level_idx < l_minus_one {
8543                    let start = level_idx * p;
8544                    dense
8545                        .slice_mut(s![i, start..start + p])
8546                        .assign(&base.row(i));
8547                } else {
8548                    for level in 0..l_minus_one {
8549                        let start = level * p;
8550                        dense
8551                            .slice_mut(s![i, start..start + p])
8552                            .assign(&base.row(i).mapv(|v| -v));
8553                    }
8554                }
8555            }
8556            let mut candidates = Vec::<PenaltyCandidate>::with_capacity(
8557                inner_built.active_penalties.len() * levels.len(),
8558            );
8559            // Replicate each marginal penalty into the sum-to-zero contrast
8560            // space. With `L-1` free deviation blocks and the reference level
8561            // `d_L = -Σ_{k<L} d_k`, the marginal penalty summed over ALL `L`
8562            // levels, `Σ_{k=1}^{L} d_kᵀ S d_k`, expands to the `(I + 11ᵀ) ⊗ S`
8563            // contrast form (factor 2 on the diagonal blocks, 1 off-diagonal).
8564            //
8565            // PER-GROUP SMOOTHING PARAMETERS (#1074). mgcv's `bs="sz"` does NOT
8566            // pool that sum under one λ: `smooth.construct.sz` emits ONE penalty
8567            // matrix per factor level (here 6 separate `S`s, each with its own
8568            // smoothing parameter), so REML can shrink a low-amplitude group's
8569            // deviation curve hard while leaving a high-amplitude group nearly
8570            // unpenalized. A single shared wiggliness λ (the old construction)
8571            // forces every group to the SAME curvature budget, so a group whose
8572            // true curve is flat drags curvature into the noise of the busy
8573            // groups and vice-versa — systematic truth-recovery loss even when
8574            // the pooled total edf matches mgcv's (the observed `sz` 1.23× gap).
8575            //
8576            // We mirror mgcv exactly by splitting the per-marginal penalty
8577            // `Σ_{k=1}^{L} d_kᵀ S d_k` back into its `L` independent
8578            // rank-controlled summands BEFORE mapping to the contrast space, each
8579            // carrying its own λ:
8580            //   * level k < L (free block):  `d_kᵀ S d_k` → block-diagonal
8581            //     `(e_k e_kᵀ) ⊗ S`  (only the (k,k) block is `S`).
8582            //   * level L (reference):       `d_Lᵀ S d_L = (Σ_{j<L} d_j)ᵀ S (·)`
8583            //     → the fully-coupled `(11ᵀ) ⊗ S` block.
8584            // Summed at equal λ these `L` blocks recover the old `(I + 11ᵀ) ⊗ S`
8585            // exactly (`Σ_k e_k e_kᵀ = I`), so this is a strict generalization:
8586            // the pooled fit is still reachable, REML only GAINS the freedom to
8587            // spend curvature per group. The zero-sum reparameterization (hence
8588            // the `sz` vs `fs` identifiability) is untouched.
8589            //
8590            // `which_level ∈ 0..=l_minus_one`: `< l_minus_one` selects the single
8591            // free deviation block; `== l_minus_one` selects the reference-level
8592            // coupling block.
8593            let stz_per_group_penalty =
8594                |s_inner: &Array2<f64>, which_level: usize| -> Array2<f64> {
8595                    let mut s_big = Array2::<f64>::zeros((p * l_minus_one, p * l_minus_one));
8596                    if which_level < l_minus_one {
8597                        // (e_k e_kᵀ) ⊗ S: a single diagonal block.
8598                        let k = which_level;
8599                        let mut block = s_big.slice_mut(s![k * p..(k + 1) * p, k * p..(k + 1) * p]);
8600                        block.assign(s_inner);
8601                    } else {
8602                        // (11ᵀ) ⊗ S: every block (diagonal and off-diagonal) is S.
8603                        for a in 0..l_minus_one {
8604                            for b in 0..l_minus_one {
8605                                let mut block =
8606                                    s_big.slice_mut(s![a * p..(a + 1) * p, b * p..(b + 1) * p]);
8607                                block.assign(s_inner);
8608                            }
8609                        }
8610                    }
8611                    s_big
8612                };
8613            for base_penalty in &inner_built.active_penalties {
8614                // Emit `L` independent per-level blocks for this marginal penalty.
8615                for which_level in 0..=l_minus_one {
8616                    let raw = stz_per_group_penalty(&base_penalty.matrix, which_level);
8617                    let (s_big, group_scale) = normalize_penalty_in_constrained_space(&raw);
8618                    candidates.push(PenaltyCandidate {
8619                        matrix: ConstructiveQuadratic::try_from_dense_psd(
8620                            s_big,
8621                            "grouped factor-smooth penalty",
8622                        )?,
8623                        source: base_penalty.info.source.clone(),
8624                        normalization_scale: base_penalty.info.normalization_scale * group_scale,
8625                        kronecker_factors: None,
8626                        op: None,
8627                    });
8628                }
8629            }
8630
8631            // Null-space ridge, mirroring the `bs="fs"` double-penalty
8632            // construction (#1605, same defect class as #700/#712/#713). The
8633            // marginal wiggliness penalty `S` shapes curvature but leaves the
8634            // {const, linear} null space of each deviation curve COMPLETELY
8635            // unpenalized. With that null space free, the single combined
8636            // wiggliness smoothing parameter cannot separate the per-group
8637            // intercept/slope variance from the curvature variance, so REML
8638            // parks the wiggliness `λ` high — over-smoothing (under-fitting) the
8639            // deviation blocks even when the truth lives in their span (the `sz`
8640            // recovery gap vs the `fs` superset). mgcv's `bs="fs"` fixes the
8641            // analogous gap by penalizing each null-space dimension SEPARATELY
8642            // under its own shared variance; we mirror that here while keeping
8643            // the zero-sum reparameterization, so the constraint (and the
8644            // identifiability of `sz` vs `fs`) is preserved. For each orthonormal
8645            // canonical null direction `z_k` of the marginal penalty, add the
8646            // rank-1 marginal penalty `z_k z_kᵀ` mapped into the SAME `(I + 11ᵀ)`
8647            // sum-to-zero contrast space, each carrying its own `λ`.
8648            if let Some(z) = inner_null_eigenvectors.as_ref()
8649                && z.nrows() == p
8650            {
8651                let z = canonical_nullspace_directions(z)?;
8652                for k in 0..z.ncols() {
8653                    let zk = z.column(k);
8654                    let mut p_k = Array2::<f64>::zeros((p, p));
8655                    for a in 0..p {
8656                        for b in 0..p {
8657                            p_k[[a, b]] = zk[a] * zk[b];
8658                        }
8659                    }
8660                    // Null ridges stay POOLED (the `(I + 11ᵀ) ⊗ z_k z_kᵀ` form):
8661                    // they govern the per-group intercept/slope shrinkage, which
8662                    // mgcv pools under one variance even for `sz`; only the
8663                    // curvature (wiggliness) penalty is split per group above.
8664                    let stz_pooled_null = {
8665                        let mut s_big = Array2::<f64>::zeros((p * l_minus_one, p * l_minus_one));
8666                        for a in 0..l_minus_one {
8667                            for b in 0..l_minus_one {
8668                                let factor = if a == b { 2.0 } else { 1.0 };
8669                                let mut block =
8670                                    s_big.slice_mut(s![a * p..(a + 1) * p, b * p..(b + 1) * p]);
8671                                block.assign(&p_k.mapv(|v| v * factor));
8672                            }
8673                        }
8674                        s_big
8675                    };
8676                    let (s_null, null_scale) =
8677                        normalize_penalty_in_constrained_space(&stz_pooled_null);
8678                    candidates.push(PenaltyCandidate {
8679                        matrix: ConstructiveQuadratic::try_from_dense_psd(
8680                            s_null,
8681                            "grouped factor-smooth null penalty",
8682                        )?,
8683                        source: PenaltySource::DoublePenaltyNullspace,
8684                        normalization_scale: null_scale,
8685                        kronecker_factors: None,
8686                        op: None,
8687                    });
8688                }
8689            }
8690            let filtered = crate::basis::filter_penalty_candidates(candidates)?;
8691            let mut dropped_penalties = std::mem::take(&mut inner_built.dropped_penalties);
8692            dropped_penalties.extend(filtered.dropped);
8693            inner_built.dim = p * l_minus_one;
8694            inner_built.design =
8695                DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(dense));
8696            inner_built.active_penalties = filtered.active;
8697            inner_built.dropped_penalties = dropped_penalties;
8698            inner_built.joint_null_rotation =
8699                crate::basis::compute_joint_null_rotation(&inner_built.active_penalties)?;
8700            inner_built.kronecker_factored = None;
8701            return Ok(inner_built);
8702        }
8703        SmoothBasisSpec::BSpline1D { feature_col, spec } => {
8704            if *feature_col >= data.ncols() {
8705                crate::bail_dim_basis!(
8706                    "term '{}' feature column {} out of bounds for {} columns",
8707                    term.name,
8708                    feature_col,
8709                    data.ncols()
8710                );
8711            }
8712            let mut spec_local = spec.clone();
8713            if term.shape != ShapeConstraint::None {
8714                // Shape-constrained B-splines are anchored by construction.
8715                // Sum-to-zero side constraints conflict with monotonic/convex cones.
8716                spec_local.identifiability = BSplineIdentifiability::None;
8717            }
8718            // Endpoint boundary conditions are structural for B-splines: the
8719            // basis builder bakes their homogeneous nullspace transform into
8720            // the design, penalties, and stored raw-basis transform.
8721            build_bspline_basis_1d(data.column(*feature_col), &spec_local)?
8722        }
8723        SmoothBasisSpec::ThinPlate {
8724            feature_cols,
8725            spec,
8726            input_scale,
8727        } => {
8728            if term.shape != ShapeConstraint::None {
8729                if feature_cols.len() != 1 {
8730                    crate::bail_invalid_basis!(
8731                        "ShapeConstraint::{:?} for term '{}' on ThinPlate basis requires exactly 1 feature axis; found {}",
8732                        term.shape,
8733                        term.name,
8734                        feature_cols.len()
8735                    );
8736                }
8737            }
8738            let frame = term.basis.scale_contract().normalize_euclidean_frame(
8739                select_columns(data, feature_cols)?,
8740                *input_scale,
8741                Some(spec.length_scale),
8742            )?;
8743            let x = frame.coordinates;
8744            let realized_input_scale = frame.input_scale;
8745            let length_scale_eff = frame
8746                .length_scale
8747                .expect("ThinPlate declares a required length-scale coordinate");
8748            let mut spec_local = spec.clone();
8749            spec_local.length_scale = length_scale_eff;
8750            if matches!(
8751                spec_local.identifiability,
8752                SpatialIdentifiability::OrthogonalToParametric
8753            ) {
8754                spec_local.identifiability = SpatialIdentifiability::None;
8755            }
8756            let mut result = build_thin_plate_basis(x.view(), &spec_local).map_err(|err| {
8757                rewrite_thin_plate_knots_error(err, &term.name, feature_cols.len(), spec)
8758            })?;
8759            // Inject the input scale into metadata; also restore the user's
8760            // original length_scale (not the σ_geom-compensated one) so a
8761            // metadata-driven rebuild that re-applies compensation does not
8762            // double-divide. The build may auto-promote to Duchon when
8763            // canonical TPS is infeasible (k < polynomial-nullspace size);
8764            // in that case patch the Duchon metadata variant so predict-time
8765            // round-trips through the same standardized data path.
8766            match &mut result.metadata {
8767                BasisMetadata::ThinPlate {
8768                    input_scale: metadata_scale,
8769                    length_scale,
8770                    ..
8771                } => {
8772                    *metadata_scale = realized_input_scale;
8773                    *length_scale = spec.length_scale;
8774                }
8775                BasisMetadata::Duchon {
8776                    input_scale: metadata_scale,
8777                    length_scale,
8778                    ..
8779                } => {
8780                    // Auto-promotion (canonical TPS infeasible at this (d, k)).
8781                    // Since #1091 the promotion does NOT forward the incoming
8782                    // σ_geom-compensated `spec_local.length_scale` to the Duchon
8783                    // builder — it DISCARDS it and substitutes the geometric mean
8784                    // of the center pairwise distances (`promotion_length_scale`,
8785                    // the natural radial-kernel scale where κ·r ≈ O(1)). So the
8786                    // realized kernel bandwidth recorded in this metadata bears no
8787                    // fixed relation to the user's `spec.length_scale`; clobbering
8788                    // it to the user-facing value (the pre-#1091 behavior) makes
8789                    // freeze→replay re-derive `compensate(spec.length_scale, σ) ≠
8790                    // promotion_length_scale`, evaluating the kernel at the wrong
8791                    // bandwidth and corrupting the replayed design (#1091 broke the
8792                    // e7ff5ed83 freeze contract for the auto-promoted path).
8793                    //
8794                    // The freeze→replay round trip rebuilds through the Duchon arm,
8795                    // which re-applies σ_geom compensation:
8796                    //   replay_eff = compensate(metadata.length_scale, σ)
8797                    //              = metadata.length_scale / σ_geom.
8798                    // For replay_eff to reproduce the realized `promotion_length_scale`
8799                    // we must store the UN-compensated value `promotion_length_scale
8800                    // · σ_geom`. `compensate(1.0, σ) = 1/σ_geom`, so divide the
8801                    // realized scale by it to multiply back through σ_geom. With no
8802                    // standardization. Restore original units before freezing.
8803                    if let Some(realized) = *length_scale {
8804                        *length_scale = Some(realized * realized_input_scale.get());
8805                    }
8806                    *metadata_scale = realized_input_scale;
8807                }
8808                _ => {}
8809            }
8810            result
8811        }
8812        SmoothBasisSpec::Sphere { feature_cols, spec } => {
8813            if term.shape != ShapeConstraint::None {
8814                crate::bail_invalid_basis!(
8815                    "ShapeConstraint::{:?} for term '{}' is not supported on spherical splines",
8816                    term.shape,
8817                    term.name
8818                );
8819            }
8820            let x = select_columns(data, feature_cols)?;
8821            build_spherical_spline_basis(x.view(), spec)?
8822        }
8823        SmoothBasisSpec::ConstantCurvature { feature_cols, spec } => {
8824            if term.shape != ShapeConstraint::None {
8825                crate::bail_invalid_basis!(
8826                    "ShapeConstraint::{:?} for term '{}' is not supported on constant-curvature smooths",
8827                    term.shape,
8828                    term.name
8829                );
8830            }
8831            // Chart coordinates are consumed verbatim: NO auto-standardization.
8832            // Rescaling axes would change the chart gauge `1 + κ‖x‖²` and
8833            // silently redefine which curvature κ refers to (the same point
8834            // cloud at a different chart scale has a different κ̂); the user's
8835            // coordinates ARE the geometry here, exactly as for the sphere
8836            // smooth's (lat, lon).
8837            let x = select_columns(data, feature_cols)?;
8838            build_constant_curvature_basis(x.view(), spec)?
8839        }
8840        SmoothBasisSpec::MeasureJet {
8841            feature_cols,
8842            spec,
8843            input_scale,
8844        } => {
8845            if term.shape != ShapeConstraint::None {
8846                crate::bail_invalid_basis!(
8847                    "ShapeConstraint::{:?} for term '{}' is not supported on measure-jet smooths",
8848                    term.shape,
8849                    term.name
8850                );
8851            }
8852            // The typed scale contract owns the intentionally asymmetric
8853            // fresh/replay rule: fresh explicit ranges are in original units,
8854            // while a frozen MeasureJet range is already in its realized frame.
8855            let frame = term.basis.scale_contract().normalize_euclidean_frame(
8856                select_columns(data, feature_cols)?,
8857                *input_scale,
8858                Some(spec.length_scale),
8859            )?;
8860            let x = frame.coordinates;
8861            let realized_input_scale = frame.input_scale;
8862            let length_scale_eff = frame
8863                .length_scale
8864                .expect("MeasureJet declares a required length-scale coordinate");
8865            let mut spec_local = spec.clone();
8866            spec_local.length_scale = length_scale_eff;
8867            let mut result = build_measure_jet_basis(x.view(), &spec_local)?;
8868            if let BasisMetadata::MeasureJet {
8869                input_scale: metadata_scale,
8870                ..
8871            } = &mut result.metadata
8872            {
8873                *metadata_scale = realized_input_scale;
8874            }
8875            result
8876        }
8877        SmoothBasisSpec::Matern {
8878            feature_cols,
8879            spec,
8880            input_scale,
8881        } => {
8882            if term.shape != ShapeConstraint::None {
8883                if feature_cols.len() != 1 {
8884                    crate::bail_invalid_basis!(
8885                        "ShapeConstraint::{:?} for term '{}' on Matern basis requires exactly 1 feature axis; found {}",
8886                        term.shape,
8887                        term.name,
8888                        feature_cols.len()
8889                    );
8890                }
8891            }
8892            let original_length_scale = spec.length_scale.resolved().ok_or_else(|| {
8893                BasisError::InvalidInput(format!(
8894                    "term '{}' reached Matérn construction before its Auto length scale was resolved",
8895                    term.name
8896                ))
8897            })?;
8898            let frame = term.basis.scale_contract().normalize_euclidean_frame(
8899                select_columns(data, feature_cols)?,
8900                *input_scale,
8901                Some(original_length_scale),
8902            )?;
8903            let x = frame.coordinates;
8904            let realized_input_scale = frame.input_scale;
8905            let length_scale_eff = frame
8906                .length_scale
8907                .expect("Matérn declares a required length-scale coordinate");
8908            let mut spec_local = spec.clone();
8909            spec_local.length_scale.set_resolved(length_scale_eff);
8910            let mut result = build_matern_basiswithworkspace(x.view(), &spec_local, workspace)?;
8911            if let BasisMetadata::Matern {
8912                input_scale: metadata_scale,
8913                length_scale,
8914                ..
8915            } = &mut result.metadata
8916            {
8917                *metadata_scale = realized_input_scale;
8918                *length_scale = original_length_scale;
8919            }
8920            result
8921        }
8922        SmoothBasisSpec::Duchon {
8923            feature_cols,
8924            spec,
8925            input_scale,
8926        } => {
8927            if term.shape != ShapeConstraint::None {
8928                if feature_cols.len() != 1 {
8929                    crate::bail_invalid_basis!(
8930                        "ShapeConstraint::{:?} for term '{}' on Duchon basis requires exactly 1 feature axis; found {}",
8931                        term.shape,
8932                        term.name,
8933                        feature_cols.len()
8934                    );
8935                }
8936            }
8937            let frame = term.basis.scale_contract().normalize_euclidean_frame(
8938                select_columns(data, feature_cols)?,
8939                *input_scale,
8940                spec.length_scale,
8941            )?;
8942            let x = frame.coordinates;
8943            let realized_input_scale = frame.input_scale;
8944            let length_scale_eff = frame.length_scale;
8945            let mut spec_local = spec.clone();
8946            spec_local.length_scale = length_scale_eff;
8947            // The Duchon input axis is standardized in place above (`x → x/σ`,
8948            // scale-only, no centering). A 1-D cyclic boundary `[start, end)`
8949            // declared in ORIGINAL covariate units must move into that same
8950            // standardized frame, or the periodic wrap in
8951            // `build_periodic_duchon_basis_1d` (which the cyclic-boundary
8952            // dispatch in `build_duchon_basis_uncached` normalizes onto) folds
8953            // the standardized coordinate against an original-unit period: the
8954            // seam never closes and the basis silently degrades to
8955            // non-periodic (#1074: `duchon(x, periodic=true)` predictions
8956            // diverged across the wrap, f(0) ≠ f(2π)). Rescale by the same
8957            // 1/σ applied to the data so training and predict share one
8958            // periodic geometry.
8959            if let crate::basis::OneDimensionalBoundary::Cyclic { start, end } =
8960                spec_local.boundary.clone()
8961            {
8962                spec_local.boundary = crate::basis::OneDimensionalBoundary::Cyclic {
8963                    start: realized_input_scale.to_standardized_units(start),
8964                    end: realized_input_scale.to_standardized_units(end),
8965                };
8966            }
8967            // The SAME original-units-vs-standardized-frame reasoning applies
8968            // to `spec.periodic` (the per-axis period vector the position API
8969            // and mixed-periodicity tensor paths use): each declared period is
8970            // in original covariate units and must be divided by that axis's
8971            // uniform input scale, or the wrap folds standardized coordinates against an
8972            // original-unit period (the #1074 seam failure, previously fixed
8973            // only for the 1-D `boundary` spelling above).
8974            if let Some(periods) = spec_local.periodic.as_mut() {
8975                for axis_period in periods {
8976                    if let Some(period) = axis_period.as_mut() {
8977                        *period = realized_input_scale.to_standardized_units(*period);
8978                    }
8979                }
8980            }
8981            if matches!(
8982                spec_local.identifiability,
8983                SpatialIdentifiability::OrthogonalToParametric
8984            ) {
8985                spec_local.identifiability = SpatialIdentifiability::None;
8986            }
8987            let mut result = build_duchon_basiswithworkspace(x.view(), &spec_local, workspace)?;
8988            if let BasisMetadata::Duchon {
8989                input_scale: metadata_scale,
8990                length_scale,
8991                periodic,
8992                ..
8993            } = &mut result.metadata
8994            {
8995                *metadata_scale = realized_input_scale;
8996                *length_scale = spec.length_scale;
8997                // Same convention as `length_scale`: metadata (and hence the
8998                // frozen replay spec design_freezing copies it into) always
8999                // stores the period in ORIGINAL covariate units, and the
9000                // standardization rescale above recomputes the standardized
9001                // period fresh from `input_scale` on EVERY build — fresh fit
9002                // and frozen replay alike — so the compensation stays
9003                // idempotent with no fit-vs-replay branch. Leaving the
9004                // builder-resolved (standardized-frame) period here would
9005                // double-divide on replay. Invariant this relies on: every
9006                // producer that sets a Cyclic `boundary` also sets
9007                // `spec.periodic` from the same original-units source (the
9008                // formula DSL does; see `parse_periodic_axes_option` /
9009                // `parse_cyclic_boundary` in term_builder.rs), so the pristine
9010                // `spec.periodic` is a valid original-units record for the
9011                // boundary spelling too.
9012                if spec.periodic.is_some() || spec.boundary.period().is_some() {
9013                    *periodic = spec
9014                        .periodic
9015                        .clone()
9016                        .or_else(|| spec.boundary.period().map(|(_, _, p)| vec![Some(p)]));
9017                }
9018            }
9019            result
9020        }
9021        SmoothBasisSpec::Pca {
9022            feature_cols,
9023            basis_matrix,
9024            centered,
9025            smooth_penalty,
9026            center_mean,
9027            pca_basis_path,
9028            chunk_size,
9029        } => {
9030            if term.shape != ShapeConstraint::None {
9031                crate::bail_invalid_basis!(
9032                    "ShapeConstraint::{:?} for term '{}' is not supported on Pca basis",
9033                    term.shape,
9034                    term.name
9035                );
9036            }
9037            build_pca_smooth_basis(
9038                data,
9039                feature_cols,
9040                basis_matrix,
9041                *centered,
9042                *smooth_penalty,
9043                center_mean.as_ref(),
9044                pca_basis_path.as_ref(),
9045                *chunk_size,
9046            )?
9047        }
9048        SmoothBasisSpec::TensorBSpline { feature_cols, spec } => {
9049            build_tensor_bspline_basis(data, feature_cols, spec)?
9050        }
9051        SmoothBasisSpec::ByVariable { .. } => {
9052            crate::bail_invalid_basis!(
9053                "internal: ByVariable smooths must return before inner basis dispatch"
9054            );
9055        }
9056        SmoothBasisSpec::BySmooth { .. } => {
9057            crate::bail_invalid_basis!("internal: BySmooth smooths must be lowered to ByVariable before inner basis dispatch"
9058                    .to_string(),);
9059        }
9060        SmoothBasisSpec::FactorSmooth { spec } => {
9061            if term.shape != ShapeConstraint::None {
9062                crate::bail_invalid_basis!(
9063                    "ShapeConstraint::{:?} is unsupported for factor smooth term '{}'",
9064                    term.shape,
9065                    term.name
9066                );
9067            }
9068            return build_factor_smooth(data, spec, &term.name, workspace);
9069        }
9070    };
9071
9072    // The Matérn design ALWAYS uses the operator-collocation {mass, tension,
9073    // stiffness} penalty triplet, overriding whatever penalty
9074    // `build_matern_basis_seeded` produced for the `double_penalty` flag.
9075    //
9076    // #1074 investigated swapping this for the genuine RKHS kernel penalty
9077    // `β' K_CC β` (mgcv `bs="gp"` / fields kriging) on the theory that the
9078    // operator triplet under-smooths the rougher half-integer kernels. MSI
9079    // truth-recovery measurement REFUTED that: the kernel penalty did NOT
9080    // improve ν=3/2 recovery (`matern(x,nu=1.5)` RMSE-vs-truth stayed 0.0554)
9081    // and it REGRESSED the high-frequency-init guard — `matern(x,nu≥5/2)` on
9082    // sin(2π·8·x) collapsed (span 0.53, RMSE 0.70) because the single RKHS
9083    // norm over-smooths a high-frequency truth where the Sobolev-order operator
9084    // dials do not. The operator triplet is therefore retained as the Matérn
9085    // penalty, and the κ-optimizer re-key / ψ-derivative paths route through the
9086    // same triplet builder so the block count stays ψ-stable (#1270).
9087    if let SmoothBasisSpec::Matern { .. } = &term.basis {
9088        let filtered = matern_operator_penalty_triplet_from_metadata(&built.metadata)?;
9089        built.active_penalties = filtered.active;
9090        built.dropped_penalties = filtered.dropped;
9091    }
9092
9093    if built.affine_offset.is_some() && term.shape != ShapeConstraint::None {
9094        crate::bail_invalid_basis!(
9095            "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",
9096            term.shape,
9097            term.name
9098        );
9099    }
9100    let p_local = built.design.ncols();
9101    let affine_offset = built.affine_offset;
9102    let mut metadata = built.metadata.clone();
9103    // Extract factored Kronecker representation before consuming fields.
9104    // Invalidate it if shape transforms will be applied (they break structure).
9105    let kron_factored = if term.shape == ShapeConstraint::None {
9106        built.kronecker_factored
9107    } else {
9108        None
9109    };
9110    let mut design_t = built.design;
9111    let mut penalties_t = built.active_penalties;
9112    let mut dropped_penalties_t = built.dropped_penalties;
9113    if matches!(
9114        spatial_identifiability_policy(term),
9115        Some(SpatialIdentifiability::OrthogonalToParametric)
9116    ) {
9117        metadata = freeze_raw_spatial_metadata(metadata, design_t.ncols());
9118    }
9119
9120    let use_box_reparam =
9121        term.shape != ShapeConstraint::None && shape_uses_box_reparameterization(&term.basis);
9122    if let Some((order, sign)) = shape_order_and_sign(term.shape)
9123        && use_box_reparam
9124    {
9125        // Order 1 (monotone): the plain first-difference cone θ_{i+1}−θ_i ≥ 0 is
9126        // the control-polygon monotonicity criterion, which is independent of
9127        // Greville-abscissa spacing (it only fixes the *sign* of consecutive
9128        // control-point gaps), so the integer-difference transform is exact.
9129        //
9130        // Order 2 (convex/concave): the plain second-difference cone is only
9131        // correct for evenly spaced Greville abscissae. gam's B-splines are
9132        // clamped (and may use quantile knots), so the abscissae are not
9133        // uniform and the geometrically-correct cone is the second *divided*
9134        // difference. Build the knot-span-scaled transform so γ_{≥2} ≥ 0
9135        // certifies convexity of the function, not of the raw coefficient
9136        // index. Periodic splines are rejected by the exact-support gate: their
9137        // cyclic coefficient chart cannot use this open divided-difference cone.
9138        let t = if order == 2 {
9139            let (knots, degree) = match &metadata {
9140                BasisMetadata::BSpline1D {
9141                    knots,
9142                    degree: Some(degree),
9143                    periodic,
9144                    ..
9145                } if periodic.is_none() => (knots, *degree),
9146                _ => {
9147                    crate::bail_invalid_basis!(
9148                        "shape-constrained convex/concave term '{}' requires realized open B-spline knot and degree metadata",
9149                        term.name
9150                    );
9151                }
9152            };
9153            let spans = bspline_first_derivative_control_spans(knots.view(), degree)?;
9154            if spans.len() + 1 != p_local {
9155                crate::bail_invalid_basis!(
9156                    "shape-constraint derivative-control span count {} does not match basis dim {} for term '{}'",
9157                    spans.len(),
9158                    p_local,
9159                    term.name
9160                );
9161            }
9162            convex_derivative_control_transform_matrix(&spans, sign)?
9163        } else {
9164            cumulative_sum_transform_matrix(p_local, order, sign)
9165        };
9166        // Coefficient-side transform: wrap the design in an operator that
9167        // applies T on the coefficient side, preserving sparsity/operator
9168        // structure of the inner design.
9169        let inner_dense = match design_t {
9170            DesignMatrix::Dense(d) => d,
9171            DesignMatrix::Sparse(sp) => gam_linalg::matrix::DenseDesignMatrix::from(
9172                sp.try_to_dense_arc("shape-constrained coefficient transform")
9173                    .map_err(BasisError::InvalidInput)?,
9174            ),
9175        };
9176        let coeff_op =
9177            gam_linalg::matrix::CoefficientTransformOperator::new(inner_dense, t.clone()).map_err(
9178                |e| BasisError::InvalidInput(format!("CoefficientTransformOperator: {e}")),
9179            )?;
9180        design_t = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(
9181            coeff_op,
9182        )));
9183        // `β = Tγ` is an invertible change of coefficient chart. Every
9184        // physical quadratic functional, including the function-space
9185        // null-component penalty, therefore transforms by the same congruence
9186        // `S_γ = Tᵀ S_β T`. Rebuilding `ZZᵀ` in the γ Euclidean metric
9187        // would change the represented functional under this harmless chart
9188        // change and violate SPEC 5.
9189        for penalty in &mut penalties_t {
9190            let tt_s = fast_atb(&t, &penalty.matrix);
9191            penalty.matrix = fast_ab(&tt_s, &t);
9192            penalty.op = None;
9193            penalty.info.kronecker_factors = None;
9194        }
9195    }
9196    let penalty_candidates = penalties_t
9197        .into_iter()
9198        .map(|penalty| -> Result<PenaltyCandidate, BasisError> {
9199            let ActivePenalty {
9200                matrix,
9201                op: op_in,
9202                info,
9203                ..
9204            } = penalty;
9205            let (matrix, c_new) = normalize_penalty_in_constrained_space(&matrix);
9206            let normalization_scale = info.normalization_scale * c_new;
9207            let op_scale = 1.0 / c_new;
9208            let kronecker_scale = 1.0 / c_new;
9209            // Frobenius rescale: wrap inner op in `ScaledPenaltyOp(1/c_new)`
9210            // so `op.as_dense() == matrix` post-normalization.
9211            let scaled_op = if op_scale > 0.0 && op_scale.is_finite() {
9212                op_in.map(|op| {
9213                    std::sync::Arc::new(crate::analytic_penalties::ScaledPenaltyOp::new(
9214                        op, op_scale,
9215                    ))
9216                        as std::sync::Arc<dyn crate::analytic_penalties::PenaltyOp>
9217                })
9218            } else {
9219                None
9220            };
9221            let kronecker_factors = info.kronecker_factors.map(|mut factors| {
9222                if let Some(first) = factors.first_mut() {
9223                    first.mapv_inplace(|v| v * kronecker_scale);
9224                }
9225                factors
9226            });
9227            Ok(PenaltyCandidate {
9228                matrix: ConstructiveQuadratic::try_from_dense_psd(
9229                    matrix,
9230                    "shape-constrained transformed penalty",
9231                )?,
9232                source: info.source,
9233                normalization_scale,
9234                kronecker_factors,
9235                op: scaled_op,
9236            })
9237        })
9238        .collect::<Result<Vec<_>, _>>()?;
9239    let filtered = crate::basis::filter_penalty_candidates(penalty_candidates)?;
9240    dropped_penalties_t.extend(filtered.dropped);
9241    // Joint-null absorption rotation. Fresh fit specs compute Q from the final
9242    // per-smooth penalty set (after all in-smooth reparameterizations have
9243    // already been applied). Frozen specs already carry the complete realized
9244    // coefficient chart in their `FrozenTransform`; recomputing Q there would
9245    // rotate an already-frozen chart a second time and desynchronize value
9246    // rebuilds from derivative operators.
9247    //
9248    // Kronecker-factored smooths (tensor B-splines under `TensorBSplineIdentifiability::None`)
9249    // carry their joint penalty as `Σ_d S_d` with `S_d = I ⊗ … ⊗ S_d^{1D} ⊗ … ⊗ I`.
9250    // The joint null space is the tensor of marginal nulls and is handled directly
9251    // by the REML runtime's `kronecker_penalty_system` path (see
9252    // `runtime.rs:8334-8344`). Applying a dense (p × p) Q here would densify
9253    // `X_raw = mx ⊗ my` into `X_raw · Q`, destroying the Kronecker product
9254    // structure that the runtime relies on for fast log-det/derivative
9255    // assembly — and the rotation block at the wrapper site also unconditionally
9256    // wipes `kronecker_factored`, leaving the runtime to fall back to the
9257    // dense per-block log-det. Skip the rotation for Kronecker-factored terms
9258    // so the factored representation survives end-to-end.
9259    let joint_null_rotation = match term.joint_null_rotation.clone() {
9260        Some(persisted) => Some(persisted),
9261        None if smooth_has_frozen_identifiability(term) => None,
9262        None if kron_factored.is_some() => None,
9263        None => crate::basis::compute_joint_null_rotation(&filtered.active)?,
9264    };
9265
9266    Ok(LocalSmoothTermBuild {
9267        dim: p_local,
9268        design: design_t,
9269        affine_offset,
9270        active_penalties: filtered.active,
9271        joint_null_rotation,
9272        dropped_penalties: dropped_penalties_t,
9273        metadata,
9274        linear_constraints: None,
9275        box_reparam: use_box_reparam,
9276        kronecker_factored: kron_factored,
9277    })
9278}
9279
9280pub fn build_smooth_design(
9281    data: ArrayView2<'_, f64>,
9282    terms: &[SmoothTermSpec],
9283) -> Result<RawSmoothDesign, BasisError> {
9284    let mut ws = crate::basis::BasisWorkspace::new();
9285    build_smooth_design_withworkspace(data, terms, &mut ws)
9286}
9287
9288/// Like `build_smooth_design`, but honors the caller workspace policy while
9289/// building each planned smooth term with an independent per-term workspace.
9290///
9291/// Independent workspaces avoid shared mutable distance-cache state during the
9292/// parallel term build; the final design, penalties, and metadata are assembled
9293/// in the original smooth-term order.
9294pub fn build_smooth_design_withworkspace(
9295    data: ArrayView2<'_, f64>,
9296    terms: &[SmoothTermSpec],
9297    workspace: &mut crate::basis::BasisWorkspace,
9298) -> Result<RawSmoothDesign, BasisError> {
9299    validate_smooth_terms_finite_inputs(data, terms)?;
9300    build_smooth_design_withworkspace_unvalidated(data, terms, workspace)
9301}
9302
9303pub fn build_smooth_design_withworkspace_unvalidated(
9304    data: ArrayView2<'_, f64>,
9305    terms: &[SmoothTermSpec],
9306    workspace: &mut crate::basis::BasisWorkspace,
9307) -> Result<RawSmoothDesign, BasisError> {
9308    let mut planned_blocks = plan_joint_spatial_centers_for_term_blocks(data, &[terms.to_vec()])?;
9309    let planned_terms = planned_blocks.pop().ok_or_else(|| {
9310        BasisError::InvalidInput(
9311            "joint spatial center planner returned no smooth blocks".to_string(),
9312        )
9313    })?;
9314    let policy = workspace.policy().clone();
9315    let local_builds: Vec<LocalSmoothTermBuild> = {
9316        use rayon::iter::{IntoParallelIterator, ParallelIterator};
9317        planned_terms
9318            .into_par_iter()
9319            .map(|term| {
9320                let mut term_workspace = crate::basis::BasisWorkspace::with_policy(policy.clone());
9321                build_single_local_smooth_term(data, &term, &mut term_workspace)
9322            })
9323            .collect::<Result<Vec<_>, _>>()?
9324    };
9325
9326    let total_p: usize = local_builds.iter().map(|built| built.dim).sum();
9327
9328    let mut local_designs: Vec<DesignMatrix> = Vec::with_capacity(local_builds.len());
9329    let mut affine_offset = Array1::<f64>::zeros(data.nrows());
9330    let mut terms_out = Vec::<SmoothTerm>::with_capacity(terms.len());
9331    let mut penalties_global = Vec::<BlockwisePenalty>::new();
9332    let mut nullspace_dims_global = Vec::<usize>::new();
9333    let mut penaltyinfo_global = Vec::<PenaltyBlockInfo>::new();
9334    let mut dropped_penaltyinfo_global = Vec::<DroppedPenaltyBlockInfo>::new();
9335    let mut coefficient_lower_bounds = Array1::<f64>::from_elem(total_p, f64::NEG_INFINITY);
9336    let mut any_bounds = false;
9337    // Each linear-constraint row only touches the current term's column slice.
9338    // Track `(col_start, col_end, local_row_values)` and assemble the final
9339    // dense `Array2` in one pass, avoiding per-row `Array1::zeros(total_p)`
9340    // allocation plus a row-by-row copy at the end.
9341    let mut linear_constraintsrows: Vec<(usize, usize, Array1<f64>)> = Vec::new();
9342    let mut linear_constraints_b: Vec<f64> = Vec::new();
9343
9344    let mut col_start = 0usize;
9345    for (term, mut built) in terms.iter().zip(local_builds.into_iter()) {
9346        let p_local = built.dim;
9347        let col_end = col_start + p_local;
9348        let lb_local = if built.box_reparam {
9349            shape_lower_bounds_local(term.shape, p_local)
9350        } else {
9351            None
9352        };
9353
9354        // Stage-2 joint-null absorption rotation. Fired *before* the
9355        // penalty / design / global aggregation loops below so that every
9356        // subsequent reference to `built.active_penalties` and `built.design`
9357        // sees the post-rotation values.
9358        //
9359        // The math: when the smooth's joint penalty `Σ_k S_k` has a
9360        // non-trivial null space, eigh selects `Q = [U_range | U_null]`
9361        // with null columns at the tail. Setting `β_raw = Q · γ` and
9362        // applying:
9363        //     design        ← X · Q
9364        //     penalties[k]  ← Qᵀ · S_k · Q   (block-diag, zero null tail)
9365        // yields a model whose fitted γ is invariant to the rotation
9366        // (since likelihood depends only on `X · β_raw = X · Q · γ`), but
9367        // whose penalty is full-rank on the range columns. The large-scale
9368        // failing case (cert refusal in the joint-Newton inner solve)
9369        // resolves because `H_pen = H_loglik + S` becomes full rank on
9370        // the smooth's range columns.
9371        //
9372        // Rotation is suppressed when the smooth carries coordinate-wise
9373        // shape constraints (`lb_local` or `built.linear_constraints`):
9374        // those encode a cone in the original coordinate system and a
9375        // general orthogonal rotation breaks the cone geometry. Smooths
9376        // with shape constraints typically have full-rank joint penalty
9377        // (their structural shape comes from the cone, not from null
9378        // directions in the penalty), so suppression is rarely a loss.
9379        //
9380        // `applied_rotation` carries the Q that was applied (or `None`
9381        // if no rotation fired). It is persisted onto `SmoothTerm` below
9382        // so prediction-side `X_new_raw · Q` replay can reproduce the
9383        // exact rotation. Persistence through the saved-model artifact
9384        // is a follow-up — see the doc on `SmoothTerm.joint_null_rotation`.
9385        let applied_rotation: Option<crate::basis::JointNullRotation> = match (
9386            built.joint_null_rotation.take(),
9387            lb_local.is_some(),
9388            built.linear_constraints.is_some(),
9389        ) {
9390            (Some(rot), false, false) => {
9391                let q = &rot.rotation;
9392                built.design =
9393                    apply_smooth_transform_to_design(built.design.clone(), q, &term.name)?;
9394                for penalty in &mut built.active_penalties {
9395                    let qt_s = gam_linalg::faer_ndarray::fast_atb(q, &penalty.matrix);
9396                    penalty.matrix = gam_linalg::faer_ndarray::fast_ab(&qt_s, q);
9397                    penalty.null_eigenvectors = penalty
9398                        .null_eigenvectors
9399                        .as_ref()
9400                        .map(|basis| gam_linalg::faer_ndarray::fast_atb(q, basis));
9401                    penalty.op = None;
9402                    penalty.info.kronecker_factors = None;
9403                }
9404                built.kronecker_factored = None;
9405                Some(rot)
9406            }
9407            (Some(_), _, _) => None,
9408            (None, _, _) => None,
9409        };
9410
9411        for active_penalty in &built.active_penalties {
9412            let global_index = penalties_global.len();
9413            penalties_global.push(
9414                BlockwisePenalty::new(col_start..col_end, active_penalty.matrix.clone())
9415                    .with_op(active_penalty.op.clone()),
9416            );
9417            nullspace_dims_global.push(active_penalty.nullity);
9418            penaltyinfo_global.push(PenaltyBlockInfo {
9419                global_index,
9420                termname: Some(term.name.clone()),
9421                penalty: active_penalty.info.clone(),
9422            });
9423        }
9424        for info in &built.dropped_penalties {
9425            dropped_penaltyinfo_global.push(DroppedPenaltyBlockInfo {
9426                termname: Some(term.name.clone()),
9427                penalty: info.clone(),
9428            });
9429        }
9430
9431        if let Some(lin_local) = &built.linear_constraints {
9432            for r in 0..lin_local.a.nrows() {
9433                linear_constraintsrows.push((col_start, col_end, lin_local.a.row(r).to_owned()));
9434                linear_constraints_b.push(lin_local.b[r]);
9435            }
9436        }
9437        if let Some(lb_local) = &lb_local {
9438            coefficient_lower_bounds
9439                .slice_mut(s![col_start..col_end])
9440                .assign(lb_local);
9441            any_bounds = true;
9442        }
9443
9444        if let Some(term_offset) = built.affine_offset.as_ref() {
9445            if term_offset.len() != data.nrows() {
9446                crate::bail_dim_basis!(
9447                    "smooth term '{}' affine offset has {} rows but the realized data has {}",
9448                    term.name,
9449                    term_offset.len(),
9450                    data.nrows()
9451                );
9452            }
9453            affine_offset += term_offset;
9454        }
9455
9456        // Move the per-term design out of `built` rather than cloning it.
9457        local_designs.push(built.design);
9458
9459        terms_out.push(SmoothTerm {
9460            name: term.name.clone(),
9461            coeff_range: col_start..col_end,
9462            shape: term.shape,
9463            active_penalties: built.active_penalties,
9464            dropped_penalties: built.dropped_penalties,
9465            metadata: built.metadata,
9466            lower_bounds_local: lb_local,
9467            linear_constraints_local: built.linear_constraints,
9468            kronecker_factored: built.kronecker_factored.take(),
9469            joint_null_rotation: applied_rotation,
9470            unabsorbed_global_orthogonality: None,
9471        });
9472
9473        col_start = col_end;
9474    }
9475
9476    assert_eq!(
9477        penalties_global.len(),
9478        nullspace_dims_global.len(),
9479        "global smooth penalty/nullspace bookkeeping diverged"
9480    );
9481    assert_eq!(
9482        penalties_global.len(),
9483        penaltyinfo_global.len(),
9484        "global smooth penalty metadata bookkeeping diverged"
9485    );
9486
9487    Ok(RawSmoothDesign {
9488        term_designs: local_designs,
9489        affine_offset,
9490        penalties: penalties_global,
9491        nullspace_dims: nullspace_dims_global,
9492        penaltyinfo: penaltyinfo_global,
9493        dropped_penaltyinfo: dropped_penaltyinfo_global,
9494        terms: terms_out,
9495        coefficient_lower_bounds: if any_bounds {
9496            Some(coefficient_lower_bounds)
9497        } else {
9498            None
9499        },
9500        linear_constraints: if linear_constraintsrows.is_empty() {
9501            None
9502        } else {
9503            let mut a = Array2::<f64>::zeros((linear_constraintsrows.len(), total_p));
9504            for (i, (cs, ce, values)) in linear_constraintsrows.iter().enumerate() {
9505                a.row_mut(i).slice_mut(s![*cs..*ce]).assign(values);
9506            }
9507            Some(LinearInequalityConstraints {
9508                a,
9509                b: Array1::from_vec(linear_constraints_b),
9510            })
9511        },
9512    })
9513}
9514
9515#[cfg(test)]
9516mod factor_smooth_heldout_group_tests {
9517    use super::*;
9518    use crate::basis::BasisWorkspace;
9519    use ndarray::{Array1, array};
9520
9521    fn pinned_marginal() -> BSplineBasisSpec {
9522        BSplineBasisSpec {
9523            degree: 3,
9524            penalty_order: 2,
9525            knotspec: BSplineKnotSpec::Provided(Array1::from(vec![
9526                0.0, 0.0, 0.0, 0.0, 0.25, 0.6, 1.0, 1.0, 1.0, 1.0,
9527            ])),
9528            double_penalty: false,
9529            identifiability: BSplineIdentifiability::None,
9530            boundary: crate::basis::OneDimensionalBoundary::Open,
9531            boundary_conditions: crate::basis::BSplineBoundaryConditions::default(),
9532        }
9533    }
9534
9535    fn factor_smooth_term(flavour: FactorSmoothFlavour, frozen: Option<Vec<u64>>) -> SmoothTermSpec {
9536        SmoothTermSpec {
9537            name: "fs_heldout".to_string(),
9538            basis: SmoothBasisSpec::FactorSmooth {
9539                spec: FactorSmoothSpec {
9540                    continuous_cols: vec![0],
9541                    group_col: 1,
9542                    marginal: pinned_marginal(),
9543                    flavour,
9544                    group_frozen_levels: frozen,
9545                    frozen_global_orthogonality: None,
9546                },
9547            },
9548            shape: ShapeConstraint::None,
9549            joint_null_rotation: None,
9550        }
9551    }
9552
9553    const FROZEN_01: [f64; 2] = [0.0, 1.0];
9554
9555    fn frozen_bits() -> Vec<u64> {
9556        FROZEN_01.iter().map(|v| v.to_bits()).collect()
9557    }
9558
9559    /// #2365: in the frozen (predict/replay) context, a `bs="re"` row whose
9560    /// group is outside the training vocabulary must build with an all-zero
9561    /// row — zero fitted deviation, population prediction — instead of
9562    /// erroring before the random-effect operator can apply its held-out-group
9563    /// contract.
9564    #[test]
9565    fn re_heldout_group_row_is_zero_deviation() {
9566        let data = array![[0.1, 0.0], [0.5, 1.0], [0.9, 7.0]];
9567        let term = factor_smooth_term(FactorSmoothFlavour::Re, Some(frozen_bits()));
9568        let mut workspace = BasisWorkspace::default();
9569        let build = build_single_local_smooth_term(data.view(), &term, &mut workspace)
9570            .expect("a held-out group must not fail the bs=\"re\" design build");
9571        let dense = build
9572            .design
9573            .try_to_dense_by_chunks("heldout test")
9574            .expect("dense");
9575        assert!(
9576            dense.row(2).iter().all(|&v| v == 0.0),
9577            "unseen-group row must carry zero deviation across every group block, got {:?}",
9578            dense.row(2)
9579        );
9580        assert!(
9581            dense.row(0).iter().any(|&v| v != 0.0) && dense.row(1).iter().any(|&v| v != 0.0),
9582            "in-vocabulary rows must still populate their group blocks"
9583        );
9584    }
9585
9586    /// The `fs` flavour estimates a per-level deviation FUNCTION — an unseen
9587    /// level has no zero-deviation population fallback — so the frozen-context
9588    /// build must stay strict (#2102/#2137 must not regress through #2365).
9589    #[test]
9590    fn fs_heldout_group_stays_strict() {
9591        let data = array![[0.1, 0.0], [0.5, 1.0], [0.9, 7.0]];
9592        let term = factor_smooth_term(
9593            FactorSmoothFlavour::Fs {
9594                m_null_penalty_orders: vec![1],
9595            },
9596            Some(frozen_bits()),
9597        );
9598        let mut workspace = BasisWorkspace::default();
9599        let err = match build_single_local_smooth_term(data.view(), &term, &mut workspace) {
9600            Ok(_) => panic!("fs must reject an unseen grouping level"),
9601            Err(err) => err,
9602        };
9603        assert!(
9604            err.to_string().contains("unseen grouping level"),
9605            "fs unseen-level refusal must name the defect, got: {err}"
9606        );
9607    }
9608}