1use super::*;
15
16use super::shape_constraints::{
17 linear_constraints_from_lower_bounds_global, merge_linear_constraints_global,
18};
19use super::structure_analysis::smooth_has_frozen_identifiability;
20use crate::basis::{
21 ConstantCurvatureIdentifiability, MaternIdentifiability, MeasureJetIdentifiability,
22 SphericalSplineIdentifiability, orthogonality_transform_for_design,
23};
24use gam_linalg::matrix::{CoefficientTransformOperator, RandomEffectOperator};
25use ndarray::ArrayView1;
26
27fn linear_function_mass(column: ArrayView1<'_, f64>, term_name: &str) -> Result<f64, BasisError> {
32 if column.is_empty() {
33 crate::bail_invalid_basis!(
34 "linear term '{term_name}' cannot define a function-space penalty on zero rows"
35 );
36 }
37 let scale = column.iter().copied().map(f64::abs).fold(0.0_f64, f64::max);
38 if !scale.is_finite() {
39 crate::bail_invalid_basis!(
40 "linear term '{term_name}' has a non-finite realized design column"
41 );
42 }
43 if scale == 0.0 {
44 crate::bail_invalid_basis!(
45 "linear term '{term_name}' is identically zero and cannot carry a recoverable effect"
46 );
47 }
48 let scaled_mean_square = column
49 .iter()
50 .map(|&value| {
51 let normalized = value / scale;
52 normalized * normalized
53 })
54 .sum::<f64>()
55 / column.len() as f64;
56 let mass = scale * scale * scaled_mean_square;
57 if !mass.is_finite() || mass <= 0.0 {
58 crate::bail_invalid_basis!(
59 "linear term '{term_name}' has an invalid empirical function mass {mass}"
60 );
61 }
62 Ok(mass)
63}
64
65pub fn build_term_collection_design_inner(
66 data: ArrayView2<'_, f64>,
67 spec: &TermCollectionSpec,
68) -> Result<TermCollectionDesign, BasisError> {
69 let policy = gam_runtime::resource::ResourcePolicy::default_library();
70 build_term_collection_design_inner_with_policy(data, spec, &policy)
71}
72
73pub fn build_term_collection_design_inner_with_policy(
79 data: ArrayView2<'_, f64>,
80 spec: &TermCollectionSpec,
81 policy: &gam_runtime::resource::ResourcePolicy,
82) -> Result<TermCollectionDesign, BasisError> {
83 use rayon::iter::{IntoParallelIterator, IntoParallelRefIterator, ParallelIterator};
84
85 let n = data.nrows();
86 let p_intercept = usize::from(!term_collection_has_anchored_bspline(spec));
87 let p_lin = spec.linear_terms.len();
88
89 let (smooth_raw_result, (random_blocks_result, linear_block_result)) = rayon::join(
94 || {
95 let mut ws = crate::basis::BasisWorkspace::with_policy(policy.clone());
96 build_smooth_design_withworkspace_unvalidated(data, &spec.smooth_terms, &mut ws)
97 },
98 || {
99 rayon::join(
100 || {
101 spec.random_effect_terms
102 .par_iter()
103 .map(|term| build_random_effect_block(data, term))
104 .collect::<Result<Vec<_>, _>>()
105 },
106 || -> Result<Option<Array2<f64>>, BasisError> {
107 if p_lin == 0 {
108 return Ok(None);
109 }
110
111 let linear_columns = (0..p_lin)
112 .into_par_iter()
113 .map(|j| {
114 let linear = &spec.linear_terms[j];
115 linear
126 .realized_design_column(data)
127 .map_err(BasisError::InvalidInput)
128 })
129 .collect::<Result<Vec<_>, _>>()?;
130
131 let mut out = Array2::<f64>::zeros((n, p_lin));
132 for (j, column) in linear_columns.iter().enumerate() {
133 out.column_mut(j).assign(column);
134 }
135 Ok(Some(out))
136 },
137 )
138 },
139 );
140
141 let smooth_raw = smooth_raw_result?;
142 let random_blocks = random_blocks_result?;
143 let linear_block = linear_block_result?;
144 let linear_function_masses = match linear_block.as_ref() {
155 Some(block) => spec
156 .linear_terms
157 .iter()
158 .enumerate()
159 .map(|(j, term)| -> Result<Option<f64>, BasisError> {
160 if !term.double_penalty {
161 return Ok(None);
162 }
163 if let Some(frozen_mass) = term.frozen_function_mass {
164 return Ok(Some(frozen_mass));
165 }
166 linear_function_mass(block.column(j), &term.name).map(Some)
167 })
168 .collect::<Result<Vec<_>, _>>()?,
169 None => Vec::new(),
170 };
171
172 let (smooth, affine_offset) = apply_global_smooth_identifiability(
173 smooth_raw,
174 data,
175 &spec.linear_terms,
176 &spec.smooth_terms,
177 )?;
178
179 let p_rand: usize = random_blocks.iter().map(|b| b.num_groups).sum();
180 let p_smooth = smooth.total_smooth_cols();
181 let p_total = p_intercept + p_lin + p_rand + p_smooth;
182
183 let mut linear_ranges = Vec::<(String, Range<usize>)>::with_capacity(p_lin);
184 for (j, linear) in spec.linear_terms.iter().enumerate() {
185 let col = p_intercept + j;
186 linear_ranges.push((linear.name.clone(), col..(col + 1)));
189 }
190
191 let mut random_effect_ranges =
194 Vec::<(String, Range<usize>)>::with_capacity(random_blocks.len());
195 let mut random_effect_levels = Vec::<(String, Vec<u64>)>::with_capacity(random_blocks.len());
196 let mut col_cursor = p_intercept + p_lin;
197 for block in &random_blocks {
198 let q = block.num_groups;
199 let end = col_cursor + q;
200 random_effect_ranges.push((block.name.clone(), col_cursor..end));
201 random_effect_levels.push((block.name.clone(), block.kept_levels.clone()));
202 col_cursor = end;
203 }
204
205 let mut blocks = Vec::<DesignBlock>::new();
221
222 if p_intercept == 1 {
227 blocks.push(DesignBlock::Intercept(n));
228 }
229
230 if let Some(lin_block) = linear_block {
232 blocks.push(DesignBlock::Dense(
233 gam_linalg::matrix::DenseDesignMatrix::from(lin_block),
234 ));
235 }
236
237 for block in &random_blocks {
239 let re_op = RandomEffectOperator::new(block.group_ids.clone(), block.num_groups);
240 blocks.push(DesignBlock::RandomEffect(Arc::new(re_op)));
241 }
242
243 if p_smooth > 0 {
247 for term_design in &smooth.term_designs {
248 match term_design {
249 DesignMatrix::Dense(dense) => blocks.push(DesignBlock::Dense(dense.clone())),
250 DesignMatrix::Sparse(sparse) => blocks.push(DesignBlock::Sparse(sparse.clone())),
251 }
252 }
253 }
254
255 let design = assemble_term_collection_design_matrix(blocks)?;
256
257 let mut penalties = Vec::<BlockwisePenalty>::new();
258 let mut nullspace_dims = Vec::<usize>::new();
259 let mut penaltyinfo = Vec::<PenaltyBlockInfo>::new();
260 let mut dropped_penaltyinfo = Vec::<DroppedPenaltyBlockInfo>::new();
261 let mut coefficient_lower_bounds = Array1::<f64>::from_elem(p_total, f64::NEG_INFINITY);
262 let mut any_bounds = false;
263 let mut linear_constraintrows = Vec::<Array1<f64>>::new();
264 let mut linear_constraint_b = Vec::<f64>::new();
265
266 for (j, linear) in spec.linear_terms.iter().enumerate() {
267 let col = p_intercept + j;
268 if let Some(lb) = linear.coefficient_min {
269 let mut row = Array1::<f64>::zeros(p_total);
270 row[col] = 1.0;
271 linear_constraintrows.push(row);
272 linear_constraint_b.push(lb);
273 }
274 if let Some(ub) = linear.coefficient_max {
275 let mut row = Array1::<f64>::zeros(p_total);
276 row[col] = -1.0;
277 linear_constraintrows.push(row);
278 linear_constraint_b.push(-ub);
279 }
280 }
281
282 for (j, linear) in spec.linear_terms.iter().enumerate() {
290 let Some(function_mass) = linear_function_masses.get(j).copied().flatten() else {
291 continue;
292 };
293 let col = p_intercept + j;
294 let global_index = penalties.len();
295 penalties.push(BlockwisePenalty::new(
296 col..(col + 1),
297 Array2::from_elem((1, 1), function_mass),
298 ));
299 nullspace_dims.push(0);
300 penaltyinfo.push(PenaltyBlockInfo {
301 global_index,
302 termname: Some(linear.name.clone()),
303 penalty: ActivePenaltyInfo {
304 source: PenaltySource::Other("LinearTermRidge".to_string()),
305 original_index: j,
306 effective_rank: 1,
307 normalization_scale: 1.0,
308 kronecker_factors: None,
309 structural_null_frame: None,
310 },
311 });
312 }
313
314 for (re_idx, (name, range)) in random_effect_ranges.iter().enumerate() {
315 if range.is_empty() || !spec.random_effect_terms[re_idx].penalized {
316 continue;
317 }
318 let block_size = range.len();
319 let global_index = penalties.len();
320 penalties.push(BlockwisePenalty::ridge(range.clone(), 1.0));
321 nullspace_dims.push(0);
322 penaltyinfo.push(PenaltyBlockInfo {
323 global_index,
324 termname: Some(name.clone()),
325 penalty: ActivePenaltyInfo {
326 source: PenaltySource::Other(format!("RandomEffectRidge({name})")),
327 original_index: re_idx,
328 effective_rank: block_size,
329 normalization_scale: 1.0,
330 kronecker_factors: None,
331 structural_null_frame: None,
332 },
333 });
334 }
335
336 if smooth.penaltyinfo.len() != smooth.penalties.len() {
337 gam_problem::bail_invalid_basis!(
338 "smooth penalty metadata mismatch: penalties={}, metadata={}",
339 smooth.penalties.len(),
340 smooth.penaltyinfo.len()
341 );
342 }
343 let smooth_start = p_intercept + p_lin + p_rand;
344 for ((bp_smooth, &ns), localinfo) in smooth
345 .penalties
346 .iter()
347 .zip(smooth.nullspace_dims.iter())
348 .zip(smooth.penaltyinfo.iter())
349 {
350 let global_index = penalties.len();
351 let offset_range =
353 (bp_smooth.col_range.start + smooth_start)..(bp_smooth.col_range.end + smooth_start);
354 let bp = if let Some(factors) = localinfo.penalty.kronecker_factors.as_ref() {
355 BlockwisePenalty::kronecker(offset_range, bp_smooth.local.clone(), factors.clone())
356 .with_op(bp_smooth.op.clone())
357 } else if matches!(
358 localinfo.penalty.source,
359 PenaltySource::Other(ref s) if s.starts_with("RandomEffectRidge")
360 ) {
361 BlockwisePenalty::ridge(offset_range, 1.0)
362 } else {
363 BlockwisePenalty::new(offset_range, bp_smooth.local.clone())
364 .with_op(bp_smooth.op.clone())
365 };
366 penalties.push(bp);
367 nullspace_dims.push(ns);
368 penaltyinfo.push(PenaltyBlockInfo {
369 global_index,
370 termname: localinfo.termname.clone(),
371 penalty: localinfo.penalty.clone(),
372 });
373 }
374 dropped_penaltyinfo.extend(smooth.dropped_penaltyinfo.iter().cloned());
375
376 assert_eq!(
377 penalties.len(),
378 nullspace_dims.len(),
379 "term-collection penalty/nullspace bookkeeping diverged"
380 );
381 assert_eq!(
382 penalties.len(),
383 penaltyinfo.len(),
384 "term-collection penalty metadata bookkeeping diverged"
385 );
386
387 if let Some(lb_smooth) = smooth.coefficient_lower_bounds.as_ref() {
388 let start = p_intercept + p_lin + p_rand;
389 coefficient_lower_bounds
390 .slice_mut(s![start..(start + p_smooth)])
391 .assign(lb_smooth);
392 any_bounds = true;
393 }
394 if let Some(lin_smooth) = smooth.linear_constraints.as_ref() {
395 let mut a_global = Array2::<f64>::zeros((lin_smooth.a.nrows(), p_total));
396 let start = p_intercept + p_lin + p_rand;
397 a_global
398 .slice_mut(s![.., start..(start + p_smooth)])
399 .assign(&lin_smooth.a);
400 for r in 0..a_global.nrows() {
401 linear_constraintrows.push(a_global.row(r).to_owned());
402 linear_constraint_b.push(lin_smooth.b[r]);
403 }
404 }
405
406 let lower_bound_constraints = if any_bounds {
410 linear_constraints_from_lower_bounds_global(&coefficient_lower_bounds)
411 } else {
412 None
413 };
414 let explicit_linear_constraints = if linear_constraintrows.is_empty() {
415 None
416 } else {
417 let mut a = Array2::<f64>::zeros((linear_constraintrows.len(), p_total));
418 for (i, row) in linear_constraintrows.iter().enumerate() {
419 a.row_mut(i).assign(row);
420 }
421 Some(LinearInequalityConstraints {
422 a,
423 b: Array1::from_vec(linear_constraint_b),
424 })
425 };
426 let linear_constraints =
427 merge_linear_constraints_global(explicit_linear_constraints, lower_bound_constraints)?;
428
429 Ok(TermCollectionDesign {
430 design,
431 affine_offset,
432 penalties,
433 nullspace_dims,
434 penaltyinfo,
435 dropped_penaltyinfo,
436 coefficient_lower_bounds: if any_bounds {
437 Some(coefficient_lower_bounds)
438 } else {
439 None
440 },
441 linear_constraints,
442 intercept_range: 0..p_intercept,
443 linear_ranges,
444 linear_function_masses,
445 random_effect_ranges,
446 random_effect_levels,
447 smooth,
448 })
449}
450
451pub fn term_collection_has_anchored_bspline(spec: &TermCollectionSpec) -> bool {
456 spec.smooth_terms
457 .iter()
458 .any(|term| smooth_basis_has_anchored_bspline(&term.basis))
459}
460
461pub fn term_collection_has_nonzero_anchor(spec: &TermCollectionSpec) -> bool {
464 spec.smooth_terms
465 .iter()
466 .any(|term| smooth_basis_has_nonzero_anchor(&term.basis))
467}
468
469fn smooth_basis_has_nonzero_anchor(basis: &SmoothBasisSpec) -> bool {
470 match basis {
471 SmoothBasisSpec::ByVariable { inner, .. }
472 | SmoothBasisSpec::FactorSumToZero { inner, .. } => smooth_basis_has_nonzero_anchor(inner),
473 SmoothBasisSpec::BSpline1D { spec, .. } => spec.boundary_conditions.has_nonzero_anchor(),
474 SmoothBasisSpec::BySmooth { smooth, .. } => smooth_basis_has_nonzero_anchor(smooth),
475 SmoothBasisSpec::TensorBSpline { spec, .. } => spec
476 .marginalspecs
477 .iter()
478 .any(|marginal| marginal.boundary_conditions.has_nonzero_anchor()),
479 SmoothBasisSpec::FactorSmooth { spec } => {
480 spec.marginal.boundary_conditions.has_nonzero_anchor()
481 }
482 SmoothBasisSpec::ThinPlate { .. }
483 | SmoothBasisSpec::Sphere { .. }
484 | SmoothBasisSpec::ConstantCurvature { .. }
485 | SmoothBasisSpec::Matern { .. }
486 | SmoothBasisSpec::MeasureJet { .. }
487 | SmoothBasisSpec::Duchon { .. }
488 | SmoothBasisSpec::Pca { .. } => false,
489 }
490}
491
492fn smooth_basis_has_anchored_bspline(basis: &SmoothBasisSpec) -> bool {
493 match basis {
494 SmoothBasisSpec::ByVariable { inner, .. }
495 | SmoothBasisSpec::FactorSumToZero { inner, .. } => {
496 smooth_basis_has_anchored_bspline(inner)
497 }
498 SmoothBasisSpec::BSpline1D { spec, .. } => {
499 bspline_conditions_have_anchor(&spec.boundary_conditions)
500 }
501 SmoothBasisSpec::BySmooth { smooth, .. } => smooth_basis_has_anchored_bspline(smooth),
502 SmoothBasisSpec::TensorBSpline { spec, .. } => spec
503 .marginalspecs
504 .iter()
505 .any(|marginal| bspline_conditions_have_anchor(&marginal.boundary_conditions)),
506 SmoothBasisSpec::FactorSmooth { .. }
507 | SmoothBasisSpec::ThinPlate { .. }
508 | SmoothBasisSpec::Sphere { .. }
509 | SmoothBasisSpec::ConstantCurvature { .. }
510 | SmoothBasisSpec::Matern { .. }
511 | SmoothBasisSpec::MeasureJet { .. }
512 | SmoothBasisSpec::Duchon { .. }
513 | SmoothBasisSpec::Pca { .. } => false,
514 }
515}
516
517fn bspline_conditions_have_anchor(conditions: &crate::basis::BSplineBoundaryConditions) -> bool {
518 conditions.has_anchor()
519}
520
521pub fn build_term_collection_design(
522 data: ArrayView2<'_, f64>,
523 spec: &TermCollectionSpec,
524) -> Result<TermCollectionDesign, BasisError> {
525 let policy = gam_runtime::resource::ResourcePolicy::default_library();
526 build_term_collection_design_with_policy(data, spec, &policy)
527}
528
529pub fn build_term_collection_design_with_policy(
533 data: ArrayView2<'_, f64>,
534 spec: &TermCollectionSpec,
535 policy: &gam_runtime::resource::ResourcePolicy,
536) -> Result<TermCollectionDesign, BasisError> {
537 validate_term_collection_finite_inputs(data, spec)?;
538 let mut planned_specs =
539 plan_joint_spatial_centers_for_term_blocks(data, &[spec.smooth_terms.clone()])?;
540 let planned_smooth_terms = planned_specs.pop().ok_or_else(|| {
541 BasisError::InvalidInput(
542 "joint spatial center planner returned no smooth terms for single-spec build"
543 .to_string(),
544 )
545 })?;
546 let mut planned_spec = spec.clone();
547 planned_spec.smooth_terms = planned_smooth_terms;
548 build_term_collection_design_inner_with_policy(data, &planned_spec, policy)
549}
550
551#[derive(Debug, Clone)]
553pub struct TermCollectionDerivativeDesign {
554 pub design: Array2<f64>,
556 pub affine_offset: Array1<f64>,
558}
559
560impl TermCollectionDerivativeDesign {
561 pub fn apply(&self, beta: ArrayView1<'_, f64>) -> Result<Array1<f64>, BasisError> {
563 if beta.len() != self.design.ncols() {
564 crate::bail_dim_basis!(
565 "term-collection derivative coefficient length {} does not match design width {}",
566 beta.len(),
567 self.design.ncols()
568 );
569 }
570 if self.affine_offset.len() != self.design.nrows() {
571 crate::bail_dim_basis!(
572 "term-collection derivative affine offset has {} rows but derivative design has {}",
573 self.affine_offset.len(),
574 self.design.nrows()
575 );
576 }
577 if beta.iter().any(|value| !value.is_finite())
578 || self.affine_offset.iter().any(|value| !value.is_finite())
579 {
580 crate::bail_invalid_basis!(
581 "term-collection derivative coefficients and affine offset must be finite"
582 );
583 }
584 Ok(self.design.dot(&beta.to_owned()) + &self.affine_offset)
585 }
586}
587
588pub fn build_term_collection_derivative_design(
621 data: ArrayView2<'_, f64>,
622 spec: &TermCollectionSpec,
623 deriv_col: usize,
624) -> Result<TermCollectionDerivativeDesign, BasisError> {
625 if deriv_col >= data.ncols() {
626 return Err(BasisError::InvalidInput(format!(
627 "average-derivative column {deriv_col} out of range for data with {} columns",
628 data.ncols()
629 )));
630 }
631
632 let value = build_term_collection_design(data, spec)?;
636 let n = data.nrows();
637 let p_total = value.design.ncols();
638 let mut d = Array2::<f64>::zeros((n, p_total));
639 let mut affine_derivative = Array1::<f64>::zeros(n);
640
641 let p_intercept = value.intercept_range.len();
643 let p_lin = spec.linear_terms.len();
644 let p_rand: usize = value
645 .random_effect_ranges
646 .iter()
647 .map(|(_, range)| range.len())
648 .sum();
649
650 for (j, linear) in spec.linear_terms.iter().enumerate() {
655 let col = p_intercept + j;
656 let derivative = linear_term_derivative_column(data, linear, deriv_col)?;
657 if let Some(column) = derivative {
658 d.column_mut(col).assign(&column);
659 }
660 }
661
662 let smooth_start = p_intercept + p_lin + p_rand;
664 if value.smooth.terms.len() != spec.smooth_terms.len() {
665 return Err(BasisError::InvalidInput(format!(
666 "average-derivative design: value build produced {} smooth terms but spec has {}",
667 value.smooth.terms.len(),
668 spec.smooth_terms.len()
669 )));
670 }
671 for (idx, termspec) in spec.smooth_terms.iter().enumerate() {
672 let term_value = &value.smooth.terms[idx];
673 let feature_cols = smooth_term_feature_cols(termspec);
674 if !feature_cols.contains(&deriv_col) {
675 continue;
677 }
678 let (block, term_affine_derivative) =
679 smooth_term_first_derivative_block(data, termspec, term_value, deriv_col)?;
680 let range = (term_value.coeff_range.start + smooth_start)
681 ..(term_value.coeff_range.end + smooth_start);
682 if block.ncols() != range.len() {
683 return Err(BasisError::DimensionMismatch(format!(
684 "average-derivative design: smooth term '{}' derivative block has {} columns \
685 but the fitted block spans {}",
686 termspec.name,
687 block.ncols(),
688 range.len()
689 )));
690 }
691 d.slice_mut(s![.., range]).assign(&block);
692 if let Some(term_offset) = term_affine_derivative {
693 if term_offset.len() != n {
694 return Err(BasisError::DimensionMismatch(format!(
695 "average-derivative design: smooth term '{}' affine derivative has {} rows but the data has {n}",
696 termspec.name,
697 term_offset.len()
698 )));
699 }
700 affine_derivative += &term_offset;
701 }
702 }
703
704 Ok(TermCollectionDerivativeDesign {
705 design: d,
706 affine_offset: affine_derivative,
707 })
708}
709
710fn linear_term_derivative_column(
718 data: ArrayView2<'_, f64>,
719 linear: &LinearTermSpec,
720 deriv_col: usize,
721) -> Result<Option<Array1<f64>>, BasisError> {
722 let numeric_cols: Vec<usize> = if linear.categorical_levels.is_empty() {
723 linear.effective_feature_cols()
724 } else {
725 linear.feature_cols.clone()
726 };
727 let occurrences = numeric_cols.iter().filter(|&&c| c == deriv_col).count();
728 if occurrences == 0 {
729 return Ok(None);
730 }
731 let n = data.nrows();
732 let p = data.ncols();
733 for &c in &numeric_cols {
734 if c >= p {
735 return Err(BasisError::InvalidInput(format!(
736 "linear term '{}' feature column {c} out of bounds for {p} columns",
737 linear.name
738 )));
739 }
740 }
741
742 let mut gate = Array1::<f64>::ones(n);
744 for &(col, level_bits) in &linear.categorical_levels {
745 if col >= p {
746 return Err(BasisError::InvalidInput(format!(
747 "linear term '{}' categorical column {col} out of bounds for {p} columns",
748 linear.name
749 )));
750 }
751 let level_bits = gam_data::canonical_level_bits(f64::from_bits(level_bits));
752 for (row, g) in gate.iter_mut().enumerate() {
753 if gam_data::canonical_level_bits(data[[row, col]]) != level_bits {
754 *g = 0.0;
755 }
756 }
757 }
758
759 let mut derivative = Array1::<f64>::zeros(n);
761 for (j, &c_j) in numeric_cols.iter().enumerate() {
762 if c_j != deriv_col {
763 continue;
764 }
765 let mut term = gate.clone();
766 for (k, &c_k) in numeric_cols.iter().enumerate() {
767 if k != j {
768 term *= &data.column(c_k);
769 }
770 }
771 derivative += &term;
772 }
773 Ok(Some(derivative))
774}
775
776fn smooth_term_first_derivative_block(
785 data: ArrayView2<'_, f64>,
786 termspec: &SmoothTermSpec,
787 term_value: &SmoothTerm,
788 deriv_col: usize,
789) -> Result<(Array2<f64>, Option<Array1<f64>>), BasisError> {
790 let feature_col = match &termspec.basis {
791 SmoothBasisSpec::BSpline1D { feature_col, .. } => *feature_col,
792 other => {
793 return Err(BasisError::InvalidInput(format!(
794 "analytic average-derivative design only supports non-periodic 1-D B-spline \
795 smooths over the differentiated covariate; term '{}' uses unsupported basis {}",
796 termspec.name,
797 smooth_basis_kind_label(other)
798 )));
799 }
800 };
801 if feature_col != deriv_col {
802 return Err(BasisError::InvalidInput(format!(
806 "analytic average-derivative design: B-spline term '{}' is over column {feature_col}, \
807 not the differentiated column {deriv_col}",
808 termspec.name
809 )));
810 }
811
812 let (knots, degree, transform, periodic, anchor_offset_coeffs) = match &term_value.metadata {
813 BasisMetadata::BSpline1D {
814 knots,
815 degree,
816 identifiability_transform,
817 periodic,
818 anchor_offset_coeffs,
819 ..
820 } => (
821 knots,
822 *degree,
823 identifiability_transform.as_ref(),
824 periodic,
825 anchor_offset_coeffs.as_ref(),
826 ),
827 other => {
828 return Err(BasisError::InvalidInput(format!(
829 "analytic average-derivative design expected B-spline metadata for term '{}', \
830 found {other:?}",
831 termspec.name
832 )));
833 }
834 };
835 if periodic.is_some() {
836 return Err(BasisError::InvalidInput(format!(
837 "analytic average-derivative design does not support periodic/cyclic B-spline \
838 term '{}'",
839 termspec.name
840 )));
841 }
842 let degree = degree.ok_or_else(|| {
843 BasisError::InvalidInput(format!(
844 "B-spline term '{}' metadata is missing its effective degree",
845 termspec.name
846 ))
847 })?;
848
849 let (deriv_basis_arc, _) = crate::basis::create_basis::<crate::basis::Dense>(
851 data.column(deriv_col),
852 crate::basis::KnotSource::Provided(knots.view()),
853 degree,
854 crate::basis::BasisOptions::first_derivative(),
855 )?;
856 let deriv_basis = deriv_basis_arc.as_ref();
857
858 let affine_derivative = match anchor_offset_coeffs {
859 Some(beta_p) => {
860 if deriv_basis.ncols() != beta_p.len() {
861 return Err(BasisError::DimensionMismatch(format!(
862 "B-spline term '{}': raw derivative basis has {} columns but the affine anchor lift has {} coefficients",
863 termspec.name,
864 deriv_basis.ncols(),
865 beta_p.len()
866 )));
867 }
868 Some(deriv_basis.dot(beta_p))
869 }
870 None => None,
871 };
872
873 let block = match transform {
877 Some(z) => {
878 if deriv_basis.ncols() != z.nrows() {
879 return Err(BasisError::DimensionMismatch(format!(
880 "B-spline term '{}': raw derivative basis has {} columns but the frozen \
881 identifiability transform has {} rows",
882 termspec.name,
883 deriv_basis.ncols(),
884 z.nrows()
885 )));
886 }
887 gam_linalg::faer_ndarray::fast_ab(deriv_basis, z)
888 }
889 None => deriv_basis.to_owned(),
890 };
891 Ok((block, affine_derivative))
892}
893
894fn smooth_basis_kind_label(basis: &SmoothBasisSpec) -> &'static str {
897 match basis {
898 SmoothBasisSpec::BSpline1D { .. } => "BSpline1D",
899 SmoothBasisSpec::TensorBSpline { .. } => "TensorBSpline",
900 SmoothBasisSpec::ByVariable { .. } => "ByVariable",
901 SmoothBasisSpec::FactorSumToZero { .. } => "FactorSumToZero",
902 SmoothBasisSpec::FactorSmooth { .. } => "FactorSmooth",
903 SmoothBasisSpec::BySmooth { .. } => "BySmooth",
904 SmoothBasisSpec::ThinPlate { .. } => "ThinPlate",
905 SmoothBasisSpec::Duchon { .. } => "Duchon",
906 SmoothBasisSpec::Matern { .. } => "Matern",
907 SmoothBasisSpec::Sphere { .. } => "Sphere",
908 SmoothBasisSpec::ConstantCurvature { .. } => "ConstantCurvature",
909 SmoothBasisSpec::MeasureJet { .. } => "MeasureJet",
910 SmoothBasisSpec::Pca { .. } => "Pca",
911 }
912}
913
914enum GlobalIdentifiabilityPlan {
922 Absent,
924 Delete { block: Array2<f64> },
927 Residualize { block: Array2<f64> },
930}
931
932impl GlobalIdentifiabilityPlan {
933 fn as_gauge(
935 &self,
936 owner_terms: &[usize],
937 has_parametric_block: bool,
938 local_identifiability_transform: Option<Array2<f64>>,
939 coefficient_transform: Array2<f64>,
940 local_columns: usize,
941 joint_null_rotation: Option<crate::basis::JointNullRotation>,
942 ) -> Option<SmoothCollectionGauge> {
943 let (arm, block) = match self {
944 Self::Absent => return None,
945 Self::Delete { block } => (SmoothCollectionGaugeArm::Delete, block),
946 Self::Residualize { block } => (SmoothCollectionGaugeArm::Residualize, block),
947 };
948 Some(SmoothCollectionGauge {
949 arm,
950 constraint_block: block.clone(),
951 owner_terms: owner_terms.to_vec(),
952 has_parametric_block,
953 local_identifiability_transform,
954 coefficient_transform,
955 local_columns,
956 joint_null_rotation,
957 })
958 }
959}
960
961fn basis_local_identifiability_transform(metadata: &BasisMetadata) -> Option<Array2<f64>> {
973 match metadata {
974 BasisMetadata::BSpline1D {
975 identifiability_transform,
976 ..
977 }
978 | BasisMetadata::CubicRegression1D {
979 identifiability_transform,
980 ..
981 }
982 | BasisMetadata::ThinPlate {
983 identifiability_transform,
984 ..
985 }
986 | BasisMetadata::Matern {
987 identifiability_transform,
988 ..
989 }
990 | BasisMetadata::Duchon {
991 identifiability_transform,
992 ..
993 }
994 | BasisMetadata::TensorBSpline {
995 identifiability_transform,
996 ..
997 } => identifiability_transform.clone(),
998 BasisMetadata::Sphere {
999 constraint_transform,
1000 ..
1001 }
1002 | BasisMetadata::ConstantCurvature {
1003 constraint_transform,
1004 ..
1005 }
1006 | BasisMetadata::MeasureJet {
1007 constraint_transform,
1008 ..
1009 } => constraint_transform.clone(),
1010 BasisMetadata::Pca { .. }
1011 | BasisMetadata::SphereHarmonics { .. }
1012 | BasisMetadata::BySmooth { .. }
1013 | BasisMetadata::FactorSmooth { .. } => None,
1014 }
1015}
1016
1017pub struct RealizedCollectionGauge {
1019 pub design: DesignMatrix,
1021 pub coefficient_transform: Array2<f64>,
1024 pub residualization: crate::basis::ParametricResidualization,
1028}
1029
1030fn derive_smooth_collection_coefficient_transform(
1039 design_local: &DesignMatrix,
1040 arm: SmoothCollectionGaugeArm,
1041 block: ArrayView2<'_, f64>,
1042 has_owner_terms: bool,
1043) -> Result<Array2<f64>, BasisError> {
1044 match arm {
1045 SmoothCollectionGaugeArm::Delete => {
1046 match orthogonality_transform_for_design(design_local, block, None) {
1047 Ok(transform) => Ok(transform),
1048 Err(BasisError::ConstraintNullspaceCollapsed { .. }) if has_owner_terms => {
1052 Ok(Array2::zeros((design_local.ncols(), 0)))
1053 }
1054 Err(error) => Err(error),
1055 }
1056 }
1057 SmoothCollectionGaugeArm::Residualize => {
1058 Ok(crate::basis::parametric_residualization_for_design(
1059 design_local,
1060 block,
1061 None, )?
1063 .coefficient_transform)
1064 }
1065 }
1066}
1067
1068pub fn realize_smooth_collection_gauge(
1087 design_local: DesignMatrix,
1088 gauge: &SmoothCollectionGauge,
1089 termname: &str,
1090) -> Result<RealizedCollectionGauge, BasisError> {
1091 let block = gauge.constraint_block.view();
1092 if block.nrows() != design_local.nrows() {
1093 gam_problem::bail_dim_basis!(
1094 "collection gauge row mismatch for term '{termname}': the design has {} rows and the frozen constraint block has {}",
1095 design_local.nrows(),
1096 block.nrows()
1097 );
1098 }
1099 if gauge.coefficient_transform.nrows() != gauge.local_columns {
1100 gam_problem::bail_dim_basis!(
1101 "collection gauge for term '{termname}' declares {} local columns but its fixed coefficient chart has {} rows",
1102 gauge.local_columns,
1103 gauge.coefficient_transform.nrows()
1104 );
1105 }
1106 if design_local.ncols() != gauge.local_columns {
1107 gam_problem::bail_dim_basis!(
1108 "collection gauge local width mismatch for term '{termname}': the design has {} columns but the fixed chart was derived on {}",
1109 design_local.ncols(),
1110 gauge.local_columns
1111 );
1112 }
1113 let coefficient_transform = gauge.coefficient_transform.clone();
1114 let design = apply_smooth_transform_to_design(
1115 design_local,
1116 &coefficient_transform,
1117 termname,
1118 )?;
1119 let projector = crate::basis::FixedRowSpaceProjector::from_constraint_block(block)?;
1120 let (design, row_space_correction) = projector.project_design(design, termname)?;
1121 let residualization = crate::basis::ParametricResidualization {
1122 coefficient_transform: coefficient_transform.clone(),
1123 row_space_correction,
1124 };
1125 assert_orthogonal_to_constraint_block(&design, block, termname)?;
1126 Ok(RealizedCollectionGauge {
1127 design,
1128 coefficient_transform,
1129 residualization,
1130 })
1131}
1132
1133pub struct LocalTermRealization<'a> {
1138 pub design: DesignMatrix,
1140 pub metadata: &'a BasisMetadata,
1142 pub active_penalties: &'a [ActivePenalty],
1143 pub dropped_penalties: Vec<DroppedPenaltyInfo>,
1144 pub linear_constraints_local: Option<&'a gam_problem::LinearInequalityConstraints>,
1145 pub joint_null_rotation: Option<&'a crate::basis::JointNullRotation>,
1148 pub termname: &'a str,
1149}
1150
1151pub struct CollectionGaugedTerm {
1153 pub design: DesignMatrix,
1154 pub metadata: BasisMetadata,
1155 pub active_penalties: Vec<ActivePenalty>,
1156 pub dropped_penalties: Vec<DroppedPenaltyInfo>,
1157 pub linear_constraints_local: Option<gam_problem::LinearInequalityConstraints>,
1158 pub parametric_residualization: Option<ParametricResidualizationChart>,
1159}
1160
1161pub fn place_term_in_collection_gauge(
1181 gauge: &SmoothCollectionGauge,
1182 local: LocalTermRealization<'_>,
1183) -> Result<CollectionGaugedTerm, BasisError> {
1184 let LocalTermRealization {
1185 design,
1186 metadata,
1187 active_penalties,
1188 dropped_penalties,
1189 linear_constraints_local,
1190 joint_null_rotation,
1191 termname,
1192 } = local;
1193 let realized = realize_smooth_collection_gauge(design, gauge, termname)?;
1194 let coefficient_gauge =
1195 gam_problem::Gauge::from_block_transforms(&[realized.coefficient_transform.clone()]);
1196 let candidates = penalty_candidates_under_collection_gauge(
1197 active_penalties,
1198 Some(&coefficient_gauge),
1199 termname,
1200 )?;
1201 let filtered = filter_penalty_candidates(candidates)?;
1202 let mut dropped_penalties = dropped_penalties;
1203 dropped_penalties.extend(filtered.dropped);
1204 let linear_constraints_local = linear_constraints_local.map(|lin| {
1205 gam_problem::LinearInequalityConstraints {
1206 a: lin.a.dot(&coefficient_gauge.block_transform(0)),
1207 b: lin.b.clone(),
1208 }
1209 });
1210 let realized_transform = match joint_null_rotation {
1211 Some(rotation) => {
1212 gam_linalg::faer_ndarray::fast_ab(&rotation.rotation, &realized.coefficient_transform)
1213 }
1214 None => realized.coefficient_transform.clone(),
1215 };
1216 let metadata = with_identifiability_transform(metadata, Some(&realized_transform))?;
1217 let parametric_residualization = Some(ParametricResidualizationChart {
1218 owner_terms: gauge.owner_terms.clone(),
1219 has_parametric_block: gauge.has_parametric_block,
1220 correction: realized.residualization.row_space_correction.clone(),
1221 });
1222 Ok(CollectionGaugedTerm {
1223 design: realized.design,
1224 metadata,
1225 active_penalties: filtered.active,
1226 dropped_penalties,
1227 linear_constraints_local,
1228 parametric_residualization,
1229 })
1230}
1231
1232const ORTHOGONALITY_REL_RESIDUAL_TOL: f64 = 1e-8;
1235
1236fn assert_orthogonal_to_constraint_block(
1237 design: &DesignMatrix,
1238 constraint: ArrayView2<'_, f64>,
1239 termname: &str,
1240) -> Result<(), BasisError> {
1241 let rel = orthogonality_relative_residual_for_design(design, constraint)?;
1242 if rel > ORTHOGONALITY_REL_RESIDUAL_TOL {
1243 gam_problem::bail_invalid_basis!(
1244 "smooth orthogonality residual too large for term '{}': {:.3e} > {:.1e}",
1245 termname,
1246 rel,
1247 ORTHOGONALITY_REL_RESIDUAL_TOL
1248 );
1249 }
1250 Ok(())
1251}
1252
1253fn subtract_row_space_correction(
1258 design: DesignMatrix,
1259 constraint: ArrayView2<'_, f64>,
1260 correction: ArrayView2<'_, f64>,
1261 termname: &str,
1262) -> Result<DesignMatrix, BasisError> {
1263 use gam_linalg::matrix::{BlockDesignOperator, DesignBlock};
1264 let p = design.ncols();
1265 let q = constraint.ncols();
1266 let k = correction.ncols();
1267 if correction.nrows() != q || p != k {
1268 return Err(BasisError::InvalidInput(format!(
1269 "row-space correction shape mismatch for term '{termname}': design is {}x{p}, \
1270 constraint is {}x{q}, correction is {}x{k}",
1271 design.nrows(),
1272 constraint.nrows(),
1273 correction.nrows(),
1274 )));
1275 }
1276 if q == 0 {
1277 return Ok(design);
1278 }
1279 let design_block = match design {
1280 DesignMatrix::Dense(inner) => DesignBlock::Dense(inner),
1281 DesignMatrix::Sparse(inner) => DesignBlock::Sparse(inner),
1282 };
1283 let stacked = BlockDesignOperator::new(vec![
1284 design_block,
1285 DesignBlock::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
1286 constraint.to_owned(),
1287 )),
1288 ])
1289 .map_err(BasisError::InvalidInput)?;
1290 let mut transform = Array2::<f64>::zeros((p + q, k));
1293 for i in 0..p {
1294 transform[[i, i]] = 1.0;
1295 }
1296 for i in 0..q {
1297 for j in 0..k {
1298 transform[[p + i, j]] = -correction[[i, j]];
1299 }
1300 }
1301 let operator = gam_linalg::matrix::CoefficientTransformOperator::new(
1302 gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(stacked)),
1303 transform,
1304 )
1305 .map_err(|e| {
1306 BasisError::InvalidInput(format!(
1307 "row-space correction failed for term '{termname}': {e}"
1308 ))
1309 })?;
1310 Ok(DesignMatrix::Dense(
1311 gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(operator)),
1312 ))
1313}
1314
1315fn build_constraint_block(
1316 n: usize,
1317 parametric_block: Option<&Array2<f64>>,
1318 owner_blocks: &[&DesignMatrix],
1319) -> Result<Array2<f64>, BasisError> {
1320 let param_cols = parametric_block.map_or(0, |mat| mat.ncols());
1321 let owner_cols: usize = owner_blocks.iter().map(|design| design.ncols()).sum();
1322 let mut block = Array2::<f64>::zeros((n, param_cols + owner_cols));
1323 let mut col_start = 0usize;
1324 if let Some(parametric) = parametric_block {
1325 let col_end = col_start + parametric.ncols();
1326 block
1327 .slice_mut(s![.., col_start..col_end])
1328 .assign(parametric);
1329 col_start = col_end;
1330 }
1331 const CHUNK: usize = 1024;
1332 for owner in owner_blocks {
1333 let col_end = col_start + owner.ncols();
1334 for row_start in (0..n).step_by(CHUNK) {
1335 let row_end = (row_start + CHUNK).min(n);
1336 let chunk = (*owner)
1337 .try_row_chunk(row_start..row_end)
1338 .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
1339 block
1340 .slice_mut(s![row_start..row_end, col_start..col_end])
1341 .assign(&chunk);
1342 }
1343 col_start = col_end;
1344 }
1345 Ok(block)
1346}
1347
1348fn design_cross_relative_residual(
1349 lhs: &DesignMatrix,
1350 rhs: &DesignMatrix,
1351) -> Result<f64, BasisError> {
1352 let n = lhs.nrows();
1353 if rhs.nrows() != n {
1354 return Err(BasisError::ConstraintMatrixRowMismatch {
1355 basisrows: n,
1356 constraintrows: rhs.nrows(),
1357 });
1358 }
1359 const CHUNK: usize = 1024;
1360 let mut cross = Array2::<f64>::zeros((lhs.ncols(), rhs.ncols()));
1361 let mut lhs_sumsq = 0.0;
1362 let mut rhs_sumsq = 0.0;
1363 for start in (0..n).step_by(CHUNK) {
1364 let end = (start + CHUNK).min(n);
1365 let lhs_chunk = lhs
1366 .try_row_chunk(start..end)
1367 .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
1368 let rhs_chunk = rhs
1369 .try_row_chunk(start..end)
1370 .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
1371 cross += &gam_linalg::faer_ndarray::fast_atb(&lhs_chunk, &rhs_chunk);
1372 lhs_sumsq += lhs_chunk.iter().map(|v| v * v).sum::<f64>();
1373 rhs_sumsq += rhs_chunk.iter().map(|v| v * v).sum::<f64>();
1374 }
1375 let num = cross.iter().map(|v| v * v).sum::<f64>().sqrt();
1376 let denom = lhs_sumsq.sqrt() * rhs_sumsq.sqrt();
1377 if denom == 0.0 {
1380 return Ok(0.0);
1381 }
1382 Ok(num / denom)
1383}
1384
1385fn smooth_has_overlapping_linear_terms(
1386 linear_terms: &[LinearTermSpec],
1387 termspec: &SmoothTermSpec,
1388) -> bool {
1389 let feature_cols = smooth_term_feature_cols(termspec);
1390 linear_terms
1391 .iter()
1392 .any(|linear| feature_cols.contains(&linear.feature_col))
1393}
1394
1395pub fn smooth_intrinsic_parametric_feature_cols(
1399 linear_terms: &[LinearTermSpec],
1400 term: &SmoothTermSpec,
1401) -> Vec<usize> {
1402 let feature_cols = smooth_term_feature_cols(term);
1417 let mut owned = Vec::new();
1418 for linear in linear_terms {
1419 if feature_cols.contains(&linear.feature_col) && !owned.contains(&linear.feature_col) {
1420 owned.push(linear.feature_col);
1421 }
1422 }
1423 owned
1424}
1425
1426fn apply_global_smooth_identifiability(
1427 smooth: RawSmoothDesign,
1428 data: ArrayView2<'_, f64>,
1429 linear_terms: &[LinearTermSpec],
1430 smoothspecs: &[SmoothTermSpec],
1431) -> Result<(SmoothDesign, Array1<f64>), BasisError> {
1432 if smoothspecs.len() != smooth.terms.len() {
1443 gam_problem::bail_dim_basis!(
1444 "smooth spec count ({}) does not match built term count ({})",
1445 smoothspecs.len(),
1446 smooth.terms.len()
1447 );
1448 }
1449
1450 if smooth.terms.is_empty() {
1451 let RawSmoothDesign {
1452 term_designs,
1453 affine_offset,
1454 penalties,
1455 nullspace_dims,
1456 penaltyinfo,
1457 dropped_penaltyinfo,
1458 terms,
1459 coefficient_lower_bounds,
1460 linear_constraints,
1461 } = smooth;
1462 return Ok((
1463 SmoothDesign {
1464 term_designs,
1465 penalties,
1466 nullspace_dims,
1467 penaltyinfo,
1468 dropped_penaltyinfo,
1469 terms,
1470 coefficient_lower_bounds,
1471 linear_constraints,
1472 },
1473 affine_offset,
1474 ));
1475 }
1476
1477 let mut local_designs = vec![None; smooth.terms.len()];
1478 let mut local_active_penalties = vec![Vec::<ActivePenalty>::new(); smooth.terms.len()];
1479 let mut local_dropped_penalties = vec![Vec::<DroppedPenaltyInfo>::new(); smooth.terms.len()];
1480 let mut local_metadata = vec![None; smooth.terms.len()];
1481 let mut local_dims = vec![0usize; smooth.terms.len()];
1482 let mut local_linear_constraints = vec![None; smooth.terms.len()];
1483 let mut local_unabsorbed_z = vec![None::<Array2<f64>>; smooth.terms.len()];
1484 let mut local_residualization =
1485 vec![None::<ParametricResidualizationChart>; smooth.terms.len()];
1486 let mut local_collection_gauge = vec![None::<SmoothCollectionGauge>; smooth.terms.len()];
1487
1488 let SmoothStructureAnalysis {
1489 ownership_order,
1490 term_owners,
1491 ..
1492 } = analyze_smooth_ownership(smoothspecs);
1493
1494 use rayon::iter::{IntoParallelIterator, ParallelIterator};
1495
1496 for &idx in &ownership_order {
1497 let term = &smooth.terms[idx];
1498 let termspec = &smoothspecs[idx];
1499 let design_local = smooth.term_designs[idx].clone();
1500 let replay_z = frozen_global_orthogonality(termspec);
1508 let skip_global_transform = replay_z.is_none()
1509 && (smooth_has_frozen_identifiability(termspec) || term.lower_bounds_local.is_some());
1510 let owner_indices = if replay_z.is_some()
1524 || skip_global_transform
1525 || termspec.basis.is_marginally_centered_tensor()
1526 || termspec.basis.is_sum_to_zero_factor_smooth()
1527 {
1528 Vec::new()
1529 } else {
1530 const OVERLAP_REL_RESIDUAL_TOL: f64 = 1e-10;
1534 let owner_cross_checks = term_owners[idx]
1535 .clone()
1536 .into_par_iter()
1537 .map(|owner_idx| {
1538 let owner_design = local_designs[owner_idx]
1539 .as_ref()
1540 .expect("owner design must be available before dependent smooth");
1541 design_cross_relative_residual(&design_local, owner_design)
1542 .map(|rel| (owner_idx, rel))
1543 })
1544 .collect::<Vec<_>>();
1545 let mut out = Vec::new();
1546 for check in owner_cross_checks {
1547 let (owner_idx, rel) = check?;
1548 if rel > OVERLAP_REL_RESIDUAL_TOL {
1549 out.push(owner_idx);
1550 }
1551 }
1552 out
1553 };
1554 let owner_blocks = owner_indices
1555 .iter()
1556 .map(|owner_idx| {
1557 local_designs[*owner_idx]
1558 .as_ref()
1559 .expect("owner design must be available before dependent smooth")
1560 })
1561 .collect::<Vec<_>>();
1562 let replay_correction = frozen_parametric_residualization(termspec);
1567 let needs_parametric_block = match replay_correction {
1568 Some(chart) => chart.has_parametric_block,
1569 None => {
1570 replay_z.is_none()
1571 && !skip_global_transform
1572 && (smooth_has_overlapping_linear_terms(linear_terms, termspec)
1573 || !smooth_intrinsic_parametric_feature_cols(linear_terms, termspec)
1574 .is_empty()
1575 || smooth_requires_parametric_orthogonality(termspec)
1576 || factor_by_level_gate(termspec).is_some())
1582 }
1583 };
1584 let parametric_block = if !needs_parametric_block {
1585 None
1586 } else {
1587 Some(build_parametric_constraint_block_for_term(
1588 data,
1589 linear_terms,
1590 termspec,
1591 )?)
1592 };
1593 let replay_owner_blocks = match replay_correction {
1597 Some(chart) => chart
1598 .owner_terms
1599 .iter()
1600 .map(|owner_idx| {
1601 local_designs.get(*owner_idx).and_then(|slot| slot.as_ref()).ok_or_else(|| {
1602 BasisError::InvalidInput(format!(
1603 "term '{}' replays a parametric residualization against owner term {owner_idx}, which is not available at this point of the rebuild",
1604 termspec.name
1605 ))
1606 })
1607 })
1608 .collect::<Result<Vec<_>, _>>()?,
1609 None => Vec::new(),
1610 };
1611 let plan = if replay_correction.is_some() {
1637 GlobalIdentifiabilityPlan::Absent
1639 } else if skip_global_transform
1640 || (parametric_block.is_none() && owner_blocks.is_empty())
1641 {
1642 GlobalIdentifiabilityPlan::Absent
1643 } else {
1644 let raw =
1645 build_constraint_block(data.nrows(), parametric_block.as_ref(), &owner_blocks)?;
1646 let contained =
1647 crate::basis::contained_constraint_directions(&design_local, raw.view(), None)?;
1648 if raw.ncols() == 0 {
1649 GlobalIdentifiabilityPlan::Absent
1650 } else if contained.ncols() == raw.ncols() {
1651 GlobalIdentifiabilityPlan::Delete { block: contained }
1655 } else {
1656 GlobalIdentifiabilityPlan::Residualize { block: raw }
1657 }
1658 };
1659 let collection_coefficient_transform = match &plan {
1669 GlobalIdentifiabilityPlan::Absent => None,
1670 GlobalIdentifiabilityPlan::Delete { block } => Some(
1671 derive_smooth_collection_coefficient_transform(
1672 &design_local,
1673 SmoothCollectionGaugeArm::Delete,
1674 block.view(),
1675 !owner_indices.is_empty(),
1676 )?,
1677 ),
1678 GlobalIdentifiabilityPlan::Residualize { block } => Some(
1679 derive_smooth_collection_coefficient_transform(
1680 &design_local,
1681 SmoothCollectionGaugeArm::Residualize,
1682 block.view(),
1683 !owner_indices.is_empty(),
1684 )?,
1685 ),
1686 };
1687 let collection_gauge = collection_coefficient_transform.map(|transform| {
1688 plan.as_gauge(
1689 &owner_indices,
1690 parametric_block.is_some(),
1691 basis_local_identifiability_transform(&term.metadata),
1692 transform,
1693 design_local.ncols(),
1694 term.joint_null_rotation.clone(),
1699 )
1700 .expect("a derived collection coefficient chart implies a present gauge")
1701 });
1702 let mut residualization: Option<crate::basis::ParametricResidualization> = None;
1703 let (design_constrained, z_opt) = if let Some(gauge) = collection_gauge.as_ref() {
1704 let realized = realize_smooth_collection_gauge(design_local, gauge, &term.name)?;
1712 residualization = Some(realized.residualization);
1713 (realized.design, Some(realized.coefficient_transform))
1714 } else {
1715 let z_opt = if let Some(z) = replay_z {
1719 if design_local.ncols() != z.nrows() {
1720 gam_problem::bail_dim_basis!(
1721 "frozen global-orthogonality transform mismatch for term '{}': rebuilt design has {} columns but the persisted fit-time transform has {} rows",
1722 term.name,
1723 design_local.ncols(),
1724 z.nrows()
1725 );
1726 }
1727 Some(z.clone())
1728 } else if skip_global_transform {
1729 None
1730 } else {
1731 maybe_smooth_identifiability_transform(termspec, &design_local, None)?
1734 };
1735 let design_transformed = match z_opt.as_ref() {
1736 Some(z) => apply_smooth_transform_to_design(design_local, z, &term.name)?,
1737 None => design_local,
1738 };
1739 let design_constrained = match replay_correction {
1740 Some(chart) => {
1741 let block = build_constraint_block(
1746 data.nrows(),
1747 parametric_block.as_ref(),
1748 &replay_owner_blocks,
1749 )?;
1750 if block.ncols() != chart.correction.nrows() {
1751 gam_problem::bail_dim_basis!(
1752 "frozen parametric residualization mismatch for term '{}': rebuilt constraint block has {} columns but the persisted fit-time correction has {} rows",
1753 term.name,
1754 block.ncols(),
1755 chart.correction.nrows()
1756 );
1757 }
1758 subtract_row_space_correction(
1759 design_transformed,
1760 block.view(),
1761 chart.correction.view(),
1762 &term.name,
1763 )?
1764 }
1765 None => design_transformed,
1766 };
1767 (design_constrained, z_opt)
1768 };
1769 let coefficient_gauge = z_opt
1770 .as_ref()
1771 .map(|z| gam_problem::Gauge::from_block_transforms(&[z.clone()]));
1772
1773 let penalty_candidates = penalty_candidates_under_collection_gauge(
1774 &term.active_penalties,
1775 coefficient_gauge.as_ref(),
1776 &term.name,
1777 )?;
1778 let filtered = filter_penalty_candidates(penalty_candidates)?;
1779 let linear_constraints_constrained =
1780 if let Some(lin_local) = term.linear_constraints_local.as_ref() {
1781 if let Some(gauge) = coefficient_gauge.as_ref() {
1782 Some(LinearInequalityConstraints {
1783 a: lin_local.a.dot(&gauge.block_transform(0)),
1784 b: lin_local.b.clone(),
1785 })
1786 } else {
1787 Some(lin_local.clone())
1788 }
1789 } else {
1790 None
1791 };
1792
1793 local_residualization[idx] = residualization
1796 .as_ref()
1797 .map(|plan| ParametricResidualizationChart {
1798 owner_terms: owner_indices.clone(),
1799 has_parametric_block: parametric_block.is_some(),
1800 correction: plan.row_space_correction.clone(),
1801 })
1802 .or_else(|| replay_correction.cloned());
1803 local_collection_gauge[idx] = collection_gauge;
1804 local_dims[idx] = design_constrained.ncols();
1805 local_designs[idx] = Some(design_constrained);
1806 local_active_penalties[idx] = filtered.active;
1807 local_dropped_penalties[idx] = term.dropped_penalties.clone();
1808 local_dropped_penalties[idx].extend(filtered.dropped);
1809 local_linear_constraints[idx] = linear_constraints_constrained;
1810 let realized_transform = match (term.joint_null_rotation.as_ref(), z_opt.as_ref()) {
1811 (Some(rotation), Some(z)) => {
1812 Some(gam_linalg::faer_ndarray::fast_ab(&rotation.rotation, z))
1813 }
1814 (Some(rotation), None) => Some(rotation.rotation.clone()),
1815 (None, Some(z)) => Some(z.clone()),
1816 (None, None) => None,
1817 };
1818 match &termspec.basis {
1843 SmoothBasisSpec::FactorSumToZero { .. } | SmoothBasisSpec::FactorSmooth { .. } => {
1844 local_metadata[idx] = Some(term.metadata.clone());
1845 local_unabsorbed_z[idx] = z_opt.clone();
1846 }
1847 _ => {
1848 local_metadata[idx] = Some(with_identifiability_transform(
1849 &term.metadata,
1850 realized_transform.as_ref(),
1851 )?);
1852 }
1853 }
1854 }
1855
1856 let total_p: usize = local_dims.iter().sum();
1857 let mut terms_out = Vec::<SmoothTerm>::with_capacity(smooth.terms.len());
1858 let mut penalties_global = Vec::<BlockwisePenalty>::new();
1859 let mut nullspace_dims_global = Vec::<usize>::new();
1860 let mut penaltyinfo_global = Vec::<PenaltyBlockInfo>::new();
1861 let mut dropped_penaltyinfo_global = Vec::<DroppedPenaltyBlockInfo>::new();
1862 let mut coefficient_lower_bounds = Array1::<f64>::from_elem(total_p, f64::NEG_INFINITY);
1863 let mut any_bounds = false;
1864 let mut linear_constraintsrows: Vec<Array1<f64>> = Vec::new();
1865 let mut linear_constraints_b: Vec<f64> = Vec::new();
1866
1867 let mut col_start = 0usize;
1868 for idx in 0..smooth.terms.len() {
1869 let p_local = local_dims[idx];
1870 let col_end = col_start + p_local;
1871
1872 for active_penalty in &local_active_penalties[idx] {
1873 let global_index = penalties_global.len();
1874 penalties_global.push(BlockwisePenalty::new(
1875 col_start..col_end,
1876 active_penalty.matrix.clone(),
1877 ));
1878 nullspace_dims_global.push(active_penalty.nullity);
1879 penaltyinfo_global.push(PenaltyBlockInfo {
1880 global_index,
1881 termname: Some(smooth.terms[idx].name.clone()),
1882 penalty: active_penalty.info.clone(),
1883 });
1884 }
1885 for info in &local_dropped_penalties[idx] {
1886 dropped_penaltyinfo_global.push(DroppedPenaltyBlockInfo {
1887 termname: Some(smooth.terms[idx].name.clone()),
1888 penalty: info.clone(),
1889 });
1890 }
1891
1892 terms_out.push(SmoothTerm {
1893 name: smooth.terms[idx].name.clone(),
1894 coeff_range: col_start..col_end,
1895 shape: smooth.terms[idx].shape,
1896 active_penalties: local_active_penalties[idx].clone(),
1897 dropped_penalties: local_dropped_penalties[idx].clone(),
1898 metadata: local_metadata[idx]
1899 .clone()
1900 .expect("local metadata must exist for every smooth term"),
1901 lower_bounds_local: smooth.terms[idx].lower_bounds_local.clone(),
1902 linear_constraints_local: local_linear_constraints[idx].clone(),
1903 kronecker_factored: None,
1905 joint_null_rotation: None,
1911 unabsorbed_global_orthogonality: local_unabsorbed_z[idx].clone(),
1914 parametric_residualization: local_residualization[idx].clone(),
1915 collection_gauge: local_collection_gauge[idx].clone(),
1918 });
1919 if let Some(lin_local) = &local_linear_constraints[idx] {
1920 for r in 0..lin_local.a.nrows() {
1921 let mut row = Array1::<f64>::zeros(total_p);
1922 row.slice_mut(s![col_start..col_end])
1923 .assign(&lin_local.a.row(r));
1924 linear_constraintsrows.push(row);
1925 linear_constraints_b.push(lin_local.b[r]);
1926 }
1927 }
1928 if let Some(lb_local) = smooth.terms[idx].lower_bounds_local.as_ref()
1929 && lb_local.len() == p_local
1930 {
1931 coefficient_lower_bounds
1932 .slice_mut(s![col_start..col_end])
1933 .assign(lb_local);
1934 any_bounds = true;
1935 }
1936
1937 col_start = col_end;
1938 }
1939
1940 assert_eq!(
1941 penalties_global.len(),
1942 nullspace_dims_global.len(),
1943 "globally reparameterized smooth penalty/nullspace bookkeeping diverged"
1944 );
1945 assert_eq!(
1946 penalties_global.len(),
1947 penaltyinfo_global.len(),
1948 "globally reparameterized smooth penalty metadata bookkeeping diverged"
1949 );
1950
1951 Ok((
1952 SmoothDesign {
1953 term_designs: local_designs
1954 .into_iter()
1955 .map(|design| design.expect("local design must exist for every smooth term"))
1956 .collect(),
1957 penalties: penalties_global,
1958 nullspace_dims: nullspace_dims_global,
1959 penaltyinfo: penaltyinfo_global,
1960 dropped_penaltyinfo: dropped_penaltyinfo_global,
1961 terms: terms_out,
1962 coefficient_lower_bounds: if any_bounds {
1963 Some(coefficient_lower_bounds)
1964 } else {
1965 None
1966 },
1967 linear_constraints: if linear_constraintsrows.is_empty() {
1968 None
1969 } else {
1970 let mut a = Array2::<f64>::zeros((linear_constraintsrows.len(), total_p));
1971 for (i, row) in linear_constraintsrows.iter().enumerate() {
1972 a.row_mut(i).assign(row);
1973 }
1974 Some(LinearInequalityConstraints {
1975 a,
1976 b: Array1::from_vec(linear_constraints_b),
1977 })
1978 },
1979 },
1980 smooth.affine_offset,
1981 ))
1982}
1983
1984fn factor_by_level_gate(termspec: &SmoothTermSpec) -> Option<(usize, u64)> {
2007 match &termspec.basis {
2008 SmoothBasisSpec::ByVariable {
2009 by_col,
2010 by: ByVariableSpec::Level { value_bits, .. },
2011 ..
2012 } => Some((*by_col, *value_bits)),
2013 _ => None,
2014 }
2015}
2016
2017fn build_parametric_constraint_block_for_term(
2018 data: ArrayView2<'_, f64>,
2019 linear_terms: &[LinearTermSpec],
2020 termspec: &SmoothTermSpec,
2021) -> Result<Array2<f64>, BasisError> {
2022 let n = data.nrows();
2023 let p_data = data.ncols();
2024
2025 if let Some((by_col, value_bits)) = factor_by_level_gate(termspec) {
2029 if by_col >= p_data {
2030 gam_problem::bail_dim_basis!(
2031 "factor-by smooth term '{}' by column {by_col} out of bounds for {p_data} columns",
2032 termspec.name
2033 );
2034 }
2035 let mut c = Array2::<f64>::zeros((n, 1));
2036 let by = data.column(by_col);
2037 let value_bits = gam_data::canonical_level_bits(f64::from_bits(value_bits));
2038 for (row, &value) in by.iter().enumerate() {
2039 if gam_data::canonical_level_bits(value) == value_bits {
2040 c[[row, 0]] = 1.0;
2041 }
2042 }
2043 return Ok(c);
2044 }
2045
2046 let feature_cols = smooth_term_feature_cols(termspec);
2047 let mut parametric_cols = smooth_intrinsic_parametric_feature_cols(linear_terms, termspec);
2048 for &feature_col in ¶metric_cols {
2049 if feature_col >= p_data {
2050 gam_problem::bail_dim_basis!(
2051 "smooth term feature column {feature_col} out of bounds for {p_data} columns"
2052 );
2053 }
2054 }
2055 for linear in linear_terms
2056 .iter()
2057 .filter(|linear| feature_cols.contains(&linear.feature_col))
2058 {
2059 if linear.feature_col >= p_data {
2060 gam_problem::bail_dim_basis!(
2061 "linear term '{}' feature column {} out of bounds for {} columns",
2062 linear.name,
2063 linear.feature_col,
2064 p_data
2065 );
2066 }
2067 if !parametric_cols.contains(&linear.feature_col) {
2068 parametric_cols.push(linear.feature_col);
2069 }
2070 }
2071
2072 let mut c = Array2::<f64>::zeros((n, 1 + parametric_cols.len()));
2073 c.column_mut(0).fill(1.0);
2074 for (j, &feature_col) in parametric_cols.iter().enumerate() {
2075 c.column_mut(j + 1).assign(&data.column(feature_col));
2076 }
2077 Ok(c)
2078}
2079
2080pub fn apply_smooth_transform_to_design(
2081 design_local: DesignMatrix,
2082 transform: &Array2<f64>,
2083 termname: &str,
2084) -> Result<DesignMatrix, BasisError> {
2085 match design_local {
2086 DesignMatrix::Dense(inner) => {
2087 let op = CoefficientTransformOperator::new(inner, transform.clone()).map_err(|e| {
2088 BasisError::InvalidInput(format!(
2089 "smooth identifiability transform failed for term '{termname}': {e}"
2090 ))
2091 })?;
2092 Ok(DesignMatrix::Dense(
2093 gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(op)),
2094 ))
2095 }
2096 DesignMatrix::Sparse(inner) => {
2097 let dense = inner
2098 .try_to_dense_arc("smooth identifiability sparse transform")
2099 .map_err(BasisError::InvalidInput)?
2100 .as_ref()
2101 .dot(transform);
2102 Ok(DesignMatrix::Dense(
2103 gam_linalg::matrix::DenseDesignMatrix::from(dense),
2104 ))
2105 }
2106 }
2107}
2108
2109fn design_constraint_cross(
2110 design: &DesignMatrix,
2111 constraint_matrix: ArrayView2<'_, f64>,
2112) -> Result<Array2<f64>, BasisError> {
2113 let n = design.nrows();
2114 if constraint_matrix.nrows() != n {
2115 return Err(BasisError::ConstraintMatrixRowMismatch {
2116 basisrows: n,
2117 constraintrows: constraint_matrix.nrows(),
2118 });
2119 }
2120 let mut cross = Array2::<f64>::zeros((design.ncols(), constraint_matrix.ncols()));
2121 const CHUNK: usize = 1024;
2122 for start in (0..n).step_by(CHUNK) {
2123 let end = (start + CHUNK).min(n);
2124 let design_chunk = design
2125 .try_row_chunk(start..end)
2126 .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
2127 let constraint_chunk = constraint_matrix.slice(s![start..end, ..]).to_owned();
2128 cross += &gam_linalg::faer_ndarray::fast_atb(&design_chunk, &constraint_chunk);
2129 }
2130 Ok(cross)
2131}
2132
2133fn design_frobenius_norm(design: &DesignMatrix) -> Result<f64, BasisError> {
2134 let n = design.nrows();
2135 const CHUNK: usize = 1024;
2136 let mut sumsq = 0.0;
2137 for start in (0..n).step_by(CHUNK) {
2138 let end = (start + CHUNK).min(n);
2139 let chunk = design
2140 .try_row_chunk(start..end)
2141 .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
2142 sumsq += chunk.iter().map(|v| v * v).sum::<f64>();
2143 }
2144 Ok(sumsq.sqrt())
2145}
2146
2147fn frozen_parametric_residualization(
2155 termspec: &SmoothTermSpec,
2156) -> Option<&ParametricResidualizationChart> {
2157 termspec.frozen_parametric_residualization.as_ref()
2158}
2159
2160fn frozen_global_orthogonality(termspec: &SmoothTermSpec) -> Option<&Array2<f64>> {
2166 match &termspec.basis {
2167 SmoothBasisSpec::FactorSumToZero {
2168 frozen_global_orthogonality,
2169 ..
2170 } => frozen_global_orthogonality.as_ref(),
2171 SmoothBasisSpec::FactorSmooth { spec } => spec.frozen_global_orthogonality.as_ref(),
2172 _ => None,
2173 }
2174}
2175
2176fn penalty_candidates_under_collection_gauge(
2189 active_penalties: &[ActivePenalty],
2190 coefficient_gauge: Option<&gam_problem::Gauge>,
2191 term_name: &str,
2192) -> Result<Vec<PenaltyCandidate>, BasisError> {
2193 use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
2194 let penalty_candidates = active_penalties
2195 .par_iter()
2196 .map(|penalty| -> Result<PenaltyCandidate, BasisError> {
2197 let raw = ConstructiveQuadratic::try_from_dense_psd(
2198 penalty.matrix.clone(),
2199 "global smooth source penalty",
2200 )?;
2201 let raw = match penalty.info.structural_null_frame.as_ref() {
2209 Some(frame) => raw.with_structural_null_frame(
2210 frame.clone(),
2211 "global smooth source penalty structural frame",
2212 )?,
2213 None => raw,
2214 };
2215 let restricted = if let Some(gauge) = coefficient_gauge {
2216 raw.restricted(gauge, "global smooth identifiability restriction")?
2217 } else {
2218 raw
2219 };
2220 let (_, c_new) = normalize_penalty_in_constrained_space(restricted.dense());
2221 let matrix = restricted.scaled(1.0 / c_new, "normalized global smooth penalty")?;
2222 Ok(PenaltyCandidate {
2223 matrix,
2224 source: penalty.info.source.clone(),
2225 normalization_scale: penalty.info.normalization_scale * c_new,
2226 kronecker_factors: None,
2227 op: None,
2228 })
2229 })
2230 .collect::<Result<Vec<_>, _>>()?;
2231 let mut penalty_candidates = penalty_candidates;
2251 if coefficient_gauge.is_some()
2252 && penalty_candidates
2253 .iter()
2254 .any(|c| matches!(c.source, PenaltySource::DoublePenaltyNullspace))
2255 {
2256 const SUPPORT_TOL: f64 = 0.0;
2265 let support_rows = |m: &Array2<f64>| -> (usize, usize) {
2266 let n = m.nrows();
2267 let mut lo = n;
2268 let mut hi = 0usize;
2269 for i in 0..n {
2270 let any = (0..m.ncols()).any(|j| m[[i, j]].abs() > SUPPORT_TOL);
2271 if any {
2272 lo = lo.min(i);
2273 hi = hi.max(i + 1);
2274 }
2275 }
2276 (lo, hi)
2277 };
2278 let primaries: Vec<((usize, usize), ConstructiveQuadratic)> = penalty_candidates
2281 .iter()
2282 .filter(|c| matches!(c.source, PenaltySource::Primary))
2283 .map(|c| -> Result<_, BasisError> {
2284 Ok((
2285 support_rows(&c.matrix),
2286 c.matrix
2287 .scaled(c.normalization_scale, "physical global smooth primary")?,
2288 ))
2289 })
2290 .collect::<Result<Vec<_>, _>>()?;
2291 for candidate in &mut penalty_candidates {
2292 if !matches!(candidate.source, PenaltySource::DoublePenaltyNullspace) {
2293 continue;
2294 }
2295 let q = candidate.matrix.nrows();
2296 let (rlo, rhi) = support_rows(&candidate.matrix);
2297 let owner = primaries
2301 .iter()
2302 .find(|((plo, phi), _)| *plo <= rlo && rhi <= *phi)
2303 .or_else(|| (primaries.len() == 1).then(|| &primaries[0]))
2304 .ok_or_else(|| {
2305 BasisError::InvalidInput(format!(
2306 "double-penalty ridge for smooth '{}' has no co-located primary penalty",
2307 term_name
2308 ))
2309 })?;
2310 let ((plo, phi), s_full) = owner;
2311 let block = ConstructiveQuadratic::from_energy_factor(
2317 s_full.factor().slice(s![.., *plo..*phi]).to_owned(),
2318 "owned global smooth primary block",
2319 )?;
2320 let block = match s_full.structural_null_frame_block(*plo, *phi) {
2326 Some(frame) => block.with_structural_null_frame(
2327 frame,
2328 "owned global smooth primary block structural frame",
2329 )?,
2330 None => block,
2331 };
2332 let ridge_full = candidate.matrix.scaled(
2333 candidate.normalization_scale,
2334 "physical global smooth null ridge",
2335 )?;
2336 let ridge_block = ConstructiveQuadratic::from_energy_factor(
2337 ridge_full.factor().slice(s![.., *plo..*phi]).to_owned(),
2338 "owned global smooth null-ridge block",
2339 )?;
2340 let rebuilt_block =
2341 crate::basis::rebuild_metric_consistent_ridge(&block, &ridge_block)?;
2342 match rebuilt_block {
2343 Some(ridge_block) => {
2344 let mut full_factor =
2345 Array2::<f64>::zeros((ridge_block.factor().nrows(), q));
2346 full_factor
2347 .slice_mut(s![.., *plo..*phi])
2348 .assign(ridge_block.factor());
2349 let full = ConstructiveQuadratic::from_energy_factor(
2350 full_factor,
2351 "embedded global smooth null ridge",
2352 )?;
2353 let (_, scale) = normalize_penalty_in_constrained_space(full.dense());
2354 candidate.matrix = full
2355 .scaled(1.0 / scale, "normalized embedded global smooth null ridge")?;
2356 candidate.normalization_scale = scale;
2357 candidate.kronecker_factors = None;
2358 candidate.op = None;
2359 }
2360 None => {
2363 candidate.matrix = ConstructiveQuadratic::zero(q);
2364 candidate.normalization_scale = 1.0;
2365 candidate.kronecker_factors = None;
2366 candidate.op = None;
2367 }
2368 }
2369 }
2370 }
2371 Ok(penalty_candidates)
2372}
2373
2374fn maybe_smooth_identifiability_transform(
2375 termspec: &SmoothTermSpec,
2376 design_local: &DesignMatrix,
2377 constraint_block: Option<ArrayView2<'_, f64>>,
2378) -> Result<Option<Array2<f64>>, BasisError> {
2379 if let Some(SpatialIdentifiability::FrozenTransform { transform }) =
2380 spatial_identifiability_policy(termspec)
2381 {
2382 if design_local.ncols() != transform.nrows() {
2383 gam_problem::bail_dim_basis!(
2384 "frozen spatial identifiability transform mismatch: design has {} columns but transform has {} rows",
2385 design_local.ncols(),
2386 transform.nrows()
2387 );
2388 }
2389 return Ok(Some(transform.clone()));
2390 }
2391
2392 if let Some(c) = constraint_block {
2393 if c.ncols() == 0 {
2394 Ok(None)
2395 } else {
2396 Ok(Some(orthogonality_transform_for_design(
2397 design_local,
2398 c,
2399 None, )?))
2401 }
2402 } else {
2403 Ok(None)
2404 }
2405}
2406
2407fn smooth_requires_parametric_orthogonality(termspec: &SmoothTermSpec) -> bool {
2476 match &termspec.basis {
2477 SmoothBasisSpec::ByVariable { inner, .. }
2478 | SmoothBasisSpec::FactorSumToZero { inner, .. } => {
2479 smooth_requires_parametric_orthogonality(&SmoothTermSpec {
2480 frozen_parametric_residualization: None,
2481 name: termspec.name.clone(),
2482 basis: (**inner).clone(),
2483 shape: termspec.shape,
2484 joint_null_rotation: None,
2485 })
2486 }
2487 SmoothBasisSpec::BySmooth { smooth, .. } => {
2488 smooth_requires_parametric_orthogonality(&SmoothTermSpec {
2489 frozen_parametric_residualization: None,
2490 name: termspec.name.clone(),
2491 basis: (**smooth).clone(),
2492 shape: termspec.shape,
2493 joint_null_rotation: None,
2494 })
2495 }
2496 SmoothBasisSpec::ThinPlate { spec, .. } => {
2497 matches!(
2498 spec.identifiability,
2499 SpatialIdentifiability::OrthogonalToParametric
2500 )
2501 }
2502 SmoothBasisSpec::Duchon { spec, .. } => {
2503 matches!(
2504 spec.identifiability,
2505 SpatialIdentifiability::OrthogonalToParametric
2506 )
2507 }
2508 SmoothBasisSpec::Matern { spec, .. } => matches!(
2509 spec.identifiability,
2510 MaternIdentifiability::CenterSumToZero | MaternIdentifiability::CenterLinearOrthogonal
2511 ),
2512 SmoothBasisSpec::Sphere { spec, .. } => {
2520 matches!(spec.method, crate::basis::SphereMethod::Wahba)
2521 && !matches!(spec.wahba_kernel, crate::basis::SphereWahbaKernel::Pseudo)
2522 && matches!(
2523 spec.identifiability,
2524 SphericalSplineIdentifiability::CenterSumToZero
2525 )
2526 }
2527 SmoothBasisSpec::ConstantCurvature { spec, .. } => matches!(
2535 spec.identifiability,
2536 ConstantCurvatureIdentifiability::CenterSumToZero
2537 ),
2538 SmoothBasisSpec::MeasureJet { spec, .. } => matches!(
2544 spec.identifiability,
2545 MeasureJetIdentifiability::CenterSumToZero
2546 ),
2547 SmoothBasisSpec::BSpline1D { .. }
2548 | SmoothBasisSpec::TensorBSpline { .. }
2549 | SmoothBasisSpec::Pca { .. }
2550 | SmoothBasisSpec::FactorSmooth { .. } => false,
2551 }
2552}
2553
2554fn compose_identifiability_transforms(
2555 existing: Option<&Array2<f64>>,
2556 extra: Option<&Array2<f64>>,
2557) -> Result<Option<Array2<f64>>, BasisError> {
2558 match (existing, extra) {
2559 (Some(lhs), Some(rhs)) => {
2560 if lhs.ncols() == rhs.nrows() {
2561 Ok(Some(lhs.dot(rhs)))
2562 } else if lhs.nrows() == rhs.nrows() && lhs.ncols() == rhs.ncols() {
2563 Ok(Some(rhs.clone()))
2567 } else {
2568 Err(BasisError::DimensionMismatch(format!(
2569 "identifiability transform mismatch: existing is {}x{}, extra is {}x{}",
2570 lhs.nrows(),
2571 lhs.ncols(),
2572 rhs.nrows(),
2573 rhs.ncols(),
2574 )))
2575 }
2576 }
2577 (Some(lhs), None) => Ok(Some(lhs.clone())),
2578 (None, Some(rhs)) => Ok(Some(rhs.clone())),
2579 (None, None) => Ok(None),
2580 }
2581}
2582
2583fn with_identifiability_transform(
2584 metadata: &BasisMetadata,
2585 transform: Option<&Array2<f64>>,
2586) -> Result<BasisMetadata, BasisError> {
2587 match metadata {
2588 BasisMetadata::BSpline1D {
2589 knots,
2590 identifiability_transform,
2591 periodic,
2592 degree,
2593 auto_shrink_note,
2594 anchor_offset_coeffs,
2595 } => Ok(BasisMetadata::BSpline1D {
2596 knots: knots.clone(),
2597 periodic: *periodic,
2598 identifiability_transform: compose_identifiability_transforms(
2599 identifiability_transform.as_ref(),
2600 transform,
2601 )?,
2602 degree: *degree,
2603 auto_shrink_note: auto_shrink_note.clone(),
2604 anchor_offset_coeffs: anchor_offset_coeffs.clone(),
2608 }),
2609 BasisMetadata::CubicRegression1D {
2610 knots,
2611 identifiability_transform,
2612 } => Ok(BasisMetadata::CubicRegression1D {
2613 knots: knots.clone(),
2614 identifiability_transform: compose_identifiability_transforms(
2615 identifiability_transform.as_ref(),
2616 transform,
2617 )?,
2618 }),
2619 BasisMetadata::ThinPlate {
2620 centers,
2621 length_scale,
2622 periodic,
2623 identifiability_transform,
2624 input_scale,
2625 radial_reparam,
2626 } => Ok(BasisMetadata::ThinPlate {
2627 centers: centers.clone(),
2628 length_scale: *length_scale,
2629 periodic: periodic.clone(),
2630 identifiability_transform: compose_identifiability_transforms(
2631 identifiability_transform.as_ref(),
2632 transform,
2633 )?,
2634 input_scale: *input_scale,
2635 radial_reparam: radial_reparam.clone(),
2636 }),
2637 BasisMetadata::Sphere {
2638 centers,
2639 penalty_order,
2640 method,
2641 max_degree,
2642 wahba_kernel,
2643 constraint_transform,
2644 } => Ok(BasisMetadata::Sphere {
2645 centers: centers.clone(),
2646 penalty_order: *penalty_order,
2647 method: *method,
2648 max_degree: *max_degree,
2649 wahba_kernel: *wahba_kernel,
2650 constraint_transform: compose_identifiability_transforms(
2651 constraint_transform.as_ref(),
2652 transform,
2653 )?,
2654 }),
2655 BasisMetadata::ConstantCurvature {
2656 centers,
2657 kappa,
2658 length_scale,
2659 constraint_transform,
2660 } => Ok(BasisMetadata::ConstantCurvature {
2661 centers: centers.clone(),
2662 kappa: *kappa,
2663 length_scale: *length_scale,
2664 constraint_transform: compose_identifiability_transforms(
2665 constraint_transform.as_ref(),
2666 transform,
2667 )?,
2668 }),
2669 BasisMetadata::MeasureJet {
2670 centers,
2671 input_scale,
2672 length_scale,
2673 eps_band,
2674 order_s,
2675 alpha,
2676 tau0,
2677 masses,
2678 support_means,
2679 penalty_normalization_scales,
2680 raw_penalty_normalization_scales,
2681 fused_penalty_normalization_scale,
2682 constraint_transform,
2683 sigma_coord,
2684 } => Ok(BasisMetadata::MeasureJet {
2685 centers: centers.clone(),
2686 input_scale: *input_scale,
2687 length_scale: *length_scale,
2688 eps_band: eps_band.clone(),
2689 order_s: *order_s,
2690 alpha: *alpha,
2691 tau0: *tau0,
2692 masses: masses.clone(),
2693 support_means: support_means.clone(),
2694 penalty_normalization_scales: penalty_normalization_scales.clone(),
2695 raw_penalty_normalization_scales: raw_penalty_normalization_scales.clone(),
2696 fused_penalty_normalization_scale: *fused_penalty_normalization_scale,
2697 constraint_transform: compose_identifiability_transforms(
2698 constraint_transform.as_ref(),
2699 transform,
2700 )?,
2701 sigma_coord: *sigma_coord,
2702 }),
2703 BasisMetadata::Matern {
2704 centers,
2705 length_scale,
2706 periodic,
2707 nu,
2708 include_intercept,
2709 identifiability_transform,
2710 input_scale,
2711 aniso_log_scales,
2712 } => Ok(BasisMetadata::Matern {
2713 centers: centers.clone(),
2714 length_scale: *length_scale,
2715 periodic: periodic.clone(),
2716 nu: *nu,
2717 include_intercept: *include_intercept,
2718 identifiability_transform: compose_identifiability_transforms(
2719 identifiability_transform.as_ref(),
2720 transform,
2721 )?,
2722 input_scale: *input_scale,
2723 aniso_log_scales: aniso_log_scales.clone(),
2724 }),
2725 BasisMetadata::Duchon {
2726 centers,
2727 length_scale,
2728 periodic,
2729 power,
2730 nullspace_order,
2731 identifiability_transform,
2732 input_scale,
2733 aniso_log_scales,
2734 operator_collocation_points,
2735 radial_reparam,
2736 spectral_basis,
2737 } => Ok(BasisMetadata::Duchon {
2738 centers: centers.clone(),
2739 length_scale: *length_scale,
2740 periodic: periodic.clone(),
2741 power: *power,
2742 nullspace_order: *nullspace_order,
2743 input_scale: *input_scale,
2744 aniso_log_scales: aniso_log_scales.clone(),
2745 operator_collocation_points: operator_collocation_points.clone(),
2746 radial_reparam: radial_reparam.clone(),
2747 spectral_basis: spectral_basis.clone(),
2748 identifiability_transform: compose_identifiability_transforms(
2749 identifiability_transform.as_ref(),
2750 transform,
2751 )?,
2752 }),
2753 BasisMetadata::SphereHarmonics {
2754 max_degree,
2755 radians,
2756 } => Ok(BasisMetadata::SphereHarmonics {
2757 max_degree: *max_degree,
2758 radians: *radians,
2759 }),
2760 BasisMetadata::TensorBSpline {
2761 feature_cols,
2762 knots,
2763 degrees,
2764 periods,
2765 is_cr,
2766 identifiability_transform,
2767 } => Ok(BasisMetadata::TensorBSpline {
2768 feature_cols: feature_cols.clone(),
2769 knots: knots.clone(),
2770 degrees: degrees.clone(),
2771 periods: periods.clone(),
2772 is_cr: is_cr.clone(),
2773 identifiability_transform: compose_identifiability_transforms(
2774 identifiability_transform.as_ref(),
2775 transform,
2776 )?,
2777 }),
2778 BasisMetadata::BySmooth {
2779 inner,
2780 by_col,
2781 levels,
2782 ordered,
2783 } => Ok(BasisMetadata::BySmooth {
2784 inner: Box::new(with_identifiability_transform(inner, transform)?),
2785 by_col: *by_col,
2786 levels: levels.clone(),
2787 ordered: *ordered,
2788 }),
2789 BasisMetadata::FactorSmooth {
2790 continuous_cols,
2791 group_col,
2792 knots,
2793 degree,
2794 periodic,
2795 group_levels,
2796 flavour,
2797 marginal_is_cr,
2798 } => {
2799 if transform.is_some() {
2806 gam_problem::bail_invalid_basis!(
2807 "FactorSmooth metadata cannot absorb an identifiability transform; \
2808 route it through the term-level frozen_global_orthogonality carrier"
2809 );
2810 }
2811 Ok(BasisMetadata::FactorSmooth {
2812 continuous_cols: continuous_cols.clone(),
2813 group_col: *group_col,
2814 knots: knots.clone(),
2815 degree: *degree,
2816 periodic: *periodic,
2817 group_levels: group_levels.clone(),
2818 flavour: flavour.clone(),
2819 marginal_is_cr: *marginal_is_cr,
2820 })
2821 }
2822 BasisMetadata::Pca {
2823 feature_cols,
2824 basis_matrix,
2825 centered,
2826 smooth_penalty,
2827 center_mean,
2828 pca_basis_path,
2829 chunk_size,
2830 } => {
2831 if transform.is_some() {
2837 gam_problem::bail_invalid_basis!(
2838 "PCA bases do not expose a composable identifiability transform"
2839 );
2840 }
2841 Ok(BasisMetadata::Pca {
2842 feature_cols: feature_cols.clone(),
2843 basis_matrix: basis_matrix.clone(),
2844 centered: *centered,
2845 smooth_penalty: *smooth_penalty,
2846 center_mean: center_mean.clone(),
2847 pca_basis_path: pca_basis_path.clone(),
2848 chunk_size: *chunk_size,
2849 })
2850 }
2851 }
2852}
2853
2854pub fn orthogonality_relative_residual_for_design(
2858 design: &DesignMatrix,
2859 constraint_matrix: ArrayView2<'_, f64>,
2860) -> Result<f64, BasisError> {
2861 let cross = design_constraint_cross(design, constraint_matrix)?;
2862 let num = cross.iter().map(|v| v * v).sum::<f64>().sqrt();
2863 let b_norm = design_frobenius_norm(design)?;
2864 let c_norm = constraint_matrix.iter().map(|v| v * v).sum::<f64>().sqrt();
2865 let denom = b_norm * c_norm;
2866 if denom == 0.0 {
2868 return Ok(0.0);
2869 }
2870 Ok(num / denom)
2871}
2872
2873#[cfg(test)]
2874mod frozen_linear_term_mass_rebuild_tests {
2875 use super::*;
2876
2877 fn one_linear_term_spec() -> TermCollectionSpec {
2881 TermCollectionSpec {
2882 linear_terms: vec![LinearTermSpec {
2883 name: "x".to_string(),
2884 feature_col: 0,
2885 feature_cols: vec![0],
2886 categorical_levels: vec![],
2887 double_penalty: true,
2888 coefficient_geometry: LinearCoefficientGeometry::Unconstrained,
2889 coefficient_min: None,
2890 coefficient_max: None,
2891 frozen_function_mass: None,
2892 }],
2893 random_effect_terms: Vec::new(),
2894 smooth_terms: Vec::new(),
2895 }
2896 }
2897
2898 fn training_data_varying_x(n: usize) -> Array2<f64> {
2899 let mut data = Array2::<f64>::zeros((n, 1));
2900 for i in 0..n {
2901 data[[i, 0]] = 1.0 + i as f64;
2903 }
2904 data
2905 }
2906
2907 fn constant_zero_x(n_rows: usize) -> Array2<f64> {
2908 Array2::<f64>::zeros((n_rows, 1))
2909 }
2910
2911 #[test]
2917 fn unfrozen_spec_still_rejects_a_genuinely_zero_training_column() {
2918 let spec = one_linear_term_spec();
2919 let degenerate_training_data = constant_zero_x(20);
2920 let err = build_term_collection_design(degenerate_training_data.view(), &spec)
2921 .expect_err("an unfrozen spec fit directly on an all-zero column must still fail");
2922 let message = err.to_string();
2923 assert!(
2924 message.contains("identically zero"),
2925 "expected the identifiability guard's message, got: {message}"
2926 );
2927 }
2928
2929 #[test]
2940 fn frozen_spec_rebuilds_at_a_constant_evaluation_column_using_the_training_mass() {
2941 let spec = one_linear_term_spec();
2942 let training_data = training_data_varying_x(40);
2943
2944 let training_design = build_term_collection_design(training_data.view(), &spec)
2945 .expect("fit-time build over a genuinely varying column must succeed");
2946 let training_mass = training_design
2947 .linear_function_masses
2948 .first()
2949 .copied()
2950 .flatten()
2951 .expect("a double_penalty=true term must report its fit-time function mass");
2952 assert!(
2953 training_mass > 0.0,
2954 "training mass for a genuinely varying column must be strictly positive, got {training_mass}"
2955 );
2956
2957 let frozen_spec = freeze_term_collection_from_design(&spec, &training_design)
2958 .expect("freezing the spec against its own fit-time design must succeed");
2959 assert_eq!(
2960 frozen_spec.linear_terms[0].frozen_function_mass,
2961 Some(training_mass),
2962 "freezing must persist the exact fit-time mass onto the term"
2963 );
2964
2965 let evaluation_grid = constant_zero_x(3);
2969 let rebuilt_design = build_term_collection_design(evaluation_grid.view(), &frozen_spec)
2970 .expect(
2971 "rebuilding a FROZEN spec's design at a constant-covariate evaluation grid must \
2972 succeed — the training-time mass is reused, never recomputed from these rows",
2973 );
2974 assert_eq!(
2975 rebuilt_design
2976 .linear_function_masses
2977 .first()
2978 .copied()
2979 .flatten(),
2980 Some(training_mass),
2981 "the rebuilt design must carry the REUSED training-time mass, not a value \
2982 recomputed from the (all-zero) evaluation rows"
2983 );
2984 }
2985}