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, y, weights);
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 _ => {}
1248 }
1249 }
1250 let (
1263 Some(mass_local),
1264 Some(tension_local),
1265 Some(stiffness_local),
1266 Some(mass_scale),
1267 Some(tension_scale),
1268 Some(stiffness_scale),
1269 ) = (
1270 mass_local_idx,
1271 tension_local_idx,
1272 stiffness_local_idx,
1273 mass_norm,
1274 tension_norm,
1275 stiffness_norm,
1276 )
1277 else {
1278 continue;
1279 };
1280 let mass_global_idx = global_base_idx + mass_local;
1281 let tension_global_idx = global_base_idx + tension_local;
1282 let stiffness_global_idx = global_base_idx + stiffness_local;
1283
1284 let (feature_cols, mut d0, mut d1, mut d2, collocation_points, dim, center_mass_rows) =
1285 match (&termspec.basis, &term_fit.metadata) {
1286 (
1287 SmoothBasisSpec::Matern { feature_cols, .. },
1288 BasisMetadata::Matern {
1289 centers,
1290 length_scale,
1291 nu,
1292 include_intercept,
1293 identifiability_transform,
1294 aniso_log_scales,
1295 input_scale,
1296 ..
1297 },
1298 ) => {
1299 let collocation_length_scale =
1305 input_scale.to_standardized_units(*length_scale);
1306 let ops = build_matern_collocation_operator_matrices(
1307 centers.view(),
1308 None,
1309 collocation_length_scale,
1310 *nu,
1311 *include_intercept,
1312 identifiability_transform.as_ref().map(|z| z.view()),
1313 aniso_log_scales.as_deref(),
1314 )?;
1315 (
1316 feature_cols.clone(),
1317 ops.d0,
1318 ops.d1,
1319 ops.d2,
1320 ops.collocation_points,
1321 centers.ncols(),
1322 false,
1323 )
1324 }
1325 (
1326 SmoothBasisSpec::Duchon { feature_cols, .. },
1327 BasisMetadata::Duchon {
1328 centers,
1329 length_scale,
1330 power,
1331 nullspace_order,
1332 identifiability_transform,
1333 input_scale,
1334 aniso_log_scales,
1335 operator_collocation_points: Some(collocation_points),
1336 radial_reparam,
1337 ..
1338 },
1339 ) => {
1340 let collocation_length_scale = (*length_scale)
1341 .map(|length| input_scale.to_standardized_units(length));
1342 let ops =
1343 gam_terms::basis::build_duchon_collocation_operator_matriceswithworkspace(
1344 centers.view(),
1345 collocation_points.view(),
1346 None,
1347 collocation_length_scale,
1348 *power,
1349 *nullspace_order,
1350 aniso_log_scales.as_deref(),
1351 identifiability_transform.as_ref().map(|z| z.view()),
1352 2,
1353 radial_reparam.as_ref().map(|v| v.view()),
1354 &mut BasisWorkspace::default(),
1355 )?;
1356 (
1357 feature_cols.clone(),
1358 ops.d0,
1359 ops.d1,
1360 ops.d2,
1361 ops.collocation_points,
1362 centers.ncols(),
1363 true,
1364 )
1365 }
1366 _ => continue,
1367 };
1368 if center_mass_rows && d0.nrows() > 0 && d0.ncols() > 0 {
1369 let means = d0.sum_axis(Axis(0)).mapv(|v| v / d0.nrows() as f64);
1370 for mut row in d0.rows_mut() {
1371 row -= &means;
1372 }
1373 }
1374
1375 let mass_scale = mass_scale.max(1e-12).sqrt();
1393 let tension_scale = tension_scale.max(1e-12).sqrt();
1394 let stiffness_scale = stiffness_scale.max(1e-12).sqrt();
1395 d0.mapv_inplace(|v| v / mass_scale);
1396 d1.mapv_inplace(|v| v / tension_scale);
1397 d2.mapv_inplace(|v| v / stiffness_scale);
1398
1399 let coeff_global_range =
1400 (smooth_start + term_fit.coeff_range.start)..(smooth_start + term_fit.coeff_range.end);
1401 if d0.ncols() != coeff_global_range.len()
1402 || d1.ncols() != coeff_global_range.len()
1403 || d2.ncols() != coeff_global_range.len()
1404 {
1405 crate::bail_invalid_estim!(
1406 "spatial operator dimension mismatch for term '{}': D0 cols={}, D1 cols={}, D2 cols={}, coeffs={}",
1407 term_fit.name,
1408 d0.ncols(),
1409 d1.ncols(),
1410 d2.ncols(),
1411 coeff_global_range.len()
1412 );
1413 }
1414 out.push(SpatialOperatorRuntimeCache {
1415 termname: term_fit.name.clone(),
1416 feature_cols,
1417 coeff_global_range,
1418 mass_penalty_global_idx: mass_global_idx,
1419 tension_penalty_global_idx: tension_global_idx,
1420 stiffness_penalty_global_idx: stiffness_global_idx,
1421 d0,
1422 d1,
1423 d2,
1424 collocation_points,
1425 dimension: dim,
1426 });
1427 }
1428 Ok(out)
1429}
1430
1431fn scalar_operator_response_variance(
1443 operator: &Array2<f64>,
1444 cov_local: &Array2<f64>,
1445) -> Array1<f64> {
1446 Array1::from_iter(operator.rows().into_iter().map(|row| {
1447 let s = cov_local.dot(&row);
1448 row.dot(&s).max(0.0)
1449 }))
1450}
1451
1452fn grouped_operator_response_variance(
1463 operator: &Array2<f64>,
1464 block_dim: usize,
1465 cov_local: &Array2<f64>,
1466) -> Result<Array1<f64>, EstimationError> {
1467 if block_dim == 0 || !operator.nrows().is_multiple_of(block_dim) {
1468 crate::bail_invalid_estim!(
1469 "grouped variance row layout invalid: rows={}, block_dim={block_dim}",
1470 operator.nrows()
1471 );
1472 }
1473 let p = operator.nrows() / block_dim;
1474 let mut out = Array1::<f64>::zeros(p);
1475 for k in 0..p {
1476 let mut acc = 0.0;
1477 for axis in 0..block_dim {
1478 let row = operator.row(k * block_dim + axis);
1479 let s = cov_local.dot(&row);
1480 acc += row.dot(&s);
1481 }
1482 out[k] = acc.max(0.0);
1483 }
1484 Ok(out)
1485}
1486
1487fn compute_spatial_adaptiveweights_for_beta(
1488 beta: &Array1<f64>,
1489 caches: &[SpatialOperatorRuntimeCache],
1490 epsilon_0: f64,
1491 epsilon_g: f64,
1492 epsilon_c: f64,
1493 weight_floor: f64,
1494 weight_ceiling: f64,
1495 beta_covariance: Option<&Array2<f64>>,
1496) -> Result<Vec<SpatialAdaptiveWeights>, EstimationError> {
1497 caches
1529 .iter()
1530 .map(|cache| {
1531 let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
1532 let exact = SpatialPenaltyExactState::from_beta_local(
1533 beta_local,
1534 cache,
1535 [epsilon_0, epsilon_g, epsilon_c],
1536 )?;
1537 let cov_local = beta_covariance.map(|cov| {
1538 cov.slice(s![
1539 cache.coeff_global_range.clone(),
1540 cache.coeff_global_range.clone()
1541 ])
1542 .to_owned()
1543 });
1544 let dim = cache.dimension;
1545 let (var_0, var_g, var_c) = match cov_local.as_ref() {
1546 Some(cov) => (
1547 scalar_operator_response_variance(&cache.d0, cov),
1548 grouped_operator_response_variance(&cache.d1, dim, cov)?,
1549 grouped_operator_response_variance(&cache.d2, dim * dim, cov)?,
1550 ),
1551 None => (
1552 Array1::<f64>::zeros(exact.magnitude.signal.len()),
1553 Array1::<f64>::zeros(exact.gradient.norm.len()),
1554 Array1::<f64>::zeros(exact.curvature.norm.len()),
1555 ),
1556 };
1557 let (_, inv_0) = exact.magnitude.surrogateweights_posterior_snr(
1558 &var_0,
1559 weight_floor,
1560 weight_ceiling,
1561 );
1562 let (_, inv_g) =
1563 exact
1564 .gradient
1565 .surrogateweights_posterior_snr(&var_g, weight_floor, weight_ceiling);
1566 let (_, inv_c) = exact.curvature.surrogateweights_posterior_snr(
1567 &var_c,
1568 weight_floor,
1569 weight_ceiling,
1570 );
1571 Ok(SpatialAdaptiveWeights {
1572 inv_magweight: inv_0,
1573 invgradweight: inv_g,
1574 inv_lapweight: inv_c,
1575 })
1576 })
1577 .collect()
1578}
1579
1580fn compute_initial_epsilons(
1581 beta: &Array1<f64>,
1582 caches: &[SpatialOperatorRuntimeCache],
1583 min_epsilon: f64,
1584) -> Result<(f64, f64, f64), EstimationError> {
1585 let mut fvals = Vec::<f64>::new();
1586 let mut gvals = Vec::<f64>::new();
1587 let mut cvals = Vec::<f64>::new();
1588 for cache in caches {
1589 let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
1590 let exact = SpatialPenaltyExactState::from_beta_local(
1591 beta_local,
1592 cache,
1593 [min_epsilon, min_epsilon, min_epsilon],
1594 )?;
1595 let (f, g, c) = exact.absolute_collocation_magnitudes();
1596 fvals.extend(f.iter().copied());
1597 gvals.extend(g.iter().copied());
1598 cvals.extend(c.iter().copied());
1599 }
1600 let eps_0 = robust_epsilon_from_samples(&fvals, min_epsilon);
1606 let eps_g = robust_epsilon_from_samples(&gvals, min_epsilon);
1607 let eps_c = robust_epsilon_from_samples(&cvals, min_epsilon);
1608 Ok((eps_0, eps_g, eps_c))
1609}
1610
1611fn exact_spatial_adaptive_penalty_index_set(
1612 caches: &[SpatialOperatorRuntimeCache],
1613) -> BTreeSet<usize> {
1614 let mut out = BTreeSet::new();
1615 for cache in caches {
1616 out.insert(cache.mass_penalty_global_idx);
1617 out.insert(cache.tension_penalty_global_idx);
1618 out.insert(cache.stiffness_penalty_global_idx);
1619 }
1620 out
1621}
1622
1623fn checked_fit_log_lambdas(
1624 lambdas: &Array1<f64>,
1625 context: &str,
1626) -> Result<Array1<f64>, EstimationError> {
1627 let values = lambdas
1628 .iter()
1629 .copied()
1630 .enumerate()
1631 .map(|(coordinate, lambda)| {
1632 gam_problem::checked_log_strength(lambda).map_err(|error| {
1633 EstimationError::InvalidInput(format!(
1634 "{context} lambda coordinate {coordinate} is outside the canonical physical-strength domain: {error}"
1635 ))
1636 })
1637 })
1638 .collect::<Result<Vec<_>, _>>()?;
1639 Ok(Array1::from_vec(values))
1640}
1641
1642fn build_spatial_adaptive_hyperspecs(cache_count: usize) -> Vec<SpatialAdaptiveHyperSpec> {
1643 let mut out = Vec::with_capacity(cache_count * 3 + 3);
1644 for cache_index in 0..cache_count {
1645 out.push(SpatialAdaptiveHyperSpec {
1646 cache_index,
1647 kind: SpatialAdaptiveHyperKind::LogLambdaMagnitude,
1648 });
1649 out.push(SpatialAdaptiveHyperSpec {
1650 cache_index,
1651 kind: SpatialAdaptiveHyperKind::LogLambdaGradient,
1652 });
1653 out.push(SpatialAdaptiveHyperSpec {
1654 cache_index,
1655 kind: SpatialAdaptiveHyperKind::LogLambdaCurvature,
1656 });
1657 }
1658 out.push(SpatialAdaptiveHyperSpec {
1659 cache_index: 0,
1660 kind: SpatialAdaptiveHyperKind::LogEpsilonMagnitude,
1661 });
1662 out.push(SpatialAdaptiveHyperSpec {
1663 cache_index: 0,
1664 kind: SpatialAdaptiveHyperKind::LogEpsilonGradient,
1665 });
1666 out.push(SpatialAdaptiveHyperSpec {
1667 cache_index: 0,
1668 kind: SpatialAdaptiveHyperKind::LogEpsilonCurvature,
1669 });
1670 out
1671}
1672
1673fn penalty_matrixwith_local_block(
1674 total_dim: usize,
1675 coeff_range: Range<usize>,
1676 local: &Array2<f64>,
1677) -> Array2<f64> {
1678 let mut out = Array2::<f64>::zeros((total_dim, total_dim));
1679 out.slice_mut(s![coeff_range.clone(), coeff_range])
1680 .assign(local);
1681 out
1682}
1683
1684fn fit_term_collectionwith_exact_spatial_adaptive_regularization(
1685 baseline: FittedTermCollection,
1686 y: ArrayView1<'_, f64>,
1687 weights: ArrayView1<'_, f64>,
1688 offset: ArrayView1<'_, f64>,
1689 family: LikelihoodSpec,
1690 options: &FitOptions,
1691 runtime_caches: &[SpatialOperatorRuntimeCache],
1692) -> Result<FittedTermCollection, EstimationError> {
1693 let adaptive_opts = options.adaptive_regularization.clone().unwrap_or_default();
1722 let adaptive_penalty_indices = exact_spatial_adaptive_penalty_index_set(runtime_caches);
1723 let p_total = baseline.design.design.ncols();
1724 if baseline.fit.lambdas.len() != baseline.design.penalties.len() {
1725 crate::bail_invalid_estim!(
1726 "exact spatial adaptive fit received {} baseline lambdas for {} penalties",
1727 baseline.fit.lambdas.len(),
1728 baseline.design.penalties.len(),
1729 );
1730 }
1731 let baseline_log_lambdas =
1732 checked_fit_log_lambdas(&baseline.fit.lambdas, "exact spatial adaptive baseline")?;
1733 for (cache_idx, cache) in runtime_caches.iter().enumerate() {
1734 for (operator, penalty_idx) in [
1735 ("mass", cache.mass_penalty_global_idx),
1736 ("tension", cache.tension_penalty_global_idx),
1737 ("stiffness", cache.stiffness_penalty_global_idx),
1738 ] {
1739 if penalty_idx >= baseline.fit.lambdas.len() {
1740 crate::bail_invalid_estim!(
1741 "exact spatial adaptive cache {cache_idx} {operator} penalty index {penalty_idx} is out of bounds for {} baseline lambdas",
1742 baseline.fit.lambdas.len(),
1743 );
1744 }
1745 }
1746 }
1747 struct RetainedPenaltySetup {
1748 global_idx: usize,
1749 global_penalty: Array2<f64>,
1750 nullspace_dim: usize,
1751 log_lambda: f64,
1752 }
1753 use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
1754 let retained_setups = baseline
1755 .design
1756 .penalties
1757 .par_iter()
1758 .enumerate()
1759 .map(|(idx, bp)| {
1760 if adaptive_penalty_indices.contains(&idx) {
1761 return None;
1762 }
1763 Some(RetainedPenaltySetup {
1764 global_idx: idx,
1765 global_penalty: bp.to_global(p_total),
1766 nullspace_dim: baseline
1767 .design
1768 .nullspace_dims
1769 .get(idx)
1770 .copied()
1771 .unwrap_or(0),
1772 log_lambda: baseline_log_lambdas[idx],
1773 })
1774 })
1775 .collect::<Vec<_>>();
1776 let retained_count = retained_setups
1777 .iter()
1778 .filter(|setup| setup.is_some())
1779 .count();
1780 let mut retained_penalties = Vec::<Array2<f64>>::with_capacity(retained_count);
1781 let mut retained_nullspace_dims = Vec::<usize>::with_capacity(retained_count);
1782 let mut retained_log_lambdas = Vec::<f64>::with_capacity(retained_count);
1783 let mut retained_global_indices = Vec::<usize>::with_capacity(retained_count);
1784 for setup in retained_setups.into_iter().flatten() {
1785 retained_penalties.push(setup.global_penalty);
1786 retained_nullspace_dims.push(setup.nullspace_dim);
1787 retained_log_lambdas.push(setup.log_lambda);
1788 retained_global_indices.push(setup.global_idx);
1789 }
1790
1791 let (eps_0_init, eps_g_init, eps_c_init) = compute_initial_epsilons(
1792 &baseline.fit.beta,
1793 runtime_caches,
1794 adaptive_opts.min_epsilon,
1795 )?;
1796 let mut initial_theta =
1797 Array1::<f64>::zeros(retained_penalties.len() + runtime_caches.len() * 3 + 3);
1798 for (idx, value) in retained_log_lambdas.iter().enumerate() {
1799 initial_theta[idx] = *value;
1800 }
1801 let adaptive_log_lambda_components = runtime_caches
1802 .par_iter()
1803 .map(|cache| {
1804 [
1805 baseline_log_lambdas[cache.mass_penalty_global_idx],
1806 baseline_log_lambdas[cache.tension_penalty_global_idx],
1807 baseline_log_lambdas[cache.stiffness_penalty_global_idx],
1808 ]
1809 })
1810 .collect::<Vec<_>>();
1811 let mut at = retained_penalties.len();
1812 for logs in &adaptive_log_lambda_components {
1813 initial_theta[at] = logs[0];
1814 initial_theta[at + 1] = logs[1];
1815 initial_theta[at + 2] = logs[2];
1816 at += 3;
1817 }
1818 let minimum_log_epsilon = gam_problem::checked_log_strength(adaptive_opts.min_epsilon)
1819 .map_err(|error| {
1820 EstimationError::InvalidInput(format!(
1821 "adaptive minimum epsilon is outside the canonical positive-strength domain: {error}"
1822 ))
1823 })?;
1824 for (slot, epsilon) in [eps_0_init, eps_g_init, eps_c_init].into_iter().enumerate() {
1825 initial_theta[at + slot] =
1826 gam_problem::checked_log_strength(epsilon.max(adaptive_opts.min_epsilon)).map_err(
1827 |error| {
1828 EstimationError::InvalidInput(format!(
1829 "adaptive initial epsilon coordinate {slot} is outside the canonical positive-strength domain: {error}"
1830 ))
1831 },
1832 )?;
1833 }
1834
1835 let hyperspecs = build_spatial_adaptive_hyperspecs(runtime_caches.len());
1836 let zero_psi_op: std::sync::Arc<dyn gam_custom_family::CustomFamilyPsiDerivativeOperator> =
1837 std::sync::Arc::new(gam_custom_family::ZeroPsiDerivativeOperator::new(
1838 baseline.design.design.nrows(),
1839 baseline.design.design.ncols(),
1840 ));
1841 let derivative_blocks = vec![
1842 hyperspecs
1843 .par_iter()
1844 .map(|_| CustomFamilyBlockPsiDerivative {
1845 penalty_index: None,
1846 x_psi: Array2::<f64>::zeros((0, 0)),
1847 s_psi: Array2::<f64>::zeros((0, 0)),
1848 s_psi_components: None,
1849 s_psi_penalty_components: None,
1850 x_psi_psi: None,
1851 s_psi_psi: None,
1852 s_psi_psi_components: None,
1853 s_psi_psi_penalty_components: None,
1854 implicit_operator: Some(std::sync::Arc::clone(&zero_psi_op)),
1855 implicit_axis: 0,
1856 implicit_group_id: None,
1857 })
1858 .collect::<Vec<_>>(),
1859 ];
1860
1861 let mixture_link_state = options
1862 .mixture_link
1863 .clone()
1864 .as_ref()
1865 .map(state_fromspec)
1866 .transpose()
1867 .map_err(EstimationError::InvalidInput)?;
1868 let sas_link_state = options
1869 .sas_link
1870 .map(|spec| {
1871 if family.is_binomial_beta_logistic() {
1872 state_from_beta_logisticspec(spec)
1873 } else {
1874 state_from_sasspec(spec)
1875 }
1876 })
1877 .transpose()
1878 .map_err(EstimationError::InvalidInput)?;
1879 let latent_cloglog_state = options.latent_cloglog;
1880 let shared_y = Arc::new(y.to_owned());
1881 let sharedweights = Arc::new(weights.to_owned());
1882 let shared_design = baseline
1883 .design
1884 .design
1885 .try_to_dense_arc("spatial adaptive exact hyperfit design")
1886 .map_err(EstimationError::InvalidInput)?;
1887 let shared_offset = Arc::new(offset.to_owned());
1888 let shared_runtime_caches = Arc::new(runtime_caches.to_vec());
1889 let shared_hyperspecs = Arc::new(hyperspecs.clone());
1890 let zero_quadratic = ValidatedFixedQuadraticHessian::zero(
1891 baseline.design.design.ncols(),
1892 )
1893 .map_err(EstimationError::InvalidInput)?;
1894 let base_family = SpatialAdaptiveExactFamily {
1895 family: family.clone(),
1896 latent_cloglog_state,
1897 mixture_link_state: mixture_link_state.clone(),
1898 sas_link_state,
1899 y: shared_y.clone(),
1900 weights: sharedweights.clone(),
1901 design: shared_design.clone(),
1902 offset: shared_offset.clone(),
1903 linear_constraints: baseline.design.linear_constraints.clone(),
1904 runtime_caches: shared_runtime_caches.clone(),
1905 adaptive_params: Vec::new(),
1906 fixed_quadratic_hessian: zero_quadratic.clone(),
1907 hyperspecs: shared_hyperspecs.clone(),
1908 exact_eval_cache: Arc::new(Mutex::new(None)),
1909 };
1910
1911 let rho_dim = retained_penalties.len();
1912 let operator_slots_end = rho_dim + runtime_caches.len() * 3;
1913 const UNIFIED_LOG_WINDOW: f64 = 6.0;
1923 const RETAINED_LAMBDA_LOG_LOWER_FLOOR: f64 = -30.0;
1924 const RETAINED_LAMBDA_LOG_UPPER_CAP: f64 = 30.0;
1925 const OPERATOR_LAMBDA_LOG_LOWER_FLOOR: f64 = -10.0;
1926 const OPERATOR_LAMBDA_LOG_UPPER_CAP: f64 = 30.0;
1927 let epsilon_floor_log = minimum_log_epsilon;
1928 let anchored_bound = |idx: usize, sign: f64| -> f64 {
1929 let raw = initial_theta[idx] + sign * UNIFIED_LOG_WINDOW;
1930 if idx < rho_dim {
1931 raw.clamp(
1932 RETAINED_LAMBDA_LOG_LOWER_FLOOR,
1933 RETAINED_LAMBDA_LOG_UPPER_CAP,
1934 )
1935 } else if idx < operator_slots_end {
1936 raw.clamp(
1937 OPERATOR_LAMBDA_LOG_LOWER_FLOOR,
1938 OPERATOR_LAMBDA_LOG_UPPER_CAP,
1939 )
1940 } else {
1941 raw.clamp(epsilon_floor_log, gam_problem::LOG_STRENGTH_MAX)
1942 }
1943 };
1944 let eps_lower =
1945 Array1::from_iter((0..initial_theta.len()).map(|idx| anchored_bound(idx, -1.0)));
1946 let eps_upper = Array1::from_iter((0..initial_theta.len()).map(|idx| anchored_bound(idx, 1.0)));
1947 let blockspec = ParameterBlockSpec {
1948 name: "eta".to_string(),
1949 design: baseline.design.design.clone(),
1950 offset: offset.to_owned(),
1951 penalties: retained_penalties
1952 .iter()
1953 .cloned()
1954 .map(PenaltyMatrix::Dense)
1955 .collect(),
1956 nullspace_dims: retained_nullspace_dims.clone(),
1957 initial_log_lambdas: Array1::from_vec(retained_log_lambdas.clone()),
1958 initial_beta: Some(baseline.fit.beta.clone()),
1959 gauge_priority: 100,
1960 jacobian_callback: None,
1961 stacked_design: None,
1962 stacked_offset: None,
1963 };
1964 let screening_cap = Arc::new(AtomicUsize::new(0));
1965 let outer_opts = BlockwiseFitOptions {
1966 inner_max_cycles: options.max_iter,
1967 inner_tol: options.tol,
1968 outer_max_iter: options.max_iter,
1969 outer_tol: options.tol,
1970 compute_covariance: false,
1971 screening_max_inner_iterations: Some(Arc::clone(&screening_cap)),
1972 ..BlockwiseFitOptions::default()
1973 };
1974
1975 use gam_problem::{DeclaredHessianForm, Derivative, HessianValue, OuterEval};
1976 use gam_solve::rho_optimizer::OuterProblem;
1977
1978 struct SpatialAdaptiveOuterState {
1979 warm_cache: Option<CustomFamilyWarmStart>,
1980 terminal_mode: Option<(Array1<f64>, f64, CustomFamilyOwnedMode)>,
1981 last_eval: Option<(
1982 Array1<f64>,
1983 f64,
1984 Array1<f64>,
1985 HessianValue,
1986 CustomFamilyWarmStart,
1987 )>,
1988 }
1989
1990 struct DecodedSpatialAdaptiveTheta {
1991 rho: Array1<f64>,
1992 retained_lambdas: Array1<f64>,
1993 adaptive_params: Vec<SpatialAdaptiveTermHyperParams>,
1994 epsilon: [f64; 3],
1995 }
1996
1997 let n_theta = initial_theta.len();
1998
1999 let theta_bounds = Some((eps_lower.clone(), eps_upper.clone()));
2000
2001 let decode_theta =
2002 |theta: &Array1<f64>| -> Result<DecodedSpatialAdaptiveTheta, EstimationError> {
2003 let physical = gam_problem::checked_exp_log_strengths(theta.iter().copied()).map_err(
2004 |error| {
2005 EstimationError::InvalidInput(format!(
2006 "exact spatial adaptive outer coordinate is outside the canonical log-strength domain: {error}"
2007 ))
2008 },
2009 )?;
2010 let rho = theta.slice(s![..rho_dim]).to_owned();
2011 let retained_lambdas = Array1::from_vec(physical[..rho_dim].to_vec());
2012 let adaptive_lambda_start = rho_dim;
2013 let adaptive_lambda_end = adaptive_lambda_start + runtime_caches.len() * 3;
2014 let eps = [
2015 physical[adaptive_lambda_end],
2016 physical[adaptive_lambda_end + 1],
2017 physical[adaptive_lambda_end + 2],
2018 ];
2019 let adaptive_params = runtime_caches
2020 .iter()
2021 .enumerate()
2022 .map(|(cache_idx, _)| SpatialAdaptiveTermHyperParams {
2023 lambda: [
2024 physical[adaptive_lambda_start + cache_idx * 3],
2025 physical[adaptive_lambda_start + cache_idx * 3 + 1],
2026 physical[adaptive_lambda_start + cache_idx * 3 + 2],
2027 ],
2028 epsilon: eps,
2029 })
2030 .collect::<Vec<_>>();
2031 Ok(DecodedSpatialAdaptiveTheta {
2032 rho,
2033 retained_lambdas,
2034 adaptive_params,
2035 epsilon: eps,
2036 })
2037 };
2038 let clamp_theta = |theta: &Array1<f64>| -> Array1<f64> {
2044 Array1::from_shape_fn(theta.len(), |i| theta[i].clamp(eps_lower[i], eps_upper[i]))
2045 };
2046 let realize_hyper_layout = |theta: &Array1<f64>| {
2047 gam_custom_family::CustomFamilyHyperLayout::new(
2048 derivative_blocks.clone(),
2049 Vec::new(),
2050 theta.slice(s![rho_dim..]).to_owned(),
2051 )
2052 .map_err(EstimationError::InvalidInput)
2053 };
2054 let analytic_outer_hessian_available =
2055 gam_custom_family::joint_exact_analytic_outer_hessian_available()
2056 && base_family
2057 .exact_outer_derivative_order(std::slice::from_ref(&blockspec), &outer_opts)
2058 .has_hessian()
2059 && gam_custom_family::exact_newton_outer_geometry_supports_second_order_solver(
2060 &base_family,
2061 );
2062 let problem = OuterProblem::new(n_theta)
2068 .with_gradient(Derivative::Analytic)
2069 .with_hessian(if analytic_outer_hessian_available {
2070 DeclaredHessianForm::Either
2071 } else {
2072 DeclaredHessianForm::Unavailable
2073 })
2074 .with_fallback_policy(gam_solve::rho_optimizer::FallbackPolicy::Disabled)
2075 .with_psi_dim(n_theta.saturating_sub(rho_dim))
2076 .with_tolerance(options.tol)
2077 .with_rel_cost_tolerance(Some(options.tol))
2082 .with_max_iter(options.max_iter)
2083 .with_seed_config(gam_problem::SeedConfig::default())
2084 .with_screening_cap(Arc::clone(&screening_cap))
2085 .with_initial_rho(initial_theta.clone());
2086 let problem = if let Some((lo, hi)) = theta_bounds {
2087 problem.with_bounds(lo, hi)
2088 } else {
2089 problem
2090 };
2091
2092 let eval_outer = |st: &mut SpatialAdaptiveOuterState,
2093 theta: &Array1<f64>,
2094 order: gam_solve::rho_optimizer::OuterEvalOrder|
2095 -> Result<OuterEval, EstimationError> {
2096 let decoded = decode_theta(theta)?;
2097
2098 if let Some((cached_theta, cached_cost, cached_grad, cached_hess, cached_warm)) =
2099 &st.last_eval
2100 && cached_theta.len() == theta.len()
2101 && cached_theta
2102 .iter()
2103 .zip(theta.iter())
2104 .all(|(&a, &b)| a.to_bits() == b.to_bits())
2105 && st
2106 .terminal_mode
2107 .as_ref()
2108 .is_some_and(|(mode_theta, mode_objective, _)| {
2109 mode_theta.len() == theta.len()
2110 && mode_theta
2111 .iter()
2112 .zip(theta.iter())
2113 .all(|(&a, &b)| a.to_bits() == b.to_bits())
2114 && mode_objective.to_bits() == cached_cost.to_bits()
2115 })
2116 && (!matches!(
2117 order,
2118 gam_solve::rho_optimizer::OuterEvalOrder::ValueGradientHessian
2119 ) || analytic_outer_hessian_available)
2120 {
2121 st.warm_cache = Some(cached_warm.clone());
2122 return Ok(OuterEval {
2123 cost: *cached_cost,
2124 gradient: cached_grad.clone(),
2125 hessian: if matches!(
2126 order,
2127 gam_solve::rho_optimizer::OuterEvalOrder::ValueGradientHessian
2128 ) && analytic_outer_hessian_available
2129 {
2130 cached_hess.clone()
2131 } else {
2132 HessianValue::Unavailable
2133 },
2134 inner_beta_hint: None,
2135 });
2136 }
2137
2138 let family_eval =
2139 base_family.with_adaptive_params(decoded.adaptive_params, zero_quadratic.clone());
2140 let hyper_layout = realize_hyper_layout(theta)?;
2141 let need_hessian = matches!(
2142 order,
2143 gam_solve::rho_optimizer::OuterEvalOrder::ValueGradientHessian
2144 ) && analytic_outer_hessian_available;
2145 let owned = evaluate_custom_family_joint_hyper_owned(
2146 &family_eval,
2147 std::slice::from_ref(&blockspec),
2148 &outer_opts,
2149 &decoded.rho,
2150 &hyper_layout,
2151 st.warm_cache.as_ref(),
2152 if need_hessian {
2153 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueGradientHessian
2154 } else {
2155 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueAndGradient
2156 },
2157 )
2158 .map_err(|e| {
2159 EstimationError::RemlOptimizationFailed(format!("spatial adaptive eval failed: {e}"))
2160 })?;
2161 if !owned.result.inner_converged {
2162 st.warm_cache = Some(owned.result.warm_start.clone());
2163 return Err(EstimationError::RemlOptimizationFailed(
2164 "exact spatial adaptive inner solve did not converge".to_string(),
2165 ));
2166 }
2167 if !owned.result.objective.is_finite()
2168 || owned.result.gradient.iter().any(|v| !v.is_finite())
2169 {
2170 return Err(EstimationError::RemlOptimizationFailed(
2171 "exact spatial adaptive objective returned non-finite values".to_string(),
2172 ));
2173 }
2174 let hessian_result = if need_hessian {
2175 if !owned.result.outer_hessian.is_analytic() {
2176 return Err(EstimationError::RemlOptimizationFailed(
2177 "exact spatial adaptive objective did not return an exact outer Hessian"
2178 .to_string(),
2179 ));
2180 }
2181 match owned.result.outer_hessian.dim() {
2182 Some(dim) if dim == theta.len() => {}
2183 Some(dim) => {
2184 return Err(EstimationError::RemlOptimizationFailed(format!(
2185 "exact spatial adaptive outer Hessian dimension mismatch: got {dim}, expected {}",
2186 theta.len(),
2187 )));
2188 }
2189 None => {
2190 return Err(EstimationError::RemlOptimizationFailed(
2191 "exact spatial adaptive objective did not report an outer Hessian dimension"
2192 .to_string(),
2193 ));
2194 }
2195 }
2196 st.last_eval = Some((
2197 theta.to_owned(),
2198 owned.result.objective,
2199 owned.result.gradient.clone(),
2200 owned.result.outer_hessian.clone(),
2201 owned.result.warm_start.clone(),
2202 ));
2203 owned.result.outer_hessian
2204 } else {
2205 HessianValue::Unavailable
2206 };
2207 let objective = owned.result.objective;
2208 let gradient = owned.result.gradient;
2209 st.warm_cache = Some(owned.result.warm_start);
2210 st.terminal_mode = Some((theta.to_owned(), objective, owned.mode));
2211 Ok(OuterEval {
2212 cost: objective,
2213 gradient,
2214 hessian: hessian_result,
2215 inner_beta_hint: None,
2216 })
2217 };
2218
2219 let mut obj = problem.build_objective_with_screening_proxy(
2220 SpatialAdaptiveOuterState {
2221 warm_cache: None,
2222 terminal_mode: None,
2223 last_eval: None,
2224 },
2225 |st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
2226 let theta = clamp_theta(theta);
2227 let DecodedSpatialAdaptiveTheta {
2228 rho,
2229 adaptive_params,
2230 ..
2231 } = decode_theta(&theta)?;
2232 let family_eval =
2233 base_family.with_adaptive_params(adaptive_params, zero_quadratic.clone());
2234 let hyper_layout = realize_hyper_layout(&theta)?;
2235 let owned = evaluate_custom_family_joint_hyper_owned(
2236 &family_eval,
2237 std::slice::from_ref(&blockspec),
2238 &outer_opts,
2239 &rho,
2240 &hyper_layout,
2241 st.warm_cache.as_ref(),
2242 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueOnly,
2243 )
2244 .map_err(|e| {
2245 EstimationError::RemlOptimizationFailed(format!(
2246 "spatial adaptive cost eval failed: {e}"
2247 ))
2248 })?;
2249 if !owned.result.inner_converged {
2250 st.warm_cache = Some(owned.result.warm_start);
2251 return Err(EstimationError::RemlOptimizationFailed(
2252 "exact spatial adaptive cost inner solve did not converge".to_string(),
2253 ));
2254 }
2255 let objective = owned.result.objective;
2256 st.warm_cache = Some(owned.result.warm_start);
2257 st.terminal_mode = Some((theta, objective, owned.mode));
2258 Ok(objective)
2259 },
2260 |st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
2261 eval_outer(
2262 st,
2263 theta,
2264 if analytic_outer_hessian_available {
2265 gam_solve::rho_optimizer::OuterEvalOrder::ValueGradientHessian
2266 } else {
2267 gam_solve::rho_optimizer::OuterEvalOrder::ValueAndGradient
2268 },
2269 )
2270 },
2271 |st: &mut SpatialAdaptiveOuterState,
2272 theta: &Array1<f64>,
2273 order: gam_solve::rho_optimizer::OuterEvalOrder| { eval_outer(st, theta, order) },
2274 Some(|st: &mut SpatialAdaptiveOuterState| {
2275 st.warm_cache = None;
2276 st.terminal_mode = None;
2277 st.last_eval = None;
2278 }),
2279 Some(|st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
2280 let theta = clamp_theta(theta);
2281 let DecodedSpatialAdaptiveTheta {
2282 rho,
2283 adaptive_params,
2284 ..
2285 } = decode_theta(&theta)?;
2286 let family_eval =
2287 base_family.with_adaptive_params(adaptive_params, zero_quadratic.clone());
2288 let hyper_layout = realize_hyper_layout(&theta)?;
2289 let owned = evaluate_custom_family_joint_hyper_efs_owned(
2290 &family_eval,
2291 std::slice::from_ref(&blockspec),
2292 &outer_opts,
2293 &rho,
2294 &hyper_layout,
2295 st.warm_cache.as_ref(),
2296 )
2297 .map_err(|e| {
2298 EstimationError::RemlOptimizationFailed(format!(
2299 "spatial adaptive EFS eval failed: {e}"
2300 ))
2301 })?;
2302 if !owned.result.inner_converged {
2303 st.warm_cache = Some(owned.result.warm_start);
2304 return Err(EstimationError::RemlOptimizationFailed(
2305 "exact spatial adaptive EFS inner solve did not converge".to_string(),
2306 ));
2307 }
2308 let objective = owned.result.efs_eval.cost;
2309 st.warm_cache = Some(owned.result.warm_start);
2310 st.terminal_mode = Some((theta, objective, owned.mode));
2311 Ok(owned.result.efs_eval)
2312 }),
2313 |st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
2325 let theta = clamp_theta(theta);
2326 let DecodedSpatialAdaptiveTheta {
2327 rho,
2328 adaptive_params,
2329 ..
2330 } = decode_theta(&theta)?;
2331 let family_eval =
2332 base_family.with_adaptive_params(adaptive_params, zero_quadratic.clone());
2333 let hyper_layout = realize_hyper_layout(&theta)?;
2334 let owned = evaluate_custom_family_joint_hyper_owned(
2335 &family_eval,
2336 std::slice::from_ref(&blockspec),
2337 &outer_opts,
2338 &rho,
2339 &hyper_layout,
2340 st.warm_cache.as_ref(),
2341 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueOnly,
2342 )
2343 .map_err(|e| {
2344 EstimationError::RemlOptimizationFailed(format!(
2345 "spatial adaptive screening eval failed: {e}"
2346 ))
2347 })?;
2348 st.warm_cache = Some(owned.result.warm_start);
2349 Ok(owned.result.objective)
2350 },
2351 );
2352
2353 let certified_outer = problem
2354 .run_certified(&mut obj, "exact spatial adaptive regularization")
2355 .map_err(|e| {
2356 EstimationError::InvalidInput(format!(
2357 "exact spatial adaptive outer optimization failed: {e}"
2358 ))
2359 })?;
2360 let outer_iterations = certified_outer.iterations();
2361 let outer_grad_norm = certified_outer.final_grad_norm();
2362 let theta_star = certified_outer.rho().clone();
2363 let (mode_theta, mode_objective, terminal_mode) =
2364 obj.state.terminal_mode.take().ok_or_else(|| {
2365 EstimationError::InvalidInput(
2366 "exact spatial adaptive optimization certified without retaining its terminal coefficient mode"
2367 .to_string(),
2368 )
2369 })?;
2370 if mode_theta.len() != theta_star.len()
2371 || mode_theta
2372 .iter()
2373 .zip(theta_star.iter())
2374 .any(|(mode, certified)| mode.to_bits() != certified.to_bits())
2375 {
2376 return Err(EstimationError::InvalidInput(
2377 "exact spatial adaptive terminal coefficient mode does not bitwise match the certified hyperparameter vector"
2378 .to_string(),
2379 ));
2380 }
2381 if mode_objective.to_bits() != certified_outer.final_value().to_bits() {
2382 return Err(EstimationError::InvalidInput(format!(
2383 "exact spatial adaptive terminal coefficient mode objective does not bitwise match the certified objective: mode={mode_objective:.17e}, certified={:.17e}",
2384 certified_outer.final_value(),
2385 )));
2386 }
2387 let DecodedSpatialAdaptiveTheta {
2388 rho: _,
2389 retained_lambdas,
2390 adaptive_params,
2391 epsilon: eps_star,
2392 } = decode_theta(&theta_star)?;
2393 let mut fixed_total = Array2::<f64>::zeros((
2394 baseline.design.design.ncols(),
2395 baseline.design.design.ncols(),
2396 ));
2397 for (idx, penalty) in retained_penalties.iter().enumerate() {
2398 fixed_total.scaled_add(retained_lambdas[idx], penalty);
2399 }
2400 let certified_final_family = base_family.with_adaptive_params(
2406 adaptive_params.clone(),
2407 zero_quadratic.clone(),
2408 );
2409 let fixed_total = ValidatedFixedQuadraticHessian::try_from_dense(
2410 fixed_total,
2411 baseline.design.design.ncols(),
2412 )
2413 .map_err(|error| {
2414 EstimationError::InvalidInput(format!(
2415 "optimized spatial adaptive fixed quadratic Hessian is invalid: {error}"
2416 ))
2417 })?;
2418 let final_family =
2419 base_family.with_adaptive_params(adaptive_params.clone(), fixed_total.clone());
2420 let final_blockspec = ParameterBlockSpec {
2421 name: "eta".to_string(),
2422 design: baseline.design.design.clone(),
2423 offset: offset.to_owned(),
2424 penalties: retained_penalties
2425 .iter()
2426 .cloned()
2427 .map(PenaltyMatrix::Dense)
2428 .collect(),
2429 nullspace_dims: retained_nullspace_dims.clone(),
2430 initial_log_lambdas: theta_star.slice(s![..rho_dim]).to_owned(),
2431 initial_beta: Some(baseline.fit.beta.clone()),
2432 gauge_priority: 100,
2433 jacobian_callback: None,
2434 stacked_design: None,
2435 stacked_offset: None,
2436 };
2437 let final_fit = fit_custom_family_fixed_log_lambdas_from_owned_mode(
2438 &certified_final_family,
2439 &[final_blockspec],
2440 &BlockwiseFitOptions {
2441 inner_max_cycles: options.max_iter,
2442 inner_tol: options.tol,
2443 outer_max_iter: 1,
2444 outer_tol: options.tol,
2445 compute_covariance: true,
2446 ..BlockwiseFitOptions::default()
2447 },
2448 terminal_mode,
2449 &theta_star,
2450 &certified_outer,
2451 )
2452 .map_err(EstimationError::CustomFamily)?;
2453 let beta = final_fit.block_states[0].beta.clone();
2454 let final_eval = final_family
2455 .exact_evaluation(&beta)
2456 .map_err(EstimationError::InvalidInput)?;
2457 let penalized_hessian = final_eval
2458 .totalobjectivehessian(&final_family.design)
2459 .map_err(EstimationError::InvalidInput)?;
2460 let beta_covariance = final_fit.covariance_conditional.clone();
2461 let beta_standard_errors = beta_covariance
2462 .as_ref()
2463 .map(|cov| Array1::from_iter((0..cov.nrows()).map(|i| cov[[i, i]].max(0.0).sqrt())));
2464
2465 let mut full_lambdas = baseline.fit.lambdas.clone();
2466 for (idx, &global_idx) in retained_global_indices.iter().enumerate() {
2467 full_lambdas[global_idx] = retained_lambdas[idx];
2468 }
2469 for (cache_idx, cache) in runtime_caches.iter().enumerate() {
2470 full_lambdas[cache.mass_penalty_global_idx] = adaptive_params[cache_idx].lambda[0];
2471 full_lambdas[cache.tension_penalty_global_idx] = adaptive_params[cache_idx].lambda[1];
2472 full_lambdas[cache.stiffness_penalty_global_idx] = adaptive_params[cache_idx].lambda[2];
2473 }
2474
2475 let deviance = -2.0 * final_eval.obs.log_likelihood;
2476 let mut local_penalty_blocks =
2477 Vec::<PenaltySpec>::with_capacity(baseline.design.penalties.len());
2478 for (global_idx, bp) in baseline.design.penalties.iter().enumerate() {
2479 if adaptive_penalty_indices.contains(&global_idx) {
2480 let cache = runtime_caches
2481 .iter()
2482 .find(|cache| {
2483 cache.mass_penalty_global_idx == global_idx
2484 || cache.tension_penalty_global_idx == global_idx
2485 || cache.stiffness_penalty_global_idx == global_idx
2486 })
2487 .ok_or_else(|| {
2488 EstimationError::InvalidInput(format!(
2489 "missing runtime cache for adaptive penalty index {global_idx}"
2490 ))
2491 })?;
2492 let cache_idx = runtime_caches
2493 .iter()
2494 .position(|c| {
2495 c.mass_penalty_global_idx == global_idx
2496 || c.tension_penalty_global_idx == global_idx
2497 || c.stiffness_penalty_global_idx == global_idx
2498 })
2499 .ok_or_else(|| {
2500 EstimationError::InvalidInput(format!(
2501 "missing adaptive cache position for penalty index {global_idx}"
2502 ))
2503 })?;
2504 let state = &final_eval.adaptive_states[cache_idx];
2505 let local = if cache.mass_penalty_global_idx == global_idx {
2506 scalar_operatorhessian(&cache.d0, &state.magnitude.betahessian_diag())
2507 .mapv(|v| adaptive_params[cache_idx].lambda[0] * v)
2508 } else if cache.tension_penalty_global_idx == global_idx {
2509 grouped_operatorhessian(
2510 &cache.d1,
2511 cache.dimension,
2512 &state.gradient.betahessian_blocks(),
2513 )?
2514 .mapv(|v| adaptive_params[cache_idx].lambda[1] * v)
2515 } else {
2516 grouped_operatorhessian(
2517 &cache.d2,
2518 cache.dimension * cache.dimension,
2519 &state.curvature.betahessian_blocks(),
2520 )?
2521 .mapv(|v| adaptive_params[cache_idx].lambda[2] * v)
2522 };
2523 local_penalty_blocks.push(PenaltySpec::Dense(penalty_matrixwith_local_block(
2525 baseline.design.design.ncols(),
2526 cache.coeff_global_range.clone(),
2527 &local,
2528 )));
2529 } else {
2530 local_penalty_blocks.push(PenaltySpec::Dense(
2531 bp.to_global(p_total).mapv(|v| v * full_lambdas[global_idx]),
2532 ));
2533 }
2534 }
2535 let (edf_by_block, penalty_block_trace, edf_total) = if let Some(cov) = beta_covariance.as_ref()
2536 {
2537 exact_bounded_edf(
2538 &local_penalty_blocks,
2539 &Array1::from_elem(local_penalty_blocks.len(), 1.0),
2540 cov,
2541 )?
2542 } else {
2543 (
2544 vec![0.0; local_penalty_blocks.len()],
2545 vec![0.0; local_penalty_blocks.len()],
2546 0.0,
2547 )
2548 };
2549 let stable_penalty_term = 2.0 * final_eval.adaptive_penalty_value
2550 + beta.dot(&fixed_total.as_dense().dot(&beta));
2551 let standard_deviation = if family.is_gaussian_identity() {
2552 let denom = (y.len() as f64 - edf_total).max(1.0);
2553 (deviance / denom).sqrt()
2554 } else {
2555 1.0
2556 };
2557 let maps = compute_spatial_adaptiveweights_for_beta(
2558 &beta,
2559 runtime_caches,
2560 eps_star[0],
2561 eps_star[1],
2562 eps_star[2],
2563 adaptive_opts.weight_floor,
2564 adaptive_opts.weight_ceiling,
2565 beta_covariance.as_ref(),
2569 )?
2570 .into_iter()
2571 .zip(runtime_caches.iter())
2572 .map(|(w, cache)| AdaptiveSpatialMap {
2573 termname: cache.termname.clone(),
2574 feature_cols: cache.feature_cols.clone(),
2575 collocation_points: cache.collocation_points.clone(),
2576 inv_magweight: w.inv_magweight,
2577 invgradweight: w.invgradweight,
2578 inv_lapweight: w.inv_lapweight,
2579 })
2580 .collect::<Vec<_>>();
2581 let fitted_link = if family.is_latent_cloglog() {
2582 FittedLinkState::LatentCLogLog {
2583 state: latent_cloglog_state
2584 .expect("BinomialLatentCLogLog requires an explicit latent-cloglog state"),
2585 }
2586 } else if family.is_binomial_mixture() {
2587 mixture_link_state
2588 .clone()
2589 .map(|state| FittedLinkState::Mixture {
2590 state,
2591 covariance: None,
2592 })
2593 .unwrap_or(FittedLinkState::Standard(None))
2594 } else if family.is_binomial_sas() {
2595 sas_link_state
2596 .map(|state| FittedLinkState::Sas {
2597 state,
2598 covariance: None,
2599 })
2600 .unwrap_or(FittedLinkState::Standard(None))
2601 } else if family.is_binomial_beta_logistic() {
2602 sas_link_state
2603 .map(|state| FittedLinkState::BetaLogistic {
2604 state,
2605 covariance: None,
2606 })
2607 .unwrap_or(FittedLinkState::Standard(None))
2608 } else {
2609 FittedLinkState::Standard(None)
2610 };
2611 let max_abs_eta = final_eval
2612 .obs
2613 .eta
2614 .iter()
2615 .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
2616 let fitted = FittedTermCollection {
2617 fit: {
2618 let log_lambdas =
2619 checked_fit_log_lambdas(&full_lambdas, "final exact spatial adaptive fit")?;
2620 let working = gam_solve::estimate::WorkingGeometry {
2621 weights: final_eval.obs.fisherweight.clone(),
2622 response: exact_standard_working_response(&final_eval.obs)?,
2623 };
2624 let inf = FitInference {
2625 edf_by_block,
2626 penalty_block_trace,
2627 edf_total,
2628 smoothing_correction: None,
2629 smoothing_correction_method: None,
2630 smoothing_correction_first_order: None,
2631 smoothing_correction_method_first_order: None,
2632 penalized_hessian: penalized_hessian.clone().into(),
2635 reparam_qs: None,
2636 dispersion: gam_solve::estimate::Dispersion::UNIT,
2637 beta_covariance: beta_covariance
2638 .clone()
2639 .map(gam_problem::dispersion_cov::PhiScaledCovariance::from),
2640 beta_standard_errors,
2641 beta_covariance_corrected: None,
2642 beta_standard_errors_corrected: None,
2643 beta_covariance_frequentist: None,
2644 coefficient_influence: None,
2645 weighted_gram: None,
2646 bias_correction_beta: None,
2647 bias_correction_jacobian: None,
2648 };
2649 let geometry = Some(gam_solve::estimate::FitGeometry {
2650 coefficient_gauge: gam_problem::gauge::Gauge::identity(&[beta.len()]),
2651 penalized_hessian: penalized_hessian.into(),
2652 working: Some(working),
2653 });
2654 let covariance_conditional = beta_covariance;
2655 let convergence = final_fit.convergence_evidence();
2656 let pirls_status_val = convergence.inner_status();
2657 let certified_outer_present = convergence.outer_certificate().is_some();
2658 UnifiedFitResult::try_from_parts(UnifiedFitResultParts {
2659 blocks: vec![gam_solve::estimate::FittedBlock {
2660 beta: beta.clone(),
2661 role: gam_problem::BlockRole::Mean,
2662 edf: edf_total,
2663 lambdas: full_lambdas.clone(),
2664 }],
2665 log_lambdas,
2666 lambdas: full_lambdas,
2667 likelihood_scale: family.default_scale_metadata(),
2668 likelihood_family: Some(family),
2669 log_likelihood_normalization: gam_spec::LogLikelihoodNormalization::UserProvided,
2670 log_likelihood: final_eval.obs.log_likelihood,
2671 deviance,
2672 reml_score: final_fit.penalized_objective,
2673 stable_penalty_term,
2674 penalized_objective: final_fit.penalized_objective,
2675 used_device: false,
2676 outer_iterations,
2677 outer_converged: certified_outer_present,
2678 outer_gradient_norm: outer_grad_norm,
2679 standard_deviation,
2680 covariance_conditional,
2681 covariance_corrected: None,
2682 inference: Some(inf),
2683 fitted_link,
2684 geometry,
2685 block_states: Vec::new(),
2686 pirls_status: pirls_status_val,
2687 max_abs_eta,
2688 constraint_kkt: None,
2689 artifacts: gam_solve::estimate::FitArtifacts {
2690 pirls: None,
2691 criterion_certificate: final_fit.artifacts.criterion_certificate.clone(),
2692 ..Default::default()
2693 },
2694 inner_cycles: 0,
2695 })?
2696 },
2697 design: baseline.design,
2698 adaptive_diagnostics: Some(AdaptiveRegularizationDiagnostics {
2699 epsilon_0: eps_star[0],
2700 epsilon_g: eps_star[1],
2701 epsilon_c: eps_star[2],
2702 epsilon_outer_iterations: outer_iterations,
2703 mm_iterations: 0,
2704 converged: true,
2705 maps,
2706 }),
2707 };
2708 enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
2709 Ok(fitted)
2710}
2711
2712fn relax_smoothing_rho_prior(
2744 options: &FitOptions,
2745 design: &TermCollectionDesign,
2746 y: ArrayView1<'_, f64>,
2747 weights: ArrayView1<'_, f64>,
2748) -> gam_spec::RhoPrior {
2749 use gam_terms::basis::BasisMetadata;
2750 let base = &options.rho_prior;
2751 if matches!(
2754 base,
2755 gam_spec::RhoPrior::Flat | gam_spec::RhoPrior::Independent(_)
2756 ) {
2757 return base.clone();
2758 }
2759 let has_link_aux = options.sas_link.is_some()
2779 || options.optimize_sas
2780 || options.mixture_link.is_some()
2781 || options.optimize_mixture;
2782 let has_moving_kappa = design.smooth.terms.iter().any(|t| {
2783 if let BasisMetadata::Duchon {
2796 length_scale,
2797 aniso_log_scales,
2798 ..
2799 } = &t.metadata
2800 {
2801 return length_scale.is_some() || aniso_log_scales.is_some();
2802 }
2803 matches!(
2804 t.metadata,
2805 BasisMetadata::Matern { .. }
2806 | BasisMetadata::Sphere { .. }
2807 | BasisMetadata::SphereHarmonics { .. }
2808 | BasisMetadata::ConstantCurvature { .. }
2809 | BasisMetadata::MeasureJet { .. }
2810 )
2811 });
2812 let length_safe = !has_link_aux && !has_moving_kappa;
2819 if !length_safe {
2820 return base.clone();
2821 }
2822 let coords = &design.penaltyinfo;
2823 if coords.is_empty() {
2824 return base.clone();
2825 }
2826 let n_obs = design.design.nrows();
2837 let p_total = design.design.ncols();
2838 let underdetermined = n_obs < 2 * p_total;
2869 let relaxable_terms: std::collections::HashSet<&str> = design
2881 .smooth
2882 .terms
2883 .iter()
2884 .filter(|t| {
2885 (matches!(
2886 t.metadata,
2887 BasisMetadata::BSpline1D { .. }
2888 | BasisMetadata::ThinPlate { .. }
2889 | BasisMetadata::TensorBSpline { .. }
2890 )
2891 || matches!(
2903 t.metadata,
2904 BasisMetadata::Duchon {
2905 length_scale: None,
2906 aniso_log_scales: None,
2907 ..
2908 }
2909 ))
2910 && matches!(t.shape, gam_terms::smooth::ShapeConstraint::None)
2924 })
2925 .map(|t| t.name.as_str())
2926 .collect();
2927 let any_relaxed = coords.iter().any(|info| {
2928 info.termname
2929 .as_deref()
2930 .is_some_and(|name| relaxable_terms.contains(name))
2931 });
2932 if !any_relaxed {
2933 return base.clone();
2934 }
2935 let relaxed_prior = if underdetermined {
2940 gam_spec::RhoPrior::Normal {
2941 mean: 0.0,
2942 sd: RELAX_UNDERDETERMINED_RHO_SD,
2943 }
2944 } else {
2945 gam_spec::RhoPrior::Flat
2946 };
2947 let nullspace_select_prior = gam_spec::RhoPrior::PenalizedComplexity {
2974 upper: NULLSPACE_SELECT_PC_UPPER,
2975 tail_prob: NULLSPACE_SELECT_PC_TAIL_PROB,
2976 };
2977 let nullspace_degeneracy_prior = gam_spec::RhoPrior::Normal {
3004 mean: 0.0,
3005 sd: NULLSPACE_WELLDET_DEGENERACY_RHO_SD,
3006 };
3007 let per_coord = coords
3008 .iter()
3009 .enumerate()
3010 .map(|(coord_idx, info)| {
3011 let relax = info
3012 .termname
3013 .as_deref()
3014 .is_some_and(|name| relaxable_terms.contains(name));
3015 if !relax {
3016 return base.clone();
3017 }
3018 let is_nullspace = matches!(info.penalty.source, PenaltySource::DoublePenaltyNullspace);
3019 if is_nullspace {
3058 if underdetermined
3079 && !nullspace_directions_are_supported(design, coord_idx, y, weights)
3080 {
3081 nullspace_select_prior.clone()
3082 } else {
3083 nullspace_degeneracy_prior.clone()
3084 }
3085 } else {
3086 relaxed_prior.clone()
3087 }
3088 })
3089 .collect::<Vec<_>>();
3090 gam_spec::RhoPrior::Independent(per_coord)
3091}
3092
3093const NULLSPACE_SUPPORT_FRACTION_THRESHOLD: f64 = 0.5;
3100
3101fn nullspace_directions_are_supported(
3123 design: &TermCollectionDesign,
3124 penalty_idx: usize,
3125 y: ArrayView1<'_, f64>,
3126 weights: ArrayView1<'_, f64>,
3127) -> bool {
3128 use gam_linalg::faer_ndarray::FaerEigh;
3129
3130 let Some(pen) = design.penalties.get(penalty_idx) else {
3131 return false;
3132 };
3133 let col_range = pen.col_range.clone();
3134 if col_range.is_empty() {
3135 return false;
3136 }
3137 let x = design.design.to_dense();
3138 let n = x.nrows();
3139 if n == 0 || y.len() != n || weights.len() != n || x.ncols() < col_range.end {
3140 return false;
3141 }
3142 let mut resp = y.to_owned();
3146 if design.affine_offset.len() == n {
3147 resp -= &design.affine_offset;
3148 }
3149
3150 let Ok((evals, evecs)) = pen.local.eigh(faer::Side::Lower) else {
3153 return false;
3154 };
3155 let max_eig = evals.iter().cloned().fold(0.0_f64, |m, v| m.max(v));
3156 if !(max_eig > 0.0) {
3157 return false;
3158 }
3159 let tol = 1.0e-9 * max_eig;
3160 let pos_cols: Vec<usize> = (0..evals.len()).filter(|&j| evals[j] > tol).collect();
3161 if pos_cols.is_empty() {
3162 return false;
3163 }
3164 let xblock = x.slice(s![.., col_range.clone()]);
3165 let mut z = Array2::<f64>::zeros((n, pos_cols.len()));
3166 for (out_j, &j) in pos_cols.iter().enumerate() {
3167 let v = evecs.column(j);
3168 if v.len() != xblock.ncols() {
3170 return false;
3171 }
3172 z.column_mut(out_j).assign(&xblock.dot(&v));
3173 }
3174
3175 let mut control: Vec<usize> = design.intercept_range.clone().collect();
3180 for (_, r) in &design.linear_ranges {
3181 control.extend(r.clone());
3182 }
3183 control.retain(|&c| c < x.ncols());
3184 let mut cmat = Array2::<f64>::zeros((n, control.len().max(1)));
3185 if control.is_empty() {
3186 cmat.column_mut(0).fill(1.0);
3188 } else {
3189 for (out_j, &c) in control.iter().enumerate() {
3190 cmat.column_mut(out_j).assign(&x.column(c));
3191 }
3192 }
3193
3194 let rss_c = weighted_regression_rss(cmat.view(), resp.view(), weights);
3195 let Some(rss_c) = rss_c else { return false };
3196 let base_scale = weighted_total_ss(resp.view(), weights);
3199 if !(rss_c > 1.0e-12 * base_scale.max(f64::MIN_POSITIVE)) {
3200 return false;
3201 }
3202 let mut cz = Array2::<f64>::zeros((n, cmat.ncols() + z.ncols()));
3203 cz.slice_mut(s![.., ..cmat.ncols()]).assign(&cmat);
3204 cz.slice_mut(s![.., cmat.ncols()..]).assign(&z);
3205 let Some(rss_cz) = weighted_regression_rss(cz.view(), resp.view(), weights) else {
3206 return false;
3207 };
3208
3209 let support = (rss_c - rss_cz) / rss_c;
3210 support.is_finite() && support > NULLSPACE_SUPPORT_FRACTION_THRESHOLD
3211}
3212
3213fn weighted_total_ss(y: ArrayView1<'_, f64>, w: ArrayView1<'_, f64>) -> f64 {
3215 let mut sw = 0.0;
3216 let mut swy = 0.0;
3217 for (&yi, &wi) in y.iter().zip(w.iter()) {
3218 if wi > 0.0 && yi.is_finite() {
3219 sw += wi;
3220 swy += wi * yi;
3221 }
3222 }
3223 if sw <= 0.0 {
3224 return 0.0;
3225 }
3226 let mean = swy / sw;
3227 let mut ss = 0.0;
3228 for (&yi, &wi) in y.iter().zip(w.iter()) {
3229 if wi > 0.0 && yi.is_finite() {
3230 ss += wi * (yi - mean) * (yi - mean);
3231 }
3232 }
3233 ss
3234}
3235
3236fn weighted_regression_rss(
3243 d: ArrayView2<'_, f64>,
3244 y: ArrayView1<'_, f64>,
3245 w: ArrayView1<'_, f64>,
3246) -> Option<f64> {
3247 use gam_linalg::faer_ndarray::FaerCholesky;
3248
3249 let m = d.ncols();
3250 if m == 0 {
3251 return Some(weighted_total_ss(y, w));
3252 }
3253 let mut gram = Array2::<f64>::zeros((m, m));
3254 let mut rhs = Array1::<f64>::zeros(m);
3255 for row in 0..d.nrows() {
3256 let wi = w[row];
3257 if !(wi > 0.0) || !y[row].is_finite() {
3258 continue;
3259 }
3260 let dr = d.row(row);
3261 for a in 0..m {
3262 let wda = wi * dr[a];
3263 rhs[a] += wda * y[row];
3264 for b in a..m {
3265 gram[[a, b]] += wda * dr[b];
3266 }
3267 }
3268 }
3269 for a in 0..m {
3270 for b in (a + 1)..m {
3271 gram[[b, a]] = gram[[a, b]];
3272 }
3273 }
3274 let trace = (0..m).map(|i| gram[[i, i]]).sum::<f64>();
3275 if !(trace > 0.0) {
3276 return Some(weighted_total_ss(y, w));
3277 }
3278 let ridge = 1.0e-10 * trace / (m as f64);
3279 for i in 0..m {
3280 gram[[i, i]] += ridge;
3281 }
3282 let chol = gram.cholesky(faer::Side::Lower).ok()?;
3283 let beta = chol.solvevec(&rhs);
3284 let mut rss = 0.0;
3285 for row in 0..d.nrows() {
3286 let wi = w[row];
3287 if !(wi > 0.0) || !y[row].is_finite() {
3288 continue;
3289 }
3290 let fitted = d.row(row).dot(&beta);
3291 let resid = y[row] - fitted;
3292 rss += wi * resid * resid;
3293 }
3294 Some(rss)
3295}
3296
3297const RELAX_UNDERDETERMINED_RHO_SD: f64 = 15.0;
3310
3311const NULLSPACE_SELECT_PC_UPPER: f64 = 0.05;
3329
3330const NULLSPACE_SELECT_PC_TAIL_PROB: f64 = 0.01;
3340
3341fn adaptive_fit_options_base(options: &FitOptions, design: &TermCollectionDesign) -> FitOptions {
3342 FitOptions {
3343 resource_policy: options.resource_policy.clone(),
3344 latent_cloglog: options.latent_cloglog,
3345 mixture_link: options.mixture_link.clone(),
3346 optimize_mixture: options.optimize_mixture,
3347 sas_link: options.sas_link,
3348 optimize_sas: options.optimize_sas,
3349 compute_inference: options.compute_inference,
3350 skip_rho_posterior_inference: options.skip_rho_posterior_inference,
3351 max_iter: options.max_iter,
3352 tol: options.tol,
3353 nullspace_dims: design.nullspace_dims.clone(),
3354 linear_constraints: design.linear_constraints.clone(),
3355 firth_bias_reduction: options.firth_bias_reduction,
3356 adaptive_regularization: None,
3357 penalty_shrinkage_floor: options.penalty_shrinkage_floor,
3358 rho_prior: options.rho_prior.clone(),
3361 kronecker_penalty_system: design.kronecker_penalty_system(),
3362 kronecker_factored: design
3363 .smooth
3364 .terms
3365 .iter()
3366 .find_map(|t| t.kronecker_factored.clone()),
3367 persist_warm_start_disk: options.persist_warm_start_disk,
3368 }
3369}
3370
3371fn superseded_fit_options(options: &FitOptions) -> FitOptions {
3372 let mut fit_options = options.clone();
3373 fit_options.skip_rho_posterior_inference = true;
3374 fit_options
3375}
3376
3377#[derive(Clone)]
3378struct BoundedLinearTermMeta {
3379 col_idx: usize,
3380 min: f64,
3381 max: f64,
3382 prior: BoundedCoefficientPriorSpec,
3383}
3384
3385struct BoundedEffectiveJacobian {
3409 design: Array2<f64>,
3410 bounded_terms: Vec<BoundedLinearTermMeta>,
3411}
3412
3413impl BlockEffectiveJacobian for BoundedEffectiveJacobian {
3414 fn effective_jacobian_rows(
3415 &self,
3416 state: &FamilyLinearizationState<'_>,
3417 rows: std::ops::Range<usize>,
3418 ) -> Result<Array2<f64>, String> {
3419 let p = self.design.ncols();
3420 let n = self.design.nrows();
3421 let rows = rows.start.min(n)..rows.end.min(n);
3422 if !state.beta.is_empty() {
3423 if state.beta.len() != p {
3424 return Err(format!(
3425 "BoundedEffectiveJacobian::effective_jacobian_at: beta length {} != design \
3426 ncols {p}",
3427 state.beta.len(),
3428 ));
3429 }
3430 if state.beta.iter().any(|v| !v.is_finite()) {
3431 return Err(
3432 "BoundedEffectiveJacobian::effective_jacobian_at: beta contains a non-finite value"
3433 .to_string(),
3434 );
3435 }
3436 }
3437 let mut jac = self
3438 .design
3439 .slice(ndarray::s![rows.start..rows.end, ..])
3440 .to_owned();
3441 for term in &self.bounded_terms {
3442 if term.col_idx >= p {
3443 return Err(format!(
3444 "BoundedEffectiveJacobian::effective_jacobian_at: bounded column {} is outside {p} columns",
3445 term.col_idx
3446 ));
3447 }
3448 let theta = if state.beta.is_empty() {
3449 0.0
3450 } else {
3451 state.beta[term.col_idx]
3452 };
3453 let (_, _, db_dtheta, _, _) = bounded_latent_derivatives(theta, term.min, term.max);
3454 if !(db_dtheta.is_finite() && db_dtheta > 0.0) {
3455 return Err(format!(
3456 "BoundedEffectiveJacobian::effective_jacobian_at: bounded column {} has unrepresentable derivative {db_dtheta} at theta={theta}",
3457 term.col_idx
3458 ));
3459 }
3460 jac.column_mut(term.col_idx).mapv_inplace(|v| v * db_dtheta);
3461 }
3462 Ok(jac)
3463 }
3464}
3465
3466#[derive(Clone)]
3467struct BoundedLinearFamily {
3468 likelihood: gam_spec::GlmLikelihoodSpec,
3469 latent_cloglog_state: Option<LatentCLogLogState>,
3470 mixture_link_state: Option<MixtureLinkState>,
3471 sas_link_state: Option<SasLinkState>,
3472 y: Array1<f64>,
3473 weights: Array1<f64>,
3474 design: Array2<f64>,
3475 designzeroed: Array2<f64>,
3476 offset: Array1<f64>,
3477 bounded_terms: Vec<BoundedLinearTermMeta>,
3478}
3479
3480#[derive(Clone, Debug)]
3481struct StandardFamilyObservationState {
3482 eta: Array1<f64>,
3483 score: Array1<f64>,
3484 fisherweight: Array1<f64>,
3485 neghessian_eta: Array1<f64>,
3486 neghessian_eta_derivative: Array1<f64>,
3487 log_likelihood: f64,
3488}
3489
3490fn bounded_latent_to_user(theta: f64, min: f64, max: f64) -> (f64, f64, f64) {
3491 let jet = logit_inverse_link_jet5(theta);
3492 let z = jet.mu;
3493 let width = max - min;
3494 let beta = min + width * z;
3495 let db_dtheta = width * jet.d1;
3496 (beta, z, db_dtheta)
3497}
3498
3499fn bounded_user_to_latent(beta: f64, min: f64, max: f64) -> f64 {
3508 (beta - min).ln() - (max - beta).ln()
3509}
3510
3511#[derive(Debug, Clone, Copy)]
3515pub struct BoundedSampleColumn {
3516 pub col_idx: usize,
3518 pub min: f64,
3520 pub max: f64,
3522}
3523
3524pub fn sample_bounded_latent_posterior_internal(
3562 beta_user: &Array1<f64>,
3563 user_hessian: &Array2<f64>,
3564 bounded_columns: &[BoundedSampleColumn],
3565 n_draws: usize,
3566 sqrt_cov_scale: f64,
3567 base_seed: u64,
3568) -> Result<Array2<f64>, EstimationError> {
3569 let p = beta_user.len();
3570 if user_hessian.nrows() != p || user_hessian.ncols() != p {
3571 crate::bail_invalid_estim!(
3572 "bounded posterior sampling dimension mismatch: mode has {p} entries, user Hessian is {}x{}",
3573 user_hessian.nrows(),
3574 user_hessian.ncols()
3575 );
3576 }
3577 if beta_user.iter().any(|value| !value.is_finite()) {
3578 crate::bail_invalid_estim!("bounded posterior sampling requires a finite mode");
3579 }
3580 if user_hessian.iter().any(|value| !value.is_finite()) {
3581 crate::bail_invalid_estim!("bounded posterior sampling requires a finite Hessian");
3582 }
3583 if !(sqrt_cov_scale.is_finite() && sqrt_cov_scale >= 0.0) {
3584 crate::bail_invalid_estim!(
3585 "bounded posterior sampling covariance scale must be finite and non-negative, got {sqrt_cov_scale}"
3586 );
3587 }
3588
3589 let mut theta_mode = beta_user.clone();
3591 let mut jac_diag = Array1::<f64>::ones(p);
3592 for bc in bounded_columns {
3593 if bc.col_idx >= p {
3594 crate::bail_invalid_estim!(
3595 "bounded posterior sampling: bounded column index {} out of range for {p} coefficients",
3596 bc.col_idx
3597 );
3598 }
3599 if !(bc.min.is_finite()
3600 && bc.max.is_finite()
3601 && (bc.max - bc.min).is_finite()
3602 && bc.min < beta_user[bc.col_idx]
3603 && beta_user[bc.col_idx] < bc.max)
3604 {
3605 crate::bail_invalid_estim!(
3606 "bounded posterior sampling column {} requires finite bounds with a finite width and a mode strictly inside ({}, {}); got {}",
3607 bc.col_idx,
3608 bc.min,
3609 bc.max,
3610 beta_user[bc.col_idx]
3611 );
3612 }
3613 let theta_i = bounded_user_to_latent(beta_user[bc.col_idx], bc.min, bc.max);
3614 let (_, _, db_dtheta) = bounded_latent_to_user(theta_i, bc.min, bc.max);
3615 if !(theta_i.is_finite() && db_dtheta.is_finite() && db_dtheta > 0.0) {
3616 crate::bail_invalid_estim!(
3617 "bounded posterior sampling column {} has unrepresentable latent geometry: theta={theta_i}, d_beta/d_theta={db_dtheta}",
3618 bc.col_idx
3619 );
3620 }
3621 theta_mode[bc.col_idx] = theta_i;
3622 jac_diag[bc.col_idx] = db_dtheta;
3623 }
3624
3625 let mut h_latent = user_hessian.clone();
3628 for i in 0..p {
3629 let ji = jac_diag[i];
3630 if ji != 1.0 {
3631 h_latent.row_mut(i).mapv_inplace(|v| v * ji);
3632 h_latent.column_mut(i).mapv_inplace(|v| v * ji);
3633 }
3634 }
3635
3636 use gam_linalg::faer_ndarray::FaerCholesky as _;
3639 use rand::SeedableRng as _;
3640 let chol = h_latent.cholesky(faer::Side::Lower).map_err(|err| {
3641 EstimationError::InvalidInput(format!(
3642 "bounded posterior sampling: Cholesky of the latent penalized Hessian failed: {err:?}"
3643 ))
3644 })?;
3645 let l = chol.lower_triangular();
3646
3647 let mut draws = Array2::<f64>::zeros((n_draws, p));
3648 let mut eps = Array1::<f64>::zeros(p);
3649 let mut delta = Array1::<f64>::zeros(p);
3650 let mut rng = rand::rngs::StdRng::seed_from_u64(base_seed);
3651 for k in 0..n_draws {
3652 for e in eps.iter_mut() {
3653 *e = standard_normal_draw(&mut rng);
3654 }
3655 solve_lower_transpose_into(&l, &eps, &mut delta)?;
3656 for i in 0..p {
3657 draws[(k, i)] = theta_mode[i] + sqrt_cov_scale * delta[i];
3660 }
3661 for bc in bounded_columns {
3666 let (beta_draw, _, _) = bounded_latent_to_user(draws[(k, bc.col_idx)], bc.min, bc.max);
3667 draws[(k, bc.col_idx)] = beta_draw;
3668 }
3669 }
3670
3671 Ok(draws)
3672}
3673
3674#[inline]
3677fn standard_normal_draw<R: rand::Rng + ?Sized>(rng: &mut R) -> f64 {
3678 use rand::RngExt as _;
3679 let u1 = loop {
3680 let candidate = rng.random::<f64>();
3681 if candidate > 0.0 {
3682 break candidate;
3683 }
3684 };
3685 let u2 = rng.random::<f64>();
3686 (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos()
3687}
3688
3689fn solve_lower_transpose_into(
3693 l: &Array2<f64>,
3694 b: &Array1<f64>,
3695 out: &mut Array1<f64>,
3696) -> Result<(), EstimationError> {
3697 let p = l.nrows();
3698 if l.ncols() != p || b.len() != p || out.len() != p {
3699 crate::bail_invalid_estim!(
3700 "bounded triangular solve dimension mismatch: L={}x{}, b={}, out={}",
3701 l.nrows(),
3702 l.ncols(),
3703 b.len(),
3704 out.len()
3705 );
3706 }
3707 for i in (0..p).rev() {
3708 let mut acc = b[i];
3709 for j in (i + 1)..p {
3710 acc -= l[(j, i)] * out[j];
3711 }
3712 let diag = l[(i, i)];
3713 if !(diag.is_finite() && diag > 0.0 && acc.is_finite()) {
3714 crate::bail_invalid_estim!(
3715 "bounded triangular solve has invalid row {i}: diagonal={diag}, residual={acc}"
3716 );
3717 }
3718 let value = acc / diag;
3719 if !value.is_finite() {
3720 crate::bail_invalid_estim!(
3721 "bounded triangular solve produced a non-finite value at row {i}: {acc}/{diag}"
3722 );
3723 }
3724 out[i] = value;
3725 }
3726 Ok(())
3727}
3728
3729fn bounded_latent_derivatives(theta: f64, min: f64, max: f64) -> (f64, f64, f64, f64, f64) {
3730 let jet = logit_inverse_link_jet5(theta);
3731 let z = jet.mu;
3732 let width = max - min;
3733 let beta = min + width * z;
3734 let db_dtheta = width * jet.d1;
3735 let d2b_dtheta2 = width * jet.d2;
3736 let d3b_dtheta3 = width * jet.d3;
3737 (beta, z, db_dtheta, d2b_dtheta2, d3b_dtheta3)
3738}
3739
3740fn bounded_prior_terms(
3741 theta: f64,
3742 prior: &BoundedCoefficientPriorSpec,
3743) -> Result<(f64, f64, f64, f64), String> {
3744 if !theta.is_finite() {
3745 return Err(format!(
3746 "bounded coefficient prior requires a finite latent coordinate, got {theta}"
3747 ));
3748 }
3749 let (a, b) = match prior {
3750 BoundedCoefficientPriorSpec::None => return Ok((0.0, 0.0, 0.0, 0.0)),
3752 BoundedCoefficientPriorSpec::Uniform => (1.0, 1.0),
3755 BoundedCoefficientPriorSpec::Beta { a, b } => (*a, *b),
3756 };
3757 if !(a.is_finite() && a > 0.0 && b.is_finite() && b > 0.0) {
3758 return Err(format!(
3759 "bounded coefficient Beta prior requires finite positive shapes, got ({a}, {b})"
3760 ));
3761 }
3762 let jet = logit_inverse_link_jet5(theta);
3763 let z = jet.mu;
3764 let logp = -a * gam_linalg::utils::stable_softplus(-theta)
3769 - b * gam_linalg::utils::stable_softplus(theta);
3770 let grad = a - (a + b) * z;
3771 let neghess = (a + b) * jet.d1;
3772 let neghess_derivative = (a + b) * jet.d2;
3773 let terms = (logp, grad, neghess, neghess_derivative);
3774 if [terms.0, terms.1, terms.2, terms.3]
3775 .iter()
3776 .any(|value| !value.is_finite())
3777 {
3778 return Err(format!(
3779 "bounded coefficient prior geometry is not representable at theta={theta}: {terms:?}"
3780 ));
3781 }
3782 Ok(terms)
3783}
3784
3785#[derive(Clone, Copy)]
3786struct ExactStandardObservationRow {
3787 mu: f64,
3788 score: f64,
3789 fisherweight: f64,
3790 neghessian_eta: f64,
3791 neghessian_eta_derivative: f64,
3792 log_likelihood: f64,
3793}
3794
3795impl ExactStandardObservationRow {
3796 #[inline]
3797 fn zero_weight(mu: f64) -> Self {
3798 Self {
3799 mu,
3800 score: 0.0,
3801 fisherweight: 0.0,
3802 neghessian_eta: 0.0,
3803 neghessian_eta_derivative: 0.0,
3804 log_likelihood: 0.0,
3805 }
3806 }
3807}
3808
3809#[inline]
3810fn bounded_row_error(row: usize, quantity: &'static str, eta: f64, value: f64) -> EstimationError {
3811 EstimationError::PirlsRowGeometryUnrepresentable {
3812 row,
3813 quantity,
3814 eta,
3815 value,
3816 }
3817}
3818
3819#[inline]
3820fn certify_bounded_row(
3821 row: usize,
3822 eta: f64,
3823 state: ExactStandardObservationRow,
3824) -> Result<ExactStandardObservationRow, EstimationError> {
3825 for (quantity, value) in [
3826 ("bounded-family mean", state.mu),
3827 ("bounded-family score", state.score),
3828 ("bounded-family Fisher weight", state.fisherweight),
3829 ("bounded-family observed Hessian", state.neghessian_eta),
3830 (
3831 "bounded-family observed Hessian derivative",
3832 state.neghessian_eta_derivative,
3833 ),
3834 ("bounded-family log likelihood", state.log_likelihood),
3835 ] {
3836 if !value.is_finite() {
3837 return Err(bounded_row_error(row, quantity, eta, value));
3838 }
3839 }
3840 if state.fisherweight < 0.0 {
3841 return Err(bounded_row_error(
3842 row,
3843 "bounded-family Fisher weight",
3844 eta,
3845 state.fisherweight,
3846 ));
3847 }
3848 Ok(state)
3849}
3850
3851#[inline]
3852fn weighted_positive_from_log(weight: f64, log_value: f64) -> f64 {
3853 if weight == 0.0 {
3854 return 0.0;
3855 }
3856 (weight.ln() + log_value).exp()
3857}
3858
3859#[inline]
3860fn weighted_product3(a: f64, b: f64, c: f64) -> f64 {
3861 crate::gamlss::scaled_signed_product3(a, b, c)
3862}
3863
3864#[inline]
3865fn convex_combination(y: f64, left: f64, right: f64) -> f64 {
3866 if y == 0.0 {
3867 right
3868 } else if y == 1.0 {
3869 left
3870 } else {
3871 y.mul_add(left, (1.0 - y) * right)
3872 }
3873}
3874
3875#[derive(Clone, Copy)]
3881struct BernoulliNaturalJet {
3882 mu: f64,
3883 log_mu: [f64; 4],
3884 log_one_minus_mu: [f64; 4],
3885 log_fisher: f64,
3886}
3887
3888#[inline]
3889fn probit_natural_jet(eta: f64) -> BernoulliNaturalJet {
3890 let left = gam_math::probability::normal_logcdf_derivatives(eta);
3891 let right_at_neg_eta = gam_math::probability::normal_logcdf_derivatives(-eta);
3892 let log_pdf = if eta.abs() <= f64::MAX.sqrt() {
3893 -0.5 * eta * eta - 0.5 * (2.0 * std::f64::consts::PI).ln()
3894 } else {
3895 f64::NEG_INFINITY
3896 };
3897 BernoulliNaturalJet {
3898 mu: left[0].exp(),
3899 log_mu: [left[0], left[1], left[2], left[3]],
3900 log_one_minus_mu: [
3901 right_at_neg_eta[0],
3902 -right_at_neg_eta[1],
3903 right_at_neg_eta[2],
3904 -right_at_neg_eta[3],
3905 ],
3906 log_fisher: 2.0 * log_pdf - left[0] - right_at_neg_eta[0],
3907 }
3908}
3909
3910#[inline]
3911fn cloglog_natural_jet(eta: f64) -> BernoulliNaturalJet {
3912 let x = eta.exp();
3913 if x == f64::INFINITY {
3914 return BernoulliNaturalJet {
3915 mu: 1.0,
3916 log_mu: [0.0; 4],
3917 log_one_minus_mu: [f64::NEG_INFINITY; 4],
3918 log_fisher: f64::NEG_INFINITY,
3919 };
3920 }
3921 if x == 0.0 {
3922 return BernoulliNaturalJet {
3923 mu: 0.0,
3924 log_mu: [eta, 1.0, 0.0, 0.0],
3925 log_one_minus_mu: [0.0; 4],
3926 log_fisher: eta,
3927 };
3928 }
3929 let mu = -(-x).exp_m1();
3930 let log_mu = if x < 0.5 {
3931 eta + (mu / x).ln()
3932 } else {
3933 mu.ln()
3934 };
3935 let h = if x < 1.0 {
3936 x / x.exp_m1()
3937 } else {
3938 let exp_neg_x = (-x).exp();
3939 x * exp_neg_x / (1.0 - exp_neg_x)
3940 };
3941 let a = 1.0 - x - h;
3942 let d2_log_mu = h * a;
3943 let d3_log_mu = h * (a * a - x - h * a);
3944 BernoulliNaturalJet {
3945 mu,
3946 log_mu: [log_mu, h, d2_log_mu, d3_log_mu],
3947 log_one_minus_mu: [-x, -x, -x, -x],
3948 log_fisher: 2.0 * eta - x - log_mu,
3949 }
3950}
3951
3952#[inline]
3953fn loglog_natural_jet(eta: f64) -> BernoulliNaturalJet {
3954 let mirrored = cloglog_natural_jet(-eta);
3955 BernoulliNaturalJet {
3956 mu: mirrored.log_one_minus_mu[0].exp(),
3957 log_mu: [
3958 mirrored.log_one_minus_mu[0],
3959 -mirrored.log_one_minus_mu[1],
3960 mirrored.log_one_minus_mu[2],
3961 -mirrored.log_one_minus_mu[3],
3962 ],
3963 log_one_minus_mu: [
3964 mirrored.log_mu[0],
3965 -mirrored.log_mu[1],
3966 mirrored.log_mu[2],
3967 -mirrored.log_mu[3],
3968 ],
3969 log_fisher: mirrored.log_fisher,
3970 }
3971}
3972
3973#[inline]
3974fn cauchit_natural_jet(eta: f64) -> BernoulliNaturalJet {
3975 let (mu, one_minus_mu) = if eta > 0.0 {
3976 let q = (eta.recip()).atan() / std::f64::consts::PI;
3977 (1.0 - q, q)
3978 } else if eta < 0.0 {
3979 let p = (-eta.recip()).atan() / std::f64::consts::PI;
3980 (p, 1.0 - p)
3981 } else {
3982 (0.5, 0.5)
3983 };
3984 let abs_eta = eta.abs();
3985 let log_one_plus_eta_sq = if abs_eta <= f64::MAX.sqrt() {
3986 (eta * eta).ln_1p()
3987 } else {
3988 2.0 * abs_eta.ln() + eta.recip().powi(2).ln_1p()
3989 };
3990 let log_d1 = -std::f64::consts::PI.ln() - log_one_plus_eta_sq;
3991 let ratio = if abs_eta <= 1.0 {
3992 eta / (1.0 + eta * eta)
3993 } else {
3994 1.0 / (eta + eta.recip())
3995 };
3996 let d2_over_d1 = -2.0 * ratio;
3997 let inv_one_plus_sq = if abs_eta <= 1.0 {
3998 1.0 / (1.0 + eta * eta)
3999 } else {
4000 let inv = eta.recip();
4001 inv * inv / (1.0 + inv * inv)
4002 };
4003 let d3_over_d1 = inv_one_plus_sq * (6.0 * (eta * ratio) - 2.0 * inv_one_plus_sq);
4004 let d1_over_mu = (log_d1 - mu.ln()).exp();
4005 let d1_over_q = (log_d1 - one_minus_mu.ln()).exp();
4006 let left_d2_ratio = d2_over_d1 * d1_over_mu;
4007 let right_d2_ratio = d2_over_d1 * d1_over_q;
4008 BernoulliNaturalJet {
4009 mu,
4010 log_mu: [
4011 mu.ln(),
4012 d1_over_mu,
4013 left_d2_ratio - d1_over_mu * d1_over_mu,
4014 d3_over_d1 * d1_over_mu - 3.0 * d1_over_mu * left_d2_ratio + 2.0 * d1_over_mu.powi(3),
4015 ],
4016 log_one_minus_mu: [
4017 one_minus_mu.ln(),
4018 -d1_over_q,
4019 -right_d2_ratio - d1_over_q * d1_over_q,
4020 -d3_over_d1 * d1_over_q - 3.0 * d1_over_q * right_d2_ratio - 2.0 * d1_over_q.powi(3),
4021 ],
4022 log_fisher: 2.0 * log_d1 - mu.ln() - one_minus_mu.ln(),
4023 }
4024}
4025
4026#[inline]
4027fn generic_bernoulli_natural_jet(
4028 row: usize,
4029 eta: f64,
4030 link: &InverseLink,
4031) -> Result<BernoulliNaturalJet, EstimationError> {
4032 let jet = inverse_link_jet_for_inverse_link(link, eta)?;
4033 if !(jet.mu.is_finite()
4034 && jet.mu > 0.0
4035 && jet.mu < 1.0
4036 && jet.d1.is_finite()
4037 && jet.d1 > 0.0
4038 && jet.d2.is_finite()
4039 && jet.d3.is_finite())
4040 {
4041 return Err(bounded_row_error(
4042 row,
4043 "bounded-family inverse-link jet",
4044 eta,
4045 jet.mu,
4046 ));
4047 }
4048 let mu = jet.mu;
4049 let q = 1.0 - mu;
4050 let r1 = jet.d1 / mu;
4051 let r2 = jet.d2 / mu;
4052 let r3 = jet.d3 / mu;
4053 let s1 = jet.d1 / q;
4054 let s2 = jet.d2 / q;
4055 let s3 = jet.d3 / q;
4056 Ok(BernoulliNaturalJet {
4057 mu,
4058 log_mu: [
4059 mu.ln(),
4060 r1,
4061 r2 - r1 * r1,
4062 r3 - 3.0 * r1 * r2 + 2.0 * r1.powi(3),
4063 ],
4064 log_one_minus_mu: [
4065 (-mu).ln_1p(),
4066 -s1,
4067 -s2 - s1 * s1,
4068 -s3 - 3.0 * s1 * s2 - 2.0 * s1.powi(3),
4069 ],
4070 log_fisher: 2.0 * jet.d1.ln() - mu.ln() - q.ln(),
4071 })
4072}
4073
4074fn resolved_bounded_binomial_link(
4075 family: &LikelihoodSpec,
4076 latent_cloglog_state: Option<&LatentCLogLogState>,
4077 mixture_link_state: Option<&MixtureLinkState>,
4078 sas_link_state: Option<&SasLinkState>,
4079) -> InverseLink {
4080 match &family.link {
4081 InverseLink::LatentCLogLog(_) => latent_cloglog_state
4082 .copied()
4083 .map(InverseLink::LatentCLogLog)
4084 .unwrap_or_else(|| family.link.clone()),
4085 InverseLink::Mixture(_) => mixture_link_state
4086 .cloned()
4087 .map(InverseLink::Mixture)
4088 .unwrap_or_else(|| family.link.clone()),
4089 InverseLink::Sas(_) => sas_link_state
4090 .copied()
4091 .map(InverseLink::Sas)
4092 .unwrap_or_else(|| family.link.clone()),
4093 InverseLink::BetaLogistic(_) => sas_link_state
4094 .copied()
4095 .map(InverseLink::BetaLogistic)
4096 .unwrap_or_else(|| family.link.clone()),
4097 InverseLink::Standard(_) => family.link.clone(),
4098 }
4099}
4100
4101fn binomial_natural_jet(
4102 row: usize,
4103 eta: f64,
4104 link: &InverseLink,
4105) -> Result<BernoulliNaturalJet, EstimationError> {
4106 match link {
4107 InverseLink::Standard(StandardLink::Probit) => Ok(probit_natural_jet(eta)),
4108 InverseLink::Standard(StandardLink::CLogLog) => Ok(cloglog_natural_jet(eta)),
4109 InverseLink::Standard(StandardLink::LogLog) => Ok(loglog_natural_jet(eta)),
4110 InverseLink::Standard(StandardLink::Cauchit) => Ok(cauchit_natural_jet(eta)),
4111 _ => generic_bernoulli_natural_jet(row, eta, link),
4112 }
4113}
4114
4115fn exact_logit_observation_row(
4116 row: usize,
4117 y: f64,
4118 weight: f64,
4119 eta: f64,
4120) -> Result<ExactStandardObservationRow, EstimationError> {
4121 let tail = (-eta.abs()).exp();
4122 let (mu, one_minus_mu) = if eta >= 0.0 {
4123 let q = tail / (1.0 + tail);
4124 (1.0 - q, q)
4125 } else {
4126 let p = tail / (1.0 + tail);
4127 (p, 1.0 - p)
4128 };
4129 if weight == 0.0 {
4130 return Ok(ExactStandardObservationRow::zero_weight(mu));
4131 }
4132 let log_fisher =
4133 -gam_linalg::utils::stable_softplus(eta) - gam_linalg::utils::stable_softplus(-eta);
4134 let fisherweight = weighted_positive_from_log(weight, log_fisher);
4135 if !(fisherweight.is_finite() && fisherweight > 0.0) {
4136 return Err(bounded_row_error(
4137 row,
4138 "bounded logit Fisher weight",
4139 eta,
4140 fisherweight,
4141 ));
4142 }
4143 let residual = if eta >= 0.0 {
4144 if y == 1.0 {
4145 one_minus_mu
4146 } else {
4147 (y - 1.0) + one_minus_mu
4148 }
4149 } else {
4150 y - mu
4151 };
4152 let log_likelihood_unit = if eta >= 0.0 {
4153 -(1.0 - y) * eta - gam_linalg::utils::stable_softplus(-eta)
4154 } else {
4155 y * eta - gam_linalg::utils::stable_softplus(eta)
4156 };
4157 certify_bounded_row(
4158 row,
4159 eta,
4160 ExactStandardObservationRow {
4161 mu,
4162 score: weight * residual,
4163 fisherweight,
4164 neghessian_eta: fisherweight,
4165 neghessian_eta_derivative: fisherweight * (one_minus_mu - mu),
4166 log_likelihood: weight * log_likelihood_unit,
4167 },
4168 )
4169}
4170
4171fn exact_noncanonical_binomial_observation_row(
4172 row: usize,
4173 y: f64,
4174 weight: f64,
4175 eta: f64,
4176 link: &InverseLink,
4177) -> Result<ExactStandardObservationRow, EstimationError> {
4178 let jet = binomial_natural_jet(row, eta, link)?;
4179 if weight == 0.0 {
4180 return Ok(ExactStandardObservationRow::zero_weight(jet.mu));
4181 }
4182 let fisherweight = weighted_positive_from_log(weight, jet.log_fisher);
4183 if !(fisherweight.is_finite() && fisherweight > 0.0) {
4184 return Err(bounded_row_error(
4185 row,
4186 "bounded binomial Fisher weight",
4187 eta,
4188 fisherweight,
4189 ));
4190 }
4191 let log_likelihood = weight * convex_combination(y, jet.log_mu[0], jet.log_one_minus_mu[0]);
4192 let score = weight * convex_combination(y, jet.log_mu[1], jet.log_one_minus_mu[1]);
4193 let neghessian_eta = -weight * convex_combination(y, jet.log_mu[2], jet.log_one_minus_mu[2]);
4194 let neghessian_eta_derivative =
4195 -weight * convex_combination(y, jet.log_mu[3], jet.log_one_minus_mu[3]);
4196 certify_bounded_row(
4197 row,
4198 eta,
4199 ExactStandardObservationRow {
4200 mu: jet.mu,
4201 score,
4202 fisherweight,
4203 neghessian_eta,
4204 neghessian_eta_derivative,
4205 log_likelihood,
4206 },
4207 )
4208}
4209
4210#[inline]
4211fn eta_exprel(rate: f64, eta: f64) -> f64 {
4212 (rate * eta).exp_m1() / rate
4213}
4214
4215fn validate_bounded_observation_inputs(
4216 likelihood: &gam_spec::GlmLikelihoodSpec,
4217 y: &Array1<f64>,
4218 weights: &Array1<f64>,
4219 eta: &Array1<f64>,
4220) -> Result<gam_spec::ResolvedLikelihoodScale, EstimationError> {
4221 let family = &likelihood.spec;
4222 if weights.len() != y.len() || eta.len() != y.len() {
4223 crate::bail_invalid_estim!(
4224 "bounded family observation size mismatch: y={}, weights={}, eta={}",
4225 y.len(),
4226 weights.len(),
4227 eta.len()
4228 );
4229 }
4230 if !LikelihoodSpec::is_legal_cell(&family.response, &family.link) {
4231 crate::bail_invalid_estim!(
4232 "bounded family received illegal likelihood cell response={} link={}",
4233 family.response.name(),
4234 family.link.link_function().name()
4235 );
4236 }
4237 let resolved_scale = likelihood
4238 .resolved_scale()
4239 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
4240 match &family.response {
4241 ResponseFamily::Tweedie { p } if !(p.is_finite() && *p > 1.0 && *p < 2.0) => {
4242 crate::bail_invalid_estim!(
4243 "bounded Tweedie power must be finite and strictly inside (1, 2), got {p}"
4244 );
4245 }
4246 ResponseFamily::NegativeBinomial { theta, .. } if !(theta.is_finite() && *theta > 0.0) => {
4247 crate::bail_invalid_estim!(
4248 "bounded negative-binomial theta must be finite and positive, got {theta}"
4249 );
4250 }
4251 _ => {}
4252 }
4253 for (i, &wi) in weights.iter().enumerate() {
4256 if !(wi.is_finite() && wi >= 0.0) {
4257 return Err(EstimationError::InvalidInput(format!(
4258 "bounded-family row {} has invalid prior weight {wi:?}; expected finite weight >= 0",
4259 i + 1
4260 )));
4261 }
4262 }
4263 for i in 0..y.len() {
4264 let wi = weights[i];
4265 if wi == 0.0 {
4266 continue;
4267 }
4268 if !eta[i].is_finite() {
4269 return Err(bounded_row_error(i, "linear predictor", eta[i], eta[i]));
4270 }
4271 if !y[i].is_finite() {
4272 return Err(bounded_row_error(
4273 i,
4274 "bounded-family response",
4275 eta[i],
4276 y[i],
4277 ));
4278 }
4279 let yi = y[i];
4280 let valid = match &family.response {
4281 ResponseFamily::Gaussian => yi.is_finite(),
4282 ResponseFamily::Binomial => yi.is_finite() && (0.0..=1.0).contains(&yi),
4283 ResponseFamily::Poisson | ResponseFamily::NegativeBinomial { .. } => {
4284 yi.is_finite() && yi >= 0.0 && (yi - yi.round()).abs() <= 1e-9
4285 }
4286 ResponseFamily::Tweedie { .. } => yi.is_finite() && yi >= 0.0,
4287 ResponseFamily::Gamma => yi.is_finite() && yi > 0.0,
4288 ResponseFamily::Beta { .. } | ResponseFamily::RoystonParmar => false,
4289 };
4290 if !valid {
4291 return Err(bounded_row_error(i, "bounded-family response", eta[i], yi));
4292 }
4293 }
4294 Ok(resolved_scale)
4295}
4296
4297fn exact_standard_observation_row(
4298 likelihood: &gam_spec::GlmLikelihoodSpec,
4299 resolved_scale: gam_spec::ResolvedLikelihoodScale,
4300 binomial_link: &InverseLink,
4301 row: usize,
4302 y: f64,
4303 weight: f64,
4304 eta: f64,
4305) -> Result<ExactStandardObservationRow, EstimationError> {
4306 if weight == 0.0 {
4307 return Ok(ExactStandardObservationRow::zero_weight(0.0));
4308 }
4309 let family = &likelihood.spec;
4310 match &family.response {
4311 ResponseFamily::Gaussian => {
4312 let scaled_weight = match resolved_scale {
4313 gam_spec::ResolvedLikelihoodScale::ProfiledGaussian => weight,
4314 gam_spec::ResolvedLikelihoodScale::FixedGaussian { phi } => {
4315 crate::gamlss::scaled_positive_product_quotient(weight, 1.0, 1.0, phi.value())
4316 }
4317 _ => {
4318 crate::bail_invalid_estim!(
4319 "bounded Gaussian received a non-Gaussian resolved scale"
4320 );
4321 }
4322 };
4323 if !(scaled_weight.is_finite() && scaled_weight > 0.0) {
4324 return Err(bounded_row_error(
4325 row,
4326 "bounded Gaussian dispersion-scaled weight",
4327 eta,
4328 scaled_weight,
4329 ));
4330 }
4331 let residual = y - eta;
4332 let loss = if residual == 0.0 {
4333 0.0
4334 } else {
4335 crate::gamlss::scaled_positive_product_quotient(
4336 scaled_weight,
4337 residual.abs(),
4338 residual.abs(),
4339 2.0,
4340 )
4341 };
4342 certify_bounded_row(
4343 row,
4344 eta,
4345 ExactStandardObservationRow {
4346 mu: eta,
4347 score: scaled_weight * residual,
4348 fisherweight: scaled_weight,
4349 neghessian_eta: scaled_weight,
4350 neghessian_eta_derivative: 0.0,
4351 log_likelihood: -loss,
4352 },
4353 )
4354 }
4355 ResponseFamily::Binomial
4356 if matches!(binomial_link, InverseLink::Standard(StandardLink::Logit)) =>
4357 {
4358 exact_logit_observation_row(row, y, weight, eta)
4359 }
4360 ResponseFamily::Binomial => {
4361 exact_noncanonical_binomial_observation_row(row, y, weight, eta, binomial_link)
4362 }
4363 ResponseFamily::Poisson => {
4364 let mu = inverse_link_jet_for_inverse_link(&family.link, eta)?.mu;
4365 let fisherweight = weight * mu;
4366 let score = weight * (y - mu);
4367 let raw_log_likelihood = y.mul_add(eta, -mu);
4368 let log_likelihood = if raw_log_likelihood.is_finite() {
4369 weight * raw_log_likelihood
4370 } else {
4371 weighted_product3(weight, y, eta) - weight * mu
4372 };
4373 if !(fisherweight.is_finite() && fisherweight > 0.0) {
4374 return Err(bounded_row_error(
4375 row,
4376 "bounded Poisson Fisher weight",
4377 eta,
4378 fisherweight,
4379 ));
4380 }
4381 certify_bounded_row(
4382 row,
4383 eta,
4384 ExactStandardObservationRow {
4385 mu,
4386 score,
4387 fisherweight,
4388 neghessian_eta: fisherweight,
4389 neghessian_eta_derivative: fisherweight,
4390 log_likelihood,
4391 },
4392 )
4393 }
4394 ResponseFamily::Gamma => {
4395 let mu = inverse_link_jet_for_inverse_link(&family.link, eta)?.mu;
4396 let shape = resolved_scale
4397 .gamma_shape()
4398 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
4399 let weighted_shape = weight * shape;
4400 if !(weighted_shape.is_finite() && weighted_shape > 0.0) {
4401 return Err(bounded_row_error(
4402 row,
4403 "bounded Gamma shape-scaled weight",
4404 eta,
4405 weighted_shape,
4406 ));
4407 }
4408 let weighted_ratio =
4409 crate::gamlss::scaled_positive_product_quotient(weight, y, shape, mu);
4410 if !(weighted_ratio.is_finite() && weighted_ratio > 0.0) {
4411 return Err(bounded_row_error(
4412 row,
4413 "bounded Gamma observed Hessian",
4414 eta,
4415 weighted_ratio,
4416 ));
4417 }
4418 certify_bounded_row(
4419 row,
4420 eta,
4421 ExactStandardObservationRow {
4422 mu,
4423 score: weighted_ratio - weighted_shape,
4424 fisherweight: weighted_shape,
4425 neghessian_eta: weighted_ratio,
4426 neghessian_eta_derivative: -weighted_ratio,
4427 log_likelihood: -weighted_ratio - weighted_shape * eta,
4428 },
4429 )
4430 }
4431 ResponseFamily::Tweedie { p } => {
4432 let p = *p;
4433 let mu = inverse_link_jet_for_inverse_link(&family.link, eta)?.mu;
4434 let phi = resolved_scale
4435 .tweedie_phi()
4436 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
4437 let weight = crate::gamlss::scaled_positive_product_quotient(weight, 1.0, 1.0, phi);
4438 if !(weight.is_finite() && weight > 0.0) {
4439 return Err(bounded_row_error(
4440 row,
4441 "bounded Tweedie dispersion-scaled weight",
4442 eta,
4443 weight,
4444 ));
4445 }
4446 let a = ((1.0 - p) * eta).exp();
4447 let b = ((2.0 - p) * eta).exp();
4448 let score_unit = y.mul_add(a, -b);
4449 let score = if score_unit.is_finite() {
4450 weight * score_unit
4451 } else {
4452 weighted_product3(weight, y, a) - weight * b
4453 };
4454 let fisherweight = weight * b;
4455 let observed_unit = (p - 1.0) * y * a + (2.0 - p) * b;
4456 let neghessian_eta = if observed_unit.is_finite() {
4457 weight * observed_unit
4458 } else {
4459 weighted_product3(weight * (p - 1.0), y, a) + weight * (2.0 - p) * b
4460 };
4461 let observed_derivative_unit = -(p - 1.0).powi(2) * y * a + (2.0 - p).powi(2) * b;
4462 let neghessian_eta_derivative = if observed_derivative_unit.is_finite() {
4463 weight * observed_derivative_unit
4464 } else {
4465 -weighted_product3(weight * (p - 1.0).powi(2), y, a)
4466 + weight * (2.0 - p).powi(2) * b
4467 };
4468 let q_left = eta_exprel(1.0 - p, eta);
4471 let q_right = eta_exprel(2.0 - p, eta);
4472 let q = y.mul_add(q_left, -q_right);
4473 let log_likelihood = if q.is_finite() {
4474 weight * q
4475 } else {
4476 weighted_product3(weight, y, q_left) - weight * q_right
4477 };
4478 if !(fisherweight.is_finite() && fisherweight > 0.0) {
4479 return Err(bounded_row_error(
4480 row,
4481 "bounded Tweedie Fisher weight",
4482 eta,
4483 fisherweight,
4484 ));
4485 }
4486 certify_bounded_row(
4487 row,
4488 eta,
4489 ExactStandardObservationRow {
4490 mu,
4491 score,
4492 fisherweight,
4493 neghessian_eta,
4494 neghessian_eta_derivative,
4495 log_likelihood,
4496 },
4497 )
4498 }
4499 ResponseFamily::NegativeBinomial { .. } => {
4500 let theta = resolved_scale
4501 .negative_binomial_theta()
4502 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
4503 let mu = inverse_link_jet_for_inverse_link(&family.link, eta)?.mu;
4504 let log_theta = theta.ln();
4505 let delta = eta - log_theta;
4506 let log_q = -gam_linalg::utils::stable_softplus(-delta);
4507 let log_r = -gam_linalg::utils::stable_softplus(delta);
4508 let q = log_q.exp();
4509 let r = log_r.exp();
4510 let y_r = if y == 0.0 {
4511 0.0
4512 } else {
4513 (y.ln() + log_r).exp()
4514 };
4515 let theta_q = (log_theta + log_q).exp();
4516 let score = weight * (y_r - theta_q);
4517 let fisherweight = weighted_positive_from_log(weight, log_theta + log_q);
4518 let log_qr = log_q + log_r;
4519 let observed_y = if y == 0.0 {
4520 0.0
4521 } else {
4522 weighted_positive_from_log(weight, y.ln() + log_qr)
4523 };
4524 let observed_theta = weighted_positive_from_log(weight, log_theta + log_qr);
4525 let neghessian_eta = observed_y + observed_theta;
4526 let neghessian_eta_derivative = neghessian_eta * (r - q);
4527 let softplus_tail = if delta >= 0.0 {
4528 gam_linalg::utils::stable_softplus(-delta)
4529 } else {
4530 gam_linalg::utils::stable_softplus(delta)
4531 };
4532 let log_likelihood = if delta >= 0.0 {
4533 -weighted_product3(weight, theta, delta)
4534 - weighted_product3(weight, y, softplus_tail)
4535 - weighted_product3(weight, theta, softplus_tail)
4536 } else {
4537 weighted_product3(weight, y, delta)
4538 - weighted_product3(weight, y, softplus_tail)
4539 - weighted_product3(weight, theta, softplus_tail)
4540 };
4541 if !(fisherweight.is_finite() && fisherweight > 0.0) {
4542 return Err(bounded_row_error(
4543 row,
4544 "bounded negative-binomial Fisher weight",
4545 eta,
4546 fisherweight,
4547 ));
4548 }
4549 certify_bounded_row(
4550 row,
4551 eta,
4552 ExactStandardObservationRow {
4553 mu,
4554 score,
4555 fisherweight,
4556 neghessian_eta,
4557 neghessian_eta_derivative,
4558 log_likelihood,
4559 },
4560 )
4561 }
4562 ResponseFamily::Beta { .. } => {
4563 crate::bail_invalid_estim!("bounded linear terms are not supported for BetaLogit fits");
4564 }
4565 ResponseFamily::RoystonParmar => {
4566 crate::bail_invalid_estim!(
4567 "bounded linear terms are not supported for survival model fits"
4568 );
4569 }
4570 }
4571}
4572
4573fn evaluate_resolved_standard_family_observations(
4574 likelihood: &gam_spec::GlmLikelihoodSpec,
4575 latent_cloglog_state: Option<&LatentCLogLogState>,
4576 mixture_link_state: Option<&MixtureLinkState>,
4577 sas_link_state: Option<&SasLinkState>,
4578 y: &Array1<f64>,
4579 weights: &Array1<f64>,
4580 eta: &Array1<f64>,
4581) -> Result<StandardFamilyObservationState, EstimationError> {
4582 let n = y.len();
4583 let resolved_scale = validate_bounded_observation_inputs(likelihood, y, weights, eta)?;
4584 let family = &likelihood.spec;
4585 let binomial_link = resolved_bounded_binomial_link(
4586 &family,
4587 latent_cloglog_state,
4588 mixture_link_state,
4589 sas_link_state,
4590 );
4591
4592 let mut score = Array1::<f64>::zeros(n);
4593 let mut fisherweight = Array1::<f64>::zeros(n);
4594 let mut neghessian_eta = Array1::<f64>::zeros(n);
4595 let mut neghessian_eta_derivative = Array1::<f64>::zeros(n);
4596 let mut log_likelihood = 0.0;
4597 let mut log_likelihood_compensation = 0.0;
4598
4599 for i in 0..n {
4600 let row = exact_standard_observation_row(
4601 likelihood,
4602 resolved_scale,
4603 &binomial_link,
4604 i,
4605 y[i],
4606 weights[i],
4607 eta[i],
4608 )?;
4609 score[i] = row.score;
4610 fisherweight[i] = row.fisherweight;
4611 neghessian_eta[i] = row.neghessian_eta;
4612 neghessian_eta_derivative[i] = row.neghessian_eta_derivative;
4613 let adjusted = row.log_likelihood - log_likelihood_compensation;
4614 let updated = log_likelihood + adjusted;
4615 log_likelihood_compensation = (updated - log_likelihood) - adjusted;
4616 log_likelihood = updated;
4617 if !log_likelihood.is_finite() {
4618 return Err(bounded_row_error(
4619 i,
4620 "bounded-family cumulative log likelihood",
4621 eta[i],
4622 log_likelihood,
4623 ));
4624 }
4625 }
4626
4627 Ok(StandardFamilyObservationState {
4628 eta: eta.clone(),
4629 score,
4630 fisherweight,
4631 neghessian_eta,
4632 neghessian_eta_derivative,
4633 log_likelihood,
4634 })
4635}
4636
4637fn evaluate_standard_familyobservations(
4642 family: LikelihoodSpec,
4643 latent_cloglog_state: Option<&LatentCLogLogState>,
4644 mixture_link_state: Option<&MixtureLinkState>,
4645 sas_link_state: Option<&SasLinkState>,
4646 y: &Array1<f64>,
4647 weights: &Array1<f64>,
4648 eta: &Array1<f64>,
4649) -> Result<StandardFamilyObservationState, EstimationError> {
4650 let likelihood = gam_spec::GlmLikelihoodSpec::canonical(family);
4651 evaluate_resolved_standard_family_observations(
4652 &likelihood,
4653 latent_cloglog_state,
4654 mixture_link_state,
4655 sas_link_state,
4656 y,
4657 weights,
4658 eta,
4659 )
4660}
4661
4662fn exact_standard_working_response(
4663 state: &StandardFamilyObservationState,
4664) -> Result<Array1<f64>, EstimationError> {
4665 let mut out = state.eta.clone();
4666 for i in 0..out.len() {
4667 let weight = state.fisherweight[i];
4668 let score = state.score[i];
4669 if weight == 0.0 {
4670 if score != 0.0 {
4671 return Err(bounded_row_error(
4672 i,
4673 "zero-Fisher row with nonzero score",
4674 state.eta[i],
4675 score,
4676 ));
4677 }
4678 continue;
4679 }
4680 let increment = score / weight;
4681 let value = out[i] + increment;
4682 if !increment.is_finite() || !value.is_finite() {
4683 return Err(bounded_row_error(
4684 i,
4685 "bounded-family working response",
4686 state.eta[i],
4687 value,
4688 ));
4689 }
4690 out[i] = value;
4691 }
4692 Ok(out)
4693}
4694
4695#[derive(Clone, Copy, Debug, PartialEq, Eq)]
4696enum SpatialAdaptiveHyperKind {
4697 LogLambdaMagnitude,
4698 LogLambdaGradient,
4699 LogLambdaCurvature,
4700 LogEpsilonMagnitude,
4701 LogEpsilonGradient,
4702 LogEpsilonCurvature,
4703}
4704
4705impl SpatialAdaptiveHyperKind {
4706 fn component_index(self) -> usize {
4707 match self {
4708 SpatialAdaptiveHyperKind::LogLambdaMagnitude
4709 | SpatialAdaptiveHyperKind::LogEpsilonMagnitude => 0,
4710 SpatialAdaptiveHyperKind::LogLambdaGradient
4711 | SpatialAdaptiveHyperKind::LogEpsilonGradient => 1,
4712 SpatialAdaptiveHyperKind::LogLambdaCurvature
4713 | SpatialAdaptiveHyperKind::LogEpsilonCurvature => 2,
4714 }
4715 }
4716
4717 fn is_log_lambda(self) -> bool {
4718 matches!(
4719 self,
4720 SpatialAdaptiveHyperKind::LogLambdaMagnitude
4721 | SpatialAdaptiveHyperKind::LogLambdaGradient
4722 | SpatialAdaptiveHyperKind::LogLambdaCurvature
4723 )
4724 }
4725
4726 fn is_log_epsilon(self) -> bool {
4727 matches!(
4728 self,
4729 SpatialAdaptiveHyperKind::LogEpsilonMagnitude
4730 | SpatialAdaptiveHyperKind::LogEpsilonGradient
4731 | SpatialAdaptiveHyperKind::LogEpsilonCurvature
4732 )
4733 }
4734}
4735
4736#[derive(Clone, Copy, Debug)]
4737struct SpatialAdaptiveHyperSpec {
4738 cache_index: usize,
4739 kind: SpatialAdaptiveHyperKind,
4740}
4741
4742#[derive(Clone, Copy, Debug, PartialEq, Eq)]
4743enum SpatialAdaptiveExplicitSecondOrderKind {
4744 StructuralZero,
4745 LocalAlphaAlpha,
4746 LocalAlphaEta,
4747 SharedEtaEta,
4748}
4749
4750#[derive(Clone, Copy, Debug, PartialEq, Eq)]
4755enum AdaptiveComponent {
4756 Magnitude,
4757 Gradient,
4758 Curvature,
4759}
4760
4761impl AdaptiveComponent {
4762 fn from_index(index: usize) -> Result<Self, String> {
4763 match index {
4764 0 => Ok(AdaptiveComponent::Magnitude),
4765 1 => Ok(AdaptiveComponent::Gradient),
4766 2 => Ok(AdaptiveComponent::Curvature),
4767 other => Err(SmoothError::invalid_index(format!(
4768 "invalid adaptive component index {}",
4769 other
4770 ))
4771 .into()),
4772 }
4773 }
4774}
4775
4776#[derive(Clone, Copy, Debug, PartialEq, Eq)]
4781enum HyperDerivativeKind {
4782 Rho,
4784 LogEpsilonFirst,
4786 LogEpsilonSecond,
4788}
4789
4790#[derive(Clone, Copy, Debug, PartialEq, Eq)]
4795enum HyperDriftKind {
4796 Rho,
4797 LogEpsilon,
4798}
4799
4800impl SpatialAdaptiveHyperSpec {
4801 fn component_index(self) -> usize {
4802 self.kind.component_index()
4803 }
4804
4805 fn explicit_second_order_kind(self, other: Self) -> SpatialAdaptiveExplicitSecondOrderKind {
4806 if self.component_index() != other.component_index() {
4807 return SpatialAdaptiveExplicitSecondOrderKind::StructuralZero;
4808 }
4809 match (
4810 self.kind.is_log_lambda(),
4811 other.kind.is_log_lambda(),
4812 self.kind.is_log_epsilon(),
4813 other.kind.is_log_epsilon(),
4814 ) {
4815 (true, true, false, false) if self.cache_index == other.cache_index => {
4816 SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaAlpha
4817 }
4818 (true, false, false, true) | (false, true, true, false) => {
4819 SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaEta
4820 }
4821 (false, false, true, true) => SpatialAdaptiveExplicitSecondOrderKind::SharedEtaEta,
4822 _ => SpatialAdaptiveExplicitSecondOrderKind::StructuralZero,
4823 }
4824 }
4825}
4826
4827#[derive(Clone, Debug)]
4828struct SpatialAdaptiveTermHyperParams {
4829 lambda: [f64; 3],
4830 epsilon: [f64; 3],
4831}
4832
4833#[derive(Clone, Debug)]
4842struct ValidatedFixedQuadraticHessian {
4843 dense: Arc<Array2<f64>>,
4844}
4845
4846impl ValidatedFixedQuadraticHessian {
4847 fn try_from_dense(dense: Array2<f64>, coefficient_dim: usize) -> Result<Self, String> {
4848 gam_linalg::utils::validate_finite_symmetric_matrix(
4849 &dense,
4850 "spatial adaptive fixed quadratic Hessian",
4851 )
4852 .map_err(|error| error.to_string())?;
4853 PenaltyMatrix::Dense(dense.clone())
4854 .validate(coefficient_dim)
4855 .map_err(|error| {
4856 format!(
4857 "spatial adaptive fixed quadratic Hessian failed quadratic-form validation: {error}"
4858 )
4859 })?;
4860 Ok(Self {
4861 dense: Arc::new(dense),
4862 })
4863 }
4864
4865 fn zero(coefficient_dim: usize) -> Result<Self, String> {
4866 Self::try_from_dense(
4867 Array2::<f64>::zeros((coefficient_dim, coefficient_dim)),
4868 coefficient_dim,
4869 )
4870 }
4871
4872 fn as_dense(&self) -> &Array2<f64> {
4873 self.dense.as_ref()
4874 }
4875
4876 fn quadratic_terms(&self, beta: &Array1<f64>) -> Result<(f64, Array1<f64>), String> {
4877 if beta.len() != self.dense.ncols() {
4878 return Err(format!(
4879 "spatial adaptive fixed quadratic beta length {} does not match validated Hessian dimension {}",
4880 beta.len(),
4881 self.dense.ncols()
4882 ));
4883 }
4884 let gradient = self.dense.dot(beta);
4885 let value = 0.5 * beta.dot(&gradient);
4886 Ok((value, gradient))
4887 }
4888}
4889
4890#[derive(Clone)]
4891struct SpatialAdaptiveExactEvaluation {
4892 obs: StandardFamilyObservationState,
4893 adaptive_states: Vec<SpatialPenaltyExactState>,
4894 adaptive_penalty_value: f64,
4895 adaptive_penaltygradient: Array1<f64>,
4896 adaptive_penaltyhessian: Array2<f64>,
4897 fixed_quadraticvalue: f64,
4898 fixed_quadraticgradient: Array1<f64>,
4899 fixed_quadratic_hessian: ValidatedFixedQuadraticHessian,
4900}
4901
4902#[derive(Clone)]
4903struct CachedSpatialAdaptiveExactEvaluation {
4904 beta: Array1<f64>,
4905 eval: Arc<SpatialAdaptiveExactEvaluation>,
4906}
4907
4908impl SpatialAdaptiveExactEvaluation {
4909 fn total_penalty_value(&self) -> f64 {
4910 self.adaptive_penalty_value + self.fixed_quadraticvalue
4911 }
4912
4913 fn total_penaltygradient(&self) -> Array1<f64> {
4914 &self.adaptive_penaltygradient + &self.fixed_quadraticgradient
4915 }
4916
4917 fn total_penaltyhessian(&self) -> Array2<f64> {
4918 &self.adaptive_penaltyhessian + self.fixed_quadratic_hessian.as_dense()
4919 }
4920
4921 fn totalobjectivehessian(&self, design: &Array2<f64>) -> Result<Array2<f64>, String> {
4922 let mut out = xt_diag_x_dense(design.view(), self.obs.neghessian_eta.view())?;
4923 out += &self.total_penaltyhessian();
4924 Ok(out)
4925 }
4926}
4927
4928#[derive(Clone)]
4929struct SpatialAdaptiveExactFamily {
4930 family: LikelihoodSpec,
4931 latent_cloglog_state: Option<LatentCLogLogState>,
4932 mixture_link_state: Option<MixtureLinkState>,
4933 sas_link_state: Option<SasLinkState>,
4934 y: Arc<Array1<f64>>,
4935 weights: Arc<Array1<f64>>,
4936 design: Arc<Array2<f64>>,
4937 offset: Arc<Array1<f64>>,
4938 linear_constraints: Option<LinearInequalityConstraints>,
4939 runtime_caches: Arc<Vec<SpatialOperatorRuntimeCache>>,
4940 adaptive_params: Vec<SpatialAdaptiveTermHyperParams>,
4941 fixed_quadratic_hessian: ValidatedFixedQuadraticHessian,
4942 hyperspecs: Arc<Vec<SpatialAdaptiveHyperSpec>>,
4943 exact_eval_cache: Arc<Mutex<Option<CachedSpatialAdaptiveExactEvaluation>>>,
4944}
4945
4946impl SpatialAdaptiveExactFamily {
4947 fn with_adaptive_params(
4948 &self,
4949 adaptive_params: Vec<SpatialAdaptiveTermHyperParams>,
4950 fixed_quadratic_hessian: ValidatedFixedQuadraticHessian,
4951 ) -> Self {
4952 Self {
4953 family: self.family.clone(),
4954 latent_cloglog_state: self.latent_cloglog_state,
4955 mixture_link_state: self.mixture_link_state.clone(),
4956 sas_link_state: self.sas_link_state,
4957 y: self.y.clone(),
4958 weights: self.weights.clone(),
4959 design: self.design.clone(),
4960 offset: self.offset.clone(),
4961 linear_constraints: self.linear_constraints.clone(),
4962 runtime_caches: self.runtime_caches.clone(),
4963 adaptive_params,
4964 fixed_quadratic_hessian,
4965 hyperspecs: self.hyperspecs.clone(),
4966 exact_eval_cache: Arc::new(Mutex::new(None)),
4967 }
4968 }
4969
4970 fn total_eta(&self, beta: &Array1<f64>) -> Array1<f64> {
4971 gam_linalg::faer_ndarray::fast_av(self.design.as_ref(), beta) + self.offset.as_ref()
4972 }
4973
4974 fn fixed_quadratic_terms(
4975 &self,
4976 beta: &Array1<f64>,
4977 ) -> Result<(f64, Array1<f64>), String> {
4978 self.fixed_quadratic_hessian.quadratic_terms(beta)
4979 }
4980
4981 fn adaptive_penalty_value_only(&self, beta: &Array1<f64>) -> Result<f64, String> {
4982 let mut penalty_value = 0.0;
4983 for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
4984 let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
4985 format!(
4986 "missing adaptive parameter block for cache {}",
4987 cache.termname
4988 )
4989 })?;
4990 let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
4991 let state =
4992 SpatialPenaltyExactState::from_beta_local(beta_local, cache, params.epsilon)
4993 .map_err(|e| e.to_string())?;
4994 penalty_value += params.lambda[0] * state.magnitude.penalty_value();
4995 penalty_value += params.lambda[1] * state.gradient.penalty_value();
4996 penalty_value += params.lambda[2] * state.curvature.penalty_value();
4997 }
4998 Ok(penalty_value)
4999 }
5000
5001 fn zero_hyper_parts(&self) -> (Array1<f64>, Array2<f64>) {
5002 let total_dim = self.design.ncols();
5003 (
5004 Array1::<f64>::zeros(total_dim),
5005 Array2::<f64>::zeros((total_dim, total_dim)),
5006 )
5007 }
5008
5009 fn embed_local_hyper_parts(
5010 &self,
5011 coeff_range: &Range<usize>,
5012 local_grad: &Array1<f64>,
5013 local_hess: &Array2<f64>,
5014 ) -> (Array1<f64>, Array2<f64>) {
5015 let (mut beta_mixed, mut betahessian) = self.zero_hyper_parts();
5016 beta_mixed
5017 .slice_mut(s![coeff_range.clone()])
5018 .assign(local_grad);
5019 betahessian
5020 .slice_mut(s![coeff_range.clone(), coeff_range.clone()])
5021 .assign(local_hess);
5022 (beta_mixed, betahessian)
5023 }
5024
5025 fn embed_local_hyper_hessian(
5026 &self,
5027 coeff_range: &Range<usize>,
5028 local_hess: &Array2<f64>,
5029 ) -> Array2<f64> {
5030 let total_dim = self.design.ncols();
5031 let mut out = Array2::<f64>::zeros((total_dim, total_dim));
5032 out.slice_mut(s![coeff_range.clone(), coeff_range.clone()])
5033 .assign(local_hess);
5034 out
5035 }
5036
5037 fn adaptive_block_eval(
5046 &self,
5047 eval: &SpatialAdaptiveExactEvaluation,
5048 cache_idx: usize,
5049 component: AdaptiveComponent,
5050 derivative: HyperDerivativeKind,
5051 ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
5052 let cache = self
5053 .runtime_caches
5054 .get(cache_idx)
5055 .ok_or_else(|| format!("adaptive cache index {} out of bounds", cache_idx))?;
5056 let params = self
5057 .adaptive_params
5058 .get(cache_idx)
5059 .ok_or_else(|| format!("adaptive hyperparameter block {} out of bounds", cache_idx))?;
5060 let state = eval
5061 .adaptive_states
5062 .get(cache_idx)
5063 .ok_or_else(|| format!("adaptive exact state index {} out of bounds", cache_idx))?;
5064
5065 let (objective_local, beta_mixed_local, betahessian_local) = match component {
5066 AdaptiveComponent::Magnitude => {
5067 let lambda = params.lambda[0];
5068 let mag = &state.magnitude;
5069 let (objective, gradient_coeff, hessian_diag) = match derivative {
5070 HyperDerivativeKind::Rho => (
5071 mag.penalty_value(),
5072 mag.betagradient_coeff(),
5073 mag.betahessian_diag(),
5074 ),
5075 HyperDerivativeKind::LogEpsilonFirst => (
5076 mag.log_epsilon_gradient_terms().sum(),
5077 mag.log_epsilon_betagradient_coeff(),
5078 mag.log_epsilon_betahessian_diag(),
5079 ),
5080 HyperDerivativeKind::LogEpsilonSecond => (
5081 mag.log_epsilon_hessian_terms().sum(),
5082 mag.log_epsilon_beta_mixed_second_coeff(),
5083 mag.log_epsilon_betahessian_second_diag(),
5084 ),
5085 };
5086 (
5087 lambda * objective,
5088 lambda * scalar_operatorgradient(&cache.d0, &gradient_coeff),
5089 lambda * scalar_operatorhessian(&cache.d0, &hessian_diag),
5090 )
5091 }
5092 AdaptiveComponent::Gradient => {
5093 let lambda = params.lambda[1];
5094 let grad = &state.gradient;
5095 let (objective, gradient_blocks, hessian_blocks) = match derivative {
5096 HyperDerivativeKind::Rho => (
5097 grad.penalty_value(),
5098 grad.betagradient_blocks(),
5099 grad.betahessian_blocks(),
5100 ),
5101 HyperDerivativeKind::LogEpsilonFirst => (
5102 grad.log_epsilon_gradient_terms().sum(),
5103 grad.log_epsilon_betagradient_blocks(),
5104 grad.log_epsilon_betahessian_blocks(),
5105 ),
5106 HyperDerivativeKind::LogEpsilonSecond => (
5107 grad.log_epsilon_hessian_terms().sum(),
5108 grad.log_epsilon_beta_mixed_second_blocks(),
5109 grad.log_epsilon_betahessian_second_blocks(),
5110 ),
5111 };
5112 (
5113 lambda * objective,
5114 lambda
5115 * grouped_operatorgradient(&cache.d1, cache.dimension, &gradient_blocks)
5116 .map_err(|e| e.to_string())?,
5117 lambda
5118 * grouped_operatorhessian(&cache.d1, cache.dimension, &hessian_blocks)
5119 .map_err(|e| e.to_string())?,
5120 )
5121 }
5122 AdaptiveComponent::Curvature => {
5123 let lambda = params.lambda[2];
5124 let group = cache.dimension * cache.dimension;
5125 let curv = &state.curvature;
5126 let (objective, gradient_blocks, hessian_blocks) = match derivative {
5127 HyperDerivativeKind::Rho => (
5128 curv.penalty_value(),
5129 curv.betagradient_blocks(),
5130 curv.betahessian_blocks(),
5131 ),
5132 HyperDerivativeKind::LogEpsilonFirst => (
5133 curv.log_epsilon_gradient_terms().sum(),
5134 curv.log_epsilon_betagradient_blocks(),
5135 curv.log_epsilon_betahessian_blocks(),
5136 ),
5137 HyperDerivativeKind::LogEpsilonSecond => (
5138 curv.log_epsilon_hessian_terms().sum(),
5139 curv.log_epsilon_beta_mixed_second_blocks(),
5140 curv.log_epsilon_betahessian_second_blocks(),
5141 ),
5142 };
5143 (
5144 lambda * objective,
5145 lambda
5146 * grouped_operatorgradient(&cache.d2, group, &gradient_blocks)
5147 .map_err(|e| e.to_string())?,
5148 lambda
5149 * grouped_operatorhessian(&cache.d2, group, &hessian_blocks)
5150 .map_err(|e| e.to_string())?,
5151 )
5152 }
5153 };
5154
5155 let (beta_mixed, betahessian) = self.embed_local_hyper_parts(
5156 &cache.coeff_global_range,
5157 &beta_mixed_local,
5158 &betahessian_local,
5159 );
5160 Ok((objective_local, beta_mixed, betahessian))
5161 }
5162
5163 fn adaptive_shared_log_epsilon_parts(
5164 &self,
5165 eval: &SpatialAdaptiveExactEvaluation,
5166 component: usize,
5167 ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
5168 self.adaptive_shared_block_eval(eval, component, HyperDerivativeKind::LogEpsilonFirst)
5174 }
5175
5176 fn adaptive_shared_log_epsilon_second_parts(
5177 &self,
5178 eval: &SpatialAdaptiveExactEvaluation,
5179 component: usize,
5180 ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
5181 self.adaptive_shared_block_eval(eval, component, HyperDerivativeKind::LogEpsilonSecond)
5187 }
5188
5189 fn adaptive_shared_block_eval(
5194 &self,
5195 eval: &SpatialAdaptiveExactEvaluation,
5196 component: usize,
5197 derivative: HyperDerivativeKind,
5198 ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
5199 let component = AdaptiveComponent::from_index(component)?;
5200 let (mut score, mut hessian) = self.zero_hyper_parts();
5201 let mut objective = 0.0;
5202 for cache_idx in 0..self.runtime_caches.len() {
5203 let (local_objective, local_score, local_hessian) =
5204 self.adaptive_block_eval(eval, cache_idx, component, derivative)?;
5205 objective += local_objective;
5206 score += &local_score;
5207 hessian += &local_hessian;
5208 }
5209 Ok((objective, score, hessian))
5210 }
5211
5212 fn adaptive_shared_log_epsilon_drift(
5213 &self,
5214 eval: &SpatialAdaptiveExactEvaluation,
5215 component: usize,
5216 direction: &Array1<f64>,
5217 ) -> Result<Array2<f64>, String> {
5218 let component = AdaptiveComponent::from_index(component)?;
5222 let total_dim = self.design.ncols();
5223 let mut total = Array2::<f64>::zeros((total_dim, total_dim));
5224 for cache_idx in 0..self.runtime_caches.len() {
5225 total += &self.adaptive_block_drift_eval(
5226 eval,
5227 cache_idx,
5228 component,
5229 HyperDriftKind::LogEpsilon,
5230 direction,
5231 )?;
5232 }
5233 Ok(total)
5234 }
5235
5236 fn adaptive_explicit_second_order_parts(
5237 &self,
5238 eval: &SpatialAdaptiveExactEvaluation,
5239 left: SpatialAdaptiveHyperSpec,
5240 right: SpatialAdaptiveHyperSpec,
5241 ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
5242 match left.explicit_second_order_kind(right) {
5251 SpatialAdaptiveExplicitSecondOrderKind::StructuralZero => {
5252 let (score, hessian) = self.zero_hyper_parts();
5253 Ok((0.0, score, hessian))
5254 }
5255 SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaAlpha => self.adaptive_block_eval(
5256 eval,
5257 left.cache_index,
5258 AdaptiveComponent::from_index(left.component_index())?,
5259 HyperDerivativeKind::Rho,
5260 ),
5261 SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaEta => {
5262 let local_alpha = if left.kind.is_log_lambda() {
5263 left
5264 } else {
5265 right
5266 };
5267 self.adaptive_block_eval(
5268 eval,
5269 local_alpha.cache_index,
5270 AdaptiveComponent::from_index(local_alpha.component_index())?,
5271 HyperDerivativeKind::LogEpsilonFirst,
5272 )
5273 }
5274 SpatialAdaptiveExplicitSecondOrderKind::SharedEtaEta => {
5275 self.adaptive_shared_log_epsilon_second_parts(eval, left.component_index())
5276 }
5277 }
5278 }
5279
5280 fn adaptive_block_drift_eval(
5288 &self,
5289 eval: &SpatialAdaptiveExactEvaluation,
5290 cache_idx: usize,
5291 component: AdaptiveComponent,
5292 drift: HyperDriftKind,
5293 direction: &Array1<f64>,
5294 ) -> Result<Array2<f64>, String> {
5295 let cache = self
5296 .runtime_caches
5297 .get(cache_idx)
5298 .ok_or_else(|| format!("adaptive cache index {} out of bounds", cache_idx))?;
5299 let params = self
5300 .adaptive_params
5301 .get(cache_idx)
5302 .ok_or_else(|| format!("adaptive hyperparameter block {} out of bounds", cache_idx))?;
5303 let state = eval
5304 .adaptive_states
5305 .get(cache_idx)
5306 .ok_or_else(|| format!("adaptive exact state index {} out of bounds", cache_idx))?;
5307 let direction_local = direction.slice(s![cache.coeff_global_range.clone()]);
5308
5309 let local_hessian = match component {
5310 AdaptiveComponent::Magnitude => {
5311 let d0_u = cache.d0.dot(&direction_local);
5312 let mag = &state.magnitude;
5313 let diag = match drift {
5314 HyperDriftKind::Rho => mag.directionalhessian_diag(&d0_u),
5315 HyperDriftKind::LogEpsilon => {
5316 mag.log_epsilon_betahessian_directional_diag(&d0_u)
5317 }
5318 };
5319 params.lambda[0] * scalar_operatorhessian(&cache.d0, &diag)
5320 }
5321 AdaptiveComponent::Gradient => {
5322 let d1_u = cache.d1.dot(&direction_local);
5323 let direction_blocks = collocationgradient_blocks(&d1_u, cache.dimension)
5324 .map_err(|e| e.to_string())?;
5325 let grad = &state.gradient;
5326 let blocks = match drift {
5327 HyperDriftKind::Rho => grad.directionalhessian_blocks(&direction_blocks),
5328 HyperDriftKind::LogEpsilon => {
5329 grad.log_epsilon_betahessian_directional_blocks(&direction_blocks)
5330 }
5331 };
5332 params.lambda[1]
5333 * grouped_operatorhessian(&cache.d1, cache.dimension, &blocks)
5334 .map_err(|e| e.to_string())?
5335 }
5336 AdaptiveComponent::Curvature => {
5337 let group = cache.dimension * cache.dimension;
5338 let d2_u = cache.d2.dot(&direction_local);
5339 let direction_blocks =
5340 collocationhessian_blocks(&d2_u, cache.dimension).map_err(|e| e.to_string())?;
5341 let curv = &state.curvature;
5342 let blocks = match drift {
5343 HyperDriftKind::Rho => curv.directionalhessian_blocks(&direction_blocks),
5344 HyperDriftKind::LogEpsilon => {
5345 curv.log_epsilon_betahessian_directional_blocks(&direction_blocks)
5346 }
5347 };
5348 params.lambda[2]
5349 * grouped_operatorhessian(&cache.d2, group, &blocks)
5350 .map_err(|e| e.to_string())?
5351 }
5352 };
5353
5354 Ok(self.embed_local_hyper_hessian(&cache.coeff_global_range, &local_hessian))
5355 }
5356
5357 fn adaptive_hyper_parts(
5358 &self,
5359 eval: &SpatialAdaptiveExactEvaluation,
5360 hyper: SpatialAdaptiveHyperSpec,
5361 ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
5362 match hyper.kind {
5363 SpatialAdaptiveHyperKind::LogLambdaMagnitude
5366 | SpatialAdaptiveHyperKind::LogLambdaGradient
5367 | SpatialAdaptiveHyperKind::LogLambdaCurvature => self.adaptive_block_eval(
5368 eval,
5369 hyper.cache_index,
5370 AdaptiveComponent::from_index(hyper.component_index())?,
5371 HyperDerivativeKind::Rho,
5372 ),
5373 SpatialAdaptiveHyperKind::LogEpsilonMagnitude
5375 | SpatialAdaptiveHyperKind::LogEpsilonGradient
5376 | SpatialAdaptiveHyperKind::LogEpsilonCurvature => {
5377 self.adaptive_shared_log_epsilon_parts(eval, hyper.component_index())
5378 }
5379 }
5380 }
5381
5382 fn exact_evaluation_uncached(
5383 &self,
5384 beta: &Array1<f64>,
5385 ) -> Result<SpatialAdaptiveExactEvaluation, String> {
5386 let eta = self.total_eta(beta);
5387 let obs = evaluate_standard_familyobservations(
5388 self.family.clone(),
5389 self.latent_cloglog_state.as_ref(),
5390 self.mixture_link_state.as_ref(),
5391 self.sas_link_state.as_ref(),
5392 &self.y,
5393 &self.weights,
5394 &eta,
5395 )
5396 .map_err(|e| e.to_string())?;
5397 let p = beta.len();
5398 let mut penalty_value = 0.0;
5399 let mut penaltygradient = Array1::<f64>::zeros(p);
5400 let mut penaltyhessian = Array2::<f64>::zeros((p, p));
5401 let mut adaptive_states = Vec::with_capacity(self.runtime_caches.len());
5402
5403 for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
5404 let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
5405 format!(
5406 "missing adaptive parameter block for cache {}",
5407 cache.termname
5408 )
5409 })?;
5410 let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
5411 let state =
5412 SpatialPenaltyExactState::from_beta_local(beta_local, cache, params.epsilon)
5413 .map_err(|e| e.to_string())?;
5414
5415 let g0 = scalar_operatorgradient(&cache.d0, &state.magnitude.betagradient_coeff());
5416 let gg = grouped_operatorgradient(
5417 &cache.d1,
5418 cache.dimension,
5419 &state.gradient.betagradient_blocks(),
5420 )
5421 .map_err(|e| e.to_string())?;
5422 let gc = grouped_operatorgradient(
5423 &cache.d2,
5424 cache.dimension * cache.dimension,
5425 &state.curvature.betagradient_blocks(),
5426 )
5427 .map_err(|e| e.to_string())?;
5428 let h0 = scalar_operatorhessian(&cache.d0, &state.magnitude.betahessian_diag());
5429 let hg = grouped_operatorhessian(
5430 &cache.d1,
5431 cache.dimension,
5432 &state.gradient.betahessian_blocks(),
5433 )
5434 .map_err(|e| e.to_string())?;
5435 let hc = grouped_operatorhessian(
5436 &cache.d2,
5437 cache.dimension * cache.dimension,
5438 &state.curvature.betahessian_blocks(),
5439 )
5440 .map_err(|e| e.to_string())?;
5441
5442 let lambda0 = params.lambda[0];
5443 let lambdag = params.lambda[1];
5444 let lambdac = params.lambda[2];
5445
5446 penalty_value += lambda0 * state.magnitude.penalty_value();
5447 penalty_value += lambdag * state.gradient.penalty_value();
5448 penalty_value += lambdac * state.curvature.penalty_value();
5449
5450 let range = cache.coeff_global_range.clone();
5451 {
5452 let mut grad_local = penaltygradient.slice_mut(s![range.clone()]);
5453 grad_local += &(g0.mapv(|v| lambda0 * v));
5454 grad_local += &(gg.mapv(|v| lambdag * v));
5455 grad_local += &(gc.mapv(|v| lambdac * v));
5456 }
5457 {
5458 let mut h_local = penaltyhessian.slice_mut(s![range.clone(), range]);
5459 h_local += &h0.mapv(|v| lambda0 * v);
5460 h_local += &hg.mapv(|v| lambdag * v);
5461 h_local += &hc.mapv(|v| lambdac * v);
5462 }
5463
5464 adaptive_states.push(state);
5465 }
5466
5467 let (fixed_quadraticvalue, fixed_quadraticgradient) =
5468 self.fixed_quadratic_terms(beta)?;
5469 Ok(SpatialAdaptiveExactEvaluation {
5470 obs,
5471 adaptive_states,
5472 adaptive_penalty_value: penalty_value,
5473 adaptive_penaltygradient: penaltygradient,
5474 adaptive_penaltyhessian: penaltyhessian,
5475 fixed_quadraticvalue,
5476 fixed_quadraticgradient,
5477 fixed_quadratic_hessian: self.fixed_quadratic_hessian.clone(),
5478 })
5479 }
5480
5481 fn exact_evaluation(
5482 &self,
5483 beta: &Array1<f64>,
5484 ) -> Result<Arc<SpatialAdaptiveExactEvaluation>, String> {
5485 {
5486 let cache = self
5487 .exact_eval_cache
5488 .lock()
5489 .map_err(|_| "spatial adaptive exact-evaluation cache lock poisoned".to_string())?;
5490 if let Some(cached) = cache.as_ref()
5491 && cached.beta.len() == beta.len()
5492 && cached
5493 .beta
5494 .iter()
5495 .zip(beta.iter())
5496 .all(|(&left, &right)| left == right)
5497 {
5498 return Ok(Arc::clone(&cached.eval));
5499 }
5500 }
5501
5502 let eval = Arc::new(self.exact_evaluation_uncached(beta)?);
5503 let mut cache = self
5504 .exact_eval_cache
5505 .lock()
5506 .map_err(|_| "spatial adaptive exact-evaluation cache lock poisoned".to_string())?;
5507 *cache = Some(CachedSpatialAdaptiveExactEvaluation {
5508 beta: beta.clone(),
5509 eval: Arc::clone(&eval),
5510 });
5511 Ok(eval)
5512 }
5513
5514 fn exacthessian_directional_derivative_from_evaluation(
5515 &self,
5516 beta: &Array1<f64>,
5517 eval: &SpatialAdaptiveExactEvaluation,
5518 direction: &Array1<f64>,
5519 ) -> Result<Array2<f64>, String> {
5520 assert_eq!(
5521 beta.len(),
5522 direction.len(),
5523 "beta/direction length mismatch",
5524 );
5525 let d_eta = gam_linalg::faer_ndarray::fast_av(self.design.as_ref(), direction);
5526 let mut total = xt_diag_x_dense(
5527 self.design.view(),
5528 (&eval.obs.neghessian_eta_derivative * &d_eta).view(),
5529 )?;
5530 for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
5531 let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
5532 format!(
5533 "missing adaptive parameter block for cache {}",
5534 cache.termname
5535 )
5536 })?;
5537 let state = eval
5538 .adaptive_states
5539 .get(cache_idx)
5540 .ok_or_else(|| format!("missing adaptive state for cache {}", cache.termname))?;
5541 let direction_local = direction.slice(s![cache.coeff_global_range.clone()]);
5542 let d0_u = cache.d0.dot(&direction_local);
5543 let d1_u = cache.d1.dot(&direction_local);
5544 let d2_u = cache.d2.dot(&direction_local);
5545 let h0 =
5546 scalar_operatorhessian(&cache.d0, &state.magnitude.directionalhessian_diag(&d0_u))
5547 .mapv(|v| params.lambda[0] * v);
5548 let hg = grouped_operatorhessian(
5549 &cache.d1,
5550 cache.dimension,
5551 &state.gradient.directionalhessian_blocks(
5552 &collocationgradient_blocks(&d1_u, cache.dimension)
5553 .map_err(|e| e.to_string())?,
5554 ),
5555 )
5556 .map_err(|e| e.to_string())?
5557 .mapv(|v| params.lambda[1] * v);
5558 let hc = grouped_operatorhessian(
5559 &cache.d2,
5560 cache.dimension * cache.dimension,
5561 &state.curvature.directionalhessian_blocks(
5562 &collocationhessian_blocks(&d2_u, cache.dimension)
5563 .map_err(|e| e.to_string())?,
5564 ),
5565 )
5566 .map_err(|e| e.to_string())?
5567 .mapv(|v| params.lambda[2] * v);
5568 let range = cache.coeff_global_range.clone();
5569 let mut local = total.slice_mut(s![range.clone(), range]);
5570 local += &h0;
5571 local += &hg;
5572 local += &hc;
5573 }
5574 Ok(total)
5575 }
5576
5577 fn exacthessian_second_directional_derivative_from_evaluation(
5598 &self,
5599 eval: &SpatialAdaptiveExactEvaluation,
5600 direction_u: &Array1<f64>,
5601 direction_v: &Array1<f64>,
5602 ) -> Result<Option<Array2<f64>>, String> {
5603 let p = self.design.ncols();
5604 if eval.obs.neghessian_eta_derivative.iter().any(|&w| w != 0.0) {
5606 return Ok(None);
5607 }
5608 let mut total = Array2::<f64>::zeros((p, p));
5609 for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
5610 let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
5611 format!(
5612 "missing adaptive parameter block for cache {}",
5613 cache.termname
5614 )
5615 })?;
5616 let state = eval
5617 .adaptive_states
5618 .get(cache_idx)
5619 .ok_or_else(|| format!("missing adaptive state for cache {}", cache.termname))?;
5620 let u_local = direction_u.slice(s![cache.coeff_global_range.clone()]);
5621 let v_local = direction_v.slice(s![cache.coeff_global_range.clone()]);
5622
5623 let q0_u = cache.d0.dot(&u_local);
5625 let q0_v = cache.d0.dot(&v_local);
5626 let h0 = scalar_operatorhessian(
5627 &cache.d0,
5628 &state.magnitude.second_directionalhessian_diag(&q0_u, &q0_v),
5629 )
5630 .mapv(|x| params.lambda[0] * x);
5631
5632 let a1 = collocationgradient_blocks(&cache.d1.dot(&u_local), cache.dimension)
5634 .map_err(|e| e.to_string())?;
5635 let b1 = collocationgradient_blocks(&cache.d1.dot(&v_local), cache.dimension)
5636 .map_err(|e| e.to_string())?;
5637 let hg = grouped_operatorhessian(
5638 &cache.d1,
5639 cache.dimension,
5640 &state.gradient.second_directionalhessian_blocks(&a1, &b1),
5641 )
5642 .map_err(|e| e.to_string())?
5643 .mapv(|x| params.lambda[1] * x);
5644
5645 let a2 = collocationhessian_blocks(&cache.d2.dot(&u_local), cache.dimension)
5647 .map_err(|e| e.to_string())?;
5648 let b2 = collocationhessian_blocks(&cache.d2.dot(&v_local), cache.dimension)
5649 .map_err(|e| e.to_string())?;
5650 let hc = grouped_operatorhessian(
5651 &cache.d2,
5652 cache.dimension * cache.dimension,
5653 &state.curvature.second_directionalhessian_blocks(&a2, &b2),
5654 )
5655 .map_err(|e| e.to_string())?
5656 .mapv(|x| params.lambda[2] * x);
5657
5658 let range = cache.coeff_global_range.clone();
5659 let mut local = total.slice_mut(s![range.clone(), range]);
5660 local += &h0;
5661 local += &hg;
5662 local += &hc;
5663 }
5664 Ok(Some(total))
5665 }
5666}
5667
5668impl CustomFamily for SpatialAdaptiveExactFamily {
5669 fn joint_jeffreys_term_required(&self) -> bool {
5673 true
5674 }
5675
5676 fn joint_jeffreys_information_with_specs(
5713 &self,
5714 block_states: &[ParameterBlockState],
5715 specs: &[ParameterBlockSpec],
5716 ) -> Result<Option<Array2<f64>>, String> {
5717 let spec = expect_single_blockspec(specs, "spatial adaptive exact family")?;
5718 let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
5719 if spec.design.ncols() != beta.len() {
5720 return Err(SmoothError::dimension_mismatch(format!(
5721 "spatial adaptive Jeffreys information: spec design has {} columns, beta has {}",
5722 spec.design.ncols(),
5723 beta.len()
5724 ))
5725 .into());
5726 }
5727 let eval = self.exact_evaluation(beta)?;
5728 Ok(Some(xt_diag_x_dense(
5729 self.design.view(),
5730 eval.obs.neghessian_eta.view(),
5731 )?))
5732 }
5733
5734 fn joint_jeffreys_information_directional_derivative_with_specs(
5735 &self,
5736 block_states: &[ParameterBlockState],
5737 specs: &[ParameterBlockSpec],
5738 d_beta_flat: &Array1<f64>,
5739 ) -> Result<Option<Array2<f64>>, String> {
5740 let spec = expect_single_blockspec(specs, "spatial adaptive exact family")?;
5746 let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
5747 if spec.design.ncols() != d_beta_flat.len() {
5748 return Err(SmoothError::dimension_mismatch(format!(
5749 "spatial adaptive Jeffreys directional derivative: spec design has {} columns, direction has {}",
5750 spec.design.ncols(),
5751 d_beta_flat.len()
5752 ))
5753 .into());
5754 }
5755 let eval = self.exact_evaluation(beta)?;
5756 let d_eta = gam_linalg::faer_ndarray::fast_av(self.design.as_ref(), d_beta_flat);
5757 Ok(Some(xt_diag_x_dense(
5758 self.design.view(),
5759 (&eval.obs.neghessian_eta_derivative * &d_eta).view(),
5760 )?))
5761 }
5762
5763 fn joint_jeffreys_information_second_directional_derivative_with_specs(
5764 &self,
5765 block_states: &[ParameterBlockState],
5766 specs: &[ParameterBlockSpec],
5767 d_beta_u_flat: &Array1<f64>,
5768 d_betav_flat: &Array1<f64>,
5769 ) -> Result<Option<Array2<f64>>, String> {
5770 let spec = expect_single_blockspec(specs, "spatial adaptive exact family")?;
5777 let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
5778 if spec.design.ncols() != beta.len()
5779 || d_beta_u_flat.len() != beta.len()
5780 || d_betav_flat.len() != beta.len()
5781 {
5782 return Err(SmoothError::dimension_mismatch(format!(
5783 "spatial adaptive Jeffreys second-direction length mismatch: spec cols={}, dirs=({}, {}), expected {}",
5784 spec.design.ncols(),
5785 d_beta_u_flat.len(),
5786 d_betav_flat.len(),
5787 beta.len()
5788 ))
5789 .into());
5790 }
5791 let eval = self.exact_evaluation(beta)?;
5792 if eval.obs.neghessian_eta_derivative.iter().any(|&w| w != 0.0) {
5793 return Ok(None);
5794 }
5795 Ok(Some(Array2::<f64>::zeros((beta.len(), beta.len()))))
5796 }
5797
5798 fn joint_jeffreys_information_matches_observed_hessian(&self) -> bool {
5799 false
5804 }
5805
5806 fn joint_jeffreys_information_depends_on_psi(&self) -> bool {
5807 false
5816 }
5817
5818 fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
5819 let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
5820 let eval = self.exact_evaluation(beta)?;
5821 let mut gradient = fast_atv(&self.design, &eval.obs.score);
5822 gradient -= &eval.total_penaltygradient();
5823 let mut hessian = xt_diag_x_dense(self.design.view(), eval.obs.neghessian_eta.view())?;
5824 hessian += &eval.total_penaltyhessian();
5825 Ok(FamilyEvaluation {
5826 log_likelihood: eval.obs.log_likelihood - eval.total_penalty_value(),
5827 blockworking_sets: vec![BlockWorkingSet::ExactNewton {
5828 gradient,
5829 hessian: SymmetricMatrix::Dense(hessian),
5830 }],
5831 })
5832 }
5833
5834 fn log_likelihood_only(&self, block_states: &[ParameterBlockState]) -> Result<f64, String> {
5835 let state = expect_single_block_state(block_states, "spatial adaptive exact family")?;
5836 let beta = &state.beta;
5837 let obs = evaluate_standard_familyobservations(
5838 self.family.clone(),
5839 self.latent_cloglog_state.as_ref(),
5840 self.mixture_link_state.as_ref(),
5841 self.sas_link_state.as_ref(),
5842 &self.y,
5843 &self.weights,
5844 &state.eta,
5845 )
5846 .map_err(|e| e.to_string())?;
5847 let adaptive_penalty = self.adaptive_penalty_value_only(beta)?;
5848 let (fixed_quadratic, _) = self.fixed_quadratic_terms(beta)?;
5849 Ok(obs.log_likelihood - adaptive_penalty - fixed_quadratic)
5850 }
5851
5852 fn exact_newton_outerobjective(&self) -> ExactNewtonOuterObjective {
5853 ExactNewtonOuterObjective::StrictPseudoLaplace
5854 }
5855
5856 fn exact_newton_joint_hessian(
5857 &self,
5858 block_states: &[ParameterBlockState],
5859 ) -> Result<Option<Array2<f64>>, String> {
5860 let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
5861 let eval = self.exact_evaluation(beta)?;
5862 Ok(Some(eval.totalobjectivehessian(&self.design)?))
5863 }
5864
5865 fn exact_newton_hessian_directional_derivative(
5866 &self,
5867 block_states: &[ParameterBlockState],
5868 block_idx: usize,
5869 d_beta: &Array1<f64>,
5870 ) -> Result<Option<Array2<f64>>, String> {
5871 expect_block_idx_zero(block_idx, "spatial adaptive exact family", "")?;
5872 self.exact_newton_joint_hessian_directional_derivative(block_states, d_beta)
5873 }
5874
5875 fn exact_newton_joint_hessian_directional_derivative(
5876 &self,
5877 block_states: &[ParameterBlockState],
5878 d_beta_flat: &Array1<f64>,
5879 ) -> Result<Option<Array2<f64>>, String> {
5880 let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
5881 if d_beta_flat.len() != beta.len() {
5882 return Err(SmoothError::dimension_mismatch(format!(
5883 "spatial adaptive exact family direction length mismatch: got {}, expected {}",
5884 d_beta_flat.len(),
5885 beta.len()
5886 ))
5887 .into());
5888 }
5889 let eval = self.exact_evaluation(beta)?;
5890 Ok(Some(
5891 self.exacthessian_directional_derivative_from_evaluation(beta, &eval, d_beta_flat)?,
5892 ))
5893 }
5894
5895 fn exact_newton_joint_hessiansecond_directional_derivative(
5896 &self,
5897 block_states: &[ParameterBlockState],
5898 d_beta_u_flat: &Array1<f64>,
5899 d_betav_flat: &Array1<f64>,
5900 ) -> Result<Option<Array2<f64>>, String> {
5901 let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
5902 if d_beta_u_flat.len() != beta.len() || d_betav_flat.len() != beta.len() {
5903 return Err(SmoothError::dimension_mismatch(format!(
5904 "spatial adaptive exact family second-direction length mismatch: got ({}, {}), expected {}",
5905 d_beta_u_flat.len(),
5906 d_betav_flat.len(),
5907 beta.len()
5908 ))
5909 .into());
5910 }
5911 let eval = self.exact_evaluation(beta)?;
5912 self.exacthessian_second_directional_derivative_from_evaluation(
5913 &eval,
5914 d_beta_u_flat,
5915 d_betav_flat,
5916 )
5917 }
5918
5919 fn block_linear_constraints(
5920 &self,
5921 block_states: &[ParameterBlockState],
5922 block_idx: usize,
5923 block_spec: &ParameterBlockSpec,
5924 ) -> Result<Option<ConstraintSet>, String> {
5925 assert!(!block_states.is_empty(), "block_states must be non-empty");
5926 assert!(
5927 !block_spec.name.is_empty(),
5928 "block spec name must be non-empty",
5929 );
5930 expect_block_idx_zero(block_idx, "spatial adaptive exact family", "")?;
5931 Ok(self.linear_constraints.clone().map(ConstraintSet::Dense))
5932 }
5933
5934 fn exact_newton_joint_psi_terms(
5935 &self,
5936 block_states: &[ParameterBlockState],
5937 specs: &[ParameterBlockSpec],
5938 hyper_layout: &crate::custom_family::CustomFamilyHyperLayout,
5939 psi_index: usize,
5940 ) -> Result<Option<ExactNewtonJointPsiTerms>, String> {
5941 if hyper_layout.family_axis_count() != 0 {
5942 return Err(
5943 "spatial adaptive exact family does not declare family-owned hyper axes"
5944 .to_string(),
5945 );
5946 }
5947 let derivative_blocks = hyper_layout.design_derivative_blocks();
5948 if block_states.len() != 1 || specs.len() != 1 || derivative_blocks.len() != 1 {
5949 return Err(SmoothError::dimension_mismatch(format!(
5950 "spatial adaptive exact family expects one block/state/spec/psi payload, got states={} specs={} deriv_blocks={}",
5951 block_states.len(),
5952 specs.len(),
5953 derivative_blocks.len()
5954 ))
5955 .into());
5956 }
5957 derivative_blocks[0]
5958 .get(psi_index)
5959 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
5960 let hyper = self
5961 .hyperspecs
5962 .get(psi_index)
5963 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
5964 let beta = &block_states[0].beta;
5965 let eval = self.exact_evaluation(beta)?;
5966 let (direct, beta_mixed, betahessian_explicit) =
5967 self.adaptive_hyper_parts(&eval, *hyper)?;
5968
5969 Ok(Some(ExactNewtonJointPsiTerms {
5990 objective_psi: direct,
5991 score_psi: beta_mixed,
5992 hessian_psi: betahessian_explicit,
5993 hessian_psi_operator: None,
5994 }))
5995 }
5996
5997 fn exact_newton_joint_psisecond_order_terms(
5998 &self,
5999 block_states: &[ParameterBlockState],
6000 specs: &[ParameterBlockSpec],
6001 hyper_layout: &crate::custom_family::CustomFamilyHyperLayout,
6002 psi_i: usize,
6003 psi_j: usize,
6004 ) -> Result<Option<gam_problem::ExactNewtonJointPsiSecondOrderTerms>, String> {
6005 if hyper_layout.family_axis_count() != 0 {
6006 return Err(
6007 "spatial adaptive exact family does not declare family-owned hyper axes"
6008 .to_string(),
6009 );
6010 }
6011 let derivative_blocks = hyper_layout.design_derivative_blocks();
6012 if block_states.len() != 1 || specs.len() != 1 || derivative_blocks.len() != 1 {
6013 return Err(SmoothError::dimension_mismatch(format!(
6014 "spatial adaptive exact family expects one block/state/spec/psi payload, got states={} specs={} deriv_blocks={}",
6015 block_states.len(),
6016 specs.len(),
6017 derivative_blocks.len()
6018 ))
6019 .into());
6020 }
6021 derivative_blocks[0]
6022 .get(psi_i)
6023 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_i))?;
6024 derivative_blocks[0]
6025 .get(psi_j)
6026 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_j))?;
6027 let hyper_i = self
6028 .hyperspecs
6029 .get(psi_i)
6030 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_i))?;
6031 let hyper_j = self
6032 .hyperspecs
6033 .get(psi_j)
6034 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_j))?;
6035 let beta = &block_states[0].beta;
6036 let eval = self.exact_evaluation(beta)?;
6037 let (objective_psi_psi, score_psi_psi, hessian_psi_psi) =
6038 self.adaptive_explicit_second_order_parts(&eval, *hyper_i, *hyper_j)?;
6039
6040 Ok(Some(gam_problem::ExactNewtonJointPsiSecondOrderTerms {
6041 objective_psi_psi,
6042 score_psi_psi,
6043 hessian_psi_psi,
6044 hessian_psi_psi_operator: None,
6045 }))
6046 }
6047
6048 fn exact_newton_joint_psihessian_directional_derivative(
6049 &self,
6050 block_states: &[ParameterBlockState],
6051 specs: &[ParameterBlockSpec],
6052 hyper_layout: &crate::custom_family::CustomFamilyHyperLayout,
6053 psi_index: usize,
6054 direction: &Array1<f64>,
6055 ) -> Result<Option<Array2<f64>>, String> {
6056 if hyper_layout.family_axis_count() != 0 {
6057 return Err(
6058 "spatial adaptive exact family does not declare family-owned hyper axes"
6059 .to_string(),
6060 );
6061 }
6062 let derivative_blocks = hyper_layout.design_derivative_blocks();
6063 if block_states.len() != 1 || specs.len() != 1 || derivative_blocks.len() != 1 {
6064 return Err(SmoothError::dimension_mismatch(format!(
6065 "spatial adaptive exact family expects one block/state/spec/psi payload, got states={} specs={} deriv_blocks={}",
6066 block_states.len(),
6067 specs.len(),
6068 derivative_blocks.len()
6069 ))
6070 .into());
6071 }
6072 let beta = &block_states[0].beta;
6073 if direction.len() != beta.len() {
6074 return Err(SmoothError::dimension_mismatch(format!(
6075 "spatial adaptive exact family direction length mismatch: got {}, expected {}",
6076 direction.len(),
6077 beta.len()
6078 ))
6079 .into());
6080 }
6081 derivative_blocks[0]
6082 .get(psi_index)
6083 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
6084 let hyper = self
6085 .hyperspecs
6086 .get(psi_index)
6087 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
6088 let eval = self.exact_evaluation(beta)?;
6089 let drift = match hyper.kind {
6090 SpatialAdaptiveHyperKind::LogLambdaMagnitude
6091 | SpatialAdaptiveHyperKind::LogLambdaGradient
6092 | SpatialAdaptiveHyperKind::LogLambdaCurvature => self.adaptive_block_drift_eval(
6093 &eval,
6094 hyper.cache_index,
6095 AdaptiveComponent::from_index(hyper.kind.component_index())?,
6096 HyperDriftKind::Rho,
6097 direction,
6098 )?,
6099 SpatialAdaptiveHyperKind::LogEpsilonMagnitude
6100 | SpatialAdaptiveHyperKind::LogEpsilonGradient
6101 | SpatialAdaptiveHyperKind::LogEpsilonCurvature => self
6102 .adaptive_shared_log_epsilon_drift(
6103 &eval,
6104 hyper.kind.component_index(),
6105 direction,
6106 )?,
6107 };
6108 Ok(Some(drift))
6109 }
6110}
6111
6112fn expect_single_block_state<'a>(
6113 block_states: &'a [ParameterBlockState],
6114 family_name: &str,
6115) -> Result<&'a ParameterBlockState, String> {
6116 crate::block_layout::block_count::validate_block_count::<SmoothError>(
6117 family_name,
6118 1,
6119 block_states.len(),
6120 )?;
6121 Ok(&block_states[0])
6122}
6123
6124fn expect_single_blockspec<'a>(
6125 specs: &'a [ParameterBlockSpec],
6126 family_name: &str,
6127) -> Result<&'a ParameterBlockSpec, String> {
6128 crate::block_layout::block_count::validate_block_count::<SmoothError>(
6129 family_name,
6130 1,
6131 specs.len(),
6132 )?;
6133 Ok(&specs[0])
6134}
6135
6136fn expect_block_idx_zero(block_idx: usize, family_name: &str, context: &str) -> Result<(), String> {
6137 if block_idx != 0 {
6138 return Err(SmoothError::invalid_index(format!(
6139 "{family_name} expects block_idx 0{context}, got {block_idx}"
6140 ))
6141 .into());
6142 }
6143 Ok::<(), _>(())
6144}
6145
6146impl BoundedLinearFamily {
6147 fn bounded_term_derivative_data(
6148 &self,
6149 latent_beta: &Array1<f64>,
6150 ) -> Result<
6151 (
6152 Array1<f64>,
6153 Array1<f64>,
6154 Array1<f64>,
6155 Array1<f64>,
6156 Array1<f64>,
6157 ),
6158 String,
6159 > {
6160 let p = latent_beta.len();
6161 if p != self.design.ncols() || latent_beta.iter().any(|value| !value.is_finite()) {
6162 return Err(format!(
6163 "bounded coefficient geometry requires {} finite latent coefficients, got {}",
6164 self.design.ncols(),
6165 p
6166 ));
6167 }
6168 let mut beta_user = latent_beta.clone();
6169 let mut jac_diag = Array1::<f64>::ones(p);
6170 let mut second_diag = Array1::<f64>::zeros(p);
6171 let mut third_diag = Array1::<f64>::zeros(p);
6172 let mut priorthird = Array1::<f64>::zeros(p);
6173 for term in &self.bounded_terms {
6174 let width = term.max - term.min;
6175 if term.col_idx >= p
6176 || !term.min.is_finite()
6177 || !term.max.is_finite()
6178 || !(width.is_finite() && width > 0.0)
6179 {
6180 return Err(format!(
6181 "bounded coefficient geometry has invalid column/bounds: col={}, p={p}, bounds=({}, {})",
6182 term.col_idx, term.min, term.max
6183 ));
6184 }
6185 let (beta, _, db_dtheta, d2b_dtheta2, d3b_dtheta3) =
6186 bounded_latent_derivatives(latent_beta[term.col_idx], term.min, term.max);
6187 if [beta, db_dtheta, d2b_dtheta2, d3b_dtheta3]
6188 .iter()
6189 .any(|value| !value.is_finite())
6190 {
6191 return Err(format!(
6192 "bounded coefficient transform is not representable at column {} and theta={}",
6193 term.col_idx, latent_beta[term.col_idx]
6194 ));
6195 }
6196 beta_user[term.col_idx] = beta;
6197 jac_diag[term.col_idx] = db_dtheta;
6198 second_diag[term.col_idx] = d2b_dtheta2;
6199 third_diag[term.col_idx] = d3b_dtheta3;
6200 let (_, _, _, prior_neghess_derivative) =
6201 bounded_prior_terms(latent_beta[term.col_idx], &term.prior)?;
6202 priorthird[term.col_idx] = prior_neghess_derivative;
6203 }
6204 Ok((beta_user, jac_diag, second_diag, third_diag, priorthird))
6205 }
6206
6207 fn user_beta_and_jacobian(
6208 &self,
6209 latent_beta: &Array1<f64>,
6210 ) -> Result<(Array1<f64>, Array1<f64>), String> {
6211 let (beta_user, jac_diag, _, _, _) = self.bounded_term_derivative_data(latent_beta)?;
6212 Ok((beta_user, jac_diag))
6213 }
6214
6215 fn nonlinear_offset_from_latent(
6216 &self,
6217 latent_beta: &Array1<f64>,
6218 ) -> Result<Array1<f64>, String> {
6219 self.bounded_term_derivative_data(latent_beta)?;
6220 let mut offset = self.offset.clone();
6221 for term in &self.bounded_terms {
6222 let (beta, _, _) =
6223 bounded_latent_to_user(latent_beta[term.col_idx], term.min, term.max);
6224 offset.scaled_add(beta, &self.design.column(term.col_idx));
6225 }
6226 if offset.iter().any(|value| !value.is_finite()) {
6227 return Err("bounded nonlinear offset is not representable".to_string());
6228 }
6229 Ok(offset)
6230 }
6231
6232 fn effective_design_for_latent(&self, jac_diag: &Array1<f64>) -> Array2<f64> {
6233 let mut x_eff = self.design.clone();
6234 for term in &self.bounded_terms {
6235 x_eff
6236 .column_mut(term.col_idx)
6237 .mapv_inplace(|v| v * jac_diag[term.col_idx]);
6238 }
6239 x_eff
6240 }
6241
6242 fn exacthessian_andgradient(
6243 &self,
6244 latent_beta: &Array1<f64>,
6245 ) -> Result<
6246 (
6247 StandardFamilyObservationState,
6248 Array2<f64>,
6249 Array1<f64>,
6250 f64,
6251 Array1<f64>,
6252 Array1<f64>,
6253 Array1<f64>,
6254 ),
6255 String,
6256 > {
6257 let (_, jac_diag, second_diag, third_diag, priorthird) =
6258 self.bounded_term_derivative_data(latent_beta)?;
6259 let x_eff = self.effective_design_for_latent(&jac_diag);
6260 let eta =
6261 self.designzeroed.dot(latent_beta) + self.nonlinear_offset_from_latent(latent_beta)?;
6262 let obs = evaluate_resolved_standard_family_observations(
6263 &self.likelihood,
6264 self.latent_cloglog_state.as_ref(),
6265 self.mixture_link_state.as_ref(),
6266 self.sas_link_state.as_ref(),
6267 &self.y,
6268 &self.weights,
6269 &eta,
6270 )
6271 .map_err(|e| e.to_string())?;
6272
6273 let mut priorgrad = Array1::<f64>::zeros(latent_beta.len());
6274 let mut prior_neghess = Array2::<f64>::zeros((latent_beta.len(), latent_beta.len()));
6275 let mut prior_loglik = 0.0;
6276 for term in &self.bounded_terms {
6277 let (logp, grad, neghess, _) =
6278 bounded_prior_terms(latent_beta[term.col_idx], &term.prior)?;
6279 prior_loglik += logp;
6280 priorgrad[term.col_idx] += grad;
6281 prior_neghess[[term.col_idx, term.col_idx]] += neghess;
6282 }
6283
6284 let mut hessian = xt_diag_x_dense(x_eff.view(), obs.neghessian_eta.view())?;
6285 let mut gradient = fast_atv(&x_eff, &obs.score);
6286 for term in &self.bounded_terms {
6287 let score_beta = self.design.column(term.col_idx).dot(&obs.score);
6288 hessian[[term.col_idx, term.col_idx]] -= score_beta * second_diag[term.col_idx];
6289 }
6290 hessian += &prior_neghess;
6291 gradient += &priorgrad;
6292
6293 Ok((
6294 obs,
6295 hessian,
6296 gradient,
6297 prior_loglik,
6298 second_diag,
6299 third_diag,
6300 priorthird,
6301 ))
6302 }
6303
6304 fn evaluation_from_latent(
6305 &self,
6306 latent_beta: &Array1<f64>,
6307 ) -> Result<
6308 (
6309 StandardFamilyObservationState,
6310 Array2<f64>,
6311 Array1<f64>,
6312 f64,
6313 ),
6314 String,
6315 > {
6316 let (obs, hessian, gradient, prior_loglik, _, _, _) =
6317 self.exacthessian_andgradient(latent_beta)?;
6318 Ok((obs, hessian, gradient, prior_loglik))
6319 }
6320}
6321
6322impl CustomFamily for BoundedLinearFamily {
6323 fn joint_jeffreys_term_required(&self) -> bool {
6327 true
6328 }
6329
6330 fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
6331 let latent_beta = &expect_single_block_state(block_states, "bounded linear family")?.beta;
6332 let (obs, hessian, gradient, prior_loglik) = self.evaluation_from_latent(latent_beta)?;
6333 Ok(FamilyEvaluation {
6334 log_likelihood: obs.log_likelihood + prior_loglik,
6335 blockworking_sets: vec![BlockWorkingSet::ExactNewton {
6336 gradient,
6337 hessian: SymmetricMatrix::Dense(hessian),
6338 }],
6339 })
6340 }
6341
6342 fn exact_newton_joint_hessian(
6343 &self,
6344 block_states: &[ParameterBlockState],
6345 ) -> Result<Option<Array2<f64>>, String> {
6346 let latent_beta = &expect_single_block_state(block_states, "bounded linear family")?.beta;
6347 let (_, hessian, _, _) = self.evaluation_from_latent(latent_beta)?;
6348 Ok(Some(hessian))
6349 }
6350
6351 fn exact_newton_hessian_directional_derivative(
6352 &self,
6353 block_states: &[ParameterBlockState],
6354 block_idx: usize,
6355 d_beta: &Array1<f64>,
6356 ) -> Result<Option<Array2<f64>>, String> {
6357 expect_block_idx_zero(block_idx, "bounded linear family", "")?;
6358 self.exact_newton_joint_hessian_directional_derivative(block_states, d_beta)
6359 }
6360
6361 fn exact_newton_joint_hessian_directional_derivative(
6362 &self,
6363 block_states: &[ParameterBlockState],
6364 d_beta_flat: &Array1<f64>,
6365 ) -> Result<Option<Array2<f64>>, String> {
6366 let latent_beta = &expect_single_block_state(block_states, "bounded linear family")?.beta;
6367 if d_beta_flat.len() != latent_beta.len() {
6368 return Err(SmoothError::dimension_mismatch(format!(
6369 "bounded linear family directional derivative length mismatch: got {}, expected {}",
6370 d_beta_flat.len(),
6371 latent_beta.len()
6372 ))
6373 .into());
6374 }
6375
6376 let (obs, _, _, _, second_diag, third_diag, priorthird) =
6377 self.exacthessian_andgradient(latent_beta)?;
6378
6379 let (_, jac_diag, _, _, _) = self.bounded_term_derivative_data(latent_beta)?;
6380 let x_eff = self.effective_design_for_latent(&jac_diag);
6381 let deta = x_eff.dot(d_beta_flat);
6382 let d_neghess_eta = &obs.neghessian_eta_derivative * &deta;
6383
6384 let mut dx_eff = Array2::<f64>::zeros(x_eff.raw_dim());
6385 for term in &self.bounded_terms {
6386 let scale = second_diag[term.col_idx] * d_beta_flat[term.col_idx];
6387 if scale != 0.0 {
6388 let mut col = dx_eff.column_mut(term.col_idx);
6389 col.assign(&self.design.column(term.col_idx));
6390 col.mapv_inplace(|v| v * scale);
6391 }
6392 }
6393
6394 let mut dhessian = xt_diag_x_dense(x_eff.view(), d_neghess_eta.view())?;
6395 let mut wxdx = Array2::<f64>::zeros((x_eff.ncols(), x_eff.ncols()));
6396 for i in 0..x_eff.nrows() {
6397 let wi = obs.neghessian_eta[i];
6398 if wi == 0.0 {
6399 continue;
6400 }
6401 for a in 0..x_eff.ncols() {
6402 let xa = x_eff[[i, a]];
6403 for b in 0..x_eff.ncols() {
6404 wxdx[[a, b]] += wi * (dx_eff[[i, a]] * x_eff[[i, b]] + xa * dx_eff[[i, b]]);
6405 }
6406 }
6407 }
6408 dhessian += &wxdx;
6409
6410 let d_score = -&obs.neghessian_eta * &deta;
6411 for term in &self.bounded_terms {
6412 let score_beta = self.design.column(term.col_idx).dot(&obs.score);
6413 let d_score_beta = self.design.column(term.col_idx).dot(&d_score);
6414 dhessian[[term.col_idx, term.col_idx]] -= d_score_beta * second_diag[term.col_idx]
6415 + score_beta * third_diag[term.col_idx] * d_beta_flat[term.col_idx];
6416 dhessian[[term.col_idx, term.col_idx]] +=
6417 priorthird[term.col_idx] * d_beta_flat[term.col_idx];
6418 }
6419
6420 Ok(Some(dhessian))
6421 }
6422
6423 fn block_geometry(
6424 &self,
6425 block_states: &[ParameterBlockState],
6426 spec: &ParameterBlockSpec,
6427 ) -> Result<(DesignMatrix, Array1<f64>), String> {
6428 if block_states.is_empty() {
6429 return Ok((
6430 DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
6431 self.designzeroed.clone(),
6432 )),
6433 self.offset.clone(),
6434 ));
6435 }
6436 let offset = self.nonlinear_offset_from_latent(
6437 &expect_single_block_state(block_states, "bounded linear family")?.beta,
6438 )?;
6439 let x = if spec.design.ncols() == self.designzeroed.ncols() {
6440 self.designzeroed.clone()
6441 } else {
6442 return Err(SmoothError::dimension_mismatch(
6443 "bounded linear family design column mismatch",
6444 )
6445 .into());
6446 };
6447 Ok((
6448 DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(x)),
6449 offset,
6450 ))
6451 }
6452
6453 fn block_geometry_is_dynamic(&self) -> bool {
6454 true
6455 }
6456
6457 fn block_geometry_directional_derivative(
6458 &self,
6459 block_states: &[ParameterBlockState],
6460 block_idx: usize,
6461 spec: &ParameterBlockSpec,
6462 d_beta: &Array1<f64>,
6463 ) -> Result<Option<BlockGeometryDirectionalDerivative>, String> {
6464 expect_block_idx_zero(
6465 block_idx,
6466 "bounded linear family",
6467 " for geometry derivative",
6468 )?;
6469 expect_single_block_state(block_states, "bounded linear family")?;
6470 if d_beta.len() != spec.design.ncols() {
6471 return Err(SmoothError::dimension_mismatch(format!(
6472 "bounded linear family geometry derivative direction mismatch: got {}, expected {}",
6473 d_beta.len(),
6474 spec.design.ncols()
6475 ))
6476 .into());
6477 }
6478 let (_, jac_diag, _, _, _) = self.bounded_term_derivative_data(&block_states[0].beta)?;
6479 let mut d_offset = Array1::<f64>::zeros(self.offset.len());
6480 let has_drift = self
6481 .bounded_terms
6482 .iter()
6483 .any(|term| jac_diag[term.col_idx] != 0.0 && d_beta[term.col_idx] != 0.0);
6484 if !has_drift {
6485 return Ok(Some(BlockGeometryDirectionalDerivative {
6486 d_design: None,
6487 d_offset,
6488 }));
6489 }
6490 for term in &self.bounded_terms {
6491 let col = term.col_idx;
6492 let drift = jac_diag[col] * d_beta[col];
6493 if drift != 0.0 {
6494 d_offset.scaled_add(drift, &self.design.column(col));
6495 }
6496 }
6497 Ok(Some(BlockGeometryDirectionalDerivative {
6498 d_design: None,
6499 d_offset,
6500 }))
6501 }
6502}
6503
6504#[inline]
6505fn dense_diag_gram_chunkrows(p: usize) -> usize {
6506 const MIN_ROWS: usize = 512;
6507 const MAX_ROWS: usize = 2048;
6508 const TARGET_BYTES: usize = 2 * 1024 * 1024;
6509 let bytes_per_row = p.max(1) * std::mem::size_of::<f64>();
6510 (TARGET_BYTES / bytes_per_row).clamp(MIN_ROWS, MAX_ROWS)
6511}
6512
6513fn xt_diag_x_dense(x: ArrayView2<'_, f64>, w: ArrayView1<'_, f64>) -> Result<Array2<f64>, String> {
6514 if x.nrows() != w.len() {
6515 return Err(SmoothError::dimension_mismatch("xt_diag_x_dense row mismatch").into());
6516 }
6517 let (n, p) = x.dim();
6518 if n == 0 || p == 0 {
6519 return Ok(Array2::<f64>::zeros((p, p)));
6520 }
6521
6522 const STREAMING_BYTES_THRESHOLD: usize = 8 * 1024 * 1024;
6523 let dense_work_bytes = n
6524 .checked_mul(p)
6525 .and_then(|cells| cells.checked_mul(std::mem::size_of::<f64>()))
6526 .unwrap_or(usize::MAX);
6527 if dense_work_bytes <= STREAMING_BYTES_THRESHOLD {
6528 let mut weighted = x.to_owned();
6529 ndarray::Zip::from(weighted.rows_mut())
6530 .and(w)
6531 .par_for_each(|mut row, wi| row *= *wi);
6532 return Ok(fast_atb(&x, &weighted));
6533 }
6534
6535 let chunkrows = dense_diag_gram_chunkrows(p).min(n);
6536 let mut weighted_chunk = Array2::<f64>::zeros((chunkrows, p));
6537 let mut out = Array2::<f64>::zeros((p, p));
6538 for row_start in (0..n).step_by(chunkrows) {
6539 let rows = (n - row_start).min(chunkrows);
6540 let x_chunk = x.slice(s![row_start..row_start + rows, ..]);
6541 {
6542 let mut chunk = weighted_chunk.slice_mut(s![0..rows, ..]);
6543 for local_row in 0..rows {
6544 let scale = w[row_start + local_row];
6545 if scale == 0.0 {
6546 chunk.row_mut(local_row).fill(0.0);
6547 continue;
6548 }
6549 for col in 0..p {
6550 chunk[[local_row, col]] = x_chunk[[local_row, col]] * scale;
6551 }
6552 }
6553 }
6554 out += &fast_atb(&x_chunk, &weighted_chunk.slice(s![0..rows, ..]));
6555 }
6556 Ok(out)
6557}
6558
6559fn trace_of_dense_product(a: &Array2<f64>, b: &Array2<f64>) -> Result<f64, String> {
6560 if a.nrows() != a.ncols() || b.nrows() != b.ncols() || a.nrows() != b.nrows() {
6561 return Err(
6562 SmoothError::dimension_mismatch("trace_of_dense_product dimension mismatch").into(),
6563 );
6564 }
6565 if a.iter().chain(b.iter()).any(|value| !value.is_finite()) {
6566 return Err("trace_of_dense_product requires finite matrices".to_string());
6567 }
6568 let mut trace = gam_linalg::utils::KahanSum::default();
6569 for i in 0..a.nrows() {
6570 for j in 0..a.ncols() {
6571 let term = a[[i, j]] * b[[j, i]];
6572 if !term.is_finite() {
6573 return Err(format!(
6574 "trace_of_dense_product term ({i}, {j}) is not representable"
6575 ));
6576 }
6577 trace.add(term);
6578 }
6579 }
6580 let trace = trace.sum();
6581 if !trace.is_finite() {
6582 return Err("trace_of_dense_product sum is not representable".to_string());
6583 }
6584 Ok(trace)
6585}
6586
6587fn certify_bounded_edf_interval(
6588 value: f64,
6589 lower: f64,
6590 upper: f64,
6591 dimension: usize,
6592 label: &str,
6593) -> Result<f64, EstimationError> {
6594 if !(value.is_finite() && lower.is_finite() && upper.is_finite() && lower <= upper) {
6595 crate::bail_invalid_estim!(
6596 "{label} has invalid EDF interval/value: value={value}, interval=[{lower}, {upper}]"
6597 );
6598 }
6599 let scale = 1.0_f64.max(value.abs()).max(lower.abs()).max(upper.abs());
6600 let allowed = 256.0 * f64::EPSILON * (dimension.max(1) as f64).powi(2) * scale;
6604 if value < lower {
6605 if lower - value <= allowed {
6606 return Ok(lower);
6607 }
6608 } else if value > upper {
6609 if value - upper <= allowed {
6610 return Ok(upper);
6611 }
6612 } else {
6613 return Ok(value);
6614 }
6615 crate::bail_invalid_estim!(
6616 "{label}={value} lies outside [{lower}, {upper}] by more than the dense-trace backward-error allowance {allowed}"
6617 )
6618}
6619
6620fn exact_bounded_edf(
6621 penalties: &[PenaltySpec],
6622 lambdas: &Array1<f64>,
6623 latent_cov: &Array2<f64>,
6624) -> Result<(Vec<f64>, Vec<f64>, f64), EstimationError> {
6625 if penalties.len() != lambdas.len() {
6626 crate::bail_invalid_estim!(
6627 "bounded EDF penalty/lambda mismatch: {} penalties vs {} lambdas",
6628 penalties.len(),
6629 lambdas.len()
6630 );
6631 }
6632 if latent_cov.nrows() != latent_cov.ncols() {
6633 crate::bail_invalid_estim!("bounded EDF covariance must be square");
6634 }
6635
6636 let p = latent_cov.nrows();
6637 let mut s_lambda = Array2::<f64>::zeros((p, p));
6638 let mut edf_by_block = Vec::with_capacity(penalties.len());
6639 let mut penalty_block_trace = Vec::with_capacity(penalties.len());
6641 let mut trace_sum = gam_linalg::utils::KahanSum::default();
6642
6643 for (k, ps) in penalties.iter().enumerate() {
6644 let lambda_k = lambdas[k];
6645 if !(lambda_k.is_finite() && lambda_k >= 0.0) {
6646 crate::bail_invalid_estim!(
6647 "bounded EDF smoothing strength at block {k} must be finite and non-negative, got {lambda_k}"
6648 );
6649 }
6650 match ps {
6651 PenaltySpec::Block {
6652 local, col_range, ..
6653 } => {
6654 s_lambda
6655 .slice_mut(ndarray::s![col_range.clone(), col_range.clone()])
6656 .scaled_add(lambda_k, local);
6657 let penalty_rank =
6659 local
6660 .nrows()
6661 .saturating_sub(estimate_penalty_nullity(local).map_err(|e| {
6662 EstimationError::InvalidInput(format!("bounded EDF rank failed: {e}"))
6663 })?);
6664 let cov_block = latent_cov.slice(ndarray::s![col_range.clone(), col_range.clone()]);
6666 let trace_k = lambda_k
6667 * trace_of_dense_product(&cov_block.to_owned(), local)
6668 .map_err(EstimationError::InvalidInput)?;
6669 trace_sum.add(trace_k);
6670 penalty_block_trace.push(trace_k);
6671 let p_k = penalty_rank as f64;
6672 edf_by_block.push(certify_bounded_edf_interval(
6673 p_k - trace_k,
6674 0.0,
6675 p_k,
6676 p,
6677 &format!("bounded EDF block {k}"),
6678 )?);
6679 }
6680 PenaltySpec::Dense(m) | PenaltySpec::DenseWithMean { matrix: m, .. } => {
6681 s_lambda.scaled_add(lambda_k, m);
6682 let penalty_rank = p.saturating_sub(estimate_penalty_nullity(m).map_err(|e| {
6683 EstimationError::InvalidInput(format!("bounded EDF rank failed: {e}"))
6684 })?);
6685 let trace_k = lambda_k
6686 * trace_of_dense_product(latent_cov, m)
6687 .map_err(EstimationError::InvalidInput)?;
6688 trace_sum.add(trace_k);
6689 penalty_block_trace.push(trace_k);
6690 let p_k = penalty_rank as f64;
6691 edf_by_block.push(certify_bounded_edf_interval(
6692 p_k - trace_k,
6693 0.0,
6694 p_k,
6695 p,
6696 &format!("bounded EDF block {k}"),
6697 )?);
6698 }
6699 }
6700 }
6701
6702 let nullity_total = estimate_penalty_nullity(&s_lambda)
6703 .map_err(|e| EstimationError::InvalidInput(format!("bounded EDF nullity failed: {e}")))?
6704 as f64;
6705 let trace_sum = trace_sum.sum();
6706 let edf_total = certify_bounded_edf_interval(
6707 p as f64 - trace_sum,
6708 nullity_total,
6709 p as f64,
6710 p,
6711 "bounded total EDF",
6712 )?;
6713 Ok((edf_by_block, penalty_block_trace, edf_total))
6714}
6715
6716fn certified_bounded_posterior_covariance(
6722 precision: &Array2<f64>,
6723 label: &'static str,
6724) -> Result<Array2<f64>, EstimationError> {
6725 gam_linalg::utils::certified_spd_inverse(precision, label)
6726 .map(gam_linalg::utils::CertifiedSpdInverse::into_inverse)
6727 .map_err(|error| {
6728 EstimationError::InvalidInput(format!(
6729 "bounded posterior covariance requires an exact SPD precision: {error}"
6730 ))
6731 })
6732}
6733
6734fn transform_bounded_latent_precision_to_user_internal(
6735 latent_precision: &Array2<f64>,
6736 jac_diag: &Array1<f64>,
6737) -> Result<Array2<f64>, EstimationError> {
6738 let p = latent_precision.nrows();
6739 if latent_precision.ncols() != p || jac_diag.len() != p {
6740 crate::bail_invalid_estim!(
6741 "bounded precision transform dimension mismatch: precision is {}x{}, jacobian has {} entries",
6742 latent_precision.nrows(),
6743 latent_precision.ncols(),
6744 jac_diag.len()
6745 );
6746 }
6747 let mut out = latent_precision.clone();
6748 for i in 0..p {
6749 let scale = jac_diag[i];
6750 if !scale.is_finite() || scale <= 0.0 {
6751 crate::bail_invalid_estim!(
6752 "bounded precision transform requires a positive finite coefficient jacobian; column {i} has {scale}"
6753 );
6754 }
6755 if scale != 1.0 {
6756 out.row_mut(i).mapv_inplace(|v| v / scale);
6757 out.column_mut(i).mapv_inplace(|v| v / scale);
6758 }
6759 }
6760 Ok(out)
6761}
6762
6763fn fit_bounded_term_collection_with_design(
6764 y: ArrayView1<'_, f64>,
6765 weights: ArrayView1<'_, f64>,
6766 offset: ArrayView1<'_, f64>,
6767 spec: &TermCollectionSpec,
6768 design: &TermCollectionDesign,
6769 heuristic_lambdas: Option<&[f64]>,
6770 family: LikelihoodSpec,
6771 options: &FitOptions,
6772) -> Result<FittedTermCollection, EstimationError> {
6773 let conditioning_cols: Vec<usize> = spec
6774 .linear_terms
6775 .iter()
6776 .enumerate()
6777 .filter_map(|(j, linear)| {
6778 (!linear.double_penalty).then_some(design.intercept_range.end + j)
6779 })
6780 .collect();
6781 let conditioning = LinearFitConditioning::from_columns(design, &conditioning_cols);
6782 let dense_design = design.design.to_dense_cow();
6783 let fit_design = conditioning.apply_to_design(&dense_design);
6784 let fit_penalties = conditioning
6785 .transform_blockwise_penalties_to_internal(&design.penalties, design.design.ncols());
6786 if design.linear_constraints.is_some() {
6787 crate::bail_invalid_estim!(
6788 "bounded() terms are not yet compatible with explicit linear constraints"
6789 );
6790 }
6791 let mut bounded_terms = Vec::<BoundedLinearTermMeta>::new();
6792 for (j, term) in spec.linear_terms.iter().enumerate() {
6793 if term.double_penalty
6794 && matches!(
6795 term.coefficient_geometry,
6796 LinearCoefficientGeometry::Bounded { .. }
6797 )
6798 {
6799 crate::bail_invalid_estim!(
6800 "bounded linear term '{}' cannot also use double_penalty",
6801 term.name
6802 );
6803 }
6804 if let LinearCoefficientGeometry::Bounded { min, max, prior } =
6805 term.coefficient_geometry.clone()
6806 {
6807 let col_idx = design.intercept_range.end + j;
6808 let (min_internal, max_internal) = conditioning.internal_bounds_for(col_idx, min, max);
6809 bounded_terms.push(BoundedLinearTermMeta {
6810 col_idx,
6811 min: min_internal,
6812 max: max_internal,
6813 prior,
6814 });
6815 }
6816 }
6817 if bounded_terms.is_empty() {
6818 crate::bail_invalid_estim!("internal bounded fit path called with no bounded terms");
6819 }
6820
6821 let mut designzeroed = fit_design.clone();
6822 let mut initial_beta = Array1::<f64>::zeros(fit_design.ncols());
6823 for term in &bounded_terms {
6824 designzeroed.column_mut(term.col_idx).fill(0.0);
6825 initial_beta[term.col_idx] = 0.0;
6826 }
6827
6828 let initial_log_lambdas = heuristic_lambdas
6829 .map(|vals| Array1::from_vec(vals.to_vec()))
6830 .unwrap_or_else(|| Array1::zeros(fit_penalties.len()));
6831 if initial_log_lambdas.len() != fit_penalties.len() {
6832 crate::bail_invalid_estim!(
6833 "heuristic lambda length mismatch for bounded model: got {}, expected {}",
6834 initial_log_lambdas.len(),
6835 fit_penalties.len()
6836 );
6837 }
6838
6839 let glm_likelihood = gam_spec::GlmLikelihoodSpec::canonical(family);
6840 let resolved_likelihood_scale = glm_likelihood
6841 .resolved_scale()
6842 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
6843 let is_beta_logistic = glm_likelihood.spec.is_binomial_beta_logistic();
6844 let family_adapter = BoundedLinearFamily {
6845 likelihood: glm_likelihood.clone(),
6846 latent_cloglog_state: options.latent_cloglog,
6847 mixture_link_state: options
6848 .mixture_link
6849 .clone()
6850 .as_ref()
6851 .map(state_fromspec)
6852 .transpose()
6853 .map_err(EstimationError::InvalidInput)?,
6854 sas_link_state: options
6855 .sas_link
6856 .map(|spec| {
6857 if is_beta_logistic {
6858 state_from_beta_logisticspec(spec)
6859 } else {
6860 state_from_sasspec(spec)
6861 }
6862 })
6863 .transpose()
6864 .map_err(EstimationError::InvalidInput)?,
6865 y: y.to_owned(),
6866 weights: weights.to_owned(),
6867 design: fit_design.clone(),
6868 designzeroed: designzeroed.clone(),
6869 offset: offset.to_owned(),
6870 bounded_terms: bounded_terms.clone(),
6871 };
6872 let blockspec = ParameterBlockSpec {
6873 name: "eta".to_string(),
6874 design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(designzeroed)),
6875 offset: offset.to_owned(),
6876 penalties: fit_penalties
6877 .iter()
6878 .map(|ps| match ps {
6879 PenaltySpec::Block {
6880 local, col_range, ..
6881 } => PenaltyMatrix::Blockwise {
6882 local: local.clone(),
6883 col_range: col_range.clone(),
6884 total_dim: design.design.ncols(),
6885 },
6886 PenaltySpec::Dense(m) | PenaltySpec::DenseWithMean { matrix: m, .. } => {
6887 PenaltyMatrix::Dense(m.clone())
6888 }
6889 })
6890 .collect(),
6891 nullspace_dims: design.nullspace_dims.clone(),
6892 initial_log_lambdas,
6893 initial_beta: Some(initial_beta),
6894 gauge_priority: 100,
6895 jacobian_callback: Some(Arc::new(BoundedEffectiveJacobian {
6901 design: fit_design.clone(),
6902 bounded_terms: bounded_terms.clone(),
6903 })),
6904 stacked_design: None,
6905 stacked_offset: None,
6906 };
6907 let fit = fit_custom_family(
6908 &family_adapter,
6909 &[blockspec],
6910 &BlockwiseFitOptions {
6911 inner_max_cycles: options.max_iter,
6912 inner_tol: options.tol,
6913 outer_max_iter: options.max_iter,
6914 outer_tol: options.tol,
6915 compute_covariance: false,
6925 ..BlockwiseFitOptions::default()
6926 },
6927 )
6928 .map_err(EstimationError::CustomFamily)?;
6929
6930 let latent_beta = fit.block_states[0].beta.clone();
6931 let (beta_user_internal, jac_diag) = family_adapter
6932 .user_beta_and_jacobian(&latent_beta)
6933 .map_err(EstimationError::InvalidInput)?;
6934 let beta_user = conditioning.backtransform_beta(&beta_user_internal);
6935
6936 let (eta_state, h_data, _, _) = family_adapter
6937 .evaluation_from_latent(&latent_beta)
6938 .map_err(EstimationError::InvalidInput)?;
6939 let p_fit = fit_design.ncols();
6940 let mut s_lambda_internal = Array2::<f64>::zeros((p_fit, p_fit));
6941 for (k, penalty) in fit_penalties.iter().enumerate() {
6942 match penalty {
6943 PenaltySpec::Block {
6944 local, col_range, ..
6945 } => {
6946 s_lambda_internal
6947 .slice_mut(ndarray::s![col_range.clone(), col_range.clone()])
6948 .scaled_add(fit.lambdas[k], local);
6949 }
6950 PenaltySpec::Dense(m) | PenaltySpec::DenseWithMean { matrix: m, .. } => {
6951 s_lambda_internal.scaled_add(fit.lambdas[k], m);
6952 }
6953 }
6954 }
6955 let mut latent_precision = h_data.clone();
6956 latent_precision += &s_lambda_internal;
6957 let user_precision_internal =
6958 transform_bounded_latent_precision_to_user_internal(&latent_precision, &jac_diag)?;
6959 let penalized_hessian =
6960 conditioning.transform_penalized_hessian_to_original(&user_precision_internal);
6961
6962 let beta_covariance_unscaled = if options.compute_inference {
6989 Some(certified_bounded_posterior_covariance(
6990 &penalized_hessian,
6991 "bounded user-scale posterior precision",
6992 )?)
6993 } else {
6994 None
6995 };
6996 let latent_cov = if options.compute_inference {
7002 Some(certified_bounded_posterior_covariance(
7003 &latent_precision,
7004 "bounded latent posterior precision",
7005 )?)
7006 } else {
7007 None
7008 };
7009 let s_lambda_original = weighted_blockwise_penalty_sum(
7010 &design.penalties,
7011 fit.lambdas.as_slice().unwrap(),
7012 design.design.ncols(),
7013 );
7014 let penalty_term = beta_user.dot(&s_lambda_original.dot(&beta_user));
7015 let deviance = -2.0 * eta_state.log_likelihood;
7016 let (edf_by_block, penalty_block_trace, edf_total) = if let Some(cov) = latent_cov.as_ref() {
7017 exact_bounded_edf(&fit_penalties, &fit.lambdas, cov)?
7018 } else {
7019 (
7020 vec![0.0; fit_penalties.len()],
7021 vec![0.0; fit_penalties.len()],
7022 0.0,
7023 )
7024 };
7025
7026 let profiled_gaussian_standard_deviation = if matches!(
7038 resolved_likelihood_scale,
7039 gam_spec::ResolvedLikelihoodScale::ProfiledGaussian
7040 ) {
7041 let residual_dof = if options.compute_inference {
7042 y.len() as f64 - edf_total
7043 } else {
7044 y.len() as f64
7045 };
7046 if !(residual_dof.is_finite() && residual_dof > 0.0) {
7047 return Err(EstimationError::InvalidInput(format!(
7048 "bounded Gaussian residual degrees of freedom must be finite and positive, got n={} minus edf={edf_total} = {residual_dof}",
7049 y.len()
7050 )));
7051 }
7052 if !(deviance.is_finite() && deviance >= 0.0) {
7053 return Err(EstimationError::InvalidInput(format!(
7054 "bounded Gaussian deviance must be finite and non-negative, got {deviance}"
7055 )));
7056 }
7057 let variance = deviance / residual_dof;
7058 if !variance.is_finite() {
7059 return Err(EstimationError::InvalidInput(format!(
7060 "bounded Gaussian residual variance is not representable: {deviance}/{residual_dof}"
7061 )));
7062 }
7063 Some(variance.sqrt())
7064 } else {
7065 None
7066 };
7067 let dispersion = gam_solve::estimate::dispersion_from_likelihood(
7068 &glm_likelihood,
7069 profiled_gaussian_standard_deviation,
7070 )?;
7071 let standard_deviation = dispersion.phi().sqrt();
7072 let cov_scale = glm_likelihood
7073 .coefficient_covariance_scale(dispersion.phi())
7074 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
7075 let beta_covariance = beta_covariance_unscaled.map(|mut cov| {
7081 if cov_scale != 1.0 {
7082 cov.mapv_inplace(|v| v * cov_scale);
7083 }
7084 cov
7085 });
7086 if let Some(covariance) = beta_covariance.as_ref()
7087 && covariance.iter().any(|value| !value.is_finite())
7088 {
7089 return Err(EstimationError::InvalidInput(
7090 "bounded coefficient covariance scaling produced a non-finite value".to_string(),
7091 ));
7092 }
7093 let beta_standard_errors = beta_covariance
7094 .as_ref()
7095 .map(gam_problem::se_from_covariance)
7096 .transpose()
7097 .map_err(|err| {
7098 EstimationError::InvalidInput(format!(
7099 "bounded coefficient covariance cannot produce standard errors: {err}"
7100 ))
7101 })?;
7102 let working_response = exact_standard_working_response(&eta_state)?;
7103
7104 let geometry = Some(gam_solve::estimate::FitGeometry {
7105 coefficient_gauge: gam_problem::gauge::Gauge::identity(&[beta_user.len()]),
7106 penalized_hessian: penalized_hessian.clone().into(),
7107 working: Some(gam_solve::estimate::WorkingGeometry {
7108 weights: eta_state.fisherweight.clone(),
7109 response: working_response,
7110 }),
7111 });
7112 let max_abs_eta = eta_state
7113 .eta
7114 .iter()
7115 .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
7116 Ok(FittedTermCollection {
7117 fit: {
7118 let log_lambdas =
7119 checked_fit_log_lambdas(&fit.lambdas, "final fitted term collection")?;
7120 let inf = FitInference {
7121 edf_by_block,
7122 penalty_block_trace,
7123 edf_total,
7124 smoothing_correction: None,
7125 smoothing_correction_method: None,
7126 smoothing_correction_first_order: None,
7127 smoothing_correction_method_first_order: None,
7128 penalized_hessian: penalized_hessian.clone().into(),
7131 reparam_qs: None,
7132 dispersion,
7133 beta_covariance: beta_covariance
7134 .clone()
7135 .map(gam_problem::dispersion_cov::PhiScaledCovariance::from),
7136 beta_standard_errors,
7137 beta_covariance_corrected: None,
7138 beta_standard_errors_corrected: None,
7139 beta_covariance_frequentist: None,
7140 coefficient_influence: None,
7141 weighted_gram: None,
7142 bias_correction_beta: None,
7143 bias_correction_jacobian: None,
7144 };
7145 let covariance_conditional = beta_covariance;
7146 let pirls_status_val = gam_solve::pirls::PirlsStatus::Converged;
7149 UnifiedFitResult::try_from_parts(UnifiedFitResultParts {
7150 blocks: vec![gam_solve::estimate::FittedBlock {
7151 beta: beta_user.clone(),
7152 role: gam_problem::BlockRole::Mean,
7153 edf: edf_total,
7154 lambdas: fit.lambdas.clone(),
7155 }],
7156 log_lambdas,
7157 lambdas: fit.lambdas,
7158 likelihood_scale: glm_likelihood.scale,
7159 likelihood_family: Some(glm_likelihood.spec),
7160 log_likelihood_normalization: gam_spec::LogLikelihoodNormalization::UserProvided,
7161 log_likelihood: eta_state.log_likelihood,
7162 deviance,
7163 reml_score: fit.penalized_objective,
7164 stable_penalty_term: penalty_term,
7165 penalized_objective: fit.penalized_objective,
7166 used_device: false,
7167 outer_iterations: fit.outer_iterations,
7168 outer_converged: true,
7170 outer_gradient_norm: fit.outer_gradient_norm,
7171 standard_deviation,
7172 covariance_conditional,
7173 covariance_corrected: None,
7174 inference: Some(inf),
7175 fitted_link: gam_solve::estimate::FittedLinkState::Standard(None),
7176 geometry,
7177 block_states: Vec::new(),
7178 pirls_status: pirls_status_val,
7179 max_abs_eta,
7180 constraint_kkt: None,
7181 artifacts: gam_solve::estimate::FitArtifacts {
7182 pirls: None,
7183 ..Default::default()
7184 },
7185 inner_cycles: 0,
7186 })?
7187 },
7188 design: design.clone(),
7189 adaptive_diagnostics: None,
7190 })
7191}
7192
7193fn enforce_term_constraint_feasibility(
7194 design: &TermCollectionDesign,
7195 fit: &UnifiedFitResult,
7196) -> Result<(), EstimationError> {
7197 const CONSTRAINT_FEASIBILITY_RAW_TOL: f64 = 1e-7;
7211 let tol = CONSTRAINT_FEASIBILITY_RAW_TOL;
7212 let smooth_start = design
7213 .design
7214 .ncols()
7215 .saturating_sub(design.smooth.total_smooth_cols());
7216 let mut violations: Vec<String> = Vec::new();
7217 for term in &design.smooth.terms {
7218 let gr = (smooth_start + term.coeff_range.start)..(smooth_start + term.coeff_range.end);
7219 let beta_local = fit.beta.slice(s![gr.clone()]).to_owned();
7220 if let Some(lb) = term.lower_bounds_local.as_ref() {
7221 let mut worst = 0.0_f64;
7222 let mut worst_idx = 0usize;
7223 for i in 0..lb.len().min(beta_local.len()) {
7224 if lb[i].is_finite() {
7225 let viol = (lb[i] - beta_local[i]).max(0.0);
7226 if viol > worst {
7227 worst = viol;
7228 worst_idx = i;
7229 }
7230 }
7231 }
7232 if worst > tol {
7233 violations.push(format!(
7234 "term='{}' kind=lower-bound maxviolation={:.3e} coeff_index={}",
7235 term.name, worst, worst_idx
7236 ));
7237 }
7238 }
7239 if let Some(lin) = term.linear_constraints_local.as_ref() {
7240 let mut worst = 0.0_f64;
7241 let mut worstrow = 0usize;
7242 for i in 0..lin.a.nrows() {
7243 let norm = lin.a.row(i).dot(&lin.a.row(i)).sqrt();
7244 let inv = if norm > 0.0 { 1.0 / norm } else { 0.0 };
7245 let s = (lin.a.row(i).dot(&beta_local) - lin.b[i]) * inv;
7246 let viol = (-s).max(0.0);
7247 if viol > worst {
7248 worst = viol;
7249 worstrow = i;
7250 }
7251 }
7252 if worst > tol {
7253 violations.push(format!(
7254 "term='{}' kind=linear-inequality maxviolation={:.3e} row={}",
7255 term.name, worst, worstrow
7256 ));
7257 }
7258 }
7259 }
7260
7261 if !violations.is_empty() {
7262 let mut msg = format!(
7263 "constraint violation after fit ({} violating term constraints): {}",
7264 violations.len(),
7265 violations.join(" | ")
7266 );
7267 if let Some(kkt) = fit.constraint_kkt.as_ref() {
7268 msg.push_str(&format!(
7269 "; KKT[primal={:.3e}, dual={:.3e}, comp={:.3e}, stat={:.3e}]",
7270 kkt.primal_feasibility, kkt.dual_feasibility, kkt.complementarity, kkt.stationarity
7271 ));
7272 }
7273 return Err(EstimationError::ParameterConstraintViolation(msg));
7274 }
7275 Ok(())
7276}
7277
7278fn stratified_spatial_subsample(
7279 data: ArrayView2<'_, f64>,
7280 spec: &TermCollectionSpec,
7281 target_size: usize,
7282) -> Vec<usize> {
7283 use rand::SeedableRng;
7284 use rand::rngs::StdRng;
7285 use rand::seq::SliceRandom;
7286
7287 let n = data.nrows();
7288 if n <= target_size {
7289 return (0..n).collect();
7290 }
7291
7292 let spatial_cols: Option<Vec<usize>> =
7293 spec.smooth_terms.iter().find_map(|term| match &term.basis {
7294 SmoothBasisSpec::ThinPlate { feature_cols, .. }
7295 | SmoothBasisSpec::Matern { feature_cols, .. }
7296 | SmoothBasisSpec::Duchon { feature_cols, .. } => {
7297 if !feature_cols.is_empty() {
7298 Some(feature_cols.clone())
7299 } else {
7300 None
7301 }
7302 }
7303 _ => None,
7304 });
7305
7306 let cols = match spatial_cols {
7307 Some(c) if !c.is_empty() => c,
7308 _ => {
7309 let mut rng = StdRng::seed_from_u64(spatial_subsample_seed(data, &[], target_size));
7310 let mut indices: Vec<usize> = (0..n).collect();
7311 indices.shuffle(&mut rng);
7312 indices.truncate(target_size);
7313 indices.sort_unstable();
7314 return indices;
7315 }
7316 };
7317 let mut rng = StdRng::seed_from_u64(spatial_subsample_seed(data, &cols, target_size));
7318
7319 let d = cols.len();
7320 let mut mins = vec![f64::INFINITY; d];
7321 let mut maxs = vec![f64::NEG_INFINITY; d];
7322 for i in 0..n {
7323 for (ax, &col) in cols.iter().enumerate() {
7324 let v = data[[i, col]];
7325 if v < mins[ax] {
7326 mins[ax] = v;
7327 }
7328 if v > maxs[ax] {
7329 maxs[ax] = v;
7330 }
7331 }
7332 }
7333
7334 const TARGET_POINTS_PER_CELL: usize = 5;
7338 let total_cells_target = (target_size / TARGET_POINTS_PER_CELL).max(1);
7339 let cells_per_axis = ((total_cells_target as f64).powf(1.0 / d as f64)).ceil() as usize;
7340 let cells_per_axis = cells_per_axis.max(1);
7341
7342 let mut cell_members: std::collections::HashMap<Vec<usize>, Vec<usize>> =
7343 std::collections::HashMap::new();
7344 for i in 0..n {
7345 let mut cell_key = Vec::with_capacity(d);
7346 for (ax, &col) in cols.iter().enumerate() {
7347 let range = maxs[ax] - mins[ax];
7348 let cell = if range <= 0.0 {
7349 0
7350 } else {
7351 let frac = (data[[i, col]] - mins[ax]) / range;
7352 (frac * cells_per_axis as f64).floor() as usize
7353 };
7354 cell_key.push(cell.min(cells_per_axis - 1));
7355 }
7356 cell_members.entry(cell_key).or_default().push(i);
7357 }
7358
7359 let mut selected: Vec<usize> = Vec::with_capacity(target_size);
7360 let mut remaining_budget = target_size;
7361 let mut remaining_population = n;
7362
7363 let mut cells: Vec<(Vec<usize>, Vec<usize>)> = cell_members.into_iter().collect();
7364 cells.sort_by(|a, b| a.0.cmp(&b.0));
7365
7366 for (_, members) in &mut cells {
7367 if remaining_budget == 0 {
7368 break;
7369 }
7370 let alloc = ((members.len() as f64 / remaining_population as f64) * remaining_budget as f64)
7371 .round() as usize;
7372 let alloc = alloc.max(1).min(members.len()).min(remaining_budget);
7373 members.shuffle(&mut rng);
7374 selected.extend_from_slice(&members[..alloc]);
7375 remaining_budget = remaining_budget.saturating_sub(alloc);
7376 remaining_population = remaining_population.saturating_sub(members.len());
7377 }
7378
7379 if selected.len() > target_size {
7380 selected.shuffle(&mut rng);
7381 selected.truncate(target_size);
7382 }
7383
7384 selected.sort_unstable();
7385 selected
7386}
7387
7388fn spatial_subsample_seed(
7389 data: ArrayView2<'_, f64>,
7390 spatial_cols: &[usize],
7391 target_size: usize,
7392) -> u64 {
7393 let mut state = 0x5350_4154_4941_4C53_u64;
7394 spatial_seed_mix(&mut state, data.nrows() as u64);
7395 spatial_seed_mix(&mut state, data.ncols() as u64);
7396 spatial_seed_mix(&mut state, target_size as u64);
7397 spatial_seed_mix(&mut state, spatial_cols.len() as u64);
7398 for &col in spatial_cols {
7399 spatial_seed_mix(&mut state, col as u64);
7400 }
7401
7402 if data.nrows() > 0 {
7403 let mid = data.nrows() / 2;
7404 let last = data.nrows() - 1;
7405 for &row in &[0usize, mid, last] {
7406 for &col in spatial_cols {
7407 let value = data[[row, col]];
7408 spatial_seed_mix(&mut state, value.to_bits());
7409 }
7410 }
7411 }
7412 state
7413}
7414
7415#[inline]
7416fn spatial_seed_mix(state: &mut u64, value: u64) {
7417 let mut s = value.wrapping_add(*state);
7420 let z = gam_linalg::utils::splitmix64(&mut s);
7421 *state ^= z;
7422 *state = (*state).rotate_left(27).wrapping_mul(0x3C79_AC49_2BA7_B653);
7423}
7424
7425fn sampled_rows(data: ArrayView2<'_, f64>, indices: &[usize]) -> Array2<f64> {
7426 let mut sampled = Array2::<f64>::zeros((indices.len(), data.ncols()));
7427 for (new_row, &orig_row) in indices.iter().enumerate() {
7428 sampled.row_mut(new_row).assign(&data.row(orig_row));
7429 }
7430 sampled
7431}
7432
7433fn spatial_term_user_centers(term: &SmoothTermSpec) -> Option<ArrayView2<'_, f64>> {
7434 match spatial_term_center_strategy(term) {
7435 Some(CenterStrategy::UserProvided(centers)) => Some(centers.view()),
7436 _ => None,
7437 }
7438}
7439
7440fn finite_centered_axis_contrasts(values: &[f64], expected_dim: usize) -> Option<Vec<f64>> {
7441 if values.len() != expected_dim || expected_dim <= 1 {
7442 return None;
7443 }
7444 if values.iter().any(|value| !value.is_finite()) {
7445 return None;
7446 }
7447 Some(center_aniso_log_scales(values))
7448}
7449
7450fn blended_pilot_axis_contrasts(
7451 pilot_data: ArrayView2<'_, f64>,
7452 term: &SmoothTermSpec,
7453 centers: ArrayView2<'_, f64>,
7454) -> Result<Option<Vec<f64>>, BasisError> {
7455 let d = centers.ncols();
7456 if d <= 1 {
7457 return Ok(None);
7458 }
7459 let center_eta = initial_aniso_contrasts(centers);
7460 let standardized_data = standardized_spatial_term_data(pilot_data, term)?;
7461 let data_eta = finite_centered_axis_contrasts(
7462 &initial_aniso_contrasts(standardized_data.view()),
7463 d,
7464 );
7465 let Some(center_eta) = finite_centered_axis_contrasts(¢er_eta, d) else {
7466 return Ok(None);
7467 };
7468 let blended = match data_eta {
7469 Some(data_eta) => center_eta
7470 .iter()
7471 .zip(data_eta.iter())
7472 .map(|(&from_centers, &from_data)| 0.5 * (from_centers + from_data))
7473 .collect::<Vec<_>>(),
7474 None => center_eta,
7475 };
7476 Ok(finite_centered_axis_contrasts(&blended, d))
7477}
7478
7479fn apply_pilot_spatial_psi_reseed(
7480 pilot_data: ArrayView2<'_, f64>,
7481 spec: &TermCollectionSpec,
7482 spatial_terms: &[usize],
7483 kappa_options: &SpatialLengthScaleOptimizationOptions,
7484) -> Result<TermCollectionSpec, EstimationError> {
7485 let dims_per_term = spatial_dims_per_term(spec, spatial_terms);
7486 let use_aniso = has_aniso_terms(spec, spatial_terms);
7487 let log_kappa0 = if use_aniso {
7488 SpatialLogKappaCoords::from_length_scales_aniso(spec, spatial_terms, kappa_options)
7489 } else {
7490 SpatialLogKappaCoords::from_length_scales(spec, spatial_terms, kappa_options)
7491 };
7492 let log_kappa0 = log_kappa0
7493 .reseed_from_data(pilot_data, spec, spatial_terms, kappa_options)
7494 .map_err(EstimationError::BasisError)?;
7495 let log_kappa_lower = if use_aniso {
7496 SpatialLogKappaCoords::lower_bounds_aniso_from_data(
7497 pilot_data,
7498 spec,
7499 spatial_terms,
7500 &dims_per_term,
7501 kappa_options,
7502 )
7503 } else {
7504 SpatialLogKappaCoords::lower_bounds_from_data(
7505 pilot_data,
7506 spec,
7507 spatial_terms,
7508 kappa_options,
7509 )
7510 }
7511 .map_err(EstimationError::BasisError)?;
7512 let log_kappa_upper = if use_aniso {
7513 SpatialLogKappaCoords::upper_bounds_aniso_from_data(
7514 pilot_data,
7515 spec,
7516 spatial_terms,
7517 &dims_per_term,
7518 kappa_options,
7519 )
7520 } else {
7521 SpatialLogKappaCoords::upper_bounds_from_data(
7522 pilot_data,
7523 spec,
7524 spatial_terms,
7525 kappa_options,
7526 )
7527 }
7528 .map_err(EstimationError::BasisError)?;
7529 log_kappa0
7530 .clamp_to_bounds(&log_kappa_lower, &log_kappa_upper)
7531 .apply_tospec(spec, spatial_terms)
7532}
7533
7534pub(crate) fn apply_spatial_anisotropy_pilot_initializer(
7535 data: ArrayView2<'_, f64>,
7536 spec: &mut TermCollectionSpec,
7537 spatial_terms: &[usize],
7538 target_size: usize,
7539 kappa_options: &SpatialLengthScaleOptimizationOptions,
7540) -> Result<usize, EstimationError> {
7541 if target_size == 0 || data.nrows() <= target_size.saturating_mul(2) || spatial_terms.is_empty()
7542 {
7543 return Ok(0);
7544 }
7545 if !has_aniso_terms(spec, spatial_terms) {
7546 return Ok(0);
7547 }
7548 let indices = stratified_spatial_subsample(data, spec, target_size);
7549 let pilot_data = sampled_rows(data, &indices);
7550 let mut working = spec.clone();
7551 let mut updated_terms = 0usize;
7552 const GEOMETRY_UPDATES: usize = 2;
7553
7554 for pass in 0..GEOMETRY_UPDATES {
7555 let planned_terms = plan_joint_spatial_centers_for_term_blocks(
7556 pilot_data.view(),
7557 &[working.smooth_terms.clone()],
7558 )
7559 .and_then(|mut blocks| {
7560 blocks.pop().ok_or_else(|| {
7561 BasisError::InvalidInput(
7562 "pilot geometry initializer produced no smooth-term block".to_string(),
7563 )
7564 })
7565 })
7566 .map_err(EstimationError::BasisError)?;
7567
7568 for &term_idx in spatial_terms {
7569 let Some(current_eta) = get_spatial_aniso_log_scales(&working, term_idx) else {
7570 continue;
7571 };
7572 let Some(d) = get_spatial_feature_dim(&working, term_idx) else {
7573 continue;
7574 };
7575 if d <= 1 || current_eta.len() != d {
7576 continue;
7577 }
7578 let Some(planned_term) = planned_terms.get(term_idx) else {
7579 continue;
7580 };
7581 let Some(centers) = spatial_term_user_centers(planned_term) else {
7582 continue;
7583 };
7584 let Some(eta) = blended_pilot_axis_contrasts(
7585 pilot_data.view(),
7586 planned_term,
7587 centers,
7588 )
7589 .map_err(EstimationError::BasisError)?
7590 else {
7591 continue;
7592 };
7593 set_spatial_aniso_log_scales(&mut working, term_idx, eta)?;
7594 updated_terms += usize::from(pass == 0);
7595 }
7596
7597 working = apply_pilot_spatial_psi_reseed(
7598 pilot_data.view(),
7599 &working,
7600 spatial_terms,
7601 kappa_options,
7602 )?;
7603 }
7604
7605 if updated_terms > 0 {
7606 log::info!(
7607 "[spatial-kappa] initialized anisotropy from {}-row pilot geometry for {} spatial term(s); proceeding to full-data optimization",
7608 indices.len(),
7609 updated_terms
7610 );
7611 *spec = working;
7612 }
7613 Ok(updated_terms)
7614}
7615
7616pub(crate) fn spatial_length_scale_term_indices(spec: &TermCollectionSpec) -> Vec<usize> {
7617 spec.smooth_terms
7618 .iter()
7619 .enumerate()
7620 .filter_map(|(idx, _)| spatial_term_supports_hyper_optimization(spec, idx).then_some(idx))
7621 .collect()
7622}
7623
7624fn fit_score(fit: &UnifiedFitResult) -> f64 {
7636 if fit.reml_score.is_finite() {
7637 return fit.reml_score;
7638 }
7639 let score = 0.5 * fit.deviance + 0.5 * fit.stable_penalty_term;
7640 if score.is_finite() {
7641 score
7642 } else {
7643 f64::INFINITY
7644 }
7645}
7646
7647fn is_recoverable_trial_point_error(err: &EstimationError) -> bool {
7669 matches!(err, EstimationError::BasisError(_))
7670 || err.is_inner_solve_retreat()
7671 || is_recoverable_fit_inference_finiteness_error(err)
7672}
7673
7674fn is_recoverable_fit_inference_finiteness_error(err: &EstimationError) -> bool {
7675 let EstimationError::InvalidInput(message) = err else {
7676 return false;
7677 };
7678
7679 message.contains("must be finite")
7680 && [
7681 "fit_result.beta_covariance_frequentist",
7682 "fit_result.coefficient_influence",
7683 "fit_result.weighted_gram",
7684 ]
7685 .iter()
7686 .any(|field| message.contains(field))
7687}
7688
7689#[cfg(test)]
7690mod spatial_trial_recovery_tests {
7691 use super::*;
7692
7693 #[test]
7694 fn nonfinite_frequentist_covariance_is_recoverable_trial_point() {
7695 let err = EstimationError::InvalidInput(
7696 "fit_result.beta_covariance_frequentist[0] must be finite, got NaN".to_string(),
7697 );
7698
7699 assert!(
7700 is_recoverable_trial_point_error(&err),
7701 "singular trial-point curvature should make spatial κ retreat, not abort"
7702 );
7703 }
7704
7705 #[test]
7706 fn arbitrary_invalid_input_remains_fatal_trial_point_error() {
7707 let err = EstimationError::InvalidInput("outer rho bounds are invalid".to_string());
7708
7709 assert!(
7710 !is_recoverable_trial_point_error(&err),
7711 "the spatial κ recovery gate must not mask unrelated invalid inputs"
7712 );
7713 }
7714}
7715
7716fn require_successful_spatial_optimization_result<T>(
7717 initial_score: f64,
7718 result: Result<Option<(T, f64)>, EstimationError>,
7719) -> Result<T, EstimationError> {
7720 match result {
7721 Ok(Some((value, exact_score))) => {
7722 const SCORE_DRIFT_ABS_TOL: f64 = 1e-6;
7731 const SCORE_DRIFT_REL_TOL: f64 = 1e-8;
7732 let tol = SCORE_DRIFT_ABS_TOL.max(initial_score.abs() * SCORE_DRIFT_REL_TOL);
7733 if exact_score <= initial_score + tol {
7734 Ok(value)
7735 } else {
7736 Err(EstimationError::RemlOptimizationFailed(format!(
7737 "spatial kappa optimization made REML score worse ({initial_score:.6e} -> {exact_score:.6e})"
7738 )))
7739 }
7740 }
7741 Ok(None) => Err(EstimationError::RemlOptimizationFailed(
7742 "spatial kappa optimization is unavailable for one or more eligible spatial terms"
7743 .to_string(),
7744 )),
7745 Err(err) => Err(EstimationError::RemlOptimizationFailed(format!(
7746 "spatial kappa optimization failed: {err}"
7747 ))),
7748 }
7749}
7750
7751fn external_opts_for_design(
7752 family: &LikelihoodSpec,
7753 design: &TermCollectionDesign,
7754 options: &FitOptions,
7755) -> ExternalOptimOptions {
7756 ExternalOptimOptions {
7757 family: family.clone(),
7758 latent_cloglog: options.latent_cloglog,
7759 mixture_link: options.mixture_link.clone(),
7760 optimize_mixture: options.optimize_mixture,
7761 sas_link: options.sas_link,
7762 optimize_sas: options.optimize_sas,
7763 compute_inference: options.compute_inference,
7764 skip_rho_posterior_inference: options.skip_rho_posterior_inference,
7765 max_iter: options.max_iter,
7766 tol: options.tol,
7767 nullspace_dims: design.nullspace_dims.clone(),
7768 linear_constraints: design.linear_constraints.clone(),
7769 firth_bias_reduction: Some(options.firth_bias_reduction),
7770 penalty_shrinkage_floor: options.penalty_shrinkage_floor,
7771 rho_prior: options.rho_prior.clone(),
7772 kronecker_penalty_system: design.kronecker_penalty_system(),
7775 kronecker_factored: design
7776 .smooth
7777 .terms
7778 .iter()
7779 .find_map(|t| t.kronecker_factored.clone()),
7780 persist_warm_start_disk: options.persist_warm_start_disk,
7781 }
7782}
7783
7784fn evaluate_joint_reml_outer_eval_at_theta(
7792 evaluator: &mut gam_solve::estimate::ExternalJointHyperEvaluator<'_>,
7793 design: &TermCollectionDesign,
7794 theta: &Array1<f64>,
7795 rho_dim: usize,
7796 hyper_dirs: Vec<gam_solve::estimate::reml::DirectionalHyperParam>,
7797 warm_start_beta: Option<ArrayView1<'_, f64>>,
7798 order: gam_solve::rho_optimizer::OuterEvalOrder,
7799 design_revision: Option<u64>,
7800) -> Result<(f64, Array1<f64>, gam_problem::HessianValue), EstimationError> {
7801 evaluator.evaluate_with_order(
7802 &design.design,
7803 &design.penalties,
7804 &design.nullspace_dims,
7805 design.linear_constraints.clone(),
7806 theta,
7807 rho_dim,
7808 hyper_dirs,
7809 warm_start_beta,
7810 "evaluate_joint_reml_outer_eval_at_theta",
7811 order,
7812 design_revision,
7813 )
7814}
7815
7816fn evaluate_joint_reml_efs_at_theta(
7817 evaluator: &mut gam_solve::estimate::ExternalJointHyperEvaluator<'_>,
7818 design: &TermCollectionDesign,
7819 theta: &Array1<f64>,
7820 rho_dim: usize,
7821 hyper_dirs: Vec<gam_solve::estimate::reml::DirectionalHyperParam>,
7822 warm_start_beta: Option<ArrayView1<'_, f64>>,
7823 design_revision: Option<u64>,
7824) -> Result<gam_problem::EfsEval, EstimationError> {
7825 evaluator.evaluate_efs(
7826 &design.design,
7827 &design.penalties,
7828 &design.nullspace_dims,
7829 design.linear_constraints.clone(),
7830 theta,
7831 rho_dim,
7832 hyper_dirs,
7833 warm_start_beta,
7834 "evaluate_joint_reml_efs_at_theta",
7835 design_revision,
7836 )
7837}
7838
7839fn exact_joint_spatial_outer_hessian_available(
7840 family: &LikelihoodSpec,
7841 design: &TermCollectionDesign,
7842) -> bool {
7843 let family_supported = match &family.response {
7866 ResponseFamily::Gaussian
7867 | ResponseFamily::Binomial
7868 | ResponseFamily::Poisson
7869 | ResponseFamily::Tweedie { .. }
7870 | ResponseFamily::NegativeBinomial { .. }
7871 | ResponseFamily::Beta { .. }
7872 | ResponseFamily::Gamma
7873 | ResponseFamily::RoystonParmar => true,
7874 };
7875 family_supported && design.design.ncols() > 0
7878}
7879
7880fn try_build_spatial_term_log_kappa_derivativeinfo(
7881 data: ArrayView2<'_, f64>,
7882 resolvedspec: &TermCollectionSpec,
7883 design: &TermCollectionDesign,
7884 term_idx: usize,
7885) -> Result<Option<SpatialPsiDerivative>, EstimationError> {
7886 let Some((
7887 global_range,
7888 total_p,
7889 x_psi_local,
7890 s_psi_local_check,
7891 x_psi_psi_local,
7892 s_psi_psi_local,
7893 s_psi_components_local,
7894 s_psi_psi_components_local,
7895 implicit_operator,
7896 )) = try_build_spatial_term_log_kappa_derivative(data, resolvedspec, design, term_idx)?
7897 else {
7898 return Ok(None);
7899 };
7900 let Some(penalty_range) = design
7901 .smooth_term_penalty_range(term_idx)
7902 .map_err(EstimationError::InvalidInput)?
7903 else {
7904 return Ok(None);
7905 };
7906 let penalty_start = penalty_range.start;
7907 if s_psi_components_local.is_empty() || s_psi_psi_components_local.is_empty() {
7908 return Ok(None);
7909 }
7910 if s_psi_components_local.len() != s_psi_psi_components_local.len() {
7911 return Ok(None);
7912 }
7913 let penalty_indices = (0..s_psi_components_local.len())
7914 .map(|j| penalty_start + j)
7915 .collect::<Vec<_>>();
7916 let penalty_index = penalty_indices[0];
7917 if s_psi_local_check.nrows() == 0 || s_psi_psi_local.nrows() == 0 {
7918 return Ok(None);
7919 }
7920 Ok(Some(SpatialPsiDerivative {
7921 penalty_index,
7922 penalty_indices,
7923 global_range,
7924 total_p,
7925 x_psi_local,
7926 s_psi_components_local,
7927 x_psi_psi_local,
7928 s_psi_psi_components_local,
7929 aniso_group_id: None,
7930 aniso_cross_designs: None,
7931 aniso_cross_penalty_provider: None,
7932 implicit_operator,
7933 implicit_axis: 0,
7934 }))
7935}
7936
7937pub(crate) fn try_build_spatial_log_kappa_derivativeinfo_list(
7938 data: ArrayView2<'_, f64>,
7939 resolvedspec: &TermCollectionSpec,
7940 design: &TermCollectionDesign,
7941 spatial_terms: &[usize],
7942) -> Result<Option<Vec<SpatialPsiDerivative>>, EstimationError> {
7943 let mut out = Vec::new();
7944 let mut aniso_gid = 0usize;
7945 for &term_idx in spatial_terms {
7946 if spatial_term_uses_per_axis_psi(resolvedspec, term_idx) {
7947 if let Some(entries) = try_build_spatial_term_log_kappa_aniso_derivativeinfos(
7948 data,
7949 resolvedspec,
7950 design,
7951 term_idx,
7952 aniso_gid,
7953 )? {
7954 aniso_gid += 1;
7955 out.extend(entries);
7956 continue;
7957 } else {
7958 return Ok(None);
7959 }
7960 }
7961 let Some(info) =
7962 try_build_spatial_term_log_kappa_derivativeinfo(data, resolvedspec, design, term_idx)?
7963 else {
7964 return Ok(None);
7965 };
7966 out.push(info);
7967 }
7968 Ok(Some(out))
7969}
7970
7971fn try_build_spatial_term_log_kappa_aniso_derivativeinfos(
7973 data: ArrayView2<'_, f64>,
7974 resolvedspec: &TermCollectionSpec,
7975 design: &TermCollectionDesign,
7976 term_idx: usize,
7977 aniso_group_id: usize,
7978) -> Result<Option<Vec<SpatialPsiDerivative>>, EstimationError> {
7979 let Some(smooth_term) = design.smooth.terms.get(term_idx) else {
7980 return Ok(None);
7981 };
7982 let Some(termspec) = resolvedspec.smooth_terms.get(term_idx) else {
7983 return Ok(None);
7984 };
7985 let mut aniso_result = match &termspec.basis {
7986 SmoothBasisSpec::Sphere { .. } => return Ok(None),
7987 SmoothBasisSpec::Matern {
7988 feature_cols,
7989 spec,
7990 input_scale,
7991 } => {
7992 let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
7993 if let Some(scale) = input_scale {
7994 scale.standardize(&mut x);
7995 }
7996 let mut spec_operator = spec.clone();
8005 spec_operator.double_penalty = false;
8006 build_matern_basis_log_kappa_aniso_derivatives(x.view(), &spec_operator)
8007 .map_err(EstimationError::from)?
8008 }
8009 SmoothBasisSpec::MeasureJet {
8015 feature_cols,
8016 spec,
8017 input_scale,
8018 } => {
8019 let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
8020 if let Some(scale) = input_scale {
8021 scale.standardize(&mut x);
8022 }
8023 build_measure_jet_basis_psi_derivatives(x.view(), spec)
8024 .map_err(EstimationError::from)?
8025 }
8026 _ => return Ok(None),
8027 };
8028 let d = if let Some(ref op) = aniso_result.implicit_operator {
8031 op.n_axes()
8032 } else if !aniso_result.design_first.is_empty() {
8033 aniso_result.design_first.len()
8034 } else {
8035 0
8036 };
8037 if d == 0 {
8038 return Ok(None);
8039 }
8040 let Some(penalty_range) = design
8041 .smooth_term_penalty_range(term_idx)
8042 .map_err(EstimationError::InvalidInput)?
8043 else {
8044 return Ok(None);
8045 };
8046 let penalty_start = penalty_range.start;
8047 let p_total = design.design.ncols();
8048 let smooth_start = p_total.saturating_sub(design.smooth.total_smooth_cols());
8049 let global_range = (smooth_start + smooth_term.coeff_range.start)
8050 ..(smooth_start + smooth_term.coeff_range.end);
8051 let num_penalties = aniso_result.penalties_first[0].len();
8052 let penalty_indices: Vec<usize> = (0..num_penalties).map(|j| penalty_start + j).collect();
8053 let penalties_cross_provider = aniso_result.penalties_cross_provider.clone();
8054
8055 let use_implicit_design = aniso_result.design_first.is_empty();
8059 let implicit_op_arc = aniso_result
8060 .implicit_operator
8061 .as_ref()
8062 .map(|op| std::sync::Arc::new(op.clone()));
8063
8064 let mut entries = Vec::with_capacity(d);
8065 for a in 0..d {
8066 let (x_psi_local, x_psi_psi_local) = if use_implicit_design {
8067 (Array2::<f64>::zeros((0, 0)), Array2::<f64>::zeros((0, 0)))
8073 } else {
8074 let x_first = std::mem::take(&mut aniso_result.design_first[a]);
8079 let x_second = std::mem::take(&mut aniso_result.design_second_diag[a]);
8080 if x_first.ncols() != smooth_term.coeff_range.len() {
8081 return Ok(None);
8082 }
8083 (x_first, x_second)
8084 };
8085 let s_psi_components = std::mem::take(&mut aniso_result.penalties_first[a]);
8086 let s_psi_psi_components = std::mem::take(&mut aniso_result.penalties_second_diag[a]);
8087 let cross_designs = if implicit_op_arc.is_some() {
8093 let mut cd = Vec::with_capacity(d - 1);
8094 for b in 0..d {
8095 if b == a {
8096 continue;
8097 }
8098 cd.push((b, Array2::<f64>::zeros((0, 0))));
8099 }
8100 cd
8101 } else if !aniso_result.design_second_cross.is_empty() {
8102 let mut cd = Vec::new();
8103 for (cross_idx, &(pa, pb)) in aniso_result.design_second_cross_pairs.iter().enumerate()
8104 {
8105 if pa == a {
8106 cd.push((pb, aniso_result.design_second_cross[cross_idx].clone()));
8107 } else if pb == a {
8108 cd.push((pa, aniso_result.design_second_cross[cross_idx].clone()));
8109 }
8110 }
8111 cd
8112 } else {
8113 Vec::new()
8114 };
8115 let cross_penalty_provider = if d > 1 {
8116 let penalties_cross_provider = penalties_cross_provider.clone();
8117 Some(std::sync::Arc::new(
8118 move |b_axis: usize| -> Result<Vec<Array2<f64>>, EstimationError> {
8119 if b_axis == a {
8120 return Ok(Vec::new());
8121 }
8122 let (axis_lo, axis_hi) = if a < b_axis { (a, b_axis) } else { (b_axis, a) };
8123 if let Some(provider) = penalties_cross_provider.as_ref() {
8124 provider
8125 .evaluate(axis_lo, axis_hi)
8126 .map_err(EstimationError::from)
8127 } else {
8128 Ok(Vec::new())
8132 }
8133 },
8134 )
8135 as std::sync::Arc<
8136 dyn Fn(usize) -> Result<Vec<Array2<f64>>, EstimationError>
8137 + Send
8138 + Sync
8139 + 'static,
8140 >)
8141 } else {
8142 None
8143 };
8144
8145 entries.push(SpatialPsiDerivative {
8146 penalty_index: penalty_indices[0],
8147 penalty_indices: penalty_indices.clone(),
8148 global_range: global_range.clone(),
8149 total_p: p_total,
8150 x_psi_local,
8151 s_psi_components_local: s_psi_components,
8152 x_psi_psi_local,
8153 s_psi_psi_components_local: s_psi_psi_components,
8154 aniso_group_id: Some(aniso_group_id),
8155 aniso_cross_designs: if cross_designs.is_empty() {
8156 None
8157 } else {
8158 Some(cross_designs)
8159 },
8160 aniso_cross_penalty_provider: cross_penalty_provider,
8161 implicit_operator: implicit_op_arc.clone(),
8162 implicit_axis: a,
8163 });
8164 }
8165 Ok(Some(entries))
8166}
8167
8168#[cfg(test)]
8169mod glm_eta_observation_fd_tests {
8170 use super::*;
8176 use ndarray::array;
8177
8178 fn one_obs_weight(
8179 spec: &LikelihoodSpec,
8180 y: f64,
8181 weight: f64,
8182 eta: f64,
8183 ) -> StandardFamilyObservationState {
8184 let yv = Array1::from_vec(vec![y]);
8185 let wv = Array1::from_vec(vec![weight]);
8186 let ev = Array1::from_vec(vec![eta]);
8187 evaluate_standard_familyobservations(spec.clone(), None, None, None, &yv, &wv, &ev)
8188 .expect("standard family observation state assembles")
8189 }
8190
8191 fn one_obs(spec: &LikelihoodSpec, y: f64, eta: f64) -> StandardFamilyObservationState {
8192 one_obs_weight(spec, y, 1.0, eta)
8193 }
8194
8195 fn one_obs_resolved(
8196 likelihood: &gam_spec::GlmLikelihoodSpec,
8197 y: f64,
8198 weight: f64,
8199 eta: f64,
8200 ) -> StandardFamilyObservationState {
8201 evaluate_resolved_standard_family_observations(
8202 likelihood,
8203 None,
8204 None,
8205 None,
8206 &array![y],
8207 &array![weight],
8208 &array![eta],
8209 )
8210 .expect("resolved standard family observation state assembles")
8211 }
8212
8213 #[test]
8214 fn bounded_gamma_and_tweedie_use_the_resolved_likelihood_scale() {
8215 let gamma_unit = gam_spec::GlmLikelihoodSpec {
8216 spec: LikelihoodSpec::gamma_log(),
8217 scale: gam_spec::LikelihoodScaleMetadata::FixedGammaShape { shape: 1.0 },
8218 };
8219 let gamma_scaled = gam_spec::GlmLikelihoodSpec {
8220 spec: LikelihoodSpec::gamma_log(),
8221 scale: gam_spec::LikelihoodScaleMetadata::FixedGammaShape { shape: 8.0 },
8222 };
8223 let unit = one_obs_resolved(&gamma_unit, 2.3, 0.7, 0.2);
8224 let scaled = one_obs_resolved(&gamma_scaled, 2.3, 0.7, 0.2);
8225 for (label, actual, base) in [
8226 ("Gamma score", scaled.score[0], unit.score[0]),
8227 (
8228 "Gamma Fisher weight",
8229 scaled.fisherweight[0],
8230 unit.fisherweight[0],
8231 ),
8232 (
8233 "Gamma observed Hessian",
8234 scaled.neghessian_eta[0],
8235 unit.neghessian_eta[0],
8236 ),
8237 (
8238 "Gamma Hessian derivative",
8239 scaled.neghessian_eta_derivative[0],
8240 unit.neghessian_eta_derivative[0],
8241 ),
8242 (
8243 "Gamma log likelihood",
8244 scaled.log_likelihood,
8245 unit.log_likelihood,
8246 ),
8247 ] {
8248 let expected = 8.0 * base;
8249 assert!(
8250 (actual - expected).abs() <= 32.0 * f64::EPSILON * expected.abs().max(1.0),
8251 "{label} scale mismatch: actual={actual}, expected={expected}"
8252 );
8253 }
8254
8255 let tweedie_unit = gam_spec::GlmLikelihoodSpec {
8256 spec: LikelihoodSpec::tweedie_log(1.5),
8257 scale: gam_spec::LikelihoodScaleMetadata::FixedDispersion { phi: 1.0 },
8258 };
8259 let tweedie_scaled = gam_spec::GlmLikelihoodSpec {
8260 spec: LikelihoodSpec::tweedie_log(1.5),
8261 scale: gam_spec::LikelihoodScaleMetadata::FixedDispersion { phi: 0.25 },
8262 };
8263 let unit = one_obs_resolved(&tweedie_unit, 1.7, 0.8, -0.1);
8264 let scaled = one_obs_resolved(&tweedie_scaled, 1.7, 0.8, -0.1);
8265 for (actual, base) in [
8266 (scaled.score[0], unit.score[0]),
8267 (scaled.fisherweight[0], unit.fisherweight[0]),
8268 (scaled.neghessian_eta[0], unit.neghessian_eta[0]),
8269 (
8270 scaled.neghessian_eta_derivative[0],
8271 unit.neghessian_eta_derivative[0],
8272 ),
8273 (scaled.log_likelihood, unit.log_likelihood),
8274 ] {
8275 let expected = 4.0 * base;
8276 assert!((actual - expected).abs() <= 32.0 * f64::EPSILON * expected.abs().max(1.0));
8277 }
8278 }
8279
8280 #[test]
8281 fn bounded_zero_rows_are_dormant_and_weight_preflight_is_atomic() {
8282 let likelihood = gam_spec::GlmLikelihoodSpec::canonical(LikelihoodSpec::poisson_log());
8283 let dormant = evaluate_resolved_standard_family_observations(
8284 &likelihood,
8285 None,
8286 None,
8287 None,
8288 &array![f64::NAN, 2.0],
8289 &array![0.0, 1.0],
8290 &array![f64::NAN, 0.2],
8291 )
8292 .expect("zero-weight response and predictor are dormant");
8293 assert_eq!(dormant.score[0], 0.0);
8294 assert_eq!(dormant.fisherweight[0], 0.0);
8295
8296 let error = evaluate_resolved_standard_family_observations(
8297 &likelihood,
8298 None,
8299 None,
8300 None,
8301 &array![f64::NAN, 2.0],
8302 &array![1.0, f64::NAN],
8303 &array![f64::NAN, 0.2],
8304 )
8305 .expect_err("later invalid weight must refuse before row evaluation");
8306 assert!(
8307 error.to_string().contains("row 2 has invalid prior weight"),
8308 "unexpected atomic preflight error: {error}"
8309 );
8310 }
8311
8312 fn check_fd(label: &str, spec: &LikelihoodSpec, y: f64, eta: f64) {
8313 let h = 1e-5;
8314 let s0 = one_obs(spec, y, eta);
8315 let sp = one_obs(spec, y, eta + h);
8316 let sm = one_obs(spec, y, eta - h);
8317
8318 let score_fd = (sp.log_likelihood - sm.log_likelihood) / (2.0 * h);
8320 let score = s0.score[0];
8321 assert!(
8322 (score - score_fd).abs() <= 1e-4 * (1.0 + score.abs()),
8323 "{label}: score {score} vs FD {score_fd}"
8324 );
8325
8326 let neghess_fd = -(sp.score[0] - sm.score[0]) / (2.0 * h);
8328 let neghess = s0.neghessian_eta[0];
8329 assert!(
8330 (neghess - neghess_fd).abs() <= 1e-3 * (1.0 + neghess.abs()),
8331 "{label}: neghessian_eta {neghess} vs FD {neghess_fd}"
8332 );
8333
8334 let nhd_fd = (sp.neghessian_eta[0] - sm.neghessian_eta[0]) / (2.0 * h);
8336 let nhd = s0.neghessian_eta_derivative[0];
8337 assert!(
8338 (nhd - nhd_fd).abs() <= 1e-2 * (1.0 + nhd.abs()),
8339 "{label}: neghessian_eta_derivative {nhd} vs FD {nhd_fd}"
8340 );
8341 }
8342
8343 #[test]
8344 fn poisson_gamma_nb_tweedie_arms_match_finite_differences_1615_1616() {
8345 let log = InverseLink::Standard(StandardLink::Log);
8346 let poisson = LikelihoodSpec {
8347 response: ResponseFamily::Poisson,
8348 link: log.clone(),
8349 };
8350 check_fd("poisson y=3", &poisson, 3.0, 0.4);
8351 check_fd("poisson y=0", &poisson, 0.0, -0.2);
8352
8353 let gamma = LikelihoodSpec {
8354 response: ResponseFamily::Gamma,
8355 link: log.clone(),
8356 };
8357 check_fd("gamma y=2.5", &gamma, 2.5, 0.3);
8358 check_fd("gamma y=0.7", &gamma, 0.7, -0.1);
8359
8360 let nb = LikelihoodSpec {
8361 response: ResponseFamily::NegativeBinomial {
8362 theta: 1.5,
8363 theta_fixed: true,
8364 },
8365 link: log.clone(),
8366 };
8367 check_fd("negbin y=4", &nb, 4.0, 0.5);
8368 check_fd("negbin y=0", &nb, 0.0, -0.3);
8369
8370 let tweedie = LikelihoodSpec {
8371 response: ResponseFamily::Tweedie { p: 1.5 },
8372 link: log.clone(),
8373 };
8374 check_fd("tweedie y=2", &tweedie, 2.0, 0.25);
8375 check_fd("tweedie y=0.5", &tweedie, 0.5, -0.15);
8376 }
8377
8378 #[test]
8379 fn binomial_natural_coordinate_towers_match_finite_differences() {
8380 for (label, family, eta) in [
8381 ("logit", LikelihoodSpec::binomial_logit(), 0.7),
8382 ("probit", LikelihoodSpec::binomial_probit(), -1.1),
8383 ("cloglog", LikelihoodSpec::binomial_cloglog(), 0.4),
8384 (
8385 "loglog",
8386 LikelihoodSpec::try_new(
8387 ResponseFamily::Binomial,
8388 InverseLink::Standard(StandardLink::LogLog),
8389 )
8390 .unwrap(),
8391 -0.35,
8392 ),
8393 (
8394 "cauchit",
8395 LikelihoodSpec::try_new(
8396 ResponseFamily::Binomial,
8397 InverseLink::Standard(StandardLink::Cauchit),
8398 )
8399 .unwrap(),
8400 1.25,
8401 ),
8402 ] {
8403 check_fd(label, &family, 0.37, eta);
8404 }
8405 }
8406
8407 #[test]
8408 fn logit_observation_geometry_carries_the_prior_weight_everywhere() {
8409 let eta = 1.75;
8410 let y = 0.3;
8411 let weight = 7.25;
8412 let state = one_obs_weight(&LikelihoodSpec::binomial_logit(), y, weight, eta);
8413 let jet = logit_inverse_link_jet5(eta);
8414 for (got, expected) in [
8415 (state.fisherweight[0], weight * jet.d1),
8416 (state.neghessian_eta[0], weight * jet.d1),
8417 (state.neghessian_eta_derivative[0], weight * jet.d2),
8418 (state.score[0], weight * (y - jet.mu)),
8419 ] {
8420 assert!((got - expected).abs() <= 4.0 * f64::EPSILON * (1.0 + expected.abs()));
8421 }
8422 }
8423
8424 #[test]
8425 fn tiny_positive_and_zero_weights_are_not_projected() {
8426 let tiny = 1e-200;
8427 let logit = one_obs_weight(&LikelihoodSpec::binomial_logit(), 0.4, tiny, 0.0);
8428 assert!((logit.fisherweight[0] / tiny - 0.25).abs() <= 2.0 * f64::EPSILON);
8429 assert!(logit.fisherweight[0] < 1e-190);
8430
8431 let zero = one_obs_weight(&LikelihoodSpec::gaussian_identity(), 3.0, 0.0, -2.0);
8432 assert_eq!(zero.score[0], 0.0);
8433 assert_eq!(zero.fisherweight[0], 0.0);
8434 assert_eq!(zero.neghessian_eta[0], 0.0);
8435 assert_eq!(zero.neghessian_eta_derivative[0], 0.0);
8436 assert_eq!(zero.log_likelihood, 0.0);
8437 assert_eq!(exact_standard_working_response(&zero).unwrap()[0], -2.0);
8438 }
8439
8440 #[test]
8441 fn log_link_tails_balance_tiny_weights_before_certification() {
8442 let poisson = one_obs_weight(&LikelihoodSpec::poisson_log(), 0.0, 1e-300, 700.0);
8443 assert!(poisson.fisherweight[0].is_finite() && poisson.fisherweight[0] > 1.0);
8444 assert!(poisson.score[0].is_finite());
8445 assert!(poisson.log_likelihood.is_finite());
8446
8447 let gamma = one_obs_weight(&LikelihoodSpec::gamma_log(), 1.0, 1e-300, -700.0);
8448 assert!(gamma.neghessian_eta[0].is_finite() && gamma.neghessian_eta[0] > 1.0);
8449 assert!(gamma.score[0].is_finite());
8450 assert!(gamma.log_likelihood.is_finite());
8451 }
8452
8453 #[test]
8454 fn invalid_weights_and_nonfinite_inputs_are_refused_in_row_order() {
8455 let family = LikelihoodSpec::gaussian_identity();
8456 let y = array![1.0, 2.0];
8457 let eta = array![0.0, 0.0];
8458 for weights in [array![-1.0, 1.0], array![f64::NAN, 1.0]] {
8459 let err = evaluate_standard_familyobservations(
8460 family.clone(),
8461 None,
8462 None,
8463 None,
8464 &y,
8465 &weights,
8466 &eta,
8467 )
8468 .expect_err("invalid prior weight must be refused");
8469 assert!(err.to_string().contains("row 0"), "{err}");
8470 }
8471
8472 let err = evaluate_standard_familyobservations(
8473 family,
8474 None,
8475 None,
8476 None,
8477 &array![f64::NAN],
8478 &array![0.0],
8479 &array![0.0],
8480 )
8481 .expect_err("a non-finite response may not hide behind zero weight");
8482 assert!(err.to_string().contains("row 0"), "{err}");
8483 }
8484
8485 #[test]
8486 fn unrepresentable_cloglog_curvature_is_refused_without_a_floor() {
8487 let err = evaluate_standard_familyobservations(
8488 LikelihoodSpec::binomial_cloglog(),
8489 None,
8490 None,
8491 None,
8492 &array![1.0],
8493 &array![1.0],
8494 &array![18.0],
8495 )
8496 .expect_err("mathematically sub-f64 Fisher information must be refused");
8497 assert!(err.to_string().contains("Fisher weight"), "{err}");
8498 }
8499
8500 #[test]
8501 fn bounded_covariance_requires_a_certified_strict_spd_precision() {
8502 let covariance = certified_bounded_posterior_covariance(
8503 &array![[4.0, 1.0], [1.0, 3.0]],
8504 "bounded covariance regression",
8505 )
8506 .expect("strict SPD precision");
8507 assert!((covariance[[0, 0]] - 3.0 / 11.0).abs() < 1e-14);
8508 assert!((covariance[[0, 1]] + 1.0 / 11.0).abs() < 1e-14);
8509 assert!((covariance[[1, 1]] - 4.0 / 11.0).abs() < 1e-14);
8510
8511 for invalid in [
8512 array![[1.0, 1.0], [1.0, 1.0]],
8513 array![[1.0, 2.0], [2.0, 1.0]],
8514 ] {
8515 assert!(
8516 certified_bounded_posterior_covariance(
8517 &invalid,
8518 "invalid bounded covariance regression"
8519 )
8520 .is_err(),
8521 "singular/indefinite precision must not become a pseudo-covariance"
8522 );
8523 }
8524 }
8525}