1use coefficient_transforms::{
2 convex_derivative_control_transform_matrix, cumulative_exp, cumulative_sum_transform_matrix,
3 second_cumulative_exp,
4};
5
6pub use error::SmoothError;
7
8use input_standardization::{
9 apply_input_standardization, compensate_length_scale_for_standardization,
10 compensate_optional_length_scale_for_standardization, compute_spatial_input_scales,
11};
12
13use shape_constraints::{
14 bspline_first_derivative_control_spans, shape_lower_bounds_local, shape_order_and_sign,
15 shape_supports_basis, shape_uses_box_reparameterization,
16};
17
18pub fn describe_thin_plate_center_request(strategy: &CenterStrategy) -> String {
19 match strategy {
20 CenterStrategy::Auto(inner) => describe_thin_plate_center_request(inner),
21 CenterStrategy::UserProvided(centers) => format!("{} centers", centers.nrows()),
22 CenterStrategy::EqualMass { num_centers }
23 | CenterStrategy::EqualMassCovarRepresentative { num_centers }
24 | CenterStrategy::FarthestPoint { num_centers }
25 | CenterStrategy::KMeans { num_centers, .. } => format!("{num_centers} centers"),
26 CenterStrategy::UniformGrid { points_per_dim } => {
27 format!("uniform grid with {points_per_dim} points per dimension")
28 }
29 }
30}
31
32pub fn rewrite_thin_plate_knots_error(
33 err: BasisError,
34 termname: &str,
35 feature_count: usize,
36 spec: &ThinPlateBasisSpec,
37) -> BasisError {
38 match err {
39 BasisError::InvalidInput(msg)
42 if msg.contains("thin-plate spline requires at least")
43 && (msg.contains("centers to span") || msg.contains("knots to span")) =>
44 {
45 let min_centers = crate::basis::thin_plate_polynomial_basis_dimension(feature_count);
46 let requested = describe_thin_plate_center_request(&spec.center_strategy);
47 BasisError::InvalidInput(format!(
48 "joint TPS term '{termname}' over {feature_count} covariates with {requested} is invalid; minimum centers is {min_centers}"
49 ))
50 }
51 BasisError::InvalidInput(msg)
56 if msg.starts_with("requested ") && msg.contains(" knots but only ") =>
57 {
58 let min_centers = crate::basis::thin_plate_polynomial_basis_dimension(feature_count);
59 let requested = describe_thin_plate_center_request(&spec.center_strategy);
60 BasisError::InvalidInput(format!(
61 "joint TPS term '{termname}' over {feature_count} covariates with {requested} is invalid; minimum centers is {min_centers}"
62 ))
63 }
64 other => other,
65 }
66}
67
68#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
69pub enum ShapeConstraint {
70 None,
71 MonotoneIncreasing,
72 MonotoneDecreasing,
73 Convex,
74 Concave,
75}
76
77pub fn parse_shape_constraint(raw: &str) -> Result<ShapeConstraint, String> {
88 let normalized = raw.trim().to_ascii_lowercase().replace('-', "_");
89 match normalized.as_str() {
90 "" | "none" => Ok(ShapeConstraint::None),
91 "monotone_increasing" | "monotonic_increasing" | "increasing" | "mono_inc" | "mpi" => {
92 Ok(ShapeConstraint::MonotoneIncreasing)
93 }
94 "monotone_decreasing" | "monotonic_decreasing" | "decreasing" | "mono_dec" | "mpd" => {
95 Ok(ShapeConstraint::MonotoneDecreasing)
96 }
97 "convex" | "cvx" => Ok(ShapeConstraint::Convex),
98 "concave" | "ccv" => Ok(ShapeConstraint::Concave),
99 other => Err(format!(
100 "unknown shape constraint {other:?}; expected one of \
101 \"none\", \"monotone_increasing\", \"monotone_decreasing\", \
102 \"convex\", \"concave\""
103 )),
104 }
105}
106
107impl ShapeConstraint {
108 pub fn dsl_str(&self) -> &'static str {
111 match self {
112 ShapeConstraint::None => "none",
113 ShapeConstraint::MonotoneIncreasing => "monotone_increasing",
114 ShapeConstraint::MonotoneDecreasing => "monotone_decreasing",
115 ShapeConstraint::Convex => "convex",
116 ShapeConstraint::Concave => "concave",
117 }
118 }
119}
120
121pub const SMOOTH_HEAD_KEYWORDS: [&str; 11] = [
124 "s",
125 "smooth",
126 "te",
127 "tensor",
128 "thinplate",
129 "tps",
130 "duchon",
131 "matern",
132 "sphere",
133 "bs",
134 "bspline",
135];
136
137pub fn apply_shape_constraints_to_formula(
150 formula: &str,
151 constraints: &[(String, String)],
152) -> Result<String, String> {
153 use std::collections::{BTreeMap, BTreeSet};
154
155 if constraints.is_empty() {
156 return Ok(formula.to_string());
157 }
158 let strip_ws = |s: &str| -> String { s.chars().filter(|c| !c.is_whitespace()).collect() };
159
160 let mut wanted: BTreeMap<String, &'static str> = BTreeMap::new();
162 let mut originals: BTreeMap<String, String> = BTreeMap::new();
164 for (key, kind_raw) in constraints {
165 let kind = parse_shape_constraint(kind_raw)?;
166 let nk = strip_ws(key);
167 originals.entry(nk.clone()).or_insert_with(|| key.clone());
168 if kind != ShapeConstraint::None {
169 wanted.insert(nk, kind.dsl_str());
170 }
171 }
172 if wanted.is_empty() {
173 return Ok(formula.to_string());
174 }
175
176 let chars: Vec<char> = formula.chars().collect();
177 let n = chars.len();
178 let is_ident = |c: char| c.is_ascii_alphanumeric() || c == '_';
179
180 let mut out = String::with_capacity(formula.len() + 32);
181 let mut matched: BTreeSet<String> = BTreeSet::new();
182 let mut i = 0usize;
183 while i < n {
184 let mut head: Option<(usize, usize)> = None; let mut p = i;
188 while p < n {
189 let boundary = p == 0 || !is_ident(chars[p - 1]);
190 if boundary {
191 for kw in SMOOTH_HEAD_KEYWORDS.iter() {
192 let klen = kw.chars().count();
193 if p + klen > n || chars[p..p + klen].iter().collect::<String>() != **kw {
194 continue;
195 }
196 let mut q = p + klen;
197 while q < n && chars[q].is_whitespace() {
198 q += 1;
199 }
200 if q < n && chars[q] == '(' {
201 head = Some((p, q));
202 break;
203 }
204 }
205 }
206 if head.is_some() {
207 break;
208 }
209 p += 1;
210 }
211 let (head_start, paren_open) = match head {
212 Some(h) => h,
213 None => {
214 out.extend(chars[i..].iter());
215 break;
216 }
217 };
218 out.extend(chars[i..head_start].iter());
219
220 let body_start = paren_open + 1;
222 let mut depth = 1i32;
223 let mut j = body_start;
224 let mut in_str: Option<char> = None;
225 let mut closed = false;
226 while j < n {
227 let ch = chars[j];
228 if let Some(quote) = in_str {
229 if ch == quote {
230 in_str = None;
231 }
232 } else if ch == '\'' || ch == '"' {
233 in_str = Some(ch);
234 } else if ch == '(' {
235 depth += 1;
236 } else if ch == ')' {
237 depth -= 1;
238 if depth == 0 {
239 closed = true;
240 break;
241 }
242 }
243 j += 1;
244 }
245
246 if !closed {
247 out.extend(chars[head_start..].iter());
250 break;
251 }
252
253 let term_text: String = chars[head_start..=j].iter().collect();
254
255 let key_norm = strip_ws(&term_text);
256
257 match wanted.get(&key_norm) {
258 None => out.extend(chars[head_start..=j].iter()),
259 Some(kind) => {
260 let head_paren: String = chars[head_start..body_start].iter().collect();
261 let inside: String = chars[body_start..j].iter().collect();
262 let inside = inside.trim();
263 if inside.is_empty() {
264 out.push_str(&format!("{head_paren}shape={kind})"));
265 } else {
266 out.push_str(&format!("{head_paren}{inside}, shape={kind})"));
267 }
268 matched.insert(key_norm);
269 }
270 }
271
272 i = j + 1;
273 }
274
275 let mut missing: Vec<String> = wanted
276 .keys()
277 .filter(|k| !matched.contains(*k))
278 .map(|k| originals.get(k).cloned().unwrap_or_else(|| k.clone()))
279 .collect();
280
281 if !missing.is_empty() {
282 missing.sort();
283 return Err(format!(
284 "shape constraints referenced smooth term(s) not found in formula: {}",
285 missing.join(", ")
286 ));
287 }
288
289 Ok(out)
290}
291
292#[derive(Debug, Clone, Serialize, Deserialize)]
293pub enum BySmoothKind {
294 Numeric,
295 Level { level_bits: u64 },
296}
297
298#[derive(Debug, Clone, Serialize, Deserialize)]
299pub enum SmoothBasisSpec {
300 ByVariable {
310 inner: Box<SmoothBasisSpec>,
311 by_col: usize,
312 kind: BySmoothKind,
313 by: ByVariableSpec,
314 },
315 FactorSumToZero {
319 inner: Box<SmoothBasisSpec>,
320 by_col: usize,
321 levels: Vec<u64>,
322 #[serde(default)]
333 frozen_global_orthogonality: Option<Array2<f64>>,
334 },
335 BSpline1D {
336 feature_col: usize,
337 spec: BSplineBasisSpec,
338 },
339 BySmooth {
342 smooth: Box<SmoothBasisSpec>,
343 by_kind: ByVarKind,
344 },
345 FactorSmooth { spec: FactorSmoothSpec },
348 ThinPlate {
349 feature_cols: Vec<usize>,
350 spec: ThinPlateBasisSpec,
351 #[serde(default)]
355 input_scales: Option<Vec<f64>>,
356 },
357 Sphere {
358 feature_cols: Vec<usize>,
359 spec: SphericalSplineBasisSpec,
360 },
361 ConstantCurvature {
367 feature_cols: Vec<usize>,
368 spec: ConstantCurvatureBasisSpec,
369 },
370 Matern {
371 feature_cols: Vec<usize>,
372 spec: MaternBasisSpec,
373 #[serde(default)]
374 input_scales: Option<Vec<f64>>,
375 },
376 MeasureJet {
382 feature_cols: Vec<usize>,
383 spec: MeasureJetBasisSpec,
384 #[serde(default)]
385 input_scales: Option<Vec<f64>>,
386 },
387 Duchon {
388 feature_cols: Vec<usize>,
389 spec: DuchonBasisSpec,
390 #[serde(default)]
391 input_scales: Option<Vec<f64>>,
392 },
393 Pca {
394 feature_cols: Vec<usize>,
395 basis_matrix: Array2<f64>,
396 centered: bool,
397 #[serde(default = "default_pca_smooth_penalty")]
398 smooth_penalty: f64,
399 #[serde(default)]
400 center_mean: Option<Array1<f64>>,
401 #[serde(default)]
402 pca_basis_path: Option<PathBuf>,
403 #[serde(default = "default_pca_chunk_size")]
404 chunk_size: usize,
405 },
406 TensorBSpline {
411 feature_cols: Vec<usize>,
412 spec: TensorBSplineSpec,
413 },
414}
415
416impl SmoothBasisSpec {
417 pub fn min_sample_rows(&self) -> usize {
434 const RADIAL_FLOOR: usize = 5;
439
440 match self {
441 Self::ByVariable { inner, .. } => inner.min_sample_rows(),
442 Self::FactorSumToZero { inner, levels, .. } => {
443 let inner_min = inner.min_sample_rows();
447 let lvls = levels.len().saturating_sub(1).max(1);
448 inner_min.saturating_mul(lvls)
449 }
450 Self::BSpline1D { spec, .. } => bspline_basis_min_rows(spec),
451 Self::BySmooth { smooth, .. } => smooth.min_sample_rows(),
452 Self::FactorSmooth { spec } => {
453 bspline_basis_min_rows(&spec.marginal)
457 }
458 Self::ThinPlate { .. }
459 | Self::Sphere { .. }
460 | Self::ConstantCurvature { .. }
461 | Self::Matern { .. }
462 | Self::MeasureJet { .. }
463 | Self::Duchon { .. } => RADIAL_FLOOR,
464 Self::Pca { basis_matrix, .. } => basis_matrix.ncols().max(1),
465 Self::TensorBSpline { spec, .. } => {
466 let mut total: usize = 0;
512 for marginal in &spec.marginalspecs {
513 let m = bspline_basis_min_rows(marginal);
514 total = total.saturating_add(m.max(1));
515 }
516 total.max(RADIAL_FLOOR)
517 }
518 }
519 }
520
521 pub fn structural_kind(&self) -> &'static str {
532 match self {
533 Self::ByVariable { .. } => "by_variable",
534 Self::FactorSumToZero { .. } => "factor_sum_to_zero",
535 Self::BSpline1D { .. } => "bspline_1d",
536 Self::BySmooth { .. } => "by_smooth",
537 Self::FactorSmooth { .. } => "factor_smooth",
538 Self::ThinPlate { .. } => "thin_plate",
539 Self::Sphere { .. } => "sphere",
540 Self::ConstantCurvature { .. } => "constant_curvature",
541 Self::Matern { .. } => "matern",
542 Self::MeasureJet { .. } => "measurejet",
543 Self::Duchon { .. } => "duchon",
544 Self::Pca { .. } => "pca",
545 Self::TensorBSpline { .. } => "tensor_bspline",
546 }
547 }
548
549 pub fn is_marginally_centered_tensor(&self) -> bool {
558 matches!(
559 self,
560 Self::TensorBSpline { spec, .. }
561 if matches!(spec.identifiability, TensorBSplineIdentifiability::MarginalSumToZero)
562 )
563 }
564
565 pub fn is_sum_to_zero_factor_smooth(&self) -> bool {
582 matches!(
583 self,
584 Self::FactorSumToZero { .. }
585 | Self::FactorSmooth {
586 spec: FactorSmoothSpec {
587 flavour: FactorSmoothFlavour::Sz,
588 ..
589 }
590 }
591 )
592 }
593
594 pub fn structural_feature_cols(&self) -> Vec<usize> {
598 match self {
599 Self::ByVariable { inner, .. } | Self::FactorSumToZero { inner, .. } => {
600 inner.structural_feature_cols()
601 }
602 Self::BySmooth { smooth, .. } => smooth.structural_feature_cols(),
603 Self::FactorSmooth { .. } => Vec::new(),
604 Self::BSpline1D { feature_col, .. } => vec![*feature_col],
605 Self::ThinPlate { feature_cols, .. }
606 | Self::Sphere { feature_cols, .. }
607 | Self::ConstantCurvature { feature_cols, .. }
608 | Self::Matern { feature_cols, .. }
609 | Self::MeasureJet { feature_cols, .. }
610 | Self::Duchon { feature_cols, .. }
611 | Self::Pca { feature_cols, .. }
612 | Self::TensorBSpline { feature_cols, .. } => feature_cols.clone(),
613 }
614 }
615}
616
617pub fn bspline_basis_min_rows(spec: &crate::basis::BSplineBasisSpec) -> usize {
642 use crate::basis::BSplineKnotSpec;
643 let columns = match &spec.knotspec {
644 BSplineKnotSpec::Generate {
645 num_internal_knots, ..
646 } => *num_internal_knots + spec.degree + 1,
647 BSplineKnotSpec::Automatic {
648 num_internal_knots: Some(k),
649 ..
650 } => *k + spec.degree + 1,
651 BSplineKnotSpec::Automatic {
652 num_internal_knots: None,
653 ..
654 } => {
655 spec.degree + 2
659 }
660 BSplineKnotSpec::Provided(knots) => knots.len().saturating_sub(spec.degree + 1).max(1),
661 BSplineKnotSpec::NaturalCubicRegression { knots } => knots.len(),
663 BSplineKnotSpec::PeriodicUniform { num_basis, .. } => *num_basis,
664 };
665 let columns = columns.max(spec.degree + 2);
666
667 if spec.double_penalty {
668 const DOUBLE_PENALTY_FLOOR: usize = 2;
671 DOUBLE_PENALTY_FLOOR.min(columns).max(1)
672 } else {
673 columns
674 }
675}
676
677#[derive(Debug, Clone, Serialize, Deserialize)]
678pub enum ByVariableSpec {
679 Numeric,
680 Level { value_bits: u64, label: String },
681}
682
683#[derive(Debug, Clone, Serialize, Deserialize)]
684pub enum ByVarKind {
685 Numeric {
686 feature_col: usize,
687 },
688 Factor {
689 feature_col: usize,
690 ordered: bool,
691 frozen_levels: Option<Vec<u64>>,
692 },
693}
694
695#[derive(Debug, Clone, Serialize, Deserialize)]
696pub struct FactorSmoothSpec {
697 pub continuous_cols: Vec<usize>,
698 pub group_col: usize,
699 pub marginal: BSplineBasisSpec,
700 pub flavour: FactorSmoothFlavour,
701 pub group_frozen_levels: Option<Vec<u64>>,
702 #[serde(default)]
708 pub frozen_global_orthogonality: Option<Array2<f64>>,
709}
710
711#[derive(Debug, Clone, Serialize, Deserialize)]
712pub enum FactorSmoothFlavour {
713 Fs { m_null_penalty_orders: Vec<usize> },
714 Sz,
715 Re,
716}
717
718#[derive(Debug, Clone, Serialize, Deserialize)]
719pub struct TensorBSplineSpec {
720 pub marginalspecs: Vec<BSplineBasisSpec>,
721 #[serde(default)]
722 pub periods: Vec<Option<f64>>,
723 #[serde(default = "default_tensor_double_penalty")]
724 pub double_penalty: bool,
725 #[serde(default)]
726 pub identifiability: TensorBSplineIdentifiability,
727 #[serde(default)]
728 pub penalty_decomposition: TensorBSplinePenaltyDecomposition,
729}
730
731pub const fn default_tensor_double_penalty() -> bool {
732 true
733}
734
735impl Default for TensorBSplineSpec {
736 fn default() -> Self {
737 Self {
738 marginalspecs: Vec::new(),
739 periods: Vec::new(),
740 double_penalty: default_tensor_double_penalty(),
741 identifiability: TensorBSplineIdentifiability::default(),
742 penalty_decomposition: TensorBSplinePenaltyDecomposition::default(),
743 }
744 }
745}
746
747#[derive(Debug, Default, Clone, Serialize, Deserialize)]
748pub enum TensorBSplineIdentifiability {
749 None,
750 #[default]
751 SumToZero,
752 MarginalSumToZero,
762 FrozenTransform {
763 transform: Array2<f64>,
764 },
765}
766
767#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
768pub enum TensorBSplinePenaltyDecomposition {
769 #[default]
772 MarginalKroneckerSum,
773 Separable,
777}
778
779#[derive(Debug, Clone, Serialize, Deserialize)]
780pub struct SmoothTermSpec {
781 pub name: String,
782 pub basis: SmoothBasisSpec,
783 pub shape: ShapeConstraint,
784 #[serde(default)]
793 pub joint_null_rotation: Option<crate::basis::JointNullRotation>,
794}
795
796#[derive(Debug, Clone)]
797pub struct SmoothTerm {
798 pub name: String,
799 pub coeff_range: Range<usize>,
800 pub shape: ShapeConstraint,
801 pub penalties_local: Vec<Array2<f64>>,
802 pub nullspace_dims: Vec<usize>,
803 pub penaltyinfo_local: Vec<PenaltyInfo>,
804 pub metadata: BasisMetadata,
805 pub lower_bounds_local: Option<Array1<f64>>,
808 pub linear_constraints_local: Option<LinearInequalityConstraints>,
811 pub kronecker_factored: Option<KroneckerFactoredBasis>,
814 pub joint_null_rotation: Option<crate::basis::JointNullRotation>,
837 pub unabsorbed_global_orthogonality: Option<Array2<f64>>,
847}
848
849impl SmoothTerm {
850 pub fn apply_rotation_to_predict(
866 &self,
867 x_new_raw: Array2<f64>,
868 ) -> Result<Array2<f64>, BasisError> {
869 let Some(rot) = self.joint_null_rotation.as_ref() else {
870 return Ok(x_new_raw);
871 };
872 let p_local = rot.rotation.nrows();
873 if x_new_raw.ncols() != p_local {
874 crate::bail_dim_basis!(
875 "joint-null rotation replay for term '{}': raw design has {} columns, \
876 rotation expects {} (the raw basis builder must emit the same column \
877 count as at fit time)",
878 self.name,
879 x_new_raw.ncols(),
880 p_local,
881 );
882 }
883 Ok(gam_linalg::faer_ndarray::fast_ab(&x_new_raw, &rot.rotation))
884 }
885
886 pub fn wald_unpenalized_dim(&self) -> usize {
909 joint_unpenalized_dim(
910 self.coeff_range.len(),
911 &self.penalties_local,
912 &self.nullspace_dims,
913 )
914 }
915}
916
917pub fn joint_unpenalized_dim(
922 p_local: usize,
923 penalties_local: &[Array2<f64>],
924 nullspace_dims: &[usize],
925) -> usize {
926 use gam_linalg::faer_ndarray::FaerEigh;
927 if p_local == 0 {
928 return 0;
929 }
930 if penalties_local.is_empty() {
931 return p_local;
933 }
934 let mut s_total = Array2::<f64>::zeros((p_local, p_local));
939 let mut materialized = 0usize;
940 for s in penalties_local {
941 if s.nrows() == p_local && s.ncols() == p_local {
942 s_total += s;
943 materialized += 1;
944 }
945 }
946 if materialized == penalties_local.len() {
947 let symmetric = {
948 let transpose = s_total.t().to_owned();
949 (&s_total + &transpose) * 0.5
950 };
951 if let Ok((evals, _)) = symmetric.eigh(faer::Side::Lower) {
952 let max_abs = evals.iter().fold(0.0_f64, |acc, &v| acc.max(v.abs()));
953 if max_abs == 0.0 {
954 return p_local;
956 }
957 let tol = max_abs * (p_local as f64) * 1e-12;
958 let rank = evals.iter().filter(|&&v| v > tol).count();
959 return p_local.saturating_sub(rank);
960 }
961 }
962 if penalties_local.len() >= 2 {
967 0
968 } else {
969 nullspace_dims
970 .iter()
971 .copied()
972 .min()
973 .unwrap_or(0)
974 .min(p_local)
975 }
976}
977
978#[derive(Debug, Clone, Serialize, Deserialize)]
979pub struct PenaltyBlockInfo {
980 pub global_index: usize,
981 pub termname: Option<String>,
982 pub penalty: PenaltyInfo,
983}
984
985#[derive(Debug, Clone, Serialize, Deserialize)]
986pub struct DroppedPenaltyBlockInfo {
987 pub termname: Option<String>,
988 pub penalty: PenaltyInfo,
989}
990
991#[derive(Debug, Clone)]
992pub struct SmoothDesign {
993 pub term_designs: Vec<DesignMatrix>,
994 pub penalties: Vec<BlockwisePenalty>,
997 pub nullspace_dims: Vec<usize>,
998 pub penaltyinfo: Vec<PenaltyBlockInfo>,
999 pub dropped_penaltyinfo: Vec<DroppedPenaltyBlockInfo>,
1000 pub terms: Vec<SmoothTerm>,
1001 pub coefficient_lower_bounds: Option<Array1<f64>>,
1004 pub linear_constraints: Option<LinearInequalityConstraints>,
1007}
1008
1009impl SmoothDesign {
1010 pub fn total_smooth_cols(&self) -> usize {
1011 self.term_designs.iter().map(DesignMatrix::ncols).sum()
1012 }
1013 pub fn nrows(&self) -> usize {
1014 self.term_designs.first().map_or(0, DesignMatrix::nrows)
1015 }
1016}
1017
1018#[derive(Debug, Clone)]
1019pub struct RawSmoothDesign {
1020 pub term_designs: Vec<DesignMatrix>,
1021 pub penalties: Vec<BlockwisePenalty>,
1024 pub nullspace_dims: Vec<usize>,
1025 pub penaltyinfo: Vec<PenaltyBlockInfo>,
1026 pub dropped_penaltyinfo: Vec<DroppedPenaltyBlockInfo>,
1027 pub terms: Vec<SmoothTerm>,
1028 pub coefficient_lower_bounds: Option<Array1<f64>>,
1029 pub linear_constraints: Option<LinearInequalityConstraints>,
1030}
1031
1032impl RawSmoothDesign {
1033 pub fn total_smooth_cols(&self) -> usize {
1034 self.term_designs.iter().map(DesignMatrix::ncols).sum()
1035 }
1036 pub fn nrows(&self) -> usize {
1037 self.term_designs.first().map_or(0, DesignMatrix::nrows)
1038 }
1039}
1040
1041impl From<RawSmoothDesign> for SmoothDesign {
1042 fn from(value: RawSmoothDesign) -> Self {
1043 Self {
1044 term_designs: value.term_designs,
1045 penalties: value.penalties,
1046 nullspace_dims: value.nullspace_dims,
1047 penaltyinfo: value.penaltyinfo,
1048 dropped_penaltyinfo: value.dropped_penaltyinfo,
1049 terms: value.terms,
1050 coefficient_lower_bounds: value.coefficient_lower_bounds,
1051 linear_constraints: value.linear_constraints,
1052 }
1053 }
1054}
1055
1056#[derive(Debug, Default, Clone, Serialize, Deserialize)]
1057pub enum BoundedCoefficientPriorSpec {
1058 #[default]
1059 None,
1060 Uniform,
1061 Beta {
1062 a: f64,
1063 b: f64,
1064 },
1065}
1066
1067#[derive(Debug, Clone, Serialize, Deserialize, Default)]
1068pub enum LinearCoefficientGeometry {
1069 #[default]
1070 Unconstrained,
1071 Bounded {
1072 min: f64,
1073 max: f64,
1074 #[serde(default)]
1075 prior: BoundedCoefficientPriorSpec,
1076 },
1077}
1078
1079#[derive(Debug, Clone, Serialize, Deserialize)]
1080pub struct LinearTermSpec {
1081 pub name: String,
1082 pub feature_col: usize,
1088 #[serde(default)]
1091 pub feature_cols: Vec<usize>,
1092 #[serde(default)]
1106 pub categorical_levels: Vec<(usize, u64)>,
1107 #[serde(default = "default_linear_term_double_penalty")]
1112 pub double_penalty: bool,
1113 #[serde(default)]
1114 pub coefficient_geometry: LinearCoefficientGeometry,
1115 #[serde(default)]
1116 pub coefficient_min: Option<f64>,
1117 #[serde(default)]
1118 pub coefficient_max: Option<f64>,
1119}
1120
1121impl LinearTermSpec {
1122 pub fn effective_feature_cols(&self) -> Vec<usize> {
1125 if self.feature_cols.is_empty() {
1126 vec![self.feature_col]
1127 } else {
1128 self.feature_cols.clone()
1129 }
1130 }
1131
1132 pub fn is_interaction(&self) -> bool {
1134 self.feature_cols.len() > 1 || !self.categorical_levels.is_empty()
1135 }
1136
1137 pub fn realized_design_column(&self, data: ArrayView2<'_, f64>) -> Result<Array1<f64>, String> {
1150 let n = data.nrows();
1151 let p = data.ncols();
1152 let bounds = |col: usize| -> Result<(), String> {
1153 if col >= p {
1154 Err(format!(
1155 "linear term '{}' feature column {} out of bounds for {} columns",
1156 self.name, col, p
1157 ))
1158 } else {
1159 Ok(())
1160 }
1161 };
1162
1163 let mut column = if self.categorical_levels.is_empty() {
1168 let cols = self.effective_feature_cols();
1169 for &c in &cols {
1170 bounds(c)?;
1171 }
1172 let mut acc = data.column(cols[0]).to_owned();
1173 for &c in cols.iter().skip(1) {
1174 acc *= &data.column(c);
1175 }
1176 acc
1177 } else {
1178 let mut acc = Array1::<f64>::ones(n);
1179 for &c in &self.feature_cols {
1180 bounds(c)?;
1181 acc *= &data.column(c);
1182 }
1183 acc
1184 };
1185
1186 for &(col, level_bits) in &self.categorical_levels {
1187 bounds(col)?;
1188 let level_bits = gam_data::canonical_level_bits(f64::from_bits(level_bits));
1192 let gate = data.column(col);
1193 for (out, &v) in column.iter_mut().zip(gate.iter()) {
1194 if gam_data::canonical_level_bits(v) != level_bits {
1195 *out = 0.0;
1196 }
1197 }
1198 }
1199
1200 Ok(column)
1201 }
1202}
1203
1204pub const fn default_linear_term_double_penalty() -> bool {
1205 true
1206}
1207
1208pub const fn default_pca_smooth_penalty() -> f64 {
1209 1.0
1210}
1211
1212pub const fn default_pca_chunk_size() -> usize {
1213 4096
1214}
1215
1216#[derive(Debug, Clone, Serialize, Deserialize)]
1222pub struct RandomEffectTermSpec {
1223 pub name: String,
1224 pub feature_col: usize,
1225 pub drop_first_level: bool,
1228 #[serde(default = "default_random_effect_penalized")]
1232 pub penalized: bool,
1233 #[serde(default)]
1236 pub frozen_levels: Option<Vec<u64>>,
1237 #[serde(default = "default_random_effect_lenient_unseen")]
1254 pub lenient_unseen: bool,
1255}
1256
1257pub fn default_random_effect_penalized() -> bool {
1258 true
1259}
1260
1261pub fn default_random_effect_lenient_unseen() -> bool {
1262 true
1263}
1264
1265pub fn validate_measure_jet_positive_vec_len(
1266 label: &str,
1267 term_name: &str,
1268 field: &str,
1269 values: &[f64],
1270 expected: usize,
1271) -> Result<(), String> {
1272 if values.len() != expected {
1273 return Err(SmoothError::invalid_config(format!(
1274 "{label} term '{term_name}' frozen MeasureJet {field} has length {}, expected {expected}",
1275 values.len()
1276 ))
1277 .into());
1278 }
1279 if values
1280 .iter()
1281 .any(|value| !(value.is_finite() && *value > 0.0))
1282 {
1283 return Err(SmoothError::invalid_config(format!(
1284 "{label} term '{term_name}' frozen MeasureJet {field} values must be positive and finite"
1285 ))
1286 .into());
1287 }
1288 Ok(())
1289}
1290
1291#[derive(Debug, Clone, Serialize, Deserialize)]
1292pub struct TermCollectionSpec {
1293 pub linear_terms: Vec<LinearTermSpec>,
1294 pub random_effect_terms: Vec<RandomEffectTermSpec>,
1295 pub smooth_terms: Vec<SmoothTermSpec>,
1296}
1297
1298pub fn validate_smooth_basis_frozen(
1299 basis: &SmoothBasisSpec,
1300 label: &str,
1301 term_name: &str,
1302) -> Result<(), String> {
1303 match basis {
1304 SmoothBasisSpec::ByVariable { inner, .. }
1305 | SmoothBasisSpec::FactorSumToZero { inner, .. } => {
1306 validate_smooth_basis_frozen(inner, label, term_name)
1307 }
1308 SmoothBasisSpec::BSpline1D { spec, .. } => {
1309 if !matches!(
1310 spec.knotspec,
1311 BSplineKnotSpec::Provided(_)
1312 | BSplineKnotSpec::PeriodicUniform { .. }
1313 | BSplineKnotSpec::NaturalCubicRegression { .. }
1314 ) {
1315 return Err(format!(
1316 "{label} term '{term_name}' is not frozen: BSpline knotspec must be Provided, PeriodicUniform, or NaturalCubicRegression"
1317 ));
1318 }
1319 Ok(())
1320 }
1321 SmoothBasisSpec::ThinPlate { spec, .. } => {
1322 if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
1323 return Err(format!(
1324 "{label} term '{term_name}' is not frozen: ThinPlate centers must be UserProvided"
1325 ));
1326 }
1327 if matches!(
1328 spec.identifiability,
1329 SpatialIdentifiability::OrthogonalToParametric
1330 ) {
1331 return Err(format!(
1332 "{label} term '{term_name}' is not frozen: ThinPlate identifiability must be FrozenTransform or None"
1333 ));
1334 }
1335 Ok(())
1336 }
1337 _ => Ok(()),
1338 }
1339}
1340
1341impl TermCollectionSpec {
1342 pub fn write_structural_shape_hash(&self, h: &mut gam_runtime::warm_start::Fingerprinter) {
1356 h.write_str("term-collection");
1357 h.write_usize(self.linear_terms.len());
1358 for linear in &self.linear_terms {
1359 h.write_str(&linear.name);
1360 }
1361 h.write_usize(self.random_effect_terms.len());
1362 h.write_usize(self.smooth_terms.len());
1363 for smooth in &self.smooth_terms {
1364 h.write_str(&smooth.name);
1365 h.write_str(smooth.basis.structural_kind());
1366 for col in smooth.basis.structural_feature_cols() {
1367 h.write_usize(col);
1368 }
1369 }
1370 }
1371
1372 pub fn validate_frozen(&self, label: &str) -> Result<(), String> {
1376 for linear in &self.linear_terms {
1377 if let (Some(min), Some(max)) = (linear.coefficient_min, linear.coefficient_max)
1378 && (!min.is_finite() || !max.is_finite() || min > max)
1379 {
1380 return Err(SmoothError::invalid_config(format!(
1381 "{label} linear term '{}' has invalid coefficient constraint [{min}, {max}]",
1382 linear.name
1383 ))
1384 .into());
1385 }
1386 if let Some(min) = linear.coefficient_min
1387 && !min.is_finite()
1388 {
1389 return Err(SmoothError::invalid_config(format!(
1390 "{label} linear term '{}' has non-finite coefficient minimum {min}",
1391 linear.name
1392 ))
1393 .into());
1394 }
1395 if let Some(max) = linear.coefficient_max
1396 && !max.is_finite()
1397 {
1398 return Err(SmoothError::invalid_config(format!(
1399 "{label} linear term '{}' has non-finite coefficient maximum {max}",
1400 linear.name
1401 ))
1402 .into());
1403 }
1404 if let LinearCoefficientGeometry::Bounded { min, max, prior } =
1405 &linear.coefficient_geometry
1406 {
1407 if !min.is_finite() || !max.is_finite() || min >= max {
1408 return Err(SmoothError::invalid_config(format!(
1409 "{label} bounded term '{}' has invalid bounds [{min}, {max}]",
1410 linear.name
1411 ))
1412 .into());
1413 }
1414 match prior {
1415 BoundedCoefficientPriorSpec::None | BoundedCoefficientPriorSpec::Uniform => {}
1416 BoundedCoefficientPriorSpec::Beta { a, b } => {
1417 if !a.is_finite() || !b.is_finite() || *a < 1.0 || *b < 1.0 {
1418 return Err(SmoothError::invalid_config(format!(
1419 "{label} bounded term '{}' has invalid Beta prior ({a}, {b})",
1420 linear.name
1421 ))
1422 .into());
1423 }
1424 }
1425 }
1426 }
1427 }
1428 for st in &self.smooth_terms {
1429 match &st.basis {
1430 SmoothBasisSpec::ByVariable { inner, .. } => {
1431 validate_smooth_basis_frozen(inner, label, &st.name)?;
1432 let nested = SmoothTermSpec {
1433 name: st.name.clone(),
1434 basis: (**inner).clone(),
1435 shape: st.shape,
1436 joint_null_rotation: None,
1437 };
1438 TermCollectionSpec {
1439 linear_terms: Vec::new(),
1440 random_effect_terms: Vec::new(),
1441 smooth_terms: vec![nested],
1442 }
1443 .validate_frozen(label)?;
1444 }
1445 SmoothBasisSpec::FactorSumToZero { inner, levels, .. } => {
1446 if levels.len() < 2 {
1447 return Err(format!(
1448 "{label} term '{}' has invalid frozen sz levels",
1449 st.name
1450 ));
1451 }
1452 validate_smooth_basis_frozen(inner, label, &st.name)?;
1453 }
1454 SmoothBasisSpec::BSpline1D { spec, .. } => {
1455 if !matches!(
1456 spec.knotspec,
1457 BSplineKnotSpec::Provided(_)
1458 | BSplineKnotSpec::PeriodicUniform { .. }
1459 | BSplineKnotSpec::NaturalCubicRegression { .. }
1460 ) {
1461 return Err(SmoothError::invalid_config(format!(
1462 "{label} term '{}' is not frozen: BSpline knotspec must be Provided, PeriodicUniform, or NaturalCubicRegression",
1463 st.name
1464 ))
1465 .into());
1466 }
1467 }
1468 SmoothBasisSpec::ThinPlate { spec, .. } => {
1469 if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
1470 return Err(SmoothError::invalid_config(format!(
1471 "{label} term '{}' is not frozen: ThinPlate centers must be UserProvided",
1472 st.name
1473 ))
1474 .into());
1475 }
1476 if matches!(
1477 spec.identifiability,
1478 SpatialIdentifiability::OrthogonalToParametric
1479 ) {
1480 return Err(SmoothError::invalid_config(format!(
1481 "{label} term '{}' is not frozen: ThinPlate identifiability must be FrozenTransform or None",
1482 st.name
1483 ))
1484 .into());
1485 }
1486 }
1487 SmoothBasisSpec::Sphere { spec, .. } => {
1488 if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
1489 return Err(SmoothError::invalid_config(format!(
1490 "{label} term '{}' is not frozen: Sphere centers must be UserProvided",
1491 st.name
1492 ))
1493 .into());
1494 }
1495 if matches!(spec.method, crate::basis::SphereMethod::Harmonic)
1496 && spec.max_degree.is_none_or(|d| d == 0)
1497 {
1498 return Err(format!(
1499 "{label} term '{}' is not frozen: sphere max_degree must be positive",
1500 st.name
1501 ));
1502 }
1503 }
1504 SmoothBasisSpec::ConstantCurvature { spec, .. } => {
1505 if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
1506 return Err(SmoothError::invalid_config(format!(
1507 "{label} term '{}' is not frozen: ConstantCurvature centers must be UserProvided",
1508 st.name
1509 ))
1510 .into());
1511 }
1512 if !(spec.length_scale.is_finite() && spec.length_scale > 0.0) {
1513 return Err(SmoothError::invalid_config(format!(
1514 "{label} term '{}' is not frozen: ConstantCurvature length_scale must be the realized positive value",
1515 st.name
1516 ))
1517 .into());
1518 }
1519 }
1520 SmoothBasisSpec::MeasureJet { spec, .. } => {
1521 let centers = match &spec.center_strategy {
1522 CenterStrategy::UserProvided(centers) => centers,
1523 _ => {
1524 return Err(SmoothError::invalid_config(format!(
1525 "{label} term '{}' is not frozen: MeasureJet centers must be UserProvided",
1526 st.name
1527 ))
1528 .into());
1529 }
1530 };
1531 if centers.nrows() == 0 {
1532 return Err(SmoothError::invalid_config(format!(
1533 "{label} term '{}' is not frozen: MeasureJet centers are empty",
1534 st.name
1535 ))
1536 .into());
1537 }
1538 if !(spec.length_scale.is_finite() && spec.length_scale > 0.0) {
1539 return Err(SmoothError::invalid_config(format!(
1540 "{label} term '{}' is not frozen: MeasureJet length_scale must be the realized positive value",
1541 st.name
1542 ))
1543 .into());
1544 }
1545 let frozen = spec.frozen_quadrature.as_ref().ok_or_else(|| {
1548 SmoothError::invalid_config(format!(
1549 "{label} term '{}' is not frozen: MeasureJet frozen_quadrature payload is missing",
1550 st.name
1551 ))
1552 })?;
1553 if frozen.masses.len() != centers.nrows() {
1554 return Err(SmoothError::invalid_config(format!(
1555 "{label} term '{}' frozen MeasureJet has {} masses for {} centers",
1556 st.name,
1557 frozen.masses.len(),
1558 centers.nrows()
1559 ))
1560 .into());
1561 }
1562 let total_mass = frozen.masses.sum();
1563 if frozen
1564 .masses
1565 .iter()
1566 .any(|mass| !(mass.is_finite() && *mass >= 0.0))
1567 || !(total_mass.is_finite() && total_mass > 0.0)
1568 {
1569 return Err(SmoothError::invalid_config(format!(
1570 "{label} term '{}' frozen MeasureJet masses must be finite, nonnegative, and have positive total mass",
1571 st.name
1572 ))
1573 .into());
1574 }
1575 let n_levels = frozen.eps_band.len();
1576 if n_levels == 0
1577 || frozen
1578 .eps_band
1579 .iter()
1580 .any(|eps| !(eps.is_finite() && *eps > 0.0))
1581 {
1582 return Err(SmoothError::invalid_config(format!(
1583 "{label} term '{}' frozen MeasureJet eps_band must be nonempty, finite, and positive",
1584 st.name
1585 ))
1586 .into());
1587 }
1588 for (idx, pair) in frozen.eps_band.windows(2).enumerate() {
1589 if pair[1] <= pair[0] {
1590 return Err(SmoothError::invalid_config(format!(
1591 "{label} term '{}' frozen MeasureJet eps_band is not strictly ascending at {idx}: {} then {}",
1592 st.name,
1593 pair[0],
1594 pair[1]
1595 ))
1596 .into());
1597 }
1598 }
1599 validate_measure_jet_positive_vec_len(
1600 label,
1601 &st.name,
1602 "support_means",
1603 &frozen.support_means,
1604 n_levels,
1605 )?;
1606 let per_level = crate::basis::measure_jet_multiscale_mode(spec);
1614 if per_level {
1615 validate_measure_jet_positive_vec_len(
1616 label,
1617 &st.name,
1618 "penalty_normalization_scales",
1619 &frozen.penalty_normalization_scales,
1620 n_levels,
1621 )?;
1622 validate_measure_jet_positive_vec_len(
1623 label,
1624 &st.name,
1625 "raw_penalty_normalization_scales",
1626 &frozen.raw_penalty_normalization_scales,
1627 n_levels,
1628 )?;
1629 if frozen.fused_penalty_normalization_scale.is_some() {
1630 return Err(SmoothError::invalid_config(format!(
1631 "{label} term '{}' per-level MeasureJet must not carry a fused penalty normalization scale",
1632 st.name
1633 ))
1634 .into());
1635 }
1636 } else {
1637 if !frozen.penalty_normalization_scales.is_empty()
1638 || !frozen.raw_penalty_normalization_scales.is_empty()
1639 {
1640 return Err(SmoothError::invalid_config(format!(
1641 "{label} term '{}' fused MeasureJet must not carry per-level penalty normalization scales",
1642 st.name
1643 ))
1644 .into());
1645 }
1646 match frozen.fused_penalty_normalization_scale {
1647 Some(scale) if scale.is_finite() && scale > 0.0 => {}
1648 Some(scale) => {
1649 return Err(SmoothError::invalid_config(format!(
1650 "{label} term '{}' fused MeasureJet penalty normalization scale must be positive and finite, got {scale}",
1651 st.name
1652 ))
1653 .into());
1654 }
1655 None => {
1656 return Err(SmoothError::invalid_config(format!(
1657 "{label} term '{}' fused MeasureJet is missing its penalty normalization scale",
1658 st.name
1659 ))
1660 .into());
1661 }
1662 }
1663 }
1664 }
1665 SmoothBasisSpec::Matern { spec, .. } => {
1666 if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
1667 return Err(SmoothError::invalid_config(format!(
1668 "{label} term '{}' is not frozen: Matern centers must be UserProvided",
1669 st.name
1670 ))
1671 .into());
1672 }
1673 }
1674 SmoothBasisSpec::Duchon { spec, .. } => {
1675 if !matches!(spec.center_strategy, CenterStrategy::UserProvided(_)) {
1676 return Err(SmoothError::invalid_config(format!(
1677 "{label} term '{}' is not frozen: Duchon centers must be UserProvided",
1678 st.name
1679 ))
1680 .into());
1681 }
1682 if matches!(
1683 spec.identifiability,
1684 SpatialIdentifiability::OrthogonalToParametric
1685 ) {
1686 return Err(SmoothError::invalid_config(format!(
1687 "{label} term '{}' is not frozen: Duchon identifiability must be FrozenTransform or None",
1688 st.name
1689 ))
1690 .into());
1691 }
1692 }
1693 SmoothBasisSpec::Pca {
1694 centered,
1695 center_mean,
1696 pca_basis_path,
1697 ..
1698 } => {
1699 if *centered && center_mean.is_none() && pca_basis_path.is_none() {
1700 return Err(SmoothError::invalid_config(format!(
1701 "{label} term '{}' is not frozen: centered Pca missing center_mean",
1702 st.name
1703 ))
1704 .into());
1705 }
1706 }
1707 SmoothBasisSpec::BySmooth { smooth, by_kind } => {
1708 if let SmoothBasisSpec::BySmooth { .. } = smooth.as_ref() {
1709 return Err(format!("{label} term '{}' has nested by-smooths", st.name));
1710 }
1711 match by_kind {
1712 ByVarKind::Numeric { .. } => {}
1713 ByVarKind::Factor { frozen_levels, .. } if frozen_levels.is_none() => {
1714 return Err(format!(
1715 "{label} term '{}' is not frozen: by-factor levels missing",
1716 st.name
1717 ));
1718 }
1719 ByVarKind::Factor { .. } => {}
1720 }
1721 let nested = TermCollectionSpec {
1722 linear_terms: vec![],
1723 random_effect_terms: vec![],
1724 smooth_terms: vec![SmoothTermSpec {
1725 name: st.name.clone(),
1726 basis: (**smooth).clone(),
1727 shape: st.shape,
1728 joint_null_rotation: None,
1729 }],
1730 };
1731 nested.validate_frozen(label)?;
1732 }
1733 SmoothBasisSpec::FactorSmooth { spec } => {
1734 if spec.group_frozen_levels.is_none() {
1735 return Err(format!(
1736 "{label} term '{}' is not frozen: factor-smooth levels missing",
1737 st.name
1738 ));
1739 }
1740 if !matches!(
1741 spec.marginal.knotspec,
1742 BSplineKnotSpec::Provided(_)
1743 | BSplineKnotSpec::PeriodicUniform { .. }
1744 | BSplineKnotSpec::NaturalCubicRegression { .. }
1756 ) {
1757 return Err(format!(
1758 "{label} term '{}' is not frozen: factor-smooth marginal knots missing",
1759 st.name
1760 ));
1761 }
1762 }
1763 SmoothBasisSpec::TensorBSpline { spec, .. } => {
1764 for (dim, marginal) in spec.marginalspecs.iter().enumerate() {
1765 if !matches!(
1766 marginal.knotspec,
1767 BSplineKnotSpec::Provided(_)
1768 | BSplineKnotSpec::PeriodicUniform { .. }
1769 | BSplineKnotSpec::NaturalCubicRegression { .. }
1770 ) {
1771 return Err(SmoothError::invalid_config(format!(
1772 "{label} term '{}' dim {} is not frozen: tensor marginal knotspec must be Provided, PeriodicUniform, or NaturalCubicRegression",
1773 st.name, dim
1774 ))
1775 .into());
1776 }
1777 }
1778 if matches!(
1779 spec.identifiability,
1780 TensorBSplineIdentifiability::SumToZero
1781 | TensorBSplineIdentifiability::MarginalSumToZero
1782 ) {
1783 return Err(SmoothError::invalid_config(format!(
1784 "{label} term '{}' is not frozen: tensor identifiability must be FrozenTransform or None",
1785 st.name
1786 ))
1787 .into());
1788 }
1789 }
1790 }
1791 }
1792
1793 for rt in &self.random_effect_terms {
1794 if rt.frozen_levels.is_none() {
1795 return Err(SmoothError::invalid_config(format!(
1796 "{label} random-effect term '{}' is not frozen: missing frozen_levels",
1797 rt.name
1798 ))
1799 .into());
1800 }
1801 }
1802
1803 Ok(())
1804 }
1805
1806 pub fn remap_feature_columns<E, F>(&self, mut remap: F) -> Result<TermCollectionSpec, E>
1825 where
1826 F: FnMut(usize) -> Result<usize, E>,
1827 {
1828 let mut out = self.clone();
1829 for lt in &mut out.linear_terms {
1830 lt.feature_col = remap(lt.feature_col)?;
1831 for fc in lt.feature_cols.iter_mut() {
1841 *fc = remap(*fc)?;
1842 }
1843 for (col, _bits) in lt.categorical_levels.iter_mut() {
1848 *col = remap(*col)?;
1849 }
1850 }
1851 for rt in &mut out.random_effect_terms {
1852 rt.feature_col = remap(rt.feature_col)?;
1853 }
1854 for st in &mut out.smooth_terms {
1855 remap_smooth_basis_feature_columns(&mut st.basis, &mut remap)?;
1856 }
1857 Ok(out)
1858 }
1859}
1860
1861pub fn remap_smooth_basis_feature_columns<E, F>(
1866 basis: &mut SmoothBasisSpec,
1867 remap: &mut F,
1868) -> Result<(), E>
1869where
1870 F: FnMut(usize) -> Result<usize, E>,
1871{
1872 match basis {
1873 SmoothBasisSpec::ByVariable { inner, by_col, .. }
1874 | SmoothBasisSpec::FactorSumToZero { inner, by_col, .. } => {
1875 *by_col = remap(*by_col)?;
1876 remap_smooth_basis_feature_columns(inner, remap)?;
1877 }
1878 SmoothBasisSpec::BSpline1D { feature_col, .. } => {
1879 *feature_col = remap(*feature_col)?;
1880 }
1881 SmoothBasisSpec::BySmooth { smooth, by_kind } => {
1882 let by_feature_col = match by_kind {
1883 ByVarKind::Numeric { feature_col } | ByVarKind::Factor { feature_col, .. } => {
1884 feature_col
1885 }
1886 };
1887 *by_feature_col = remap(*by_feature_col)?;
1888 remap_smooth_basis_feature_columns(smooth, remap)?;
1889 }
1890 SmoothBasisSpec::FactorSmooth { spec } => {
1891 for fc in spec.continuous_cols.iter_mut() {
1892 *fc = remap(*fc)?;
1893 }
1894 spec.group_col = remap(spec.group_col)?;
1895 }
1896 SmoothBasisSpec::ThinPlate { feature_cols, .. }
1897 | SmoothBasisSpec::Sphere { feature_cols, .. }
1898 | SmoothBasisSpec::ConstantCurvature { feature_cols, .. }
1899 | SmoothBasisSpec::Matern { feature_cols, .. }
1900 | SmoothBasisSpec::MeasureJet { feature_cols, .. }
1901 | SmoothBasisSpec::Duchon { feature_cols, .. }
1902 | SmoothBasisSpec::Pca { feature_cols, .. }
1903 | SmoothBasisSpec::TensorBSpline { feature_cols, .. } => {
1904 for fc in feature_cols.iter_mut() {
1905 *fc = remap(*fc)?;
1906 }
1907 }
1908 }
1909 Ok(())
1910}
1911
1912#[derive(Debug, Clone)]
1913pub enum PenaltyStructureHint {
1914 Ridge(f64),
1915 Kronecker(Vec<Array2<f64>>),
1916}
1917
1918#[derive(Clone)]
1925pub struct BlockwisePenalty {
1926 pub col_range: Range<usize>,
1928 pub local: Array2<f64>,
1931 pub prior_mean: gam_problem::CoefficientPriorMean,
1933 pub structure_hint: Option<PenaltyStructureHint>,
1936 pub op: Option<std::sync::Arc<dyn crate::analytic_penalties::PenaltyOp>>,
1941}
1942
1943impl std::fmt::Debug for BlockwisePenalty {
1944 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1945 f.debug_struct("BlockwisePenalty")
1946 .field("col_range", &self.col_range)
1947 .field(
1948 "local",
1949 &format_args!("{}×{}", self.local.nrows(), self.local.ncols()),
1950 )
1951 .field("prior_mean", &self.prior_mean)
1952 .field("structure_hint", &self.structure_hint)
1953 .field("op", &self.op.as_ref().map(|o| o.dim()))
1954 .finish()
1955 }
1956}
1957
1958impl BlockwisePenalty {
1959 pub fn new(col_range: Range<usize>, local: Array2<f64>) -> Self {
1961 assert_eq!(col_range.len(), local.nrows());
1962 assert_eq!(col_range.len(), local.ncols());
1963 Self {
1964 col_range,
1965 local,
1966 prior_mean: gam_problem::CoefficientPriorMean::Zero,
1967 structure_hint: None,
1968 op: None,
1969 }
1970 }
1971
1972 pub fn with_prior_mean(mut self, prior_mean: gam_problem::CoefficientPriorMean) -> Self {
1973 self.prior_mean = prior_mean;
1974 self
1975 }
1976
1977 pub fn with_op(
1979 mut self,
1980 op: Option<std::sync::Arc<dyn crate::analytic_penalties::PenaltyOp>>,
1981 ) -> Self {
1982 self.op = op;
1983 self
1984 }
1985
1986 pub fn ridge(col_range: Range<usize>, scale: f64) -> Self {
1987 let block_size = col_range.len();
1988 let mut local = Array2::<f64>::zeros((block_size, block_size));
1989 for i in 0..block_size {
1990 local[[i, i]] = scale;
1991 }
1992 Self {
1993 col_range,
1994 local,
1995 prior_mean: gam_problem::CoefficientPriorMean::Zero,
1996 structure_hint: Some(PenaltyStructureHint::Ridge(scale)),
1997 op: None,
1998 }
1999 }
2000
2001 pub fn kronecker(
2002 col_range: Range<usize>,
2003 local: Array2<f64>,
2004 factors: Vec<Array2<f64>>,
2005 ) -> Self {
2006 assert_eq!(col_range.len(), local.nrows());
2007 assert_eq!(col_range.len(), local.ncols());
2008 Self {
2009 col_range,
2010 local,
2011 prior_mean: gam_problem::CoefficientPriorMean::Zero,
2012 structure_hint: Some(PenaltyStructureHint::Kronecker(factors)),
2013 op: None,
2014 }
2015 }
2016
2017 pub fn to_global(&self, p_total: usize) -> Array2<f64> {
2021 let mut g = Array2::<f64>::zeros((p_total, p_total));
2022 let r = &self.col_range;
2023 assert!(
2024 r.end <= p_total && self.local.nrows() == r.len() && self.local.ncols() == r.len(),
2025 "BlockwisePenalty::to_global shape invariant violated: \
2026 col_range={}..{}, local={}x{}, p_total={}",
2027 r.start,
2028 r.end,
2029 self.local.nrows(),
2030 self.local.ncols(),
2031 p_total,
2032 );
2033 g.slice_mut(s![r.start..r.end, r.start..r.end])
2034 .assign(&self.local);
2035 g
2036 }
2037
2038 pub fn to_penalty_matrix(&self, total_dim: usize) -> gam_problem::PenaltyMatrix {
2041 gam_problem::PenaltyMatrix::Blockwise {
2042 local: self.local.clone(),
2043 col_range: self.col_range.clone(),
2044 total_dim,
2045 }
2046 }
2047
2048 #[inline]
2050 pub fn block_size(&self) -> usize {
2051 self.col_range.len()
2052 }
2053}
2054
2055pub fn weighted_blockwise_penalty_sum(
2059 penalties: &[BlockwisePenalty],
2060 lambdas: &[f64],
2061 p_total: usize,
2062) -> Array2<f64> {
2063 assert_eq!(penalties.len(), lambdas.len());
2064 for (idx, &lam) in lambdas.iter().enumerate() {
2071 assert!(
2072 lam.is_finite() && lam >= 0.0,
2073 "weighted_blockwise_penalty_sum: lambdas[{idx}] = {lam} is invalid (must be finite and non-negative; negative smoothing parameters violate S_λ ⪰ 0)",
2074 );
2075 }
2076 for (idx, bp) in penalties.iter().enumerate() {
2080 let r = &bp.col_range;
2081 assert!(
2082 r.end <= p_total,
2083 "weighted_blockwise_penalty_sum: penalties[{idx}] col_range {:?} exceeds p_total = {p_total}",
2084 r,
2085 );
2086 }
2087 let mut out = Array2::<f64>::zeros((p_total, p_total));
2088 for (bp, &lam) in penalties.iter().zip(lambdas.iter()) {
2089 let r = &bp.col_range;
2090 let mut slice = out.slice_mut(s![r.start..r.end, r.start..r.end]);
2091 slice.scaled_add(lam, &bp.local);
2092 }
2093 out
2094}
2095
2096#[derive(Debug, Clone)]
2103pub struct KroneckerPenaltySystem {
2104 pub marginal_penalties: Vec<Array2<f64>>,
2106 pub marginal_eigensystems: Vec<(Array1<f64>, Array2<f64>)>,
2108 pub marginal_dims: Vec<usize>,
2110 pub has_double_penalty: bool,
2112}
2113
2114impl KroneckerPenaltySystem {
2115 pub fn new(
2116 marginal_penalties: Vec<Array2<f64>>,
2117 marginal_dims: Vec<usize>,
2118 has_double_penalty: bool,
2119 ) -> Result<Self, BasisError> {
2120 if marginal_penalties.len() != marginal_dims.len() {
2121 crate::bail_dim_basis!(
2122 "KroneckerPenaltySystem: {} penalties vs {} dims",
2123 marginal_penalties.len(),
2124 marginal_dims.len()
2125 );
2126 }
2127 let eigensystems =
2128 kronecker_marginal_eigensystems(&marginal_penalties, "KroneckerPenaltySystem")
2129 .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
2130 Ok(Self {
2131 marginal_penalties,
2132 marginal_eigensystems: eigensystems,
2133 marginal_dims,
2134 has_double_penalty,
2135 })
2136 }
2137
2138 pub fn p_total(&self) -> usize {
2139 self.marginal_dims.iter().copied().product()
2140 }
2141
2142 pub fn ndim(&self) -> usize {
2143 self.marginal_dims.len()
2144 }
2145
2146 pub fn num_penalties(&self) -> usize {
2147 self.marginal_dims.len() + if self.has_double_penalty { 1 } else { 0 }
2148 }
2149
2150 pub fn logdet_and_derivatives(
2154 &self,
2155 lambdas: &[f64],
2156 ridge: f64,
2157 ) -> (f64, Array1<f64>, Array2<f64>) {
2158 let n_pen = self.num_penalties();
2159 assert_eq!(lambdas.len(), n_pen, "lambda count mismatch");
2160 let marginal_evals: Vec<_> = self
2161 .marginal_eigensystems
2162 .iter()
2163 .map(|(evals, _)| evals.view())
2164 .collect();
2165 kronecker_logdet_and_derivatives(
2166 &marginal_evals,
2167 &self.marginal_dims,
2168 lambdas,
2169 self.has_double_penalty,
2170 ridge,
2171 )
2172 }
2173
2174 pub fn logdet_rank_and_derivatives(
2175 &self,
2176 lambdas: &[f64],
2177 ridge: f64,
2178 ) -> (f64, usize, Array1<f64>, Array2<f64>) {
2179 let n_pen = self.num_penalties();
2180 assert_eq!(lambdas.len(), n_pen, "lambda count mismatch");
2181 let d = self.marginal_dims.len();
2182 let mut logdet = 0.0;
2183 let mut rank = 0usize;
2184 let mut grad = Array1::<f64>::zeros(n_pen);
2185 let mut hess = Array2::<f64>::zeros((n_pen, n_pen));
2186 const EIGENVALUE_POSITIVITY_FLOOR: f64 = 1e-12;
2190 const STRUCTURAL_ZERO_FLOOR: f64 = 1e-12;
2194 let mut multi_idx = vec![0usize; d];
2195 loop {
2196 let mut sigma = 0.0;
2197 let mut structural_sigma = 0.0;
2198 for k in 0..d {
2199 let marginal_eigenvalue = self.marginal_eigensystems[k].0[multi_idx[k]];
2200 structural_sigma += marginal_eigenvalue;
2201 sigma += lambdas[k] * marginal_eigenvalue;
2202 }
2203 let joint_null = structural_sigma <= STRUCTURAL_ZERO_FLOOR;
2204 if self.has_double_penalty && joint_null {
2205 sigma += lambdas[d];
2206 }
2207 if structural_sigma > STRUCTURAL_ZERO_FLOOR {
2208 sigma += ridge;
2209 }
2210
2211 if sigma > EIGENVALUE_POSITIVITY_FLOOR {
2212 rank += 1;
2213 logdet += sigma.ln();
2214 let inv_sigma = 1.0 / sigma;
2215 let inv_sigma2 = inv_sigma * inv_sigma;
2216 for k in 0..n_pen {
2217 let ck = if k < d {
2218 lambdas[k] * self.marginal_eigensystems[k].0[multi_idx[k]]
2219 } else if joint_null {
2220 lambdas[d]
2221 } else {
2222 0.0
2223 };
2224 grad[k] += ck * inv_sigma;
2225 hess[[k, k]] += ck * inv_sigma - ck * ck * inv_sigma2;
2226 for l in (k + 1)..n_pen {
2227 let cl = if l < d {
2228 lambdas[l] * self.marginal_eigensystems[l].0[multi_idx[l]]
2229 } else if joint_null {
2230 lambdas[d]
2231 } else {
2232 0.0
2233 };
2234 let off = -ck * cl * inv_sigma2;
2235 hess[[k, l]] += off;
2236 hess[[l, k]] += off;
2237 }
2238 }
2239 }
2240
2241 let mut carry = true;
2242 for dim in (0..d).rev() {
2243 if carry {
2244 multi_idx[dim] += 1;
2245 if multi_idx[dim] < self.marginal_dims[dim] {
2246 carry = false;
2247 } else {
2248 multi_idx[dim] = 0;
2249 }
2250 }
2251 }
2252 if carry {
2253 break;
2254 }
2255 }
2256 (logdet, rank, grad, hess)
2257 }
2258}
2259
2260#[cfg(test)]
2261mod joint_unpenalized_dim_tests {
2262 use super::joint_unpenalized_dim;
2263 use ndarray::{Array2, array};
2264
2265 #[test]
2266 fn no_penalty_is_fully_unpenalized() {
2267 assert_eq!(joint_unpenalized_dim(4, &[], &[]), 4);
2268 }
2269
2270 #[test]
2271 fn single_penalty_returns_its_own_null_space() {
2272 let s = array![[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 5.0]];
2275 assert_eq!(joint_unpenalized_dim(3, std::slice::from_ref(&s), &[2]), 2);
2276 }
2277
2278 #[test]
2279 fn complementary_double_penalty_has_empty_joint_null_space() {
2280 let bending = array![[0.0, 0.0, 0.0], [0.0, 4.0, 0.0], [0.0, 0.0, 4.0]];
2287 let ridge = array![[2.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]];
2288 assert_eq!(joint_unpenalized_dim(3, &[bending, ridge], &[1, 2]), 0);
2289 }
2290
2291 #[test]
2292 fn partial_overlap_keeps_shared_null_direction() {
2293 let a = array![[0.0, 0.0, 0.0], [0.0, 3.0, 0.0], [0.0, 0.0, 0.0]];
2297 let b = array![[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 3.0]];
2298 assert_eq!(joint_unpenalized_dim(3, &[a, b], &[2, 2]), 1);
2299 }
2300
2301 #[test]
2302 fn non_materialized_penalty_falls_back_conservatively() {
2303 let full: Array2<f64> = array![[0.0, 0.0], [0.0, 1.0]];
2307 let factor: Array2<f64> = array![[1.0]]; assert_eq!(
2309 joint_unpenalized_dim(2, &[full, factor.clone()], &[1, 0]),
2310 0
2311 );
2312 assert_eq!(
2314 joint_unpenalized_dim(4, std::slice::from_ref(&factor), &[2]),
2315 2
2316 );
2317 }
2318}
2319
2320#[cfg(test)]
2321mod kronecker_penalty_system_tests {
2322 use super::KroneckerPenaltySystem;
2323 use ndarray::array;
2324
2325 #[test]
2326 fn double_penalty_rank_derivatives_use_only_joint_null_space() {
2327 let penalties = vec![
2328 array![[0.0, 0.0], [0.0, 2.0]],
2329 array![[0.0, 0.0], [0.0, 3.0]],
2330 ];
2331 let system = KroneckerPenaltySystem::new(penalties, vec![2usize, 2usize], true).unwrap();
2332 let lambdas = vec![5.0, 7.0, 11.0];
2333
2334 let (logdet, rank, grad, hess) = system.logdet_rank_and_derivatives(&lambdas, 0.0);
2335
2336 let expected_diag = [11.0_f64, 21.0, 10.0, 31.0];
2337 let expected_logdet: f64 = expected_diag.iter().map(|v| v.ln()).sum();
2338 assert_eq!(rank, 4);
2339 assert!((logdet - expected_logdet).abs() <= 1e-12);
2340 assert!(
2341 (grad[2] - 1.0).abs() <= 1e-12,
2342 "double-penalty rank derivative must count only the joint null mode, got {}",
2343 grad[2]
2344 );
2345 assert!(hess[[2, 2]].abs() <= 1e-12);
2346 }
2347}
2348
2349#[derive(Clone, Debug)]
2350pub struct TermCollectionDesign {
2351 pub design: DesignMatrix,
2360 pub penalties: Vec<BlockwisePenalty>,
2361 pub nullspace_dims: Vec<usize>,
2362 pub penaltyinfo: Vec<PenaltyBlockInfo>,
2363 pub dropped_penaltyinfo: Vec<DroppedPenaltyBlockInfo>,
2364 pub coefficient_lower_bounds: Option<Array1<f64>>,
2367 pub linear_constraints: Option<LinearInequalityConstraints>,
2370 pub intercept_range: Range<usize>,
2371 pub linear_ranges: Vec<(String, Range<usize>)>,
2372 pub random_effect_ranges: Vec<(String, Range<usize>)>,
2373 pub random_effect_levels: Vec<(String, Vec<u64>)>,
2374 pub smooth: SmoothDesign,
2375}
2376
2377impl TermCollectionDesign {
2378 pub fn leading_penalty_blocks_before_smooth(&self) -> usize {
2386 self.penaltyinfo
2387 .iter()
2388 .take_while(|info| {
2389 matches!(
2390 &info.penalty.source,
2391 crate::basis::PenaltySource::Other(source)
2392 if source == "LinearTermRidge"
2393 || source.starts_with("RandomEffectRidge(")
2394 )
2395 })
2396 .count()
2397 }
2398
2399 pub fn penalties_as_penalty_matrix(&self) -> Vec<gam_problem::PenaltyMatrix> {
2403 let p = self.design.ncols();
2404 self.penalties
2405 .iter()
2406 .map(|bp| bp.to_penalty_matrix(p))
2407 .collect()
2408 }
2409
2410 #[inline]
2412 pub fn num_penalties(&self) -> usize {
2413 self.penalties.len()
2414 }
2415
2416 pub fn realize_coefficient_groups(
2419 &self,
2420 groups: &[CoefficientGroupSpec],
2421 base_prior: &gam_spec::RhoPrior,
2422 ) -> Result<RealizedCoefficientGroups, BasisError> {
2423 realize_coefficient_groups(self, groups, base_prior)
2424 }
2425
2426 pub fn kronecker_penalty_system(&self) -> Option<KroneckerPenaltySystem> {
2437 let [only_term] = self.smooth.terms.as_slice() else {
2438 return None;
2439 };
2440 let kron = only_term.kronecker_factored.as_ref()?;
2441 if kron.marginal_dims.len() < 2
2447 || kron.marginal_penalties.len() != kron.marginal_dims.len()
2448 || kron.marginal_designs.len() != kron.marginal_dims.len()
2449 {
2450 return None;
2451 }
2452 KroneckerPenaltySystem::new(
2453 kron.marginal_penalties.clone(),
2454 kron.marginal_dims.clone(),
2455 kron.has_double_penalty,
2456 )
2457 .ok()
2458 }
2459}
2460
2461#[derive(Clone)]
2467pub struct StandardLatentCoordConfig {
2468 pub values: std::sync::Arc<crate::latent::LatentCoordValues>,
2469 pub term_index: gam_problem::types::SmoothTermIdx,
2470 pub feature_cols: Vec<usize>,
2471 pub manifold: crate::latent::LatentManifold,
2472 pub manifold_auto: bool,
2473 pub retraction_registry: gam_problem::LatentRetractionRegistry,
2474 pub analytic_penalties: Option<std::sync::Arc<crate::AnalyticPenaltyRegistry>>,
2475}
2476
2477#[derive(Clone, Debug, Serialize, Deserialize)]
2478pub struct AdaptiveSpatialMap {
2479 pub termname: String,
2480 pub feature_cols: Vec<usize>,
2481 pub collocation_points: Array2<f64>,
2482 pub inv_magweight: Array1<f64>,
2483 pub invgradweight: Array1<f64>,
2484 pub inv_lapweight: Array1<f64>,
2485}
2486
2487#[derive(Clone, Debug, Serialize, Deserialize)]
2488pub struct AdaptiveRegularizationDiagnostics {
2489 pub epsilon_0: f64,
2490 pub epsilon_g: f64,
2491 pub epsilon_c: f64,
2492 pub epsilon_outer_iterations: usize,
2493 pub mm_iterations: usize,
2494 pub converged: bool,
2495 pub maps: Vec<AdaptiveSpatialMap>,
2496}
2497
2498#[derive(Debug, Clone)]
2499pub struct LinearColumnConditioning {
2500 col_idx: usize,
2501 mean: f64,
2502 scale: f64,
2503}
2504
2505#[derive(Debug, Clone, Default)]
2506pub struct LinearFitConditioning {
2507 pub intercept_idx: usize,
2508 pub columns: Vec<LinearColumnConditioning>,
2509}
2510
2511#[derive(Clone)]
2512pub struct SpatialPsiDerivative {
2513 pub penalty_index: usize,
2515 pub penalty_indices: Vec<usize>,
2516 pub global_range: Range<usize>,
2517 pub total_p: usize,
2518 pub x_psi_local: Array2<f64>,
2519 pub s_psi_components_local: Vec<Array2<f64>>,
2520 pub x_psi_psi_local: Array2<f64>,
2521 pub s_psi_psi_components_local: Vec<Array2<f64>>,
2522 pub aniso_group_id: Option<usize>,
2523 pub aniso_cross_designs: Option<Vec<(usize, Array2<f64>)>>,
2526 pub aniso_cross_penalty_provider: Option<
2530 std::sync::Arc<
2531 dyn Fn(usize) -> Result<Vec<Array2<f64>>, EstimationError> + Send + Sync + 'static,
2532 >,
2533 >,
2534 pub implicit_operator: Option<std::sync::Arc<crate::basis::ImplicitDesignPsiDerivative>>,
2539 pub implicit_axis: usize,
2541}
2542
2543#[derive(Debug, Clone)]
2544pub struct SpatialLogKappaCoords {
2545 pub values: Array1<f64>,
2548 pub dims_per_term: Vec<usize>,
2550}
2551
2552#[derive(Clone, Copy)]
2557pub enum AnisoBoundEnd {
2558 Lower,
2559 Upper,
2560}
2561
2562impl SpatialLogKappaCoords {
2563 pub fn new_with_dims(values: Array1<f64>, dims_per_term: Vec<usize>) -> Self {
2565 assert_eq!(
2566 values.len(),
2567 dims_per_term.iter().sum::<usize>(),
2568 "SpatialLogKappaCoords: values length {} != sum of dims_per_term {}",
2569 values.len(),
2570 dims_per_term.iter().sum::<usize>(),
2571 );
2572 Self {
2573 values,
2574 dims_per_term,
2575 }
2576 }
2577
2578 pub fn from_length_scales(
2580 spec: &TermCollectionSpec,
2581 term_indices: &[usize],
2582 options: &SpatialLengthScaleOptimizationOptions,
2583 ) -> Self {
2584 let mut out = Array1::<f64>::zeros(term_indices.len());
2585 for (slot, &term_idx) in term_indices.iter().enumerate() {
2586 if let Some(cc) = constant_curvature_term_spec(spec, term_idx) {
2592 out[slot] = cc.kappa;
2593 continue;
2594 }
2595 let length_scale = get_spatial_length_scale(spec, term_idx)
2596 .unwrap_or(options.min_length_scale)
2597 .clamp(options.min_length_scale, options.max_length_scale);
2598 out[slot] = -length_scale.ln();
2599 }
2600 Self {
2601 values: out,
2602 dims_per_term: vec![1; term_indices.len()],
2603 }
2604 }
2605
2606 pub fn from_length_scales_aniso(
2624 spec: &TermCollectionSpec,
2625 term_indices: &[usize],
2626 options: &SpatialLengthScaleOptimizationOptions,
2627 ) -> Self {
2628 let mut vals = Vec::new();
2629 let mut dims = Vec::new();
2630 for &term_idx in term_indices {
2631 if let Some(mj) = measure_jet_term_spec(spec, term_idx) {
2635 let seed = measure_jet_psi_seed(mj);
2636 dims.push(seed.len());
2637 vals.extend(seed);
2638 continue;
2639 }
2640 if let Some(cc) = constant_curvature_term_spec(spec, term_idx) {
2646 vals.push(cc.kappa);
2647 dims.push(1);
2648 continue;
2649 }
2650 let length_scale = get_spatial_length_scale(spec, term_idx)
2651 .unwrap_or(options.min_length_scale)
2652 .clamp(options.min_length_scale, options.max_length_scale);
2653 let psi_bar = -length_scale.ln(); if spatial_term_uses_per_axis_psi(spec, term_idx) {
2656 let d = get_spatial_feature_dim(spec, term_idx).unwrap_or(1);
2661 let eta_raw = get_spatial_aniso_log_scales(spec, term_idx)
2662 .expect("predicate guarantees aniso_log_scales is Some");
2663 let eta = center_aniso_log_scales(&eta_raw);
2664 for &eta_a in &eta {
2665 vals.push(psi_bar + eta_a);
2666 }
2667 dims.push(d);
2668 } else {
2669 vals.push(psi_bar);
2676 dims.push(1);
2677 }
2678 }
2679 Self {
2680 values: Array1::from_vec(vals),
2681 dims_per_term: dims,
2682 }
2683 }
2684
2685 pub fn lower_bounds_from_data(
2689 data: ArrayView2<'_, f64>,
2690 spec: &TermCollectionSpec,
2691 term_indices: &[usize],
2692 options: &SpatialLengthScaleOptimizationOptions,
2693 ) -> Self {
2694 let mut values = Array1::<f64>::zeros(term_indices.len());
2695 for (slot, &term_idx) in term_indices.iter().enumerate() {
2696 values[slot] = spatial_term_psi_bounds(data, spec, term_idx, options).0;
2697 }
2698 Self {
2699 values,
2700 dims_per_term: vec![1; term_indices.len()],
2701 }
2702 }
2703
2704 pub fn upper_bounds_from_data(
2706 data: ArrayView2<'_, f64>,
2707 spec: &TermCollectionSpec,
2708 term_indices: &[usize],
2709 options: &SpatialLengthScaleOptimizationOptions,
2710 ) -> Self {
2711 let mut values = Array1::<f64>::zeros(term_indices.len());
2712 for (slot, &term_idx) in term_indices.iter().enumerate() {
2713 values[slot] = spatial_term_psi_bounds(data, spec, term_idx, options).1;
2714 }
2715 Self {
2716 values,
2717 dims_per_term: vec![1; term_indices.len()],
2718 }
2719 }
2720
2721 pub fn lower_bounds_aniso_from_data(
2738 data: ArrayView2<'_, f64>,
2739 spec: &TermCollectionSpec,
2740 term_indices: &[usize],
2741 dims_per_term: &[usize],
2742 options: &SpatialLengthScaleOptimizationOptions,
2743 ) -> Self {
2744 Self::aniso_bounds_from_data(
2745 data,
2746 spec,
2747 term_indices,
2748 dims_per_term,
2749 options,
2750 AnisoBoundEnd::Lower,
2751 )
2752 }
2753
2754 pub fn upper_bounds_aniso_from_data(
2758 data: ArrayView2<'_, f64>,
2759 spec: &TermCollectionSpec,
2760 term_indices: &[usize],
2761 dims_per_term: &[usize],
2762 options: &SpatialLengthScaleOptimizationOptions,
2763 ) -> Self {
2764 Self::aniso_bounds_from_data(
2765 data,
2766 spec,
2767 term_indices,
2768 dims_per_term,
2769 options,
2770 AnisoBoundEnd::Upper,
2771 )
2772 }
2773
2774 fn aniso_bounds_from_data(
2780 data: ArrayView2<'_, f64>,
2781 spec: &TermCollectionSpec,
2782 term_indices: &[usize],
2783 dims_per_term: &[usize],
2784 options: &SpatialLengthScaleOptimizationOptions,
2785 end: AnisoBoundEnd,
2786 ) -> Self {
2787 assert_eq!(term_indices.len(), dims_per_term.len());
2788 let total: usize = dims_per_term.iter().sum();
2789 let mut values = Array1::<f64>::zeros(total);
2790 let mut cursor = 0;
2791 for (slot, &term_idx) in term_indices.iter().enumerate() {
2792 let d = dims_per_term[slot];
2793 if let Some(mj) = measure_jet_term_spec(spec, term_idx) {
2796 let bounds = measure_jet_psi_bound_values(mj, matches!(end, AnisoBoundEnd::Upper));
2797 for (offset, bound) in bounds.into_iter().enumerate() {
2798 if offset < d {
2799 values[cursor + offset] = bound;
2800 }
2801 }
2802 cursor += d;
2803 continue;
2804 }
2805 if constant_curvature_term_spec(spec, term_idx).is_some() {
2808 let (lo, hi) = constant_curvature_kappa_bounds(data, spec, term_idx);
2809 if d >= 1 {
2810 values[cursor] = match end {
2811 AnisoBoundEnd::Lower => lo,
2812 AnisoBoundEnd::Upper => hi,
2813 };
2814 }
2815 cursor += d;
2816 continue;
2817 }
2818 let psi_bound = {
2819 let (lo, hi) = spatial_term_psi_bounds(data, spec, term_idx, options);
2820 match end {
2821 AnisoBoundEnd::Lower => lo,
2822 AnisoBoundEnd::Upper => hi,
2823 }
2824 };
2825 let axis_offsets = if d <= 1 {
2826 vec![0.0; d]
2827 } else {
2828 get_spatial_aniso_log_scales(spec, term_idx)
2829 .filter(|eta| eta.len() == d)
2830 .map(|eta| center_aniso_log_scales(&eta))
2831 .unwrap_or_else(|| vec![0.0; d])
2832 };
2833 for offset in 0..d {
2834 values[cursor + offset] = psi_bound + axis_offsets[offset];
2835 }
2836 cursor += d;
2837 }
2838 Self {
2839 values,
2840 dims_per_term: dims_per_term.to_vec(),
2841 }
2842 }
2843
2844 pub fn reseed_from_data(
2853 mut self,
2854 data: ArrayView2<'_, f64>,
2855 spec: &TermCollectionSpec,
2856 term_indices: &[usize],
2857 options: &SpatialLengthScaleOptimizationOptions,
2858 ) -> Self {
2859 assert_eq!(term_indices.len(), self.dims_per_term.len());
2860 let mut cursor = 0;
2861 for (slot, &term_idx) in term_indices.iter().enumerate() {
2862 let d = self.dims_per_term[slot];
2863 if measure_jet_term_spec(spec, term_idx).is_some() {
2866 cursor += d;
2867 continue;
2868 }
2869 if constant_curvature_term_spec(spec, term_idx).is_some() {
2873 cursor += d;
2874 continue;
2875 }
2876 let Some(psi_bar_new) = spatial_term_psi_seed(data, spec, term_idx, options) else {
2877 cursor += d;
2878 continue;
2879 };
2880 if d == 0 {
2881 continue;
2882 }
2883 let current: Vec<f64> = self.values.slice(s![cursor..cursor + d]).to_vec();
2884 let psi_bar_old = current.iter().sum::<f64>() / d as f64;
2885 for (offset, &old_value) in current.iter().enumerate() {
2886 self.values[cursor + offset] = psi_bar_new + (old_value - psi_bar_old);
2887 }
2888 cursor += d;
2889 }
2890 self
2891 }
2892
2893 pub fn clamp_to_bounds(
2904 mut self,
2905 lower: &SpatialLogKappaCoords,
2906 upper: &SpatialLogKappaCoords,
2907 ) -> Self {
2908 assert_eq!(self.values.len(), lower.values.len());
2909 assert_eq!(self.values.len(), upper.values.len());
2910 let mut n_projected = 0usize;
2911 let mut worst_delta = 0.0_f64;
2912 for idx in 0..self.values.len() {
2913 let lo = lower.values[idx];
2914 let hi = upper.values[idx];
2915 if !(lo.is_finite() && hi.is_finite()) {
2916 continue;
2917 }
2918 let v = self.values[idx];
2919 if v < lo {
2920 worst_delta = worst_delta.max(lo - v);
2921 self.values[idx] = lo;
2922 n_projected += 1;
2923 } else if v > hi {
2924 worst_delta = worst_delta.max(v - hi);
2925 self.values[idx] = hi;
2926 n_projected += 1;
2927 }
2928 }
2929 if n_projected > 0 {
2930 log::info!(
2931 "[spatial-kappa] projected {n_projected}/{} ψ seed coords into data-derived bounds \
2932 (worst excess={worst_delta:.3} log units); user length_scale falls outside \
2933 [{KERNEL_RANGE_MIN_DIAMETER_FRACTION}/r_max, {KERNEL_RANGE_MAX_SPACING_MULTIPLE}/r_min] geometry window",
2934 self.values.len()
2935 );
2936 }
2937 self
2938 }
2939
2940 pub fn from_theta_tail_with_dims(
2942 theta: &Array1<f64>,
2943 start: usize,
2944 dims_per_term: Vec<usize>,
2945 ) -> Self {
2946 let total: usize = dims_per_term.iter().sum();
2947 Self {
2948 values: theta.slice(s![start..start + total]).to_owned(),
2949 dims_per_term,
2950 }
2951 }
2952
2953 pub fn len(&self) -> usize {
2955 self.values.len()
2956 }
2957
2958 pub fn dims_per_term(&self) -> &[usize] {
2960 &self.dims_per_term
2961 }
2962
2963 fn term_offset(&self, term_idx: usize) -> usize {
2965 self.dims_per_term[..term_idx].iter().sum()
2966 }
2967
2968 pub fn term_slice(&self, term_idx: usize) -> &[f64] {
2970 let offset = self.term_offset(term_idx);
2971 let d = self.dims_per_term[term_idx];
2972 &self.values.as_slice().unwrap()[offset..offset + d]
2973 }
2974
2975 pub fn as_array(&self) -> &Array1<f64> {
2976 &self.values
2977 }
2978
2979 pub fn set_scalar_slot(&mut self, slot: usize, value: f64) -> bool {
2985 if slot >= self.dims_per_term.len() || self.dims_per_term[slot] != 1 {
2986 return false;
2987 }
2988 let offset = self.term_offset(slot);
2989 self.values[offset] = value;
2990 true
2991 }
2992
2993 pub fn split_at(&self, mid: usize) -> (Self, Self) {
2996 let flat_mid: usize = self.dims_per_term[..mid].iter().sum();
2997 (
2998 Self {
2999 values: self.values.slice(s![0..flat_mid]).to_owned(),
3000 dims_per_term: self.dims_per_term[..mid].to_vec(),
3001 },
3002 Self {
3003 values: self.values.slice(s![flat_mid..]).to_owned(),
3004 dims_per_term: self.dims_per_term[mid..].to_vec(),
3005 },
3006 )
3007 }
3008
3009 pub fn apply_tospec(
3016 &self,
3017 spec: &TermCollectionSpec,
3018 term_indices: &[usize],
3019 ) -> Result<TermCollectionSpec, EstimationError> {
3020 if term_indices.len() != self.dims_per_term.len() {
3021 crate::bail_invalid_estim!(
3022 "SpatialLogKappaCoords::apply_tospec: term count mismatch: \
3023 term_indices={} dims_per_term={}",
3024 term_indices.len(),
3025 self.dims_per_term.len()
3026 );
3027 }
3028 let mut updated = spec.clone();
3029 for (slot, &term_idx) in term_indices.iter().enumerate() {
3030 let psi = self.term_slice(slot);
3031 let d = self.dims_per_term[slot];
3032 if measure_jet_term_spec(&updated, term_idx).is_some() {
3035 set_measure_jet_psi_dials(&mut updated, term_idx, psi)?;
3036 continue;
3037 }
3038 if constant_curvature_term_spec(&updated, term_idx).is_some() {
3042 set_constant_curvature_kappa(&mut updated, term_idx, psi)?;
3043 continue;
3044 }
3045 let (next_length_scale, next_aniso) = spatial_term_psi_to_length_scale_and_aniso(psi);
3046 if (d == 1 || next_length_scale.is_some())
3047 && let Some(length_scale) = next_length_scale
3048 {
3049 set_spatial_length_scale(&mut updated, term_idx, length_scale)?;
3050 }
3051 if let Some(eta) = next_aniso {
3052 set_spatial_aniso_log_scales(&mut updated, term_idx, eta)?;
3053 }
3054 }
3055 Ok(updated)
3056 }
3057}
3058
3059pub fn center_aniso_log_scales(eta: &[f64]) -> Vec<f64> {
3060 if eta.len() <= 1 {
3061 return eta.to_vec();
3062 }
3063 let mean = eta.iter().sum::<f64>() / eta.len() as f64;
3064 eta.iter()
3065 .map(|&v| {
3066 let centered = v - mean;
3067 if centered.abs() <= 1e-15 {
3068 0.0
3069 } else {
3070 centered
3071 }
3072 })
3073 .collect()
3074}
3075
3076pub fn spatial_term_uses_per_axis_psi(resolvedspec: &TermCollectionSpec, term_idx: usize) -> bool {
3079 if let Some(mj) = measure_jet_term_spec(resolvedspec, term_idx) {
3080 return measure_jet_enrolls_psi(mj);
3081 }
3082 let Some(d) = get_spatial_feature_dim(resolvedspec, term_idx) else {
3083 return false;
3084 };
3085 if d <= 1 {
3086 return false;
3087 }
3088 let Some(eta) = get_spatial_aniso_log_scales(resolvedspec, term_idx) else {
3089 return false;
3090 };
3091 if eta.len() != d {
3092 return false;
3093 }
3094 !matches!(
3095 resolvedspec
3096 .smooth_terms
3097 .get(term_idx)
3098 .map(|term| &term.basis),
3099 Some(SmoothBasisSpec::Duchon { .. })
3100 )
3101}
3102
3103pub fn set_spatial_length_scale(
3104 spec: &mut TermCollectionSpec,
3105 term_idx: usize,
3106 length_scale: f64,
3107) -> Result<(), EstimationError> {
3108 let Some(term) = spec.smooth_terms.get_mut(term_idx) else {
3109 crate::bail_invalid_estim!("spatial length-scale term index {term_idx} out of range");
3110 };
3111 match &mut term.basis {
3112 SmoothBasisSpec::ThinPlate { spec, .. } => {
3113 spec.length_scale = length_scale;
3114 Ok(())
3115 }
3116 SmoothBasisSpec::Matern { spec, .. } => {
3117 spec.length_scale = length_scale;
3118 Ok(())
3119 }
3120 SmoothBasisSpec::Duchon { spec, .. } => {
3121 spec.length_scale = Some(length_scale);
3122 Ok(())
3123 }
3124 _ => Err(EstimationError::InvalidInput(format!(
3125 "term '{}' does not expose a spatial length scale",
3126 term.name
3127 ))),
3128 }
3129}
3130
3131pub fn get_spatial_length_scale(spec: &TermCollectionSpec, term_idx: usize) -> Option<f64> {
3132 spec.smooth_terms
3133 .get(term_idx)
3134 .and_then(|term| match &term.basis {
3135 SmoothBasisSpec::ThinPlate { spec, .. } => Some(spec.length_scale),
3136 SmoothBasisSpec::Matern { spec, .. } => Some(spec.length_scale),
3137 SmoothBasisSpec::Duchon { spec, .. } => spec.length_scale,
3138 _ => None,
3139 })
3140}
3141
3142pub fn spatial_term_supports_hyper_optimization(
3143 spec: &TermCollectionSpec,
3144 term_idx: usize,
3145) -> bool {
3146 if let Some(term) = spec.smooth_terms.get(term_idx)
3152 && let SmoothBasisSpec::ThinPlate { .. } = &term.basis
3153 {
3154 return false;
3155 }
3156
3157 if let Some(term) = spec.smooth_terms.get(term_idx)
3182 && let SmoothBasisSpec::Matern { .. } = &term.basis
3183 {
3184 return true;
3185 }
3186
3187 if let Some(mj) = measure_jet_term_spec(spec, term_idx) {
3190 return measure_jet_enrolls_psi(mj);
3191 }
3192
3193 if constant_curvature_term_spec(spec, term_idx).is_some() {
3200 return true;
3201 }
3202
3203 get_spatial_length_scale(spec, term_idx).is_some()
3204}
3205
3206pub fn measure_jet_term_spec(
3209 spec: &TermCollectionSpec,
3210 term_idx: usize,
3211) -> Option<&crate::basis::MeasureJetBasisSpec> {
3212 spec.smooth_terms
3213 .get(term_idx)
3214 .and_then(|term| match &term.basis {
3215 SmoothBasisSpec::MeasureJet { spec, .. } => Some(spec),
3216 _ => None,
3217 })
3218}
3219
3220pub fn measure_jet_enrolls_psi(mj: &crate::basis::MeasureJetBasisSpec) -> bool {
3227 measure_jet_learns_length_scale(mj)
3236 || (mj.tau0 > 0.0 && crate::basis::measure_jet_multiscale_mode(mj))
3237}
3238
3239pub fn measure_jet_learns_length_scale(mj: &crate::basis::MeasureJetBasisSpec) -> bool {
3242 mj.learn_length_scale
3243}
3244
3245pub fn freeze_measure_jet_length_scale_learning(spec: &mut TermCollectionSpec) -> usize {
3246 let mut frozen = 0;
3247 for term in spec.smooth_terms.iter_mut() {
3248 if let SmoothBasisSpec::MeasureJet { spec: mj, .. } = &mut term.basis
3249 && mj.learn_length_scale
3250 {
3251 mj.learn_length_scale = false;
3252 frozen += 1;
3253 }
3254 }
3255 frozen
3256}
3257
3258pub const MEASURE_JET_PSI_ALPHA_BOUNDS: (f64, f64) = (-1.0, 3.0);
3266
3267pub const MEASURE_JET_PSI_LN_TAU_BOUNDS: (f64, f64) = (-18.420680743952367, 4.605170185988092);
3268
3269pub const MEASURE_JET_PSI_LN_LENGTH_SCALE_BOUNDS: (f64, f64) =
3275 (-6.907755278982137, 4.605170185988092);
3276
3277pub fn measure_jet_penalty_psi_dim(mj: &crate::basis::MeasureJetBasisSpec) -> usize {
3285 if crate::basis::measure_jet_multiscale_mode(mj) {
3286 2
3287 } else {
3288 0
3289 }
3290}
3291
3292pub fn measure_jet_psi_dim(mj: &crate::basis::MeasureJetBasisSpec) -> usize {
3296 usize::from(measure_jet_learns_length_scale(mj)) + measure_jet_penalty_psi_dim(mj)
3297}
3298
3299pub fn measure_jet_psi_seed(mj: &crate::basis::MeasureJetBasisSpec) -> Vec<f64> {
3304 let mut seed = Vec::with_capacity(measure_jet_psi_dim(mj));
3305 if measure_jet_learns_length_scale(mj) {
3306 let ell = if mj.length_scale > 0.0 {
3310 mj.length_scale
3311 } else {
3312 1.0
3313 };
3314 seed.push(ell.ln());
3315 }
3316 if measure_jet_penalty_psi_dim(mj) > 0 {
3317 let ln_tau = mj.tau0.max(f64::MIN_POSITIVE).ln();
3319 seed.extend_from_slice(&[mj.alpha, ln_tau]);
3320 }
3321 seed
3322}
3323
3324pub fn measure_jet_psi_bound_values(
3327 mj: &crate::basis::MeasureJetBasisSpec,
3328 upper: bool,
3329) -> Vec<f64> {
3330 let pick = |b: (f64, f64)| if upper { b.1 } else { b.0 };
3331 let mut bounds = Vec::with_capacity(measure_jet_psi_dim(mj));
3332 if measure_jet_learns_length_scale(mj) {
3333 bounds.push(pick(MEASURE_JET_PSI_LN_LENGTH_SCALE_BOUNDS));
3334 }
3335 if measure_jet_penalty_psi_dim(mj) > 0 {
3336 bounds.push(pick(MEASURE_JET_PSI_ALPHA_BOUNDS));
3338 bounds.push(pick(MEASURE_JET_PSI_LN_TAU_BOUNDS));
3339 }
3340 bounds
3341}
3342
3343pub fn apply_measure_jet_psi(
3348 mj: &mut crate::basis::MeasureJetBasisSpec,
3349 psi: &[f64],
3350) -> Result<bool, EstimationError> {
3351 if psi.len() != measure_jet_psi_dim(mj) {
3352 crate::bail_invalid_estim!(
3353 "measure-jet ψ write-back dimension mismatch: got {} values for a {}-dial term",
3354 psi.len(),
3355 measure_jet_psi_dim(mj)
3356 );
3357 }
3358 let mut changed = false;
3359 let mut cursor = 0usize;
3363 if measure_jet_learns_length_scale(mj) {
3364 let next_ell = psi[cursor].exp();
3365 cursor += 1;
3366 if !(next_ell.is_finite() && next_ell > 0.0) {
3367 crate::bail_invalid_estim!(
3368 "measure-jet ψ write-back produced a non-finite/non-positive length_scale (ℓ={next_ell})"
3369 );
3370 }
3371 if next_ell != mj.length_scale {
3372 mj.length_scale = next_ell;
3373 changed = true;
3374 }
3375 }
3376 if measure_jet_penalty_psi_dim(mj) > 0 {
3377 let next_alpha = psi[cursor];
3380 let next_tau = psi[cursor + 1].exp();
3381 if !(next_alpha.is_finite() && next_tau.is_finite() && next_tau > 0.0) {
3382 crate::bail_invalid_estim!(
3383 "measure-jet ψ write-back produced non-finite dials (alpha={next_alpha}, tau={next_tau})"
3384 );
3385 }
3386 if next_alpha != mj.alpha {
3387 mj.alpha = next_alpha;
3388 changed = true;
3389 }
3390 if next_tau != mj.tau0 {
3391 mj.tau0 = next_tau;
3392 changed = true;
3393 }
3394 }
3395 Ok(changed)
3396}
3397
3398pub fn set_measure_jet_psi_dials(
3401 spec: &mut TermCollectionSpec,
3402 term_idx: usize,
3403 psi: &[f64],
3404) -> Result<bool, EstimationError> {
3405 let Some(term) = spec.smooth_terms.get_mut(term_idx) else {
3406 crate::bail_invalid_estim!("measure-jet ψ write-back: term index {term_idx} out of range");
3407 };
3408 set_single_term_measure_jet_psi_dials(term, psi)
3409}
3410
3411pub fn set_single_term_measure_jet_psi_dials(
3416 term: &mut SmoothTermSpec,
3417 psi: &[f64],
3418) -> Result<bool, EstimationError> {
3419 let SmoothBasisSpec::MeasureJet { spec: mj, .. } = &mut term.basis else {
3420 crate::bail_invalid_estim!("measure-jet ψ write-back targeted a non-measure-jet term");
3421 };
3422 apply_measure_jet_psi(mj, psi)
3423}
3424
3425pub fn constant_curvature_term_spec(
3428 spec: &TermCollectionSpec,
3429 term_idx: usize,
3430) -> Option<&crate::basis::ConstantCurvatureBasisSpec> {
3431 spec.smooth_terms
3432 .get(term_idx)
3433 .and_then(|term| match &term.basis {
3434 SmoothBasisSpec::ConstantCurvature { spec, .. } => Some(spec),
3435 _ => None,
3436 })
3437}
3438
3439pub const CONSTANT_CURVATURE_KAPPA_CHART_FRACTION: f64 = 0.5;
3447
3448pub const CONSTANT_CURVATURE_MIN_CHART_RADIUS2: f64 = 1e-8;
3452
3453pub fn constant_curvature_kappa_bounds(
3458 data: ArrayView2<'_, f64>,
3459 spec: &TermCollectionSpec,
3460 term_idx: usize,
3461) -> (f64, f64) {
3462 let feature_cols = match spec.smooth_terms.get(term_idx).map(|t| &t.basis) {
3463 Some(SmoothBasisSpec::ConstantCurvature { feature_cols, .. }) => feature_cols,
3464 _ => return (-1.0, 1.0),
3465 };
3466 let mut max_r2 = CONSTANT_CURVATURE_MIN_CHART_RADIUS2;
3467 for row in data.outer_iter() {
3468 let mut r2 = 0.0_f64;
3469 for &c in feature_cols.iter() {
3470 if let Some(&v) = row.get(c)
3471 && v.is_finite()
3472 {
3473 r2 += v * v;
3474 }
3475 }
3476 if r2 > max_r2 {
3477 max_r2 = r2;
3478 }
3479 }
3480 let half = CONSTANT_CURVATURE_KAPPA_CHART_FRACTION / max_r2;
3481 (-half, half)
3482}
3483
3484pub fn set_constant_curvature_kappa(
3488 spec: &mut TermCollectionSpec,
3489 term_idx: usize,
3490 psi: &[f64],
3491) -> Result<bool, EstimationError> {
3492 let Some(term) = spec.smooth_terms.get_mut(term_idx) else {
3493 crate::bail_invalid_estim!(
3494 "constant-curvature κ write-back: term index {term_idx} out of range"
3495 );
3496 };
3497 set_single_term_constant_curvature_kappa(term, psi)
3498}
3499
3500pub fn set_single_term_constant_curvature_kappa(
3505 term: &mut SmoothTermSpec,
3506 psi: &[f64],
3507) -> Result<bool, EstimationError> {
3508 if psi.len() != 1 {
3509 crate::bail_invalid_estim!(
3510 "constant-curvature κ write-back expects exactly one value, got {}",
3511 psi.len()
3512 );
3513 }
3514 let next_kappa = psi[0];
3515 if !next_kappa.is_finite() {
3516 crate::bail_invalid_estim!(
3517 "constant-curvature κ write-back produced a non-finite κ = {next_kappa}"
3518 );
3519 }
3520 let SmoothBasisSpec::ConstantCurvature { spec: cc, .. } = &mut term.basis else {
3521 crate::bail_invalid_estim!(
3522 "constant-curvature κ write-back targeted a non-constant-curvature term"
3523 );
3524 };
3525 if cc.kappa != next_kappa {
3526 cc.kappa = next_kappa;
3527 Ok(true)
3528 } else {
3529 Ok(false)
3530 }
3531}
3532
3533pub fn spatial_term_has_locked_kappa(spec: &TermCollectionSpec, term_idx: usize) -> bool {
3544 get_spatial_length_scale(spec, term_idx).is_some()
3545 && !spatial_term_uses_per_axis_psi(spec, term_idx)
3546}
3547
3548pub fn all_spatial_terms_kappa_fixed(spec: &TermCollectionSpec) -> bool {
3549 spec.smooth_terms.iter().enumerate().all(|(idx, _)| {
3550 !spatial_term_supports_hyper_optimization(spec, idx)
3551 || spatial_term_has_locked_kappa(spec, idx)
3552 })
3553}
3554
3555pub fn spatial_identifiability_policy(
3556 termspec: &SmoothTermSpec,
3557) -> Option<&SpatialIdentifiability> {
3558 match &termspec.basis {
3559 SmoothBasisSpec::ThinPlate { spec, .. } => Some(&spec.identifiability),
3560 SmoothBasisSpec::Duchon { spec, .. } => Some(&spec.identifiability),
3561 _ => None,
3562 }
3563}
3564
3565pub const NULLSPACE_WELLDET_DEGENERACY_RHO_SD: f64 = 15.0;
3569
3570pub fn is_nullspace_degeneracy_prior(prior: &gam_spec::RhoPrior) -> bool {
3573 matches!(
3574 prior,
3575 gam_spec::RhoPrior::Normal { mean, sd }
3576 if *mean == 0.0 && *sd == NULLSPACE_WELLDET_DEGENERACY_RHO_SD
3577 )
3578}
3579
3580pub const KERNEL_RANGE_MIN_DIAMETER_FRACTION: f64 = 2.0;
3592
3593pub const KERNEL_RANGE_MAX_SPACING_MULTIPLE: f64 = 1e2;
3598
3599fn spatial_term_stored_input_scales(term: &SmoothTermSpec) -> Option<Vec<f64>> {
3600 match &term.basis {
3601 SmoothBasisSpec::ThinPlate { input_scales, .. }
3602 | SmoothBasisSpec::Matern { input_scales, .. }
3603 | SmoothBasisSpec::Duchon { input_scales, .. } => input_scales.clone(),
3604 _ => None,
3605 }
3606}
3607
3608fn spatial_term_realized_input_scales(
3609 data: ArrayView2<'_, f64>,
3610 term: &SmoothTermSpec,
3611) -> Option<Vec<f64>> {
3612 let (feature_cols, stored) = match &term.basis {
3613 SmoothBasisSpec::ThinPlate {
3614 feature_cols,
3615 input_scales,
3616 ..
3617 }
3618 | SmoothBasisSpec::Matern {
3619 feature_cols,
3620 input_scales,
3621 ..
3622 }
3623 | SmoothBasisSpec::Duchon {
3624 feature_cols,
3625 input_scales,
3626 ..
3627 } => (feature_cols, input_scales),
3628 _ => return None,
3629 };
3630 if let Some(scales) = stored {
3631 return Some(scales.clone());
3632 }
3633 let x = select_columns(data, feature_cols).ok()?;
3634 compute_spatial_input_scales(x.view())
3635}
3636
3637pub fn spatial_term_psi_bounds(
3646 data: ArrayView2<'_, f64>,
3647 spec: &TermCollectionSpec,
3648 term_idx: usize,
3649 options: &SpatialLengthScaleOptimizationOptions,
3650) -> (f64, f64) {
3651 let fallback = (
3652 -options.max_length_scale.ln(),
3653 -options.min_length_scale.ln(),
3654 );
3655 if constant_curvature_term_spec(spec, term_idx).is_some() {
3660 return constant_curvature_kappa_bounds(data, spec, term_idx);
3661 }
3662 let Some(term) = spec.smooth_terms.get(term_idx) else {
3663 return fallback;
3664 };
3665 let aniso = get_spatial_aniso_log_scales(spec, term_idx);
3678 let (r_bounds, input_scales) = match spatial_term_center_strategy(term) {
3679 Some(CenterStrategy::UserProvided(centers)) if centers.nrows() >= 2 => {
3680 let bounds = match aniso.as_deref() {
3681 Some(eta) if eta.len() == centers.ncols() => {
3682 let y = points_in_aniso_y_space(centers.view(), eta);
3683 pairwise_distance_bounds(y.view())
3684 }
3685 _ => pairwise_distance_bounds(centers.view()),
3686 };
3687 (bounds, spatial_term_stored_input_scales(term))
3693 }
3694 _ => {
3695 let input_scales = spatial_term_realized_input_scales(data, term);
3696 let bounds = standardized_spatial_term_data(data, term)
3697 .ok()
3698 .and_then(|x| match aniso.as_deref() {
3699 Some(eta) if eta.len() == x.ncols() => {
3700 let y = points_in_aniso_y_space(x.view(), eta);
3701 pairwise_distance_bounds_sampled(y.view())
3702 }
3703 _ => pairwise_distance_bounds_sampled(x.view()),
3704 });
3705 (bounds, input_scales)
3706 }
3707 };
3708 let Some((r_min, r_max)) = r_bounds else {
3709 return fallback;
3710 };
3711 let inverse_sigma_geom = input_scales
3728 .as_deref()
3729 .map(|scales| compensate_length_scale_for_standardization(1.0, scales))
3730 .unwrap_or(1.0);
3731 let psi_chart_offset = inverse_sigma_geom.ln();
3732 let psi_lo_data =
3733 (KERNEL_RANGE_MIN_DIAMETER_FRACTION / r_max).ln() + psi_chart_offset;
3734 let psi_hi_data =
3735 (KERNEL_RANGE_MAX_SPACING_MULTIPLE / r_min).ln() + psi_chart_offset;
3736 let psi_lo = psi_lo_data.max(fallback.0);
3746 let psi_hi = psi_hi_data.min(fallback.1);
3747 if psi_lo >= psi_hi {
3748 return fallback;
3751 }
3752 (psi_lo, psi_hi)
3753}
3754
3755#[cfg(test)]
3756mod spatial_psi_bound_coordinate_tests {
3757 use super::*;
3758 use crate::basis::{MaternIdentifiability, MaternNu};
3759 use ndarray::array;
3760
3761 fn frozen_matern_bounds(theta: f64, dilation: f64) -> (f64, f64) {
3762 let source = array![
3763 [-1.7, -0.4],
3764 [-1.1, 0.8],
3765 [-0.2, -1.3],
3766 [0.5, 1.6],
3767 [1.4, -0.7],
3768 [2.1, 0.5],
3769 ];
3770 let (cos_theta, sin_theta) = (theta.cos(), theta.sin());
3771 let mut data = Array2::<f64>::zeros(source.raw_dim());
3772 for row in 0..source.nrows() {
3773 let x = source[[row, 0]];
3774 let y = source[[row, 1]];
3775 data[[row, 0]] = dilation * (cos_theta * x - sin_theta * y);
3776 data[[row, 1]] = dilation * (sin_theta * x + cos_theta * y);
3777 }
3778 let input_scales = compute_spatial_input_scales(data.view()).expect("input scales");
3779 let mut centers = data.clone();
3780 apply_input_standardization(&mut centers, &input_scales);
3781 let spec = TermCollectionSpec {
3782 linear_terms: Vec::new(),
3783 random_effect_terms: Vec::new(),
3784 smooth_terms: vec![SmoothTermSpec {
3785 name: "matern".to_string(),
3786 basis: SmoothBasisSpec::Matern {
3787 feature_cols: vec![0, 1],
3788 spec: MaternBasisSpec {
3789 periodic: None,
3790 center_strategy: CenterStrategy::UserProvided(centers),
3791 length_scale: 1.0,
3792 nu: MaternNu::FiveHalves,
3793 include_intercept: false,
3794 double_penalty: true,
3795 identifiability: MaternIdentifiability::CenterSumToZero,
3796 aniso_log_scales: None,
3797 nullspace_shrinkage_survived: None,
3798 },
3799 input_scales: Some(input_scales),
3800 },
3801 shape: ShapeConstraint::None,
3802 joint_null_rotation: None,
3803 }],
3804 };
3805 spatial_term_psi_bounds(
3806 data.view(),
3807 &spec,
3808 0,
3809 &SpatialLengthScaleOptimizationOptions::default(),
3810 )
3811 }
3812
3813 fn assert_close(left: f64, right: f64) {
3814 assert!(
3815 (left - right).abs() <= 1e-12,
3816 "coordinate-equivalent bounds differ: left={left:.16e}, right={right:.16e}"
3817 );
3818 }
3819
3820 #[test]
3821 fn standardized_center_bounds_return_to_original_units_under_rotation_and_scaling() {
3822 let base = frozen_matern_bounds(0.0, 1.0);
3823 let rotated = frozen_matern_bounds(0.61, 1.0);
3824 assert_close(rotated.0, base.0);
3825 assert_close(rotated.1, base.1);
3826
3827 let dilation = 4.0_f64;
3828 let rotated_scaled = frozen_matern_bounds(0.61, dilation);
3829 let expected_shift = dilation.ln();
3830 assert_close(rotated_scaled.0, base.0 - expected_shift);
3831 assert_close(rotated_scaled.1, base.1 - expected_shift);
3832 }
3833}
3834
3835pub fn spatial_term_psi_seed(
3839 data: ArrayView2<'_, f64>,
3840 spec: &TermCollectionSpec,
3841 term_idx: usize,
3842 options: &SpatialLengthScaleOptimizationOptions,
3843) -> Option<f64> {
3844 if get_spatial_length_scale(spec, term_idx).is_some() {
3845 return None; }
3847 let (psi_lo, psi_hi) = spatial_term_psi_bounds(data, spec, term_idx, options);
3848 Some(0.5 * (psi_lo + psi_hi))
3849}
3850
3851pub fn spatial_term_psi_to_length_scale_and_aniso(psi: &[f64]) -> (Option<f64>, Option<Vec<f64>>) {
3852 if psi.len() <= 1 {
3853 (Some((-psi.first().copied().unwrap_or(0.0)).exp()), None)
3854 } else {
3855 let psi_bar = psi.iter().sum::<f64>() / psi.len() as f64;
3856 (
3857 Some((-psi_bar).exp()),
3858 Some(psi.iter().map(|&value| value - psi_bar).collect()),
3859 )
3860 }
3861}
3862
3863pub fn get_spatial_aniso_log_scales(
3865 spec: &TermCollectionSpec,
3866 term_idx: usize,
3867) -> Option<Vec<f64>> {
3868 spec.smooth_terms
3869 .get(term_idx)
3870 .and_then(|term| match &term.basis {
3871 SmoothBasisSpec::Matern { spec, .. } => spec.aniso_log_scales.clone(),
3872 SmoothBasisSpec::Duchon { spec, .. } => spec.aniso_log_scales.clone(),
3873 _ => None,
3874 })
3875}
3876
3877pub fn response_aware_axis_contrasts(
3897 x: ndarray::ArrayView2<'_, f64>,
3898 y: ndarray::ArrayView1<'_, f64>,
3899) -> Option<Vec<f64>> {
3900 let n = x.nrows();
3901 let d = x.ncols();
3902 if d <= 1 || n < 4 || y.len() != n {
3903 return None;
3904 }
3905 if x.iter().any(|v| !v.is_finite()) || y.iter().any(|v| !v.is_finite()) {
3906 return None;
3907 }
3908 let mut scores = Vec::with_capacity(d);
3909 for a in 0..d {
3910 let mut order: Vec<usize> = (0..n).collect();
3911 let col = x.column(a);
3912 order.sort_by(|&i, &j| {
3913 col[i]
3914 .partial_cmp(&col[j])
3915 .unwrap_or(std::cmp::Ordering::Equal)
3916 });
3917 let mut tv = 0.0_f64;
3918 for w in order.windows(2) {
3919 let diff = y[w[1]] - y[w[0]];
3920 tv += diff * diff;
3921 }
3922 scores.push(-0.5 * (tv + 1e-12).ln());
3924 }
3925 if scores.iter().any(|v| !v.is_finite()) {
3926 return None;
3927 }
3928 let mean = scores.iter().sum::<f64>() / d as f64;
3929 let centered: Vec<f64> = scores.iter().map(|&s| s - mean).collect();
3930 if centered.iter().all(|&v| v.abs() < 1e-9) {
3933 return None;
3934 }
3935 Some(centered)
3936}
3937
3938pub fn apply_response_aware_anisotropy_seed(
3947 data: ArrayView2<'_, f64>,
3948 y: ndarray::ArrayView1<'_, f64>,
3949 spec: &mut TermCollectionSpec,
3950 spatial_terms: &[usize],
3951) {
3952 const MAX_NUDGE: f64 = std::f64::consts::LN_2;
3957 for &term_idx in spatial_terms {
3958 let Some(current_eta) = get_spatial_aniso_log_scales(spec, term_idx) else {
3959 continue;
3960 };
3961 let d = current_eta.len();
3962 if d <= 1 {
3963 continue;
3964 }
3965 let Some(term) = spec.smooth_terms.get(term_idx) else {
3966 continue;
3967 };
3968 let feature_cols = term.basis.structural_feature_cols();
3969 if feature_cols.len() != d {
3970 continue;
3971 }
3972 let Ok(x) = select_columns(data, &feature_cols) else {
3973 continue;
3974 };
3975 let Some(contrast) = response_aware_axis_contrasts(x.view(), y) else {
3976 continue;
3977 };
3978 let nudged: Vec<f64> = current_eta
3979 .iter()
3980 .zip(contrast.iter())
3981 .map(|(&eta_a, &c_a)| eta_a + c_a.clamp(-MAX_NUDGE, MAX_NUDGE))
3982 .collect();
3983 if let Err(err) = set_spatial_aniso_log_scales(spec, term_idx, nudged) {
3986 log::debug!(
3987 "[spatial-kappa] response-aware anisotropy seed skipped for term {term_idx}: {err}"
3988 );
3989 }
3990 }
3991}
3992
3993pub fn get_spatial_feature_dim(spec: &TermCollectionSpec, term_idx: usize) -> Option<usize> {
3995 spec.smooth_terms
3996 .get(term_idx)
3997 .and_then(|term| match &term.basis {
3998 SmoothBasisSpec::ThinPlate { feature_cols, .. } => Some(feature_cols.len()),
3999 SmoothBasisSpec::Matern { feature_cols, .. } => Some(feature_cols.len()),
4000 SmoothBasisSpec::Duchon { feature_cols, .. } => Some(feature_cols.len()),
4001 _ => None,
4002 })
4003}
4004
4005pub fn log_spatial_aniso_scales(spec: &TermCollectionSpec) {
4012 for (term_idx, term) in spec.smooth_terms.iter().enumerate() {
4013 let (aniso, length_scale) = match &term.basis {
4014 SmoothBasisSpec::Matern { spec, .. } => {
4015 (spec.aniso_log_scales.as_ref(), Some(spec.length_scale))
4016 }
4017 SmoothBasisSpec::Duchon { spec, .. } => {
4018 (spec.aniso_log_scales.as_ref(), spec.length_scale)
4019 }
4020 _ => (None, None),
4021 };
4022 let Some(eta) = aniso else { continue };
4023 if eta.is_empty() {
4024 continue;
4025 }
4026 let mut lines = match length_scale {
4027 Some(ls) => format!(
4028 "[spatial-kappa] term {} (\"{}\"): anisotropic length scales optimized (global length_scale={:.4})",
4029 term_idx, term.name, ls
4030 ),
4031 None => format!(
4032 "[spatial-kappa] term {} (\"{}\"): pure Duchon shape anisotropy optimized",
4033 term_idx, term.name
4034 ),
4035 };
4036 for (a, &eta_a) in eta.iter().enumerate() {
4037 if let Some(ls) = length_scale {
4038 let length_a = ls * (-eta_a).exp();
4039 let kappa_a = (1.0 / ls) * eta_a.exp();
4040 lines.push_str(&format!(
4041 "\n axis {}: eta={:+.4}, length={:.4}, kappa={:.4}",
4042 a, eta_a, length_a, kappa_a
4043 ));
4044 } else {
4045 lines.push_str(&format!("\n axis {}: eta={:+.4}", a, eta_a));
4046 }
4047 }
4048 log::info!("{}", lines);
4049 }
4050}
4051
4052pub fn set_spatial_aniso_log_scales(
4054 spec: &mut TermCollectionSpec,
4055 term_idx: usize,
4056 eta: Vec<f64>,
4057) -> Result<(), EstimationError> {
4058 let eta = center_aniso_log_scales(&eta);
4059 let Some(term) = spec.smooth_terms.get_mut(term_idx) else {
4060 crate::bail_invalid_estim!("spatial aniso_log_scales term index {term_idx} out of range");
4061 };
4062 match &mut term.basis {
4063 SmoothBasisSpec::Matern { spec, .. } => {
4064 spec.aniso_log_scales = Some(eta);
4065 Ok(())
4066 }
4067 SmoothBasisSpec::Duchon { spec, .. } => {
4068 spec.aniso_log_scales = Some(eta);
4069 Ok(())
4070 }
4071 _ => Err(EstimationError::InvalidInput(format!(
4072 "term '{}' does not support aniso_log_scales",
4073 term.name
4074 ))),
4075 }
4076}
4077
4078pub fn sync_aniso_contrasts_from_metadata(spec: &mut TermCollectionSpec, design: &SmoothDesign) {
4085 for (term_idx, term) in design.terms.iter().enumerate() {
4086 let meta_aniso = match &term.metadata {
4087 BasisMetadata::Matern {
4088 aniso_log_scales, ..
4089 } => aniso_log_scales.clone(),
4090 BasisMetadata::Duchon {
4091 aniso_log_scales, ..
4092 } => aniso_log_scales.clone(),
4093 _ => None,
4094 };
4095 if let Some(eta) = meta_aniso
4096 && eta.len() > 1
4097 {
4098 set_spatial_aniso_log_scales(spec, term_idx, eta).ok();
4099 }
4100 }
4101}
4102
4103#[derive(Debug, Clone)]
4104pub struct SpatialLengthScaleOptimizationOptions {
4105 pub enabled: bool,
4109 pub max_outer_iter: usize,
4111 pub rel_tol: f64,
4113 pub log_step: f64,
4115 pub min_length_scale: f64,
4117 pub max_length_scale: f64,
4119 pub pilot_subsample_threshold: usize,
4132}
4133
4134impl Default for SpatialLengthScaleOptimizationOptions {
4135 fn default() -> Self {
4136 Self {
4137 enabled: true,
4138 max_outer_iter: 80,
4139 rel_tol: 1e-4,
4140 log_step: std::f64::consts::LN_2,
4141 min_length_scale: 1e-3,
4142 max_length_scale: 1e3,
4143 pilot_subsample_threshold: 10_000,
4144 }
4145 }
4146}
4147
4148impl SpatialLengthScaleOptimizationOptions {
4149 pub fn validate(&self) -> Result<(), String> {
4167 if !self.min_length_scale.is_finite() || self.min_length_scale <= 0.0 {
4168 return Err(SmoothError::invalid_config(format!(
4169 "SpatialLengthScaleOptimizationOptions::min_length_scale must be > 0 and finite, got {}",
4170 self.min_length_scale
4171 ))
4172 .into());
4173 }
4174 if !self.max_length_scale.is_finite() || self.max_length_scale <= 0.0 {
4175 return Err(SmoothError::invalid_config(format!(
4176 "SpatialLengthScaleOptimizationOptions::max_length_scale must be > 0 and finite, got {}",
4177 self.max_length_scale
4178 ))
4179 .into());
4180 }
4181 if self.min_length_scale >= self.max_length_scale {
4182 return Err(SmoothError::invalid_config(format!(
4183 "SpatialLengthScaleOptimizationOptions requires min_length_scale < max_length_scale, got min={} max={}",
4184 self.min_length_scale, self.max_length_scale
4185 ))
4186 .into());
4187 }
4188 if !self.rel_tol.is_finite() || self.rel_tol <= 0.0 {
4189 return Err(SmoothError::invalid_config(format!(
4190 "SpatialLengthScaleOptimizationOptions::rel_tol must be > 0 and finite, got {}",
4191 self.rel_tol
4192 ))
4193 .into());
4194 }
4195 if !self.log_step.is_finite() || self.log_step <= 0.0 {
4196 return Err(SmoothError::invalid_config(format!(
4197 "SpatialLengthScaleOptimizationOptions::log_step must be > 0 and finite, got {}",
4198 self.log_step
4199 ))
4200 .into());
4201 }
4202 Ok(())
4203 }
4204}
4205
4206#[derive(Debug, Clone)]
4207pub struct RandomEffectBlock {
4208 pub name: String,
4209 pub group_ids: Vec<Option<usize>>,
4212 pub num_groups: usize,
4213 pub kept_levels: Vec<u64>,
4214}
4215
4216pub const BLOCK_SPARSE_ZERO_EPS: f64 = 1e-12;
4217
4218pub const BLOCK_SPARSE_MAX_DENSITY: f64 = 0.20;
4219
4220pub fn blocks_have_intrinsic_sparse_structure(blocks: &[DesignBlock]) -> bool {
4221 blocks
4222 .iter()
4223 .any(|block| matches!(block, DesignBlock::Sparse(_) | DesignBlock::RandomEffect(_)))
4224}
4225
4226pub fn sparse_compatible_block_nnz(block: &DesignBlock) -> Option<usize> {
4227 match block {
4228 DesignBlock::Intercept(n) => Some(*n),
4229 DesignBlock::RandomEffect(op) => {
4230 Some(op.group_ids.iter().filter(|gid| gid.is_some()).count())
4231 }
4232 DesignBlock::Sparse(sparse) => Some(sparse.val().len()),
4233 DesignBlock::Dense(dense) => dense.as_dense_ref().map(|matrix| {
4234 matrix
4235 .iter()
4236 .filter(|&&value| value.abs() > BLOCK_SPARSE_ZERO_EPS)
4237 .count()
4238 }),
4239 }
4240}
4241
4242pub fn try_build_sparse_design_from_blocks(
4243 blocks: &[DesignBlock],
4244) -> Result<Option<DesignMatrix>, BasisError> {
4245 if blocks.is_empty() {
4246 return Ok(None);
4247 }
4248 let nrows = blocks[0].nrows();
4249 let ncols: usize = blocks.iter().map(DesignBlock::ncols).sum();
4250 if nrows == 0 || ncols == 0 || ncols <= 32 {
4251 return Ok(None);
4252 }
4253
4254 let preserve_sparse_storage = blocks_have_intrinsic_sparse_structure(blocks);
4255 let sparse_nnz_limit = if preserve_sparse_storage {
4256 usize::MAX
4257 } else {
4258 let total_cells = nrows.saturating_mul(ncols);
4259 ((total_cells as f64) * BLOCK_SPARSE_MAX_DENSITY).floor() as usize
4260 };
4261 let mut nnz = 0usize;
4262 for block in blocks {
4263 let block_nnz = if let Some(block_nnz) = sparse_compatible_block_nnz(block) {
4264 block_nnz
4265 } else {
4266 return Ok(None);
4267 };
4268 nnz = nnz.saturating_add(block_nnz);
4269 if nnz > sparse_nnz_limit {
4270 return Ok(None);
4271 }
4272 }
4273
4274 let mut triplets = Vec::<Triplet<usize, usize, f64>>::with_capacity(nnz);
4275 let mut col_offset = 0usize;
4276 for block in blocks {
4277 match block {
4278 DesignBlock::Intercept(n) => {
4279 for row in 0..*n {
4280 triplets.push(Triplet::new(row, col_offset, 1.0));
4281 }
4282 }
4283 DesignBlock::RandomEffect(op) => {
4284 for (row, group_id) in op.group_ids.iter().enumerate() {
4285 if let Some(group) = group_id {
4286 triplets.push(Triplet::new(row, col_offset + group, 1.0));
4287 }
4288 }
4289 }
4290 DesignBlock::Sparse(sparse) => {
4291 let (symbolic, values) = sparse.parts();
4292 let col_ptr = symbolic.col_ptr();
4293 let row_idx = symbolic.row_idx();
4294 for col in 0..sparse.ncols() {
4295 for idx in col_ptr[col]..col_ptr[col + 1] {
4296 let value = values[idx];
4297 if value.abs() > BLOCK_SPARSE_ZERO_EPS {
4298 triplets.push(Triplet::new(row_idx[idx], col_offset + col, value));
4299 }
4300 }
4301 }
4302 }
4303 DesignBlock::Dense(dense) => {
4304 let matrix = dense.as_dense_ref().ok_or_else(|| {
4305 BasisError::InvalidInput(
4306 "sparse-compatible block assembly requires materialized dense blocks"
4307 .to_string(),
4308 )
4309 })?;
4310 for row in 0..matrix.nrows() {
4311 for col in 0..matrix.ncols() {
4312 let value = matrix[[row, col]];
4313 if value.abs() > BLOCK_SPARSE_ZERO_EPS {
4314 triplets.push(Triplet::new(row, col_offset + col, value));
4315 }
4316 }
4317 }
4318 }
4319 }
4320 col_offset += block.ncols();
4321 }
4322
4323 let sparse = SparseColMat::try_new_from_triplets(nrows, ncols, &triplets).map_err(|_| {
4324 BasisError::SparseCreation("failed to assemble sparse term-collection design".to_string())
4325 })?;
4326 Ok(Some(DesignMatrix::Sparse(
4327 gam_linalg::matrix::SparseDesignMatrix::new(sparse),
4328 )))
4329}
4330
4331pub fn assemble_term_collection_design_matrix(
4332 blocks: Vec<DesignBlock>,
4333) -> Result<DesignMatrix, BasisError> {
4334 if let Some(sparse) = try_build_sparse_design_from_blocks(&blocks)? {
4335 return Ok(sparse);
4336 }
4337 let block_op = BlockDesignOperator::new(blocks).map_err(|e| {
4338 BasisError::InvalidInput(format!("failed to build block design operator: {e}"))
4339 })?;
4340 Ok(DesignMatrix::Dense(
4341 gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(block_op)),
4342 ))
4343}
4344
4345pub fn select_columns(
4346 data: ArrayView2<'_, f64>,
4347 cols: &[usize],
4348) -> Result<Array2<f64>, BasisError> {
4349 let n = data.nrows();
4350 let p = data.ncols();
4351 for &c in cols {
4352 if c >= p {
4353 crate::bail_dim_basis!("feature column {c} is out of bounds for data with {p} columns");
4354 }
4355 }
4356 let mut out = Array2::<f64>::zeros((n, cols.len()));
4357 for (j, &c) in cols.iter().enumerate() {
4358 out.column_mut(j).assign(&data.column(c));
4359 }
4360 Ok(out)
4361}
4362
4363pub fn nonfinite_value_label(value: f64) -> &'static str {
4364 if value.is_nan() {
4365 "NaN"
4366 } else if value.is_sign_positive() {
4367 "+Inf"
4368 } else {
4369 "-Inf"
4370 }
4371}
4372
4373pub fn validate_term_feature_column_finite(
4374 data: ArrayView2<'_, f64>,
4375 term_kind: &str,
4376 term_name: &str,
4377 feature_col: usize,
4378) -> Result<(), BasisError> {
4379 let p = data.ncols();
4380 if feature_col >= p {
4381 crate::bail_dim_basis!(
4382 "{term_kind} term '{term_name}' feature column {feature_col} out of bounds for {p} columns"
4383 );
4384 }
4385 for (row, &value) in data.column(feature_col).iter().enumerate() {
4386 if !value.is_finite() {
4387 crate::bail_invalid_basis!(
4388 "{term_kind} term '{term_name}' feature column {feature_col} row {row} contains non-finite value {}",
4389 nonfinite_value_label(value)
4390 );
4391 }
4392 }
4393 Ok(())
4394}
4395
4396pub fn validate_smooth_terms_finite_inputs(
4397 data: ArrayView2<'_, f64>,
4398 terms: &[SmoothTermSpec],
4399) -> Result<(), BasisError> {
4400 for term in terms {
4401 for feature_col in smooth_term_feature_cols(term) {
4402 validate_term_feature_column_finite(data, "smooth", &term.name, feature_col)?;
4403 }
4404 }
4405 Ok(())
4406}
4407
4408pub fn validate_term_collection_finite_inputs(
4409 data: ArrayView2<'_, f64>,
4410 spec: &TermCollectionSpec,
4411) -> Result<(), BasisError> {
4412 for term in &spec.linear_terms {
4413 validate_term_feature_column_finite(data, "linear", &term.name, term.feature_col)?;
4414 }
4415 for term in &spec.random_effect_terms {
4416 validate_term_feature_column_finite(data, "random-effect", &term.name, term.feature_col)?;
4417 }
4418 validate_smooth_terms_finite_inputs(data, &spec.smooth_terms)
4419}
4420
4421#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
4422pub struct JointSpatialCenterGroupKey {
4423 feature_cols: Vec<usize>,
4424 strategy_kind: CenterStrategyKind,
4425 strategy_aux: usize,
4426 requested_num_centers: usize,
4427 input_scale_bits: Option<Vec<u64>>,
4428}
4429
4430pub fn spatial_term_min_center_count(term: &SmoothTermSpec) -> usize {
4431 match &term.basis {
4432 SmoothBasisSpec::ThinPlate { feature_cols, .. } => feature_cols.len() + 1,
4433 SmoothBasisSpec::Duchon {
4434 feature_cols, spec, ..
4435 } => match spec.nullspace_order {
4436 crate::basis::DuchonNullspaceOrder::Zero => 1,
4437 crate::basis::DuchonNullspaceOrder::Linear => feature_cols.len() + 1,
4438 crate::basis::DuchonNullspaceOrder::Degree(degree) => {
4439 crate::basis::duchon_nullspace_dimension(feature_cols.len(), degree)
4440 }
4441 },
4442 SmoothBasisSpec::Matern { .. } => 1,
4443 _ => 1,
4444 }
4445}
4446
4447pub fn spatial_term_group_key(term: &SmoothTermSpec) -> Option<JointSpatialCenterGroupKey> {
4448 let (feature_cols, strategy, input_scales) = match &term.basis {
4449 SmoothBasisSpec::ThinPlate {
4450 feature_cols,
4451 spec,
4452 input_scales,
4453 } => (feature_cols, &spec.center_strategy, input_scales.as_ref()),
4454 SmoothBasisSpec::Matern {
4455 feature_cols,
4456 spec,
4457 input_scales,
4458 } => (feature_cols, &spec.center_strategy, input_scales.as_ref()),
4459 SmoothBasisSpec::Duchon {
4460 feature_cols,
4461 spec,
4462 input_scales,
4463 } => (feature_cols, &spec.center_strategy, input_scales.as_ref()),
4464 _ => return None,
4465 };
4466 let strategy_kind = center_strategy_kind(strategy);
4467 let strategy_aux = match strategy {
4468 CenterStrategy::Auto(inner) => match inner.as_ref() {
4469 CenterStrategy::KMeans { max_iter, .. } => *max_iter,
4470 CenterStrategy::UniformGrid { points_per_dim } => *points_per_dim,
4471 _ => 0,
4472 },
4473 CenterStrategy::KMeans { max_iter, .. } => *max_iter,
4474 CenterStrategy::UniformGrid { points_per_dim } => *points_per_dim,
4475 _ => 0,
4476 };
4477 Some(JointSpatialCenterGroupKey {
4478 feature_cols: feature_cols.clone(),
4479 strategy_kind,
4480 strategy_aux,
4481 requested_num_centers: strategy.planned_num_centers(feature_cols.len()),
4482 input_scale_bits: input_scales
4483 .map(|values| values.iter().map(|value| value.to_bits()).collect()),
4484 })
4485}
4486
4487pub fn spatial_term_center_strategy(term: &SmoothTermSpec) -> Option<&CenterStrategy> {
4488 match &term.basis {
4489 SmoothBasisSpec::ThinPlate { spec, .. } => Some(&spec.center_strategy),
4490 SmoothBasisSpec::Matern { spec, .. } => Some(&spec.center_strategy),
4491 SmoothBasisSpec::Duchon { spec, .. } => Some(&spec.center_strategy),
4492 _ => None,
4493 }
4494}
4495
4496pub fn set_spatial_term_centers(
4497 term: &mut SmoothTermSpec,
4498 centers: Array2<f64>,
4499) -> Result<(), BasisError> {
4500 match &mut term.basis {
4501 SmoothBasisSpec::ThinPlate { spec, .. } => {
4502 spec.center_strategy = CenterStrategy::UserProvided(centers);
4503 Ok(())
4504 }
4505 SmoothBasisSpec::Matern { spec, .. } => {
4506 spec.center_strategy = CenterStrategy::UserProvided(centers);
4507 Ok(())
4508 }
4509 SmoothBasisSpec::Duchon { spec, .. } => {
4510 spec.center_strategy = CenterStrategy::UserProvided(centers);
4511 Ok(())
4512 }
4513 _ => Err(BasisError::InvalidInput(format!(
4514 "term '{}' does not support spatial center planning",
4515 term.name
4516 ))),
4517 }
4518}
4519
4520pub fn standardized_spatial_term_data(
4521 data: ArrayView2<'_, f64>,
4522 term: &SmoothTermSpec,
4523) -> Result<Array2<f64>, BasisError> {
4524 let (feature_cols, input_scales) = match &term.basis {
4525 SmoothBasisSpec::ThinPlate {
4526 feature_cols,
4527 input_scales,
4528 ..
4529 }
4530 | SmoothBasisSpec::Matern {
4531 feature_cols,
4532 input_scales,
4533 ..
4534 }
4535 | SmoothBasisSpec::Duchon {
4536 feature_cols,
4537 input_scales,
4538 ..
4539 } => (feature_cols, input_scales.as_ref()),
4540 _ => {
4541 crate::bail_invalid_basis!("term '{}' is not a spatial smooth", term.name);
4542 }
4543 };
4544 let mut x = select_columns(data, feature_cols)?;
4545 if let Some(scales) = input_scales {
4546 apply_input_standardization(&mut x, scales);
4547 } else if let Some(scales) = compute_spatial_input_scales(x.view()) {
4548 apply_input_standardization(&mut x, &scales);
4549 }
4550 Ok(x)
4551}
4552
4553pub fn plan_joint_spatial_centers_for_term_blocks(
4554 data: ArrayView2<'_, f64>,
4555 term_blocks: &[Vec<SmoothTermSpec>],
4556) -> Result<Vec<Vec<SmoothTermSpec>>, BasisError> {
4557 let mut planned_blocks = term_blocks.to_vec();
4558 let n = data.nrows();
4559 let mut groups: BTreeMap<JointSpatialCenterGroupKey, Vec<(usize, usize)>> = BTreeMap::new();
4560
4561 for (block_idx, terms) in planned_blocks.iter().enumerate() {
4562 for (term_idx, term) in terms.iter().enumerate() {
4563 let Some(strategy) = spatial_term_center_strategy(term) else {
4564 continue;
4565 };
4566 if !center_strategy_is_auto(strategy) {
4567 continue;
4568 }
4569 let Some(group_key) = spatial_term_group_key(term) else {
4570 continue;
4571 };
4572 if !matches!(
4573 group_key.strategy_kind,
4574 CenterStrategyKind::EqualMass
4575 | CenterStrategyKind::EqualMassCovarRepresentative
4576 | CenterStrategyKind::FarthestPoint
4577 | CenterStrategyKind::KMeans
4578 | CenterStrategyKind::UniformGrid
4579 ) {
4580 continue;
4581 }
4582 groups
4583 .entry(group_key)
4584 .or_default()
4585 .push((block_idx, term_idx));
4586 }
4587 }
4588
4589 for (group_key, members) in groups {
4590 if members.len() < 2 {
4591 continue;
4592 }
4593 let min_required = members
4594 .iter()
4595 .map(|&(block_idx, term_idx)| {
4596 spatial_term_min_center_count(&planned_blocks[block_idx][term_idx])
4597 })
4598 .max()
4599 .unwrap_or(1);
4600 let joint_centers = group_key
4601 .requested_num_centers
4602 .max(min_required)
4603 .min(n.max(1));
4604 let (first_block_idx, first_term_idx) = members[0];
4605 let prototype = &planned_blocks[first_block_idx][first_term_idx];
4606 let standardized = standardized_spatial_term_data(data, prototype)?;
4607 let strategy = spatial_term_center_strategy(prototype).ok_or_else(|| {
4608 BasisError::InvalidInput(format!(
4609 "term '{}' lost its spatial center strategy during joint planning",
4610 prototype.name
4611 ))
4612 })?;
4613 let joint_strategy = center_strategy_with_num_centers(
4614 strategy,
4615 joint_centers,
4616 group_key.feature_cols.len(),
4617 )?;
4618 let shared_centers = select_centers_by_strategy(standardized.view(), &joint_strategy)?;
4619 log::info!(
4620 "sharing {} spatial centers across {} smooth terms over columns {:?} (requested {} centers)",
4621 shared_centers.nrows(),
4622 members.len(),
4623 group_key.feature_cols,
4624 group_key.requested_num_centers,
4625 );
4626 for (block_idx, term_idx) in members {
4627 set_spatial_term_centers(
4628 &mut planned_blocks[block_idx][term_idx],
4629 shared_centers.clone(),
4630 )?;
4631 }
4632 }
4633
4634 for block in planned_blocks.iter_mut() {
4641 for term in block.iter_mut() {
4642 auto_init_length_scale_in_place(data, term);
4643 }
4644 }
4645
4646 Ok(planned_blocks)
4647}
4648
4649const AUTO_LENGTH_SCALE_FLOOR: f64 = 1e-6;
4652
4653fn feature_columns_max_range(data: ArrayView2<'_, f64>, feature_cols: &[usize]) -> Option<f64> {
4656 let mut max_range = 0.0_f64;
4657 for &c in feature_cols {
4658 if c >= data.ncols() {
4659 continue;
4660 }
4661 let col = data.column(c);
4662 let mut lo = f64::INFINITY;
4663 let mut hi = f64::NEG_INFINITY;
4664 for &v in col.iter() {
4665 if v.is_finite() {
4666 if v < lo {
4667 lo = v;
4668 }
4669 if v > hi {
4670 hi = v;
4671 }
4672 }
4673 }
4674 if hi > lo {
4675 let r = hi - lo;
4676 if r > max_range {
4677 max_range = r;
4678 }
4679 }
4680 }
4681 if max_range.is_finite() && max_range > 0.0 {
4682 Some(max_range)
4683 } else {
4684 None
4685 }
4686}
4687
4688fn feature_columns_rotation_invariant_range(
4699 data: ArrayView2<'_, f64>,
4700 feature_cols: &[usize],
4701) -> Option<f64> {
4702 let cols: Vec<usize> = feature_cols
4703 .iter()
4704 .copied()
4705 .filter(|&c| c < data.ncols())
4706 .collect();
4707 if cols.is_empty() {
4708 return None;
4709 }
4710 let mut points: Vec<Vec<f64>> = data
4711 .rows()
4712 .into_iter()
4713 .filter_map(|row| {
4714 let point: Vec<f64> = cols.iter().map(|&column| row[column]).collect();
4715 point.iter().all(|value| value.is_finite()).then_some(point)
4716 })
4717 .collect();
4718 if points.is_empty() {
4719 return None;
4720 }
4721 points.sort_by(|left, right| {
4722 left.iter()
4723 .zip(right)
4724 .find_map(|(a, b)| {
4725 let ordering = a.total_cmp(b);
4726 ordering.is_ne().then_some(ordering)
4727 })
4728 .unwrap_or(std::cmp::Ordering::Equal)
4729 });
4730
4731 let dimensions = cols.len();
4732 let count = points.len() as f64;
4733 let mut centroid = vec![0.0_f64; dimensions];
4734 for point in &points {
4735 for (coordinate, value) in centroid.iter_mut().zip(point) {
4736 *coordinate += *value;
4737 }
4738 }
4739 for coordinate in &mut centroid {
4740 *coordinate /= count;
4741 }
4742
4743 let mut covariance = Array2::<f64>::zeros((dimensions, dimensions));
4744 for point in &points {
4745 for row in 0..dimensions {
4746 let centered_row = point[row] - centroid[row];
4747 for column in 0..=row {
4748 covariance[[row, column]] += centered_row * (point[column] - centroid[column]);
4749 }
4750 }
4751 }
4752 for row in 0..dimensions {
4753 for column in 0..=row {
4754 let value = covariance[[row, column]] / count;
4755 covariance[[row, column]] = value;
4756 covariance[[column, row]] = value;
4757 }
4758 }
4759
4760 use gam_linalg::faer_ndarray::FaerEigh;
4761 let (eigenvalues, _) = covariance
4762 .eigh(faer::Side::Lower)
4763 .expect("finite covariance must have a symmetric eigendecomposition");
4764 let leading_variance = eigenvalues[eigenvalues.len() - 1];
4765 let extent = (12.0 * leading_variance).sqrt();
4766 if extent.is_finite() && extent > 0.0 {
4767 Some(extent)
4768 } else {
4769 None
4770 }
4771}
4772
4773pub fn auto_initial_length_scale(data: ArrayView2<'_, f64>, feature_cols: &[usize]) -> f64 {
4780 let n = data.nrows();
4781 if n == 0 || feature_cols.is_empty() {
4782 return 1.0;
4783 }
4784 let Some(max_range) = feature_columns_max_range(data, feature_cols) else {
4785 return 1.0;
4786 };
4787 let init = max_range / (n as f64).sqrt();
4788 init.max(AUTO_LENGTH_SCALE_FLOOR).min(max_range)
4789}
4790
4791pub fn auto_initial_length_scale_for_centers(
4814 data: ArrayView2<'_, f64>,
4815 feature_cols: &[usize],
4816 num_centers: usize,
4817) -> f64 {
4818 let n = data.nrows();
4819 if n == 0 || feature_cols.is_empty() {
4820 return 1.0;
4821 }
4822 let Some(max_range) = feature_columns_rotation_invariant_range(data, feature_cols) else {
4833 return 1.0;
4834 };
4835 let resolution_points = n.max(num_centers).max(1) as f64;
4841 let spacing = max_range / resolution_points.sqrt();
4842 spacing.max(AUTO_LENGTH_SCALE_FLOOR).min(max_range)
4843}
4844
4845pub fn matern_low_rank_center_resolution_length_scale(
4855 data: ArrayView2<'_, f64>,
4856 feature_cols: &[usize],
4857 num_centers: usize,
4858) -> Option<f64> {
4859 if data.nrows() == 0 || feature_cols.is_empty() || num_centers == 0 {
4860 return None;
4861 }
4862 let extent = feature_columns_rotation_invariant_range(data, feature_cols)?;
4863 let length_scale = extent / (num_centers as f64).sqrt();
4864 Some(
4865 length_scale
4866 .max(AUTO_LENGTH_SCALE_FLOOR)
4867 .min(extent),
4868 )
4869}
4870
4871pub fn auto_initial_length_scale_for_low_rank_centers(
4881 data: ArrayView2<'_, f64>,
4882 feature_cols: &[usize],
4883 num_centers: usize,
4884) -> f64 {
4885 if data.nrows() == 0 || feature_cols.is_empty() {
4886 return 1.0;
4887 }
4888 let Some(max_range) = feature_columns_max_range(data, feature_cols) else {
4889 return 1.0;
4890 };
4891 let resolution_points = num_centers.max(1) as f64;
4892 let spacing = max_range / resolution_points.sqrt();
4893 spacing.max(AUTO_LENGTH_SCALE_FLOOR).min(max_range)
4894}
4895
4896fn center_strategy_requested_count(strategy: &CenterStrategy) -> Option<usize> {
4899 match strategy {
4900 CenterStrategy::Auto(inner) => center_strategy_requested_count(inner),
4901 CenterStrategy::UserProvided(centers) => Some(centers.nrows()),
4902 CenterStrategy::EqualMass { num_centers }
4903 | CenterStrategy::EqualMassCovarRepresentative { num_centers }
4904 | CenterStrategy::FarthestPoint { num_centers }
4905 | CenterStrategy::KMeans { num_centers, .. } => Some(*num_centers),
4906 CenterStrategy::UniformGrid { .. } => None,
4907 }
4908}
4909
4910pub fn auto_init_length_scale_in_place(data: ArrayView2<'_, f64>, term: &mut SmoothTermSpec) {
4914 auto_init_length_scale_in_basis(data, &mut term.basis);
4915}
4916
4917pub fn auto_init_length_scale_in_basis(data: ArrayView2<'_, f64>, basis: &mut SmoothBasisSpec) {
4930 match basis {
4931 SmoothBasisSpec::Matern {
4932 feature_cols, spec, ..
4933 } => {
4934 if spec.length_scale == 0.0 {
4935 spec.length_scale = match center_strategy_requested_count(&spec.center_strategy) {
4944 Some(k) => auto_initial_length_scale_for_centers(data, feature_cols, k),
4945 None => auto_initial_length_scale(data, feature_cols),
4946 };
4947 }
4948 }
4949 SmoothBasisSpec::ThinPlate {
4950 feature_cols, spec, ..
4951 } => {
4952 if spec.length_scale == 0.0 {
4953 spec.length_scale = match center_strategy_requested_count(&spec.center_strategy) {
4954 Some(k) => {
4955 auto_initial_length_scale_for_low_rank_centers(data, feature_cols, k)
4956 }
4957 None => auto_initial_length_scale(data, feature_cols),
4958 };
4959 }
4960 }
4961 SmoothBasisSpec::ByVariable { inner, .. }
4962 | SmoothBasisSpec::FactorSumToZero { inner, .. } => {
4963 auto_init_length_scale_in_basis(data, inner);
4964 }
4965 SmoothBasisSpec::BySmooth { smooth, .. } => {
4966 auto_init_length_scale_in_basis(data, smooth);
4967 }
4968 _ => {}
4969 }
4970}
4971
4972impl LinearFitConditioning {
4973 pub fn from_columns(design: &TermCollectionDesign, selected_cols: &[usize]) -> Self {
4974 const SCALE_EPS: f64 = 1e-12;
4975 let n = design.design.nrows();
4976 let p = design.design.ncols();
4977 let mut columns = Vec::with_capacity(selected_cols.len());
4978 if n == 0 || selected_cols.is_empty() {
4979 return Self {
4980 intercept_idx: design.intercept_range.start,
4981 columns,
4982 };
4983 }
4984 let chunk_rows = gam_linalg::utils::row_chunk_for_byte_budget(n, p);
4985 let mut sums = vec![0.0_f64; selected_cols.len()];
4991 for start in (0..n).step_by(chunk_rows) {
4992 let end = (start + chunk_rows).min(n);
4993 let chunk = design
4994 .design
4995 .try_row_chunk(start..end)
4996 .expect("LinearFitConditioning::from_columns row chunk failed");
4997 for (k, &col_idx) in selected_cols.iter().enumerate() {
4998 let column = chunk.column(col_idx);
4999 for &v in column.iter() {
5000 sums[k] += v;
5001 }
5002 }
5003 }
5004 let inv_n = 1.0_f64 / n as f64;
5005 let means: Vec<f64> = sums.iter().map(|&s| s * inv_n).collect();
5006 let mut sq_devs = vec![0.0_f64; selected_cols.len()];
5007 for start in (0..n).step_by(chunk_rows) {
5008 let end = (start + chunk_rows).min(n);
5009 let chunk = design
5010 .design
5011 .try_row_chunk(start..end)
5012 .expect("LinearFitConditioning::from_columns row chunk failed");
5013 for (k, &col_idx) in selected_cols.iter().enumerate() {
5014 let mean_k = means[k];
5015 let column = chunk.column(col_idx);
5016 for &v in column.iter() {
5017 let d = v - mean_k;
5018 sq_devs[k] += d * d;
5019 }
5020 }
5021 }
5022 for (k, &col_idx) in selected_cols.iter().enumerate() {
5023 let mean = means[k];
5024 let var = sq_devs[k] * inv_n;
5025 let (mean, scale) = if var.is_finite() && var > SCALE_EPS * SCALE_EPS {
5026 (mean, var.sqrt())
5027 } else {
5028 (0.0, 1.0)
5031 };
5032 columns.push(LinearColumnConditioning {
5033 col_idx,
5034 mean,
5035 scale,
5036 });
5037 }
5038 Self {
5039 intercept_idx: design.intercept_range.start,
5040 columns,
5041 }
5042 }
5043
5044 pub fn apply_to_design(&self, design: &Array2<f64>) -> Array2<f64> {
5045 let mut out = design.clone();
5046 for col in &self.columns {
5047 {
5048 let mut dst = out.column_mut(col.col_idx);
5049 dst -= col.mean;
5050 }
5051 if col.scale != 1.0 {
5052 out.column_mut(col.col_idx).mapv_inplace(|v| v / col.scale);
5053 }
5054 }
5055 out
5056 }
5057
5058 fn transform_matrix_columnswith_a(&self, mat: &Array2<f64>) -> Array2<f64> {
5059 let mut out = mat.clone();
5060 let intercept = self.intercept_idx;
5061 for col in &self.columns {
5062 let intercept_col = out.column(intercept).to_owned();
5063 let mut target = out.column_mut(col.col_idx);
5064 target -= &(intercept_col * col.mean);
5065 if col.scale != 1.0 {
5066 target.mapv_inplace(|v| v / col.scale);
5067 }
5068 }
5069 out
5070 }
5071
5072 fn transform_matrixrowswith_a_transpose(&self, mat: &Array2<f64>) -> Array2<f64> {
5073 let mut out = mat.clone();
5074 let intercept = self.intercept_idx;
5075 for col in &self.columns {
5076 let interceptrow = out.row(intercept).to_owned();
5077 let mut target = out.row_mut(col.col_idx);
5078 target -= &(interceptrow * col.mean);
5079 if col.scale != 1.0 {
5080 target.mapv_inplace(|v| v / col.scale);
5081 }
5082 }
5083 out
5084 }
5085
5086 fn left_multiply_by_m_inv_transpose(&self, mat_internal: &Array2<f64>) -> Array2<f64> {
5091 let mut out = mat_internal.clone();
5092 let intercept = self.intercept_idx;
5093 let interceptrow_snapshot = mat_internal.row(intercept).to_owned();
5094 for col in &self.columns {
5095 if col.scale != 1.0 {
5096 out.row_mut(col.col_idx).mapv_inplace(|v| v * col.scale);
5097 }
5098 if col.mean != 0.0 {
5099 let mut target = out.row_mut(col.col_idx);
5100 target += &(&interceptrow_snapshot * col.mean);
5101 }
5102 }
5103 out
5104 }
5105
5106 fn right_multiply_by_m_inv(&self, mat_internal: &Array2<f64>) -> Array2<f64> {
5109 let mut out = mat_internal.clone();
5110 let intercept = self.intercept_idx;
5111 let intercept_col_snapshot = mat_internal.column(intercept).to_owned();
5112 for col in &self.columns {
5113 if col.scale != 1.0 {
5114 out.column_mut(col.col_idx).mapv_inplace(|v| v * col.scale);
5115 }
5116 if col.mean != 0.0 {
5117 let mut target = out.column_mut(col.col_idx);
5118 target += &(&intercept_col_snapshot * col.mean);
5119 }
5120 }
5121 out
5122 }
5123
5124 pub fn transform_blockwise_penalties_to_internal(
5131 &self,
5132 penalties: &[BlockwisePenalty],
5133 p: usize,
5134 ) -> Vec<crate::penalty_spec::PenaltySpec> {
5135 let conditioning_cols: std::collections::HashSet<usize> =
5136 self.columns.iter().map(|c| c.col_idx).collect();
5137 penalties
5138 .iter()
5139 .map(|bp| {
5140 let overlaps =
5141 (bp.col_range.start..bp.col_range.end).any(|j| conditioning_cols.contains(&j));
5142 if overlaps {
5143 let global = bp.to_global(p);
5146 let right = self.transform_matrix_columnswith_a(&global);
5147 let transformed = self.transform_matrixrowswith_a_transpose(&right);
5148 crate::penalty_spec::PenaltySpec::Dense(transformed)
5149 } else {
5150 crate::penalty_spec::PenaltySpec::from_blockwise(bp.clone())
5153 }
5154 })
5155 .collect()
5156 }
5157
5158 pub fn backtransform_beta(&self, beta_internal: &Array1<f64>) -> Array1<f64> {
5159 let mut beta = beta_internal.clone();
5160 let intercept = self.intercept_idx;
5161 for col in &self.columns {
5162 beta[intercept] -= beta_internal[col.col_idx] * col.mean / col.scale;
5163 beta[col.col_idx] = beta_internal[col.col_idx] / col.scale;
5164 }
5165 beta
5166 }
5167
5168 pub fn transform_penalized_hessian_to_original(&self, h_internal: &Array2<f64>) -> Array2<f64> {
5171 let right = self.right_multiply_by_m_inv(h_internal);
5172 self.left_multiply_by_m_inv_transpose(&right)
5173 }
5174
5175 pub fn internal_bounds_for(&self, col_idx: usize, min: f64, max: f64) -> (f64, f64) {
5176 if let Some(col) = self.columns.iter().find(|c| c.col_idx == col_idx) {
5177 (min * col.scale, max * col.scale)
5178 } else {
5179 (min, max)
5180 }
5181 }
5182}
5183
5184pub fn freeze_raw_spatial_metadata(metadata: BasisMetadata, raw_cols: usize) -> BasisMetadata {
5185 match metadata {
5186 BasisMetadata::ThinPlate {
5187 centers,
5188 length_scale,
5189 periodic,
5190 identifiability_transform: None,
5191 input_scales,
5192 radial_reparam,
5193 } => BasisMetadata::ThinPlate {
5194 centers,
5195 length_scale,
5196 periodic,
5197 identifiability_transform: Some(Array2::eye(raw_cols)),
5198 input_scales,
5199 radial_reparam,
5200 },
5201 BasisMetadata::Duchon {
5202 centers,
5203 length_scale,
5204 periodic,
5205 power,
5206 nullspace_order,
5207 identifiability_transform: None,
5208 input_scales,
5209 aniso_log_scales,
5210 operator_collocation_points,
5211 radial_reparam,
5212 } => BasisMetadata::Duchon {
5213 centers,
5214 length_scale,
5215 periodic,
5216 power,
5217 nullspace_order,
5218 identifiability_transform: Some(Array2::eye(raw_cols)),
5219 input_scales,
5220 aniso_log_scales,
5221 operator_collocation_points,
5222 radial_reparam,
5223 },
5224 other => other,
5225 }
5226}
5227
5228pub fn matern_operator_penalty_triplet_from_metadata(
5229 metadata: &BasisMetadata,
5230) -> Result<(Vec<Array2<f64>>, Vec<usize>, Vec<PenaltyInfo>), BasisError> {
5231 let BasisMetadata::Matern {
5232 centers,
5233 length_scale,
5234 periodic,
5235 nu,
5236 include_intercept,
5237 identifiability_transform,
5238 aniso_log_scales,
5239 input_scales,
5240 ..
5241 } = metadata
5242 else {
5243 crate::bail_invalid_basis!("Matérn operator penalties require Matérn metadata");
5244 };
5245 let penalty_length_scale = match input_scales.as_deref() {
5257 Some(scales) => compensate_length_scale_for_standardization(*length_scale, scales),
5258 None => *length_scale,
5259 };
5260 matern_operator_penalty_triplet_at_length_scale(
5261 centers.view(),
5262 periodic.as_deref(),
5263 identifiability_transform.as_ref(),
5264 *nu,
5265 *include_intercept,
5266 aniso_log_scales.as_deref(),
5267 penalty_length_scale,
5268 )
5269}
5270
5271pub fn matern_operator_penalty_triplet_at_length_scale(
5289 centers: ArrayView2<'_, f64>,
5290 periodic: Option<&[Option<f64>]>,
5291 identifiability_transform: Option<&Array2<f64>>,
5292 nu: crate::basis::MaternNu,
5293 include_intercept: bool,
5294 aniso_log_scales: Option<&[f64]>,
5295 effective_length_scale: f64,
5296) -> Result<(Vec<Array2<f64>>, Vec<usize>, Vec<PenaltyInfo>), BasisError> {
5297 let penalty_centers = crate::basis::expand_periodic_centers(¢ers.to_owned(), periodic)?;
5298 let ops = build_matern_collocation_operator_matrices(
5299 penalty_centers.view(),
5300 None,
5301 effective_length_scale,
5302 nu,
5303 include_intercept,
5304 identifiability_transform.map(|z| z.view()),
5305 aniso_log_scales,
5306 )?;
5307 const ORDER_EPS: f64 = 1e-9;
5314 let d = penalty_centers.ncols();
5315 let m = nu.half_integer_value() + 0.5 * d as f64;
5316 let mut candidates = Vec::with_capacity(3);
5317 for (raw, source, min_order) in [
5318 (ops.d0.t().dot(&ops.d0), PenaltySource::OperatorMass, 0.0),
5319 (ops.d1.t().dot(&ops.d1), PenaltySource::OperatorTension, 1.0),
5320 (
5321 ops.d2.t().dot(&ops.d2),
5322 PenaltySource::OperatorStiffness,
5323 2.0,
5324 ),
5325 ] {
5326 let nondifferentiable_ou = matches!(nu, crate::basis::MaternNu::Half);
5327 if min_order > 0.0 && (nondifferentiable_ou || m + ORDER_EPS < min_order) {
5328 continue;
5329 }
5330 let sym = (&raw + &raw.t()) * 0.5;
5331 let (matrix, normalization_scale) = normalize_penalty_in_constrained_space(&sym);
5332 candidates.push(PenaltyCandidate {
5333 matrix,
5334 nullspace_dim_hint: 0,
5335 source,
5336 normalization_scale,
5337 kronecker_factors: None,
5338 op: None,
5339 });
5340 }
5341 filter_active_penalty_candidates(candidates)
5342}
5343
5344pub fn normalize_penalty_in_constrained_space(matrix: &Array2<f64>) -> (Array2<f64>, f64) {
5345 let matrix = (matrix + &matrix.t().to_owned()) * 0.5;
5350 let matrix = crate::basis::project_penalty_to_psd_cone(&matrix);
5352 let c = matrix.iter().map(|v| v * v).sum::<f64>().sqrt();
5353 if c.is_finite() && c > 0.0 {
5354 (matrix.mapv(|v| v / c), c)
5355 } else {
5356 (matrix, 1.0)
5357 }
5358}
5359
5360pub fn tensor_product_design_from_sparse_marginals(
5361 marginal_sparse: &[&SparseColMat<usize, f64>],
5362) -> Result<SparseColMat<usize, f64>, BasisError> {
5363 if marginal_sparse.is_empty() {
5364 crate::bail_invalid_basis!("TensorBSpline requires at least one marginal basis");
5365 }
5366 let n = marginal_sparse[0].nrows();
5367 for (i, m) in marginal_sparse.iter().enumerate().skip(1) {
5368 if m.nrows() != n {
5369 crate::bail_dim_basis!(
5370 "tensor sparse marginal row mismatch at dim {i}: expected {n}, got {}",
5371 m.nrows()
5372 );
5373 }
5374 }
5375 let dims: Vec<usize> = marginal_sparse.iter().map(|m| m.ncols()).collect();
5376 let total_cols = dims.iter().try_fold(1usize, |acc, &q| {
5377 acc.checked_mul(q)
5378 .ok_or_else(|| BasisError::DimensionMismatch("tensor basis too large".to_string()))
5379 })?;
5380 let mut strides = vec![1usize; dims.len()];
5381 for d in (0..dims.len().saturating_sub(1)).rev() {
5382 strides[d] = strides[d + 1]
5383 .checked_mul(dims[d + 1])
5384 .ok_or_else(|| BasisError::DimensionMismatch("tensor basis too large".to_string()))?;
5385 }
5386
5387 use faer::sparse::SparseRowMat;
5388 let csrs: Vec<SparseRowMat<usize, f64>> = marginal_sparse
5389 .iter()
5390 .enumerate()
5391 .map(|(d, m)| {
5392 m.as_ref().to_row_major().map_err(|e| {
5393 BasisError::SparseCreation(format!(
5394 "tensor sparse marginal {d} CSR conversion failed: {e:?}"
5395 ))
5396 })
5397 })
5398 .collect::<Result<Vec<_>, _>>()?;
5399 let row_ptrs: Vec<&[usize]> = csrs.iter().map(|c| c.symbolic().row_ptr()).collect();
5400 let col_idxs: Vec<&[usize]> = csrs.iter().map(|c| c.symbolic().col_idx()).collect();
5401 let vals: Vec<&[f64]> = csrs.iter().map(|c| c.val()).collect();
5402
5403 use rayon::prelude::*;
5404 const CHUNK: usize = 1024;
5405 let num_chunks = n.div_ceil(CHUNK);
5406 let per_chunk: Vec<Vec<Triplet<usize, usize, f64>>> = (0..num_chunks)
5407 .into_par_iter()
5408 .map(|chunk_idx| {
5409 let row_start = chunk_idx * CHUNK;
5410 let row_end = (row_start + CHUNK).min(n);
5411 let mut chunk_triplets = Vec::<Triplet<usize, usize, f64>>::new();
5412 let mut cur_cols = Vec::<usize>::with_capacity(64);
5413 let mut cur_vals = Vec::<f64>::with_capacity(64);
5414 let mut next_cols = Vec::<usize>::with_capacity(64);
5415 let mut next_vals = Vec::<f64>::with_capacity(64);
5416 for i in row_start..row_end {
5417 cur_cols.clear();
5418 cur_vals.clear();
5419 cur_cols.push(0);
5420 cur_vals.push(1.0);
5421 let mut row_is_zero = false;
5422 for d in 0..dims.len() {
5423 let row_start_d = row_ptrs[d][i];
5424 let row_end_d = row_ptrs[d][i + 1];
5425 if row_start_d == row_end_d {
5426 row_is_zero = true;
5427 break;
5428 }
5429 let stride = strides[d];
5430 next_cols.clear();
5431 next_vals.clear();
5432 next_cols.reserve(cur_cols.len() * (row_end_d - row_start_d));
5433 next_vals.reserve(cur_vals.len() * (row_end_d - row_start_d));
5434 for (&prev_col, &prev_val) in cur_cols.iter().zip(cur_vals.iter()) {
5435 for ptr in row_start_d..row_end_d {
5436 let cj = col_idxs[d][ptr];
5437 let vj = vals[d][ptr];
5438 next_cols.push(prev_col + cj * stride);
5439 next_vals.push(prev_val * vj);
5440 }
5441 }
5442 std::mem::swap(&mut cur_cols, &mut next_cols);
5443 std::mem::swap(&mut cur_vals, &mut next_vals);
5444 }
5445 if row_is_zero {
5446 continue;
5447 }
5448 for (&col, &val) in cur_cols.iter().zip(cur_vals.iter()) {
5449 chunk_triplets.push(Triplet::new(i, col, val));
5450 }
5451 }
5452 chunk_triplets
5453 })
5454 .collect();
5455 let total_nnz: usize = per_chunk.iter().map(Vec::len).sum();
5456 let mut triplets = Vec::<Triplet<usize, usize, f64>>::with_capacity(total_nnz);
5457 for chunk in per_chunk {
5458 triplets.extend(chunk);
5459 }
5460 SparseColMat::try_new_from_triplets(n, total_cols, &triplets).map_err(|e| {
5461 BasisError::SparseCreation(format!(
5462 "failed to assemble sparse tensor product design: {e:?}"
5463 ))
5464 })
5465}
5466
5467pub fn dense_local_margin_to_sparse(
5468 dense: &Array2<f64>,
5469) -> Result<SparseColMat<usize, f64>, BasisError> {
5470 let expected_row_nnz = dense.ncols().min(4);
5471 let mut triplets =
5472 Vec::<Triplet<usize, usize, f64>>::with_capacity(dense.nrows() * expected_row_nnz);
5473 for ((row, col), &value) in dense.indexed_iter() {
5474 if value != 0.0 {
5475 triplets.push(Triplet::new(row, col, value));
5476 }
5477 }
5478 SparseColMat::try_new_from_triplets(dense.nrows(), dense.ncols(), &triplets).map_err(|e| {
5479 BasisError::SparseCreation(format!(
5480 "failed to convert tensor marginal design to sparse form: {e:?}"
5481 ))
5482 })
5483}
5484
5485pub struct TensorMarginRangeNullProjectors {
5486 range: Array2<f64>,
5487 null: Array2<f64>,
5488}
5489
5490pub fn projector_from_columns(columns: &Array2<f64>, indices: &[usize]) -> Array2<f64> {
5491 if indices.is_empty() {
5492 return Array2::<f64>::zeros((columns.nrows(), columns.nrows()));
5493 }
5494 let basis = columns.select(Axis(1), indices);
5495 basis.dot(&basis.t())
5496}
5497
5498pub fn tensor_margin_range_null_projectors(
5499 normalized_marginal_penalties: &[(Array2<f64>, f64)],
5500) -> Result<Vec<TensorMarginRangeNullProjectors>, BasisError> {
5501 normalized_marginal_penalties
5502 .iter()
5503 .enumerate()
5504 .map(|(dim, (penalty, _))| {
5505 let analysis = crate::basis::analyze_penalty_block(penalty)?;
5506 if analysis.rank == 0 {
5507 crate::bail_invalid_basis!(
5508 "t2 separable tensor penalty margin {dim} has rank-zero penalty; \
5509 cannot split penalized and null subspaces"
5510 );
5511 }
5512 let mut range_idx = Vec::<usize>::new();
5513 let mut null_idx = Vec::<usize>::new();
5514 for (idx, &ev) in analysis.eigenvalues.iter().enumerate() {
5515 if ev > analysis.tol {
5516 range_idx.push(idx);
5517 } else {
5518 null_idx.push(idx);
5519 }
5520 }
5521 Ok(TensorMarginRangeNullProjectors {
5522 range: projector_from_columns(&analysis.eigenvectors, &range_idx),
5523 null: projector_from_columns(&analysis.eigenvectors, &null_idx),
5524 })
5525 })
5526 .collect()
5527}
5528
5529pub fn build_tensor_bspline_basis(
5530 data: ArrayView2<'_, f64>,
5531 feature_cols: &[usize],
5532 spec: &TensorBSplineSpec,
5533) -> Result<BasisBuildResult, BasisError> {
5534 if feature_cols.is_empty() {
5535 crate::bail_invalid_basis!("TensorBSpline requires at least one feature column");
5536 }
5537 if feature_cols.len() != spec.marginalspecs.len() {
5538 crate::bail_dim_basis!(
5539 "TensorBSpline feature/spec mismatch: feature_cols={}, marginalspecs={}",
5540 feature_cols.len(),
5541 spec.marginalspecs.len()
5542 );
5543 }
5544 if !spec.periods.is_empty() && spec.periods.len() != feature_cols.len() {
5545 crate::bail_dim_basis!(
5546 "TensorBSpline periods length {} does not match feature count {}",
5547 spec.periods.len(),
5548 feature_cols.len()
5549 );
5550 }
5551 let p = data.ncols();
5552 for &c in feature_cols {
5553 if c >= p {
5554 crate::bail_dim_basis!(
5555 "tensor feature column {c} is out of bounds for data with {p} columns"
5556 );
5557 }
5558 }
5559
5560 let mut marginal_knots = Vec::<Array1<f64>>::with_capacity(feature_cols.len());
5561 let mut marginal_is_cr_flags = Vec::<bool>::with_capacity(feature_cols.len());
5564 let mut marginal_degrees = Vec::<usize>::with_capacity(feature_cols.len());
5565 let mut marginalnum_basis = Vec::<usize>::with_capacity(feature_cols.len());
5566 let mut marginal_penalties = Vec::<Array2<f64>>::with_capacity(feature_cols.len());
5567 let mut marginal_function_grams = Vec::<Array2<f64>>::with_capacity(feature_cols.len());
5568 let mut marginal_designs = Vec::<Array2<f64>>::with_capacity(feature_cols.len());
5569 let mut marginal_effective_periods = Vec::<Option<f64>>::with_capacity(feature_cols.len());
5577 let mut marginal_sparse =
5585 Vec::<Option<SparseColMat<usize, f64>>>::with_capacity(feature_cols.len());
5586
5587 for (dim, (&col, marginalspec)) in feature_cols
5590 .iter()
5591 .zip(spec.marginalspecs.iter())
5592 .enumerate()
5593 {
5594 let mut marginal_unconstrained = marginalspec.clone();
5599 marginal_unconstrained.identifiability = BSplineIdentifiability::None;
5600 let built = build_bspline_basis_1d(data.column(col), &marginal_unconstrained)?;
5601 let (knots, marginal_is_cr, effective_degree, function_gram) = match built.metadata {
5606 BasisMetadata::BSpline1D {
5607 knots,
5608 periodic,
5609 degree,
5610 ..
5611 } => {
5612 let effective_degree = degree.unwrap_or(marginal_unconstrained.degree);
5613 let gram = if spec.double_penalty {
5614 Some(match periodic {
5615 Some((start, period, num_basis)) => {
5616 crate::basis::periodic_bspline_function_gram(
5617 start,
5618 start + period,
5619 effective_degree,
5620 num_basis,
5621 )?
5622 }
5623 None => crate::basis::bspline_function_gram(&knots, effective_degree)?,
5624 })
5625 } else {
5626 None
5627 };
5628 (knots, false, effective_degree, gram)
5629 }
5630 BasisMetadata::CubicRegression1D { knots, .. } => {
5631 let gram = spec
5632 .double_penalty
5633 .then(|| crate::basis::cubic_regression_function_gram(&knots))
5634 .transpose()?;
5635 (knots, true, marginalspec.degree, gram)
5636 }
5637 _ => {
5638 crate::bail_invalid_basis!(
5639 "internal TensorBSpline error at dim {dim}: expected BSpline1D or CubicRegression1D metadata"
5640 );
5641 }
5642 };
5643 let metadata_knots = match marginalspec.knotspec {
5644 BSplineKnotSpec::PeriodicUniform {
5645 data_range,
5646 num_basis,
5647 } => Array1::linspace(data_range.0, data_range.1, num_basis),
5648 _ => knots,
5649 };
5650 if let Some(function_gram) = function_gram {
5651 if function_gram.dim() != (built.design.ncols(), built.design.ncols()) {
5652 crate::bail_dim_basis!(
5653 "internal TensorBSpline error at dim {dim}: function Gram is {:?}, basis has {} columns",
5654 function_gram.dim(),
5655 built.design.ncols()
5656 );
5657 }
5658 marginal_function_grams.push(function_gram);
5659 }
5660 marginal_knots.push(metadata_knots);
5661 marginal_is_cr_flags.push(marginal_is_cr);
5662 marginal_degrees.push(effective_degree);
5663 marginalnum_basis.push(built.design.ncols());
5664 let dense_marginal = built.design.to_dense();
5669 let sparse_view: Option<SparseColMat<usize, f64>> = match built.design.as_sparse() {
5670 Some(sd) => {
5671 let inner: &SparseColMat<usize, f64> = sd;
5672 Some(inner.clone())
5673 }
5674 None => match marginalspec.knotspec {
5675 BSplineKnotSpec::PeriodicUniform { .. } => {
5676 Some(dense_local_margin_to_sparse(&dense_marginal)?)
5677 }
5678 _ => None,
5679 },
5680 };
5681 marginal_sparse.push(sparse_view);
5682 marginal_designs.push(dense_marginal);
5683 marginal_penalties.push(
5684 built
5685 .penalties
5686 .first()
5687 .ok_or_else(|| {
5688 BasisError::InvalidInput(format!(
5689 "internal TensorBSpline error at dim {dim}: missing marginal penalty"
5690 ))
5691 })?
5692 .clone(),
5693 );
5694 built.nullspace_dims.first().ok_or_else(|| {
5695 BasisError::InvalidInput(format!(
5696 "internal TensorBSpline error at dim {dim}: missing marginal nullspace dim"
5697 ))
5698 })?;
5699 let implied_period = match marginalspec.knotspec {
5707 BSplineKnotSpec::PeriodicUniform { data_range, .. } => {
5708 Some(data_range.1 - data_range.0)
5709 }
5710 _ => spec.periods.get(dim).and_then(|p| *p),
5711 };
5712 marginal_effective_periods.push(implied_period);
5713 }
5714
5715 let total_cols: usize = marginalnum_basis.iter().product();
5716 let mut dense_design = (!matches!(spec.identifiability, TensorBSplineIdentifiability::None))
5717 .then(|| tensor_product_design_from_marginals(&marginal_designs))
5718 .transpose()?;
5719 let mut candidates = Vec::<PenaltyCandidate>::with_capacity(
5720 match spec.penalty_decomposition {
5721 TensorBSplinePenaltyDecomposition::MarginalKroneckerSum => marginal_penalties.len(),
5722 TensorBSplinePenaltyDecomposition::Separable => marginal_penalties.len() * 2,
5723 } + if spec.double_penalty { 1 } else { 0 },
5724 );
5725
5726 let normalized_marginal_penalties: Vec<(Array2<f64>, f64)> = marginal_penalties
5734 .iter()
5735 .map(normalize_penalty_in_constrained_space)
5736 .collect();
5737 let tensor_function_gram = if spec.double_penalty {
5738 if marginal_function_grams.len() != marginalnum_basis.len() {
5739 crate::bail_dim_basis!(
5740 "TensorBSpline double penalty requires one function Gram per margin; got {} for {} margins",
5741 marginal_function_grams.len(),
5742 marginalnum_basis.len()
5743 );
5744 }
5745 let mut gram = Array2::<f64>::eye(1);
5746 for marginal_gram in &marginal_function_grams {
5747 gram = kronecker_product(&gram, marginal_gram);
5748 }
5749 Some(gram)
5750 } else {
5751 None
5752 };
5753 let joint_wiggliness = if spec.double_penalty {
5758 let mut sum = Array2::<f64>::zeros((total_cols, total_cols));
5759 for dim in 0..normalized_marginal_penalties.len() {
5760 let mut embedded = Array2::<f64>::eye(1);
5761 for (margin, &width) in marginalnum_basis.iter().enumerate() {
5762 let factor = if margin == dim {
5763 normalized_marginal_penalties[margin].0.clone()
5764 } else {
5765 Array2::<f64>::eye(width)
5766 };
5767 embedded = kronecker_product(&embedded, &factor);
5768 }
5769 sum += &embedded;
5770 }
5771 Some(sum)
5772 } else {
5773 None
5774 };
5775 let mut kronecker_marginal_penalties =
5776 Vec::<Array2<f64>>::with_capacity(normalized_marginal_penalties.len());
5777
5778 match spec.penalty_decomposition {
5779 TensorBSplinePenaltyDecomposition::MarginalKroneckerSum => {
5780 for dim in 0..normalized_marginal_penalties.len() {
5786 let mut s_dim = Array2::<f64>::eye(1);
5787 let mut factors = Vec::<Array2<f64>>::with_capacity(marginalnum_basis.len());
5788 for (j, &qj) in marginalnum_basis.iter().enumerate() {
5789 let factor = if j == dim {
5790 normalized_marginal_penalties[j].0.clone()
5791 } else {
5792 Array2::<f64>::eye(qj)
5793 };
5794 factors.push(factor.clone());
5795 s_dim = kronecker_product(&s_dim, &factor);
5796 }
5797 if dim == kronecker_marginal_penalties.len() {
5798 kronecker_marginal_penalties.push(normalized_marginal_penalties[dim].0.clone());
5799 }
5800 candidates.push(PenaltyCandidate {
5801 matrix: s_dim,
5802 nullspace_dim_hint: 0,
5803 source: PenaltySource::TensorMarginal { dim },
5804 normalization_scale: normalized_marginal_penalties[dim].1,
5805 kronecker_factors: Some(factors),
5806 op: None,
5807 });
5808 }
5809
5810 if let (Some(primary), Some(gram)) =
5811 (joint_wiggliness.as_ref(), tensor_function_gram.as_ref())
5812 && let Some(shrink) =
5813 crate::basis::function_space_nullspace_shrinkage(primary, gram)?
5814 {
5815 let (matrix, normalization_scale) = normalize_penalty_in_constrained_space(&shrink);
5816 candidates.push(PenaltyCandidate {
5817 matrix,
5818 nullspace_dim_hint: 0,
5819 source: PenaltySource::TensorGlobalRidge,
5820 normalization_scale,
5821 kronecker_factors: None,
5822 op: None,
5823 });
5824 }
5825 }
5826 TensorBSplinePenaltyDecomposition::Separable => {
5827 let projectors = tensor_margin_range_null_projectors(&normalized_marginal_penalties)?;
5828 let n_masks = 1usize.checked_shl(projectors.len() as u32).ok_or_else(|| {
5829 BasisError::InvalidInput(format!(
5830 "t2 separable tensor penalty supports at most {} margins, got {}",
5831 usize::BITS - 1,
5832 projectors.len()
5833 ))
5834 })?;
5835 for mask in 1..n_masks {
5836 let mut matrix = Array2::<f64>::eye(1);
5837 let mut factors = Vec::<Array2<f64>>::with_capacity(projectors.len());
5838 let mut penalized_margins = Vec::<usize>::new();
5839 for (dim, projector) in projectors.iter().enumerate() {
5840 let use_range = ((mask >> dim) & 1) == 1;
5841 let factor = if use_range {
5842 penalized_margins.push(dim);
5843 projector.range.clone()
5844 } else {
5845 projector.null.clone()
5846 };
5847 matrix = kronecker_product(&matrix, &factor);
5848 factors.push(factor);
5849 }
5850 let (matrix, normalization_scale) = normalize_penalty_in_constrained_space(&matrix);
5851 candidates.push(PenaltyCandidate {
5852 matrix,
5853 nullspace_dim_hint: 0,
5854 source: PenaltySource::TensorSeparable { penalized_margins },
5855 normalization_scale,
5856 kronecker_factors: Some(factors),
5857 op: None,
5858 });
5859 }
5860
5861 if let (Some(primary), Some(gram)) =
5862 (joint_wiggliness.as_ref(), tensor_function_gram.as_ref())
5863 && let Some(matrix) =
5864 crate::basis::function_space_nullspace_shrinkage(primary, gram)?
5865 {
5866 let (matrix, normalization_scale) = normalize_penalty_in_constrained_space(&matrix);
5867 candidates.push(PenaltyCandidate {
5868 matrix,
5869 nullspace_dim_hint: 0,
5870 source: PenaltySource::TensorGlobalRidge,
5871 normalization_scale,
5872 kronecker_factors: None,
5873 op: None,
5874 });
5875 }
5876 }
5877 }
5878
5879 let z_opt = match &spec.identifiability {
5880 TensorBSplineIdentifiability::None => None,
5881 TensorBSplineIdentifiability::SumToZero => {
5882 if total_cols < 2 {
5883 crate::bail_invalid_basis!(
5884 "TensorBSpline requires at least 2 basis coefficients to enforce sum-to-zero identifiability"
5885 );
5886 }
5887 let dense_design_ref = dense_design.as_ref().ok_or_else(|| {
5888 BasisError::InvalidInput(
5889 "tensor sum-to-zero identifiability requires a realized basis".to_string(),
5890 )
5891 })?;
5892 let (_, z) = apply_sum_to_zero_constraint(dense_design_ref.view(), None)?;
5893 let gauge = gam_problem::Gauge::sum_to_zero(z);
5894 Some(gauge.block_transform(0))
5895 }
5896 TensorBSplineIdentifiability::MarginalSumToZero => {
5897 if marginal_designs.len() < 2 {
5908 crate::bail_invalid_basis!(
5909 "tensor interaction (ti) identifiability requires at least 2 margins"
5910 );
5911 }
5912 let mut z = Array2::<f64>::eye(1);
5913 for (dim, marginal) in marginal_designs.iter().enumerate() {
5914 if marginal.ncols() < 2 {
5915 crate::bail_invalid_basis!(
5916 "tensor interaction (ti) margin {dim} has fewer than 2 basis functions; \
5917 cannot remove its marginal main effect"
5918 );
5919 }
5920 let (_, z_dim) = apply_sum_to_zero_constraint(marginal.view(), None)?;
5921 let gauge_dim = gam_problem::Gauge::sum_to_zero(z_dim);
5922 let z_dim = gauge_dim.block_transform(0);
5923 z = kronecker_product(&z, &z_dim);
5924 }
5925 Some(z)
5926 }
5927 TensorBSplineIdentifiability::FrozenTransform { transform } => {
5928 if transform.nrows() != total_cols {
5929 crate::bail_dim_basis!(
5930 "frozen tensor identifiability transform mismatch: design has {} columns but transform has {} rows",
5931 total_cols,
5932 transform.nrows()
5933 );
5934 }
5935 Some(transform.clone())
5936 }
5937 };
5938
5939 if let Some(z) = z_opt.as_ref() {
5940 let gauge = gam_problem::Gauge::from_block_transforms(&[z.clone()]);
5941 let dense = dense_design.as_mut().ok_or_else(|| {
5942 BasisError::InvalidInput(
5943 "tensor identifiability transform requires a realized basis".to_string(),
5944 )
5945 })?;
5946 let restricted_design = gauge.restrict_design(dense);
5947 *dense = restricted_design;
5948 candidates = candidates
5949 .into_iter()
5950 .map(|candidate| -> Result<PenaltyCandidate, BasisError> {
5951 let matrix = gauge.restrict_penalty(&candidate.matrix);
5952 let (matrix, c_new) = normalize_penalty_in_constrained_space(&matrix);
5960 Ok(PenaltyCandidate {
5961 nullspace_dim_hint: candidate.nullspace_dim_hint,
5962 matrix,
5963 source: candidate.source,
5964 normalization_scale: candidate.normalization_scale * c_new,
5965 kronecker_factors: None,
5971 op: candidate.op.clone(),
5972 })
5973 })
5974 .collect::<Result<Vec<_>, _>>()?;
5975
5976 if candidates
5977 .iter()
5978 .any(|candidate| matches!(candidate.source, PenaltySource::TensorGlobalRidge))
5979 {
5980 let width = candidates
5981 .first()
5982 .ok_or_else(|| {
5983 BasisError::InvalidInput(
5984 "TensorBSpline global ridge has no penalty candidates".to_string(),
5985 )
5986 })?
5987 .matrix
5988 .nrows();
5989 let mut joint_primary = Array2::<f64>::zeros((width, width));
5990 for candidate in &candidates {
5991 if !matches!(candidate.source, PenaltySource::TensorGlobalRidge) {
5992 joint_primary += &candidate
5993 .matrix
5994 .mapv(|value| value * candidate.normalization_scale);
5995 }
5996 }
5997 for candidate in &mut candidates {
5998 if !matches!(candidate.source, PenaltySource::TensorGlobalRidge) {
5999 continue;
6000 }
6001 let physical_ridge = candidate
6002 .matrix
6003 .mapv(|value| value * candidate.normalization_scale);
6004 match crate::basis::rebuild_metric_consistent_ridge(
6005 &joint_primary,
6006 &physical_ridge,
6007 )? {
6008 Some(rebuilt) => {
6009 let (matrix, scale) = normalize_penalty_in_constrained_space(&rebuilt);
6010 candidate.matrix = matrix;
6011 candidate.normalization_scale = scale;
6012 }
6013 None => {
6014 candidate.matrix = Array2::<f64>::zeros((width, width));
6015 candidate.normalization_scale = 1.0;
6016 }
6017 }
6018 candidate.kronecker_factors = None;
6019 candidate.op = None;
6020 }
6021 }
6022 }
6023
6024 let (penalties, nullspace_dims, penaltyinfo, null_eigenvectors, ops) =
6025 filter_active_penalty_candidates_with_ops(candidates)?;
6026 let identifiability_is_none =
6027 matches!(spec.identifiability, TensorBSplineIdentifiability::None);
6028 let all_marginals_sparse = marginal_sparse.iter().all(Option::is_some);
6036 let design = if let Some(dense_design) = dense_design {
6037 DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(dense_design))
6038 } else if identifiability_is_none && all_marginals_sparse {
6039 let sparse_marginals: Vec<&SparseColMat<usize, f64>> = marginal_sparse
6045 .iter()
6046 .map(|m| m.as_ref().expect("all_marginals_sparse just verified"))
6047 .collect();
6048 let sparse_design = tensor_product_design_from_sparse_marginals(&sparse_marginals)?;
6049 DesignMatrix::Sparse(gam_linalg::matrix::SparseDesignMatrix::new(sparse_design))
6050 } else {
6051 let marginals: Vec<Arc<Array2<f64>>> = marginal_designs
6052 .iter()
6053 .map(|m| Arc::new(m.clone()))
6054 .collect();
6055 let op = TensorProductDesignOperator::new(marginals).map_err(|e| {
6056 BasisError::InvalidInput(format!("TensorProductDesignOperator build failed: {e}"))
6057 })?;
6058 DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(op)))
6059 };
6060
6061 Ok(BasisBuildResult {
6062 design,
6063 penalties,
6064 nullspace_dims,
6065 penaltyinfo,
6066 ops,
6067 null_eigenvectors,
6068 joint_null_rotation: None,
6069 metadata: BasisMetadata::TensorBSpline {
6070 feature_cols: feature_cols.to_vec(),
6071 knots: marginal_knots,
6072 degrees: marginal_degrees,
6073 periods: marginal_effective_periods,
6080 is_cr: marginal_is_cr_flags,
6081 identifiability_transform: z_opt,
6082 },
6083 kronecker_factored: if !spec.double_penalty
6090 && matches!(spec.identifiability, TensorBSplineIdentifiability::None)
6091 && matches!(
6092 spec.penalty_decomposition,
6093 TensorBSplinePenaltyDecomposition::MarginalKroneckerSum
6094 ) {
6095 Some(KroneckerFactoredBasis::new(
6096 marginal_designs,
6097 kronecker_marginal_penalties,
6098 marginalnum_basis.clone(),
6099 spec.double_penalty,
6100 ))
6101 } else {
6102 None
6103 },
6104 })
6105}
6106
6107#[cfg(test)]
6108mod tensor_function_space_runtime_tests {
6109 use super::*;
6110 use crate::basis::{BSplineBoundaryConditions, OneDimensionalBoundary};
6111 use ndarray::array;
6112
6113 fn marginal() -> BSplineBasisSpec {
6114 BSplineBasisSpec {
6115 degree: 2,
6116 penalty_order: 1,
6117 knotspec: BSplineKnotSpec::Generate {
6118 data_range: (0.0, 1.0),
6119 num_internal_knots: 2,
6120 },
6121 double_penalty: false,
6122 identifiability: BSplineIdentifiability::None,
6123 boundary: OneDimensionalBoundary::Open,
6124 boundary_conditions: BSplineBoundaryConditions::default(),
6125 }
6126 }
6127
6128 #[test]
6129 fn function_space_tensor_ridge_uses_exact_canonical_runtime() {
6130 let data = array![
6131 [0.00, 0.13],
6132 [0.15, 0.82],
6133 [0.29, 0.37],
6134 [0.43, 0.95],
6135 [0.58, 0.21],
6136 [0.71, 0.66],
6137 [0.86, 0.48],
6138 [1.00, 0.04]
6139 ];
6140 let mut spec = TensorBSplineSpec {
6141 marginalspecs: vec![marginal(), marginal()],
6142 periods: Vec::new(),
6143 double_penalty: true,
6144 identifiability: TensorBSplineIdentifiability::None,
6145 penalty_decomposition: TensorBSplinePenaltyDecomposition::MarginalKroneckerSum,
6146 };
6147 let built = build_tensor_bspline_basis(data.view(), &[0, 1], &spec)
6148 .expect("double-penalty tensor basis");
6149 assert!(built.penaltyinfo.iter().any(|info| {
6150 info.active && matches!(info.source, PenaltySource::TensorGlobalRidge)
6151 }));
6152 assert!(
6153 built.kronecker_factored.is_none(),
6154 "the legacy factored runtime cannot represent a function-metric global ridge"
6155 );
6156
6157 spec.double_penalty = false;
6158 let singly_penalized = build_tensor_bspline_basis(data.view(), &[0, 1], &spec)
6159 .expect("single-penalty tensor basis");
6160 assert!(
6161 singly_penalized.kronecker_factored.is_some(),
6162 "the exact marginal-only fast path must remain available"
6163 );
6164 }
6165}
6166
6167pub fn tensor_product_design_from_marginals(
6168 marginal_designs: &[Array2<f64>],
6169) -> Result<Array2<f64>, BasisError> {
6170 if marginal_designs.is_empty() {
6171 crate::bail_invalid_basis!("TensorBSpline requires at least one marginal basis");
6172 }
6173 let n = marginal_designs[0].nrows();
6174 for (i, b) in marginal_designs.iter().enumerate().skip(1) {
6175 if b.nrows() != n {
6176 crate::bail_dim_basis!(
6177 "tensor marginal row mismatch at dim {i}: expected {n}, got {}",
6178 b.nrows()
6179 );
6180 }
6181 }
6182 let total_cols = marginal_designs.iter().try_fold(1usize, |acc, b| {
6183 acc.checked_mul(b.ncols())
6184 .ok_or_else(|| BasisError::DimensionMismatch("tensor basis too large".to_string()))
6185 })?;
6186 use ndarray::parallel::prelude::*;
6192 use rayon::iter::{IntoParallelIterator, ParallelIterator};
6193 let mut design = Array2::<f64>::zeros((n, total_cols));
6194 design
6195 .axis_chunks_iter_mut(ndarray::Axis(0), 1024)
6196 .into_par_iter()
6197 .enumerate()
6198 .for_each(|(chunk_idx, mut block)| {
6199 let row_offset = chunk_idx * 1024;
6200 let mut cur = Vec::<f64>::with_capacity(total_cols);
6202 let mut next = Vec::<f64>::with_capacity(total_cols);
6203 for (local_i, mut out_row) in block.outer_iter_mut().enumerate() {
6204 let i = row_offset + local_i;
6205 cur.clear();
6206 cur.push(1.0);
6207 for b in marginal_designs {
6208 let q = b.ncols();
6209 next.clear();
6210 next.resize(cur.len() * q, 0.0);
6211 let b_row = b.row(i);
6215 let b_slice = b_row
6216 .as_slice()
6217 .expect("Array2 row from outer_iter is contiguous");
6218 for (a_idx, &aval) in cur.iter().enumerate() {
6219 let off = a_idx * q;
6220 let dst = &mut next[off..off + q];
6221 for col in 0..q {
6222 dst[col] = aval * b_slice[col];
6223 }
6224 }
6225 std::mem::swap(&mut cur, &mut next);
6226 }
6227 let out_slice = out_row
6232 .as_slice_mut()
6233 .expect("design row is contiguous in C-major Array2");
6234 out_slice.copy_from_slice(&cur);
6235 }
6236 });
6237 Ok(design)
6238}
6239
6240fn fmt_level_value(v: f64) -> String {
6244 if v.is_finite() && v.fract() == 0.0 && v.abs() < 1e15 {
6245 format!("{}", v as i64)
6246 } else {
6247 format!("{v}")
6248 }
6249}
6250
6251pub fn build_random_effect_block(
6252 data: ArrayView2<'_, f64>,
6253 spec: &RandomEffectTermSpec,
6254) -> Result<RandomEffectBlock, BasisError> {
6255 let n = data.nrows();
6256 let p = data.ncols();
6257 if spec.feature_col >= p {
6258 crate::bail_dim_basis!(
6259 "random-effect term '{}' feature column {} out of bounds for {} columns",
6260 spec.name,
6261 spec.feature_col,
6262 p
6263 );
6264 }
6265
6266 let col = data.column(spec.feature_col);
6267 if col.iter().any(|v| !v.is_finite()) {
6268 crate::bail_invalid_basis!(
6269 "random-effect term '{}' contains non-finite group values",
6270 spec.name
6271 );
6272 }
6273
6274 let kept_levels: Vec<u64> = if let Some(levels) = spec.frozen_levels.as_ref() {
6275 if levels.is_empty() {
6276 crate::bail_invalid_basis!(
6277 "random-effect term '{}' has empty frozen_levels",
6278 spec.name
6279 );
6280 }
6281 levels
6285 .iter()
6286 .map(|&b| gam_data::canonical_level_bits(f64::from_bits(b)))
6287 .collect()
6288 } else {
6289 let mut seen = BTreeSet::<u64>::new();
6290 let mut levels = Vec::<u64>::new();
6291 for &v in col {
6292 let bits = gam_data::canonical_level_bits(v);
6293 if seen.insert(bits) {
6294 levels.push(bits);
6295 }
6296 }
6297 if levels.is_empty() {
6298 crate::bail_invalid_basis!("random-effect term '{}' has no observed levels", spec.name);
6299 }
6300 let start_idx = if spec.drop_first_level && levels.len() > 1 {
6301 1usize
6302 } else {
6303 0usize
6304 };
6305 levels[start_idx..].to_vec()
6306 };
6307
6308 if kept_levels.is_empty() {
6309 crate::bail_invalid_basis!(
6310 "random-effect term '{}' drops all levels; keep at least one level",
6311 spec.name
6312 );
6313 }
6314
6315 let q = kept_levels.len();
6316 let mut level_to_col = BTreeMap::<u64, usize>::new();
6317 for (idx, &bits) in kept_levels.iter().enumerate() {
6318 if level_to_col.insert(bits, idx).is_some() {
6319 crate::bail_invalid_basis!(
6320 "random-effect term '{}' has duplicate frozen level bits {bits}",
6321 spec.name
6322 );
6323 }
6324 }
6325 let strict_unseen =
6339 !spec.lenient_unseen && !spec.drop_first_level && spec.frozen_levels.is_some();
6340 let mut group_ids = Vec::with_capacity(n);
6341 for (row, &v) in col.iter().enumerate() {
6342 let bits = gam_data::canonical_level_bits(v);
6343 let group_id = level_to_col.get(&bits).copied();
6344 if strict_unseen && group_id.is_none() {
6345 crate::bail_invalid_basis!(
6346 "unseen level '{}' in fixed factor column '{}' at row {}; the factor's levels \
6347 were fixed at fit time and an out-of-vocabulary level cannot be predicted \
6348 (use group({}) for a random effect that tolerates held-out levels)",
6349 fmt_level_value(v),
6350 spec.name,
6351 row,
6352 spec.name
6353 );
6354 }
6355 group_ids.push(group_id);
6356 }
6357
6358 Ok(RandomEffectBlock {
6359 name: spec.name.clone(),
6360 group_ids,
6361 num_groups: q,
6362 kept_levels,
6363 })
6364}
6365
6366#[cfg(test)]
6367mod random_effect_signed_zero_tests {
6368 use super::{RandomEffectTermSpec, build_random_effect_block};
6369 use ndarray::array;
6370
6371 fn spec() -> RandomEffectTermSpec {
6372 RandomEffectTermSpec {
6373 name: "g".to_string(),
6374 feature_col: 0,
6375 drop_first_level: false,
6376 penalized: true,
6377 frozen_levels: None,
6378 lenient_unseen: true,
6379 }
6380 }
6381
6382 #[test]
6383 fn signed_zero_rows_share_one_group() {
6384 let data = array![[-0.0_f64], [0.0], [1.0], [-0.0], [1.0]];
6388 let block = build_random_effect_block(data.view(), &spec()).unwrap();
6389 assert_eq!(
6390 block.num_groups, 2,
6391 "0.0/-0.0 must not split into two groups"
6392 );
6393 assert_eq!(block.group_ids[0], block.group_ids[1]);
6395 assert_eq!(block.group_ids[0], block.group_ids[3]);
6396 assert_eq!(block.group_ids[2], block.group_ids[4]);
6397 assert_ne!(block.group_ids[0], block.group_ids[2]);
6398 }
6399
6400 #[test]
6401 fn frozen_positive_zero_matches_negative_zero_row() {
6402 let mut s = spec();
6405 s.frozen_levels = Some(vec![0.0_f64.to_bits(), 1.0_f64.to_bits()]);
6406 let data = array![[-0.0_f64], [1.0]];
6407 let block = build_random_effect_block(data.view(), &s).unwrap();
6408 assert_eq!(
6409 block.group_ids[0],
6410 Some(0),
6411 "-0.0 must match the +0.0 column"
6412 );
6413 assert_eq!(block.group_ids[1], Some(1));
6414 }
6415
6416 #[test]
6417 fn frozen_negative_zero_matches_positive_zero_row() {
6418 let mut s = spec();
6421 s.frozen_levels = Some(vec![(-0.0_f64).to_bits(), 1.0_f64.to_bits()]);
6422 let data = array![[0.0_f64], [1.0]];
6423 let block = build_random_effect_block(data.view(), &s).unwrap();
6424 assert_eq!(
6425 block.group_ids[0],
6426 Some(0),
6427 "+0.0 must match the -0.0 column"
6428 );
6429 }
6430
6431 fn fixed_factor_spec() -> RandomEffectTermSpec {
6434 let mut s = spec();
6437 s.name = "year".to_string();
6438 s.lenient_unseen = false;
6439 s
6440 }
6441
6442 #[test]
6443 fn fixed_factor_rejects_unseen_numeric_level_at_predict() {
6444 let mut s = fixed_factor_spec();
6449 s.frozen_levels = Some(vec![2000.0_f64.to_bits(), 2001.0_f64.to_bits()]);
6450 let data = array![[2000.0_f64], [1999.0]];
6451 let err = build_random_effect_block(data.view(), &s)
6452 .expect_err("an unseen fixed-factor level must be rejected");
6453 let msg = format!("{err}");
6454 assert!(
6455 msg.contains("unseen level"),
6456 "message must name the defect: {msg}"
6457 );
6458 assert!(
6459 msg.contains("1999"),
6460 "message must name the integer level (not 1999.0): {msg}"
6461 );
6462 assert!(msg.contains("year"), "message must name the column: {msg}");
6463 }
6464
6465 #[test]
6466 fn fixed_factor_accepts_seen_numeric_levels_at_predict() {
6467 let mut s = fixed_factor_spec();
6470 s.frozen_levels = Some(vec![2000.0_f64.to_bits(), 2001.0_f64.to_bits()]);
6471 let data = array![[2001.0_f64], [2000.0]];
6472 let block = build_random_effect_block(data.view(), &s).unwrap();
6473 assert_eq!(block.group_ids[0], Some(1));
6474 assert_eq!(block.group_ids[1], Some(0));
6475 }
6476
6477 #[test]
6478 fn fixed_factor_at_fit_time_derives_vocabulary_and_never_false_rejects() {
6479 let mut s = fixed_factor_spec();
6483 s.frozen_levels = None;
6484 let data = array![[2000.0_f64], [2001.0], [2002.0], [2000.0]];
6485 let block = build_random_effect_block(data.view(), &s)
6486 .expect("fit-time build must not reject its own levels");
6487 assert_eq!(block.num_groups, 3);
6488 }
6489
6490 #[test]
6491 fn random_effect_still_tolerates_unseen_numeric_level() {
6492 let mut s = spec(); s.frozen_levels = Some(vec![2000.0_f64.to_bits(), 2001.0_f64.to_bits()]);
6497 let data = array![[2000.0_f64], [1999.0]];
6498 let block = build_random_effect_block(data.view(), &s)
6499 .expect("a random effect tolerates unseen levels");
6500 assert_eq!(block.group_ids[0], Some(0));
6501 assert_eq!(
6502 block.group_ids[1], None,
6503 "unseen level → population mean, not a reject"
6504 );
6505 }
6506}
6507
6508impl SmoothDesign {
6509 pub fn map_term_coefficients(
6512 unconstrained: &Array1<f64>,
6513 shape: ShapeConstraint,
6514 ) -> Result<Array1<f64>, BasisError> {
6515 if unconstrained.is_empty() {
6516 crate::bail_invalid_basis!("unconstrained coefficient vector cannot be empty");
6517 }
6518 let mapped = match shape {
6519 ShapeConstraint::None => unconstrained.clone(),
6520 ShapeConstraint::MonotoneIncreasing => cumulative_exp(unconstrained, 1.0),
6521 ShapeConstraint::MonotoneDecreasing => cumulative_exp(unconstrained, -1.0),
6522 ShapeConstraint::Convex => second_cumulative_exp(unconstrained, 1.0),
6523 ShapeConstraint::Concave => second_cumulative_exp(unconstrained, -1.0),
6524 };
6525 Ok(mapped)
6526 }
6527}
6528
6529pub struct LocalSmoothTermBuild {
6530 pub dim: usize,
6531 pub design: DesignMatrix,
6532 pub penalties: Vec<Array2<f64>>,
6533 pub ops: Vec<Option<std::sync::Arc<dyn crate::analytic_penalties::PenaltyOp>>>,
6534 pub nullspaces: Vec<usize>,
6535 pub null_eigenvectors: Vec<Option<Array2<f64>>>,
6543 pub joint_null_rotation: Option<crate::basis::JointNullRotation>,
6550 pub penaltyinfo: Vec<PenaltyInfo>,
6551 pub pre_dropped_penaltyinfo: Vec<PenaltyInfo>,
6552 pub metadata: BasisMetadata,
6553 pub linear_constraints: Option<LinearInequalityConstraints>,
6554 pub box_reparam: bool,
6555 pub kronecker_factored: Option<KroneckerFactoredBasis>,
6556}
6557
6558#[derive(Clone)]
6559pub struct PcaScoresMemmapDesignOperator {
6560 mmap: Arc<memmap2::Mmap>,
6561 data_offset: usize,
6562 nrows: usize,
6563 ncols: usize,
6564 chunk_size: usize,
6565}
6566
6567impl PcaScoresMemmapDesignOperator {
6568 fn open(path: PathBuf, chunk_size: usize) -> Result<Self, BasisError> {
6569 let file = File::open(&path).map_err(|err| {
6570 BasisError::InvalidInput(format!(
6571 "failed to open lazy Pca .npy scores '{}': {err}",
6572 path.display()
6573 ))
6574 })?;
6575 let mmap = unsafe {
6581 memmap2::Mmap::map(&file).map_err(|err| {
6582 BasisError::InvalidInput(format!(
6583 "failed to memmap lazy Pca .npy scores '{}': {err}",
6584 path.display()
6585 ))
6586 })?
6587 };
6588 let (data_offset, nrows, ncols) = parse_f64_2d_npy_header(&mmap, &path)?;
6589 let expected = data_offset
6590 .checked_add(nrows.saturating_mul(ncols).saturating_mul(8))
6591 .ok_or_else(|| {
6592 BasisError::InvalidInput(format!(
6593 "lazy Pca .npy scores '{}' shape is too large",
6594 path.display()
6595 ))
6596 })?;
6597 if mmap.len() < expected {
6598 crate::bail_invalid_basis!(
6599 "lazy Pca .npy scores '{}' is truncated: header expects {} bytes, file has {}",
6600 path.display(),
6601 expected,
6602 mmap.len()
6603 );
6604 }
6605 Ok(Self {
6606 mmap: Arc::new(mmap),
6607 data_offset,
6608 nrows,
6609 ncols,
6610 chunk_size: chunk_size.max(1),
6611 })
6612 }
6613
6614 fn value(&self, row: usize, col: usize) -> f64 {
6615 let offset = self.data_offset + (row * self.ncols + col) * 8;
6616 let mut bytes = [0_u8; 8];
6617 bytes.copy_from_slice(&self.mmap[offset..offset + 8]);
6618 f64::from_le_bytes(bytes)
6619 }
6620
6621 fn chunk_rows(&self) -> usize {
6622 self.chunk_size.min(self.nrows.max(1))
6623 }
6624}
6625
6626impl LinearOperator for PcaScoresMemmapDesignOperator {
6627 fn nrows(&self) -> usize {
6628 self.nrows
6629 }
6630
6631 fn ncols(&self) -> usize {
6632 self.ncols
6633 }
6634
6635 fn apply(&self, vector: &Array1<f64>) -> Array1<f64> {
6636 assert_eq!(
6637 vector.len(),
6638 self.ncols,
6639 "lazy Pca apply vector length mismatch"
6640 );
6641 let mut out = Array1::<f64>::zeros(self.nrows);
6642 for start in (0..self.nrows).step_by(self.chunk_rows()) {
6643 let end = (start + self.chunk_rows()).min(self.nrows);
6644 for row in start..end {
6645 let mut acc = 0.0;
6646 for col in 0..self.ncols {
6647 acc += self.value(row, col) * vector[col];
6648 }
6649 out[row] = acc;
6650 }
6651 }
6652 out
6653 }
6654
6655 fn apply_transpose(&self, vector: &Array1<f64>) -> Array1<f64> {
6656 assert_eq!(
6657 vector.len(),
6658 self.nrows,
6659 "lazy Pca apply_transpose vector length mismatch"
6660 );
6661 let mut out = Array1::<f64>::zeros(self.ncols);
6662 for start in (0..self.nrows).step_by(self.chunk_rows()) {
6663 let end = (start + self.chunk_rows()).min(self.nrows);
6664 for row in start..end {
6665 let scale = vector[row];
6666 if scale == 0.0 {
6667 continue;
6668 }
6669 for col in 0..self.ncols {
6670 out[col] += scale * self.value(row, col);
6671 }
6672 }
6673 }
6674 out
6675 }
6676
6677 fn diag_xtw_x(&self, weights: &Array1<f64>) -> Result<Array2<f64>, String> {
6678 if weights.len() != self.nrows {
6679 return Err(format!(
6680 "lazy Pca diag_xtw_x weight length mismatch: weights={}, nrows={}",
6681 weights.len(),
6682 self.nrows
6683 ));
6684 }
6685 let mut gram = Array2::<f64>::zeros((self.ncols, self.ncols));
6686 for start in (0..self.nrows).step_by(self.chunk_rows()) {
6687 let end = (start + self.chunk_rows()).min(self.nrows);
6688 for row in start..end {
6689 let w = weights[row];
6690 if w == 0.0 {
6691 continue;
6692 }
6693 for a in 0..self.ncols {
6694 let xa = self.value(row, a);
6695 if xa == 0.0 {
6696 continue;
6697 }
6698 for b in a..self.ncols {
6699 gram[[a, b]] += w * xa * self.value(row, b);
6700 }
6701 }
6702 }
6703 }
6704 for a in 0..self.ncols {
6705 for b in 0..a {
6706 gram[[a, b]] = gram[[b, a]];
6707 }
6708 }
6709 Ok(gram)
6710 }
6711
6712 fn apply_weighted_normal(
6713 &self,
6714 weights: &Array1<f64>,
6715 vector: &Array1<f64>,
6716 penalty: Option<&Array2<f64>>,
6717 ridge: f64,
6718 ) -> Array1<f64> {
6719 assert_eq!(
6720 weights.len(),
6721 self.nrows,
6722 "lazy Pca weighted-normal weight mismatch"
6723 );
6724 assert_eq!(
6725 vector.len(),
6726 self.ncols,
6727 "lazy Pca weighted-normal vector mismatch"
6728 );
6729 let mut out = Array1::<f64>::zeros(self.ncols);
6730 for start in (0..self.nrows).step_by(self.chunk_rows()) {
6731 let end = (start + self.chunk_rows()).min(self.nrows);
6732 for row in start..end {
6733 let w = weights[row].max(0.0);
6734 if w == 0.0 {
6735 continue;
6736 }
6737 let mut row_dot = 0.0;
6738 for col in 0..self.ncols {
6739 row_dot += self.value(row, col) * vector[col];
6740 }
6741 if row_dot == 0.0 {
6742 continue;
6743 }
6744 let scaled = w * row_dot;
6745 for col in 0..self.ncols {
6746 out[col] += scaled * self.value(row, col);
6747 }
6748 }
6749 }
6750 if let Some(pen) = penalty {
6751 out += &pen.dot(vector);
6752 }
6753 if ridge > 0.0 {
6754 out += &vector.mapv(|x| ridge * x);
6755 }
6756 out
6757 }
6758}
6759
6760impl DenseDesignOperator for PcaScoresMemmapDesignOperator {
6761 fn compute_xtwy(&self, weights: &Array1<f64>, y: &Array1<f64>) -> Result<Array1<f64>, String> {
6762 if weights.len() != self.nrows || y.len() != self.nrows {
6763 return Err(format!(
6764 "lazy Pca compute_xtwy dimension mismatch: weights={}, y={}, nrows={}",
6765 weights.len(),
6766 y.len(),
6767 self.nrows
6768 ));
6769 }
6770 let mut out = Array1::<f64>::zeros(self.ncols);
6771 for start in (0..self.nrows).step_by(self.chunk_rows()) {
6772 let end = (start + self.chunk_rows()).min(self.nrows);
6773 for row in start..end {
6774 let scale = weights[row] * y[row];
6775 if scale == 0.0 {
6776 continue;
6777 }
6778 for col in 0..self.ncols {
6779 out[col] += scale * self.value(row, col);
6780 }
6781 }
6782 }
6783 Ok(out)
6784 }
6785
6786 fn row_chunk_into(
6787 &self,
6788 rows: Range<usize>,
6789 mut out: ArrayViewMut2<'_, f64>,
6790 ) -> Result<(), MatrixMaterializationError> {
6791 if rows.end > self.nrows || rows.start > rows.end {
6792 return Err(MatrixMaterializationError::MissingRowChunk {
6793 context: "lazy Pca row range out of bounds",
6794 });
6795 }
6796 if out.nrows() != rows.end - rows.start || out.ncols() != self.ncols {
6797 return Err(MatrixMaterializationError::MissingRowChunk {
6798 context: "lazy Pca row_chunk_into shape mismatch",
6799 });
6800 }
6801 for (local, row) in (rows.start..rows.end).enumerate() {
6802 for col in 0..self.ncols {
6803 out[[local, col]] = self.value(row, col);
6804 }
6805 }
6806 Ok(())
6807 }
6808
6809 fn to_dense(&self) -> Array2<f64> {
6810 let mut out = Array2::<f64>::zeros((self.nrows, self.ncols));
6811 self.row_chunk_into(0..self.nrows, out.view_mut())
6812 .expect("lazy Pca full materialization failed");
6813 out
6814 }
6815}
6816
6817pub fn parse_f64_2d_npy_header(
6818 bytes: &[u8],
6819 path: &PathBuf,
6820) -> Result<(usize, usize, usize), BasisError> {
6821 if bytes.len() < 10 || &bytes[0..6] != b"\x93NUMPY" {
6822 crate::bail_invalid_basis!("lazy Pca scores '{}' is not a .npy file", path.display());
6823 }
6824 let major = bytes[6];
6825 let header_len = match major {
6826 1 => u16::from_le_bytes([bytes[8], bytes[9]]) as usize,
6827 2 | 3 => {
6828 if bytes.len() < 12 {
6829 crate::bail_invalid_basis!(
6830 "lazy Pca scores '{}' has a truncated .npy header",
6831 path.display()
6832 );
6833 }
6834 u32::from_le_bytes([bytes[8], bytes[9], bytes[10], bytes[11]]) as usize
6835 }
6836 other => {
6837 crate::bail_invalid_basis!(
6838 "lazy Pca scores '{}' uses unsupported .npy version {}",
6839 path.display(),
6840 other
6841 );
6842 }
6843 };
6844 let header_start = if major == 1 { 10 } else { 12 };
6845 let data_offset = header_start + header_len;
6846 if bytes.len() < data_offset {
6847 crate::bail_invalid_basis!(
6848 "lazy Pca scores '{}' has a truncated .npy header",
6849 path.display()
6850 );
6851 }
6852 let header = std::str::from_utf8(&bytes[header_start..data_offset]).map_err(|err| {
6853 BasisError::InvalidInput(format!(
6854 "lazy Pca scores '{}' has a non-UTF8 .npy header: {err}",
6855 path.display()
6856 ))
6857 })?;
6858 if !(header.contains("'descr': '<f8'")
6859 || header.contains("\"descr\": \"<f8\"")
6860 || header.contains("'descr': '|f8'")
6861 || header.contains("\"descr\": \"|f8\""))
6862 {
6863 crate::bail_invalid_basis!(
6864 "lazy Pca scores '{}' must be float64 little-endian .npy",
6865 path.display()
6866 );
6867 }
6868 if header.contains("True") {
6869 crate::bail_invalid_basis!(
6870 "lazy Pca scores '{}' must be C-contiguous, not Fortran-ordered",
6871 path.display()
6872 );
6873 }
6874 let shape_pos = header.find("shape").ok_or_else(|| {
6875 BasisError::InvalidInput(format!(
6876 "lazy Pca scores '{}' .npy header is missing shape",
6877 path.display()
6878 ))
6879 })?;
6880 let open = header[shape_pos..].find('(').ok_or_else(|| {
6881 BasisError::InvalidInput(format!(
6882 "lazy Pca scores '{}' .npy header has malformed shape",
6883 path.display()
6884 ))
6885 })? + shape_pos;
6886 let close = header[open..].find(')').ok_or_else(|| {
6887 BasisError::InvalidInput(format!(
6888 "lazy Pca scores '{}' .npy header has malformed shape",
6889 path.display()
6890 ))
6891 })? + open;
6892 let dims = header[open + 1..close]
6893 .split(',')
6894 .map(str::trim)
6895 .filter(|part| !part.is_empty())
6896 .map(|part| part.parse::<usize>())
6897 .collect::<Result<Vec<_>, _>>()
6898 .map_err(|err| {
6899 BasisError::InvalidInput(format!(
6900 "lazy Pca scores '{}' .npy shape is not integral: {err}",
6901 path.display()
6902 ))
6903 })?;
6904 if dims.len() != 2 {
6905 crate::bail_invalid_basis!(
6906 "lazy Pca scores '{}' must have shape (N, K), got {:?}",
6907 path.display(),
6908 dims
6909 );
6910 }
6911 Ok((data_offset, dims[0], dims[1]))
6912}
6913
6914pub fn pca_center_mean(x: ArrayView2<'_, f64>) -> Result<Array1<f64>, BasisError> {
6915 if x.nrows() == 0 {
6916 crate::bail_invalid_basis!("Pca basis requires at least one row to compute center mean");
6917 }
6918 let mut mean = Array1::<f64>::zeros(x.ncols());
6919 for row in x.rows() {
6920 mean += &row;
6921 }
6922 mean.mapv_inplace(|v| v / x.nrows() as f64);
6923 Ok(mean)
6924}
6925
6926fn pca_function_mass_penalty(
6937 mut raw_score_gram: Array2<f64>,
6938 n_rows: usize,
6939 smooth_penalty: f64,
6940) -> Result<Array2<f64>, BasisError> {
6941 let k = raw_score_gram.ncols();
6942 if raw_score_gram.nrows() != k {
6943 crate::bail_dim_basis!(
6944 "Pca score Gram must be square, got {}x{}",
6945 raw_score_gram.nrows(),
6946 k
6947 );
6948 }
6949 if n_rows == 0 {
6950 crate::bail_invalid_basis!("Pca basis requires at least one score row");
6951 }
6952 if k == 0 {
6953 crate::bail_invalid_basis!("Pca basis requires at least one score column");
6954 }
6955 if k > n_rows {
6956 crate::bail_invalid_basis!(
6957 "Pca score design is rank deficient: {} score columns cannot have full column rank with only {} rows; remove redundant components",
6958 k,
6959 n_rows
6960 );
6961 }
6962 if raw_score_gram.iter().any(|value| !value.is_finite()) {
6963 crate::bail_invalid_basis!("Pca score design produced a non-finite function Gram");
6964 }
6965
6966 let rrqr = gam_linalg::faer_ndarray::rrqr_from_gram_with_permutation(
6970 &raw_score_gram,
6971 n_rows,
6972 gam_linalg::faer_ndarray::default_rrqr_rank_alpha(),
6973 )
6974 .map_err(BasisError::LinalgError)?;
6975 if rrqr.rank != k {
6976 let redundant_columns = &rrqr.column_permutation[rrqr.rank..];
6977 crate::bail_invalid_basis!(
6978 "Pca score design is rank deficient under canonical RRQR: rank {} < {} (tolerance {:.6e}); redundant score columns {:?}; remove zero or dependent components instead of stabilizing them with a coefficient ridge",
6979 rrqr.rank,
6980 k,
6981 rrqr.rank_tol,
6982 redundant_columns
6983 );
6984 }
6985
6986 raw_score_gram.mapv_inplace(|value| value * smooth_penalty / n_rows as f64);
6987 Ok(raw_score_gram)
6988}
6989
6990pub fn build_pca_smooth_basis(
6991 data: ArrayView2<'_, f64>,
6992 feature_cols: &[usize],
6993 basis_matrix: &Array2<f64>,
6994 centered: bool,
6995 smooth_penalty: f64,
6996 center_mean: Option<&Array1<f64>>,
6997 pca_basis_path: Option<&PathBuf>,
6998 chunk_size: usize,
6999) -> Result<BasisBuildResult, BasisError> {
7000 if !smooth_penalty.is_finite() || smooth_penalty < 0.0 {
7001 crate::bail_invalid_basis!(
7002 "Pca smooth_penalty must be finite and non-negative, got {}",
7003 smooth_penalty
7004 );
7005 }
7006 if data.nrows() == 0 {
7007 crate::bail_invalid_basis!("Pca basis requires at least one data row");
7008 }
7009
7010 if let Some(path) = pca_basis_path {
7011 let op = PcaScoresMemmapDesignOperator::open(path.clone(), chunk_size)?;
7012 if op.nrows != data.nrows() {
7013 crate::bail_dim_basis!(
7014 "lazy Pca scores row mismatch: .npy has {}, data has {}",
7015 op.nrows,
7016 data.nrows()
7017 );
7018 }
7019 let raw_score_gram = op
7022 .diag_xtw_x(&Array1::<f64>::ones(op.nrows))
7023 .map_err(|err| {
7024 BasisError::InvalidInput(format!(
7025 "lazy Pca function-mass Gram construction failed: {err}"
7026 ))
7027 })?;
7028 let penalty = pca_function_mass_penalty(raw_score_gram, op.nrows, smooth_penalty)?;
7029 let (penalties, nullspace_dims, penaltyinfo, null_eigenvectors, ops) =
7030 filter_active_penalty_candidates_with_ops(vec![PenaltyCandidate {
7031 matrix: penalty,
7032 nullspace_dim_hint: 0,
7033 source: PenaltySource::OperatorMass,
7034 normalization_scale: 1.0,
7035 kronecker_factors: None,
7036 op: None,
7037 }])?;
7038 return Ok(BasisBuildResult {
7039 design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(op))),
7040 penalties,
7041 nullspace_dims,
7042 penaltyinfo,
7043 ops,
7044 null_eigenvectors,
7045 joint_null_rotation: None,
7046 metadata: BasisMetadata::Pca {
7047 feature_cols: feature_cols.to_vec(),
7048 basis_matrix: basis_matrix.clone(),
7049 centered,
7050 smooth_penalty,
7051 center_mean: center_mean.cloned(),
7052 pca_basis_path: Some(path.clone()),
7053 chunk_size: chunk_size.max(1),
7054 },
7055 kronecker_factored: None,
7056 });
7057 }
7058 if basis_matrix.nrows() != feature_cols.len() {
7059 crate::bail_dim_basis!(
7060 "Pca basis row mismatch: basis rows={}, feature columns={}",
7061 basis_matrix.nrows(),
7062 feature_cols.len()
7063 );
7064 }
7065 let mut x = select_columns(data, feature_cols)?;
7066 let mean = if centered {
7067 match center_mean {
7068 Some(mean) => mean.clone(),
7069 None => pca_center_mean(x.view())?,
7070 }
7071 } else {
7072 Array1::<f64>::zeros(feature_cols.len())
7073 };
7074 if centered {
7075 for mut row in x.rows_mut() {
7076 row -= &mean;
7077 }
7078 }
7079 let design = fast_ab(&x, basis_matrix);
7080 let raw_score_gram = gam_linalg::faer_ndarray::fast_ata(&design);
7081 let penalty = pca_function_mass_penalty(raw_score_gram, design.nrows(), smooth_penalty)?;
7082 let (penalties, nullspace_dims, penaltyinfo, null_eigenvectors, ops) =
7083 filter_active_penalty_candidates_with_ops(vec![PenaltyCandidate {
7084 matrix: penalty,
7085 nullspace_dim_hint: 0,
7086 source: PenaltySource::OperatorMass,
7087 normalization_scale: 1.0,
7088 kronecker_factors: None,
7089 op: None,
7090 }])?;
7091 Ok(BasisBuildResult {
7092 design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(design)),
7093 penalties,
7094 nullspace_dims,
7095 penaltyinfo,
7096 ops,
7097 null_eigenvectors,
7098 joint_null_rotation: None,
7099 metadata: BasisMetadata::Pca {
7100 feature_cols: feature_cols.to_vec(),
7101 basis_matrix: basis_matrix.clone(),
7102 centered,
7103 smooth_penalty,
7104 center_mean: centered.then_some(mean),
7105 pca_basis_path: None,
7106 chunk_size: chunk_size.max(1),
7107 },
7108 kronecker_factored: None,
7109 })
7110}
7111
7112#[cfg(test)]
7113mod pca_function_mass_tests {
7114 use super::{PenaltySource, build_pca_smooth_basis};
7115 use ndarray::{Array1, Array2, array};
7116 use std::io::Write;
7117 use std::path::PathBuf;
7118
7119 fn quadratic_form(matrix: &Array2<f64>, coefficients: &Array1<f64>) -> f64 {
7120 coefficients.dot(&matrix.dot(coefficients))
7121 }
7122
7123 fn assert_close(left: f64, right: f64) {
7124 let scale = left.abs().max(right.abs()).max(1.0);
7125 assert!(
7126 (left - right).abs() <= 1e-11 * scale,
7127 "values differ: left={left:.16e}, right={right:.16e}"
7128 );
7129 }
7130
7131 fn write_f64_npy(scores: &Array2<f64>) -> PathBuf {
7132 let path = std::env::temp_dir().join(format!(
7133 "gam_terms_pca_function_mass_{}.npy",
7134 std::process::id()
7135 ));
7136 let mut header = format!(
7137 "{{'descr': '<f8', 'fortran_order': False, 'shape': ({}, {}), }}",
7138 scores.nrows(),
7139 scores.ncols()
7140 );
7141 while (10 + header.len() + 1) % 16 != 0 {
7142 header.push(' ');
7143 }
7144 header.push('\n');
7145 let header_len = u16::try_from(header.len()).expect("test .npy header fits u16");
7146
7147 let mut file = std::fs::File::create(&path).expect("create test .npy");
7148 file.write_all(b"\x93NUMPY").expect("write .npy magic");
7149 file.write_all(&[1, 0]).expect("write .npy version");
7150 file.write_all(&header_len.to_le_bytes())
7151 .expect("write .npy header length");
7152 file.write_all(header.as_bytes())
7153 .expect("write .npy header");
7154 for &value in scores {
7155 file.write_all(&value.to_le_bytes())
7156 .expect("write .npy score");
7157 }
7158 path
7159 }
7160
7161 #[test]
7162 fn pca_penalty_quadratic_equals_empirical_fitted_function_norm() {
7163 let data = array![[1.0, 2.0], [-1.0, 0.5], [2.0, -0.5], [0.25, -1.5]];
7164 let basis = array![[1.0, 0.5], [-0.25, 2.0]];
7165 let smooth_penalty = 2.5;
7166 let built = build_pca_smooth_basis(
7167 data.view(),
7168 &[0, 1],
7169 &basis,
7170 false,
7171 smooth_penalty,
7172 None,
7173 None,
7174 2,
7175 )
7176 .expect("full-rank PCA basis");
7177 let coefficients = array![0.7, -1.2];
7178 let design = built.design.to_dense();
7179 let fitted = design.dot(&coefficients);
7180 let expected = smooth_penalty * fitted.dot(&fitted) / fitted.len() as f64;
7181 let actual = quadratic_form(&built.penalties[0], &coefficients);
7182
7183 assert_close(actual, expected);
7184 assert_eq!(built.nullspace_dims, vec![0]);
7185 assert_eq!(built.penaltyinfo[0].source, PenaltySource::OperatorMass);
7186 }
7187
7188 #[test]
7189 fn pca_function_mass_is_invariant_to_nonorthogonal_score_reparameterization() {
7190 let scores = array![[1.0, 2.0], [-1.0, 0.5], [2.0, -0.5], [0.25, -1.5]];
7191 let identity = Array2::<f64>::eye(2);
7192 let transform = array![[2.0, 0.5], [0.0, 0.25]];
7194 let base_coefficients = array![0.8, -1.1];
7195 let transformed_coefficients = array![1.5, -4.4];
7197 let smooth_penalty = 1.7;
7198
7199 let base = build_pca_smooth_basis(
7200 scores.view(),
7201 &[0, 1],
7202 &identity,
7203 false,
7204 smooth_penalty,
7205 None,
7206 None,
7207 2,
7208 )
7209 .expect("base PCA chart");
7210 let transformed = build_pca_smooth_basis(
7211 scores.view(),
7212 &[0, 1],
7213 &transform,
7214 false,
7215 smooth_penalty,
7216 None,
7217 None,
7218 2,
7219 )
7220 .expect("reparameterized PCA chart");
7221
7222 let fitted_base = base.design.to_dense().dot(&base_coefficients);
7223 let fitted_transformed = transformed.design.to_dense().dot(&transformed_coefficients);
7224 for (&left, &right) in fitted_base.iter().zip(fitted_transformed.iter()) {
7225 assert_close(left, right);
7226 }
7227 assert_close(
7228 quadratic_form(&base.penalties[0], &base_coefficients),
7229 quadratic_form(&transformed.penalties[0], &transformed_coefficients),
7230 );
7231 }
7232
7233 #[test]
7234 fn rank_deficient_pca_score_design_is_rejected() {
7235 let scores = array![[1.0, 0.0], [2.0, 0.0], [3.0, 0.0], [4.0, 0.0]];
7236 let result = build_pca_smooth_basis(
7237 scores.view(),
7238 &[0, 1],
7239 &Array2::<f64>::eye(2),
7240 false,
7241 1.0,
7242 None,
7243 None,
7244 2,
7245 );
7246 let err = result.err().expect("zero score column must be rejected");
7247 let message = err.to_string();
7248 assert!(
7249 message.contains("rank deficient"),
7250 "unexpected error: {message}"
7251 );
7252 assert!(
7253 message.contains("rank 1 < 2"),
7254 "missing RRQR evidence: {message}"
7255 );
7256 }
7257
7258 #[test]
7259 fn lazy_and_dense_pca_function_mass_penalties_match() {
7260 let scores = array![[1.0, 2.0], [-1.0, 0.5], [2.0, -0.5], [0.25, -1.5]];
7261 let smooth_penalty = 2.25;
7262 let path = write_f64_npy(&scores);
7263 let dense = build_pca_smooth_basis(
7264 scores.view(),
7265 &[0, 1],
7266 &Array2::<f64>::eye(2),
7267 false,
7268 smooth_penalty,
7269 None,
7270 None,
7271 2,
7272 )
7273 .expect("dense PCA basis");
7274 let lazy_data = Array2::<f64>::zeros((scores.nrows(), 0));
7275 let lazy = build_pca_smooth_basis(
7276 lazy_data.view(),
7277 &[],
7278 &Array2::<f64>::zeros((0, scores.ncols())),
7279 false,
7280 smooth_penalty,
7281 None,
7282 Some(&path),
7283 2,
7284 )
7285 .expect("lazy PCA basis");
7286 std::fs::remove_file(&path).expect("remove test .npy");
7287
7288 for (&left, &right) in dense.penalties[0].iter().zip(lazy.penalties[0].iter()) {
7289 assert_close(left, right);
7290 }
7291 for (&left, &right) in dense
7292 .design
7293 .to_dense()
7294 .iter()
7295 .zip(lazy.design.to_dense().iter())
7296 {
7297 assert_close(left, right);
7298 }
7299 }
7300}
7301
7302pub fn defer_inner_model_centering_to_factor_level_wrapper(basis: &mut SmoothBasisSpec) {
7318 if let SmoothBasisSpec::BSpline1D { spec, .. } = basis
7319 && matches!(
7320 spec.identifiability,
7321 BSplineIdentifiability::WeightedSumToZero { .. }
7322 )
7323 {
7324 spec.identifiability = BSplineIdentifiability::None;
7325 }
7326}
7327
7328pub fn apply_by_variable_to_local_build(
7329 mut built: LocalSmoothTermBuild,
7330 data: ArrayView2<'_, f64>,
7331 by_col: usize,
7332 by: &ByVariableSpec,
7333 term_name: &str,
7334) -> Result<LocalSmoothTermBuild, BasisError> {
7335 if by_col >= data.ncols() {
7336 crate::bail_dim_basis!(
7337 "by-variable smooth term '{term_name}' references column {by_col}, but data has {} columns",
7338 data.ncols()
7339 );
7340 }
7341 let weights = match by {
7342 ByVariableSpec::Numeric => data.column(by_col).to_owned(),
7343 ByVariableSpec::Level { value_bits, .. } => {
7344 let value_bits = gam_data::canonical_level_bits(f64::from_bits(*value_bits));
7345 data.column(by_col).mapv(|value| {
7346 if gam_data::canonical_level_bits(value) == value_bits {
7347 1.0
7348 } else {
7349 0.0
7350 }
7351 })
7352 }
7353 };
7354 if weights.iter().any(|value| !value.is_finite()) {
7355 crate::bail_invalid_basis!(
7356 "by-variable smooth term '{term_name}' has non-finite by-column values"
7357 );
7358 }
7359
7360 let mut dense = built
7361 .design
7362 .try_to_dense_by_chunks("by-variable smooth row gating")
7363 .map_err(BasisError::InvalidInput)?;
7364 for (mut row, &weight) in dense.rows_mut().into_iter().zip(weights.iter()) {
7365 row.mapv_inplace(|value| value * weight);
7366 }
7367 built.design = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(dense));
7368 built.kronecker_factored = None;
7369 Ok(built)
7370}
7371
7372pub fn build_by_smooth_local(
7383 data: ArrayView2<'_, f64>,
7384 term: &SmoothTermSpec,
7385 smooth: &SmoothBasisSpec,
7386 by_kind: &ByVarKind,
7387 workspace: &mut crate::basis::BasisWorkspace,
7388) -> Result<LocalSmoothTermBuild, BasisError> {
7389 let inner_term = SmoothTermSpec {
7390 name: term.name.clone(),
7391 basis: (*smooth).clone(),
7392 shape: term.shape,
7393 joint_null_rotation: None,
7394 };
7395 let inner = build_single_local_smooth_term(data, &inner_term, workspace)?;
7396
7397 match by_kind {
7398 ByVarKind::Numeric { feature_col } => {
7399 let inner_meta = inner.metadata.clone();
7400 let mut built = apply_by_variable_to_local_build(
7401 inner,
7402 data,
7403 *feature_col,
7404 &ByVariableSpec::Numeric,
7405 &term.name,
7406 )?;
7407 built.metadata = BasisMetadata::BySmooth {
7408 inner: Box::new(inner_meta),
7409 by_col: *feature_col,
7410 levels: None,
7411 ordered: false,
7412 };
7413 Ok(built)
7414 }
7415 ByVarKind::Factor {
7416 feature_col,
7417 frozen_levels,
7418 ordered,
7419 } => {
7420 let level_bits: Vec<u64> = if let Some(fl) = frozen_levels {
7423 fl.iter()
7424 .map(|&b| gam_data::canonical_level_bits(f64::from_bits(b)))
7425 .collect()
7426 } else {
7427 let col = data.column(*feature_col);
7428 let mut seen = BTreeSet::<u64>::new();
7429 for &v in col.iter() {
7430 if v.is_finite() {
7431 seen.insert(gam_data::canonical_level_bits(v));
7432 }
7433 }
7434 seen.into_iter().collect()
7435 };
7436 let n_levels = level_bits.len();
7437 if n_levels == 0 {
7438 crate::bail_invalid_basis!(
7439 "by-factor smooth term '{}': factor column {} has no observed levels",
7440 term.name,
7441 feature_col
7442 );
7443 }
7444 let p = inner.dim;
7445 let q = n_levels * p;
7446 let n = data.nrows();
7447
7448 let inner_dense = inner
7449 .design
7450 .try_to_dense_by_chunks("by-factor smooth design gating")
7451 .map_err(BasisError::InvalidInput)?;
7452
7453 let mut combined = Array2::<f64>::zeros((n, q));
7455 for (lvl_idx, &bits) in level_bits.iter().enumerate() {
7456 let col_start = lvl_idx * p;
7457 for row in 0..n {
7458 if gam_data::canonical_level_bits(data[[row, *feature_col]]) == bits {
7459 combined
7460 .slice_mut(s![row, col_start..col_start + p])
7461 .assign(&inner_dense.row(row));
7462 }
7463 }
7464 }
7465
7466 let inner_meta = inner.metadata.clone();
7478 let n_penalties = inner.penalties.len();
7479 let n_blocks = n_penalties.saturating_mul(n_levels);
7480 let mut penalties = Vec::<Array2<f64>>::with_capacity(n_blocks);
7481 let mut penaltyinfo = Vec::<PenaltyInfo>::with_capacity(n_blocks);
7482 let mut nullspaces = Vec::<usize>::with_capacity(n_blocks);
7483 for (pen_pos, s_inner) in inner.penalties.iter().enumerate() {
7484 for lvl in 0..n_levels {
7485 let off = lvl * p;
7486 let mut s_big = Array2::<f64>::zeros((q, q));
7487 s_big
7488 .slice_mut(s![off..off + p, off..off + p])
7489 .assign(s_inner);
7490 let (s_big, scale) = normalize_penalty_in_constrained_space(&s_big);
7491 let mut info = inner.penaltyinfo[pen_pos].clone();
7492 info.original_index = pen_pos * n_levels + lvl;
7495 info.normalization_scale *= scale;
7496 info.kronecker_factors = None;
7499 penalties.push(s_big);
7500 penaltyinfo.push(info);
7501 nullspaces.push(inner.nullspaces[pen_pos]);
7502 }
7503 }
7504
7505 let null_eigenvectors = vec![None; penalties.len()];
7506 let ops = vec![None; penalties.len()];
7507
7508 Ok(LocalSmoothTermBuild {
7509 dim: q,
7510 design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(combined)),
7511 penalties,
7512 ops,
7513 nullspaces,
7514 null_eigenvectors,
7515 joint_null_rotation: None,
7516 penaltyinfo,
7517 pre_dropped_penaltyinfo: inner.pre_dropped_penaltyinfo,
7518 metadata: BasisMetadata::BySmooth {
7519 inner: Box::new(inner_meta),
7520 by_col: *feature_col,
7521 levels: Some(level_bits),
7522 ordered: *ordered,
7523 },
7524 linear_constraints: None,
7525 box_reparam: false,
7526 kronecker_factored: None,
7527 })
7528 }
7529 }
7530}
7531
7532pub fn ensure_by_variable_specs_match(
7533 kind: &BySmoothKind,
7534 by: &ByVariableSpec,
7535 term_name: &str,
7536) -> Result<(), BasisError> {
7537 match (kind, by) {
7538 (BySmoothKind::Numeric, ByVariableSpec::Numeric) => Ok(()),
7539 (BySmoothKind::Level { level_bits }, ByVariableSpec::Level { value_bits, .. })
7540 if level_bits == value_bits =>
7541 {
7542 Ok(())
7543 }
7544 _ => Err(BasisError::InvalidInput(format!(
7545 "by-variable smooth term '{term_name}' has inconsistent by-variable specifications"
7546 ))),
7547 }
7548}
7549
7550pub fn build_factor_smooth(
7578 data: ArrayView2<'_, f64>,
7579 spec: &FactorSmoothSpec,
7580 term_name: &str,
7581 workspace: &mut crate::basis::BasisWorkspace,
7582) -> Result<LocalSmoothTermBuild, BasisError> {
7583 if spec.continuous_cols.len() != 1 {
7584 crate::bail_invalid_basis!(
7585 "factor smooth term '{}' currently supports exactly one continuous covariate; found {}",
7586 term_name,
7587 spec.continuous_cols.len()
7588 );
7589 }
7590 let feature_col = spec.continuous_cols[0];
7591 let group_col = spec.group_col;
7592 if feature_col >= data.ncols() || group_col >= data.ncols() {
7593 crate::bail_dim_basis!(
7594 "factor smooth term '{}' references columns ({}, {}) out of bounds for {} columns",
7595 term_name,
7596 feature_col,
7597 group_col,
7598 data.ncols()
7599 );
7600 }
7601
7602 if matches!(spec.flavour, FactorSmoothFlavour::Sz) {
7605 let levels = resolve_factor_smooth_levels(data, group_col, spec, term_name)?;
7606 let inner = SmoothBasisSpec::BSpline1D {
7607 feature_col,
7608 spec: factor_smooth_marginal_for_replay(&spec.marginal),
7609 };
7610 let sz_term = SmoothTermSpec {
7611 name: term_name.to_string(),
7612 basis: SmoothBasisSpec::FactorSumToZero {
7613 inner: Box::new(inner),
7614 by_col: group_col,
7615 levels: levels.clone(),
7616 frozen_global_orthogonality: None,
7617 },
7618 shape: ShapeConstraint::None,
7619 joint_null_rotation: None,
7620 };
7621 let mut built = build_single_local_smooth_term(data, &sz_term, workspace)?;
7622 let (knots, degree, periodic, marginal_is_cr) = match &built.metadata {
7643 BasisMetadata::BSpline1D {
7644 knots,
7645 periodic,
7646 degree,
7647 ..
7648 } => (
7649 knots.clone(),
7650 degree.unwrap_or(spec.marginal.degree),
7651 *periodic,
7652 false,
7653 ),
7654 BasisMetadata::CubicRegression1D { knots, .. } => {
7655 (knots.clone(), spec.marginal.degree, None, true)
7656 }
7657 other => {
7658 crate::bail_invalid_basis!(
7659 "sz factor smooth term '{}' produced an unexpected marginal metadata variant {:?}",
7660 term_name,
7661 other
7662 );
7663 }
7664 };
7665 built.metadata = BasisMetadata::FactorSmooth {
7666 continuous_cols: spec.continuous_cols.clone(),
7667 group_col,
7668 knots,
7669 degree,
7670 periodic,
7671 group_levels: levels,
7672 flavour: "sz".to_string(),
7673 marginal_is_cr,
7674 };
7675 return Ok(built);
7676 }
7677
7678 let levels = resolve_factor_smooth_levels(data, group_col, spec, term_name)?;
7679 let n_levels = levels.len();
7680 if n_levels < 2 {
7681 crate::bail_invalid_basis!(
7682 "factor smooth term '{}' requires at least two grouping levels; found {}",
7683 term_name,
7684 n_levels
7685 );
7686 }
7687
7688 let use_per_dim_null = matches!(
7696 &spec.flavour,
7697 FactorSmoothFlavour::Fs { m_null_penalty_orders }
7698 if m_null_penalty_orders.iter().copied().max().unwrap_or(0) >= 1
7699 );
7700
7701 let mut marginal_spec = factor_smooth_marginal_for_replay(&spec.marginal);
7707 if use_per_dim_null {
7708 marginal_spec.double_penalty = false;
7709 }
7710 let inner_term = SmoothTermSpec {
7711 name: format!("{term_name}::marginal"),
7712 basis: SmoothBasisSpec::BSpline1D {
7713 feature_col,
7714 spec: marginal_spec,
7715 },
7716 shape: ShapeConstraint::None,
7717 joint_null_rotation: None,
7718 };
7719 let inner = build_single_local_smooth_term(data, &inner_term, workspace)?;
7720 let mut base = inner
7721 .design
7722 .try_to_dense_by_chunks("factor smooth marginal")
7723 .map_err(BasisError::InvalidInput)?;
7724 if matches!(spec.flavour, FactorSmoothFlavour::Re) {
7725 let center = match &inner.metadata {
7735 BasisMetadata::BSpline1D { knots, .. } if !knots.is_empty() => {
7736 0.5 * (knots[0] + knots[knots.len() - 1])
7737 }
7738 _ => 0.0,
7739 };
7740 let mut linear = Array2::<f64>::ones((data.nrows(), 2));
7741 linear
7742 .column_mut(1)
7743 .assign(&data.column(feature_col).mapv(|x| x - center));
7744 base = linear;
7745 }
7746 let n = base.nrows();
7747 let p = base.ncols();
7748 let q = p * n_levels;
7749
7750 let mut dense = Array2::<f64>::zeros((n, q));
7753 for i in 0..n {
7754 let bits = gam_data::canonical_level_bits(data[[i, group_col]]);
7755 let level_idx = levels.iter().position(|b| *b == bits).ok_or_else(|| {
7756 BasisError::InvalidInput(format!(
7757 "factor smooth term '{term_name}' saw an unseen grouping level at row {}",
7758 i + 1
7759 ))
7760 })?;
7761 let start = level_idx * p;
7762 dense
7763 .slice_mut(s![i, start..start + p])
7764 .assign(&base.row(i));
7765 }
7766
7767 let marginal_penalties: Vec<Array2<f64>> = if matches!(spec.flavour, FactorSmoothFlavour::Re) {
7773 (0..p)
7774 .map(|j| {
7775 let mut s = Array2::<f64>::zeros((p, p));
7776 s[[j, j]] = 1.0;
7777 s
7778 })
7779 .collect()
7780 } else {
7781 inner.penalties.clone()
7782 };
7783 let marginal_penaltyinfo: Vec<PenaltyInfo> = if matches!(spec.flavour, FactorSmoothFlavour::Re)
7784 {
7785 (0..p)
7786 .map(|j| PenaltyInfo {
7787 source: PenaltySource::Primary,
7788 original_index: j,
7789 active: true,
7790 effective_rank: 1,
7791 dropped_reason: None,
7792 nullspace_dim_hint: p.saturating_sub(1),
7793 normalization_scale: 1.0,
7794 kronecker_factors: None,
7795 })
7796 .collect()
7797 } else {
7798 inner.penaltyinfo.clone()
7799 };
7800 if marginal_penalties.len() != marginal_penaltyinfo.len() {
7801 crate::bail_invalid_basis!(
7802 "internal factor-smooth penalty metadata mismatch for term '{}': penalties={}, infos={}",
7803 term_name,
7804 marginal_penalties.len(),
7805 marginal_penaltyinfo.len()
7806 );
7807 }
7808
7809 let mut penalties = Vec::<Array2<f64>>::with_capacity(marginal_penalties.len());
7810 let mut penaltyinfo = Vec::<PenaltyInfo>::with_capacity(marginal_penalties.len());
7811 for (penalty_pos, s_inner) in marginal_penalties.iter().enumerate() {
7812 let mut s_big = Array2::<f64>::zeros((q, q));
7813 for level in 0..n_levels {
7814 let start = level * p;
7815 s_big
7816 .slice_mut(s![start..start + p, start..start + p])
7817 .assign(s_inner);
7818 }
7819 let (s_big, factor_smooth_scale) = normalize_penalty_in_constrained_space(&s_big);
7820 let mut info = marginal_penaltyinfo[penalty_pos].clone();
7821 info.original_index = penalty_pos;
7822 info.normalization_scale *= factor_smooth_scale;
7823 info.nullspace_dim_hint = info.nullspace_dim_hint.saturating_mul(n_levels);
7824 info.kronecker_factors = None;
7825 penalties.push(s_big);
7826 penaltyinfo.push(info);
7827 }
7828
7829 let mut nullspaces: Vec<usize> = if matches!(spec.flavour, FactorSmoothFlavour::Re) {
7830 vec![q.saturating_sub(n_levels); p]
7831 } else {
7832 inner
7833 .nullspaces
7834 .iter()
7835 .map(|ns| ns.saturating_mul(n_levels))
7836 .collect()
7837 };
7838
7839 if use_per_dim_null
7869 && let Some(Some(z)) = inner.null_eigenvectors.first()
7870 && z.nrows() == p
7871 {
7872 for k in 0..z.ncols() {
7873 let zk = z.column(k);
7878 let mut p_k = Array2::<f64>::zeros((p, p));
7879 for a in 0..p {
7880 for b in 0..p {
7881 p_k[[a, b]] = zk[a] * zk[b];
7882 }
7883 }
7884 let mut s_null = Array2::<f64>::zeros((q, q));
7885 for level in 0..n_levels {
7886 let start = level * p;
7887 s_null
7888 .slice_mut(s![start..start + p, start..start + p])
7889 .assign(&p_k);
7890 }
7891 let (s_null, null_scale) = normalize_penalty_in_constrained_space(&s_null);
7892 let null_block = crate::basis::analyze_penalty_block_with_op(&s_null, None)?;
7893 if null_block.rank > 0 {
7894 let original_index = penalties.len();
7895 penalties.push(null_block.sym_penalty);
7896 nullspaces.push(null_block.nullity);
7897 penaltyinfo.push(PenaltyInfo {
7898 source: PenaltySource::Primary,
7899 original_index,
7900 active: true,
7901 effective_rank: null_block.rank,
7902 dropped_reason: None,
7903 nullspace_dim_hint: null_block.nullity,
7904 normalization_scale: null_scale,
7905 kronecker_factors: None,
7906 });
7907 }
7908 }
7909 }
7910 let null_eigenvectors = crate::basis::recompute_null_eigenvectors(&penalties)?;
7911 let joint_null_rotation = crate::basis::compute_joint_null_rotation(&penalties)?;
7912
7913 let (knots, degree, periodic) = match &inner.metadata {
7916 BasisMetadata::BSpline1D {
7917 knots,
7918 periodic,
7919 degree,
7920 ..
7921 } => (
7922 knots.clone(),
7923 degree.unwrap_or(spec.marginal.degree),
7924 *periodic,
7925 ),
7926 other => {
7927 crate::bail_invalid_basis!(
7928 "factor smooth term '{}' produced an unexpected marginal metadata variant {:?}",
7929 term_name,
7930 other
7931 );
7932 }
7933 };
7934 let flavour_tag = match &spec.flavour {
7935 FactorSmoothFlavour::Fs { .. } => "fs",
7936 FactorSmoothFlavour::Sz => "sz",
7937 FactorSmoothFlavour::Re => "re",
7938 }
7939 .to_string();
7940 let metadata = BasisMetadata::FactorSmooth {
7941 continuous_cols: spec.continuous_cols.clone(),
7942 group_col,
7943 knots,
7944 degree,
7945 periodic,
7946 group_levels: levels,
7947 flavour: flavour_tag,
7948 marginal_is_cr: false,
7951 };
7952
7953 let ops = vec![None; penalties.len()];
7954 Ok(LocalSmoothTermBuild {
7955 dim: q,
7956 design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(dense)),
7957 penalties,
7958 ops,
7959 nullspaces,
7960 null_eigenvectors,
7961 joint_null_rotation,
7962 penaltyinfo,
7963 pre_dropped_penaltyinfo: Vec::new(),
7964 metadata,
7965 linear_constraints: None,
7966 box_reparam: false,
7967 kronecker_factored: None,
7968 })
7969}
7970
7971pub fn resolve_factor_smooth_levels(
7975 data: ArrayView2<'_, f64>,
7976 group_col: usize,
7977 spec: &FactorSmoothSpec,
7978 term_name: &str,
7979) -> Result<Vec<u64>, BasisError> {
7980 if let Some(frozen) = &spec.group_frozen_levels {
7981 if frozen.is_empty() {
7982 crate::bail_invalid_basis!(
7983 "factor smooth term '{}' has an empty frozen level list",
7984 term_name
7985 );
7986 }
7987 return Ok(frozen
7988 .iter()
7989 .map(|&b| gam_data::canonical_level_bits(f64::from_bits(b)))
7990 .collect());
7991 }
7992 let mut bits: Vec<u64> = data
7993 .column(group_col)
7994 .iter()
7995 .map(|v| gam_data::canonical_level_bits(*v))
7996 .collect();
7997 bits.sort_by(|a, b| {
7998 f64::from_bits(*a)
7999 .partial_cmp(&f64::from_bits(*b))
8000 .unwrap_or(std::cmp::Ordering::Equal)
8001 });
8002 bits.dedup();
8003 Ok(bits)
8004}
8005
8006pub fn factor_smooth_marginal_for_replay(marginal: &BSplineBasisSpec) -> BSplineBasisSpec {
8013 let mut m = marginal.clone();
8014 m.identifiability = BSplineIdentifiability::None;
8015 m
8016}
8017
8018pub fn build_single_local_smooth_term(
8019 data: ArrayView2<'_, f64>,
8020 term: &SmoothTermSpec,
8021 workspace: &mut crate::basis::BasisWorkspace,
8022) -> Result<LocalSmoothTermBuild, BasisError> {
8023 if term.shape != ShapeConstraint::None && !shape_supports_basis(term) {
8024 crate::bail_invalid_basis!(
8025 "ShapeConstraint::{:?} is unsupported for term '{}'",
8026 term.shape,
8027 term.name
8028 );
8029 }
8030 if let SmoothBasisSpec::ByVariable {
8031 inner,
8032 by_col,
8033 kind,
8034 by,
8035 } = &term.basis
8036 {
8037 ensure_by_variable_specs_match(kind, by, &term.name)?;
8038 let mut inner_basis = (**inner).clone();
8039 if matches!(by, ByVariableSpec::Level { .. }) {
8046 defer_inner_model_centering_to_factor_level_wrapper(&mut inner_basis);
8047 }
8048 let inner_term = SmoothTermSpec {
8049 name: term.name.clone(),
8050 basis: inner_basis,
8051 shape: term.shape,
8052 joint_null_rotation: None,
8053 };
8054 let built = build_single_local_smooth_term(data, &inner_term, workspace)?;
8055 return apply_by_variable_to_local_build(built, data, *by_col, by, &term.name);
8056 }
8057
8058 if let SmoothBasisSpec::BySmooth { smooth, by_kind } = &term.basis {
8061 return build_by_smooth_local(data, term, smooth, by_kind, workspace);
8062 }
8063
8064 let mut built: BasisBuildResult = match &term.basis {
8065 SmoothBasisSpec::FactorSumToZero {
8066 inner,
8067 by_col,
8068 levels,
8069 ..
8070 } => {
8071 if *by_col >= data.ncols() {
8072 crate::bail_dim_basis!(
8073 "term '{}' by column {} out of bounds for {} columns",
8074 term.name,
8075 by_col,
8076 data.ncols()
8077 );
8078 }
8079 if levels.len() < 2 {
8080 crate::bail_invalid_basis!(
8081 "sum-to-zero factor smooth term '{}' requires at least two levels",
8082 term.name
8083 );
8084 }
8085 if term.shape != ShapeConstraint::None {
8086 crate::bail_invalid_basis!(
8087 "ShapeConstraint::{:?} is unsupported for sum-to-zero factor smooth term '{}'",
8088 term.shape,
8089 term.name
8090 );
8091 }
8092 let inner_term = SmoothTermSpec {
8093 name: format!("{}::inner", term.name),
8094 basis: (**inner).clone(),
8095 shape: ShapeConstraint::None,
8096 joint_null_rotation: None,
8097 };
8098 let mut inner_built = build_single_local_smooth_term(data, &inner_term, workspace)?;
8099 let inner_null_eigenvectors = inner_built.null_eigenvectors.clone();
8103 let base = inner_built
8104 .design
8105 .try_to_dense_by_chunks("sum-to-zero factor smooth")
8106 .map_err(BasisError::InvalidInput)?;
8107 let n = base.nrows();
8108 let p = base.ncols();
8109 let l_minus_one = levels.len() - 1;
8110 let canon_levels: Vec<u64> = levels
8113 .iter()
8114 .map(|&b| gam_data::canonical_level_bits(f64::from_bits(b)))
8115 .collect();
8116 let mut dense = Array2::<f64>::zeros((n, p * l_minus_one));
8117 for i in 0..n {
8118 let bits = gam_data::canonical_level_bits(data[[i, *by_col]]);
8119 let level_idx = canon_levels
8120 .iter()
8121 .position(|b| *b == bits)
8122 .ok_or_else(|| {
8123 BasisError::InvalidInput(format!(
8124 "sum-to-zero factor smooth term '{}' saw an unseen level at row {}",
8125 term.name,
8126 i + 1
8127 ))
8128 })?;
8129 if level_idx < l_minus_one {
8130 let start = level_idx * p;
8131 dense
8132 .slice_mut(s![i, start..start + p])
8133 .assign(&base.row(i));
8134 } else {
8135 for level in 0..l_minus_one {
8136 let start = level * p;
8137 dense
8138 .slice_mut(s![i, start..start + p])
8139 .assign(&base.row(i).mapv(|v| -v));
8140 }
8141 }
8142 }
8143 let mut penalties = Vec::<Array2<f64>>::with_capacity(inner_built.penalties.len());
8144 let active_penalty_indices = inner_built
8145 .penaltyinfo
8146 .iter()
8147 .enumerate()
8148 .filter_map(|(idx, info)| info.active.then_some(idx))
8149 .collect::<Vec<_>>();
8150 if active_penalty_indices.len() != inner_built.penalties.len() {
8151 crate::bail_invalid_basis!(
8152 "internal sz penalty metadata mismatch: activeinfos={}, penalties={}",
8153 active_penalty_indices.len(),
8154 inner_built.penalties.len()
8155 );
8156 }
8157 let stz_per_group_penalty =
8192 |s_inner: &Array2<f64>, which_level: usize| -> Array2<f64> {
8193 let mut s_big = Array2::<f64>::zeros((p * l_minus_one, p * l_minus_one));
8194 if which_level < l_minus_one {
8195 let k = which_level;
8197 let mut block = s_big.slice_mut(s![k * p..(k + 1) * p, k * p..(k + 1) * p]);
8198 block.assign(s_inner);
8199 } else {
8200 for a in 0..l_minus_one {
8202 for b in 0..l_minus_one {
8203 let mut block =
8204 s_big.slice_mut(s![a * p..(a + 1) * p, b * p..(b + 1) * p]);
8205 block.assign(s_inner);
8206 }
8207 }
8208 }
8209 s_big
8210 };
8211 let mut nullspaces = Vec::<usize>::with_capacity(penalties.capacity());
8217 for (penalty_pos, s_inner) in inner_built.penalties.iter().enumerate() {
8218 let info_idx = active_penalty_indices[penalty_pos];
8219 let base_info = inner_built.penaltyinfo[info_idx].clone();
8220 let marginal_nullity = inner_built
8221 .nullspaces
8222 .get(penalty_pos)
8223 .copied()
8224 .unwrap_or(0);
8225 for which_level in 0..=l_minus_one {
8227 let raw = stz_per_group_penalty(s_inner, which_level);
8228 let (s_big, group_scale) = normalize_penalty_in_constrained_space(&raw);
8229 let block = crate::basis::analyze_penalty_block_with_op(&s_big, None)?;
8230 if block.rank == 0 {
8231 continue;
8232 }
8233 if which_level == 0 {
8234 inner_built.penaltyinfo[info_idx].normalization_scale *= group_scale;
8237 inner_built.penaltyinfo[info_idx].original_index = penalties.len();
8238 inner_built.penaltyinfo[info_idx].effective_rank = block.rank;
8239 inner_built.penaltyinfo[info_idx].nullspace_dim_hint = block.nullity;
8240 } else {
8241 let mut info = base_info.clone();
8242 info.original_index = penalties.len();
8243 info.normalization_scale = base_info.normalization_scale * group_scale;
8244 info.effective_rank = block.rank;
8245 info.nullspace_dim_hint = block.nullity;
8246 info.kronecker_factors = None;
8247 inner_built.penaltyinfo.push(info);
8248 }
8249 penalties.push(block.sym_penalty);
8250 nullspaces.push(marginal_nullity);
8256 }
8257 }
8258
8259 if let Some(Some(z)) = inner_null_eigenvectors.first()
8277 && z.nrows() == p
8278 {
8279 for k in 0..z.ncols() {
8280 let zk = z.column(k);
8281 let mut p_k = Array2::<f64>::zeros((p, p));
8282 for a in 0..p {
8283 for b in 0..p {
8284 p_k[[a, b]] = zk[a] * zk[b];
8285 }
8286 }
8287 let stz_pooled_null = {
8292 let mut s_big = Array2::<f64>::zeros((p * l_minus_one, p * l_minus_one));
8293 for a in 0..l_minus_one {
8294 for b in 0..l_minus_one {
8295 let factor = if a == b { 2.0 } else { 1.0 };
8296 let mut block =
8297 s_big.slice_mut(s![a * p..(a + 1) * p, b * p..(b + 1) * p]);
8298 block.assign(&p_k.mapv(|v| v * factor));
8299 }
8300 }
8301 s_big
8302 };
8303 let (s_null, null_scale) =
8304 normalize_penalty_in_constrained_space(&stz_pooled_null);
8305 let null_block = crate::basis::analyze_penalty_block_with_op(&s_null, None)?;
8306 if null_block.rank > 0 {
8307 let original_index = penalties.len();
8308 penalties.push(null_block.sym_penalty);
8309 nullspaces.push(null_block.nullity);
8310 inner_built.penaltyinfo.push(PenaltyInfo {
8311 source: PenaltySource::Primary,
8312 original_index,
8313 active: true,
8314 effective_rank: null_block.rank,
8315 dropped_reason: None,
8316 nullspace_dim_hint: null_block.nullity,
8317 normalization_scale: null_scale,
8318 kronecker_factors: None,
8319 });
8320 }
8321 }
8322 }
8323 inner_built.dim = p * l_minus_one;
8324 inner_built.design =
8325 DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(dense));
8326 inner_built.penalties = penalties;
8327 inner_built.ops = vec![None; inner_built.penalties.len()];
8328 inner_built.nullspaces = nullspaces;
8329 inner_built.null_eigenvectors =
8336 crate::basis::recompute_null_eigenvectors(&inner_built.penalties)?;
8337 inner_built.joint_null_rotation =
8338 crate::basis::compute_joint_null_rotation(&inner_built.penalties)?;
8339 inner_built.kronecker_factored = None;
8340 return Ok(inner_built);
8341 }
8342 SmoothBasisSpec::BSpline1D { feature_col, spec } => {
8343 if *feature_col >= data.ncols() {
8344 crate::bail_dim_basis!(
8345 "term '{}' feature column {} out of bounds for {} columns",
8346 term.name,
8347 feature_col,
8348 data.ncols()
8349 );
8350 }
8351 let mut spec_local = spec.clone();
8352 if term.shape != ShapeConstraint::None {
8353 spec_local.identifiability = BSplineIdentifiability::None;
8356 }
8357 build_bspline_basis_1d(data.column(*feature_col), &spec_local)?
8361 }
8362 SmoothBasisSpec::ThinPlate {
8363 feature_cols,
8364 spec,
8365 input_scales,
8366 } => {
8367 if term.shape != ShapeConstraint::None {
8368 if feature_cols.len() != 1 {
8369 crate::bail_invalid_basis!(
8370 "ShapeConstraint::{:?} for term '{}' on ThinPlate basis requires exactly 1 feature axis; found {}",
8371 term.shape,
8372 term.name,
8373 feature_cols.len()
8374 );
8375 }
8376 }
8377 let mut x = select_columns(data, feature_cols)?;
8378 let (scales, length_scale_eff) = if let Some(s) = input_scales {
8384 apply_input_standardization(&mut x, s);
8385 (
8386 Some(s.clone()),
8387 compensate_length_scale_for_standardization(spec.length_scale, s),
8388 )
8389 } else if let Some(s) = compute_spatial_input_scales(x.view()) {
8390 apply_input_standardization(&mut x, &s);
8391 let l_eff = compensate_length_scale_for_standardization(spec.length_scale, &s);
8392 (Some(s), l_eff)
8393 } else {
8394 (None, spec.length_scale)
8395 };
8396 let mut spec_local = spec.clone();
8397 spec_local.length_scale = length_scale_eff;
8398 if matches!(
8399 spec_local.identifiability,
8400 SpatialIdentifiability::OrthogonalToParametric
8401 ) {
8402 spec_local.identifiability = SpatialIdentifiability::None;
8403 }
8404 let mut result = build_thin_plate_basis(x.view(), &spec_local).map_err(|err| {
8405 rewrite_thin_plate_knots_error(err, &term.name, feature_cols.len(), spec)
8406 })?;
8407 match &mut result.metadata {
8415 BasisMetadata::ThinPlate {
8416 input_scales: ms,
8417 length_scale,
8418 ..
8419 } => {
8420 *ms = scales;
8421 *length_scale = spec.length_scale;
8422 }
8423 BasisMetadata::Duchon {
8424 input_scales: ms,
8425 length_scale,
8426 ..
8427 } => {
8428 if let (Some(s), Some(realized)) = (scales.as_ref(), *length_scale) {
8453 let inv_sigma_geom = compensate_length_scale_for_standardization(1.0, s);
8454 if inv_sigma_geom.is_finite() && inv_sigma_geom > 0.0 {
8455 *length_scale = Some(realized / inv_sigma_geom);
8456 }
8457 }
8458 *ms = scales;
8459 }
8460 _ => {}
8461 }
8462 result
8463 }
8464 SmoothBasisSpec::Sphere { feature_cols, spec } => {
8465 if term.shape != ShapeConstraint::None {
8466 crate::bail_invalid_basis!(
8467 "ShapeConstraint::{:?} for term '{}' is not supported on spherical splines",
8468 term.shape,
8469 term.name
8470 );
8471 }
8472 let x = select_columns(data, feature_cols)?;
8473 build_spherical_spline_basis(x.view(), spec)?
8474 }
8475 SmoothBasisSpec::ConstantCurvature { feature_cols, spec } => {
8476 if term.shape != ShapeConstraint::None {
8477 crate::bail_invalid_basis!(
8478 "ShapeConstraint::{:?} for term '{}' is not supported on constant-curvature smooths",
8479 term.shape,
8480 term.name
8481 );
8482 }
8483 let x = select_columns(data, feature_cols)?;
8490 build_constant_curvature_basis(x.view(), spec)?
8491 }
8492 SmoothBasisSpec::MeasureJet {
8493 feature_cols,
8494 spec,
8495 input_scales,
8496 } => {
8497 if term.shape != ShapeConstraint::None {
8498 crate::bail_invalid_basis!(
8499 "ShapeConstraint::{:?} for term '{}' is not supported on measure-jet smooths",
8500 term.shape,
8501 term.name
8502 );
8503 }
8504 let mut x = select_columns(data, feature_cols)?;
8505 let (scales, length_scale_eff) = if let Some(s) = input_scales {
8517 apply_input_standardization(&mut x, s);
8518 (Some(s.clone()), spec.length_scale)
8519 } else if let Some(s) = compute_spatial_input_scales(x.view()) {
8520 apply_input_standardization(&mut x, &s);
8521 let l_eff = if spec.length_scale > 0.0 {
8522 compensate_length_scale_for_standardization(spec.length_scale, &s)
8523 } else {
8524 spec.length_scale
8525 };
8526 (Some(s), l_eff)
8527 } else {
8528 (None, spec.length_scale)
8529 };
8530 let mut spec_local = spec.clone();
8531 spec_local.length_scale = length_scale_eff;
8532 let mut result = build_measure_jet_basis(x.view(), &spec_local)?;
8533 if let BasisMetadata::MeasureJet {
8534 input_scales: ms, ..
8535 } = &mut result.metadata
8536 {
8537 *ms = scales;
8538 }
8539 result
8540 }
8541 SmoothBasisSpec::Matern {
8542 feature_cols,
8543 spec,
8544 input_scales,
8545 } => {
8546 if term.shape != ShapeConstraint::None {
8547 if feature_cols.len() != 1 {
8548 crate::bail_invalid_basis!(
8549 "ShapeConstraint::{:?} for term '{}' on Matern basis requires exactly 1 feature axis; found {}",
8550 term.shape,
8551 term.name,
8552 feature_cols.len()
8553 );
8554 }
8555 }
8556 let mut x = select_columns(data, feature_cols)?;
8557 let (scales, length_scale_eff) = if let Some(s) = input_scales {
8572 apply_input_standardization(&mut x, s);
8573 (
8574 Some(s.clone()),
8575 compensate_length_scale_for_standardization(spec.length_scale, s),
8576 )
8577 } else if let Some(s) = compute_spatial_input_scales(x.view()) {
8578 apply_input_standardization(&mut x, &s);
8579 let l_eff = compensate_length_scale_for_standardization(spec.length_scale, &s);
8580 (Some(s), l_eff)
8581 } else {
8582 (None, spec.length_scale)
8583 };
8584 let mut spec_local = spec.clone();
8585 spec_local.length_scale = length_scale_eff;
8586 let mut result = build_matern_basiswithworkspace(x.view(), &spec_local, workspace)?;
8587 if let BasisMetadata::Matern {
8588 input_scales,
8589 length_scale,
8590 ..
8591 } = &mut result.metadata
8592 {
8593 *input_scales = scales;
8594 *length_scale = spec.length_scale;
8595 }
8596 result
8597 }
8598 SmoothBasisSpec::Duchon {
8599 feature_cols,
8600 spec,
8601 input_scales,
8602 } => {
8603 if term.shape != ShapeConstraint::None {
8604 if feature_cols.len() != 1 {
8605 crate::bail_invalid_basis!(
8606 "ShapeConstraint::{:?} for term '{}' on Duchon basis requires exactly 1 feature axis; found {}",
8607 term.shape,
8608 term.name,
8609 feature_cols.len()
8610 );
8611 }
8612 }
8613 let mut x = select_columns(data, feature_cols)?;
8614 let (scales, length_scale_eff) = if let Some(s) = input_scales {
8625 apply_input_standardization(&mut x, s);
8626 (
8627 Some(s.clone()),
8628 compensate_optional_length_scale_for_standardization(spec.length_scale, s),
8629 )
8630 } else if let Some(s) = compute_spatial_input_scales(x.view()) {
8631 apply_input_standardization(&mut x, &s);
8632 let l_eff =
8633 compensate_optional_length_scale_for_standardization(spec.length_scale, &s);
8634 (Some(s), l_eff)
8635 } else {
8636 (None, spec.length_scale)
8637 };
8638 let mut spec_local = spec.clone();
8639 spec_local.length_scale = length_scale_eff;
8640 if let (Some(s), crate::basis::OneDimensionalBoundary::Cyclic { start, end }) =
8653 (scales.as_ref(), spec_local.boundary.clone())
8654 && s.len() == 1
8655 && s[0] > 0.0
8656 {
8657 spec_local.boundary = crate::basis::OneDimensionalBoundary::Cyclic {
8658 start: start / s[0],
8659 end: end / s[0],
8660 };
8661 }
8662 if let (Some(s), Some(periods)) = (scales.as_ref(), spec_local.periodic.as_mut())
8670 && s.len() == periods.len()
8671 {
8672 for (axis_period, &sigma) in periods.iter_mut().zip(s.iter()) {
8673 if sigma > 0.0
8674 && let Some(p) = axis_period.as_mut()
8675 {
8676 *p /= sigma;
8677 }
8678 }
8679 }
8680 if matches!(
8681 spec_local.identifiability,
8682 SpatialIdentifiability::OrthogonalToParametric
8683 ) {
8684 spec_local.identifiability = SpatialIdentifiability::None;
8685 }
8686 let mut result = build_duchon_basiswithworkspace(x.view(), &spec_local, workspace)?;
8687 if let BasisMetadata::Duchon {
8688 input_scales,
8689 length_scale,
8690 periodic,
8691 ..
8692 } = &mut result.metadata
8693 {
8694 *input_scales = scales;
8695 *length_scale = spec.length_scale;
8696 if spec.periodic.is_some() || spec.boundary.period().is_some() {
8712 *periodic = spec
8713 .periodic
8714 .clone()
8715 .or_else(|| spec.boundary.period().map(|(_, _, p)| vec![Some(p)]));
8716 }
8717 }
8718 result
8719 }
8720 SmoothBasisSpec::Pca {
8721 feature_cols,
8722 basis_matrix,
8723 centered,
8724 smooth_penalty,
8725 center_mean,
8726 pca_basis_path,
8727 chunk_size,
8728 } => {
8729 if term.shape != ShapeConstraint::None {
8730 crate::bail_invalid_basis!(
8731 "ShapeConstraint::{:?} for term '{}' is not supported on Pca basis",
8732 term.shape,
8733 term.name
8734 );
8735 }
8736 build_pca_smooth_basis(
8737 data,
8738 feature_cols,
8739 basis_matrix,
8740 *centered,
8741 *smooth_penalty,
8742 center_mean.as_ref(),
8743 pca_basis_path.as_ref(),
8744 *chunk_size,
8745 )?
8746 }
8747 SmoothBasisSpec::TensorBSpline { feature_cols, spec } => {
8748 build_tensor_bspline_basis(data, feature_cols, spec)?
8749 }
8750 SmoothBasisSpec::ByVariable { .. } => {
8751 crate::bail_invalid_basis!(
8752 "internal: ByVariable smooths must return before inner basis dispatch"
8753 );
8754 }
8755 SmoothBasisSpec::BySmooth { .. } => {
8756 crate::bail_invalid_basis!("internal: BySmooth smooths must be lowered to ByVariable before inner basis dispatch"
8757 .to_string(),);
8758 }
8759 SmoothBasisSpec::FactorSmooth { spec } => {
8760 if term.shape != ShapeConstraint::None {
8761 crate::bail_invalid_basis!(
8762 "ShapeConstraint::{:?} is unsupported for factor smooth term '{}'",
8763 term.shape,
8764 term.name
8765 );
8766 }
8767 return build_factor_smooth(data, spec, &term.name, workspace);
8768 }
8769 };
8770
8771 if let SmoothBasisSpec::Matern { .. } = &term.basis {
8787 let (penalties, nullspace_dims, penaltyinfo) =
8788 matern_operator_penalty_triplet_from_metadata(&built.metadata)?;
8789 built.penalties = penalties;
8790 built.nullspace_dims = nullspace_dims;
8791 built.penaltyinfo = penaltyinfo;
8792 }
8793
8794 let p_local = built.design.ncols();
8795 let mut metadata = built.metadata.clone();
8796 let kron_factored = if term.shape == ShapeConstraint::None {
8799 built.kronecker_factored
8800 } else {
8801 None
8802 };
8803 let mut design_t = built.design;
8804 let mut penalties_t: Vec<Array2<f64>> = built.penalties;
8805 let mut ops_t: Vec<Option<std::sync::Arc<dyn crate::analytic_penalties::PenaltyOp>>> =
8810 built.ops;
8811 if matches!(
8812 spatial_identifiability_policy(term),
8813 Some(SpatialIdentifiability::OrthogonalToParametric)
8814 ) {
8815 metadata = freeze_raw_spatial_metadata(metadata, design_t.ncols());
8816 }
8817
8818 let active_penaltyinfo_t = built
8819 .penaltyinfo
8820 .iter()
8821 .filter(|info| info.active)
8822 .cloned()
8823 .collect::<Vec<_>>();
8824 let pre_dropped_penaltyinfo_t = built
8825 .penaltyinfo
8826 .iter()
8827 .filter(|info| !info.active)
8828 .cloned()
8829 .collect::<Vec<_>>();
8830 let use_box_reparam =
8831 term.shape != ShapeConstraint::None && shape_uses_box_reparameterization(&term.basis);
8832 if let Some((order, sign)) = shape_order_and_sign(term.shape)
8833 && use_box_reparam
8834 {
8835 let t = if order == 2 {
8849 let (knots, degree) = match &metadata {
8850 BasisMetadata::BSpline1D {
8851 knots,
8852 degree: Some(degree),
8853 periodic,
8854 ..
8855 } if periodic.is_none() => (knots, *degree),
8856 _ => {
8857 crate::bail_invalid_basis!(
8858 "shape-constrained convex/concave term '{}' requires realized open B-spline knot and degree metadata",
8859 term.name
8860 );
8861 }
8862 };
8863 let spans = bspline_first_derivative_control_spans(knots.view(), degree)?;
8864 if spans.len() + 1 != p_local {
8865 crate::bail_invalid_basis!(
8866 "shape-constraint derivative-control span count {} does not match basis dim {} for term '{}'",
8867 spans.len(),
8868 p_local,
8869 term.name
8870 );
8871 }
8872 convex_derivative_control_transform_matrix(&spans, sign)?
8873 } else {
8874 cumulative_sum_transform_matrix(p_local, order, sign)
8875 };
8876 let inner_dense = match design_t {
8880 DesignMatrix::Dense(d) => d,
8881 DesignMatrix::Sparse(sp) => gam_linalg::matrix::DenseDesignMatrix::from(
8882 sp.try_to_dense_arc("shape-constrained coefficient transform")
8883 .map_err(BasisError::InvalidInput)?,
8884 ),
8885 };
8886 let coeff_op =
8887 gam_linalg::matrix::CoefficientTransformOperator::new(inner_dense, t.clone()).map_err(
8888 |e| BasisError::InvalidInput(format!("CoefficientTransformOperator: {e}")),
8889 )?;
8890 design_t = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(
8891 coeff_op,
8892 )));
8893 if penalties_t.len() != active_penaltyinfo_t.len() {
8894 crate::bail_invalid_basis!(
8895 "internal box-reparam penalty/info mismatch for term '{}': penalties={}, infos={}",
8896 term.name,
8897 penalties_t.len(),
8898 active_penaltyinfo_t.len()
8899 );
8900 }
8901 let mut rebuilt = Vec::with_capacity(penalties_t.len());
8908 for s_local in &penalties_t {
8909 let tt_s = fast_atb(&t, s_local);
8910 rebuilt.push(fast_ab(&tt_s, &t));
8911 }
8912 penalties_t = rebuilt;
8913 ops_t = vec![None; penalties_t.len()];
8916 }
8917 if penalties_t.len() != active_penaltyinfo_t.len() {
8918 crate::bail_invalid_basis!(
8919 "internal penalty metadata mismatch for term '{}': active penalties={}, active infos={}",
8920 term.name,
8921 penalties_t.len(),
8922 active_penaltyinfo_t.len()
8923 );
8924 }
8925 if ops_t.len() != penalties_t.len() {
8926 ops_t = vec![None; penalties_t.len()];
8927 }
8928 let penalty_candidates = penalties_t
8929 .into_iter()
8930 .zip(active_penaltyinfo_t.into_iter())
8931 .zip(ops_t.into_iter())
8932 .map(
8933 |((matrix, info), op_in)| -> Result<PenaltyCandidate, BasisError> {
8934 let (matrix, c_new) = normalize_penalty_in_constrained_space(&matrix);
8935 let normalization_scale = info.normalization_scale * c_new;
8936 let op_scale = 1.0 / c_new;
8937 let kronecker_scale = 1.0 / c_new;
8938 let scaled_op = if op_scale > 0.0 && op_scale.is_finite() {
8941 op_in.map(|op| {
8942 std::sync::Arc::new(crate::analytic_penalties::ScaledPenaltyOp::new(
8943 op, op_scale,
8944 ))
8945 as std::sync::Arc<dyn crate::analytic_penalties::PenaltyOp>
8946 })
8947 } else {
8948 None
8949 };
8950 let kronecker_factors = info.kronecker_factors.map(|mut factors| {
8951 if let Some(first) = factors.first_mut() {
8952 first.mapv_inplace(|v| v * kronecker_scale);
8953 }
8954 factors
8955 });
8956 Ok(PenaltyCandidate {
8957 nullspace_dim_hint: info.nullspace_dim_hint,
8958 matrix,
8959 source: info.source,
8960 normalization_scale,
8961 kronecker_factors,
8962 op: scaled_op,
8963 })
8964 },
8965 )
8966 .collect::<Result<Vec<_>, _>>()?;
8967 let (penalties_t, nullspaces_t, penaltyinfo_t, null_eigenvectors_t, ops_t) =
8968 crate::basis::filter_active_penalty_candidates_with_ops(penalty_candidates)?;
8969 let joint_null_rotation = match term.joint_null_rotation.clone() {
8988 Some(persisted) => Some(persisted),
8989 None if smooth_has_frozen_identifiability(term) => None,
8990 None if kron_factored.is_some() => None,
8991 None => crate::basis::compute_joint_null_rotation(&penalties_t)?,
8992 };
8993
8994 Ok(LocalSmoothTermBuild {
8995 dim: p_local,
8996 design: design_t,
8997 penalties: penalties_t,
8998 ops: ops_t,
8999 nullspaces: nullspaces_t,
9000 null_eigenvectors: null_eigenvectors_t,
9001 joint_null_rotation,
9002 penaltyinfo: penaltyinfo_t,
9003 pre_dropped_penaltyinfo: pre_dropped_penaltyinfo_t,
9004 metadata,
9005 linear_constraints: None,
9006 box_reparam: use_box_reparam,
9007 kronecker_factored: kron_factored,
9008 })
9009}
9010
9011pub fn build_smooth_design(
9012 data: ArrayView2<'_, f64>,
9013 terms: &[SmoothTermSpec],
9014) -> Result<RawSmoothDesign, BasisError> {
9015 let mut ws = crate::basis::BasisWorkspace::new();
9016 build_smooth_design_withworkspace(data, terms, &mut ws)
9017}
9018
9019pub fn build_smooth_design_withworkspace(
9026 data: ArrayView2<'_, f64>,
9027 terms: &[SmoothTermSpec],
9028 workspace: &mut crate::basis::BasisWorkspace,
9029) -> Result<RawSmoothDesign, BasisError> {
9030 validate_smooth_terms_finite_inputs(data, terms)?;
9031 build_smooth_design_withworkspace_unvalidated(data, terms, workspace)
9032}
9033
9034pub fn build_smooth_design_withworkspace_unvalidated(
9035 data: ArrayView2<'_, f64>,
9036 terms: &[SmoothTermSpec],
9037 workspace: &mut crate::basis::BasisWorkspace,
9038) -> Result<RawSmoothDesign, BasisError> {
9039 let mut planned_blocks = plan_joint_spatial_centers_for_term_blocks(data, &[terms.to_vec()])?;
9040 let planned_terms = planned_blocks.pop().ok_or_else(|| {
9041 BasisError::InvalidInput(
9042 "joint spatial center planner returned no smooth blocks".to_string(),
9043 )
9044 })?;
9045 let policy = workspace.policy().clone();
9046 let local_builds: Vec<LocalSmoothTermBuild> = {
9047 use rayon::iter::{IntoParallelIterator, ParallelIterator};
9048 planned_terms
9049 .into_par_iter()
9050 .map(|term| {
9051 let mut term_workspace = crate::basis::BasisWorkspace::with_policy(policy.clone());
9052 build_single_local_smooth_term(data, &term, &mut term_workspace)
9053 })
9054 .collect::<Result<Vec<_>, _>>()?
9055 };
9056
9057 let total_p: usize = local_builds.iter().map(|built| built.dim).sum();
9058
9059 let mut local_designs: Vec<DesignMatrix> = Vec::with_capacity(local_builds.len());
9060 let mut terms_out = Vec::<SmoothTerm>::with_capacity(terms.len());
9061 let mut penalties_global = Vec::<BlockwisePenalty>::new();
9062 let mut nullspace_dims_global = Vec::<usize>::new();
9063 let mut penaltyinfo_global = Vec::<PenaltyBlockInfo>::new();
9064 let mut dropped_penaltyinfo_global = Vec::<DroppedPenaltyBlockInfo>::new();
9065 let mut coefficient_lower_bounds = Array1::<f64>::from_elem(total_p, f64::NEG_INFINITY);
9066 let mut any_bounds = false;
9067 let mut linear_constraintsrows: Vec<(usize, usize, Array1<f64>)> = Vec::new();
9072 let mut linear_constraints_b: Vec<f64> = Vec::new();
9073
9074 let mut col_start = 0usize;
9075 for (term, mut built) in terms.iter().zip(local_builds.into_iter()) {
9076 let p_local = built.dim;
9077 let col_end = col_start + p_local;
9078 let lb_local = if built.box_reparam {
9079 shape_lower_bounds_local(term.shape, p_local)
9080 } else {
9081 None
9082 };
9083
9084 let applied_rotation: Option<crate::basis::JointNullRotation> = match (
9116 built.joint_null_rotation.take(),
9117 lb_local.is_some(),
9118 built.linear_constraints.is_some(),
9119 ) {
9120 (Some(rot), false, false) => {
9121 let q = &rot.rotation;
9122 built.design =
9123 apply_smooth_transform_to_design(built.design.clone(), q, &term.name)?;
9124 built.penalties = built
9125 .penalties
9126 .into_iter()
9127 .map(|s_local| {
9128 let qt_s = gam_linalg::faer_ndarray::fast_atb(q, &s_local);
9129 gam_linalg::faer_ndarray::fast_ab(&qt_s, q)
9130 })
9131 .collect();
9132 built.ops = vec![None; built.penalties.len()];
9133 built.kronecker_factored = None;
9134 Some(rot)
9135 }
9136 (Some(_), _, _) => None,
9137 (None, _, _) => None,
9138 };
9139
9140 let activeinfos = built
9141 .penaltyinfo
9142 .iter()
9143 .filter(|info| info.active)
9144 .collect::<Vec<_>>();
9145 if activeinfos.len() != built.penalties.len() {
9146 crate::bail_invalid_basis!(
9147 "internal penalty info mismatch for term '{}': activeinfos={}, penalties={}",
9148 term.name,
9149 activeinfos.len(),
9150 built.penalties.len()
9151 );
9152 }
9153 for (((s_local, &ns), info), op_local) in built
9154 .penalties
9155 .iter()
9156 .zip(built.nullspaces.iter())
9157 .zip(activeinfos.into_iter())
9158 .zip(built.ops.iter())
9159 {
9160 let global_index = penalties_global.len();
9161 penalties_global.push(
9162 BlockwisePenalty::new(col_start..col_end, s_local.clone())
9163 .with_op(op_local.clone()),
9164 );
9165 nullspace_dims_global.push(ns);
9166 let mut penalty = info.clone();
9167 penalty.nullspace_dim_hint = ns;
9168 penaltyinfo_global.push(PenaltyBlockInfo {
9169 global_index,
9170 termname: Some(term.name.clone()),
9171 penalty,
9172 });
9173 }
9174 for info in built.penaltyinfo.iter().filter(|info| !info.active) {
9175 dropped_penaltyinfo_global.push(DroppedPenaltyBlockInfo {
9176 termname: Some(term.name.clone()),
9177 penalty: info.clone(),
9178 });
9179 }
9180 for info in &built.pre_dropped_penaltyinfo {
9181 dropped_penaltyinfo_global.push(DroppedPenaltyBlockInfo {
9182 termname: Some(term.name.clone()),
9183 penalty: info.clone(),
9184 });
9185 }
9186
9187 if let Some(lin_local) = &built.linear_constraints {
9188 for r in 0..lin_local.a.nrows() {
9189 linear_constraintsrows.push((col_start, col_end, lin_local.a.row(r).to_owned()));
9190 linear_constraints_b.push(lin_local.b[r]);
9191 }
9192 }
9193 if let Some(lb_local) = &lb_local {
9194 coefficient_lower_bounds
9195 .slice_mut(s![col_start..col_end])
9196 .assign(lb_local);
9197 any_bounds = true;
9198 }
9199
9200 local_designs.push(built.design);
9202
9203 terms_out.push(SmoothTerm {
9204 name: term.name.clone(),
9205 coeff_range: col_start..col_end,
9206 shape: term.shape,
9207 penalties_local: built.penalties,
9208 nullspace_dims: built.nullspaces,
9209 penaltyinfo_local: built.penaltyinfo,
9210 metadata: built.metadata,
9211 lower_bounds_local: lb_local,
9212 linear_constraints_local: built.linear_constraints,
9213 kronecker_factored: built.kronecker_factored.take(),
9214 joint_null_rotation: applied_rotation,
9215 unabsorbed_global_orthogonality: None,
9216 });
9217
9218 col_start = col_end;
9219 }
9220
9221 assert_eq!(
9222 penalties_global.len(),
9223 nullspace_dims_global.len(),
9224 "global smooth penalty/nullspace bookkeeping diverged"
9225 );
9226 assert_eq!(
9227 penalties_global.len(),
9228 penaltyinfo_global.len(),
9229 "global smooth penalty metadata bookkeeping diverged"
9230 );
9231
9232 Ok(RawSmoothDesign {
9233 term_designs: local_designs,
9234 penalties: penalties_global,
9235 nullspace_dims: nullspace_dims_global,
9236 penaltyinfo: penaltyinfo_global,
9237 dropped_penaltyinfo: dropped_penaltyinfo_global,
9238 terms: terms_out,
9239 coefficient_lower_bounds: if any_bounds {
9240 Some(coefficient_lower_bounds)
9241 } else {
9242 None
9243 },
9244 linear_constraints: if linear_constraintsrows.is_empty() {
9245 None
9246 } else {
9247 let mut a = Array2::<f64>::zeros((linear_constraintsrows.len(), total_p));
9248 for (i, (cs, ce, values)) in linear_constraintsrows.iter().enumerate() {
9249 a.row_mut(i).slice_mut(s![*cs..*ce]).assign(values);
9250 }
9251 Some(LinearInequalityConstraints {
9252 a,
9253 b: Array1::from_vec(linear_constraints_b),
9254 })
9255 },
9256 })
9257}