gam_terms/basis/types.rs
1use super::*;
2
3/// Wrapper to send a raw pointer across thread boundaries for parallel buffer fills.
4/// SAFETY: every `SendPtr` value must be built from live, properly aligned `f64`
5/// storage whose mutable borrow is held until all worker threads finish; callers
6/// may only dereference offsets that are in-bounds and disjoint across workers.
7#[derive(Clone, Copy)]
8pub(crate) struct SendPtr(pub(crate) *mut f64);
9
10// SAFETY: SendPtr only grants raw-pointer transport. Actual dereferences occur
11// at call sites after row-chunk partitioning proves each thread writes a
12// distinct in-bounds element of the backing Array/Vec allocation.
13unsafe impl Send for SendPtr {}
14
15// SAFETY: shared references to SendPtr are sound because the pointee is never
16// accessed through the wrapper without the call-site disjoint-offset proof.
17unsafe impl Sync for SendPtr {}
18
19impl SendPtr {
20 #[inline(always)]
21 pub(crate) fn add(self, offset: usize) -> *mut f64 {
22 // SAFETY: callers pass offsets within the backing allocation and only
23 // dereference the returned pointer after proving the target element is
24 // uniquely owned by that worker's chunk for the whole parallel region.
25 unsafe { self.0.add(offset) }
26 }
27}
28
29/// Re-export of the neutral basis-error contract. #1521: `BasisError` lives
30/// in `gam-problem` so `EstimationError` can wrap it (`#[from]`) without a
31/// back-edge; gam-terms re-exports it to preserve `gam_terms::basis::BasisError`.
32pub use gam_problem::BasisError;
33
34// ============================================================================
35// Unified Basis Generation API
36// ============================================================================
37
38/// Options for basis generation, controlling derivative order.
39#[derive(Clone, Copy, Debug, Default)]
40pub struct BasisOptions {
41 /// Derivative order: 0 = value (default), 1 = first derivative, 2 = second derivative
42 pub derivative_order: usize,
43 /// Basis family to evaluate.
44 pub basis_family: BasisFamily,
45}
46
47impl BasisOptions {
48 /// Create options for evaluating basis functions (no derivative).
49 pub const fn value() -> Self {
50 Self {
51 derivative_order: 0,
52 basis_family: BasisFamily::BSpline,
53 }
54 }
55
56 /// Create options for evaluating first derivatives of basis functions.
57 pub const fn first_derivative() -> Self {
58 Self {
59 derivative_order: 1,
60 basis_family: BasisFamily::BSpline,
61 }
62 }
63
64 /// Create options for evaluating second derivatives of basis functions.
65 pub const fn second_derivative() -> Self {
66 Self {
67 derivative_order: 2,
68 basis_family: BasisFamily::BSpline,
69 }
70 }
71
72 /// Create options for evaluating M-spline basis values.
73 pub const fn m_spline() -> Self {
74 Self {
75 derivative_order: 0,
76 basis_family: BasisFamily::MSpline,
77 }
78 }
79
80 /// Create options for evaluating I-spline basis values.
81 pub const fn i_spline() -> Self {
82 Self {
83 derivative_order: 0,
84 basis_family: BasisFamily::ISpline,
85 }
86 }
87}
88
89/// Basis-family selector for 1D spline evaluation.
90#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
91pub enum BasisFamily {
92 /// Standard B-splines.
93 #[default]
94 BSpline,
95 /// M-splines: normalized B-splines, M_i = ((k+1)/(t_{i+k+1}-t_i)) B_i.
96 MSpline,
97 /// I-splines: integrated M-splines, implemented by right-cumulative
98 /// sums of B-splines at degree k+1.
99 ISpline,
100}
101
102/// Specifies the source of knots for basis generation.
103#[derive(Clone, Debug)]
104pub enum KnotSource<'a> {
105 /// Use a pre-computed knot vector.
106 Provided(ArrayView1<'a, f64>),
107 /// Generate uniformly spaced knots based on data range.
108 Generate {
109 /// Data range (min, max) for knot placement.
110 data_range: (f64, f64),
111 /// Number of internal knots to place between boundaries.
112 num_internal_knots: usize,
113 },
114}
115/// Thin-plate regression spline basis and penalty (order m=2).
116///
117/// The returned basis has columns `[K_c | P]` where:
118/// - `K_c` is the constrained radial basis block (`K * Z`) with
119/// `P(knots)^T * α = 0` enforced via nullspace projection
120/// - `P` is the TPS polynomial null-space block containing all monomials of
121/// total degree `< m`, where `m = thin_plate_penalty_order(d)` (so `P` is
122/// just `[1, x_1, ..., x_d]` for `d <= 3`)
123///
124/// The returned penalty matrix is block-diagonal with:
125/// - upper-left `Omega_c = Z^T Omega Z` for the constrained radial block
126/// - zero lower-right block for unpenalized polynomial terms.
127///
128/// For double-penalty GAMs, a second ridge penalty `I` is also returned so the
129/// caller can optimize `(lambda_bending, lambdaridge)` jointly.
130#[derive(Debug, Clone)]
131pub struct ThinPlateSplineBasis {
132 pub basis: Array2<f64>,
133 pub penalty_bending: Array2<f64>,
134 pub penalty_ridge: Array2<f64>,
135 pub num_kernel_basis: usize,
136 pub num_polynomial_basis: usize,
137 pub dimension: usize,
138 /// Wood-TPRS radial reparameterization matrix `V`.
139 ///
140 /// Rows live in the side-constrained radial coefficient space. Columns are
141 /// the retained positive bending eigendirections of `Z' Ω Z`; numerically
142 /// near-null radial directions are dropped before the basis is exposed.
143 /// Therefore `V` can be rectangular: design columns are `Φ Z V`, and the
144 /// radial penalty is `diag(Λ_retained)`.
145 pub radial_reparam: Array2<f64>,
146}
147
148/// Matérn smoothness parameter `nu` (half-integer variants with closed forms).
149#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
150pub enum MaternNu {
151 Half,
152 ThreeHalves,
153 FiveHalves,
154 SevenHalves,
155 NineHalves,
156}
157
158impl MaternNu {
159 /// The half-integer smoothness value ν as an `f64` (0.5, 1.5, …).
160 pub const fn half_integer_value(self) -> f64 {
161 match self {
162 MaternNu::Half => 0.5,
163 MaternNu::ThreeHalves => 1.5,
164 MaternNu::FiveHalves => 2.5,
165 MaternNu::SevenHalves => 3.5,
166 MaternNu::NineHalves => 4.5,
167 }
168 }
169}
170
171/// Matérn radial basis and penalties.
172#[derive(Debug, Clone)]
173pub struct MaternSplineBasis {
174 pub basis: Array2<f64>,
175 pub penalty_kernel: Array2<f64>,
176 pub penalty_ridge: Array2<f64>,
177 pub num_kernel_basis: usize,
178 pub num_polynomial_basis: usize,
179 pub dimension: usize,
180}
181
182#[derive(Debug, Clone)]
183pub(crate) struct DuchonBasisDesign {
184 pub(crate) basis: Array2<f64>,
185}
186
187/// Boundary-condition policy for one-dimensional smooth bases.
188#[derive(Debug, Clone, Serialize, Deserialize, Default)]
189pub enum OneDimensionalBoundary {
190 /// Ordinary open interval basis with clamped endpoint behavior.
191 #[default]
192 Open,
193 /// Periodic/cyclic basis over the half-open interval `[start, end)`.
194 ///
195 /// Values are evaluated modulo `period = end - start`; the basis and its
196 /// first `degree - 1` derivatives agree at the two endpoints for B-splines.
197 Cyclic { start: f64, end: f64 },
198}
199
200impl OneDimensionalBoundary {
201 pub(crate) fn period(&self) -> Option<(f64, f64, f64)> {
202 match *self {
203 OneDimensionalBoundary::Open => None,
204 OneDimensionalBoundary::Cyclic { start, end } if end > start => {
205 Some((start, end, end - start))
206 }
207 OneDimensionalBoundary::Cyclic { .. } => None,
208 }
209 }
210}
211
212/// Which knot strategy to use for 1D B-spline bases.
213#[derive(Debug, Clone, Serialize, Deserialize)]
214pub enum BSplineKnotSpec {
215 Generate {
216 data_range: (f64, f64),
217 num_internal_knots: usize,
218 },
219 /// Uniform cyclic B-spline basis on `[data_range.0, data_range.1)`.
220 ///
221 /// The first and last endpoints are identified, so evaluating at `x` and
222 /// `x + m * period` gives identical rows. `num_basis` is the number of
223 /// periodic control sites around the loop and must be at least
224 /// `degree + 1` for an unaliased local support stencil.
225 PeriodicUniform {
226 data_range: (f64, f64),
227 num_basis: usize,
228 },
229 Automatic {
230 num_internal_knots: Option<usize>,
231 placement: BSplineKnotPlacement,
232 },
233 Provided(Array1<f64>),
234 /// Natural cubic regression spline (`bs="cr"`/`"cs"`) knot set (#1074).
235 ///
236 /// Unlike the open-spline variants above, these `knots` are the `k`
237 /// Lancaster–Salkauskas knots `x*_1 < … < x*_k` that *directly* index the
238 /// basis values `β_i = f(x*_i)` — the basis dimension equals `knots.len()`
239 /// (not `knots.len() - degree - 1`). The 1-D builder routes this variant to
240 /// the cubic-regression builder; the cr identity therefore round-trips
241 /// through freeze/reload by virtue of the variant itself (no separate
242 /// metadata marker is required), and tensor margins inherit cr by carrying
243 /// this knotspec into `build_bspline_basis_1d`.
244 NaturalCubicRegression {
245 knots: Array1<f64>,
246 },
247}
248
249/// Internal-knot placement strategy when knots are automatically inferred.
250#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
251pub enum BSplineKnotPlacement {
252 Uniform,
253 Quantile,
254}
255
256/// 1D B-spline basis configuration.
257#[derive(Debug, Clone, Serialize, Deserialize)]
258pub struct BSplineBasisSpec {
259 pub degree: usize,
260 pub penalty_order: usize,
261 pub knotspec: BSplineKnotSpec,
262 pub double_penalty: bool,
263 pub identifiability: BSplineIdentifiability,
264 #[serde(default)]
265 pub boundary: OneDimensionalBoundary,
266 /// Optional endpoint boundary constraints (Hermite-style pin of value and/or
267 /// derivative at the left/right knot extents). Default = `Free` on both
268 /// sides which is a no-op.
269 #[serde(default)]
270 pub boundary_conditions: BSplineBoundaryConditions,
271}
272
273/// Per-endpoint boundary constraint policy for B-spline 1D bases.
274#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize, Default)]
275pub enum BSplineEndpointBoundaryCondition {
276 /// No endpoint constraint.
277 #[default]
278 Free,
279 /// Pin the first derivative to zero at this endpoint.
280 Clamped,
281 /// Pin the value at this endpoint to `value` (currently only `value == 0`
282 /// is accepted in the builder; non-zero anchors require an affine offset).
283 Anchored { value: f64 },
284}
285
286/// Left/right pair of B-spline endpoint constraints.
287#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize, Default)]
288pub struct BSplineBoundaryConditions {
289 #[serde(default)]
290 pub left: BSplineEndpointBoundaryCondition,
291 #[serde(default)]
292 pub right: BSplineEndpointBoundaryCondition,
293}
294
295impl BSplineBoundaryConditions {
296 pub const fn is_free(&self) -> bool {
297 matches!(self.left, BSplineEndpointBoundaryCondition::Free)
298 && matches!(self.right, BSplineEndpointBoundaryCondition::Free)
299 }
300
301 /// Whether exactly one endpoint fixes the function's absolute level.
302 ///
303 /// A one-sided anchor replaces the global intercept as the level-setting
304 /// constraint. Centering that same smooth to zero would impose a second,
305 /// incompatible level constraint and exclude every non-zero-mean anchored
306 /// function from the model space.
307 pub const fn has_one_sided_anchor(&self) -> bool {
308 let left = matches!(self.left, BSplineEndpointBoundaryCondition::Anchored { .. });
309 let right = matches!(
310 self.right,
311 BSplineEndpointBoundaryCondition::Anchored { .. }
312 );
313 left != right
314 }
315}
316
317/// Per-smooth identifiability policy for 1D B-spline bases.
318///
319/// These constraints are applied directly in the builder via a reparameterization
320/// `B_constrained = B * Z`, and every penalty matrix is projected as
321/// `S_constrained = Z' S Z`, so solver geometry stays consistent.
322#[derive(Debug, Clone, Serialize, Deserialize)]
323pub enum BSplineIdentifiability {
324 /// Keep unconstrained basis columns.
325 None,
326 /// Enforce weighted sum-to-zero: `B' w = 0` (or unweighted when `weights=None`).
327 // Smooth terms are centered by default to avoid intercept confounding.
328 WeightedSumToZero { weights: Option<Array1<f64>> },
329 /// Remove intercept + linear trend in coefficient space using Greville geometry.
330 RemoveLinearTrend,
331 /// Enforce orthogonality to supplied design columns `C` (n x q):
332 /// `B_c' W C = 0` (or unweighted when `weights=None`).
333 ///
334 /// To enforce `[intercept, x, ...]`, provide `columns` with those columns.
335 OrthogonalToDesignColumns {
336 columns: Array2<f64>,
337 weights: Option<Array1<f64>>,
338 },
339 /// Apply an explicit coefficient-space transform `Z` learned at fit time.
340 ///
341 /// This freezes identifiability behavior so prediction cannot drift based on
342 /// new-data distribution. The constrained basis is `B * Z`.
343 FrozenTransform { transform: Array2<f64> },
344}
345
346impl Default for BSplineIdentifiability {
347 fn default() -> Self {
348 BSplineIdentifiability::WeightedSumToZero { weights: None }
349 }
350}
351
352/// Spatial center selection strategy.
353///
354/// `num_centers` is the exact number of knot/center rows selected by the
355/// strategy. Polynomial nullspace columns are added separately by each basis
356/// builder and must never be folded into this count.
357#[derive(Debug, Clone, Serialize, Deserialize)]
358pub enum CenterStrategy {
359 Auto(Box<CenterStrategy>),
360 UserProvided(Array2<f64>),
361 /// Joint multidimensional equal-mass partitioning in the full smooth space.
362 EqualMass {
363 num_centers: usize,
364 },
365 /// Covariate-representative equal-mass partitioning along one selected axis.
366 EqualMassCovarRepresentative {
367 num_centers: usize,
368 },
369 FarthestPoint {
370 num_centers: usize,
371 },
372 KMeans {
373 num_centers: usize,
374 max_iter: usize,
375 },
376 UniformGrid {
377 points_per_dim: usize,
378 },
379}
380
381impl CenterStrategy {
382 /// The number of centers this strategy will select, computed from the
383 /// strategy alone (no data pass). `d` is the smooth's covariate
384 /// dimensionality, needed only by `UniformGrid` whose count is
385 /// `points_per_dim^d`. Adaptive-fit provenance consults this before freeze,
386 /// because the frozen center matrix can contain periodic image expansion
387 /// and therefore is not the requested resolution for the next refit.
388 pub fn planned_num_centers(&self, d: usize) -> usize {
389 match self {
390 Self::Auto(inner) => inner.planned_num_centers(d),
391 Self::UserProvided(centers) => centers.nrows(),
392 Self::EqualMass { num_centers }
393 | Self::EqualMassCovarRepresentative { num_centers }
394 | Self::FarthestPoint { num_centers }
395 | Self::KMeans { num_centers, .. } => *num_centers,
396 Self::UniformGrid { points_per_dim } => {
397 points_per_dim.saturating_pow(d.clamp(1, u32::MAX as usize) as u32)
398 }
399 }
400 }
401}
402
403#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
404pub enum CenterStrategyKind {
405 UserProvided,
406 EqualMass,
407 EqualMassCovarRepresentative,
408 FarthestPoint,
409 KMeans,
410 UniformGrid,
411}
412
413/// Adaptive default center count for spatial smooths (TPS, Duchon, Matérn).
414///
415/// Use this when the user has not explicitly specified a knot/center count.
416/// The basis size is the sub-linear `ceil(8 * d_factor * n^0.4)`, clamped above
417/// at `K_MAX = 2000` and below at a *data-proportional* floor `min(200, n/8)` so
418/// the floor only engages once there are enough observations to support a rich
419/// basis. The result is additionally capped at `n/4` so the penalty matrices
420/// stay well-conditioned relative to the data:
421///
422/// | n | d=1 | d=2 | d=5 |
423/// |--------|------|------|------|
424/// | 800 | 116 | 134 | 186 |
425/// | 1 000 | 127 | 146 | 200 |
426/// | 2 000 | 200 | 200 | 268 |
427/// | 10 000 | 319 | 367 | 510 |
428/// | 100 000| 801 | 921 | 1281 |
429/// | 400 000| 1393 | 1602 | 2000 |
430/// | 1 000 000| 2000 | 2000 | 2000 |
431///
432/// The flat `200` floor used to inflate moderate-`n` spatial smooths (a few
433/// hundred to ~2000 rows) up to a dense 200-column design even though the raw
434/// sub-linear count — and the mesh/knot density that mgcv and R-INLA use on the
435/// same data — is far smaller. On ~800 rows that turned a single 2-D thin-plate
436/// REML fit into an `O(n·p² + p³)` grind at `p ≈ 200` (#718). Smoothness is
437/// already controlled by REML's penalty weight λ, not by the center count, so a
438/// data-proportional floor recovers the same surface at a fraction of the cost.
439///
440/// # Arguments
441/// * `n` - sample size (number of observations)
442/// * `d` - covariate dimensionality (number of input variables in the smooth)
443pub fn default_num_centers(n: usize, d: usize) -> usize {
444 const K_MIN: usize = 200;
445 const K_MAX: usize = 2000;
446 const ALPHA: f64 = 0.4;
447 const C: f64 = 8.0;
448 /// Per-extra-dimension growth in the center count: each covariate axis
449 /// beyond the first widens the basis by 15% to keep the per-axis mesh
450 /// density roughly constant as the smooth's domain dimensionality grows.
451 const PER_DIM_GROWTH: f64 = 0.15;
452 /// Divisor for the data-proportional floor: the `K_MIN` floor only engages
453 /// once `n` exceeds `K_MIN * FLOOR_N_DIVISOR`, so small samples are not
454 /// forced up to a dense `K_MIN`-column design.
455 const FLOOR_N_DIVISOR: usize = 8;
456 /// Divisor for the conditioning cap: the center count never exceeds `n /
457 /// COND_N_DIVISOR`, keeping the penalty matrices well-conditioned relative
458 /// to the data.
459 const COND_N_DIVISOR: usize = 4;
460
461 let d_factor = 1.0 + PER_DIM_GROWTH * (d.max(1) - 1) as f64;
462 let raw = (C * d_factor * (n as f64).powf(ALPHA)).ceil() as usize;
463
464 // Data-proportional floor: never inflate beyond n/FLOOR_N_DIVISOR, so the
465 // K_MIN-center floor only takes effect once n is large enough (~1600) to
466 // genuinely support that many basis columns.
467 let floor = K_MIN.min(n / FLOOR_N_DIVISOR);
468 let k = raw.clamp(floor, K_MAX);
469
470 // Never exceed n itself; cap at n/COND_N_DIVISOR to keep the penalty
471 // matrices well-conditioned relative to the data.
472 k.min(n).min(n / COND_N_DIVISOR)
473}
474
475/// Conservative center count for a *secondary* (distributional) predictor's
476/// spatial smooth — e.g. the log-σ scale model in a Gaussian location-scale
477/// fit.
478///
479/// The mean is identified directly by the response, so it warrants the
480/// generous [`default_num_centers`] basis. A scale/shape predictor is
481/// identified only through (noisy) squared residuals: handing it a basis sized
482/// for the mean lets REML/LAML smoothing selection over-fit it, because where
483/// the fitted scale is driven small the *observed* information collapses and
484/// the determinant penalty stops holding the wiggle down (#501). This mirrors
485/// standard GAMLSS/mgcv practice of giving distribution parameters a modest
486/// default (mgcv's modest default basis for a 1-D `s()`), grown gently with
487/// dimensionality and never exceeding the generous primary-predictor default.
488pub fn conservative_secondary_centers(n: usize, d: usize) -> usize {
489 const BASE_1D_CENTERS: usize = 15;
490 let modest = BASE_1D_CENTERS.saturating_mul(d.max(1));
491 default_num_centers(n, d).min(modest).max(1)
492}
493
494/// Low-rank starting center count for saturation-driven spatial fitting.
495///
496/// The structural minimum (`d + 1` polynomial directions plus one radial
497/// direction) is only enough to make the algebra identifiable. It is not an
498/// adequate pilot function space: structure orthogonal to that single radial
499/// direction is absorbed into the residual, so REML can legitimately shrink
500/// the direction and report EDF below its ceiling even when the surface is
501/// badly under-resolved (#1689). Start from the project's established
502/// thin-plate-style low-rank resolution `10 * 3^(d - 1)` instead. This is the
503/// same dimension rule already used by the automatic Duchon builder, capped by
504/// [`default_num_centers`] so the pilot never exceeds the validated production
505/// basis at small sample sizes.
506pub fn starting_num_centers(n: usize, d: usize) -> usize {
507 let low_rank_resolution = 10usize
508 .saturating_mul(3usize.saturating_pow(d.saturating_sub(1).min(u32::MAX as usize) as u32));
509 low_rank_resolution
510 .min(default_num_centers(n, d))
511 .min(n)
512 .max(1)
513}
514
515/// Next evidence-backed center count for a saturated spatial basis, bounded by
516/// the already validated production-default resolution.
517///
518/// Growth is geometric so the number of certified refits is logarithmic. The
519/// ceiling is supplied by the owning workflow because it depends on the
520/// spatial family/dimension and resource plan; the standard formula workflow
521/// uses [`default_num_centers`]. Adaptive resolution may therefore avoid work
522/// below the previous default, but can never turn an ordinary fit into an
523/// unvalidated row-rank dense basis. `None` means the validated function-space
524/// ceiling has been reached.
525pub fn expanded_num_centers(current: usize, ceiling: usize) -> Option<usize> {
526 if current >= ceiling {
527 return None;
528 }
529 let expanded = current.saturating_mul(2).min(ceiling);
530 (expanded > current).then_some(expanded)
531}
532
533/// Is a fitted spatial smooth's basis SATURATED — i.e. does its own evidence say
534/// the data wants more resolution than its realized coefficient span provides (#1689)?
535///
536/// The penalizable capacity is `realized_width − nullspace_dim`: the unpenalized
537/// polynomial null space is always fully used, so it is excluded from the "is the
538/// PENALIZED part maxed out?" test. The supplied `edf` is the total term EDF;
539/// subtracting `nullspace_dim` yields its penalized contribution, which rises
540/// toward that capacity exactly as REML drives the penalty
541/// λ toward its floor to chase structure the basis cannot resolve. Saturated ⟺
542/// `edf ≥ capacity − ε`, with the margin `ε` DERIVED from the outer REML
543/// numerical resolution (`ε = capacity · resolution_tol`, floored at
544/// `resolution_tol` so a tiny-capacity block still has a positive margin) rather
545/// than a tuned knob. The workflow derives `resolution_tol` from the maximum of
546/// its outer convergence tolerance and any rho-independent penalty shrinkage
547/// floor, because that floor bounds how closely EDF can approach the algebraic
548/// ceiling even as lambda tends to zero. Non-positive capacity (a block whose
549/// null space already exhausts its columns) is never saturated. The absolute
550/// scale of `ε` is what the MSI truth-recovery sweep
551/// (sin8/kappa/large_scale + #1074) validates — the criterion SHAPE
552/// (edf-vs-capacity, nullspace excluded, tol-tied margin) is the load-bearing
553/// contract this function pins.
554pub fn basis_is_saturated(
555 edf: f64,
556 realized_width: usize,
557 nullspace_dim: usize,
558 resolution_tol: f64,
559) -> bool {
560 let capacity = realized_width.saturating_sub(nullspace_dim) as f64;
561 if !(capacity > 0.0) || !edf.is_finite() {
562 return false;
563 }
564 let penalized_edf = (edf - nullspace_dim as f64).clamp(0.0, capacity);
565 let margin = (capacity * resolution_tol).max(resolution_tol);
566 penalized_edf >= capacity - margin
567}
568
569/// Resource-aware plan for a spatial smooth (Duchon / Matérn / TPS).
570///
571/// Returned by [`plan_spatial_basis`]. Captures the resolved center count,
572/// final basis dimension `p`, the dense byte cost for the value matrix and
573/// each derivative tier, and a recommended storage mode that is consistent
574/// with the supplied [`gam_runtime::resource::ResourcePolicy`].
575#[derive(Clone, Debug)]
576pub struct SpatialBasisPlan {
577 pub n: usize,
578 pub d: usize,
579 pub centers: usize,
580 pub p_final_estimate: usize,
581 pub dense_design_bytes: usize,
582 pub first_derivative_dense_bytes: usize,
583 pub second_derivative_dense_bytes: usize,
584 pub recommended_storage: SpatialStorageMode,
585}
586
587/// Storage mode recommended by [`plan_spatial_basis`].
588///
589/// * `DenseValueDenseDerivatives` — both the value design and its derivative
590/// matrices fit under the policy's single-materialization budget.
591/// * `LazyValueImplicitDerivatives` — the value design fits dense but the
592/// derivative matrices do not; switch derivatives to the implicit operator.
593/// * `OperatorOnly` — neither the design nor its derivatives fit; everything
594/// must be operator-backed.
595#[derive(Clone, Copy, Debug, PartialEq, Eq)]
596pub enum SpatialStorageMode {
597 DenseValueDenseDerivatives,
598 LazyValueImplicitDerivatives,
599 OperatorOnly,
600}
601
602/// How [`plan_spatial_basis`] should pick the spatial center count.
603#[derive(Clone, Copy, Debug)]
604pub enum CenterCountRequest {
605 /// Use the heuristic [`default_num_centers`].
606 Default,
607 /// Use the caller-supplied count exactly.
608 Explicit(usize),
609 /// Use [`default_num_centers`] but cap at `cap` to bound dense cost.
610 HeuristicCapped { cap: usize },
611}
612
613/// Build a resource-aware plan for a spatial smooth basis.
614///
615/// Computes the resolved center count, final basis dimension, dense byte
616/// estimates for the value design and first/second derivative tiers, and a
617/// recommended [`SpatialStorageMode`] derived from `policy`. This is the
618/// resource-aware replacement for ad-hoc calls to [`default_num_centers`] /
619/// [`heuristic_centers`](crate::term_builder::heuristic_centers).
620pub fn plan_spatial_basis(
621 n: usize,
622 d: usize,
623 requested_centers: CenterCountRequest,
624 nullspace_order: DuchonNullspaceOrder,
625 scale_dims: bool,
626 policy: &gam_runtime::resource::ResourcePolicy,
627) -> Result<SpatialBasisPlan, BasisError> {
628 if n == 0 {
629 crate::bail_invalid_basis!("plan_spatial_basis: n must be >= 1");
630 }
631 if d == 0 {
632 crate::bail_invalid_basis!("plan_spatial_basis: d must be >= 1");
633 }
634
635 // 1. Resolve center count.
636 let centers = match requested_centers {
637 CenterCountRequest::Default => default_num_centers(n, d),
638 CenterCountRequest::Explicit(k) => k,
639 CenterCountRequest::HeuristicCapped { cap } => default_num_centers(n, d).min(cap),
640 };
641
642 // 2. Nullspace dimension (Duchon polynomial null space of degree p-1).
643 // `duchon_p_from_nullspace_order` returns m such that the null space is
644 // polynomials of total degree < m, matching `duchon_nullspace_dimension`'s
645 // `max_total_degree = m - 1` argument.
646 let m = duchon_p_from_nullspace_order(nullspace_order);
647 let nullspace_dim = if m == 0 {
648 0
649 } else {
650 duchon_nullspace_dimension(d, m - 1)
651 };
652
653 let p = centers.saturating_add(nullspace_dim);
654
655 // 3. Dense byte estimates.
656 let derivative_axes = if scale_dims { d } else { 0 };
657 let bytes_per_f64 = std::mem::size_of::<f64>();
658 let dense_design_bytes = bytes_per_f64.saturating_mul(n).saturating_mul(p);
659 let first_derivative_dense_bytes = dense_design_bytes.saturating_mul(derivative_axes);
660 // Diagonal second derivatives are also (D × n × p); off-diagonal cross terms
661 // would scale as D^2 but the planner reports the diagonal tier here.
662 let second_derivative_dense_bytes = first_derivative_dense_bytes;
663
664 // 4. Pick storage mode based on policy.
665 let recommended_storage = match policy.derivative_storage_mode {
666 gam_runtime::resource::DerivativeStorageMode::AnalyticOperatorRequired => {
667 SpatialStorageMode::OperatorOnly
668 }
669 gam_runtime::resource::DerivativeStorageMode::MaterializeIfSmall => {
670 let budget = policy.max_single_materialization_bytes;
671 if derivative_axes == 0 {
672 if dense_design_bytes <= budget {
673 SpatialStorageMode::DenseValueDenseDerivatives
674 } else {
675 SpatialStorageMode::LazyValueImplicitDerivatives
676 }
677 } else {
678 let total = dense_design_bytes
679 .saturating_add(first_derivative_dense_bytes)
680 .saturating_add(second_derivative_dense_bytes);
681 if total <= budget {
682 SpatialStorageMode::DenseValueDenseDerivatives
683 } else if dense_design_bytes <= budget {
684 SpatialStorageMode::LazyValueImplicitDerivatives
685 } else {
686 SpatialStorageMode::OperatorOnly
687 }
688 }
689 }
690 gam_runtime::resource::DerivativeStorageMode::DiagnosticsOnly => {
691 // Diagnostic mode still prefers analytic storage for correctness.
692 SpatialStorageMode::OperatorOnly
693 }
694 };
695
696 Ok(SpatialBasisPlan {
697 n,
698 d,
699 centers,
700 p_final_estimate: p,
701 dense_design_bytes,
702 first_derivative_dense_bytes,
703 second_derivative_dense_bytes,
704 recommended_storage,
705 })
706}
707
708pub const fn default_spatial_center_strategy(num_centers: usize, d: usize) -> CenterStrategy {
709 if d <= 3 {
710 CenterStrategy::FarthestPoint { num_centers }
711 } else {
712 CenterStrategy::EqualMassCovarRepresentative { num_centers }
713 }
714}
715
716pub fn auto_spatial_center_strategy(num_centers: usize, d: usize) -> CenterStrategy {
717 let strategy = if d == 1 {
718 // In one dimension, farthest-point selection is the deterministic
719 // maximin grid over the observed domain. Equal-mass midpoints leave the
720 // low-frequency Duchon radial block slightly under-resolved at the
721 // boundaries, and REML then compensates with an over-smooth λ on
722 // low-noise signals (#504). The maximin grid matches the native
723 // reproducing-kernel interpolation geometry. The default strategy below
724 // extends the same space-filling contract to low-dimensional spatial
725 // GP bases, where kriging accuracy is governed by fill distance rather
726 // than marginal quantile balance.
727 CenterStrategy::FarthestPoint { num_centers }
728 } else {
729 default_spatial_center_strategy(num_centers, d)
730 };
731 CenterStrategy::Auto(Box::new(strategy))
732}
733
734pub const fn center_strategy_is_auto(strategy: &CenterStrategy) -> bool {
735 matches!(strategy, CenterStrategy::Auto(_))
736}
737
738pub(crate) fn realized_center_strategy(strategy: &CenterStrategy) -> &CenterStrategy {
739 match strategy {
740 CenterStrategy::Auto(inner) => inner.as_ref(),
741 other => other,
742 }
743}
744
745pub fn center_strategy_kind(strategy: &CenterStrategy) -> CenterStrategyKind {
746 match strategy {
747 CenterStrategy::Auto(inner) => center_strategy_kind(inner.as_ref()),
748 CenterStrategy::UserProvided(_) => CenterStrategyKind::UserProvided,
749 CenterStrategy::EqualMass { .. } => CenterStrategyKind::EqualMass,
750 CenterStrategy::EqualMassCovarRepresentative { .. } => {
751 CenterStrategyKind::EqualMassCovarRepresentative
752 }
753 CenterStrategy::FarthestPoint { .. } => CenterStrategyKind::FarthestPoint,
754 CenterStrategy::KMeans { .. } => CenterStrategyKind::KMeans,
755 CenterStrategy::UniformGrid { .. } => CenterStrategyKind::UniformGrid,
756 }
757}
758
759pub fn center_strategy_num_centers(strategy: &CenterStrategy) -> Option<usize> {
760 match strategy {
761 CenterStrategy::Auto(inner) => center_strategy_num_centers(inner.as_ref()),
762 CenterStrategy::UserProvided(centers) => Some(centers.nrows()),
763 CenterStrategy::EqualMass { num_centers }
764 | CenterStrategy::EqualMassCovarRepresentative { num_centers }
765 | CenterStrategy::FarthestPoint { num_centers }
766 | CenterStrategy::KMeans { num_centers, .. } => Some(*num_centers),
767 CenterStrategy::UniformGrid { .. } => None,
768 }
769}
770
771pub fn center_strategy_with_num_centers(
772 strategy: &CenterStrategy,
773 num_centers: usize,
774 d: usize,
775) -> Result<CenterStrategy, BasisError> {
776 validate_center_count(num_centers)?;
777 fn rebuild_inner(
778 strategy: &CenterStrategy,
779 num_centers: usize,
780 d: usize,
781 ) -> Result<CenterStrategy, BasisError> {
782 match strategy {
783 CenterStrategy::Auto(inner) => rebuild_inner(inner.as_ref(), num_centers, d),
784 CenterStrategy::EqualMass { .. } => Ok(CenterStrategy::EqualMass { num_centers }),
785 CenterStrategy::EqualMassCovarRepresentative { .. } => {
786 Ok(CenterStrategy::EqualMassCovarRepresentative { num_centers })
787 }
788 CenterStrategy::FarthestPoint { .. } => {
789 Ok(CenterStrategy::FarthestPoint { num_centers })
790 }
791 CenterStrategy::KMeans { max_iter, .. } => Ok(CenterStrategy::KMeans {
792 num_centers,
793 max_iter: *max_iter,
794 }),
795 CenterStrategy::UniformGrid { .. } if d == 1 => Ok(CenterStrategy::UniformGrid {
796 points_per_dim: num_centers,
797 }),
798 CenterStrategy::UserProvided(_) | CenterStrategy::UniformGrid { .. } => {
799 Err(BasisError::InvalidInput(format!(
800 "cannot replace center count for {:?} strategy",
801 center_strategy_kind(strategy)
802 )))
803 }
804 }
805 }
806 let rebuilt = rebuild_inner(strategy, num_centers, d)?;
807 Ok(match strategy {
808 CenterStrategy::Auto(_) => CenterStrategy::Auto(Box::new(rebuilt)),
809 _ => rebuilt,
810 })
811}
812
813/// Thin-plate basis configuration.
814#[derive(Debug, Clone, Serialize, Deserialize)]
815pub struct ThinPlateBasisSpec {
816 pub center_strategy: CenterStrategy,
817 #[serde(default)]
818 pub periodic: Option<Vec<Option<f64>>>,
819 pub length_scale: f64,
820 pub double_penalty: bool,
821 #[serde(default)]
822 pub identifiability: SpatialIdentifiability,
823 /// Frozen Wood-TPRS radial reparameterization. When `Some`, the builder
824 /// reuses this `(raw_radial_cols) × (kept_radial_cols)` matrix instead of
825 /// recomputing it from the constrained kernel penalty eigensystem. The
826 /// rectangular case is the truncated regression-spline path; carrying it
827 /// into prediction guarantees identical radial modes to fit-time.
828 #[serde(default)]
829 pub radial_reparam: Option<Array2<f64>>,
830}
831
832/// Per-smooth identifiability policy for spatial (TPS / Duchon) bases.
833///
834/// For a raw local basis `B` and parametric design block `C`, the orthogonalized
835/// basis is `B_c = B Z` where columns of `Z` span `null((B^T C)^T)`. This enforces:
836/// `B_c^T C = 0`
837/// in the unweighted inner product, so spatial effects cannot absorb parametric
838/// directions that actually exist in the model. The standalone basis builder has
839/// only an implicit intercept available, so it centers smooths against that
840/// intercept. The term-collection builder augments `C` with explicit linear
841/// terms when those terms are present in the formula.
842#[derive(Debug, Default, Clone, Serialize, Deserialize)]
843pub enum SpatialIdentifiability {
844 /// Keep unconstrained basis columns.
845 None,
846 /// Orthogonalize the smooth against model-owned parametric columns.
847 // "Magic" default for modular GAMs with explicit parametric block:
848 // keep spatial smooth orthogonal to intercept/linear terms.
849 // ApproxKind: Exact (orthogonalization is an exact projection).
850 #[default]
851 OrthogonalToParametric,
852 /// Freeze a fit-time transform `Z`; prediction uses `B_new * Z` unchanged.
853 FrozenTransform { transform: Array2<f64> },
854}
855
856pub(crate) use sphere_kernels::{
857 wahba_sphere_kernel_derivative_dcos_kind, wahba_sphere_kernel_from_cos_kind,
858 wahba_sphere_kernel_from_cos_simd_kind, wahba_sphere_kernel_sobolev_derivative_dcos,
859};
860
861pub use sphere_spectral::{
862 pseudo_s2_truncated_coefficients, sobolev_s2_truncated_coefficients,
863 sphere_truncated_spectral_eval,
864};
865
866/// Matérn basis configuration.
867#[derive(Debug, Clone, Serialize, Deserialize)]
868pub struct MaternBasisSpec {
869 pub center_strategy: CenterStrategy,
870 #[serde(default)]
871 pub periodic: Option<Vec<Option<f64>>>,
872 pub length_scale: f64,
873 pub nu: MaternNu,
874 #[serde(default)]
875 pub include_intercept: bool,
876 pub double_penalty: bool,
877 #[serde(default)]
878 pub identifiability: MaternIdentifiability,
879 /// Per-axis anisotropy log-scales η_a (contrasts with Ση_a = 0).
880 ///
881 /// This implements geometric anisotropy: Λ = κA where A = diag(exp(η_a)),
882 /// det(A) = 1. The kernel is evaluated at r = κ|Ah| instead of r = κ|h|.
883 /// The decomposition preserves the isotropic scaling law for global κ
884 /// and adds d−1 shape parameters for directional relevance.
885 ///
886 /// Conditional positive definiteness is preserved under any invertible
887 /// linear coordinate transform (Schoenberg), so the kernel remains valid.
888 ///
889 /// When Some, the distance is r = √(Σ_a exp(2η_a) · (x_a - c_a)²).
890 /// When None, isotropic distance r = ‖x - c‖ is used.
891 #[serde(default)]
892 pub aniso_log_scales: Option<Vec<f64>>,
893 /// Frozen double-penalty nullspace-shrinkage decision (gam#787/#860).
894 ///
895 /// `None` (the default, and the cold-build value) = decide whether to emit
896 /// the `DoublePenaltyNullspace` candidate via the κ-dependent spectral test in
897 /// `build_nullspace_shrinkage_penalty`. `Some(b)` = force the decision (set by
898 /// the freeze step from the bootstrap-κ build, mirrored from
899 /// `MaternIdentifiability::FrozenTransform`) so the learned-penalty count stays
900 /// invariant as the κ-optimizer rebuilds the design at each trial length-scale.
901 /// Only consulted when `double_penalty` is true.
902 #[serde(default)]
903 pub nullspace_shrinkage_survived: Option<bool>,
904}
905
906/// Per-smooth identifiability policy for Matérn kernel coefficients.
907///
908/// These constraints are geometric (center-based), so they are stable across
909/// train/predict and do not depend on response weights.
910#[derive(Debug, Default, Clone, Serialize, Deserialize)]
911pub enum MaternIdentifiability {
912 /// Keep the unconstrained kernel coefficient space.
913 None,
914 /// Enforce `1^T alpha = 0` at center locations (removes constant drift).
915 // Safe default with model intercepts: prevent kernel block from absorbing
916 // a global mean level.
917 #[default]
918 CenterSumToZero,
919 /// Enforce orthogonality to `[1, c_1, ..., c_d]` at centers.
920 /// Use this when explicit linear terms should own global trends.
921 CenterLinearOrthogonal,
922 /// Freeze a fit-time transform `Z` so prediction cannot drift.
923 ///
924 /// `nullspace_shrinkage_survived` freezes the double-penalty
925 /// nullspace-shrinkage decision alongside the transform (gam#787/#860). The
926 /// matern double-penalty path emits a `DoublePenaltyNullspace` candidate iff
927 /// `build_nullspace_shrinkage_penalty(&projected_kernel)` finds a near-zero
928 /// eigenvalue — but that spectral test is κ-DEPENDENT (its tolerance scales
929 /// with `λ_max`), so a near-zero eigenvalue can cross the threshold as the
930 /// κ-optimizer rebuilds the design at each trial length-scale. That flips the
931 /// learned-penalty count 6↔7 across the rebuild and the rebuilt design's ρ
932 /// dimension then disagrees with the frozen joint setup ("joint hyper rho
933 /// dimension mismatch" → every κ seed fails startup validation). Freezing the
934 /// bootstrap-κ decision here (`Some(true)` = always emit the shrinkage
935 /// candidate, `Some(false)` = never) keeps the penalty count INVARIANT across
936 /// the κ rebuild so κ actually optimizes. `None` = decide via the spectral
937 /// test (the non-frozen / cold-build behavior; also the serde back-compat
938 /// default for transforms frozen before this field existed).
939 FrozenTransform {
940 transform: Array2<f64>,
941 #[serde(default)]
942 nullspace_shrinkage_survived: Option<bool>,
943 },
944}
945
946/// Duchon null-space polynomial degree.
947///
948/// Controls the polynomial null space of the Duchon / polyharmonic spline. The
949/// Duchon seminorm `‖D^m f‖²` annihilates all polynomials of total degree
950/// `< m`, so those polynomials must be handled as explicit unpenalized columns.
951///
952/// The user-facing `order` knob selects the polynomial degree cutoff `r`, and
953/// the resulting polynomial null space has dimension `C(d + r, r)` where `d`
954/// is the covariate dimension. In the `duchon(...)` formula DSL:
955///
956/// | `order=` | Variant | max total degree | null-space dim |
957/// |----------|-----------------|------------------|-----------------|
958/// | `0` | `Zero` | 0 | `C(d+0,0) = 1` |
959/// | `1` | `Linear` | 1 | `C(d+1,1) = d+1`|
960/// | `k≥2` | `Degree(k)` | k | `C(d+k,k)` |
961///
962/// **How the polynomial null space is consumed during basis construction:**
963///
964/// 1. `polynomial_block_from_order` materialises an `(n, C(d+r,r))` block `P`
965/// of monomials up to total degree `r` at the selected `centers`.
966/// 2. `kernel_constraint_nullspace` computes `Z = null(P_centers^T)`, a
967/// `(k, k − C(d+r,r))` matrix. Reparameterising the radial kernel
968/// coefficients as `α = Z γ` enforces the side condition `P_centers^T α = 0`
969/// and yields `k − C(d+r,r)` free kernel parameters.
970/// 3. The polynomial block `P_data` evaluated at the data rows is appended to
971/// the kernel block `Φ Z`, giving a total of
972/// `(k − C(d+r,r)) + C(d+r,r) = k` columns before the spatial
973/// identifiability transform. Crucially, the total width equals the
974/// requested center count `k`, **not** `k + C(d+r,r)`.
975///
976/// **Example — `duchon(PC1, PC2, PC3, centers=10, order=1)` (d=3):**
977///
978/// - Polynomial null space: `C(3+1,1) = 4` monomials `{1, x₁, x₂, x₃}`.
979/// - Kernel columns after constraint: `10 − 4 = 6`.
980/// - Appended polynomial block: 4 columns.
981/// - Pre-identifiability total: `6 + 4 = 10` columns, i.e. exactly `centers`.
982///
983/// The variant naming matches the Duchon `m` parameter:
984/// `Zero` → `m=1`, `Linear` → `m=2`, `Degree(k)` → `m=k+1`.
985#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
986pub enum DuchonNullspaceOrder {
987 Zero,
988 Linear,
989 Degree(usize),
990}
991
992/// Duchon-like basis configuration with explicit low-frequency null-space
993/// control and explicit spectral power.
994#[derive(Debug, Clone, Serialize, Deserialize)]
995#[serde(deny_unknown_fields)]
996pub struct DuchonBasisSpec {
997 pub center_strategy: CenterStrategy,
998 #[serde(default)]
999 pub periodic: Option<Vec<Option<f64>>>,
1000 /// Optional hybrid Matérn width. `None` means pure scale-free Duchon with
1001 /// spectrum `||w||^(2p + 2s)`. `Some(length_scale)` enables the hybrid
1002 /// spectrum `||w||^(2p) * (kappa^2 + ||w||^2)^s`, `kappa = 1/length_scale`.
1003 pub length_scale: Option<f64>,
1004 /// Literal Duchon spectral power `s` (`f64`, fractional values fully
1005 /// threaded end-to-end). The pure-Duchon kernel exponent is `2(p + s) − d`,
1006 /// so this is the knob that sets `φ(r)`: `s = 0` is the integer-order Duchon
1007 /// kernel `r^{2p−d}` (its `r²·log r` log case in even `d`, ≡ the thin-plate
1008 /// kernel); `s = (d − 1)/2` gives the cubic `r³` in every dimension.
1009 ///
1010 /// This field is taken LITERALLY by the basis builder — `power = 0` means
1011 /// `s = 0`, NOT "use a default". The magic cubic default (applied when the
1012 /// user gives no explicit power) is a request-layer choice resolved by the
1013 /// formula / CLI / pyffi front-ends via [`duchon_cubic_default`]; by the time
1014 /// a spec reaches the builder this value is the final intended `s`. The
1015 /// hybrid Duchon–Matérn path (`length_scale = Some`) still requires an
1016 /// integer `s` (read via `spec.power_as_usize()`).
1017 pub power: f64,
1018 pub nullspace_order: DuchonNullspaceOrder,
1019 #[serde(default)]
1020 pub identifiability: SpatialIdentifiability,
1021 /// Per-axis anisotropy log-scales η_a.
1022 ///
1023 /// For hybrid Duchon (`length_scale=Some`), these are centered contrasts in
1024 /// the decomposition Λ = κA with det(A)=1. For pure Duchon
1025 /// (`length_scale=None`), they parameterize shape-only axis warping on the
1026 /// public path and are centered before basis evaluation/writeback so no
1027 /// global length scale is introduced.
1028 ///
1029 /// When Some, the distance is r = √(Σ_a exp(2η_a) · (x_a - c_a)²).
1030 /// When None, isotropic distance r = ‖x - c‖ is used.
1031 #[serde(default)]
1032 pub aniso_log_scales: Option<Vec<f64>>,
1033 #[serde(default)]
1034 pub operator_penalties: DuchonOperatorPenaltySpec,
1035 #[serde(default)]
1036 pub boundary: OneDimensionalBoundary,
1037 /// Data-metric radial reparameterization `V` (#1355), mirroring the
1038 /// thin-plate Wood-TPRS reparam. When `Some`, the constrained kernel
1039 /// transform is folded to `Z·V` so the realized design columns rotate into
1040 /// the `G_c`-orthonormal generalized eigenbasis of `Ω_c v = μ G_c v` and the
1041 /// native penalty becomes the diagonal curvature-per-unit-data-variance
1042 /// spectrum (mgcv's cliff), preventing the REML over-smoothing collapse to
1043 /// EDF = 1. Frozen at the cold dense build and replayed verbatim by the
1044 /// predict / κ-trial / ψ-derivative paths so they stay bit-consistent with
1045 /// the fit-time design. `None` on the lazy/streaming path (huge `n`), which
1046 /// retains the original constrained basis.
1047 #[serde(default)]
1048 pub radial_reparam: Option<Array2<f64>>,
1049}
1050
1051impl DuchonBasisSpec {
1052 /// Integer view of `power` for the existing integer-only downstream chain.
1053 /// Non-finite or non-integer values fall back to `0` (the integer-only
1054 /// validators downstream already reject this case with a clear message).
1055 pub fn power_as_usize(&self) -> usize {
1056 duchon_power_to_usize(self.power)
1057 }
1058}
1059
1060/// Convert a Duchon spectral-power `f64` into the integer view used by the
1061/// closed-form code paths. Non-finite, negative, or fractional values clamp to
1062/// `0` so the validator downstream emits the canonical error.
1063pub fn duchon_power_to_usize(power: f64) -> usize {
1064 if !power.is_finite() || power < 0.0 {
1065 return 0;
1066 }
1067 let rounded = power.round();
1068 if (rounded - power).abs() > 1e-9 {
1069 return 0;
1070 }
1071 rounded as usize
1072}
1073
1074#[derive(Clone, Debug, Serialize, Deserialize)]
1075pub struct DuchonOperatorPenaltySpec {
1076 pub mass: OperatorPenaltySpec,
1077 pub tension: OperatorPenaltySpec,
1078 pub stiffness: OperatorPenaltySpec,
1079}
1080
1081#[derive(Clone, Debug, Serialize, Deserialize)]
1082pub enum OperatorPenaltySpec {
1083 Active {
1084 initial_log_lambda: f64,
1085 prior: Option<RhoPrior>,
1086 },
1087 Disabled,
1088}
1089
1090impl Default for DuchonOperatorPenaltySpec {
1091 fn default() -> Self {
1092 // ALL ON. The Duchon penalty is a Hilbert scale: curvature is the
1093 // always-on exact RKHS `Primary` Gram and the trend ridge is always on;
1094 // the lower orders — mass (amplitude `Σ(f−f̄)²`) and tension (first-order
1095 // roughness `Σ‖∇f‖²`) — are active here, collocated on a density-blind
1096 // data-support sample. REML deselects any the data don't support (SPEC:
1097 // recover the null by default, opt INTO overfitting). Stiffness (`D2`)
1098 // stays off — `Primary` is the exact, superior curvature. (The Matérn
1099 // collocation overlay builds its own `all_active()`; SAE atoms, which
1100 // ship only `Primary`, use `all_disabled()`.)
1101 Self {
1102 mass: OperatorPenaltySpec::Active {
1103 initial_log_lambda: 0.0,
1104 prior: None,
1105 },
1106 tension: OperatorPenaltySpec::Active {
1107 initial_log_lambda: 0.0,
1108 prior: None,
1109 },
1110 stiffness: OperatorPenaltySpec::Disabled,
1111 }
1112 }
1113}
1114
1115impl DuchonOperatorPenaltySpec {
1116 pub fn has_active_operator_penalty(&self) -> bool {
1117 matches!(self.mass, OperatorPenaltySpec::Active { .. })
1118 || matches!(self.tension, OperatorPenaltySpec::Active { .. })
1119 || matches!(self.stiffness, OperatorPenaltySpec::Active { .. })
1120 }
1121
1122 pub fn all_disabled() -> Self {
1123 Self {
1124 mass: OperatorPenaltySpec::Disabled,
1125 tension: OperatorPenaltySpec::Disabled,
1126 stiffness: OperatorPenaltySpec::Disabled,
1127 }
1128 }
1129
1130 /// All three operator dials active — used by the Matérn collocation overlay.
1131 pub fn all_active() -> Self {
1132 let active = || OperatorPenaltySpec::Active {
1133 initial_log_lambda: 0.0,
1134 prior: None,
1135 };
1136 Self {
1137 mass: active(),
1138 tension: active(),
1139 stiffness: active(),
1140 }
1141 }
1142
1143 /// Operator-penalty dials appropriate for a Matérn-ν kernel in dimension `d`.
1144 ///
1145 /// The Matérn-ν RKHS is the Sobolev space `H^m` with `m = ν + d/2`: its
1146 /// squared norm controls the order-`j` derivative in L2 exactly when
1147 /// `j ≤ m`. The collocation overlay penalizes the squared L2 norms of the
1148 /// value (mass, `D0`, j=0), gradient (tension, `D1`, j=1) and Hessian
1149 /// (stiffness, `D2`, j=2). Activating a penalty whose derivative order
1150 /// exceeds the RKHS smoothness (`j > m`) imposes a roughness constraint the
1151 /// true kernel does NOT — it over-smooths the reduced-rank fit relative to
1152 /// the exact GP (mgcv `bs="gp"`, GpGp).
1153 ///
1154 /// The ν=1/2 Ornstein–Uhlenbeck kernel is the sole exception: its cusp at a
1155 /// center makes the collocated gradient/Hessian undefined, so it retains
1156 /// mass only (#707). Every differentiable order uses the inclusive Sobolev
1157 /// boundary. In particular ν=3/2 in d=1 has `m=2`, and its finite `D2`
1158 /// stiffness energy belongs to H². Omitting that block leaves the rough
1159 /// kernel with only mass+tension, inflates EDF, and changes the REML model
1160 /// class relative to its stated RKHS.
1161 pub fn matern_for_smoothness(nu: MaternNu, d: usize) -> Self {
1162 let m = nu.half_integer_value() + 0.5 * d as f64;
1163 // Tolerance keeps the mathematically inclusive `j ≤ m` boundary stable
1164 // under floating-point representation of half-integer orders.
1165 const ORDER_EPS: f64 = 1e-9;
1166 let active = || OperatorPenaltySpec::Active {
1167 initial_log_lambda: 0.0,
1168 prior: None,
1169 };
1170 let gate = |order: f64| {
1171 if !matches!(nu, MaternNu::Half) && m + ORDER_EPS >= order {
1172 active()
1173 } else {
1174 OperatorPenaltySpec::Disabled
1175 }
1176 };
1177 Self {
1178 mass: active(),
1179 tension: gate(1.0),
1180 stiffness: gate(2.0),
1181 }
1182 }
1183}
1184
1185pub fn minimum_duchon_power_for_operator_penalties(
1186 dim: usize,
1187 nullspace_order: DuchonNullspaceOrder,
1188 max_operator_derivative_order: usize,
1189) -> usize {
1190 let p = duchon_p_from_nullspace_order(nullspace_order);
1191 let mut s = 0usize;
1192 while 2 * (p + s) <= dim + max_operator_derivative_order {
1193 s += 1;
1194 }
1195 s
1196}
1197
1198/// Resolve a fully admissible Duchon `(nullspace_order, power)` pair.
1199///
1200/// Three constraints fold into one resolution:
1201/// (a) operator collocation up to `max_op`: `2(p + s) > d + max_op`
1202/// (b) pure-mode CPD vs polynomial nullspace P_p: `2s < d`
1203/// (Wendland Thm 8.17: pure polyharmonic kernel of order m = p+s in
1204/// R^d is CPD of order `m − ⌊d/2⌋ + 1[d even, log] / m − (d−1)/2
1205/// [d odd]`, and Duchon interpolation against P_p is well-posed iff
1206/// CPD order ≤ p, which collapses to `2s < d` since 2s, d are
1207/// integers and 2s is even.)
1208/// (a) implies the kernel-existence condition `2(p + s) > d`.
1209/// (b) is dropped when `length_scale` is `Some` (hybrid Matérn-blended
1210/// kernel is strictly PD, CPD order 0).
1211///
1212/// Strategy: at the requested `nullspace_order`, take the smallest `s`
1213/// satisfying (a). If that `s` violates (b) in pure mode, escalate the
1214/// nullspace order by one and retry. Termination: at `p ≥ ⌈(d+max_op)/2⌉ + 1`
1215/// the operator constraint (a) admits `s = 0`, and `0 < d` satisfies (b)
1216/// for any `d ≥ 1`, so escalation always converges.
1217///
1218/// The returned nullspace order is monotone in the request: it never
1219/// decreases the user's requested order — only strengthens it when pure-mode
1220/// CPD requires a richer polynomial absorption space.
1221pub fn resolve_duchon_orders(
1222 dim: usize,
1223 requested_nullspace_order: DuchonNullspaceOrder,
1224 max_operator_derivative_order: usize,
1225 length_scale: Option<f64>,
1226) -> (DuchonNullspaceOrder, usize) {
1227 assert!(dim >= 1, "Duchon basis requires dim >= 1");
1228 let pure = length_scale.is_none();
1229 let mut nullspace = requested_nullspace_order;
1230 // Bounded loop: escalation terminates by the argument above.
1231 for _ in 0..=(dim + max_operator_derivative_order + 1) {
1232 let p = duchon_p_from_nullspace_order(nullspace);
1233 // Smallest s with 2(p + s) > d + max_op:
1234 // 2p > d + max_op ⇒ s = 0
1235 // else s = ⌈(d + max_op + 1 − 2p) / 2⌉ = (d + max_op + 2 − 2p) / 2
1236 let s_op = if 2 * p > dim + max_operator_derivative_order {
1237 0
1238 } else {
1239 (dim + max_operator_derivative_order + 2 - 2 * p) / 2
1240 };
1241 if !pure || 2 * s_op < dim {
1242 return (nullspace, s_op);
1243 }
1244 nullspace = duchon_next_nullspace_order(nullspace);
1245 }
1246 // Bounded-loop fallback: by the analysis in the docstring, for
1247 // `p >= ceil((dim + max_op) / 2) + 1` the operator constraint admits
1248 // `s = 0` and (in pure mode) `0 < dim` satisfies the kernel-existence
1249 // condition. The loop above always reaches that regime within the bound,
1250 // so returning the last `nullspace` with `s = 0` is a valid answer.
1251 (nullspace, 0)
1252}
1253
1254#[inline]
1255pub(crate) fn duchon_next_nullspace_order(order: DuchonNullspaceOrder) -> DuchonNullspaceOrder {
1256 match order {
1257 DuchonNullspaceOrder::Zero => DuchonNullspaceOrder::Linear,
1258 DuchonNullspaceOrder::Linear => DuchonNullspaceOrder::Degree(2),
1259 DuchonNullspaceOrder::Degree(k) => DuchonNullspaceOrder::Degree(k + 1),
1260 }
1261}
1262
1263pub(crate) fn duchon_previous_nullspace_order(order: DuchonNullspaceOrder) -> DuchonNullspaceOrder {
1264 match order {
1265 DuchonNullspaceOrder::Zero => DuchonNullspaceOrder::Zero,
1266 DuchonNullspaceOrder::Linear => DuchonNullspaceOrder::Zero,
1267 DuchonNullspaceOrder::Degree(2) => DuchonNullspaceOrder::Linear,
1268 DuchonNullspaceOrder::Degree(k) => DuchonNullspaceOrder::Degree(k - 1),
1269 }
1270}
1271
1272/// Returns the maximum derivative order required by the *active* operator
1273/// penalties: 2 if stiffness is Active, else 1 if tension is Active, else 0.
1274/// Mass-only (or no active operator) penalties only require kernel validity
1275/// (`2(p+s) > d`), tension requires D1 collocation (`2(p+s) > d+1`), and
1276/// stiffness requires D2 collocation (`2(p+s) > d+2`).
1277pub fn duchon_max_active_operator_derivative_order(
1278 operator_penalties: &DuchonOperatorPenaltySpec,
1279) -> usize {
1280 if matches!(
1281 operator_penalties.stiffness,
1282 OperatorPenaltySpec::Active { .. }
1283 ) {
1284 2
1285 } else if matches!(
1286 operator_penalties.tension,
1287 OperatorPenaltySpec::Active { .. }
1288 ) {
1289 1
1290 } else {
1291 0
1292 }
1293}
1294
1295/// Metadata returned by generic basis builders.
1296#[derive(Debug, Clone)]
1297pub enum BasisMetadata {
1298 BSpline1D {
1299 knots: Array1<f64>,
1300 identifiability_transform: Option<Array2<f64>>,
1301 periodic: Option<(f64, f64, usize)>,
1302 /// Effective B-spline polynomial degree carried by `knots`.
1303 ///
1304 /// Persisted alongside `knots` so prediction can reconstruct an
1305 /// evaluator that matches fit-time geometry, even when the fit-time
1306 /// auto-shrink (issue #340) reduced the user's requested degree to
1307 /// fit the available data (`n` too small for cubic ⇒ quadratic ⇒
1308 /// linear). When `None` the consumer should fall back to the
1309 /// upstream `BSplineBasisSpec.degree` (legacy / non-shrunk path).
1310 degree: Option<usize>,
1311 /// Human-readable description of an automatic basis shrink (issue #340)
1312 /// when the user's requested `(degree, num_internal_knots)` exceeded the
1313 /// available evaluation count `n`. `Some(note)` records the before→after
1314 /// configuration; `None` means no auto-shrink occurred for this basis.
1315 auto_shrink_note: Option<String>,
1316 },
1317 /// Natural cubic regression spline (`bs="cr"`/`"cs"`) metadata (#1074).
1318 ///
1319 /// `knots` are the `k` Lancaster–Salkauskas knots that index the basis
1320 /// values directly (basis dim = `knots.len()`). Predict-time rebuilds
1321 /// reconstruct the cr geometry from `knots` and replay the captured
1322 /// `identifiability_transform` exactly, mirroring `BSpline1D`.
1323 CubicRegression1D {
1324 knots: Array1<f64>,
1325 identifiability_transform: Option<Array2<f64>>,
1326 },
1327 ThinPlate {
1328 centers: Array2<f64>,
1329 length_scale: f64,
1330 periodic: Option<Vec<Option<f64>>>,
1331 identifiability_transform: Option<Array2<f64>>,
1332 /// Per-column standard deviations used for input standardization (d > 1).
1333 input_scales: Option<Vec<f64>>,
1334 /// Wood-TPRS radial reparameterization carried into prediction so the
1335 /// rotated radial basis at predict-time matches fit-time exactly. `None`
1336 /// in the lazy/streaming path which retains the original basis.
1337 radial_reparam: Option<Array2<f64>>,
1338 },
1339 Sphere {
1340 centers: Array2<f64>,
1341 penalty_order: usize,
1342 method: SphereMethod,
1343 max_degree: Option<usize>,
1344 wahba_kernel: SphereWahbaKernel,
1345 constraint_transform: Option<Array2<f64>>,
1346 },
1347 /// Constant-curvature (`M_κ`) geodesic-kernel smooth (#944). `kappa` and
1348 /// the realized `length_scale` are persisted so predict-time (and the
1349 /// future ψ-channel per-trial) rebuilds replay the exact fit-time
1350 /// geometry; `constraint_transform` is the composed `z · z_parametric`
1351 /// frozen by the global identifiability pipeline (#532 pattern).
1352 ConstantCurvature {
1353 centers: Array2<f64>,
1354 kappa: f64,
1355 length_scale: f64,
1356 constraint_transform: Option<Array2<f64>>,
1357 },
1358 /// Measure-jet spline smooth: multiscale local-jet-residual energy of the
1359 /// empirical measure, quadratured on the center set. `centers` are the
1360 /// REALIZED barycenter nodes; `order_s` stores the spec's order sentinel
1361 /// verbatim as the mode marker (0.0 = per-level/spectral, > 0 = fused
1362 /// pin — persisting a realized default would flip the rebuilt mode). The
1363 /// penalty depends on the FIT data through `masses`, the realized
1364 /// `eps_band`, the support anchors, and the normalization scales, so all
1365 /// are persisted and replayed verbatim by
1366 /// predict-time (and per-ψ-trial) rebuilds — recomputing either from
1367 /// predict rows would change the penalty the coefficients were estimated
1368 /// under. `constraint_transform` is the composed `z · z_parametric`
1369 /// frozen by the global identifiability pipeline (#532 pattern).
1370 MeasureJet {
1371 centers: Array2<f64>,
1372 input_scales: Option<Vec<f64>>,
1373 length_scale: f64,
1374 eps_band: Vec<f64>,
1375 order_s: f64,
1376 alpha: f64,
1377 tau0: f64,
1378 masses: Array1<f64>,
1379 support_means: Vec<f64>,
1380 penalty_normalization_scales: Vec<f64>,
1381 raw_penalty_normalization_scales: Vec<f64>,
1382 fused_penalty_normalization_scale: Option<f64>,
1383 constraint_transform: Option<Array2<f64>>,
1384 /// Ambient input-measurement-error scale `σ_coord` (issue #2225): the
1385 /// perpendicular off-manifold residual spread of the fit rows, in the
1386 /// centers' (standardized) frame. `None` when it could not be estimated.
1387 /// Carried into `MeasureJetFrozenQuadrature::sigma_coord` at freeze time.
1388 sigma_coord: Option<f64>,
1389 },
1390 Matern {
1391 centers: Array2<f64>,
1392 length_scale: f64,
1393 periodic: Option<Vec<Option<f64>>>,
1394 nu: MaternNu,
1395 include_intercept: bool,
1396 identifiability_transform: Option<Array2<f64>>,
1397 /// Per-column standard deviations used for input standardization (d > 1).
1398 input_scales: Option<Vec<f64>>,
1399 /// Per-axis anisotropy log-scales η_a for geometric anisotropy.
1400 /// When Some, distance is r = √(Σ_a exp(2η_a) · (x_a - c_a)²).
1401 aniso_log_scales: Option<Vec<f64>>,
1402 /// Realized double-penalty nullspace-shrinkage decision at this build
1403 /// (gam#787/#860). The freeze step pins this into
1404 /// `MaternIdentifiability::FrozenTransform::nullspace_shrinkage_survived`
1405 /// so the κ-optimizer's per-trial rebuilds keep the learned-penalty count
1406 /// invariant (otherwise the κ-dependent spectral test flips it 6↔7 → "joint
1407 /// hyper rho dimension mismatch").
1408 nullspace_shrinkage_survived: bool,
1409 },
1410 Duchon {
1411 centers: Array2<f64>,
1412 length_scale: Option<f64>,
1413 periodic: Option<Vec<Option<f64>>>,
1414 power: f64,
1415 nullspace_order: DuchonNullspaceOrder,
1416 identifiability_transform: Option<Array2<f64>>,
1417 /// Per-column standard deviations used for input standardization (d > 1).
1418 input_scales: Option<Vec<f64>>,
1419 /// Per-axis anisotropy log-scales η_a, stored for prediction.
1420 aniso_log_scales: Option<Vec<f64>>,
1421 /// Support points used to build the active lower-order operator
1422 /// penalties (mass/tension/stiffness). Stored so runtime adaptive
1423 /// caches can rebuild the exact same operator rows instead of guessing
1424 /// from centers.
1425 operator_collocation_points: Option<Array2<f64>>,
1426 /// Data-metric radial reparameterization `V` (#1355). When `Some`, the
1427 /// constrained kernel transform is folded to `Z·V` so predict-time and
1428 /// κ-trial rebuilds replay the exact fit-time rotated radial basis.
1429 /// `None` on the lazy/streaming path (original constrained basis).
1430 radial_reparam: Option<Array2<f64>>,
1431 },
1432 Pca {
1433 feature_cols: Vec<usize>,
1434 basis_matrix: Array2<f64>,
1435 centered: bool,
1436 smooth_penalty: f64,
1437 center_mean: Option<Array1<f64>>,
1438 pca_basis_path: Option<std::path::PathBuf>,
1439 chunk_size: usize,
1440 },
1441 TensorBSpline {
1442 feature_cols: Vec<usize>,
1443 knots: Vec<Array1<f64>>,
1444 degrees: Vec<usize>,
1445 periods: Vec<Option<f64>>,
1446 /// Per-margin flag: `true` when that margin is a natural cubic
1447 /// regression spline (`NaturalCubicRegression` knotspec) rather than an
1448 /// open/periodic B-spline (#1074). Persisted so the tensor freeze
1449 /// rebuilds the cr marginal knotspec (value-at-knot) instead of an open
1450 /// `Provided(knots)` B-spline, keeping predict-time marginals identical
1451 /// to the fit-time cr margins. Defaults to all-`false` (legacy B-spline
1452 /// tensors) when deserialized from an older persisted model (the
1453 /// older-model default is applied on the persisted `SmoothBasisSpec`
1454 /// side; `BasisMetadata` itself is transient builder output and is not
1455 /// serde-serialized, so it carries no `#[serde]` attributes).
1456 is_cr: Vec<bool>,
1457 identifiability_transform: Option<Array2<f64>>,
1458 },
1459 SphereHarmonics {
1460 max_degree: usize,
1461 radians: bool,
1462 },
1463 /// Wrap an inner basis metadata to record a multiplicative `by` (continuous or
1464 /// factor) along a column of the dataset.
1465 BySmooth {
1466 inner: Box<BasisMetadata>,
1467 by_col: usize,
1468 levels: Option<Vec<u64>>,
1469 ordered: bool,
1470 },
1471 /// Factor-by-smooth (mgcv-style `s(x, by=g, bs="fs"|"sz"|"re")`).
1472 FactorSmooth {
1473 continuous_cols: Vec<usize>,
1474 group_col: usize,
1475 knots: Array1<f64>,
1476 degree: usize,
1477 periodic: Option<(f64, f64, usize)>,
1478 group_levels: Vec<u64>,
1479 flavour: String,
1480 /// `true` when the per-level marginal is a cubic regression spline
1481 /// (`NaturalCubicRegression` knotspec, mgcv's `bs="sz"` default marginal,
1482 /// #1074). Predict-time freeze must then restore a cr knotspec from the
1483 /// stored value-knots rather than treating them as a B-spline knot
1484 /// vector. Defaults to `false` (B-spline marginal) for backward compat.
1485 marginal_is_cr: bool,
1486 },
1487}
1488
1489/// Standardized basis build result for engine-level composition.
1490#[derive(Clone)]
1491pub struct BasisBuildResult {
1492 pub design: DesignMatrix,
1493 pub penalties: Vec<Array2<f64>>,
1494 pub nullspace_dims: Vec<usize>,
1495 pub penaltyinfo: Vec<PenaltyInfo>,
1496 pub metadata: BasisMetadata,
1497 /// Optional factored rowwise-Kronecker representation for tensor-product
1498 /// bases. When present, downstream code can keep the design operator-backed
1499 /// instead of forcing a fully materialized `n x prod(q_j)` block.
1500 pub kronecker_factored: Option<KroneckerFactoredBasis>,
1501 /// Per-active-penalty operator handles (parallel to `penalties`). Each
1502 /// entry is `Some(op)` when the closed-form factory emitted an op-form
1503 /// penalty bit-equivalent to the dense matrix, `None` for ordinary dense
1504 /// penalties. Downstream consumers route through the `Some` entries to
1505 /// avoid materializing dense `p x p` Grams in exact operator algebra.
1506 pub ops: Vec<Option<std::sync::Arc<dyn crate::analytic_penalties::PenaltyOp>>>,
1507 /// Per-active-penalty null-space eigenvector matrices (parallel to
1508 /// `penalties`). Each entry is `Some(U_null)` with `U_null.ncols() ==
1509 /// nullspace_dims[k]` when the active block has a non-trivial null space
1510 /// (eigenvalues ≤ spectral tolerance), and `None` when the block is
1511 /// already full-rank. The columns of `U_null` are the eigenvectors of
1512 /// `sym_penalty` at the (near-)zero eigenvalues — i.e., an orthonormal
1513 /// basis of `null(S_block)` in the block's own coordinate system.
1514 ///
1515 /// This is the raw spectral data that the construction pipeline uses to
1516 /// absorb each smooth's penalty null space into the parametric block
1517 /// (reparameterize-and-split). Without absorption the inner Newton solve
1518 /// cannot converge on data whose unpenalized signal lies along a null
1519 /// direction of `S` (phantom-multiplier refusal at the KKT certificate).
1520 pub null_eigenvectors: Vec<Option<Array2<f64>>>,
1521 /// Joint-null absorption rotation for this basis, when the basis carries
1522 /// any penalties with a non-trivial joint null space.
1523 ///
1524 /// `Some(rotation)` records `Q = [U_range | U_null]` where `U_null` spans
1525 /// the joint null space `null(Σ_k S_k)` over this basis's active
1526 /// penalties (unscaled — the structural joint null is independent of
1527 /// `λ`). After the basis pipeline applies this rotation, the design
1528 /// becomes `X · Q` and each penalty becomes `Qᵀ S_k Q`, block-diagonal
1529 /// with a guaranteed-zero null tail. The same `Q` must be replayed at
1530 /// prediction time, so it is persisted in the fitted model. `None`
1531 /// indicates either no penalties on this basis, or a full-rank joint
1532 /// penalty (joint nullity = 0). A `Some` value is never recorded with
1533 /// `joint_nullity == 0` — the `None` discriminant is canonical for
1534 /// "nothing to absorb".
1535 ///
1536 /// Stage-2 commit A: this field is plumbed into the struct but neither
1537 /// computed nor applied yet. Stage-2 commit B populates it; Stage-2
1538 /// commit D applies the rotation to `design` and `penalties`.
1539 pub joint_null_rotation: Option<JointNullRotation>,
1540}
1541
1542/// Joint-null absorption rotation, attached to a smooth's basis when the
1543/// basis's joint penalty `Σ_k S_k` has a non-trivial null space.
1544///
1545/// The `rotation` field stores the orthonormal eigenvector matrix
1546/// `Q = [U_range | U_null]` of the symmetric joint penalty: the first
1547/// `range_dim = rotation.ncols() - joint_nullity` columns span
1548/// `range(Σ_k S_k)`; the remaining `joint_nullity` columns span
1549/// `null(Σ_k S_k)`. After the pipeline applies the rotation, the smooth's
1550/// coefficient vector satisfies `β = Q · γ`, the design becomes `X · Q`,
1551/// and each per-block penalty `S_k` becomes `Qᵀ S_k Q`, which is guaranteed
1552/// block-diagonal with a zero `(joint_nullity × joint_nullity)` tail
1553/// (because the joint null annihilates every active `S_k`).
1554#[derive(Clone, Serialize, Deserialize)]
1555pub struct JointNullRotation {
1556 /// `(p_smooth × p_smooth)` orthonormal matrix; range columns first,
1557 /// joint-null columns last.
1558 pub rotation: Array2<f64>,
1559 /// Number of columns at the tail of `rotation` that span the joint
1560 /// null space. Always `> 0` when wrapped in `Some` — the value `0`
1561 /// is encoded as `None`.
1562 pub joint_nullity: usize,
1563}
1564
1565impl std::fmt::Debug for JointNullRotation {
1566 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1567 f.debug_struct("JointNullRotation")
1568 .field(
1569 "rotation",
1570 &format_args!("{}×{}", self.rotation.nrows(), self.rotation.ncols()),
1571 )
1572 .field("joint_nullity", &self.joint_nullity)
1573 .finish()
1574 }
1575}
1576
1577impl std::fmt::Debug for BasisBuildResult {
1578 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1579 f.debug_struct("BasisBuildResult")
1580 .field("design", &self.design)
1581 .field("penalties", &self.penalties)
1582 .field("nullspace_dims", &self.nullspace_dims)
1583 .field("penaltyinfo", &self.penaltyinfo)
1584 .field("metadata", &self.metadata)
1585 .field("kronecker_factored", &self.kronecker_factored)
1586 .field(
1587 "ops",
1588 &format_args!(
1589 "[{}]",
1590 self.ops
1591 .iter()
1592 .map(|o| if o.is_some() { "Some" } else { "None" })
1593 .collect::<Vec<_>>()
1594 .join(", ")
1595 ),
1596 )
1597 .field(
1598 "null_eigenvectors",
1599 &format_args!(
1600 "[{}]",
1601 self.null_eigenvectors
1602 .iter()
1603 .map(|u| match u {
1604 Some(m) => format!("Some({}x{})", m.nrows(), m.ncols()),
1605 None => "None".to_string(),
1606 })
1607 .collect::<Vec<_>>()
1608 .join(", ")
1609 ),
1610 )
1611 .field("joint_null_rotation", &self.joint_null_rotation)
1612 .finish()
1613 }
1614}
1615
1616/// Factored tensor-product basis metadata for operator-backed downstream use.
1617#[derive(Debug)]
1618pub struct KroneckerFactoredBasis {
1619 /// Marginal design matrices: `marginal_designs[j]` is `(n, q_j)`.
1620 pub marginal_designs: Vec<Array2<f64>>,
1621 /// Marginal penalty matrices: `marginal_penalties[k]` is `(q_k, q_k)`.
1622 pub marginal_penalties: Vec<Array2<f64>>,
1623 /// Marginal basis dimensions: `[q_0, ..., q_{d-1}]`.
1624 pub marginal_dims: Vec<usize>,
1625 /// Whether the system includes a global ridge (double) penalty.
1626 pub has_double_penalty: bool,
1627 /// λ-invariant tensor structure (marginal eigensystems, reparameterized
1628 /// marginals, shrinkage scale), memoized once per fit. The marginal
1629 /// designs/penalties are fixed for the whole fit, so the expensive marginal
1630 /// `eigh()` and `B_k·U_k` GEMMs only need to run once — every outer REML
1631 /// iterate (50+ on the #1082 tensor cases) then reuses this. Filled lazily
1632 /// on first use via [`Self::invariant_structure`]. NOT serialized and reset
1633 /// to empty on `Clone` (it is purely a within-fit performance cache; a fresh
1634 /// owner recomputes on first demand, keeping every result bit-identical).
1635 invariant: std::sync::OnceLock<std::sync::Arc<crate::kronecker::KroneckerInvariantStructure>>,
1636}
1637
1638impl Clone for KroneckerFactoredBasis {
1639 fn clone(&self) -> Self {
1640 Self {
1641 marginal_designs: self.marginal_designs.clone(),
1642 marginal_penalties: self.marginal_penalties.clone(),
1643 marginal_dims: self.marginal_dims.clone(),
1644 has_double_penalty: self.has_double_penalty,
1645 // Propagate the memoized structure when present so a clone made
1646 // mid-fit keeps the hoist; otherwise start empty (recomputed on
1647 // first demand, identical result).
1648 invariant: match self.invariant.get() {
1649 Some(s) => {
1650 let cell = std::sync::OnceLock::new();
1651 cell.get_or_init(|| std::sync::Arc::clone(s));
1652 cell
1653 }
1654 None => std::sync::OnceLock::new(),
1655 },
1656 }
1657 }
1658}
1659
1660impl KroneckerFactoredBasis {
1661 /// Construct from the fixed marginal data with an empty invariant cache.
1662 pub fn new(
1663 marginal_designs: Vec<Array2<f64>>,
1664 marginal_penalties: Vec<Array2<f64>>,
1665 marginal_dims: Vec<usize>,
1666 has_double_penalty: bool,
1667 ) -> Self {
1668 Self {
1669 marginal_designs,
1670 marginal_penalties,
1671 marginal_dims,
1672 has_double_penalty,
1673 invariant: std::sync::OnceLock::new(),
1674 }
1675 }
1676
1677 /// Lazily compute (once) and return the λ-invariant tensor structure
1678 /// (marginal eigensystems, reparameterized marginals, shrinkage scale).
1679 ///
1680 /// Computed from the fixed marginal designs/penalties, so the result is the
1681 /// same on every call within a fit; the first call pays the `eigh()` cost
1682 /// and every later call is a pointer load. Because the cache is keyed on the
1683 /// fixed marginal data and `marginal_penalties`/`marginal_designs` are
1684 /// immutable for the fit's lifetime, no invalidation is needed.
1685 pub fn invariant_structure(
1686 &self,
1687 ) -> Result<std::sync::Arc<crate::kronecker::KroneckerInvariantStructure>, BasisError> {
1688 // Fast path: already memoized.
1689 if let Some(s) = self.invariant.get() {
1690 return Ok(std::sync::Arc::clone(s));
1691 }
1692 // Compute outside the cell (fallible) and install via `get_or_init`. If a
1693 // concurrent racer already won, `get_or_init` drops our `computed` and
1694 // returns the stored one; either way the value is the unique function of
1695 // the fixed marginal data, so the returned Arc is correct.
1696 let computed = std::sync::Arc::new(crate::kronecker::KroneckerInvariantStructure::compute(
1697 &self.marginal_designs,
1698 &self.marginal_penalties,
1699 &self.marginal_dims,
1700 )?);
1701 let installed = self.invariant.get_or_init(|| computed);
1702 Ok(std::sync::Arc::clone(installed))
1703 }
1704}
1705
1706#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1707pub enum PenaltySource {
1708 Primary,
1709 DoublePenaltyNullspace,
1710 OperatorMass,
1711 OperatorTension,
1712 OperatorStiffness,
1713 /// One per input axis `a` of a multivariate Duchon smooth: the gradient
1714 /// energy along axis `a`, `Σ(∂f/∂x_a)²`, each with its own REML λ_a. REML
1715 /// shrinks an axis's contribution toward flat only when it does not earn
1716 /// its keep — penalty-based ARD / variable relevance, the replacement for
1717 /// brittle kernel-η optimization. Emitted when `scale_dims` is on.
1718 OperatorRelevance {
1719 axis: usize,
1720 },
1721 TensorMarginal {
1722 dim: usize,
1723 },
1724 TensorSeparable {
1725 penalized_margins: Vec<usize>,
1726 },
1727 TensorGlobalRidge,
1728 Other(String),
1729}
1730
1731#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1732pub enum PenaltyDropReason {
1733 ZeroMatrix,
1734 NumericalRankZero,
1735}
1736
1737fn default_normalization_scale() -> f64 {
1738 1.0
1739}
1740
1741#[derive(Debug, Clone, Serialize, Deserialize)]
1742pub struct PenaltyInfo {
1743 pub source: PenaltySource,
1744 pub original_index: usize,
1745 pub active: bool,
1746 pub effective_rank: usize,
1747 pub dropped_reason: Option<PenaltyDropReason>,
1748 pub nullspace_dim_hint: usize,
1749 #[serde(default = "default_normalization_scale")]
1750 pub normalization_scale: f64,
1751 /// Kronecker factors preserved from tensor penalty construction.
1752 /// When present, spectral decomposition can use per-factor eigendecomposition.
1753 #[serde(skip)]
1754 pub kronecker_factors: Option<Vec<Array2<f64>>>,
1755}
1756
1757#[derive(Clone)]
1758pub struct PenaltyCandidate {
1759 pub matrix: Array2<f64>,
1760 pub nullspace_dim_hint: usize,
1761 pub source: PenaltySource,
1762 pub normalization_scale: f64,
1763 /// Optional Kronecker factors whose product equals `matrix`.
1764 /// When present, spectral decomposition can be done per-factor
1765 /// (O(Σ q_j³) instead of O((Π q_j)³)).
1766 pub kronecker_factors: Option<Vec<Array2<f64>>>,
1767 /// Optional operator-form handle whose `as_dense()` matches `matrix`. When
1768 /// populated by the closed-form factories, this is propagated through to
1769 /// `CanonicalPenaltyBlock` so downstream consumers can use exact matvec
1770 /// algebra without rebuilding the dense Gram. When `None`, only the dense
1771 /// `matrix` path is available.
1772 pub op: Option<std::sync::Arc<dyn crate::analytic_penalties::PenaltyOp>>,
1773}
1774
1775impl std::fmt::Debug for PenaltyCandidate {
1776 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1777 f.debug_struct("PenaltyCandidate")
1778 .field(
1779 "matrix",
1780 &format_args!("{}×{}", self.matrix.nrows(), self.matrix.ncols()),
1781 )
1782 .field("nullspace_dim_hint", &self.nullspace_dim_hint)
1783 .field("source", &self.source)
1784 .field("normalization_scale", &self.normalization_scale)
1785 .field(
1786 "kronecker_factors",
1787 &self.kronecker_factors.as_ref().map(|v| v.len()),
1788 )
1789 .field("op", &self.op.as_ref().map(|o| o.dim()))
1790 .finish()
1791 }
1792}
1793
1794#[derive(Clone)]
1795pub struct CanonicalPenaltyBlock {
1796 pub sym_penalty: Array2<f64>,
1797 /// Eigenvalues from spectral decomposition (retained to avoid recomputation).
1798 pub eigenvalues: Array1<f64>,
1799 /// Eigenvectors from spectral decomposition (retained to avoid recomputation).
1800 pub eigenvectors: Array2<f64>,
1801 pub rank: usize,
1802 pub nullity: usize,
1803 /// Number of genuine negative-curvature eigendirections (`ev < -tol`).
1804 /// A non-PSD penalty has `negative_dim > 0`; these directions are
1805 /// neither range nor null and are never absorbed as unpenalized (#1425).
1806 pub negative_dim: usize,
1807 pub tol: f64,
1808 pub iszero: bool,
1809 /// Optional operator-form handle that is bit-equivalent to `sym_penalty`.
1810 /// Propagated from `PenaltyCandidate.op` when present so downstream
1811 /// consumers can use matvec without rebuilding the dense Gram.
1812 pub op: Option<std::sync::Arc<dyn crate::analytic_penalties::PenaltyOp>>,
1813}
1814
1815impl std::fmt::Debug for CanonicalPenaltyBlock {
1816 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1817 f.debug_struct("CanonicalPenaltyBlock")
1818 .field(
1819 "sym_penalty",
1820 &format_args!("{}×{}", self.sym_penalty.nrows(), self.sym_penalty.ncols()),
1821 )
1822 .field("eigenvalues", &self.eigenvalues)
1823 .field(
1824 "eigenvectors",
1825 &format_args!(
1826 "{}×{}",
1827 self.eigenvectors.nrows(),
1828 self.eigenvectors.ncols()
1829 ),
1830 )
1831 .field("rank", &self.rank)
1832 .field("nullity", &self.nullity)
1833 .field("negative_dim", &self.negative_dim)
1834 .field("tol", &self.tol)
1835 .field("iszero", &self.iszero)
1836 .field("op", &self.op.as_ref().map(|o| o.dim()))
1837 .finish()
1838 }
1839}
1840
1841#[derive(Debug)]
1842pub struct BasisPsiDerivativeResult {
1843 pub design_derivative: Array2<f64>,
1844 pub penalties_derivative: Vec<Array2<f64>>,
1845 /// Operator-backed design derivative for standalone first-derivative
1846 /// callers. Bundled first+second callers receive the shared operator on
1847 /// `BasisPsiDerivativeBundle` instead.
1848 pub implicit_operator: Option<ImplicitDesignPsiDerivative>,
1849}
1850
1851#[derive(Debug)]
1852pub struct BasisPsiSecondDerivativeResult {
1853 pub designsecond_derivative: Array2<f64>,
1854 pub penaltiessecond_derivative: Vec<Array2<f64>>,
1855 /// Operator-backed design derivative for standalone second-derivative
1856 /// callers. Bundled first+second callers receive the shared operator on
1857 /// `BasisPsiDerivativeBundle` instead.
1858 pub implicit_operator: Option<ImplicitDesignPsiDerivative>,
1859}
1860
1861#[derive(Debug)]
1862pub struct BasisPsiDerivativeBundle {
1863 pub first: BasisPsiDerivativeResult,
1864 pub second: BasisPsiSecondDerivativeResult,
1865 /// Shared operator-backed design derivative for the first and second
1866 /// psi derivatives. Bundled callers consume this once instead of storing
1867 /// duplicate materialized/streaming operators in both derivative payloads.
1868 pub implicit_operator: Option<ImplicitDesignPsiDerivative>,
1869}
1870
1871/// Per-axis psi_a derivative package for anisotropic spatial terms.
1872///
1873/// For a d-dimensional anisotropic term, the kernel phi(r) depends on
1874/// the anisotropic distance r = |Lambda h| where Lambda = diag(kappa_a). Each axis a
1875/// has its own log-scale psi_a = log(kappa_a), yielding d first derivatives,
1876/// d diagonal second derivatives, and d*(d-1)/2 cross second derivatives.
1877///
1878/// The cross second derivative d2 phi/(d psi_a d psi_b) = t * s_a * s_b (a != b)
1879/// is rank-1, so we store the t_values and s_components vectors rather
1880/// than materializing d^2 matrices.
1881#[derive(Clone)]
1882pub struct AnisoBasisPsiDerivatives {
1883 /// d matrices, each (n x p_smooth): dX/d psi_a.
1884 pub design_first: Vec<Array2<f64>>,
1885 /// d matrices, each (n x p_smooth): d2X/d psi_a^2 (diagonal second derivatives).
1886 pub design_second_diag: Vec<Array2<f64>>,
1887 /// Cross second derivatives d2X/(d psi_a d psi_b) for a < b.
1888 pub design_second_cross: Vec<Array2<f64>>,
1889 /// Axis-pair indices corresponding to `design_second_cross`.
1890 pub design_second_cross_pairs: Vec<(usize, usize)>,
1891 /// d x num_penalties: dS_m/d psi_a for each axis a and penalty m.
1892 pub penalties_first: Vec<Vec<Array2<f64>>>,
1893 /// d x num_penalties: d2S_m/d psi_a^2 for each axis a and penalty m.
1894 pub penalties_second_diag: Vec<Vec<Array2<f64>>>,
1895 /// The (a, b) axis pairs supported by the on-demand cross-penalty
1896 /// provider. Only the upper triangle (a < b) is stored.
1897 pub penalties_cross_pairs: Vec<(usize, usize)>,
1898 /// On-demand cross-penalty second-derivative provider. Exact anisotropic
1899 /// cross-axis penalty seconds are streamed one pair at a time rather than
1900 /// stored as a dense upper triangle of blocks.
1901 pub penalties_cross_provider: Option<AnisoPenaltyCrossProvider>,
1902 /// Shared operator-backed representation of the anisotropic kernel-side
1903 /// design derivatives. When `design_first` / `design_second_diag` are empty,
1904 /// callers must use this operator directly; when they are present, this
1905 /// operator still provides exact cross-axis second derivatives without
1906 /// duplicating separate `t` / `s_a` storage layouts.
1907 pub implicit_operator: Option<ImplicitDesignPsiDerivative>,
1908}
1909
1910#[derive(Clone)]
1911pub struct AnisoPenaltyCrossProvider(
1912 std::sync::Arc<
1913 dyn Fn(usize, usize) -> Result<Vec<Array2<f64>>, BasisError> + Send + Sync + 'static,
1914 >,
1915);
1916
1917impl AnisoPenaltyCrossProvider {
1918 pub(crate) fn new<F>(f: F) -> Self
1919 where
1920 F: Fn(usize, usize) -> Result<Vec<Array2<f64>>, BasisError> + Send + Sync + 'static,
1921 {
1922 Self(std::sync::Arc::new(f))
1923 }
1924
1925 pub fn evaluate(&self, axis_a: usize, axis_b: usize) -> Result<Vec<Array2<f64>>, BasisError> {
1926 (self.0)(axis_a, axis_b)
1927 }
1928}
1929
1930// ═══════════════════════════════════════════════════════════════════════════
1931// Implicit derivative operator for scalable anisotropic REML gradients
1932// ═══════════════════════════════════════════════════════════════════════════
1933
1934pub(crate) const SPATIAL_CENTER_CENTER_MAX_BYTES: usize = 512 * 1024 * 1024; // 512 MiB
1935pub(crate) const DESIGN_CROSS_CHUNK_SIZE: usize = 1024;
1936
1937/// Determine whether implicit operators should be used based on problem size
1938/// and the supplied [`ResourcePolicy`].
1939///
1940/// Returns `true` when the dense materialization of D first-derivative
1941/// matrices would exceed `policy.max_single_materialization_bytes`.
1942///
1943/// For D axes with n data points and p_smooth basis columns, the dense path
1944/// allocates D * n * p_smooth * 8 bytes for first-derivative matrices alone
1945/// (plus a similar amount for second derivatives). The implicit path stores
1946/// only the compact (n * n_knots) radial jets plus (n * n_knots * D) axis
1947/// fractions, which is O(n * k * D) instead of O(n * p * D).
1948pub fn should_use_implicit_operators_with_policy(
1949 n: usize,
1950 p: usize,
1951 d: usize,
1952 policy: &gam_runtime::resource::ResourcePolicy,
1953) -> bool {
1954 // Each first-derivative matrix is (n x p) f64 → n*p*8 bytes.
1955 // We need D of them for first derivatives, D for second diag, plus
1956 // the cross-t matrix and s_components. Conservative estimate: 3*D matrices.
1957 let dense_bytes = 3usize
1958 .saturating_mul(n)
1959 .saturating_mul(p)
1960 .saturating_mul(d)
1961 .saturating_mul(8);
1962 dense_bytes > policy.max_single_materialization_bytes
1963}
1964
1965pub(crate) fn implicit_radial_cache_bytes(n: usize, k: usize, n_axes: usize) -> usize {
1966 n.saturating_mul(k)
1967 .saturating_mul(n_axes.saturating_add(3))
1968 .saturating_mul(8)
1969}
1970
1971pub(crate) fn should_cache_implicit_radial_components(
1972 n: usize,
1973 k: usize,
1974 n_axes: usize,
1975 policy: &gam_runtime::resource::ResourcePolicy,
1976) -> bool {
1977 implicit_radial_cache_bytes(n, k, n_axes) <= policy.max_operator_cache_bytes
1978}
1979
1980pub fn assert_no_dense_derivative_materialization(n: usize, p: usize, d_pc: usize) {
1981 let first = dense_design_bytes(n, p).saturating_mul(d_pc);
1982 let second = dense_design_bytes(n, p).saturating_mul(d_pc.saturating_mul(d_pc));
1983 // Consult the library default ResourcePolicy. Production large-scale runs
1984 // configure `AnalyticOperatorRequired`, which still refuses every dense
1985 // materialization here. The default `MaterializeIfSmall` mode lets tiny
1986 // problems (and small-data/test usage) materialize as long as the combined
1987 // first- and second-order dense bytes fit under the single-materialization
1988 // byte budget. `DiagnosticsOnly` is treated like `MaterializeIfSmall` for
1989 // this guard: it permits dense materialization under the same byte cap.
1990 let policy = gam_runtime::resource::ResourcePolicy::default_library();
1991 let budget = policy.max_single_materialization_bytes;
1992 let needed = first.saturating_add(second);
1993 match policy.derivative_storage_mode {
1994 gam_runtime::resource::DerivativeStorageMode::AnalyticOperatorRequired => {
1995 // SAFETY: this assertion helper exists specifically to enforce
1996 // the large-scale invariant that spatial-PC Duchon derivative
1997 // designs never persist as dense `Array2<f64>` storage. When the
1998 // resource policy is `AnalyticOperatorRequired`, any caller that
1999 // reached this point has materialized something the strict
2000 // operator contract forbids.
2001 // SAFETY: AnalyticOperatorRequired forbids dense derivative materialization.
2002 panic!(
2003 "spatial PC Duchon derivative designs must remain operator-backed; refused persistent dense derivative materialization (n={n}, p={p}, d_pc={d_pc}, first_order={:.1} MiB, second_order={:.1} MiB)",
2004 first as f64 / (1024.0 * 1024.0),
2005 second as f64 / (1024.0 * 1024.0),
2006 );
2007 }
2008 gam_runtime::resource::DerivativeStorageMode::MaterializeIfSmall
2009 | gam_runtime::resource::DerivativeStorageMode::DiagnosticsOnly => {
2010 // SAFETY: exceeding the single-materialization budget here is a
2011 // contract violation by an upstream caller that must route through
2012 // the operator-backed path; failing loudly surfaces it rather than
2013 // silently materializing an oversized dense derivative design.
2014 assert!(
2015 needed <= budget,
2016 "spatial PC Duchon derivative designs would exceed the single-materialization budget; refused persistent dense derivative materialization (n={n}, p={p}, d_pc={d_pc}, first_order={:.1} MiB, second_order={:.1} MiB, budget={:.1} MiB)",
2017 first as f64 / (1024.0 * 1024.0),
2018 second as f64 / (1024.0 * 1024.0),
2019 budget as f64 / (1024.0 * 1024.0),
2020 );
2021 }
2022 }
2023}
2024
2025pub fn assert_spatial_centers_below_large_scale_cap(
2026 d_pc: usize,
2027 centers: ArrayView2<'_, f64>,
2028) -> Result<(), BasisError> {
2029 if centers.ncols() != d_pc {
2030 crate::bail_dim_basis!(
2031 "spatial PC center dimension mismatch: centers have {} columns, expected {d_pc}",
2032 centers.ncols()
2033 );
2034 }
2035 let k = centers.nrows();
2036 let centers_bytes = dense_design_bytes(k, d_pc);
2037 let center_center_bytes = dense_design_bytes(k, k);
2038 if centers_bytes > SPATIAL_CENTER_CENTER_MAX_BYTES {
2039 crate::bail_invalid_basis!(
2040 "spatial PC centers exceed center storage cap: K={k}, d_pc={d_pc}, centers={:.1} MiB, cap={:.1} MiB",
2041 centers_bytes as f64 / (1024.0 * 1024.0),
2042 SPATIAL_CENTER_CENTER_MAX_BYTES as f64 / (1024.0 * 1024.0),
2043 );
2044 }
2045 if center_center_bytes > SPATIAL_CENTER_CENTER_MAX_BYTES {
2046 crate::bail_invalid_basis!(
2047 "spatial PC centers exceed center-center large-scale cap: K={k}, d_pc={d_pc}, KxK={:.1} MiB, cap={:.1} MiB",
2048 center_center_bytes as f64 / (1024.0 * 1024.0),
2049 SPATIAL_CENTER_CENTER_MAX_BYTES as f64 / (1024.0 * 1024.0),
2050 );
2051 }
2052 Ok(())
2053}
2054
2055pub(crate) fn dense_design_bytes(n: usize, p: usize) -> usize {
2056 n.saturating_mul(p)
2057 .saturating_mul(std::mem::size_of::<f64>())
2058}
2059
2060pub(crate) fn should_use_lazy_spatial_design(
2061 n: usize,
2062 p: usize,
2063 policy: &gam_runtime::resource::ResourcePolicy,
2064) -> bool {
2065 matches!(
2066 policy.derivative_storage_mode,
2067 gam_runtime::resource::DerivativeStorageMode::AnalyticOperatorRequired
2068 ) || dense_design_bytes(n, p) > policy.max_single_materialization_bytes
2069}
2070
2071pub(crate) fn wrap_dense_design_with_transform(
2072 design: DesignMatrix,
2073 transform: &Array2<f64>,
2074 label: &str,
2075) -> Result<DesignMatrix, BasisError> {
2076 match design {
2077 DesignMatrix::Dense(inner) => {
2078 let op = CoefficientTransformOperator::new(inner, transform.clone()).map_err(|e| {
2079 BasisError::InvalidInput(format!("{label} coefficient transform failed: {e}"))
2080 })?;
2081 Ok(DesignMatrix::Dense(
2082 gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(op)),
2083 ))
2084 }
2085 DesignMatrix::Sparse(_) => Err(BasisError::InvalidInput(format!(
2086 "{label} coefficient transform requires a dense/operator-backed design"
2087 ))),
2088 }
2089}
2090
2091/// Single-pass `(Bᵀ(W·C), BᵀB)` accumulation over the streamed design.
2092///
2093/// Materialises each row chunk of the design **once** and reuses it for both
2094/// the constraint cross `Bᵀ(W·C)` and the Gram `BᵀB`. On the lazy chunked
2095/// spatial path each `try_row_chunk` re-evaluates all kernel columns for the
2096/// chunk, so accumulating both products in a single sweep halves the per-build
2097/// kernel re-evaluation work (the dominant cost at large scale) versus two
2098/// independent streaming passes — without changing the result beyond
2099/// floating-point reassociation. The cross is masked off (`q == 0`) by the
2100/// caller, which never invokes this when there is no constraint block.
2101pub(crate) fn design_cross_and_gram(
2102 design: &DesignMatrix,
2103 constraint_matrix: ArrayView2<'_, f64>,
2104 weights: Option<ArrayView1<'_, f64>>,
2105) -> Result<(Array2<f64>, Array2<f64>), BasisError> {
2106 let n = design.nrows();
2107 let k = design.ncols();
2108 if constraint_matrix.nrows() != n {
2109 return Err(BasisError::ConstraintMatrixRowMismatch {
2110 basisrows: n,
2111 constraintrows: constraint_matrix.nrows(),
2112 });
2113 }
2114 if let Some(w) = weights
2115 && w.len() != n
2116 {
2117 return Err(BasisError::WeightsDimensionMismatch {
2118 expected: n,
2119 found: w.len(),
2120 });
2121 }
2122 let q = constraint_matrix.ncols();
2123 let mut cross = Array2::<f64>::zeros((k, q));
2124 let mut gram = Array2::<f64>::zeros((k, k));
2125 for start in (0..n).step_by(DESIGN_CROSS_CHUNK_SIZE) {
2126 let end = (start + DESIGN_CROSS_CHUNK_SIZE).min(n);
2127 let basis_chunk = design
2128 .try_row_chunk(start..end)
2129 .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
2130 let mut constraint_chunk = constraint_matrix.slice(s![start..end, ..]).to_owned();
2131 if let Some(w) = weights {
2132 for (mut row, &weight) in constraint_chunk
2133 .axis_iter_mut(Axis(0))
2134 .zip(w.slice(s![start..end]).iter())
2135 {
2136 row *= weight;
2137 }
2138 }
2139 cross += &fast_atb(&basis_chunk, &constraint_chunk);
2140 gram += &fast_atb(&basis_chunk, &basis_chunk);
2141 }
2142 Ok((cross, gram))
2143}
2144
2145pub(crate) fn positive_spectral_whitener_from_gram(
2146 gram: &Array2<f64>,
2147) -> Result<Array2<f64>, BasisError> {
2148 // Inverse-square-root for the positive part of `gram`. Eigenvalues at or
2149 // below the relative rank tolerance `α·ε·n·max_eval` are *dropped*: the
2150 // returned whitener has shape `(n × keep)` where `keep` counts strictly
2151 // positive eigendirections of `gram`.
2152 //
2153 // Dropping (rather than ridging) is what makes the result a true
2154 // square-root inverse on the column space of `gram`. This whitener is
2155 // used by `stabilized_orthogonality_transform_from_gram` to make a
2156 // pre-existing transform `K_raw` orthonormal under the W-inner product:
2157 // when some columns of `K_raw` map to zero (or near-zero) under `B`, the
2158 // constrained Gram `K_raw^T G K_raw` is rank-deficient. Ridging those
2159 // tail directions with `1/sqrt(ε)` produced spurious basis columns
2160 // whose coefficient norms blew up to `~1/sqrt(ε)` while their image in
2161 // `B` was floating-point zero, contaminating downstream linear algebra
2162 // (in particular it forced `smooth.rs` to widen the post-transform
2163 // orthogonality residual tolerance to absorb a `cond ≈ 1/sqrt(ε)`
2164 // rounding floor). Dropping these directions is the right behavior:
2165 // they contribute nothing to `B`'s column space, and removing them
2166 // tightens the orthogonality residual back down to the genuine
2167 // floating-point limit.
2168 let (eigenvalues, eigenvectors) = gram.eigh(Side::Lower).map_err(BasisError::LinalgError)?;
2169 let n = gram.nrows();
2170 let max_eval = eigenvalues.iter().copied().fold(0.0_f64, f64::max);
2171 // Scale-invariant rank tolerance: the cutoff must track the Gram's own
2172 // spectrum (`α·ε·n·max_eval`), not an absolute floor. An earlier `max_eval
2173 // .max(1.0)` clamped the reference scale to 1.0, which is only harmless when
2174 // `max_eval ≥ 1`; for a genuinely well-conditioned but small-magnitude Gram
2175 // (e.g. a Duchon hybrid whose evaluated kernel sits far below unit scale in
2176 // moderate-to-high d) it inflated the tolerance to an absolute `α·ε·n` floor
2177 // that swallows the entire — perfectly valid — spectrum, spuriously reporting
2178 // `keep == 0`. Using the true `max_eval` makes `keep` invariant to a uniform
2179 // rescaling of the Gram (which scales every eigenvalue and the cutoff
2180 // identically). The residual `.max(f64::EPSILON)` only guards the degenerate
2181 // all-zero Gram so that numerical-zero roundoff directions are still dropped.
2182 let tol =
2183 (default_rrqr_rank_alpha() * f64::EPSILON * (n.max(1) as f64) * max_eval).max(f64::EPSILON);
2184 let keep = eigenvalues.iter().filter(|&&ev| ev > tol).count();
2185 if keep == 0 {
2186 let min_ev = eigenvalues.iter().copied().fold(f64::INFINITY, f64::min);
2187 return Err(BasisError::ConstraintNullspaceCollapsed {
2188 site: "positive_spectral_whitener_from_gram",
2189 cross_rank: 0,
2190 coeff_dim: gram.nrows(),
2191 cross_frobenius: gram.iter().map(|v| v * v).sum::<f64>().sqrt(),
2192 gram_spectrum: format!(
2193 "max eigenvalue {max_eval:.3e} (min {min_ev:.3e}, spectral tolerance {tol:.3e})"
2194 ),
2195 });
2196 }
2197 // `eigh` returns eigenvalues in ascending order, so the largest `keep`
2198 // eigenvalues live at the tail.
2199 let eig_start = eigenvalues.len() - keep;
2200 let kept_vectors = eigenvectors.slice(s![.., eig_start..]).to_owned();
2201 let mut inv_sqrt = Array2::<f64>::zeros((keep, keep));
2202 for (out_i, eig_i) in (eig_start..eigenvalues.len()).enumerate() {
2203 inv_sqrt[[out_i, out_i]] = 1.0 / eigenvalues[eig_i].sqrt();
2204 }
2205 Ok(fast_ab(&kept_vectors, &inv_sqrt))
2206}
2207
2208pub(crate) fn stabilized_orthogonality_transform_from_gram(
2209 gram: &Array2<f64>,
2210 transform: &Array2<f64>,
2211) -> Result<Array2<f64>, BasisError> {
2212 let constrained_gram = {
2213 let gt = fast_ab(gram, transform);
2214 fast_atb(transform, >)
2215 };
2216 let whitening = positive_spectral_whitener_from_gram(&constrained_gram)?;
2217 Ok(fast_ab(transform, &whitening))
2218}
2219
2220pub(crate) fn orthogonality_transform_from_cross_and_gram(
2221 constraint_cross: &Array2<f64>,
2222 gram: &Array2<f64>,
2223) -> Result<Array2<f64>, BasisError> {
2224 // Compute null(M^T) directly on M = B^T W C (k × q) via column-pivoted QR.
2225 // Working in the original k-dim coefficient space rather than first
2226 // whitening by B^T B avoids a fundamental failure mode: when B is heavily
2227 // collinear, `positive_spectral_whitener_from_gram` truncates the design
2228 // column-space to a `keep`-dim subspace, and if `keep <= q` the subsequent
2229 // nullspace search has no room — even though dim null(M^T) = k - rank(M)
2230 // ≥ k - q is always positive when k > q. The constraint nullspace is a
2231 // property of M alone; conditioning of the design only matters for the
2232 // downstream stabilization of B*K_raw.
2233 let k = constraint_cross.nrows();
2234 if k == 0 {
2235 return Err(BasisError::InsufficientColumnsForConstraint { found: 0 });
2236 }
2237 let (transform_raw, rank) = rrqr_nullspace_basis(constraint_cross, default_rrqr_rank_alpha())
2238 .map_err(BasisError::LinalgError)?;
2239 if rank >= k || transform_raw.ncols() == 0 {
2240 return Err(BasisError::ConstraintNullspaceCollapsed {
2241 site: "orthogonality_transform_from_cross_and_gram",
2242 cross_rank: rank,
2243 coeff_dim: k,
2244 cross_frobenius: constraint_cross.iter().map(|v| v * v).sum::<f64>().sqrt(),
2245 gram_spectrum: "not computed (structural cross-rank collapse: null(Mᵀ) is empty, \
2246 so no constrained design exists to eigendecompose)"
2247 .to_string(),
2248 });
2249 }
2250
2251 // Make the constrained design B*K_raw orthonormal under the W-inner product.
2252 // If the constrained Gram K_raw^T G K_raw is rank-deficient (because some
2253 // directions in null(M^T) collapse under B), the spectral whitener drops
2254 // them — that is the right behavior: a degenerate column never contributes
2255 // to B's column space and shouldn't appear in the reparameterized basis.
2256 stabilized_orthogonality_transform_from_gram(gram, &transform_raw)
2257}
2258
2259pub fn orthogonality_transform_for_design(
2260 design: &DesignMatrix,
2261 constraint_matrix: ArrayView2<'_, f64>,
2262 weights: Option<ArrayView1<'_, f64>>,
2263) -> Result<Array2<f64>, BasisError> {
2264 let k = design.ncols();
2265 if k == 0 {
2266 return Err(BasisError::InsufficientColumnsForConstraint { found: 0 });
2267 }
2268 let q = constraint_matrix.ncols();
2269 if q == 0 {
2270 return Ok(Array2::eye(k));
2271 }
2272 // Scale every constraint column to unit (weighted) L2 norm before forming the
2273 // design/constraint cross `M = Bᵀ W C`. The downstream rank detection
2274 // (`rrqr_nullspace_basis`) decides HOW MANY parametric directions the smooth
2275 // genuinely spans by testing the pivoted magnitudes of `M` against an
2276 // essentially absolute floor `α·ε·max(k,q)` (the `max(|R₀₀|, 1)` reference
2277 // clamps to 1 whenever the cross is sub-unit). With a RAW constraint column
2278 // that floor is scale-wrong: the all-ones intercept has norm √n, so `M`
2279 // carries a √n factor while the tolerance is referenced to 1. For a
2280 // kernel/radial smooth whose realized design is already (numerically)
2281 // orthogonal to the constant, `‖Bᵀ1‖` is pure floating-point roundoff
2282 // (~ε·‖B‖·√n); the √n inflation lands it right at the floor, so a rigid
2283 // rotation of the covariates — which only perturbs that roundoff — flips the
2284 // detected rank between 0 and 1. A spurious rank 1 then removes an ARBITRARY
2285 // real smooth direction (the pivot of a noise vector), and the fitted
2286 // surface, its EDF, and the REML-selected λ all drift under rotation
2287 // (gam#1818). Measuring the design's overlap with UNIT constraint directions
2288 // turns the test into a genuine rotation-invariant cosine: a real overlap is
2289 // O(1) and always detected, while roundoff-level overlap stays consistently
2290 // below the floor (rank 0). Column scaling of `C` leaves `null(Mᵀ)` — hence
2291 // the constrained-design span and the emitted transform — unchanged wherever
2292 // the rank is unchanged; it only removes the roundoff-driven rank flip.
2293 let normalized_constraint = unit_normalize_constraint_columns(constraint_matrix, weights);
2294 let (constraint_cross, gram) =
2295 design_cross_and_gram(design, normalized_constraint.view(), weights)?;
2296 orthogonality_transform_from_cross_and_gram(&constraint_cross, &gram)
2297}
2298
2299/// Scale each column of a constraint block to unit L2 norm under the inner
2300/// product used to form the identifiability cross — the `weights`-weighted
2301/// product when `Some`, the plain product otherwise. A column that is already
2302/// numerically zero (norm 0 or non-finite) is left untouched: its cross entries
2303/// are zero regardless, so it contributes no rank. The returned owned copy is
2304/// used only to build the cross; the emitted transform and the realized
2305/// constrained design are unaffected (column scaling of `C` preserves
2306/// `null(Mᵀ)`).
2307fn unit_normalize_constraint_columns(
2308 constraint_matrix: ArrayView2<'_, f64>,
2309 weights: Option<ArrayView1<'_, f64>>,
2310) -> Array2<f64> {
2311 let mut c = constraint_matrix.to_owned();
2312 let (n, q) = c.dim();
2313 for col in 0..q {
2314 let mut norm_sq = 0.0_f64;
2315 for row in 0..n {
2316 let v = c[[row, col]];
2317 let w = weights.map_or(1.0, |ws| ws[row]);
2318 norm_sq += w * v * v;
2319 }
2320 let norm = norm_sq.sqrt();
2321 if norm > 0.0 && norm.is_finite() {
2322 let inv = 1.0 / norm;
2323 for row in 0..n {
2324 c[[row, col]] *= inv;
2325 }
2326 }
2327 }
2328 c
2329}
2330
2331#[cfg(test)]
2332mod saturation_escalation_tests {
2333 use super::*;
2334
2335 #[test]
2336 fn starting_count_is_a_supported_low_rank_pilot_capped_by_default() {
2337 assert_eq!(starting_num_centers(800, 2), 30);
2338 assert_eq!(starting_num_centers(100_000, 1), 10);
2339 // The generic conditioning ceiling is `n / 4` and therefore reports
2340 // zero below four rows; the pilot retains the basis-wide one-center
2341 // degenerate minimum, which materialization subsequently raises to the
2342 // exact polynomial floor for the requested family.
2343 assert_eq!(starting_num_centers(3, 5), 1);
2344 assert_eq!(starting_num_centers(1, 2), 1);
2345 }
2346
2347 #[test]
2348 fn saturated_expansion_doubles_then_pins_at_validated_ceiling() {
2349 assert_eq!(expanded_num_centers(30, 157), Some(60));
2350 assert_eq!(expanded_num_centers(120, 157), Some(157));
2351 assert_eq!(expanded_num_centers(157, 157), None);
2352 assert_eq!(
2353 expanded_num_centers(usize::MAX - 1, usize::MAX),
2354 Some(usize::MAX)
2355 );
2356 }
2357
2358 #[test]
2359 fn saturation_excludes_the_nullspace_and_tracks_edf() {
2360 let tol = 1e-4;
2361 // Total term EDF includes the three-dimensional nullspace. Saturation
2362 // means its penalized component spends all 97 remaining directions.
2363 assert!(basis_is_saturated(100.0, 100, 3, tol));
2364 // Half-used basis is NOT saturated.
2365 assert!(!basis_is_saturated(48.5, 100, 3, tol));
2366 // Just below capacity by more than the derived margin: not saturated.
2367 assert!(!basis_is_saturated(90.0, 100, 3, tol));
2368 // A block whose null space already exhausts its columns has no penalizable
2369 // capacity and is never saturated.
2370 assert!(!basis_is_saturated(3.0, 3, 3, tol));
2371 assert!(!basis_is_saturated(f64::NAN, 100, 3, tol));
2372 }
2373
2374 #[test]
2375 fn saturation_is_monotone_in_edf() {
2376 let tol = 1e-3;
2377 let (k, null) = (60usize, 3usize);
2378 let full_width = k as f64;
2379 // Once saturated at some edf, any larger edf stays saturated.
2380 let mut first_true: Option<f64> = None;
2381 let mut e = full_width - 5.0;
2382 while e <= full_width {
2383 let sat = basis_is_saturated(e, k, null, tol);
2384 if sat && first_true.is_none() {
2385 first_true = Some(e);
2386 }
2387 if let Some(t) = first_true {
2388 assert!(
2389 basis_is_saturated(e.max(t), k, null, tol),
2390 "saturation must not flip back to false as edf grows"
2391 );
2392 }
2393 e += 0.25;
2394 }
2395 assert!(first_true.is_some(), "edf reaching capacity must saturate");
2396 }
2397}