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 build_term_collection_design_inner_with_policy_and_plan(data, spec, policy, false)
84}
85
86pub fn build_planned_term_collection_design_inner_with_policy(
88 data: ArrayView2<'_, f64>,
89 spec: &TermCollectionSpec,
90 policy: &gam_runtime::resource::ResourcePolicy,
91) -> Result<TermCollectionDesign, BasisError> {
92 build_term_collection_design_inner_with_policy_and_plan(data, spec, policy, true)
93}
94
95fn build_term_collection_design_inner_with_policy_and_plan(
96 data: ArrayView2<'_, f64>,
97 spec: &TermCollectionSpec,
98 policy: &gam_runtime::resource::ResourcePolicy,
99 spatial_plan_is_resolved: bool,
100) -> Result<TermCollectionDesign, BasisError> {
101 use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
102
103 let n = data.nrows();
104 let p_intercept = usize::from(!term_collection_has_anchored_bspline(spec));
105 let p_lin = spec.linear_terms.len();
106
107 let (smooth_raw_result, (random_blocks_result, linear_block_result)) = rayon::join(
112 || {
113 let mut ws = crate::basis::BasisWorkspace::with_policy(policy.clone());
114 if spatial_plan_is_resolved {
115 build_smooth_design_from_planned_terms(data, &spec.smooth_terms, &mut ws)
116 } else {
117 build_smooth_design_withworkspace_unvalidated(data, &spec.smooth_terms, &mut ws)
118 }
119 },
120 || {
121 rayon::join(
122 || {
123 spec.random_effect_terms
124 .par_iter()
125 .map(|term| build_random_effect_block(data, term))
126 .collect::<Result<Vec<_>, _>>()
127 },
128 || -> Result<Option<Array2<f64>>, BasisError> {
129 if p_lin == 0 {
130 return Ok(None);
131 }
132
133 let mut out = Array2::<f64>::zeros((n, p_lin));
137 for (j, linear) in spec.linear_terms.iter().enumerate() {
138 let column = linear
149 .realized_design_column(data)
150 .map_err(BasisError::InvalidInput)?;
151 out.column_mut(j).assign(&column);
152 }
153 Ok(Some(out))
154 },
155 )
156 },
157 );
158
159 let smooth_raw = smooth_raw_result?;
160 let random_blocks = random_blocks_result?;
161 let linear_block = linear_block_result?;
162 let linear_function_masses = match linear_block.as_ref() {
173 Some(block) => spec
174 .linear_terms
175 .iter()
176 .enumerate()
177 .map(|(j, term)| -> Result<Option<f64>, BasisError> {
178 if !term.double_penalty {
179 return Ok(None);
180 }
181 if let Some(frozen_mass) = term.frozen_function_mass {
182 return Ok(Some(frozen_mass));
183 }
184 linear_function_mass(block.column(j), &term.name).map(Some)
185 })
186 .collect::<Result<Vec<_>, _>>()?,
187 None => Vec::new(),
188 };
189
190 let (smooth, affine_offset) = apply_global_smooth_identifiability(
191 smooth_raw,
192 data,
193 &spec.linear_terms,
194 &spec.smooth_terms,
195 )?;
196
197 let p_rand: usize = random_blocks.iter().map(|b| b.num_groups).sum();
198 let p_smooth = smooth.total_smooth_cols();
199 let p_total = p_intercept + p_lin + p_rand + p_smooth;
200
201 let mut linear_ranges = Vec::<(String, Range<usize>)>::with_capacity(p_lin);
202 for (j, linear) in spec.linear_terms.iter().enumerate() {
203 let col = p_intercept + j;
204 linear_ranges.push((linear.name.clone(), col..(col + 1)));
207 }
208
209 let mut random_effect_ranges =
212 Vec::<(String, Range<usize>)>::with_capacity(random_blocks.len());
213 let mut random_effect_levels = Vec::<(String, Vec<u64>)>::with_capacity(random_blocks.len());
214 let mut col_cursor = p_intercept + p_lin;
215 for block in &random_blocks {
216 let q = block.num_groups;
217 let end = col_cursor + q;
218 random_effect_ranges.push((block.name.clone(), col_cursor..end));
219 random_effect_levels.push((block.name.clone(), block.kept_levels.clone()));
220 col_cursor = end;
221 }
222
223 let mut blocks = Vec::<DesignBlock>::new();
239
240 if p_intercept == 1 {
245 blocks.push(DesignBlock::Intercept(n));
246 }
247
248 if let Some(lin_block) = linear_block {
250 blocks.push(DesignBlock::Dense(
251 gam_linalg::matrix::DenseDesignMatrix::from(lin_block),
252 ));
253 }
254
255 for block in &random_blocks {
257 let re_op = RandomEffectOperator::new(block.group_ids.clone(), block.num_groups);
258 blocks.push(DesignBlock::RandomEffect(Arc::new(re_op)));
259 }
260
261 if p_smooth > 0 {
265 for term_design in &smooth.term_designs {
266 match term_design {
267 DesignMatrix::Dense(dense) => blocks.push(DesignBlock::Dense(dense.clone())),
268 DesignMatrix::Sparse(sparse) => blocks.push(DesignBlock::Sparse(sparse.clone())),
269 }
270 }
271 }
272
273 let design = assemble_term_collection_design_matrix(blocks)?;
274
275 let mut penalties = Vec::<BlockwisePenalty>::new();
276 let mut nullspace_dims = Vec::<usize>::new();
277 let mut penaltyinfo = Vec::<PenaltyBlockInfo>::new();
278 let mut dropped_penaltyinfo = Vec::<DroppedPenaltyBlockInfo>::new();
279 let mut coefficient_lower_bounds = Array1::<f64>::from_elem(p_total, f64::NEG_INFINITY);
280 let mut any_bounds = false;
281 let mut linear_constraintrows = Vec::<Array1<f64>>::new();
282 let mut linear_constraint_b = Vec::<f64>::new();
283
284 for (j, linear) in spec.linear_terms.iter().enumerate() {
285 let col = p_intercept + j;
286 if let Some(lb) = linear.coefficient_min {
287 let mut row = Array1::<f64>::zeros(p_total);
288 row[col] = 1.0;
289 linear_constraintrows.push(row);
290 linear_constraint_b.push(lb);
291 }
292 if let Some(ub) = linear.coefficient_max {
293 let mut row = Array1::<f64>::zeros(p_total);
294 row[col] = -1.0;
295 linear_constraintrows.push(row);
296 linear_constraint_b.push(-ub);
297 }
298 }
299
300 for (j, linear) in spec.linear_terms.iter().enumerate() {
308 let Some(function_mass) = linear_function_masses.get(j).copied().flatten() else {
309 continue;
310 };
311 let col = p_intercept + j;
312 let global_index = penalties.len();
313 penalties.push(BlockwisePenalty::new(
314 col..(col + 1),
315 Array2::from_elem((1, 1), function_mass),
316 ));
317 nullspace_dims.push(0);
318 penaltyinfo.push(PenaltyBlockInfo {
319 global_index,
320 termname: Some(linear.name.clone()),
321 penalty: ActivePenaltyInfo {
322 source: PenaltySource::Other("LinearTermRidge".to_string()),
323 original_index: j,
324 effective_rank: 1,
325 normalization_scale: 1.0,
326 kronecker_factors: None,
327 structural_null_frame: None,
328 },
329 });
330 }
331
332 for (re_idx, (name, range)) in random_effect_ranges.iter().enumerate() {
333 if range.is_empty() || !spec.random_effect_terms[re_idx].penalized {
334 continue;
335 }
336 let block_size = range.len();
337 let global_index = penalties.len();
338 penalties.push(BlockwisePenalty::ridge(range.clone(), 1.0));
339 nullspace_dims.push(0);
340 penaltyinfo.push(PenaltyBlockInfo {
341 global_index,
342 termname: Some(name.clone()),
343 penalty: ActivePenaltyInfo {
344 source: PenaltySource::Other(format!("RandomEffectRidge({name})")),
345 original_index: re_idx,
346 effective_rank: block_size,
347 normalization_scale: 1.0,
348 kronecker_factors: None,
349 structural_null_frame: None,
350 },
351 });
352 }
353
354 if smooth.penaltyinfo.len() != smooth.penalties.len() {
355 gam_problem::bail_invalid_basis!(
356 "smooth penalty metadata mismatch: penalties={}, metadata={}",
357 smooth.penalties.len(),
358 smooth.penaltyinfo.len()
359 );
360 }
361 let smooth_start = p_intercept + p_lin + p_rand;
362 for ((bp_smooth, &ns), localinfo) in smooth
363 .penalties
364 .iter()
365 .zip(smooth.nullspace_dims.iter())
366 .zip(smooth.penaltyinfo.iter())
367 {
368 let global_index = penalties.len();
369 let offset_range =
371 (bp_smooth.col_range.start + smooth_start)..(bp_smooth.col_range.end + smooth_start);
372 let bp = if let Some(factors) = localinfo.penalty.kronecker_factors.as_ref() {
373 BlockwisePenalty::kronecker(offset_range, bp_smooth.local.clone(), factors.clone())
374 .with_op(bp_smooth.op.clone())
375 } else if matches!(
376 localinfo.penalty.source,
377 PenaltySource::Other(ref s) if s.starts_with("RandomEffectRidge")
378 ) {
379 BlockwisePenalty::ridge(offset_range, 1.0)
380 } else {
381 BlockwisePenalty::new(offset_range, bp_smooth.local.clone())
382 .with_op(bp_smooth.op.clone())
383 };
384 penalties.push(bp);
385 nullspace_dims.push(ns);
386 penaltyinfo.push(PenaltyBlockInfo {
387 global_index,
388 termname: localinfo.termname.clone(),
389 penalty: localinfo.penalty.clone(),
390 });
391 }
392 dropped_penaltyinfo.extend(smooth.dropped_penaltyinfo.iter().cloned());
393
394 assert_eq!(
395 penalties.len(),
396 nullspace_dims.len(),
397 "term-collection penalty/nullspace bookkeeping diverged"
398 );
399 assert_eq!(
400 penalties.len(),
401 penaltyinfo.len(),
402 "term-collection penalty metadata bookkeeping diverged"
403 );
404
405 if let Some(lb_smooth) = smooth.coefficient_lower_bounds.as_ref() {
406 let start = p_intercept + p_lin + p_rand;
407 coefficient_lower_bounds
408 .slice_mut(s![start..(start + p_smooth)])
409 .assign(lb_smooth);
410 any_bounds = true;
411 }
412 if let Some(lin_smooth) = smooth.linear_constraints.as_ref() {
413 let mut a_global = Array2::<f64>::zeros((lin_smooth.a.nrows(), p_total));
414 let start = p_intercept + p_lin + p_rand;
415 a_global
416 .slice_mut(s![.., start..(start + p_smooth)])
417 .assign(&lin_smooth.a);
418 for r in 0..a_global.nrows() {
419 linear_constraintrows.push(a_global.row(r).to_owned());
420 linear_constraint_b.push(lin_smooth.b[r]);
421 }
422 }
423
424 let lower_bound_constraints = if any_bounds {
428 linear_constraints_from_lower_bounds_global(&coefficient_lower_bounds)
429 } else {
430 None
431 };
432 let explicit_linear_constraints = if linear_constraintrows.is_empty() {
433 None
434 } else {
435 let mut a = Array2::<f64>::zeros((linear_constraintrows.len(), p_total));
436 for (i, row) in linear_constraintrows.iter().enumerate() {
437 a.row_mut(i).assign(row);
438 }
439 Some(LinearInequalityConstraints {
440 a,
441 b: Array1::from_vec(linear_constraint_b),
442 })
443 };
444 let linear_constraints =
445 merge_linear_constraints_global(explicit_linear_constraints, lower_bound_constraints)?;
446
447 Ok(TermCollectionDesign {
448 design,
449 affine_offset,
450 penalties,
451 nullspace_dims,
452 penaltyinfo,
453 dropped_penaltyinfo,
454 coefficient_lower_bounds: if any_bounds {
455 Some(coefficient_lower_bounds)
456 } else {
457 None
458 },
459 linear_constraints,
460 intercept_range: 0..p_intercept,
461 linear_ranges,
462 linear_function_masses,
463 random_effect_ranges,
464 random_effect_levels,
465 smooth,
466 })
467}
468
469pub fn term_collection_has_anchored_bspline(spec: &TermCollectionSpec) -> bool {
474 spec.smooth_terms
475 .iter()
476 .any(|term| smooth_basis_has_anchored_bspline(&term.basis))
477}
478
479pub fn term_collection_has_nonzero_anchor(spec: &TermCollectionSpec) -> bool {
482 spec.smooth_terms
483 .iter()
484 .any(|term| smooth_basis_has_nonzero_anchor(&term.basis))
485}
486
487fn smooth_basis_has_nonzero_anchor(basis: &SmoothBasisSpec) -> bool {
488 match basis {
489 SmoothBasisSpec::ByVariable { inner, .. }
490 | SmoothBasisSpec::FactorSumToZero { inner, .. } => smooth_basis_has_nonzero_anchor(inner),
491 SmoothBasisSpec::BSpline1D { spec, .. } => spec.boundary_conditions.has_nonzero_anchor(),
492 SmoothBasisSpec::BySmooth { smooth, .. } => smooth_basis_has_nonzero_anchor(smooth),
493 SmoothBasisSpec::TensorBSpline { spec, .. } => spec
494 .marginalspecs
495 .iter()
496 .any(|marginal| marginal.boundary_conditions.has_nonzero_anchor()),
497 SmoothBasisSpec::FactorSmooth { spec } => {
498 spec.marginal.boundary_conditions.has_nonzero_anchor()
499 }
500 SmoothBasisSpec::ThinPlate { .. }
501 | SmoothBasisSpec::Sphere { .. }
502 | SmoothBasisSpec::ConstantCurvature { .. }
503 | SmoothBasisSpec::Matern { .. }
504 | SmoothBasisSpec::MeasureJet { .. }
505 | SmoothBasisSpec::Duchon { .. }
506 | SmoothBasisSpec::Pca { .. } => false,
507 }
508}
509
510fn smooth_basis_has_anchored_bspline(basis: &SmoothBasisSpec) -> bool {
511 match basis {
512 SmoothBasisSpec::ByVariable { inner, .. }
513 | SmoothBasisSpec::FactorSumToZero { inner, .. } => {
514 smooth_basis_has_anchored_bspline(inner)
515 }
516 SmoothBasisSpec::BSpline1D { spec, .. } => {
517 bspline_conditions_have_anchor(&spec.boundary_conditions)
518 }
519 SmoothBasisSpec::BySmooth { smooth, .. } => smooth_basis_has_anchored_bspline(smooth),
520 SmoothBasisSpec::TensorBSpline { spec, .. } => spec
521 .marginalspecs
522 .iter()
523 .any(|marginal| bspline_conditions_have_anchor(&marginal.boundary_conditions)),
524 SmoothBasisSpec::FactorSmooth { .. }
525 | SmoothBasisSpec::ThinPlate { .. }
526 | SmoothBasisSpec::Sphere { .. }
527 | SmoothBasisSpec::ConstantCurvature { .. }
528 | SmoothBasisSpec::Matern { .. }
529 | SmoothBasisSpec::MeasureJet { .. }
530 | SmoothBasisSpec::Duchon { .. }
531 | SmoothBasisSpec::Pca { .. } => false,
532 }
533}
534
535fn bspline_conditions_have_anchor(conditions: &crate::basis::BSplineBoundaryConditions) -> bool {
536 conditions.has_anchor()
537}
538
539pub fn build_term_collection_design(
540 data: ArrayView2<'_, f64>,
541 spec: &TermCollectionSpec,
542) -> Result<TermCollectionDesign, BasisError> {
543 let policy = gam_runtime::resource::ResourcePolicy::default_library();
544 build_term_collection_design_with_policy(data, spec, &policy)
545}
546
547pub fn build_term_collection_design_with_policy(
551 data: ArrayView2<'_, f64>,
552 spec: &TermCollectionSpec,
553 policy: &gam_runtime::resource::ResourcePolicy,
554) -> Result<TermCollectionDesign, BasisError> {
555 validate_term_collection_finite_inputs(data, spec)?;
556 let mut planned_specs =
557 plan_joint_spatial_centers_for_term_blocks(data, &[spec.smooth_terms.clone()])?;
558 let planned_smooth_terms = planned_specs.pop().ok_or_else(|| {
559 BasisError::InvalidInput(
560 "joint spatial center planner returned no smooth terms for single-spec build"
561 .to_string(),
562 )
563 })?;
564 let mut planned_spec = spec.clone();
565 planned_spec.smooth_terms = planned_smooth_terms;
566 build_term_collection_design_inner_with_policy(data, &planned_spec, policy)
567}
568
569#[derive(Debug, Clone)]
571pub struct TermCollectionDerivativeDesign {
572 pub design: Array2<f64>,
574 pub affine_offset: Array1<f64>,
576}
577
578impl TermCollectionDerivativeDesign {
579 pub fn apply(&self, beta: ArrayView1<'_, f64>) -> Result<Array1<f64>, BasisError> {
581 if beta.len() != self.design.ncols() {
582 crate::bail_dim_basis!(
583 "term-collection derivative coefficient length {} does not match design width {}",
584 beta.len(),
585 self.design.ncols()
586 );
587 }
588 if self.affine_offset.len() != self.design.nrows() {
589 crate::bail_dim_basis!(
590 "term-collection derivative affine offset has {} rows but derivative design has {}",
591 self.affine_offset.len(),
592 self.design.nrows()
593 );
594 }
595 if beta.iter().any(|value| !value.is_finite())
596 || self.affine_offset.iter().any(|value| !value.is_finite())
597 {
598 crate::bail_invalid_basis!(
599 "term-collection derivative coefficients and affine offset must be finite"
600 );
601 }
602 Ok(self.design.dot(&beta.to_owned()) + &self.affine_offset)
603 }
604}
605
606pub fn build_term_collection_derivative_design(
639 data: ArrayView2<'_, f64>,
640 spec: &TermCollectionSpec,
641 deriv_col: usize,
642) -> Result<TermCollectionDerivativeDesign, BasisError> {
643 if deriv_col >= data.ncols() {
644 return Err(BasisError::InvalidInput(format!(
645 "average-derivative column {deriv_col} out of range for data with {} columns",
646 data.ncols()
647 )));
648 }
649
650 let value = build_term_collection_design(data, spec)?;
654 let n = data.nrows();
655 let p_total = value.design.ncols();
656 let mut d = Array2::<f64>::zeros((n, p_total));
657 let mut affine_derivative = Array1::<f64>::zeros(n);
658
659 let p_intercept = value.intercept_range.len();
661 let p_lin = spec.linear_terms.len();
662 let p_rand: usize = value
663 .random_effect_ranges
664 .iter()
665 .map(|(_, range)| range.len())
666 .sum();
667
668 for (j, linear) in spec.linear_terms.iter().enumerate() {
673 let col = p_intercept + j;
674 let derivative = linear_term_derivative_column(data, linear, deriv_col)?;
675 if let Some(column) = derivative {
676 d.column_mut(col).assign(&column);
677 }
678 }
679
680 let smooth_start = p_intercept + p_lin + p_rand;
682 if value.smooth.terms.len() != spec.smooth_terms.len() {
683 return Err(BasisError::InvalidInput(format!(
684 "average-derivative design: value build produced {} smooth terms but spec has {}",
685 value.smooth.terms.len(),
686 spec.smooth_terms.len()
687 )));
688 }
689 for (idx, termspec) in spec.smooth_terms.iter().enumerate() {
690 let term_value = &value.smooth.terms[idx];
691 let feature_cols = smooth_term_feature_cols(termspec);
692 if !feature_cols.contains(&deriv_col) {
693 continue;
695 }
696 let (block, term_affine_derivative) =
697 smooth_term_first_derivative_block(data, termspec, term_value, deriv_col)?;
698 let range = (term_value.coeff_range.start + smooth_start)
699 ..(term_value.coeff_range.end + smooth_start);
700 if block.ncols() != range.len() {
701 return Err(BasisError::DimensionMismatch(format!(
702 "average-derivative design: smooth term '{}' derivative block has {} columns \
703 but the fitted block spans {}",
704 termspec.name,
705 block.ncols(),
706 range.len()
707 )));
708 }
709 d.slice_mut(s![.., range]).assign(&block);
710 if let Some(term_offset) = term_affine_derivative {
711 if term_offset.len() != n {
712 return Err(BasisError::DimensionMismatch(format!(
713 "average-derivative design: smooth term '{}' affine derivative has {} rows but the data has {n}",
714 termspec.name,
715 term_offset.len()
716 )));
717 }
718 affine_derivative += &term_offset;
719 }
720 }
721
722 Ok(TermCollectionDerivativeDesign {
723 design: d,
724 affine_offset: affine_derivative,
725 })
726}
727
728fn linear_term_derivative_column(
736 data: ArrayView2<'_, f64>,
737 linear: &LinearTermSpec,
738 deriv_col: usize,
739) -> Result<Option<Array1<f64>>, BasisError> {
740 let numeric_cols: Vec<usize> = if linear.categorical_levels.is_empty() {
741 linear.effective_feature_cols()
742 } else {
743 linear.feature_cols.clone()
744 };
745 let occurrences = numeric_cols.iter().filter(|&&c| c == deriv_col).count();
746 if occurrences == 0 {
747 return Ok(None);
748 }
749 let n = data.nrows();
750 let p = data.ncols();
751 for &c in &numeric_cols {
752 if c >= p {
753 return Err(BasisError::InvalidInput(format!(
754 "linear term '{}' feature column {c} out of bounds for {p} columns",
755 linear.name
756 )));
757 }
758 }
759
760 let mut gate = Array1::<f64>::ones(n);
762 for &(col, level_bits) in &linear.categorical_levels {
763 if col >= p {
764 return Err(BasisError::InvalidInput(format!(
765 "linear term '{}' categorical column {col} out of bounds for {p} columns",
766 linear.name
767 )));
768 }
769 let level_bits = gam_data::canonical_level_bits(f64::from_bits(level_bits));
770 for (row, g) in gate.iter_mut().enumerate() {
771 if gam_data::canonical_level_bits(data[[row, col]]) != level_bits {
772 *g = 0.0;
773 }
774 }
775 }
776
777 let mut derivative = Array1::<f64>::zeros(n);
779 for (j, &c_j) in numeric_cols.iter().enumerate() {
780 if c_j != deriv_col {
781 continue;
782 }
783 let mut term = gate.clone();
784 for (k, &c_k) in numeric_cols.iter().enumerate() {
785 if k != j {
786 term *= &data.column(c_k);
787 }
788 }
789 derivative += &term;
790 }
791 Ok(Some(derivative))
792}
793
794fn smooth_term_first_derivative_block(
803 data: ArrayView2<'_, f64>,
804 termspec: &SmoothTermSpec,
805 term_value: &SmoothTerm,
806 deriv_col: usize,
807) -> Result<(Array2<f64>, Option<Array1<f64>>), BasisError> {
808 let feature_col = match &termspec.basis {
809 SmoothBasisSpec::BSpline1D { feature_col, .. } => *feature_col,
810 other => {
811 return Err(BasisError::InvalidInput(format!(
812 "analytic average-derivative design only supports non-periodic 1-D B-spline \
813 smooths over the differentiated covariate; term '{}' uses unsupported basis {}",
814 termspec.name,
815 smooth_basis_kind_label(other)
816 )));
817 }
818 };
819 if feature_col != deriv_col {
820 return Err(BasisError::InvalidInput(format!(
824 "analytic average-derivative design: B-spline term '{}' is over column {feature_col}, \
825 not the differentiated column {deriv_col}",
826 termspec.name
827 )));
828 }
829
830 let (knots, degree, transform, periodic, anchor_offset_coeffs) = match &term_value.metadata {
831 BasisMetadata::BSpline1D {
832 knots,
833 degree,
834 identifiability_transform,
835 periodic,
836 anchor_offset_coeffs,
837 ..
838 } => (
839 knots,
840 *degree,
841 identifiability_transform.as_ref(),
842 periodic,
843 anchor_offset_coeffs.as_ref(),
844 ),
845 other => {
846 return Err(BasisError::InvalidInput(format!(
847 "analytic average-derivative design expected B-spline metadata for term '{}', \
848 found {other:?}",
849 termspec.name
850 )));
851 }
852 };
853 if periodic.is_some() {
854 return Err(BasisError::InvalidInput(format!(
855 "analytic average-derivative design does not support periodic/cyclic B-spline \
856 term '{}'",
857 termspec.name
858 )));
859 }
860 let degree = degree.ok_or_else(|| {
861 BasisError::InvalidInput(format!(
862 "B-spline term '{}' metadata is missing its effective degree",
863 termspec.name
864 ))
865 })?;
866
867 let (deriv_basis_arc, _) = crate::basis::create_basis::<crate::basis::Dense>(
869 data.column(deriv_col),
870 crate::basis::KnotSource::Provided(knots.view()),
871 degree,
872 crate::basis::BasisOptions::first_derivative(),
873 )?;
874 let deriv_basis = deriv_basis_arc.as_ref();
875
876 let affine_derivative = match anchor_offset_coeffs {
877 Some(beta_p) => {
878 if deriv_basis.ncols() != beta_p.len() {
879 return Err(BasisError::DimensionMismatch(format!(
880 "B-spline term '{}': raw derivative basis has {} columns but the affine anchor lift has {} coefficients",
881 termspec.name,
882 deriv_basis.ncols(),
883 beta_p.len()
884 )));
885 }
886 Some(deriv_basis.dot(beta_p))
887 }
888 None => None,
889 };
890
891 let block = match transform {
895 Some(z) => {
896 if deriv_basis.ncols() != z.nrows() {
897 return Err(BasisError::DimensionMismatch(format!(
898 "B-spline term '{}': raw derivative basis has {} columns but the frozen \
899 identifiability transform has {} rows",
900 termspec.name,
901 deriv_basis.ncols(),
902 z.nrows()
903 )));
904 }
905 gam_linalg::faer_ndarray::fast_ab(deriv_basis, z)
906 }
907 None => deriv_basis.to_owned(),
908 };
909 Ok((block, affine_derivative))
910}
911
912fn smooth_basis_kind_label(basis: &SmoothBasisSpec) -> &'static str {
915 match basis {
916 SmoothBasisSpec::BSpline1D { .. } => "BSpline1D",
917 SmoothBasisSpec::TensorBSpline { .. } => "TensorBSpline",
918 SmoothBasisSpec::ByVariable { .. } => "ByVariable",
919 SmoothBasisSpec::FactorSumToZero { .. } => "FactorSumToZero",
920 SmoothBasisSpec::FactorSmooth { .. } => "FactorSmooth",
921 SmoothBasisSpec::BySmooth { .. } => "BySmooth",
922 SmoothBasisSpec::ThinPlate { .. } => "ThinPlate",
923 SmoothBasisSpec::Duchon { .. } => "Duchon",
924 SmoothBasisSpec::Matern { .. } => "Matern",
925 SmoothBasisSpec::Sphere { .. } => "Sphere",
926 SmoothBasisSpec::ConstantCurvature { .. } => "ConstantCurvature",
927 SmoothBasisSpec::MeasureJet { .. } => "MeasureJet",
928 SmoothBasisSpec::Pca { .. } => "Pca",
929 }
930}
931
932enum GlobalIdentifiabilityPlan {
940 Absent,
942 Delete { block: Array2<f64> },
945 Residualize { block: Array2<f64> },
948}
949
950impl GlobalIdentifiabilityPlan {
951 fn as_gauge(
953 &self,
954 owner_terms: &[usize],
955 has_parametric_block: bool,
956 local_identifiability_transform: Option<Array2<f64>>,
957 coefficient_transform: Array2<f64>,
958 local_columns: usize,
959 joint_null_rotation: Option<crate::basis::JointNullRotation>,
960 ) -> Option<SmoothCollectionGauge> {
961 let (arm, block) = match self {
962 Self::Absent => return None,
963 Self::Delete { block } => (SmoothCollectionGaugeArm::Delete, block),
964 Self::Residualize { block } => (SmoothCollectionGaugeArm::Residualize, block),
965 };
966 Some(SmoothCollectionGauge {
967 arm,
968 constraint_block: block.clone(),
969 owner_terms: owner_terms.to_vec(),
970 has_parametric_block,
971 local_identifiability_transform,
972 coefficient_transform,
973 local_columns,
974 joint_null_rotation,
975 })
976 }
977}
978
979fn basis_local_identifiability_transform(metadata: &BasisMetadata) -> Option<Array2<f64>> {
991 match metadata {
992 BasisMetadata::BSpline1D {
993 identifiability_transform,
994 ..
995 }
996 | BasisMetadata::CubicRegression1D {
997 identifiability_transform,
998 ..
999 }
1000 | BasisMetadata::ThinPlate {
1001 identifiability_transform,
1002 ..
1003 }
1004 | BasisMetadata::Matern {
1005 identifiability_transform,
1006 ..
1007 }
1008 | BasisMetadata::Duchon {
1009 identifiability_transform,
1010 ..
1011 }
1012 | BasisMetadata::TensorBSpline {
1013 identifiability_transform,
1014 ..
1015 } => identifiability_transform.clone(),
1016 BasisMetadata::Sphere {
1017 constraint_transform,
1018 ..
1019 }
1020 | BasisMetadata::ConstantCurvature {
1021 constraint_transform,
1022 ..
1023 }
1024 | BasisMetadata::MeasureJet {
1025 constraint_transform,
1026 ..
1027 } => constraint_transform.clone(),
1028 BasisMetadata::Pca { .. }
1029 | BasisMetadata::SphereHarmonics { .. }
1030 | BasisMetadata::BySmooth { .. }
1031 | BasisMetadata::FactorSmooth { .. } => None,
1032 }
1033}
1034
1035pub struct RealizedCollectionGauge {
1037 pub design: DesignMatrix,
1039 pub coefficient_transform: Array2<f64>,
1042 pub residualization: crate::basis::ParametricResidualization,
1046}
1047
1048fn derive_smooth_collection_coefficient_transform(
1057 design_local: &DesignMatrix,
1058 arm: SmoothCollectionGaugeArm,
1059 block: ArrayView2<'_, f64>,
1060 has_owner_terms: bool,
1061) -> Result<Array2<f64>, BasisError> {
1062 match arm {
1063 SmoothCollectionGaugeArm::Delete => {
1064 match orthogonality_transform_for_design(design_local, block, None) {
1065 Ok(transform) => Ok(transform),
1066 Err(BasisError::ConstraintNullspaceCollapsed { .. }) if has_owner_terms => {
1070 Ok(Array2::zeros((design_local.ncols(), 0)))
1071 }
1072 Err(error) => Err(error),
1073 }
1074 }
1075 SmoothCollectionGaugeArm::Residualize => {
1076 Ok(crate::basis::parametric_residualization_for_design(
1077 design_local,
1078 block,
1079 None, )?
1081 .coefficient_transform)
1082 }
1083 }
1084}
1085
1086pub fn realize_smooth_collection_gauge(
1105 design_local: DesignMatrix,
1106 gauge: &SmoothCollectionGauge,
1107 termname: &str,
1108) -> Result<RealizedCollectionGauge, BasisError> {
1109 let block = gauge.constraint_block.view();
1110 if block.nrows() != design_local.nrows() {
1111 gam_problem::bail_dim_basis!(
1112 "collection gauge row mismatch for term '{termname}': the design has {} rows and the frozen constraint block has {}",
1113 design_local.nrows(),
1114 block.nrows()
1115 );
1116 }
1117 if gauge.coefficient_transform.nrows() != gauge.local_columns {
1118 gam_problem::bail_dim_basis!(
1119 "collection gauge for term '{termname}' declares {} local columns but its fixed coefficient chart has {} rows",
1120 gauge.local_columns,
1121 gauge.coefficient_transform.nrows()
1122 );
1123 }
1124 if design_local.ncols() != gauge.local_columns {
1125 gam_problem::bail_dim_basis!(
1126 "collection gauge local width mismatch for term '{termname}': the design has {} columns but the fixed chart was derived on {}",
1127 design_local.ncols(),
1128 gauge.local_columns
1129 );
1130 }
1131 let coefficient_transform = gauge.coefficient_transform.clone();
1132 let design = apply_smooth_transform_to_design(
1133 design_local,
1134 &coefficient_transform,
1135 termname,
1136 )?;
1137 let projector = crate::basis::FixedRowSpaceProjector::from_constraint_block(block)?;
1138 let (design, row_space_correction) = projector.project_design(design, termname)?;
1139 let residualization = crate::basis::ParametricResidualization {
1140 coefficient_transform: coefficient_transform.clone(),
1141 row_space_correction,
1142 };
1143 assert_orthogonal_to_constraint_block(&design, block, termname)?;
1144 Ok(RealizedCollectionGauge {
1145 design,
1146 coefficient_transform,
1147 residualization,
1148 })
1149}
1150
1151pub struct LocalTermRealization<'a> {
1156 pub design: DesignMatrix,
1158 pub metadata: &'a BasisMetadata,
1160 pub active_penalties: &'a [ActivePenalty],
1161 pub dropped_penalties: Vec<DroppedPenaltyInfo>,
1162 pub linear_constraints_local: Option<&'a gam_problem::LinearInequalityConstraints>,
1163 pub joint_null_rotation: Option<&'a crate::basis::JointNullRotation>,
1166 pub termname: &'a str,
1167}
1168
1169pub struct CollectionGaugedTerm {
1171 pub design: DesignMatrix,
1172 pub metadata: BasisMetadata,
1173 pub active_penalties: Vec<ActivePenalty>,
1174 pub dropped_penalties: Vec<DroppedPenaltyInfo>,
1175 pub linear_constraints_local: Option<gam_problem::LinearInequalityConstraints>,
1176 pub parametric_residualization: Option<ParametricResidualizationChart>,
1177}
1178
1179pub fn place_term_in_collection_gauge(
1199 gauge: &SmoothCollectionGauge,
1200 local: LocalTermRealization<'_>,
1201) -> Result<CollectionGaugedTerm, BasisError> {
1202 let LocalTermRealization {
1203 design,
1204 metadata,
1205 active_penalties,
1206 dropped_penalties,
1207 linear_constraints_local,
1208 joint_null_rotation,
1209 termname,
1210 } = local;
1211 let realized = realize_smooth_collection_gauge(design, gauge, termname)?;
1212 let coefficient_gauge =
1213 gam_problem::Gauge::from_block_transforms(&[realized.coefficient_transform.clone()]);
1214 let candidates = penalty_candidates_under_collection_gauge(
1215 active_penalties,
1216 Some(&coefficient_gauge),
1217 termname,
1218 )?;
1219 let filtered = filter_penalty_candidates(candidates)?;
1220 let mut dropped_penalties = dropped_penalties;
1221 dropped_penalties.extend(filtered.dropped);
1222 let linear_constraints_local = linear_constraints_local.map(|lin| {
1223 gam_problem::LinearInequalityConstraints {
1224 a: lin.a.dot(&coefficient_gauge.block_transform(0)),
1225 b: lin.b.clone(),
1226 }
1227 });
1228 let realized_transform = match joint_null_rotation {
1229 Some(rotation) => {
1230 gam_linalg::faer_ndarray::fast_ab(&rotation.rotation, &realized.coefficient_transform)
1231 }
1232 None => realized.coefficient_transform.clone(),
1233 };
1234 let metadata = with_identifiability_transform(metadata, Some(&realized_transform))?;
1235 let parametric_residualization = Some(ParametricResidualizationChart {
1236 owner_terms: gauge.owner_terms.clone(),
1237 has_parametric_block: gauge.has_parametric_block,
1238 correction: realized.residualization.row_space_correction.clone(),
1239 });
1240 Ok(CollectionGaugedTerm {
1241 design: realized.design,
1242 metadata,
1243 active_penalties: filtered.active,
1244 dropped_penalties,
1245 linear_constraints_local,
1246 parametric_residualization,
1247 })
1248}
1249
1250const ORTHOGONALITY_REL_RESIDUAL_TOL: f64 = 1e-8;
1253
1254fn assert_orthogonal_to_constraint_block(
1255 design: &DesignMatrix,
1256 constraint: ArrayView2<'_, f64>,
1257 termname: &str,
1258) -> Result<(), BasisError> {
1259 let rel = orthogonality_relative_residual_for_design(design, constraint)?;
1260 if rel > ORTHOGONALITY_REL_RESIDUAL_TOL {
1261 gam_problem::bail_invalid_basis!(
1262 "smooth orthogonality residual too large for term '{}': {:.3e} > {:.1e}",
1263 termname,
1264 rel,
1265 ORTHOGONALITY_REL_RESIDUAL_TOL
1266 );
1267 }
1268 Ok(())
1269}
1270
1271fn subtract_row_space_correction(
1276 design: DesignMatrix,
1277 constraint: ArrayView2<'_, f64>,
1278 correction: ArrayView2<'_, f64>,
1279 termname: &str,
1280) -> Result<DesignMatrix, BasisError> {
1281 use gam_linalg::matrix::{BlockDesignOperator, DesignBlock};
1282 let p = design.ncols();
1283 let q = constraint.ncols();
1284 let k = correction.ncols();
1285 if correction.nrows() != q || p != k {
1286 return Err(BasisError::InvalidInput(format!(
1287 "row-space correction shape mismatch for term '{termname}': design is {}x{p}, \
1288 constraint is {}x{q}, correction is {}x{k}",
1289 design.nrows(),
1290 constraint.nrows(),
1291 correction.nrows(),
1292 )));
1293 }
1294 if q == 0 {
1295 return Ok(design);
1296 }
1297 let design_block = match design {
1298 DesignMatrix::Dense(inner) => DesignBlock::Dense(inner),
1299 DesignMatrix::Sparse(inner) => DesignBlock::Sparse(inner),
1300 };
1301 let stacked = BlockDesignOperator::new(vec![
1302 design_block,
1303 DesignBlock::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
1304 constraint.to_owned(),
1305 )),
1306 ])
1307 .map_err(BasisError::InvalidInput)?;
1308 let mut transform = Array2::<f64>::zeros((p + q, k));
1311 for i in 0..p {
1312 transform[[i, i]] = 1.0;
1313 }
1314 for i in 0..q {
1315 for j in 0..k {
1316 transform[[p + i, j]] = -correction[[i, j]];
1317 }
1318 }
1319 let operator = gam_linalg::matrix::CoefficientTransformOperator::new(
1320 gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(stacked)),
1321 transform,
1322 )
1323 .map_err(|e| {
1324 BasisError::InvalidInput(format!(
1325 "row-space correction failed for term '{termname}': {e}"
1326 ))
1327 })?;
1328 Ok(DesignMatrix::Dense(
1329 gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(operator)),
1330 ))
1331}
1332
1333fn build_constraint_block(
1334 n: usize,
1335 parametric_block: Option<&Array2<f64>>,
1336 owner_blocks: &[&DesignMatrix],
1337) -> Result<Array2<f64>, BasisError> {
1338 let param_cols = parametric_block.map_or(0, |mat| mat.ncols());
1339 let owner_cols: usize = owner_blocks.iter().map(|design| design.ncols()).sum();
1340 let mut block = Array2::<f64>::zeros((n, param_cols + owner_cols));
1341 let mut col_start = 0usize;
1342 if let Some(parametric) = parametric_block {
1343 let col_end = col_start + parametric.ncols();
1344 block
1345 .slice_mut(s![.., col_start..col_end])
1346 .assign(parametric);
1347 col_start = col_end;
1348 }
1349 const CHUNK: usize = 1024;
1350 for owner in owner_blocks {
1351 let col_end = col_start + owner.ncols();
1352 for row_start in (0..n).step_by(CHUNK) {
1353 let row_end = (row_start + CHUNK).min(n);
1354 let chunk = (*owner)
1355 .try_row_chunk(row_start..row_end)
1356 .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
1357 block
1358 .slice_mut(s![row_start..row_end, col_start..col_end])
1359 .assign(&chunk);
1360 }
1361 col_start = col_end;
1362 }
1363 Ok(block)
1364}
1365
1366fn design_cross_relative_residual(
1367 lhs: &DesignMatrix,
1368 rhs: &DesignMatrix,
1369) -> Result<f64, BasisError> {
1370 let n = lhs.nrows();
1371 if rhs.nrows() != n {
1372 return Err(BasisError::ConstraintMatrixRowMismatch {
1373 basisrows: n,
1374 constraintrows: rhs.nrows(),
1375 });
1376 }
1377 const CHUNK: usize = 1024;
1378 let mut cross = Array2::<f64>::zeros((lhs.ncols(), rhs.ncols()));
1379 let mut lhs_sumsq = 0.0;
1380 let mut rhs_sumsq = 0.0;
1381 for start in (0..n).step_by(CHUNK) {
1382 let end = (start + CHUNK).min(n);
1383 let lhs_chunk = lhs
1384 .try_row_chunk(start..end)
1385 .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
1386 let rhs_chunk = rhs
1387 .try_row_chunk(start..end)
1388 .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
1389 cross += &gam_linalg::faer_ndarray::fast_atb(&lhs_chunk, &rhs_chunk);
1390 lhs_sumsq += lhs_chunk.iter().map(|v| v * v).sum::<f64>();
1391 rhs_sumsq += rhs_chunk.iter().map(|v| v * v).sum::<f64>();
1392 }
1393 let num = cross.iter().map(|v| v * v).sum::<f64>().sqrt();
1394 let denom = lhs_sumsq.sqrt() * rhs_sumsq.sqrt();
1395 if denom == 0.0 {
1398 return Ok(0.0);
1399 }
1400 Ok(num / denom)
1401}
1402
1403fn smooth_has_overlapping_linear_terms(
1404 linear_terms: &[LinearTermSpec],
1405 termspec: &SmoothTermSpec,
1406) -> bool {
1407 let feature_cols = smooth_term_feature_cols(termspec);
1408 linear_terms
1409 .iter()
1410 .any(|linear| feature_cols.contains(&linear.feature_col))
1411}
1412
1413pub fn smooth_intrinsic_parametric_feature_cols(
1417 linear_terms: &[LinearTermSpec],
1418 term: &SmoothTermSpec,
1419) -> Vec<usize> {
1420 let feature_cols = smooth_term_feature_cols(term);
1435 let mut owned = Vec::new();
1436 for linear in linear_terms {
1437 if feature_cols.contains(&linear.feature_col) && !owned.contains(&linear.feature_col) {
1438 owned.push(linear.feature_col);
1439 }
1440 }
1441 owned
1442}
1443
1444fn apply_global_smooth_identifiability(
1445 smooth: RawSmoothDesign,
1446 data: ArrayView2<'_, f64>,
1447 linear_terms: &[LinearTermSpec],
1448 smoothspecs: &[SmoothTermSpec],
1449) -> Result<(SmoothDesign, Array1<f64>), BasisError> {
1450 if smoothspecs.len() != smooth.terms.len() {
1461 gam_problem::bail_dim_basis!(
1462 "smooth spec count ({}) does not match built term count ({})",
1463 smoothspecs.len(),
1464 smooth.terms.len()
1465 );
1466 }
1467
1468 if smooth.terms.is_empty() {
1469 let RawSmoothDesign {
1470 term_designs,
1471 affine_offset,
1472 penalties,
1473 nullspace_dims,
1474 penaltyinfo,
1475 dropped_penaltyinfo,
1476 terms,
1477 coefficient_lower_bounds,
1478 linear_constraints,
1479 } = smooth;
1480 return Ok((
1481 SmoothDesign {
1482 term_designs,
1483 penalties,
1484 nullspace_dims,
1485 penaltyinfo,
1486 dropped_penaltyinfo,
1487 terms,
1488 coefficient_lower_bounds,
1489 linear_constraints,
1490 },
1491 affine_offset,
1492 ));
1493 }
1494
1495 let mut local_designs = vec![None; smooth.terms.len()];
1496 let mut local_active_penalties = vec![Vec::<ActivePenalty>::new(); smooth.terms.len()];
1497 let mut local_dropped_penalties = vec![Vec::<DroppedPenaltyInfo>::new(); smooth.terms.len()];
1498 let mut local_metadata = vec![None; smooth.terms.len()];
1499 let mut local_dims = vec![0usize; smooth.terms.len()];
1500 let mut local_linear_constraints = vec![None; smooth.terms.len()];
1501 let mut local_unabsorbed_z = vec![None::<Array2<f64>>; smooth.terms.len()];
1502 let mut local_residualization =
1503 vec![None::<ParametricResidualizationChart>; smooth.terms.len()];
1504 let mut local_collection_gauge = vec![None::<SmoothCollectionGauge>; smooth.terms.len()];
1505
1506 let SmoothStructureAnalysis {
1507 ownership_order,
1508 term_owners,
1509 ..
1510 } = analyze_smooth_ownership(smoothspecs);
1511
1512 use rayon::iter::{IntoParallelIterator, ParallelIterator};
1513
1514 for &idx in &ownership_order {
1515 let term = &smooth.terms[idx];
1516 let termspec = &smoothspecs[idx];
1517 let design_local = smooth.term_designs[idx].clone();
1518 let replay_z = frozen_global_orthogonality(termspec);
1526 let skip_global_transform = replay_z.is_none()
1527 && (smooth_has_frozen_identifiability(termspec) || term.lower_bounds_local.is_some());
1528 let owner_indices = if replay_z.is_some()
1542 || skip_global_transform
1543 || termspec.basis.is_marginally_centered_tensor()
1544 || termspec.basis.is_sum_to_zero_factor_smooth()
1545 {
1546 Vec::new()
1547 } else {
1548 const OVERLAP_REL_RESIDUAL_TOL: f64 = 1e-10;
1552 let owner_cross_checks = term_owners[idx]
1553 .clone()
1554 .into_par_iter()
1555 .map(|owner_idx| {
1556 let owner_design = local_designs[owner_idx]
1557 .as_ref()
1558 .expect("owner design must be available before dependent smooth");
1559 design_cross_relative_residual(&design_local, owner_design)
1560 .map(|rel| (owner_idx, rel))
1561 })
1562 .collect::<Vec<_>>();
1563 let mut out = Vec::new();
1564 for check in owner_cross_checks {
1565 let (owner_idx, rel) = check?;
1566 if rel > OVERLAP_REL_RESIDUAL_TOL {
1567 out.push(owner_idx);
1568 }
1569 }
1570 out
1571 };
1572 let owner_blocks = owner_indices
1573 .iter()
1574 .map(|owner_idx| {
1575 local_designs[*owner_idx]
1576 .as_ref()
1577 .expect("owner design must be available before dependent smooth")
1578 })
1579 .collect::<Vec<_>>();
1580 let replay_correction = frozen_parametric_residualization(termspec);
1585 let needs_parametric_block = match replay_correction {
1586 Some(chart) => chart.has_parametric_block,
1587 None => {
1588 replay_z.is_none()
1589 && !skip_global_transform
1590 && (smooth_has_overlapping_linear_terms(linear_terms, termspec)
1591 || !smooth_intrinsic_parametric_feature_cols(linear_terms, termspec)
1592 .is_empty()
1593 || smooth_requires_parametric_orthogonality(termspec)
1594 || factor_by_level_gate(termspec).is_some())
1600 }
1601 };
1602 let parametric_block = if !needs_parametric_block {
1603 None
1604 } else {
1605 Some(build_parametric_constraint_block_for_term(
1606 data,
1607 linear_terms,
1608 termspec,
1609 )?)
1610 };
1611 let replay_owner_blocks = match replay_correction {
1615 Some(chart) => chart
1616 .owner_terms
1617 .iter()
1618 .map(|owner_idx| {
1619 local_designs.get(*owner_idx).and_then(|slot| slot.as_ref()).ok_or_else(|| {
1620 BasisError::InvalidInput(format!(
1621 "term '{}' replays a parametric residualization against owner term {owner_idx}, which is not available at this point of the rebuild",
1622 termspec.name
1623 ))
1624 })
1625 })
1626 .collect::<Result<Vec<_>, _>>()?,
1627 None => Vec::new(),
1628 };
1629 let plan = if replay_correction.is_some() {
1655 GlobalIdentifiabilityPlan::Absent
1657 } else if skip_global_transform
1658 || (parametric_block.is_none() && owner_blocks.is_empty())
1659 {
1660 GlobalIdentifiabilityPlan::Absent
1661 } else {
1662 let raw =
1663 build_constraint_block(data.nrows(), parametric_block.as_ref(), &owner_blocks)?;
1664 let contained =
1665 crate::basis::contained_constraint_directions(&design_local, raw.view(), None)?;
1666 if raw.ncols() == 0 {
1667 GlobalIdentifiabilityPlan::Absent
1668 } else if contained.ncols() == raw.ncols() {
1669 GlobalIdentifiabilityPlan::Delete { block: contained }
1673 } else {
1674 GlobalIdentifiabilityPlan::Residualize { block: raw }
1675 }
1676 };
1677 let collection_coefficient_transform = match &plan {
1687 GlobalIdentifiabilityPlan::Absent => None,
1688 GlobalIdentifiabilityPlan::Delete { block } => Some(
1689 derive_smooth_collection_coefficient_transform(
1690 &design_local,
1691 SmoothCollectionGaugeArm::Delete,
1692 block.view(),
1693 !owner_indices.is_empty(),
1694 )?,
1695 ),
1696 GlobalIdentifiabilityPlan::Residualize { block } => Some(
1697 derive_smooth_collection_coefficient_transform(
1698 &design_local,
1699 SmoothCollectionGaugeArm::Residualize,
1700 block.view(),
1701 !owner_indices.is_empty(),
1702 )?,
1703 ),
1704 };
1705 let collection_gauge = collection_coefficient_transform.map(|transform| {
1706 plan.as_gauge(
1707 &owner_indices,
1708 parametric_block.is_some(),
1709 basis_local_identifiability_transform(&term.metadata),
1710 transform,
1711 design_local.ncols(),
1712 term.joint_null_rotation.clone(),
1717 )
1718 .expect("a derived collection coefficient chart implies a present gauge")
1719 });
1720 let mut residualization: Option<crate::basis::ParametricResidualization> = None;
1721 let (design_constrained, z_opt) = if let Some(gauge) = collection_gauge.as_ref() {
1722 let realized = realize_smooth_collection_gauge(design_local, gauge, &term.name)?;
1730 residualization = Some(realized.residualization);
1731 (realized.design, Some(realized.coefficient_transform))
1732 } else {
1733 let z_opt = if let Some(z) = replay_z {
1737 if design_local.ncols() != z.nrows() {
1738 gam_problem::bail_dim_basis!(
1739 "frozen global-orthogonality transform mismatch for term '{}': rebuilt design has {} columns but the persisted fit-time transform has {} rows",
1740 term.name,
1741 design_local.ncols(),
1742 z.nrows()
1743 );
1744 }
1745 Some(z.clone())
1746 } else if skip_global_transform {
1747 None
1748 } else {
1749 maybe_smooth_identifiability_transform(termspec, &design_local, None)?
1752 };
1753 let design_transformed = match z_opt.as_ref() {
1754 Some(z) => apply_smooth_transform_to_design(design_local, z, &term.name)?,
1755 None => design_local,
1756 };
1757 let design_constrained = match replay_correction {
1758 Some(chart) => {
1759 let block = build_constraint_block(
1764 data.nrows(),
1765 parametric_block.as_ref(),
1766 &replay_owner_blocks,
1767 )?;
1768 if block.ncols() != chart.correction.nrows() {
1769 gam_problem::bail_dim_basis!(
1770 "frozen parametric residualization mismatch for term '{}': rebuilt constraint block has {} columns but the persisted fit-time correction has {} rows",
1771 term.name,
1772 block.ncols(),
1773 chart.correction.nrows()
1774 );
1775 }
1776 subtract_row_space_correction(
1777 design_transformed,
1778 block.view(),
1779 chart.correction.view(),
1780 &term.name,
1781 )?
1782 }
1783 None => design_transformed,
1784 };
1785 (design_constrained, z_opt)
1786 };
1787 let coefficient_gauge = z_opt
1788 .as_ref()
1789 .map(|z| gam_problem::Gauge::from_block_transforms(&[z.clone()]));
1790
1791 let penalty_candidates = penalty_candidates_under_collection_gauge(
1792 &term.active_penalties,
1793 coefficient_gauge.as_ref(),
1794 &term.name,
1795 )?;
1796 let filtered = filter_penalty_candidates(penalty_candidates)?;
1797 let linear_constraints_constrained =
1798 if let Some(lin_local) = term.linear_constraints_local.as_ref() {
1799 if let Some(gauge) = coefficient_gauge.as_ref() {
1800 Some(LinearInequalityConstraints {
1801 a: lin_local.a.dot(&gauge.block_transform(0)),
1802 b: lin_local.b.clone(),
1803 })
1804 } else {
1805 Some(lin_local.clone())
1806 }
1807 } else {
1808 None
1809 };
1810
1811 local_residualization[idx] = residualization
1814 .as_ref()
1815 .map(|plan| ParametricResidualizationChart {
1816 owner_terms: owner_indices.clone(),
1817 has_parametric_block: parametric_block.is_some(),
1818 correction: plan.row_space_correction.clone(),
1819 })
1820 .or_else(|| replay_correction.cloned());
1821 local_collection_gauge[idx] = collection_gauge;
1822 local_dims[idx] = design_constrained.ncols();
1823 local_designs[idx] = Some(design_constrained);
1824 local_active_penalties[idx] = filtered.active;
1825 local_dropped_penalties[idx] = term.dropped_penalties.clone();
1826 local_dropped_penalties[idx].extend(filtered.dropped);
1827 local_linear_constraints[idx] = linear_constraints_constrained;
1828 let realized_transform = match (term.joint_null_rotation.as_ref(), z_opt.as_ref()) {
1829 (Some(rotation), Some(z)) => {
1830 Some(gam_linalg::faer_ndarray::fast_ab(&rotation.rotation, z))
1831 }
1832 (Some(rotation), None) => Some(rotation.rotation.clone()),
1833 (None, Some(z)) => Some(z.clone()),
1834 (None, None) => None,
1835 };
1836 match &termspec.basis {
1861 SmoothBasisSpec::FactorSumToZero { .. } | SmoothBasisSpec::FactorSmooth { .. } => {
1862 local_metadata[idx] = Some(term.metadata.clone());
1863 local_unabsorbed_z[idx] = z_opt.clone();
1864 }
1865 _ => {
1866 local_metadata[idx] = Some(with_identifiability_transform(
1867 &term.metadata,
1868 realized_transform.as_ref(),
1869 )?);
1870 }
1871 }
1872 }
1873
1874 let total_p: usize = local_dims.iter().sum();
1875 let mut terms_out = Vec::<SmoothTerm>::with_capacity(smooth.terms.len());
1876 let mut penalties_global = Vec::<BlockwisePenalty>::new();
1877 let mut nullspace_dims_global = Vec::<usize>::new();
1878 let mut penaltyinfo_global = Vec::<PenaltyBlockInfo>::new();
1879 let mut dropped_penaltyinfo_global = Vec::<DroppedPenaltyBlockInfo>::new();
1880 let mut coefficient_lower_bounds = Array1::<f64>::from_elem(total_p, f64::NEG_INFINITY);
1881 let mut any_bounds = false;
1882 let mut linear_constraintsrows: Vec<Array1<f64>> = Vec::new();
1883 let mut linear_constraints_b: Vec<f64> = Vec::new();
1884
1885 let mut col_start = 0usize;
1886 for idx in 0..smooth.terms.len() {
1887 let p_local = local_dims[idx];
1888 let col_end = col_start + p_local;
1889
1890 for active_penalty in &local_active_penalties[idx] {
1891 let global_index = penalties_global.len();
1892 penalties_global.push(BlockwisePenalty::new(
1893 col_start..col_end,
1894 active_penalty.matrix.clone(),
1895 ));
1896 nullspace_dims_global.push(active_penalty.nullity);
1897 penaltyinfo_global.push(PenaltyBlockInfo {
1898 global_index,
1899 termname: Some(smooth.terms[idx].name.clone()),
1900 penalty: active_penalty.info.clone(),
1901 });
1902 }
1903 for info in &local_dropped_penalties[idx] {
1904 dropped_penaltyinfo_global.push(DroppedPenaltyBlockInfo {
1905 termname: Some(smooth.terms[idx].name.clone()),
1906 penalty: info.clone(),
1907 });
1908 }
1909
1910 terms_out.push(SmoothTerm {
1911 name: smooth.terms[idx].name.clone(),
1912 coeff_range: col_start..col_end,
1913 shape: smooth.terms[idx].shape,
1914 active_penalties: local_active_penalties[idx].clone(),
1915 dropped_penalties: local_dropped_penalties[idx].clone(),
1916 metadata: local_metadata[idx]
1917 .clone()
1918 .expect("local metadata must exist for every smooth term"),
1919 lower_bounds_local: smooth.terms[idx].lower_bounds_local.clone(),
1920 linear_constraints_local: local_linear_constraints[idx].clone(),
1921 kronecker_factored: None,
1923 joint_null_rotation: None,
1929 unabsorbed_global_orthogonality: local_unabsorbed_z[idx].clone(),
1932 parametric_residualization: local_residualization[idx].clone(),
1933 collection_gauge: local_collection_gauge[idx].clone(),
1936 });
1937 if let Some(lin_local) = &local_linear_constraints[idx] {
1938 for r in 0..lin_local.a.nrows() {
1939 let mut row = Array1::<f64>::zeros(total_p);
1940 row.slice_mut(s![col_start..col_end])
1941 .assign(&lin_local.a.row(r));
1942 linear_constraintsrows.push(row);
1943 linear_constraints_b.push(lin_local.b[r]);
1944 }
1945 }
1946 if let Some(lb_local) = smooth.terms[idx].lower_bounds_local.as_ref()
1947 && lb_local.len() == p_local
1948 {
1949 coefficient_lower_bounds
1950 .slice_mut(s![col_start..col_end])
1951 .assign(lb_local);
1952 any_bounds = true;
1953 }
1954
1955 col_start = col_end;
1956 }
1957
1958 assert_eq!(
1959 penalties_global.len(),
1960 nullspace_dims_global.len(),
1961 "globally reparameterized smooth penalty/nullspace bookkeeping diverged"
1962 );
1963 assert_eq!(
1964 penalties_global.len(),
1965 penaltyinfo_global.len(),
1966 "globally reparameterized smooth penalty metadata bookkeeping diverged"
1967 );
1968
1969 Ok((
1970 SmoothDesign {
1971 term_designs: local_designs
1972 .into_iter()
1973 .map(|design| design.expect("local design must exist for every smooth term"))
1974 .collect(),
1975 penalties: penalties_global,
1976 nullspace_dims: nullspace_dims_global,
1977 penaltyinfo: penaltyinfo_global,
1978 dropped_penaltyinfo: dropped_penaltyinfo_global,
1979 terms: terms_out,
1980 coefficient_lower_bounds: if any_bounds {
1981 Some(coefficient_lower_bounds)
1982 } else {
1983 None
1984 },
1985 linear_constraints: if linear_constraintsrows.is_empty() {
1986 None
1987 } else {
1988 let mut a = Array2::<f64>::zeros((linear_constraintsrows.len(), total_p));
1989 for (i, row) in linear_constraintsrows.iter().enumerate() {
1990 a.row_mut(i).assign(row);
1991 }
1992 Some(LinearInequalityConstraints {
1993 a,
1994 b: Array1::from_vec(linear_constraints_b),
1995 })
1996 },
1997 },
1998 smooth.affine_offset,
1999 ))
2000}
2001
2002fn factor_by_level_gate(termspec: &SmoothTermSpec) -> Option<(usize, u64)> {
2025 match &termspec.basis {
2026 SmoothBasisSpec::ByVariable {
2027 by_col,
2028 by: ByVariableSpec::Level { value_bits, .. },
2029 ..
2030 } => Some((*by_col, *value_bits)),
2031 _ => None,
2032 }
2033}
2034
2035fn build_parametric_constraint_block_for_term(
2036 data: ArrayView2<'_, f64>,
2037 linear_terms: &[LinearTermSpec],
2038 termspec: &SmoothTermSpec,
2039) -> Result<Array2<f64>, BasisError> {
2040 let n = data.nrows();
2041 let p_data = data.ncols();
2042
2043 if let Some((by_col, value_bits)) = factor_by_level_gate(termspec) {
2047 if by_col >= p_data {
2048 gam_problem::bail_dim_basis!(
2049 "factor-by smooth term '{}' by column {by_col} out of bounds for {p_data} columns",
2050 termspec.name
2051 );
2052 }
2053 let mut c = Array2::<f64>::zeros((n, 1));
2054 let by = data.column(by_col);
2055 let value_bits = gam_data::canonical_level_bits(f64::from_bits(value_bits));
2056 for (row, &value) in by.iter().enumerate() {
2057 if gam_data::canonical_level_bits(value) == value_bits {
2058 c[[row, 0]] = 1.0;
2059 }
2060 }
2061 return Ok(c);
2062 }
2063
2064 let feature_cols = smooth_term_feature_cols(termspec);
2065 let mut parametric_cols = smooth_intrinsic_parametric_feature_cols(linear_terms, termspec);
2066 for &feature_col in ¶metric_cols {
2067 if feature_col >= p_data {
2068 gam_problem::bail_dim_basis!(
2069 "smooth term feature column {feature_col} out of bounds for {p_data} columns"
2070 );
2071 }
2072 }
2073 for linear in linear_terms
2074 .iter()
2075 .filter(|linear| feature_cols.contains(&linear.feature_col))
2076 {
2077 if linear.feature_col >= p_data {
2078 gam_problem::bail_dim_basis!(
2079 "linear term '{}' feature column {} out of bounds for {} columns",
2080 linear.name,
2081 linear.feature_col,
2082 p_data
2083 );
2084 }
2085 if !parametric_cols.contains(&linear.feature_col) {
2086 parametric_cols.push(linear.feature_col);
2087 }
2088 }
2089
2090 let mut c = Array2::<f64>::zeros((n, 1 + parametric_cols.len()));
2091 c.column_mut(0).fill(1.0);
2092 for (j, &feature_col) in parametric_cols.iter().enumerate() {
2093 c.column_mut(j + 1).assign(&data.column(feature_col));
2094 }
2095 Ok(c)
2096}
2097
2098pub fn apply_smooth_transform_to_design(
2099 design_local: DesignMatrix,
2100 transform: &Array2<f64>,
2101 termname: &str,
2102) -> Result<DesignMatrix, BasisError> {
2103 match design_local {
2104 DesignMatrix::Dense(inner) => {
2105 let op = CoefficientTransformOperator::new(inner, transform.clone()).map_err(|e| {
2106 BasisError::InvalidInput(format!(
2107 "smooth identifiability transform failed for term '{termname}': {e}"
2108 ))
2109 })?;
2110 Ok(DesignMatrix::Dense(
2111 gam_linalg::matrix::DenseDesignMatrix::from(Arc::new(op)),
2112 ))
2113 }
2114 DesignMatrix::Sparse(inner) => {
2115 let dense = inner
2116 .try_to_dense_arc("smooth identifiability sparse transform")
2117 .map_err(BasisError::InvalidInput)?
2118 .as_ref()
2119 .dot(transform);
2120 Ok(DesignMatrix::Dense(
2121 gam_linalg::matrix::DenseDesignMatrix::from(dense),
2122 ))
2123 }
2124 }
2125}
2126
2127fn design_constraint_cross(
2128 design: &DesignMatrix,
2129 constraint_matrix: ArrayView2<'_, f64>,
2130) -> Result<Array2<f64>, BasisError> {
2131 let n = design.nrows();
2132 if constraint_matrix.nrows() != n {
2133 return Err(BasisError::ConstraintMatrixRowMismatch {
2134 basisrows: n,
2135 constraintrows: constraint_matrix.nrows(),
2136 });
2137 }
2138 let mut cross = Array2::<f64>::zeros((design.ncols(), constraint_matrix.ncols()));
2139 const CHUNK: usize = 1024;
2140 for start in (0..n).step_by(CHUNK) {
2141 let end = (start + CHUNK).min(n);
2142 let design_chunk = design
2143 .try_row_chunk(start..end)
2144 .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
2145 let constraint_chunk = constraint_matrix.slice(s![start..end, ..]).to_owned();
2146 cross += &gam_linalg::faer_ndarray::fast_atb(&design_chunk, &constraint_chunk);
2147 }
2148 Ok(cross)
2149}
2150
2151fn design_frobenius_norm(design: &DesignMatrix) -> Result<f64, BasisError> {
2152 let n = design.nrows();
2153 const CHUNK: usize = 1024;
2154 let mut sumsq = 0.0;
2155 for start in (0..n).step_by(CHUNK) {
2156 let end = (start + CHUNK).min(n);
2157 let chunk = design
2158 .try_row_chunk(start..end)
2159 .map_err(|e| BasisError::InvalidInput(e.to_string()))?;
2160 sumsq += chunk.iter().map(|v| v * v).sum::<f64>();
2161 }
2162 Ok(sumsq.sqrt())
2163}
2164
2165fn frozen_parametric_residualization(
2173 termspec: &SmoothTermSpec,
2174) -> Option<&ParametricResidualizationChart> {
2175 termspec.frozen_parametric_residualization.as_ref()
2176}
2177
2178fn frozen_global_orthogonality(termspec: &SmoothTermSpec) -> Option<&Array2<f64>> {
2184 match &termspec.basis {
2185 SmoothBasisSpec::FactorSumToZero {
2186 frozen_global_orthogonality,
2187 ..
2188 } => frozen_global_orthogonality.as_ref(),
2189 SmoothBasisSpec::FactorSmooth { spec } => spec.frozen_global_orthogonality.as_ref(),
2190 _ => None,
2191 }
2192}
2193
2194fn penalty_candidates_under_collection_gauge(
2207 active_penalties: &[ActivePenalty],
2208 coefficient_gauge: Option<&gam_problem::Gauge>,
2209 term_name: &str,
2210) -> Result<Vec<PenaltyCandidate>, BasisError> {
2211 use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
2212 let penalty_candidates = active_penalties
2213 .par_iter()
2214 .map(|penalty| -> Result<PenaltyCandidate, BasisError> {
2215 let raw = ConstructiveQuadratic::try_from_dense_psd(
2216 penalty.matrix.clone(),
2217 "global smooth source penalty",
2218 )?;
2219 let raw = match penalty.info.structural_null_frame.as_ref() {
2227 Some(frame) => raw.with_structural_null_frame(
2228 frame.clone(),
2229 "global smooth source penalty structural frame",
2230 )?,
2231 None => raw,
2232 };
2233 let restricted = if let Some(gauge) = coefficient_gauge {
2234 raw.restricted(gauge, "global smooth identifiability restriction")?
2235 } else {
2236 raw
2237 };
2238 let (_, c_new) = normalize_penalty_in_constrained_space(restricted.dense());
2239 let matrix = restricted.scaled(1.0 / c_new, "normalized global smooth penalty")?;
2240 Ok(PenaltyCandidate {
2241 matrix,
2242 source: penalty.info.source.clone(),
2243 normalization_scale: penalty.info.normalization_scale * c_new,
2244 kronecker_factors: None,
2245 op: None,
2246 })
2247 })
2248 .collect::<Result<Vec<_>, _>>()?;
2249 let mut penalty_candidates = penalty_candidates;
2269 if coefficient_gauge.is_some()
2270 && penalty_candidates
2271 .iter()
2272 .any(|c| matches!(c.source, PenaltySource::DoublePenaltyNullspace))
2273 {
2274 const SUPPORT_TOL: f64 = 0.0;
2283 let support_rows = |m: &Array2<f64>| -> (usize, usize) {
2284 let n = m.nrows();
2285 let mut lo = n;
2286 let mut hi = 0usize;
2287 for i in 0..n {
2288 let any = (0..m.ncols()).any(|j| m[[i, j]].abs() > SUPPORT_TOL);
2289 if any {
2290 lo = lo.min(i);
2291 hi = hi.max(i + 1);
2292 }
2293 }
2294 (lo, hi)
2295 };
2296 let primaries: Vec<((usize, usize), ConstructiveQuadratic)> = penalty_candidates
2299 .iter()
2300 .filter(|c| matches!(c.source, PenaltySource::Primary))
2301 .map(|c| -> Result<_, BasisError> {
2302 Ok((
2303 support_rows(&c.matrix),
2304 c.matrix
2305 .scaled(c.normalization_scale, "physical global smooth primary")?,
2306 ))
2307 })
2308 .collect::<Result<Vec<_>, _>>()?;
2309 for candidate in &mut penalty_candidates {
2310 if !matches!(candidate.source, PenaltySource::DoublePenaltyNullspace) {
2311 continue;
2312 }
2313 let q = candidate.matrix.nrows();
2314 let (rlo, rhi) = support_rows(&candidate.matrix);
2315 let owner = primaries
2319 .iter()
2320 .find(|((plo, phi), _)| *plo <= rlo && rhi <= *phi)
2321 .or_else(|| (primaries.len() == 1).then(|| &primaries[0]))
2322 .ok_or_else(|| {
2323 BasisError::InvalidInput(format!(
2324 "double-penalty ridge for smooth '{}' has no co-located primary penalty",
2325 term_name
2326 ))
2327 })?;
2328 let ((plo, phi), s_full) = owner;
2329 let block = ConstructiveQuadratic::from_energy_factor(
2335 s_full.factor().slice(s![.., *plo..*phi]).to_owned(),
2336 "owned global smooth primary block",
2337 )?;
2338 let block = match s_full.structural_null_frame_block(*plo, *phi) {
2344 Some(frame) => block.with_structural_null_frame(
2345 frame,
2346 "owned global smooth primary block structural frame",
2347 )?,
2348 None => block,
2349 };
2350 let ridge_full = candidate.matrix.scaled(
2351 candidate.normalization_scale,
2352 "physical global smooth null ridge",
2353 )?;
2354 let ridge_block = ConstructiveQuadratic::from_energy_factor(
2355 ridge_full.factor().slice(s![.., *plo..*phi]).to_owned(),
2356 "owned global smooth null-ridge block",
2357 )?;
2358 let rebuilt_block =
2359 crate::basis::rebuild_metric_consistent_ridge(&block, &ridge_block)?;
2360 match rebuilt_block {
2361 Some(ridge_block) => {
2362 let mut full_factor =
2363 Array2::<f64>::zeros((ridge_block.factor().nrows(), q));
2364 full_factor
2365 .slice_mut(s![.., *plo..*phi])
2366 .assign(ridge_block.factor());
2367 let full = ConstructiveQuadratic::from_energy_factor(
2368 full_factor,
2369 "embedded global smooth null ridge",
2370 )?;
2371 let (_, scale) = normalize_penalty_in_constrained_space(full.dense());
2372 candidate.matrix = full
2373 .scaled(1.0 / scale, "normalized embedded global smooth null ridge")?;
2374 candidate.normalization_scale = scale;
2375 candidate.kronecker_factors = None;
2376 candidate.op = None;
2377 }
2378 None => {
2381 candidate.matrix = ConstructiveQuadratic::zero(q);
2382 candidate.normalization_scale = 1.0;
2383 candidate.kronecker_factors = None;
2384 candidate.op = None;
2385 }
2386 }
2387 }
2388 }
2389 Ok(penalty_candidates)
2390}
2391
2392fn maybe_smooth_identifiability_transform(
2393 termspec: &SmoothTermSpec,
2394 design_local: &DesignMatrix,
2395 constraint_block: Option<ArrayView2<'_, f64>>,
2396) -> Result<Option<Array2<f64>>, BasisError> {
2397 if let Some(SpatialIdentifiability::FrozenTransform { transform }) =
2398 spatial_identifiability_policy(termspec)
2399 {
2400 if design_local.ncols() != transform.nrows() {
2401 gam_problem::bail_dim_basis!(
2402 "frozen spatial identifiability transform mismatch: design has {} columns but transform has {} rows",
2403 design_local.ncols(),
2404 transform.nrows()
2405 );
2406 }
2407 return Ok(Some(transform.clone()));
2408 }
2409
2410 if let Some(c) = constraint_block {
2411 if c.ncols() == 0 {
2412 Ok(None)
2413 } else {
2414 Ok(Some(orthogonality_transform_for_design(
2415 design_local,
2416 c,
2417 None, )?))
2419 }
2420 } else {
2421 Ok(None)
2422 }
2423}
2424
2425fn smooth_requires_parametric_orthogonality(termspec: &SmoothTermSpec) -> bool {
2494 match &termspec.basis {
2495 SmoothBasisSpec::ByVariable { inner, .. }
2496 | SmoothBasisSpec::FactorSumToZero { inner, .. } => {
2497 smooth_requires_parametric_orthogonality(&SmoothTermSpec {
2498 frozen_parametric_residualization: None,
2499 name: termspec.name.clone(),
2500 basis: (**inner).clone(),
2501 shape: termspec.shape,
2502 joint_null_rotation: None,
2503 })
2504 }
2505 SmoothBasisSpec::BySmooth { smooth, .. } => {
2506 smooth_requires_parametric_orthogonality(&SmoothTermSpec {
2507 frozen_parametric_residualization: None,
2508 name: termspec.name.clone(),
2509 basis: (**smooth).clone(),
2510 shape: termspec.shape,
2511 joint_null_rotation: None,
2512 })
2513 }
2514 SmoothBasisSpec::ThinPlate { spec, .. } => {
2515 matches!(
2516 spec.identifiability,
2517 SpatialIdentifiability::OrthogonalToParametric
2518 )
2519 }
2520 SmoothBasisSpec::Duchon { spec, .. } => {
2521 matches!(
2522 spec.identifiability,
2523 SpatialIdentifiability::OrthogonalToParametric
2524 )
2525 }
2526 SmoothBasisSpec::Matern { spec, .. } => matches!(
2527 spec.identifiability,
2528 MaternIdentifiability::CenterSumToZero | MaternIdentifiability::CenterLinearOrthogonal
2529 ),
2530 SmoothBasisSpec::Sphere { spec, .. } => {
2538 matches!(spec.method, crate::basis::SphereMethod::Wahba)
2539 && !matches!(spec.wahba_kernel, crate::basis::SphereWahbaKernel::Pseudo)
2540 && matches!(
2541 spec.identifiability,
2542 SphericalSplineIdentifiability::CenterSumToZero
2543 )
2544 }
2545 SmoothBasisSpec::ConstantCurvature { spec, .. } => matches!(
2553 spec.identifiability,
2554 ConstantCurvatureIdentifiability::CenterSumToZero
2555 ),
2556 SmoothBasisSpec::MeasureJet { spec, .. } => matches!(
2562 spec.identifiability,
2563 MeasureJetIdentifiability::CenterSumToZero
2564 ),
2565 SmoothBasisSpec::BSpline1D { .. }
2566 | SmoothBasisSpec::TensorBSpline { .. }
2567 | SmoothBasisSpec::Pca { .. }
2568 | SmoothBasisSpec::FactorSmooth { .. } => false,
2569 }
2570}
2571
2572fn compose_identifiability_transforms(
2573 existing: Option<&Array2<f64>>,
2574 extra: Option<&Array2<f64>>,
2575) -> Result<Option<Array2<f64>>, BasisError> {
2576 match (existing, extra) {
2577 (Some(lhs), Some(rhs)) => {
2578 if lhs.ncols() == rhs.nrows() {
2579 Ok(Some(lhs.dot(rhs)))
2580 } else if lhs.nrows() == rhs.nrows() && lhs.ncols() == rhs.ncols() {
2581 Ok(Some(rhs.clone()))
2585 } else {
2586 Err(BasisError::DimensionMismatch(format!(
2587 "identifiability transform mismatch: existing is {}x{}, extra is {}x{}",
2588 lhs.nrows(),
2589 lhs.ncols(),
2590 rhs.nrows(),
2591 rhs.ncols(),
2592 )))
2593 }
2594 }
2595 (Some(lhs), None) => Ok(Some(lhs.clone())),
2596 (None, Some(rhs)) => Ok(Some(rhs.clone())),
2597 (None, None) => Ok(None),
2598 }
2599}
2600
2601fn with_identifiability_transform(
2602 metadata: &BasisMetadata,
2603 transform: Option<&Array2<f64>>,
2604) -> Result<BasisMetadata, BasisError> {
2605 match metadata {
2606 BasisMetadata::BSpline1D {
2607 knots,
2608 identifiability_transform,
2609 periodic,
2610 degree,
2611 auto_shrink_note,
2612 anchor_offset_coeffs,
2613 } => Ok(BasisMetadata::BSpline1D {
2614 knots: knots.clone(),
2615 periodic: *periodic,
2616 identifiability_transform: compose_identifiability_transforms(
2617 identifiability_transform.as_ref(),
2618 transform,
2619 )?,
2620 degree: *degree,
2621 auto_shrink_note: auto_shrink_note.clone(),
2622 anchor_offset_coeffs: anchor_offset_coeffs.clone(),
2626 }),
2627 BasisMetadata::CubicRegression1D {
2628 knots,
2629 identifiability_transform,
2630 } => Ok(BasisMetadata::CubicRegression1D {
2631 knots: knots.clone(),
2632 identifiability_transform: compose_identifiability_transforms(
2633 identifiability_transform.as_ref(),
2634 transform,
2635 )?,
2636 }),
2637 BasisMetadata::ThinPlate {
2638 centers,
2639 length_scale,
2640 periodic,
2641 identifiability_transform,
2642 input_scale,
2643 radial_reparam,
2644 } => Ok(BasisMetadata::ThinPlate {
2645 centers: centers.clone(),
2646 length_scale: *length_scale,
2647 periodic: periodic.clone(),
2648 identifiability_transform: compose_identifiability_transforms(
2649 identifiability_transform.as_ref(),
2650 transform,
2651 )?,
2652 input_scale: *input_scale,
2653 radial_reparam: radial_reparam.clone(),
2654 }),
2655 BasisMetadata::Sphere {
2656 centers,
2657 penalty_order,
2658 method,
2659 max_degree,
2660 wahba_kernel,
2661 constraint_transform,
2662 } => Ok(BasisMetadata::Sphere {
2663 centers: centers.clone(),
2664 penalty_order: *penalty_order,
2665 method: *method,
2666 max_degree: *max_degree,
2667 wahba_kernel: *wahba_kernel,
2668 constraint_transform: compose_identifiability_transforms(
2669 constraint_transform.as_ref(),
2670 transform,
2671 )?,
2672 }),
2673 BasisMetadata::ConstantCurvature {
2674 centers,
2675 kappa,
2676 length_scale,
2677 constraint_transform,
2678 } => Ok(BasisMetadata::ConstantCurvature {
2679 centers: centers.clone(),
2680 kappa: *kappa,
2681 length_scale: *length_scale,
2682 constraint_transform: compose_identifiability_transforms(
2683 constraint_transform.as_ref(),
2684 transform,
2685 )?,
2686 }),
2687 BasisMetadata::MeasureJet {
2688 centers,
2689 input_scale,
2690 length_scale,
2691 eps_band,
2692 order_s,
2693 alpha,
2694 tau0,
2695 masses,
2696 support_means,
2697 penalty_normalization_scales,
2698 raw_penalty_normalization_scales,
2699 fused_penalty_normalization_scale,
2700 constraint_transform,
2701 sigma_coord,
2702 } => Ok(BasisMetadata::MeasureJet {
2703 centers: centers.clone(),
2704 input_scale: *input_scale,
2705 length_scale: *length_scale,
2706 eps_band: eps_band.clone(),
2707 order_s: *order_s,
2708 alpha: *alpha,
2709 tau0: *tau0,
2710 masses: masses.clone(),
2711 support_means: support_means.clone(),
2712 penalty_normalization_scales: penalty_normalization_scales.clone(),
2713 raw_penalty_normalization_scales: raw_penalty_normalization_scales.clone(),
2714 fused_penalty_normalization_scale: *fused_penalty_normalization_scale,
2715 constraint_transform: compose_identifiability_transforms(
2716 constraint_transform.as_ref(),
2717 transform,
2718 )?,
2719 sigma_coord: *sigma_coord,
2720 }),
2721 BasisMetadata::Matern {
2722 centers,
2723 length_scale,
2724 periodic,
2725 nu,
2726 include_intercept,
2727 identifiability_transform,
2728 input_scale,
2729 aniso_log_scales,
2730 } => Ok(BasisMetadata::Matern {
2731 centers: centers.clone(),
2732 length_scale: *length_scale,
2733 periodic: periodic.clone(),
2734 nu: *nu,
2735 include_intercept: *include_intercept,
2736 identifiability_transform: compose_identifiability_transforms(
2737 identifiability_transform.as_ref(),
2738 transform,
2739 )?,
2740 input_scale: *input_scale,
2741 aniso_log_scales: aniso_log_scales.clone(),
2742 }),
2743 BasisMetadata::Duchon {
2744 centers,
2745 length_scale,
2746 periodic,
2747 power,
2748 nullspace_order,
2749 identifiability_transform,
2750 input_scale,
2751 aniso_log_scales,
2752 operator_collocation_points,
2753 radial_reparam,
2754 spectral_basis,
2755 } => Ok(BasisMetadata::Duchon {
2756 centers: centers.clone(),
2757 length_scale: *length_scale,
2758 periodic: periodic.clone(),
2759 power: *power,
2760 nullspace_order: *nullspace_order,
2761 input_scale: *input_scale,
2762 aniso_log_scales: aniso_log_scales.clone(),
2763 operator_collocation_points: operator_collocation_points.clone(),
2764 radial_reparam: radial_reparam.clone(),
2765 spectral_basis: spectral_basis.clone(),
2766 identifiability_transform: compose_identifiability_transforms(
2767 identifiability_transform.as_ref(),
2768 transform,
2769 )?,
2770 }),
2771 BasisMetadata::SphereHarmonics {
2772 max_degree,
2773 radians,
2774 } => Ok(BasisMetadata::SphereHarmonics {
2775 max_degree: *max_degree,
2776 radians: *radians,
2777 }),
2778 BasisMetadata::TensorBSpline {
2779 feature_cols,
2780 knots,
2781 degrees,
2782 periods,
2783 is_cr,
2784 identifiability_transform,
2785 } => Ok(BasisMetadata::TensorBSpline {
2786 feature_cols: feature_cols.clone(),
2787 knots: knots.clone(),
2788 degrees: degrees.clone(),
2789 periods: periods.clone(),
2790 is_cr: is_cr.clone(),
2791 identifiability_transform: compose_identifiability_transforms(
2792 identifiability_transform.as_ref(),
2793 transform,
2794 )?,
2795 }),
2796 BasisMetadata::BySmooth {
2797 inner,
2798 by_col,
2799 levels,
2800 ordered,
2801 } => Ok(BasisMetadata::BySmooth {
2802 inner: Box::new(with_identifiability_transform(inner, transform)?),
2803 by_col: *by_col,
2804 levels: levels.clone(),
2805 ordered: *ordered,
2806 }),
2807 BasisMetadata::FactorSmooth {
2808 continuous_cols,
2809 group_col,
2810 knots,
2811 degree,
2812 periodic,
2813 group_levels,
2814 flavour,
2815 marginal_is_cr,
2816 } => {
2817 if transform.is_some() {
2824 gam_problem::bail_invalid_basis!(
2825 "FactorSmooth metadata cannot absorb an identifiability transform; \
2826 route it through the term-level frozen_global_orthogonality carrier"
2827 );
2828 }
2829 Ok(BasisMetadata::FactorSmooth {
2830 continuous_cols: continuous_cols.clone(),
2831 group_col: *group_col,
2832 knots: knots.clone(),
2833 degree: *degree,
2834 periodic: *periodic,
2835 group_levels: group_levels.clone(),
2836 flavour: flavour.clone(),
2837 marginal_is_cr: *marginal_is_cr,
2838 })
2839 }
2840 BasisMetadata::Pca {
2841 feature_cols,
2842 basis_matrix,
2843 centered,
2844 smooth_penalty,
2845 center_mean,
2846 pca_basis_path,
2847 chunk_size,
2848 } => {
2849 if transform.is_some() {
2855 gam_problem::bail_invalid_basis!(
2856 "PCA bases do not expose a composable identifiability transform"
2857 );
2858 }
2859 Ok(BasisMetadata::Pca {
2860 feature_cols: feature_cols.clone(),
2861 basis_matrix: basis_matrix.clone(),
2862 centered: *centered,
2863 smooth_penalty: *smooth_penalty,
2864 center_mean: center_mean.clone(),
2865 pca_basis_path: pca_basis_path.clone(),
2866 chunk_size: *chunk_size,
2867 })
2868 }
2869 }
2870}
2871
2872pub fn orthogonality_relative_residual_for_design(
2876 design: &DesignMatrix,
2877 constraint_matrix: ArrayView2<'_, f64>,
2878) -> Result<f64, BasisError> {
2879 let cross = design_constraint_cross(design, constraint_matrix)?;
2880 let num = cross.iter().map(|v| v * v).sum::<f64>().sqrt();
2881 let b_norm = design_frobenius_norm(design)?;
2882 let c_norm = constraint_matrix.iter().map(|v| v * v).sum::<f64>().sqrt();
2883 let denom = b_norm * c_norm;
2884 if denom == 0.0 {
2886 return Ok(0.0);
2887 }
2888 Ok(num / denom)
2889}
2890
2891#[cfg(test)]
2892mod frozen_linear_term_mass_rebuild_tests {
2893 use super::*;
2894
2895 #[cfg(target_os = "linux")]
2899 #[test]
2900 fn million_row_linear_design_stays_below_derived_peak_rss() {
2901 fn hwm() -> usize {
2902 std::fs::read_to_string("/proc/self/status")
2903 .expect("Linux status")
2904 .lines()
2905 .find_map(|line| line.strip_prefix("VmHWM:")?.split_whitespace().next()?.parse::<usize>().ok())
2906 .expect("VmHWM") * 1024
2907 }
2908 let (n, p) = (1_000_000usize, 16usize);
2909 let data = Array2::from_shape_fn((n, p), |(i, j)| ((i + j) % 19) as f64);
2910 let spec = TermCollectionSpec {
2911 linear_terms: (0..p).map(|j| LinearTermSpec {
2912 name: format!("x{j}"), feature_col: j, feature_cols: vec![j],
2913 categorical_levels: vec![], double_penalty: false,
2914 coefficient_geometry: LinearCoefficientGeometry::Unconstrained,
2915 coefficient_min: None, coefficient_max: None, frozen_function_mass: None,
2916 }).collect(),
2917 random_effect_terms: vec![], smooth_terms: vec![],
2918 };
2919 let before = hwm();
2920 let built = build_term_collection_design(data.view(), &spec).expect("million-row design");
2921 let growth = hwm().saturating_sub(before);
2922 let column_bytes = n * std::mem::size_of::<f64>();
2923 let derived_limit = n * p * std::mem::size_of::<f64>() + 2 * column_bytes;
2924 assert!(growth <= derived_limit, "peak RSS growth {growth} exceeded derived {derived_limit}");
2925 assert_eq!(built.design.nrows(), n);
2926 }
2927
2928 fn one_linear_term_spec() -> TermCollectionSpec {
2932 TermCollectionSpec {
2933 linear_terms: vec![LinearTermSpec {
2934 name: "x".to_string(),
2935 feature_col: 0,
2936 feature_cols: vec![0],
2937 categorical_levels: vec![],
2938 double_penalty: true,
2939 coefficient_geometry: LinearCoefficientGeometry::Unconstrained,
2940 coefficient_min: None,
2941 coefficient_max: None,
2942 frozen_function_mass: None,
2943 }],
2944 random_effect_terms: Vec::new(),
2945 smooth_terms: Vec::new(),
2946 }
2947 }
2948
2949 fn training_data_varying_x(n: usize) -> Array2<f64> {
2950 let mut data = Array2::<f64>::zeros((n, 1));
2951 for i in 0..n {
2952 data[[i, 0]] = 1.0 + i as f64;
2954 }
2955 data
2956 }
2957
2958 fn constant_zero_x(n_rows: usize) -> Array2<f64> {
2959 Array2::<f64>::zeros((n_rows, 1))
2960 }
2961
2962 #[test]
2968 fn unfrozen_spec_still_rejects_a_genuinely_zero_training_column() {
2969 let spec = one_linear_term_spec();
2970 let degenerate_training_data = constant_zero_x(20);
2971 let err = build_term_collection_design(degenerate_training_data.view(), &spec)
2972 .expect_err("an unfrozen spec fit directly on an all-zero column must still fail");
2973 let message = err.to_string();
2974 assert!(
2975 message.contains("identically zero"),
2976 "expected the identifiability guard's message, got: {message}"
2977 );
2978 }
2979
2980 #[test]
2991 fn frozen_spec_rebuilds_at_a_constant_evaluation_column_using_the_training_mass() {
2992 let spec = one_linear_term_spec();
2993 let training_data = training_data_varying_x(40);
2994
2995 let training_design = build_term_collection_design(training_data.view(), &spec)
2996 .expect("fit-time build over a genuinely varying column must succeed");
2997 let training_mass = training_design
2998 .linear_function_masses
2999 .first()
3000 .copied()
3001 .flatten()
3002 .expect("a double_penalty=true term must report its fit-time function mass");
3003 assert!(
3004 training_mass > 0.0,
3005 "training mass for a genuinely varying column must be strictly positive, got {training_mass}"
3006 );
3007
3008 let frozen_spec = freeze_term_collection_from_design(&spec, &training_design)
3009 .expect("freezing the spec against its own fit-time design must succeed");
3010 assert_eq!(
3011 frozen_spec.linear_terms[0].frozen_function_mass,
3012 Some(training_mass),
3013 "freezing must persist the exact fit-time mass onto the term"
3014 );
3015
3016 let evaluation_grid = constant_zero_x(3);
3020 let rebuilt_design = build_term_collection_design(evaluation_grid.view(), &frozen_spec)
3021 .expect(
3022 "rebuilding a FROZEN spec's design at a constant-covariate evaluation grid must \
3023 succeed — the training-time mass is reused, never recomputed from these rows",
3024 );
3025 assert_eq!(
3026 rebuilt_design
3027 .linear_function_masses
3028 .first()
3029 .copied()
3030 .flatten(),
3031 Some(training_mass),
3032 "the rebuilt design must carry the REUSED training-time mass, not a value \
3033 recomputed from the (all-zero) evaluation rows"
3034 );
3035 }
3036}