Skip to main content

gam_terms/smooth/
term_specs.rs

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