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 local_columns: usize,
940 ) -> Option<SmoothCollectionGauge> {
941 let (arm, block) = match self {
942 Self::Absent => return None,
943 Self::Delete { block } => (SmoothCollectionGaugeArm::Delete, block),
944 Self::Residualize { block } => (SmoothCollectionGaugeArm::Residualize, block),
945 };
946 Some(SmoothCollectionGauge {
947 arm,
948 constraint_block: block.clone(),
949 owner_terms: owner_terms.to_vec(),
950 has_parametric_block,
951 local_identifiability_transform,
952 local_columns,
953 })
954 }
955}
956
957fn basis_local_identifiability_transform(metadata: &BasisMetadata) -> Option<Array2<f64>> {
969 match metadata {
970 BasisMetadata::BSpline1D {
971 identifiability_transform,
972 ..
973 }
974 | BasisMetadata::CubicRegression1D {
975 identifiability_transform,
976 ..
977 }
978 | BasisMetadata::ThinPlate {
979 identifiability_transform,
980 ..
981 }
982 | BasisMetadata::Matern {
983 identifiability_transform,
984 ..
985 }
986 | BasisMetadata::Duchon {
987 identifiability_transform,
988 ..
989 }
990 | BasisMetadata::TensorBSpline {
991 identifiability_transform,
992 ..
993 } => identifiability_transform.clone(),
994 BasisMetadata::Sphere {
995 constraint_transform,
996 ..
997 }
998 | BasisMetadata::ConstantCurvature {
999 constraint_transform,
1000 ..
1001 }
1002 | BasisMetadata::MeasureJet {
1003 constraint_transform,
1004 ..
1005 } => constraint_transform.clone(),
1006 BasisMetadata::Pca { .. }
1007 | BasisMetadata::SphereHarmonics { .. }
1008 | BasisMetadata::BySmooth { .. }
1009 | BasisMetadata::FactorSmooth { .. } => None,
1010 }
1011}
1012
1013pub struct RealizedCollectionGauge {
1015 pub design: DesignMatrix,
1017 pub coefficient_transform: Array2<f64>,
1020 pub residualization: Option<crate::basis::ParametricResidualization>,
1022}
1023
1024pub fn realize_smooth_collection_gauge(
1042 design_local: DesignMatrix,
1043 gauge: &SmoothCollectionGauge,
1044 termname: &str,
1045) -> Result<RealizedCollectionGauge, BasisError> {
1046 let block = gauge.constraint_block.view();
1047 if block.nrows() != design_local.nrows() {
1048 gam_problem::bail_dim_basis!(
1049 "collection gauge row mismatch for term '{termname}': the design has {} rows and the frozen constraint block has {}",
1050 design_local.nrows(),
1051 block.nrows()
1052 );
1053 }
1054 let (design, coefficient_transform, residualization) = match gauge.arm {
1055 SmoothCollectionGaugeArm::Delete => {
1056 let z = match orthogonality_transform_for_design(&design_local, block, None) {
1057 Ok(z) => z,
1058 Err(BasisError::ConstraintNullspaceCollapsed { .. })
1063 if !gauge.owner_terms.is_empty() =>
1064 {
1065 Array2::zeros((design_local.ncols(), 0))
1066 }
1067 Err(err) => return Err(err),
1068 };
1069 let design = apply_smooth_transform_to_design(design_local, &z, termname)?;
1070 (design, z, None)
1071 }
1072 SmoothCollectionGaugeArm::Residualize => {
1073 let plan = crate::basis::parametric_residualization_for_design(
1074 &design_local,
1075 block,
1076 None, )?;
1078 let transform = plan.coefficient_transform.clone();
1079 let design = apply_smooth_transform_to_design(design_local, &transform, termname)?;
1080 let design = subtract_row_space_correction(
1081 design,
1082 block,
1083 plan.row_space_correction.view(),
1084 termname,
1085 )?;
1086 (design, transform, Some(plan))
1087 }
1088 };
1089 assert_orthogonal_to_constraint_block(&design, block, termname)?;
1090 Ok(RealizedCollectionGauge {
1091 design,
1092 coefficient_transform,
1093 residualization,
1094 })
1095}
1096
1097pub struct LocalTermRealization<'a> {
1102 pub design: DesignMatrix,
1104 pub metadata: &'a BasisMetadata,
1106 pub active_penalties: &'a [ActivePenalty],
1107 pub dropped_penalties: Vec<DroppedPenaltyInfo>,
1108 pub linear_constraints_local: Option<&'a gam_problem::LinearInequalityConstraints>,
1109 pub joint_null_rotation: Option<&'a crate::basis::JointNullRotation>,
1112 pub termname: &'a str,
1113}
1114
1115pub struct CollectionGaugedTerm {
1117 pub design: DesignMatrix,
1118 pub metadata: BasisMetadata,
1119 pub active_penalties: Vec<ActivePenalty>,
1120 pub dropped_penalties: Vec<DroppedPenaltyInfo>,
1121 pub linear_constraints_local: Option<gam_problem::LinearInequalityConstraints>,
1122 pub parametric_residualization: Option<ParametricResidualizationChart>,
1123}
1124
1125pub fn place_term_in_collection_gauge(
1145 gauge: &SmoothCollectionGauge,
1146 local: LocalTermRealization<'_>,
1147) -> Result<CollectionGaugedTerm, BasisError> {
1148 let LocalTermRealization {
1149 design,
1150 metadata,
1151 active_penalties,
1152 dropped_penalties,
1153 linear_constraints_local,
1154 joint_null_rotation,
1155 termname,
1156 } = local;
1157 let realized = realize_smooth_collection_gauge(design, gauge, termname)?;
1158 let coefficient_gauge =
1159 gam_problem::Gauge::from_block_transforms(&[realized.coefficient_transform.clone()]);
1160 let candidates = penalty_candidates_under_collection_gauge(
1161 active_penalties,
1162 Some(&coefficient_gauge),
1163 termname,
1164 )?;
1165 let filtered = filter_penalty_candidates(candidates)?;
1166 let mut dropped_penalties = dropped_penalties;
1167 dropped_penalties.extend(filtered.dropped);
1168 let linear_constraints_local = linear_constraints_local.map(|lin| {
1169 gam_problem::LinearInequalityConstraints {
1170 a: lin.a.dot(&coefficient_gauge.block_transform(0)),
1171 b: lin.b.clone(),
1172 }
1173 });
1174 let realized_transform = match joint_null_rotation {
1175 Some(rotation) => {
1176 gam_linalg::faer_ndarray::fast_ab(&rotation.rotation, &realized.coefficient_transform)
1177 }
1178 None => realized.coefficient_transform.clone(),
1179 };
1180 let metadata = with_identifiability_transform(metadata, Some(&realized_transform))?;
1181 let parametric_residualization =
1182 realized
1183 .residualization
1184 .as_ref()
1185 .map(|plan| ParametricResidualizationChart {
1186 owner_terms: gauge.owner_terms.clone(),
1187 has_parametric_block: gauge.has_parametric_block,
1188 correction: plan.row_space_correction.clone(),
1189 });
1190 Ok(CollectionGaugedTerm {
1191 design: realized.design,
1192 metadata,
1193 active_penalties: filtered.active,
1194 dropped_penalties,
1195 linear_constraints_local,
1196 parametric_residualization,
1197 })
1198}
1199
1200const ORTHOGONALITY_REL_RESIDUAL_TOL: f64 = 1e-8;
1203
1204fn assert_orthogonal_to_constraint_block(
1205 design: &DesignMatrix,
1206 constraint: ArrayView2<'_, f64>,
1207 termname: &str,
1208) -> Result<(), BasisError> {
1209 let rel = orthogonality_relative_residual_for_design(design, constraint)?;
1210 if rel > ORTHOGONALITY_REL_RESIDUAL_TOL {
1211 gam_problem::bail_invalid_basis!(
1212 "smooth orthogonality residual too large for term '{}': {:.3e} > {:.1e}",
1213 termname,
1214 rel,
1215 ORTHOGONALITY_REL_RESIDUAL_TOL
1216 );
1217 }
1218 Ok(())
1219}
1220
1221fn subtract_row_space_correction(
1226 design: DesignMatrix,
1227 constraint: ArrayView2<'_, f64>,
1228 correction: ArrayView2<'_, f64>,
1229 termname: &str,
1230) -> Result<DesignMatrix, BasisError> {
1231 use gam_linalg::matrix::{BlockDesignOperator, DesignBlock};
1232 let p = design.ncols();
1233 let q = constraint.ncols();
1234 let k = correction.ncols();
1235 if correction.nrows() != q || p != k {
1236 return Err(BasisError::InvalidInput(format!(
1237 "row-space correction shape mismatch for term '{termname}': design is {}x{p}, \
1238 constraint is {}x{q}, correction is {}x{k}",
1239 design.nrows(),
1240 constraint.nrows(),
1241 correction.nrows(),
1242 )));
1243 }
1244 if q == 0 {
1245 return Ok(design);
1246 }
1247 let design_block = match design {
1248 DesignMatrix::Dense(inner) => DesignBlock::Dense(inner),
1249 DesignMatrix::Sparse(inner) => DesignBlock::Sparse(inner),
1250 };
1251 let stacked = BlockDesignOperator::new(vec![
1252 design_block,
1253 DesignBlock::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
1254 constraint.to_owned(),
1255 )),
1256 ])
1257 .map_err(BasisError::InvalidInput)?;
1258 let mut transform = Array2::<f64>::zeros((p + q, k));
1261 for i in 0..p {
1262 transform[[i, i]] = 1.0;
1263 }
1264 for i in 0..q {
1265 for j in 0..k {
1266 transform[[p + i, j]] = -correction[[i, j]];
1267 }
1268 }
1269 let operator = gam_linalg::matrix::CoefficientTransformOperator::new(
1270 gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(stacked)),
1271 transform,
1272 )
1273 .map_err(|e| {
1274 BasisError::InvalidInput(format!(
1275 "row-space correction failed for term '{termname}': {e}"
1276 ))
1277 })?;
1278 Ok(DesignMatrix::Dense(
1279 gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(operator)),
1280 ))
1281}
1282
1283fn build_constraint_block(
1284 n: usize,
1285 parametric_block: Option<&Array2<f64>>,
1286 owner_blocks: &[&DesignMatrix],
1287) -> Result<Array2<f64>, BasisError> {
1288 let param_cols = parametric_block.map_or(0, |mat| mat.ncols());
1289 let owner_cols: usize = owner_blocks.iter().map(|design| design.ncols()).sum();
1290 let mut block = Array2::<f64>::zeros((n, param_cols + owner_cols));
1291 let mut col_start = 0usize;
1292 if let Some(parametric) = parametric_block {
1293 let col_end = col_start + parametric.ncols();
1294 block
1295 .slice_mut(s![.., col_start..col_end])
1296 .assign(parametric);
1297 col_start = col_end;
1298 }
1299 const CHUNK: usize = 1024;
1300 for owner in owner_blocks {
1301 let col_end = col_start + owner.ncols();
1302 for row_start in (0..n).step_by(CHUNK) {
1303 let row_end = (row_start + CHUNK).min(n);
1304 let chunk = (*owner)
1305 .try_row_chunk(row_start..row_end)
1306 .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
1307 block
1308 .slice_mut(s![row_start..row_end, col_start..col_end])
1309 .assign(&chunk);
1310 }
1311 col_start = col_end;
1312 }
1313 Ok(block)
1314}
1315
1316fn design_cross_relative_residual(
1317 lhs: &DesignMatrix,
1318 rhs: &DesignMatrix,
1319) -> Result<f64, BasisError> {
1320 let n = lhs.nrows();
1321 if rhs.nrows() != n {
1322 return Err(BasisError::ConstraintMatrixRowMismatch {
1323 basisrows: n,
1324 constraintrows: rhs.nrows(),
1325 });
1326 }
1327 const CHUNK: usize = 1024;
1328 let mut cross = Array2::<f64>::zeros((lhs.ncols(), rhs.ncols()));
1329 let mut lhs_sumsq = 0.0;
1330 let mut rhs_sumsq = 0.0;
1331 for start in (0..n).step_by(CHUNK) {
1332 let end = (start + CHUNK).min(n);
1333 let lhs_chunk = lhs
1334 .try_row_chunk(start..end)
1335 .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
1336 let rhs_chunk = rhs
1337 .try_row_chunk(start..end)
1338 .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
1339 cross += &gam_linalg::faer_ndarray::fast_atb(&lhs_chunk, &rhs_chunk);
1340 lhs_sumsq += lhs_chunk.iter().map(|v| v * v).sum::<f64>();
1341 rhs_sumsq += rhs_chunk.iter().map(|v| v * v).sum::<f64>();
1342 }
1343 let num = cross.iter().map(|v| v * v).sum::<f64>().sqrt();
1344 let denom = (lhs_sumsq.sqrt() * rhs_sumsq.sqrt()).max(1e-300);
1345 Ok(num / denom)
1346}
1347
1348fn smooth_has_overlapping_linear_terms(
1349 linear_terms: &[LinearTermSpec],
1350 termspec: &SmoothTermSpec,
1351) -> bool {
1352 let feature_cols = smooth_term_feature_cols(termspec);
1353 linear_terms
1354 .iter()
1355 .any(|linear| feature_cols.contains(&linear.feature_col))
1356}
1357
1358pub fn smooth_intrinsic_parametric_feature_cols(
1362 linear_terms: &[LinearTermSpec],
1363 term: &SmoothTermSpec,
1364) -> Vec<usize> {
1365 let feature_cols = smooth_term_feature_cols(term);
1380 let mut owned = Vec::new();
1381 for linear in linear_terms {
1382 if feature_cols.contains(&linear.feature_col) && !owned.contains(&linear.feature_col) {
1383 owned.push(linear.feature_col);
1384 }
1385 }
1386 owned
1387}
1388
1389fn apply_global_smooth_identifiability(
1390 smooth: RawSmoothDesign,
1391 data: ArrayView2<'_, f64>,
1392 linear_terms: &[LinearTermSpec],
1393 smoothspecs: &[SmoothTermSpec],
1394) -> Result<(SmoothDesign, Array1<f64>), BasisError> {
1395 if smoothspecs.len() != smooth.terms.len() {
1406 gam_problem::bail_dim_basis!(
1407 "smooth spec count ({}) does not match built term count ({})",
1408 smoothspecs.len(),
1409 smooth.terms.len()
1410 );
1411 }
1412
1413 if smooth.terms.is_empty() {
1414 let RawSmoothDesign {
1415 term_designs,
1416 affine_offset,
1417 penalties,
1418 nullspace_dims,
1419 penaltyinfo,
1420 dropped_penaltyinfo,
1421 terms,
1422 coefficient_lower_bounds,
1423 linear_constraints,
1424 } = smooth;
1425 return Ok((
1426 SmoothDesign {
1427 term_designs,
1428 penalties,
1429 nullspace_dims,
1430 penaltyinfo,
1431 dropped_penaltyinfo,
1432 terms,
1433 coefficient_lower_bounds,
1434 linear_constraints,
1435 },
1436 affine_offset,
1437 ));
1438 }
1439
1440 let mut local_designs = vec![None; smooth.terms.len()];
1441 let mut local_active_penalties = vec![Vec::<ActivePenalty>::new(); smooth.terms.len()];
1442 let mut local_dropped_penalties = vec![Vec::<DroppedPenaltyInfo>::new(); smooth.terms.len()];
1443 let mut local_metadata = vec![None; smooth.terms.len()];
1444 let mut local_dims = vec![0usize; smooth.terms.len()];
1445 let mut local_linear_constraints = vec![None; smooth.terms.len()];
1446 let mut local_unabsorbed_z = vec![None::<Array2<f64>>; smooth.terms.len()];
1447 let mut local_residualization =
1448 vec![None::<ParametricResidualizationChart>; smooth.terms.len()];
1449 let mut local_collection_gauge = vec![None::<SmoothCollectionGauge>; smooth.terms.len()];
1450
1451 let SmoothStructureAnalysis {
1452 ownership_order,
1453 term_owners,
1454 ..
1455 } = analyze_smooth_ownership(smoothspecs);
1456
1457 use rayon::iter::{IntoParallelIterator, ParallelIterator};
1458
1459 for &idx in &ownership_order {
1460 let term = &smooth.terms[idx];
1461 let termspec = &smoothspecs[idx];
1462 let design_local = smooth.term_designs[idx].clone();
1463 let replay_z = frozen_global_orthogonality(termspec);
1471 let skip_global_transform = replay_z.is_none()
1472 && (smooth_has_frozen_identifiability(termspec) || term.lower_bounds_local.is_some());
1473 let owner_indices = if replay_z.is_some()
1487 || skip_global_transform
1488 || termspec.basis.is_marginally_centered_tensor()
1489 || termspec.basis.is_sum_to_zero_factor_smooth()
1490 {
1491 Vec::new()
1492 } else {
1493 const OVERLAP_REL_RESIDUAL_TOL: f64 = 1e-10;
1497 let owner_cross_checks = term_owners[idx]
1498 .clone()
1499 .into_par_iter()
1500 .map(|owner_idx| {
1501 let owner_design = local_designs[owner_idx]
1502 .as_ref()
1503 .expect("owner design must be available before dependent smooth");
1504 design_cross_relative_residual(&design_local, owner_design)
1505 .map(|rel| (owner_idx, rel))
1506 })
1507 .collect::<Vec<_>>();
1508 let mut out = Vec::new();
1509 for check in owner_cross_checks {
1510 let (owner_idx, rel) = check?;
1511 if rel > OVERLAP_REL_RESIDUAL_TOL {
1512 out.push(owner_idx);
1513 }
1514 }
1515 out
1516 };
1517 let owner_blocks = owner_indices
1518 .iter()
1519 .map(|owner_idx| {
1520 local_designs[*owner_idx]
1521 .as_ref()
1522 .expect("owner design must be available before dependent smooth")
1523 })
1524 .collect::<Vec<_>>();
1525 let replay_correction = frozen_parametric_residualization(termspec);
1530 let needs_parametric_block = match replay_correction {
1531 Some(chart) => chart.has_parametric_block,
1532 None => {
1533 replay_z.is_none()
1534 && !skip_global_transform
1535 && (smooth_has_overlapping_linear_terms(linear_terms, termspec)
1536 || !smooth_intrinsic_parametric_feature_cols(linear_terms, termspec)
1537 .is_empty()
1538 || smooth_requires_parametric_orthogonality(termspec)
1539 || factor_by_level_gate(termspec).is_some())
1545 }
1546 };
1547 let parametric_block = if !needs_parametric_block {
1548 None
1549 } else {
1550 Some(build_parametric_constraint_block_for_term(
1551 data,
1552 linear_terms,
1553 termspec,
1554 )?)
1555 };
1556 let replay_owner_blocks = match replay_correction {
1560 Some(chart) => chart
1561 .owner_terms
1562 .iter()
1563 .map(|owner_idx| {
1564 local_designs.get(*owner_idx).and_then(|slot| slot.as_ref()).ok_or_else(|| {
1565 BasisError::InvalidInput(format!(
1566 "term '{}' replays a parametric residualization against owner term {owner_idx}, which is not available at this point of the rebuild",
1567 termspec.name
1568 ))
1569 })
1570 })
1571 .collect::<Result<Vec<_>, _>>()?,
1572 None => Vec::new(),
1573 };
1574 let plan = if replay_correction.is_some() {
1600 GlobalIdentifiabilityPlan::Absent
1602 } else if skip_global_transform
1603 || (parametric_block.is_none() && owner_blocks.is_empty())
1604 {
1605 GlobalIdentifiabilityPlan::Absent
1606 } else {
1607 let raw =
1608 build_constraint_block(data.nrows(), parametric_block.as_ref(), &owner_blocks)?;
1609 let contained =
1610 crate::basis::contained_constraint_directions(&design_local, raw.view(), None)?;
1611 if raw.ncols() == 0 {
1612 GlobalIdentifiabilityPlan::Absent
1613 } else if contained.ncols() == raw.ncols() {
1614 GlobalIdentifiabilityPlan::Delete { block: contained }
1618 } else {
1619 GlobalIdentifiabilityPlan::Residualize { block: raw }
1620 }
1621 };
1622 let collection_gauge = plan.as_gauge(
1631 &owner_indices,
1632 parametric_block.is_some(),
1633 basis_local_identifiability_transform(&term.metadata),
1634 design_local.ncols(),
1635 );
1636 let mut residualization: Option<crate::basis::ParametricResidualization> = None;
1637 let (design_constrained, z_opt) = if let Some(gauge) = collection_gauge.as_ref() {
1638 let realized = realize_smooth_collection_gauge(design_local, gauge, &term.name)?;
1646 residualization = realized.residualization;
1647 (realized.design, Some(realized.coefficient_transform))
1648 } else {
1649 let z_opt = if let Some(z) = replay_z {
1653 if design_local.ncols() != z.nrows() {
1654 gam_problem::bail_dim_basis!(
1655 "frozen global-orthogonality transform mismatch for term '{}': rebuilt design has {} columns but the persisted fit-time transform has {} rows",
1656 term.name,
1657 design_local.ncols(),
1658 z.nrows()
1659 );
1660 }
1661 Some(z.clone())
1662 } else if skip_global_transform {
1663 None
1664 } else {
1665 maybe_smooth_identifiability_transform(termspec, &design_local, None)?
1668 };
1669 let design_transformed = match z_opt.as_ref() {
1670 Some(z) => apply_smooth_transform_to_design(design_local, z, &term.name)?,
1671 None => design_local,
1672 };
1673 let design_constrained = match replay_correction {
1674 Some(chart) => {
1675 let block = build_constraint_block(
1680 data.nrows(),
1681 parametric_block.as_ref(),
1682 &replay_owner_blocks,
1683 )?;
1684 if block.ncols() != chart.correction.nrows() {
1685 gam_problem::bail_dim_basis!(
1686 "frozen parametric residualization mismatch for term '{}': rebuilt constraint block has {} columns but the persisted fit-time correction has {} rows",
1687 term.name,
1688 block.ncols(),
1689 chart.correction.nrows()
1690 );
1691 }
1692 subtract_row_space_correction(
1693 design_transformed,
1694 block.view(),
1695 chart.correction.view(),
1696 &term.name,
1697 )?
1698 }
1699 None => design_transformed,
1700 };
1701 (design_constrained, z_opt)
1702 };
1703 let coefficient_gauge = z_opt
1704 .as_ref()
1705 .map(|z| gam_problem::Gauge::from_block_transforms(&[z.clone()]));
1706
1707 let penalty_candidates = penalty_candidates_under_collection_gauge(
1708 &term.active_penalties,
1709 coefficient_gauge.as_ref(),
1710 &term.name,
1711 )?;
1712 let filtered = filter_penalty_candidates(penalty_candidates)?;
1713 let linear_constraints_constrained =
1714 if let Some(lin_local) = term.linear_constraints_local.as_ref() {
1715 if let Some(gauge) = coefficient_gauge.as_ref() {
1716 Some(LinearInequalityConstraints {
1717 a: lin_local.a.dot(&gauge.block_transform(0)),
1718 b: lin_local.b.clone(),
1719 })
1720 } else {
1721 Some(lin_local.clone())
1722 }
1723 } else {
1724 None
1725 };
1726
1727 local_residualization[idx] = residualization
1730 .as_ref()
1731 .map(|plan| ParametricResidualizationChart {
1732 owner_terms: owner_indices.clone(),
1733 has_parametric_block: parametric_block.is_some(),
1734 correction: plan.row_space_correction.clone(),
1735 })
1736 .or_else(|| replay_correction.cloned());
1737 local_collection_gauge[idx] = collection_gauge;
1738 local_dims[idx] = design_constrained.ncols();
1739 local_designs[idx] = Some(design_constrained);
1740 local_active_penalties[idx] = filtered.active;
1741 local_dropped_penalties[idx] = term.dropped_penalties.clone();
1742 local_dropped_penalties[idx].extend(filtered.dropped);
1743 local_linear_constraints[idx] = linear_constraints_constrained;
1744 let realized_transform = match (term.joint_null_rotation.as_ref(), z_opt.as_ref()) {
1745 (Some(rotation), Some(z)) => {
1746 Some(gam_linalg::faer_ndarray::fast_ab(&rotation.rotation, z))
1747 }
1748 (Some(rotation), None) => Some(rotation.rotation.clone()),
1749 (None, Some(z)) => Some(z.clone()),
1750 (None, None) => None,
1751 };
1752 match &termspec.basis {
1777 SmoothBasisSpec::FactorSumToZero { .. } | SmoothBasisSpec::FactorSmooth { .. } => {
1778 local_metadata[idx] = Some(term.metadata.clone());
1779 local_unabsorbed_z[idx] = z_opt.clone();
1780 }
1781 _ => {
1782 local_metadata[idx] = Some(with_identifiability_transform(
1783 &term.metadata,
1784 realized_transform.as_ref(),
1785 )?);
1786 }
1787 }
1788 }
1789
1790 let total_p: usize = local_dims.iter().sum();
1791 let mut terms_out = Vec::<SmoothTerm>::with_capacity(smooth.terms.len());
1792 let mut penalties_global = Vec::<BlockwisePenalty>::new();
1793 let mut nullspace_dims_global = Vec::<usize>::new();
1794 let mut penaltyinfo_global = Vec::<PenaltyBlockInfo>::new();
1795 let mut dropped_penaltyinfo_global = Vec::<DroppedPenaltyBlockInfo>::new();
1796 let mut coefficient_lower_bounds = Array1::<f64>::from_elem(total_p, f64::NEG_INFINITY);
1797 let mut any_bounds = false;
1798 let mut linear_constraintsrows: Vec<Array1<f64>> = Vec::new();
1799 let mut linear_constraints_b: Vec<f64> = Vec::new();
1800
1801 let mut col_start = 0usize;
1802 for idx in 0..smooth.terms.len() {
1803 let p_local = local_dims[idx];
1804 let col_end = col_start + p_local;
1805
1806 for active_penalty in &local_active_penalties[idx] {
1807 let global_index = penalties_global.len();
1808 penalties_global.push(BlockwisePenalty::new(
1809 col_start..col_end,
1810 active_penalty.matrix.clone(),
1811 ));
1812 nullspace_dims_global.push(active_penalty.nullity);
1813 penaltyinfo_global.push(PenaltyBlockInfo {
1814 global_index,
1815 termname: Some(smooth.terms[idx].name.clone()),
1816 penalty: active_penalty.info.clone(),
1817 });
1818 }
1819 for info in &local_dropped_penalties[idx] {
1820 dropped_penaltyinfo_global.push(DroppedPenaltyBlockInfo {
1821 termname: Some(smooth.terms[idx].name.clone()),
1822 penalty: info.clone(),
1823 });
1824 }
1825
1826 terms_out.push(SmoothTerm {
1827 name: smooth.terms[idx].name.clone(),
1828 coeff_range: col_start..col_end,
1829 shape: smooth.terms[idx].shape,
1830 active_penalties: local_active_penalties[idx].clone(),
1831 dropped_penalties: local_dropped_penalties[idx].clone(),
1832 metadata: local_metadata[idx]
1833 .clone()
1834 .expect("local metadata must exist for every smooth term"),
1835 lower_bounds_local: smooth.terms[idx].lower_bounds_local.clone(),
1836 linear_constraints_local: local_linear_constraints[idx].clone(),
1837 kronecker_factored: None,
1839 joint_null_rotation: None,
1845 unabsorbed_global_orthogonality: local_unabsorbed_z[idx].clone(),
1848 parametric_residualization: local_residualization[idx].clone(),
1849 collection_gauge: local_collection_gauge[idx].clone(),
1852 });
1853 if let Some(lin_local) = &local_linear_constraints[idx] {
1854 for r in 0..lin_local.a.nrows() {
1855 let mut row = Array1::<f64>::zeros(total_p);
1856 row.slice_mut(s![col_start..col_end])
1857 .assign(&lin_local.a.row(r));
1858 linear_constraintsrows.push(row);
1859 linear_constraints_b.push(lin_local.b[r]);
1860 }
1861 }
1862 if let Some(lb_local) = smooth.terms[idx].lower_bounds_local.as_ref()
1863 && lb_local.len() == p_local
1864 {
1865 coefficient_lower_bounds
1866 .slice_mut(s![col_start..col_end])
1867 .assign(lb_local);
1868 any_bounds = true;
1869 }
1870
1871 col_start = col_end;
1872 }
1873
1874 assert_eq!(
1875 penalties_global.len(),
1876 nullspace_dims_global.len(),
1877 "globally reparameterized smooth penalty/nullspace bookkeeping diverged"
1878 );
1879 assert_eq!(
1880 penalties_global.len(),
1881 penaltyinfo_global.len(),
1882 "globally reparameterized smooth penalty metadata bookkeeping diverged"
1883 );
1884
1885 Ok((
1886 SmoothDesign {
1887 term_designs: local_designs
1888 .into_iter()
1889 .map(|design| design.expect("local design must exist for every smooth term"))
1890 .collect(),
1891 penalties: penalties_global,
1892 nullspace_dims: nullspace_dims_global,
1893 penaltyinfo: penaltyinfo_global,
1894 dropped_penaltyinfo: dropped_penaltyinfo_global,
1895 terms: terms_out,
1896 coefficient_lower_bounds: if any_bounds {
1897 Some(coefficient_lower_bounds)
1898 } else {
1899 None
1900 },
1901 linear_constraints: if linear_constraintsrows.is_empty() {
1902 None
1903 } else {
1904 let mut a = Array2::<f64>::zeros((linear_constraintsrows.len(), total_p));
1905 for (i, row) in linear_constraintsrows.iter().enumerate() {
1906 a.row_mut(i).assign(row);
1907 }
1908 Some(LinearInequalityConstraints {
1909 a,
1910 b: Array1::from_vec(linear_constraints_b),
1911 })
1912 },
1913 },
1914 smooth.affine_offset,
1915 ))
1916}
1917
1918fn factor_by_level_gate(termspec: &SmoothTermSpec) -> Option<(usize, u64)> {
1941 match &termspec.basis {
1942 SmoothBasisSpec::ByVariable {
1943 by_col,
1944 by: ByVariableSpec::Level { value_bits, .. },
1945 ..
1946 } => Some((*by_col, *value_bits)),
1947 _ => None,
1948 }
1949}
1950
1951fn build_parametric_constraint_block_for_term(
1952 data: ArrayView2<'_, f64>,
1953 linear_terms: &[LinearTermSpec],
1954 termspec: &SmoothTermSpec,
1955) -> Result<Array2<f64>, BasisError> {
1956 let n = data.nrows();
1957 let p_data = data.ncols();
1958
1959 if let Some((by_col, value_bits)) = factor_by_level_gate(termspec) {
1963 if by_col >= p_data {
1964 gam_problem::bail_dim_basis!(
1965 "factor-by smooth term '{}' by column {by_col} out of bounds for {p_data} columns",
1966 termspec.name
1967 );
1968 }
1969 let mut c = Array2::<f64>::zeros((n, 1));
1970 let by = data.column(by_col);
1971 let value_bits = gam_data::canonical_level_bits(f64::from_bits(value_bits));
1972 for (row, &value) in by.iter().enumerate() {
1973 if gam_data::canonical_level_bits(value) == value_bits {
1974 c[[row, 0]] = 1.0;
1975 }
1976 }
1977 return Ok(c);
1978 }
1979
1980 let feature_cols = smooth_term_feature_cols(termspec);
1981 let mut parametric_cols = smooth_intrinsic_parametric_feature_cols(linear_terms, termspec);
1982 for &feature_col in ¶metric_cols {
1983 if feature_col >= p_data {
1984 gam_problem::bail_dim_basis!(
1985 "smooth term feature column {feature_col} out of bounds for {p_data} columns"
1986 );
1987 }
1988 }
1989 for linear in linear_terms
1990 .iter()
1991 .filter(|linear| feature_cols.contains(&linear.feature_col))
1992 {
1993 if linear.feature_col >= p_data {
1994 gam_problem::bail_dim_basis!(
1995 "linear term '{}' feature column {} out of bounds for {} columns",
1996 linear.name,
1997 linear.feature_col,
1998 p_data
1999 );
2000 }
2001 if !parametric_cols.contains(&linear.feature_col) {
2002 parametric_cols.push(linear.feature_col);
2003 }
2004 }
2005
2006 let mut c = Array2::<f64>::zeros((n, 1 + parametric_cols.len()));
2007 c.column_mut(0).fill(1.0);
2008 for (j, &feature_col) in parametric_cols.iter().enumerate() {
2009 c.column_mut(j + 1).assign(&data.column(feature_col));
2010 }
2011 Ok(c)
2012}
2013
2014pub fn apply_smooth_transform_to_design(
2015 design_local: DesignMatrix,
2016 transform: &Array2<f64>,
2017 termname: &str,
2018) -> Result<DesignMatrix, BasisError> {
2019 match design_local {
2020 DesignMatrix::Dense(inner) => {
2021 let op = CoefficientTransformOperator::new(inner, transform.clone()).map_err(|e| {
2022 BasisError::InvalidInput(format!(
2023 "smooth identifiability transform failed for term '{termname}': {e}"
2024 ))
2025 })?;
2026 Ok(DesignMatrix::Dense(
2027 gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(op)),
2028 ))
2029 }
2030 DesignMatrix::Sparse(inner) => {
2031 let dense = inner
2032 .try_to_dense_arc("smooth identifiability sparse transform")
2033 .map_err(BasisError::InvalidInput)?
2034 .as_ref()
2035 .dot(transform);
2036 Ok(DesignMatrix::Dense(
2037 gam_linalg::matrix::DenseDesignMatrix::from(dense),
2038 ))
2039 }
2040 }
2041}
2042
2043fn design_constraint_cross(
2044 design: &DesignMatrix,
2045 constraint_matrix: ArrayView2<'_, f64>,
2046) -> Result<Array2<f64>, BasisError> {
2047 let n = design.nrows();
2048 if constraint_matrix.nrows() != n {
2049 return Err(BasisError::ConstraintMatrixRowMismatch {
2050 basisrows: n,
2051 constraintrows: constraint_matrix.nrows(),
2052 });
2053 }
2054 let mut cross = Array2::<f64>::zeros((design.ncols(), constraint_matrix.ncols()));
2055 const CHUNK: usize = 1024;
2056 for start in (0..n).step_by(CHUNK) {
2057 let end = (start + CHUNK).min(n);
2058 let design_chunk = design
2059 .try_row_chunk(start..end)
2060 .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
2061 let constraint_chunk = constraint_matrix.slice(s![start..end, ..]).to_owned();
2062 cross += &gam_linalg::faer_ndarray::fast_atb(&design_chunk, &constraint_chunk);
2063 }
2064 Ok(cross)
2065}
2066
2067fn design_frobenius_norm(design: &DesignMatrix) -> Result<f64, BasisError> {
2068 let n = design.nrows();
2069 const CHUNK: usize = 1024;
2070 let mut sumsq = 0.0;
2071 for start in (0..n).step_by(CHUNK) {
2072 let end = (start + CHUNK).min(n);
2073 let chunk = design
2074 .try_row_chunk(start..end)
2075 .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
2076 sumsq += chunk.iter().map(|v| v * v).sum::<f64>();
2077 }
2078 Ok(sumsq.sqrt())
2079}
2080
2081fn frozen_parametric_residualization(
2089 termspec: &SmoothTermSpec,
2090) -> Option<&ParametricResidualizationChart> {
2091 termspec.frozen_parametric_residualization.as_ref()
2092}
2093
2094fn frozen_global_orthogonality(termspec: &SmoothTermSpec) -> Option<&Array2<f64>> {
2100 match &termspec.basis {
2101 SmoothBasisSpec::FactorSumToZero {
2102 frozen_global_orthogonality,
2103 ..
2104 } => frozen_global_orthogonality.as_ref(),
2105 SmoothBasisSpec::FactorSmooth { spec } => spec.frozen_global_orthogonality.as_ref(),
2106 _ => None,
2107 }
2108}
2109
2110fn penalty_candidates_under_collection_gauge(
2123 active_penalties: &[ActivePenalty],
2124 coefficient_gauge: Option<&gam_problem::Gauge>,
2125 term_name: &str,
2126) -> Result<Vec<PenaltyCandidate>, BasisError> {
2127 use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
2128 let penalty_candidates = active_penalties
2129 .par_iter()
2130 .map(|penalty| -> Result<PenaltyCandidate, BasisError> {
2131 let raw = ConstructiveQuadratic::try_from_dense_psd(
2132 penalty.matrix.clone(),
2133 "global smooth source penalty",
2134 )?;
2135 let raw = match penalty.info.structural_null_frame.as_ref() {
2143 Some(frame) => raw.with_structural_null_frame(
2144 frame.clone(),
2145 "global smooth source penalty structural frame",
2146 )?,
2147 None => raw,
2148 };
2149 let restricted = if let Some(gauge) = coefficient_gauge {
2150 raw.restricted(gauge, "global smooth identifiability restriction")?
2151 } else {
2152 raw
2153 };
2154 let (_, c_new) = normalize_penalty_in_constrained_space(restricted.dense());
2155 let matrix = restricted.scaled(1.0 / c_new, "normalized global smooth penalty")?;
2156 Ok(PenaltyCandidate {
2157 matrix,
2158 source: penalty.info.source.clone(),
2159 normalization_scale: penalty.info.normalization_scale * c_new,
2160 kronecker_factors: None,
2161 op: None,
2162 })
2163 })
2164 .collect::<Result<Vec<_>, _>>()?;
2165 let mut penalty_candidates = penalty_candidates;
2185 if coefficient_gauge.is_some()
2186 && penalty_candidates
2187 .iter()
2188 .any(|c| matches!(c.source, PenaltySource::DoublePenaltyNullspace))
2189 {
2190 const SUPPORT_TOL: f64 = 0.0;
2199 let support_rows = |m: &Array2<f64>| -> (usize, usize) {
2200 let n = m.nrows();
2201 let mut lo = n;
2202 let mut hi = 0usize;
2203 for i in 0..n {
2204 let any = (0..m.ncols()).any(|j| m[[i, j]].abs() > SUPPORT_TOL);
2205 if any {
2206 lo = lo.min(i);
2207 hi = hi.max(i + 1);
2208 }
2209 }
2210 (lo, hi)
2211 };
2212 let primaries: Vec<((usize, usize), ConstructiveQuadratic)> = penalty_candidates
2215 .iter()
2216 .filter(|c| matches!(c.source, PenaltySource::Primary))
2217 .map(|c| -> Result<_, BasisError> {
2218 Ok((
2219 support_rows(&c.matrix),
2220 c.matrix
2221 .scaled(c.normalization_scale, "physical global smooth primary")?,
2222 ))
2223 })
2224 .collect::<Result<Vec<_>, _>>()?;
2225 for candidate in &mut penalty_candidates {
2226 if !matches!(candidate.source, PenaltySource::DoublePenaltyNullspace) {
2227 continue;
2228 }
2229 let q = candidate.matrix.nrows();
2230 let (rlo, rhi) = support_rows(&candidate.matrix);
2231 let owner = primaries
2235 .iter()
2236 .find(|((plo, phi), _)| *plo <= rlo && rhi <= *phi)
2237 .or_else(|| (primaries.len() == 1).then(|| &primaries[0]))
2238 .ok_or_else(|| {
2239 BasisError::InvalidInput(format!(
2240 "double-penalty ridge for smooth '{}' has no co-located primary penalty",
2241 term_name
2242 ))
2243 })?;
2244 let ((plo, phi), s_full) = owner;
2245 let block = ConstructiveQuadratic::from_energy_factor(
2251 s_full.factor().slice(s![.., *plo..*phi]).to_owned(),
2252 "owned global smooth primary block",
2253 )?;
2254 let block = match s_full.structural_null_frame_block(*plo, *phi) {
2260 Some(frame) => block.with_structural_null_frame(
2261 frame,
2262 "owned global smooth primary block structural frame",
2263 )?,
2264 None => block,
2265 };
2266 let ridge_full = candidate.matrix.scaled(
2267 candidate.normalization_scale,
2268 "physical global smooth null ridge",
2269 )?;
2270 let ridge_block = ConstructiveQuadratic::from_energy_factor(
2271 ridge_full.factor().slice(s![.., *plo..*phi]).to_owned(),
2272 "owned global smooth null-ridge block",
2273 )?;
2274 let rebuilt_block =
2275 crate::basis::rebuild_metric_consistent_ridge(&block, &ridge_block)?;
2276 match rebuilt_block {
2277 Some(ridge_block) => {
2278 let mut full_factor =
2279 Array2::<f64>::zeros((ridge_block.factor().nrows(), q));
2280 full_factor
2281 .slice_mut(s![.., *plo..*phi])
2282 .assign(ridge_block.factor());
2283 let full = ConstructiveQuadratic::from_energy_factor(
2284 full_factor,
2285 "embedded global smooth null ridge",
2286 )?;
2287 let (_, scale) = normalize_penalty_in_constrained_space(full.dense());
2288 candidate.matrix = full
2289 .scaled(1.0 / scale, "normalized embedded global smooth null ridge")?;
2290 candidate.normalization_scale = scale;
2291 candidate.kronecker_factors = None;
2292 candidate.op = None;
2293 }
2294 None => {
2297 candidate.matrix = ConstructiveQuadratic::zero(q);
2298 candidate.normalization_scale = 1.0;
2299 candidate.kronecker_factors = None;
2300 candidate.op = None;
2301 }
2302 }
2303 }
2304 }
2305 Ok(penalty_candidates)
2306}
2307
2308fn maybe_smooth_identifiability_transform(
2309 termspec: &SmoothTermSpec,
2310 design_local: &DesignMatrix,
2311 constraint_block: Option<ArrayView2<'_, f64>>,
2312) -> Result<Option<Array2<f64>>, BasisError> {
2313 if let Some(SpatialIdentifiability::FrozenTransform { transform }) =
2314 spatial_identifiability_policy(termspec)
2315 {
2316 if design_local.ncols() != transform.nrows() {
2317 gam_problem::bail_dim_basis!(
2318 "frozen spatial identifiability transform mismatch: design has {} columns but transform has {} rows",
2319 design_local.ncols(),
2320 transform.nrows()
2321 );
2322 }
2323 return Ok(Some(transform.clone()));
2324 }
2325
2326 if let Some(c) = constraint_block {
2327 if c.ncols() == 0 {
2328 Ok(None)
2329 } else {
2330 Ok(Some(orthogonality_transform_for_design(
2331 design_local,
2332 c,
2333 None, )?))
2335 }
2336 } else {
2337 Ok(None)
2338 }
2339}
2340
2341fn smooth_requires_parametric_orthogonality(termspec: &SmoothTermSpec) -> bool {
2410 match &termspec.basis {
2411 SmoothBasisSpec::ByVariable { inner, .. }
2412 | SmoothBasisSpec::FactorSumToZero { inner, .. } => {
2413 smooth_requires_parametric_orthogonality(&SmoothTermSpec {
2414 frozen_parametric_residualization: None,
2415 name: termspec.name.clone(),
2416 basis: (**inner).clone(),
2417 shape: termspec.shape,
2418 joint_null_rotation: None,
2419 })
2420 }
2421 SmoothBasisSpec::BySmooth { smooth, .. } => {
2422 smooth_requires_parametric_orthogonality(&SmoothTermSpec {
2423 frozen_parametric_residualization: None,
2424 name: termspec.name.clone(),
2425 basis: (**smooth).clone(),
2426 shape: termspec.shape,
2427 joint_null_rotation: None,
2428 })
2429 }
2430 SmoothBasisSpec::ThinPlate { spec, .. } => {
2431 matches!(
2432 spec.identifiability,
2433 SpatialIdentifiability::OrthogonalToParametric
2434 )
2435 }
2436 SmoothBasisSpec::Duchon { spec, .. } => {
2437 matches!(
2438 spec.identifiability,
2439 SpatialIdentifiability::OrthogonalToParametric
2440 )
2441 }
2442 SmoothBasisSpec::Matern { spec, .. } => matches!(
2443 spec.identifiability,
2444 MaternIdentifiability::CenterSumToZero | MaternIdentifiability::CenterLinearOrthogonal
2445 ),
2446 SmoothBasisSpec::Sphere { spec, .. } => {
2454 matches!(spec.method, crate::basis::SphereMethod::Wahba)
2455 && !matches!(spec.wahba_kernel, crate::basis::SphereWahbaKernel::Pseudo)
2456 && matches!(
2457 spec.identifiability,
2458 SphericalSplineIdentifiability::CenterSumToZero
2459 )
2460 }
2461 SmoothBasisSpec::ConstantCurvature { spec, .. } => matches!(
2469 spec.identifiability,
2470 ConstantCurvatureIdentifiability::CenterSumToZero
2471 ),
2472 SmoothBasisSpec::MeasureJet { spec, .. } => matches!(
2478 spec.identifiability,
2479 MeasureJetIdentifiability::CenterSumToZero
2480 ),
2481 SmoothBasisSpec::BSpline1D { .. }
2482 | SmoothBasisSpec::TensorBSpline { .. }
2483 | SmoothBasisSpec::Pca { .. }
2484 | SmoothBasisSpec::FactorSmooth { .. } => false,
2485 }
2486}
2487
2488fn compose_identifiability_transforms(
2489 existing: Option<&Array2<f64>>,
2490 extra: Option<&Array2<f64>>,
2491) -> Result<Option<Array2<f64>>, BasisError> {
2492 match (existing, extra) {
2493 (Some(lhs), Some(rhs)) => {
2494 if lhs.ncols() == rhs.nrows() {
2495 Ok(Some(lhs.dot(rhs)))
2496 } else if lhs.nrows() == rhs.nrows() && lhs.ncols() == rhs.ncols() {
2497 Ok(Some(rhs.clone()))
2501 } else {
2502 Err(BasisError::DimensionMismatch(format!(
2503 "identifiability transform mismatch: existing is {}x{}, extra is {}x{}",
2504 lhs.nrows(),
2505 lhs.ncols(),
2506 rhs.nrows(),
2507 rhs.ncols(),
2508 )))
2509 }
2510 }
2511 (Some(lhs), None) => Ok(Some(lhs.clone())),
2512 (None, Some(rhs)) => Ok(Some(rhs.clone())),
2513 (None, None) => Ok(None),
2514 }
2515}
2516
2517fn with_identifiability_transform(
2518 metadata: &BasisMetadata,
2519 transform: Option<&Array2<f64>>,
2520) -> Result<BasisMetadata, BasisError> {
2521 match metadata {
2522 BasisMetadata::BSpline1D {
2523 knots,
2524 identifiability_transform,
2525 periodic,
2526 degree,
2527 auto_shrink_note,
2528 anchor_offset_coeffs,
2529 } => Ok(BasisMetadata::BSpline1D {
2530 knots: knots.clone(),
2531 periodic: *periodic,
2532 identifiability_transform: compose_identifiability_transforms(
2533 identifiability_transform.as_ref(),
2534 transform,
2535 )?,
2536 degree: *degree,
2537 auto_shrink_note: auto_shrink_note.clone(),
2538 anchor_offset_coeffs: anchor_offset_coeffs.clone(),
2542 }),
2543 BasisMetadata::CubicRegression1D {
2544 knots,
2545 identifiability_transform,
2546 } => Ok(BasisMetadata::CubicRegression1D {
2547 knots: knots.clone(),
2548 identifiability_transform: compose_identifiability_transforms(
2549 identifiability_transform.as_ref(),
2550 transform,
2551 )?,
2552 }),
2553 BasisMetadata::ThinPlate {
2554 centers,
2555 length_scale,
2556 periodic,
2557 identifiability_transform,
2558 input_scale,
2559 radial_reparam,
2560 } => Ok(BasisMetadata::ThinPlate {
2561 centers: centers.clone(),
2562 length_scale: *length_scale,
2563 periodic: periodic.clone(),
2564 identifiability_transform: compose_identifiability_transforms(
2565 identifiability_transform.as_ref(),
2566 transform,
2567 )?,
2568 input_scale: *input_scale,
2569 radial_reparam: radial_reparam.clone(),
2570 }),
2571 BasisMetadata::Sphere {
2572 centers,
2573 penalty_order,
2574 method,
2575 max_degree,
2576 wahba_kernel,
2577 constraint_transform,
2578 } => Ok(BasisMetadata::Sphere {
2579 centers: centers.clone(),
2580 penalty_order: *penalty_order,
2581 method: *method,
2582 max_degree: *max_degree,
2583 wahba_kernel: *wahba_kernel,
2584 constraint_transform: compose_identifiability_transforms(
2585 constraint_transform.as_ref(),
2586 transform,
2587 )?,
2588 }),
2589 BasisMetadata::ConstantCurvature {
2590 centers,
2591 kappa,
2592 length_scale,
2593 constraint_transform,
2594 } => Ok(BasisMetadata::ConstantCurvature {
2595 centers: centers.clone(),
2596 kappa: *kappa,
2597 length_scale: *length_scale,
2598 constraint_transform: compose_identifiability_transforms(
2599 constraint_transform.as_ref(),
2600 transform,
2601 )?,
2602 }),
2603 BasisMetadata::MeasureJet {
2604 centers,
2605 input_scale,
2606 length_scale,
2607 eps_band,
2608 order_s,
2609 alpha,
2610 tau0,
2611 masses,
2612 support_means,
2613 penalty_normalization_scales,
2614 raw_penalty_normalization_scales,
2615 fused_penalty_normalization_scale,
2616 constraint_transform,
2617 sigma_coord,
2618 } => Ok(BasisMetadata::MeasureJet {
2619 centers: centers.clone(),
2620 input_scale: *input_scale,
2621 length_scale: *length_scale,
2622 eps_band: eps_band.clone(),
2623 order_s: *order_s,
2624 alpha: *alpha,
2625 tau0: *tau0,
2626 masses: masses.clone(),
2627 support_means: support_means.clone(),
2628 penalty_normalization_scales: penalty_normalization_scales.clone(),
2629 raw_penalty_normalization_scales: raw_penalty_normalization_scales.clone(),
2630 fused_penalty_normalization_scale: *fused_penalty_normalization_scale,
2631 constraint_transform: compose_identifiability_transforms(
2632 constraint_transform.as_ref(),
2633 transform,
2634 )?,
2635 sigma_coord: *sigma_coord,
2636 }),
2637 BasisMetadata::Matern {
2638 centers,
2639 length_scale,
2640 periodic,
2641 nu,
2642 include_intercept,
2643 identifiability_transform,
2644 input_scale,
2645 aniso_log_scales,
2646 } => Ok(BasisMetadata::Matern {
2647 centers: centers.clone(),
2648 length_scale: *length_scale,
2649 periodic: periodic.clone(),
2650 nu: *nu,
2651 include_intercept: *include_intercept,
2652 identifiability_transform: compose_identifiability_transforms(
2653 identifiability_transform.as_ref(),
2654 transform,
2655 )?,
2656 input_scale: *input_scale,
2657 aniso_log_scales: aniso_log_scales.clone(),
2658 }),
2659 BasisMetadata::Duchon {
2660 centers,
2661 length_scale,
2662 periodic,
2663 power,
2664 nullspace_order,
2665 identifiability_transform,
2666 input_scale,
2667 aniso_log_scales,
2668 operator_collocation_points,
2669 radial_reparam,
2670 spectral_basis,
2671 } => Ok(BasisMetadata::Duchon {
2672 centers: centers.clone(),
2673 length_scale: *length_scale,
2674 periodic: periodic.clone(),
2675 power: *power,
2676 nullspace_order: *nullspace_order,
2677 input_scale: *input_scale,
2678 aniso_log_scales: aniso_log_scales.clone(),
2679 operator_collocation_points: operator_collocation_points.clone(),
2680 radial_reparam: radial_reparam.clone(),
2681 spectral_basis: spectral_basis.clone(),
2682 identifiability_transform: compose_identifiability_transforms(
2683 identifiability_transform.as_ref(),
2684 transform,
2685 )?,
2686 }),
2687 BasisMetadata::SphereHarmonics {
2688 max_degree,
2689 radians,
2690 } => Ok(BasisMetadata::SphereHarmonics {
2691 max_degree: *max_degree,
2692 radians: *radians,
2693 }),
2694 BasisMetadata::TensorBSpline {
2695 feature_cols,
2696 knots,
2697 degrees,
2698 periods,
2699 is_cr,
2700 identifiability_transform,
2701 } => Ok(BasisMetadata::TensorBSpline {
2702 feature_cols: feature_cols.clone(),
2703 knots: knots.clone(),
2704 degrees: degrees.clone(),
2705 periods: periods.clone(),
2706 is_cr: is_cr.clone(),
2707 identifiability_transform: compose_identifiability_transforms(
2708 identifiability_transform.as_ref(),
2709 transform,
2710 )?,
2711 }),
2712 BasisMetadata::BySmooth {
2713 inner,
2714 by_col,
2715 levels,
2716 ordered,
2717 } => Ok(BasisMetadata::BySmooth {
2718 inner: Box::new(with_identifiability_transform(inner, transform)?),
2719 by_col: *by_col,
2720 levels: levels.clone(),
2721 ordered: *ordered,
2722 }),
2723 BasisMetadata::FactorSmooth {
2724 continuous_cols,
2725 group_col,
2726 knots,
2727 degree,
2728 periodic,
2729 group_levels,
2730 flavour,
2731 marginal_is_cr,
2732 } => {
2733 if transform.is_some() {
2740 gam_problem::bail_invalid_basis!(
2741 "FactorSmooth metadata cannot absorb an identifiability transform; \
2742 route it through the term-level frozen_global_orthogonality carrier"
2743 );
2744 }
2745 Ok(BasisMetadata::FactorSmooth {
2746 continuous_cols: continuous_cols.clone(),
2747 group_col: *group_col,
2748 knots: knots.clone(),
2749 degree: *degree,
2750 periodic: *periodic,
2751 group_levels: group_levels.clone(),
2752 flavour: flavour.clone(),
2753 marginal_is_cr: *marginal_is_cr,
2754 })
2755 }
2756 BasisMetadata::Pca {
2757 feature_cols,
2758 basis_matrix,
2759 centered,
2760 smooth_penalty,
2761 center_mean,
2762 pca_basis_path,
2763 chunk_size,
2764 } => {
2765 if transform.is_some() {
2771 gam_problem::bail_invalid_basis!(
2772 "PCA bases do not expose a composable identifiability transform"
2773 );
2774 }
2775 Ok(BasisMetadata::Pca {
2776 feature_cols: feature_cols.clone(),
2777 basis_matrix: basis_matrix.clone(),
2778 centered: *centered,
2779 smooth_penalty: *smooth_penalty,
2780 center_mean: center_mean.clone(),
2781 pca_basis_path: pca_basis_path.clone(),
2782 chunk_size: *chunk_size,
2783 })
2784 }
2785 }
2786}
2787
2788pub fn orthogonality_relative_residual_for_design(
2792 design: &DesignMatrix,
2793 constraint_matrix: ArrayView2<'_, f64>,
2794) -> Result<f64, BasisError> {
2795 let cross = design_constraint_cross(design, constraint_matrix)?;
2796 let num = cross.iter().map(|v| v * v).sum::<f64>().sqrt();
2797 let b_norm = design_frobenius_norm(design)?;
2798 let c_norm = constraint_matrix.iter().map(|v| v * v).sum::<f64>().sqrt();
2799 let denom = (b_norm * c_norm).max(1e-300);
2800 Ok(num / denom)
2801}
2802
2803#[cfg(test)]
2804mod frozen_linear_term_mass_rebuild_tests {
2805 use super::*;
2806
2807 fn one_linear_term_spec() -> TermCollectionSpec {
2811 TermCollectionSpec {
2812 linear_terms: vec![LinearTermSpec {
2813 name: "x".to_string(),
2814 feature_col: 0,
2815 feature_cols: vec![0],
2816 categorical_levels: vec![],
2817 double_penalty: true,
2818 coefficient_geometry: LinearCoefficientGeometry::Unconstrained,
2819 coefficient_min: None,
2820 coefficient_max: None,
2821 frozen_function_mass: None,
2822 }],
2823 random_effect_terms: Vec::new(),
2824 smooth_terms: Vec::new(),
2825 }
2826 }
2827
2828 fn training_data_varying_x(n: usize) -> Array2<f64> {
2829 let mut data = Array2::<f64>::zeros((n, 1));
2830 for i in 0..n {
2831 data[[i, 0]] = 1.0 + i as f64;
2833 }
2834 data
2835 }
2836
2837 fn constant_zero_x(n_rows: usize) -> Array2<f64> {
2838 Array2::<f64>::zeros((n_rows, 1))
2839 }
2840
2841 #[test]
2847 fn unfrozen_spec_still_rejects_a_genuinely_zero_training_column() {
2848 let spec = one_linear_term_spec();
2849 let degenerate_training_data = constant_zero_x(20);
2850 let err = build_term_collection_design(degenerate_training_data.view(), &spec)
2851 .expect_err("an unfrozen spec fit directly on an all-zero column must still fail");
2852 let message = err.to_string();
2853 assert!(
2854 message.contains("identically zero"),
2855 "expected the identifiability guard's message, got: {message}"
2856 );
2857 }
2858
2859 #[test]
2870 fn frozen_spec_rebuilds_at_a_constant_evaluation_column_using_the_training_mass() {
2871 let spec = one_linear_term_spec();
2872 let training_data = training_data_varying_x(40);
2873
2874 let training_design = build_term_collection_design(training_data.view(), &spec)
2875 .expect("fit-time build over a genuinely varying column must succeed");
2876 let training_mass = training_design
2877 .linear_function_masses
2878 .first()
2879 .copied()
2880 .flatten()
2881 .expect("a double_penalty=true term must report its fit-time function mass");
2882 assert!(
2883 training_mass > 0.0,
2884 "training mass for a genuinely varying column must be strictly positive, got {training_mass}"
2885 );
2886
2887 let frozen_spec = freeze_term_collection_from_design(&spec, &training_design)
2888 .expect("freezing the spec against its own fit-time design must succeed");
2889 assert_eq!(
2890 frozen_spec.linear_terms[0].frozen_function_mass,
2891 Some(training_mass),
2892 "freezing must persist the exact fit-time mass onto the term"
2893 );
2894
2895 let evaluation_grid = constant_zero_x(3);
2899 let rebuilt_design = build_term_collection_design(evaluation_grid.view(), &frozen_spec)
2900 .expect(
2901 "rebuilding a FROZEN spec's design at a constant-covariate evaluation grid must \
2902 succeed — the training-time mass is reused, never recomputed from these rows",
2903 );
2904 assert_eq!(
2905 rebuilt_design
2906 .linear_function_masses
2907 .first()
2908 .copied()
2909 .flatten(),
2910 Some(training_mass),
2911 "the rebuilt design must carry the REUSED training-time mass, not a value \
2912 recomputed from the (all-zero) evaluation rows"
2913 );
2914 }
2915}