1use coefficient_transforms::{
2 convex_derivative_control_transform_matrix, cumulative_exp, cumulative_sum_transform_matrix,
3 second_cumulative_exp,
4};
5
6pub use error::SmoothError;
7
8use input_standardization::estimate_isotropic_scale;
9
10use shape_constraints::{
11 bspline_first_derivative_control_spans, shape_lower_bounds_local, shape_order_and_sign,
12 shape_supports_basis, shape_uses_box_reparameterization,
13};
14
15pub fn describe_thin_plate_center_request(strategy: &CenterStrategy) -> String {
16 match strategy {
17 CenterStrategy::Auto(inner) => describe_thin_plate_center_request(inner),
18 CenterStrategy::DuchonSpectral { knots, basis } => format!(
19 "{} with Duchon spectral rank {}",
20 describe_thin_plate_center_request(knots),
21 basis.rank()
22 ),
23 CenterStrategy::UserProvided(centers) => format!("{} centers", centers.nrows()),
24 CenterStrategy::EqualMass { num_centers }
25 | CenterStrategy::EqualMassCovarRepresentative { num_centers }
26 | CenterStrategy::FarthestPoint { num_centers }
27 | CenterStrategy::KMeans { num_centers, .. } => format!("{num_centers} centers"),
28 CenterStrategy::UniformGrid { points_per_dim } => {
29 format!("uniform grid with {points_per_dim} points per dimension")
30 }
31 }
32}
33
34pub fn rewrite_thin_plate_knots_error(
35 err: BasisError,
36 termname: &str,
37 feature_count: usize,
38 spec: &ThinPlateBasisSpec,
39) -> BasisError {
40 match err {
41 BasisError::InvalidInput(msg)
44 if msg.contains("thin-plate spline requires at least")
45 && (msg.contains("centers to span") || msg.contains("knots to span")) =>
46 {
47 let min_centers = crate::basis::thin_plate_polynomial_basis_dimension(feature_count);
48 let requested = describe_thin_plate_center_request(&spec.center_strategy);
49 BasisError::InvalidInput(format!(
50 "joint TPS term '{termname}' over {feature_count} covariates with {requested} is invalid; minimum centers is {min_centers}"
51 ))
52 }
53 BasisError::InvalidInput(msg)
58 if msg.starts_with("requested ") && msg.contains(" knots but only ") =>
59 {
60 let min_centers = crate::basis::thin_plate_polynomial_basis_dimension(feature_count);
61 let requested = describe_thin_plate_center_request(&spec.center_strategy);
62 BasisError::InvalidInput(format!(
63 "joint TPS term '{termname}' over {feature_count} covariates with {requested} is invalid; minimum centers is {min_centers}"
64 ))
65 }
66 other => other,
67 }
68}
69
70#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
71pub enum ShapeConstraint {
72 None,
73 MonotoneIncreasing,
74 MonotoneDecreasing,
75 Convex,
76 Concave,
77}
78
79pub fn parse_shape_constraint(raw: &str) -> Result<ShapeConstraint, String> {
90 let normalized = raw.trim().to_ascii_lowercase().replace('-', "_");
91 match normalized.as_str() {
92 "" | "none" => Ok(ShapeConstraint::None),
93 "monotone_increasing" | "monotonic_increasing" | "increasing" | "mono_inc" | "mpi" => {
94 Ok(ShapeConstraint::MonotoneIncreasing)
95 }
96 "monotone_decreasing" | "monotonic_decreasing" | "decreasing" | "mono_dec" | "mpd" => {
97 Ok(ShapeConstraint::MonotoneDecreasing)
98 }
99 "convex" | "cvx" => Ok(ShapeConstraint::Convex),
100 "concave" | "ccv" => Ok(ShapeConstraint::Concave),
101 other => Err(format!(
102 "unknown shape constraint {other:?}; expected one of \
103 \"none\", \"monotone_increasing\", \"monotone_decreasing\", \
104 \"convex\", \"concave\""
105 )),
106 }
107}
108
109impl ShapeConstraint {
110 pub fn dsl_str(&self) -> &'static str {
113 match self {
114 ShapeConstraint::None => "none",
115 ShapeConstraint::MonotoneIncreasing => "monotone_increasing",
116 ShapeConstraint::MonotoneDecreasing => "monotone_decreasing",
117 ShapeConstraint::Convex => "convex",
118 ShapeConstraint::Concave => "concave",
119 }
120 }
121}
122
123pub const SMOOTH_HEAD_KEYWORDS: [&str; 11] = [
126 "s",
127 "smooth",
128 "te",
129 "tensor",
130 "thinplate",
131 "tps",
132 "duchon",
133 "matern",
134 "sphere",
135 "bs",
136 "bspline",
137];
138
139pub fn apply_shape_constraints_to_formula(
152 formula: &str,
153 constraints: &[(String, String)],
154) -> Result<String, String> {
155 use std::collections::{BTreeMap, BTreeSet};
156
157 if constraints.is_empty() {
158 return Ok(formula.to_string());
159 }
160 let strip_ws = |s: &str| -> String { s.chars().filter(|c| !c.is_whitespace()).collect() };
161
162 let mut wanted: BTreeMap<String, &'static str> = BTreeMap::new();
164 let mut originals: BTreeMap<String, String> = BTreeMap::new();
166 for (key, kind_raw) in constraints {
167 let kind = parse_shape_constraint(kind_raw)?;
168 let nk = strip_ws(key);
169 originals.entry(nk.clone()).or_insert_with(|| key.clone());
170 if kind != ShapeConstraint::None {
171 wanted.insert(nk, kind.dsl_str());
172 }
173 }
174 if wanted.is_empty() {
175 return Ok(formula.to_string());
176 }
177
178 let chars: Vec<char> = formula.chars().collect();
179 let n = chars.len();
180 let is_ident = |c: char| c.is_ascii_alphanumeric() || c == '_';
181
182 let mut out = String::with_capacity(formula.len() + 32);
183 let mut matched: BTreeSet<String> = BTreeSet::new();
184 let mut i = 0usize;
185 while i < n {
186 let mut head: Option<(usize, usize)> = None; let mut p = i;
190 while p < n {
191 let boundary = p == 0 || !is_ident(chars[p - 1]);
192 if boundary {
193 for kw in SMOOTH_HEAD_KEYWORDS.iter() {
194 let klen = kw.chars().count();
195 if p + klen > n || chars[p..p + klen].iter().collect::<String>() != **kw {
196 continue;
197 }
198 let mut q = p + klen;
199 while q < n && chars[q].is_whitespace() {
200 q += 1;
201 }
202 if q < n && chars[q] == '(' {
203 head = Some((p, q));
204 break;
205 }
206 }
207 }
208 if head.is_some() {
209 break;
210 }
211 p += 1;
212 }
213 let (head_start, paren_open) = match head {
214 Some(h) => h,
215 None => {
216 out.extend(chars[i..].iter());
217 break;
218 }
219 };
220 out.extend(chars[i..head_start].iter());
221
222 let body_start = paren_open + 1;
224 let mut depth = 1i32;
225 let mut j = body_start;
226 let mut in_str: Option<char> = None;
227 let mut closed = false;
228 while j < n {
229 let ch = chars[j];
230 if let Some(quote) = in_str {
231 if ch == quote {
232 in_str = None;
233 }
234 } else if ch == '\'' || ch == '"' {
235 in_str = Some(ch);
236 } else if ch == '(' {
237 depth += 1;
238 } else if ch == ')' {
239 depth -= 1;
240 if depth == 0 {
241 closed = true;
242 break;
243 }
244 }
245 j += 1;
246 }
247
248 if !closed {
249 out.extend(chars[head_start..].iter());
252 break;
253 }
254
255 let term_text: String = chars[head_start..=j].iter().collect();
256
257 let key_norm = strip_ws(&term_text);
258
259 match wanted.get(&key_norm) {
260 None => out.extend(chars[head_start..=j].iter()),
261 Some(kind) => {
262 let head_paren: String = chars[head_start..body_start].iter().collect();
263 let inside: String = chars[body_start..j].iter().collect();
264 let inside = inside.trim();
265 if inside.is_empty() {
266 out.push_str(&format!("{head_paren}shape={kind})"));
267 } else {
268 out.push_str(&format!("{head_paren}{inside}, shape={kind})"));
269 }
270 matched.insert(key_norm);
271 }
272 }
273
274 i = j + 1;
275 }
276
277 let mut missing: Vec<String> = wanted
278 .keys()
279 .filter(|k| !matched.contains(*k))
280 .map(|k| originals.get(k).cloned().unwrap_or_else(|| k.clone()))
281 .collect();
282
283 if !missing.is_empty() {
284 missing.sort();
285 return Err(format!(
286 "shape constraints referenced smooth term(s) not found in formula: {}",
287 missing.join(", ")
288 ));
289 }
290
291 Ok(out)
292}
293
294#[derive(Debug, Clone, Serialize, Deserialize)]
295pub enum BySmoothKind {
296 Numeric,
297 Level { level_bits: u64 },
298}
299
300#[derive(Debug, Clone, Serialize, Deserialize)]
301#[serde(deny_unknown_fields)]
302pub enum SmoothBasisSpec {
303 ByVariable {
313 inner: Box<SmoothBasisSpec>,
314 by_col: usize,
315 kind: BySmoothKind,
316 by: ByVariableSpec,
317 },
318 FactorSumToZero {
322 inner: Box<SmoothBasisSpec>,
323 by_col: usize,
324 levels: Vec<u64>,
325 #[serde(default)]
336 frozen_global_orthogonality: Option<Array2<f64>>,
337 },
338 BSpline1D {
339 feature_col: usize,
340 spec: BSplineBasisSpec,
341 },
342 BySmooth {
345 smooth: Box<SmoothBasisSpec>,
346 by_kind: ByVarKind,
347 },
348 FactorSmooth { spec: FactorSmoothSpec },
351 ThinPlate {
352 feature_cols: Vec<usize>,
353 spec: ThinPlateBasisSpec,
354 input_scale: Option<crate::IsotropicScale>,
357 },
358 Sphere {
359 feature_cols: Vec<usize>,
360 spec: SphericalSplineBasisSpec,
361 },
362 ConstantCurvature {
368 feature_cols: Vec<usize>,
369 spec: ConstantCurvatureBasisSpec,
370 },
371 Matern {
372 feature_cols: Vec<usize>,
373 spec: MaternBasisSpec,
374 input_scale: Option<crate::IsotropicScale>,
375 },
376 MeasureJet {
382 feature_cols: Vec<usize>,
383 spec: MeasureJetBasisSpec,
384 input_scale: Option<crate::IsotropicScale>,
385 },
386 Duchon {
387 feature_cols: Vec<usize>,
388 spec: DuchonBasisSpec,
389 input_scale: Option<crate::IsotropicScale>,
390 },
391 Pca {
392 feature_cols: Vec<usize>,
393 basis_matrix: Array2<f64>,
394 centered: bool,
395 #[serde(default = "default_pca_smooth_penalty")]
396 smooth_penalty: f64,
397 #[serde(default)]
398 center_mean: Option<Array1<f64>>,
399 #[serde(default)]
400 pca_basis_path: Option<PathBuf>,
401 #[serde(default = "default_pca_chunk_size")]
402 chunk_size: usize,
403 },
404 TensorBSpline {
409 feature_cols: Vec<usize>,
410 spec: TensorBSplineSpec,
411 },
412}
413
414impl SmoothBasisSpec {
415 pub fn min_sample_rows(&self) -> usize {
432 const RADIAL_FLOOR: usize = 5;
437
438 match self {
439 Self::ByVariable { inner, .. } => inner.min_sample_rows(),
440 Self::FactorSumToZero { inner, levels, .. } => {
441 let inner_min = inner.min_sample_rows();
445 let lvls = levels.len().saturating_sub(1).max(1);
446 inner_min.saturating_mul(lvls)
447 }
448 Self::BSpline1D { spec, .. } => bspline_basis_min_rows(spec),
449 Self::BySmooth { smooth, .. } => smooth.min_sample_rows(),
450 Self::FactorSmooth { spec } => {
451 bspline_basis_min_rows(&spec.marginal)
455 }
456 Self::ThinPlate { .. }
457 | Self::Sphere { .. }
458 | Self::ConstantCurvature { .. }
459 | Self::Matern { .. }
460 | Self::MeasureJet { .. }
461 | Self::Duchon { .. } => RADIAL_FLOOR,
462 Self::Pca { basis_matrix, .. } => basis_matrix.ncols().max(1),
463 Self::TensorBSpline { spec, .. } => {
464 let mut total: usize = 0;
510 for marginal in &spec.marginalspecs {
511 let m = bspline_basis_min_rows(marginal);
512 total = total.saturating_add(m.max(1));
513 }
514 total.max(RADIAL_FLOOR)
515 }
516 }
517 }
518
519 pub fn structural_kind(&self) -> &'static str {
530 match self {
531 Self::ByVariable { .. } => "by_variable",
532 Self::FactorSumToZero { .. } => "factor_sum_to_zero",
533 Self::BSpline1D { .. } => "bspline_1d",
534 Self::BySmooth { .. } => "by_smooth",
535 Self::FactorSmooth { .. } => "factor_smooth",
536 Self::ThinPlate { .. } => "thin_plate",
537 Self::Sphere { .. } => "sphere",
538 Self::ConstantCurvature { .. } => "constant_curvature",
539 Self::Matern { .. } => "matern",
540 Self::MeasureJet { .. } => "measurejet",
541 Self::Duchon { .. } => "duchon",
542 Self::Pca { .. } => "pca",
543 Self::TensorBSpline { .. } => "tensor_bspline",
544 }
545 }
546
547 pub fn is_marginally_centered_tensor(&self) -> bool {
556 matches!(
557 self,
558 Self::TensorBSpline { spec, .. }
559 if matches!(spec.identifiability, TensorBSplineIdentifiability::MarginalSumToZero)
560 )
561 }
562
563 pub fn is_sum_to_zero_factor_smooth(&self) -> bool {
580 matches!(
581 self,
582 Self::FactorSumToZero { .. }
583 | Self::FactorSmooth {
584 spec: FactorSmoothSpec {
585 flavour: FactorSmoothFlavour::Sz,
586 ..
587 }
588 }
589 )
590 }
591
592 pub fn structural_feature_cols(&self) -> Vec<usize> {
596 match self {
597 Self::ByVariable { inner, .. } | Self::FactorSumToZero { inner, .. } => {
598 inner.structural_feature_cols()
599 }
600 Self::BySmooth { smooth, .. } => smooth.structural_feature_cols(),
601 Self::FactorSmooth { .. } => Vec::new(),
602 Self::BSpline1D { feature_col, .. } => vec![*feature_col],
603 Self::ThinPlate { feature_cols, .. }
604 | Self::Sphere { feature_cols, .. }
605 | Self::ConstantCurvature { feature_cols, .. }
606 | Self::Matern { feature_cols, .. }
607 | Self::MeasureJet { feature_cols, .. }
608 | Self::Duchon { feature_cols, .. }
609 | Self::Pca { feature_cols, .. }
610 | Self::TensorBSpline { feature_cols, .. } => feature_cols.clone(),
611 }
612 }
613}
614
615pub fn bspline_basis_min_rows(spec: &crate::basis::BSplineBasisSpec) -> usize {
640 use crate::basis::BSplineKnotSpec;
641 let columns = match &spec.knotspec {
642 BSplineKnotSpec::Generate {
643 num_internal_knots, ..
644 } => *num_internal_knots + spec.degree + 1,
645 BSplineKnotSpec::Automatic {
646 num_internal_knots: Some(k),
647 ..
648 } => *k + spec.degree + 1,
649 BSplineKnotSpec::Automatic {
650 num_internal_knots: None,
651 ..
652 } => {
653 spec.degree + 2
657 }
658 BSplineKnotSpec::Provided(knots) => knots.len().saturating_sub(spec.degree + 1).max(1),
659 BSplineKnotSpec::NaturalCubicRegression { knots } => knots.len(),
661 BSplineKnotSpec::PeriodicUniform { num_basis, .. } => *num_basis,
662 };
663 let columns = columns.max(spec.degree + 2);
664
665 if spec.double_penalty {
666 const DOUBLE_PENALTY_FLOOR: usize = 2;
669 DOUBLE_PENALTY_FLOOR.min(columns).max(1)
670 } else {
671 columns
672 }
673}
674
675#[derive(Debug, Clone, Serialize, Deserialize)]
676pub enum ByVariableSpec {
677 Numeric,
678 Level { value_bits: u64, label: String },
679}
680
681#[derive(Debug, Clone, Serialize, Deserialize)]
682pub enum ByVarKind {
683 Numeric {
684 feature_col: usize,
685 },
686 Factor {
687 feature_col: usize,
688 ordered: bool,
689 frozen_levels: Option<Vec<u64>>,
690 },
691}
692
693#[derive(Debug, Clone, Serialize, Deserialize)]
694pub struct FactorSmoothSpec {
695 pub continuous_cols: Vec<usize>,
696 pub group_col: usize,
697 pub marginal: BSplineBasisSpec,
698 pub flavour: FactorSmoothFlavour,
699 pub group_frozen_levels: Option<Vec<u64>>,
700 #[serde(default)]
706 pub frozen_global_orthogonality: Option<Array2<f64>>,
707}
708
709#[derive(Debug, Clone, Serialize, Deserialize)]
710pub enum FactorSmoothFlavour {
711 Fs { m_null_penalty_orders: Vec<usize> },
712 Sz,
713 Re,
714}
715
716#[derive(Debug, Clone, Serialize, Deserialize)]
717pub struct TensorBSplineSpec {
718 pub marginalspecs: Vec<BSplineBasisSpec>,
719 #[serde(default)]
720 pub periods: Vec<Option<f64>>,
721 #[serde(default = "default_tensor_double_penalty")]
722 pub double_penalty: bool,
723 #[serde(default)]
724 pub identifiability: TensorBSplineIdentifiability,
725 #[serde(default)]
726 pub penalty_decomposition: TensorBSplinePenaltyDecomposition,
727}
728
729pub const fn default_tensor_double_penalty() -> bool {
730 true
731}
732
733impl Default for TensorBSplineSpec {
734 fn default() -> Self {
735 Self {
736 marginalspecs: Vec::new(),
737 periods: Vec::new(),
738 double_penalty: default_tensor_double_penalty(),
739 identifiability: TensorBSplineIdentifiability::default(),
740 penalty_decomposition: TensorBSplinePenaltyDecomposition::default(),
741 }
742 }
743}
744
745#[derive(Debug, Default, Clone, Serialize, Deserialize)]
746pub enum TensorBSplineIdentifiability {
747 None,
748 #[default]
749 SumToZero,
750 MarginalSumToZero,
760 FrozenTransform {
761 transform: Array2<f64>,
762 },
763}
764
765#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
766pub enum TensorBSplinePenaltyDecomposition {
767 #[default]
770 MarginalKroneckerSum,
771 Separable,
775}
776
777#[derive(Debug, Clone, Serialize, Deserialize)]
778pub struct SmoothTermSpec {
779 pub name: String,
780 pub basis: SmoothBasisSpec,
781 pub shape: ShapeConstraint,
782 #[serde(default)]
791 pub joint_null_rotation: Option<crate::basis::JointNullRotation>,
792 #[serde(default)]
796 pub frozen_parametric_residualization: Option<ParametricResidualizationChart>,
797}
798
799#[derive(Debug, Clone, Serialize, Deserialize)]
817pub struct ParametricResidualizationChart {
818 pub owner_terms: Vec<usize>,
822 pub has_parametric_block: bool,
826 pub correction: Array2<f64>,
828}
829
830#[derive(Debug, Clone, Copy, PartialEq, Eq)]
839pub enum SmoothCollectionGaugeArm {
840 Delete,
843 Residualize,
845}
846
847#[derive(Debug, Clone)]
872pub struct SmoothCollectionGauge {
873 pub arm: SmoothCollectionGaugeArm,
875 pub constraint_block: Array2<f64>,
877 pub owner_terms: Vec<usize>,
881 pub has_parametric_block: bool,
883 pub local_identifiability_transform: Option<Array2<f64>>,
911 pub local_columns: usize,
921}
922
923#[derive(Debug, Clone)]
924pub struct SmoothTerm {
925 pub name: String,
926 pub coeff_range: Range<usize>,
927 pub shape: ShapeConstraint,
928 pub active_penalties: Vec<ActivePenalty>,
931 pub dropped_penalties: Vec<DroppedPenaltyInfo>,
932 pub metadata: BasisMetadata,
933 pub lower_bounds_local: Option<Array1<f64>>,
936 pub linear_constraints_local: Option<LinearInequalityConstraints>,
939 pub kronecker_factored: Option<KroneckerFactoredBasis>,
942 pub joint_null_rotation: Option<crate::basis::JointNullRotation>,
965 pub unabsorbed_global_orthogonality: Option<Array2<f64>>,
975 pub parametric_residualization: Option<ParametricResidualizationChart>,
982 pub collection_gauge: Option<SmoothCollectionGauge>,
993}
994
995impl SmoothTerm {
996 pub fn apply_rotation_to_predict(
1012 &self,
1013 x_new_raw: Array2<f64>,
1014 ) -> Result<Array2<f64>, BasisError> {
1015 let Some(rot) = self.joint_null_rotation.as_ref() else {
1016 return Ok(x_new_raw);
1017 };
1018 let p_local = rot.rotation.nrows();
1019 if x_new_raw.ncols() != p_local {
1020 crate::bail_dim_basis!(
1021 "joint-null rotation replay for term '{}': raw design has {} columns, \
1022 rotation expects {} (the raw basis builder must emit the same column \
1023 count as at fit time)",
1024 self.name,
1025 x_new_raw.ncols(),
1026 p_local,
1027 );
1028 }
1029 Ok(gam_linalg::faer_ndarray::fast_ab(&x_new_raw, &rot.rotation))
1030 }
1031
1032 pub fn wald_unpenalized_dim(&self) -> usize {
1055 joint_unpenalized_dim(self.coeff_range.len(), &self.active_penalties)
1056 }
1057}
1058
1059pub fn joint_unpenalized_dim(p_local: usize, active_penalties: &[ActivePenalty]) -> usize {
1064 use gam_linalg::faer_ndarray::FaerEigh;
1065 if p_local == 0 {
1066 return 0;
1067 }
1068 if active_penalties.is_empty() {
1069 return p_local;
1071 }
1072 let mut s_total = Array2::<f64>::zeros((p_local, p_local));
1077 let mut materialized = 0usize;
1078 for penalty in active_penalties {
1079 let s = &penalty.matrix;
1080 if s.nrows() == p_local && s.ncols() == p_local {
1081 s_total += s;
1082 materialized += 1;
1083 }
1084 }
1085 if materialized == active_penalties.len() {
1086 let symmetric = {
1087 let transpose = s_total.t().to_owned();
1088 (&s_total + &transpose) * 0.5
1089 };
1090 if let Ok((evals, _)) = symmetric.eigh(faer::Side::Lower) {
1091 let max_abs = evals.iter().fold(0.0_f64, |acc, &v| acc.max(v.abs()));
1092 if max_abs == 0.0 {
1093 return p_local;
1095 }
1096 let tol = max_abs * (p_local as f64) * 1e-12;
1097 let rank = evals.iter().filter(|&&v| v > tol).count();
1098 return p_local.saturating_sub(rank);
1099 }
1100 }
1101 if active_penalties.len() >= 2 {
1106 0
1107 } else {
1108 active_penalties
1109 .iter()
1110 .map(|penalty| penalty.nullity)
1111 .min()
1112 .unwrap_or(0)
1113 .min(p_local)
1114 }
1115}
1116
1117#[derive(Debug, Clone, Serialize, Deserialize)]
1118pub struct PenaltyBlockInfo {
1119 pub global_index: usize,
1120 pub termname: Option<String>,
1121 pub penalty: ActivePenaltyInfo,
1122}
1123
1124#[derive(Debug, Clone, Serialize, Deserialize)]
1125pub struct DroppedPenaltyBlockInfo {
1126 pub termname: Option<String>,
1127 pub penalty: DroppedPenaltyInfo,
1128}
1129
1130#[derive(Debug, Clone)]
1131pub struct SmoothDesign {
1132 pub term_designs: Vec<DesignMatrix>,
1133 pub penalties: Vec<BlockwisePenalty>,
1136 pub nullspace_dims: Vec<usize>,
1137 pub penaltyinfo: Vec<PenaltyBlockInfo>,
1138 pub dropped_penaltyinfo: Vec<DroppedPenaltyBlockInfo>,
1139 pub terms: Vec<SmoothTerm>,
1140 pub coefficient_lower_bounds: Option<Array1<f64>>,
1143 pub linear_constraints: Option<LinearInequalityConstraints>,
1146}
1147
1148impl SmoothDesign {
1149 pub fn total_smooth_cols(&self) -> usize {
1150 self.term_designs.iter().map(DesignMatrix::ncols).sum()
1151 }
1152 pub fn nrows(&self) -> usize {
1153 self.term_designs.first().map_or(0, DesignMatrix::nrows)
1154 }
1155}
1156
1157#[derive(Debug, Clone)]
1158pub struct RawSmoothDesign {
1159 pub term_designs: Vec<DesignMatrix>,
1160 pub affine_offset: Array1<f64>,
1162 pub penalties: Vec<BlockwisePenalty>,
1165 pub nullspace_dims: Vec<usize>,
1166 pub penaltyinfo: Vec<PenaltyBlockInfo>,
1167 pub dropped_penaltyinfo: Vec<DroppedPenaltyBlockInfo>,
1168 pub terms: Vec<SmoothTerm>,
1169 pub coefficient_lower_bounds: Option<Array1<f64>>,
1170 pub linear_constraints: Option<LinearInequalityConstraints>,
1171}
1172
1173impl RawSmoothDesign {
1174 pub fn total_smooth_cols(&self) -> usize {
1175 self.term_designs.iter().map(DesignMatrix::ncols).sum()
1176 }
1177 pub fn nrows(&self) -> usize {
1178 self.term_designs.first().map_or(0, DesignMatrix::nrows)
1179 }
1180}
1181
1182#[derive(Debug, Default, Clone, Serialize, Deserialize)]
1183pub enum BoundedCoefficientPriorSpec {
1184 #[default]
1185 None,
1186 Uniform,
1187 Beta {
1188 a: f64,
1189 b: f64,
1190 },
1191}
1192
1193#[derive(Debug, Clone, Serialize, Deserialize, Default)]
1194pub enum LinearCoefficientGeometry {
1195 #[default]
1196 Unconstrained,
1197 Bounded {
1198 min: f64,
1199 max: f64,
1200 #[serde(default)]
1201 prior: BoundedCoefficientPriorSpec,
1202 },
1203}
1204
1205#[derive(Debug, Clone, Serialize, Deserialize)]
1206pub struct LinearTermSpec {
1207 pub name: String,
1208 pub feature_col: usize,
1214 #[serde(default)]
1217 pub feature_cols: Vec<usize>,
1218 #[serde(default)]
1232 pub categorical_levels: Vec<(usize, u64)>,
1233 #[serde(default = "default_linear_term_double_penalty")]
1237 pub double_penalty: bool,
1238 #[serde(default)]
1239 pub coefficient_geometry: LinearCoefficientGeometry,
1240 #[serde(default)]
1241 pub coefficient_min: Option<f64>,
1242 #[serde(default)]
1243 pub coefficient_max: Option<f64>,
1244 #[serde(default)]
1259 pub frozen_function_mass: Option<f64>,
1260}
1261
1262impl LinearTermSpec {
1263 pub fn effective_feature_cols(&self) -> Vec<usize> {
1266 if self.feature_cols.is_empty() {
1267 vec![self.feature_col]
1268 } else {
1269 self.feature_cols.clone()
1270 }
1271 }
1272
1273 pub fn is_interaction(&self) -> bool {
1275 self.feature_cols.len() > 1 || !self.categorical_levels.is_empty()
1276 }
1277
1278 pub fn realized_design_column(&self, data: ArrayView2<'_, f64>) -> Result<Array1<f64>, String> {
1291 let n = data.nrows();
1292 let p = data.ncols();
1293 let bounds = |col: usize| -> Result<(), String> {
1294 if col >= p {
1295 Err(format!(
1296 "linear term '{}' feature column {} out of bounds for {} columns",
1297 self.name, col, p
1298 ))
1299 } else {
1300 Ok(())
1301 }
1302 };
1303
1304 let mut column = if self.categorical_levels.is_empty() {
1309 let cols = self.effective_feature_cols();
1310 for &c in &cols {
1311 bounds(c)?;
1312 }
1313 let mut acc = data.column(cols[0]).to_owned();
1314 for &c in cols.iter().skip(1) {
1315 acc *= &data.column(c);
1316 }
1317 acc
1318 } else {
1319 let mut acc = Array1::<f64>::ones(n);
1320 for &c in &self.feature_cols {
1321 bounds(c)?;
1322 acc *= &data.column(c);
1323 }
1324 acc
1325 };
1326
1327 for &(col, level_bits) in &self.categorical_levels {
1328 bounds(col)?;
1329 let level_bits = gam_data::canonical_level_bits(f64::from_bits(level_bits));
1333 let gate = data.column(col);
1334 for (out, &v) in column.iter_mut().zip(gate.iter()) {
1335 if gam_data::canonical_level_bits(v) != level_bits {
1336 *out = 0.0;
1337 }
1338 }
1339 }
1340
1341 Ok(column)
1342 }
1343}
1344
1345pub const fn default_linear_term_double_penalty() -> bool {
1346 false
1347}
1348
1349pub const fn default_pca_smooth_penalty() -> f64 {
1350 1.0
1351}
1352
1353pub const fn default_pca_chunk_size() -> usize {
1354 4096
1355}
1356
1357#[derive(Debug, Clone, Serialize, Deserialize)]
1363pub struct RandomEffectTermSpec {
1364 pub name: String,
1365 pub feature_col: usize,
1366 pub drop_first_level: bool,
1369 #[serde(default = "default_random_effect_penalized")]
1373 pub penalized: bool,
1374 #[serde(default)]
1377 pub frozen_levels: Option<Vec<u64>>,
1378 #[serde(default = "default_random_effect_lenient_unseen")]
1395 pub lenient_unseen: bool,
1396}
1397
1398pub fn default_random_effect_penalized() -> bool {
1399 true
1400}
1401
1402pub fn default_random_effect_lenient_unseen() -> bool {
1403 true
1404}
1405
1406pub fn validate_measure_jet_positive_vec_len(
1407 label: &str,
1408 term_name: &str,
1409 field: &str,
1410 values: &[f64],
1411 expected: usize,
1412) -> Result<(), String> {
1413 if values.len() != expected {
1414 return Err(SmoothError::invalid_config(format!(
1415 "{label} term '{term_name}' frozen MeasureJet {field} has length {}, expected {expected}",
1416 values.len()
1417 ))
1418 .into());
1419 }
1420 if values
1421 .iter()
1422 .any(|value| !(value.is_finite() && *value > 0.0))
1423 {
1424 return Err(SmoothError::invalid_config(format!(
1425 "{label} term '{term_name}' frozen MeasureJet {field} values must be positive and finite"
1426 ))
1427 .into());
1428 }
1429 Ok(())
1430}
1431
1432#[derive(Debug, Clone, Serialize, Deserialize)]
1433pub struct TermCollectionSpec {
1434 pub linear_terms: Vec<LinearTermSpec>,
1435 pub random_effect_terms: Vec<RandomEffectTermSpec>,
1436 pub smooth_terms: Vec<SmoothTermSpec>,
1437}
1438
1439pub fn validate_smooth_basis_frozen(
1440 basis: &SmoothBasisSpec,
1441 label: &str,
1442 term_name: &str,
1443) -> Result<(), String> {
1444 if let Err(error) = basis.validate_scale_configuration() {
1445 return Err(SmoothError::invalid_config(format!(
1446 "{label} term '{term_name}' has an invalid scale contract: {error}"
1447 ))
1448 .into());
1449 }
1450 match basis {
1451 SmoothBasisSpec::ByVariable { inner, .. }
1452 | SmoothBasisSpec::FactorSumToZero { inner, .. } => {
1453 validate_smooth_basis_frozen(inner, label, term_name)
1454 }
1455 SmoothBasisSpec::BSpline1D { spec, .. } => {
1456 if !matches!(
1457 spec.knotspec,
1458 BSplineKnotSpec::Provided(_)
1459 | BSplineKnotSpec::PeriodicUniform { .. }
1460 | BSplineKnotSpec::NaturalCubicRegression { .. }
1461 ) {
1462 return Err(format!(
1463 "{label} term '{term_name}' is not frozen: BSpline knotspec must be Provided, PeriodicUniform, or NaturalCubicRegression"
1464 ));
1465 }
1466 Ok(())
1467 }
1468 SmoothBasisSpec::ThinPlate { spec, .. } => {
1469 if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
1470 return Err(format!(
1471 "{label} term '{term_name}' is not frozen: ThinPlate centers must be UserProvided"
1472 ));
1473 }
1474 if matches!(
1475 spec.identifiability,
1476 SpatialIdentifiability::OrthogonalToParametric
1477 ) {
1478 return Err(format!(
1479 "{label} term '{term_name}' is not frozen: ThinPlate identifiability must be FrozenTransform or None"
1480 ));
1481 }
1482 Ok(())
1483 }
1484 _ => Ok(()),
1485 }
1486}
1487
1488impl TermCollectionSpec {
1489 pub fn write_structural_shape_hash(&self, h: &mut gam_runtime::warm_start::Fingerprinter) {
1503 h.write_str("term-collection");
1504 h.write_usize(self.linear_terms.len());
1505 for linear in &self.linear_terms {
1506 h.write_str(&linear.name);
1507 }
1508 h.write_usize(self.random_effect_terms.len());
1509 h.write_usize(self.smooth_terms.len());
1510 for smooth in &self.smooth_terms {
1511 h.write_str(&smooth.name);
1512 h.write_str(smooth.basis.structural_kind());
1513 for col in smooth.basis.structural_feature_cols() {
1514 h.write_usize(col);
1515 }
1516 }
1517 }
1518
1519 pub fn validate_frozen(&self, label: &str) -> Result<(), String> {
1523 for linear in &self.linear_terms {
1524 if let (Some(min), Some(max)) = (linear.coefficient_min, linear.coefficient_max)
1525 && (!min.is_finite() || !max.is_finite() || min > max)
1526 {
1527 return Err(SmoothError::invalid_config(format!(
1528 "{label} linear term '{}' has invalid coefficient constraint [{min}, {max}]",
1529 linear.name
1530 ))
1531 .into());
1532 }
1533 if let Some(min) = linear.coefficient_min
1534 && !min.is_finite()
1535 {
1536 return Err(SmoothError::invalid_config(format!(
1537 "{label} linear term '{}' has non-finite coefficient minimum {min}",
1538 linear.name
1539 ))
1540 .into());
1541 }
1542 if let Some(max) = linear.coefficient_max
1543 && !max.is_finite()
1544 {
1545 return Err(SmoothError::invalid_config(format!(
1546 "{label} linear term '{}' has non-finite coefficient maximum {max}",
1547 linear.name
1548 ))
1549 .into());
1550 }
1551 if let LinearCoefficientGeometry::Bounded { min, max, prior } =
1552 &linear.coefficient_geometry
1553 {
1554 if !min.is_finite() || !max.is_finite() || min >= max {
1555 return Err(SmoothError::invalid_config(format!(
1556 "{label} bounded term '{}' has invalid bounds [{min}, {max}]",
1557 linear.name
1558 ))
1559 .into());
1560 }
1561 match prior {
1562 BoundedCoefficientPriorSpec::None | BoundedCoefficientPriorSpec::Uniform => {}
1563 BoundedCoefficientPriorSpec::Beta { a, b } => {
1564 if !a.is_finite() || !b.is_finite() || *a < 1.0 || *b < 1.0 {
1565 return Err(SmoothError::invalid_config(format!(
1566 "{label} bounded term '{}' has invalid Beta prior ({a}, {b})",
1567 linear.name
1568 ))
1569 .into());
1570 }
1571 }
1572 }
1573 }
1574 }
1575 for st in &self.smooth_terms {
1576 if let Err(error) = st.basis.validate_scale_configuration() {
1577 return Err(SmoothError::invalid_config(format!(
1578 "{label} term '{}' has an invalid scale contract: {error}",
1579 st.name
1580 ))
1581 .into());
1582 }
1583 match &st.basis {
1584 SmoothBasisSpec::ByVariable { inner, .. } => {
1585 validate_smooth_basis_frozen(inner, label, &st.name)?;
1586 let nested = SmoothTermSpec {
1587 frozen_parametric_residualization: None,
1588 name: st.name.clone(),
1589 basis: (**inner).clone(),
1590 shape: st.shape,
1591 joint_null_rotation: None,
1592 };
1593 TermCollectionSpec {
1594 linear_terms: Vec::new(),
1595 random_effect_terms: Vec::new(),
1596 smooth_terms: vec![nested],
1597 }
1598 .validate_frozen(label)?;
1599 }
1600 SmoothBasisSpec::FactorSumToZero { inner, levels, .. } => {
1601 if levels.len() < 2 {
1602 return Err(format!(
1603 "{label} term '{}' has invalid frozen sz levels",
1604 st.name
1605 ));
1606 }
1607 validate_smooth_basis_frozen(inner, label, &st.name)?;
1608 }
1609 SmoothBasisSpec::BSpline1D { spec, .. } => {
1610 if !matches!(
1611 spec.knotspec,
1612 BSplineKnotSpec::Provided(_)
1613 | BSplineKnotSpec::PeriodicUniform { .. }
1614 | BSplineKnotSpec::NaturalCubicRegression { .. }
1615 ) {
1616 return Err(SmoothError::invalid_config(format!(
1617 "{label} term '{}' is not frozen: BSpline knotspec must be Provided, PeriodicUniform, or NaturalCubicRegression",
1618 st.name
1619 ))
1620 .into());
1621 }
1622 }
1623 SmoothBasisSpec::ThinPlate {
1624 spec, input_scale, ..
1625 } => {
1626 if input_scale.is_none() {
1627 return Err(SmoothError::invalid_config(format!(
1628 "{label} term '{}' is not frozen: ThinPlate input_scale is missing",
1629 st.name
1630 ))
1631 .into());
1632 }
1633 if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
1634 return Err(SmoothError::invalid_config(format!(
1635 "{label} term '{}' is not frozen: ThinPlate centers must be UserProvided",
1636 st.name
1637 ))
1638 .into());
1639 }
1640 if matches!(
1641 spec.identifiability,
1642 SpatialIdentifiability::OrthogonalToParametric
1643 ) {
1644 return Err(SmoothError::invalid_config(format!(
1645 "{label} term '{}' is not frozen: ThinPlate identifiability must be FrozenTransform or None",
1646 st.name
1647 ))
1648 .into());
1649 }
1650 }
1651 SmoothBasisSpec::Sphere { spec, .. } => {
1652 if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
1653 return Err(SmoothError::invalid_config(format!(
1654 "{label} term '{}' is not frozen: Sphere centers must be UserProvided",
1655 st.name
1656 ))
1657 .into());
1658 }
1659 if matches!(spec.method, crate::basis::SphereMethod::Harmonic)
1660 && spec.max_degree.is_none_or(|d| d == 0)
1661 {
1662 return Err(format!(
1663 "{label} term '{}' is not frozen: sphere max_degree must be positive",
1664 st.name
1665 ));
1666 }
1667 }
1668 SmoothBasisSpec::ConstantCurvature { spec, .. } => {
1669 if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
1670 return Err(SmoothError::invalid_config(format!(
1671 "{label} term '{}' is not frozen: ConstantCurvature centers must be UserProvided",
1672 st.name
1673 ))
1674 .into());
1675 }
1676 if !(spec.length_scale.is_finite() && spec.length_scale > 0.0) {
1677 return Err(SmoothError::invalid_config(format!(
1678 "{label} term '{}' is not frozen: ConstantCurvature length_scale must be the realized positive value",
1679 st.name
1680 ))
1681 .into());
1682 }
1683 }
1684 SmoothBasisSpec::MeasureJet {
1685 spec, input_scale, ..
1686 } => {
1687 if input_scale.is_none() {
1688 return Err(SmoothError::invalid_config(format!(
1689 "{label} term '{}' is not frozen: MeasureJet input_scale is missing",
1690 st.name
1691 ))
1692 .into());
1693 }
1694 let centers = match &spec.center_strategy {
1695 CenterStrategy::UserProvided(centers) => centers,
1696 _ => {
1697 return Err(SmoothError::invalid_config(format!(
1698 "{label} term '{}' is not frozen: MeasureJet centers must be UserProvided",
1699 st.name
1700 ))
1701 .into());
1702 }
1703 };
1704 if centers.nrows() == 0 {
1705 return Err(SmoothError::invalid_config(format!(
1706 "{label} term '{}' is not frozen: MeasureJet centers are empty",
1707 st.name
1708 ))
1709 .into());
1710 }
1711 if !(spec.length_scale.is_finite() && spec.length_scale > 0.0) {
1712 return Err(SmoothError::invalid_config(format!(
1713 "{label} term '{}' is not frozen: MeasureJet length_scale must be the realized positive value",
1714 st.name
1715 ))
1716 .into());
1717 }
1718 let frozen = spec.frozen_quadrature.as_ref().ok_or_else(|| {
1721 SmoothError::invalid_config(format!(
1722 "{label} term '{}' is not frozen: MeasureJet frozen_quadrature payload is missing",
1723 st.name
1724 ))
1725 })?;
1726 if frozen.masses.len() != centers.nrows() {
1727 return Err(SmoothError::invalid_config(format!(
1728 "{label} term '{}' frozen MeasureJet has {} masses for {} centers",
1729 st.name,
1730 frozen.masses.len(),
1731 centers.nrows()
1732 ))
1733 .into());
1734 }
1735 let total_mass = frozen.masses.sum();
1736 if frozen
1737 .masses
1738 .iter()
1739 .any(|mass| !(mass.is_finite() && *mass >= 0.0))
1740 || !(total_mass.is_finite() && total_mass > 0.0)
1741 {
1742 return Err(SmoothError::invalid_config(format!(
1743 "{label} term '{}' frozen MeasureJet masses must be finite, nonnegative, and have positive total mass",
1744 st.name
1745 ))
1746 .into());
1747 }
1748 let n_levels = frozen.eps_band.len();
1749 if n_levels == 0
1750 || frozen
1751 .eps_band
1752 .iter()
1753 .any(|eps| !(eps.is_finite() && *eps > 0.0))
1754 {
1755 return Err(SmoothError::invalid_config(format!(
1756 "{label} term '{}' frozen MeasureJet eps_band must be nonempty, finite, and positive",
1757 st.name
1758 ))
1759 .into());
1760 }
1761 for (idx, pair) in frozen.eps_band.windows(2).enumerate() {
1762 if pair[1] <= pair[0] {
1763 return Err(SmoothError::invalid_config(format!(
1764 "{label} term '{}' frozen MeasureJet eps_band is not strictly ascending at {idx}: {} then {}",
1765 st.name,
1766 pair[0],
1767 pair[1]
1768 ))
1769 .into());
1770 }
1771 }
1772 validate_measure_jet_positive_vec_len(
1773 label,
1774 &st.name,
1775 "support_means",
1776 &frozen.support_means,
1777 n_levels,
1778 )?;
1779 let per_level = crate::basis::measure_jet_multiscale_mode(spec);
1787 if per_level {
1788 validate_measure_jet_positive_vec_len(
1789 label,
1790 &st.name,
1791 "penalty_normalization_scales",
1792 &frozen.penalty_normalization_scales,
1793 n_levels,
1794 )?;
1795 validate_measure_jet_positive_vec_len(
1796 label,
1797 &st.name,
1798 "raw_penalty_normalization_scales",
1799 &frozen.raw_penalty_normalization_scales,
1800 n_levels,
1801 )?;
1802 if frozen.fused_penalty_normalization_scale.is_some() {
1803 return Err(SmoothError::invalid_config(format!(
1804 "{label} term '{}' per-level MeasureJet must not carry a fused penalty normalization scale",
1805 st.name
1806 ))
1807 .into());
1808 }
1809 } else {
1810 if !frozen.penalty_normalization_scales.is_empty()
1811 || !frozen.raw_penalty_normalization_scales.is_empty()
1812 {
1813 return Err(SmoothError::invalid_config(format!(
1814 "{label} term '{}' fused MeasureJet must not carry per-level penalty normalization scales",
1815 st.name
1816 ))
1817 .into());
1818 }
1819 match frozen.fused_penalty_normalization_scale {
1820 Some(scale) if scale.is_finite() && scale > 0.0 => {}
1821 Some(scale) => {
1822 return Err(SmoothError::invalid_config(format!(
1823 "{label} term '{}' fused MeasureJet penalty normalization scale must be positive and finite, got {scale}",
1824 st.name
1825 ))
1826 .into());
1827 }
1828 None => {
1829 return Err(SmoothError::invalid_config(format!(
1830 "{label} term '{}' fused MeasureJet is missing its penalty normalization scale",
1831 st.name
1832 ))
1833 .into());
1834 }
1835 }
1836 }
1837 }
1838 SmoothBasisSpec::Matern {
1839 spec, input_scale, ..
1840 } => {
1841 if input_scale.is_none() {
1842 return Err(SmoothError::invalid_config(format!(
1843 "{label} term '{}' is not frozen: Matern input_scale is missing",
1844 st.name
1845 ))
1846 .into());
1847 }
1848 if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
1849 return Err(SmoothError::invalid_config(format!(
1850 "{label} term '{}' is not frozen: Matern centers must be UserProvided",
1851 st.name
1852 ))
1853 .into());
1854 }
1855 if spec
1856 .length_scale
1857 .resolved()
1858 .is_none_or(|value| !value.is_finite() || value <= 0.0)
1859 {
1860 return Err(SmoothError::invalid_config(format!(
1861 "{label} term '{}' is not frozen: Matern length_scale must be resolved, finite, and positive",
1862 st.name
1863 ))
1864 .into());
1865 }
1866 }
1867 SmoothBasisSpec::Duchon {
1868 spec, input_scale, ..
1869 } => {
1870 if input_scale.is_none() {
1871 return Err(SmoothError::invalid_config(format!(
1872 "{label} term '{}' is not frozen: Duchon input_scale is missing",
1873 st.name
1874 ))
1875 .into());
1876 }
1877 if !crate::basis::duchon_center_strategy_is_frozen(&spec.center_strategy) {
1878 return Err(SmoothError::invalid_config(format!(
1879 "{label} term '{}' is not frozen: Duchon knots and spectral basis must be resolved",
1880 st.name
1881 ))
1882 .into());
1883 }
1884 if matches!(
1885 spec.identifiability,
1886 SpatialIdentifiability::OrthogonalToParametric
1887 ) {
1888 return Err(SmoothError::invalid_config(format!(
1889 "{label} term '{}' is not frozen: Duchon identifiability must be FrozenTransform or None",
1890 st.name
1891 ))
1892 .into());
1893 }
1894 }
1895 SmoothBasisSpec::Pca {
1896 centered,
1897 center_mean,
1898 pca_basis_path,
1899 ..
1900 } => {
1901 if *centered && center_mean.is_none() && pca_basis_path.is_none() {
1902 return Err(SmoothError::invalid_config(format!(
1903 "{label} term '{}' is not frozen: centered Pca missing center_mean",
1904 st.name
1905 ))
1906 .into());
1907 }
1908 }
1909 SmoothBasisSpec::BySmooth { smooth, by_kind } => {
1910 if let SmoothBasisSpec::BySmooth { .. } = smooth.as_ref() {
1911 return Err(format!("{label} term '{}' has nested by-smooths", st.name));
1912 }
1913 match by_kind {
1914 ByVarKind::Numeric { .. } => {}
1915 ByVarKind::Factor { frozen_levels, .. } if frozen_levels.is_none() => {
1916 return Err(format!(
1917 "{label} term '{}' is not frozen: by-factor levels missing",
1918 st.name
1919 ));
1920 }
1921 ByVarKind::Factor { .. } => {}
1922 }
1923 let nested = TermCollectionSpec {
1924 linear_terms: vec![],
1925 random_effect_terms: vec![],
1926 smooth_terms: vec![SmoothTermSpec {
1927 frozen_parametric_residualization: None,
1928 name: st.name.clone(),
1929 basis: (**smooth).clone(),
1930 shape: st.shape,
1931 joint_null_rotation: None,
1932 }],
1933 };
1934 nested.validate_frozen(label)?;
1935 }
1936 SmoothBasisSpec::FactorSmooth { spec } => {
1937 if spec.group_frozen_levels.is_none() {
1938 return Err(format!(
1939 "{label} term '{}' is not frozen: factor-smooth levels missing",
1940 st.name
1941 ));
1942 }
1943 if !matches!(
1944 spec.marginal.knotspec,
1945 BSplineKnotSpec::Provided(_)
1946 | BSplineKnotSpec::PeriodicUniform { .. }
1947 | BSplineKnotSpec::NaturalCubicRegression { .. }
1959 ) {
1960 return Err(format!(
1961 "{label} term '{}' is not frozen: factor-smooth marginal knots missing",
1962 st.name
1963 ));
1964 }
1965 }
1966 SmoothBasisSpec::TensorBSpline { spec, .. } => {
1967 for (dim, marginal) in spec.marginalspecs.iter().enumerate() {
1968 if !matches!(
1969 marginal.knotspec,
1970 BSplineKnotSpec::Provided(_)
1971 | BSplineKnotSpec::PeriodicUniform { .. }
1972 | BSplineKnotSpec::NaturalCubicRegression { .. }
1973 ) {
1974 return Err(SmoothError::invalid_config(format!(
1975 "{label} term '{}' dim {} is not frozen: tensor marginal knotspec must be Provided, PeriodicUniform, or NaturalCubicRegression",
1976 st.name, dim
1977 ))
1978 .into());
1979 }
1980 }
1981 if matches!(
1982 spec.identifiability,
1983 TensorBSplineIdentifiability::SumToZero
1984 | TensorBSplineIdentifiability::MarginalSumToZero
1985 ) {
1986 return Err(SmoothError::invalid_config(format!(
1987 "{label} term '{}' is not frozen: tensor identifiability must be FrozenTransform or None",
1988 st.name
1989 ))
1990 .into());
1991 }
1992 }
1993 }
1994 }
1995
1996 for rt in &self.random_effect_terms {
1997 if rt.frozen_levels.is_none() {
1998 return Err(SmoothError::invalid_config(format!(
1999 "{label} random-effect term '{}' is not frozen: missing frozen_levels",
2000 rt.name
2001 ))
2002 .into());
2003 }
2004 }
2005
2006 Ok(())
2007 }
2008
2009 pub fn remap_feature_columns<E, F>(&self, mut remap: F) -> Result<TermCollectionSpec, E>
2028 where
2029 F: FnMut(usize) -> Result<usize, E>,
2030 {
2031 let mut out = self.clone();
2032 for lt in &mut out.linear_terms {
2033 lt.feature_col = remap(lt.feature_col)?;
2034 for fc in lt.feature_cols.iter_mut() {
2044 *fc = remap(*fc)?;
2045 }
2046 for (col, _bits) in lt.categorical_levels.iter_mut() {
2051 *col = remap(*col)?;
2052 }
2053 }
2054 for rt in &mut out.random_effect_terms {
2055 rt.feature_col = remap(rt.feature_col)?;
2056 }
2057 for st in &mut out.smooth_terms {
2058 remap_smooth_basis_feature_columns(&mut st.basis, &mut remap)?;
2059 }
2060 Ok(out)
2061 }
2062}
2063
2064pub fn remap_smooth_basis_feature_columns<E, F>(
2069 basis: &mut SmoothBasisSpec,
2070 remap: &mut F,
2071) -> Result<(), E>
2072where
2073 F: FnMut(usize) -> Result<usize, E>,
2074{
2075 match basis {
2076 SmoothBasisSpec::ByVariable { inner, by_col, .. }
2077 | SmoothBasisSpec::FactorSumToZero { inner, by_col, .. } => {
2078 *by_col = remap(*by_col)?;
2079 remap_smooth_basis_feature_columns(inner, remap)?;
2080 }
2081 SmoothBasisSpec::BSpline1D { feature_col, .. } => {
2082 *feature_col = remap(*feature_col)?;
2083 }
2084 SmoothBasisSpec::BySmooth { smooth, by_kind } => {
2085 let by_feature_col = match by_kind {
2086 ByVarKind::Numeric { feature_col } | ByVarKind::Factor { feature_col, .. } => {
2087 feature_col
2088 }
2089 };
2090 *by_feature_col = remap(*by_feature_col)?;
2091 remap_smooth_basis_feature_columns(smooth, remap)?;
2092 }
2093 SmoothBasisSpec::FactorSmooth { spec } => {
2094 for fc in spec.continuous_cols.iter_mut() {
2095 *fc = remap(*fc)?;
2096 }
2097 spec.group_col = remap(spec.group_col)?;
2098 }
2099 SmoothBasisSpec::ThinPlate { feature_cols, .. }
2100 | SmoothBasisSpec::Sphere { feature_cols, .. }
2101 | SmoothBasisSpec::ConstantCurvature { feature_cols, .. }
2102 | SmoothBasisSpec::Matern { feature_cols, .. }
2103 | SmoothBasisSpec::MeasureJet { feature_cols, .. }
2104 | SmoothBasisSpec::Duchon { feature_cols, .. }
2105 | SmoothBasisSpec::Pca { feature_cols, .. }
2106 | SmoothBasisSpec::TensorBSpline { feature_cols, .. } => {
2107 for fc in feature_cols.iter_mut() {
2108 *fc = remap(*fc)?;
2109 }
2110 }
2111 }
2112 Ok(())
2113}
2114
2115#[derive(Debug, Clone)]
2116pub enum PenaltyStructureHint {
2117 Ridge(f64),
2118 Kronecker(Vec<Array2<f64>>),
2119}
2120
2121#[derive(Clone)]
2128pub struct BlockwisePenalty {
2129 pub col_range: Range<usize>,
2131 pub local: Array2<f64>,
2134 pub prior_mean: gam_problem::CoefficientPriorMean,
2136 pub structure_hint: Option<PenaltyStructureHint>,
2139 pub op: Option<std::sync::Arc<dyn crate::analytic_penalties::PenaltyOp>>,
2144}
2145
2146impl std::fmt::Debug for BlockwisePenalty {
2147 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2148 f.debug_struct("BlockwisePenalty")
2149 .field("col_range", &self.col_range)
2150 .field(
2151 "local",
2152 &format_args!("{}×{}", self.local.nrows(), self.local.ncols()),
2153 )
2154 .field("prior_mean", &self.prior_mean)
2155 .field("structure_hint", &self.structure_hint)
2156 .field("op", &self.op.as_ref().map(|o| o.dim()))
2157 .finish()
2158 }
2159}
2160
2161impl BlockwisePenalty {
2162 pub fn new(col_range: Range<usize>, local: Array2<f64>) -> Self {
2164 assert_eq!(col_range.len(), local.nrows());
2165 assert_eq!(col_range.len(), local.ncols());
2166 Self {
2167 col_range,
2168 local,
2169 prior_mean: gam_problem::CoefficientPriorMean::Zero,
2170 structure_hint: None,
2171 op: None,
2172 }
2173 }
2174
2175 pub fn with_prior_mean(mut self, prior_mean: gam_problem::CoefficientPriorMean) -> Self {
2176 self.prior_mean = prior_mean;
2177 self
2178 }
2179
2180 pub fn with_op(
2182 mut self,
2183 op: Option<std::sync::Arc<dyn crate::analytic_penalties::PenaltyOp>>,
2184 ) -> Self {
2185 self.op = op;
2186 self
2187 }
2188
2189 pub fn ridge(col_range: Range<usize>, scale: f64) -> Self {
2190 let block_size = col_range.len();
2191 let mut local = Array2::<f64>::zeros((block_size, block_size));
2192 for i in 0..block_size {
2193 local[[i, i]] = scale;
2194 }
2195 Self {
2196 col_range,
2197 local,
2198 prior_mean: gam_problem::CoefficientPriorMean::Zero,
2199 structure_hint: Some(PenaltyStructureHint::Ridge(scale)),
2200 op: None,
2201 }
2202 }
2203
2204 pub fn kronecker(
2205 col_range: Range<usize>,
2206 local: Array2<f64>,
2207 factors: Vec<Array2<f64>>,
2208 ) -> Self {
2209 assert_eq!(col_range.len(), local.nrows());
2210 assert_eq!(col_range.len(), local.ncols());
2211 Self {
2212 col_range,
2213 local,
2214 prior_mean: gam_problem::CoefficientPriorMean::Zero,
2215 structure_hint: Some(PenaltyStructureHint::Kronecker(factors)),
2216 op: None,
2217 }
2218 }
2219
2220 pub fn to_global(&self, p_total: usize) -> Array2<f64> {
2224 let mut g = Array2::<f64>::zeros((p_total, p_total));
2225 let r = &self.col_range;
2226 assert!(
2227 r.end <= p_total && self.local.nrows() == r.len() && self.local.ncols() == r.len(),
2228 "BlockwisePenalty::to_global shape invariant violated: \
2229 col_range={}..{}, local={}x{}, p_total={}",
2230 r.start,
2231 r.end,
2232 self.local.nrows(),
2233 self.local.ncols(),
2234 p_total,
2235 );
2236 g.slice_mut(s![r.start..r.end, r.start..r.end])
2237 .assign(&self.local);
2238 g
2239 }
2240
2241 pub fn to_penalty_matrix(&self, total_dim: usize) -> gam_problem::PenaltyMatrix {
2244 gam_problem::PenaltyMatrix::Blockwise {
2245 local: self.local.clone(),
2246 col_range: self.col_range.clone(),
2247 total_dim,
2248 }
2249 }
2250
2251 #[inline]
2253 pub fn block_size(&self) -> usize {
2254 self.col_range.len()
2255 }
2256}
2257
2258pub fn weighted_blockwise_penalty_sum(
2262 penalties: &[BlockwisePenalty],
2263 lambdas: &[f64],
2264 p_total: usize,
2265) -> Array2<f64> {
2266 assert_eq!(penalties.len(), lambdas.len());
2267 for (idx, &lam) in lambdas.iter().enumerate() {
2274 assert!(
2275 lam.is_finite() && lam >= 0.0,
2276 "weighted_blockwise_penalty_sum: lambdas[{idx}] = {lam} is invalid (must be finite and non-negative; negative smoothing parameters violate S_λ ⪰ 0)",
2277 );
2278 }
2279 for (idx, bp) in penalties.iter().enumerate() {
2283 let r = &bp.col_range;
2284 assert!(
2285 r.end <= p_total,
2286 "weighted_blockwise_penalty_sum: penalties[{idx}] col_range {:?} exceeds p_total = {p_total}",
2287 r,
2288 );
2289 }
2290 let mut out = Array2::<f64>::zeros((p_total, p_total));
2291 for (bp, &lam) in penalties.iter().zip(lambdas.iter()) {
2292 let r = &bp.col_range;
2293 let mut slice = out.slice_mut(s![r.start..r.end, r.start..r.end]);
2294 slice.scaled_add(lam, &bp.local);
2295 }
2296 out
2297}
2298
2299#[derive(Debug, Clone)]
2306pub struct KroneckerPenaltySystem {
2307 pub marginal_penalties: Vec<Array2<f64>>,
2309 pub marginal_eigensystems: Vec<(Array1<f64>, Array2<f64>)>,
2311 pub marginal_dims: Vec<usize>,
2313 pub has_double_penalty: bool,
2315}
2316
2317impl KroneckerPenaltySystem {
2318 pub fn new(
2319 marginal_penalties: Vec<Array2<f64>>,
2320 marginal_dims: Vec<usize>,
2321 has_double_penalty: bool,
2322 ) -> Result<Self, BasisError> {
2323 if marginal_penalties.len() != marginal_dims.len() {
2324 crate::bail_dim_basis!(
2325 "KroneckerPenaltySystem: {} penalties vs {} dims",
2326 marginal_penalties.len(),
2327 marginal_dims.len()
2328 );
2329 }
2330 let eigensystems =
2331 kronecker_marginal_eigensystems(&marginal_penalties, "KroneckerPenaltySystem")
2332 .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
2333 Ok(Self {
2334 marginal_penalties,
2335 marginal_eigensystems: eigensystems,
2336 marginal_dims,
2337 has_double_penalty,
2338 })
2339 }
2340
2341 pub fn p_total(&self) -> usize {
2342 self.marginal_dims.iter().copied().product()
2343 }
2344
2345 pub fn ndim(&self) -> usize {
2346 self.marginal_dims.len()
2347 }
2348
2349 pub fn num_penalties(&self) -> usize {
2350 self.marginal_dims.len() + if self.has_double_penalty { 1 } else { 0 }
2351 }
2352
2353 pub fn logdet_and_derivatives(
2357 &self,
2358 lambdas: &[f64],
2359 objective_ridge: f64,
2360 ) -> (f64, Array1<f64>, Array2<f64>) {
2361 let n_pen = self.num_penalties();
2362 assert_eq!(lambdas.len(), n_pen, "lambda count mismatch");
2363 let marginal_evals: Vec<_> = self
2364 .marginal_eigensystems
2365 .iter()
2366 .map(|(evals, _)| evals.view())
2367 .collect();
2368 kronecker_logdet_and_derivatives(
2369 &marginal_evals,
2370 &self.marginal_dims,
2371 lambdas,
2372 self.has_double_penalty,
2373 objective_ridge,
2374 )
2375 }
2376
2377 pub fn logdet_rank_and_derivatives(
2378 &self,
2379 lambdas: &[f64],
2380 objective_ridge: f64,
2381 ) -> (f64, usize, Array1<f64>, Array2<f64>) {
2382 let n_pen = self.num_penalties();
2383 assert_eq!(lambdas.len(), n_pen, "lambda count mismatch");
2384 let d = self.marginal_dims.len();
2385 let mut logdet = 0.0;
2386 let mut rank = 0usize;
2387 let mut grad = Array1::<f64>::zeros(n_pen);
2388 let mut hess = Array2::<f64>::zeros((n_pen, n_pen));
2389 const EIGENVALUE_POSITIVITY_FLOOR: f64 = 1e-12;
2393 const STRUCTURAL_ZERO_FLOOR: f64 = 1e-12;
2396 let mut multi_idx = vec![0usize; d];
2397 loop {
2398 let mut sigma = 0.0;
2399 let mut structural_sigma = 0.0;
2400 for k in 0..d {
2401 let marginal_eigenvalue = self.marginal_eigensystems[k].0[multi_idx[k]];
2402 structural_sigma += marginal_eigenvalue;
2403 sigma += lambdas[k] * marginal_eigenvalue;
2404 }
2405 let joint_null = structural_sigma <= STRUCTURAL_ZERO_FLOOR;
2406 if self.has_double_penalty && joint_null {
2407 sigma += lambdas[d];
2408 }
2409 if structural_sigma > STRUCTURAL_ZERO_FLOOR {
2410 sigma += objective_ridge;
2411 }
2412
2413 if sigma > EIGENVALUE_POSITIVITY_FLOOR {
2414 rank += 1;
2415 logdet += sigma.ln();
2416 let inv_sigma = 1.0 / sigma;
2417 let inv_sigma2 = inv_sigma * inv_sigma;
2418 for k in 0..n_pen {
2419 let ck = if k < d {
2420 lambdas[k] * self.marginal_eigensystems[k].0[multi_idx[k]]
2421 } else if joint_null {
2422 lambdas[d]
2423 } else {
2424 0.0
2425 };
2426 grad[k] += ck * inv_sigma;
2427 hess[[k, k]] += ck * inv_sigma - ck * ck * inv_sigma2;
2428 for l in (k + 1)..n_pen {
2429 let cl = if l < d {
2430 lambdas[l] * self.marginal_eigensystems[l].0[multi_idx[l]]
2431 } else if joint_null {
2432 lambdas[d]
2433 } else {
2434 0.0
2435 };
2436 let off = -ck * cl * inv_sigma2;
2437 hess[[k, l]] += off;
2438 hess[[l, k]] += off;
2439 }
2440 }
2441 }
2442
2443 let mut carry = true;
2444 for dim in (0..d).rev() {
2445 if carry {
2446 multi_idx[dim] += 1;
2447 if multi_idx[dim] < self.marginal_dims[dim] {
2448 carry = false;
2449 } else {
2450 multi_idx[dim] = 0;
2451 }
2452 }
2453 }
2454 if carry {
2455 break;
2456 }
2457 }
2458 (logdet, rank, grad, hess)
2459 }
2460}
2461
2462#[cfg(test)]
2463mod joint_unpenalized_dim_tests {
2464 use super::{ActivePenalty, ActivePenaltyInfo, PenaltySource, joint_unpenalized_dim};
2465 use ndarray::{Array2, array};
2466
2467 fn active_penalty(
2468 matrix: Array2<f64>,
2469 effective_rank: usize,
2470 nullity: usize,
2471 original_index: usize,
2472 source: PenaltySource,
2473 ) -> ActivePenalty {
2474 ActivePenalty {
2475 matrix,
2476 nullity,
2477 null_eigenvectors: None,
2478 op: None,
2479 info: ActivePenaltyInfo {
2480 source,
2481 original_index,
2482 effective_rank,
2483 normalization_scale: 1.0,
2484 kronecker_factors: None,
2485 structural_null_frame: None,
2486 },
2487 }
2488 }
2489
2490 #[test]
2491 fn no_penalty_is_fully_unpenalized() {
2492 assert_eq!(joint_unpenalized_dim(4, &[]), 4);
2493 }
2494
2495 #[test]
2496 fn single_penalty_returns_its_own_null_space() {
2497 let s = array![[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 5.0]];
2500 let penalties = [active_penalty(s, 1, 2, 0, PenaltySource::Primary)];
2501 assert_eq!(joint_unpenalized_dim(3, &penalties), 2);
2502 }
2503
2504 #[test]
2505 fn complementary_double_penalty_has_empty_joint_null_space() {
2506 let bending = array![[0.0, 0.0, 0.0], [0.0, 4.0, 0.0], [0.0, 0.0, 4.0]];
2513 let ridge = array![[2.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]];
2514 let penalties = [
2515 active_penalty(bending, 2, 1, 0, PenaltySource::Primary),
2516 active_penalty(ridge, 1, 2, 1, PenaltySource::DoublePenaltyNullspace),
2517 ];
2518 assert_eq!(joint_unpenalized_dim(3, &penalties), 0);
2519 }
2520
2521 #[test]
2522 fn partial_overlap_keeps_shared_null_direction() {
2523 let a = array![[0.0, 0.0, 0.0], [0.0, 3.0, 0.0], [0.0, 0.0, 0.0]];
2527 let b = array![[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 3.0]];
2528 let penalties = [
2529 active_penalty(a, 1, 2, 0, PenaltySource::Primary),
2530 active_penalty(b, 1, 2, 1, PenaltySource::OperatorStiffness),
2531 ];
2532 assert_eq!(joint_unpenalized_dim(3, &penalties), 1);
2533 }
2534
2535 #[test]
2536 fn non_materialized_penalty_falls_back_conservatively() {
2537 let full: Array2<f64> = array![[0.0, 0.0], [0.0, 1.0]];
2541 let factor: Array2<f64> = array![[1.0]]; let mixed_penalties = [
2543 active_penalty(full, 1, 1, 0, PenaltySource::Primary),
2544 active_penalty(
2545 factor.clone(),
2546 2,
2547 0,
2548 1,
2549 PenaltySource::TensorMarginal { dim: 0 },
2550 ),
2551 ];
2552 assert_eq!(joint_unpenalized_dim(2, &mixed_penalties), 0);
2553 let factor_penalties = [active_penalty(
2555 factor,
2556 2,
2557 2,
2558 0,
2559 PenaltySource::TensorMarginal { dim: 0 },
2560 )];
2561 assert_eq!(joint_unpenalized_dim(4, &factor_penalties), 2);
2562 }
2563}
2564
2565#[cfg(test)]
2566mod kronecker_penalty_system_tests {
2567 use super::KroneckerPenaltySystem;
2568 use ndarray::array;
2569
2570 #[test]
2571 fn double_penalty_rank_derivatives_use_only_joint_null_space() {
2572 let penalties = vec![
2573 array![[0.0, 0.0], [0.0, 2.0]],
2574 array![[0.0, 0.0], [0.0, 3.0]],
2575 ];
2576 let system = KroneckerPenaltySystem::new(penalties, vec![2usize, 2usize], true).unwrap();
2577 let lambdas = vec![5.0, 7.0, 11.0];
2578
2579 let (logdet, rank, grad, hess) = system.logdet_rank_and_derivatives(&lambdas, 0.0);
2580
2581 let expected_diag = [11.0_f64, 21.0, 10.0, 31.0];
2582 let expected_logdet: f64 = expected_diag.iter().map(|v| v.ln()).sum();
2583 assert_eq!(rank, 4);
2584 assert!((logdet - expected_logdet).abs() <= 1e-12);
2585 assert!(
2586 (grad[2] - 1.0).abs() <= 1e-12,
2587 "double-penalty rank derivative must count only the joint null mode, got {}",
2588 grad[2]
2589 );
2590 assert!(hess[[2, 2]].abs() <= 1e-12);
2591 }
2592}
2593
2594#[derive(Clone, Debug)]
2595pub struct TermCollectionDesign {
2596 pub design: DesignMatrix,
2605 pub affine_offset: Array1<f64>,
2612 pub penalties: Vec<BlockwisePenalty>,
2613 pub nullspace_dims: Vec<usize>,
2614 pub penaltyinfo: Vec<PenaltyBlockInfo>,
2615 pub dropped_penaltyinfo: Vec<DroppedPenaltyBlockInfo>,
2616 pub coefficient_lower_bounds: Option<Array1<f64>>,
2619 pub linear_constraints: Option<LinearInequalityConstraints>,
2622 pub intercept_range: Range<usize>,
2623 pub linear_ranges: Vec<(String, Range<usize>)>,
2624 pub linear_function_masses: Vec<Option<f64>>,
2633 pub random_effect_ranges: Vec<(String, Range<usize>)>,
2634 pub random_effect_levels: Vec<(String, Vec<u64>)>,
2635 pub smooth: SmoothDesign,
2636}
2637
2638impl TermCollectionDesign {
2639 pub fn compose_offset(
2643 &self,
2644 base: ArrayView1<'_, f64>,
2645 context: &str,
2646 ) -> Result<Array1<f64>, BasisError> {
2647 let n = self.design.nrows();
2648 if self.affine_offset.len() != n || base.len() != n {
2649 crate::bail_dim_basis!(
2650 "{context}: design rows={n}, affine offset rows={}, base offset rows={}",
2651 self.affine_offset.len(),
2652 base.len()
2653 );
2654 }
2655 if self.affine_offset.iter().any(|value| !value.is_finite())
2656 || base.iter().any(|value| !value.is_finite())
2657 {
2658 crate::bail_invalid_basis!("{context}: offsets must be finite");
2659 }
2660 Ok(base.to_owned() + &self.affine_offset)
2661 }
2662
2663 pub fn apply(&self, beta: ArrayView1<'_, f64>) -> Result<Array1<f64>, BasisError> {
2667 if beta.len() != self.design.ncols() {
2668 crate::bail_dim_basis!(
2669 "term-collection predictor coefficient length {} does not match design width {}",
2670 beta.len(),
2671 self.design.ncols()
2672 );
2673 }
2674 if beta.iter().any(|value| !value.is_finite()) {
2675 crate::bail_invalid_basis!("term-collection predictor coefficients must be finite");
2676 }
2677 if self.affine_offset.len() != self.design.nrows() {
2678 crate::bail_dim_basis!(
2679 "term-collection affine offset has {} rows but design has {}",
2680 self.affine_offset.len(),
2681 self.design.nrows()
2682 );
2683 }
2684 if self.affine_offset.iter().any(|value| !value.is_finite()) {
2685 crate::bail_invalid_basis!("term-collection affine offset must be finite");
2686 }
2687 Ok(self.design.apply(&beta.to_owned()) + &self.affine_offset)
2688 }
2689
2690 pub fn leading_penalty_blocks_before_smooth(&self) -> usize {
2698 self.penaltyinfo
2699 .iter()
2700 .take_while(|info| {
2701 matches!(
2702 &info.penalty.source,
2703 crate::basis::PenaltySource::Other(source)
2704 if source == "LinearTermRidge"
2705 || source.starts_with("RandomEffectRidge(")
2706 )
2707 })
2708 .count()
2709 }
2710
2711 pub fn smooth_term_penalty_range(
2719 &self,
2720 term_idx: usize,
2721 ) -> Result<Option<Range<usize>>, String> {
2722 let Some(term) = self.smooth.terms.get(term_idx) else {
2723 return Ok(None);
2724 };
2725 if term.active_penalties.is_empty() {
2726 return Ok(None);
2727 }
2728
2729 let leading = self.leading_penalty_blocks_before_smooth();
2730 let smooth_count = self
2731 .smooth
2732 .terms
2733 .iter()
2734 .map(|smooth| smooth.active_penalties.len())
2735 .sum::<usize>();
2736 let expected = leading
2737 .checked_add(smooth_count)
2738 .ok_or_else(|| "term-collection penalty count overflow".to_string())?;
2739 if expected != self.penalties.len() || self.penaltyinfo.len() != self.penalties.len() {
2740 return Err(format!(
2741 "term-collection penalty layout is inconsistent: {leading} leading blocks + \
2742 {smooth_count} smooth blocks = {expected}, but there are {} penalties and {} \
2743 metadata records",
2744 self.penalties.len(),
2745 self.penaltyinfo.len()
2746 ));
2747 }
2748
2749 let local_offset = self
2750 .smooth
2751 .terms
2752 .iter()
2753 .take(term_idx)
2754 .map(|smooth| smooth.active_penalties.len())
2755 .sum::<usize>();
2756 let start = leading
2757 .checked_add(local_offset)
2758 .ok_or_else(|| "smooth penalty offset overflow".to_string())?;
2759 let end = start
2760 .checked_add(term.active_penalties.len())
2761 .ok_or_else(|| "smooth penalty range overflow".to_string())?;
2762 Ok(Some(start..end))
2763 }
2764
2765 pub fn penalties_as_penalty_matrix(&self) -> Vec<gam_problem::PenaltyMatrix> {
2769 let p = self.design.ncols();
2770 self.penalties
2771 .iter()
2772 .map(|bp| bp.to_penalty_matrix(p))
2773 .collect()
2774 }
2775
2776 #[inline]
2778 pub fn num_penalties(&self) -> usize {
2779 self.penalties.len()
2780 }
2781
2782 pub fn realize_coefficient_groups(
2785 &self,
2786 groups: &[CoefficientGroupSpec],
2787 base_prior: &gam_spec::RhoPrior,
2788 ) -> Result<RealizedCoefficientGroups, BasisError> {
2789 realize_coefficient_groups(self, groups, base_prior)
2790 }
2791
2792 pub fn kronecker_penalty_system(&self) -> Option<KroneckerPenaltySystem> {
2803 let [only_term] = self.smooth.terms.as_slice() else {
2804 return None;
2805 };
2806 let kron = only_term.kronecker_factored.as_ref()?;
2807 if kron.marginal_dims.len() < 2
2813 || kron.marginal_penalties.len() != kron.marginal_dims.len()
2814 || kron.marginal_designs.len() != kron.marginal_dims.len()
2815 {
2816 return None;
2817 }
2818 KroneckerPenaltySystem::new(
2819 kron.marginal_penalties.clone(),
2820 kron.marginal_dims.clone(),
2821 kron.has_double_penalty,
2822 )
2823 .ok()
2824 }
2825}
2826
2827#[derive(Clone)]
2833pub struct StandardLatentCoordConfig {
2834 pub values: std::sync::Arc<crate::latent::LatentCoordValues>,
2835 pub term_index: gam_problem::types::SmoothTermIdx,
2836 pub feature_cols: Vec<usize>,
2837 pub manifold: crate::latent::LatentManifold,
2838 pub manifold_auto: bool,
2839 pub retraction_registry: gam_problem::LatentRetractionRegistry,
2840 pub analytic_penalties: Option<std::sync::Arc<crate::AnalyticPenaltyRegistry>>,
2841}
2842
2843#[derive(Clone, Debug, Serialize, Deserialize)]
2844pub struct AdaptiveSpatialMap {
2845 pub termname: String,
2846 pub feature_cols: Vec<usize>,
2847 pub collocation_points: Array2<f64>,
2848 pub inv_magweight: Array1<f64>,
2849 pub invgradweight: Array1<f64>,
2850 pub inv_lapweight: Array1<f64>,
2851}
2852
2853#[derive(Clone, Debug, Serialize, Deserialize)]
2854pub struct AdaptiveRegularizationDiagnostics {
2855 pub epsilon_0: f64,
2856 pub epsilon_g: f64,
2857 pub epsilon_c: f64,
2858 pub epsilon_outer_iterations: usize,
2859 pub mm_iterations: usize,
2860 pub converged: bool,
2861 pub maps: Vec<AdaptiveSpatialMap>,
2862}
2863
2864#[derive(Debug, Clone)]
2865pub struct LinearColumnConditioning {
2866 col_idx: usize,
2867 mean: f64,
2868 scale: f64,
2869}
2870
2871#[derive(Debug, Clone, Default)]
2872pub struct LinearFitConditioning {
2873 pub intercept_idx: usize,
2874 pub columns: Vec<LinearColumnConditioning>,
2875}
2876
2877#[derive(Clone)]
2878pub struct SpatialPsiDerivative {
2879 pub penalty_index: usize,
2881 pub penalty_indices: Vec<usize>,
2882 pub global_range: Range<usize>,
2883 pub total_p: usize,
2884 pub x_psi_local: Array2<f64>,
2885 pub s_psi_components_local: Vec<Array2<f64>>,
2886 pub x_psi_psi_local: Array2<f64>,
2887 pub s_psi_psi_components_local: Vec<Array2<f64>>,
2888 pub aniso_group_id: Option<usize>,
2889 pub aniso_cross_designs: Option<Vec<(usize, Array2<f64>)>>,
2892 pub aniso_cross_penalty_provider: Option<
2896 std::sync::Arc<
2897 dyn Fn(usize) -> Result<Vec<Array2<f64>>, EstimationError> + Send + Sync + 'static,
2898 >,
2899 >,
2900 pub implicit_operator: Option<std::sync::Arc<crate::basis::ImplicitDesignPsiDerivative>>,
2905 pub implicit_axis: usize,
2907}
2908
2909#[derive(Debug, Clone)]
2910pub struct SpatialLogKappaCoords {
2911 pub values: Array1<f64>,
2914 pub dims_per_term: Vec<usize>,
2916}
2917
2918#[derive(Clone, Copy)]
2922pub enum AnisoBoundEnd {
2923 Lower,
2924 Upper,
2925}
2926
2927impl SpatialLogKappaCoords {
2928 pub fn new_with_dims(values: Array1<f64>, dims_per_term: Vec<usize>) -> Self {
2930 assert_eq!(
2931 values.len(),
2932 dims_per_term.iter().sum::<usize>(),
2933 "SpatialLogKappaCoords: values length {} != sum of dims_per_term {}",
2934 values.len(),
2935 dims_per_term.iter().sum::<usize>(),
2936 );
2937 Self {
2938 values,
2939 dims_per_term,
2940 }
2941 }
2942
2943 pub fn from_length_scales(
2945 spec: &TermCollectionSpec,
2946 term_indices: &[usize],
2947 options: &SpatialLengthScaleOptimizationOptions,
2948 ) -> Self {
2949 let mut out = Array1::<f64>::zeros(term_indices.len());
2950 for (slot, &term_idx) in term_indices.iter().enumerate() {
2951 if let Some(cc) = constant_curvature_term_spec(spec, term_idx) {
2957 out[slot] = cc.kappa;
2958 continue;
2959 }
2960 out[slot] = spatial_term_seed_psi(spec, term_idx, options);
2964 }
2965 Self {
2966 values: out,
2967 dims_per_term: vec![1; term_indices.len()],
2968 }
2969 }
2970
2971 pub fn from_length_scales_aniso(
2985 spec: &TermCollectionSpec,
2986 term_indices: &[usize],
2987 options: &SpatialLengthScaleOptimizationOptions,
2988 ) -> Self {
2989 let mut vals = Vec::new();
2990 let mut dims = Vec::new();
2991 for &term_idx in term_indices {
2992 if let Some(mj) = measure_jet_term_spec(spec, term_idx) {
2996 let seed = measure_jet_psi_seed(mj);
2997 dims.push(seed.len());
2998 vals.extend(seed);
2999 continue;
3000 }
3001 if let Some(cc) = constant_curvature_term_spec(spec, term_idx) {
3007 vals.push(cc.kappa);
3008 dims.push(1);
3009 continue;
3010 }
3011 let psi_bar = spatial_term_seed_psi(spec, term_idx, options);
3015
3016 if spatial_term_uses_per_axis_psi(spec, term_idx) {
3017 let d = get_spatial_feature_dim(spec, term_idx).unwrap_or(1);
3022 let eta_raw = get_spatial_aniso_log_scales(spec, term_idx)
3023 .expect("predicate guarantees aniso_log_scales is Some");
3024 let eta = center_aniso_log_scales(&eta_raw);
3025 for &eta_a in &eta {
3026 vals.push(psi_bar + eta_a);
3027 }
3028 dims.push(d);
3029 } else {
3030 vals.push(psi_bar);
3037 dims.push(1);
3038 }
3039 }
3040 Self {
3041 values: Array1::from_vec(vals),
3042 dims_per_term: dims,
3043 }
3044 }
3045
3046 pub fn lower_bounds_from_data(
3053 data: ArrayView2<'_, f64>,
3054 spec: &TermCollectionSpec,
3055 term_indices: &[usize],
3056 options: &SpatialLengthScaleOptimizationOptions,
3057 ) -> Result<Self, BasisError> {
3058 let mut values = Array1::<f64>::zeros(term_indices.len());
3059 for (slot, &term_idx) in term_indices.iter().enumerate() {
3060 values[slot] = spatial_term_psi_search_box(data, spec, term_idx, options)?.0;
3061 }
3062 Ok(Self {
3063 values,
3064 dims_per_term: vec![1; term_indices.len()],
3065 })
3066 }
3067
3068 pub fn upper_bounds_from_data(
3071 data: ArrayView2<'_, f64>,
3072 spec: &TermCollectionSpec,
3073 term_indices: &[usize],
3074 options: &SpatialLengthScaleOptimizationOptions,
3075 ) -> Result<Self, BasisError> {
3076 let mut values = Array1::<f64>::zeros(term_indices.len());
3077 for (slot, &term_idx) in term_indices.iter().enumerate() {
3078 values[slot] = spatial_term_psi_search_box(data, spec, term_idx, options)?.1;
3079 }
3080 Ok(Self {
3081 values,
3082 dims_per_term: vec![1; term_indices.len()],
3083 })
3084 }
3085
3086 pub fn lower_bounds_aniso_from_data(
3095 data: ArrayView2<'_, f64>,
3096 spec: &TermCollectionSpec,
3097 term_indices: &[usize],
3098 dims_per_term: &[usize],
3099 options: &SpatialLengthScaleOptimizationOptions,
3100 ) -> Result<Self, BasisError> {
3101 Self::aniso_bounds_from_data(
3102 data,
3103 spec,
3104 term_indices,
3105 dims_per_term,
3106 options,
3107 AnisoBoundEnd::Lower,
3108 )
3109 }
3110
3111 pub fn upper_bounds_aniso_from_data(
3115 data: ArrayView2<'_, f64>,
3116 spec: &TermCollectionSpec,
3117 term_indices: &[usize],
3118 dims_per_term: &[usize],
3119 options: &SpatialLengthScaleOptimizationOptions,
3120 ) -> Result<Self, BasisError> {
3121 Self::aniso_bounds_from_data(
3122 data,
3123 spec,
3124 term_indices,
3125 dims_per_term,
3126 options,
3127 AnisoBoundEnd::Upper,
3128 )
3129 }
3130
3131 fn aniso_bounds_from_data(
3135 data: ArrayView2<'_, f64>,
3136 spec: &TermCollectionSpec,
3137 term_indices: &[usize],
3138 dims_per_term: &[usize],
3139 options: &SpatialLengthScaleOptimizationOptions,
3140 end: AnisoBoundEnd,
3141 ) -> Result<Self, BasisError> {
3142 assert_eq!(term_indices.len(), dims_per_term.len());
3143 let total: usize = dims_per_term.iter().sum();
3144 let mut values = Array1::<f64>::zeros(total);
3145 let mut cursor = 0;
3146 for (slot, &term_idx) in term_indices.iter().enumerate() {
3147 let d = dims_per_term[slot];
3148 if measure_jet_term_spec(spec, term_idx).is_some() {
3155 let term = spec
3156 .smooth_terms
3157 .get(term_idx)
3158 .expect("measure_jet_term_spec resolved this index");
3159 let bounds = measure_jet_psi_bound_values(
3160 data,
3161 &term.basis,
3162 matches!(end, AnisoBoundEnd::Upper),
3163 )?;
3164 for (offset, bound) in bounds.into_iter().enumerate() {
3165 if offset < d {
3166 values[cursor + offset] = bound;
3167 }
3168 }
3169 cursor += d;
3170 continue;
3171 }
3172 if constant_curvature_term_spec(spec, term_idx).is_some() {
3175 let (lo, hi) = constant_curvature_kappa_bounds(data, spec, term_idx);
3176 if d >= 1 {
3177 values[cursor] = match end {
3178 AnisoBoundEnd::Lower => lo,
3179 AnisoBoundEnd::Upper => hi,
3180 };
3181 }
3182 cursor += d;
3183 continue;
3184 }
3185 let psi_bound = {
3186 let (lo, hi) = spatial_term_psi_search_box(data, spec, term_idx, options)?;
3191 match end {
3192 AnisoBoundEnd::Lower => lo,
3193 AnisoBoundEnd::Upper => hi,
3194 }
3195 };
3196 let axis_offsets = if d <= 1 {
3197 vec![0.0; d]
3198 } else {
3199 get_spatial_aniso_log_scales(spec, term_idx)
3200 .filter(|eta| eta.len() == d)
3201 .map(|eta| center_aniso_log_scales(&eta))
3202 .unwrap_or_else(|| vec![0.0; d])
3203 };
3204 for offset in 0..d {
3205 values[cursor + offset] = psi_bound + axis_offsets[offset];
3206 }
3207 cursor += d;
3208 }
3209 Ok(Self {
3210 values,
3211 dims_per_term: dims_per_term.to_vec(),
3212 })
3213 }
3214
3215 pub fn reseed_from_data(
3224 mut self,
3225 data: ArrayView2<'_, f64>,
3226 spec: &TermCollectionSpec,
3227 term_indices: &[usize],
3228 options: &SpatialLengthScaleOptimizationOptions,
3229 ) -> Result<Self, BasisError> {
3230 assert_eq!(term_indices.len(), self.dims_per_term.len());
3231 let mut cursor = 0;
3232 for (slot, &term_idx) in term_indices.iter().enumerate() {
3233 let d = self.dims_per_term[slot];
3234 if measure_jet_term_spec(spec, term_idx).is_some() {
3237 cursor += d;
3238 continue;
3239 }
3240 if constant_curvature_term_spec(spec, term_idx).is_some() {
3244 cursor += d;
3245 continue;
3246 }
3247 let Some(psi_bar_new) = spatial_term_psi_seed(data, spec, term_idx, options)? else {
3248 cursor += d;
3249 continue;
3250 };
3251 if d == 0 {
3252 continue;
3253 }
3254 let current: Vec<f64> = self.values.slice(s![cursor..cursor + d]).to_vec();
3255 let psi_bar_old = current.iter().sum::<f64>() / d as f64;
3256 for (offset, &old_value) in current.iter().enumerate() {
3257 self.values[cursor + offset] = psi_bar_new + (old_value - psi_bar_old);
3258 }
3259 cursor += d;
3260 }
3261 Ok(self)
3262 }
3263
3264 pub fn clamp_to_bounds(
3275 mut self,
3276 lower: &SpatialLogKappaCoords,
3277 upper: &SpatialLogKappaCoords,
3278 ) -> Self {
3279 assert_eq!(self.values.len(), lower.values.len());
3280 assert_eq!(self.values.len(), upper.values.len());
3281 let mut n_projected = 0usize;
3282 let mut worst_delta = 0.0_f64;
3283 for idx in 0..self.values.len() {
3284 let lo = lower.values[idx];
3285 let hi = upper.values[idx];
3286 if !(lo.is_finite() && hi.is_finite()) {
3287 continue;
3288 }
3289 let v = self.values[idx];
3290 if v < lo {
3291 worst_delta = worst_delta.max(lo - v);
3292 self.values[idx] = lo;
3293 n_projected += 1;
3294 } else if v > hi {
3295 worst_delta = worst_delta.max(v - hi);
3296 self.values[idx] = hi;
3297 n_projected += 1;
3298 }
3299 }
3300 if n_projected > 0 {
3301 log::info!(
3302 "[spatial-kappa] projected {n_projected}/{} ψ seed coords into data-derived bounds \
3303 (worst excess={worst_delta:.3} log units); user length_scale falls outside \
3304 [{KERNEL_RANGE_MIN_DIAMETER_FRACTION}/r_max, {KERNEL_RANGE_MAX_SPACING_MULTIPLE}/r_min] geometry window",
3305 self.values.len()
3306 );
3307 }
3308 self
3309 }
3310
3311 pub fn from_theta_tail_with_dims(
3313 theta: &Array1<f64>,
3314 start: usize,
3315 dims_per_term: Vec<usize>,
3316 ) -> Self {
3317 let total: usize = dims_per_term.iter().sum();
3318 Self {
3319 values: theta.slice(s![start..start + total]).to_owned(),
3320 dims_per_term,
3321 }
3322 }
3323
3324 pub fn len(&self) -> usize {
3326 self.values.len()
3327 }
3328
3329 pub fn dims_per_term(&self) -> &[usize] {
3331 &self.dims_per_term
3332 }
3333
3334 fn term_offset(&self, term_idx: usize) -> usize {
3336 self.dims_per_term[..term_idx].iter().sum()
3337 }
3338
3339 pub fn term_slice(&self, term_idx: usize) -> &[f64] {
3341 let offset = self.term_offset(term_idx);
3342 let d = self.dims_per_term[term_idx];
3343 &self
3344 .values
3345 .as_slice()
3346 .expect("psi values are an owned contiguous Array1")[offset..offset + d]
3347 }
3348
3349 pub fn as_array(&self) -> &Array1<f64> {
3350 &self.values
3351 }
3352
3353 pub fn set_scalar_slot(&mut self, slot: usize, value: f64) -> bool {
3360 if slot >= self.dims_per_term.len() || self.dims_per_term[slot] != 1 {
3361 return false;
3362 }
3363 let offset = self.term_offset(slot);
3364 self.values[offset] = value;
3365 true
3366 }
3367
3368 pub fn split_at(&self, mid: usize) -> (Self, Self) {
3371 let flat_mid: usize = self.dims_per_term[..mid].iter().sum();
3372 (
3373 Self {
3374 values: self.values.slice(s![0..flat_mid]).to_owned(),
3375 dims_per_term: self.dims_per_term[..mid].to_vec(),
3376 },
3377 Self {
3378 values: self.values.slice(s![flat_mid..]).to_owned(),
3379 dims_per_term: self.dims_per_term[mid..].to_vec(),
3380 },
3381 )
3382 }
3383
3384 pub fn apply_tospec(
3391 &self,
3392 spec: &TermCollectionSpec,
3393 term_indices: &[usize],
3394 ) -> Result<TermCollectionSpec, EstimationError> {
3395 if term_indices.len() != self.dims_per_term.len() {
3396 crate::bail_invalid_estim!(
3397 "SpatialLogKappaCoords::apply_tospec: term count mismatch: \
3398 term_indices={} dims_per_term={}",
3399 term_indices.len(),
3400 self.dims_per_term.len()
3401 );
3402 }
3403 let mut updated = spec.clone();
3404 for (slot, &term_idx) in term_indices.iter().enumerate() {
3405 let psi = self.term_slice(slot);
3406 let d = self.dims_per_term[slot];
3407 if measure_jet_term_spec(&updated, term_idx).is_some() {
3410 set_measure_jet_psi_dials(&mut updated, term_idx, psi)?;
3411 continue;
3412 }
3413 if constant_curvature_term_spec(&updated, term_idx).is_some() {
3417 set_constant_curvature_kappa(&mut updated, term_idx, psi)?;
3418 continue;
3419 }
3420 let (next_length_scale, next_aniso) = spatial_term_psi_to_length_scale_and_aniso(psi);
3421 if (d == 1 || next_length_scale.is_some())
3422 && let Some(length_scale) = next_length_scale
3423 {
3424 set_spatial_length_scale(&mut updated, term_idx, length_scale)?;
3425 }
3426 if let Some(eta) = next_aniso {
3427 set_spatial_aniso_log_scales(&mut updated, term_idx, eta)?;
3428 }
3429 }
3430 Ok(updated)
3431 }
3432}
3433
3434pub fn center_aniso_log_scales(eta: &[f64]) -> Vec<f64> {
3435 if eta.len() <= 1 {
3436 return eta.to_vec();
3437 }
3438 let mean = eta.iter().sum::<f64>() / eta.len() as f64;
3439 eta.iter()
3440 .map(|&v| {
3441 let centered = v - mean;
3442 if centered.abs() <= 1e-15 {
3443 0.0
3444 } else {
3445 centered
3446 }
3447 })
3448 .collect()
3449}
3450
3451pub fn spatial_term_uses_per_axis_psi(resolvedspec: &TermCollectionSpec, term_idx: usize) -> bool {
3454 if let Some(mj) = measure_jet_term_spec(resolvedspec, term_idx) {
3455 return measure_jet_enrolls_psi(mj);
3456 }
3457 let Some(d) = get_spatial_feature_dim(resolvedspec, term_idx) else {
3458 return false;
3459 };
3460 if d <= 1 {
3461 return false;
3462 }
3463 let Some(eta) = get_spatial_aniso_log_scales(resolvedspec, term_idx) else {
3464 return false;
3465 };
3466 if eta.len() != d {
3467 return false;
3468 }
3469 let Some(term) = resolvedspec.smooth_terms.get(term_idx) else {
3490 return false;
3491 };
3492 match &term.basis {
3493 SmoothBasisSpec::Duchon { spec, .. } => {
3494 term.joint_null_rotation.is_none()
3502 && crate::basis::duchon_spec_supports_axis_psi(spec, d)
3503 }
3504 _ => true,
3505 }
3506}
3507
3508pub fn set_spatial_length_scale(
3509 spec: &mut TermCollectionSpec,
3510 term_idx: usize,
3511 length_scale: f64,
3512) -> Result<(), EstimationError> {
3513 let Some(term) = spec.smooth_terms.get_mut(term_idx) else {
3514 crate::bail_invalid_estim!("spatial length-scale term index {term_idx} out of range");
3515 };
3516 match &mut term.basis {
3517 SmoothBasisSpec::ThinPlate { spec, .. } => {
3518 spec.length_scale = length_scale;
3519 Ok(())
3520 }
3521 SmoothBasisSpec::Matern { spec, .. } => {
3522 spec.length_scale.set_resolved(length_scale);
3523 Ok(())
3524 }
3525 SmoothBasisSpec::Duchon { spec, .. } => {
3526 spec.length_scale = Some(length_scale);
3527 Ok(())
3528 }
3529 _ => Err(EstimationError::InvalidInput(format!(
3530 "term '{}' does not expose a spatial length scale",
3531 term.name
3532 ))),
3533 }
3534}
3535
3536pub fn get_spatial_length_scale(spec: &TermCollectionSpec, term_idx: usize) -> Option<f64> {
3537 spec.smooth_terms
3538 .get(term_idx)
3539 .and_then(|term| match &term.basis {
3540 SmoothBasisSpec::ThinPlate { spec, .. } => Some(spec.length_scale),
3541 SmoothBasisSpec::Matern { spec, .. } => spec.length_scale.resolved(),
3542 SmoothBasisSpec::Duchon { spec, .. } => spec.length_scale,
3543 _ => None,
3544 })
3545}
3546
3547pub fn spatial_term_supports_hyper_optimization(
3548 spec: &TermCollectionSpec,
3549 term_idx: usize,
3550) -> bool {
3551 if let Some(term) = spec.smooth_terms.get(term_idx)
3557 && let SmoothBasisSpec::ThinPlate { .. } = &term.basis
3558 {
3559 return false;
3560 }
3561
3562 if let Some(term) = spec.smooth_terms.get(term_idx)
3592 && let SmoothBasisSpec::Matern { .. } = &term.basis
3593 {
3594 return true;
3595 }
3596
3597 if let Some(mj) = measure_jet_term_spec(spec, term_idx) {
3600 return measure_jet_enrolls_psi(mj);
3601 }
3602
3603 if constant_curvature_term_spec(spec, term_idx).is_some() {
3615 return true;
3616 }
3617
3618 get_spatial_length_scale(spec, term_idx).is_some()
3619}
3620
3621pub fn constant_curvature_term_spec(
3624 spec: &TermCollectionSpec,
3625 term_idx: usize,
3626) -> Option<&crate::basis::ConstantCurvatureBasisSpec> {
3627 spec.smooth_terms
3628 .get(term_idx)
3629 .and_then(|term| match &term.basis {
3630 SmoothBasisSpec::ConstantCurvature { spec, .. } => Some(spec),
3631 _ => None,
3632 })
3633}
3634
3635pub const CONSTANT_CURVATURE_KAPPA_CHART_FRACTION: f64 = 0.5;
3690
3691pub const CONSTANT_CURVATURE_MIN_CHART_RADIUS2: f64 = 1e-8;
3695
3696pub fn constant_curvature_kappa_bounds(
3754 data: ArrayView2<'_, f64>,
3755 spec: &TermCollectionSpec,
3756 term_idx: usize,
3757) -> (f64, f64) {
3758 let (feature_cols, cc) = match spec.smooth_terms.get(term_idx).map(|t| &t.basis) {
3759 Some(SmoothBasisSpec::ConstantCurvature {
3760 feature_cols, spec, ..
3761 }) => (feature_cols, spec),
3762 _ => return (-1.0, 1.0),
3763 };
3764 let data_r2 = crate::basis::constant_curvature_data_chart_radius2(data, feature_cols);
3765 let center_r2 = crate::basis::constant_curvature_center_chart_radius2(
3766 data,
3767 feature_cols,
3768 &cc.center_strategy,
3769 );
3770 let max_r2 = data_r2
3771 .max(center_r2)
3772 .max(CONSTANT_CURVATURE_MIN_CHART_RADIUS2);
3773 let half = CONSTANT_CURVATURE_KAPPA_CHART_FRACTION / max_r2;
3774 (-half, half)
3775}
3776
3777pub fn set_constant_curvature_kappa(
3781 spec: &mut TermCollectionSpec,
3782 term_idx: usize,
3783 psi: &[f64],
3784) -> Result<bool, EstimationError> {
3785 let Some(term) = spec.smooth_terms.get_mut(term_idx) else {
3786 crate::bail_invalid_estim!(
3787 "constant-curvature κ write-back: term index {term_idx} out of range"
3788 );
3789 };
3790 set_single_term_constant_curvature_kappa(term, psi)
3791}
3792
3793pub fn set_single_term_constant_curvature_kappa(
3798 term: &mut SmoothTermSpec,
3799 psi: &[f64],
3800) -> Result<bool, EstimationError> {
3801 if psi.len() != 1 {
3802 crate::bail_invalid_estim!(
3803 "constant-curvature κ write-back expects exactly one value, got {}",
3804 psi.len()
3805 );
3806 }
3807 let next_kappa = psi[0];
3808 if !next_kappa.is_finite() {
3809 crate::bail_invalid_estim!(
3810 "constant-curvature κ write-back produced a non-finite κ = {next_kappa}"
3811 );
3812 }
3813 let SmoothBasisSpec::ConstantCurvature { spec: cc, .. } = &mut term.basis else {
3814 crate::bail_invalid_estim!(
3815 "constant-curvature κ write-back targeted a non-constant-curvature term"
3816 );
3817 };
3818 if cc.kappa != next_kappa {
3819 cc.kappa = next_kappa;
3820 Ok(true)
3821 } else {
3822 Ok(false)
3823 }
3824}
3825
3826pub fn spatial_term_has_locked_kappa(spec: &TermCollectionSpec, term_idx: usize) -> bool {
3837 let explicitly_fixed = spec
3838 .smooth_terms
3839 .get(term_idx)
3840 .is_some_and(|term| match &term.basis {
3841 SmoothBasisSpec::Matern { spec, .. } => spec.length_scale.is_fixed(),
3842 SmoothBasisSpec::ThinPlate { .. } => true,
3843 SmoothBasisSpec::Duchon { spec, .. } => spec.length_scale.is_some(),
3844 _ => false,
3845 });
3846 explicitly_fixed && !spatial_term_uses_per_axis_psi(spec, term_idx)
3847}
3848
3849pub fn all_spatial_terms_kappa_fixed(spec: &TermCollectionSpec) -> bool {
3861 spec.smooth_terms.iter().enumerate().all(|(idx, _)| {
3862 !spatial_term_supports_hyper_optimization(spec, idx)
3863 || spatial_term_has_locked_kappa(spec, idx)
3864 })
3865}
3866
3867pub fn spatial_identifiability_policy(
3868 termspec: &SmoothTermSpec,
3869) -> Option<&SpatialIdentifiability> {
3870 match &termspec.basis {
3871 SmoothBasisSpec::ThinPlate { spec, .. } => Some(&spec.identifiability),
3872 SmoothBasisSpec::Duchon { spec, .. } => Some(&spec.identifiability),
3873 _ => None,
3874 }
3875}
3876
3877pub const NULLSPACE_DEGENERACY_RHO_SD: f64 = 15.0;
3881
3882
3883pub const KERNEL_RANGE_MIN_DIAMETER_FRACTION: f64 = 2.0;
3895
3896pub const KERNEL_RANGE_MAX_SPACING_MULTIPLE: f64 = 1e2;
3901
3902fn spatial_term_stored_input_scale(term: &SmoothTermSpec) -> Option<crate::IsotropicScale> {
3903 match &term.basis {
3904 SmoothBasisSpec::ThinPlate { input_scale, .. }
3905 | SmoothBasisSpec::Matern { input_scale, .. }
3906 | SmoothBasisSpec::Duchon { input_scale, .. } => *input_scale,
3907 _ => None,
3908 }
3909}
3910
3911fn spatial_term_realized_input_scale(
3912 data: ArrayView2<'_, f64>,
3913 term: &SmoothTermSpec,
3914) -> Result<crate::IsotropicScale, BasisError> {
3915 let (feature_cols, stored) = match &term.basis {
3916 SmoothBasisSpec::ThinPlate {
3917 feature_cols,
3918 input_scale,
3919 ..
3920 }
3921 | SmoothBasisSpec::Matern {
3922 feature_cols,
3923 input_scale,
3924 ..
3925 }
3926 | SmoothBasisSpec::Duchon {
3927 feature_cols,
3928 input_scale,
3929 ..
3930 } => (feature_cols, input_scale),
3931 _ => {
3932 return Err(BasisError::InvalidInput(format!(
3933 "term '{}' does not have an isotropic Euclidean input frame",
3934 term.name
3935 )));
3936 }
3937 };
3938 if let Some(scale) = stored {
3939 return Ok(*scale);
3940 }
3941 let x = select_columns(data, feature_cols)?;
3942 estimate_isotropic_scale(x.view())
3943}
3944
3945pub fn spatial_term_psi_bounds(
3952 data: ArrayView2<'_, f64>,
3953 spec: &TermCollectionSpec,
3954 term_idx: usize,
3955 options: &SpatialLengthScaleOptimizationOptions,
3956) -> Result<(f64, f64), BasisError> {
3957 let options_window = (
3958 -options.max_length_scale.ln(),
3959 -options.min_length_scale.ln(),
3960 );
3961 if constant_curvature_term_spec(spec, term_idx).is_some() {
3966 return Ok(constant_curvature_kappa_bounds(data, spec, term_idx));
3967 }
3968 let term = spec.smooth_terms.get(term_idx).ok_or_else(|| {
3969 BasisError::InvalidInput(format!(
3970 "spatial term index {term_idx} is out of bounds for {} smooth terms",
3971 spec.smooth_terms.len()
3972 ))
3973 })?;
3974 let aniso = get_spatial_aniso_log_scales(spec, term_idx);
3987 let stored_input_scale = spatial_term_stored_input_scale(term);
3988 let input_scale = spatial_term_realized_input_scale(data, term)?;
3989 let r_bounds = match spatial_term_center_strategy(term) {
3990 Some(CenterStrategy::UserProvided(centers)) if centers.nrows() >= 2 => {
3991 let mut centers_in_frame = centers.clone();
3992 if stored_input_scale.is_none() {
3993 input_scale.standardize(&mut centers_in_frame);
3994 }
3995 let bounds = match aniso.as_deref() {
3996 Some(eta) if eta.len() == centers_in_frame.ncols() => {
3997 let y = points_in_aniso_y_space(centers_in_frame.view(), eta);
3998 pairwise_distance_bounds(y.view())
3999 }
4000 _ => pairwise_distance_bounds(centers_in_frame.view()),
4001 };
4002 bounds
4003 }
4004 _ => {
4005 let x = standardized_spatial_term_data(data, term)?;
4006 match aniso.as_deref() {
4007 Some(eta) if eta.len() == x.ncols() => {
4008 let y = points_in_aniso_y_space(x.view(), eta);
4009 pairwise_distance_bounds_sampled(y.view())
4010 }
4011 _ => pairwise_distance_bounds_sampled(x.view()),
4012 }
4013 }
4014 };
4015 let (r_min, r_max) = r_bounds.ok_or_else(|| {
4016 BasisError::InvalidInput(format!(
4017 "term '{}' has no positive finite pairwise-distance range",
4018 term.name
4019 ))
4020 })?;
4021 let inverse_sigma = input_scale.reciprocal();
4038 let psi_chart_offset = inverse_sigma.ln();
4039 let psi_lo_data = (KERNEL_RANGE_MIN_DIAMETER_FRACTION / r_max).ln() + psi_chart_offset;
4040 let psi_hi_data = (KERNEL_RANGE_MAX_SPACING_MULTIPLE / r_min).ln() + psi_chart_offset;
4041 let psi_lo = psi_lo_data.max(options_window.0);
4051 let psi_hi = psi_hi_data.min(options_window.1);
4052 if psi_lo >= psi_hi {
4053 return Err(BasisError::InvalidInput(format!(
4054 "term '{}' has an empty spatial ψ window after intersecting data bounds [{psi_lo_data}, {psi_hi_data}] with configured bounds [{}, {}]",
4055 term.name, options_window.0, options_window.1
4056 )));
4057 }
4058 Ok((psi_lo, psi_hi))
4059}
4060
4061pub fn spatial_term_psi_search_box(
4098 data: ArrayView2<'_, f64>,
4099 spec: &TermCollectionSpec,
4100 term_idx: usize,
4101 options: &SpatialLengthScaleOptimizationOptions,
4102) -> Result<(f64, f64), BasisError> {
4103 let (mut psi_lo, mut psi_hi) = spatial_term_psi_bounds(data, spec, term_idx, options)?;
4104 if constant_curvature_term_spec(spec, term_idx).is_some() {
4108 return Ok((psi_lo, psi_hi));
4109 }
4110 let options_window = (
4111 -options.max_length_scale.ln(),
4112 -options.min_length_scale.ln(),
4113 );
4114 if let Some(length_scale) = get_spatial_length_scale(spec, term_idx)
4115 && length_scale.is_finite()
4116 && length_scale > 0.0
4117 {
4118 let psi_incumbent = -length_scale.ln();
4119 if psi_incumbent.is_finite() {
4120 psi_lo = psi_lo.min(psi_incumbent.max(options_window.0));
4121 psi_hi = psi_hi.max(psi_incumbent.min(options_window.1));
4122 }
4123 }
4124 Ok((psi_lo, psi_hi))
4125}
4126
4127#[cfg(test)]
4128mod spatial_psi_bound_coordinate_tests {
4129 use super::*;
4130 use crate::basis::{MaternIdentifiability, MaternNu};
4131 use ndarray::array;
4132
4133 fn frozen_matern_bounds(theta: f64, dilation: f64) -> (f64, f64) {
4134 let source = array![
4135 [-1.7, -0.4],
4136 [-1.1, 0.8],
4137 [-0.2, -1.3],
4138 [0.5, 1.6],
4139 [1.4, -0.7],
4140 [2.1, 0.5],
4141 ];
4142 let (cos_theta, sin_theta) = (theta.cos(), theta.sin());
4143 let mut data = Array2::<f64>::zeros(source.raw_dim());
4144 for row in 0..source.nrows() {
4145 let x = source[[row, 0]];
4146 let y = source[[row, 1]];
4147 data[[row, 0]] = dilation * (cos_theta * x - sin_theta * y);
4148 data[[row, 1]] = dilation * (sin_theta * x + cos_theta * y);
4149 }
4150 let input_scale = estimate_isotropic_scale(data.view()).expect("isotropic input scale");
4151 let mut centers = data.clone();
4152 input_scale.standardize(&mut centers);
4153 let spec = TermCollectionSpec {
4154 linear_terms: Vec::new(),
4155 random_effect_terms: Vec::new(),
4156 smooth_terms: vec![SmoothTermSpec {
4157 frozen_parametric_residualization: None,
4158 name: "matern".to_string(),
4159 basis: SmoothBasisSpec::Matern {
4160 feature_cols: vec![0, 1],
4161 spec: MaternBasisSpec {
4162 periodic: None,
4163 center_strategy: CenterStrategy::UserProvided(centers),
4164 length_scale: crate::basis::MaternLengthScale::fixed(1.0),
4165 nu: MaternNu::FiveHalves,
4166 include_intercept: false,
4167 double_penalty: true,
4168 identifiability: MaternIdentifiability::CenterSumToZero,
4169 aniso_log_scales: None,
4170 },
4171 input_scale: Some(input_scale),
4172 },
4173 shape: ShapeConstraint::None,
4174 joint_null_rotation: None,
4175 }],
4176 };
4177 spatial_term_psi_bounds(
4178 data.view(),
4179 &spec,
4180 0,
4181 &SpatialLengthScaleOptimizationOptions::default(),
4182 )
4183 .expect("finite spatial ψ bounds")
4184 }
4185
4186 fn assert_close(left: f64, right: f64) {
4187 assert!(
4188 (left - right).abs() <= 1e-12,
4189 "coordinate-equivalent bounds differ: left={left:.16e}, right={right:.16e}"
4190 );
4191 }
4192
4193 #[test]
4207 fn psi_search_box_contains_the_incumbent_length_scale_2454() {
4208 let source = array![
4209 [-1.7, -0.4],
4210 [-1.1, 0.8],
4211 [-0.2, -1.3],
4212 [0.5, 1.6],
4213 [1.4, -0.7],
4214 [2.1, 0.5],
4215 ];
4216 let options = SpatialLengthScaleOptimizationOptions::default();
4217 let box_for = |length_scale: f64| -> ((f64, f64), (f64, f64)) {
4218 let input_scale =
4219 estimate_isotropic_scale(source.view()).expect("isotropic input scale");
4220 let mut centers = source.clone();
4221 input_scale.standardize(&mut centers);
4222 let spec = TermCollectionSpec {
4223 linear_terms: Vec::new(),
4224 random_effect_terms: Vec::new(),
4225 smooth_terms: vec![SmoothTermSpec {
4226 frozen_parametric_residualization: None,
4227 name: "matern".to_string(),
4228 basis: SmoothBasisSpec::Matern {
4229 feature_cols: vec![0, 1],
4230 spec: MaternBasisSpec {
4231 periodic: None,
4232 center_strategy: CenterStrategy::UserProvided(centers),
4233 length_scale: crate::basis::MaternLengthScale::fixed(length_scale),
4234 nu: MaternNu::FiveHalves,
4235 include_intercept: false,
4236 double_penalty: true,
4237 identifiability: MaternIdentifiability::CenterSumToZero,
4238 aniso_log_scales: None,
4239 },
4240 input_scale: Some(input_scale),
4241 },
4242 shape: ShapeConstraint::None,
4243 joint_null_rotation: None,
4244 }],
4245 };
4246 let geometry = spatial_term_psi_bounds(source.view(), &spec, 0, &options)
4247 .expect("finite geometry window");
4248 let search = spatial_term_psi_search_box(source.view(), &spec, 0, &options)
4249 .expect("finite search box");
4250 (geometry, search)
4251 };
4252
4253 let far = 1.0e3_f64;
4257 let (geometry, search) = box_for(far);
4258 let psi_far = -far.ln();
4259 assert!(
4260 psi_far < geometry.0,
4261 "fixture must place the incumbent OUTSIDE the geometry window, got \
4262 psi={psi_far} against [{}, {}]",
4263 geometry.0,
4264 geometry.1
4265 );
4266 assert!(
4267 search.0 <= psi_far && psi_far <= search.1,
4268 "the search box [{}, {}] must contain the incumbent psi={psi_far}; a seed \
4269 the box excludes is projected onto its edge and the optimum is then taken \
4270 over a set that does not contain the point it is graded against (#2454)",
4271 search.0,
4272 search.1
4273 );
4274 assert!(
4275 search.1 == geometry.1 && search.0 <= geometry.0,
4276 "the search box must be the geometry window WIDENED, never narrowed: \
4277 geometry=[{}, {}] search=[{}, {}]",
4278 geometry.0,
4279 geometry.1,
4280 search.0,
4281 search.1
4282 );
4283
4284 let (geometry_mid, search_mid) = box_for((-0.5 * (geometry.0 + geometry.1)).exp());
4286 assert!(
4287 search_mid == geometry_mid,
4288 "an incumbent inside the geometry window must not move the search box: \
4289 geometry=[{}, {}] search=[{}, {}]",
4290 geometry_mid.0,
4291 geometry_mid.1,
4292 search_mid.0,
4293 search_mid.1
4294 );
4295 }
4296
4297 #[test]
4298 fn standardized_center_bounds_return_to_original_units_under_rotation_and_scaling() {
4299 let base = frozen_matern_bounds(0.0, 1.0);
4300 let rotated = frozen_matern_bounds(0.61, 1.0);
4301 assert_close(rotated.0, base.0);
4302 assert_close(rotated.1, base.1);
4303
4304 let dilation = 4.0_f64;
4305 let rotated_scaled = frozen_matern_bounds(0.61, dilation);
4306 let expected_shift = dilation.ln();
4307 assert_close(rotated_scaled.0, base.0 - expected_shift);
4308 assert_close(rotated_scaled.1, base.1 - expected_shift);
4309 }
4310}
4311
4312pub fn spatial_term_psi_seed(
4316 data: ArrayView2<'_, f64>,
4317 spec: &TermCollectionSpec,
4318 term_idx: usize,
4319 options: &SpatialLengthScaleOptimizationOptions,
4320) -> Result<Option<f64>, BasisError> {
4321 if get_spatial_length_scale(spec, term_idx).is_some() {
4322 return Ok(None); }
4324 let (psi_lo, psi_hi) = spatial_term_psi_bounds(data, spec, term_idx, options)?;
4325 Ok(Some(0.5 * (psi_lo + psi_hi)))
4326}
4327
4328pub fn spatial_term_psi_to_length_scale_and_aniso(psi: &[f64]) -> (Option<f64>, Option<Vec<f64>>) {
4329 if psi.len() <= 1 {
4330 (Some((-psi.first().copied().unwrap_or(0.0)).exp()), None)
4331 } else {
4332 let psi_bar = psi.iter().sum::<f64>() / psi.len() as f64;
4333 (
4334 Some((-psi_bar).exp()),
4335 Some(psi.iter().map(|&value| value - psi_bar).collect()),
4336 )
4337 }
4338}
4339
4340pub fn get_spatial_aniso_log_scales(
4342 spec: &TermCollectionSpec,
4343 term_idx: usize,
4344) -> Option<Vec<f64>> {
4345 spec.smooth_terms
4346 .get(term_idx)
4347 .and_then(|term| match &term.basis {
4348 SmoothBasisSpec::Matern { spec, .. } => spec.aniso_log_scales.clone(),
4349 SmoothBasisSpec::Duchon { spec, .. } => spec.aniso_log_scales.clone(),
4350 _ => None,
4351 })
4352}
4353
4354pub fn response_aware_axis_contrasts(
4374 x: ndarray::ArrayView2<'_, f64>,
4375 y: ndarray::ArrayView1<'_, f64>,
4376) -> Option<Vec<f64>> {
4377 let n = x.nrows();
4378 let d = x.ncols();
4379 if d <= 1 || n < 4 || y.len() != n {
4380 return None;
4381 }
4382 if x.iter().any(|v| !v.is_finite()) || y.iter().any(|v| !v.is_finite()) {
4383 return None;
4384 }
4385 let mut scores = Vec::with_capacity(d);
4386 for a in 0..d {
4387 let mut order: Vec<usize> = (0..n).collect();
4388 let col = x.column(a);
4389 order.sort_by(|&i, &j| {
4390 col[i]
4391 .partial_cmp(&col[j])
4392 .unwrap_or(std::cmp::Ordering::Equal)
4393 });
4394 let mut tv = 0.0_f64;
4395 for w in order.windows(2) {
4396 let diff = y[w[1]] - y[w[0]];
4397 tv += diff * diff;
4398 }
4399 scores.push(-0.5 * (tv + 1e-12).ln());
4401 }
4402 if scores.iter().any(|v| !v.is_finite()) {
4403 return None;
4404 }
4405 let mean = scores.iter().sum::<f64>() / d as f64;
4406 let centered: Vec<f64> = scores.iter().map(|&s| s - mean).collect();
4407 if centered.iter().all(|&v| v.abs() < 1e-9) {
4410 return None;
4411 }
4412 Some(centered)
4413}
4414
4415pub fn apply_response_aware_anisotropy_seed(
4424 data: ArrayView2<'_, f64>,
4425 y: ndarray::ArrayView1<'_, f64>,
4426 spec: &mut TermCollectionSpec,
4427 spatial_terms: &[usize],
4428) {
4429 const MAX_NUDGE: f64 = std::f64::consts::LN_2;
4434 for &term_idx in spatial_terms {
4435 let Some(current_eta) = get_spatial_aniso_log_scales(spec, term_idx) else {
4436 continue;
4437 };
4438 let d = current_eta.len();
4439 if d <= 1 {
4440 continue;
4441 }
4442 let Some(term) = spec.smooth_terms.get(term_idx) else {
4443 continue;
4444 };
4445 let feature_cols = term.basis.structural_feature_cols();
4446 if feature_cols.len() != d {
4447 continue;
4448 }
4449 let Ok(x) = select_columns(data, &feature_cols) else {
4450 continue;
4451 };
4452 let Some(contrast) = response_aware_axis_contrasts(x.view(), y) else {
4453 continue;
4454 };
4455 let nudged: Vec<f64> = current_eta
4456 .iter()
4457 .zip(contrast.iter())
4458 .map(|(&eta_a, &c_a)| eta_a + c_a.clamp(-MAX_NUDGE, MAX_NUDGE))
4459 .collect();
4460 if let Err(err) = set_spatial_aniso_log_scales(spec, term_idx, nudged) {
4463 log::debug!(
4464 "[spatial-kappa] response-aware anisotropy seed skipped for term {term_idx}: {err}"
4465 );
4466 }
4467 }
4468}
4469
4470pub fn get_spatial_feature_dim(spec: &TermCollectionSpec, term_idx: usize) -> Option<usize> {
4472 spec.smooth_terms
4473 .get(term_idx)
4474 .and_then(|term| match &term.basis {
4475 SmoothBasisSpec::ThinPlate { feature_cols, .. } => Some(feature_cols.len()),
4476 SmoothBasisSpec::Matern { feature_cols, .. } => Some(feature_cols.len()),
4477 SmoothBasisSpec::Duchon { feature_cols, .. } => Some(feature_cols.len()),
4478 _ => None,
4479 })
4480}
4481
4482pub fn log_spatial_aniso_scales(spec: &TermCollectionSpec) {
4489 for (term_idx, term) in spec.smooth_terms.iter().enumerate() {
4490 let (aniso, length_scale) = match &term.basis {
4491 SmoothBasisSpec::Matern { spec, .. } => {
4492 (spec.aniso_log_scales.as_ref(), spec.length_scale.resolved())
4493 }
4494 SmoothBasisSpec::Duchon { spec, .. } => {
4495 (spec.aniso_log_scales.as_ref(), spec.length_scale)
4496 }
4497 _ => (None, None),
4498 };
4499 let Some(eta) = aniso else { continue };
4500 if eta.is_empty() {
4501 continue;
4502 }
4503 let mut lines = match length_scale {
4504 Some(ls) => format!(
4505 "[spatial-kappa] term {} (\"{}\"): anisotropic length scales optimized (global length_scale={:.4})",
4506 term_idx, term.name, ls
4507 ),
4508 None => format!(
4509 "[spatial-kappa] term {} (\"{}\"): pure Duchon shape anisotropy optimized",
4510 term_idx, term.name
4511 ),
4512 };
4513 for (a, &eta_a) in eta.iter().enumerate() {
4514 if let Some(ls) = length_scale {
4515 let length_a = ls * (-eta_a).exp();
4516 let kappa_a = (1.0 / ls) * eta_a.exp();
4517 lines.push_str(&format!(
4518 "\n axis {}: eta={:+.4}, length={:.4}, kappa={:.4}",
4519 a, eta_a, length_a, kappa_a
4520 ));
4521 } else {
4522 lines.push_str(&format!("\n axis {}: eta={:+.4}", a, eta_a));
4523 }
4524 }
4525 log::info!("{}", lines);
4526 }
4527}
4528
4529pub fn set_spatial_aniso_log_scales(
4531 spec: &mut TermCollectionSpec,
4532 term_idx: usize,
4533 eta: Vec<f64>,
4534) -> Result<(), EstimationError> {
4535 let eta = center_aniso_log_scales(&eta);
4536 let Some(term) = spec.smooth_terms.get_mut(term_idx) else {
4537 crate::bail_invalid_estim!("spatial aniso_log_scales term index {term_idx} out of range");
4538 };
4539 match &mut term.basis {
4540 SmoothBasisSpec::Matern { spec, .. } => {
4541 spec.aniso_log_scales = Some(eta);
4542 Ok(())
4543 }
4544 SmoothBasisSpec::Duchon { spec, .. } => {
4545 spec.aniso_log_scales = Some(eta);
4546 Ok(())
4547 }
4548 _ => Err(EstimationError::InvalidInput(format!(
4549 "term '{}' does not support aniso_log_scales",
4550 term.name
4551 ))),
4552 }
4553}
4554
4555pub fn sync_aniso_contrasts_from_metadata(spec: &mut TermCollectionSpec, design: &SmoothDesign) {
4562 for (term_idx, term) in design.terms.iter().enumerate() {
4563 let meta_aniso = match &term.metadata {
4564 BasisMetadata::Matern {
4565 aniso_log_scales, ..
4566 } => aniso_log_scales.clone(),
4567 BasisMetadata::Duchon {
4568 aniso_log_scales, ..
4569 } => aniso_log_scales.clone(),
4570 _ => None,
4571 };
4572 if let Some(eta) = meta_aniso
4573 && eta.len() > 1
4574 {
4575 if let Err(err) = set_spatial_aniso_log_scales(spec, term_idx, eta) {
4576 log::debug!(
4577 "term {term_idx}: anisotropic log-scale sync skipped, keeping the existing scales: {err}"
4578 );
4579 }
4580 }
4581 }
4582}
4583
4584#[derive(Debug, Clone)]
4585pub struct SpatialLengthScaleOptimizationOptions {
4586 pub enabled: bool,
4590 pub max_outer_iter: usize,
4592 pub rel_tol: f64,
4594 pub log_step: f64,
4596 pub min_length_scale: f64,
4598 pub max_length_scale: f64,
4600 pub pilot_subsample_threshold: usize,
4613}
4614
4615impl Default for SpatialLengthScaleOptimizationOptions {
4616 fn default() -> Self {
4617 Self {
4618 enabled: true,
4619 max_outer_iter: 80,
4620 rel_tol: 1e-4,
4621 log_step: std::f64::consts::LN_2,
4622 min_length_scale: 1e-3,
4623 max_length_scale: 1e3,
4624 pilot_subsample_threshold: 10_000,
4625 }
4626 }
4627}
4628
4629impl SpatialLengthScaleOptimizationOptions {
4630 pub fn validate(&self) -> Result<(), String> {
4648 if !self.min_length_scale.is_finite() || self.min_length_scale <= 0.0 {
4649 return Err(SmoothError::invalid_config(format!(
4650 "SpatialLengthScaleOptimizationOptions::min_length_scale must be > 0 and finite, got {}",
4651 self.min_length_scale
4652 ))
4653 .into());
4654 }
4655 if !self.max_length_scale.is_finite() || self.max_length_scale <= 0.0 {
4656 return Err(SmoothError::invalid_config(format!(
4657 "SpatialLengthScaleOptimizationOptions::max_length_scale must be > 0 and finite, got {}",
4658 self.max_length_scale
4659 ))
4660 .into());
4661 }
4662 if self.min_length_scale >= self.max_length_scale {
4663 return Err(SmoothError::invalid_config(format!(
4664 "SpatialLengthScaleOptimizationOptions requires min_length_scale < max_length_scale, got min={} max={}",
4665 self.min_length_scale, self.max_length_scale
4666 ))
4667 .into());
4668 }
4669 if !self.rel_tol.is_finite() || self.rel_tol <= 0.0 {
4670 return Err(SmoothError::invalid_config(format!(
4671 "SpatialLengthScaleOptimizationOptions::rel_tol must be > 0 and finite, got {}",
4672 self.rel_tol
4673 ))
4674 .into());
4675 }
4676 if !self.log_step.is_finite() || self.log_step <= 0.0 {
4677 return Err(SmoothError::invalid_config(format!(
4678 "SpatialLengthScaleOptimizationOptions::log_step must be > 0 and finite, got {}",
4679 self.log_step
4680 ))
4681 .into());
4682 }
4683 Ok(())
4684 }
4685}
4686
4687#[derive(Debug, Clone)]
4688pub struct RandomEffectBlock {
4689 pub name: String,
4690 pub group_ids: Vec<Option<usize>>,
4693 pub num_groups: usize,
4694 pub kept_levels: Vec<u64>,
4695}
4696
4697pub const BLOCK_SPARSE_ZERO_EPS: f64 = 1e-12;
4698
4699pub const BLOCK_SPARSE_MAX_DENSITY: f64 = 0.20;
4700
4701pub fn blocks_have_intrinsic_sparse_structure(blocks: &[DesignBlock]) -> bool {
4702 blocks
4703 .iter()
4704 .any(|block| matches!(block, DesignBlock::Sparse(_) | DesignBlock::RandomEffect(_)))
4705}
4706
4707pub fn sparse_compatible_block_nnz(block: &DesignBlock) -> Option<usize> {
4708 match block {
4709 DesignBlock::Intercept(n) => Some(*n),
4710 DesignBlock::RandomEffect(op) => {
4711 Some(op.group_ids.iter().filter(|gid| gid.is_some()).count())
4712 }
4713 DesignBlock::Sparse(sparse) => Some(sparse.val().len()),
4714 DesignBlock::Dense(dense) => dense.as_dense_ref().map(|matrix| {
4715 matrix
4716 .iter()
4717 .filter(|&&value| value.abs() > BLOCK_SPARSE_ZERO_EPS)
4718 .count()
4719 }),
4720 }
4721}
4722
4723pub fn try_build_sparse_design_from_blocks(
4724 blocks: &[DesignBlock],
4725) -> Result<Option<DesignMatrix>, BasisError> {
4726 if blocks.is_empty() {
4727 return Ok(None);
4728 }
4729 let nrows = blocks[0].nrows();
4730 let ncols: usize = blocks.iter().map(DesignBlock::ncols).sum();
4731 if nrows == 0 || ncols == 0 || ncols <= 32 {
4732 return Ok(None);
4733 }
4734
4735 let preserve_sparse_storage = blocks_have_intrinsic_sparse_structure(blocks);
4736 let sparse_nnz_limit = if preserve_sparse_storage {
4737 usize::MAX
4738 } else {
4739 let total_cells = nrows.saturating_mul(ncols);
4740 ((total_cells as f64) * BLOCK_SPARSE_MAX_DENSITY).floor() as usize
4741 };
4742 let mut nnz = 0usize;
4743 for block in blocks {
4744 let block_nnz = if let Some(block_nnz) = sparse_compatible_block_nnz(block) {
4745 block_nnz
4746 } else {
4747 return Ok(None);
4748 };
4749 nnz = nnz.saturating_add(block_nnz);
4750 if nnz > sparse_nnz_limit {
4751 return Ok(None);
4752 }
4753 }
4754
4755 let mut triplets = Vec::<Triplet<usize, usize, f64>>::with_capacity(nnz);
4756 let mut col_offset = 0usize;
4757 for block in blocks {
4758 match block {
4759 DesignBlock::Intercept(n) => {
4760 for row in 0..*n {
4761 triplets.push(Triplet::new(row, col_offset, 1.0));
4762 }
4763 }
4764 DesignBlock::RandomEffect(op) => {
4765 for (row, group_id) in op.group_ids.iter().enumerate() {
4766 if let Some(group) = group_id {
4767 triplets.push(Triplet::new(row, col_offset + group, 1.0));
4768 }
4769 }
4770 }
4771 DesignBlock::Sparse(sparse) => {
4772 let (symbolic, values) = sparse.parts();
4773 let col_ptr = symbolic.col_ptr();
4774 let row_idx = symbolic.row_idx();
4775 for col in 0..sparse.ncols() {
4776 for idx in col_ptr[col]..col_ptr[col + 1] {
4777 let value = values[idx];
4778 if value.abs() > BLOCK_SPARSE_ZERO_EPS {
4779 triplets.push(Triplet::new(row_idx[idx], col_offset + col, value));
4780 }
4781 }
4782 }
4783 }
4784 DesignBlock::Dense(dense) => {
4785 let matrix = dense.as_dense_ref().ok_or_else(|| {
4786 BasisError::InvalidInput(
4787 "sparse-compatible block assembly requires materialized dense blocks"
4788 .to_string(),
4789 )
4790 })?;
4791 for row in 0..matrix.nrows() {
4792 for col in 0..matrix.ncols() {
4793 let value = matrix[[row, col]];
4794 if value.abs() > BLOCK_SPARSE_ZERO_EPS {
4795 triplets.push(Triplet::new(row, col_offset + col, value));
4796 }
4797 }
4798 }
4799 }
4800 }
4801 col_offset += block.ncols();
4802 }
4803
4804 let sparse = SparseColMat::try_new_from_triplets(nrows, ncols, &triplets).map_err(|_| {
4805 BasisError::SparseCreation("failed to assemble sparse term-collection design".to_string())
4806 })?;
4807 Ok(Some(DesignMatrix::Sparse(
4808 gam_linalg::matrix::SparseDesignMatrix::new(sparse),
4809 )))
4810}
4811
4812pub fn assemble_term_collection_design_matrix(
4813 blocks: Vec<DesignBlock>,
4814) -> Result<DesignMatrix, BasisError> {
4815 if let Some(sparse) = try_build_sparse_design_from_blocks(&blocks)? {
4816 return Ok(sparse);
4817 }
4818 let block_op = BlockDesignOperator::new(blocks).map_err(|e| {
4819 BasisError::InvalidInput(format!("failed to build block design operator: {e}"))
4820 })?;
4821 Ok(DesignMatrix::Dense(
4822 gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(block_op)),
4823 ))
4824}
4825
4826pub fn select_columns(
4827 data: ArrayView2<'_, f64>,
4828 cols: &[usize],
4829) -> Result<Array2<f64>, BasisError> {
4830 let n = data.nrows();
4831 let p = data.ncols();
4832 for &c in cols {
4833 if c >= p {
4834 crate::bail_dim_basis!("feature column {c} is out of bounds for data with {p} columns");
4835 }
4836 }
4837 let mut out = Array2::<f64>::zeros((n, cols.len()));
4838 for (j, &c) in cols.iter().enumerate() {
4839 out.column_mut(j).assign(&data.column(c));
4840 }
4841 Ok(out)
4842}
4843
4844pub fn nonfinite_value_label(value: f64) -> &'static str {
4845 if value.is_nan() {
4846 "NaN"
4847 } else if value.is_sign_positive() {
4848 "+Inf"
4849 } else {
4850 "-Inf"
4851 }
4852}
4853
4854pub fn validate_term_feature_column_finite(
4855 data: ArrayView2<'_, f64>,
4856 term_kind: &str,
4857 term_name: &str,
4858 feature_col: usize,
4859) -> Result<(), BasisError> {
4860 let p = data.ncols();
4861 if feature_col >= p {
4862 crate::bail_dim_basis!(
4863 "{term_kind} term '{term_name}' feature column {feature_col} out of bounds for {p} columns"
4864 );
4865 }
4866 for (row, &value) in data.column(feature_col).iter().enumerate() {
4867 if !value.is_finite() {
4868 crate::bail_invalid_basis!(
4869 "{term_kind} term '{term_name}' feature column {feature_col} row {row} contains non-finite value {}",
4870 nonfinite_value_label(value)
4871 );
4872 }
4873 }
4874 Ok(())
4875}
4876
4877pub fn validate_smooth_terms_finite_inputs(
4878 data: ArrayView2<'_, f64>,
4879 terms: &[SmoothTermSpec],
4880) -> Result<(), BasisError> {
4881 for term in terms {
4882 for feature_col in smooth_term_feature_cols(term) {
4883 validate_term_feature_column_finite(data, "smooth", &term.name, feature_col)?;
4884 }
4885 }
4886 Ok(())
4887}
4888
4889pub fn validate_term_collection_finite_inputs(
4890 data: ArrayView2<'_, f64>,
4891 spec: &TermCollectionSpec,
4892) -> Result<(), BasisError> {
4893 for term in &spec.linear_terms {
4894 validate_term_feature_column_finite(data, "linear", &term.name, term.feature_col)?;
4895 }
4896 for term in &spec.random_effect_terms {
4897 validate_term_feature_column_finite(data, "random-effect", &term.name, term.feature_col)?;
4898 }
4899 validate_smooth_terms_finite_inputs(data, &spec.smooth_terms)
4900}
4901
4902#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
4903pub struct JointSpatialCenterGroupKey {
4904 feature_cols: Vec<usize>,
4905 strategy_kind: CenterStrategyKind,
4906 strategy_aux: usize,
4907 requested_num_centers: usize,
4908 input_scale_bits: Option<u64>,
4909}
4910
4911pub fn spatial_term_min_center_count(term: &SmoothTermSpec) -> usize {
4912 match &term.basis {
4913 SmoothBasisSpec::ThinPlate { feature_cols, .. } => feature_cols.len() + 1,
4914 SmoothBasisSpec::Duchon {
4915 feature_cols, spec, ..
4916 } => match spec.nullspace_order {
4917 crate::basis::DuchonNullspaceOrder::Zero => 1,
4918 crate::basis::DuchonNullspaceOrder::Linear => feature_cols.len() + 1,
4919 crate::basis::DuchonNullspaceOrder::Degree(degree) => {
4920 crate::basis::duchon_nullspace_dimension(feature_cols.len(), degree)
4921 }
4922 },
4923 SmoothBasisSpec::Matern { .. } => 1,
4924 _ => 1,
4925 }
4926}
4927
4928pub fn spatial_term_group_key(term: &SmoothTermSpec) -> Option<JointSpatialCenterGroupKey> {
4929 let (feature_cols, strategy, input_scale) = match &term.basis {
4930 SmoothBasisSpec::ThinPlate {
4931 feature_cols,
4932 spec,
4933 input_scale,
4934 } => (feature_cols, &spec.center_strategy, *input_scale),
4935 SmoothBasisSpec::Matern {
4936 feature_cols,
4937 spec,
4938 input_scale,
4939 } => (feature_cols, &spec.center_strategy, *input_scale),
4940 SmoothBasisSpec::Duchon {
4941 feature_cols,
4942 spec,
4943 input_scale,
4944 } => (feature_cols, &spec.center_strategy, *input_scale),
4945 _ => return None,
4946 };
4947 let strategy_kind = center_strategy_kind(strategy);
4948 let strategy_aux = match strategy {
4949 CenterStrategy::Auto(inner) => match inner.as_ref() {
4950 CenterStrategy::KMeans { max_iter, .. } => *max_iter,
4951 CenterStrategy::UniformGrid { points_per_dim } => *points_per_dim,
4952 _ => 0,
4953 },
4954 CenterStrategy::KMeans { max_iter, .. } => *max_iter,
4955 CenterStrategy::UniformGrid { points_per_dim } => *points_per_dim,
4956 _ => 0,
4957 };
4958 Some(JointSpatialCenterGroupKey {
4959 feature_cols: feature_cols.clone(),
4960 strategy_kind,
4961 strategy_aux,
4962 requested_num_centers: strategy.planned_num_centers(feature_cols.len()),
4963 input_scale_bits: input_scale.map(crate::IsotropicScale::to_bits),
4964 })
4965}
4966
4967pub fn spatial_term_center_strategy(term: &SmoothTermSpec) -> Option<&CenterStrategy> {
4968 match &term.basis {
4969 SmoothBasisSpec::ThinPlate { spec, .. } => Some(&spec.center_strategy),
4970 SmoothBasisSpec::Matern { spec, .. } => Some(&spec.center_strategy),
4971 SmoothBasisSpec::Duchon { spec, .. } => Some(&spec.center_strategy),
4972 _ => None,
4973 }
4974}
4975
4976pub fn set_spatial_term_centers(
4977 term: &mut SmoothTermSpec,
4978 centers: Array2<f64>,
4979) -> Result<(), BasisError> {
4980 match &mut term.basis {
4981 SmoothBasisSpec::ThinPlate { spec, .. } => {
4982 spec.center_strategy = CenterStrategy::UserProvided(centers);
4983 Ok(())
4984 }
4985 SmoothBasisSpec::Matern { spec, .. } => {
4986 spec.center_strategy = CenterStrategy::UserProvided(centers);
4987 Ok(())
4988 }
4989 SmoothBasisSpec::Duchon { spec, .. } => {
4990 spec.center_strategy = CenterStrategy::UserProvided(centers);
4991 Ok(())
4992 }
4993 _ => Err(BasisError::InvalidInput(format!(
4994 "term '{}' does not support spatial center planning",
4995 term.name
4996 ))),
4997 }
4998}
4999
5000pub fn standardized_spatial_term_data(
5001 data: ArrayView2<'_, f64>,
5002 term: &SmoothTermSpec,
5003) -> Result<Array2<f64>, BasisError> {
5004 let (feature_cols, input_scale) = match &term.basis {
5005 SmoothBasisSpec::ThinPlate {
5006 feature_cols,
5007 input_scale,
5008 ..
5009 }
5010 | SmoothBasisSpec::Matern {
5011 feature_cols,
5012 input_scale,
5013 ..
5014 }
5015 | SmoothBasisSpec::Duchon {
5016 feature_cols,
5017 input_scale,
5018 ..
5019 } => (feature_cols, *input_scale),
5020 _ => {
5021 crate::bail_invalid_basis!("term '{}' is not a spatial smooth", term.name);
5022 }
5023 };
5024 let mut x = select_columns(data, feature_cols)?;
5025 input_scale
5026 .map_or_else(|| estimate_isotropic_scale(x.view()), Ok)?
5027 .standardize(&mut x);
5028 Ok(x)
5029}
5030
5031pub fn plan_joint_spatial_centers_for_term_blocks(
5032 data: ArrayView2<'_, f64>,
5033 term_blocks: &[Vec<SmoothTermSpec>],
5034) -> Result<Vec<Vec<SmoothTermSpec>>, BasisError> {
5035 let mut planned_blocks = term_blocks.to_vec();
5036 let n = data.nrows();
5037 let mut groups: BTreeMap<JointSpatialCenterGroupKey, Vec<(usize, usize)>> = BTreeMap::new();
5038
5039 for (block_idx, terms) in planned_blocks.iter().enumerate() {
5040 for (term_idx, term) in terms.iter().enumerate() {
5041 let Some(strategy) = spatial_term_center_strategy(term) else {
5042 continue;
5043 };
5044 if !center_strategy_is_auto(strategy) {
5045 continue;
5046 }
5047 let Some(group_key) = spatial_term_group_key(term) else {
5048 continue;
5049 };
5050 if !matches!(
5051 group_key.strategy_kind,
5052 CenterStrategyKind::EqualMass
5053 | CenterStrategyKind::EqualMassCovarRepresentative
5054 | CenterStrategyKind::FarthestPoint
5055 | CenterStrategyKind::KMeans
5056 | CenterStrategyKind::UniformGrid
5057 ) {
5058 continue;
5059 }
5060 groups
5061 .entry(group_key)
5062 .or_default()
5063 .push((block_idx, term_idx));
5064 }
5065 }
5066
5067 for (group_key, members) in groups {
5068 if members.len() < 2 {
5069 continue;
5070 }
5071 let min_required = members
5072 .iter()
5073 .map(|&(block_idx, term_idx)| {
5074 spatial_term_min_center_count(&planned_blocks[block_idx][term_idx])
5075 })
5076 .max()
5077 .unwrap_or(1);
5078 let joint_centers = group_key
5079 .requested_num_centers
5080 .max(min_required)
5081 .min(n.max(1));
5082 let (first_block_idx, first_term_idx) = members[0];
5083 let prototype = &planned_blocks[first_block_idx][first_term_idx];
5084 let standardized = standardized_spatial_term_data(data, prototype)?;
5085 let strategy = spatial_term_center_strategy(prototype).ok_or_else(|| {
5086 BasisError::InvalidInput(format!(
5087 "term '{}' lost its spatial center strategy during joint planning",
5088 prototype.name
5089 ))
5090 })?;
5091 let joint_strategy = center_strategy_with_num_centers(
5092 strategy,
5093 joint_centers,
5094 group_key.feature_cols.len(),
5095 )?;
5096 let shared_centers = select_centers_by_strategy(standardized.view(), &joint_strategy)?;
5097 log::info!(
5098 "sharing {} spatial centers across {} smooth terms over columns {:?} (requested {} centers)",
5099 shared_centers.nrows(),
5100 members.len(),
5101 group_key.feature_cols,
5102 group_key.requested_num_centers,
5103 );
5104 for (block_idx, term_idx) in members {
5105 set_spatial_term_centers(
5106 &mut planned_blocks[block_idx][term_idx],
5107 shared_centers.clone(),
5108 )?;
5109 }
5110 }
5111
5112 for block in planned_blocks.iter_mut() {
5117 for term in block.iter_mut() {
5118 auto_init_length_scale_in_place(data, term);
5119 }
5120 }
5121
5122 Ok(planned_blocks)
5123}
5124
5125const AUTO_LENGTH_SCALE_FLOOR: f64 = 1e-6;
5128
5129fn feature_columns_max_range(data: ArrayView2<'_, f64>, feature_cols: &[usize]) -> Option<f64> {
5132 let mut max_range = 0.0_f64;
5133 for &c in feature_cols {
5134 if c >= data.ncols() {
5135 continue;
5136 }
5137 let col = data.column(c);
5138 let mut lo = f64::INFINITY;
5139 let mut hi = f64::NEG_INFINITY;
5140 for &v in col.iter() {
5141 if v.is_finite() {
5142 if v < lo {
5143 lo = v;
5144 }
5145 if v > hi {
5146 hi = v;
5147 }
5148 }
5149 }
5150 if hi > lo {
5151 let r = hi - lo;
5152 if r > max_range {
5153 max_range = r;
5154 }
5155 }
5156 }
5157 if max_range.is_finite() && max_range > 0.0 {
5158 Some(max_range)
5159 } else {
5160 None
5161 }
5162}
5163
5164fn feature_columns_rotation_invariant_range(
5175 data: ArrayView2<'_, f64>,
5176 feature_cols: &[usize],
5177) -> Option<f64> {
5178 let cols: Vec<usize> = feature_cols
5179 .iter()
5180 .copied()
5181 .filter(|&c| c < data.ncols())
5182 .collect();
5183 if cols.is_empty() {
5184 return None;
5185 }
5186 let mut points: Vec<Vec<f64>> = data
5187 .rows()
5188 .into_iter()
5189 .filter_map(|row| {
5190 let point: Vec<f64> = cols.iter().map(|&column| row[column]).collect();
5191 point.iter().all(|value| value.is_finite()).then_some(point)
5192 })
5193 .collect();
5194 if points.is_empty() {
5195 return None;
5196 }
5197 points.sort_by(|left, right| {
5198 left.iter()
5199 .zip(right)
5200 .find_map(|(a, b)| {
5201 let ordering = a.total_cmp(b);
5202 ordering.is_ne().then_some(ordering)
5203 })
5204 .unwrap_or(std::cmp::Ordering::Equal)
5205 });
5206
5207 let dimensions = cols.len();
5208 let count = points.len() as f64;
5209 let mut centroid = vec![0.0_f64; dimensions];
5210 for point in &points {
5211 for (coordinate, value) in centroid.iter_mut().zip(point) {
5212 *coordinate += *value;
5213 }
5214 }
5215 for coordinate in &mut centroid {
5216 *coordinate /= count;
5217 }
5218
5219 let mut covariance = Array2::<f64>::zeros((dimensions, dimensions));
5220 for point in &points {
5221 for row in 0..dimensions {
5222 let centered_row = point[row] - centroid[row];
5223 for column in 0..=row {
5224 covariance[[row, column]] += centered_row * (point[column] - centroid[column]);
5225 }
5226 }
5227 }
5228 for row in 0..dimensions {
5229 for column in 0..=row {
5230 let value = covariance[[row, column]] / count;
5231 covariance[[row, column]] = value;
5232 covariance[[column, row]] = value;
5233 }
5234 }
5235
5236 use gam_linalg::faer_ndarray::FaerEigh;
5237 let (eigenvalues, _) = covariance
5238 .eigh(faer::Side::Lower)
5239 .expect("finite covariance must have a symmetric eigendecomposition");
5240 let leading_variance = eigenvalues[eigenvalues.len() - 1];
5241 let extent = (12.0 * leading_variance).sqrt();
5242 if extent.is_finite() && extent > 0.0 {
5243 Some(extent)
5244 } else {
5245 None
5246 }
5247}
5248
5249pub fn auto_initial_length_scale(data: ArrayView2<'_, f64>, feature_cols: &[usize]) -> f64 {
5256 let n = data.nrows();
5257 if n == 0 || feature_cols.is_empty() {
5258 return 1.0;
5259 }
5260 let Some(max_range) = feature_columns_max_range(data, feature_cols) else {
5261 return 1.0;
5262 };
5263 let init = max_range / (n as f64).sqrt();
5264 init.max(AUTO_LENGTH_SCALE_FLOOR).min(max_range)
5265}
5266
5267pub fn auto_initial_length_scale_for_centers(
5290 data: ArrayView2<'_, f64>,
5291 feature_cols: &[usize],
5292 num_centers: usize,
5293) -> f64 {
5294 let n = data.nrows();
5295 if n == 0 || feature_cols.is_empty() {
5296 return 1.0;
5297 }
5298 let Some(max_range) = feature_columns_rotation_invariant_range(data, feature_cols) else {
5309 return 1.0;
5310 };
5311 let resolution_points = n.max(num_centers).max(1) as f64;
5317 let spacing = max_range / resolution_points.sqrt();
5318 spacing.max(AUTO_LENGTH_SCALE_FLOOR).min(max_range)
5319}
5320
5321pub fn matern_low_rank_center_resolution_length_scale(
5331 data: ArrayView2<'_, f64>,
5332 feature_cols: &[usize],
5333 num_centers: usize,
5334) -> Option<f64> {
5335 if data.nrows() == 0 || feature_cols.is_empty() || num_centers == 0 {
5336 return None;
5337 }
5338 let extent = feature_columns_rotation_invariant_range(data, feature_cols)?;
5339 let length_scale = extent / (num_centers as f64).sqrt();
5340 Some(length_scale.max(AUTO_LENGTH_SCALE_FLOOR).min(extent))
5341}
5342
5343pub fn auto_initial_length_scale_for_low_rank_centers(
5353 data: ArrayView2<'_, f64>,
5354 feature_cols: &[usize],
5355 num_centers: usize,
5356) -> f64 {
5357 if data.nrows() == 0 || feature_cols.is_empty() {
5358 return 1.0;
5359 }
5360 let Some(max_range) = feature_columns_max_range(data, feature_cols) else {
5361 return 1.0;
5362 };
5363 let resolution_points = num_centers.max(1) as f64;
5364 let spacing = max_range / resolution_points.sqrt();
5365 spacing.max(AUTO_LENGTH_SCALE_FLOOR).min(max_range)
5366}
5367
5368fn center_strategy_requested_count(strategy: &CenterStrategy) -> Option<usize> {
5371 match strategy {
5372 CenterStrategy::Auto(inner) => center_strategy_requested_count(inner),
5373 CenterStrategy::DuchonSpectral { knots, .. } => center_strategy_requested_count(knots),
5374 CenterStrategy::UserProvided(centers) => Some(centers.nrows()),
5375 CenterStrategy::EqualMass { num_centers }
5376 | CenterStrategy::EqualMassCovarRepresentative { num_centers }
5377 | CenterStrategy::FarthestPoint { num_centers }
5378 | CenterStrategy::KMeans { num_centers, .. } => Some(*num_centers),
5379 CenterStrategy::UniformGrid { .. } => None,
5380 }
5381}
5382
5383pub fn auto_init_length_scale_in_place(data: ArrayView2<'_, f64>, term: &mut SmoothTermSpec) {
5387 auto_init_length_scale_in_basis(data, &mut term.basis);
5388}
5389
5390pub fn auto_init_length_scale_in_basis(data: ArrayView2<'_, f64>, basis: &mut SmoothBasisSpec) {
5402 match basis {
5403 SmoothBasisSpec::Matern {
5404 feature_cols, spec, ..
5405 } => {
5406 if spec.length_scale.resolved().is_none() {
5407 let resolved = match center_strategy_requested_count(&spec.center_strategy) {
5416 Some(k) => auto_initial_length_scale_for_centers(data, feature_cols, k),
5417 None => auto_initial_length_scale(data, feature_cols),
5418 };
5419 spec.length_scale.resolve_auto_once(resolved);
5420 }
5421 }
5422 SmoothBasisSpec::ThinPlate {
5423 feature_cols, spec, ..
5424 } => {
5425 if spec.length_scale == 0.0 {
5426 spec.length_scale = match center_strategy_requested_count(&spec.center_strategy) {
5427 Some(k) => {
5428 auto_initial_length_scale_for_low_rank_centers(data, feature_cols, k)
5429 }
5430 None => auto_initial_length_scale(data, feature_cols),
5431 };
5432 }
5433 }
5434 SmoothBasisSpec::ByVariable { inner, .. }
5435 | SmoothBasisSpec::FactorSumToZero { inner, .. } => {
5436 auto_init_length_scale_in_basis(data, inner);
5437 }
5438 SmoothBasisSpec::BySmooth { smooth, .. } => {
5439 auto_init_length_scale_in_basis(data, smooth);
5440 }
5441 SmoothBasisSpec::BSpline1D { .. }
5448 | SmoothBasisSpec::FactorSmooth { .. }
5449 | SmoothBasisSpec::Sphere { .. }
5450 | SmoothBasisSpec::ConstantCurvature { .. }
5451 | SmoothBasisSpec::MeasureJet { .. }
5452 | SmoothBasisSpec::Duchon { .. }
5453 | SmoothBasisSpec::Pca { .. }
5454 | SmoothBasisSpec::TensorBSpline { .. } => {}
5455 }
5456}
5457
5458impl LinearFitConditioning {
5459 pub fn from_columns(design: &TermCollectionDesign, selected_cols: &[usize]) -> Self {
5460 const SCALE_EPS: f64 = 1e-12;
5461 let n = design.design.nrows();
5462 let p = design.design.ncols();
5463 let mut columns = Vec::with_capacity(selected_cols.len());
5464 if n == 0 || selected_cols.is_empty() {
5465 return Self {
5466 intercept_idx: design.intercept_range.start,
5467 columns,
5468 };
5469 }
5470 let chunk_rows = gam_linalg::utils::row_chunk_for_byte_budget(n, p);
5471 let mut sums = vec![0.0_f64; selected_cols.len()];
5477 for start in (0..n).step_by(chunk_rows) {
5478 let end = (start + chunk_rows).min(n);
5479 let chunk = design
5480 .design
5481 .try_row_chunk(start..end)
5482 .expect("LinearFitConditioning::from_columns row chunk failed");
5483 for (k, &col_idx) in selected_cols.iter().enumerate() {
5484 let column = chunk.column(col_idx);
5485 for &v in column.iter() {
5486 sums[k] += v;
5487 }
5488 }
5489 }
5490 let inv_n = 1.0_f64 / n as f64;
5491 let means: Vec<f64> = sums.iter().map(|&s| s * inv_n).collect();
5492 let mut sq_devs = vec![0.0_f64; selected_cols.len()];
5493 for start in (0..n).step_by(chunk_rows) {
5494 let end = (start + chunk_rows).min(n);
5495 let chunk = design
5496 .design
5497 .try_row_chunk(start..end)
5498 .expect("LinearFitConditioning::from_columns row chunk failed");
5499 for (k, &col_idx) in selected_cols.iter().enumerate() {
5500 let mean_k = means[k];
5501 let column = chunk.column(col_idx);
5502 for &v in column.iter() {
5503 let d = v - mean_k;
5504 sq_devs[k] += d * d;
5505 }
5506 }
5507 }
5508 for (k, &col_idx) in selected_cols.iter().enumerate() {
5509 let mean = means[k];
5510 let var = sq_devs[k] * inv_n;
5511 let (mean, scale) = if var.is_finite() && var > SCALE_EPS * SCALE_EPS {
5512 (mean, var.sqrt())
5513 } else {
5514 (0.0, 1.0)
5517 };
5518 columns.push(LinearColumnConditioning {
5519 col_idx,
5520 mean,
5521 scale,
5522 });
5523 }
5524 Self {
5525 intercept_idx: design.intercept_range.start,
5526 columns,
5527 }
5528 }
5529
5530 pub fn apply_to_design(&self, design: &Array2<f64>) -> Array2<f64> {
5531 let mut out = design.clone();
5532 for col in &self.columns {
5533 {
5534 let mut dst = out.column_mut(col.col_idx);
5535 dst -= col.mean;
5536 }
5537 if col.scale != 1.0 {
5538 out.column_mut(col.col_idx).mapv_inplace(|v| v / col.scale);
5539 }
5540 }
5541 out
5542 }
5543
5544 fn transform_matrix_columnswith_a(&self, mat: &Array2<f64>) -> Array2<f64> {
5545 let mut out = mat.clone();
5546 let intercept = self.intercept_idx;
5547 for col in &self.columns {
5548 let intercept_col = out.column(intercept).to_owned();
5549 let mut target = out.column_mut(col.col_idx);
5550 target -= &(intercept_col * col.mean);
5551 if col.scale != 1.0 {
5552 target.mapv_inplace(|v| v / col.scale);
5553 }
5554 }
5555 out
5556 }
5557
5558 fn transform_matrixrowswith_a_transpose(&self, mat: &Array2<f64>) -> Array2<f64> {
5559 let mut out = mat.clone();
5560 let intercept = self.intercept_idx;
5561 for col in &self.columns {
5562 let interceptrow = out.row(intercept).to_owned();
5563 let mut target = out.row_mut(col.col_idx);
5564 target -= &(interceptrow * col.mean);
5565 if col.scale != 1.0 {
5566 target.mapv_inplace(|v| v / col.scale);
5567 }
5568 }
5569 out
5570 }
5571
5572 fn left_multiply_by_m_inv_transpose(&self, mat_internal: &Array2<f64>) -> Array2<f64> {
5577 let mut out = mat_internal.clone();
5578 let intercept = self.intercept_idx;
5579 let interceptrow_snapshot = mat_internal.row(intercept).to_owned();
5580 for col in &self.columns {
5581 if col.scale != 1.0 {
5582 out.row_mut(col.col_idx).mapv_inplace(|v| v * col.scale);
5583 }
5584 if col.mean != 0.0 {
5585 let mut target = out.row_mut(col.col_idx);
5586 target += &(&interceptrow_snapshot * col.mean);
5587 }
5588 }
5589 out
5590 }
5591
5592 fn right_multiply_by_m_inv(&self, mat_internal: &Array2<f64>) -> Array2<f64> {
5595 let mut out = mat_internal.clone();
5596 let intercept = self.intercept_idx;
5597 let intercept_col_snapshot = mat_internal.column(intercept).to_owned();
5598 for col in &self.columns {
5599 if col.scale != 1.0 {
5600 out.column_mut(col.col_idx).mapv_inplace(|v| v * col.scale);
5601 }
5602 if col.mean != 0.0 {
5603 let mut target = out.column_mut(col.col_idx);
5604 target += &(&intercept_col_snapshot * col.mean);
5605 }
5606 }
5607 out
5608 }
5609
5610 pub fn transform_blockwise_penalties_to_internal(
5617 &self,
5618 penalties: &[BlockwisePenalty],
5619 p: usize,
5620 ) -> Vec<crate::penalty_spec::PenaltySpec> {
5621 let conditioning_cols: std::collections::HashSet<usize> =
5622 self.columns.iter().map(|c| c.col_idx).collect();
5623 penalties
5624 .iter()
5625 .map(|bp| {
5626 let overlaps =
5627 (bp.col_range.start..bp.col_range.end).any(|j| conditioning_cols.contains(&j));
5628 if overlaps {
5629 let global = bp.to_global(p);
5632 let right = self.transform_matrix_columnswith_a(&global);
5633 let transformed = self.transform_matrixrowswith_a_transpose(&right);
5634 crate::penalty_spec::PenaltySpec::Dense(transformed)
5635 } else {
5636 crate::penalty_spec::PenaltySpec::from_blockwise(bp.clone())
5639 }
5640 })
5641 .collect()
5642 }
5643
5644 pub fn backtransform_beta(&self, beta_internal: &Array1<f64>) -> Array1<f64> {
5645 let mut beta = beta_internal.clone();
5646 let intercept = self.intercept_idx;
5647 for col in &self.columns {
5648 beta[intercept] -= beta_internal[col.col_idx] * col.mean / col.scale;
5649 beta[col.col_idx] = beta_internal[col.col_idx] / col.scale;
5650 }
5651 beta
5652 }
5653
5654 pub fn transform_penalized_hessian_to_original(&self, h_internal: &Array2<f64>) -> Array2<f64> {
5657 let right = self.right_multiply_by_m_inv(h_internal);
5658 self.left_multiply_by_m_inv_transpose(&right)
5659 }
5660
5661 pub fn internal_bounds_for(&self, col_idx: usize, min: f64, max: f64) -> (f64, f64) {
5662 if let Some(col) = self.columns.iter().find(|c| c.col_idx == col_idx) {
5663 (min * col.scale, max * col.scale)
5664 } else {
5665 (min, max)
5666 }
5667 }
5668}
5669
5670pub fn freeze_raw_spatial_metadata(metadata: BasisMetadata, raw_cols: usize) -> BasisMetadata {
5671 match metadata {
5672 BasisMetadata::ThinPlate {
5673 centers,
5674 length_scale,
5675 periodic,
5676 identifiability_transform: None,
5677 input_scale,
5678 radial_reparam,
5679 } => BasisMetadata::ThinPlate {
5680 centers,
5681 length_scale,
5682 periodic,
5683 identifiability_transform: Some(Array2::eye(raw_cols)),
5684 input_scale,
5685 radial_reparam,
5686 },
5687 BasisMetadata::Duchon {
5688 centers,
5689 length_scale,
5690 periodic,
5691 power,
5692 nullspace_order,
5693 identifiability_transform: None,
5694 input_scale,
5695 aniso_log_scales,
5696 operator_collocation_points,
5697 radial_reparam,
5698 spectral_basis,
5699 } => BasisMetadata::Duchon {
5700 centers,
5701 length_scale,
5702 periodic,
5703 power,
5704 nullspace_order,
5705 identifiability_transform: Some(Array2::eye(raw_cols)),
5706 input_scale,
5707 aniso_log_scales,
5708 operator_collocation_points,
5709 radial_reparam,
5710 spectral_basis,
5711 },
5712 other => other,
5713 }
5714}
5715
5716pub fn matern_operator_penalty_triplet_from_metadata(
5717 metadata: &BasisMetadata,
5718) -> Result<crate::basis::FilteredPenalties, BasisError> {
5719 let BasisMetadata::Matern {
5720 centers,
5721 length_scale,
5722 periodic,
5723 nu,
5724 include_intercept,
5725 identifiability_transform,
5726 aniso_log_scales,
5727 input_scale,
5728 ..
5729 } = metadata
5730 else {
5731 crate::bail_invalid_basis!("Matérn operator penalties require Matérn metadata");
5732 };
5733 let penalty_length_scale = input_scale
5747 .to_standardized_units(*length_scale)
5748 .standardized_value();
5749 matern_operator_penalty_triplet_at_length_scale(
5750 centers.view(),
5751 periodic.as_deref(),
5752 identifiability_transform.as_ref(),
5753 *nu,
5754 *include_intercept,
5755 aniso_log_scales.as_deref(),
5756 penalty_length_scale,
5757 )
5758}
5759
5760pub fn matern_operator_penalty_triplet_at_length_scale(
5778 centers: ArrayView2<'_, f64>,
5779 periodic: Option<&[Option<f64>]>,
5780 identifiability_transform: Option<&Array2<f64>>,
5781 nu: crate::basis::MaternNu,
5782 include_intercept: bool,
5783 aniso_log_scales: Option<&[f64]>,
5784 effective_length_scale: f64,
5785) -> Result<crate::basis::FilteredPenalties, BasisError> {
5786 let penalty_centers = crate::basis::expand_periodic_centers(¢ers.to_owned(), periodic)?;
5787 let ops = build_matern_collocation_operator_matrices(
5788 penalty_centers.view(),
5789 None,
5790 effective_length_scale,
5791 nu,
5792 include_intercept,
5793 identifiability_transform.map(|z| z.view()),
5794 aniso_log_scales,
5795 )?;
5796 const ORDER_EPS: f64 = 1e-9;
5803 let d = penalty_centers.ncols();
5804 let m = nu.half_integer_value() + 0.5 * d as f64;
5805 let mut candidates = Vec::with_capacity(3);
5806 for (raw, source, min_order) in [
5807 (ops.d0.t().dot(&ops.d0), PenaltySource::OperatorMass, 0.0),
5808 (ops.d1.t().dot(&ops.d1), PenaltySource::OperatorTension, 1.0),
5809 (
5810 ops.d2.t().dot(&ops.d2),
5811 PenaltySource::OperatorStiffness,
5812 2.0,
5813 ),
5814 ] {
5815 let nondifferentiable_ou = matches!(nu, crate::basis::MaternNu::Half);
5816 if min_order > 0.0 && (nondifferentiable_ou || m + ORDER_EPS < min_order) {
5817 continue;
5818 }
5819 let sym = (&raw + &raw.t()) * 0.5;
5820 let (matrix, normalization_scale) = normalize_penalty_in_constrained_space(&sym);
5821 candidates.push(PenaltyCandidate {
5822 matrix: ConstructiveQuadratic::try_from_dense_psd(matrix, "Matérn operator penalty")?,
5823 source,
5824 normalization_scale,
5825 kronecker_factors: None,
5826 op: None,
5827 });
5828 }
5829 filter_penalty_candidates(candidates)
5830}
5831
5832pub fn normalize_penalty_in_constrained_space(matrix: &Array2<f64>) -> (Array2<f64>, f64) {
5833 let matrix = (matrix + &matrix.t().to_owned()) * 0.5;
5838 let matrix = crate::basis::project_penalty_to_psd_cone(&matrix);
5840 let c = matrix.iter().map(|v| v * v).sum::<f64>().sqrt();
5841 if c.is_finite() && c > 0.0 {
5842 (matrix.mapv(|v| v / c), c)
5843 } else {
5844 (matrix, 1.0)
5845 }
5846}
5847
5848pub fn tensor_product_design_from_sparse_marginals(
5849 marginal_sparse: &[&SparseColMat<usize, f64>],
5850) -> Result<SparseColMat<usize, f64>, BasisError> {
5851 if marginal_sparse.is_empty() {
5852 crate::bail_invalid_basis!("TensorBSpline requires at least one marginal basis");
5853 }
5854 let n = marginal_sparse[0].nrows();
5855 for (i, m) in marginal_sparse.iter().enumerate().skip(1) {
5856 if m.nrows() != n {
5857 crate::bail_dim_basis!(
5858 "tensor sparse marginal row mismatch at dim {i}: expected {n}, got {}",
5859 m.nrows()
5860 );
5861 }
5862 }
5863 let dims: Vec<usize> = marginal_sparse.iter().map(|m| m.ncols()).collect();
5864 let total_cols = dims.iter().try_fold(1usize, |acc, &q| {
5865 acc.checked_mul(q)
5866 .ok_or_else(|| BasisError::DimensionMismatch("tensor basis too large".to_string()))
5867 })?;
5868 let mut strides = vec![1usize; dims.len()];
5869 for d in (0..dims.len().saturating_sub(1)).rev() {
5870 strides[d] = strides[d + 1]
5871 .checked_mul(dims[d + 1])
5872 .ok_or_else(|| BasisError::DimensionMismatch("tensor basis too large".to_string()))?;
5873 }
5874
5875 use faer::sparse::SparseRowMat;
5876 let csrs: Vec<SparseRowMat<usize, f64>> = marginal_sparse
5877 .iter()
5878 .enumerate()
5879 .map(|(d, m)| {
5880 m.as_ref().to_row_major().map_err(|e| {
5881 BasisError::SparseCreation(format!(
5882 "tensor sparse marginal {d} CSR conversion failed: {e:?}"
5883 ))
5884 })
5885 })
5886 .collect::<Result<Vec<_>, _>>()?;
5887 let row_ptrs: Vec<&[usize]> = csrs.iter().map(|c| c.symbolic().row_ptr()).collect();
5888 let col_idxs: Vec<&[usize]> = csrs.iter().map(|c| c.symbolic().col_idx()).collect();
5889 let vals: Vec<&[f64]> = csrs.iter().map(|c| c.val()).collect();
5890
5891 use rayon::prelude::*;
5892 const CHUNK: usize = 1024;
5893 let num_chunks = n.div_ceil(CHUNK);
5894 let per_chunk: Vec<Vec<Triplet<usize, usize, f64>>> = (0..num_chunks)
5895 .into_par_iter()
5896 .map(|chunk_idx| {
5897 let row_start = chunk_idx * CHUNK;
5898 let row_end = (row_start + CHUNK).min(n);
5899 let mut chunk_triplets = Vec::<Triplet<usize, usize, f64>>::new();
5900 let mut cur_cols = Vec::<usize>::with_capacity(64);
5901 let mut cur_vals = Vec::<f64>::with_capacity(64);
5902 let mut next_cols = Vec::<usize>::with_capacity(64);
5903 let mut next_vals = Vec::<f64>::with_capacity(64);
5904 for i in row_start..row_end {
5905 cur_cols.clear();
5906 cur_vals.clear();
5907 cur_cols.push(0);
5908 cur_vals.push(1.0);
5909 let mut row_is_zero = false;
5910 for d in 0..dims.len() {
5911 let row_start_d = row_ptrs[d][i];
5912 let row_end_d = row_ptrs[d][i + 1];
5913 if row_start_d == row_end_d {
5914 row_is_zero = true;
5915 break;
5916 }
5917 let stride = strides[d];
5918 next_cols.clear();
5919 next_vals.clear();
5920 next_cols.reserve(cur_cols.len() * (row_end_d - row_start_d));
5921 next_vals.reserve(cur_vals.len() * (row_end_d - row_start_d));
5922 for (&prev_col, &prev_val) in cur_cols.iter().zip(cur_vals.iter()) {
5923 for ptr in row_start_d..row_end_d {
5924 let cj = col_idxs[d][ptr];
5925 let vj = vals[d][ptr];
5926 next_cols.push(prev_col + cj * stride);
5927 next_vals.push(prev_val * vj);
5928 }
5929 }
5930 std::mem::swap(&mut cur_cols, &mut next_cols);
5931 std::mem::swap(&mut cur_vals, &mut next_vals);
5932 }
5933 if row_is_zero {
5934 continue;
5935 }
5936 for (&col, &val) in cur_cols.iter().zip(cur_vals.iter()) {
5937 chunk_triplets.push(Triplet::new(i, col, val));
5938 }
5939 }
5940 chunk_triplets
5941 })
5942 .collect();
5943 let total_nnz: usize = per_chunk.iter().map(Vec::len).sum();
5944 let mut triplets = Vec::<Triplet<usize, usize, f64>>::with_capacity(total_nnz);
5945 for chunk in per_chunk {
5946 triplets.extend(chunk);
5947 }
5948 SparseColMat::try_new_from_triplets(n, total_cols, &triplets).map_err(|e| {
5949 BasisError::SparseCreation(format!(
5950 "failed to assemble sparse tensor product design: {e:?}"
5951 ))
5952 })
5953}
5954
5955pub fn dense_local_margin_to_sparse(
5956 dense: &Array2<f64>,
5957) -> Result<SparseColMat<usize, f64>, BasisError> {
5958 let expected_row_nnz = dense.ncols().min(4);
5959 let mut triplets =
5960 Vec::<Triplet<usize, usize, f64>>::with_capacity(dense.nrows() * expected_row_nnz);
5961 for ((row, col), &value) in dense.indexed_iter() {
5962 if value != 0.0 {
5963 triplets.push(Triplet::new(row, col, value));
5964 }
5965 }
5966 SparseColMat::try_new_from_triplets(dense.nrows(), dense.ncols(), &triplets).map_err(|e| {
5967 BasisError::SparseCreation(format!(
5968 "failed to convert tensor marginal design to sparse form: {e:?}"
5969 ))
5970 })
5971}
5972
5973pub struct TensorMarginRangeNullProjectors {
5974 range: Array2<f64>,
5975 null: Array2<f64>,
5976}
5977
5978pub fn projector_from_columns(columns: &Array2<f64>, indices: &[usize]) -> Array2<f64> {
5979 if indices.is_empty() {
5980 return Array2::<f64>::zeros((columns.nrows(), columns.nrows()));
5981 }
5982 let basis = columns.select(Axis(1), indices);
5983 basis.dot(&basis.t())
5984}
5985
5986pub fn tensor_margin_range_null_projectors(
5987 normalized_marginal_penalties: &[(Array2<f64>, f64)],
5988) -> Result<Vec<TensorMarginRangeNullProjectors>, BasisError> {
5989 normalized_marginal_penalties
5990 .iter()
5991 .enumerate()
5992 .map(|(dim, (penalty, _))| {
5993 let analysis = crate::basis::analyze_penalty_block(penalty)?;
5994 if analysis.rank == 0 {
5995 crate::bail_invalid_basis!(
5996 "t2 separable tensor penalty margin {dim} has rank-zero penalty; \
5997 cannot split penalized and null subspaces"
5998 );
5999 }
6000 let mut range_idx = Vec::<usize>::new();
6001 let mut null_idx = Vec::<usize>::new();
6002 for (idx, &ev) in analysis.eigenvalues.iter().enumerate() {
6003 if ev > analysis.rank_tol {
6004 range_idx.push(idx);
6005 } else {
6006 null_idx.push(idx);
6007 }
6008 }
6009 Ok(TensorMarginRangeNullProjectors {
6010 range: projector_from_columns(&analysis.eigenvectors, &range_idx),
6011 null: projector_from_columns(&analysis.eigenvectors, &null_idx),
6012 })
6013 })
6014 .collect()
6015}
6016
6017pub fn build_tensor_bspline_basis(
6018 data: ArrayView2<'_, f64>,
6019 feature_cols: &[usize],
6020 spec: &TensorBSplineSpec,
6021) -> Result<BasisBuildResult, BasisError> {
6022 if feature_cols.is_empty() {
6023 crate::bail_invalid_basis!("TensorBSpline requires at least one feature column");
6024 }
6025 if feature_cols.len() != spec.marginalspecs.len() {
6026 crate::bail_dim_basis!(
6027 "TensorBSpline feature/spec mismatch: feature_cols={}, marginalspecs={}",
6028 feature_cols.len(),
6029 spec.marginalspecs.len()
6030 );
6031 }
6032 if let Some((margin, _)) = spec
6033 .marginalspecs
6034 .iter()
6035 .enumerate()
6036 .find(|(_, marginal)| marginal.boundary_conditions.has_nonzero_anchor())
6037 {
6038 crate::bail_invalid_basis!(
6039 "TensorBSpline margin {margin} has a non-zero endpoint anchor. An inhomogeneous \
6040 marginal constraint cannot be represented by the tensor's homogeneous coefficient \
6041 chart plus one scalar row offset; use a separate anchored 1-D smooth or an explicit \
6042 model offset"
6043 );
6044 }
6045 if !spec.periods.is_empty() && spec.periods.len() != feature_cols.len() {
6046 crate::bail_dim_basis!(
6047 "TensorBSpline periods length {} does not match feature count {}",
6048 spec.periods.len(),
6049 feature_cols.len()
6050 );
6051 }
6052 let p = data.ncols();
6053 for &c in feature_cols {
6054 if c >= p {
6055 crate::bail_dim_basis!(
6056 "tensor feature column {c} is out of bounds for data with {p} columns"
6057 );
6058 }
6059 }
6060
6061 let mut marginal_knots = Vec::<Array1<f64>>::with_capacity(feature_cols.len());
6062 let mut marginal_is_cr_flags = Vec::<bool>::with_capacity(feature_cols.len());
6065 let mut marginal_degrees = Vec::<usize>::with_capacity(feature_cols.len());
6066 let mut marginalnum_basis = Vec::<usize>::with_capacity(feature_cols.len());
6067 let mut marginal_penalties = Vec::<Array2<f64>>::with_capacity(feature_cols.len());
6068 let mut marginal_function_grams = Vec::<Array2<f64>>::with_capacity(feature_cols.len());
6069 let mut marginal_designs = Vec::<Array2<f64>>::with_capacity(feature_cols.len());
6070 let mut marginal_effective_periods = Vec::<Option<f64>>::with_capacity(feature_cols.len());
6078 let mut marginal_sparse =
6086 Vec::<Option<SparseColMat<usize, f64>>>::with_capacity(feature_cols.len());
6087
6088 for (dim, (&col, marginalspec)) in feature_cols
6091 .iter()
6092 .zip(spec.marginalspecs.iter())
6093 .enumerate()
6094 {
6095 let mut marginal_unconstrained = marginalspec.clone();
6100 marginal_unconstrained.identifiability = BSplineIdentifiability::None;
6101 let built = build_bspline_basis_1d(data.column(col), &marginal_unconstrained)?;
6102 let (knots, marginal_is_cr, effective_degree, function_gram) = match built.metadata {
6107 BasisMetadata::BSpline1D {
6108 knots,
6109 periodic,
6110 degree,
6111 ..
6112 } => {
6113 let effective_degree = degree.unwrap_or(marginal_unconstrained.degree);
6114 let gram = if spec.double_penalty {
6115 Some(match periodic {
6116 Some((start, period, num_basis)) => {
6117 crate::basis::periodic_bspline_function_gram(
6118 start,
6119 start + period,
6120 effective_degree,
6121 num_basis,
6122 )?
6123 }
6124 None => crate::basis::bspline_function_gram(&knots, effective_degree)?,
6125 })
6126 } else {
6127 None
6128 };
6129 (knots, false, effective_degree, gram)
6130 }
6131 BasisMetadata::CubicRegression1D { knots, .. } => {
6132 let gram = spec
6133 .double_penalty
6134 .then(|| crate::basis::cubic_regression_function_gram(&knots))
6135 .transpose()?;
6136 (knots, true, marginalspec.degree, gram)
6137 }
6138 _ => {
6139 crate::bail_invalid_basis!(
6140 "internal TensorBSpline error at dim {dim}: expected BSpline1D or CubicRegression1D metadata"
6141 );
6142 }
6143 };
6144 let metadata_knots = match marginalspec.knotspec {
6145 BSplineKnotSpec::PeriodicUniform {
6146 data_range,
6147 num_basis,
6148 } => Array1::linspace(data_range.0, data_range.1, num_basis),
6149 _ => knots,
6150 };
6151 if let Some(function_gram) = function_gram {
6152 if function_gram.dim() != (built.design.ncols(), built.design.ncols()) {
6153 crate::bail_dim_basis!(
6154 "internal TensorBSpline error at dim {dim}: function Gram is {:?}, basis has {} columns",
6155 function_gram.dim(),
6156 built.design.ncols()
6157 );
6158 }
6159 marginal_function_grams.push(function_gram);
6160 }
6161 marginal_knots.push(metadata_knots);
6162 marginal_is_cr_flags.push(marginal_is_cr);
6163 marginal_degrees.push(effective_degree);
6164 marginalnum_basis.push(built.design.ncols());
6165 let dense_marginal = built.design.to_dense();
6170 let sparse_view: Option<SparseColMat<usize, f64>> = match built.design.as_sparse() {
6171 Some(sd) => {
6172 let inner: &SparseColMat<usize, f64> = sd;
6173 Some(inner.clone())
6174 }
6175 None => match marginalspec.knotspec {
6176 BSplineKnotSpec::PeriodicUniform { .. } => {
6177 Some(dense_local_margin_to_sparse(&dense_marginal)?)
6178 }
6179 _ => None,
6180 },
6181 };
6182 marginal_sparse.push(sparse_view);
6183 marginal_designs.push(dense_marginal);
6184 marginal_penalties.push(
6185 built
6186 .active_penalties
6187 .first()
6188 .ok_or_else(|| {
6189 BasisError::InvalidInput(format!(
6190 "internal TensorBSpline error at dim {dim}: missing marginal penalty"
6191 ))
6192 })?
6193 .matrix
6194 .clone(),
6195 );
6196 built.active_penalties.first().ok_or_else(|| {
6197 BasisError::InvalidInput(format!(
6198 "internal TensorBSpline error at dim {dim}: missing marginal nullspace dim"
6199 ))
6200 })?;
6201 let implied_period = match marginalspec.knotspec {
6209 BSplineKnotSpec::PeriodicUniform { data_range, .. } => {
6210 Some(data_range.1 - data_range.0)
6211 }
6212 _ => spec.periods.get(dim).and_then(|p| *p),
6213 };
6214 marginal_effective_periods.push(implied_period);
6215 }
6216
6217 let total_cols: usize = marginalnum_basis.iter().product();
6218 let mut dense_design = (!matches!(spec.identifiability, TensorBSplineIdentifiability::None))
6219 .then(|| tensor_product_design_from_marginals(&marginal_designs))
6220 .transpose()?;
6221 let mut candidates = Vec::<PenaltyCandidate>::with_capacity(
6222 match spec.penalty_decomposition {
6223 TensorBSplinePenaltyDecomposition::MarginalKroneckerSum => marginal_penalties.len(),
6224 TensorBSplinePenaltyDecomposition::Separable => marginal_penalties.len() * 2,
6225 } + if spec.double_penalty { 1 } else { 0 },
6226 );
6227
6228 let normalized_marginal_penalties: Vec<(Array2<f64>, f64)> = marginal_penalties
6236 .iter()
6237 .map(normalize_penalty_in_constrained_space)
6238 .collect();
6239 let tensor_function_gram = if spec.double_penalty {
6240 if marginal_function_grams.len() != marginalnum_basis.len() {
6241 crate::bail_dim_basis!(
6242 "TensorBSpline double penalty requires one function Gram per margin; got {} for {} margins",
6243 marginal_function_grams.len(),
6244 marginalnum_basis.len()
6245 );
6246 }
6247 let mut gram = Array2::<f64>::eye(1);
6248 for marginal_gram in &marginal_function_grams {
6249 gram = kronecker_product(&gram, marginal_gram);
6250 }
6251 Some(gram)
6252 } else {
6253 None
6254 };
6255 let joint_wiggliness = if spec.double_penalty {
6260 let mut sum = Array2::<f64>::zeros((total_cols, total_cols));
6261 for dim in 0..normalized_marginal_penalties.len() {
6262 let mut embedded = Array2::<f64>::eye(1);
6263 for (margin, &width) in marginalnum_basis.iter().enumerate() {
6264 let factor = if margin == dim {
6265 normalized_marginal_penalties[margin].0.clone()
6266 } else {
6267 Array2::<f64>::eye(width)
6268 };
6269 embedded = kronecker_product(&embedded, &factor);
6270 }
6271 sum += &embedded;
6272 }
6273 Some(sum)
6274 } else {
6275 None
6276 };
6277 let mut kronecker_marginal_penalties =
6278 Vec::<Array2<f64>>::with_capacity(normalized_marginal_penalties.len());
6279
6280 match spec.penalty_decomposition {
6281 TensorBSplinePenaltyDecomposition::MarginalKroneckerSum => {
6282 for dim in 0..normalized_marginal_penalties.len() {
6288 let mut s_dim = Array2::<f64>::eye(1);
6289 let mut factors = Vec::<Array2<f64>>::with_capacity(marginalnum_basis.len());
6290 for (j, &qj) in marginalnum_basis.iter().enumerate() {
6291 let factor = if j == dim {
6292 normalized_marginal_penalties[j].0.clone()
6293 } else {
6294 Array2::<f64>::eye(qj)
6295 };
6296 factors.push(factor.clone());
6297 s_dim = kronecker_product(&s_dim, &factor);
6298 }
6299 if dim == kronecker_marginal_penalties.len() {
6300 kronecker_marginal_penalties.push(normalized_marginal_penalties[dim].0.clone());
6301 }
6302 candidates.push(PenaltyCandidate {
6303 matrix: ConstructiveQuadratic::try_from_dense_psd(
6304 s_dim,
6305 "tensor marginal penalty",
6306 )?,
6307 source: PenaltySource::TensorMarginal { dim },
6308 normalization_scale: normalized_marginal_penalties[dim].1,
6309 kronecker_factors: Some(factors),
6310 op: None,
6311 });
6312 }
6313
6314 if let (Some(primary), Some(gram)) =
6315 (joint_wiggliness.as_ref(), tensor_function_gram.as_ref())
6316 && let Some(shrink) =
6317 crate::basis::function_space_nullspace_shrinkage(primary, gram)?
6318 {
6319 let (matrix, normalization_scale) = normalize_penalty_in_constrained_space(&shrink);
6320 candidates.push(PenaltyCandidate {
6321 matrix: ConstructiveQuadratic::try_from_dense_psd(
6322 matrix,
6323 "tensor global null-function ridge",
6324 )?,
6325 source: PenaltySource::TensorGlobalRidge,
6326 normalization_scale,
6327 kronecker_factors: None,
6328 op: None,
6329 });
6330 }
6331 }
6332 TensorBSplinePenaltyDecomposition::Separable => {
6333 let projectors = tensor_margin_range_null_projectors(&normalized_marginal_penalties)?;
6334 let n_masks = 1usize.checked_shl(projectors.len() as u32).ok_or_else(|| {
6335 BasisError::InvalidInput(format!(
6336 "t2 separable tensor penalty supports at most {} margins, got {}",
6337 usize::BITS - 1,
6338 projectors.len()
6339 ))
6340 })?;
6341 for mask in 1..n_masks {
6342 let mut matrix = Array2::<f64>::eye(1);
6343 let mut factors = Vec::<Array2<f64>>::with_capacity(projectors.len());
6344 let mut penalized_margins = Vec::<usize>::new();
6345 for (dim, projector) in projectors.iter().enumerate() {
6346 let use_range = ((mask >> dim) & 1) == 1;
6347 let factor = if use_range {
6348 penalized_margins.push(dim);
6349 projector.range.clone()
6350 } else {
6351 projector.null.clone()
6352 };
6353 matrix = kronecker_product(&matrix, &factor);
6354 factors.push(factor);
6355 }
6356 let (matrix, normalization_scale) = normalize_penalty_in_constrained_space(&matrix);
6357 candidates.push(PenaltyCandidate {
6358 matrix: ConstructiveQuadratic::try_from_dense_psd(
6359 matrix,
6360 "tensor separable penalty",
6361 )?,
6362 source: PenaltySource::TensorSeparable { penalized_margins },
6363 normalization_scale,
6364 kronecker_factors: Some(factors),
6365 op: None,
6366 });
6367 }
6368
6369 if let (Some(primary), Some(gram)) =
6370 (joint_wiggliness.as_ref(), tensor_function_gram.as_ref())
6371 && let Some(matrix) =
6372 crate::basis::function_space_nullspace_shrinkage(primary, gram)?
6373 {
6374 let (matrix, normalization_scale) = normalize_penalty_in_constrained_space(&matrix);
6375 candidates.push(PenaltyCandidate {
6376 matrix: ConstructiveQuadratic::try_from_dense_psd(
6377 matrix,
6378 "separable tensor global null-function ridge",
6379 )?,
6380 source: PenaltySource::TensorGlobalRidge,
6381 normalization_scale,
6382 kronecker_factors: None,
6383 op: None,
6384 });
6385 }
6386 }
6387 }
6388
6389 let z_opt = match &spec.identifiability {
6390 TensorBSplineIdentifiability::None => None,
6391 TensorBSplineIdentifiability::SumToZero => {
6392 if total_cols < 2 {
6393 crate::bail_invalid_basis!(
6394 "TensorBSpline requires at least 2 basis coefficients to enforce sum-to-zero identifiability"
6395 );
6396 }
6397 let dense_design_ref = dense_design.as_ref().ok_or_else(|| {
6398 BasisError::InvalidInput(
6399 "tensor sum-to-zero identifiability requires a realized basis".to_string(),
6400 )
6401 })?;
6402 let (_, z) = apply_sum_to_zero_constraint(dense_design_ref.view(), None)?;
6403 let gauge = gam_problem::Gauge::sum_to_zero(z);
6404 Some(gauge.block_transform(0))
6405 }
6406 TensorBSplineIdentifiability::MarginalSumToZero => {
6407 if marginal_designs.len() < 2 {
6418 crate::bail_invalid_basis!(
6419 "tensor interaction (ti) identifiability requires at least 2 margins"
6420 );
6421 }
6422 let mut z = Array2::<f64>::eye(1);
6423 for (dim, marginal) in marginal_designs.iter().enumerate() {
6424 if marginal.ncols() < 2 {
6425 crate::bail_invalid_basis!(
6426 "tensor interaction (ti) margin {dim} has fewer than 2 basis functions; \
6427 cannot remove its marginal main effect"
6428 );
6429 }
6430 let (_, z_dim) = apply_sum_to_zero_constraint(marginal.view(), None)?;
6431 let gauge_dim = gam_problem::Gauge::sum_to_zero(z_dim);
6432 let z_dim = gauge_dim.block_transform(0);
6433 z = kronecker_product(&z, &z_dim);
6434 }
6435 Some(z)
6436 }
6437 TensorBSplineIdentifiability::FrozenTransform { transform } => {
6438 if transform.nrows() != total_cols {
6439 crate::bail_dim_basis!(
6440 "frozen tensor identifiability transform mismatch: design has {} columns but transform has {} rows",
6441 total_cols,
6442 transform.nrows()
6443 );
6444 }
6445 Some(transform.clone())
6446 }
6447 };
6448
6449 if let Some(z) = z_opt.as_ref() {
6450 let gauge = gam_problem::Gauge::from_block_transforms(&[z.clone()]);
6451 let dense = dense_design.as_mut().ok_or_else(|| {
6452 BasisError::InvalidInput(
6453 "tensor identifiability transform requires a realized basis".to_string(),
6454 )
6455 })?;
6456 let restricted_design = gauge.restrict_design(dense);
6457 *dense = restricted_design;
6458 candidates = candidates
6459 .into_iter()
6460 .map(|candidate| -> Result<PenaltyCandidate, BasisError> {
6461 let restricted = candidate
6462 .matrix
6463 .restricted(&gauge, "tensor identifiability restriction")?;
6464 let (_, c_new) = normalize_penalty_in_constrained_space(restricted.dense());
6472 let matrix = restricted.scaled(
6473 1.0 / c_new,
6474 "normalized tensor penalty after identifiability",
6475 )?;
6476 Ok(PenaltyCandidate {
6477 matrix,
6478 source: candidate.source,
6479 normalization_scale: candidate.normalization_scale * c_new,
6480 kronecker_factors: None,
6486 op: candidate.op.clone(),
6487 })
6488 })
6489 .collect::<Result<Vec<_>, _>>()?;
6490
6491 if candidates
6492 .iter()
6493 .any(|candidate| matches!(candidate.source, PenaltySource::TensorGlobalRidge))
6494 {
6495 let width = candidates
6496 .first()
6497 .ok_or_else(|| {
6498 BasisError::InvalidInput(
6499 "TensorBSpline global ridge has no penalty candidates".to_string(),
6500 )
6501 })?
6502 .matrix
6503 .nrows();
6504 let physical_primary_terms = candidates
6505 .iter()
6506 .filter(|candidate| !matches!(candidate.source, PenaltySource::TensorGlobalRidge))
6507 .map(|candidate| {
6508 candidate.matrix.scaled(
6509 candidate.normalization_scale,
6510 "physical tensor primary penalty",
6511 )
6512 })
6513 .collect::<Result<Vec<_>, _>>()?;
6514 let joint_primary = ConstructiveQuadratic::sum(
6515 &physical_primary_terms,
6516 "joint tensor primary penalty",
6517 )?;
6518 for candidate in &mut candidates {
6519 if !matches!(candidate.source, PenaltySource::TensorGlobalRidge) {
6520 continue;
6521 }
6522 let physical_ridge = candidate
6523 .matrix
6524 .scaled(candidate.normalization_scale, "physical tensor null ridge")?;
6525 match crate::basis::rebuild_metric_consistent_ridge(
6526 &joint_primary,
6527 &physical_ridge,
6528 )? {
6529 Some(rebuilt) => {
6530 let (_, scale) = normalize_penalty_in_constrained_space(rebuilt.dense());
6531 candidate.matrix =
6532 rebuilt.scaled(1.0 / scale, "normalized rebuilt tensor null ridge")?;
6533 candidate.normalization_scale = scale;
6534 }
6535 None => {
6536 candidate.matrix = ConstructiveQuadratic::zero(width);
6537 candidate.normalization_scale = 1.0;
6538 }
6539 }
6540 candidate.kronecker_factors = None;
6541 candidate.op = None;
6542 }
6543 }
6544 }
6545
6546 let filtered = filter_penalty_candidates(candidates)?;
6547 let identifiability_is_none =
6548 matches!(spec.identifiability, TensorBSplineIdentifiability::None);
6549 let all_marginals_sparse = marginal_sparse.iter().all(Option::is_some);
6557 let design = if let Some(dense_design) = dense_design {
6558 DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(dense_design))
6559 } else if identifiability_is_none && all_marginals_sparse {
6560 let sparse_marginals: Vec<&SparseColMat<usize, f64>> = marginal_sparse
6566 .iter()
6567 .map(|m| m.as_ref().expect("all_marginals_sparse just verified"))
6568 .collect();
6569 let sparse_design = tensor_product_design_from_sparse_marginals(&sparse_marginals)?;
6570 DesignMatrix::Sparse(gam_linalg::matrix::SparseDesignMatrix::new(sparse_design))
6571 } else {
6572 let marginals: Vec<Arc<Array2<f64>>> = marginal_designs
6573 .iter()
6574 .map(|m| Arc::new(m.clone()))
6575 .collect();
6576 let op = TensorProductDesignOperator::new(marginals).map_err(|e| {
6577 BasisError::InvalidInput(format!("TensorProductDesignOperator build failed: {e}"))
6578 })?;
6579 DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(op)))
6580 };
6581
6582 Ok(BasisBuildResult {
6583 design,
6584 affine_offset: None,
6585 active_penalties: filtered.active,
6586 dropped_penalties: filtered.dropped,
6587 joint_null_rotation: None,
6588 metadata: BasisMetadata::TensorBSpline {
6589 feature_cols: feature_cols.to_vec(),
6590 knots: marginal_knots,
6591 degrees: marginal_degrees,
6592 periods: marginal_effective_periods,
6599 is_cr: marginal_is_cr_flags,
6600 identifiability_transform: z_opt,
6601 },
6602 kronecker_factored: if !spec.double_penalty
6609 && matches!(spec.identifiability, TensorBSplineIdentifiability::None)
6610 && matches!(
6611 spec.penalty_decomposition,
6612 TensorBSplinePenaltyDecomposition::MarginalKroneckerSum
6613 ) {
6614 Some(KroneckerFactoredBasis::new(
6615 marginal_designs,
6616 kronecker_marginal_penalties,
6617 marginalnum_basis.clone(),
6618 spec.double_penalty,
6619 ))
6620 } else {
6621 None
6622 },
6623 })
6624}
6625
6626#[cfg(test)]
6627mod tensor_function_space_runtime_tests {
6628 use super::*;
6629 use crate::basis::{
6630 BSplineBoundaryConditions, BSplineEndpointBoundaryCondition, OneDimensionalBoundary,
6631 };
6632 use ndarray::array;
6633
6634 fn marginal() -> BSplineBasisSpec {
6635 BSplineBasisSpec {
6636 degree: 2,
6637 penalty_order: 1,
6638 knotspec: BSplineKnotSpec::Generate {
6639 data_range: (0.0, 1.0),
6640 num_internal_knots: 2,
6641 },
6642 double_penalty: false,
6643 identifiability: BSplineIdentifiability::None,
6644 boundary: OneDimensionalBoundary::Open,
6645 boundary_conditions: BSplineBoundaryConditions::default(),
6646 }
6647 }
6648
6649 #[test]
6650 fn function_space_tensor_ridge_uses_exact_canonical_runtime() {
6651 let data = array![
6652 [0.00, 0.13],
6653 [0.15, 0.82],
6654 [0.29, 0.37],
6655 [0.43, 0.95],
6656 [0.58, 0.21],
6657 [0.71, 0.66],
6658 [0.86, 0.48],
6659 [1.00, 0.04]
6660 ];
6661 let mut spec = TensorBSplineSpec {
6662 marginalspecs: vec![marginal(), marginal()],
6663 periods: Vec::new(),
6664 double_penalty: true,
6665 identifiability: TensorBSplineIdentifiability::None,
6666 penalty_decomposition: TensorBSplinePenaltyDecomposition::MarginalKroneckerSum,
6667 };
6668 let built = build_tensor_bspline_basis(data.view(), &[0, 1], &spec)
6669 .expect("double-penalty tensor basis");
6670 assert!(
6671 built
6672 .active_penalties
6673 .iter()
6674 .any(|penalty| { matches!(penalty.info.source, PenaltySource::TensorGlobalRidge) })
6675 );
6676 assert!(
6677 built.kronecker_factored.is_none(),
6678 "the legacy factored runtime cannot represent a function-metric global ridge"
6679 );
6680
6681 spec.double_penalty = false;
6682 let singly_penalized = build_tensor_bspline_basis(data.view(), &[0, 1], &spec)
6683 .expect("single-penalty tensor basis");
6684 assert!(
6685 singly_penalized.kronecker_factored.is_some(),
6686 "the exact marginal-only fast path must remain available"
6687 );
6688 }
6689
6690 #[test]
6691 fn tensor_nonzero_anchor_is_rejected_before_its_affine_lift_can_be_dropped() {
6692 let data = array![[0.0, 0.0], [0.25, 0.75], [0.75, 0.25], [1.0, 1.0]];
6693 let mut anchored = marginal();
6694 anchored.boundary_conditions.left =
6695 BSplineEndpointBoundaryCondition::Anchored { value: 1.25 };
6696 let spec = TensorBSplineSpec {
6697 marginalspecs: vec![anchored, marginal()],
6698 periods: Vec::new(),
6699 double_penalty: false,
6700 identifiability: TensorBSplineIdentifiability::None,
6701 penalty_decomposition: TensorBSplinePenaltyDecomposition::MarginalKroneckerSum,
6702 };
6703
6704 let error = build_tensor_bspline_basis(data.view(), &[0, 1], &spec)
6705 .expect_err("a tensor margin cannot silently discard an inhomogeneous lift");
6706 let message = error.to_string();
6707 assert!(message.contains("TensorBSpline margin 0"));
6708 assert!(message.contains("non-zero endpoint anchor"));
6709 assert!(message.contains("explicit model offset"));
6710 }
6711}
6712
6713pub fn tensor_product_design_from_marginals(
6714 marginal_designs: &[Array2<f64>],
6715) -> Result<Array2<f64>, BasisError> {
6716 if marginal_designs.is_empty() {
6717 crate::bail_invalid_basis!("TensorBSpline requires at least one marginal basis");
6718 }
6719 let n = marginal_designs[0].nrows();
6720 for (i, b) in marginal_designs.iter().enumerate().skip(1) {
6721 if b.nrows() != n {
6722 crate::bail_dim_basis!(
6723 "tensor marginal row mismatch at dim {i}: expected {n}, got {}",
6724 b.nrows()
6725 );
6726 }
6727 }
6728 let total_cols = marginal_designs.iter().try_fold(1usize, |acc, b| {
6729 acc.checked_mul(b.ncols())
6730 .ok_or_else(|| BasisError::DimensionMismatch("tensor basis too large".to_string()))
6731 })?;
6732 use ndarray::parallel::prelude::*;
6738 use rayon::iter::{IntoParallelIterator, ParallelIterator};
6739 let mut design = Array2::<f64>::zeros((n, total_cols));
6740 design
6741 .axis_chunks_iter_mut(ndarray::Axis(0), 1024)
6742 .into_par_iter()
6743 .enumerate()
6744 .for_each(|(chunk_idx, mut block)| {
6745 let row_offset = chunk_idx * 1024;
6746 let mut cur = Vec::<f64>::with_capacity(total_cols);
6748 let mut next = Vec::<f64>::with_capacity(total_cols);
6749 for (local_i, mut out_row) in block.outer_iter_mut().enumerate() {
6750 let i = row_offset + local_i;
6751 cur.clear();
6752 cur.push(1.0);
6753 for b in marginal_designs {
6754 let q = b.ncols();
6755 next.clear();
6756 next.resize(cur.len() * q, 0.0);
6757 let b_row = b.row(i);
6761 let b_slice = b_row
6762 .as_slice()
6763 .expect("Array2 row from outer_iter is contiguous");
6764 for (a_idx, &aval) in cur.iter().enumerate() {
6765 let off = a_idx * q;
6766 let dst = &mut next[off..off + q];
6767 for col in 0..q {
6768 dst[col] = aval * b_slice[col];
6769 }
6770 }
6771 std::mem::swap(&mut cur, &mut next);
6772 }
6773 let out_slice = out_row
6778 .as_slice_mut()
6779 .expect("design row is contiguous in C-major Array2");
6780 out_slice.copy_from_slice(&cur);
6781 }
6782 });
6783 Ok(design)
6784}
6785
6786fn fmt_level_value(v: f64) -> String {
6790 if v.is_finite() && v.fract() == 0.0 && v.abs() < 1e15 {
6791 format!("{}", v as i64)
6792 } else {
6793 format!("{v}")
6794 }
6795}
6796
6797pub fn build_random_effect_block(
6798 data: ArrayView2<'_, f64>,
6799 spec: &RandomEffectTermSpec,
6800) -> Result<RandomEffectBlock, BasisError> {
6801 let n = data.nrows();
6802 let p = data.ncols();
6803 if spec.feature_col >= p {
6804 crate::bail_dim_basis!(
6805 "random-effect term '{}' feature column {} out of bounds for {} columns",
6806 spec.name,
6807 spec.feature_col,
6808 p
6809 );
6810 }
6811
6812 let col = data.column(spec.feature_col);
6813 if col.iter().any(|v| !v.is_finite()) {
6814 crate::bail_invalid_basis!(
6815 "random-effect term '{}' contains non-finite group values",
6816 spec.name
6817 );
6818 }
6819
6820 let kept_levels: Vec<u64> = if let Some(levels) = spec.frozen_levels.as_ref() {
6821 if levels.is_empty() {
6822 crate::bail_invalid_basis!(
6823 "random-effect term '{}' has empty frozen_levels",
6824 spec.name
6825 );
6826 }
6827 levels
6831 .iter()
6832 .map(|&b| gam_data::canonical_level_bits(f64::from_bits(b)))
6833 .collect()
6834 } else {
6835 let mut seen = BTreeSet::<u64>::new();
6836 let mut levels = Vec::<u64>::new();
6837 for &v in col {
6838 let bits = gam_data::canonical_level_bits(v);
6839 if seen.insert(bits) {
6840 levels.push(bits);
6841 }
6842 }
6843 if levels.is_empty() {
6844 crate::bail_invalid_basis!("random-effect term '{}' has no observed levels", spec.name);
6845 }
6846 let start_idx = if spec.drop_first_level && levels.len() > 1 {
6847 1usize
6848 } else {
6849 0usize
6850 };
6851 levels[start_idx..].to_vec()
6852 };
6853
6854 if kept_levels.is_empty() {
6855 crate::bail_invalid_basis!(
6856 "random-effect term '{}' drops all levels; keep at least one level",
6857 spec.name
6858 );
6859 }
6860
6861 let q = kept_levels.len();
6862 let mut level_to_col = BTreeMap::<u64, usize>::new();
6863 for (idx, &bits) in kept_levels.iter().enumerate() {
6864 if level_to_col.insert(bits, idx).is_some() {
6865 crate::bail_invalid_basis!(
6866 "random-effect term '{}' has duplicate frozen level bits {bits}",
6867 spec.name
6868 );
6869 }
6870 }
6871 let strict_unseen =
6885 !spec.lenient_unseen && !spec.drop_first_level && spec.frozen_levels.is_some();
6886 let mut group_ids = Vec::with_capacity(n);
6887 for (row, &v) in col.iter().enumerate() {
6888 let bits = gam_data::canonical_level_bits(v);
6889 let group_id = level_to_col.get(&bits).copied();
6890 if strict_unseen && group_id.is_none() {
6891 crate::bail_invalid_basis!(
6892 "unseen level '{}' in fixed factor column '{}' at row {}; the factor's levels \
6893 were fixed at fit time and an out-of-vocabulary level cannot be predicted \
6894 (use group({}) for a random effect that tolerates held-out levels)",
6895 fmt_level_value(v),
6896 spec.name,
6897 row,
6898 spec.name
6899 );
6900 }
6901 group_ids.push(group_id);
6902 }
6903
6904 Ok(RandomEffectBlock {
6905 name: spec.name.clone(),
6906 group_ids,
6907 num_groups: q,
6908 kept_levels,
6909 })
6910}
6911
6912#[cfg(test)]
6913mod random_effect_signed_zero_tests {
6914 use super::{RandomEffectTermSpec, build_random_effect_block};
6915 use ndarray::array;
6916
6917 fn spec() -> RandomEffectTermSpec {
6918 RandomEffectTermSpec {
6919 name: "g".to_string(),
6920 feature_col: 0,
6921 drop_first_level: false,
6922 penalized: true,
6923 frozen_levels: None,
6924 lenient_unseen: true,
6925 }
6926 }
6927
6928 #[test]
6929 fn signed_zero_rows_share_one_group() {
6930 let data = array![[-0.0_f64], [0.0], [1.0], [-0.0], [1.0]];
6934 let block = build_random_effect_block(data.view(), &spec()).unwrap();
6935 assert_eq!(
6936 block.num_groups, 2,
6937 "0.0/-0.0 must not split into two groups"
6938 );
6939 assert_eq!(block.group_ids[0], block.group_ids[1]);
6941 assert_eq!(block.group_ids[0], block.group_ids[3]);
6942 assert_eq!(block.group_ids[2], block.group_ids[4]);
6943 assert_ne!(block.group_ids[0], block.group_ids[2]);
6944 }
6945
6946 #[test]
6947 fn frozen_positive_zero_matches_negative_zero_row() {
6948 let mut s = spec();
6951 s.frozen_levels = Some(vec![0.0_f64.to_bits(), 1.0_f64.to_bits()]);
6952 let data = array![[-0.0_f64], [1.0]];
6953 let block = build_random_effect_block(data.view(), &s).unwrap();
6954 assert_eq!(
6955 block.group_ids[0],
6956 Some(0),
6957 "-0.0 must match the +0.0 column"
6958 );
6959 assert_eq!(block.group_ids[1], Some(1));
6960 }
6961
6962 #[test]
6963 fn frozen_negative_zero_matches_positive_zero_row() {
6964 let mut s = spec();
6967 s.frozen_levels = Some(vec![(-0.0_f64).to_bits(), 1.0_f64.to_bits()]);
6968 let data = array![[0.0_f64], [1.0]];
6969 let block = build_random_effect_block(data.view(), &s).unwrap();
6970 assert_eq!(
6971 block.group_ids[0],
6972 Some(0),
6973 "+0.0 must match the -0.0 column"
6974 );
6975 }
6976
6977 fn fixed_factor_spec() -> RandomEffectTermSpec {
6980 let mut s = spec();
6983 s.name = "year".to_string();
6984 s.lenient_unseen = false;
6985 s
6986 }
6987
6988 #[test]
6989 fn fixed_factor_rejects_unseen_numeric_level_at_predict() {
6990 let mut s = fixed_factor_spec();
6995 s.frozen_levels = Some(vec![2000.0_f64.to_bits(), 2001.0_f64.to_bits()]);
6996 let data = array![[2000.0_f64], [1999.0]];
6997 let err = build_random_effect_block(data.view(), &s)
6998 .expect_err("an unseen fixed-factor level must be rejected");
6999 let msg = format!("{err}");
7000 assert!(
7001 msg.contains("unseen level"),
7002 "message must name the defect: {msg}"
7003 );
7004 assert!(
7005 msg.contains("1999"),
7006 "message must name the integer level (not 1999.0): {msg}"
7007 );
7008 assert!(msg.contains("year"), "message must name the column: {msg}");
7009 }
7010
7011 #[test]
7012 fn fixed_factor_accepts_seen_numeric_levels_at_predict() {
7013 let mut s = fixed_factor_spec();
7016 s.frozen_levels = Some(vec![2000.0_f64.to_bits(), 2001.0_f64.to_bits()]);
7017 let data = array![[2001.0_f64], [2000.0]];
7018 let block = build_random_effect_block(data.view(), &s).unwrap();
7019 assert_eq!(block.group_ids[0], Some(1));
7020 assert_eq!(block.group_ids[1], Some(0));
7021 }
7022
7023 #[test]
7024 fn fixed_factor_at_fit_time_derives_vocabulary_and_never_false_rejects() {
7025 let mut s = fixed_factor_spec();
7029 s.frozen_levels = None;
7030 let data = array![[2000.0_f64], [2001.0], [2002.0], [2000.0]];
7031 let block = build_random_effect_block(data.view(), &s)
7032 .expect("fit-time build must not reject its own levels");
7033 assert_eq!(block.num_groups, 3);
7034 }
7035
7036 #[test]
7037 fn random_effect_still_tolerates_unseen_numeric_level() {
7038 let mut s = spec(); s.frozen_levels = Some(vec![2000.0_f64.to_bits(), 2001.0_f64.to_bits()]);
7043 let data = array![[2000.0_f64], [1999.0]];
7044 let block = build_random_effect_block(data.view(), &s)
7045 .expect("a random effect tolerates unseen levels");
7046 assert_eq!(block.group_ids[0], Some(0));
7047 assert_eq!(
7048 block.group_ids[1], None,
7049 "unseen level → population mean, not a reject"
7050 );
7051 }
7052}
7053
7054impl SmoothDesign {
7055 pub fn map_term_coefficients(
7058 unconstrained: &Array1<f64>,
7059 shape: ShapeConstraint,
7060 ) -> Result<Array1<f64>, BasisError> {
7061 if unconstrained.is_empty() {
7062 crate::bail_invalid_basis!("unconstrained coefficient vector cannot be empty");
7063 }
7064 let mapped = match shape {
7065 ShapeConstraint::None => unconstrained.clone(),
7066 ShapeConstraint::MonotoneIncreasing => cumulative_exp(unconstrained, 1.0),
7067 ShapeConstraint::MonotoneDecreasing => cumulative_exp(unconstrained, -1.0),
7068 ShapeConstraint::Convex => second_cumulative_exp(unconstrained, 1.0),
7069 ShapeConstraint::Concave => second_cumulative_exp(unconstrained, -1.0),
7070 };
7071 Ok(mapped)
7072 }
7073}
7074
7075pub struct LocalSmoothTermBuild {
7076 pub dim: usize,
7077 pub design: DesignMatrix,
7078 pub affine_offset: Option<Array1<f64>>,
7080 pub active_penalties: Vec<ActivePenalty>,
7081 pub joint_null_rotation: Option<crate::basis::JointNullRotation>,
7088 pub dropped_penalties: Vec<DroppedPenaltyInfo>,
7089 pub metadata: BasisMetadata,
7090 pub linear_constraints: Option<LinearInequalityConstraints>,
7091 pub box_reparam: bool,
7092 pub kronecker_factored: Option<KroneckerFactoredBasis>,
7093}
7094
7095#[derive(Clone)]
7096pub struct PcaScoresMemmapDesignOperator {
7097 mmap: Arc<memmap2::Mmap>,
7098 data_offset: usize,
7099 nrows: usize,
7100 ncols: usize,
7101 chunk_size: usize,
7102}
7103
7104impl PcaScoresMemmapDesignOperator {
7105 fn open(path: PathBuf, chunk_size: usize) -> Result<Self, BasisError> {
7106 let file = File::open(&path).map_err(|err| {
7107 BasisError::InvalidInput(format!(
7108 "failed to open lazy Pca .npy scores '{}': {err}",
7109 path.display()
7110 ))
7111 })?;
7112 let mmap = unsafe {
7118 memmap2::Mmap::map(&file).map_err(|err| {
7119 BasisError::InvalidInput(format!(
7120 "failed to memmap lazy Pca .npy scores '{}': {err}",
7121 path.display()
7122 ))
7123 })?
7124 };
7125 let (data_offset, nrows, ncols) = parse_f64_2d_npy_header(&mmap, &path)?;
7126 let expected = data_offset
7127 .checked_add(nrows.saturating_mul(ncols).saturating_mul(8))
7128 .ok_or_else(|| {
7129 BasisError::InvalidInput(format!(
7130 "lazy Pca .npy scores '{}' shape is too large",
7131 path.display()
7132 ))
7133 })?;
7134 if mmap.len() < expected {
7135 crate::bail_invalid_basis!(
7136 "lazy Pca .npy scores '{}' is truncated: header expects {} bytes, file has {}",
7137 path.display(),
7138 expected,
7139 mmap.len()
7140 );
7141 }
7142 Ok(Self {
7143 mmap: Arc::new(mmap),
7144 data_offset,
7145 nrows,
7146 ncols,
7147 chunk_size: chunk_size.max(1),
7148 })
7149 }
7150
7151 fn value(&self, row: usize, col: usize) -> f64 {
7152 let offset = self.data_offset + (row * self.ncols + col) * 8;
7153 let mut bytes = [0_u8; 8];
7154 bytes.copy_from_slice(&self.mmap[offset..offset + 8]);
7155 f64::from_le_bytes(bytes)
7156 }
7157
7158 fn chunk_rows(&self) -> usize {
7159 self.chunk_size.min(self.nrows.max(1))
7160 }
7161}
7162
7163impl LinearOperator for PcaScoresMemmapDesignOperator {
7164 fn nrows(&self) -> usize {
7165 self.nrows
7166 }
7167
7168 fn ncols(&self) -> usize {
7169 self.ncols
7170 }
7171
7172 fn apply(&self, vector: &Array1<f64>) -> Array1<f64> {
7173 assert_eq!(
7174 vector.len(),
7175 self.ncols,
7176 "lazy Pca apply vector length mismatch"
7177 );
7178 let mut out = Array1::<f64>::zeros(self.nrows);
7179 for start in (0..self.nrows).step_by(self.chunk_rows()) {
7180 let end = (start + self.chunk_rows()).min(self.nrows);
7181 for row in start..end {
7182 let mut acc = 0.0;
7183 for col in 0..self.ncols {
7184 acc += self.value(row, col) * vector[col];
7185 }
7186 out[row] = acc;
7187 }
7188 }
7189 out
7190 }
7191
7192 fn apply_transpose(&self, vector: &Array1<f64>) -> Array1<f64> {
7193 assert_eq!(
7194 vector.len(),
7195 self.nrows,
7196 "lazy Pca apply_transpose vector length mismatch"
7197 );
7198 let mut out = Array1::<f64>::zeros(self.ncols);
7199 for start in (0..self.nrows).step_by(self.chunk_rows()) {
7200 let end = (start + self.chunk_rows()).min(self.nrows);
7201 for row in start..end {
7202 let scale = vector[row];
7203 if scale == 0.0 {
7204 continue;
7205 }
7206 for col in 0..self.ncols {
7207 out[col] += scale * self.value(row, col);
7208 }
7209 }
7210 }
7211 out
7212 }
7213
7214 fn diag_xtw_x(&self, weights: &Array1<f64>) -> Result<Array2<f64>, String> {
7215 if weights.len() != self.nrows {
7216 return Err(format!(
7217 "lazy Pca diag_xtw_x weight length mismatch: weights={}, nrows={}",
7218 weights.len(),
7219 self.nrows
7220 ));
7221 }
7222 FiniteSignedWeightsView::try_from_array(weights)
7223 .map_err(|reason| format!("lazy Pca diag_xtw_x: {reason}"))?;
7224 let mut gram = Array2::<f64>::zeros((self.ncols, self.ncols));
7225 for start in (0..self.nrows).step_by(self.chunk_rows()) {
7226 let end = (start + self.chunk_rows()).min(self.nrows);
7227 for row in start..end {
7228 let w = weights[row];
7229 if w == 0.0 {
7230 continue;
7231 }
7232 for a in 0..self.ncols {
7233 let xa = self.value(row, a);
7234 if xa == 0.0 {
7235 continue;
7236 }
7237 for b in a..self.ncols {
7238 gram[[a, b]] += w * xa * self.value(row, b);
7239 }
7240 }
7241 }
7242 }
7243 for a in 0..self.ncols {
7244 for b in 0..a {
7245 gram[[a, b]] = gram[[b, a]];
7246 }
7247 }
7248 Ok(gram)
7249 }
7250
7251 fn apply_weighted_normal(
7252 &self,
7253 weights: FiniteSignedWeightsView<'_>,
7254 vector: &Array1<f64>,
7255 penalty: Option<&Array2<f64>>,
7256 ridge: f64,
7257 ) -> Array1<f64> {
7258 assert_eq!(
7259 weights.len(),
7260 self.nrows,
7261 "lazy Pca weighted-normal weight mismatch"
7262 );
7263 assert_eq!(
7264 vector.len(),
7265 self.ncols,
7266 "lazy Pca weighted-normal vector mismatch"
7267 );
7268 let weights = weights.view();
7269 let mut out = Array1::<f64>::zeros(self.ncols);
7270 for start in (0..self.nrows).step_by(self.chunk_rows()) {
7271 let end = (start + self.chunk_rows()).min(self.nrows);
7272 for row in start..end {
7273 let w = weights[row];
7274 if w == 0.0 {
7275 continue;
7276 }
7277 let mut row_dot = 0.0;
7278 for col in 0..self.ncols {
7279 row_dot += self.value(row, col) * vector[col];
7280 }
7281 if row_dot == 0.0 {
7282 continue;
7283 }
7284 let scaled = w * row_dot;
7285 for col in 0..self.ncols {
7286 out[col] += scaled * self.value(row, col);
7287 }
7288 }
7289 }
7290 if let Some(pen) = penalty {
7291 out += &pen.dot(vector);
7292 }
7293 if ridge > 0.0 {
7294 out += &vector.mapv(|x| ridge * x);
7295 }
7296 out
7297 }
7298}
7299
7300impl DenseDesignOperator for PcaScoresMemmapDesignOperator {
7301 fn compute_xtwy(&self, weights: &Array1<f64>, y: &Array1<f64>) -> Result<Array1<f64>, String> {
7302 if weights.len() != self.nrows || y.len() != self.nrows {
7303 return Err(format!(
7304 "lazy Pca compute_xtwy dimension mismatch: weights={}, y={}, nrows={}",
7305 weights.len(),
7306 y.len(),
7307 self.nrows
7308 ));
7309 }
7310 FiniteSignedWeightsView::try_from_array(weights)
7311 .map_err(|reason| format!("lazy Pca compute_xtwy: {reason}"))?;
7312 let mut out = Array1::<f64>::zeros(self.ncols);
7313 for start in (0..self.nrows).step_by(self.chunk_rows()) {
7314 let end = (start + self.chunk_rows()).min(self.nrows);
7315 for row in start..end {
7316 let scale = weights[row] * y[row];
7317 if scale == 0.0 {
7318 continue;
7319 }
7320 for col in 0..self.ncols {
7321 out[col] += scale * self.value(row, col);
7322 }
7323 }
7324 }
7325 Ok(out)
7326 }
7327
7328 fn row_chunk_into(
7329 &self,
7330 rows: Range<usize>,
7331 mut out: ArrayViewMut2<'_, f64>,
7332 ) -> Result<(), MatrixMaterializationError> {
7333 if rows.end > self.nrows || rows.start > rows.end {
7334 return Err(MatrixMaterializationError::MissingRowChunk {
7335 context: "lazy Pca row range out of bounds",
7336 });
7337 }
7338 if out.nrows() != rows.end - rows.start || out.ncols() != self.ncols {
7339 return Err(MatrixMaterializationError::MissingRowChunk {
7340 context: "lazy Pca row_chunk_into shape mismatch",
7341 });
7342 }
7343 for (local, row) in (rows.start..rows.end).enumerate() {
7344 for col in 0..self.ncols {
7345 out[[local, col]] = self.value(row, col);
7346 }
7347 }
7348 Ok(())
7349 }
7350
7351 fn to_dense(&self) -> Array2<f64> {
7352 let mut out = Array2::<f64>::zeros((self.nrows, self.ncols));
7353 self.row_chunk_into(0..self.nrows, out.view_mut())
7354 .expect("lazy Pca full materialization failed");
7355 out
7356 }
7357}
7358
7359pub fn parse_f64_2d_npy_header(
7360 bytes: &[u8],
7361 path: &PathBuf,
7362) -> Result<(usize, usize, usize), BasisError> {
7363 let mut reader = std::io::Cursor::new(bytes);
7364 let header = npyz::NpyHeader::from_reader(&mut reader).map_err(|err| {
7365 BasisError::InvalidInput(format!(
7366 "lazy Pca scores '{}' has an invalid .npy header: {err}",
7367 path.display()
7368 ))
7369 })?;
7370 let is_little_endian_f64 = matches!(
7371 header.dtype(),
7372 npyz::DType::Plain(ref dtype)
7373 if dtype.type_char() == npyz::TypeChar::Float
7374 && dtype.size_field() == 8
7375 && dtype.endianness() == npyz::Endianness::Little
7376 );
7377 if !is_little_endian_f64 {
7378 crate::bail_invalid_basis!(
7379 "lazy Pca scores '{}' must be scalar little-endian float64 .npy, got {}",
7380 path.display(),
7381 header.dtype().descr()
7382 );
7383 }
7384 if header.order() != npyz::Order::C {
7385 crate::bail_invalid_basis!(
7386 "lazy Pca scores '{}' must be C-contiguous, not Fortran-ordered",
7387 path.display()
7388 );
7389 }
7390 if header.shape().len() != 2 {
7391 crate::bail_invalid_basis!(
7392 "lazy Pca scores '{}' must have shape (N, K), got {:?}",
7393 path.display(),
7394 header.shape()
7395 );
7396 }
7397 let nrows = usize::try_from(header.shape()[0]).map_err(|_| {
7398 BasisError::InvalidInput(format!(
7399 "lazy Pca scores '{}' row count {} exceeds this platform's address space",
7400 path.display(),
7401 header.shape()[0]
7402 ))
7403 })?;
7404 let ncols = usize::try_from(header.shape()[1]).map_err(|_| {
7405 BasisError::InvalidInput(format!(
7406 "lazy Pca scores '{}' column count {} exceeds this platform's address space",
7407 path.display(),
7408 header.shape()[1]
7409 ))
7410 })?;
7411 let data_offset = usize::try_from(reader.position()).map_err(|_| {
7412 BasisError::InvalidInput(format!(
7413 "lazy Pca scores '{}' header offset exceeds this platform's address space",
7414 path.display()
7415 ))
7416 })?;
7417 Ok((data_offset, nrows, ncols))
7418}
7419
7420pub fn pca_center_mean(x: ArrayView2<'_, f64>) -> Result<Array1<f64>, BasisError> {
7421 if x.nrows() == 0 {
7422 crate::bail_invalid_basis!("Pca basis requires at least one row to compute center mean");
7423 }
7424 let mut mean = Array1::<f64>::zeros(x.ncols());
7425 for row in x.rows() {
7426 mean += &row;
7427 }
7428 mean.mapv_inplace(|v| v / x.nrows() as f64);
7429 Ok(mean)
7430}
7431
7432fn pca_function_mass_penalty(
7443 mut raw_score_gram: Array2<f64>,
7444 n_rows: usize,
7445 smooth_penalty: f64,
7446) -> Result<Array2<f64>, BasisError> {
7447 let k = raw_score_gram.ncols();
7448 if raw_score_gram.nrows() != k {
7449 crate::bail_dim_basis!(
7450 "Pca score Gram must be square, got {}x{}",
7451 raw_score_gram.nrows(),
7452 k
7453 );
7454 }
7455 if n_rows == 0 {
7456 crate::bail_invalid_basis!("Pca basis requires at least one score row");
7457 }
7458 if k == 0 {
7459 crate::bail_invalid_basis!("Pca basis requires at least one score column");
7460 }
7461 if k > n_rows {
7462 crate::bail_invalid_basis!(
7463 "Pca score design is rank deficient: {} score columns cannot have full column rank with only {} rows; remove redundant components",
7464 k,
7465 n_rows
7466 );
7467 }
7468 if raw_score_gram.iter().any(|value| !value.is_finite()) {
7469 crate::bail_invalid_basis!("Pca score design produced a non-finite function Gram");
7470 }
7471
7472 let rrqr = gam_linalg::faer_ndarray::rrqr_from_gram_with_permutation(
7476 &raw_score_gram,
7477 n_rows,
7478 gam_linalg::faer_ndarray::default_rrqr_rank_alpha(),
7479 )
7480 .map_err(BasisError::LinalgError)?;
7481 if rrqr.rank != k {
7482 let redundant_columns = &rrqr.column_permutation[rrqr.rank..];
7483 crate::bail_invalid_basis!(
7484 "Pca score design is rank deficient under canonical RRQR: rank {} < {} (tolerance {:.6e}); redundant score columns {:?}; remove zero or dependent components instead of stabilizing them with a coefficient ridge",
7485 rrqr.rank,
7486 k,
7487 rrqr.rank_tol,
7488 redundant_columns
7489 );
7490 }
7491
7492 raw_score_gram.mapv_inplace(|value| value * smooth_penalty / n_rows as f64);
7493 Ok(raw_score_gram)
7494}
7495
7496pub fn build_pca_smooth_basis(
7497 data: ArrayView2<'_, f64>,
7498 feature_cols: &[usize],
7499 basis_matrix: &Array2<f64>,
7500 centered: bool,
7501 smooth_penalty: f64,
7502 center_mean: Option<&Array1<f64>>,
7503 pca_basis_path: Option<&PathBuf>,
7504 chunk_size: usize,
7505) -> Result<BasisBuildResult, BasisError> {
7506 if !smooth_penalty.is_finite() || smooth_penalty < 0.0 {
7507 crate::bail_invalid_basis!(
7508 "Pca smooth_penalty must be finite and non-negative, got {}",
7509 smooth_penalty
7510 );
7511 }
7512 if data.nrows() == 0 {
7513 crate::bail_invalid_basis!("Pca basis requires at least one data row");
7514 }
7515
7516 if let Some(path) = pca_basis_path {
7517 let op = PcaScoresMemmapDesignOperator::open(path.clone(), chunk_size)?;
7518 if op.nrows != data.nrows() {
7519 crate::bail_dim_basis!(
7520 "lazy Pca scores row mismatch: .npy has {}, data has {}",
7521 op.nrows,
7522 data.nrows()
7523 );
7524 }
7525 let raw_score_gram = op
7528 .diag_xtw_x(&Array1::<f64>::ones(op.nrows))
7529 .map_err(|err| {
7530 BasisError::InvalidInput(format!(
7531 "lazy Pca function-mass Gram construction failed: {err}"
7532 ))
7533 })?;
7534 let penalty = pca_function_mass_penalty(raw_score_gram, op.nrows, smooth_penalty)?;
7535 let filtered = filter_penalty_candidates(vec![PenaltyCandidate {
7536 matrix: ConstructiveQuadratic::try_from_dense_psd(
7537 penalty,
7538 "lazy PCA function-mass penalty",
7539 )?,
7540 source: PenaltySource::OperatorMass,
7541 normalization_scale: 1.0,
7542 kronecker_factors: None,
7543 op: None,
7544 }])?;
7545 return Ok(BasisBuildResult {
7546 design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(op))),
7547 affine_offset: None,
7548 active_penalties: filtered.active,
7549 dropped_penalties: filtered.dropped,
7550 joint_null_rotation: None,
7551 metadata: BasisMetadata::Pca {
7552 feature_cols: feature_cols.to_vec(),
7553 basis_matrix: basis_matrix.clone(),
7554 centered,
7555 smooth_penalty,
7556 center_mean: center_mean.cloned(),
7557 pca_basis_path: Some(path.clone()),
7558 chunk_size: chunk_size.max(1),
7559 },
7560 kronecker_factored: None,
7561 });
7562 }
7563 if basis_matrix.nrows() != feature_cols.len() {
7564 crate::bail_dim_basis!(
7565 "Pca basis row mismatch: basis rows={}, feature columns={}",
7566 basis_matrix.nrows(),
7567 feature_cols.len()
7568 );
7569 }
7570 let mut x = select_columns(data, feature_cols)?;
7571 let mean = if centered {
7572 match center_mean {
7573 Some(mean) => mean.clone(),
7574 None => pca_center_mean(x.view())?,
7575 }
7576 } else {
7577 Array1::<f64>::zeros(feature_cols.len())
7578 };
7579 if centered {
7580 for mut row in x.rows_mut() {
7581 row -= &mean;
7582 }
7583 }
7584 let design = fast_ab(&x, basis_matrix);
7585 let raw_score_gram = gam_linalg::faer_ndarray::fast_ata(&design);
7586 let penalty = pca_function_mass_penalty(raw_score_gram, design.nrows(), smooth_penalty)?;
7587 let filtered = filter_penalty_candidates(vec![PenaltyCandidate {
7588 matrix: ConstructiveQuadratic::try_from_dense_psd(penalty, "PCA function-mass penalty")?,
7589 source: PenaltySource::OperatorMass,
7590 normalization_scale: 1.0,
7591 kronecker_factors: None,
7592 op: None,
7593 }])?;
7594 Ok(BasisBuildResult {
7595 design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(design)),
7596 affine_offset: None,
7597 active_penalties: filtered.active,
7598 dropped_penalties: filtered.dropped,
7599 joint_null_rotation: None,
7600 metadata: BasisMetadata::Pca {
7601 feature_cols: feature_cols.to_vec(),
7602 basis_matrix: basis_matrix.clone(),
7603 centered,
7604 smooth_penalty,
7605 center_mean: centered.then_some(mean),
7606 pca_basis_path: None,
7607 chunk_size: chunk_size.max(1),
7608 },
7609 kronecker_factored: None,
7610 })
7611}
7612
7613#[cfg(test)]
7614mod pca_function_mass_tests {
7615 use super::{PenaltySource, build_pca_smooth_basis, parse_f64_2d_npy_header};
7616 use ndarray::{Array1, Array2, array};
7617 use std::io::Write;
7618 use std::path::PathBuf;
7619
7620 fn quadratic_form(matrix: &Array2<f64>, coefficients: &Array1<f64>) -> f64 {
7621 coefficients.dot(&matrix.dot(coefficients))
7622 }
7623
7624 fn assert_close(left: f64, right: f64) {
7625 let scale = left.abs().max(right.abs()).max(1.0);
7626 assert!(
7627 (left - right).abs() <= 1e-11 * scale,
7628 "values differ: left={left:.16e}, right={right:.16e}"
7629 );
7630 }
7631
7632 fn write_f64_npy(scores: &Array2<f64>) -> PathBuf {
7633 let path = std::env::temp_dir().join(format!(
7634 "gam_terms_pca_function_mass_{}.npy",
7635 std::process::id()
7636 ));
7637 let mut header = format!(
7638 "{{'descr': '<f8', 'fortran_order': False, 'shape': ({}, {}), }}",
7639 scores.nrows(),
7640 scores.ncols()
7641 );
7642 while (10 + header.len() + 1) % 16 != 0 {
7643 header.push(' ');
7644 }
7645 header.push('\n');
7646 let header_len = u16::try_from(header.len()).expect("test .npy header fits u16");
7647
7648 let mut file = std::fs::File::create(&path).expect("create test .npy");
7649 file.write_all(b"\x93NUMPY").expect("write .npy magic");
7650 file.write_all(&[1, 0]).expect("write .npy version");
7651 file.write_all(&header_len.to_le_bytes())
7652 .expect("write .npy header length");
7653 file.write_all(header.as_bytes())
7654 .expect("write .npy header");
7655 for &value in scores {
7656 file.write_all(&value.to_le_bytes())
7657 .expect("write .npy score");
7658 }
7659 path
7660 }
7661
7662 fn npy_v1_bytes(mut header: String) -> Vec<u8> {
7663 while (10 + header.len() + 1) % 16 != 0 {
7664 header.push(' ');
7665 }
7666 header.push('\n');
7667 let header_len = u16::try_from(header.len()).expect("test header fits v1");
7668 let mut bytes = b"\x93NUMPY".to_vec();
7669 bytes.extend_from_slice(&[1, 0]);
7670 bytes.extend_from_slice(&header_len.to_le_bytes());
7671 bytes.extend_from_slice(header.as_bytes());
7672 bytes
7673 }
7674
7675 #[test]
7676 fn npy_header_parser_uses_exact_ast_fields_2293() {
7677 let path = PathBuf::from("scores.npy");
7678 let bytes = npy_v1_bytes(
7679 "{'shape':(3, 2), 'note':'True', 'descr':'<f8', 'fortran_order':False,}".to_string(),
7680 );
7681 let (offset, rows, cols) =
7682 parse_f64_2d_npy_header(&bytes, &path).expect("valid reordered header");
7683 assert_eq!((rows, cols), (3, 2));
7684 assert_eq!(offset, bytes.len());
7685
7686 for header in [
7687 "{'descr':'<f8','fortran_order':True,'shape':(3,2),}",
7688 "{'descr':'>f8','fortran_order':False,'shape':(3,2),}",
7689 "{'descr':'<f8','shape':(3,2),}",
7690 "{'descr':'<f8','fortran_order':'False','shape':(3,2),}",
7691 "{'descr':'<f8','fortran_order':False,'shape':(6,),}",
7692 ] {
7693 let invalid = npy_v1_bytes(header.to_string());
7694 assert!(
7695 parse_f64_2d_npy_header(&invalid, &path).is_err(),
7696 "{header}"
7697 );
7698 }
7699 }
7700
7701 #[test]
7702 fn pca_penalty_quadratic_equals_empirical_fitted_function_norm() {
7703 let data = array![[1.0, 2.0], [-1.0, 0.5], [2.0, -0.5], [0.25, -1.5]];
7704 let basis = array![[1.0, 0.5], [-0.25, 2.0]];
7705 let smooth_penalty = 2.5;
7706 let built = build_pca_smooth_basis(
7707 data.view(),
7708 &[0, 1],
7709 &basis,
7710 false,
7711 smooth_penalty,
7712 None,
7713 None,
7714 2,
7715 )
7716 .expect("full-rank PCA basis");
7717 let coefficients = array![0.7, -1.2];
7718 let design = built.design.to_dense();
7719 let fitted = design.dot(&coefficients);
7720 let expected = smooth_penalty * fitted.dot(&fitted) / fitted.len() as f64;
7721 let actual = quadratic_form(&built.active_penalties[0].matrix, &coefficients);
7722
7723 assert_close(actual, expected);
7724 assert_eq!(built.active_penalties[0].nullity, 0);
7725 assert_eq!(
7726 built.active_penalties[0].info.source,
7727 PenaltySource::OperatorMass
7728 );
7729 }
7730
7731 #[test]
7732 fn pca_function_mass_is_invariant_to_nonorthogonal_score_reparameterization() {
7733 let scores = array![[1.0, 2.0], [-1.0, 0.5], [2.0, -0.5], [0.25, -1.5]];
7734 let identity = Array2::<f64>::eye(2);
7735 let transform = array![[2.0, 0.5], [0.0, 0.25]];
7737 let base_coefficients = array![0.8, -1.1];
7738 let transformed_coefficients = array![1.5, -4.4];
7740 let smooth_penalty = 1.7;
7741
7742 let base = build_pca_smooth_basis(
7743 scores.view(),
7744 &[0, 1],
7745 &identity,
7746 false,
7747 smooth_penalty,
7748 None,
7749 None,
7750 2,
7751 )
7752 .expect("base PCA chart");
7753 let transformed = build_pca_smooth_basis(
7754 scores.view(),
7755 &[0, 1],
7756 &transform,
7757 false,
7758 smooth_penalty,
7759 None,
7760 None,
7761 2,
7762 )
7763 .expect("reparameterized PCA chart");
7764
7765 let fitted_base = base.design.to_dense().dot(&base_coefficients);
7766 let fitted_transformed = transformed.design.to_dense().dot(&transformed_coefficients);
7767 for (&left, &right) in fitted_base.iter().zip(fitted_transformed.iter()) {
7768 assert_close(left, right);
7769 }
7770 assert_close(
7771 quadratic_form(&base.active_penalties[0].matrix, &base_coefficients),
7772 quadratic_form(
7773 &transformed.active_penalties[0].matrix,
7774 &transformed_coefficients,
7775 ),
7776 );
7777 }
7778
7779 #[test]
7780 fn rank_deficient_pca_score_design_is_rejected() {
7781 let scores = array![[1.0, 0.0], [2.0, 0.0], [3.0, 0.0], [4.0, 0.0]];
7782 let result = build_pca_smooth_basis(
7783 scores.view(),
7784 &[0, 1],
7785 &Array2::<f64>::eye(2),
7786 false,
7787 1.0,
7788 None,
7789 None,
7790 2,
7791 );
7792 let err = result.err().expect("zero score column must be rejected");
7793 let message = err.to_string();
7794 assert!(
7795 message.contains("rank deficient"),
7796 "unexpected error: {message}"
7797 );
7798 assert!(
7799 message.contains("rank 1 < 2"),
7800 "missing RRQR evidence: {message}"
7801 );
7802 }
7803
7804 #[test]
7805 fn lazy_and_dense_pca_function_mass_penalties_match() {
7806 let scores = array![[1.0, 2.0], [-1.0, 0.5], [2.0, -0.5], [0.25, -1.5]];
7807 let smooth_penalty = 2.25;
7808 let path = write_f64_npy(&scores);
7809 let dense = build_pca_smooth_basis(
7810 scores.view(),
7811 &[0, 1],
7812 &Array2::<f64>::eye(2),
7813 false,
7814 smooth_penalty,
7815 None,
7816 None,
7817 2,
7818 )
7819 .expect("dense PCA basis");
7820 let lazy_data = Array2::<f64>::zeros((scores.nrows(), 0));
7821 let lazy = build_pca_smooth_basis(
7822 lazy_data.view(),
7823 &[],
7824 &Array2::<f64>::zeros((0, scores.ncols())),
7825 false,
7826 smooth_penalty,
7827 None,
7828 Some(&path),
7829 2,
7830 )
7831 .expect("lazy PCA basis");
7832 std::fs::remove_file(&path).expect("remove test .npy");
7833
7834 for (&left, &right) in dense.active_penalties[0]
7835 .matrix
7836 .iter()
7837 .zip(lazy.active_penalties[0].matrix.iter())
7838 {
7839 assert_close(left, right);
7840 }
7841 for (&left, &right) in dense
7842 .design
7843 .to_dense()
7844 .iter()
7845 .zip(lazy.design.to_dense().iter())
7846 {
7847 assert_close(left, right);
7848 }
7849 }
7850}
7851
7852pub fn defer_inner_model_centering_to_factor_level_wrapper(basis: &mut SmoothBasisSpec) {
7868 if let SmoothBasisSpec::BSpline1D { spec, .. } = basis
7869 && matches!(
7870 spec.identifiability,
7871 BSplineIdentifiability::WeightedSumToZero { .. }
7872 )
7873 {
7874 spec.identifiability = BSplineIdentifiability::None;
7875 }
7876}
7877
7878pub fn apply_by_variable_to_local_build(
7879 mut built: LocalSmoothTermBuild,
7880 data: ArrayView2<'_, f64>,
7881 by_col: usize,
7882 by: &ByVariableSpec,
7883 term_name: &str,
7884) -> Result<LocalSmoothTermBuild, BasisError> {
7885 if by_col >= data.ncols() {
7886 crate::bail_dim_basis!(
7887 "by-variable smooth term '{term_name}' references column {by_col}, but data has {} columns",
7888 data.ncols()
7889 );
7890 }
7891 let weights = match by {
7892 ByVariableSpec::Numeric => data.column(by_col).to_owned(),
7893 ByVariableSpec::Level { value_bits, .. } => {
7894 let value_bits = gam_data::canonical_level_bits(f64::from_bits(*value_bits));
7895 data.column(by_col).mapv(|value| {
7896 if gam_data::canonical_level_bits(value) == value_bits {
7897 1.0
7898 } else {
7899 0.0
7900 }
7901 })
7902 }
7903 };
7904 if weights.iter().any(|value| !value.is_finite()) {
7905 crate::bail_invalid_basis!(
7906 "by-variable smooth term '{term_name}' has non-finite by-column values"
7907 );
7908 }
7909
7910 let mut dense = built
7911 .design
7912 .try_to_dense_by_chunks("by-variable smooth row gating")
7913 .map_err(BasisError::InvalidInput)?;
7914 for (mut row, &weight) in dense.rows_mut().into_iter().zip(weights.iter()) {
7915 row.mapv_inplace(|value| value * weight);
7916 }
7917 if let Some(offset) = built.affine_offset.as_mut() {
7918 if offset.len() != weights.len() {
7919 crate::bail_dim_basis!(
7920 "by-variable smooth term '{term_name}' affine offset has {} rows but the by-variable has {}",
7921 offset.len(),
7922 weights.len()
7923 );
7924 }
7925 *offset *= &weights;
7926 }
7927 built.design = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(dense));
7928 built.kronecker_factored = None;
7929 Ok(built)
7930}
7931
7932pub fn build_by_smooth_local(
7943 data: ArrayView2<'_, f64>,
7944 term: &SmoothTermSpec,
7945 smooth: &SmoothBasisSpec,
7946 by_kind: &ByVarKind,
7947 workspace: &mut crate::basis::BasisWorkspace,
7948) -> Result<LocalSmoothTermBuild, BasisError> {
7949 let inner_term = SmoothTermSpec {
7950 frozen_parametric_residualization: None,
7951 name: term.name.clone(),
7952 basis: (*smooth).clone(),
7953 shape: term.shape,
7954 joint_null_rotation: None,
7955 };
7956 let inner = build_single_local_smooth_term(data, &inner_term, workspace)?;
7957
7958 match by_kind {
7959 ByVarKind::Numeric { feature_col } => {
7960 let inner_meta = inner.metadata.clone();
7961 let mut built = apply_by_variable_to_local_build(
7962 inner,
7963 data,
7964 *feature_col,
7965 &ByVariableSpec::Numeric,
7966 &term.name,
7967 )?;
7968 built.metadata = BasisMetadata::BySmooth {
7969 inner: Box::new(inner_meta),
7970 by_col: *feature_col,
7971 levels: None,
7972 ordered: false,
7973 };
7974 Ok(built)
7975 }
7976 ByVarKind::Factor {
7977 feature_col,
7978 frozen_levels,
7979 ordered,
7980 } => {
7981 let level_bits: Vec<u64> = if let Some(fl) = frozen_levels {
7984 fl.iter()
7985 .map(|&b| gam_data::canonical_level_bits(f64::from_bits(b)))
7986 .collect()
7987 } else {
7988 let col = data.column(*feature_col);
7989 let mut seen = BTreeSet::<u64>::new();
7990 for &v in col.iter() {
7991 if v.is_finite() {
7992 seen.insert(gam_data::canonical_level_bits(v));
7993 }
7994 }
7995 seen.into_iter().collect()
7996 };
7997 let n_levels = level_bits.len();
7998 if n_levels == 0 {
7999 crate::bail_invalid_basis!(
8000 "by-factor smooth term '{}': factor column {} has no observed levels",
8001 term.name,
8002 feature_col
8003 );
8004 }
8005 let p = inner.dim;
8006 let q = n_levels * p;
8007 let n = data.nrows();
8008
8009 let inner_dense = inner
8010 .design
8011 .try_to_dense_by_chunks("by-factor smooth design gating")
8012 .map_err(BasisError::InvalidInput)?;
8013
8014 let mut combined = Array2::<f64>::zeros((n, q));
8016 for (lvl_idx, &bits) in level_bits.iter().enumerate() {
8017 let col_start = lvl_idx * p;
8018 for row in 0..n {
8019 if gam_data::canonical_level_bits(data[[row, *feature_col]]) == bits {
8020 combined
8021 .slice_mut(s![row, col_start..col_start + p])
8022 .assign(&inner_dense.row(row));
8023 }
8024 }
8025 }
8026
8027 let inner_meta = inner.metadata.clone();
8039 let n_penalties = inner.active_penalties.len();
8040 let n_blocks = n_penalties.saturating_mul(n_levels);
8041 let mut candidates = Vec::<PenaltyCandidate>::with_capacity(n_blocks);
8042 for base_penalty in &inner.active_penalties {
8043 for lvl in 0..n_levels {
8044 let off = lvl * p;
8045 let mut s_big = Array2::<f64>::zeros((q, q));
8046 s_big
8047 .slice_mut(s![off..off + p, off..off + p])
8048 .assign(&base_penalty.matrix);
8049 let (s_big, scale) = normalize_penalty_in_constrained_space(&s_big);
8050 candidates.push(PenaltyCandidate {
8051 matrix: ConstructiveQuadratic::try_from_dense_psd(
8052 s_big,
8053 "factor-smooth replicated penalty",
8054 )?,
8055 source: base_penalty.info.source.clone(),
8056 normalization_scale: base_penalty.info.normalization_scale * scale,
8057 kronecker_factors: None,
8058 op: None,
8059 });
8060 }
8061 }
8062
8063 let filtered = crate::basis::filter_penalty_candidates(candidates)?;
8069 let joint_null_rotation = crate::basis::compute_joint_null_rotation(&filtered.active)?;
8070 let mut dropped_penalties = inner.dropped_penalties;
8071 dropped_penalties.extend(filtered.dropped);
8072
8073 Ok(LocalSmoothTermBuild {
8074 dim: q,
8075 design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(combined)),
8076 affine_offset: inner.affine_offset,
8080 active_penalties: filtered.active,
8081 joint_null_rotation,
8082 dropped_penalties,
8083 metadata: BasisMetadata::BySmooth {
8084 inner: Box::new(inner_meta),
8085 by_col: *feature_col,
8086 levels: Some(level_bits),
8087 ordered: *ordered,
8088 },
8089 linear_constraints: None,
8090 box_reparam: false,
8091 kronecker_factored: None,
8092 })
8093 }
8094 }
8095}
8096
8097pub fn ensure_by_variable_specs_match(
8098 kind: &BySmoothKind,
8099 by: &ByVariableSpec,
8100 term_name: &str,
8101) -> Result<(), BasisError> {
8102 match (kind, by) {
8103 (BySmoothKind::Numeric, ByVariableSpec::Numeric) => Ok(()),
8104 (BySmoothKind::Level { level_bits }, ByVariableSpec::Level { value_bits, .. })
8105 if level_bits == value_bits =>
8106 {
8107 Ok(())
8108 }
8109 _ => Err(BasisError::InvalidInput(format!(
8110 "by-variable smooth term '{term_name}' has inconsistent by-variable specifications"
8111 ))),
8112 }
8113}
8114
8115fn canonical_nullspace_directions(z: &Array2<f64>) -> Result<Array2<f64>, BasisError> {
8126 let (coefficient_dim, nullity) = z.dim();
8127 if nullity == 0 {
8128 return Ok(Array2::zeros((coefficient_dim, 0)));
8129 }
8130 if coefficient_dim < nullity || z.iter().any(|value| !value.is_finite()) {
8131 crate::bail_invalid_basis!(
8132 "null-space basis must be finite with rows >= columns, got {}x{}",
8133 coefficient_dim,
8134 nullity
8135 );
8136 }
8137
8138 let tolerance = 128.0 * f64::EPSILON * coefficient_dim.max(1) as f64;
8139 let mut canonical = Array2::<f64>::zeros((coefficient_dim, nullity));
8140 for accepted in 0..nullity {
8141 let mut best_coordinate = usize::MAX;
8142 let mut best_norm = 0.0_f64;
8143 let mut best = Array1::<f64>::zeros(coefficient_dim);
8144
8145 for coordinate in 0..coefficient_dim {
8146 let mut candidate = Array1::<f64>::zeros(coefficient_dim);
8148 for row in 0..coefficient_dim {
8149 candidate[row] = (0..nullity)
8150 .map(|axis| z[[row, axis]] * z[[coordinate, axis]])
8151 .sum();
8152 }
8153 for _ in 0..2 {
8156 for axis in 0..accepted {
8157 let direction = canonical.column(axis);
8158 let projection = direction.dot(&candidate);
8159 candidate.scaled_add(-projection, &direction);
8160 }
8161 }
8162 let norm = candidate.dot(&candidate).sqrt();
8163 let tie_band = tolerance * best_norm.max(1.0);
8164 if best_coordinate == usize::MAX || norm > best_norm + tie_band {
8165 best_coordinate = coordinate;
8166 best_norm = norm;
8167 best = candidate;
8168 }
8169 }
8170
8171 if best_coordinate == usize::MAX || best_norm <= tolerance {
8172 crate::bail_invalid_basis!(
8173 "null-space projector exposed only {} of {} independent directions",
8174 accepted,
8175 nullity
8176 );
8177 }
8178 best.mapv_inplace(|value| value / best_norm);
8179 let sign_anchor = best
8181 .iter()
8182 .enumerate()
8183 .max_by(|(left_index, left), (right_index, right)| {
8184 left.abs()
8185 .partial_cmp(&right.abs())
8186 .unwrap_or(std::cmp::Ordering::Equal)
8187 .then_with(|| right_index.cmp(left_index))
8188 })
8189 .map(|(_, value)| *value)
8190 .unwrap_or(1.0);
8191 if sign_anchor < 0.0 {
8192 best.mapv_inplace(|value| -value);
8193 }
8194 canonical.column_mut(accepted).assign(&best);
8195 }
8196 Ok(canonical)
8197}
8198
8199#[cfg(test)]
8200mod canonical_nullspace_direction_tests {
8201 use super::*;
8202 use ndarray::array;
8203
8204 #[test]
8205 fn per_axis_null_penalties_are_invariant_to_eigensolver_gauge_2315() {
8206 let inv_sqrt_two = 0.5_f64.sqrt();
8207 let z = array![
8208 [inv_sqrt_two, 0.0],
8209 [inv_sqrt_two, 0.0],
8210 [0.0, 1.0],
8211 [0.0, 0.0]
8212 ];
8213 let rotation = array![[0.6, -0.8], [0.8, 0.6]];
8214 let rotated = z.dot(&rotation);
8215 let reference = canonical_nullspace_directions(&z).expect("canonical null basis");
8216 let actual =
8217 canonical_nullspace_directions(&rotated).expect("rotated canonical null basis");
8218 for axis in 0..reference.ncols() {
8219 let reference_penalty = reference
8220 .column(axis)
8221 .to_owned()
8222 .insert_axis(Axis(1))
8223 .dot(&reference.column(axis).insert_axis(Axis(0)));
8224 let actual_penalty = actual
8225 .column(axis)
8226 .to_owned()
8227 .insert_axis(Axis(1))
8228 .dot(&actual.column(axis).insert_axis(Axis(0)));
8229 let max_error = reference_penalty
8230 .iter()
8231 .zip(actual_penalty.iter())
8232 .map(|(left, right)| (left - right).abs())
8233 .fold(0.0_f64, f64::max);
8234 assert!(
8235 max_error <= 256.0 * f64::EPSILON,
8236 "axis {axis} changed by {max_error:e}"
8237 );
8238 }
8239 }
8240}
8241
8242pub fn build_factor_smooth(
8270 data: ArrayView2<'_, f64>,
8271 spec: &FactorSmoothSpec,
8272 term_name: &str,
8273 workspace: &mut crate::basis::BasisWorkspace,
8274) -> Result<LocalSmoothTermBuild, BasisError> {
8275 if spec.continuous_cols.len() != 1 {
8276 crate::bail_invalid_basis!(
8277 "factor smooth term '{}' currently supports exactly one continuous covariate; found {}",
8278 term_name,
8279 spec.continuous_cols.len()
8280 );
8281 }
8282 let feature_col = spec.continuous_cols[0];
8283 let group_col = spec.group_col;
8284 if feature_col >= data.ncols() || group_col >= data.ncols() {
8285 crate::bail_dim_basis!(
8286 "factor smooth term '{}' references columns ({}, {}) out of bounds for {} columns",
8287 term_name,
8288 feature_col,
8289 group_col,
8290 data.ncols()
8291 );
8292 }
8293
8294 if matches!(spec.flavour, FactorSmoothFlavour::Sz) {
8297 let levels = resolve_factor_smooth_levels(data, group_col, spec, term_name)?;
8298 let inner = SmoothBasisSpec::BSpline1D {
8299 feature_col,
8300 spec: factor_smooth_marginal_for_replay(&spec.marginal),
8301 };
8302 let sz_term = SmoothTermSpec {
8303 frozen_parametric_residualization: None,
8304 name: term_name.to_string(),
8305 basis: SmoothBasisSpec::FactorSumToZero {
8306 inner: Box::new(inner),
8307 by_col: group_col,
8308 levels: levels.clone(),
8309 frozen_global_orthogonality: None,
8310 },
8311 shape: ShapeConstraint::None,
8312 joint_null_rotation: None,
8313 };
8314 let mut built = build_single_local_smooth_term(data, &sz_term, workspace)?;
8315 let (knots, degree, periodic, marginal_is_cr) = match &built.metadata {
8336 BasisMetadata::BSpline1D {
8337 knots,
8338 periodic,
8339 degree,
8340 ..
8341 } => (
8342 knots.clone(),
8343 degree.unwrap_or(spec.marginal.degree),
8344 *periodic,
8345 false,
8346 ),
8347 BasisMetadata::CubicRegression1D { knots, .. } => {
8348 (knots.clone(), spec.marginal.degree, None, true)
8349 }
8350 other => {
8351 crate::bail_invalid_basis!(
8352 "sz factor smooth term '{}' produced an unexpected marginal metadata variant {:?}",
8353 term_name,
8354 other
8355 );
8356 }
8357 };
8358 built.metadata = BasisMetadata::FactorSmooth {
8359 continuous_cols: spec.continuous_cols.clone(),
8360 group_col,
8361 knots,
8362 degree,
8363 periodic,
8364 group_levels: levels,
8365 flavour: "sz".to_string(),
8366 marginal_is_cr,
8367 };
8368 return Ok(built);
8369 }
8370
8371 let levels = resolve_factor_smooth_levels(data, group_col, spec, term_name)?;
8372 let n_levels = levels.len();
8373 if n_levels < 2 {
8374 crate::bail_invalid_basis!(
8375 "factor smooth term '{}' requires at least two grouping levels; found {}",
8376 term_name,
8377 n_levels
8378 );
8379 }
8380
8381 let use_per_dim_null = matches!(
8389 &spec.flavour,
8390 FactorSmoothFlavour::Fs { m_null_penalty_orders }
8391 if m_null_penalty_orders.iter().copied().max().unwrap_or(0) >= 1
8392 );
8393
8394 let mut marginal_spec = factor_smooth_marginal_for_replay(&spec.marginal);
8400 if use_per_dim_null {
8401 marginal_spec.double_penalty = false;
8402 }
8403 let inner_term = SmoothTermSpec {
8404 frozen_parametric_residualization: None,
8405 name: format!("{term_name}::marginal"),
8406 basis: SmoothBasisSpec::BSpline1D {
8407 feature_col,
8408 spec: marginal_spec,
8409 },
8410 shape: ShapeConstraint::None,
8411 joint_null_rotation: None,
8412 };
8413 let inner = build_single_local_smooth_term(data, &inner_term, workspace)?;
8414 let mut base = inner
8415 .design
8416 .try_to_dense_by_chunks("factor smooth marginal")
8417 .map_err(BasisError::InvalidInput)?;
8418 if matches!(spec.flavour, FactorSmoothFlavour::Re) {
8419 let center = match &inner.metadata {
8429 BasisMetadata::BSpline1D { knots, .. } if !knots.is_empty() => {
8430 0.5 * (knots[0] + knots[knots.len() - 1])
8431 }
8432 _ => 0.0,
8433 };
8434 let mut linear = Array2::<f64>::ones((data.nrows(), 2));
8435 linear
8436 .column_mut(1)
8437 .assign(&data.column(feature_col).mapv(|x| x - center));
8438 base = linear;
8439 }
8440 let n = base.nrows();
8441 let p = base.ncols();
8442 let q = p * n_levels;
8443
8444 let mut dense = Array2::<f64>::zeros((n, q));
8447 for i in 0..n {
8448 let bits = gam_data::canonical_level_bits(data[[i, group_col]]);
8449 let Some(level_idx) = levels.iter().position(|b| *b == bits) else {
8450 if matches!(spec.flavour, FactorSmoothFlavour::Re) && spec.group_frozen_levels.is_some()
8461 {
8462 continue;
8463 }
8464 return Err(BasisError::InvalidInput(format!(
8465 "factor smooth term '{term_name}' saw an unseen grouping level at row {}",
8466 i + 1
8467 )));
8468 };
8469 let start = level_idx * p;
8470 dense
8471 .slice_mut(s![i, start..start + p])
8472 .assign(&base.row(i));
8473 }
8474
8475 let marginal_penalties: Vec<(Array2<f64>, PenaltySource, f64)> =
8481 if matches!(spec.flavour, FactorSmoothFlavour::Re) {
8482 (0..p)
8483 .map(|j| {
8484 let mut matrix = Array2::<f64>::zeros((p, p));
8485 matrix[[j, j]] = 1.0;
8486 (matrix, PenaltySource::Primary, 1.0)
8487 })
8488 .collect()
8489 } else {
8490 inner
8491 .active_penalties
8492 .iter()
8493 .map(|penalty| {
8494 (
8495 penalty.matrix.clone(),
8496 penalty.info.source.clone(),
8497 penalty.info.normalization_scale,
8498 )
8499 })
8500 .collect()
8501 };
8502
8503 let mut candidates = Vec::<PenaltyCandidate>::with_capacity(marginal_penalties.len());
8504 for (s_inner, source, base_scale) in marginal_penalties {
8505 let mut s_big = Array2::<f64>::zeros((q, q));
8506 for level in 0..n_levels {
8507 let start = level * p;
8508 s_big
8509 .slice_mut(s![start..start + p, start..start + p])
8510 .assign(&s_inner);
8511 }
8512 let (s_big, factor_smooth_scale) = normalize_penalty_in_constrained_space(&s_big);
8513 candidates.push(PenaltyCandidate {
8514 matrix: ConstructiveQuadratic::try_from_dense_psd(
8515 s_big,
8516 "factor-smooth shared penalty",
8517 )?,
8518 source,
8519 normalization_scale: base_scale * factor_smooth_scale,
8520 kronecker_factors: None,
8521 op: None,
8522 });
8523 }
8524
8525 if use_per_dim_null
8555 && let Some(Some(z)) = inner
8556 .active_penalties
8557 .first()
8558 .map(|penalty| &penalty.null_eigenvectors)
8559 && z.nrows() == p
8560 {
8561 let z = canonical_nullspace_directions(z)?;
8562 for k in 0..z.ncols() {
8563 let zk = z.column(k);
8568 let mut p_k = Array2::<f64>::zeros((p, p));
8569 for a in 0..p {
8570 for b in 0..p {
8571 p_k[[a, b]] = zk[a] * zk[b];
8572 }
8573 }
8574 let mut s_null = Array2::<f64>::zeros((q, q));
8575 for level in 0..n_levels {
8576 let start = level * p;
8577 s_null
8578 .slice_mut(s![start..start + p, start..start + p])
8579 .assign(&p_k);
8580 }
8581 let (s_null, null_scale) = normalize_penalty_in_constrained_space(&s_null);
8582 candidates.push(PenaltyCandidate {
8583 matrix: ConstructiveQuadratic::try_from_dense_psd(
8584 s_null,
8585 "factor-smooth null-function penalty",
8586 )?,
8587 source: PenaltySource::Primary,
8588 normalization_scale: null_scale,
8589 kronecker_factors: None,
8590 op: None,
8591 });
8592 }
8593 }
8594 let filtered = crate::basis::filter_penalty_candidates(candidates)?;
8595 let joint_null_rotation = crate::basis::compute_joint_null_rotation(&filtered.active)?;
8596 let mut dropped_penalties = inner.dropped_penalties;
8597 dropped_penalties.extend(filtered.dropped);
8598
8599 let (knots, degree, periodic) = match &inner.metadata {
8602 BasisMetadata::BSpline1D {
8603 knots,
8604 periodic,
8605 degree,
8606 ..
8607 } => (
8608 knots.clone(),
8609 degree.unwrap_or(spec.marginal.degree),
8610 *periodic,
8611 ),
8612 other => {
8613 crate::bail_invalid_basis!(
8614 "factor smooth term '{}' produced an unexpected marginal metadata variant {:?}",
8615 term_name,
8616 other
8617 );
8618 }
8619 };
8620 let flavour_tag = match &spec.flavour {
8621 FactorSmoothFlavour::Fs { .. } => "fs",
8622 FactorSmoothFlavour::Sz => "sz",
8623 FactorSmoothFlavour::Re => "re",
8624 }
8625 .to_string();
8626 let metadata = BasisMetadata::FactorSmooth {
8627 continuous_cols: spec.continuous_cols.clone(),
8628 group_col,
8629 knots,
8630 degree,
8631 periodic,
8632 group_levels: levels,
8633 flavour: flavour_tag,
8634 marginal_is_cr: false,
8637 };
8638
8639 Ok(LocalSmoothTermBuild {
8640 dim: q,
8641 design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(dense)),
8642 affine_offset: inner.affine_offset,
8645 active_penalties: filtered.active,
8646 joint_null_rotation,
8647 dropped_penalties,
8648 metadata,
8649 linear_constraints: None,
8650 box_reparam: false,
8651 kronecker_factored: None,
8652 })
8653}
8654
8655pub fn resolve_factor_smooth_levels(
8659 data: ArrayView2<'_, f64>,
8660 group_col: usize,
8661 spec: &FactorSmoothSpec,
8662 term_name: &str,
8663) -> Result<Vec<u64>, BasisError> {
8664 if let Some(frozen) = &spec.group_frozen_levels {
8665 if frozen.is_empty() {
8666 crate::bail_invalid_basis!(
8667 "factor smooth term '{}' has an empty frozen level list",
8668 term_name
8669 );
8670 }
8671 return Ok(frozen
8672 .iter()
8673 .map(|&b| gam_data::canonical_level_bits(f64::from_bits(b)))
8674 .collect());
8675 }
8676 let mut bits: Vec<u64> = data
8677 .column(group_col)
8678 .iter()
8679 .map(|v| gam_data::canonical_level_bits(*v))
8680 .collect();
8681 bits.sort_by(|a, b| {
8682 f64::from_bits(*a)
8683 .partial_cmp(&f64::from_bits(*b))
8684 .unwrap_or(std::cmp::Ordering::Equal)
8685 });
8686 bits.dedup();
8687 Ok(bits)
8688}
8689
8690pub fn factor_smooth_marginal_for_replay(marginal: &BSplineBasisSpec) -> BSplineBasisSpec {
8697 let mut m = marginal.clone();
8698 m.identifiability = BSplineIdentifiability::None;
8699 m
8700}
8701
8702pub fn build_single_local_smooth_term(
8703 data: ArrayView2<'_, f64>,
8704 term: &SmoothTermSpec,
8705 workspace: &mut crate::basis::BasisWorkspace,
8706) -> Result<LocalSmoothTermBuild, BasisError> {
8707 term.basis.validate_scale_configuration()?;
8708 if term.shape != ShapeConstraint::None && !shape_supports_basis(term) {
8709 crate::bail_invalid_basis!(
8710 "ShapeConstraint::{:?} is unsupported for term '{}'",
8711 term.shape,
8712 term.name
8713 );
8714 }
8715 if let SmoothBasisSpec::ByVariable {
8716 inner,
8717 by_col,
8718 kind,
8719 by,
8720 } = &term.basis
8721 {
8722 ensure_by_variable_specs_match(kind, by, &term.name)?;
8723 let mut inner_basis = (**inner).clone();
8724 if matches!(by, ByVariableSpec::Level { .. }) {
8731 defer_inner_model_centering_to_factor_level_wrapper(&mut inner_basis);
8732 }
8733 let inner_term = SmoothTermSpec {
8734 frozen_parametric_residualization: None,
8735 name: term.name.clone(),
8736 basis: inner_basis,
8737 shape: term.shape,
8738 joint_null_rotation: None,
8739 };
8740 let built = build_single_local_smooth_term(data, &inner_term, workspace)?;
8741 return apply_by_variable_to_local_build(built, data, *by_col, by, &term.name);
8742 }
8743
8744 if let SmoothBasisSpec::BySmooth { smooth, by_kind } = &term.basis {
8747 return build_by_smooth_local(data, term, smooth, by_kind, workspace);
8748 }
8749
8750 let mut built: BasisBuildResult = match &term.basis {
8751 SmoothBasisSpec::FactorSumToZero {
8752 inner,
8753 by_col,
8754 levels,
8755 ..
8756 } => {
8757 if *by_col >= data.ncols() {
8758 crate::bail_dim_basis!(
8759 "term '{}' by column {} out of bounds for {} columns",
8760 term.name,
8761 by_col,
8762 data.ncols()
8763 );
8764 }
8765 if levels.len() < 2 {
8766 crate::bail_invalid_basis!(
8767 "sum-to-zero factor smooth term '{}' requires at least two levels",
8768 term.name
8769 );
8770 }
8771 if term.shape != ShapeConstraint::None {
8772 crate::bail_invalid_basis!(
8773 "ShapeConstraint::{:?} is unsupported for sum-to-zero factor smooth term '{}'",
8774 term.shape,
8775 term.name
8776 );
8777 }
8778 let inner_term = SmoothTermSpec {
8779 frozen_parametric_residualization: None,
8780 name: format!("{}::inner", term.name),
8781 basis: (**inner).clone(),
8782 shape: ShapeConstraint::None,
8783 joint_null_rotation: None,
8784 };
8785 let mut inner_built = build_single_local_smooth_term(data, &inner_term, workspace)?;
8786 if inner_built.affine_offset.is_some() {
8787 crate::bail_invalid_basis!(
8788 "sum-to-zero factor smooth term '{}' cannot contain a non-zero endpoint anchor: a shared fixed affine lift would violate the per-covariate zero-sum deviation identity",
8789 term.name
8790 );
8791 }
8792 let inner_null_eigenvectors = inner_built
8796 .active_penalties
8797 .first()
8798 .and_then(|penalty| penalty.null_eigenvectors.clone());
8799 let base = inner_built
8800 .design
8801 .try_to_dense_by_chunks("sum-to-zero factor smooth")
8802 .map_err(BasisError::InvalidInput)?;
8803 let n = base.nrows();
8804 let p = base.ncols();
8805 let l_minus_one = levels.len() - 1;
8806 let canon_levels: Vec<u64> = levels
8809 .iter()
8810 .map(|&b| gam_data::canonical_level_bits(f64::from_bits(b)))
8811 .collect();
8812 let mut dense = Array2::<f64>::zeros((n, p * l_minus_one));
8813 for i in 0..n {
8814 let bits = gam_data::canonical_level_bits(data[[i, *by_col]]);
8815 let level_idx = canon_levels
8816 .iter()
8817 .position(|b| *b == bits)
8818 .ok_or_else(|| {
8819 BasisError::InvalidInput(format!(
8820 "sum-to-zero factor smooth term '{}' saw an unseen level at row {}",
8821 term.name,
8822 i + 1
8823 ))
8824 })?;
8825 if level_idx < l_minus_one {
8826 let start = level_idx * p;
8827 dense
8828 .slice_mut(s![i, start..start + p])
8829 .assign(&base.row(i));
8830 } else {
8831 for level in 0..l_minus_one {
8832 let start = level * p;
8833 dense
8834 .slice_mut(s![i, start..start + p])
8835 .assign(&base.row(i).mapv(|v| -v));
8836 }
8837 }
8838 }
8839 let mut candidates = Vec::<PenaltyCandidate>::with_capacity(
8840 inner_built.active_penalties.len() * levels.len(),
8841 );
8842 let stz_per_group_penalty =
8877 |s_inner: &Array2<f64>, which_level: usize| -> Array2<f64> {
8878 let mut s_big = Array2::<f64>::zeros((p * l_minus_one, p * l_minus_one));
8879 if which_level < l_minus_one {
8880 let k = which_level;
8882 let mut block = s_big.slice_mut(s![k * p..(k + 1) * p, k * p..(k + 1) * p]);
8883 block.assign(s_inner);
8884 } else {
8885 for a in 0..l_minus_one {
8887 for b in 0..l_minus_one {
8888 let mut block =
8889 s_big.slice_mut(s![a * p..(a + 1) * p, b * p..(b + 1) * p]);
8890 block.assign(s_inner);
8891 }
8892 }
8893 }
8894 s_big
8895 };
8896 for base_penalty in &inner_built.active_penalties {
8897 for which_level in 0..=l_minus_one {
8899 let raw = stz_per_group_penalty(&base_penalty.matrix, which_level);
8900 let (s_big, group_scale) = normalize_penalty_in_constrained_space(&raw);
8901 candidates.push(PenaltyCandidate {
8902 matrix: ConstructiveQuadratic::try_from_dense_psd(
8903 s_big,
8904 "grouped factor-smooth penalty",
8905 )?,
8906 source: base_penalty.info.source.clone(),
8907 normalization_scale: base_penalty.info.normalization_scale * group_scale,
8908 kronecker_factors: None,
8909 op: None,
8910 });
8911 }
8912 }
8913
8914 if let Some(z) = inner_null_eigenvectors.as_ref()
8932 && z.nrows() == p
8933 {
8934 let z = canonical_nullspace_directions(z)?;
8935 for k in 0..z.ncols() {
8936 let zk = z.column(k);
8937 let mut p_k = Array2::<f64>::zeros((p, p));
8938 for a in 0..p {
8939 for b in 0..p {
8940 p_k[[a, b]] = zk[a] * zk[b];
8941 }
8942 }
8943 let stz_pooled_null = {
8948 let mut s_big = Array2::<f64>::zeros((p * l_minus_one, p * l_minus_one));
8949 for a in 0..l_minus_one {
8950 for b in 0..l_minus_one {
8951 let factor = if a == b { 2.0 } else { 1.0 };
8952 let mut block =
8953 s_big.slice_mut(s![a * p..(a + 1) * p, b * p..(b + 1) * p]);
8954 block.assign(&p_k.mapv(|v| v * factor));
8955 }
8956 }
8957 s_big
8958 };
8959 let (s_null, null_scale) =
8960 normalize_penalty_in_constrained_space(&stz_pooled_null);
8961 candidates.push(PenaltyCandidate {
8962 matrix: ConstructiveQuadratic::try_from_dense_psd(
8963 s_null,
8964 "grouped factor-smooth null penalty",
8965 )?,
8966 source: PenaltySource::DoublePenaltyNullspace,
8967 normalization_scale: null_scale,
8968 kronecker_factors: None,
8969 op: None,
8970 });
8971 }
8972 }
8973 let filtered = crate::basis::filter_penalty_candidates(candidates)?;
8974 let mut dropped_penalties = std::mem::take(&mut inner_built.dropped_penalties);
8975 dropped_penalties.extend(filtered.dropped);
8976 inner_built.dim = p * l_minus_one;
8977 inner_built.design =
8978 DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(dense));
8979 inner_built.active_penalties = filtered.active;
8980 inner_built.dropped_penalties = dropped_penalties;
8981 inner_built.joint_null_rotation =
8982 crate::basis::compute_joint_null_rotation(&inner_built.active_penalties)?;
8983 inner_built.kronecker_factored = None;
8984 return Ok(inner_built);
8985 }
8986 SmoothBasisSpec::BSpline1D { feature_col, spec } => {
8987 if *feature_col >= data.ncols() {
8988 crate::bail_dim_basis!(
8989 "term '{}' feature column {} out of bounds for {} columns",
8990 term.name,
8991 feature_col,
8992 data.ncols()
8993 );
8994 }
8995 let mut spec_local = spec.clone();
8996 if term.shape != ShapeConstraint::None {
8997 spec_local.identifiability = BSplineIdentifiability::None;
9000 }
9001 build_bspline_basis_1d(data.column(*feature_col), &spec_local)?
9005 }
9006 SmoothBasisSpec::ThinPlate {
9007 feature_cols,
9008 spec,
9009 input_scale,
9010 } => {
9011 if term.shape != ShapeConstraint::None {
9012 if feature_cols.len() != 1 {
9013 crate::bail_invalid_basis!(
9014 "ShapeConstraint::{:?} for term '{}' on ThinPlate basis requires exactly 1 feature axis; found {}",
9015 term.shape,
9016 term.name,
9017 feature_cols.len()
9018 );
9019 }
9020 }
9021 let mut spec_local = spec.clone();
9022 let frame = term.basis.scale_contract().normalize_euclidean_frame(
9023 select_columns(data, feature_cols)?,
9024 *input_scale,
9025 Some(spec.length_scale),
9026 &mut spec_local.center_strategy,
9027 )?;
9028 let x = frame.coordinates;
9029 let realized_input_scale = frame.input_scale;
9030 let length_scale_eff = frame
9031 .length_scale
9032 .expect("ThinPlate declares a required length-scale coordinate");
9033 spec_local.length_scale = length_scale_eff.standardized_value();
9034 if matches!(
9035 spec_local.identifiability,
9036 SpatialIdentifiability::OrthogonalToParametric
9037 ) {
9038 spec_local.identifiability = SpatialIdentifiability::None;
9039 }
9040 let mut result = build_thin_plate_basis(x.view(), &spec_local).map_err(|err| {
9041 rewrite_thin_plate_knots_error(err, &term.name, feature_cols.len(), spec)
9042 })?;
9043 match &mut result.metadata {
9051 BasisMetadata::ThinPlate {
9052 input_scale: metadata_scale,
9053 length_scale,
9054 ..
9055 } => {
9056 *metadata_scale = realized_input_scale;
9057 *length_scale = crate::OriginalUnits::new(spec.length_scale);
9058 }
9059 BasisMetadata::Duchon {
9060 input_scale: metadata_scale,
9061 length_scale,
9062 ..
9063 } => {
9064 if let Some(promoted) = *length_scale {
9088 let promoted = crate::StandardizedUnits::new(promoted.original_value());
9097 *length_scale = Some(realized_input_scale.to_original_units(promoted));
9098 }
9099 *metadata_scale = realized_input_scale;
9100 }
9101 _ => {
9107 crate::bail_invalid_basis!(
9108 "term '{}' thin-plate build produced metadata that is neither ThinPlate \
9109 nor its Duchon auto-promotion, so the realized input scale cannot be frozen",
9110 term.name
9111 );
9112 }
9113 }
9114 result
9115 }
9116 SmoothBasisSpec::Sphere { feature_cols, spec } => {
9117 if term.shape != ShapeConstraint::None {
9118 crate::bail_invalid_basis!(
9119 "ShapeConstraint::{:?} for term '{}' is not supported on spherical splines",
9120 term.shape,
9121 term.name
9122 );
9123 }
9124 let x = select_columns(data, feature_cols)?;
9125 build_spherical_spline_basis(x.view(), spec)?
9126 }
9127 SmoothBasisSpec::ConstantCurvature { feature_cols, spec } => {
9128 if term.shape != ShapeConstraint::None {
9129 crate::bail_invalid_basis!(
9130 "ShapeConstraint::{:?} for term '{}' is not supported on constant-curvature smooths",
9131 term.shape,
9132 term.name
9133 );
9134 }
9135 let x = select_columns(data, feature_cols)?;
9142 build_constant_curvature_basis(x.view(), spec)?
9143 }
9144 SmoothBasisSpec::MeasureJet {
9145 feature_cols,
9146 spec,
9147 input_scale,
9148 } => {
9149 if term.shape != ShapeConstraint::None {
9150 crate::bail_invalid_basis!(
9151 "ShapeConstraint::{:?} for term '{}' is not supported on measure-jet smooths",
9152 term.shape,
9153 term.name
9154 );
9155 }
9156 let mut spec_local = spec.clone();
9160 let frame = term.basis.scale_contract().normalize_euclidean_frame(
9161 select_columns(data, feature_cols)?,
9162 *input_scale,
9163 Some(spec.length_scale),
9164 &mut spec_local.center_strategy,
9165 )?;
9166 let x = frame.coordinates;
9167 let realized_input_scale = frame.input_scale;
9168 let length_scale_eff = frame
9169 .length_scale
9170 .expect("MeasureJet declares a required length-scale coordinate");
9171 spec_local.length_scale = length_scale_eff.standardized_value();
9172 let mut result = build_measure_jet_basis(x.view(), &spec_local)?;
9173 if let BasisMetadata::MeasureJet {
9174 input_scale: metadata_scale,
9175 ..
9176 } = &mut result.metadata
9177 {
9178 *metadata_scale = realized_input_scale;
9179 }
9180 result
9181 }
9182 SmoothBasisSpec::Matern {
9183 feature_cols,
9184 spec,
9185 input_scale,
9186 } => {
9187 if term.shape != ShapeConstraint::None {
9188 if feature_cols.len() != 1 {
9189 crate::bail_invalid_basis!(
9190 "ShapeConstraint::{:?} for term '{}' on Matern basis requires exactly 1 feature axis; found {}",
9191 term.shape,
9192 term.name,
9193 feature_cols.len()
9194 );
9195 }
9196 }
9197 let original_length_scale = spec.length_scale.resolved().ok_or_else(|| {
9198 BasisError::InvalidInput(format!(
9199 "term '{}' reached Matérn construction before its Auto length scale was resolved",
9200 term.name
9201 ))
9202 })?;
9203 let mut spec_local = spec.clone();
9204 let frame = term.basis.scale_contract().normalize_euclidean_frame(
9205 select_columns(data, feature_cols)?,
9206 *input_scale,
9207 Some(original_length_scale),
9208 &mut spec_local.center_strategy,
9209 )?;
9210 let x = frame.coordinates;
9211 let realized_input_scale = frame.input_scale;
9212 let length_scale_eff = frame
9213 .length_scale
9214 .expect("Matérn declares a required length-scale coordinate");
9215 spec_local
9216 .length_scale
9217 .set_resolved(length_scale_eff.standardized_value());
9218 let mut result = build_matern_basiswithworkspace(x.view(), &spec_local, workspace)?;
9219 if let BasisMetadata::Matern {
9220 input_scale: metadata_scale,
9221 length_scale,
9222 ..
9223 } = &mut result.metadata
9224 {
9225 *metadata_scale = realized_input_scale;
9226 *length_scale = crate::OriginalUnits::new(original_length_scale);
9227 }
9228 result
9229 }
9230 SmoothBasisSpec::Duchon {
9231 feature_cols,
9232 spec,
9233 input_scale,
9234 } => {
9235 if term.shape != ShapeConstraint::None {
9236 if feature_cols.len() != 1 {
9237 crate::bail_invalid_basis!(
9238 "ShapeConstraint::{:?} for term '{}' on Duchon basis requires exactly 1 feature axis; found {}",
9239 term.shape,
9240 term.name,
9241 feature_cols.len()
9242 );
9243 }
9244 }
9245 let mut spec_local = spec.clone();
9246 let frame = term.basis.scale_contract().normalize_euclidean_frame(
9247 select_columns(data, feature_cols)?,
9248 *input_scale,
9249 spec.length_scale,
9250 &mut spec_local.center_strategy,
9251 )?;
9252 let x = frame.coordinates;
9253 let realized_input_scale = frame.input_scale;
9254 let length_scale_eff = frame.length_scale;
9255 spec_local.length_scale =
9256 length_scale_eff.map(crate::StandardizedUnits::standardized_value);
9257 if let crate::basis::OneDimensionalBoundary::Cyclic { start, end } =
9270 spec_local.boundary.clone()
9271 {
9272 spec_local.boundary = crate::basis::OneDimensionalBoundary::Cyclic {
9273 start: realized_input_scale
9274 .to_standardized_units(crate::OriginalUnits::new(start))
9275 .standardized_value(),
9276 end: realized_input_scale
9277 .to_standardized_units(crate::OriginalUnits::new(end))
9278 .standardized_value(),
9279 };
9280 }
9281 if let Some(periods) = spec_local.periodic.as_mut() {
9289 for axis_period in periods {
9290 if let Some(period) = axis_period.as_mut() {
9291 *period = realized_input_scale
9292 .to_standardized_units(crate::OriginalUnits::new(*period))
9293 .standardized_value();
9294 }
9295 }
9296 }
9297 if matches!(
9298 spec_local.identifiability,
9299 SpatialIdentifiability::OrthogonalToParametric
9300 ) {
9301 spec_local.identifiability = SpatialIdentifiability::None;
9302 }
9303 let mut result = build_duchon_basiswithworkspace(x.view(), &spec_local, workspace)?;
9304 if let BasisMetadata::Duchon {
9305 input_scale: metadata_scale,
9306 length_scale,
9307 periodic,
9308 ..
9309 } = &mut result.metadata
9310 {
9311 *metadata_scale = realized_input_scale;
9312 *length_scale = spec.length_scale.map(crate::OriginalUnits::new);
9313 if spec.periodic.is_some() || spec.boundary.period().is_some() {
9329 *periodic = spec
9330 .periodic
9331 .clone()
9332 .or_else(|| spec.boundary.period().map(|(_, _, p)| vec![Some(p)]));
9333 }
9334 }
9335 result
9336 }
9337 SmoothBasisSpec::Pca {
9338 feature_cols,
9339 basis_matrix,
9340 centered,
9341 smooth_penalty,
9342 center_mean,
9343 pca_basis_path,
9344 chunk_size,
9345 } => {
9346 if term.shape != ShapeConstraint::None {
9347 crate::bail_invalid_basis!(
9348 "ShapeConstraint::{:?} for term '{}' is not supported on Pca basis",
9349 term.shape,
9350 term.name
9351 );
9352 }
9353 build_pca_smooth_basis(
9354 data,
9355 feature_cols,
9356 basis_matrix,
9357 *centered,
9358 *smooth_penalty,
9359 center_mean.as_ref(),
9360 pca_basis_path.as_ref(),
9361 *chunk_size,
9362 )?
9363 }
9364 SmoothBasisSpec::TensorBSpline { feature_cols, spec } => {
9365 build_tensor_bspline_basis(data, feature_cols, spec)?
9366 }
9367 SmoothBasisSpec::ByVariable { .. } => {
9368 crate::bail_invalid_basis!(
9369 "internal: ByVariable smooths must return before inner basis dispatch"
9370 );
9371 }
9372 SmoothBasisSpec::BySmooth { .. } => {
9373 crate::bail_invalid_basis!("internal: BySmooth smooths must be lowered to ByVariable before inner basis dispatch"
9374 .to_string(),);
9375 }
9376 SmoothBasisSpec::FactorSmooth { spec } => {
9377 if term.shape != ShapeConstraint::None {
9378 crate::bail_invalid_basis!(
9379 "ShapeConstraint::{:?} is unsupported for factor smooth term '{}'",
9380 term.shape,
9381 term.name
9382 );
9383 }
9384 return build_factor_smooth(data, spec, &term.name, workspace);
9385 }
9386 };
9387
9388 if let SmoothBasisSpec::Matern { .. } = &term.basis {
9404 let filtered = matern_operator_penalty_triplet_from_metadata(&built.metadata)?;
9405 built.active_penalties = filtered.active;
9406 built.dropped_penalties = filtered.dropped;
9407 }
9408
9409 if built.affine_offset.is_some() && term.shape != ShapeConstraint::None {
9410 crate::bail_invalid_basis!(
9411 "non-zero endpoint anchors cannot be combined with ShapeConstraint::{:?} on term '{}': the coefficient cone constrains only the homogeneous spline and would not certify the final affine function",
9412 term.shape,
9413 term.name
9414 );
9415 }
9416 let p_local = built.design.ncols();
9417 let affine_offset = built.affine_offset;
9418 let mut metadata = built.metadata.clone();
9419 let kron_factored = if term.shape == ShapeConstraint::None {
9422 built.kronecker_factored
9423 } else {
9424 None
9425 };
9426 let mut design_t = built.design;
9427 let mut penalties_t = built.active_penalties;
9428 let mut dropped_penalties_t = built.dropped_penalties;
9429 if matches!(
9430 spatial_identifiability_policy(term),
9431 Some(SpatialIdentifiability::OrthogonalToParametric)
9432 ) {
9433 metadata = freeze_raw_spatial_metadata(metadata, design_t.ncols());
9434 }
9435
9436 let use_box_reparam =
9437 term.shape != ShapeConstraint::None && shape_uses_box_reparameterization(&term.basis);
9438 if let Some((order, sign)) = shape_order_and_sign(term.shape)
9439 && use_box_reparam
9440 {
9441 let t = if order == 2 {
9455 let (knots, degree) = match &metadata {
9456 BasisMetadata::BSpline1D {
9457 knots,
9458 degree: Some(degree),
9459 periodic,
9460 ..
9461 } if periodic.is_none() => (knots, *degree),
9462 _ => {
9463 crate::bail_invalid_basis!(
9464 "shape-constrained convex/concave term '{}' requires realized open B-spline knot and degree metadata",
9465 term.name
9466 );
9467 }
9468 };
9469 let spans = bspline_first_derivative_control_spans(knots.view(), degree)?;
9470 if spans.len() + 1 != p_local {
9471 crate::bail_invalid_basis!(
9472 "shape-constraint derivative-control span count {} does not match basis dim {} for term '{}'",
9473 spans.len(),
9474 p_local,
9475 term.name
9476 );
9477 }
9478 convex_derivative_control_transform_matrix(&spans, sign)?
9479 } else {
9480 cumulative_sum_transform_matrix(p_local, order, sign)
9481 };
9482 let inner_dense = match design_t {
9486 DesignMatrix::Dense(d) => d,
9487 DesignMatrix::Sparse(sp) => gam_linalg::matrix::DenseDesignMatrix::from(
9488 sp.try_to_dense_arc("shape-constrained coefficient transform")
9489 .map_err(BasisError::InvalidInput)?,
9490 ),
9491 };
9492 let coeff_op =
9493 gam_linalg::matrix::CoefficientTransformOperator::new(inner_dense, t.clone()).map_err(
9494 |e| BasisError::InvalidInput(format!("CoefficientTransformOperator: {e}")),
9495 )?;
9496 design_t = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(
9497 coeff_op,
9498 )));
9499 for penalty in &mut penalties_t {
9506 let tt_s = fast_atb(&t, &penalty.matrix);
9507 penalty.matrix = fast_ab(&tt_s, &t);
9508 penalty.op = None;
9509 penalty.info.kronecker_factors = None;
9510 penalty.info.structural_null_frame = None;
9519 }
9520 }
9521 let penalty_candidates = penalties_t
9522 .into_iter()
9523 .map(|penalty| -> Result<PenaltyCandidate, BasisError> {
9524 let ActivePenalty {
9525 matrix,
9526 op: op_in,
9527 info,
9528 ..
9529 } = penalty;
9530 let (matrix, c_new) = normalize_penalty_in_constrained_space(&matrix);
9531 let normalization_scale = info.normalization_scale * c_new;
9532 let op_scale = 1.0 / c_new;
9533 let kronecker_scale = 1.0 / c_new;
9534 let scaled_op = if op_scale > 0.0 && op_scale.is_finite() {
9537 op_in.map(|op| {
9538 std::sync::Arc::new(crate::analytic_penalties::ScaledPenaltyOp::new(
9539 op, op_scale,
9540 ))
9541 as std::sync::Arc<dyn crate::analytic_penalties::PenaltyOp>
9542 })
9543 } else {
9544 None
9545 };
9546 let kronecker_factors = info.kronecker_factors.map(|mut factors| {
9547 if let Some(first) = factors.first_mut() {
9548 first.mapv_inplace(|v| v * kronecker_scale);
9549 }
9550 factors
9551 });
9552 let structural_null_frame = info.structural_null_frame;
9565 let matrix = ConstructiveQuadratic::try_from_dense_psd(
9566 matrix,
9567 "shape-constrained transformed penalty",
9568 )?;
9569 let matrix = match structural_null_frame {
9570 Some(frame) => matrix.with_structural_null_frame(
9571 frame,
9572 "renormalized smooth penalty structural frame",
9573 )?,
9574 None => matrix,
9575 };
9576 Ok(PenaltyCandidate {
9577 matrix,
9578 source: info.source,
9579 normalization_scale,
9580 kronecker_factors,
9581 op: scaled_op,
9582 })
9583 })
9584 .collect::<Result<Vec<_>, _>>()?;
9585 let filtered = crate::basis::filter_penalty_candidates(penalty_candidates)?;
9586 dropped_penalties_t.extend(filtered.dropped);
9587 let joint_null_rotation = match term.joint_null_rotation.clone() {
9606 Some(persisted) => Some(persisted),
9607 None if smooth_has_frozen_identifiability(term) => None,
9608 None if kron_factored.is_some() => None,
9609 None => crate::basis::compute_joint_null_rotation(&filtered.active)?,
9610 };
9611
9612 Ok(LocalSmoothTermBuild {
9613 dim: p_local,
9614 design: design_t,
9615 affine_offset,
9616 active_penalties: filtered.active,
9617 joint_null_rotation,
9618 dropped_penalties: dropped_penalties_t,
9619 metadata,
9620 linear_constraints: None,
9621 box_reparam: use_box_reparam,
9622 kronecker_factored: kron_factored,
9623 })
9624}
9625
9626pub fn build_smooth_design(
9627 data: ArrayView2<'_, f64>,
9628 terms: &[SmoothTermSpec],
9629) -> Result<RawSmoothDesign, BasisError> {
9630 let mut ws = crate::basis::BasisWorkspace::new();
9631 build_smooth_design_withworkspace(data, terms, &mut ws)
9632}
9633
9634pub fn build_smooth_design_withworkspace(
9641 data: ArrayView2<'_, f64>,
9642 terms: &[SmoothTermSpec],
9643 workspace: &mut crate::basis::BasisWorkspace,
9644) -> Result<RawSmoothDesign, BasisError> {
9645 validate_smooth_terms_finite_inputs(data, terms)?;
9646 build_smooth_design_withworkspace_unvalidated(data, terms, workspace)
9647}
9648
9649pub fn build_smooth_design_withworkspace_unvalidated(
9650 data: ArrayView2<'_, f64>,
9651 terms: &[SmoothTermSpec],
9652 workspace: &mut crate::basis::BasisWorkspace,
9653) -> Result<RawSmoothDesign, BasisError> {
9654 let mut planned_blocks = plan_joint_spatial_centers_for_term_blocks(data, &[terms.to_vec()])?;
9655 let planned_terms = planned_blocks.pop().ok_or_else(|| {
9656 BasisError::InvalidInput(
9657 "joint spatial center planner returned no smooth blocks".to_string(),
9658 )
9659 })?;
9660 let policy = workspace.policy().clone();
9661 let local_builds: Vec<LocalSmoothTermBuild> = {
9662 use rayon::iter::{IntoParallelIterator, ParallelIterator};
9663 planned_terms
9664 .into_par_iter()
9665 .map(|term| {
9666 let mut term_workspace = crate::basis::BasisWorkspace::with_policy(policy.clone());
9667 build_single_local_smooth_term(data, &term, &mut term_workspace)
9668 })
9669 .collect::<Result<Vec<_>, _>>()?
9670 };
9671
9672 let total_p: usize = local_builds.iter().map(|built| built.dim).sum();
9673
9674 let mut local_designs: Vec<DesignMatrix> = Vec::with_capacity(local_builds.len());
9675 let mut affine_offset = Array1::<f64>::zeros(data.nrows());
9676 let mut terms_out = Vec::<SmoothTerm>::with_capacity(terms.len());
9677 let mut penalties_global = Vec::<BlockwisePenalty>::new();
9678 let mut nullspace_dims_global = Vec::<usize>::new();
9679 let mut penaltyinfo_global = Vec::<PenaltyBlockInfo>::new();
9680 let mut dropped_penaltyinfo_global = Vec::<DroppedPenaltyBlockInfo>::new();
9681 let mut coefficient_lower_bounds = Array1::<f64>::from_elem(total_p, f64::NEG_INFINITY);
9682 let mut any_bounds = false;
9683 let mut linear_constraintsrows: Vec<(usize, usize, Array1<f64>)> = Vec::new();
9688 let mut linear_constraints_b: Vec<f64> = Vec::new();
9689
9690 let mut col_start = 0usize;
9691 for (term, mut built) in terms.iter().zip(local_builds.into_iter()) {
9692 let p_local = built.dim;
9693 let col_end = col_start + p_local;
9694 let lb_local = if built.box_reparam {
9695 shape_lower_bounds_local(term.shape, p_local)
9696 } else {
9697 None
9698 };
9699
9700 let applied_rotation: Option<crate::basis::JointNullRotation> = match (
9732 built.joint_null_rotation.take(),
9733 lb_local.is_some(),
9734 built.linear_constraints.is_some(),
9735 ) {
9736 (Some(rot), false, false) => {
9737 let q = &rot.rotation;
9738 built.design =
9739 apply_smooth_transform_to_design(built.design.clone(), q, &term.name)?;
9740 for penalty in &mut built.active_penalties {
9741 let qt_s = gam_linalg::faer_ndarray::fast_atb(q, &penalty.matrix);
9742 penalty.matrix = gam_linalg::faer_ndarray::fast_ab(&qt_s, q);
9743 penalty.null_eigenvectors = penalty
9744 .null_eigenvectors
9745 .as_ref()
9746 .map(|basis| gam_linalg::faer_ndarray::fast_atb(q, basis));
9747 penalty.info.structural_null_frame = penalty
9757 .info
9758 .structural_null_frame
9759 .as_ref()
9760 .map(|frame| gam_linalg::faer_ndarray::fast_atb(q, frame));
9761 penalty.op = None;
9762 penalty.info.kronecker_factors = None;
9763 }
9764 built.kronecker_factored = None;
9765 Some(rot)
9766 }
9767 (Some(_), _, _) => None,
9768 (None, _, _) => None,
9769 };
9770
9771 for active_penalty in &built.active_penalties {
9772 let global_index = penalties_global.len();
9773 penalties_global.push(
9774 BlockwisePenalty::new(col_start..col_end, active_penalty.matrix.clone())
9775 .with_op(active_penalty.op.clone()),
9776 );
9777 nullspace_dims_global.push(active_penalty.nullity);
9778 penaltyinfo_global.push(PenaltyBlockInfo {
9779 global_index,
9780 termname: Some(term.name.clone()),
9781 penalty: active_penalty.info.clone(),
9782 });
9783 }
9784 for info in &built.dropped_penalties {
9785 dropped_penaltyinfo_global.push(DroppedPenaltyBlockInfo {
9786 termname: Some(term.name.clone()),
9787 penalty: info.clone(),
9788 });
9789 }
9790
9791 if let Some(lin_local) = &built.linear_constraints {
9792 for r in 0..lin_local.a.nrows() {
9793 linear_constraintsrows.push((col_start, col_end, lin_local.a.row(r).to_owned()));
9794 linear_constraints_b.push(lin_local.b[r]);
9795 }
9796 }
9797 if let Some(lb_local) = &lb_local {
9798 coefficient_lower_bounds
9799 .slice_mut(s![col_start..col_end])
9800 .assign(lb_local);
9801 any_bounds = true;
9802 }
9803
9804 if let Some(term_offset) = built.affine_offset.as_ref() {
9805 if term_offset.len() != data.nrows() {
9806 crate::bail_dim_basis!(
9807 "smooth term '{}' affine offset has {} rows but the realized data has {}",
9808 term.name,
9809 term_offset.len(),
9810 data.nrows()
9811 );
9812 }
9813 affine_offset += term_offset;
9814 }
9815
9816 local_designs.push(built.design);
9818
9819 terms_out.push(SmoothTerm {
9820 parametric_residualization: None,
9821 name: term.name.clone(),
9822 coeff_range: col_start..col_end,
9823 shape: term.shape,
9824 active_penalties: built.active_penalties,
9825 dropped_penalties: built.dropped_penalties,
9826 metadata: built.metadata,
9827 lower_bounds_local: lb_local,
9828 linear_constraints_local: built.linear_constraints,
9829 kronecker_factored: built.kronecker_factored.take(),
9830 joint_null_rotation: applied_rotation,
9831 unabsorbed_global_orthogonality: None,
9832 collection_gauge: None,
9835 });
9836
9837 col_start = col_end;
9838 }
9839
9840 assert_eq!(
9841 penalties_global.len(),
9842 nullspace_dims_global.len(),
9843 "global smooth penalty/nullspace bookkeeping diverged"
9844 );
9845 assert_eq!(
9846 penalties_global.len(),
9847 penaltyinfo_global.len(),
9848 "global smooth penalty metadata bookkeeping diverged"
9849 );
9850
9851 Ok(RawSmoothDesign {
9852 term_designs: local_designs,
9853 affine_offset,
9854 penalties: penalties_global,
9855 nullspace_dims: nullspace_dims_global,
9856 penaltyinfo: penaltyinfo_global,
9857 dropped_penaltyinfo: dropped_penaltyinfo_global,
9858 terms: terms_out,
9859 coefficient_lower_bounds: if any_bounds {
9860 Some(coefficient_lower_bounds)
9861 } else {
9862 None
9863 },
9864 linear_constraints: if linear_constraintsrows.is_empty() {
9865 None
9866 } else {
9867 let mut a = Array2::<f64>::zeros((linear_constraintsrows.len(), total_p));
9868 for (i, (cs, ce, values)) in linear_constraintsrows.iter().enumerate() {
9869 a.row_mut(i).slice_mut(s![*cs..*ce]).assign(values);
9870 }
9871 Some(LinearInequalityConstraints {
9872 a,
9873 b: Array1::from_vec(linear_constraints_b),
9874 })
9875 },
9876 })
9877}
9878
9879#[cfg(test)]
9880mod factor_smooth_heldout_group_tests {
9881 use super::*;
9882 use crate::basis::BasisWorkspace;
9883 use ndarray::{Array1, array};
9884
9885 fn pinned_marginal() -> BSplineBasisSpec {
9886 BSplineBasisSpec {
9887 degree: 3,
9888 penalty_order: 2,
9889 knotspec: BSplineKnotSpec::Provided(Array1::from(vec![
9890 0.0, 0.0, 0.0, 0.0, 0.25, 0.6, 1.0, 1.0, 1.0, 1.0,
9891 ])),
9892 double_penalty: false,
9893 identifiability: BSplineIdentifiability::None,
9894 boundary: crate::basis::OneDimensionalBoundary::Open,
9895 boundary_conditions: crate::basis::BSplineBoundaryConditions::default(),
9896 }
9897 }
9898
9899 fn factor_smooth_term(
9900 flavour: FactorSmoothFlavour,
9901 frozen: Option<Vec<u64>>,
9902 ) -> SmoothTermSpec {
9903 SmoothTermSpec {
9904 frozen_parametric_residualization: None,
9905 name: "fs_heldout".to_string(),
9906 basis: SmoothBasisSpec::FactorSmooth {
9907 spec: FactorSmoothSpec {
9908 continuous_cols: vec![0],
9909 group_col: 1,
9910 marginal: pinned_marginal(),
9911 flavour,
9912 group_frozen_levels: frozen,
9913 frozen_global_orthogonality: None,
9914 },
9915 },
9916 shape: ShapeConstraint::None,
9917 joint_null_rotation: None,
9918 }
9919 }
9920
9921 const FROZEN_01: [f64; 2] = [0.0, 1.0];
9922
9923 fn frozen_bits() -> Vec<u64> {
9924 FROZEN_01.iter().map(|v| v.to_bits()).collect()
9925 }
9926
9927 #[test]
9933 fn re_heldout_group_row_is_zero_deviation() {
9934 let data = array![[0.1, 0.0], [0.5, 1.0], [0.9, 7.0]];
9935 let term = factor_smooth_term(FactorSmoothFlavour::Re, Some(frozen_bits()));
9936 let mut workspace = BasisWorkspace::default();
9937 let build = build_single_local_smooth_term(data.view(), &term, &mut workspace)
9938 .expect("a held-out group must not fail the bs=\"re\" design build");
9939 let dense = build
9940 .design
9941 .try_to_dense_by_chunks("heldout test")
9942 .expect("dense");
9943 assert!(
9944 dense.row(2).iter().all(|&v| v == 0.0),
9945 "unseen-group row must carry zero deviation across every group block, got {:?}",
9946 dense.row(2)
9947 );
9948 assert!(
9949 dense.row(0).iter().any(|&v| v != 0.0) && dense.row(1).iter().any(|&v| v != 0.0),
9950 "in-vocabulary rows must still populate their group blocks"
9951 );
9952 }
9953
9954 #[test]
9958 fn fs_heldout_group_stays_strict() {
9959 let data = array![[0.1, 0.0], [0.5, 1.0], [0.9, 7.0]];
9960 let term = factor_smooth_term(
9961 FactorSmoothFlavour::Fs {
9962 m_null_penalty_orders: vec![1],
9963 },
9964 Some(frozen_bits()),
9965 );
9966 let mut workspace = BasisWorkspace::default();
9967 let err = match build_single_local_smooth_term(data.view(), &term, &mut workspace) {
9968 Ok(_) => panic!("fs must reject an unseen grouping level"),
9969 Err(err) => err,
9970 };
9971 assert!(
9972 err.to_string().contains("unseen grouping level"),
9973 "fs unseen-level refusal must name the defect, got: {err}"
9974 );
9975 }
9976}
9977
9978#[cfg(test)]
9979mod linear_term_contract_tests {
9980 use super::LinearTermSpec;
9981
9982 #[test]
9983 fn missing_linear_double_penalty_deserializes_to_unpenalized_mle() {
9984 let term: LinearTermSpec = serde_json::from_str(r#"{"name":"x","feature_col":0}"#)
9985 .expect("minimal saved linear term");
9986 assert!(
9987 !term.double_penalty,
9988 "descriptor and formula defaults must both preserve parametric MLE semantics"
9989 );
9990 }
9991}