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