1pub fn build_term_collection_designs_joint(
9 data: ArrayView2<'_, f64>,
10 specs: &[TermCollectionSpec],
11) -> Result<Vec<TermCollectionDesign>, BasisError> {
12 for spec in specs {
13 validate_term_collection_finite_inputs(data, spec)?;
14 }
15 let smooth_blocks = specs
16 .iter()
17 .map(|spec| spec.smooth_terms.clone())
18 .collect::<Vec<_>>();
19 let planned_blocks = plan_joint_spatial_centers_for_term_blocks(data, &smooth_blocks)?;
20 let mut out = Vec::with_capacity(specs.len());
21 for (spec, planned_terms) in specs.iter().zip(planned_blocks.into_iter()) {
22 let mut planned_spec = spec.clone();
23 planned_spec.smooth_terms = planned_terms;
24 out.push(build_term_collection_design_inner(data, &planned_spec)?);
25 }
26 Ok(out)
27}
28
29pub fn build_term_collection_designs_and_freeze_joint(
30 data: ArrayView2<'_, f64>,
31 specs: &[TermCollectionSpec],
32) -> Result<(Vec<TermCollectionDesign>, Vec<TermCollectionSpec>), EstimationError> {
33 let designs = build_term_collection_designs_joint(data, specs)?;
34 let mut resolved_specs = Vec::with_capacity(specs.len());
35 for (spec, design) in specs.iter().zip(designs.iter()) {
36 resolved_specs.push(freeze_term_collection_from_design(spec, design)?);
37 }
38 Ok((designs, resolved_specs))
39}
40
41pub fn fit_term_collection_forspec(
42 data: ArrayView2<'_, f64>,
43 y: ArrayView1<'_, f64>,
44 weights: ArrayView1<'_, f64>,
45 offset: ArrayView1<'_, f64>,
46 spec: &TermCollectionSpec,
47 family: LikelihoodSpec,
48 options: &FitOptions,
49) -> Result<FittedTermCollection, EstimationError> {
50 fit_term_collection_forspecwith_heuristic_lambdas(
51 data, y, weights, offset, spec, None, family, options,
52 )
53}
54
55pub fn fit_term_collection_with_coefficient_groups(
56 data: ArrayView2<'_, f64>,
57 y: ArrayView1<'_, f64>,
58 weights: ArrayView1<'_, f64>,
59 offset: ArrayView1<'_, f64>,
60 spec: &TermCollectionSpec,
61 groups: &[CoefficientGroupSpec],
62 family: LikelihoodSpec,
63 options: &FitOptions,
64) -> Result<FittedTermCollection, EstimationError> {
65 if groups.is_empty() {
66 return fit_term_collection_forspec(data, y, weights, offset, spec, family, options);
67 }
68 let design = build_term_collection_design_with_policy(data, spec, &options.resource_policy)?;
69 let base_fit_opts = adaptive_fit_options_base(options, &design);
70 let realized = design
71 .realize_coefficient_groups(groups, &base_fit_opts.rho_prior)
72 .map_err(EstimationError::BasisError)?;
73 let effective_offset = design
74 .compose_offset(offset, "coefficient-group fit")
75 .map_err(EstimationError::BasisError)?;
76 let mut grouped_options = base_fit_opts.clone();
77 grouped_options.rho_prior = realized.rho_prior;
78 let fitted = FittedTermCollection {
79 fit: gam_solve::estimate::fit_gam_with_penalty_specs(
80 design.design.clone(),
81 y,
82 weights,
83 effective_offset.view(),
84 realized.penalty_specs,
85 realized.nullspace_dims,
86 family.clone(),
87 &grouped_options,
88 )?,
89 design,
90 adaptive_diagnostics: None,
91 };
92 enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
93 Ok(fitted)
94}
95
96pub fn fit_term_collection_with_penalty_block_gamma_prior_callback<F>(
97 data: ArrayView2<'_, f64>,
98 y: ArrayView1<'_, f64>,
99 weights: ArrayView1<'_, f64>,
100 offset: ArrayView1<'_, f64>,
101 spec: &TermCollectionSpec,
102 callback: F,
103 family: LikelihoodSpec,
104 options: &FitOptions,
105) -> Result<FittedTermCollection, EstimationError>
106where
107 F: FnMut(&PenaltyBlockGammaPriorMetadata<'_>) -> Option<(f64, f64)>,
108{
109 let design = build_term_collection_design_with_policy(data, spec, &options.resource_policy)?;
110 let effective_offset = design
111 .compose_offset(offset, "penalty-prior callback fit")
112 .map_err(EstimationError::BasisError)?;
113 let mut fit_opts = adaptive_fit_options_base(options, &design);
114 fit_opts.rho_prior = realize_penalty_block_gamma_priors(&design, callback)
115 .map_err(EstimationError::BasisError)?;
116 let fitted = FittedTermCollection {
117 fit: fit_gamwith_heuristic_lambdas(
118 design.design.clone(),
119 y,
120 weights,
121 effective_offset.view(),
122 &design.penalties,
123 None,
124 family.clone(),
125 &fit_opts,
126 )?,
127 design,
128 adaptive_diagnostics: None,
129 };
130 enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
131 Ok(fitted)
132}
133
134pub fn fit_term_collection_with_penalty_block_gamma_priors(
135 data: ArrayView2<'_, f64>,
136 y: ArrayView1<'_, f64>,
137 weights: ArrayView1<'_, f64>,
138 offset: ArrayView1<'_, f64>,
139 spec: &TermCollectionSpec,
140 priors: &[(String, f64, f64)],
141 family: LikelihoodSpec,
142 options: &FitOptions,
143) -> Result<FittedTermCollection, EstimationError> {
144 let design = build_term_collection_design_with_policy(data, spec, &options.resource_policy)?;
145 let effective_offset = design
146 .compose_offset(offset, "penalty-prior fit")
147 .map_err(EstimationError::BasisError)?;
148 let mut fit_opts = adaptive_fit_options_base(options, &design);
149 fit_opts.rho_prior = realize_keyed_penalty_block_gamma_priors(&design, priors)
150 .map_err(EstimationError::BasisError)?;
151 let fitted = FittedTermCollection {
152 fit: fit_gamwith_heuristic_lambdas(
153 design.design.clone(),
154 y,
155 weights,
156 effective_offset.view(),
157 &design.penalties,
158 None,
159 family.clone(),
160 &fit_opts,
161 )?,
162 design,
163 adaptive_diagnostics: None,
164 };
165 enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
166 Ok(fitted)
167}
168
169pub fn fit_term_collection_with_coefficient_groups_and_penalty_block_gamma_priors(
170 data: ArrayView2<'_, f64>,
171 y: ArrayView1<'_, f64>,
172 weights: ArrayView1<'_, f64>,
173 offset: ArrayView1<'_, f64>,
174 spec: &TermCollectionSpec,
175 groups: &[CoefficientGroupSpec],
176 priors: &[(String, f64, f64)],
177 family: LikelihoodSpec,
178 options: &FitOptions,
179) -> Result<FittedTermCollection, EstimationError> {
180 if groups.is_empty() {
181 return fit_term_collection_with_penalty_block_gamma_priors(
182 data, y, weights, offset, spec, priors, family, options,
183 );
184 }
185 if priors.is_empty() {
186 return fit_term_collection_with_coefficient_groups(
187 data, y, weights, offset, spec, groups, family, options,
188 );
189 }
190
191 let design = build_term_collection_design_with_policy(data, spec, &options.resource_policy)?;
195 let base_fit_opts = adaptive_fit_options_base(options, &design);
196 let base_rho_prior = realize_keyed_penalty_block_gamma_priors(&design, priors)
197 .map_err(EstimationError::BasisError)?;
198 let realized = design
199 .realize_coefficient_groups(groups, &base_rho_prior)
200 .map_err(EstimationError::BasisError)?;
201 let effective_offset = design
202 .compose_offset(offset, "coefficient-group and penalty-prior fit")
203 .map_err(EstimationError::BasisError)?;
204 let mut grouped_options = base_fit_opts.clone();
205 grouped_options.rho_prior = realized.rho_prior;
206 let fitted = FittedTermCollection {
207 fit: gam_solve::estimate::fit_gam_with_penalty_specs(
208 design.design.clone(),
209 y,
210 weights,
211 effective_offset.view(),
212 realized.penalty_specs,
213 realized.nullspace_dims,
214 family.clone(),
215 &grouped_options,
216 )?,
217 design,
218 adaptive_diagnostics: None,
219 };
220 enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
221 Ok(fitted)
222}
223
224fn fit_term_collection_forspecwith_heuristic_lambdas(
225 data: ArrayView2<'_, f64>,
226 y: ArrayView1<'_, f64>,
227 weights: ArrayView1<'_, f64>,
228 offset: ArrayView1<'_, f64>,
229 spec: &TermCollectionSpec,
230 heuristic_lambdas: Option<&[f64]>,
231 family: LikelihoodSpec,
232 options: &FitOptions,
233) -> Result<FittedTermCollection, EstimationError> {
234 let adaptive_opts = options.adaptive_regularization.clone().unwrap_or_default();
235 let resolved_spec;
236 let design_spec = if adaptive_opts.enabled {
237 resolved_spec = ensure_matern_adaptive_center_resolution(spec, data.nrows());
238 &resolved_spec
239 } else {
240 spec
241 };
242 let base_design =
243 build_term_collection_design_with_policy(data, design_spec, &options.resource_policy)?;
244 fit_term_collection_on_realized_design(
245 y,
246 weights,
247 offset,
248 design_spec,
249 &base_design,
250 heuristic_lambdas,
251 family,
252 options,
253 )
254}
255
256fn ensure_matern_adaptive_center_resolution(
257 spec: &TermCollectionSpec,
258 n_rows: usize,
259) -> TermCollectionSpec {
260 let mut out = spec.clone();
261 for term in &mut out.smooth_terms {
262 let gam_terms::smooth::SmoothBasisSpec::Matern {
263 feature_cols,
264 spec: matern,
265 ..
266 } = &mut term.basis
267 else {
268 continue;
269 };
270 if let gam_terms::basis::CenterStrategy::FarthestPoint { num_centers } =
271 &mut matern.center_strategy
272 {
273 let min_centers = (4 * feature_cols.len()).min(n_rows).max(*num_centers);
286 *num_centers = min_centers;
287 }
288 }
289 out
290}
291
292fn has_bounded_linear_terms(spec: &TermCollectionSpec) -> bool {
293 spec.linear_terms.iter().any(|term| {
294 matches!(
295 term.coefficient_geometry,
296 LinearCoefficientGeometry::Bounded { .. }
297 )
298 })
299}
300
301fn fit_term_collection_on_realized_design(
302 y: ArrayView1<'_, f64>,
303 weights: ArrayView1<'_, f64>,
304 offset: ArrayView1<'_, f64>,
305 spec: &TermCollectionSpec,
306 design: &TermCollectionDesign,
307 heuristic_lambdas: Option<&[f64]>,
308 family: LikelihoodSpec,
309 options: &FitOptions,
310) -> Result<FittedTermCollection, EstimationError> {
311 let effective_offset = design
312 .compose_offset(offset, "term-collection fit")
313 .map_err(EstimationError::BasisError)?;
314 let offset = effective_offset.view();
315 if has_bounded_linear_terms(spec) {
316 return fit_bounded_term_collection_with_design(
317 y,
318 weights,
319 offset,
320 spec,
321 design,
322 heuristic_lambdas,
323 family,
324 options,
325 );
326 }
327 let mut base_fit_opts = adaptive_fit_options_base(options, design);
328 base_fit_opts.rho_prior = relax_smoothing_rho_prior(options, design);
335 let fitted = FittedTermCollection {
336 fit: fit_gamwith_heuristic_lambdas(
337 design.design.clone(),
338 y,
339 weights,
340 offset,
341 &design.penalties,
342 heuristic_lambdas,
343 family.clone(),
344 &base_fit_opts,
345 )?,
346 design: design.clone(),
347 adaptive_diagnostics: None,
348 };
349 enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
350
351 let adaptive_opts = options.adaptive_regularization.clone().unwrap_or_default();
352 if !adaptive_opts.enabled {
353 return Ok(fitted);
354 }
355 let runtime_caches = extract_spatial_operator_runtime_caches(spec, &fitted.design)?;
356 if runtime_caches.is_empty() {
357 return Ok(fitted);
358 }
359 fit_term_collectionwith_exact_spatial_adaptive_regularization(
366 fitted,
367 y,
368 weights,
369 offset,
370 family,
371 options,
372 &runtime_caches,
373 )
374}
375
376#[derive(Clone)]
377struct SpatialOperatorRuntimeCache {
378 termname: String,
379 feature_cols: Vec<usize>,
380 coeff_global_range: Range<usize>,
381 mass_penalty_global_idx: usize,
382 tension_penalty_global_idx: usize,
383 stiffness_penalty_global_idx: usize,
384 d0: Array2<f64>,
385 d1: Array2<f64>,
386 d2: Array2<f64>,
387 collocation_points: Array2<f64>,
388 dimension: usize,
389}
390
391#[derive(Clone)]
392struct SpatialAdaptiveWeights {
393 inv_magweight: Array1<f64>,
394 invgradweight: Array1<f64>,
395 inv_lapweight: Array1<f64>,
396}
397
398#[derive(Clone)]
399struct CharbonnierScalarBlockState {
400 signal: Array1<f64>,
401 radius: Array1<f64>,
402 epsilon: f64,
403}
404
405impl CharbonnierScalarBlockState {
406 fn from_signal(signal: Array1<f64>, epsilon: f64) -> Self {
407 let eps = epsilon.max(1e-12);
408 let radius = signal.mapv(|t| (t * t + eps * eps).sqrt());
409 Self {
410 signal,
411 radius,
412 epsilon: eps,
413 }
414 }
415
416 fn absolute_signal(&self) -> Array1<f64> {
417 self.signal.mapv(f64::abs)
418 }
419
420 fn penalty_value(&self) -> f64 {
421 self.radius.iter().map(|r| r - self.epsilon).sum::<f64>()
422 }
423
424 fn betagradient_coeff(&self) -> Array1<f64> {
425 Array1::from_iter(
426 self.signal
427 .iter()
428 .zip(self.radius.iter())
429 .map(|(t, r)| t / r),
430 )
431 }
432
433 fn betahessian_diag(&self) -> Array1<f64> {
434 let eps2 = self.epsilon * self.epsilon;
435 self.radius.mapv(|r| eps2 / r.powi(3))
436 }
437
438 fn log_epsilon_gradient_terms(&self) -> Array1<f64> {
439 let epsilon = self.epsilon;
440 let eps2 = epsilon * epsilon;
441 self.radius.mapv(|r| eps2 / r - epsilon)
442 }
443
444 fn log_epsilon_betagradient_coeff(&self) -> Array1<f64> {
445 let eps2 = self.epsilon * self.epsilon;
446 Array1::from_iter(
447 self.signal
448 .iter()
449 .zip(self.radius.iter())
450 .map(|(t, r)| -eps2 * t / r.powi(3)),
451 )
452 }
453
454 fn log_epsilon_hessian_terms(&self) -> Array1<f64> {
455 let epsilon = self.epsilon;
456 let eps2 = epsilon * epsilon;
457 let eps4 = eps2 * eps2;
458 self.radius
459 .mapv(|r| 2.0 * eps2 / r - eps4 / r.powi(3) - epsilon)
460 }
461
462 fn surrogateweights_posterior_snr(
463 &self,
464 variance: &Array1<f64>,
465 weight_floor: f64,
466 weight_ceiling: f64,
467 ) -> (Array1<f64>, Array1<f64>) {
468 let eps2 = self.epsilon * self.epsilon;
526 let weight = Array1::from_iter(self.signal.iter().zip(variance.iter()).map(|(&t, &v)| {
527 let credible2 = (t * t - v.max(0.0)).max(0.0);
528 let r = (credible2 + eps2).sqrt();
529 (1.0 / r).clamp(weight_floor, weight_ceiling)
530 }));
531 let invweight = weight.mapv(|u| 1.0 / u);
532 (weight, invweight)
533 }
534
535 fn directionalhessian_diag(&self, direction_signal: &Array1<f64>) -> Array1<f64> {
536 let eps2 = self.epsilon * self.epsilon;
551 Array1::from_iter(
552 self.signal
553 .iter()
554 .zip(direction_signal.iter())
555 .zip(self.radius.iter())
556 .map(|((t, q), r)| -3.0 * eps2 * t * q / r.powi(5)),
557 )
558 }
559
560 fn second_directionalhessian_diag(
567 &self,
568 direction1_signal: &Array1<f64>,
569 direction2_signal: &Array1<f64>,
570 ) -> Array1<f64> {
571 let eps2 = self.epsilon * self.epsilon;
572 Array1::from_iter(
573 self.signal
574 .iter()
575 .zip(direction1_signal.iter())
576 .zip(direction2_signal.iter())
577 .zip(self.radius.iter())
578 .map(|(((t, q1), q2), r)| {
579 let r2 = r * r;
580 let psi4 = -3.0 * eps2 / r.powi(5) + 15.0 * eps2 * t * t / (r.powi(5) * r2);
581 psi4 * q1 * q2
582 }),
583 )
584 }
585
586 fn log_epsilon_betahessian_diag(&self) -> Array1<f64> {
587 let eps2 = self.epsilon * self.epsilon;
588 let eps4 = eps2 * eps2;
589 Array1::from_iter(
590 self.signal
591 .iter()
592 .zip(self.radius.iter())
593 .map(|(_, r)| 2.0 * eps2 / r.powi(3) - 3.0 * eps4 / r.powi(5)),
594 )
595 }
596
597 fn log_epsilon_beta_mixed_second_coeff(&self) -> Array1<f64> {
598 let eps2 = self.epsilon * self.epsilon;
599 Array1::from_iter(
600 self.signal
601 .iter()
602 .zip(self.radius.iter())
603 .map(|(t, r)| eps2 * t * (eps2 - 2.0 * t * t) / r.powi(5)),
604 )
605 }
606
607 fn log_epsilon_betahessian_second_diag(&self) -> Array1<f64> {
608 let eps2 = self.epsilon * self.epsilon;
609 let eps4 = eps2 * eps2;
610 let eps6 = eps4 * eps2;
611 Array1::from_iter(
612 self.radius.iter().map(|r| {
613 4.0 * eps2 / r.powi(3) - 18.0 * eps4 / r.powi(5) + 15.0 * eps6 / r.powi(7)
614 }),
615 )
616 }
617
618 fn log_epsilon_betahessian_directional_diag(
619 &self,
620 direction_signal: &Array1<f64>,
621 ) -> Array1<f64> {
622 let eps2 = self.epsilon * self.epsilon;
623 let eps4 = eps2 * eps2;
624 Array1::from_iter(
625 self.signal
626 .iter()
627 .zip(direction_signal.iter())
628 .zip(self.radius.iter())
629 .map(|((t, q), r)| (-6.0 * eps2 * t / r.powi(5) + 15.0 * eps4 * t / r.powi(7)) * q),
630 )
631 }
632}
633
634#[derive(Clone)]
635struct CharbonnierGroupedBlockState {
636 norm: Array1<f64>,
637 radius: Array1<f64>,
638 signal_blocks: Array2<f64>,
639 epsilon: f64,
640}
641
642impl CharbonnierGroupedBlockState {
643 fn from_signal_blocks(signal_blocks: Array2<f64>, epsilon: f64) -> Self {
644 let eps = epsilon.max(1e-12);
645 let norm = Array1::from_iter(
646 signal_blocks
647 .rows()
648 .into_iter()
649 .map(|row| row.iter().map(|v| v * v).sum::<f64>().sqrt()),
650 );
651 let radius = norm.mapv(|g| (g * g + eps * eps).sqrt());
652 Self {
653 norm,
654 radius,
655 signal_blocks,
656 epsilon: eps,
657 }
658 }
659
660 fn penalty_value(&self) -> f64 {
661 self.radius.iter().map(|r| r - self.epsilon).sum::<f64>()
662 }
663
664 fn norm_signal(&self) -> Array1<f64> {
665 self.norm.clone()
666 }
667
668 fn betagradient_blocks(&self) -> Array2<f64> {
669 let mut out = self.signal_blocks.clone();
670 for (k, mut row) in out.rows_mut().into_iter().enumerate() {
671 let scale = 1.0 / self.radius[k];
672 row.mapv_inplace(|v| v * scale);
673 }
674 out
675 }
676
677 fn betahessian_blocks(&self) -> Vec<Array2<f64>> {
678 let mut out = Vec::with_capacity(self.signal_blocks.nrows());
679 for (k, row) in self.signal_blocks.rows().into_iter().enumerate() {
680 let dim = row.len();
681 let mut block = Array2::<f64>::eye(dim);
682 block.mapv_inplace(|v| v / self.radius[k]);
683 for i in 0..dim {
684 for j in 0..dim {
685 block[[i, j]] -= row[i] * row[j] / self.radius[k].powi(3);
686 }
687 }
688 out.push(block);
689 }
690 out
691 }
692
693 fn log_epsilon_gradient_terms(&self) -> Array1<f64> {
694 let epsilon = self.epsilon;
695 let eps2 = epsilon * epsilon;
696 self.radius.mapv(|r| eps2 / r - epsilon)
697 }
698
699 fn log_epsilon_betagradient_blocks(&self) -> Array2<f64> {
700 let mut out = self.signal_blocks.clone();
701 let eps2 = self.epsilon * self.epsilon;
702 for (k, mut row) in out.rows_mut().into_iter().enumerate() {
703 let scale = -eps2 / self.radius[k].powi(3);
704 row.mapv_inplace(|v| v * scale);
705 }
706 out
707 }
708
709 fn log_epsilon_hessian_terms(&self) -> Array1<f64> {
710 let epsilon = self.epsilon;
711 let eps2 = epsilon * epsilon;
712 let eps4 = eps2 * eps2;
713 self.radius
714 .mapv(|r| 2.0 * eps2 / r - eps4 / r.powi(3) - epsilon)
715 }
716
717 fn surrogateweights_posterior_snr(
718 &self,
719 variance: &Array1<f64>,
720 weight_floor: f64,
721 weight_ceiling: f64,
722 ) -> (Array1<f64>, Array1<f64>) {
723 let eps2 = self.epsilon * self.epsilon;
765 let weight = Array1::from_iter(self.norm.iter().zip(variance.iter()).map(|(&g, &v)| {
766 let credible2 = (g * g - v.max(0.0)).max(0.0);
767 let r = (credible2 + eps2).sqrt();
768 (1.0 / r).clamp(weight_floor, weight_ceiling)
769 }));
770 let invweight = weight.mapv(|u| 1.0 / u);
771 (weight, invweight)
772 }
773
774 fn directionalhessian_blocks(&self, direction_blocks: &Array2<f64>) -> Vec<Array2<f64>> {
775 let mut out = Vec::with_capacity(self.signal_blocks.nrows());
800 for (k, (v, q)) in self
801 .signal_blocks
802 .rows()
803 .into_iter()
804 .zip(direction_blocks.rows().into_iter())
805 .enumerate()
806 {
807 let dim = v.len();
808 let dot = v.iter().zip(q.iter()).map(|(a, b)| a * b).sum::<f64>();
809 let r3 = self.radius[k].powi(3);
810 let r5 = self.radius[k].powi(5);
811 let mut block = Array2::<f64>::eye(dim);
812 block.mapv_inplace(|x| -dot * x / r3);
813 for i in 0..dim {
814 for j in 0..dim {
815 block[[i, j]] -= (q[i] * v[j] + v[i] * q[j]) / r3;
816 block[[i, j]] += 3.0 * dot * v[i] * v[j] / r5;
817 }
818 }
819 out.push(block);
820 }
821 out
822 }
823
824 fn second_directionalhessian_blocks(
841 &self,
842 direction1_blocks: &Array2<f64>,
843 direction2_blocks: &Array2<f64>,
844 ) -> Vec<Array2<f64>> {
845 let mut out = Vec::with_capacity(self.signal_blocks.nrows());
846 for ((k, v), (a, b)) in self.signal_blocks.rows().into_iter().enumerate().zip(
847 direction1_blocks
848 .rows()
849 .into_iter()
850 .zip(direction2_blocks.rows().into_iter()),
851 ) {
852 let dim = v.len();
853 let dot = |x: ndarray::ArrayView1<'_, f64>, y: ndarray::ArrayView1<'_, f64>| {
854 x.iter().zip(y.iter()).map(|(p, q)| p * q).sum::<f64>()
855 };
856 let sa = dot(v, a);
857 let sb = dot(v, b);
858 let ab = dot(a, b);
859 let r = self.radius[k];
860 let r3 = r.powi(3);
861 let r5 = r.powi(5);
862 let r7 = r5 * r * r;
863 let diag = -ab / r3 + 3.0 * sa * sb / r5;
864 let mut block = Array2::<f64>::eye(dim);
865 block.mapv_inplace(|x| diag * x);
866 for i in 0..dim {
867 for j in 0..dim {
868 block[[i, j]] -= (a[i] * b[j] + b[i] * a[j]) / r3;
869 block[[i, j]] += 3.0 * sb * (a[i] * v[j] + v[i] * a[j]) / r5;
870 block[[i, j]] += 3.0 * ab * v[i] * v[j] / r5;
871 block[[i, j]] += 3.0 * sa * (b[i] * v[j] + v[i] * b[j]) / r5;
872 block[[i, j]] -= 15.0 * sa * sb * v[i] * v[j] / r7;
873 }
874 }
875 out.push(block);
876 }
877 out
878 }
879
880 fn log_epsilon_betahessian_blocks(&self) -> Vec<Array2<f64>> {
881 let mut out = Vec::with_capacity(self.signal_blocks.nrows());
882 for (k, row) in self.signal_blocks.rows().into_iter().enumerate() {
883 let dim = row.len();
884 let r3 = self.radius[k].powi(3);
885 let r5 = self.radius[k].powi(5);
886 let mut block = Array2::<f64>::eye(dim);
887 let eps2 = self.epsilon * self.epsilon;
888 block.mapv_inplace(|v| -eps2 * v / r3);
889 for i in 0..dim {
890 for j in 0..dim {
891 block[[i, j]] += 3.0 * eps2 * row[i] * row[j] / r5;
892 }
893 }
894 out.push(block);
895 }
896 out
897 }
898
899 fn log_epsilon_beta_mixed_second_blocks(&self) -> Array2<f64> {
900 let mut out = self.signal_blocks.clone();
901 let eps2 = self.epsilon * self.epsilon;
902 for (k, mut row) in out.rows_mut().into_iter().enumerate() {
903 let norm2 = self.norm[k] * self.norm[k];
904 let scale = eps2 * (eps2 - 2.0 * norm2) / self.radius[k].powi(5);
905 row.mapv_inplace(|v| v * scale);
906 }
907 out
908 }
909
910 fn log_epsilon_betahessian_second_blocks(&self) -> Vec<Array2<f64>> {
911 let mut out = Vec::with_capacity(self.signal_blocks.nrows());
912 let eps2 = self.epsilon * self.epsilon;
913 for (k, row) in self.signal_blocks.rows().into_iter().enumerate() {
914 let dim = row.len();
915 let norm2 = self.norm[k] * self.norm[k];
916 let r5 = self.radius[k].powi(5);
917 let r7 = self.radius[k].powi(7);
918 let mut block = Array2::<f64>::eye(dim);
919 block.mapv_inplace(|v| eps2 * (eps2 - 2.0 * norm2) * v / r5);
920 for i in 0..dim {
921 for j in 0..dim {
922 block[[i, j]] += 3.0 * eps2 * (2.0 * norm2 - 3.0 * eps2) * row[i] * row[j] / r7;
923 }
924 }
925 out.push(block);
926 }
927 out
928 }
929
930 fn log_epsilon_betahessian_directional_blocks(
931 &self,
932 direction_blocks: &Array2<f64>,
933 ) -> Vec<Array2<f64>> {
934 let mut out = Vec::with_capacity(self.signal_blocks.nrows());
935 let eps2 = self.epsilon * self.epsilon;
936 for (k, (v, q)) in self
937 .signal_blocks
938 .rows()
939 .into_iter()
940 .zip(direction_blocks.rows().into_iter())
941 .enumerate()
942 {
943 let dim = v.len();
944 let dot = v.iter().zip(q.iter()).map(|(a, b)| a * b).sum::<f64>();
945 let r5 = self.radius[k].powi(5);
946 let r7 = self.radius[k].powi(7);
947 let mut block = Array2::<f64>::eye(dim);
948 block.mapv_inplace(|x| 3.0 * eps2 * dot * x / r5);
949 for i in 0..dim {
950 for j in 0..dim {
951 block[[i, j]] += 3.0 * eps2 * (q[i] * v[j] + v[i] * q[j]) / r5;
952 block[[i, j]] -= 15.0 * eps2 * dot * v[i] * v[j] / r7;
953 }
954 }
955 out.push(block);
956 }
957 out
958 }
959}
960
961fn scalar_operatorgradient(operator: &Array2<f64>, coeff: &Array1<f64>) -> Array1<f64> {
962 operator.t().dot(coeff)
963}
964
965fn scalar_operatorhessian(operator: &Array2<f64>, diag: &Array1<f64>) -> Array2<f64> {
966 let mut weighted = operator.clone();
967 for (k, &w) in diag.iter().enumerate() {
968 weighted.row_mut(k).mapv_inplace(|v| v * w);
969 }
970 let gram = operator.t().dot(&weighted);
971 (&gram + &gram.t().to_owned()) * 0.5
972}
973
974fn grouped_operatorgradient(
975 d1: &Array2<f64>,
976 dimension: usize,
977 blocks: &Array2<f64>,
978) -> Result<Array1<f64>, EstimationError> {
979 if blocks.ncols() != dimension {
980 crate::bail_invalid_estim!(
981 "grouped gradient block dimension mismatch: got {}, expected {dimension}",
982 blocks.ncols()
983 );
984 }
985 if d1.nrows() != blocks.nrows() * dimension {
986 crate::bail_invalid_estim!(
987 "grouped gradient row mismatch: D1 has {} rows, blocks imply {}",
988 d1.nrows(),
989 blocks.nrows() * dimension
990 );
991 }
992 let mut out = Array1::<f64>::zeros(d1.ncols());
993 for k in 0..blocks.nrows() {
994 let gk = d1
995 .slice(s![k * dimension..(k + 1) * dimension, ..])
996 .to_owned();
997 out += &gk.t().dot(&blocks.row(k));
998 }
999 Ok(out)
1000}
1001
1002fn grouped_operatorhessian(
1003 d1: &Array2<f64>,
1004 dimension: usize,
1005 blocks: &[Array2<f64>],
1006) -> Result<Array2<f64>, EstimationError> {
1007 if d1.nrows() != blocks.len() * dimension {
1008 crate::bail_invalid_estim!(
1009 "grouped Hessian row mismatch: D1 has {} rows, blocks imply {}",
1010 d1.nrows(),
1011 blocks.len() * dimension
1012 );
1013 }
1014 let p = d1.ncols();
1015 let mut out = Array2::<f64>::zeros((p, p));
1016 for (k, block) in blocks.iter().enumerate() {
1017 if block.nrows() != dimension || block.ncols() != dimension {
1018 crate::bail_invalid_estim!(
1019 "grouped Hessian block {k} has shape {}x{}, expected {}x{}",
1020 block.nrows(),
1021 block.ncols(),
1022 dimension,
1023 dimension
1024 );
1025 }
1026 let gk = d1
1027 .slice(s![k * dimension..(k + 1) * dimension, ..])
1028 .to_owned();
1029 out += &gk.t().dot(&block.dot(&gk));
1030 }
1031 Ok((&out + &out.t().to_owned()) * 0.5)
1032}
1033
1034#[derive(Clone)]
1035struct SpatialPenaltyExactState {
1036 magnitude: CharbonnierScalarBlockState,
1037 gradient: CharbonnierGroupedBlockState,
1038 curvature: CharbonnierGroupedBlockState,
1039}
1040
1041fn collocationgradient_blocks(
1042 gradrows: &Array1<f64>,
1043 dimension: usize,
1044) -> Result<Array2<f64>, EstimationError> {
1045 if dimension == 0 || !gradrows.len().is_multiple_of(dimension) {
1046 crate::bail_invalid_estim!(
1047 "invalid collocation gradient layout: rows={}, dimension={dimension}",
1048 gradrows.len()
1049 );
1050 }
1051 let p = gradrows.len() / dimension;
1052 let mut out = Array2::<f64>::zeros((p, dimension));
1053 for k in 0..p {
1054 for axis in 0..dimension {
1055 out[[k, axis]] = gradrows[k * dimension + axis];
1056 }
1057 }
1058 Ok(out)
1059}
1060
1061fn collocationhessian_blocks(
1062 hessianrows: &Array1<f64>,
1063 dimension: usize,
1064) -> Result<Array2<f64>, EstimationError> {
1065 let block_dim = dimension.checked_mul(dimension).ok_or_else(|| {
1066 EstimationError::InvalidInput("invalid collocation Hessian dimension overflow".to_string())
1067 })?;
1068 if block_dim == 0 || !hessianrows.len().is_multiple_of(block_dim) {
1069 crate::bail_invalid_estim!(
1070 "invalid collocation Hessian layout: rows={}, dimension={dimension}",
1071 hessianrows.len()
1072 );
1073 }
1074 let p = hessianrows.len() / block_dim;
1075 let mut out = Array2::<f64>::zeros((p, block_dim));
1076 for k in 0..p {
1077 for idx in 0..block_dim {
1078 out[[k, idx]] = hessianrows[k * block_dim + idx];
1079 }
1080 }
1081 Ok(out)
1082}
1083
1084impl SpatialPenaltyExactState {
1085 fn from_beta_local(
1086 beta_local: ArrayView1<'_, f64>,
1087 cache: &SpatialOperatorRuntimeCache,
1088 epsilons: [f64; 3],
1089 ) -> Result<Self, EstimationError> {
1090 let gradientrows = cache.d1.dot(&beta_local);
1120 let hessianrows = cache.d2.dot(&beta_local);
1121 Ok(Self {
1122 magnitude: CharbonnierScalarBlockState::from_signal(
1123 cache.d0.dot(&beta_local),
1124 epsilons[0],
1125 ),
1126 gradient: CharbonnierGroupedBlockState::from_signal_blocks(
1127 collocationgradient_blocks(&gradientrows, cache.dimension)?,
1128 epsilons[1],
1129 ),
1130 curvature: CharbonnierGroupedBlockState::from_signal_blocks(
1131 collocationhessian_blocks(&hessianrows, cache.dimension)?,
1132 epsilons[2],
1133 ),
1134 })
1135 }
1136
1137 fn absolute_collocation_magnitudes(&self) -> (Array1<f64>, Array1<f64>, Array1<f64>) {
1138 (
1139 self.magnitude.absolute_signal(),
1140 self.gradient.norm_signal(),
1141 self.curvature.norm_signal(),
1142 )
1143 }
1144}
1145
1146fn robust_epsilon_from_samples(values: &[f64], min_epsilon_cfg: f64) -> f64 {
1147 if values.is_empty() {
1148 return min_epsilon_cfg.max(1e-12);
1149 }
1150 let mut clean = values
1151 .iter()
1152 .copied()
1153 .filter(|v| v.is_finite() && *v >= 0.0)
1154 .collect::<Vec<_>>();
1155 if clean.is_empty() {
1156 return min_epsilon_cfg.max(1e-12);
1157 }
1158 clean.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
1159
1160 let n = clean.len();
1161 let median = quantile_from_sorted(&clean, 0.5);
1162 let q75 = quantile_from_sorted(&clean, 0.75);
1163 let q95 = quantile_from_sorted(&clean, 0.95);
1164
1165 let mut abs_dev = clean
1166 .iter()
1167 .map(|v| (v - median).abs())
1168 .filter(|v| v.is_finite())
1169 .collect::<Vec<_>>();
1170 abs_dev.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
1171 let mad = 1.4826 * quantile_from_sorted(&abs_dev, 0.5);
1172
1173 let mut scale = median.max(mad).max(q75);
1183
1184 let delta = (f64::EPSILON.sqrt() * q95.max(1.0))
1186 .max(min_epsilon_cfg)
1187 .max(1e-12);
1188 let s_min = min_epsilon_cfg.max(1e-12);
1189
1190 if scale <= delta {
1192 let rms = (clean.iter().map(|v| v * v).sum::<f64>() / n as f64).sqrt();
1193 scale = q95.max(rms);
1194 }
1195 if scale <= delta {
1196 scale = s_min;
1197 }
1198
1199 let kappa = 1.0_f64;
1202 (kappa * scale).max(s_min)
1203}
1204
1205fn extract_spatial_operator_runtime_caches(
1206 spec: &TermCollectionSpec,
1207 design: &TermCollectionDesign,
1208) -> Result<Vec<SpatialOperatorRuntimeCache>, EstimationError> {
1209 let smooth_start = design
1210 .design
1211 .ncols()
1212 .saturating_sub(design.smooth.total_smooth_cols());
1213 let mut out = Vec::<SpatialOperatorRuntimeCache>::new();
1214 for (term_idx, (termspec, term_fit)) in spec
1215 .smooth_terms
1216 .iter()
1217 .zip(design.smooth.terms.iter())
1218 .enumerate()
1219 {
1220 let Some(global_range) = design
1221 .smooth_term_penalty_range(term_idx)
1222 .map_err(EstimationError::InvalidInput)?
1223 else {
1224 continue;
1225 };
1226 let global_base_idx = global_range.start;
1227 let mut mass_local_idx = None;
1228 let mut tension_local_idx = None;
1229 let mut stiffness_local_idx = None;
1230 let mut mass_norm = None;
1231 let mut tension_norm = None;
1232 let mut stiffness_norm = None;
1233 for (active_local_idx, penalty) in term_fit.active_penalties.iter().enumerate() {
1234 match penalty.info.source {
1235 PenaltySource::OperatorMass => {
1236 mass_local_idx = Some(active_local_idx);
1237 mass_norm = Some(penalty.info.normalization_scale);
1238 }
1239 PenaltySource::OperatorTension => {
1240 tension_local_idx = Some(active_local_idx);
1241 tension_norm = Some(penalty.info.normalization_scale);
1242 }
1243 PenaltySource::OperatorStiffness => {
1244 stiffness_local_idx = Some(active_local_idx);
1245 stiffness_norm = Some(penalty.info.normalization_scale);
1246 }
1247 PenaltySource::Primary
1250 | PenaltySource::DoublePenaltyNullspace
1251 | PenaltySource::OperatorRelevance { .. }
1252 | PenaltySource::TensorMarginal { .. }
1253 | PenaltySource::TensorSeparable { .. }
1254 | PenaltySource::TensorGlobalRidge
1255 | PenaltySource::Other(_) => {}
1256 }
1257 }
1258 let (
1271 Some(mass_local),
1272 Some(tension_local),
1273 Some(stiffness_local),
1274 Some(mass_scale),
1275 Some(tension_scale),
1276 Some(stiffness_scale),
1277 ) = (
1278 mass_local_idx,
1279 tension_local_idx,
1280 stiffness_local_idx,
1281 mass_norm,
1282 tension_norm,
1283 stiffness_norm,
1284 )
1285 else {
1286 continue;
1287 };
1288 let mass_global_idx = global_base_idx + mass_local;
1289 let tension_global_idx = global_base_idx + tension_local;
1290 let stiffness_global_idx = global_base_idx + stiffness_local;
1291
1292 let (feature_cols, mut d0, mut d1, mut d2, collocation_points, dim, center_mass_rows) =
1293 match (&termspec.basis, &term_fit.metadata) {
1294 (
1295 SmoothBasisSpec::Matern { feature_cols, .. },
1296 BasisMetadata::Matern {
1297 centers,
1298 length_scale,
1299 nu,
1300 include_intercept,
1301 identifiability_transform,
1302 aniso_log_scales,
1303 input_scale,
1304 ..
1305 },
1306 ) => {
1307 let collocation_length_scale = input_scale
1313 .to_standardized_units(*length_scale)
1314 .standardized_value();
1315 let ops = build_matern_collocation_operator_matrices(
1316 centers.view(),
1317 None,
1318 collocation_length_scale,
1319 *nu,
1320 *include_intercept,
1321 identifiability_transform.as_ref().map(|z| z.view()),
1322 aniso_log_scales.as_deref(),
1323 )?;
1324 (
1325 feature_cols.clone(),
1326 ops.d0,
1327 ops.d1,
1328 ops.d2,
1329 ops.collocation_points,
1330 centers.ncols(),
1331 false,
1332 )
1333 }
1334 (
1335 SmoothBasisSpec::Duchon { feature_cols, .. },
1336 BasisMetadata::Duchon {
1337 centers,
1338 length_scale,
1339 power,
1340 nullspace_order,
1341 identifiability_transform,
1342 input_scale,
1343 aniso_log_scales,
1344 operator_collocation_points: Some(collocation_points),
1345 radial_reparam,
1346 ..
1347 },
1348 ) => {
1349 let collocation_length_scale = (*length_scale).map(|length| {
1350 input_scale.to_standardized_units(length).standardized_value()
1351 });
1352 let ops =
1353 gam_terms::basis::build_duchon_collocation_operator_matriceswithworkspace(
1354 centers.view(),
1355 collocation_points.view(),
1356 None,
1357 collocation_length_scale,
1358 *power,
1359 *nullspace_order,
1360 aniso_log_scales.as_deref(),
1361 identifiability_transform.as_ref().map(|z| z.view()),
1362 2,
1363 radial_reparam.as_ref().map(|v| v.view()),
1364 &mut BasisWorkspace::default(),
1365 )?;
1366 (
1367 feature_cols.clone(),
1368 ops.d0,
1369 ops.d1,
1370 ops.d2,
1371 ops.collocation_points,
1372 centers.ncols(),
1373 true,
1374 )
1375 }
1376 _ => continue,
1377 };
1378 if center_mass_rows && d0.nrows() > 0 && d0.ncols() > 0 {
1379 let means = d0.sum_axis(Axis(0)).mapv(|v| v / d0.nrows() as f64);
1380 for mut row in d0.rows_mut() {
1381 row -= &means;
1382 }
1383 }
1384
1385 let mass_scale = mass_scale.max(1e-12).sqrt();
1403 let tension_scale = tension_scale.max(1e-12).sqrt();
1404 let stiffness_scale = stiffness_scale.max(1e-12).sqrt();
1405 d0.mapv_inplace(|v| v / mass_scale);
1406 d1.mapv_inplace(|v| v / tension_scale);
1407 d2.mapv_inplace(|v| v / stiffness_scale);
1408
1409 let coeff_global_range =
1410 (smooth_start + term_fit.coeff_range.start)..(smooth_start + term_fit.coeff_range.end);
1411 if d0.ncols() != coeff_global_range.len()
1412 || d1.ncols() != coeff_global_range.len()
1413 || d2.ncols() != coeff_global_range.len()
1414 {
1415 crate::bail_invalid_estim!(
1416 "spatial operator dimension mismatch for term '{}': D0 cols={}, D1 cols={}, D2 cols={}, coeffs={}",
1417 term_fit.name,
1418 d0.ncols(),
1419 d1.ncols(),
1420 d2.ncols(),
1421 coeff_global_range.len()
1422 );
1423 }
1424 out.push(SpatialOperatorRuntimeCache {
1425 termname: term_fit.name.clone(),
1426 feature_cols,
1427 coeff_global_range,
1428 mass_penalty_global_idx: mass_global_idx,
1429 tension_penalty_global_idx: tension_global_idx,
1430 stiffness_penalty_global_idx: stiffness_global_idx,
1431 d0,
1432 d1,
1433 d2,
1434 collocation_points,
1435 dimension: dim,
1436 });
1437 }
1438 Ok(out)
1439}
1440
1441fn scalar_operator_response_variance(
1453 operator: &Array2<f64>,
1454 cov_local: &Array2<f64>,
1455) -> Array1<f64> {
1456 Array1::from_iter(operator.rows().into_iter().map(|row| {
1457 let s = cov_local.dot(&row);
1458 row.dot(&s).max(0.0)
1459 }))
1460}
1461
1462fn grouped_operator_response_variance(
1473 operator: &Array2<f64>,
1474 block_dim: usize,
1475 cov_local: &Array2<f64>,
1476) -> Result<Array1<f64>, EstimationError> {
1477 if block_dim == 0 || !operator.nrows().is_multiple_of(block_dim) {
1478 crate::bail_invalid_estim!(
1479 "grouped variance row layout invalid: rows={}, block_dim={block_dim}",
1480 operator.nrows()
1481 );
1482 }
1483 let p = operator.nrows() / block_dim;
1484 let mut out = Array1::<f64>::zeros(p);
1485 for k in 0..p {
1486 let mut acc = 0.0;
1487 for axis in 0..block_dim {
1488 let row = operator.row(k * block_dim + axis);
1489 let s = cov_local.dot(&row);
1490 acc += row.dot(&s);
1491 }
1492 out[k] = acc.max(0.0);
1493 }
1494 Ok(out)
1495}
1496
1497fn compute_spatial_adaptiveweights_for_beta(
1498 beta: &Array1<f64>,
1499 caches: &[SpatialOperatorRuntimeCache],
1500 epsilon_0: f64,
1501 epsilon_g: f64,
1502 epsilon_c: f64,
1503 weight_floor: f64,
1504 weight_ceiling: f64,
1505 beta_covariance: Option<&Array2<f64>>,
1506) -> Result<Vec<SpatialAdaptiveWeights>, EstimationError> {
1507 caches
1539 .iter()
1540 .map(|cache| {
1541 let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
1542 let exact = SpatialPenaltyExactState::from_beta_local(
1543 beta_local,
1544 cache,
1545 [epsilon_0, epsilon_g, epsilon_c],
1546 )?;
1547 let cov_local = beta_covariance.map(|cov| {
1548 cov.slice(s![
1549 cache.coeff_global_range.clone(),
1550 cache.coeff_global_range.clone()
1551 ])
1552 .to_owned()
1553 });
1554 let dim = cache.dimension;
1555 let (var_0, var_g, var_c) = match cov_local.as_ref() {
1556 Some(cov) => (
1557 scalar_operator_response_variance(&cache.d0, cov),
1558 grouped_operator_response_variance(&cache.d1, dim, cov)?,
1559 grouped_operator_response_variance(&cache.d2, dim * dim, cov)?,
1560 ),
1561 None => (
1562 Array1::<f64>::zeros(exact.magnitude.signal.len()),
1563 Array1::<f64>::zeros(exact.gradient.norm.len()),
1564 Array1::<f64>::zeros(exact.curvature.norm.len()),
1565 ),
1566 };
1567 let (_, inv_0) = exact.magnitude.surrogateweights_posterior_snr(
1568 &var_0,
1569 weight_floor,
1570 weight_ceiling,
1571 );
1572 let (_, inv_g) =
1573 exact
1574 .gradient
1575 .surrogateweights_posterior_snr(&var_g, weight_floor, weight_ceiling);
1576 let (_, inv_c) = exact.curvature.surrogateweights_posterior_snr(
1577 &var_c,
1578 weight_floor,
1579 weight_ceiling,
1580 );
1581 Ok(SpatialAdaptiveWeights {
1582 inv_magweight: inv_0,
1583 invgradweight: inv_g,
1584 inv_lapweight: inv_c,
1585 })
1586 })
1587 .collect()
1588}
1589
1590fn compute_initial_epsilons(
1591 beta: &Array1<f64>,
1592 caches: &[SpatialOperatorRuntimeCache],
1593 min_epsilon: f64,
1594) -> Result<(f64, f64, f64), EstimationError> {
1595 let mut fvals = Vec::<f64>::new();
1596 let mut gvals = Vec::<f64>::new();
1597 let mut cvals = Vec::<f64>::new();
1598 for cache in caches {
1599 let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
1600 let exact = SpatialPenaltyExactState::from_beta_local(
1601 beta_local,
1602 cache,
1603 [min_epsilon, min_epsilon, min_epsilon],
1604 )?;
1605 let (f, g, c) = exact.absolute_collocation_magnitudes();
1606 fvals.extend(f.iter().copied());
1607 gvals.extend(g.iter().copied());
1608 cvals.extend(c.iter().copied());
1609 }
1610 let eps_0 = robust_epsilon_from_samples(&fvals, min_epsilon);
1616 let eps_g = robust_epsilon_from_samples(&gvals, min_epsilon);
1617 let eps_c = robust_epsilon_from_samples(&cvals, min_epsilon);
1618 Ok((eps_0, eps_g, eps_c))
1619}
1620
1621fn exact_spatial_adaptive_penalty_index_set(
1622 caches: &[SpatialOperatorRuntimeCache],
1623) -> BTreeSet<usize> {
1624 let mut out = BTreeSet::new();
1625 for cache in caches {
1626 out.insert(cache.mass_penalty_global_idx);
1627 out.insert(cache.tension_penalty_global_idx);
1628 out.insert(cache.stiffness_penalty_global_idx);
1629 }
1630 out
1631}
1632
1633fn checked_fit_log_lambdas(
1634 lambdas: &Array1<f64>,
1635 context: &str,
1636) -> Result<Array1<f64>, EstimationError> {
1637 let values = lambdas
1638 .iter()
1639 .copied()
1640 .enumerate()
1641 .map(|(coordinate, lambda)| {
1642 gam_problem::checked_log_strength(lambda).map_err(|error| {
1643 EstimationError::InvalidInput(format!(
1644 "{context} lambda coordinate {coordinate} is outside the canonical physical-strength domain: {error}"
1645 ))
1646 })
1647 })
1648 .collect::<Result<Vec<_>, _>>()?;
1649 Ok(Array1::from_vec(values))
1650}
1651
1652fn build_spatial_adaptive_hyperspecs(cache_count: usize) -> Vec<SpatialAdaptiveHyperSpec> {
1653 let mut out = Vec::with_capacity(cache_count * 3 + 3);
1654 for cache_index in 0..cache_count {
1655 out.push(SpatialAdaptiveHyperSpec {
1656 cache_index,
1657 kind: SpatialAdaptiveHyperKind::LogLambdaMagnitude,
1658 });
1659 out.push(SpatialAdaptiveHyperSpec {
1660 cache_index,
1661 kind: SpatialAdaptiveHyperKind::LogLambdaGradient,
1662 });
1663 out.push(SpatialAdaptiveHyperSpec {
1664 cache_index,
1665 kind: SpatialAdaptiveHyperKind::LogLambdaCurvature,
1666 });
1667 }
1668 out.push(SpatialAdaptiveHyperSpec {
1669 cache_index: 0,
1670 kind: SpatialAdaptiveHyperKind::LogEpsilonMagnitude,
1671 });
1672 out.push(SpatialAdaptiveHyperSpec {
1673 cache_index: 0,
1674 kind: SpatialAdaptiveHyperKind::LogEpsilonGradient,
1675 });
1676 out.push(SpatialAdaptiveHyperSpec {
1677 cache_index: 0,
1678 kind: SpatialAdaptiveHyperKind::LogEpsilonCurvature,
1679 });
1680 out
1681}
1682
1683fn penalty_matrixwith_local_block(
1684 total_dim: usize,
1685 coeff_range: Range<usize>,
1686 local: &Array2<f64>,
1687) -> Array2<f64> {
1688 let mut out = Array2::<f64>::zeros((total_dim, total_dim));
1689 out.slice_mut(s![coeff_range.clone(), coeff_range])
1690 .assign(local);
1691 out
1692}
1693
1694fn fit_term_collectionwith_exact_spatial_adaptive_regularization(
1695 baseline: FittedTermCollection,
1696 y: ArrayView1<'_, f64>,
1697 weights: ArrayView1<'_, f64>,
1698 offset: ArrayView1<'_, f64>,
1699 family: LikelihoodSpec,
1700 options: &FitOptions,
1701 runtime_caches: &[SpatialOperatorRuntimeCache],
1702) -> Result<FittedTermCollection, EstimationError> {
1703 let adaptive_opts = options.adaptive_regularization.clone().unwrap_or_default();
1732 let adaptive_penalty_indices = exact_spatial_adaptive_penalty_index_set(runtime_caches);
1733 let p_total = baseline.design.design.ncols();
1734 if baseline.fit.lambdas.len() != baseline.design.penalties.len() {
1735 crate::bail_invalid_estim!(
1736 "exact spatial adaptive fit received {} baseline lambdas for {} penalties",
1737 baseline.fit.lambdas.len(),
1738 baseline.design.penalties.len(),
1739 );
1740 }
1741 let baseline_log_lambdas =
1742 checked_fit_log_lambdas(&baseline.fit.lambdas, "exact spatial adaptive baseline")?;
1743 for (cache_idx, cache) in runtime_caches.iter().enumerate() {
1744 for (operator, penalty_idx) in [
1745 ("mass", cache.mass_penalty_global_idx),
1746 ("tension", cache.tension_penalty_global_idx),
1747 ("stiffness", cache.stiffness_penalty_global_idx),
1748 ] {
1749 if penalty_idx >= baseline.fit.lambdas.len() {
1750 crate::bail_invalid_estim!(
1751 "exact spatial adaptive cache {cache_idx} {operator} penalty index {penalty_idx} is out of bounds for {} baseline lambdas",
1752 baseline.fit.lambdas.len(),
1753 );
1754 }
1755 }
1756 }
1757 struct RetainedPenaltySetup {
1758 global_idx: usize,
1759 global_penalty: Array2<f64>,
1760 nullspace_dim: usize,
1761 log_lambda: f64,
1762 }
1763 use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
1764 let retained_setups = baseline
1765 .design
1766 .penalties
1767 .par_iter()
1768 .enumerate()
1769 .map(|(idx, bp)| {
1770 if adaptive_penalty_indices.contains(&idx) {
1771 return None;
1772 }
1773 Some(RetainedPenaltySetup {
1774 global_idx: idx,
1775 global_penalty: bp.to_global(p_total),
1776 nullspace_dim: baseline
1777 .design
1778 .nullspace_dims
1779 .get(idx)
1780 .copied()
1781 .unwrap_or(0),
1782 log_lambda: baseline_log_lambdas[idx],
1783 })
1784 })
1785 .collect::<Vec<_>>();
1786 let retained_count = retained_setups
1787 .iter()
1788 .filter(|setup| setup.is_some())
1789 .count();
1790 let mut retained_penalties = Vec::<Array2<f64>>::with_capacity(retained_count);
1791 let mut retained_nullspace_dims = Vec::<usize>::with_capacity(retained_count);
1792 let mut retained_log_lambdas = Vec::<f64>::with_capacity(retained_count);
1793 let mut retained_global_indices = Vec::<usize>::with_capacity(retained_count);
1794 for setup in retained_setups.into_iter().flatten() {
1795 retained_penalties.push(setup.global_penalty);
1796 retained_nullspace_dims.push(setup.nullspace_dim);
1797 retained_log_lambdas.push(setup.log_lambda);
1798 retained_global_indices.push(setup.global_idx);
1799 }
1800
1801 let (eps_0_init, eps_g_init, eps_c_init) = compute_initial_epsilons(
1802 &baseline.fit.beta,
1803 runtime_caches,
1804 adaptive_opts.min_epsilon,
1805 )?;
1806 let mut initial_theta =
1807 Array1::<f64>::zeros(retained_penalties.len() + runtime_caches.len() * 3 + 3);
1808 for (idx, value) in retained_log_lambdas.iter().enumerate() {
1809 initial_theta[idx] = *value;
1810 }
1811 let adaptive_log_lambda_components = runtime_caches
1812 .par_iter()
1813 .map(|cache| {
1814 [
1815 baseline_log_lambdas[cache.mass_penalty_global_idx],
1816 baseline_log_lambdas[cache.tension_penalty_global_idx],
1817 baseline_log_lambdas[cache.stiffness_penalty_global_idx],
1818 ]
1819 })
1820 .collect::<Vec<_>>();
1821 let mut at = retained_penalties.len();
1822 for logs in &adaptive_log_lambda_components {
1823 initial_theta[at] = logs[0];
1824 initial_theta[at + 1] = logs[1];
1825 initial_theta[at + 2] = logs[2];
1826 at += 3;
1827 }
1828 let minimum_log_epsilon = gam_problem::checked_log_strength(adaptive_opts.min_epsilon)
1829 .map_err(|error| {
1830 EstimationError::InvalidInput(format!(
1831 "adaptive minimum epsilon is outside the canonical positive-strength domain: {error}"
1832 ))
1833 })?;
1834 for (slot, epsilon) in [eps_0_init, eps_g_init, eps_c_init].into_iter().enumerate() {
1835 initial_theta[at + slot] =
1836 gam_problem::checked_log_strength(epsilon.max(adaptive_opts.min_epsilon)).map_err(
1837 |error| {
1838 EstimationError::InvalidInput(format!(
1839 "adaptive initial epsilon coordinate {slot} is outside the canonical positive-strength domain: {error}"
1840 ))
1841 },
1842 )?;
1843 }
1844
1845 let hyperspecs = build_spatial_adaptive_hyperspecs(runtime_caches.len());
1846 let zero_psi_op: std::sync::Arc<dyn gam_custom_family::CustomFamilyPsiDerivativeOperator> =
1847 std::sync::Arc::new(gam_custom_family::ZeroPsiDerivativeOperator::new(
1848 baseline.design.design.nrows(),
1849 baseline.design.design.ncols(),
1850 hyperspecs.len(),
1851 ));
1852 let derivative_blocks = vec![
1853 hyperspecs
1854 .par_iter()
1855 .map(|_| CustomFamilyBlockPsiDerivative {
1856 penalty_index: None,
1857 x_psi: Array2::<f64>::zeros((0, 0)),
1858 s_psi: Array2::<f64>::zeros((0, 0)),
1859 s_psi_components: None,
1860 s_psi_penalty_components: None,
1861 x_psi_psi: None,
1862 s_psi_psi: None,
1863 s_psi_psi_components: None,
1864 s_psi_psi_penalty_components: None,
1865 implicit_operator: Some(std::sync::Arc::clone(&zero_psi_op)),
1866 implicit_axis: 0,
1867 implicit_group_id: None,
1868 })
1869 .collect::<Vec<_>>(),
1870 ];
1871
1872 let mixture_link_state = options
1873 .mixture_link
1874 .clone()
1875 .as_ref()
1876 .map(state_fromspec)
1877 .transpose()
1878 .map_err(EstimationError::InvalidInput)?;
1879 let sas_link_state = options
1880 .sas_link
1881 .map(|spec| {
1882 if family.is_binomial_beta_logistic() {
1883 state_from_beta_logisticspec(spec)
1884 } else {
1885 state_from_sasspec(spec)
1886 }
1887 })
1888 .transpose()
1889 .map_err(EstimationError::InvalidInput)?;
1890 let latent_cloglog_state = options.latent_cloglog;
1891 let shared_y = Arc::new(y.to_owned());
1892 let sharedweights = Arc::new(weights.to_owned());
1893 let shared_design = baseline
1894 .design
1895 .design
1896 .try_to_dense_arc("spatial adaptive exact hyperfit design")
1897 .map_err(EstimationError::InvalidInput)?;
1898 let shared_offset = Arc::new(offset.to_owned());
1899 let shared_runtime_caches = Arc::new(runtime_caches.to_vec());
1900 let shared_hyperspecs = Arc::new(hyperspecs.clone());
1901 let zero_quadratic = ValidatedFixedQuadraticHessian::zero(
1902 baseline.design.design.ncols(),
1903 )
1904 .map_err(EstimationError::InvalidInput)?;
1905 let base_family = SpatialAdaptiveExactFamily {
1906 family: family.clone(),
1907 latent_cloglog_state,
1908 mixture_link_state: mixture_link_state.clone(),
1909 sas_link_state,
1910 y: shared_y.clone(),
1911 weights: sharedweights.clone(),
1912 design: shared_design.clone(),
1913 offset: shared_offset.clone(),
1914 linear_constraints: baseline.design.linear_constraints.clone(),
1915 runtime_caches: shared_runtime_caches.clone(),
1916 adaptive_params: Vec::new(),
1917 fixed_quadratic_hessian: zero_quadratic.clone(),
1918 hyperspecs: shared_hyperspecs.clone(),
1919 exact_eval_cache: Arc::new(Mutex::new(None)),
1920 };
1921
1922 let rho_dim = retained_penalties.len();
1923 let operator_slots_end = rho_dim + runtime_caches.len() * 3;
1924 const UNIFIED_LOG_WINDOW: f64 = 6.0;
1934 const RETAINED_LAMBDA_LOG_LOWER_FLOOR: f64 = -30.0;
1935 const RETAINED_LAMBDA_LOG_UPPER_CAP: f64 = 30.0;
1936 const OPERATOR_LAMBDA_LOG_LOWER_FLOOR: f64 = -10.0;
1937 const OPERATOR_LAMBDA_LOG_UPPER_CAP: f64 = 30.0;
1938 let epsilon_floor_log = minimum_log_epsilon;
1939 let anchored_bound = |idx: usize, sign: f64| -> f64 {
1940 let raw = initial_theta[idx] + sign * UNIFIED_LOG_WINDOW;
1941 if idx < rho_dim {
1942 raw.clamp(
1943 RETAINED_LAMBDA_LOG_LOWER_FLOOR,
1944 RETAINED_LAMBDA_LOG_UPPER_CAP,
1945 )
1946 } else if idx < operator_slots_end {
1947 raw.clamp(
1948 OPERATOR_LAMBDA_LOG_LOWER_FLOOR,
1949 OPERATOR_LAMBDA_LOG_UPPER_CAP,
1950 )
1951 } else {
1952 raw.clamp(epsilon_floor_log, gam_problem::LOG_STRENGTH_MAX)
1953 }
1954 };
1955 let eps_lower =
1956 Array1::from_iter((0..initial_theta.len()).map(|idx| anchored_bound(idx, -1.0)));
1957 let eps_upper = Array1::from_iter((0..initial_theta.len()).map(|idx| anchored_bound(idx, 1.0)));
1958 let blockspec = ParameterBlockSpec {
1959 name: "eta".to_string(),
1960 design: baseline.design.design.clone(),
1961 offset: offset.to_owned(),
1962 penalties: retained_penalties
1963 .iter()
1964 .cloned()
1965 .map(PenaltyMatrix::Dense)
1966 .collect(),
1967 nullspace_dims: retained_nullspace_dims.clone(),
1968 initial_log_lambdas: Array1::from_vec(retained_log_lambdas.clone()),
1969 initial_beta: Some(baseline.fit.beta.clone()),
1970 gauge_priority: 100,
1971 jacobian_callback: None,
1972 stacked_design: None,
1973 stacked_offset: None,
1974 };
1975 let screening_cap = Arc::new(AtomicUsize::new(0));
1976 let outer_opts = BlockwiseFitOptions {
1977 inner_max_cycles: options.max_iter,
1978 inner_tol: options.tol,
1979 outer_max_iter: options.max_iter,
1980 outer_tol: options.tol,
1981 compute_covariance: false,
1982 screening_max_inner_iterations: Some(Arc::clone(&screening_cap)),
1983 ..BlockwiseFitOptions::default()
1984 };
1985
1986 use gam_problem::{DeclaredHessianForm, Derivative, HessianValue, OuterEval};
1987 use gam_solve::rho_optimizer::OuterProblem;
1988
1989 struct SpatialAdaptiveOuterState {
1990 warm_cache: Option<CustomFamilyWarmStart>,
1991 terminal_mode: Option<(Array1<f64>, f64, CustomFamilyOwnedMode)>,
1992 last_eval: Option<(
1993 Array1<f64>,
1994 f64,
1995 Array1<f64>,
1996 HessianValue,
1997 CustomFamilyWarmStart,
1998 )>,
1999 }
2000
2001 struct DecodedSpatialAdaptiveTheta {
2002 rho: Array1<f64>,
2003 retained_lambdas: Array1<f64>,
2004 adaptive_params: Vec<SpatialAdaptiveTermHyperParams>,
2005 epsilon: [f64; 3],
2006 }
2007
2008 let n_theta = initial_theta.len();
2009
2010 let theta_bounds = Some((eps_lower.clone(), eps_upper.clone()));
2011
2012 let decode_theta =
2013 |theta: &Array1<f64>| -> Result<DecodedSpatialAdaptiveTheta, EstimationError> {
2014 let physical = gam_problem::checked_exp_log_strengths(theta.iter().copied()).map_err(
2015 |error| {
2016 EstimationError::InvalidInput(format!(
2017 "exact spatial adaptive outer coordinate is outside the canonical log-strength domain: {error}"
2018 ))
2019 },
2020 )?;
2021 let rho = theta.slice(s![..rho_dim]).to_owned();
2022 let retained_lambdas = Array1::from_vec(physical[..rho_dim].to_vec());
2023 let adaptive_lambda_start = rho_dim;
2024 let adaptive_lambda_end = adaptive_lambda_start + runtime_caches.len() * 3;
2025 let eps = [
2026 physical[adaptive_lambda_end],
2027 physical[adaptive_lambda_end + 1],
2028 physical[adaptive_lambda_end + 2],
2029 ];
2030 let adaptive_params = runtime_caches
2031 .iter()
2032 .enumerate()
2033 .map(|(cache_idx, _)| SpatialAdaptiveTermHyperParams {
2034 lambda: [
2035 physical[adaptive_lambda_start + cache_idx * 3],
2036 physical[adaptive_lambda_start + cache_idx * 3 + 1],
2037 physical[adaptive_lambda_start + cache_idx * 3 + 2],
2038 ],
2039 epsilon: eps,
2040 })
2041 .collect::<Vec<_>>();
2042 Ok(DecodedSpatialAdaptiveTheta {
2043 rho,
2044 retained_lambdas,
2045 adaptive_params,
2046 epsilon: eps,
2047 })
2048 };
2049 let clamp_theta = |theta: &Array1<f64>| -> Array1<f64> {
2055 Array1::from_shape_fn(theta.len(), |i| theta[i].clamp(eps_lower[i], eps_upper[i]))
2056 };
2057 let realize_hyper_layout = |theta: &Array1<f64>| {
2058 gam_custom_family::CustomFamilyHyperLayout::new(
2059 derivative_blocks.clone(),
2060 Vec::new(),
2061 theta.slice(s![rho_dim..]).to_owned(),
2062 )
2063 .map_err(EstimationError::InvalidInput)
2064 };
2065 let analytic_outer_hessian_available =
2066 gam_custom_family::joint_exact_analytic_outer_hessian_available()
2067 && base_family
2068 .exact_outer_derivative_order(std::slice::from_ref(&blockspec), &outer_opts)
2069 .has_hessian()
2070 && gam_custom_family::exact_newton_outer_geometry_supports_second_order_solver(
2071 &base_family,
2072 );
2073 let problem = OuterProblem::new(n_theta)
2079 .with_gradient(Derivative::Analytic)
2080 .with_hessian(if analytic_outer_hessian_available {
2081 DeclaredHessianForm::Either
2082 } else {
2083 DeclaredHessianForm::Unavailable
2084 })
2085 .with_prefer_gradient_only(true)
2086 .with_fallback_policy(gam_solve::rho_optimizer::FallbackPolicy::Disabled)
2087 .with_psi_dim(n_theta.saturating_sub(rho_dim))
2088 .with_tolerance(options.tol)
2089 .with_rel_cost_tolerance(Some(options.tol))
2094 .with_max_iter(options.max_iter)
2095 .with_seed_config(gam_problem::SeedConfig::default())
2096 .with_screening_cap(Arc::clone(&screening_cap))
2097 .with_initial_rho(initial_theta.clone());
2098 let problem = if let Some((lo, hi)) = theta_bounds {
2099 problem.with_bounds(lo, hi)
2100 } else {
2101 problem
2102 };
2103
2104 let eval_outer = |st: &mut SpatialAdaptiveOuterState,
2105 theta: &Array1<f64>,
2106 order: gam_solve::rho_optimizer::OuterEvalOrder|
2107 -> Result<OuterEval, EstimationError> {
2108 let decoded = decode_theta(theta)?;
2109
2110 if let Some((cached_theta, cached_cost, cached_grad, cached_hess, cached_warm)) =
2111 &st.last_eval
2112 && cached_theta.len() == theta.len()
2113 && cached_theta
2114 .iter()
2115 .zip(theta.iter())
2116 .all(|(&a, &b)| a.to_bits() == b.to_bits())
2117 && st
2118 .terminal_mode
2119 .as_ref()
2120 .is_some_and(|(mode_theta, mode_objective, _)| {
2121 mode_theta.len() == theta.len()
2122 && mode_theta
2123 .iter()
2124 .zip(theta.iter())
2125 .all(|(&a, &b)| a.to_bits() == b.to_bits())
2126 && mode_objective.to_bits() == cached_cost.to_bits()
2127 })
2128 && (!matches!(
2129 order,
2130 gam_solve::rho_optimizer::OuterEvalOrder::ValueGradientHessian
2131 ) || analytic_outer_hessian_available)
2132 {
2133 st.warm_cache = Some(cached_warm.clone());
2134 return Ok(OuterEval {
2135 cost: *cached_cost,
2136 gradient: cached_grad.clone(),
2137 hessian: if matches!(
2138 order,
2139 gam_solve::rho_optimizer::OuterEvalOrder::ValueGradientHessian
2140 ) && analytic_outer_hessian_available
2141 {
2142 cached_hess.clone()
2143 } else {
2144 HessianValue::Unavailable
2145 },
2146 inner_beta_hint: None,
2147 });
2148 }
2149
2150 let family_eval =
2151 base_family.with_adaptive_params(decoded.adaptive_params, zero_quadratic.clone());
2152 let hyper_layout = realize_hyper_layout(theta)?;
2153 let need_hessian = matches!(
2154 order,
2155 gam_solve::rho_optimizer::OuterEvalOrder::ValueGradientHessian
2156 ) && analytic_outer_hessian_available;
2157 let owned = evaluate_custom_family_joint_hyper_owned(
2158 &family_eval,
2159 std::slice::from_ref(&blockspec),
2160 &outer_opts,
2161 &decoded.rho,
2162 &hyper_layout,
2163 st.warm_cache.as_ref(),
2164 if need_hessian {
2165 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueGradientHessian
2166 } else {
2167 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueAndGradient
2168 },
2169 )
2170 .map_err(|e| {
2179 EstimationError::CustomFamily(e)
2180 .wrap_preserving_trial_point("spatial adaptive eval failed")
2181 })?;
2182 if !owned.result.inner_converged {
2183 st.warm_cache = Some(owned.result.warm_start.clone());
2184 return Err(EstimationError::TrialPointRefused {
2187 reason: "exact spatial adaptive inner solve did not converge".to_string(),
2188 });
2189 }
2190 if !owned.result.objective.is_finite()
2191 || owned.result.gradient.iter().any(|v| !v.is_finite())
2192 {
2193 return Err(EstimationError::TrialPointRefused {
2196 reason: "exact spatial adaptive objective returned non-finite values".to_string(),
2197 });
2198 }
2199 let hessian_result = if need_hessian {
2200 if !owned.result.outer_hessian.is_analytic() {
2201 return Err(EstimationError::RemlOptimizationFailed(
2202 "exact spatial adaptive objective did not return an exact outer Hessian"
2203 .to_string(),
2204 ));
2205 }
2206 match owned.result.outer_hessian.dim() {
2207 Some(dim) if dim == theta.len() => {}
2208 Some(dim) => {
2209 return Err(EstimationError::RemlOptimizationFailed(format!(
2210 "exact spatial adaptive outer Hessian dimension mismatch: got {dim}, expected {}",
2211 theta.len(),
2212 )));
2213 }
2214 None => {
2215 return Err(EstimationError::RemlOptimizationFailed(
2216 "exact spatial adaptive objective did not report an outer Hessian dimension"
2217 .to_string(),
2218 ));
2219 }
2220 }
2221 st.last_eval = Some((
2222 theta.to_owned(),
2223 owned.result.objective,
2224 owned.result.gradient.clone(),
2225 owned.result.outer_hessian.clone(),
2226 owned.result.warm_start.clone(),
2227 ));
2228 owned.result.outer_hessian
2229 } else {
2230 HessianValue::Unavailable
2231 };
2232 let objective = owned.result.objective;
2233 let gradient = owned.result.gradient;
2234 st.warm_cache = Some(owned.result.warm_start);
2235 st.terminal_mode = Some((theta.to_owned(), objective, owned.mode));
2236 Ok(OuterEval {
2237 cost: objective,
2238 gradient,
2239 hessian: hessian_result,
2240 inner_beta_hint: None,
2241 })
2242 };
2243
2244 let mut obj = problem.build_objective_with_screening_proxy(
2245 SpatialAdaptiveOuterState {
2246 warm_cache: None,
2247 terminal_mode: None,
2248 last_eval: None,
2249 },
2250 |st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
2251 let theta = clamp_theta(theta);
2252 let DecodedSpatialAdaptiveTheta {
2253 rho,
2254 adaptive_params,
2255 ..
2256 } = decode_theta(&theta)?;
2257 let family_eval =
2258 base_family.with_adaptive_params(adaptive_params, zero_quadratic.clone());
2259 let hyper_layout = realize_hyper_layout(&theta)?;
2260 let owned = evaluate_custom_family_joint_hyper_owned(
2261 &family_eval,
2262 std::slice::from_ref(&blockspec),
2263 &outer_opts,
2264 &rho,
2265 &hyper_layout,
2266 st.warm_cache.as_ref(),
2267 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueOnly,
2268 )
2269 .map_err(|e| {
2270 EstimationError::CustomFamily(e)
2271 .wrap_preserving_trial_point("spatial adaptive cost eval failed")
2272 })?;
2273 if !owned.result.inner_converged {
2274 st.warm_cache = Some(owned.result.warm_start);
2275 return Err(EstimationError::TrialPointRefused {
2276 reason: "exact spatial adaptive cost inner solve did not converge".to_string(),
2277 });
2278 }
2279 let objective = owned.result.objective;
2280 st.warm_cache = Some(owned.result.warm_start);
2281 st.terminal_mode = Some((theta, objective, owned.mode));
2282 Ok(objective)
2283 },
2284 |st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
2285 eval_outer(
2286 st,
2287 theta,
2288 if analytic_outer_hessian_available {
2289 gam_solve::rho_optimizer::OuterEvalOrder::ValueGradientHessian
2290 } else {
2291 gam_solve::rho_optimizer::OuterEvalOrder::ValueAndGradient
2292 },
2293 )
2294 },
2295 |st: &mut SpatialAdaptiveOuterState,
2296 theta: &Array1<f64>,
2297 order: gam_solve::rho_optimizer::OuterEvalOrder| { eval_outer(st, theta, order) },
2298 Some(|st: &mut SpatialAdaptiveOuterState| {
2299 st.warm_cache = None;
2300 st.terminal_mode = None;
2301 st.last_eval = None;
2302 }),
2303 Some(|st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
2304 let theta = clamp_theta(theta);
2305 let DecodedSpatialAdaptiveTheta {
2306 rho,
2307 adaptive_params,
2308 ..
2309 } = decode_theta(&theta)?;
2310 let family_eval =
2311 base_family.with_adaptive_params(adaptive_params, zero_quadratic.clone());
2312 let hyper_layout = realize_hyper_layout(&theta)?;
2313 let owned = evaluate_custom_family_joint_hyper_efs_owned(
2314 &family_eval,
2315 std::slice::from_ref(&blockspec),
2316 &outer_opts,
2317 &rho,
2318 &hyper_layout,
2319 st.warm_cache.as_ref(),
2320 )
2321 .map_err(|e| {
2322 EstimationError::CustomFamily(e)
2323 .wrap_preserving_trial_point("spatial adaptive EFS eval failed")
2324 })?;
2325 if !owned.result.inner_converged {
2326 st.warm_cache = Some(owned.result.warm_start);
2327 return Err(EstimationError::TrialPointRefused {
2328 reason: "exact spatial adaptive EFS inner solve did not converge".to_string(),
2329 });
2330 }
2331 let objective = owned.result.efs_eval.cost;
2332 st.warm_cache = Some(owned.result.warm_start);
2333 st.terminal_mode = Some((theta, objective, owned.mode));
2334 Ok(owned.result.efs_eval)
2335 }),
2336 |st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
2348 let theta = clamp_theta(theta);
2349 let DecodedSpatialAdaptiveTheta {
2350 rho,
2351 adaptive_params,
2352 ..
2353 } = decode_theta(&theta)?;
2354 let family_eval =
2355 base_family.with_adaptive_params(adaptive_params, zero_quadratic.clone());
2356 let hyper_layout = realize_hyper_layout(&theta)?;
2357 let owned = evaluate_custom_family_joint_hyper_owned(
2358 &family_eval,
2359 std::slice::from_ref(&blockspec),
2360 &outer_opts,
2361 &rho,
2362 &hyper_layout,
2363 st.warm_cache.as_ref(),
2364 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueOnly,
2365 )
2366 .map_err(|e| {
2367 EstimationError::CustomFamily(e)
2368 .wrap_preserving_trial_point("spatial adaptive screening eval failed")
2369 })?;
2370 st.warm_cache = Some(owned.result.warm_start);
2371 Ok(owned.result.objective)
2372 },
2373 );
2374
2375 let certified_outer = problem
2376 .run_certified(&mut obj, "exact spatial adaptive regularization")
2377 .map_err(|e| {
2378 EstimationError::InvalidInput(format!(
2379 "exact spatial adaptive outer optimization failed: {e}"
2380 ))
2381 })?;
2382 let outer_iterations = certified_outer.iterations();
2383 let outer_grad_norm = certified_outer.final_grad_norm();
2384 let theta_star = certified_outer.rho().clone();
2385 let (mode_theta, mode_objective, terminal_mode) =
2386 obj.state.terminal_mode.take().ok_or_else(|| {
2387 EstimationError::InvalidInput(
2388 "exact spatial adaptive optimization certified without retaining its terminal coefficient mode"
2389 .to_string(),
2390 )
2391 })?;
2392 if mode_theta.len() != theta_star.len()
2393 || mode_theta
2394 .iter()
2395 .zip(theta_star.iter())
2396 .any(|(mode, certified)| mode.to_bits() != certified.to_bits())
2397 {
2398 return Err(EstimationError::InvalidInput(
2399 "exact spatial adaptive terminal coefficient mode does not bitwise match the certified hyperparameter vector"
2400 .to_string(),
2401 ));
2402 }
2403 if mode_objective.to_bits() != certified_outer.final_value().to_bits() {
2404 return Err(EstimationError::InvalidInput(format!(
2405 "exact spatial adaptive terminal coefficient mode objective does not bitwise match the certified objective: mode={mode_objective:.17e}, certified={:.17e}",
2406 certified_outer.final_value(),
2407 )));
2408 }
2409 let DecodedSpatialAdaptiveTheta {
2410 rho: _,
2411 retained_lambdas,
2412 adaptive_params,
2413 epsilon: eps_star,
2414 } = decode_theta(&theta_star)?;
2415 let mut fixed_total = Array2::<f64>::zeros((
2416 baseline.design.design.ncols(),
2417 baseline.design.design.ncols(),
2418 ));
2419 for (idx, penalty) in retained_penalties.iter().enumerate() {
2420 fixed_total.scaled_add(retained_lambdas[idx], penalty);
2421 }
2422 let certified_final_family = base_family.with_adaptive_params(
2428 adaptive_params.clone(),
2429 zero_quadratic.clone(),
2430 );
2431 let fixed_total = ValidatedFixedQuadraticHessian::try_from_dense(
2432 fixed_total,
2433 baseline.design.design.ncols(),
2434 )
2435 .map_err(|error| {
2436 EstimationError::InvalidInput(format!(
2437 "optimized spatial adaptive fixed quadratic Hessian is invalid: {error}"
2438 ))
2439 })?;
2440 let final_family =
2441 base_family.with_adaptive_params(adaptive_params.clone(), fixed_total.clone());
2442 let final_blockspec = ParameterBlockSpec {
2443 name: "eta".to_string(),
2444 design: baseline.design.design.clone(),
2445 offset: offset.to_owned(),
2446 penalties: retained_penalties
2447 .iter()
2448 .cloned()
2449 .map(PenaltyMatrix::Dense)
2450 .collect(),
2451 nullspace_dims: retained_nullspace_dims.clone(),
2452 initial_log_lambdas: theta_star.slice(s![..rho_dim]).to_owned(),
2453 initial_beta: Some(baseline.fit.beta.clone()),
2454 gauge_priority: 100,
2455 jacobian_callback: None,
2456 stacked_design: None,
2457 stacked_offset: None,
2458 };
2459 let final_fit = fit_custom_family_fixed_log_lambdas_from_owned_mode(
2460 &certified_final_family,
2461 &[final_blockspec],
2462 &BlockwiseFitOptions {
2463 inner_max_cycles: options.max_iter,
2464 inner_tol: options.tol,
2465 outer_max_iter: 1,
2466 outer_tol: options.tol,
2467 compute_covariance: true,
2468 ..BlockwiseFitOptions::default()
2469 },
2470 terminal_mode,
2471 &theta_star,
2472 &certified_outer,
2473 )
2474 .map_err(EstimationError::CustomFamily)?;
2475 let beta = final_fit.block_states[0].beta.clone();
2476 let final_eval = final_family
2477 .exact_evaluation(&beta)
2478 .map_err(EstimationError::InvalidInput)?;
2479 let penalized_hessian = final_eval
2480 .totalobjectivehessian(&final_family.design)
2481 .map_err(EstimationError::InvalidInput)?;
2482 let beta_covariance = final_fit.covariance_conditional.clone();
2483 let beta_standard_errors = beta_covariance
2484 .as_ref()
2485 .map(|cov| Array1::from_iter((0..cov.nrows()).map(|i| cov[[i, i]].max(0.0).sqrt())));
2486
2487 let mut full_lambdas = baseline.fit.lambdas.clone();
2488 for (idx, &global_idx) in retained_global_indices.iter().enumerate() {
2489 full_lambdas[global_idx] = retained_lambdas[idx];
2490 }
2491 for (cache_idx, cache) in runtime_caches.iter().enumerate() {
2492 full_lambdas[cache.mass_penalty_global_idx] = adaptive_params[cache_idx].lambda[0];
2493 full_lambdas[cache.tension_penalty_global_idx] = adaptive_params[cache_idx].lambda[1];
2494 full_lambdas[cache.stiffness_penalty_global_idx] = adaptive_params[cache_idx].lambda[2];
2495 }
2496
2497 let deviance = -2.0 * final_eval.obs.log_likelihood;
2498 let mut local_penalty_blocks =
2499 Vec::<PenaltySpec>::with_capacity(baseline.design.penalties.len());
2500 for (global_idx, bp) in baseline.design.penalties.iter().enumerate() {
2501 if adaptive_penalty_indices.contains(&global_idx) {
2502 let cache = runtime_caches
2503 .iter()
2504 .find(|cache| {
2505 cache.mass_penalty_global_idx == global_idx
2506 || cache.tension_penalty_global_idx == global_idx
2507 || cache.stiffness_penalty_global_idx == global_idx
2508 })
2509 .ok_or_else(|| {
2510 EstimationError::InvalidInput(format!(
2511 "missing runtime cache for adaptive penalty index {global_idx}"
2512 ))
2513 })?;
2514 let cache_idx = runtime_caches
2515 .iter()
2516 .position(|c| {
2517 c.mass_penalty_global_idx == global_idx
2518 || c.tension_penalty_global_idx == global_idx
2519 || c.stiffness_penalty_global_idx == global_idx
2520 })
2521 .ok_or_else(|| {
2522 EstimationError::InvalidInput(format!(
2523 "missing adaptive cache position for penalty index {global_idx}"
2524 ))
2525 })?;
2526 let state = &final_eval.adaptive_states[cache_idx];
2527 let local = if cache.mass_penalty_global_idx == global_idx {
2528 scalar_operatorhessian(&cache.d0, &state.magnitude.betahessian_diag())
2529 .mapv(|v| adaptive_params[cache_idx].lambda[0] * v)
2530 } else if cache.tension_penalty_global_idx == global_idx {
2531 grouped_operatorhessian(
2532 &cache.d1,
2533 cache.dimension,
2534 &state.gradient.betahessian_blocks(),
2535 )?
2536 .mapv(|v| adaptive_params[cache_idx].lambda[1] * v)
2537 } else {
2538 grouped_operatorhessian(
2539 &cache.d2,
2540 cache.dimension * cache.dimension,
2541 &state.curvature.betahessian_blocks(),
2542 )?
2543 .mapv(|v| adaptive_params[cache_idx].lambda[2] * v)
2544 };
2545 local_penalty_blocks.push(PenaltySpec::Dense(penalty_matrixwith_local_block(
2547 baseline.design.design.ncols(),
2548 cache.coeff_global_range.clone(),
2549 &local,
2550 )));
2551 } else {
2552 local_penalty_blocks.push(PenaltySpec::Dense(
2553 bp.to_global(p_total).mapv(|v| v * full_lambdas[global_idx]),
2554 ));
2555 }
2556 }
2557 let (edf_by_block, penalty_block_trace, edf_total) = if let Some(cov) = beta_covariance.as_ref()
2558 {
2559 exact_bounded_edf(
2560 &local_penalty_blocks,
2561 &Array1::from_elem(local_penalty_blocks.len(), 1.0),
2562 cov,
2563 )?
2564 } else {
2565 (
2566 vec![0.0; local_penalty_blocks.len()],
2567 vec![0.0; local_penalty_blocks.len()],
2568 0.0,
2569 )
2570 };
2571 let stable_penalty_term = 2.0 * final_eval.adaptive_penalty_value
2572 + beta.dot(&fixed_total.as_dense().dot(&beta));
2573 let standard_deviation = if family.is_gaussian_identity() {
2574 certified_profiled_gaussian_scale(
2575 deviance,
2576 y.len() as f64 - edf_total,
2577 "exact spatial-adaptive Gaussian",
2578 )?
2579 } else {
2580 1.0
2581 };
2582 let maps = compute_spatial_adaptiveweights_for_beta(
2583 &beta,
2584 runtime_caches,
2585 eps_star[0],
2586 eps_star[1],
2587 eps_star[2],
2588 adaptive_opts.weight_floor,
2589 adaptive_opts.weight_ceiling,
2590 beta_covariance.as_ref(),
2594 )?
2595 .into_iter()
2596 .zip(runtime_caches.iter())
2597 .map(|(w, cache)| AdaptiveSpatialMap {
2598 termname: cache.termname.clone(),
2599 feature_cols: cache.feature_cols.clone(),
2600 collocation_points: cache.collocation_points.clone(),
2601 inv_magweight: w.inv_magweight,
2602 invgradweight: w.invgradweight,
2603 inv_lapweight: w.inv_lapweight,
2604 })
2605 .collect::<Vec<_>>();
2606 let fitted_link = if family.is_latent_cloglog() {
2607 FittedLinkState::LatentCLogLog {
2608 state: latent_cloglog_state
2609 .expect("BinomialLatentCLogLog requires an explicit latent-cloglog state"),
2610 }
2611 } else if family.is_binomial_mixture() {
2612 mixture_link_state
2613 .clone()
2614 .map(|state| FittedLinkState::Mixture {
2615 state,
2616 covariance: None,
2617 })
2618 .unwrap_or(FittedLinkState::Standard(None))
2619 } else if family.is_binomial_sas() {
2620 sas_link_state
2621 .map(|state| FittedLinkState::Sas {
2622 state,
2623 covariance: None,
2624 })
2625 .unwrap_or(FittedLinkState::Standard(None))
2626 } else if family.is_binomial_beta_logistic() {
2627 sas_link_state
2628 .map(|state| FittedLinkState::BetaLogistic {
2629 state,
2630 covariance: None,
2631 })
2632 .unwrap_or(FittedLinkState::Standard(None))
2633 } else {
2634 FittedLinkState::Standard(None)
2635 };
2636 let max_abs_eta = final_eval
2637 .obs
2638 .eta
2639 .iter()
2640 .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
2641 let fitted = FittedTermCollection {
2642 fit: {
2643 let log_lambdas =
2644 checked_fit_log_lambdas(&full_lambdas, "final exact spatial adaptive fit")?;
2645 let working = gam_solve::estimate::WorkingGeometry {
2646 weights: final_eval.obs.fisherweight.clone(),
2647 response: exact_standard_working_response(&final_eval.obs)?,
2648 };
2649 let inf = FitInference {
2650 edf_by_block,
2651 penalty_block_trace,
2652 edf_total,
2653 smoothing_correction: None,
2654 smoothing_correction_method: None,
2655 smoothing_correction_first_order: None,
2656 smoothing_correction_method_first_order: None,
2657 penalized_hessian: penalized_hessian.clone().into(),
2660 reparam_qs: None,
2661 dispersion: gam_solve::estimate::Dispersion::UNIT,
2662 beta_covariance: beta_covariance
2663 .clone()
2664 .map(gam_problem::dispersion_cov::PhiScaledCovariance::from),
2665 beta_standard_errors,
2666 beta_covariance_corrected: None,
2667 beta_standard_errors_corrected: None,
2668 beta_covariance_frequentist: None,
2669 coefficient_influence: None,
2670 weighted_gram: None,
2671 bias_correction_beta: None,
2672 bias_correction_jacobian: None,
2673 };
2674 let geometry = Some(gam_solve::estimate::FitGeometry {
2675 coefficient_gauge: gam_problem::gauge::Gauge::identity(&[beta.len()]),
2676 penalized_hessian: penalized_hessian.into(),
2677 constrained_posterior: None,
2678 working: Some(working),
2679 });
2680 let covariance_conditional = beta_covariance;
2681 let convergence = final_fit.convergence_evidence();
2682 let pirls_status_val = convergence.inner_status();
2683 let certified_outer_present = convergence.outer_certificate().is_some();
2684 UnifiedFitResult::try_from_parts(UnifiedFitResultParts {
2685 blocks: vec![gam_solve::estimate::FittedBlock {
2686 beta: beta.clone(),
2687 role: gam_problem::BlockRole::Mean,
2688 edf: edf_total,
2689 lambdas: full_lambdas.clone(),
2690 }],
2691 training_sample_size: y.len(),
2692 log_lambdas,
2693 lambdas: full_lambdas,
2694 likelihood_scale: family.default_scale_metadata(),
2695 likelihood_family: Some(family),
2696 log_likelihood_normalization: gam_spec::LogLikelihoodNormalization::UserProvided,
2697 log_likelihood: final_eval.obs.log_likelihood,
2698 deviance,
2699 reml_score: final_fit.penalized_objective(),
2700 stable_penalty_term,
2701 penalized_objective: final_fit.penalized_objective(),
2702 used_device: false,
2703 outer_iterations,
2704 outer_converged: certified_outer_present,
2705 outer_gradient_norm: outer_grad_norm,
2706 standard_deviation,
2707 covariance_conditional,
2708 covariance_corrected: None,
2709 inference: Some(inf),
2710 fitted_link,
2711 geometry,
2712 block_states: Vec::new(),
2713 pirls_status: pirls_status_val,
2714 max_abs_eta,
2715 constraint_kkt: None,
2716 artifacts: gam_solve::estimate::FitArtifacts {
2717 pirls: None,
2718 criterion_certificate: final_fit.artifacts.criterion_certificate.clone(),
2719 ..Default::default()
2720 },
2721 inner_cycles: 0,
2722 })?
2723 },
2724 design: baseline.design,
2725 adaptive_diagnostics: Some(AdaptiveRegularizationDiagnostics {
2726 epsilon_0: eps_star[0],
2727 epsilon_g: eps_star[1],
2728 epsilon_c: eps_star[2],
2729 epsilon_outer_iterations: outer_iterations,
2730 mm_iterations: 0,
2731 converged: true,
2732 maps,
2733 }),
2734 };
2735 enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
2736 Ok(fitted)
2737}
2738
2739fn relax_smoothing_rho_prior(
2779 options: &FitOptions,
2780 design: &TermCollectionDesign,
2781) -> gam_spec::RhoPrior {
2782 use gam_terms::basis::BasisMetadata;
2783 let base = &options.rho_prior;
2784 if !base.is_unset() {
2812 return base.clone();
2813 }
2814 let has_link_aux = options.sas_link.is_some()
2834 || options.optimize_sas
2835 || options.mixture_link.is_some()
2836 || options.optimize_mixture;
2837 let has_moving_kappa = design.smooth.terms.iter().any(|t| {
2838 if let BasisMetadata::Duchon {
2851 length_scale,
2852 aniso_log_scales,
2853 ..
2854 } = &t.metadata
2855 {
2856 return length_scale.is_some() || aniso_log_scales.is_some();
2857 }
2858 matches!(
2859 t.metadata,
2860 BasisMetadata::Matern { .. }
2861 | BasisMetadata::Sphere { .. }
2862 | BasisMetadata::SphereHarmonics { .. }
2863 | BasisMetadata::ConstantCurvature { .. }
2864 | BasisMetadata::MeasureJet { .. }
2865 )
2866 });
2867 let length_safe = !has_link_aux && !has_moving_kappa;
2874 if !length_safe {
2875 return base.clone();
2876 }
2877 let coords = &design.penaltyinfo;
2878 if coords.is_empty() {
2879 return base.clone();
2880 }
2881 let n_obs = design.design.nrows();
2892 let p_total = design.design.ncols();
2893 let underdetermined = n_obs < 2 * p_total;
2923 let relaxable_terms: std::collections::HashSet<&str> = design
2934 .smooth
2935 .terms
2936 .iter()
2937 .filter(|t| {
2938 (matches!(
2939 t.metadata,
2940 BasisMetadata::BSpline1D { .. }
2941 | BasisMetadata::ThinPlate { .. }
2942 | BasisMetadata::TensorBSpline { .. }
2943 )
2944 || matches!(
2956 t.metadata,
2957 BasisMetadata::Duchon {
2958 length_scale: None,
2959 aniso_log_scales: None,
2960 ..
2961 }
2962 ))
2963 && matches!(t.shape, gam_terms::smooth::ShapeConstraint::None)
2977 })
2978 .map(|t| t.name.as_str())
2979 .collect();
2980 let any_relaxed = coords.iter().any(|info| {
2981 info.termname
2982 .as_deref()
2983 .is_some_and(|name| relaxable_terms.contains(name))
2984 });
2985 if !any_relaxed {
2986 return base.clone();
2987 }
2988 let relaxed_prior = if underdetermined {
2993 gam_spec::RhoPrior::Normal {
2994 mean: 0.0,
2995 sd: RELAX_UNDERDETERMINED_RHO_SD,
2996 }
2997 } else {
2998 gam_spec::RhoPrior::Flat
2999 };
3000 let nullspace_degeneracy_prior = gam_spec::RhoPrior::Normal {
3011 mean: 0.0,
3012 sd: NULLSPACE_DEGENERACY_RHO_SD,
3013 };
3014 let per_coord = coords
3015 .iter()
3016 .map(|info| {
3017 let relax = info
3018 .termname
3019 .as_deref()
3020 .is_some_and(|name| relaxable_terms.contains(name));
3021 if !relax {
3022 return base.clone();
3023 }
3024 let is_nullspace = matches!(info.penalty.source, PenaltySource::DoublePenaltyNullspace);
3025 if is_nullspace {
3026 nullspace_degeneracy_prior.clone()
3027 } else {
3028 relaxed_prior.clone()
3029 }
3030 })
3031 .collect::<Vec<_>>();
3032 gam_spec::RhoPrior::Independent(per_coord)
3033}
3034
3035
3036const RELAX_UNDERDETERMINED_RHO_SD: f64 = 15.0;
3049
3050
3051fn adaptive_fit_options_base(options: &FitOptions, design: &TermCollectionDesign) -> FitOptions {
3052 FitOptions {
3053 resource_policy: options.resource_policy.clone(),
3054 latent_cloglog: options.latent_cloglog,
3055 mixture_link: options.mixture_link.clone(),
3056 optimize_mixture: options.optimize_mixture,
3057 sas_link: options.sas_link,
3058 optimize_sas: options.optimize_sas,
3059 compute_inference: options.compute_inference,
3060 skip_rho_posterior_inference: options.skip_rho_posterior_inference,
3061 max_iter: options.max_iter,
3062 tol: options.tol,
3063 nullspace_dims: design.nullspace_dims.clone(),
3064 linear_constraints: design.linear_constraints.clone(),
3065 firth_bias_reduction: options.firth_bias_reduction,
3066 adaptive_regularization: None,
3067 rho_prior: options.rho_prior.clone(),
3070 kronecker_penalty_system: design.kronecker_penalty_system(),
3071 kronecker_factored: design
3072 .smooth
3073 .terms
3074 .iter()
3075 .find_map(|t| t.kronecker_factored.clone()),
3076 persistent_warm_start_store: options.persistent_warm_start_store.clone(),
3077 }
3078}
3079
3080fn superseded_fit_options(options: &FitOptions) -> FitOptions {
3081 let mut fit_options = options.clone();
3082 fit_options.skip_rho_posterior_inference = true;
3083 fit_options
3084}
3085
3086#[derive(Clone)]
3087struct BoundedLinearTermMeta {
3088 col_idx: usize,
3089 min: f64,
3090 max: f64,
3091 prior: BoundedCoefficientPriorSpec,
3092}
3093
3094struct BoundedEffectiveJacobian {
3118 design: Array2<f64>,
3119 bounded_terms: Vec<BoundedLinearTermMeta>,
3120}
3121
3122impl BlockEffectiveJacobian for BoundedEffectiveJacobian {
3123 fn effective_jacobian_rows(
3124 &self,
3125 state: &FamilyLinearizationState<'_>,
3126 rows: std::ops::Range<usize>,
3127 ) -> Result<Array2<f64>, String> {
3128 let p = self.design.ncols();
3129 let n = self.design.nrows();
3130 let rows = rows.start.min(n)..rows.end.min(n);
3131 if !state.beta.is_empty() {
3132 if state.beta.len() != p {
3133 return Err(format!(
3134 "BoundedEffectiveJacobian::effective_jacobian_at: beta length {} != design \
3135 ncols {p}",
3136 state.beta.len(),
3137 ));
3138 }
3139 if state.beta.iter().any(|v| !v.is_finite()) {
3140 return Err(
3141 "BoundedEffectiveJacobian::effective_jacobian_at: beta contains a non-finite value"
3142 .to_string(),
3143 );
3144 }
3145 }
3146 let mut jac = self
3147 .design
3148 .slice(ndarray::s![rows.start..rows.end, ..])
3149 .to_owned();
3150 for term in &self.bounded_terms {
3151 if term.col_idx >= p {
3152 return Err(format!(
3153 "BoundedEffectiveJacobian::effective_jacobian_at: bounded column {} is outside {p} columns",
3154 term.col_idx
3155 ));
3156 }
3157 let theta = if state.beta.is_empty() {
3158 0.0
3159 } else {
3160 state.beta[term.col_idx]
3161 };
3162 let (_, _, db_dtheta, _, _) = bounded_latent_derivatives(theta, term.min, term.max);
3163 if !(db_dtheta.is_finite() && db_dtheta > 0.0) {
3164 return Err(format!(
3165 "BoundedEffectiveJacobian::effective_jacobian_at: bounded column {} has unrepresentable derivative {db_dtheta} at theta={theta}",
3166 term.col_idx
3167 ));
3168 }
3169 jac.column_mut(term.col_idx).mapv_inplace(|v| v * db_dtheta);
3170 }
3171 Ok(jac)
3172 }
3173}
3174
3175#[derive(Clone)]
3176struct BoundedLinearFamily {
3177 likelihood: gam_spec::GlmLikelihoodSpec,
3178 latent_cloglog_state: Option<LatentCLogLogState>,
3179 mixture_link_state: Option<MixtureLinkState>,
3180 sas_link_state: Option<SasLinkState>,
3181 y: Array1<f64>,
3182 weights: Array1<f64>,
3183 design: Array2<f64>,
3184 designzeroed: Array2<f64>,
3185 offset: Array1<f64>,
3186 bounded_terms: Vec<BoundedLinearTermMeta>,
3187}
3188
3189#[derive(Clone, Debug)]
3190struct StandardFamilyObservationState {
3191 eta: Array1<f64>,
3192 score: Array1<f64>,
3193 fisherweight: Array1<f64>,
3194 neghessian_eta: Array1<f64>,
3195 neghessian_eta_derivative: Array1<f64>,
3196 log_likelihood: f64,
3197}
3198
3199fn bounded_latent_to_user(theta: f64, min: f64, max: f64) -> (f64, f64, f64) {
3200 let jet = logit_inverse_link_jet5(theta);
3201 let z = jet.mu;
3202 let width = max - min;
3203 let beta = min + width * z;
3204 let db_dtheta = width * jet.d1;
3205 (beta, z, db_dtheta)
3206}
3207
3208fn bounded_user_to_latent(beta: f64, min: f64, max: f64) -> f64 {
3217 (beta - min).ln() - (max - beta).ln()
3218}
3219
3220#[derive(Debug, Clone, Copy)]
3224pub struct BoundedSampleColumn {
3225 pub col_idx: usize,
3227 pub min: f64,
3229 pub max: f64,
3231}
3232
3233pub fn sample_bounded_latent_posterior_internal(
3271 beta_user: &Array1<f64>,
3272 user_hessian: &Array2<f64>,
3273 bounded_columns: &[BoundedSampleColumn],
3274 n_draws: usize,
3275 sqrt_cov_scale: f64,
3276 base_seed: u64,
3277) -> Result<Array2<f64>, EstimationError> {
3278 let p = beta_user.len();
3279 if user_hessian.nrows() != p || user_hessian.ncols() != p {
3280 crate::bail_invalid_estim!(
3281 "bounded posterior sampling dimension mismatch: mode has {p} entries, user Hessian is {}x{}",
3282 user_hessian.nrows(),
3283 user_hessian.ncols()
3284 );
3285 }
3286 if beta_user.iter().any(|value| !value.is_finite()) {
3287 crate::bail_invalid_estim!("bounded posterior sampling requires a finite mode");
3288 }
3289 if user_hessian.iter().any(|value| !value.is_finite()) {
3290 crate::bail_invalid_estim!("bounded posterior sampling requires a finite Hessian");
3291 }
3292 if !(sqrt_cov_scale.is_finite() && sqrt_cov_scale >= 0.0) {
3293 crate::bail_invalid_estim!(
3294 "bounded posterior sampling covariance scale must be finite and non-negative, got {sqrt_cov_scale}"
3295 );
3296 }
3297
3298 let mut theta_mode = beta_user.clone();
3300 let mut jac_diag = Array1::<f64>::ones(p);
3301 for bc in bounded_columns {
3302 if bc.col_idx >= p {
3303 crate::bail_invalid_estim!(
3304 "bounded posterior sampling: bounded column index {} out of range for {p} coefficients",
3305 bc.col_idx
3306 );
3307 }
3308 if !(bc.min.is_finite()
3309 && bc.max.is_finite()
3310 && (bc.max - bc.min).is_finite()
3311 && bc.min < beta_user[bc.col_idx]
3312 && beta_user[bc.col_idx] < bc.max)
3313 {
3314 crate::bail_invalid_estim!(
3315 "bounded posterior sampling column {} requires finite bounds with a finite width and a mode strictly inside ({}, {}); got {}",
3316 bc.col_idx,
3317 bc.min,
3318 bc.max,
3319 beta_user[bc.col_idx]
3320 );
3321 }
3322 let theta_i = bounded_user_to_latent(beta_user[bc.col_idx], bc.min, bc.max);
3323 let (_, _, db_dtheta) = bounded_latent_to_user(theta_i, bc.min, bc.max);
3324 if !(theta_i.is_finite() && db_dtheta.is_finite() && db_dtheta > 0.0) {
3325 crate::bail_invalid_estim!(
3326 "bounded posterior sampling column {} has unrepresentable latent geometry: theta={theta_i}, d_beta/d_theta={db_dtheta}",
3327 bc.col_idx
3328 );
3329 }
3330 theta_mode[bc.col_idx] = theta_i;
3331 jac_diag[bc.col_idx] = db_dtheta;
3332 }
3333
3334 let mut h_latent = user_hessian.clone();
3337 for i in 0..p {
3338 let ji = jac_diag[i];
3339 if ji != 1.0 {
3340 h_latent.row_mut(i).mapv_inplace(|v| v * ji);
3341 h_latent.column_mut(i).mapv_inplace(|v| v * ji);
3342 }
3343 }
3344
3345 use gam_linalg::faer_ndarray::FaerCholesky as _;
3348 use rand::SeedableRng as _;
3349 let chol = h_latent.cholesky(faer::Side::Lower).map_err(|err| {
3350 EstimationError::InvalidInput(format!(
3351 "bounded posterior sampling: Cholesky of the latent penalized Hessian failed: {err:?}"
3352 ))
3353 })?;
3354 let l = chol.lower_triangular();
3355
3356 let mut draws = Array2::<f64>::zeros((n_draws, p));
3357 let mut eps = Array1::<f64>::zeros(p);
3358 let mut delta = Array1::<f64>::zeros(p);
3359 let mut rng = rand::rngs::StdRng::seed_from_u64(base_seed);
3360 for k in 0..n_draws {
3361 for e in eps.iter_mut() {
3362 *e = standard_normal_draw(&mut rng);
3363 }
3364 solve_lower_transpose_into(&l, &eps, &mut delta)?;
3365 for i in 0..p {
3366 draws[(k, i)] = theta_mode[i] + sqrt_cov_scale * delta[i];
3369 }
3370 for bc in bounded_columns {
3375 let (beta_draw, _, _) = bounded_latent_to_user(draws[(k, bc.col_idx)], bc.min, bc.max);
3376 draws[(k, bc.col_idx)] = beta_draw;
3377 }
3378 }
3379
3380 Ok(draws)
3381}
3382
3383#[inline]
3386fn standard_normal_draw<R: rand::Rng + ?Sized>(rng: &mut R) -> f64 {
3387 use rand::RngExt as _;
3388 let u1 = loop {
3389 let candidate = rng.random::<f64>();
3390 if candidate > 0.0 {
3391 break candidate;
3392 }
3393 };
3394 let u2 = rng.random::<f64>();
3395 (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos()
3396}
3397
3398fn solve_lower_transpose_into(
3402 l: &Array2<f64>,
3403 b: &Array1<f64>,
3404 out: &mut Array1<f64>,
3405) -> Result<(), EstimationError> {
3406 let p = l.nrows();
3407 if l.ncols() != p || b.len() != p || out.len() != p {
3408 crate::bail_invalid_estim!(
3409 "bounded triangular solve dimension mismatch: L={}x{}, b={}, out={}",
3410 l.nrows(),
3411 l.ncols(),
3412 b.len(),
3413 out.len()
3414 );
3415 }
3416 for i in (0..p).rev() {
3417 let mut acc = b[i];
3418 for j in (i + 1)..p {
3419 acc -= l[(j, i)] * out[j];
3420 }
3421 let diag = l[(i, i)];
3422 if !(diag.is_finite() && diag > 0.0 && acc.is_finite()) {
3423 crate::bail_invalid_estim!(
3424 "bounded triangular solve has invalid row {i}: diagonal={diag}, residual={acc}"
3425 );
3426 }
3427 let value = acc / diag;
3428 if !value.is_finite() {
3429 crate::bail_invalid_estim!(
3430 "bounded triangular solve produced a non-finite value at row {i}: {acc}/{diag}"
3431 );
3432 }
3433 out[i] = value;
3434 }
3435 Ok(())
3436}
3437
3438fn bounded_latent_injective_limit() -> f64 {
3450 (2.0 / f64::EPSILON).ln()
3451}
3452
3453fn bounded_latent_derivatives(theta: f64, min: f64, max: f64) -> (f64, f64, f64, f64, f64) {
3454 let jet = logit_inverse_link_jet5(theta);
3455 let z = jet.mu;
3456 let width = max - min;
3457 let beta = min + width * z;
3458 let db_dtheta = width * jet.d1;
3459 let d2b_dtheta2 = width * jet.d2;
3460 let d3b_dtheta3 = width * jet.d3;
3461 (beta, z, db_dtheta, d2b_dtheta2, d3b_dtheta3)
3462}
3463
3464fn bounded_prior_terms(
3465 theta: f64,
3466 prior: &BoundedCoefficientPriorSpec,
3467) -> Result<(f64, f64, f64, f64), String> {
3468 if !theta.is_finite() {
3469 return Err(format!(
3470 "bounded coefficient prior requires a finite latent coordinate, got {theta}"
3471 ));
3472 }
3473 let (a, b) = match prior {
3474 BoundedCoefficientPriorSpec::None => return Ok((0.0, 0.0, 0.0, 0.0)),
3476 BoundedCoefficientPriorSpec::Uniform => (1.0, 1.0),
3479 BoundedCoefficientPriorSpec::Beta { a, b } => (*a, *b),
3480 };
3481 if !(a.is_finite() && a > 0.0 && b.is_finite() && b > 0.0) {
3482 return Err(format!(
3483 "bounded coefficient Beta prior requires finite positive shapes, got ({a}, {b})"
3484 ));
3485 }
3486 let jet = logit_inverse_link_jet5(theta);
3487 let z = jet.mu;
3488 let logp = -a * gam_linalg::utils::stable_softplus(-theta)
3493 - b * gam_linalg::utils::stable_softplus(theta);
3494 let grad = a - (a + b) * z;
3495 let neghess = (a + b) * jet.d1;
3496 let neghess_derivative = (a + b) * jet.d2;
3497 let terms = (logp, grad, neghess, neghess_derivative);
3498 if [terms.0, terms.1, terms.2, terms.3]
3499 .iter()
3500 .any(|value| !value.is_finite())
3501 {
3502 return Err(format!(
3503 "bounded coefficient prior geometry is not representable at theta={theta}: {terms:?}"
3504 ));
3505 }
3506 Ok(terms)
3507}
3508
3509#[derive(Clone, Copy)]
3510struct ExactStandardObservationRow {
3511 mu: f64,
3512 score: f64,
3513 fisherweight: f64,
3514 neghessian_eta: f64,
3515 neghessian_eta_derivative: f64,
3516 log_likelihood: f64,
3517}
3518
3519impl ExactStandardObservationRow {
3520 #[inline]
3521 fn zero_weight(mu: f64) -> Self {
3522 Self {
3523 mu,
3524 score: 0.0,
3525 fisherweight: 0.0,
3526 neghessian_eta: 0.0,
3527 neghessian_eta_derivative: 0.0,
3528 log_likelihood: 0.0,
3529 }
3530 }
3531}
3532
3533#[inline]
3534fn bounded_row_error(row: usize, quantity: &'static str, eta: f64, value: f64) -> EstimationError {
3535 EstimationError::PirlsRowGeometryUnrepresentable {
3536 row,
3537 quantity,
3538 eta,
3539 value,
3540 }
3541}
3542
3543#[inline]
3544fn certify_bounded_row(
3545 row: usize,
3546 eta: f64,
3547 state: ExactStandardObservationRow,
3548) -> Result<ExactStandardObservationRow, EstimationError> {
3549 for (quantity, value) in [
3550 ("bounded-family mean", state.mu),
3551 ("bounded-family score", state.score),
3552 ("bounded-family Fisher weight", state.fisherweight),
3553 ("bounded-family observed Hessian", state.neghessian_eta),
3554 (
3555 "bounded-family observed Hessian derivative",
3556 state.neghessian_eta_derivative,
3557 ),
3558 ("bounded-family log likelihood", state.log_likelihood),
3559 ] {
3560 if !value.is_finite() {
3561 return Err(bounded_row_error(row, quantity, eta, value));
3562 }
3563 }
3564 if state.fisherweight < 0.0 {
3565 return Err(bounded_row_error(
3566 row,
3567 "bounded-family Fisher weight",
3568 eta,
3569 state.fisherweight,
3570 ));
3571 }
3572 Ok(state)
3573}
3574
3575#[inline]
3576fn weighted_positive_from_log(weight: f64, log_value: f64) -> f64 {
3577 if weight == 0.0 {
3578 return 0.0;
3579 }
3580 (weight.ln() + log_value).exp()
3581}
3582
3583#[inline]
3584fn weighted_product3(a: f64, b: f64, c: f64) -> f64 {
3585 crate::gamlss::scaled_signed_product3(a, b, c)
3586}
3587
3588#[inline]
3589fn convex_combination(y: f64, left: f64, right: f64) -> f64 {
3590 if y == 0.0 {
3591 right
3592 } else if y == 1.0 {
3593 left
3594 } else {
3595 y.mul_add(left, (1.0 - y) * right)
3596 }
3597}
3598
3599#[derive(Clone, Copy)]
3605struct BernoulliNaturalJet {
3606 mu: f64,
3607 log_mu: [f64; 4],
3608 log_one_minus_mu: [f64; 4],
3609 log_fisher: f64,
3610}
3611
3612#[inline]
3613fn probit_natural_jet(eta: f64) -> BernoulliNaturalJet {
3614 let left = gam_math::probability::normal_logcdf_derivatives(eta);
3615 let right_at_neg_eta = gam_math::probability::normal_logcdf_derivatives(-eta);
3616 let log_pdf = if eta.abs() <= f64::MAX.sqrt() {
3617 -0.5 * eta * eta - 0.5 * (2.0 * std::f64::consts::PI).ln()
3618 } else {
3619 f64::NEG_INFINITY
3620 };
3621 BernoulliNaturalJet {
3622 mu: left[0].exp(),
3623 log_mu: [left[0], left[1], left[2], left[3]],
3624 log_one_minus_mu: [
3625 right_at_neg_eta[0],
3626 -right_at_neg_eta[1],
3627 right_at_neg_eta[2],
3628 -right_at_neg_eta[3],
3629 ],
3630 log_fisher: 2.0 * log_pdf - left[0] - right_at_neg_eta[0],
3631 }
3632}
3633
3634#[inline]
3635fn cloglog_natural_jet(eta: f64) -> BernoulliNaturalJet {
3636 let x = eta.exp();
3637 if x == f64::INFINITY {
3638 return BernoulliNaturalJet {
3639 mu: 1.0,
3640 log_mu: [0.0; 4],
3641 log_one_minus_mu: [f64::NEG_INFINITY; 4],
3642 log_fisher: f64::NEG_INFINITY,
3643 };
3644 }
3645 if x == 0.0 {
3646 return BernoulliNaturalJet {
3647 mu: 0.0,
3648 log_mu: [eta, 1.0, 0.0, 0.0],
3649 log_one_minus_mu: [0.0; 4],
3650 log_fisher: eta,
3651 };
3652 }
3653 let mu = -(-x).exp_m1();
3654 let log_mu = if x < 0.5 {
3655 eta + (mu / x).ln()
3656 } else {
3657 mu.ln()
3658 };
3659 let h = if x < 1.0 {
3660 x / x.exp_m1()
3661 } else {
3662 let exp_neg_x = (-x).exp();
3663 x * exp_neg_x / (1.0 - exp_neg_x)
3664 };
3665 let a = 1.0 - x - h;
3666 let d2_log_mu = h * a;
3667 let d3_log_mu = h * (a * a - x - h * a);
3668 BernoulliNaturalJet {
3669 mu,
3670 log_mu: [log_mu, h, d2_log_mu, d3_log_mu],
3671 log_one_minus_mu: [-x, -x, -x, -x],
3672 log_fisher: 2.0 * eta - x - log_mu,
3673 }
3674}
3675
3676#[inline]
3677fn loglog_natural_jet(eta: f64) -> BernoulliNaturalJet {
3678 let mirrored = cloglog_natural_jet(-eta);
3679 BernoulliNaturalJet {
3680 mu: mirrored.log_one_minus_mu[0].exp(),
3681 log_mu: [
3682 mirrored.log_one_minus_mu[0],
3683 -mirrored.log_one_minus_mu[1],
3684 mirrored.log_one_minus_mu[2],
3685 -mirrored.log_one_minus_mu[3],
3686 ],
3687 log_one_minus_mu: [
3688 mirrored.log_mu[0],
3689 -mirrored.log_mu[1],
3690 mirrored.log_mu[2],
3691 -mirrored.log_mu[3],
3692 ],
3693 log_fisher: mirrored.log_fisher,
3694 }
3695}
3696
3697#[inline]
3698fn cauchit_natural_jet(eta: f64) -> BernoulliNaturalJet {
3699 let (mu, one_minus_mu) = if eta > 0.0 {
3700 let q = (eta.recip()).atan() / std::f64::consts::PI;
3701 (1.0 - q, q)
3702 } else if eta < 0.0 {
3703 let p = (-eta.recip()).atan() / std::f64::consts::PI;
3704 (p, 1.0 - p)
3705 } else {
3706 (0.5, 0.5)
3707 };
3708 let abs_eta = eta.abs();
3709 let log_one_plus_eta_sq = if abs_eta <= f64::MAX.sqrt() {
3710 (eta * eta).ln_1p()
3711 } else {
3712 2.0 * abs_eta.ln() + eta.recip().powi(2).ln_1p()
3713 };
3714 let log_d1 = -std::f64::consts::PI.ln() - log_one_plus_eta_sq;
3715 let ratio = if abs_eta <= 1.0 {
3716 eta / (1.0 + eta * eta)
3717 } else {
3718 1.0 / (eta + eta.recip())
3719 };
3720 let d2_over_d1 = -2.0 * ratio;
3721 let inv_one_plus_sq = if abs_eta <= 1.0 {
3722 1.0 / (1.0 + eta * eta)
3723 } else {
3724 let inv = eta.recip();
3725 inv * inv / (1.0 + inv * inv)
3726 };
3727 let d3_over_d1 = inv_one_plus_sq * (6.0 * (eta * ratio) - 2.0 * inv_one_plus_sq);
3728 let d1_over_mu = (log_d1 - mu.ln()).exp();
3729 let d1_over_q = (log_d1 - one_minus_mu.ln()).exp();
3730 let left_d2_ratio = d2_over_d1 * d1_over_mu;
3731 let right_d2_ratio = d2_over_d1 * d1_over_q;
3732 BernoulliNaturalJet {
3733 mu,
3734 log_mu: [
3735 mu.ln(),
3736 d1_over_mu,
3737 left_d2_ratio - d1_over_mu * d1_over_mu,
3738 d3_over_d1 * d1_over_mu - 3.0 * d1_over_mu * left_d2_ratio + 2.0 * d1_over_mu.powi(3),
3739 ],
3740 log_one_minus_mu: [
3741 one_minus_mu.ln(),
3742 -d1_over_q,
3743 -right_d2_ratio - d1_over_q * d1_over_q,
3744 -d3_over_d1 * d1_over_q - 3.0 * d1_over_q * right_d2_ratio - 2.0 * d1_over_q.powi(3),
3745 ],
3746 log_fisher: 2.0 * log_d1 - mu.ln() - one_minus_mu.ln(),
3747 }
3748}
3749
3750#[inline]
3751fn generic_bernoulli_natural_jet(
3752 row: usize,
3753 eta: f64,
3754 link: &InverseLink,
3755) -> Result<BernoulliNaturalJet, EstimationError> {
3756 let jet = inverse_link_jet_for_inverse_link(link, eta)?;
3757 if !(jet.mu.is_finite()
3758 && jet.mu > 0.0
3759 && jet.mu < 1.0
3760 && jet.d1.is_finite()
3761 && jet.d1 > 0.0
3762 && jet.d2.is_finite()
3763 && jet.d3.is_finite())
3764 {
3765 return Err(bounded_row_error(
3766 row,
3767 "bounded-family inverse-link jet",
3768 eta,
3769 jet.mu,
3770 ));
3771 }
3772 let mu = jet.mu;
3773 let q = 1.0 - mu;
3774 let r1 = jet.d1 / mu;
3775 let r2 = jet.d2 / mu;
3776 let r3 = jet.d3 / mu;
3777 let s1 = jet.d1 / q;
3778 let s2 = jet.d2 / q;
3779 let s3 = jet.d3 / q;
3780 Ok(BernoulliNaturalJet {
3781 mu,
3782 log_mu: [
3783 mu.ln(),
3784 r1,
3785 r2 - r1 * r1,
3786 r3 - 3.0 * r1 * r2 + 2.0 * r1.powi(3),
3787 ],
3788 log_one_minus_mu: [
3789 (-mu).ln_1p(),
3790 -s1,
3791 -s2 - s1 * s1,
3792 -s3 - 3.0 * s1 * s2 - 2.0 * s1.powi(3),
3793 ],
3794 log_fisher: 2.0 * jet.d1.ln() - mu.ln() - q.ln(),
3795 })
3796}
3797
3798fn resolved_bounded_binomial_link(
3799 family: &LikelihoodSpec,
3800 latent_cloglog_state: Option<&LatentCLogLogState>,
3801 mixture_link_state: Option<&MixtureLinkState>,
3802 sas_link_state: Option<&SasLinkState>,
3803) -> InverseLink {
3804 match &family.link {
3805 InverseLink::LatentCLogLog(_) => latent_cloglog_state
3806 .copied()
3807 .map(InverseLink::LatentCLogLog)
3808 .unwrap_or_else(|| family.link.clone()),
3809 InverseLink::Mixture(_) => mixture_link_state
3810 .cloned()
3811 .map(InverseLink::Mixture)
3812 .unwrap_or_else(|| family.link.clone()),
3813 InverseLink::Sas(_) => sas_link_state
3814 .copied()
3815 .map(InverseLink::Sas)
3816 .unwrap_or_else(|| family.link.clone()),
3817 InverseLink::BetaLogistic(_) => sas_link_state
3818 .copied()
3819 .map(InverseLink::BetaLogistic)
3820 .unwrap_or_else(|| family.link.clone()),
3821 InverseLink::Standard(_) => family.link.clone(),
3822 }
3823}
3824
3825fn binomial_natural_jet(
3826 row: usize,
3827 eta: f64,
3828 link: &InverseLink,
3829) -> Result<BernoulliNaturalJet, EstimationError> {
3830 match link {
3831 InverseLink::Standard(StandardLink::Probit) => Ok(probit_natural_jet(eta)),
3832 InverseLink::Standard(StandardLink::CLogLog) => Ok(cloglog_natural_jet(eta)),
3833 InverseLink::Standard(StandardLink::LogLog) => Ok(loglog_natural_jet(eta)),
3834 InverseLink::Standard(StandardLink::Cauchit) => Ok(cauchit_natural_jet(eta)),
3835 _ => generic_bernoulli_natural_jet(row, eta, link),
3836 }
3837}
3838
3839fn exact_logit_observation_row(
3840 row: usize,
3841 y: f64,
3842 weight: f64,
3843 eta: f64,
3844) -> Result<ExactStandardObservationRow, EstimationError> {
3845 let tail = (-eta.abs()).exp();
3846 let (mu, one_minus_mu) = if eta >= 0.0 {
3847 let q = tail / (1.0 + tail);
3848 (1.0 - q, q)
3849 } else {
3850 let p = tail / (1.0 + tail);
3851 (p, 1.0 - p)
3852 };
3853 if weight == 0.0 {
3854 return Ok(ExactStandardObservationRow::zero_weight(mu));
3855 }
3856 let log_fisher =
3857 -gam_linalg::utils::stable_softplus(eta) - gam_linalg::utils::stable_softplus(-eta);
3858 let fisherweight = weighted_positive_from_log(weight, log_fisher);
3859 if !(fisherweight.is_finite() && fisherweight > 0.0) {
3860 return Err(bounded_row_error(
3861 row,
3862 "bounded logit Fisher weight",
3863 eta,
3864 fisherweight,
3865 ));
3866 }
3867 let residual = if eta >= 0.0 {
3868 if y == 1.0 {
3869 one_minus_mu
3870 } else {
3871 (y - 1.0) + one_minus_mu
3872 }
3873 } else {
3874 y - mu
3875 };
3876 let log_likelihood_unit = if eta >= 0.0 {
3877 -(1.0 - y) * eta - gam_linalg::utils::stable_softplus(-eta)
3878 } else {
3879 y * eta - gam_linalg::utils::stable_softplus(eta)
3880 };
3881 certify_bounded_row(
3882 row,
3883 eta,
3884 ExactStandardObservationRow {
3885 mu,
3886 score: weight * residual,
3887 fisherweight,
3888 neghessian_eta: fisherweight,
3889 neghessian_eta_derivative: fisherweight * (one_minus_mu - mu),
3890 log_likelihood: weight * log_likelihood_unit,
3891 },
3892 )
3893}
3894
3895fn exact_noncanonical_binomial_observation_row(
3896 row: usize,
3897 y: f64,
3898 weight: f64,
3899 eta: f64,
3900 link: &InverseLink,
3901) -> Result<ExactStandardObservationRow, EstimationError> {
3902 let jet = binomial_natural_jet(row, eta, link)?;
3903 if weight == 0.0 {
3904 return Ok(ExactStandardObservationRow::zero_weight(jet.mu));
3905 }
3906 let fisherweight = weighted_positive_from_log(weight, jet.log_fisher);
3907 if !(fisherweight.is_finite() && fisherweight > 0.0) {
3908 return Err(bounded_row_error(
3909 row,
3910 "bounded binomial Fisher weight",
3911 eta,
3912 fisherweight,
3913 ));
3914 }
3915 let log_likelihood = weight * convex_combination(y, jet.log_mu[0], jet.log_one_minus_mu[0]);
3916 let score = weight * convex_combination(y, jet.log_mu[1], jet.log_one_minus_mu[1]);
3917 let neghessian_eta = -weight * convex_combination(y, jet.log_mu[2], jet.log_one_minus_mu[2]);
3918 let neghessian_eta_derivative =
3919 -weight * convex_combination(y, jet.log_mu[3], jet.log_one_minus_mu[3]);
3920 certify_bounded_row(
3921 row,
3922 eta,
3923 ExactStandardObservationRow {
3924 mu: jet.mu,
3925 score,
3926 fisherweight,
3927 neghessian_eta,
3928 neghessian_eta_derivative,
3929 log_likelihood,
3930 },
3931 )
3932}
3933
3934#[inline]
3935fn eta_exprel(rate: f64, eta: f64) -> f64 {
3936 (rate * eta).exp_m1() / rate
3937}
3938
3939fn validate_bounded_observation_inputs(
3940 likelihood: &gam_spec::GlmLikelihoodSpec,
3941 y: &Array1<f64>,
3942 weights: &Array1<f64>,
3943 eta: &Array1<f64>,
3944) -> Result<gam_spec::ResolvedLikelihoodScale, EstimationError> {
3945 let family = &likelihood.spec;
3946 if weights.len() != y.len() || eta.len() != y.len() {
3947 crate::bail_invalid_estim!(
3948 "bounded family observation size mismatch: y={}, weights={}, eta={}",
3949 y.len(),
3950 weights.len(),
3951 eta.len()
3952 );
3953 }
3954 if !LikelihoodSpec::is_legal_cell(&family.response, &family.link) {
3955 crate::bail_invalid_estim!(
3956 "bounded family received illegal likelihood cell response={} link={}",
3957 family.response.name(),
3958 family.link.link_function().name()
3959 );
3960 }
3961 let resolved_scale = likelihood
3962 .resolved_scale()
3963 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
3964 if let ResponseFamily::Tweedie { p } = &family.response
3965 && !(p.is_finite() && *p > 1.0 && *p < 2.0)
3966 {
3967 crate::bail_invalid_estim!(
3968 "bounded Tweedie power must be finite and strictly inside (1, 2), got {p}"
3969 );
3970 }
3971 if let ResponseFamily::NegativeBinomial { theta, .. } = &family.response
3972 && !(theta.is_finite() && *theta > 0.0)
3973 {
3974 crate::bail_invalid_estim!(
3975 "bounded negative-binomial theta must be finite and positive, got {theta}"
3976 );
3977 }
3978 for (i, &wi) in weights.iter().enumerate() {
3987 if !(wi.is_finite() && wi >= 0.0) {
3988 return Err(EstimationError::InvalidInput(format!(
3989 "bounded-family row {i} has invalid prior weight {wi:?}; expected finite weight >= 0"
3990 )));
3991 }
3992 }
3993 for i in 0..y.len() {
3994 let wi = weights[i];
3995 if wi == 0.0 {
3996 continue;
3997 }
3998 if !eta[i].is_finite() {
3999 return Err(bounded_row_error(i, "linear predictor", eta[i], eta[i]));
4000 }
4001 if !y[i].is_finite() {
4002 return Err(bounded_row_error(
4003 i,
4004 "bounded-family response",
4005 eta[i],
4006 y[i],
4007 ));
4008 }
4009 let yi = y[i];
4010 let valid = match &family.response {
4011 ResponseFamily::Gaussian => yi.is_finite(),
4012 ResponseFamily::Binomial => yi.is_finite() && (0.0..=1.0).contains(&yi),
4013 ResponseFamily::Poisson | ResponseFamily::NegativeBinomial { .. } => {
4014 yi.is_finite() && yi >= 0.0 && (yi - yi.round()).abs() <= 1e-9
4015 }
4016 ResponseFamily::Tweedie { .. } => yi.is_finite() && yi >= 0.0,
4017 ResponseFamily::Gamma => yi.is_finite() && yi > 0.0,
4018 ResponseFamily::Beta { .. } | ResponseFamily::RoystonParmar => false,
4019 };
4020 if !valid {
4021 return Err(bounded_row_error(i, "bounded-family response", eta[i], yi));
4022 }
4023 }
4024 Ok(resolved_scale)
4025}
4026
4027fn exact_standard_observation_row(
4028 likelihood: &gam_spec::GlmLikelihoodSpec,
4029 resolved_scale: gam_spec::ResolvedLikelihoodScale,
4030 binomial_link: &InverseLink,
4031 row: usize,
4032 y: f64,
4033 weight: f64,
4034 eta: f64,
4035) -> Result<ExactStandardObservationRow, EstimationError> {
4036 if weight == 0.0 {
4037 return Ok(ExactStandardObservationRow::zero_weight(0.0));
4038 }
4039 let family = &likelihood.spec;
4040 match &family.response {
4041 ResponseFamily::Gaussian => {
4042 let scaled_weight = match resolved_scale {
4043 gam_spec::ResolvedLikelihoodScale::ProfiledGaussian => weight,
4044 gam_spec::ResolvedLikelihoodScale::FixedGaussian { phi } => {
4045 crate::gamlss::scaled_positive_product_quotient(weight, 1.0, 1.0, phi.value())
4046 }
4047 _ => {
4048 crate::bail_invalid_estim!(
4049 "bounded Gaussian received a non-Gaussian resolved scale"
4050 );
4051 }
4052 };
4053 if !(scaled_weight.is_finite() && scaled_weight > 0.0) {
4054 return Err(bounded_row_error(
4055 row,
4056 "bounded Gaussian dispersion-scaled weight",
4057 eta,
4058 scaled_weight,
4059 ));
4060 }
4061 let residual = y - eta;
4062 let loss = if residual == 0.0 {
4063 0.0
4064 } else {
4065 crate::gamlss::scaled_positive_product_quotient(
4066 scaled_weight,
4067 residual.abs(),
4068 residual.abs(),
4069 2.0,
4070 )
4071 };
4072 certify_bounded_row(
4073 row,
4074 eta,
4075 ExactStandardObservationRow {
4076 mu: eta,
4077 score: scaled_weight * residual,
4078 fisherweight: scaled_weight,
4079 neghessian_eta: scaled_weight,
4080 neghessian_eta_derivative: 0.0,
4081 log_likelihood: -loss,
4082 },
4083 )
4084 }
4085 ResponseFamily::Binomial
4086 if matches!(binomial_link, InverseLink::Standard(StandardLink::Logit)) =>
4087 {
4088 exact_logit_observation_row(row, y, weight, eta)
4089 }
4090 ResponseFamily::Binomial => {
4091 exact_noncanonical_binomial_observation_row(row, y, weight, eta, binomial_link)
4092 }
4093 ResponseFamily::Poisson => {
4094 let mu = inverse_link_jet_for_inverse_link(&family.link, eta)?.mu;
4095 let fisherweight = weight * mu;
4096 let score = weight * (y - mu);
4097 let raw_log_likelihood = y.mul_add(eta, -mu);
4098 let log_likelihood = if raw_log_likelihood.is_finite() {
4099 weight * raw_log_likelihood
4100 } else {
4101 weighted_product3(weight, y, eta) - weight * mu
4102 };
4103 if !(fisherweight.is_finite() && fisherweight > 0.0) {
4104 return Err(bounded_row_error(
4105 row,
4106 "bounded Poisson Fisher weight",
4107 eta,
4108 fisherweight,
4109 ));
4110 }
4111 certify_bounded_row(
4112 row,
4113 eta,
4114 ExactStandardObservationRow {
4115 mu,
4116 score,
4117 fisherweight,
4118 neghessian_eta: fisherweight,
4119 neghessian_eta_derivative: fisherweight,
4120 log_likelihood,
4121 },
4122 )
4123 }
4124 ResponseFamily::Gamma => {
4125 let mu = inverse_link_jet_for_inverse_link(&family.link, eta)?.mu;
4126 let shape = resolved_scale
4127 .gamma_shape()
4128 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
4129 let weighted_shape = weight * shape;
4130 if !(weighted_shape.is_finite() && weighted_shape > 0.0) {
4131 return Err(bounded_row_error(
4132 row,
4133 "bounded Gamma shape-scaled weight",
4134 eta,
4135 weighted_shape,
4136 ));
4137 }
4138 let weighted_ratio =
4139 crate::gamlss::scaled_positive_product_quotient(weight, y, shape, mu);
4140 if !(weighted_ratio.is_finite() && weighted_ratio > 0.0) {
4141 return Err(bounded_row_error(
4142 row,
4143 "bounded Gamma observed Hessian",
4144 eta,
4145 weighted_ratio,
4146 ));
4147 }
4148 certify_bounded_row(
4149 row,
4150 eta,
4151 ExactStandardObservationRow {
4152 mu,
4153 score: weighted_ratio - weighted_shape,
4154 fisherweight: weighted_shape,
4155 neghessian_eta: weighted_ratio,
4156 neghessian_eta_derivative: -weighted_ratio,
4157 log_likelihood: -weighted_ratio - weighted_shape * eta,
4158 },
4159 )
4160 }
4161 ResponseFamily::Tweedie { p } => {
4162 let p = *p;
4163 let mu = inverse_link_jet_for_inverse_link(&family.link, eta)?.mu;
4164 let phi = resolved_scale
4165 .tweedie_phi()
4166 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
4167 let weight = crate::gamlss::scaled_positive_product_quotient(weight, 1.0, 1.0, phi);
4168 if !(weight.is_finite() && weight > 0.0) {
4169 return Err(bounded_row_error(
4170 row,
4171 "bounded Tweedie dispersion-scaled weight",
4172 eta,
4173 weight,
4174 ));
4175 }
4176 let a = ((1.0 - p) * eta).exp();
4177 let b = ((2.0 - p) * eta).exp();
4178 let score_unit = y.mul_add(a, -b);
4179 let score = if score_unit.is_finite() {
4180 weight * score_unit
4181 } else {
4182 weighted_product3(weight, y, a) - weight * b
4183 };
4184 let fisherweight = weight * b;
4185 let observed_unit = (p - 1.0) * y * a + (2.0 - p) * b;
4186 let neghessian_eta = if observed_unit.is_finite() {
4187 weight * observed_unit
4188 } else {
4189 weighted_product3(weight * (p - 1.0), y, a) + weight * (2.0 - p) * b
4190 };
4191 let observed_derivative_unit = -(p - 1.0).powi(2) * y * a + (2.0 - p).powi(2) * b;
4192 let neghessian_eta_derivative = if observed_derivative_unit.is_finite() {
4193 weight * observed_derivative_unit
4194 } else {
4195 -weighted_product3(weight * (p - 1.0).powi(2), y, a)
4196 + weight * (2.0 - p).powi(2) * b
4197 };
4198 let q_left = eta_exprel(1.0 - p, eta);
4201 let q_right = eta_exprel(2.0 - p, eta);
4202 let q = y.mul_add(q_left, -q_right);
4203 let log_likelihood = if q.is_finite() {
4204 weight * q
4205 } else {
4206 weighted_product3(weight, y, q_left) - weight * q_right
4207 };
4208 if !(fisherweight.is_finite() && fisherweight > 0.0) {
4209 return Err(bounded_row_error(
4210 row,
4211 "bounded Tweedie Fisher weight",
4212 eta,
4213 fisherweight,
4214 ));
4215 }
4216 certify_bounded_row(
4217 row,
4218 eta,
4219 ExactStandardObservationRow {
4220 mu,
4221 score,
4222 fisherweight,
4223 neghessian_eta,
4224 neghessian_eta_derivative,
4225 log_likelihood,
4226 },
4227 )
4228 }
4229 ResponseFamily::NegativeBinomial { .. } => {
4230 let theta = resolved_scale
4231 .negative_binomial_theta()
4232 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
4233 let mu = inverse_link_jet_for_inverse_link(&family.link, eta)?.mu;
4234 let log_theta = theta.ln();
4235 let delta = eta - log_theta;
4236 let log_q = -gam_linalg::utils::stable_softplus(-delta);
4237 let log_r = -gam_linalg::utils::stable_softplus(delta);
4238 let q = log_q.exp();
4239 let r = log_r.exp();
4240 let y_r = if y == 0.0 {
4241 0.0
4242 } else {
4243 (y.ln() + log_r).exp()
4244 };
4245 let theta_q = (log_theta + log_q).exp();
4246 let score = weight * (y_r - theta_q);
4247 let fisherweight = weighted_positive_from_log(weight, log_theta + log_q);
4248 let log_qr = log_q + log_r;
4249 let observed_y = if y == 0.0 {
4250 0.0
4251 } else {
4252 weighted_positive_from_log(weight, y.ln() + log_qr)
4253 };
4254 let observed_theta = weighted_positive_from_log(weight, log_theta + log_qr);
4255 let neghessian_eta = observed_y + observed_theta;
4256 let neghessian_eta_derivative = neghessian_eta * (r - q);
4257 let softplus_tail = if delta >= 0.0 {
4258 gam_linalg::utils::stable_softplus(-delta)
4259 } else {
4260 gam_linalg::utils::stable_softplus(delta)
4261 };
4262 let log_likelihood = if delta >= 0.0 {
4263 -weighted_product3(weight, theta, delta)
4264 - weighted_product3(weight, y, softplus_tail)
4265 - weighted_product3(weight, theta, softplus_tail)
4266 } else {
4267 weighted_product3(weight, y, delta)
4268 - weighted_product3(weight, y, softplus_tail)
4269 - weighted_product3(weight, theta, softplus_tail)
4270 };
4271 if !(fisherweight.is_finite() && fisherweight > 0.0) {
4272 return Err(bounded_row_error(
4273 row,
4274 "bounded negative-binomial Fisher weight",
4275 eta,
4276 fisherweight,
4277 ));
4278 }
4279 certify_bounded_row(
4280 row,
4281 eta,
4282 ExactStandardObservationRow {
4283 mu,
4284 score,
4285 fisherweight,
4286 neghessian_eta,
4287 neghessian_eta_derivative,
4288 log_likelihood,
4289 },
4290 )
4291 }
4292 ResponseFamily::Beta { .. } => {
4293 crate::bail_invalid_estim!("bounded linear terms are not supported for BetaLogit fits");
4294 }
4295 ResponseFamily::RoystonParmar => {
4296 crate::bail_invalid_estim!(
4297 "bounded linear terms are not supported for survival model fits"
4298 );
4299 }
4300 }
4301}
4302
4303fn evaluate_resolved_standard_family_observations(
4304 likelihood: &gam_spec::GlmLikelihoodSpec,
4305 latent_cloglog_state: Option<&LatentCLogLogState>,
4306 mixture_link_state: Option<&MixtureLinkState>,
4307 sas_link_state: Option<&SasLinkState>,
4308 y: &Array1<f64>,
4309 weights: &Array1<f64>,
4310 eta: &Array1<f64>,
4311) -> Result<StandardFamilyObservationState, EstimationError> {
4312 let n = y.len();
4313 let resolved_scale = validate_bounded_observation_inputs(likelihood, y, weights, eta)?;
4314 let family = &likelihood.spec;
4315 let binomial_link = resolved_bounded_binomial_link(
4316 &family,
4317 latent_cloglog_state,
4318 mixture_link_state,
4319 sas_link_state,
4320 );
4321
4322 let mut score = Array1::<f64>::zeros(n);
4323 let mut fisherweight = Array1::<f64>::zeros(n);
4324 let mut neghessian_eta = Array1::<f64>::zeros(n);
4325 let mut neghessian_eta_derivative = Array1::<f64>::zeros(n);
4326 let mut log_likelihood = 0.0;
4327 let mut log_likelihood_compensation = 0.0;
4328
4329 for i in 0..n {
4330 let row = exact_standard_observation_row(
4331 likelihood,
4332 resolved_scale,
4333 &binomial_link,
4334 i,
4335 y[i],
4336 weights[i],
4337 eta[i],
4338 )?;
4339 score[i] = row.score;
4340 fisherweight[i] = row.fisherweight;
4341 neghessian_eta[i] = row.neghessian_eta;
4342 neghessian_eta_derivative[i] = row.neghessian_eta_derivative;
4343 let adjusted = row.log_likelihood - log_likelihood_compensation;
4344 let updated = log_likelihood + adjusted;
4345 log_likelihood_compensation = (updated - log_likelihood) - adjusted;
4346 log_likelihood = updated;
4347 if !log_likelihood.is_finite() {
4348 return Err(bounded_row_error(
4349 i,
4350 "bounded-family cumulative log likelihood",
4351 eta[i],
4352 log_likelihood,
4353 ));
4354 }
4355 }
4356
4357 Ok(StandardFamilyObservationState {
4358 eta: eta.clone(),
4359 score,
4360 fisherweight,
4361 neghessian_eta,
4362 neghessian_eta_derivative,
4363 log_likelihood,
4364 })
4365}
4366
4367fn evaluate_standard_familyobservations(
4372 family: LikelihoodSpec,
4373 latent_cloglog_state: Option<&LatentCLogLogState>,
4374 mixture_link_state: Option<&MixtureLinkState>,
4375 sas_link_state: Option<&SasLinkState>,
4376 y: &Array1<f64>,
4377 weights: &Array1<f64>,
4378 eta: &Array1<f64>,
4379) -> Result<StandardFamilyObservationState, EstimationError> {
4380 let likelihood = gam_spec::GlmLikelihoodSpec::canonical(family);
4381 evaluate_resolved_standard_family_observations(
4382 &likelihood,
4383 latent_cloglog_state,
4384 mixture_link_state,
4385 sas_link_state,
4386 y,
4387 weights,
4388 eta,
4389 )
4390}
4391
4392fn exact_standard_working_response(
4393 state: &StandardFamilyObservationState,
4394) -> Result<Array1<f64>, EstimationError> {
4395 let mut out = state.eta.clone();
4396 for i in 0..out.len() {
4397 let weight = state.fisherweight[i];
4398 let score = state.score[i];
4399 if weight == 0.0 {
4400 if score != 0.0 {
4401 return Err(bounded_row_error(
4402 i,
4403 "zero-Fisher row with nonzero score",
4404 state.eta[i],
4405 score,
4406 ));
4407 }
4408 continue;
4409 }
4410 let increment = score / weight;
4411 let value = out[i] + increment;
4412 if !increment.is_finite() || !value.is_finite() {
4413 return Err(bounded_row_error(
4414 i,
4415 "bounded-family working response",
4416 state.eta[i],
4417 value,
4418 ));
4419 }
4420 out[i] = value;
4421 }
4422 Ok(out)
4423}
4424
4425#[derive(Clone, Copy, Debug, PartialEq, Eq)]
4426enum SpatialAdaptiveHyperKind {
4427 LogLambdaMagnitude,
4428 LogLambdaGradient,
4429 LogLambdaCurvature,
4430 LogEpsilonMagnitude,
4431 LogEpsilonGradient,
4432 LogEpsilonCurvature,
4433}
4434
4435impl SpatialAdaptiveHyperKind {
4436 fn component_index(self) -> usize {
4437 match self {
4438 SpatialAdaptiveHyperKind::LogLambdaMagnitude
4439 | SpatialAdaptiveHyperKind::LogEpsilonMagnitude => 0,
4440 SpatialAdaptiveHyperKind::LogLambdaGradient
4441 | SpatialAdaptiveHyperKind::LogEpsilonGradient => 1,
4442 SpatialAdaptiveHyperKind::LogLambdaCurvature
4443 | SpatialAdaptiveHyperKind::LogEpsilonCurvature => 2,
4444 }
4445 }
4446
4447 fn is_log_lambda(self) -> bool {
4448 matches!(
4449 self,
4450 SpatialAdaptiveHyperKind::LogLambdaMagnitude
4451 | SpatialAdaptiveHyperKind::LogLambdaGradient
4452 | SpatialAdaptiveHyperKind::LogLambdaCurvature
4453 )
4454 }
4455
4456 fn is_log_epsilon(self) -> bool {
4457 matches!(
4458 self,
4459 SpatialAdaptiveHyperKind::LogEpsilonMagnitude
4460 | SpatialAdaptiveHyperKind::LogEpsilonGradient
4461 | SpatialAdaptiveHyperKind::LogEpsilonCurvature
4462 )
4463 }
4464}
4465
4466#[derive(Clone, Copy, Debug)]
4467struct SpatialAdaptiveHyperSpec {
4468 cache_index: usize,
4469 kind: SpatialAdaptiveHyperKind,
4470}
4471
4472#[derive(Clone, Copy, Debug, PartialEq, Eq)]
4473enum SpatialAdaptiveExplicitSecondOrderKind {
4474 StructuralZero,
4475 LocalAlphaAlpha,
4476 LocalAlphaEta,
4477 SharedEtaEta,
4478}
4479
4480#[derive(Clone, Copy, Debug, PartialEq, Eq)]
4485enum AdaptiveComponent {
4486 Magnitude,
4487 Gradient,
4488 Curvature,
4489}
4490
4491impl AdaptiveComponent {
4492 fn from_index(index: usize) -> Result<Self, String> {
4493 match index {
4494 0 => Ok(AdaptiveComponent::Magnitude),
4495 1 => Ok(AdaptiveComponent::Gradient),
4496 2 => Ok(AdaptiveComponent::Curvature),
4497 other => Err(SmoothError::invalid_index(format!(
4498 "invalid adaptive component index {}",
4499 other
4500 ))
4501 .into()),
4502 }
4503 }
4504}
4505
4506#[derive(Clone, Copy, Debug, PartialEq, Eq)]
4511enum HyperDerivativeKind {
4512 Rho,
4514 LogEpsilonFirst,
4516 LogEpsilonSecond,
4518}
4519
4520#[derive(Clone, Copy, Debug, PartialEq, Eq)]
4525enum HyperDriftKind {
4526 Rho,
4527 LogEpsilon,
4528}
4529
4530impl SpatialAdaptiveHyperSpec {
4531 fn component_index(self) -> usize {
4532 self.kind.component_index()
4533 }
4534
4535 fn explicit_second_order_kind(self, other: Self) -> SpatialAdaptiveExplicitSecondOrderKind {
4536 if self.component_index() != other.component_index() {
4537 return SpatialAdaptiveExplicitSecondOrderKind::StructuralZero;
4538 }
4539 match (
4540 self.kind.is_log_lambda(),
4541 other.kind.is_log_lambda(),
4542 self.kind.is_log_epsilon(),
4543 other.kind.is_log_epsilon(),
4544 ) {
4545 (true, true, false, false) if self.cache_index == other.cache_index => {
4546 SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaAlpha
4547 }
4548 (true, false, false, true) | (false, true, true, false) => {
4549 SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaEta
4550 }
4551 (false, false, true, true) => SpatialAdaptiveExplicitSecondOrderKind::SharedEtaEta,
4552 _ => SpatialAdaptiveExplicitSecondOrderKind::StructuralZero,
4553 }
4554 }
4555}
4556
4557#[derive(Clone, Debug)]
4558struct SpatialAdaptiveTermHyperParams {
4559 lambda: [f64; 3],
4560 epsilon: [f64; 3],
4561}
4562
4563#[derive(Clone, Debug)]
4572struct ValidatedFixedQuadraticHessian {
4573 dense: Arc<Array2<f64>>,
4574}
4575
4576impl ValidatedFixedQuadraticHessian {
4577 fn try_from_dense(dense: Array2<f64>, coefficient_dim: usize) -> Result<Self, String> {
4578 gam_linalg::utils::validate_finite_symmetric_matrix(
4579 &dense,
4580 "spatial adaptive fixed quadratic Hessian",
4581 )
4582 .map_err(|error| error.to_string())?;
4583 PenaltyMatrix::Dense(dense.clone())
4584 .validate(coefficient_dim)
4585 .map_err(|error| {
4586 format!(
4587 "spatial adaptive fixed quadratic Hessian failed quadratic-form validation: {error}"
4588 )
4589 })?;
4590 Ok(Self {
4591 dense: Arc::new(dense),
4592 })
4593 }
4594
4595 fn zero(coefficient_dim: usize) -> Result<Self, String> {
4596 Self::try_from_dense(
4597 Array2::<f64>::zeros((coefficient_dim, coefficient_dim)),
4598 coefficient_dim,
4599 )
4600 }
4601
4602 fn as_dense(&self) -> &Array2<f64> {
4603 self.dense.as_ref()
4604 }
4605
4606 fn quadratic_terms(&self, beta: &Array1<f64>) -> Result<(f64, Array1<f64>), String> {
4607 if beta.len() != self.dense.ncols() {
4608 return Err(format!(
4609 "spatial adaptive fixed quadratic beta length {} does not match validated Hessian dimension {}",
4610 beta.len(),
4611 self.dense.ncols()
4612 ));
4613 }
4614 let gradient = self.dense.dot(beta);
4615 let value = 0.5 * beta.dot(&gradient);
4616 Ok((value, gradient))
4617 }
4618}
4619
4620#[derive(Clone)]
4621struct SpatialAdaptiveExactEvaluation {
4622 obs: StandardFamilyObservationState,
4623 adaptive_states: Vec<SpatialPenaltyExactState>,
4624 adaptive_penalty_value: f64,
4625 adaptive_penaltygradient: Array1<f64>,
4626 adaptive_penaltyhessian: Array2<f64>,
4627 fixed_quadraticvalue: f64,
4628 fixed_quadraticgradient: Array1<f64>,
4629 fixed_quadratic_hessian: ValidatedFixedQuadraticHessian,
4630}
4631
4632#[derive(Clone)]
4633struct CachedSpatialAdaptiveExactEvaluation {
4634 beta: Array1<f64>,
4635 eval: Arc<SpatialAdaptiveExactEvaluation>,
4636}
4637
4638impl SpatialAdaptiveExactEvaluation {
4639 fn total_penalty_value(&self) -> f64 {
4640 self.adaptive_penalty_value + self.fixed_quadraticvalue
4641 }
4642
4643 fn total_penaltygradient(&self) -> Array1<f64> {
4644 &self.adaptive_penaltygradient + &self.fixed_quadraticgradient
4645 }
4646
4647 fn total_penaltyhessian(&self) -> Array2<f64> {
4648 &self.adaptive_penaltyhessian + self.fixed_quadratic_hessian.as_dense()
4649 }
4650
4651 fn totalobjectivehessian(&self, design: &Array2<f64>) -> Result<Array2<f64>, String> {
4652 let mut out = xt_diag_x_dense(design.view(), self.obs.neghessian_eta.view())?;
4653 out += &self.total_penaltyhessian();
4654 Ok(out)
4655 }
4656}
4657
4658#[derive(Clone)]
4659struct SpatialAdaptiveExactFamily {
4660 family: LikelihoodSpec,
4661 latent_cloglog_state: Option<LatentCLogLogState>,
4662 mixture_link_state: Option<MixtureLinkState>,
4663 sas_link_state: Option<SasLinkState>,
4664 y: Arc<Array1<f64>>,
4665 weights: Arc<Array1<f64>>,
4666 design: Arc<Array2<f64>>,
4667 offset: Arc<Array1<f64>>,
4668 linear_constraints: Option<LinearInequalityConstraints>,
4669 runtime_caches: Arc<Vec<SpatialOperatorRuntimeCache>>,
4670 adaptive_params: Vec<SpatialAdaptiveTermHyperParams>,
4671 fixed_quadratic_hessian: ValidatedFixedQuadraticHessian,
4672 hyperspecs: Arc<Vec<SpatialAdaptiveHyperSpec>>,
4673 exact_eval_cache: Arc<Mutex<Option<CachedSpatialAdaptiveExactEvaluation>>>,
4674}
4675
4676impl SpatialAdaptiveExactFamily {
4677 fn with_adaptive_params(
4678 &self,
4679 adaptive_params: Vec<SpatialAdaptiveTermHyperParams>,
4680 fixed_quadratic_hessian: ValidatedFixedQuadraticHessian,
4681 ) -> Self {
4682 Self {
4683 family: self.family.clone(),
4684 latent_cloglog_state: self.latent_cloglog_state,
4685 mixture_link_state: self.mixture_link_state.clone(),
4686 sas_link_state: self.sas_link_state,
4687 y: self.y.clone(),
4688 weights: self.weights.clone(),
4689 design: self.design.clone(),
4690 offset: self.offset.clone(),
4691 linear_constraints: self.linear_constraints.clone(),
4692 runtime_caches: self.runtime_caches.clone(),
4693 adaptive_params,
4694 fixed_quadratic_hessian,
4695 hyperspecs: self.hyperspecs.clone(),
4696 exact_eval_cache: Arc::new(Mutex::new(None)),
4697 }
4698 }
4699
4700 fn total_eta(&self, beta: &Array1<f64>) -> Array1<f64> {
4701 gam_linalg::faer_ndarray::fast_av(self.design.as_ref(), beta) + self.offset.as_ref()
4702 }
4703
4704 fn fixed_quadratic_terms(
4705 &self,
4706 beta: &Array1<f64>,
4707 ) -> Result<(f64, Array1<f64>), String> {
4708 self.fixed_quadratic_hessian.quadratic_terms(beta)
4709 }
4710
4711 fn adaptive_penalty_value_only(&self, beta: &Array1<f64>) -> Result<f64, String> {
4712 let mut penalty_value = 0.0;
4713 for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
4714 let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
4715 format!(
4716 "missing adaptive parameter block for cache {}",
4717 cache.termname
4718 )
4719 })?;
4720 let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
4721 let state =
4722 SpatialPenaltyExactState::from_beta_local(beta_local, cache, params.epsilon)
4723 .map_err(|e| e.to_string())?;
4724 penalty_value += params.lambda[0] * state.magnitude.penalty_value();
4725 penalty_value += params.lambda[1] * state.gradient.penalty_value();
4726 penalty_value += params.lambda[2] * state.curvature.penalty_value();
4727 }
4728 Ok(penalty_value)
4729 }
4730
4731 fn zero_hyper_parts(&self) -> (Array1<f64>, Array2<f64>) {
4732 let total_dim = self.design.ncols();
4733 (
4734 Array1::<f64>::zeros(total_dim),
4735 Array2::<f64>::zeros((total_dim, total_dim)),
4736 )
4737 }
4738
4739 fn embed_local_hyper_parts(
4740 &self,
4741 coeff_range: &Range<usize>,
4742 local_grad: &Array1<f64>,
4743 local_hess: &Array2<f64>,
4744 ) -> (Array1<f64>, Array2<f64>) {
4745 let (mut beta_mixed, mut betahessian) = self.zero_hyper_parts();
4746 beta_mixed
4747 .slice_mut(s![coeff_range.clone()])
4748 .assign(local_grad);
4749 betahessian
4750 .slice_mut(s![coeff_range.clone(), coeff_range.clone()])
4751 .assign(local_hess);
4752 (beta_mixed, betahessian)
4753 }
4754
4755 fn embed_local_hyper_hessian(
4756 &self,
4757 coeff_range: &Range<usize>,
4758 local_hess: &Array2<f64>,
4759 ) -> Array2<f64> {
4760 let total_dim = self.design.ncols();
4761 let mut out = Array2::<f64>::zeros((total_dim, total_dim));
4762 out.slice_mut(s![coeff_range.clone(), coeff_range.clone()])
4763 .assign(local_hess);
4764 out
4765 }
4766
4767 fn adaptive_block_eval(
4776 &self,
4777 eval: &SpatialAdaptiveExactEvaluation,
4778 cache_idx: usize,
4779 component: AdaptiveComponent,
4780 derivative: HyperDerivativeKind,
4781 ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
4782 let cache = self
4783 .runtime_caches
4784 .get(cache_idx)
4785 .ok_or_else(|| format!("adaptive cache index {} out of bounds", cache_idx))?;
4786 let params = self
4787 .adaptive_params
4788 .get(cache_idx)
4789 .ok_or_else(|| format!("adaptive hyperparameter block {} out of bounds", cache_idx))?;
4790 let state = eval
4791 .adaptive_states
4792 .get(cache_idx)
4793 .ok_or_else(|| format!("adaptive exact state index {} out of bounds", cache_idx))?;
4794
4795 let (objective_local, beta_mixed_local, betahessian_local) = match component {
4796 AdaptiveComponent::Magnitude => {
4797 let lambda = params.lambda[0];
4798 let mag = &state.magnitude;
4799 let (objective, gradient_coeff, hessian_diag) = match derivative {
4800 HyperDerivativeKind::Rho => (
4801 mag.penalty_value(),
4802 mag.betagradient_coeff(),
4803 mag.betahessian_diag(),
4804 ),
4805 HyperDerivativeKind::LogEpsilonFirst => (
4806 mag.log_epsilon_gradient_terms().sum(),
4807 mag.log_epsilon_betagradient_coeff(),
4808 mag.log_epsilon_betahessian_diag(),
4809 ),
4810 HyperDerivativeKind::LogEpsilonSecond => (
4811 mag.log_epsilon_hessian_terms().sum(),
4812 mag.log_epsilon_beta_mixed_second_coeff(),
4813 mag.log_epsilon_betahessian_second_diag(),
4814 ),
4815 };
4816 (
4817 lambda * objective,
4818 lambda * scalar_operatorgradient(&cache.d0, &gradient_coeff),
4819 lambda * scalar_operatorhessian(&cache.d0, &hessian_diag),
4820 )
4821 }
4822 AdaptiveComponent::Gradient => {
4823 let lambda = params.lambda[1];
4824 let grad = &state.gradient;
4825 let (objective, gradient_blocks, hessian_blocks) = match derivative {
4826 HyperDerivativeKind::Rho => (
4827 grad.penalty_value(),
4828 grad.betagradient_blocks(),
4829 grad.betahessian_blocks(),
4830 ),
4831 HyperDerivativeKind::LogEpsilonFirst => (
4832 grad.log_epsilon_gradient_terms().sum(),
4833 grad.log_epsilon_betagradient_blocks(),
4834 grad.log_epsilon_betahessian_blocks(),
4835 ),
4836 HyperDerivativeKind::LogEpsilonSecond => (
4837 grad.log_epsilon_hessian_terms().sum(),
4838 grad.log_epsilon_beta_mixed_second_blocks(),
4839 grad.log_epsilon_betahessian_second_blocks(),
4840 ),
4841 };
4842 (
4843 lambda * objective,
4844 lambda
4845 * grouped_operatorgradient(&cache.d1, cache.dimension, &gradient_blocks)
4846 .map_err(|e| e.to_string())?,
4847 lambda
4848 * grouped_operatorhessian(&cache.d1, cache.dimension, &hessian_blocks)
4849 .map_err(|e| e.to_string())?,
4850 )
4851 }
4852 AdaptiveComponent::Curvature => {
4853 let lambda = params.lambda[2];
4854 let group = cache.dimension * cache.dimension;
4855 let curv = &state.curvature;
4856 let (objective, gradient_blocks, hessian_blocks) = match derivative {
4857 HyperDerivativeKind::Rho => (
4858 curv.penalty_value(),
4859 curv.betagradient_blocks(),
4860 curv.betahessian_blocks(),
4861 ),
4862 HyperDerivativeKind::LogEpsilonFirst => (
4863 curv.log_epsilon_gradient_terms().sum(),
4864 curv.log_epsilon_betagradient_blocks(),
4865 curv.log_epsilon_betahessian_blocks(),
4866 ),
4867 HyperDerivativeKind::LogEpsilonSecond => (
4868 curv.log_epsilon_hessian_terms().sum(),
4869 curv.log_epsilon_beta_mixed_second_blocks(),
4870 curv.log_epsilon_betahessian_second_blocks(),
4871 ),
4872 };
4873 (
4874 lambda * objective,
4875 lambda
4876 * grouped_operatorgradient(&cache.d2, group, &gradient_blocks)
4877 .map_err(|e| e.to_string())?,
4878 lambda
4879 * grouped_operatorhessian(&cache.d2, group, &hessian_blocks)
4880 .map_err(|e| e.to_string())?,
4881 )
4882 }
4883 };
4884
4885 let (beta_mixed, betahessian) = self.embed_local_hyper_parts(
4886 &cache.coeff_global_range,
4887 &beta_mixed_local,
4888 &betahessian_local,
4889 );
4890 Ok((objective_local, beta_mixed, betahessian))
4891 }
4892
4893 fn adaptive_shared_log_epsilon_parts(
4894 &self,
4895 eval: &SpatialAdaptiveExactEvaluation,
4896 component: usize,
4897 ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
4898 self.adaptive_shared_block_eval(eval, component, HyperDerivativeKind::LogEpsilonFirst)
4904 }
4905
4906 fn adaptive_shared_log_epsilon_second_parts(
4907 &self,
4908 eval: &SpatialAdaptiveExactEvaluation,
4909 component: usize,
4910 ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
4911 self.adaptive_shared_block_eval(eval, component, HyperDerivativeKind::LogEpsilonSecond)
4917 }
4918
4919 fn adaptive_shared_block_eval(
4924 &self,
4925 eval: &SpatialAdaptiveExactEvaluation,
4926 component: usize,
4927 derivative: HyperDerivativeKind,
4928 ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
4929 let component = AdaptiveComponent::from_index(component)?;
4930 let (mut score, mut hessian) = self.zero_hyper_parts();
4931 let mut objective = 0.0;
4932 for cache_idx in 0..self.runtime_caches.len() {
4933 let (local_objective, local_score, local_hessian) =
4934 self.adaptive_block_eval(eval, cache_idx, component, derivative)?;
4935 objective += local_objective;
4936 score += &local_score;
4937 hessian += &local_hessian;
4938 }
4939 Ok((objective, score, hessian))
4940 }
4941
4942 fn adaptive_shared_log_epsilon_drift(
4943 &self,
4944 eval: &SpatialAdaptiveExactEvaluation,
4945 component: usize,
4946 direction: &Array1<f64>,
4947 ) -> Result<Array2<f64>, String> {
4948 let component = AdaptiveComponent::from_index(component)?;
4952 let total_dim = self.design.ncols();
4953 let mut total = Array2::<f64>::zeros((total_dim, total_dim));
4954 for cache_idx in 0..self.runtime_caches.len() {
4955 total += &self.adaptive_block_drift_eval(
4956 eval,
4957 cache_idx,
4958 component,
4959 HyperDriftKind::LogEpsilon,
4960 direction,
4961 )?;
4962 }
4963 Ok(total)
4964 }
4965
4966 fn adaptive_explicit_second_order_parts(
4967 &self,
4968 eval: &SpatialAdaptiveExactEvaluation,
4969 left: SpatialAdaptiveHyperSpec,
4970 right: SpatialAdaptiveHyperSpec,
4971 ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
4972 match left.explicit_second_order_kind(right) {
4981 SpatialAdaptiveExplicitSecondOrderKind::StructuralZero => {
4982 let (score, hessian) = self.zero_hyper_parts();
4983 Ok((0.0, score, hessian))
4984 }
4985 SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaAlpha => self.adaptive_block_eval(
4986 eval,
4987 left.cache_index,
4988 AdaptiveComponent::from_index(left.component_index())?,
4989 HyperDerivativeKind::Rho,
4990 ),
4991 SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaEta => {
4992 let local_alpha = if left.kind.is_log_lambda() {
4993 left
4994 } else {
4995 right
4996 };
4997 self.adaptive_block_eval(
4998 eval,
4999 local_alpha.cache_index,
5000 AdaptiveComponent::from_index(local_alpha.component_index())?,
5001 HyperDerivativeKind::LogEpsilonFirst,
5002 )
5003 }
5004 SpatialAdaptiveExplicitSecondOrderKind::SharedEtaEta => {
5005 self.adaptive_shared_log_epsilon_second_parts(eval, left.component_index())
5006 }
5007 }
5008 }
5009
5010 fn adaptive_block_drift_eval(
5018 &self,
5019 eval: &SpatialAdaptiveExactEvaluation,
5020 cache_idx: usize,
5021 component: AdaptiveComponent,
5022 drift: HyperDriftKind,
5023 direction: &Array1<f64>,
5024 ) -> Result<Array2<f64>, String> {
5025 let cache = self
5026 .runtime_caches
5027 .get(cache_idx)
5028 .ok_or_else(|| format!("adaptive cache index {} out of bounds", cache_idx))?;
5029 let params = self
5030 .adaptive_params
5031 .get(cache_idx)
5032 .ok_or_else(|| format!("adaptive hyperparameter block {} out of bounds", cache_idx))?;
5033 let state = eval
5034 .adaptive_states
5035 .get(cache_idx)
5036 .ok_or_else(|| format!("adaptive exact state index {} out of bounds", cache_idx))?;
5037 let direction_local = direction.slice(s![cache.coeff_global_range.clone()]);
5038
5039 let local_hessian = match component {
5040 AdaptiveComponent::Magnitude => {
5041 let d0_u = cache.d0.dot(&direction_local);
5042 let mag = &state.magnitude;
5043 let diag = match drift {
5044 HyperDriftKind::Rho => mag.directionalhessian_diag(&d0_u),
5045 HyperDriftKind::LogEpsilon => {
5046 mag.log_epsilon_betahessian_directional_diag(&d0_u)
5047 }
5048 };
5049 params.lambda[0] * scalar_operatorhessian(&cache.d0, &diag)
5050 }
5051 AdaptiveComponent::Gradient => {
5052 let d1_u = cache.d1.dot(&direction_local);
5053 let direction_blocks = collocationgradient_blocks(&d1_u, cache.dimension)
5054 .map_err(|e| e.to_string())?;
5055 let grad = &state.gradient;
5056 let blocks = match drift {
5057 HyperDriftKind::Rho => grad.directionalhessian_blocks(&direction_blocks),
5058 HyperDriftKind::LogEpsilon => {
5059 grad.log_epsilon_betahessian_directional_blocks(&direction_blocks)
5060 }
5061 };
5062 params.lambda[1]
5063 * grouped_operatorhessian(&cache.d1, cache.dimension, &blocks)
5064 .map_err(|e| e.to_string())?
5065 }
5066 AdaptiveComponent::Curvature => {
5067 let group = cache.dimension * cache.dimension;
5068 let d2_u = cache.d2.dot(&direction_local);
5069 let direction_blocks =
5070 collocationhessian_blocks(&d2_u, cache.dimension).map_err(|e| e.to_string())?;
5071 let curv = &state.curvature;
5072 let blocks = match drift {
5073 HyperDriftKind::Rho => curv.directionalhessian_blocks(&direction_blocks),
5074 HyperDriftKind::LogEpsilon => {
5075 curv.log_epsilon_betahessian_directional_blocks(&direction_blocks)
5076 }
5077 };
5078 params.lambda[2]
5079 * grouped_operatorhessian(&cache.d2, group, &blocks)
5080 .map_err(|e| e.to_string())?
5081 }
5082 };
5083
5084 Ok(self.embed_local_hyper_hessian(&cache.coeff_global_range, &local_hessian))
5085 }
5086
5087 fn adaptive_hyper_parts(
5088 &self,
5089 eval: &SpatialAdaptiveExactEvaluation,
5090 hyper: SpatialAdaptiveHyperSpec,
5091 ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
5092 match hyper.kind {
5093 SpatialAdaptiveHyperKind::LogLambdaMagnitude
5096 | SpatialAdaptiveHyperKind::LogLambdaGradient
5097 | SpatialAdaptiveHyperKind::LogLambdaCurvature => self.adaptive_block_eval(
5098 eval,
5099 hyper.cache_index,
5100 AdaptiveComponent::from_index(hyper.component_index())?,
5101 HyperDerivativeKind::Rho,
5102 ),
5103 SpatialAdaptiveHyperKind::LogEpsilonMagnitude
5105 | SpatialAdaptiveHyperKind::LogEpsilonGradient
5106 | SpatialAdaptiveHyperKind::LogEpsilonCurvature => {
5107 self.adaptive_shared_log_epsilon_parts(eval, hyper.component_index())
5108 }
5109 }
5110 }
5111
5112 fn exact_evaluation_uncached(
5113 &self,
5114 beta: &Array1<f64>,
5115 ) -> Result<SpatialAdaptiveExactEvaluation, String> {
5116 let eta = self.total_eta(beta);
5117 let obs = evaluate_standard_familyobservations(
5118 self.family.clone(),
5119 self.latent_cloglog_state.as_ref(),
5120 self.mixture_link_state.as_ref(),
5121 self.sas_link_state.as_ref(),
5122 &self.y,
5123 &self.weights,
5124 &eta,
5125 )
5126 .map_err(|e| e.to_string())?;
5127 let p = beta.len();
5128 let mut penalty_value = 0.0;
5129 let mut penaltygradient = Array1::<f64>::zeros(p);
5130 let mut penaltyhessian = Array2::<f64>::zeros((p, p));
5131 let mut adaptive_states = Vec::with_capacity(self.runtime_caches.len());
5132
5133 for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
5134 let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
5135 format!(
5136 "missing adaptive parameter block for cache {}",
5137 cache.termname
5138 )
5139 })?;
5140 let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
5141 let state =
5142 SpatialPenaltyExactState::from_beta_local(beta_local, cache, params.epsilon)
5143 .map_err(|e| e.to_string())?;
5144
5145 let g0 = scalar_operatorgradient(&cache.d0, &state.magnitude.betagradient_coeff());
5146 let gg = grouped_operatorgradient(
5147 &cache.d1,
5148 cache.dimension,
5149 &state.gradient.betagradient_blocks(),
5150 )
5151 .map_err(|e| e.to_string())?;
5152 let gc = grouped_operatorgradient(
5153 &cache.d2,
5154 cache.dimension * cache.dimension,
5155 &state.curvature.betagradient_blocks(),
5156 )
5157 .map_err(|e| e.to_string())?;
5158 let h0 = scalar_operatorhessian(&cache.d0, &state.magnitude.betahessian_diag());
5159 let hg = grouped_operatorhessian(
5160 &cache.d1,
5161 cache.dimension,
5162 &state.gradient.betahessian_blocks(),
5163 )
5164 .map_err(|e| e.to_string())?;
5165 let hc = grouped_operatorhessian(
5166 &cache.d2,
5167 cache.dimension * cache.dimension,
5168 &state.curvature.betahessian_blocks(),
5169 )
5170 .map_err(|e| e.to_string())?;
5171
5172 let lambda0 = params.lambda[0];
5173 let lambdag = params.lambda[1];
5174 let lambdac = params.lambda[2];
5175
5176 penalty_value += lambda0 * state.magnitude.penalty_value();
5177 penalty_value += lambdag * state.gradient.penalty_value();
5178 penalty_value += lambdac * state.curvature.penalty_value();
5179
5180 let range = cache.coeff_global_range.clone();
5181 {
5182 let mut grad_local = penaltygradient.slice_mut(s![range.clone()]);
5183 grad_local += &(g0.mapv(|v| lambda0 * v));
5184 grad_local += &(gg.mapv(|v| lambdag * v));
5185 grad_local += &(gc.mapv(|v| lambdac * v));
5186 }
5187 {
5188 let mut h_local = penaltyhessian.slice_mut(s![range.clone(), range]);
5189 h_local += &h0.mapv(|v| lambda0 * v);
5190 h_local += &hg.mapv(|v| lambdag * v);
5191 h_local += &hc.mapv(|v| lambdac * v);
5192 }
5193
5194 adaptive_states.push(state);
5195 }
5196
5197 let (fixed_quadraticvalue, fixed_quadraticgradient) =
5198 self.fixed_quadratic_terms(beta)?;
5199 Ok(SpatialAdaptiveExactEvaluation {
5200 obs,
5201 adaptive_states,
5202 adaptive_penalty_value: penalty_value,
5203 adaptive_penaltygradient: penaltygradient,
5204 adaptive_penaltyhessian: penaltyhessian,
5205 fixed_quadraticvalue,
5206 fixed_quadraticgradient,
5207 fixed_quadratic_hessian: self.fixed_quadratic_hessian.clone(),
5208 })
5209 }
5210
5211 fn exact_evaluation(
5212 &self,
5213 beta: &Array1<f64>,
5214 ) -> Result<Arc<SpatialAdaptiveExactEvaluation>, String> {
5215 {
5216 let cache = self
5217 .exact_eval_cache
5218 .lock()
5219 .map_err(|_| "spatial adaptive exact-evaluation cache lock poisoned".to_string())?;
5220 if let Some(cached) = cache.as_ref()
5221 && cached.beta.len() == beta.len()
5222 && cached
5223 .beta
5224 .iter()
5225 .zip(beta.iter())
5226 .all(|(&left, &right)| left == right)
5227 {
5228 return Ok(Arc::clone(&cached.eval));
5229 }
5230 }
5231
5232 let eval = Arc::new(self.exact_evaluation_uncached(beta)?);
5233 let mut cache = self
5234 .exact_eval_cache
5235 .lock()
5236 .map_err(|_| "spatial adaptive exact-evaluation cache lock poisoned".to_string())?;
5237 *cache = Some(CachedSpatialAdaptiveExactEvaluation {
5238 beta: beta.clone(),
5239 eval: Arc::clone(&eval),
5240 });
5241 Ok(eval)
5242 }
5243
5244 fn exacthessian_directional_derivative_from_evaluation(
5245 &self,
5246 beta: &Array1<f64>,
5247 eval: &SpatialAdaptiveExactEvaluation,
5248 direction: &Array1<f64>,
5249 ) -> Result<Array2<f64>, String> {
5250 assert_eq!(
5251 beta.len(),
5252 direction.len(),
5253 "beta/direction length mismatch",
5254 );
5255 let d_eta = gam_linalg::faer_ndarray::fast_av(self.design.as_ref(), direction);
5256 let mut total = xt_diag_x_dense(
5257 self.design.view(),
5258 (&eval.obs.neghessian_eta_derivative * &d_eta).view(),
5259 )?;
5260 for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
5261 let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
5262 format!(
5263 "missing adaptive parameter block for cache {}",
5264 cache.termname
5265 )
5266 })?;
5267 let state = eval
5268 .adaptive_states
5269 .get(cache_idx)
5270 .ok_or_else(|| format!("missing adaptive state for cache {}", cache.termname))?;
5271 let direction_local = direction.slice(s![cache.coeff_global_range.clone()]);
5272 let d0_u = cache.d0.dot(&direction_local);
5273 let d1_u = cache.d1.dot(&direction_local);
5274 let d2_u = cache.d2.dot(&direction_local);
5275 let h0 =
5276 scalar_operatorhessian(&cache.d0, &state.magnitude.directionalhessian_diag(&d0_u))
5277 .mapv(|v| params.lambda[0] * v);
5278 let hg = grouped_operatorhessian(
5279 &cache.d1,
5280 cache.dimension,
5281 &state.gradient.directionalhessian_blocks(
5282 &collocationgradient_blocks(&d1_u, cache.dimension)
5283 .map_err(|e| e.to_string())?,
5284 ),
5285 )
5286 .map_err(|e| e.to_string())?
5287 .mapv(|v| params.lambda[1] * v);
5288 let hc = grouped_operatorhessian(
5289 &cache.d2,
5290 cache.dimension * cache.dimension,
5291 &state.curvature.directionalhessian_blocks(
5292 &collocationhessian_blocks(&d2_u, cache.dimension)
5293 .map_err(|e| e.to_string())?,
5294 ),
5295 )
5296 .map_err(|e| e.to_string())?
5297 .mapv(|v| params.lambda[2] * v);
5298 let range = cache.coeff_global_range.clone();
5299 let mut local = total.slice_mut(s![range.clone(), range]);
5300 local += &h0;
5301 local += &hg;
5302 local += &hc;
5303 }
5304 Ok(total)
5305 }
5306
5307 fn exacthessian_second_directional_derivative_from_evaluation(
5328 &self,
5329 eval: &SpatialAdaptiveExactEvaluation,
5330 direction_u: &Array1<f64>,
5331 direction_v: &Array1<f64>,
5332 ) -> Result<Option<Array2<f64>>, String> {
5333 let p = self.design.ncols();
5334 if eval.obs.neghessian_eta_derivative.iter().any(|&w| w != 0.0) {
5336 return Ok(None);
5337 }
5338 let mut total = Array2::<f64>::zeros((p, p));
5339 for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
5340 let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
5341 format!(
5342 "missing adaptive parameter block for cache {}",
5343 cache.termname
5344 )
5345 })?;
5346 let state = eval
5347 .adaptive_states
5348 .get(cache_idx)
5349 .ok_or_else(|| format!("missing adaptive state for cache {}", cache.termname))?;
5350 let u_local = direction_u.slice(s![cache.coeff_global_range.clone()]);
5351 let v_local = direction_v.slice(s![cache.coeff_global_range.clone()]);
5352
5353 let q0_u = cache.d0.dot(&u_local);
5355 let q0_v = cache.d0.dot(&v_local);
5356 let h0 = scalar_operatorhessian(
5357 &cache.d0,
5358 &state.magnitude.second_directionalhessian_diag(&q0_u, &q0_v),
5359 )
5360 .mapv(|x| params.lambda[0] * x);
5361
5362 let a1 = collocationgradient_blocks(&cache.d1.dot(&u_local), cache.dimension)
5364 .map_err(|e| e.to_string())?;
5365 let b1 = collocationgradient_blocks(&cache.d1.dot(&v_local), cache.dimension)
5366 .map_err(|e| e.to_string())?;
5367 let hg = grouped_operatorhessian(
5368 &cache.d1,
5369 cache.dimension,
5370 &state.gradient.second_directionalhessian_blocks(&a1, &b1),
5371 )
5372 .map_err(|e| e.to_string())?
5373 .mapv(|x| params.lambda[1] * x);
5374
5375 let a2 = collocationhessian_blocks(&cache.d2.dot(&u_local), cache.dimension)
5377 .map_err(|e| e.to_string())?;
5378 let b2 = collocationhessian_blocks(&cache.d2.dot(&v_local), cache.dimension)
5379 .map_err(|e| e.to_string())?;
5380 let hc = grouped_operatorhessian(
5381 &cache.d2,
5382 cache.dimension * cache.dimension,
5383 &state.curvature.second_directionalhessian_blocks(&a2, &b2),
5384 )
5385 .map_err(|e| e.to_string())?
5386 .mapv(|x| params.lambda[2] * x);
5387
5388 let range = cache.coeff_global_range.clone();
5389 let mut local = total.slice_mut(s![range.clone(), range]);
5390 local += &h0;
5391 local += &hg;
5392 local += &hc;
5393 }
5394 Ok(Some(total))
5395 }
5396}
5397
5398impl CustomFamily for SpatialAdaptiveExactFamily {
5399 fn joint_jeffreys_term_required(&self) -> bool {
5403 true
5404 }
5405
5406 fn joint_jeffreys_information_with_specs(
5443 &self,
5444 block_states: &[ParameterBlockState],
5445 specs: &[ParameterBlockSpec],
5446 ) -> Result<Option<Array2<f64>>, String> {
5447 let spec = expect_single_blockspec(specs, "spatial adaptive exact family")?;
5448 let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
5449 if spec.design.ncols() != beta.len() {
5450 return Err(SmoothError::dimension_mismatch(format!(
5451 "spatial adaptive Jeffreys information: spec design has {} columns, beta has {}",
5452 spec.design.ncols(),
5453 beta.len()
5454 ))
5455 .into());
5456 }
5457 let eval = self.exact_evaluation(beta)?;
5458 Ok(Some(xt_diag_x_dense(
5459 self.design.view(),
5460 eval.obs.neghessian_eta.view(),
5461 )?))
5462 }
5463
5464 fn joint_jeffreys_information_directional_derivative_with_specs(
5465 &self,
5466 block_states: &[ParameterBlockState],
5467 specs: &[ParameterBlockSpec],
5468 d_beta_flat: &Array1<f64>,
5469 ) -> Result<Option<Array2<f64>>, String> {
5470 let spec = expect_single_blockspec(specs, "spatial adaptive exact family")?;
5476 let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
5477 if spec.design.ncols() != d_beta_flat.len() {
5478 return Err(SmoothError::dimension_mismatch(format!(
5479 "spatial adaptive Jeffreys directional derivative: spec design has {} columns, direction has {}",
5480 spec.design.ncols(),
5481 d_beta_flat.len()
5482 ))
5483 .into());
5484 }
5485 let eval = self.exact_evaluation(beta)?;
5486 let d_eta = gam_linalg::faer_ndarray::fast_av(self.design.as_ref(), d_beta_flat);
5487 Ok(Some(xt_diag_x_dense(
5488 self.design.view(),
5489 (&eval.obs.neghessian_eta_derivative * &d_eta).view(),
5490 )?))
5491 }
5492
5493 fn joint_jeffreys_information_second_directional_derivative_with_specs(
5494 &self,
5495 block_states: &[ParameterBlockState],
5496 specs: &[ParameterBlockSpec],
5497 d_beta_u_flat: &Array1<f64>,
5498 d_betav_flat: &Array1<f64>,
5499 ) -> Result<Option<Array2<f64>>, String> {
5500 let spec = expect_single_blockspec(specs, "spatial adaptive exact family")?;
5507 let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
5508 if spec.design.ncols() != beta.len()
5509 || d_beta_u_flat.len() != beta.len()
5510 || d_betav_flat.len() != beta.len()
5511 {
5512 return Err(SmoothError::dimension_mismatch(format!(
5513 "spatial adaptive Jeffreys second-direction length mismatch: spec cols={}, dirs=({}, {}), expected {}",
5514 spec.design.ncols(),
5515 d_beta_u_flat.len(),
5516 d_betav_flat.len(),
5517 beta.len()
5518 ))
5519 .into());
5520 }
5521 let eval = self.exact_evaluation(beta)?;
5522 if eval.obs.neghessian_eta_derivative.iter().any(|&w| w != 0.0) {
5523 return Ok(None);
5524 }
5525 Ok(Some(Array2::<f64>::zeros((beta.len(), beta.len()))))
5526 }
5527
5528 fn joint_jeffreys_information_matches_observed_hessian(&self) -> bool {
5529 false
5534 }
5535
5536 fn joint_jeffreys_information_depends_on_psi(&self) -> bool {
5537 false
5546 }
5547
5548 fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
5549 let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
5550 let eval = self.exact_evaluation(beta)?;
5551 let mut gradient = fast_atv(&self.design, &eval.obs.score);
5552 gradient -= &eval.total_penaltygradient();
5553 let mut hessian = xt_diag_x_dense(self.design.view(), eval.obs.neghessian_eta.view())?;
5554 hessian += &eval.total_penaltyhessian();
5555 Ok(FamilyEvaluation {
5556 log_likelihood: eval.obs.log_likelihood - eval.total_penalty_value(),
5557 blockworking_sets: vec![BlockWorkingSet::ExactNewton {
5558 gradient,
5559 hessian: SymmetricMatrix::Dense(hessian),
5560 }],
5561 })
5562 }
5563
5564 fn log_likelihood_only(&self, block_states: &[ParameterBlockState]) -> Result<f64, String> {
5565 let state = expect_single_block_state(block_states, "spatial adaptive exact family")?;
5566 let beta = &state.beta;
5567 let obs = evaluate_standard_familyobservations(
5568 self.family.clone(),
5569 self.latent_cloglog_state.as_ref(),
5570 self.mixture_link_state.as_ref(),
5571 self.sas_link_state.as_ref(),
5572 &self.y,
5573 &self.weights,
5574 &state.eta,
5575 )
5576 .map_err(|e| e.to_string())?;
5577 let adaptive_penalty = self.adaptive_penalty_value_only(beta)?;
5578 let (fixed_quadratic, _) = self.fixed_quadratic_terms(beta)?;
5579 Ok(obs.log_likelihood - adaptive_penalty - fixed_quadratic)
5580 }
5581
5582 fn exact_newton_outerobjective(&self) -> ExactNewtonOuterObjective {
5583 ExactNewtonOuterObjective::StrictPseudoLaplace
5584 }
5585
5586 fn exact_newton_joint_hessian(
5587 &self,
5588 block_states: &[ParameterBlockState],
5589 ) -> Result<Option<Array2<f64>>, String> {
5590 let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
5591 let eval = self.exact_evaluation(beta)?;
5592 Ok(Some(eval.totalobjectivehessian(&self.design)?))
5593 }
5594
5595 fn exact_newton_hessian_directional_derivative(
5596 &self,
5597 block_states: &[ParameterBlockState],
5598 block_idx: usize,
5599 d_beta: &Array1<f64>,
5600 ) -> Result<Option<Array2<f64>>, String> {
5601 expect_block_idx_zero(block_idx, "spatial adaptive exact family", "")?;
5602 self.exact_newton_joint_hessian_directional_derivative(block_states, d_beta)
5603 }
5604
5605 fn exact_newton_joint_hessian_directional_derivative(
5606 &self,
5607 block_states: &[ParameterBlockState],
5608 d_beta_flat: &Array1<f64>,
5609 ) -> Result<Option<Array2<f64>>, String> {
5610 let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
5611 if d_beta_flat.len() != beta.len() {
5612 return Err(SmoothError::dimension_mismatch(format!(
5613 "spatial adaptive exact family direction length mismatch: got {}, expected {}",
5614 d_beta_flat.len(),
5615 beta.len()
5616 ))
5617 .into());
5618 }
5619 let eval = self.exact_evaluation(beta)?;
5620 Ok(Some(
5621 self.exacthessian_directional_derivative_from_evaluation(beta, &eval, d_beta_flat)?,
5622 ))
5623 }
5624
5625 fn exact_newton_joint_hessiansecond_directional_derivative(
5626 &self,
5627 block_states: &[ParameterBlockState],
5628 d_beta_u_flat: &Array1<f64>,
5629 d_betav_flat: &Array1<f64>,
5630 ) -> Result<Option<Array2<f64>>, String> {
5631 let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
5632 if d_beta_u_flat.len() != beta.len() || d_betav_flat.len() != beta.len() {
5633 return Err(SmoothError::dimension_mismatch(format!(
5634 "spatial adaptive exact family second-direction length mismatch: got ({}, {}), expected {}",
5635 d_beta_u_flat.len(),
5636 d_betav_flat.len(),
5637 beta.len()
5638 ))
5639 .into());
5640 }
5641 let eval = self.exact_evaluation(beta)?;
5642 self.exacthessian_second_directional_derivative_from_evaluation(
5643 &eval,
5644 d_beta_u_flat,
5645 d_betav_flat,
5646 )
5647 }
5648
5649 fn block_linear_constraints(
5650 &self,
5651 block_states: &[ParameterBlockState],
5652 block_idx: usize,
5653 block_spec: &ParameterBlockSpec,
5654 ) -> Result<Option<ConstraintSet>, String> {
5655 assert!(!block_states.is_empty(), "block_states must be non-empty");
5656 assert!(
5657 !block_spec.name.is_empty(),
5658 "block spec name must be non-empty",
5659 );
5660 expect_block_idx_zero(block_idx, "spatial adaptive exact family", "")?;
5661 Ok(self.linear_constraints.clone().map(ConstraintSet::Dense))
5662 }
5663
5664 fn exact_newton_joint_psi_terms(
5665 &self,
5666 block_states: &[ParameterBlockState],
5667 specs: &[ParameterBlockSpec],
5668 hyper_layout: &crate::custom_family::CustomFamilyHyperLayout,
5669 psi_index: usize,
5670 ) -> Result<Option<ExactNewtonJointPsiTerms>, String> {
5671 if hyper_layout.family_axis_count() != 0 {
5672 return Err(
5673 "spatial adaptive exact family does not declare family-owned hyper axes"
5674 .to_string(),
5675 );
5676 }
5677 let derivative_blocks = hyper_layout.design_derivative_blocks();
5678 if block_states.len() != 1 || specs.len() != 1 || derivative_blocks.len() != 1 {
5679 return Err(SmoothError::dimension_mismatch(format!(
5680 "spatial adaptive exact family expects one block/state/spec/psi payload, got states={} specs={} deriv_blocks={}",
5681 block_states.len(),
5682 specs.len(),
5683 derivative_blocks.len()
5684 ))
5685 .into());
5686 }
5687 derivative_blocks[0]
5688 .get(psi_index)
5689 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
5690 let hyper = self
5691 .hyperspecs
5692 .get(psi_index)
5693 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
5694 let beta = &block_states[0].beta;
5695 let eval = self.exact_evaluation(beta)?;
5696 let (direct, beta_mixed, betahessian_explicit) =
5697 self.adaptive_hyper_parts(&eval, *hyper)?;
5698
5699 Ok(Some(ExactNewtonJointPsiTerms {
5720 objective_psi: direct,
5721 score_psi: beta_mixed,
5722 hessian_psi: betahessian_explicit,
5723 hessian_psi_operator: None,
5724 }))
5725 }
5726
5727 fn exact_newton_joint_psisecond_order_terms(
5728 &self,
5729 block_states: &[ParameterBlockState],
5730 specs: &[ParameterBlockSpec],
5731 hyper_layout: &crate::custom_family::CustomFamilyHyperLayout,
5732 psi_i: usize,
5733 psi_j: usize,
5734 ) -> Result<Option<gam_problem::ExactNewtonJointPsiSecondOrderTerms>, String> {
5735 if hyper_layout.family_axis_count() != 0 {
5736 return Err(
5737 "spatial adaptive exact family does not declare family-owned hyper axes"
5738 .to_string(),
5739 );
5740 }
5741 let derivative_blocks = hyper_layout.design_derivative_blocks();
5742 if block_states.len() != 1 || specs.len() != 1 || derivative_blocks.len() != 1 {
5743 return Err(SmoothError::dimension_mismatch(format!(
5744 "spatial adaptive exact family expects one block/state/spec/psi payload, got states={} specs={} deriv_blocks={}",
5745 block_states.len(),
5746 specs.len(),
5747 derivative_blocks.len()
5748 ))
5749 .into());
5750 }
5751 derivative_blocks[0]
5752 .get(psi_i)
5753 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_i))?;
5754 derivative_blocks[0]
5755 .get(psi_j)
5756 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_j))?;
5757 let hyper_i = self
5758 .hyperspecs
5759 .get(psi_i)
5760 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_i))?;
5761 let hyper_j = self
5762 .hyperspecs
5763 .get(psi_j)
5764 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_j))?;
5765 let beta = &block_states[0].beta;
5766 let eval = self.exact_evaluation(beta)?;
5767 let (objective_psi_psi, score_psi_psi, hessian_psi_psi) =
5768 self.adaptive_explicit_second_order_parts(&eval, *hyper_i, *hyper_j)?;
5769
5770 Ok(Some(gam_problem::ExactNewtonJointPsiSecondOrderTerms {
5771 objective_psi_psi,
5772 score_psi_psi,
5773 hessian_psi_psi,
5774 hessian_psi_psi_operator: None,
5775 }))
5776 }
5777
5778 fn exact_newton_joint_psihessian_directional_derivative(
5779 &self,
5780 block_states: &[ParameterBlockState],
5781 specs: &[ParameterBlockSpec],
5782 hyper_layout: &crate::custom_family::CustomFamilyHyperLayout,
5783 psi_index: usize,
5784 direction: &Array1<f64>,
5785 ) -> Result<Option<Array2<f64>>, String> {
5786 if hyper_layout.family_axis_count() != 0 {
5787 return Err(
5788 "spatial adaptive exact family does not declare family-owned hyper axes"
5789 .to_string(),
5790 );
5791 }
5792 let derivative_blocks = hyper_layout.design_derivative_blocks();
5793 if block_states.len() != 1 || specs.len() != 1 || derivative_blocks.len() != 1 {
5794 return Err(SmoothError::dimension_mismatch(format!(
5795 "spatial adaptive exact family expects one block/state/spec/psi payload, got states={} specs={} deriv_blocks={}",
5796 block_states.len(),
5797 specs.len(),
5798 derivative_blocks.len()
5799 ))
5800 .into());
5801 }
5802 let beta = &block_states[0].beta;
5803 if direction.len() != beta.len() {
5804 return Err(SmoothError::dimension_mismatch(format!(
5805 "spatial adaptive exact family direction length mismatch: got {}, expected {}",
5806 direction.len(),
5807 beta.len()
5808 ))
5809 .into());
5810 }
5811 derivative_blocks[0]
5812 .get(psi_index)
5813 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
5814 let hyper = self
5815 .hyperspecs
5816 .get(psi_index)
5817 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
5818 let eval = self.exact_evaluation(beta)?;
5819 let drift = match hyper.kind {
5820 SpatialAdaptiveHyperKind::LogLambdaMagnitude
5821 | SpatialAdaptiveHyperKind::LogLambdaGradient
5822 | SpatialAdaptiveHyperKind::LogLambdaCurvature => self.adaptive_block_drift_eval(
5823 &eval,
5824 hyper.cache_index,
5825 AdaptiveComponent::from_index(hyper.kind.component_index())?,
5826 HyperDriftKind::Rho,
5827 direction,
5828 )?,
5829 SpatialAdaptiveHyperKind::LogEpsilonMagnitude
5830 | SpatialAdaptiveHyperKind::LogEpsilonGradient
5831 | SpatialAdaptiveHyperKind::LogEpsilonCurvature => self
5832 .adaptive_shared_log_epsilon_drift(
5833 &eval,
5834 hyper.kind.component_index(),
5835 direction,
5836 )?,
5837 };
5838 Ok(Some(drift))
5839 }
5840}
5841
5842fn expect_single_block_state<'a>(
5843 block_states: &'a [ParameterBlockState],
5844 family_name: &str,
5845) -> Result<&'a ParameterBlockState, String> {
5846 crate::block_layout::block_count::validate_block_count::<SmoothError>(
5847 family_name,
5848 1,
5849 block_states.len(),
5850 )?;
5851 Ok(&block_states[0])
5852}
5853
5854fn expect_single_blockspec<'a>(
5855 specs: &'a [ParameterBlockSpec],
5856 family_name: &str,
5857) -> Result<&'a ParameterBlockSpec, String> {
5858 crate::block_layout::block_count::validate_block_count::<SmoothError>(
5859 family_name,
5860 1,
5861 specs.len(),
5862 )?;
5863 Ok(&specs[0])
5864}
5865
5866fn expect_block_idx_zero(block_idx: usize, family_name: &str, context: &str) -> Result<(), String> {
5867 if block_idx != 0 {
5868 return Err(SmoothError::invalid_index(format!(
5869 "{family_name} expects block_idx 0{context}, got {block_idx}"
5870 ))
5871 .into());
5872 }
5873 Ok::<(), _>(())
5874}
5875
5876impl BoundedLinearFamily {
5877 fn bounded_term_derivative_data(
5878 &self,
5879 latent_beta: &Array1<f64>,
5880 ) -> Result<
5881 (
5882 Array1<f64>,
5883 Array1<f64>,
5884 Array1<f64>,
5885 Array1<f64>,
5886 Array1<f64>,
5887 ),
5888 String,
5889 > {
5890 let p = latent_beta.len();
5891 if p != self.design.ncols() || latent_beta.iter().any(|value| !value.is_finite()) {
5892 return Err(format!(
5893 "bounded coefficient geometry requires {} finite latent coefficients, got {}",
5894 self.design.ncols(),
5895 p
5896 ));
5897 }
5898 let mut beta_user = latent_beta.clone();
5899 let mut jac_diag = Array1::<f64>::ones(p);
5900 let mut second_diag = Array1::<f64>::zeros(p);
5901 let mut third_diag = Array1::<f64>::zeros(p);
5902 let mut priorthird = Array1::<f64>::zeros(p);
5903 for term in &self.bounded_terms {
5904 let width = term.max - term.min;
5905 if term.col_idx >= p
5906 || !term.min.is_finite()
5907 || !term.max.is_finite()
5908 || !(width.is_finite() && width > 0.0)
5909 {
5910 return Err(format!(
5911 "bounded coefficient geometry has invalid column/bounds: col={}, p={p}, bounds=({}, {})",
5912 term.col_idx, term.min, term.max
5913 ));
5914 }
5915 let (beta, _, db_dtheta, d2b_dtheta2, d3b_dtheta3) =
5916 bounded_latent_derivatives(latent_beta[term.col_idx], term.min, term.max);
5917 if [beta, db_dtheta, d2b_dtheta2, d3b_dtheta3]
5918 .iter()
5919 .any(|value| !value.is_finite())
5920 {
5921 return Err(format!(
5922 "bounded coefficient transform is not representable at column {} and theta={}",
5923 term.col_idx, latent_beta[term.col_idx]
5924 ));
5925 }
5926 beta_user[term.col_idx] = beta;
5927 jac_diag[term.col_idx] = db_dtheta;
5928 second_diag[term.col_idx] = d2b_dtheta2;
5929 third_diag[term.col_idx] = d3b_dtheta3;
5930 let (_, _, _, prior_neghess_derivative) =
5931 bounded_prior_terms(latent_beta[term.col_idx], &term.prior)?;
5932 priorthird[term.col_idx] = prior_neghess_derivative;
5933 }
5934 Ok((beta_user, jac_diag, second_diag, third_diag, priorthird))
5935 }
5936
5937 fn user_beta_and_jacobian(
5938 &self,
5939 latent_beta: &Array1<f64>,
5940 ) -> Result<(Array1<f64>, Array1<f64>), String> {
5941 let (beta_user, jac_diag, _, _, _) = self.bounded_term_derivative_data(latent_beta)?;
5942 Ok((beta_user, jac_diag))
5943 }
5944
5945 fn nonlinear_offset_from_latent(
5946 &self,
5947 latent_beta: &Array1<f64>,
5948 ) -> Result<Array1<f64>, String> {
5949 self.bounded_term_derivative_data(latent_beta)?;
5950 let mut offset = self.offset.clone();
5951 for term in &self.bounded_terms {
5952 let (beta, _, _) =
5953 bounded_latent_to_user(latent_beta[term.col_idx], term.min, term.max);
5954 offset.scaled_add(beta, &self.design.column(term.col_idx));
5955 }
5956 if offset.iter().any(|value| !value.is_finite()) {
5957 return Err("bounded nonlinear offset is not representable".to_string());
5958 }
5959 Ok(offset)
5960 }
5961
5962 fn effective_design_for_latent(&self, jac_diag: &Array1<f64>) -> Array2<f64> {
5963 let mut x_eff = self.design.clone();
5964 for term in &self.bounded_terms {
5965 x_eff
5966 .column_mut(term.col_idx)
5967 .mapv_inplace(|v| v * jac_diag[term.col_idx]);
5968 }
5969 x_eff
5970 }
5971
5972 fn exacthessian_andgradient(
5973 &self,
5974 latent_beta: &Array1<f64>,
5975 ) -> Result<
5976 (
5977 StandardFamilyObservationState,
5978 Array2<f64>,
5979 Array1<f64>,
5980 f64,
5981 Array1<f64>,
5982 Array1<f64>,
5983 Array1<f64>,
5984 ),
5985 String,
5986 > {
5987 let (_, jac_diag, second_diag, third_diag, priorthird) =
5988 self.bounded_term_derivative_data(latent_beta)?;
5989 let x_eff = self.effective_design_for_latent(&jac_diag);
5990 let eta =
5991 self.designzeroed.dot(latent_beta) + self.nonlinear_offset_from_latent(latent_beta)?;
5992 let obs = evaluate_resolved_standard_family_observations(
5993 &self.likelihood,
5994 self.latent_cloglog_state.as_ref(),
5995 self.mixture_link_state.as_ref(),
5996 self.sas_link_state.as_ref(),
5997 &self.y,
5998 &self.weights,
5999 &eta,
6000 )
6001 .map_err(|e| e.to_string())?;
6002
6003 let mut priorgrad = Array1::<f64>::zeros(latent_beta.len());
6004 let mut prior_neghess = Array2::<f64>::zeros((latent_beta.len(), latent_beta.len()));
6005 let mut prior_loglik = 0.0;
6006 for term in &self.bounded_terms {
6007 let (logp, grad, neghess, _) =
6008 bounded_prior_terms(latent_beta[term.col_idx], &term.prior)?;
6009 prior_loglik += logp;
6010 priorgrad[term.col_idx] += grad;
6011 prior_neghess[[term.col_idx, term.col_idx]] += neghess;
6012 }
6013
6014 let mut hessian = xt_diag_x_dense(x_eff.view(), obs.neghessian_eta.view())?;
6015 let mut gradient = fast_atv(&x_eff, &obs.score);
6016 for term in &self.bounded_terms {
6017 let score_beta = self.design.column(term.col_idx).dot(&obs.score);
6018 hessian[[term.col_idx, term.col_idx]] -= score_beta * second_diag[term.col_idx];
6019 }
6020 hessian += &prior_neghess;
6021 gradient += &priorgrad;
6022
6023 Ok((
6024 obs,
6025 hessian,
6026 gradient,
6027 prior_loglik,
6028 second_diag,
6029 third_diag,
6030 priorthird,
6031 ))
6032 }
6033
6034 fn evaluation_from_latent(
6035 &self,
6036 latent_beta: &Array1<f64>,
6037 ) -> Result<
6038 (
6039 StandardFamilyObservationState,
6040 Array2<f64>,
6041 Array1<f64>,
6042 f64,
6043 ),
6044 String,
6045 > {
6046 let (obs, hessian, gradient, prior_loglik, _, _, _) =
6047 self.exacthessian_andgradient(latent_beta)?;
6048 Ok((obs, hessian, gradient, prior_loglik))
6049 }
6050}
6051
6052impl CustomFamily for BoundedLinearFamily {
6053 fn joint_jeffreys_term_required(&self) -> bool {
6057 true
6058 }
6059
6060 fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
6061 let latent_beta = &expect_single_block_state(block_states, "bounded linear family")?.beta;
6062 let (obs, hessian, gradient, prior_loglik) = self.evaluation_from_latent(latent_beta)?;
6063 Ok(FamilyEvaluation {
6064 log_likelihood: obs.log_likelihood + prior_loglik,
6065 blockworking_sets: vec![BlockWorkingSet::ExactNewton {
6066 gradient,
6067 hessian: SymmetricMatrix::Dense(hessian),
6068 }],
6069 })
6070 }
6071
6072 fn exact_newton_joint_hessian(
6073 &self,
6074 block_states: &[ParameterBlockState],
6075 ) -> Result<Option<Array2<f64>>, String> {
6076 let latent_beta = &expect_single_block_state(block_states, "bounded linear family")?.beta;
6077 let (_, hessian, _, _) = self.evaluation_from_latent(latent_beta)?;
6078 Ok(Some(hessian))
6079 }
6080
6081 fn exact_newton_hessian_directional_derivative(
6082 &self,
6083 block_states: &[ParameterBlockState],
6084 block_idx: usize,
6085 d_beta: &Array1<f64>,
6086 ) -> Result<Option<Array2<f64>>, String> {
6087 expect_block_idx_zero(block_idx, "bounded linear family", "")?;
6088 self.exact_newton_joint_hessian_directional_derivative(block_states, d_beta)
6089 }
6090
6091 fn exact_newton_joint_hessian_directional_derivative(
6092 &self,
6093 block_states: &[ParameterBlockState],
6094 d_beta_flat: &Array1<f64>,
6095 ) -> Result<Option<Array2<f64>>, String> {
6096 let latent_beta = &expect_single_block_state(block_states, "bounded linear family")?.beta;
6097 if d_beta_flat.len() != latent_beta.len() {
6098 return Err(SmoothError::dimension_mismatch(format!(
6099 "bounded linear family directional derivative length mismatch: got {}, expected {}",
6100 d_beta_flat.len(),
6101 latent_beta.len()
6102 ))
6103 .into());
6104 }
6105
6106 let (obs, _, _, _, second_diag, third_diag, priorthird) =
6107 self.exacthessian_andgradient(latent_beta)?;
6108
6109 let (_, jac_diag, _, _, _) = self.bounded_term_derivative_data(latent_beta)?;
6110 let x_eff = self.effective_design_for_latent(&jac_diag);
6111 let deta = x_eff.dot(d_beta_flat);
6112 let d_neghess_eta = &obs.neghessian_eta_derivative * &deta;
6113
6114 let mut dx_eff = Array2::<f64>::zeros(x_eff.raw_dim());
6115 for term in &self.bounded_terms {
6116 let scale = second_diag[term.col_idx] * d_beta_flat[term.col_idx];
6117 if scale != 0.0 {
6118 let mut col = dx_eff.column_mut(term.col_idx);
6119 col.assign(&self.design.column(term.col_idx));
6120 col.mapv_inplace(|v| v * scale);
6121 }
6122 }
6123
6124 let mut dhessian = xt_diag_x_dense(x_eff.view(), d_neghess_eta.view())?;
6125 let mut wxdx = Array2::<f64>::zeros((x_eff.ncols(), x_eff.ncols()));
6126 for i in 0..x_eff.nrows() {
6127 let wi = obs.neghessian_eta[i];
6128 if wi == 0.0 {
6129 continue;
6130 }
6131 for a in 0..x_eff.ncols() {
6132 let xa = x_eff[[i, a]];
6133 for b in 0..x_eff.ncols() {
6134 wxdx[[a, b]] += wi * (dx_eff[[i, a]] * x_eff[[i, b]] + xa * dx_eff[[i, b]]);
6135 }
6136 }
6137 }
6138 dhessian += &wxdx;
6139
6140 let d_score = -&obs.neghessian_eta * &deta;
6141 for term in &self.bounded_terms {
6142 let score_beta = self.design.column(term.col_idx).dot(&obs.score);
6143 let d_score_beta = self.design.column(term.col_idx).dot(&d_score);
6144 dhessian[[term.col_idx, term.col_idx]] -= d_score_beta * second_diag[term.col_idx]
6145 + score_beta * third_diag[term.col_idx] * d_beta_flat[term.col_idx];
6146 dhessian[[term.col_idx, term.col_idx]] +=
6147 priorthird[term.col_idx] * d_beta_flat[term.col_idx];
6148 }
6149
6150 Ok(Some(dhessian))
6151 }
6152
6153 fn block_geometry(
6154 &self,
6155 block_states: &[ParameterBlockState],
6156 spec: &ParameterBlockSpec,
6157 ) -> Result<(DesignMatrix, Array1<f64>), String> {
6158 if block_states.is_empty() {
6159 return Ok((
6160 DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
6161 self.designzeroed.clone(),
6162 )),
6163 self.offset.clone(),
6164 ));
6165 }
6166 let offset = self.nonlinear_offset_from_latent(
6167 &expect_single_block_state(block_states, "bounded linear family")?.beta,
6168 )?;
6169 let x = if spec.design.ncols() == self.designzeroed.ncols() {
6170 self.designzeroed.clone()
6171 } else {
6172 return Err(SmoothError::dimension_mismatch(
6173 "bounded linear family design column mismatch",
6174 )
6175 .into());
6176 };
6177 Ok((
6178 DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(x)),
6179 offset,
6180 ))
6181 }
6182
6183 fn block_geometry_is_dynamic(&self) -> bool {
6184 true
6185 }
6186
6187 fn block_coefficient_coordinate(
6199 &self,
6200 block_states: &[ParameterBlockState],
6201 block_index: usize,
6202 block_spec: &ParameterBlockSpec,
6203 ) -> gam_problem::CoefficientCoordinate {
6204 if block_index != 0 || block_spec.design.ncols() != self.designzeroed.ncols() {
6212 log::debug!(
6213 "bounded linear family: coefficient coordinate asked for block {block_index} \
6214 at spec width {} ({} block state(s) supplied) while this family carries one \
6215 block of width {}; the coordinate is structural either way",
6216 block_spec.design.ncols(),
6217 block_states.len(),
6218 self.designzeroed.ncols(),
6219 );
6220 }
6221 gam_problem::CoefficientCoordinate::Structural
6222 }
6223
6224 fn post_update_block_beta(
6247 &self,
6248 block_states: &[ParameterBlockState],
6249 block_index: usize,
6250 block_spec: &ParameterBlockSpec,
6251 beta: Array1<f64>,
6252 ) -> Result<Array1<f64>, String> {
6253 expect_block_idx_zero(
6254 block_index,
6255 "bounded linear family",
6256 " for post-update beta",
6257 )?;
6258 let current = expect_single_block_state(block_states, "bounded linear family")?;
6262 if beta.len() != block_spec.design.ncols() || beta.len() != current.beta.len() {
6263 return Err(SmoothError::dimension_mismatch(format!(
6264 "bounded linear family post-update beta width mismatch: got {}, expected {} \
6265 (spec) / {} (current state)",
6266 beta.len(),
6267 block_spec.design.ncols(),
6268 current.beta.len()
6269 ))
6270 .into());
6271 }
6272 let limit = bounded_latent_injective_limit();
6273 let mut clamped = beta;
6274 for term in &self.bounded_terms {
6275 let theta = clamped[term.col_idx];
6276 clamped[term.col_idx] = theta.clamp(-limit, limit);
6277 }
6278 Ok(clamped)
6279 }
6280
6281 fn block_geometry_directional_derivative(
6282 &self,
6283 block_states: &[ParameterBlockState],
6284 block_idx: usize,
6285 spec: &ParameterBlockSpec,
6286 d_beta: &Array1<f64>,
6287 ) -> Result<Option<BlockGeometryDirectionalDerivative>, String> {
6288 expect_block_idx_zero(
6289 block_idx,
6290 "bounded linear family",
6291 " for geometry derivative",
6292 )?;
6293 expect_single_block_state(block_states, "bounded linear family")?;
6294 if d_beta.len() != spec.design.ncols() {
6295 return Err(SmoothError::dimension_mismatch(format!(
6296 "bounded linear family geometry derivative direction mismatch: got {}, expected {}",
6297 d_beta.len(),
6298 spec.design.ncols()
6299 ))
6300 .into());
6301 }
6302 let (_, jac_diag, _, _, _) = self.bounded_term_derivative_data(&block_states[0].beta)?;
6303 let mut d_offset = Array1::<f64>::zeros(self.offset.len());
6304 let has_drift = self
6305 .bounded_terms
6306 .iter()
6307 .any(|term| jac_diag[term.col_idx] != 0.0 && d_beta[term.col_idx] != 0.0);
6308 if !has_drift {
6309 return Ok(Some(BlockGeometryDirectionalDerivative {
6310 d_design: None,
6311 d_offset,
6312 }));
6313 }
6314 for term in &self.bounded_terms {
6315 let col = term.col_idx;
6316 let drift = jac_diag[col] * d_beta[col];
6317 if drift != 0.0 {
6318 d_offset.scaled_add(drift, &self.design.column(col));
6319 }
6320 }
6321 Ok(Some(BlockGeometryDirectionalDerivative {
6322 d_design: None,
6323 d_offset,
6324 }))
6325 }
6326}
6327
6328#[inline]
6329fn dense_diag_gram_chunkrows(p: usize) -> usize {
6330 const MIN_ROWS: usize = 512;
6331 const MAX_ROWS: usize = 2048;
6332 const TARGET_BYTES: usize = 2 * 1024 * 1024;
6333 let bytes_per_row = p.max(1) * std::mem::size_of::<f64>();
6334 (TARGET_BYTES / bytes_per_row).clamp(MIN_ROWS, MAX_ROWS)
6335}
6336
6337fn xt_diag_x_dense(x: ArrayView2<'_, f64>, w: ArrayView1<'_, f64>) -> Result<Array2<f64>, String> {
6338 if x.nrows() != w.len() {
6339 return Err(SmoothError::dimension_mismatch("xt_diag_x_dense row mismatch").into());
6340 }
6341 let (n, p) = x.dim();
6342 if n == 0 || p == 0 {
6343 return Ok(Array2::<f64>::zeros((p, p)));
6344 }
6345
6346 const STREAMING_BYTES_THRESHOLD: usize = 8 * 1024 * 1024;
6347 let dense_work_bytes = n
6348 .checked_mul(p)
6349 .and_then(|cells| cells.checked_mul(std::mem::size_of::<f64>()))
6350 .unwrap_or(usize::MAX);
6351 if dense_work_bytes <= STREAMING_BYTES_THRESHOLD {
6352 let mut weighted = x.to_owned();
6353 ndarray::Zip::from(weighted.rows_mut())
6354 .and(w)
6355 .par_for_each(|mut row, wi| row *= *wi);
6356 return Ok(fast_atb(&x, &weighted));
6357 }
6358
6359 let chunkrows = dense_diag_gram_chunkrows(p).min(n);
6360 let mut weighted_chunk = Array2::<f64>::zeros((chunkrows, p));
6361 let mut out = Array2::<f64>::zeros((p, p));
6362 for row_start in (0..n).step_by(chunkrows) {
6363 let rows = (n - row_start).min(chunkrows);
6364 let x_chunk = x.slice(s![row_start..row_start + rows, ..]);
6365 {
6366 let mut chunk = weighted_chunk.slice_mut(s![0..rows, ..]);
6367 for local_row in 0..rows {
6368 let scale = w[row_start + local_row];
6369 if scale == 0.0 {
6370 chunk.row_mut(local_row).fill(0.0);
6371 continue;
6372 }
6373 for col in 0..p {
6374 chunk[[local_row, col]] = x_chunk[[local_row, col]] * scale;
6375 }
6376 }
6377 }
6378 out += &fast_atb(&x_chunk, &weighted_chunk.slice(s![0..rows, ..]));
6379 }
6380 Ok(out)
6381}
6382
6383fn trace_of_factored_product(
6412 penalty_root: &Array2<f64>,
6413 covariance: &Array2<f64>,
6414) -> Result<f64, String> {
6415 let m = covariance.nrows();
6416 if m != covariance.ncols() {
6417 return Err(
6418 SmoothError::dimension_mismatch("trace_of_factored_product needs a square covariance")
6419 .into(),
6420 );
6421 }
6422 if penalty_root.ncols() != m {
6423 return Err(SmoothError::dimension_mismatch(
6424 "trace_of_factored_product penalty root and covariance disagree on the block width",
6425 )
6426 .into());
6427 }
6428 if penalty_root
6429 .iter()
6430 .chain(covariance.iter())
6431 .any(|value| !value.is_finite())
6432 {
6433 return Err("trace_of_factored_product requires finite factors".to_string());
6434 }
6435 if m == 0 || penalty_root.nrows() == 0 {
6436 return Ok(0.0);
6437 }
6438 use gam_linalg::faer_ndarray::FaerEigh as _;
6444 let (eigenvalues, eigenvectors) = covariance.eigh(faer::Side::Lower).map_err(|err| {
6445 format!("trace_of_factored_product could not factor the covariance block: {err}")
6446 })?;
6447 let mut covariance_factor = eigenvectors;
6448 for (col, eigenvalue) in eigenvalues.iter().enumerate() {
6449 let scale = eigenvalue.max(0.0).sqrt();
6450 covariance_factor.column_mut(col).mapv_inplace(|v| v * scale);
6451 }
6452 let scaled = penalty_root.dot(&covariance_factor);
6453 let mut trace = gam_linalg::utils::KahanSum::default();
6454 for value in scaled.iter() {
6455 let term = value * value;
6456 if !term.is_finite() {
6457 return Err("trace_of_factored_product term is not representable".to_string());
6458 }
6459 trace.add(term);
6460 }
6461 let trace = trace.sum();
6462 if !trace.is_finite() {
6463 return Err("trace_of_factored_product sum is not representable".to_string());
6464 }
6465 Ok(trace)
6466}
6467
6468fn certify_bounded_edf_interval(
6469 value: f64,
6470 lower: f64,
6471 upper: f64,
6472 dimension: usize,
6473 label: &str,
6474) -> Result<f64, EstimationError> {
6475 if !(value.is_finite() && lower.is_finite() && upper.is_finite() && lower <= upper) {
6476 crate::bail_invalid_estim!(
6477 "{label} has invalid EDF interval/value: value={value}, interval=[{lower}, {upper}]"
6478 );
6479 }
6480 let scale = 1.0_f64.max(value.abs()).max(lower.abs()).max(upper.abs());
6481 let allowed = 256.0 * f64::EPSILON * (dimension.max(1) as f64).powi(2) * scale;
6485 if value < lower {
6486 if lower - value <= allowed {
6487 return Ok(lower);
6488 }
6489 } else if value > upper {
6490 if value - upper <= allowed {
6491 return Ok(upper);
6492 }
6493 } else {
6494 return Ok(value);
6495 }
6496 crate::bail_invalid_estim!(
6497 "{label}={value} lies outside [{lower}, {upper}] by more than the dense-trace backward-error allowance {allowed}"
6498 )
6499}
6500
6501fn exact_bounded_edf(
6502 penalties: &[PenaltySpec],
6503 lambdas: &Array1<f64>,
6504 latent_cov: &Array2<f64>,
6505) -> Result<(Vec<f64>, Vec<f64>, f64), EstimationError> {
6506 if penalties.len() != lambdas.len() {
6507 crate::bail_invalid_estim!(
6508 "bounded EDF penalty/lambda mismatch: {} penalties vs {} lambdas",
6509 penalties.len(),
6510 lambdas.len()
6511 );
6512 }
6513 if latent_cov.nrows() != latent_cov.ncols() {
6514 crate::bail_invalid_estim!("bounded EDF covariance must be square");
6515 }
6516
6517 let p = latent_cov.nrows();
6518 let mut s_lambda = Array2::<f64>::zeros((p, p));
6519 let mut edf_by_block = Vec::with_capacity(penalties.len());
6520 let mut penalty_block_trace = Vec::with_capacity(penalties.len());
6522 let mut trace_sum = gam_linalg::utils::KahanSum::default();
6523
6524 for (k, ps) in penalties.iter().enumerate() {
6525 let lambda_k = lambdas[k];
6526 if !(lambda_k.is_finite() && lambda_k >= 0.0) {
6527 crate::bail_invalid_estim!(
6528 "bounded EDF smoothing strength at block {k} must be finite and non-negative, got {lambda_k}"
6529 );
6530 }
6531 match ps {
6532 PenaltySpec::Block {
6533 local, col_range, ..
6534 } => {
6535 s_lambda
6536 .slice_mut(ndarray::s![col_range.clone(), col_range.clone()])
6537 .scaled_add(lambda_k, local);
6538 let penalty_rank =
6540 local
6541 .nrows()
6542 .saturating_sub(estimate_penalty_nullity(local).map_err(|e| {
6543 EstimationError::InvalidInput(format!("bounded EDF rank failed: {e}"))
6544 })?);
6545 let cov_block = latent_cov.slice(ndarray::s![col_range.clone(), col_range.clone()]);
6547 let penalty_root = gam_solve::estimate::reml::reml_outer_engine::penalty_matrix_root(local)
6554 .map_err(EstimationError::InvalidInput)?;
6555 let trace_k = lambda_k
6556 * trace_of_factored_product(&penalty_root, &cov_block.to_owned())
6557 .map_err(EstimationError::InvalidInput)?;
6558 trace_sum.add(trace_k);
6559 penalty_block_trace.push(trace_k);
6560 let p_k = penalty_rank as f64;
6561 edf_by_block.push(certify_bounded_edf_interval(
6562 p_k - trace_k,
6563 0.0,
6564 p_k,
6565 p,
6566 &format!("bounded EDF block {k}"),
6567 )?);
6568 }
6569 PenaltySpec::Dense(m) | PenaltySpec::DenseWithMean { matrix: m, .. } => {
6570 s_lambda.scaled_add(lambda_k, m);
6571 let penalty_rank = p.saturating_sub(estimate_penalty_nullity(m).map_err(|e| {
6572 EstimationError::InvalidInput(format!("bounded EDF rank failed: {e}"))
6573 })?);
6574 let penalty_root = gam_solve::estimate::reml::reml_outer_engine::penalty_matrix_root(m)
6575 .map_err(EstimationError::InvalidInput)?;
6576 let trace_k = lambda_k
6577 * trace_of_factored_product(&penalty_root, latent_cov)
6578 .map_err(EstimationError::InvalidInput)?;
6579 trace_sum.add(trace_k);
6580 penalty_block_trace.push(trace_k);
6581 let p_k = penalty_rank as f64;
6582 edf_by_block.push(certify_bounded_edf_interval(
6583 p_k - trace_k,
6584 0.0,
6585 p_k,
6586 p,
6587 &format!("bounded EDF block {k}"),
6588 )?);
6589 }
6590 }
6591 }
6592
6593 let nullity_total = estimate_penalty_nullity(&s_lambda)
6594 .map_err(|e| EstimationError::InvalidInput(format!("bounded EDF nullity failed: {e}")))?
6595 as f64;
6596 let trace_sum = trace_sum.sum();
6597 let edf_total = certify_bounded_edf_interval(
6598 p as f64 - trace_sum,
6599 nullity_total,
6600 p as f64,
6601 p,
6602 "bounded total EDF",
6603 )?;
6604 Ok((edf_by_block, penalty_block_trace, edf_total))
6605}
6606
6607fn certified_profiled_gaussian_scale(
6631 deviance: f64,
6632 residual_dof: f64,
6633 label: &str,
6634) -> Result<f64, EstimationError> {
6635 if !(residual_dof.is_finite() && residual_dof > 0.0) {
6636 crate::bail_invalid_estim!(
6637 "{label} residual degrees of freedom must be finite and positive, got {residual_dof}; \
6638 a fit with no residual information cannot report a scale"
6639 );
6640 }
6641 if !(deviance.is_finite() && deviance >= 0.0) {
6642 crate::bail_invalid_estim!(
6643 "{label} deviance must be finite and non-negative, got {deviance}"
6644 );
6645 }
6646 let variance = deviance / residual_dof;
6647 if !variance.is_finite() {
6648 crate::bail_invalid_estim!(
6649 "{label} residual variance is not representable: {deviance}/{residual_dof}"
6650 );
6651 }
6652 Ok(variance.sqrt())
6653}
6654
6655fn certified_bounded_posterior_covariance(
6656 precision: &Array2<f64>,
6657 label: &'static str,
6658) -> Result<Array2<f64>, EstimationError> {
6659 gam_linalg::utils::certified_spd_inverse(precision, label)
6660 .map(gam_linalg::utils::CertifiedSpdInverse::into_inverse)
6661 .map_err(|error| {
6662 EstimationError::InvalidInput(format!(
6663 "bounded posterior covariance requires an exact SPD precision: {error}"
6664 ))
6665 })
6666}
6667
6668fn transform_bounded_latent_precision_to_user_internal(
6669 latent_precision: &Array2<f64>,
6670 jac_diag: &Array1<f64>,
6671) -> Result<Array2<f64>, EstimationError> {
6672 let p = latent_precision.nrows();
6673 if latent_precision.ncols() != p || jac_diag.len() != p {
6674 crate::bail_invalid_estim!(
6675 "bounded precision transform dimension mismatch: precision is {}x{}, jacobian has {} entries",
6676 latent_precision.nrows(),
6677 latent_precision.ncols(),
6678 jac_diag.len()
6679 );
6680 }
6681 let mut out = latent_precision.clone();
6682 for i in 0..p {
6683 let scale = jac_diag[i];
6684 if !scale.is_finite() || scale <= 0.0 {
6685 crate::bail_invalid_estim!(
6686 "bounded precision transform requires a positive finite coefficient jacobian; column {i} has {scale}"
6687 );
6688 }
6689 if scale != 1.0 {
6690 out.row_mut(i).mapv_inplace(|v| v / scale);
6691 out.column_mut(i).mapv_inplace(|v| v / scale);
6692 }
6693 }
6694 Ok(out)
6695}
6696
6697fn fit_bounded_term_collection_with_design(
6698 y: ArrayView1<'_, f64>,
6699 weights: ArrayView1<'_, f64>,
6700 offset: ArrayView1<'_, f64>,
6701 spec: &TermCollectionSpec,
6702 design: &TermCollectionDesign,
6703 heuristic_lambdas: Option<&[f64]>,
6704 family: LikelihoodSpec,
6705 options: &FitOptions,
6706) -> Result<FittedTermCollection, EstimationError> {
6707 let conditioning_cols: Vec<usize> = spec
6708 .linear_terms
6709 .iter()
6710 .enumerate()
6711 .filter_map(|(j, linear)| {
6712 (!linear.double_penalty).then_some(design.intercept_range.end + j)
6713 })
6714 .collect();
6715 let conditioning = LinearFitConditioning::from_columns(design, &conditioning_cols);
6716 let dense_design = design.design.to_dense_cow();
6717 let fit_design = conditioning.apply_to_design(&dense_design);
6718 let fit_penalties = conditioning
6719 .transform_blockwise_penalties_to_internal(&design.penalties, design.design.ncols());
6720 if design.linear_constraints.is_some() {
6721 crate::bail_invalid_estim!(
6722 "bounded() terms are not yet compatible with explicit linear constraints"
6723 );
6724 }
6725 let mut bounded_terms = Vec::<BoundedLinearTermMeta>::new();
6726 for (j, term) in spec.linear_terms.iter().enumerate() {
6727 if term.double_penalty
6728 && matches!(
6729 term.coefficient_geometry,
6730 LinearCoefficientGeometry::Bounded { .. }
6731 )
6732 {
6733 crate::bail_invalid_estim!(
6734 "bounded linear term '{}' cannot also use double_penalty",
6735 term.name
6736 );
6737 }
6738 if let LinearCoefficientGeometry::Bounded { min, max, prior } =
6739 term.coefficient_geometry.clone()
6740 {
6741 let col_idx = design.intercept_range.end + j;
6742 let (min_internal, max_internal) = conditioning.internal_bounds_for(col_idx, min, max);
6743 bounded_terms.push(BoundedLinearTermMeta {
6744 col_idx,
6745 min: min_internal,
6746 max: max_internal,
6747 prior,
6748 });
6749 }
6750 }
6751 if bounded_terms.is_empty() {
6752 crate::bail_invalid_estim!("internal bounded fit path called with no bounded terms");
6753 }
6754
6755 let mut designzeroed = fit_design.clone();
6756 let mut initial_beta = Array1::<f64>::zeros(fit_design.ncols());
6757 for term in &bounded_terms {
6758 designzeroed.column_mut(term.col_idx).fill(0.0);
6759 initial_beta[term.col_idx] = 0.0;
6760 }
6761
6762 let initial_log_lambdas = heuristic_lambdas
6763 .map(|vals| Array1::from_vec(vals.to_vec()))
6764 .unwrap_or_else(|| Array1::zeros(fit_penalties.len()));
6765 if initial_log_lambdas.len() != fit_penalties.len() {
6766 crate::bail_invalid_estim!(
6767 "heuristic lambda length mismatch for bounded model: got {}, expected {}",
6768 initial_log_lambdas.len(),
6769 fit_penalties.len()
6770 );
6771 }
6772
6773 let glm_likelihood = gam_spec::GlmLikelihoodSpec::canonical(family);
6774 let resolved_likelihood_scale = glm_likelihood
6775 .resolved_scale()
6776 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
6777 let is_beta_logistic = glm_likelihood.spec.is_binomial_beta_logistic();
6778 let family_adapter = BoundedLinearFamily {
6779 likelihood: glm_likelihood.clone(),
6780 latent_cloglog_state: options.latent_cloglog,
6781 mixture_link_state: options
6782 .mixture_link
6783 .clone()
6784 .as_ref()
6785 .map(state_fromspec)
6786 .transpose()
6787 .map_err(EstimationError::InvalidInput)?,
6788 sas_link_state: options
6789 .sas_link
6790 .map(|spec| {
6791 if is_beta_logistic {
6792 state_from_beta_logisticspec(spec)
6793 } else {
6794 state_from_sasspec(spec)
6795 }
6796 })
6797 .transpose()
6798 .map_err(EstimationError::InvalidInput)?,
6799 y: y.to_owned(),
6800 weights: weights.to_owned(),
6801 design: fit_design.clone(),
6802 designzeroed: designzeroed.clone(),
6803 offset: offset.to_owned(),
6804 bounded_terms: bounded_terms.clone(),
6805 };
6806 let blockspec = ParameterBlockSpec {
6807 name: "eta".to_string(),
6808 design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(designzeroed)),
6809 offset: offset.to_owned(),
6810 penalties: fit_penalties
6811 .iter()
6812 .map(|ps| match ps {
6813 PenaltySpec::Block {
6814 local, col_range, ..
6815 } => PenaltyMatrix::Blockwise {
6816 local: local.clone(),
6817 col_range: col_range.clone(),
6818 total_dim: design.design.ncols(),
6819 },
6820 PenaltySpec::Dense(m) | PenaltySpec::DenseWithMean { matrix: m, .. } => {
6821 PenaltyMatrix::Dense(m.clone())
6822 }
6823 })
6824 .collect(),
6825 nullspace_dims: design.nullspace_dims.clone(),
6826 initial_log_lambdas,
6827 initial_beta: Some(initial_beta),
6828 gauge_priority: 100,
6829 jacobian_callback: Some(Arc::new(BoundedEffectiveJacobian {
6835 design: fit_design.clone(),
6836 bounded_terms: bounded_terms.clone(),
6837 })),
6838 stacked_design: None,
6839 stacked_offset: None,
6840 };
6841 let fit = fit_custom_family(
6842 &family_adapter,
6843 &[blockspec],
6844 &BlockwiseFitOptions {
6845 inner_max_cycles: options.max_iter,
6846 inner_tol: options.tol,
6847 outer_max_iter: options.max_iter,
6848 outer_tol: options.tol,
6849 compute_covariance: false,
6859 ..BlockwiseFitOptions::default()
6860 },
6861 )
6862 .map_err(EstimationError::CustomFamily)?;
6863
6864 let latent_beta = fit.block_states[0].beta.clone();
6865 let (beta_user_internal, jac_diag) = family_adapter
6866 .user_beta_and_jacobian(&latent_beta)
6867 .map_err(EstimationError::InvalidInput)?;
6868 let beta_user = conditioning.backtransform_beta(&beta_user_internal);
6869
6870 let (eta_state, h_data, _, _) = family_adapter
6871 .evaluation_from_latent(&latent_beta)
6872 .map_err(EstimationError::InvalidInput)?;
6873 let p_fit = fit_design.ncols();
6874 let mut s_lambda_internal = Array2::<f64>::zeros((p_fit, p_fit));
6875 for (k, penalty) in fit_penalties.iter().enumerate() {
6876 match penalty {
6877 PenaltySpec::Block {
6878 local, col_range, ..
6879 } => {
6880 s_lambda_internal
6881 .slice_mut(ndarray::s![col_range.clone(), col_range.clone()])
6882 .scaled_add(fit.lambdas[k], local);
6883 }
6884 PenaltySpec::Dense(m) | PenaltySpec::DenseWithMean { matrix: m, .. } => {
6885 s_lambda_internal.scaled_add(fit.lambdas[k], m);
6886 }
6887 }
6888 }
6889 let mut latent_precision = h_data.clone();
6890 latent_precision += &s_lambda_internal;
6891 let user_precision_internal =
6892 transform_bounded_latent_precision_to_user_internal(&latent_precision, &jac_diag)?;
6893 let penalized_hessian =
6894 conditioning.transform_penalized_hessian_to_original(&user_precision_internal);
6895
6896 let beta_covariance_unscaled = if options.compute_inference {
6923 Some(certified_bounded_posterior_covariance(
6924 &penalized_hessian,
6925 "bounded user-scale posterior precision",
6926 )?)
6927 } else {
6928 None
6929 };
6930 let latent_cov = if options.compute_inference {
6936 Some(certified_bounded_posterior_covariance(
6937 &latent_precision,
6938 "bounded latent posterior precision",
6939 )?)
6940 } else {
6941 None
6942 };
6943 let s_lambda_original = weighted_blockwise_penalty_sum(
6944 &design.penalties,
6945 fit.lambdas
6946 .as_slice()
6947 .expect("the fitted lambdas are a contiguous standard-layout array"),
6948 design.design.ncols(),
6949 );
6950 let penalty_term = beta_user.dot(&s_lambda_original.dot(&beta_user));
6951 let deviance = -2.0 * eta_state.log_likelihood;
6952 let (edf_by_block, penalty_block_trace, edf_total) = if let Some(cov) = latent_cov.as_ref() {
6953 exact_bounded_edf(&fit_penalties, &fit.lambdas, cov)?
6954 } else {
6955 (
6956 vec![0.0; fit_penalties.len()],
6957 vec![0.0; fit_penalties.len()],
6958 0.0,
6959 )
6960 };
6961
6962 let profiled_gaussian_standard_deviation = if matches!(
6974 resolved_likelihood_scale,
6975 gam_spec::ResolvedLikelihoodScale::ProfiledGaussian
6976 ) {
6977 let residual_dof = if options.compute_inference {
6978 y.len() as f64 - edf_total
6979 } else {
6980 y.len() as f64
6981 };
6982 Some(certified_profiled_gaussian_scale(
6983 deviance,
6984 residual_dof,
6985 "bounded Gaussian",
6986 )?)
6987 } else {
6988 None
6989 };
6990 let dispersion = gam_solve::estimate::dispersion_from_likelihood(
6991 &glm_likelihood,
6992 profiled_gaussian_standard_deviation,
6993 )?;
6994 let standard_deviation = dispersion.phi().sqrt();
6995 let cov_scale = glm_likelihood
6996 .coefficient_covariance_scale(dispersion.phi())
6997 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
6998 let beta_covariance = beta_covariance_unscaled.map(|mut cov| {
7004 if cov_scale != 1.0 {
7005 cov.mapv_inplace(|v| v * cov_scale);
7006 }
7007 cov
7008 });
7009 if let Some(covariance) = beta_covariance.as_ref()
7010 && covariance.iter().any(|value| !value.is_finite())
7011 {
7012 return Err(EstimationError::InvalidInput(
7013 "bounded coefficient covariance scaling produced a non-finite value".to_string(),
7014 ));
7015 }
7016 let beta_standard_errors = beta_covariance
7017 .as_ref()
7018 .map(gam_problem::se_from_covariance)
7019 .transpose()
7020 .map_err(|err| {
7021 EstimationError::InvalidInput(format!(
7022 "bounded coefficient covariance cannot produce standard errors: {err}"
7023 ))
7024 })?;
7025 let working_response = exact_standard_working_response(&eta_state)?;
7026
7027 let geometry = Some(gam_solve::estimate::FitGeometry {
7028 coefficient_gauge: gam_problem::gauge::Gauge::identity(&[beta_user.len()]),
7029 penalized_hessian: penalized_hessian.clone().into(),
7030 constrained_posterior: None,
7031 working: Some(gam_solve::estimate::WorkingGeometry {
7032 weights: eta_state.fisherweight.clone(),
7033 response: working_response,
7034 }),
7035 });
7036 let max_abs_eta = eta_state
7037 .eta
7038 .iter()
7039 .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
7040 Ok(FittedTermCollection {
7041 fit: {
7042 let log_lambdas =
7043 checked_fit_log_lambdas(&fit.lambdas, "final fitted term collection")?;
7044 let inf = FitInference {
7045 edf_by_block,
7046 penalty_block_trace,
7047 edf_total,
7048 smoothing_correction: None,
7049 smoothing_correction_method: None,
7050 smoothing_correction_first_order: None,
7051 smoothing_correction_method_first_order: None,
7052 penalized_hessian: penalized_hessian.clone().into(),
7055 reparam_qs: None,
7056 dispersion,
7057 beta_covariance: beta_covariance
7058 .clone()
7059 .map(gam_problem::dispersion_cov::PhiScaledCovariance::from),
7060 beta_standard_errors,
7061 beta_covariance_corrected: None,
7062 beta_standard_errors_corrected: None,
7063 beta_covariance_frequentist: None,
7064 coefficient_influence: None,
7065 weighted_gram: None,
7066 bias_correction_beta: None,
7067 bias_correction_jacobian: None,
7068 };
7069 let covariance_conditional = beta_covariance;
7070 let pirls_status_val = gam_solve::pirls::PirlsStatus::Converged;
7073 let fit_objective = fit.penalized_objective();
7074 UnifiedFitResult::try_from_parts(UnifiedFitResultParts {
7075 blocks: vec![gam_solve::estimate::FittedBlock {
7076 beta: beta_user.clone(),
7077 role: gam_problem::BlockRole::Mean,
7078 edf: edf_total,
7079 lambdas: fit.lambdas.clone(),
7080 }],
7081 training_sample_size: y.len(),
7082 log_lambdas,
7083 lambdas: fit.lambdas,
7084 likelihood_scale: glm_likelihood.scale,
7085 likelihood_family: Some(glm_likelihood.spec),
7086 log_likelihood_normalization: gam_spec::LogLikelihoodNormalization::UserProvided,
7087 log_likelihood: eta_state.log_likelihood,
7088 deviance,
7089 reml_score: fit_objective,
7090 stable_penalty_term: penalty_term,
7091 penalized_objective: fit_objective,
7092 used_device: false,
7093 outer_iterations: fit.outer_iterations,
7094 outer_converged: true,
7096 outer_gradient_norm: fit.outer_gradient_norm,
7097 standard_deviation,
7098 covariance_conditional,
7099 covariance_corrected: None,
7100 inference: Some(inf),
7101 fitted_link: gam_solve::estimate::FittedLinkState::Standard(None),
7102 geometry,
7103 block_states: Vec::new(),
7104 pirls_status: pirls_status_val,
7105 max_abs_eta,
7106 constraint_kkt: None,
7107 artifacts: gam_solve::estimate::FitArtifacts {
7108 pirls: None,
7109 ..Default::default()
7110 },
7111 inner_cycles: 0,
7112 })?
7113 },
7114 design: design.clone(),
7115 adaptive_diagnostics: None,
7116 })
7117}
7118
7119fn enforce_term_constraint_feasibility(
7120 design: &TermCollectionDesign,
7121 fit: &UnifiedFitResult,
7122) -> Result<(), EstimationError> {
7123 const CONSTRAINT_FEASIBILITY_RAW_TOL: f64 = 1e-7;
7137 let tol = CONSTRAINT_FEASIBILITY_RAW_TOL;
7138 let smooth_start = design
7139 .design
7140 .ncols()
7141 .saturating_sub(design.smooth.total_smooth_cols());
7142 let mut violations: Vec<String> = Vec::new();
7143 for term in &design.smooth.terms {
7144 let gr = (smooth_start + term.coeff_range.start)..(smooth_start + term.coeff_range.end);
7145 let beta_local = fit.beta.slice(s![gr.clone()]).to_owned();
7146 if let Some(lb) = term.lower_bounds_local.as_ref() {
7147 let mut worst = 0.0_f64;
7148 let mut worst_idx = 0usize;
7149 for i in 0..lb.len().min(beta_local.len()) {
7150 if lb[i].is_finite() {
7151 let viol = (lb[i] - beta_local[i]).max(0.0);
7152 if viol > worst {
7153 worst = viol;
7154 worst_idx = i;
7155 }
7156 }
7157 }
7158 if worst > tol {
7159 violations.push(format!(
7160 "term='{}' kind=lower-bound maxviolation={:.3e} coeff_index={}",
7161 term.name, worst, worst_idx
7162 ));
7163 }
7164 }
7165 if let Some(lin) = term.linear_constraints_local.as_ref() {
7166 let mut worst = 0.0_f64;
7167 let mut worstrow = 0usize;
7168 for i in 0..lin.a.nrows() {
7169 let norm = lin.a.row(i).dot(&lin.a.row(i)).sqrt();
7170 let inv = if norm > 0.0 { 1.0 / norm } else { 0.0 };
7171 let s = (lin.a.row(i).dot(&beta_local) - lin.b[i]) * inv;
7172 let viol = (-s).max(0.0);
7173 if viol > worst {
7174 worst = viol;
7175 worstrow = i;
7176 }
7177 }
7178 if worst > tol {
7179 violations.push(format!(
7180 "term='{}' kind=linear-inequality maxviolation={:.3e} row={}",
7181 term.name, worst, worstrow
7182 ));
7183 }
7184 }
7185 }
7186
7187 if !violations.is_empty() {
7188 let mut msg = format!(
7189 "constraint violation after fit ({} violating term constraints): {}",
7190 violations.len(),
7191 violations.join(" | ")
7192 );
7193 if let Some(kkt) = fit.constraint_kkt.as_ref() {
7194 msg.push_str(&format!(
7195 "; KKT[primal={:.3e}, dual={:.3e}, comp={:.3e}, stat={:.3e}]{}",
7196 kkt.primal_feasibility,
7197 kkt.dual_feasibility,
7198 kkt.complementarity,
7199 kkt.stationarity,
7200 kkt.cone_projection_note()
7204 ));
7205 }
7206 return Err(EstimationError::ParameterConstraintViolation(msg));
7207 }
7208 Ok(())
7209}
7210
7211fn stratified_spatial_subsample(
7212 data: ArrayView2<'_, f64>,
7213 spec: &TermCollectionSpec,
7214 target_size: usize,
7215) -> Vec<usize> {
7216 use rand::SeedableRng;
7217 use rand::rngs::StdRng;
7218 use rand::seq::SliceRandom;
7219
7220 let n = data.nrows();
7221 if n <= target_size {
7222 return (0..n).collect();
7223 }
7224
7225 let spatial_cols: Option<Vec<usize>> =
7226 spec.smooth_terms.iter().find_map(|term| match &term.basis {
7227 SmoothBasisSpec::ThinPlate { feature_cols, .. }
7228 | SmoothBasisSpec::Matern { feature_cols, .. }
7229 | SmoothBasisSpec::Duchon { feature_cols, .. } => {
7230 if !feature_cols.is_empty() {
7231 Some(feature_cols.clone())
7232 } else {
7233 None
7234 }
7235 }
7236 _ => None,
7237 });
7238
7239 let cols = match spatial_cols {
7240 Some(c) if !c.is_empty() => c,
7241 _ => {
7242 let mut rng = StdRng::seed_from_u64(spatial_subsample_seed(data, &[], target_size));
7243 let mut indices: Vec<usize> = (0..n).collect();
7244 indices.shuffle(&mut rng);
7245 indices.truncate(target_size);
7246 indices.sort_unstable();
7247 return indices;
7248 }
7249 };
7250 let mut rng = StdRng::seed_from_u64(spatial_subsample_seed(data, &cols, target_size));
7251
7252 let d = cols.len();
7253 let mut mins = vec![f64::INFINITY; d];
7254 let mut maxs = vec![f64::NEG_INFINITY; d];
7255 for i in 0..n {
7256 for (ax, &col) in cols.iter().enumerate() {
7257 let v = data[[i, col]];
7258 if v < mins[ax] {
7259 mins[ax] = v;
7260 }
7261 if v > maxs[ax] {
7262 maxs[ax] = v;
7263 }
7264 }
7265 }
7266
7267 const TARGET_POINTS_PER_CELL: usize = 5;
7271 let total_cells_target = (target_size / TARGET_POINTS_PER_CELL).max(1);
7272 let cells_per_axis = ((total_cells_target as f64).powf(1.0 / d as f64)).ceil() as usize;
7273 let cells_per_axis = cells_per_axis.max(1);
7274
7275 let mut cell_members: std::collections::HashMap<Vec<usize>, Vec<usize>> =
7276 std::collections::HashMap::new();
7277 for i in 0..n {
7278 let mut cell_key = Vec::with_capacity(d);
7279 for (ax, &col) in cols.iter().enumerate() {
7280 let range = maxs[ax] - mins[ax];
7281 let cell = if range <= 0.0 {
7282 0
7283 } else {
7284 let frac = (data[[i, col]] - mins[ax]) / range;
7285 (frac * cells_per_axis as f64).floor() as usize
7286 };
7287 cell_key.push(cell.min(cells_per_axis - 1));
7288 }
7289 cell_members.entry(cell_key).or_default().push(i);
7290 }
7291
7292 let mut selected: Vec<usize> = Vec::with_capacity(target_size);
7293 let mut remaining_budget = target_size;
7294 let mut remaining_population = n;
7295
7296 let mut cells: Vec<(Vec<usize>, Vec<usize>)> = cell_members.into_iter().collect();
7297 cells.sort_by(|a, b| a.0.cmp(&b.0));
7298
7299 for (_, members) in &mut cells {
7300 if remaining_budget == 0 {
7301 break;
7302 }
7303 let alloc = ((members.len() as f64 / remaining_population as f64) * remaining_budget as f64)
7304 .round() as usize;
7305 let alloc = alloc.max(1).min(members.len()).min(remaining_budget);
7306 members.shuffle(&mut rng);
7307 selected.extend_from_slice(&members[..alloc]);
7308 remaining_budget = remaining_budget.saturating_sub(alloc);
7309 remaining_population = remaining_population.saturating_sub(members.len());
7310 }
7311
7312 if selected.len() > target_size {
7313 selected.shuffle(&mut rng);
7314 selected.truncate(target_size);
7315 }
7316
7317 selected.sort_unstable();
7318 selected
7319}
7320
7321fn spatial_subsample_seed(
7322 data: ArrayView2<'_, f64>,
7323 spatial_cols: &[usize],
7324 target_size: usize,
7325) -> u64 {
7326 let mut state = 0x5350_4154_4941_4C53_u64;
7327 spatial_seed_mix(&mut state, data.nrows() as u64);
7328 spatial_seed_mix(&mut state, data.ncols() as u64);
7329 spatial_seed_mix(&mut state, target_size as u64);
7330 spatial_seed_mix(&mut state, spatial_cols.len() as u64);
7331 for &col in spatial_cols {
7332 spatial_seed_mix(&mut state, col as u64);
7333 }
7334
7335 if data.nrows() > 0 {
7336 let mid = data.nrows() / 2;
7337 let last = data.nrows() - 1;
7338 for &row in &[0usize, mid, last] {
7339 for &col in spatial_cols {
7340 let value = data[[row, col]];
7341 spatial_seed_mix(&mut state, value.to_bits());
7342 }
7343 }
7344 }
7345 state
7346}
7347
7348#[inline]
7349fn spatial_seed_mix(state: &mut u64, value: u64) {
7350 let mut s = value.wrapping_add(*state);
7353 let z = gam_linalg::utils::splitmix64(&mut s);
7354 *state ^= z;
7355 *state = (*state).rotate_left(27).wrapping_mul(0x3C79_AC49_2BA7_B653);
7356}
7357
7358fn sampled_rows(data: ArrayView2<'_, f64>, indices: &[usize]) -> Array2<f64> {
7359 let mut sampled = Array2::<f64>::zeros((indices.len(), data.ncols()));
7360 for (new_row, &orig_row) in indices.iter().enumerate() {
7361 sampled.row_mut(new_row).assign(&data.row(orig_row));
7362 }
7363 sampled
7364}
7365
7366fn spatial_term_user_centers(term: &SmoothTermSpec) -> Option<ArrayView2<'_, f64>> {
7367 match spatial_term_center_strategy(term) {
7368 Some(CenterStrategy::UserProvided(centers)) => Some(centers.view()),
7369 _ => None,
7370 }
7371}
7372
7373fn finite_centered_axis_contrasts(values: &[f64], expected_dim: usize) -> Option<Vec<f64>> {
7374 if values.len() != expected_dim || expected_dim <= 1 {
7375 return None;
7376 }
7377 if values.iter().any(|value| !value.is_finite()) {
7378 return None;
7379 }
7380 Some(center_aniso_log_scales(values))
7381}
7382
7383fn blended_pilot_axis_contrasts(
7384 pilot_data: ArrayView2<'_, f64>,
7385 term: &SmoothTermSpec,
7386 centers: ArrayView2<'_, f64>,
7387) -> Result<Option<Vec<f64>>, BasisError> {
7388 let d = centers.ncols();
7389 if d <= 1 {
7390 return Ok(None);
7391 }
7392 let center_eta = initial_aniso_contrasts(centers);
7393 let standardized_data = standardized_spatial_term_data(pilot_data, term)?;
7394 let data_eta = finite_centered_axis_contrasts(
7395 &initial_aniso_contrasts(standardized_data.view()),
7396 d,
7397 );
7398 let Some(center_eta) = finite_centered_axis_contrasts(¢er_eta, d) else {
7399 return Ok(None);
7400 };
7401 let blended = match data_eta {
7402 Some(data_eta) => center_eta
7403 .iter()
7404 .zip(data_eta.iter())
7405 .map(|(&from_centers, &from_data)| 0.5 * (from_centers + from_data))
7406 .collect::<Vec<_>>(),
7407 None => center_eta,
7408 };
7409 Ok(finite_centered_axis_contrasts(&blended, d))
7410}
7411
7412fn apply_pilot_spatial_psi_reseed(
7413 pilot_data: ArrayView2<'_, f64>,
7414 spec: &TermCollectionSpec,
7415 spatial_terms: &[usize],
7416 kappa_options: &SpatialLengthScaleOptimizationOptions,
7417) -> Result<TermCollectionSpec, EstimationError> {
7418 let dims_per_term = spatial_dims_per_term(spec, spatial_terms);
7419 let use_aniso = has_aniso_terms(spec, spatial_terms);
7420 let log_kappa0 = if use_aniso {
7421 SpatialLogKappaCoords::from_length_scales_aniso(spec, spatial_terms, kappa_options)
7422 } else {
7423 SpatialLogKappaCoords::from_length_scales(spec, spatial_terms, kappa_options)
7424 };
7425 let log_kappa0 = log_kappa0
7426 .reseed_from_data(pilot_data, spec, spatial_terms, kappa_options)
7427 .map_err(EstimationError::BasisError)?;
7428 let log_kappa_lower = if use_aniso {
7429 SpatialLogKappaCoords::lower_bounds_aniso_from_data(
7430 pilot_data,
7431 spec,
7432 spatial_terms,
7433 &dims_per_term,
7434 kappa_options,
7435 )
7436 } else {
7437 SpatialLogKappaCoords::lower_bounds_from_data(
7438 pilot_data,
7439 spec,
7440 spatial_terms,
7441 kappa_options,
7442 )
7443 }
7444 .map_err(EstimationError::BasisError)?;
7445 let log_kappa_upper = if use_aniso {
7446 SpatialLogKappaCoords::upper_bounds_aniso_from_data(
7447 pilot_data,
7448 spec,
7449 spatial_terms,
7450 &dims_per_term,
7451 kappa_options,
7452 )
7453 } else {
7454 SpatialLogKappaCoords::upper_bounds_from_data(
7455 pilot_data,
7456 spec,
7457 spatial_terms,
7458 kappa_options,
7459 )
7460 }
7461 .map_err(EstimationError::BasisError)?;
7462 log_kappa0
7463 .clamp_to_bounds(&log_kappa_lower, &log_kappa_upper)
7464 .apply_tospec(spec, spatial_terms)
7465}
7466
7467pub(crate) fn apply_spatial_anisotropy_pilot_initializer(
7468 data: ArrayView2<'_, f64>,
7469 spec: &mut TermCollectionSpec,
7470 spatial_terms: &[usize],
7471 target_size: usize,
7472 kappa_options: &SpatialLengthScaleOptimizationOptions,
7473) -> Result<usize, EstimationError> {
7474 if target_size == 0 || data.nrows() <= target_size.saturating_mul(2) || spatial_terms.is_empty()
7475 {
7476 return Ok(0);
7477 }
7478 if !has_aniso_terms(spec, spatial_terms) {
7479 return Ok(0);
7480 }
7481 let indices = stratified_spatial_subsample(data, spec, target_size);
7482 let pilot_data = sampled_rows(data, &indices);
7483 let mut working = spec.clone();
7484 let mut updated_terms = 0usize;
7485 const GEOMETRY_UPDATES: usize = 2;
7486
7487 for pass in 0..GEOMETRY_UPDATES {
7488 let planned_terms = plan_joint_spatial_centers_for_term_blocks(
7489 pilot_data.view(),
7490 &[working.smooth_terms.clone()],
7491 )
7492 .and_then(|mut blocks| {
7493 blocks.pop().ok_or_else(|| {
7494 BasisError::InvalidInput(
7495 "pilot geometry initializer produced no smooth-term block".to_string(),
7496 )
7497 })
7498 })
7499 .map_err(EstimationError::BasisError)?;
7500
7501 for &term_idx in spatial_terms {
7502 let Some(current_eta) = get_spatial_aniso_log_scales(&working, term_idx) else {
7503 continue;
7504 };
7505 let Some(d) = get_spatial_feature_dim(&working, term_idx) else {
7506 continue;
7507 };
7508 if d <= 1 || current_eta.len() != d {
7509 continue;
7510 }
7511 let Some(planned_term) = planned_terms.get(term_idx) else {
7512 continue;
7513 };
7514 let Some(centers) = spatial_term_user_centers(planned_term) else {
7515 continue;
7516 };
7517 let Some(eta) = blended_pilot_axis_contrasts(
7518 pilot_data.view(),
7519 planned_term,
7520 centers,
7521 )
7522 .map_err(EstimationError::BasisError)?
7523 else {
7524 continue;
7525 };
7526 set_spatial_aniso_log_scales(&mut working, term_idx, eta)?;
7527 updated_terms += usize::from(pass == 0);
7528 }
7529
7530 working = apply_pilot_spatial_psi_reseed(
7531 pilot_data.view(),
7532 &working,
7533 spatial_terms,
7534 kappa_options,
7535 )?;
7536 }
7537
7538 if updated_terms > 0 {
7539 log::info!(
7540 "[spatial-kappa] initialized anisotropy from {}-row pilot geometry for {} spatial term(s); proceeding to full-data optimization",
7541 indices.len(),
7542 updated_terms
7543 );
7544 *spec = working;
7545 }
7546 Ok(updated_terms)
7547}
7548
7549pub(crate) fn spatial_length_scale_term_indices(spec: &TermCollectionSpec) -> Vec<usize> {
7550 spec.smooth_terms
7551 .iter()
7552 .enumerate()
7553 .filter_map(|(idx, _)| spatial_term_supports_hyper_optimization(spec, idx).then_some(idx))
7554 .collect()
7555}
7556
7557fn fit_score(fit: &UnifiedFitResult) -> f64 {
7567 if let Some(score) = fit.reml_score().filter(|value| value.is_finite()) {
7568 return score;
7569 }
7570 let score = 0.5 * fit.deviance + 0.5 * fit.stable_penalty_term;
7571 if score.is_finite() {
7572 score
7573 } else {
7574 f64::INFINITY
7575 }
7576}
7577
7578fn is_recoverable_trial_point_error(err: &EstimationError) -> bool {
7600 matches!(err, EstimationError::BasisError(_)) || err.is_trial_point_infeasible()
7609}
7610
7611#[cfg(test)]
7612mod spatial_trial_recovery_tests {
7613 use super::*;
7614
7615 #[test]
7616 fn nonfinite_frequentist_covariance_is_recoverable_trial_point() {
7617 let err = EstimationError::TrialPointRefused {
7620 reason: "fit_result.beta_covariance_frequentist[0] must be finite, got NaN".to_string(),
7621 };
7622
7623 assert!(
7624 is_recoverable_trial_point_error(&err),
7625 "singular trial-point curvature should make spatial κ retreat, not abort"
7626 );
7627 assert!(
7630 !is_recoverable_trial_point_error(&EstimationError::InvalidInput(
7631 "fit_result.beta_covariance_frequentist[0] must be finite, got NaN".to_string(),
7632 )),
7633 "recoverability must come from the producer's variant, not from the message"
7634 );
7635 }
7636
7637 #[test]
7638 fn arbitrary_invalid_input_remains_fatal_trial_point_error() {
7639 let err = EstimationError::InvalidInput("outer rho bounds are invalid".to_string());
7640
7641 assert!(
7642 !is_recoverable_trial_point_error(&err),
7643 "the spatial κ recovery gate must not mask unrelated invalid inputs"
7644 );
7645 }
7646
7647 #[test]
7648 fn spatial_value_probe_classifier_matches_derivative_lane() {
7649 let nonfinite_covariance =
7666 "fit_result.beta_covariance_frequentist[0] must be finite, got NaN";
7667 let unrelated = "outer rho bounds are invalid";
7668 let cases = [
7669 EstimationError::TrialPointRefused {
7670 reason: nonfinite_covariance.to_string(),
7671 },
7672 EstimationError::BasisError(gam_problem::BasisError::DegenerateRange(8)),
7675 EstimationError::InvalidInput(nonfinite_covariance.to_string()),
7676 EstimationError::InvalidInput(unrelated.to_string()),
7677 ];
7678
7679 for error in cases {
7680 let message = error.to_string();
7681 let derivative_lane_recovers = is_recoverable_trial_point_error(&error);
7682 match classify_spatial_value_probe_failure(error) {
7683 Ok(value) => {
7684 assert!(
7685 derivative_lane_recovers,
7686 "the value probe retreated on {message:?} while the derivative lane \
7687 calls it fatal — the two lanes must classify one error the same way"
7688 );
7689 assert!(
7690 value.is_infinite() && value.is_sign_positive(),
7691 "a domain refusal must retreat to +INFINITY so the line search steps \
7692 away from it; got {value} for {message:?}"
7693 );
7694 }
7695 Err(propagated) => {
7696 assert!(
7697 !derivative_lane_recovers,
7698 "the value probe propagated {message:?} while the derivative lane \
7699 calls it a recoverable trial point"
7700 );
7701 assert_eq!(
7702 propagated.to_string(),
7703 message,
7704 "a fatal failure must be propagated unchanged, not reworded"
7705 );
7706 }
7707 }
7708 }
7709 }
7710}
7711
7712fn require_available_spatial_optimization_result<T>(
7740 result: Result<Option<T>, EstimationError>,
7741) -> Result<T, EstimationError> {
7742 match result {
7743 Ok(Some(value)) => Ok(value),
7744 Ok(None) => Err(EstimationError::RemlOptimizationFailed(
7745 "spatial kappa optimization is unavailable for one or more eligible spatial terms"
7746 .to_string(),
7747 )),
7748 Err(err) => Err(EstimationError::RemlOptimizationFailed(format!(
7749 "spatial kappa optimization failed: {err}"
7750 ))),
7751 }
7752}
7753
7754fn external_opts_for_design(
7755 family: &LikelihoodSpec,
7756 design: &TermCollectionDesign,
7757 options: &FitOptions,
7758) -> ExternalOptimOptions {
7759 ExternalOptimOptions {
7760 family: family.clone(),
7761 latent_cloglog: options.latent_cloglog,
7762 mixture_link: options.mixture_link.clone(),
7763 optimize_mixture: options.optimize_mixture,
7764 sas_link: options.sas_link,
7765 optimize_sas: options.optimize_sas,
7766 compute_inference: options.compute_inference,
7767 skip_rho_posterior_inference: options.skip_rho_posterior_inference,
7768 max_iter: options.max_iter,
7769 tol: options.tol,
7770 nullspace_dims: design.nullspace_dims.clone(),
7771 linear_constraints: design.linear_constraints.clone(),
7772 firth_bias_reduction: Some(options.firth_bias_reduction),
7773 rho_prior: options.rho_prior.clone(),
7774 kronecker_penalty_system: design.kronecker_penalty_system(),
7777 kronecker_factored: design
7778 .smooth
7779 .terms
7780 .iter()
7781 .find_map(|t| t.kronecker_factored.clone()),
7782 persistent_warm_start_store: options.persistent_warm_start_store.clone(),
7783 }
7784}
7785
7786fn evaluate_joint_reml_outer_eval_at_theta(
7794 evaluator: &mut gam_solve::estimate::ExternalJointHyperEvaluator<'_>,
7795 design: &TermCollectionDesign,
7796 theta: &Array1<f64>,
7797 rho_dim: usize,
7798 hyper_dirs: Vec<gam_solve::estimate::reml::DirectionalHyperParam>,
7799 warm_start_beta: Option<ArrayView1<'_, f64>>,
7800 order: gam_solve::rho_optimizer::OuterEvalOrder,
7801 design_revision: Option<u64>,
7802) -> Result<(f64, Array1<f64>, gam_problem::HessianValue), EstimationError> {
7803 evaluator.evaluate_with_order(
7804 &design.design,
7805 &design.penalties,
7806 &design.nullspace_dims,
7807 design.linear_constraints.clone(),
7808 theta,
7809 rho_dim,
7810 hyper_dirs,
7811 warm_start_beta,
7812 "evaluate_joint_reml_outer_eval_at_theta",
7813 order,
7814 design_revision,
7815 )
7816}
7817
7818fn evaluate_joint_reml_efs_at_theta(
7819 evaluator: &mut gam_solve::estimate::ExternalJointHyperEvaluator<'_>,
7820 design: &TermCollectionDesign,
7821 theta: &Array1<f64>,
7822 rho_dim: usize,
7823 hyper_dirs: Vec<gam_solve::estimate::reml::DirectionalHyperParam>,
7824 warm_start_beta: Option<ArrayView1<'_, f64>>,
7825 design_revision: Option<u64>,
7826) -> Result<gam_problem::EfsEval, EstimationError> {
7827 evaluator.evaluate_efs(
7828 &design.design,
7829 &design.penalties,
7830 &design.nullspace_dims,
7831 design.linear_constraints.clone(),
7832 theta,
7833 rho_dim,
7834 hyper_dirs,
7835 warm_start_beta,
7836 "evaluate_joint_reml_efs_at_theta",
7837 design_revision,
7838 )
7839}
7840
7841fn exact_joint_spatial_outer_hessian_available(
7842 family: &LikelihoodSpec,
7843 design: &TermCollectionDesign,
7844) -> bool {
7845 let family_supported = match &family.response {
7868 ResponseFamily::Gaussian
7869 | ResponseFamily::Binomial
7870 | ResponseFamily::Poisson
7871 | ResponseFamily::Tweedie { .. }
7872 | ResponseFamily::NegativeBinomial { .. }
7873 | ResponseFamily::Beta { .. }
7874 | ResponseFamily::Gamma
7875 | ResponseFamily::RoystonParmar => true,
7876 };
7877 family_supported && design.design.ncols() > 0
7880}
7881
7882fn try_build_spatial_term_log_kappa_derivativeinfo(
7883 data: ArrayView2<'_, f64>,
7884 resolvedspec: &TermCollectionSpec,
7885 design: &TermCollectionDesign,
7886 term_idx: usize,
7887) -> Result<Option<SpatialPsiDerivative>, EstimationError> {
7888 let Some((
7889 global_range,
7890 total_p,
7891 x_psi_local,
7892 s_psi_local_check,
7893 x_psi_psi_local,
7894 s_psi_psi_local,
7895 s_psi_components_local,
7896 s_psi_psi_components_local,
7897 implicit_operator,
7898 )) = try_build_spatial_term_log_kappa_derivative(data, resolvedspec, design, term_idx)?
7899 else {
7900 return Ok(None);
7901 };
7902 let Some(penalty_range) = design
7903 .smooth_term_penalty_range(term_idx)
7904 .map_err(EstimationError::InvalidInput)?
7905 else {
7906 return Ok(None);
7907 };
7908 let penalty_start = penalty_range.start;
7909 if s_psi_components_local.is_empty() || s_psi_psi_components_local.is_empty() {
7910 return Ok(None);
7911 }
7912 if s_psi_components_local.len() != s_psi_psi_components_local.len() {
7913 return Ok(None);
7914 }
7915 let penalty_indices = (0..s_psi_components_local.len())
7916 .map(|j| penalty_start + j)
7917 .collect::<Vec<_>>();
7918 let penalty_index = penalty_indices[0];
7919 if s_psi_local_check.nrows() == 0 || s_psi_psi_local.nrows() == 0 {
7920 return Ok(None);
7921 }
7922 Ok(Some(SpatialPsiDerivative {
7923 penalty_index,
7924 penalty_indices,
7925 global_range,
7926 total_p,
7927 x_psi_local,
7928 s_psi_components_local,
7929 x_psi_psi_local,
7930 s_psi_psi_components_local,
7931 aniso_group_id: None,
7932 aniso_cross_designs: None,
7933 aniso_cross_penalty_provider: None,
7934 implicit_operator,
7935 implicit_axis: 0,
7936 }))
7937}
7938
7939pub(crate) fn try_build_spatial_log_kappa_derivativeinfo_list(
7940 data: ArrayView2<'_, f64>,
7941 resolvedspec: &TermCollectionSpec,
7942 design: &TermCollectionDesign,
7943 spatial_terms: &[usize],
7944) -> Result<Option<Vec<SpatialPsiDerivative>>, EstimationError> {
7945 let mut out = Vec::new();
7946 let mut aniso_gid = 0usize;
7947 for &term_idx in spatial_terms {
7948 if spatial_term_uses_per_axis_psi(resolvedspec, term_idx) {
7949 if let Some(entries) = try_build_spatial_term_log_kappa_aniso_derivativeinfos(
7950 data,
7951 resolvedspec,
7952 design,
7953 term_idx,
7954 aniso_gid,
7955 )? {
7956 aniso_gid += 1;
7957 out.extend(entries);
7958 continue;
7959 } else {
7960 log::warn!(
7968 "[spatial-kappa] term {term_idx}: enrolled for per-axis ψ but its per-axis \
7969 derivative producer declined; the joint κ route is unavailable for this fit"
7970 );
7971 return Ok(None);
7972 }
7973 }
7974 let Some(info) =
7975 try_build_spatial_term_log_kappa_derivativeinfo(data, resolvedspec, design, term_idx)?
7976 else {
7977 log::warn!(
7978 "[spatial-kappa] term {term_idx}: isotropic ψ derivative producer declined; the \
7979 joint κ route is unavailable for this fit"
7980 );
7981 return Ok(None);
7982 };
7983 out.push(info);
7984 }
7985 Ok(Some(out))
7986}
7987
7988fn try_build_spatial_term_log_kappa_aniso_derivativeinfos(
7990 data: ArrayView2<'_, f64>,
7991 resolvedspec: &TermCollectionSpec,
7992 design: &TermCollectionDesign,
7993 term_idx: usize,
7994 aniso_group_id: usize,
7995) -> Result<Option<Vec<SpatialPsiDerivative>>, EstimationError> {
7996 let Some(smooth_term) = design.smooth.terms.get(term_idx) else {
7997 return Ok(None);
7998 };
7999 let Some(termspec) = resolvedspec.smooth_terms.get(term_idx) else {
8000 return Ok(None);
8001 };
8002 let mut aniso_result = match &termspec.basis {
8003 SmoothBasisSpec::Sphere { .. } => return Ok(None),
8004 SmoothBasisSpec::Matern {
8005 feature_cols,
8006 spec,
8007 input_scale,
8008 } => {
8009 let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
8010 let mut spec_operator = spec.clone();
8011 if let Some(scale) = input_scale {
8012 scale.standardize(&mut x);
8013 let length_scale = spec.length_scale.resolved().ok_or_else(|| {
8014 EstimationError::InvalidInput(
8015 "anisotropic Matérn Auto length_scale reached derivative construction \
8016 unresolved"
8017 .to_string(),
8018 )
8019 })?;
8020 spec_operator.length_scale.set_resolved(
8021 scale
8022 .to_standardized_units(gam_terms::OriginalUnits::new(length_scale))
8023 .standardized_value(),
8024 );
8025 }
8026 spec_operator.double_penalty = false;
8035 build_matern_basis_log_kappa_aniso_derivatives(x.view(), &spec_operator)
8036 .map_err(EstimationError::from)?
8037 }
8038 SmoothBasisSpec::MeasureJet {
8044 feature_cols,
8045 spec,
8046 input_scale,
8047 } => {
8048 let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
8049 if let Some(scale) = input_scale {
8050 scale.standardize(&mut x);
8051 }
8052 build_measure_jet_basis_psi_derivatives(x.view(), spec)
8053 .map_err(EstimationError::from)?
8054 }
8055 SmoothBasisSpec::Duchon {
8062 feature_cols,
8063 spec,
8064 input_scale,
8065 } => {
8066 let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
8067 let mut spec_local = spec.clone();
8068 if let Some(scale) = input_scale {
8069 scale.standardize(&mut x);
8070 spec_local.length_scale = spec.length_scale.map(|length| {
8071 scale
8072 .to_standardized_units(gam_terms::OriginalUnits::new(length))
8073 .standardized_value()
8074 });
8075 }
8076 let BasisMetadata::Duchon {
8077 centers,
8078 identifiability_transform,
8079 operator_collocation_points,
8080 radial_reparam,
8081 aniso_log_scales,
8082 ..
8083 } = &smooth_term.metadata
8084 else {
8085 log::warn!(
8086 "[spatial-kappa] term {term_idx}: per-axis ψ declined -- a Duchon spec whose \
8087 realized design does not carry Duchon metadata"
8088 );
8089 return Ok(None);
8090 };
8091 if spec_local.radial_reparam.is_none() {
8092 spec_local.radial_reparam = radial_reparam.clone();
8093 }
8094 if let Some(resolved) = aniso_log_scales.as_ref() {
8099 spec_local.aniso_log_scales = Some(resolved.clone());
8100 }
8101 gam_terms::basis::build_duchon_basis_log_kappa_aniso_derivativeswith_collocationwithworkspace(
8102 x.view(),
8103 &spec_local,
8104 centers.view(),
8105 identifiability_transform.as_ref(),
8106 operator_collocation_points
8107 .as_ref()
8108 .map(|points| points.view()),
8109 &mut BasisWorkspace::default(),
8110 )
8111 .map_err(EstimationError::from)?
8112 }
8113 _ => return Ok(None),
8114 };
8115 let d = if let Some(ref op) = aniso_result.implicit_operator {
8118 op.n_axes()
8119 } else if !aniso_result.design_first.is_empty() {
8120 aniso_result.design_first.len()
8121 } else {
8122 0
8123 };
8124 if d == 0 {
8125 log::warn!(
8126 "[spatial-kappa] term {term_idx}: per-axis ψ declined -- the producer reported zero \
8127 axes (no implicit operator and no dense design list)"
8128 );
8129 return Ok(None);
8130 }
8131 let Some(penalty_range) = design
8132 .smooth_term_penalty_range(term_idx)
8133 .map_err(EstimationError::InvalidInput)?
8134 else {
8135 log::warn!(
8136 "[spatial-kappa] term {term_idx}: per-axis ψ declined -- the realized design exposes \
8137 no penalty range for this term"
8138 );
8139 return Ok(None);
8140 };
8141 let penalty_start = penalty_range.start;
8142 let p_total = design.design.ncols();
8143 let smooth_start = p_total.saturating_sub(design.smooth.total_smooth_cols());
8144 let global_range = (smooth_start + smooth_term.coeff_range.start)
8145 ..(smooth_start + smooth_term.coeff_range.end);
8146 let emitted = aniso_result.penalties_first[0].len();
8166 let producer_emits_active_list = matches!(&termspec.basis, SmoothBasisSpec::Duchon { .. });
8167 let keep: Vec<usize> = if producer_emits_active_list {
8168 if emitted != smooth_term.active_penalties.len() {
8169 log::warn!(
8170 "[spatial-kappa] term {term_idx}: per-axis ψ declined -- the Duchon producer \
8171 emitted {emitted} active penalty block(s) but the realized design carries {} \
8172 ({:?}); the term falls back to its isotropic axis",
8173 smooth_term.active_penalties.len(),
8174 smooth_term
8175 .active_penalties
8176 .iter()
8177 .map(|active| active.info.source.clone())
8178 .collect::<Vec<_>>()
8179 );
8180 return Ok(None);
8181 }
8182 (0..emitted).collect()
8183 } else {
8184 smooth_term
8185 .active_penalties
8186 .iter()
8187 .map(|active| active.info.original_index)
8188 .collect()
8189 };
8190 if keep.is_empty() || keep.iter().any(|&index| index >= emitted) {
8191 log::warn!(
8192 "[spatial-kappa] term {term_idx}: per-axis ψ declined -- candidate→fitted map \
8193 {keep:?} does not index the {emitted} emitted block(s)"
8194 );
8195 return Ok(None);
8196 }
8197 let penalty_indices: Vec<usize> = (0..keep.len()).map(|j| penalty_start + j).collect();
8198 let penalties_cross_provider = aniso_result.penalties_cross_provider.clone();
8199 fn select_fitted(blocks: Vec<Array2<f64>>, keep: &[usize]) -> Vec<Array2<f64>> {
8202 let mut slots: Vec<Option<Array2<f64>>> = blocks.into_iter().map(Some).collect();
8203 keep.iter()
8204 .filter_map(|&index| slots.get_mut(index).and_then(Option::take))
8205 .collect()
8206 }
8207
8208 let use_implicit_design = aniso_result.design_first.is_empty();
8212 let implicit_op_arc = aniso_result
8213 .implicit_operator
8214 .as_ref()
8215 .map(|op| std::sync::Arc::new(op.clone()));
8216
8217 let mut entries = Vec::with_capacity(d);
8218 for a in 0..d {
8219 let (x_psi_local, x_psi_psi_local) = if use_implicit_design {
8220 (Array2::<f64>::zeros((0, 0)), Array2::<f64>::zeros((0, 0)))
8226 } else {
8227 let x_first = std::mem::take(&mut aniso_result.design_first[a]);
8232 let x_second = std::mem::take(&mut aniso_result.design_second_diag[a]);
8233 if x_first.ncols() != smooth_term.coeff_range.len() {
8234 return Ok(None);
8235 }
8236 (x_first, x_second)
8237 };
8238 let s_psi_components =
8239 select_fitted(std::mem::take(&mut aniso_result.penalties_first[a]), &keep);
8240 let s_psi_psi_components = select_fitted(
8241 std::mem::take(&mut aniso_result.penalties_second_diag[a]),
8242 &keep,
8243 );
8244 let cross_designs = if implicit_op_arc.is_some() {
8250 let mut cd = Vec::with_capacity(d - 1);
8251 for b in 0..d {
8252 if b == a {
8253 continue;
8254 }
8255 cd.push((b, Array2::<f64>::zeros((0, 0))));
8256 }
8257 cd
8258 } else if !aniso_result.design_second_cross.is_empty() {
8259 let mut cd = Vec::new();
8260 for (cross_idx, &(pa, pb)) in aniso_result.design_second_cross_pairs.iter().enumerate()
8261 {
8262 if pa == a {
8263 cd.push((pb, aniso_result.design_second_cross[cross_idx].clone()));
8264 } else if pb == a {
8265 cd.push((pa, aniso_result.design_second_cross[cross_idx].clone()));
8266 }
8267 }
8268 cd
8269 } else {
8270 Vec::new()
8271 };
8272 let cross_penalty_provider = if d > 1 {
8273 let penalties_cross_provider = penalties_cross_provider.clone();
8274 let keep_cross = keep.clone();
8275 Some(std::sync::Arc::new(
8276 move |b_axis: usize| -> Result<Vec<Array2<f64>>, EstimationError> {
8277 if b_axis == a {
8278 return Ok(Vec::new());
8279 }
8280 let (axis_lo, axis_hi) = if a < b_axis { (a, b_axis) } else { (b_axis, a) };
8281 if let Some(provider) = penalties_cross_provider.as_ref() {
8282 provider
8286 .evaluate(axis_lo, axis_hi)
8287 .map_err(EstimationError::from)
8288 .map(|blocks| select_fitted(blocks, &keep_cross))
8289 } else {
8290 Ok(Vec::new())
8294 }
8295 },
8296 )
8297 as std::sync::Arc<
8298 dyn Fn(usize) -> Result<Vec<Array2<f64>>, EstimationError>
8299 + Send
8300 + Sync
8301 + 'static,
8302 >)
8303 } else {
8304 None
8305 };
8306
8307 entries.push(SpatialPsiDerivative {
8308 penalty_index: penalty_indices[0],
8309 penalty_indices: penalty_indices.clone(),
8310 global_range: global_range.clone(),
8311 total_p: p_total,
8312 x_psi_local,
8313 s_psi_components_local: s_psi_components,
8314 x_psi_psi_local,
8315 s_psi_psi_components_local: s_psi_psi_components,
8316 aniso_group_id: Some(aniso_group_id),
8317 aniso_cross_designs: if cross_designs.is_empty() {
8318 None
8319 } else {
8320 Some(cross_designs)
8321 },
8322 aniso_cross_penalty_provider: cross_penalty_provider,
8323 implicit_operator: implicit_op_arc.clone(),
8324 implicit_axis: a,
8325 });
8326 }
8327 Ok(Some(entries))
8328}
8329
8330#[cfg(test)]
8331mod glm_eta_observation_fd_tests {
8332 use super::*;
8338 use ndarray::array;
8339
8340 fn one_obs_weight(
8341 spec: &LikelihoodSpec,
8342 y: f64,
8343 weight: f64,
8344 eta: f64,
8345 ) -> StandardFamilyObservationState {
8346 let yv = Array1::from_vec(vec![y]);
8347 let wv = Array1::from_vec(vec![weight]);
8348 let ev = Array1::from_vec(vec![eta]);
8349 evaluate_standard_familyobservations(spec.clone(), None, None, None, &yv, &wv, &ev)
8350 .expect("standard family observation state assembles")
8351 }
8352
8353 fn one_obs(spec: &LikelihoodSpec, y: f64, eta: f64) -> StandardFamilyObservationState {
8354 one_obs_weight(spec, y, 1.0, eta)
8355 }
8356
8357 fn one_obs_resolved(
8358 likelihood: &gam_spec::GlmLikelihoodSpec,
8359 y: f64,
8360 weight: f64,
8361 eta: f64,
8362 ) -> StandardFamilyObservationState {
8363 evaluate_resolved_standard_family_observations(
8364 likelihood,
8365 None,
8366 None,
8367 None,
8368 &array![y],
8369 &array![weight],
8370 &array![eta],
8371 )
8372 .expect("resolved standard family observation state assembles")
8373 }
8374
8375 #[test]
8376 fn bounded_gamma_and_tweedie_use_the_resolved_likelihood_scale() {
8377 let gamma_unit = gam_spec::GlmLikelihoodSpec {
8378 spec: LikelihoodSpec::gamma_log(),
8379 scale: gam_spec::LikelihoodScaleMetadata::FixedGammaShape { shape: 1.0 },
8380 };
8381 let gamma_scaled = gam_spec::GlmLikelihoodSpec {
8382 spec: LikelihoodSpec::gamma_log(),
8383 scale: gam_spec::LikelihoodScaleMetadata::FixedGammaShape { shape: 8.0 },
8384 };
8385 let unit = one_obs_resolved(&gamma_unit, 2.3, 0.7, 0.2);
8386 let scaled = one_obs_resolved(&gamma_scaled, 2.3, 0.7, 0.2);
8387 for (label, actual, base) in [
8388 ("Gamma score", scaled.score[0], unit.score[0]),
8389 (
8390 "Gamma Fisher weight",
8391 scaled.fisherweight[0],
8392 unit.fisherweight[0],
8393 ),
8394 (
8395 "Gamma observed Hessian",
8396 scaled.neghessian_eta[0],
8397 unit.neghessian_eta[0],
8398 ),
8399 (
8400 "Gamma Hessian derivative",
8401 scaled.neghessian_eta_derivative[0],
8402 unit.neghessian_eta_derivative[0],
8403 ),
8404 (
8405 "Gamma log likelihood",
8406 scaled.log_likelihood,
8407 unit.log_likelihood,
8408 ),
8409 ] {
8410 let expected = 8.0 * base;
8411 assert!(
8412 (actual - expected).abs() <= 32.0 * f64::EPSILON * expected.abs().max(1.0),
8413 "{label} scale mismatch: actual={actual}, expected={expected}"
8414 );
8415 }
8416
8417 let tweedie_unit = gam_spec::GlmLikelihoodSpec {
8418 spec: LikelihoodSpec::tweedie_log(1.5),
8419 scale: gam_spec::LikelihoodScaleMetadata::FixedDispersion { phi: 1.0 },
8420 };
8421 let tweedie_scaled = gam_spec::GlmLikelihoodSpec {
8422 spec: LikelihoodSpec::tweedie_log(1.5),
8423 scale: gam_spec::LikelihoodScaleMetadata::FixedDispersion { phi: 0.25 },
8424 };
8425 let unit = one_obs_resolved(&tweedie_unit, 1.7, 0.8, -0.1);
8426 let scaled = one_obs_resolved(&tweedie_scaled, 1.7, 0.8, -0.1);
8427 for (actual, base) in [
8428 (scaled.score[0], unit.score[0]),
8429 (scaled.fisherweight[0], unit.fisherweight[0]),
8430 (scaled.neghessian_eta[0], unit.neghessian_eta[0]),
8431 (
8432 scaled.neghessian_eta_derivative[0],
8433 unit.neghessian_eta_derivative[0],
8434 ),
8435 (scaled.log_likelihood, unit.log_likelihood),
8436 ] {
8437 let expected = 4.0 * base;
8438 assert!((actual - expected).abs() <= 32.0 * f64::EPSILON * expected.abs().max(1.0));
8439 }
8440 }
8441
8442 #[test]
8443 fn bounded_zero_rows_are_dormant_and_weight_preflight_is_atomic() {
8444 let likelihood = gam_spec::GlmLikelihoodSpec::canonical(LikelihoodSpec::poisson_log());
8445 let dormant = evaluate_resolved_standard_family_observations(
8446 &likelihood,
8447 None,
8448 None,
8449 None,
8450 &array![f64::NAN, 2.0],
8451 &array![0.0, 1.0],
8452 &array![f64::NAN, 0.2],
8453 )
8454 .expect("zero-weight response and predictor are dormant");
8455 assert_eq!(dormant.score[0], 0.0);
8456 assert_eq!(dormant.fisherweight[0], 0.0);
8457
8458 let error = evaluate_resolved_standard_family_observations(
8459 &likelihood,
8460 None,
8461 None,
8462 None,
8463 &array![f64::NAN, 2.0],
8464 &array![1.0, f64::NAN],
8465 &array![f64::NAN, 0.2],
8466 )
8467 .expect_err("later invalid weight must refuse before row evaluation");
8468 assert!(
8473 error.to_string().contains("row 1 has invalid prior weight"),
8474 "unexpected atomic preflight error: {error}"
8475 );
8476 }
8477
8478 fn check_fd(label: &str, spec: &LikelihoodSpec, y: f64, eta: f64) {
8479 let h = 1e-5;
8480 let s0 = one_obs(spec, y, eta);
8481 let sp = one_obs(spec, y, eta + h);
8482 let sm = one_obs(spec, y, eta - h);
8483
8484 let score_fd = (sp.log_likelihood - sm.log_likelihood) / (2.0 * h);
8486 let score = s0.score[0];
8487 assert!(
8488 (score - score_fd).abs() <= 1e-4 * (1.0 + score.abs()),
8489 "{label}: score {score} vs FD {score_fd}"
8490 );
8491
8492 let neghess_fd = -(sp.score[0] - sm.score[0]) / (2.0 * h);
8494 let neghess = s0.neghessian_eta[0];
8495 assert!(
8496 (neghess - neghess_fd).abs() <= 1e-3 * (1.0 + neghess.abs()),
8497 "{label}: neghessian_eta {neghess} vs FD {neghess_fd}"
8498 );
8499
8500 let nhd_fd = (sp.neghessian_eta[0] - sm.neghessian_eta[0]) / (2.0 * h);
8502 let nhd = s0.neghessian_eta_derivative[0];
8503 assert!(
8504 (nhd - nhd_fd).abs() <= 1e-2 * (1.0 + nhd.abs()),
8505 "{label}: neghessian_eta_derivative {nhd} vs FD {nhd_fd}"
8506 );
8507 }
8508
8509 #[test]
8510 fn poisson_gamma_nb_tweedie_arms_match_finite_differences_1615_1616() {
8511 let log = InverseLink::Standard(StandardLink::Log);
8512 let poisson = LikelihoodSpec {
8513 response: ResponseFamily::Poisson,
8514 link: log.clone(),
8515 };
8516 check_fd("poisson y=3", &poisson, 3.0, 0.4);
8517 check_fd("poisson y=0", &poisson, 0.0, -0.2);
8518
8519 let gamma = LikelihoodSpec {
8520 response: ResponseFamily::Gamma,
8521 link: log.clone(),
8522 };
8523 check_fd("gamma y=2.5", &gamma, 2.5, 0.3);
8524 check_fd("gamma y=0.7", &gamma, 0.7, -0.1);
8525
8526 let nb = LikelihoodSpec {
8527 response: ResponseFamily::NegativeBinomial {
8528 theta: 1.5,
8529 theta_fixed: true,
8530 },
8531 link: log.clone(),
8532 };
8533 check_fd("negbin y=4", &nb, 4.0, 0.5);
8534 check_fd("negbin y=0", &nb, 0.0, -0.3);
8535
8536 let tweedie = LikelihoodSpec {
8537 response: ResponseFamily::Tweedie { p: 1.5 },
8538 link: log.clone(),
8539 };
8540 check_fd("tweedie y=2", &tweedie, 2.0, 0.25);
8541 check_fd("tweedie y=0.5", &tweedie, 0.5, -0.15);
8542 }
8543
8544 #[test]
8545 fn binomial_natural_coordinate_towers_match_finite_differences() {
8546 for (label, family, eta) in [
8547 ("logit", LikelihoodSpec::binomial_logit(), 0.7),
8548 ("probit", LikelihoodSpec::binomial_probit(), -1.1),
8549 ("cloglog", LikelihoodSpec::binomial_cloglog(), 0.4),
8550 (
8551 "loglog",
8552 LikelihoodSpec::try_new(
8553 ResponseFamily::Binomial,
8554 InverseLink::Standard(StandardLink::LogLog),
8555 )
8556 .unwrap(),
8557 -0.35,
8558 ),
8559 (
8560 "cauchit",
8561 LikelihoodSpec::try_new(
8562 ResponseFamily::Binomial,
8563 InverseLink::Standard(StandardLink::Cauchit),
8564 )
8565 .unwrap(),
8566 1.25,
8567 ),
8568 ] {
8569 check_fd(label, &family, 0.37, eta);
8570 }
8571 }
8572
8573 #[test]
8574 fn logit_observation_geometry_carries_the_prior_weight_everywhere() {
8575 let eta = 1.75;
8576 let y = 0.3;
8577 let weight = 7.25;
8578 let state = one_obs_weight(&LikelihoodSpec::binomial_logit(), y, weight, eta);
8579 let jet = logit_inverse_link_jet5(eta);
8580 for (got, expected) in [
8581 (state.fisherweight[0], weight * jet.d1),
8582 (state.neghessian_eta[0], weight * jet.d1),
8583 (state.neghessian_eta_derivative[0], weight * jet.d2),
8584 (state.score[0], weight * (y - jet.mu)),
8585 ] {
8586 assert!((got - expected).abs() <= 4.0 * f64::EPSILON * (1.0 + expected.abs()));
8587 }
8588 }
8589
8590 #[test]
8591 fn tiny_positive_and_zero_weights_are_not_projected() {
8592 let tiny = 1e-200;
8593 let logit = one_obs_weight(&LikelihoodSpec::binomial_logit(), 0.4, tiny, 0.0);
8594 let log_argument = tiny.ln() + 0.25_f64.ln();
8608 let roundtrip_tolerance = 0.25 * log_argument.abs() * f64::EPSILON;
8609 let ratio = logit.fisherweight[0] / tiny;
8610 assert!(
8611 (ratio - 0.25).abs() <= roundtrip_tolerance,
8612 "logit Fisher weight at eta=0 must be w/4 up to the log/exp round trip: \
8613 fisherweight/w = {ratio:.17e}, want 0.25 within {roundtrip_tolerance:.3e}"
8614 );
8615 assert!(logit.fisherweight[0] < 1e-190);
8616
8617 let zero = one_obs_weight(&LikelihoodSpec::gaussian_identity(), 3.0, 0.0, -2.0);
8618 assert_eq!(zero.score[0], 0.0);
8619 assert_eq!(zero.fisherweight[0], 0.0);
8620 assert_eq!(zero.neghessian_eta[0], 0.0);
8621 assert_eq!(zero.neghessian_eta_derivative[0], 0.0);
8622 assert_eq!(zero.log_likelihood, 0.0);
8623 assert_eq!(exact_standard_working_response(&zero).unwrap()[0], -2.0);
8624 }
8625
8626 #[test]
8627 fn log_link_tails_balance_tiny_weights_before_certification() {
8628 let poisson = one_obs_weight(&LikelihoodSpec::poisson_log(), 0.0, 1e-300, 700.0);
8629 assert!(poisson.fisherweight[0].is_finite() && poisson.fisherweight[0] > 1.0);
8630 assert!(poisson.score[0].is_finite());
8631 assert!(poisson.log_likelihood.is_finite());
8632
8633 let gamma = one_obs_weight(&LikelihoodSpec::gamma_log(), 1.0, 1e-300, -700.0);
8634 assert!(gamma.neghessian_eta[0].is_finite() && gamma.neghessian_eta[0] > 1.0);
8635 assert!(gamma.score[0].is_finite());
8636 assert!(gamma.log_likelihood.is_finite());
8637 }
8638
8639 #[test]
8640 fn invalid_weights_are_refused_in_row_order_and_zero_weight_rows_are_excluded() {
8641 let family = LikelihoodSpec::gaussian_identity();
8642 let y = array![1.0, 2.0];
8643 let eta = array![0.0, 0.0];
8644 for weights in [array![-1.0, 1.0], array![f64::NAN, 1.0]] {
8645 let err = evaluate_standard_familyobservations(
8646 family.clone(),
8647 None,
8648 None,
8649 None,
8650 &y,
8651 &weights,
8652 &eta,
8653 )
8654 .expect_err("invalid prior weight must be refused");
8655 assert!(err.to_string().contains("row 0"), "{err}");
8656 }
8657
8658 let err = evaluate_standard_familyobservations(
8669 family.clone(),
8670 None,
8671 None,
8672 None,
8673 &array![f64::NAN],
8674 &array![1.0],
8675 &array![0.0],
8676 )
8677 .expect_err("a non-finite response on a weight-bearing row must be refused");
8678 assert!(err.to_string().contains("row 0"), "{err}");
8679
8680 let excluded = evaluate_standard_familyobservations(
8681 family,
8682 None,
8683 None,
8684 None,
8685 &array![f64::NAN],
8686 &array![0.0],
8687 &array![0.0],
8688 )
8689 .expect("a zero-weight row is excluded, so its response is never inspected");
8690 assert_eq!(excluded.score[0], 0.0);
8691 assert_eq!(excluded.fisherweight[0], 0.0);
8692 assert_eq!(excluded.neghessian_eta[0], 0.0);
8693 assert_eq!(excluded.log_likelihood, 0.0);
8694 }
8695
8696 #[test]
8697 fn unrepresentable_cloglog_curvature_is_refused_without_a_floor() {
8698 let err = evaluate_standard_familyobservations(
8699 LikelihoodSpec::binomial_cloglog(),
8700 None,
8701 None,
8702 None,
8703 &array![1.0],
8704 &array![1.0],
8705 &array![18.0],
8706 )
8707 .expect_err("mathematically sub-f64 Fisher information must be refused");
8708 assert!(err.to_string().contains("Fisher weight"), "{err}");
8709 }
8710
8711 #[test]
8712 fn bounded_covariance_requires_a_certified_strict_spd_precision() {
8713 let covariance = certified_bounded_posterior_covariance(
8714 &array![[4.0, 1.0], [1.0, 3.0]],
8715 "bounded covariance regression",
8716 )
8717 .expect("strict SPD precision");
8718 assert!((covariance[[0, 0]] - 3.0 / 11.0).abs() < 1e-14);
8719 assert!((covariance[[0, 1]] + 1.0 / 11.0).abs() < 1e-14);
8720 assert!((covariance[[1, 1]] - 4.0 / 11.0).abs() < 1e-14);
8721
8722 for invalid in [
8723 array![[1.0, 1.0], [1.0, 1.0]],
8724 array![[1.0, 2.0], [2.0, 1.0]],
8725 ] {
8726 assert!(
8727 certified_bounded_posterior_covariance(
8728 &invalid,
8729 "invalid bounded covariance regression"
8730 )
8731 .is_err(),
8732 "singular/indefinite precision must not become a pseudo-covariance"
8733 );
8734 }
8735 }
8736}