Skip to main content

gam_terms/smooth/
term_specs.rs

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