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 mut grouped_options = base_fit_opts.clone();
74 grouped_options.rho_prior = realized.rho_prior;
75 let fitted = FittedTermCollection {
76 fit: gam_solve::estimate::fit_gam_with_penalty_specs(
77 design.design.clone(),
78 y,
79 weights,
80 offset,
81 realized.penalty_specs,
82 realized.nullspace_dims,
83 family.clone(),
84 &grouped_options,
85 )?,
86 design,
87 adaptive_diagnostics: None,
88 };
89 enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
90 Ok(fitted)
91}
92
93pub fn fit_term_collection_with_penalty_block_gamma_prior_callback<F>(
94 data: ArrayView2<'_, f64>,
95 y: ArrayView1<'_, f64>,
96 weights: ArrayView1<'_, f64>,
97 offset: ArrayView1<'_, f64>,
98 spec: &TermCollectionSpec,
99 callback: F,
100 family: LikelihoodSpec,
101 options: &FitOptions,
102) -> Result<FittedTermCollection, EstimationError>
103where
104 F: FnMut(&PenaltyBlockGammaPriorMetadata<'_>) -> Option<(f64, f64)>,
105{
106 let design = build_term_collection_design_with_policy(data, spec, &options.resource_policy)?;
107 let mut fit_opts = adaptive_fit_options_base(options, &design);
108 fit_opts.rho_prior = realize_penalty_block_gamma_priors(&design, callback)
109 .map_err(EstimationError::BasisError)?;
110 let fitted = FittedTermCollection {
111 fit: fit_gamwith_heuristic_lambdas(
112 design.design.clone(),
113 y,
114 weights,
115 offset,
116 &design.penalties,
117 None,
118 family.clone(),
119 &fit_opts,
120 )?,
121 design,
122 adaptive_diagnostics: None,
123 };
124 enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
125 Ok(fitted)
126}
127
128pub fn fit_term_collection_with_penalty_block_gamma_priors(
129 data: ArrayView2<'_, f64>,
130 y: ArrayView1<'_, f64>,
131 weights: ArrayView1<'_, f64>,
132 offset: ArrayView1<'_, f64>,
133 spec: &TermCollectionSpec,
134 priors: &[(String, f64, f64)],
135 family: LikelihoodSpec,
136 options: &FitOptions,
137) -> Result<FittedTermCollection, EstimationError> {
138 let design = build_term_collection_design_with_policy(data, spec, &options.resource_policy)?;
139 let mut fit_opts = adaptive_fit_options_base(options, &design);
140 fit_opts.rho_prior = realize_keyed_penalty_block_gamma_priors(&design, priors)
141 .map_err(EstimationError::BasisError)?;
142 let fitted = FittedTermCollection {
143 fit: fit_gamwith_heuristic_lambdas(
144 design.design.clone(),
145 y,
146 weights,
147 offset,
148 &design.penalties,
149 None,
150 family.clone(),
151 &fit_opts,
152 )?,
153 design,
154 adaptive_diagnostics: None,
155 };
156 enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
157 Ok(fitted)
158}
159
160pub fn fit_term_collection_with_coefficient_groups_and_penalty_block_gamma_priors(
161 data: ArrayView2<'_, f64>,
162 y: ArrayView1<'_, f64>,
163 weights: ArrayView1<'_, f64>,
164 offset: ArrayView1<'_, f64>,
165 spec: &TermCollectionSpec,
166 groups: &[CoefficientGroupSpec],
167 priors: &[(String, f64, f64)],
168 family: LikelihoodSpec,
169 options: &FitOptions,
170) -> Result<FittedTermCollection, EstimationError> {
171 if groups.is_empty() {
172 return fit_term_collection_with_penalty_block_gamma_priors(
173 data, y, weights, offset, spec, priors, family, options,
174 );
175 }
176 if priors.is_empty() {
177 return fit_term_collection_with_coefficient_groups(
178 data, y, weights, offset, spec, groups, family, options,
179 );
180 }
181
182 let design = build_term_collection_design_with_policy(data, spec, &options.resource_policy)?;
186 let base_fit_opts = adaptive_fit_options_base(options, &design);
187 let base_rho_prior = realize_keyed_penalty_block_gamma_priors(&design, priors)
188 .map_err(EstimationError::BasisError)?;
189 let realized = design
190 .realize_coefficient_groups(groups, &base_rho_prior)
191 .map_err(EstimationError::BasisError)?;
192 let mut grouped_options = base_fit_opts.clone();
193 grouped_options.rho_prior = realized.rho_prior;
194 let fitted = FittedTermCollection {
195 fit: gam_solve::estimate::fit_gam_with_penalty_specs(
196 design.design.clone(),
197 y,
198 weights,
199 offset,
200 realized.penalty_specs,
201 realized.nullspace_dims,
202 family.clone(),
203 &grouped_options,
204 )?,
205 design,
206 adaptive_diagnostics: None,
207 };
208 enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
209 Ok(fitted)
210}
211
212fn fit_term_collection_forspecwith_heuristic_lambdas(
213 data: ArrayView2<'_, f64>,
214 y: ArrayView1<'_, f64>,
215 weights: ArrayView1<'_, f64>,
216 offset: ArrayView1<'_, f64>,
217 spec: &TermCollectionSpec,
218 heuristic_lambdas: Option<&[f64]>,
219 family: LikelihoodSpec,
220 options: &FitOptions,
221) -> Result<FittedTermCollection, EstimationError> {
222 let adaptive_opts = options.adaptive_regularization.clone().unwrap_or_default();
223 let resolved_spec;
224 let design_spec = if adaptive_opts.enabled {
225 resolved_spec = ensure_matern_adaptive_center_resolution(spec, data.nrows());
226 &resolved_spec
227 } else {
228 spec
229 };
230 let base_design =
231 build_term_collection_design_with_policy(data, design_spec, &options.resource_policy)?;
232 fit_term_collection_on_realized_design(
233 y,
234 weights,
235 offset,
236 design_spec,
237 &base_design,
238 heuristic_lambdas,
239 family,
240 options,
241 )
242}
243
244fn ensure_matern_adaptive_center_resolution(
245 spec: &TermCollectionSpec,
246 n_rows: usize,
247) -> TermCollectionSpec {
248 let mut out = spec.clone();
249 for term in &mut out.smooth_terms {
250 let gam_terms::smooth::SmoothBasisSpec::Matern {
251 feature_cols,
252 spec: matern,
253 ..
254 } = &mut term.basis
255 else {
256 continue;
257 };
258 if let gam_terms::basis::CenterStrategy::FarthestPoint { num_centers } =
259 &mut matern.center_strategy
260 {
261 let min_centers = (4 * feature_cols.len()).min(n_rows).max(*num_centers);
274 *num_centers = min_centers;
275 }
276 }
277 out
278}
279
280fn has_bounded_linear_terms(spec: &TermCollectionSpec) -> bool {
281 spec.linear_terms.iter().any(|term| {
282 matches!(
283 term.coefficient_geometry,
284 LinearCoefficientGeometry::Bounded { .. }
285 )
286 })
287}
288
289fn fit_term_collection_on_realized_design(
290 y: ArrayView1<'_, f64>,
291 weights: ArrayView1<'_, f64>,
292 offset: ArrayView1<'_, f64>,
293 spec: &TermCollectionSpec,
294 design: &TermCollectionDesign,
295 heuristic_lambdas: Option<&[f64]>,
296 family: LikelihoodSpec,
297 options: &FitOptions,
298) -> Result<FittedTermCollection, EstimationError> {
299 if has_bounded_linear_terms(spec) {
300 return fit_bounded_term_collection_with_design(
301 y,
302 weights,
303 offset,
304 spec,
305 design,
306 heuristic_lambdas,
307 family,
308 options,
309 );
310 }
311 let mut base_fit_opts = adaptive_fit_options_base(options, design);
312 base_fit_opts.rho_prior = relax_smoothing_rho_prior(options, design);
319 let fitted = FittedTermCollection {
320 fit: fit_gamwith_heuristic_lambdas(
321 design.design.clone(),
322 y,
323 weights,
324 offset,
325 &design.penalties,
326 heuristic_lambdas,
327 family.clone(),
328 &base_fit_opts,
329 )?,
330 design: design.clone(),
331 adaptive_diagnostics: None,
332 };
333 enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
334
335 let adaptive_opts = options.adaptive_regularization.clone().unwrap_or_default();
336 if !adaptive_opts.enabled {
337 return Ok(fitted);
338 }
339 let runtime_caches = extract_spatial_operator_runtime_caches(spec, &fitted.design)?;
340 if runtime_caches.is_empty() {
341 return Ok(fitted);
342 }
343 fit_term_collectionwith_exact_spatial_adaptive_regularization(
350 fitted,
351 y,
352 weights,
353 offset,
354 family,
355 options,
356 &runtime_caches,
357 )
358}
359
360#[derive(Clone)]
361struct SpatialOperatorRuntimeCache {
362 termname: String,
363 feature_cols: Vec<usize>,
364 coeff_global_range: Range<usize>,
365 mass_penalty_global_idx: usize,
366 tension_penalty_global_idx: usize,
367 stiffness_penalty_global_idx: usize,
368 d0: Array2<f64>,
369 d1: Array2<f64>,
370 d2: Array2<f64>,
371 collocation_points: Array2<f64>,
372 dimension: usize,
373}
374
375#[derive(Clone)]
376struct SpatialAdaptiveWeights {
377 inv_magweight: Array1<f64>,
378 invgradweight: Array1<f64>,
379 inv_lapweight: Array1<f64>,
380}
381
382#[derive(Clone)]
383struct CharbonnierScalarBlockState {
384 signal: Array1<f64>,
385 radius: Array1<f64>,
386 epsilon: f64,
387}
388
389impl CharbonnierScalarBlockState {
390 fn from_signal(signal: Array1<f64>, epsilon: f64) -> Self {
391 let eps = epsilon.max(1e-12);
392 let radius = signal.mapv(|t| (t * t + eps * eps).sqrt());
393 Self {
394 signal,
395 radius,
396 epsilon: eps,
397 }
398 }
399
400 fn absolute_signal(&self) -> Array1<f64> {
401 self.signal.mapv(f64::abs)
402 }
403
404 fn penalty_value(&self) -> f64 {
405 self.radius.iter().map(|r| r - self.epsilon).sum::<f64>()
406 }
407
408 fn betagradient_coeff(&self) -> Array1<f64> {
409 Array1::from_iter(
410 self.signal
411 .iter()
412 .zip(self.radius.iter())
413 .map(|(t, r)| t / r),
414 )
415 }
416
417 fn betahessian_diag(&self) -> Array1<f64> {
418 let eps2 = self.epsilon * self.epsilon;
419 self.radius.mapv(|r| eps2 / r.powi(3))
420 }
421
422 fn log_epsilon_gradient_terms(&self) -> Array1<f64> {
423 let epsilon = self.epsilon;
424 let eps2 = epsilon * epsilon;
425 self.radius.mapv(|r| eps2 / r - epsilon)
426 }
427
428 fn log_epsilon_betagradient_coeff(&self) -> Array1<f64> {
429 let eps2 = self.epsilon * self.epsilon;
430 Array1::from_iter(
431 self.signal
432 .iter()
433 .zip(self.radius.iter())
434 .map(|(t, r)| -eps2 * t / r.powi(3)),
435 )
436 }
437
438 fn log_epsilon_hessian_terms(&self) -> Array1<f64> {
439 let epsilon = self.epsilon;
440 let eps2 = epsilon * epsilon;
441 let eps4 = eps2 * eps2;
442 self.radius
443 .mapv(|r| 2.0 * eps2 / r - eps4 / r.powi(3) - epsilon)
444 }
445
446 fn surrogateweights_posterior_snr(
447 &self,
448 variance: &Array1<f64>,
449 weight_floor: f64,
450 weight_ceiling: f64,
451 ) -> (Array1<f64>, Array1<f64>) {
452 let eps2 = self.epsilon * self.epsilon;
510 let weight = Array1::from_iter(self.signal.iter().zip(variance.iter()).map(|(&t, &v)| {
511 let credible2 = (t * t - v.max(0.0)).max(0.0);
512 let r = (credible2 + eps2).sqrt();
513 (1.0 / r).clamp(weight_floor, weight_ceiling)
514 }));
515 let invweight = weight.mapv(|u| 1.0 / u);
516 (weight, invweight)
517 }
518
519 fn directionalhessian_diag(&self, direction_signal: &Array1<f64>) -> Array1<f64> {
520 let eps2 = self.epsilon * self.epsilon;
535 Array1::from_iter(
536 self.signal
537 .iter()
538 .zip(direction_signal.iter())
539 .zip(self.radius.iter())
540 .map(|((t, q), r)| -3.0 * eps2 * t * q / r.powi(5)),
541 )
542 }
543
544 fn second_directionalhessian_diag(
551 &self,
552 direction1_signal: &Array1<f64>,
553 direction2_signal: &Array1<f64>,
554 ) -> Array1<f64> {
555 let eps2 = self.epsilon * self.epsilon;
556 Array1::from_iter(
557 self.signal
558 .iter()
559 .zip(direction1_signal.iter())
560 .zip(direction2_signal.iter())
561 .zip(self.radius.iter())
562 .map(|(((t, q1), q2), r)| {
563 let r2 = r * r;
564 let psi4 = -3.0 * eps2 / r.powi(5) + 15.0 * eps2 * t * t / (r.powi(5) * r2);
565 psi4 * q1 * q2
566 }),
567 )
568 }
569
570 fn log_epsilon_betahessian_diag(&self) -> Array1<f64> {
571 let eps2 = self.epsilon * self.epsilon;
572 let eps4 = eps2 * eps2;
573 Array1::from_iter(
574 self.signal
575 .iter()
576 .zip(self.radius.iter())
577 .map(|(_, r)| 2.0 * eps2 / r.powi(3) - 3.0 * eps4 / r.powi(5)),
578 )
579 }
580
581 fn log_epsilon_beta_mixed_second_coeff(&self) -> Array1<f64> {
582 let eps2 = self.epsilon * self.epsilon;
583 Array1::from_iter(
584 self.signal
585 .iter()
586 .zip(self.radius.iter())
587 .map(|(t, r)| eps2 * t * (eps2 - 2.0 * t * t) / r.powi(5)),
588 )
589 }
590
591 fn log_epsilon_betahessian_second_diag(&self) -> Array1<f64> {
592 let eps2 = self.epsilon * self.epsilon;
593 let eps4 = eps2 * eps2;
594 let eps6 = eps4 * eps2;
595 Array1::from_iter(
596 self.radius.iter().map(|r| {
597 4.0 * eps2 / r.powi(3) - 18.0 * eps4 / r.powi(5) + 15.0 * eps6 / r.powi(7)
598 }),
599 )
600 }
601
602 fn log_epsilon_betahessian_directional_diag(
603 &self,
604 direction_signal: &Array1<f64>,
605 ) -> Array1<f64> {
606 let eps2 = self.epsilon * self.epsilon;
607 let eps4 = eps2 * eps2;
608 Array1::from_iter(
609 self.signal
610 .iter()
611 .zip(direction_signal.iter())
612 .zip(self.radius.iter())
613 .map(|((t, q), r)| (-6.0 * eps2 * t / r.powi(5) + 15.0 * eps4 * t / r.powi(7)) * q),
614 )
615 }
616}
617
618#[derive(Clone)]
619struct CharbonnierGroupedBlockState {
620 norm: Array1<f64>,
621 radius: Array1<f64>,
622 signal_blocks: Array2<f64>,
623 epsilon: f64,
624}
625
626impl CharbonnierGroupedBlockState {
627 fn from_signal_blocks(signal_blocks: Array2<f64>, epsilon: f64) -> Self {
628 let eps = epsilon.max(1e-12);
629 let norm = Array1::from_iter(
630 signal_blocks
631 .rows()
632 .into_iter()
633 .map(|row| row.iter().map(|v| v * v).sum::<f64>().sqrt()),
634 );
635 let radius = norm.mapv(|g| (g * g + eps * eps).sqrt());
636 Self {
637 norm,
638 radius,
639 signal_blocks,
640 epsilon: eps,
641 }
642 }
643
644 fn penalty_value(&self) -> f64 {
645 self.radius.iter().map(|r| r - self.epsilon).sum::<f64>()
646 }
647
648 fn norm_signal(&self) -> Array1<f64> {
649 self.norm.clone()
650 }
651
652 fn betagradient_blocks(&self) -> Array2<f64> {
653 let mut out = self.signal_blocks.clone();
654 for (k, mut row) in out.rows_mut().into_iter().enumerate() {
655 let scale = 1.0 / self.radius[k];
656 row.mapv_inplace(|v| v * scale);
657 }
658 out
659 }
660
661 fn betahessian_blocks(&self) -> Vec<Array2<f64>> {
662 let mut out = Vec::with_capacity(self.signal_blocks.nrows());
663 for (k, row) in self.signal_blocks.rows().into_iter().enumerate() {
664 let dim = row.len();
665 let mut block = Array2::<f64>::eye(dim);
666 block.mapv_inplace(|v| v / self.radius[k]);
667 for i in 0..dim {
668 for j in 0..dim {
669 block[[i, j]] -= row[i] * row[j] / self.radius[k].powi(3);
670 }
671 }
672 out.push(block);
673 }
674 out
675 }
676
677 fn log_epsilon_gradient_terms(&self) -> Array1<f64> {
678 let epsilon = self.epsilon;
679 let eps2 = epsilon * epsilon;
680 self.radius.mapv(|r| eps2 / r - epsilon)
681 }
682
683 fn log_epsilon_betagradient_blocks(&self) -> Array2<f64> {
684 let mut out = self.signal_blocks.clone();
685 let eps2 = self.epsilon * self.epsilon;
686 for (k, mut row) in out.rows_mut().into_iter().enumerate() {
687 let scale = -eps2 / self.radius[k].powi(3);
688 row.mapv_inplace(|v| v * scale);
689 }
690 out
691 }
692
693 fn log_epsilon_hessian_terms(&self) -> Array1<f64> {
694 let epsilon = self.epsilon;
695 let eps2 = epsilon * epsilon;
696 let eps4 = eps2 * eps2;
697 self.radius
698 .mapv(|r| 2.0 * eps2 / r - eps4 / r.powi(3) - epsilon)
699 }
700
701 fn surrogateweights_posterior_snr(
702 &self,
703 variance: &Array1<f64>,
704 weight_floor: f64,
705 weight_ceiling: f64,
706 ) -> (Array1<f64>, Array1<f64>) {
707 let eps2 = self.epsilon * self.epsilon;
749 let weight = Array1::from_iter(self.norm.iter().zip(variance.iter()).map(|(&g, &v)| {
750 let credible2 = (g * g - v.max(0.0)).max(0.0);
751 let r = (credible2 + eps2).sqrt();
752 (1.0 / r).clamp(weight_floor, weight_ceiling)
753 }));
754 let invweight = weight.mapv(|u| 1.0 / u);
755 (weight, invweight)
756 }
757
758 fn directionalhessian_blocks(&self, direction_blocks: &Array2<f64>) -> Vec<Array2<f64>> {
759 let mut out = Vec::with_capacity(self.signal_blocks.nrows());
784 for (k, (v, q)) in self
785 .signal_blocks
786 .rows()
787 .into_iter()
788 .zip(direction_blocks.rows().into_iter())
789 .enumerate()
790 {
791 let dim = v.len();
792 let dot = v.iter().zip(q.iter()).map(|(a, b)| a * b).sum::<f64>();
793 let r3 = self.radius[k].powi(3);
794 let r5 = self.radius[k].powi(5);
795 let mut block = Array2::<f64>::eye(dim);
796 block.mapv_inplace(|x| -dot * x / r3);
797 for i in 0..dim {
798 for j in 0..dim {
799 block[[i, j]] -= (q[i] * v[j] + v[i] * q[j]) / r3;
800 block[[i, j]] += 3.0 * dot * v[i] * v[j] / r5;
801 }
802 }
803 out.push(block);
804 }
805 out
806 }
807
808 fn second_directionalhessian_blocks(
825 &self,
826 direction1_blocks: &Array2<f64>,
827 direction2_blocks: &Array2<f64>,
828 ) -> Vec<Array2<f64>> {
829 let mut out = Vec::with_capacity(self.signal_blocks.nrows());
830 for ((k, v), (a, b)) in self.signal_blocks.rows().into_iter().enumerate().zip(
831 direction1_blocks
832 .rows()
833 .into_iter()
834 .zip(direction2_blocks.rows().into_iter()),
835 ) {
836 let dim = v.len();
837 let dot = |x: ndarray::ArrayView1<'_, f64>, y: ndarray::ArrayView1<'_, f64>| {
838 x.iter().zip(y.iter()).map(|(p, q)| p * q).sum::<f64>()
839 };
840 let sa = dot(v, a);
841 let sb = dot(v, b);
842 let ab = dot(a, b);
843 let r = self.radius[k];
844 let r3 = r.powi(3);
845 let r5 = r.powi(5);
846 let r7 = r5 * r * r;
847 let diag = -ab / r3 + 3.0 * sa * sb / r5;
848 let mut block = Array2::<f64>::eye(dim);
849 block.mapv_inplace(|x| diag * x);
850 for i in 0..dim {
851 for j in 0..dim {
852 block[[i, j]] -= (a[i] * b[j] + b[i] * a[j]) / r3;
853 block[[i, j]] += 3.0 * sb * (a[i] * v[j] + v[i] * a[j]) / r5;
854 block[[i, j]] += 3.0 * ab * v[i] * v[j] / r5;
855 block[[i, j]] += 3.0 * sa * (b[i] * v[j] + v[i] * b[j]) / r5;
856 block[[i, j]] -= 15.0 * sa * sb * v[i] * v[j] / r7;
857 }
858 }
859 out.push(block);
860 }
861 out
862 }
863
864 fn log_epsilon_betahessian_blocks(&self) -> Vec<Array2<f64>> {
865 let mut out = Vec::with_capacity(self.signal_blocks.nrows());
866 for (k, row) in self.signal_blocks.rows().into_iter().enumerate() {
867 let dim = row.len();
868 let r3 = self.radius[k].powi(3);
869 let r5 = self.radius[k].powi(5);
870 let mut block = Array2::<f64>::eye(dim);
871 let eps2 = self.epsilon * self.epsilon;
872 block.mapv_inplace(|v| -eps2 * v / r3);
873 for i in 0..dim {
874 for j in 0..dim {
875 block[[i, j]] += 3.0 * eps2 * row[i] * row[j] / r5;
876 }
877 }
878 out.push(block);
879 }
880 out
881 }
882
883 fn log_epsilon_beta_mixed_second_blocks(&self) -> Array2<f64> {
884 let mut out = self.signal_blocks.clone();
885 let eps2 = self.epsilon * self.epsilon;
886 for (k, mut row) in out.rows_mut().into_iter().enumerate() {
887 let norm2 = self.norm[k] * self.norm[k];
888 let scale = eps2 * (eps2 - 2.0 * norm2) / self.radius[k].powi(5);
889 row.mapv_inplace(|v| v * scale);
890 }
891 out
892 }
893
894 fn log_epsilon_betahessian_second_blocks(&self) -> Vec<Array2<f64>> {
895 let mut out = Vec::with_capacity(self.signal_blocks.nrows());
896 let eps2 = self.epsilon * self.epsilon;
897 for (k, row) in self.signal_blocks.rows().into_iter().enumerate() {
898 let dim = row.len();
899 let norm2 = self.norm[k] * self.norm[k];
900 let r5 = self.radius[k].powi(5);
901 let r7 = self.radius[k].powi(7);
902 let mut block = Array2::<f64>::eye(dim);
903 block.mapv_inplace(|v| eps2 * (eps2 - 2.0 * norm2) * v / r5);
904 for i in 0..dim {
905 for j in 0..dim {
906 block[[i, j]] += 3.0 * eps2 * (2.0 * norm2 - 3.0 * eps2) * row[i] * row[j] / r7;
907 }
908 }
909 out.push(block);
910 }
911 out
912 }
913
914 fn log_epsilon_betahessian_directional_blocks(
915 &self,
916 direction_blocks: &Array2<f64>,
917 ) -> Vec<Array2<f64>> {
918 let mut out = Vec::with_capacity(self.signal_blocks.nrows());
919 let eps2 = self.epsilon * self.epsilon;
920 for (k, (v, q)) in self
921 .signal_blocks
922 .rows()
923 .into_iter()
924 .zip(direction_blocks.rows().into_iter())
925 .enumerate()
926 {
927 let dim = v.len();
928 let dot = v.iter().zip(q.iter()).map(|(a, b)| a * b).sum::<f64>();
929 let r5 = self.radius[k].powi(5);
930 let r7 = self.radius[k].powi(7);
931 let mut block = Array2::<f64>::eye(dim);
932 block.mapv_inplace(|x| 3.0 * eps2 * dot * x / r5);
933 for i in 0..dim {
934 for j in 0..dim {
935 block[[i, j]] += 3.0 * eps2 * (q[i] * v[j] + v[i] * q[j]) / r5;
936 block[[i, j]] -= 15.0 * eps2 * dot * v[i] * v[j] / r7;
937 }
938 }
939 out.push(block);
940 }
941 out
942 }
943}
944
945fn scalar_operatorgradient(operator: &Array2<f64>, coeff: &Array1<f64>) -> Array1<f64> {
946 operator.t().dot(coeff)
947}
948
949fn scalar_operatorhessian(operator: &Array2<f64>, diag: &Array1<f64>) -> Array2<f64> {
950 let mut weighted = operator.clone();
951 for (k, &w) in diag.iter().enumerate() {
952 weighted.row_mut(k).mapv_inplace(|v| v * w);
953 }
954 let gram = operator.t().dot(&weighted);
955 (&gram + &gram.t().to_owned()) * 0.5
956}
957
958fn grouped_operatorgradient(
959 d1: &Array2<f64>,
960 dimension: usize,
961 blocks: &Array2<f64>,
962) -> Result<Array1<f64>, EstimationError> {
963 if blocks.ncols() != dimension {
964 crate::bail_invalid_estim!(
965 "grouped gradient block dimension mismatch: got {}, expected {dimension}",
966 blocks.ncols()
967 );
968 }
969 if d1.nrows() != blocks.nrows() * dimension {
970 crate::bail_invalid_estim!(
971 "grouped gradient row mismatch: D1 has {} rows, blocks imply {}",
972 d1.nrows(),
973 blocks.nrows() * dimension
974 );
975 }
976 let mut out = Array1::<f64>::zeros(d1.ncols());
977 for k in 0..blocks.nrows() {
978 let gk = d1
979 .slice(s![k * dimension..(k + 1) * dimension, ..])
980 .to_owned();
981 out += &gk.t().dot(&blocks.row(k));
982 }
983 Ok(out)
984}
985
986fn grouped_operatorhessian(
987 d1: &Array2<f64>,
988 dimension: usize,
989 blocks: &[Array2<f64>],
990) -> Result<Array2<f64>, EstimationError> {
991 if d1.nrows() != blocks.len() * dimension {
992 crate::bail_invalid_estim!(
993 "grouped Hessian row mismatch: D1 has {} rows, blocks imply {}",
994 d1.nrows(),
995 blocks.len() * dimension
996 );
997 }
998 let p = d1.ncols();
999 let mut out = Array2::<f64>::zeros((p, p));
1000 for (k, block) in blocks.iter().enumerate() {
1001 if block.nrows() != dimension || block.ncols() != dimension {
1002 crate::bail_invalid_estim!(
1003 "grouped Hessian block {k} has shape {}x{}, expected {}x{}",
1004 block.nrows(),
1005 block.ncols(),
1006 dimension,
1007 dimension
1008 );
1009 }
1010 let gk = d1
1011 .slice(s![k * dimension..(k + 1) * dimension, ..])
1012 .to_owned();
1013 out += &gk.t().dot(&block.dot(&gk));
1014 }
1015 Ok((&out + &out.t().to_owned()) * 0.5)
1016}
1017
1018#[derive(Clone)]
1019struct SpatialPenaltyExactState {
1020 magnitude: CharbonnierScalarBlockState,
1021 gradient: CharbonnierGroupedBlockState,
1022 curvature: CharbonnierGroupedBlockState,
1023}
1024
1025fn collocationgradient_blocks(
1026 gradrows: &Array1<f64>,
1027 dimension: usize,
1028) -> Result<Array2<f64>, EstimationError> {
1029 if dimension == 0 || !gradrows.len().is_multiple_of(dimension) {
1030 crate::bail_invalid_estim!(
1031 "invalid collocation gradient layout: rows={}, dimension={dimension}",
1032 gradrows.len()
1033 );
1034 }
1035 let p = gradrows.len() / dimension;
1036 let mut out = Array2::<f64>::zeros((p, dimension));
1037 for k in 0..p {
1038 for axis in 0..dimension {
1039 out[[k, axis]] = gradrows[k * dimension + axis];
1040 }
1041 }
1042 Ok(out)
1043}
1044
1045fn collocationhessian_blocks(
1046 hessianrows: &Array1<f64>,
1047 dimension: usize,
1048) -> Result<Array2<f64>, EstimationError> {
1049 let block_dim = dimension.checked_mul(dimension).ok_or_else(|| {
1050 EstimationError::InvalidInput("invalid collocation Hessian dimension overflow".to_string())
1051 })?;
1052 if block_dim == 0 || !hessianrows.len().is_multiple_of(block_dim) {
1053 crate::bail_invalid_estim!(
1054 "invalid collocation Hessian layout: rows={}, dimension={dimension}",
1055 hessianrows.len()
1056 );
1057 }
1058 let p = hessianrows.len() / block_dim;
1059 let mut out = Array2::<f64>::zeros((p, block_dim));
1060 for k in 0..p {
1061 for idx in 0..block_dim {
1062 out[[k, idx]] = hessianrows[k * block_dim + idx];
1063 }
1064 }
1065 Ok(out)
1066}
1067
1068impl SpatialPenaltyExactState {
1069 fn from_beta_local(
1070 beta_local: ArrayView1<'_, f64>,
1071 cache: &SpatialOperatorRuntimeCache,
1072 epsilons: [f64; 3],
1073 ) -> Result<Self, EstimationError> {
1074 let gradientrows = cache.d1.dot(&beta_local);
1104 let hessianrows = cache.d2.dot(&beta_local);
1105 Ok(Self {
1106 magnitude: CharbonnierScalarBlockState::from_signal(
1107 cache.d0.dot(&beta_local),
1108 epsilons[0],
1109 ),
1110 gradient: CharbonnierGroupedBlockState::from_signal_blocks(
1111 collocationgradient_blocks(&gradientrows, cache.dimension)?,
1112 epsilons[1],
1113 ),
1114 curvature: CharbonnierGroupedBlockState::from_signal_blocks(
1115 collocationhessian_blocks(&hessianrows, cache.dimension)?,
1116 epsilons[2],
1117 ),
1118 })
1119 }
1120
1121 fn absolute_collocation_magnitudes(&self) -> (Array1<f64>, Array1<f64>, Array1<f64>) {
1122 (
1123 self.magnitude.absolute_signal(),
1124 self.gradient.norm_signal(),
1125 self.curvature.norm_signal(),
1126 )
1127 }
1128}
1129
1130fn robust_epsilon_from_samples(values: &[f64], min_epsilon_cfg: f64) -> f64 {
1131 if values.is_empty() {
1132 return min_epsilon_cfg.max(1e-12);
1133 }
1134 let mut clean = values
1135 .iter()
1136 .copied()
1137 .filter(|v| v.is_finite() && *v >= 0.0)
1138 .collect::<Vec<_>>();
1139 if clean.is_empty() {
1140 return min_epsilon_cfg.max(1e-12);
1141 }
1142 clean.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
1143
1144 let n = clean.len();
1145 let median = quantile_from_sorted(&clean, 0.5);
1146 let q75 = quantile_from_sorted(&clean, 0.75);
1147 let q95 = quantile_from_sorted(&clean, 0.95);
1148
1149 let mut abs_dev = clean
1150 .iter()
1151 .map(|v| (v - median).abs())
1152 .filter(|v| v.is_finite())
1153 .collect::<Vec<_>>();
1154 abs_dev.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
1155 let mad = 1.4826 * quantile_from_sorted(&abs_dev, 0.5);
1156
1157 let mut scale = median.max(mad).max(q75);
1167
1168 let delta = (f64::EPSILON.sqrt() * q95.max(1.0))
1170 .max(min_epsilon_cfg)
1171 .max(1e-12);
1172 let s_min = min_epsilon_cfg.max(1e-12);
1173
1174 if scale <= delta {
1176 let rms = (clean.iter().map(|v| v * v).sum::<f64>() / n as f64).sqrt();
1177 scale = q95.max(rms);
1178 }
1179 if scale <= delta {
1180 scale = s_min;
1181 }
1182
1183 let kappa = 1.0_f64;
1186 (kappa * scale).max(s_min)
1187}
1188
1189fn extract_spatial_operator_runtime_caches(
1190 spec: &TermCollectionSpec,
1191 design: &TermCollectionDesign,
1192) -> Result<Vec<SpatialOperatorRuntimeCache>, EstimationError> {
1193 let smooth_start = design
1194 .design
1195 .ncols()
1196 .saturating_sub(design.smooth.total_smooth_cols());
1197 let mut out = Vec::<SpatialOperatorRuntimeCache>::new();
1198 for (term_idx, (termspec, term_fit)) in spec
1199 .smooth_terms
1200 .iter()
1201 .zip(design.smooth.terms.iter())
1202 .enumerate()
1203 {
1204 let Some(global_base_idx) = smooth_term_penalty_index(spec, design, term_idx) else {
1205 continue;
1206 };
1207 let mut active_local_idx = 0usize;
1208 let mut mass_local_idx = None;
1209 let mut tension_local_idx = None;
1210 let mut stiffness_local_idx = None;
1211 let mut mass_norm = None;
1212 let mut tension_norm = None;
1213 let mut stiffness_norm = None;
1214 for info in &term_fit.penaltyinfo_local {
1215 if !info.active {
1216 continue;
1217 }
1218 match info.source {
1219 PenaltySource::OperatorMass => {
1220 mass_local_idx = Some(active_local_idx);
1221 mass_norm = Some(info.normalization_scale);
1222 }
1223 PenaltySource::OperatorTension => {
1224 tension_local_idx = Some(active_local_idx);
1225 tension_norm = Some(info.normalization_scale);
1226 }
1227 PenaltySource::OperatorStiffness => {
1228 stiffness_local_idx = Some(active_local_idx);
1229 stiffness_norm = Some(info.normalization_scale);
1230 }
1231 _ => {}
1232 }
1233 active_local_idx += 1;
1234 }
1235 let (
1248 Some(mass_local),
1249 Some(tension_local),
1250 Some(stiffness_local),
1251 Some(mass_scale),
1252 Some(tension_scale),
1253 Some(stiffness_scale),
1254 ) = (
1255 mass_local_idx,
1256 tension_local_idx,
1257 stiffness_local_idx,
1258 mass_norm,
1259 tension_norm,
1260 stiffness_norm,
1261 )
1262 else {
1263 continue;
1264 };
1265 let mass_global_idx = global_base_idx + mass_local;
1266 let tension_global_idx = global_base_idx + tension_local;
1267 let stiffness_global_idx = global_base_idx + stiffness_local;
1268
1269 let (feature_cols, mut d0, mut d1, mut d2, collocation_points, dim, center_mass_rows) =
1270 match (&termspec.basis, &term_fit.metadata) {
1271 (
1272 SmoothBasisSpec::Matern { feature_cols, .. },
1273 BasisMetadata::Matern {
1274 centers,
1275 length_scale,
1276 nu,
1277 include_intercept,
1278 identifiability_transform,
1279 aniso_log_scales,
1280 input_scales,
1281 ..
1282 },
1283 ) => {
1284 let collocation_length_scale = match input_scales.as_deref() {
1290 Some(scales) => {
1291 compensate_length_scale_for_standardization(*length_scale, scales)
1292 }
1293 None => *length_scale,
1294 };
1295 let ops = build_matern_collocation_operator_matrices(
1296 centers.view(),
1297 None,
1298 collocation_length_scale,
1299 *nu,
1300 *include_intercept,
1301 identifiability_transform.as_ref().map(|z| z.view()),
1302 aniso_log_scales.as_deref(),
1303 )?;
1304 (
1305 feature_cols.clone(),
1306 ops.d0,
1307 ops.d1,
1308 ops.d2,
1309 ops.collocation_points,
1310 centers.ncols(),
1311 false,
1312 )
1313 }
1314 (
1315 SmoothBasisSpec::Duchon { feature_cols, .. },
1316 BasisMetadata::Duchon {
1317 centers,
1318 length_scale,
1319 power,
1320 nullspace_order,
1321 identifiability_transform,
1322 input_scales,
1323 aniso_log_scales,
1324 operator_collocation_points: Some(collocation_points),
1325 ..
1326 },
1327 ) => {
1328 let collocation_length_scale = match (length_scale, input_scales.as_deref()) {
1329 (Some(ls), Some(scales)) => {
1330 Some(compensate_length_scale_for_standardization(*ls, scales))
1331 }
1332 (Some(ls), None) => Some(*ls),
1333 (None, _) => None,
1334 };
1335 let ops =
1336 gam_terms::basis::build_duchon_collocation_operator_matriceswithworkspace(
1337 centers.view(),
1338 collocation_points.view(),
1339 None,
1340 collocation_length_scale,
1341 *power,
1342 *nullspace_order,
1343 aniso_log_scales.as_deref(),
1344 identifiability_transform.as_ref().map(|z| z.view()),
1345 2,
1346 None,
1347 &mut BasisWorkspace::default(),
1348 )?;
1349 (
1350 feature_cols.clone(),
1351 ops.d0,
1352 ops.d1,
1353 ops.d2,
1354 ops.collocation_points,
1355 centers.ncols(),
1356 true,
1357 )
1358 }
1359 _ => continue,
1360 };
1361 if center_mass_rows && d0.nrows() > 0 && d0.ncols() > 0 {
1362 let means = d0.sum_axis(Axis(0)).mapv(|v| v / d0.nrows() as f64);
1363 for mut row in d0.rows_mut() {
1364 row -= &means;
1365 }
1366 }
1367
1368 let mass_scale = mass_scale.max(1e-12).sqrt();
1386 let tension_scale = tension_scale.max(1e-12).sqrt();
1387 let stiffness_scale = stiffness_scale.max(1e-12).sqrt();
1388 d0.mapv_inplace(|v| v / mass_scale);
1389 d1.mapv_inplace(|v| v / tension_scale);
1390 d2.mapv_inplace(|v| v / stiffness_scale);
1391
1392 let coeff_global_range =
1393 (smooth_start + term_fit.coeff_range.start)..(smooth_start + term_fit.coeff_range.end);
1394 if d0.ncols() != coeff_global_range.len()
1395 || d1.ncols() != coeff_global_range.len()
1396 || d2.ncols() != coeff_global_range.len()
1397 {
1398 crate::bail_invalid_estim!(
1399 "spatial operator dimension mismatch for term '{}': D0 cols={}, D1 cols={}, D2 cols={}, coeffs={}",
1400 term_fit.name,
1401 d0.ncols(),
1402 d1.ncols(),
1403 d2.ncols(),
1404 coeff_global_range.len()
1405 );
1406 }
1407 out.push(SpatialOperatorRuntimeCache {
1408 termname: term_fit.name.clone(),
1409 feature_cols,
1410 coeff_global_range,
1411 mass_penalty_global_idx: mass_global_idx,
1412 tension_penalty_global_idx: tension_global_idx,
1413 stiffness_penalty_global_idx: stiffness_global_idx,
1414 d0,
1415 d1,
1416 d2,
1417 collocation_points,
1418 dimension: dim,
1419 });
1420 }
1421 Ok(out)
1422}
1423
1424fn scalar_operator_response_variance(
1436 operator: &Array2<f64>,
1437 cov_local: &Array2<f64>,
1438) -> Array1<f64> {
1439 Array1::from_iter(operator.rows().into_iter().map(|row| {
1440 let s = cov_local.dot(&row);
1441 row.dot(&s).max(0.0)
1442 }))
1443}
1444
1445fn grouped_operator_response_variance(
1456 operator: &Array2<f64>,
1457 block_dim: usize,
1458 cov_local: &Array2<f64>,
1459) -> Result<Array1<f64>, EstimationError> {
1460 if block_dim == 0 || !operator.nrows().is_multiple_of(block_dim) {
1461 crate::bail_invalid_estim!(
1462 "grouped variance row layout invalid: rows={}, block_dim={block_dim}",
1463 operator.nrows()
1464 );
1465 }
1466 let p = operator.nrows() / block_dim;
1467 let mut out = Array1::<f64>::zeros(p);
1468 for k in 0..p {
1469 let mut acc = 0.0;
1470 for axis in 0..block_dim {
1471 let row = operator.row(k * block_dim + axis);
1472 let s = cov_local.dot(&row);
1473 acc += row.dot(&s);
1474 }
1475 out[k] = acc.max(0.0);
1476 }
1477 Ok(out)
1478}
1479
1480fn compute_spatial_adaptiveweights_for_beta(
1481 beta: &Array1<f64>,
1482 caches: &[SpatialOperatorRuntimeCache],
1483 epsilon_0: f64,
1484 epsilon_g: f64,
1485 epsilon_c: f64,
1486 weight_floor: f64,
1487 weight_ceiling: f64,
1488 beta_covariance: Option<&Array2<f64>>,
1489) -> Result<Vec<SpatialAdaptiveWeights>, EstimationError> {
1490 caches
1522 .iter()
1523 .map(|cache| {
1524 let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
1525 let exact = SpatialPenaltyExactState::from_beta_local(
1526 beta_local,
1527 cache,
1528 [epsilon_0, epsilon_g, epsilon_c],
1529 )?;
1530 let cov_local = beta_covariance.map(|cov| {
1531 cov.slice(s![
1532 cache.coeff_global_range.clone(),
1533 cache.coeff_global_range.clone()
1534 ])
1535 .to_owned()
1536 });
1537 let dim = cache.dimension;
1538 let (var_0, var_g, var_c) = match cov_local.as_ref() {
1539 Some(cov) => (
1540 scalar_operator_response_variance(&cache.d0, cov),
1541 grouped_operator_response_variance(&cache.d1, dim, cov)?,
1542 grouped_operator_response_variance(&cache.d2, dim * dim, cov)?,
1543 ),
1544 None => (
1545 Array1::<f64>::zeros(exact.magnitude.signal.len()),
1546 Array1::<f64>::zeros(exact.gradient.norm.len()),
1547 Array1::<f64>::zeros(exact.curvature.norm.len()),
1548 ),
1549 };
1550 let (_, inv_0) = exact.magnitude.surrogateweights_posterior_snr(
1551 &var_0,
1552 weight_floor,
1553 weight_ceiling,
1554 );
1555 let (_, inv_g) =
1556 exact
1557 .gradient
1558 .surrogateweights_posterior_snr(&var_g, weight_floor, weight_ceiling);
1559 let (_, inv_c) = exact.curvature.surrogateweights_posterior_snr(
1560 &var_c,
1561 weight_floor,
1562 weight_ceiling,
1563 );
1564 Ok(SpatialAdaptiveWeights {
1565 inv_magweight: inv_0,
1566 invgradweight: inv_g,
1567 inv_lapweight: inv_c,
1568 })
1569 })
1570 .collect()
1571}
1572
1573fn compute_initial_epsilons(
1574 beta: &Array1<f64>,
1575 caches: &[SpatialOperatorRuntimeCache],
1576 min_epsilon: f64,
1577) -> Result<(f64, f64, f64), EstimationError> {
1578 let mut fvals = Vec::<f64>::new();
1579 let mut gvals = Vec::<f64>::new();
1580 let mut cvals = Vec::<f64>::new();
1581 for cache in caches {
1582 let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
1583 let exact = SpatialPenaltyExactState::from_beta_local(
1584 beta_local,
1585 cache,
1586 [min_epsilon, min_epsilon, min_epsilon],
1587 )?;
1588 let (f, g, c) = exact.absolute_collocation_magnitudes();
1589 fvals.extend(f.iter().copied());
1590 gvals.extend(g.iter().copied());
1591 cvals.extend(c.iter().copied());
1592 }
1593 let eps_0 = robust_epsilon_from_samples(&fvals, min_epsilon);
1599 let eps_g = robust_epsilon_from_samples(&gvals, min_epsilon);
1600 let eps_c = robust_epsilon_from_samples(&cvals, min_epsilon);
1601 Ok((eps_0, eps_g, eps_c))
1602}
1603
1604fn exact_spatial_adaptive_penalty_index_set(
1605 caches: &[SpatialOperatorRuntimeCache],
1606) -> BTreeSet<usize> {
1607 let mut out = BTreeSet::new();
1608 for cache in caches {
1609 out.insert(cache.mass_penalty_global_idx);
1610 out.insert(cache.tension_penalty_global_idx);
1611 out.insert(cache.stiffness_penalty_global_idx);
1612 }
1613 out
1614}
1615
1616fn build_spatial_adaptive_hyperspecs(cache_count: usize) -> Vec<SpatialAdaptiveHyperSpec> {
1617 let mut out = Vec::with_capacity(cache_count * 3 + 3);
1618 for cache_index in 0..cache_count {
1619 out.push(SpatialAdaptiveHyperSpec {
1620 cache_index,
1621 kind: SpatialAdaptiveHyperKind::LogLambdaMagnitude,
1622 });
1623 out.push(SpatialAdaptiveHyperSpec {
1624 cache_index,
1625 kind: SpatialAdaptiveHyperKind::LogLambdaGradient,
1626 });
1627 out.push(SpatialAdaptiveHyperSpec {
1628 cache_index,
1629 kind: SpatialAdaptiveHyperKind::LogLambdaCurvature,
1630 });
1631 }
1632 out.push(SpatialAdaptiveHyperSpec {
1633 cache_index: 0,
1634 kind: SpatialAdaptiveHyperKind::LogEpsilonMagnitude,
1635 });
1636 out.push(SpatialAdaptiveHyperSpec {
1637 cache_index: 0,
1638 kind: SpatialAdaptiveHyperKind::LogEpsilonGradient,
1639 });
1640 out.push(SpatialAdaptiveHyperSpec {
1641 cache_index: 0,
1642 kind: SpatialAdaptiveHyperKind::LogEpsilonCurvature,
1643 });
1644 out
1645}
1646
1647fn penalty_matrixwith_local_block(
1648 total_dim: usize,
1649 coeff_range: Range<usize>,
1650 local: &Array2<f64>,
1651) -> Array2<f64> {
1652 let mut out = Array2::<f64>::zeros((total_dim, total_dim));
1653 out.slice_mut(s![coeff_range.clone(), coeff_range])
1654 .assign(local);
1655 out
1656}
1657
1658fn fit_term_collectionwith_exact_spatial_adaptive_regularization(
1659 baseline: FittedTermCollection,
1660 y: ArrayView1<'_, f64>,
1661 weights: ArrayView1<'_, f64>,
1662 offset: ArrayView1<'_, f64>,
1663 family: LikelihoodSpec,
1664 options: &FitOptions,
1665 runtime_caches: &[SpatialOperatorRuntimeCache],
1666) -> Result<FittedTermCollection, EstimationError> {
1667 let adaptive_opts = options.adaptive_regularization.clone().unwrap_or_default();
1696 let adaptive_penalty_indices = exact_spatial_adaptive_penalty_index_set(runtime_caches);
1697 let p_total = baseline.design.design.ncols();
1698 struct RetainedPenaltySetup {
1699 global_idx: usize,
1700 global_penalty: Array2<f64>,
1701 nullspace_dim: usize,
1702 log_lambda: f64,
1703 col_range: Range<usize>,
1704 hessian_piece: Array2<f64>,
1705 }
1706 use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
1707 let retained_setups = baseline
1708 .design
1709 .penalties
1710 .par_iter()
1711 .enumerate()
1712 .map(|(idx, bp)| {
1713 if adaptive_penalty_indices.contains(&idx) {
1714 return None;
1715 }
1716 let lambda = baseline.fit.lambdas[idx];
1717 Some(RetainedPenaltySetup {
1718 global_idx: idx,
1719 global_penalty: bp.to_global(p_total),
1720 nullspace_dim: baseline
1721 .design
1722 .nullspace_dims
1723 .get(idx)
1724 .copied()
1725 .unwrap_or(0),
1726 log_lambda: lambda.max(1e-12).ln(),
1727 col_range: bp.col_range.clone(),
1728 hessian_piece: bp.local.mapv(|v| lambda * v),
1729 })
1730 })
1731 .collect::<Vec<_>>();
1732 let retained_count = retained_setups
1733 .iter()
1734 .filter(|setup| setup.is_some())
1735 .count();
1736 let mut retained_penalties = Vec::<Array2<f64>>::with_capacity(retained_count);
1737 let mut retained_nullspace_dims = Vec::<usize>::with_capacity(retained_count);
1738 let mut retained_log_lambdas = Vec::<f64>::with_capacity(retained_count);
1739 let mut retained_global_indices = Vec::<usize>::with_capacity(retained_count);
1740 let mut fixed_quadratichessian = Array2::<f64>::zeros((p_total, p_total));
1741 for setup in retained_setups.into_iter().flatten() {
1742 retained_penalties.push(setup.global_penalty);
1743 retained_nullspace_dims.push(setup.nullspace_dim);
1744 retained_log_lambdas.push(setup.log_lambda);
1745 retained_global_indices.push(setup.global_idx);
1746 fixed_quadratichessian
1747 .slice_mut(s![setup.col_range.clone(), setup.col_range])
1748 .scaled_add(1.0, &setup.hessian_piece);
1749 }
1750
1751 let (eps_0_init, eps_g_init, eps_c_init) = compute_initial_epsilons(
1752 &baseline.fit.beta,
1753 runtime_caches,
1754 adaptive_opts.min_epsilon,
1755 )?;
1756 let mut initial_theta =
1757 Array1::<f64>::zeros(retained_penalties.len() + runtime_caches.len() * 3 + 3);
1758 for (idx, value) in retained_log_lambdas.iter().enumerate() {
1759 initial_theta[idx] = *value;
1760 }
1761 let adaptive_log_lambda_components = runtime_caches
1762 .par_iter()
1763 .map(|cache| {
1764 [
1765 baseline.fit.lambdas[cache.mass_penalty_global_idx]
1766 .max(1e-12)
1767 .ln(),
1768 baseline.fit.lambdas[cache.tension_penalty_global_idx]
1769 .max(1e-12)
1770 .ln(),
1771 baseline.fit.lambdas[cache.stiffness_penalty_global_idx]
1772 .max(1e-12)
1773 .ln(),
1774 ]
1775 })
1776 .collect::<Vec<_>>();
1777 let mut at = retained_penalties.len();
1778 for logs in &adaptive_log_lambda_components {
1779 initial_theta[at] = logs[0];
1780 initial_theta[at + 1] = logs[1];
1781 initial_theta[at + 2] = logs[2];
1782 at += 3;
1783 }
1784 initial_theta[at] = eps_0_init.max(adaptive_opts.min_epsilon).ln();
1785 initial_theta[at + 1] = eps_g_init.max(adaptive_opts.min_epsilon).ln();
1786 initial_theta[at + 2] = eps_c_init.max(adaptive_opts.min_epsilon).ln();
1787
1788 let hyperspecs = build_spatial_adaptive_hyperspecs(runtime_caches.len());
1789 let zero_psi_op: std::sync::Arc<dyn gam_custom_family::CustomFamilyPsiDerivativeOperator> =
1790 std::sync::Arc::new(gam_custom_family::ZeroPsiDerivativeOperator::new(
1791 baseline.design.design.nrows(),
1792 baseline.design.design.ncols(),
1793 ));
1794 let derivative_blocks = vec![
1795 hyperspecs
1796 .par_iter()
1797 .map(|_| CustomFamilyBlockPsiDerivative {
1798 penalty_index: None,
1799 x_psi: Array2::<f64>::zeros((0, 0)),
1800 s_psi: Array2::<f64>::zeros((0, 0)),
1801 s_psi_components: None,
1802 s_psi_penalty_components: None,
1803 x_psi_psi: None,
1804 s_psi_psi: None,
1805 s_psi_psi_components: None,
1806 s_psi_psi_penalty_components: None,
1807 implicit_operator: Some(std::sync::Arc::clone(&zero_psi_op)),
1808 implicit_axis: 0,
1809 implicit_group_id: None,
1810 })
1811 .collect::<Vec<_>>(),
1812 ];
1813
1814 let mixture_link_state = options
1815 .mixture_link
1816 .clone()
1817 .as_ref()
1818 .map(state_fromspec)
1819 .transpose()
1820 .map_err(EstimationError::InvalidInput)?;
1821 let sas_link_state = options
1822 .sas_link
1823 .map(|spec| {
1824 if family.is_binomial_beta_logistic() {
1825 state_from_beta_logisticspec(spec)
1826 } else {
1827 state_from_sasspec(spec)
1828 }
1829 })
1830 .transpose()
1831 .map_err(EstimationError::InvalidInput)?;
1832 let latent_cloglog_state = options.latent_cloglog;
1833 let shared_y = Arc::new(y.to_owned());
1834 let sharedweights = Arc::new(weights.to_owned());
1835 let shared_design = baseline
1836 .design
1837 .design
1838 .try_to_dense_arc("spatial adaptive exact hyperfit design")
1839 .map_err(EstimationError::InvalidInput)?;
1840 let shared_offset = Arc::new(offset.to_owned());
1841 let shared_runtime_caches = Arc::new(runtime_caches.to_vec());
1842 let shared_hyperspecs = Arc::new(hyperspecs.clone());
1843 let zero_quadratic = Arc::new(Array2::<f64>::zeros((
1844 baseline.design.design.ncols(),
1845 baseline.design.design.ncols(),
1846 )));
1847 let base_family = SpatialAdaptiveExactFamily {
1848 family: family.clone(),
1849 latent_cloglog_state,
1850 mixture_link_state: mixture_link_state.clone(),
1851 sas_link_state,
1852 y: shared_y.clone(),
1853 weights: sharedweights.clone(),
1854 design: shared_design.clone(),
1855 offset: shared_offset.clone(),
1856 linear_constraints: baseline.design.linear_constraints.clone(),
1857 runtime_caches: shared_runtime_caches.clone(),
1858 adaptive_params: Vec::new(),
1859 fixed_quadratichessian: zero_quadratic.clone(),
1860 hyperspecs: shared_hyperspecs.clone(),
1861 exact_eval_cache: Arc::new(Mutex::new(None)),
1862 };
1863
1864 let rho_dim = retained_penalties.len();
1865 let operator_slots_end = rho_dim + runtime_caches.len() * 3;
1866 const UNIFIED_LOG_WINDOW: f64 = 6.0;
1876 const RETAINED_LAMBDA_LOG_LOWER_FLOOR: f64 = -30.0;
1877 const RETAINED_LAMBDA_LOG_UPPER_CAP: f64 = 30.0;
1878 const OPERATOR_LAMBDA_LOG_LOWER_FLOOR: f64 = -10.0;
1879 const OPERATOR_LAMBDA_LOG_UPPER_CAP: f64 = 30.0;
1880 let epsilon_floor_log = adaptive_opts.min_epsilon.max(1e-12).ln();
1881 let anchored_bound = |idx: usize, sign: f64| -> f64 {
1882 let raw = initial_theta[idx] + sign * UNIFIED_LOG_WINDOW;
1883 if idx < rho_dim {
1884 raw.clamp(
1885 RETAINED_LAMBDA_LOG_LOWER_FLOOR,
1886 RETAINED_LAMBDA_LOG_UPPER_CAP,
1887 )
1888 } else if idx < operator_slots_end {
1889 raw.clamp(
1890 OPERATOR_LAMBDA_LOG_LOWER_FLOOR,
1891 OPERATOR_LAMBDA_LOG_UPPER_CAP,
1892 )
1893 } else {
1894 raw.max(epsilon_floor_log)
1895 }
1896 };
1897 let eps_lower =
1898 Array1::from_iter((0..initial_theta.len()).map(|idx| anchored_bound(idx, -1.0)));
1899 let eps_upper = Array1::from_iter((0..initial_theta.len()).map(|idx| anchored_bound(idx, 1.0)));
1900 let blockspec = ParameterBlockSpec {
1901 name: "eta".to_string(),
1902 design: baseline.design.design.clone(),
1903 offset: offset.to_owned(),
1904 penalties: retained_penalties
1905 .iter()
1906 .cloned()
1907 .map(PenaltyMatrix::Dense)
1908 .collect(),
1909 nullspace_dims: retained_nullspace_dims.clone(),
1910 initial_log_lambdas: Array1::from_vec(retained_log_lambdas.clone()),
1911 initial_beta: Some(baseline.fit.beta.clone()),
1912 gauge_priority: 100,
1913 jacobian_callback: None,
1914 stacked_design: None,
1915 stacked_offset: None,
1916 };
1917 let screening_cap = Arc::new(AtomicUsize::new(0));
1918 let outer_opts = BlockwiseFitOptions {
1919 inner_max_cycles: options.max_iter,
1920 inner_tol: options.tol,
1921 outer_max_iter: options.max_iter,
1922 outer_tol: options.tol,
1923 compute_covariance: false,
1924 screening_max_inner_iterations: Some(Arc::clone(&screening_cap)),
1925 ..BlockwiseFitOptions::default()
1926 };
1927
1928 use gam_solve::rho_optimizer::OuterProblem;
1929 use gam_problem::{DeclaredHessianForm, Derivative, HessianValue, OuterEval};
1930
1931 struct SpatialAdaptiveOuterState {
1932 warm_cache: Option<CustomFamilyWarmStart>,
1933 last_eval: Option<(
1934 Array1<f64>,
1935 f64,
1936 Array1<f64>,
1937 HessianValue,
1938 CustomFamilyWarmStart,
1939 )>,
1940 }
1941
1942 let n_theta = initial_theta.len();
1943
1944 let theta_bounds = Some((eps_lower.clone(), eps_upper.clone()));
1947 let clamp_theta = {
1948 let lo = eps_lower;
1949 let hi = eps_upper;
1950 move |theta: &Array1<f64>| -> Array1<f64> {
1951 let mut clamped = theta.clone();
1952 for i in 0..clamped.len() {
1953 clamped[i] = clamped[i].clamp(lo[i], hi[i]);
1954 }
1955 clamped
1956 }
1957 };
1958
1959 let decode_theta = |theta: &Array1<f64>| -> (Array1<f64>, Vec<SpatialAdaptiveTermHyperParams>) {
1960 let rho = theta.slice(s![..rho_dim]).to_owned();
1961 let adaptive_lambda_start = rho_dim;
1962 let adaptive_lambda_end = adaptive_lambda_start + runtime_caches.len() * 3;
1963 let eps = [
1964 theta[adaptive_lambda_end].exp(),
1965 theta[adaptive_lambda_end + 1].exp(),
1966 theta[adaptive_lambda_end + 2].exp(),
1967 ];
1968 let adaptive_params = runtime_caches
1969 .iter()
1970 .enumerate()
1971 .map(|(cache_idx, _)| SpatialAdaptiveTermHyperParams {
1972 lambda: [
1973 theta[adaptive_lambda_start + cache_idx * 3].exp(),
1974 theta[adaptive_lambda_start + cache_idx * 3 + 1].exp(),
1975 theta[adaptive_lambda_start + cache_idx * 3 + 2].exp(),
1976 ],
1977 epsilon: eps,
1978 })
1979 .collect::<Vec<_>>();
1980 (rho, adaptive_params)
1981 };
1982 let analytic_outer_hessian_available =
1983 gam_custom_family::joint_exact_analytic_outer_hessian_available()
1984 && base_family
1985 .exact_outer_derivative_order(std::slice::from_ref(&blockspec), &outer_opts)
1986 .has_hessian()
1987 && gam_custom_family::exact_newton_outer_geometry_supports_second_order_solver(
1988 &base_family,
1989 );
1990 let problem = OuterProblem::new(n_theta)
1996 .with_gradient(Derivative::Analytic)
1997 .with_hessian(if analytic_outer_hessian_available {
1998 DeclaredHessianForm::Either
1999 } else {
2000 DeclaredHessianForm::Unavailable
2001 })
2002 .with_fallback_policy(gam_solve::rho_optimizer::FallbackPolicy::Disabled)
2003 .with_psi_dim(n_theta.saturating_sub(rho_dim))
2004 .with_tolerance(options.tol)
2005 .with_max_iter(options.max_iter)
2006 .with_seed_config(gam_problem::SeedConfig::default())
2007 .with_screening_cap(Arc::clone(&screening_cap))
2008 .with_initial_rho(initial_theta.clone());
2009 let problem = if let Some((lo, hi)) = theta_bounds {
2010 problem.with_bounds(lo, hi)
2011 } else {
2012 problem
2013 };
2014
2015 let eval_outer = |st: &mut SpatialAdaptiveOuterState,
2016 theta: &Array1<f64>,
2017 order: gam_solve::rho_optimizer::OuterEvalOrder|
2018 -> Result<OuterEval, EstimationError> {
2019 let theta = clamp_theta(theta);
2020
2021 if let Some((cached_theta, cached_cost, cached_grad, cached_hess, cached_warm)) =
2022 &st.last_eval
2023 && cached_theta.len() == theta.len()
2024 && cached_theta
2025 .iter()
2026 .zip(theta.iter())
2027 .all(|(&a, &b)| (a - b).abs() <= 1e-12)
2028 && (!matches!(
2029 order,
2030 gam_solve::rho_optimizer::OuterEvalOrder::ValueGradientHessian
2031 ) || analytic_outer_hessian_available)
2032 {
2033 st.warm_cache = Some(cached_warm.clone());
2034 return Ok(OuterEval {
2035 cost: *cached_cost,
2036 gradient: cached_grad.clone(),
2037 hessian: if matches!(
2038 order,
2039 gam_solve::rho_optimizer::OuterEvalOrder::ValueGradientHessian
2040 ) && analytic_outer_hessian_available
2041 {
2042 cached_hess.clone()
2043 } else {
2044 HessianValue::Unavailable
2045 },
2046 inner_beta_hint: None,
2047 });
2048 }
2049
2050 let (rho, adaptive_params) = decode_theta(&theta);
2051 let family_eval = base_family.with_adaptive_params(adaptive_params, zero_quadratic.clone());
2052 let need_hessian = matches!(
2053 order,
2054 gam_solve::rho_optimizer::OuterEvalOrder::ValueGradientHessian
2055 ) && analytic_outer_hessian_available;
2056 let result = evaluate_custom_family_joint_hyper(
2057 &family_eval,
2058 std::slice::from_ref(&blockspec),
2059 &outer_opts,
2060 &rho,
2061 &derivative_blocks,
2062 st.warm_cache.as_ref(),
2063 if need_hessian {
2064 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueGradientHessian
2065 } else {
2066 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueAndGradient
2067 },
2068 )
2069 .map_err(|e| {
2070 EstimationError::RemlOptimizationFailed(format!("spatial adaptive eval failed: {e}"))
2071 })?;
2072 if !result.inner_converged {
2073 st.warm_cache = Some(result.warm_start.clone());
2074 return Err(EstimationError::RemlOptimizationFailed(
2075 "exact spatial adaptive inner solve did not converge".to_string(),
2076 ));
2077 }
2078 if !result.objective.is_finite() || result.gradient.iter().any(|v| !v.is_finite()) {
2079 return Err(EstimationError::RemlOptimizationFailed(
2080 "exact spatial adaptive objective returned non-finite values".to_string(),
2081 ));
2082 }
2083 let hessian_result = if need_hessian {
2084 if !result.outer_hessian.is_analytic() {
2085 return Err(EstimationError::RemlOptimizationFailed(
2086 "exact spatial adaptive objective did not return an exact outer Hessian"
2087 .to_string(),
2088 ));
2089 }
2090 match result.outer_hessian.dim() {
2091 Some(dim) if dim == theta.len() => {}
2092 Some(dim) => {
2093 return Err(EstimationError::RemlOptimizationFailed(format!(
2094 "exact spatial adaptive outer Hessian dimension mismatch: got {dim}, expected {}",
2095 theta.len(),
2096 )));
2097 }
2098 None => {
2099 return Err(EstimationError::RemlOptimizationFailed(
2100 "exact spatial adaptive objective did not report an outer Hessian dimension"
2101 .to_string(),
2102 ));
2103 }
2104 }
2105 st.last_eval = Some((
2106 theta.clone(),
2107 result.objective,
2108 result.gradient.clone(),
2109 result.outer_hessian.clone(),
2110 result.warm_start.clone(),
2111 ));
2112 result.outer_hessian
2113 } else {
2114 HessianValue::Unavailable
2115 };
2116 st.warm_cache = Some(result.warm_start);
2117 Ok(OuterEval {
2118 cost: result.objective,
2119 gradient: result.gradient,
2120 hessian: hessian_result,
2121 inner_beta_hint: None,
2122 })
2123 };
2124
2125 let mut obj = problem.build_objective_with_screening_proxy(
2126 SpatialAdaptiveOuterState {
2127 warm_cache: None,
2128 last_eval: None,
2129 },
2130 |st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
2131 let theta = clamp_theta(theta);
2132 let (rho, adaptive_params) = decode_theta(&theta);
2133 let family_eval =
2134 base_family.with_adaptive_params(adaptive_params, zero_quadratic.clone());
2135 let result = evaluate_custom_family_joint_hyper(
2136 &family_eval,
2137 std::slice::from_ref(&blockspec),
2138 &outer_opts,
2139 &rho,
2140 &derivative_blocks,
2141 st.warm_cache.as_ref(),
2142 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueOnly,
2143 )
2144 .map_err(|e| {
2145 EstimationError::RemlOptimizationFailed(format!(
2146 "spatial adaptive cost eval failed: {e}"
2147 ))
2148 })?;
2149 if !result.inner_converged {
2150 st.warm_cache = Some(result.warm_start);
2151 return Err(EstimationError::RemlOptimizationFailed(
2152 "exact spatial adaptive cost inner solve did not converge".to_string(),
2153 ));
2154 }
2155 st.warm_cache = Some(result.warm_start);
2156 Ok(result.objective)
2157 },
2158 |st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
2159 eval_outer(
2160 st,
2161 theta,
2162 if analytic_outer_hessian_available {
2163 gam_solve::rho_optimizer::OuterEvalOrder::ValueGradientHessian
2164 } else {
2165 gam_solve::rho_optimizer::OuterEvalOrder::ValueAndGradient
2166 },
2167 )
2168 },
2169 |st: &mut SpatialAdaptiveOuterState,
2170 theta: &Array1<f64>,
2171 order: gam_solve::rho_optimizer::OuterEvalOrder| {
2172 eval_outer(st, theta, order)
2173 },
2174 Some(|st: &mut SpatialAdaptiveOuterState| {
2175 st.warm_cache = None;
2176 st.last_eval = None;
2177 }),
2178 Some(|st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
2179 let theta = clamp_theta(theta);
2180 let (rho, adaptive_params) = decode_theta(&theta);
2181 let family_eval =
2182 base_family.with_adaptive_params(adaptive_params, zero_quadratic.clone());
2183 let result = evaluate_custom_family_joint_hyper_efs(
2184 &family_eval,
2185 std::slice::from_ref(&blockspec),
2186 &outer_opts,
2187 &rho,
2188 &derivative_blocks,
2189 st.warm_cache.as_ref(),
2190 )
2191 .map_err(|e| {
2192 EstimationError::RemlOptimizationFailed(format!(
2193 "spatial adaptive EFS eval failed: {e}"
2194 ))
2195 })?;
2196 if !result.inner_converged {
2197 st.warm_cache = Some(result.warm_start);
2198 return Err(EstimationError::RemlOptimizationFailed(
2199 "exact spatial adaptive EFS inner solve did not converge".to_string(),
2200 ));
2201 }
2202 st.warm_cache = Some(result.warm_start);
2203 Ok(result.efs_eval)
2204 }),
2205 |st: &mut SpatialAdaptiveOuterState, theta: &Array1<f64>| {
2217 let theta = clamp_theta(theta);
2218 let (rho, adaptive_params) = decode_theta(&theta);
2219 let family_eval =
2220 base_family.with_adaptive_params(adaptive_params, zero_quadratic.clone());
2221 let result = evaluate_custom_family_joint_hyper(
2222 &family_eval,
2223 std::slice::from_ref(&blockspec),
2224 &outer_opts,
2225 &rho,
2226 &derivative_blocks,
2227 st.warm_cache.as_ref(),
2228 gam_solve::estimate::reml::reml_outer_engine::EvalMode::ValueOnly,
2229 )
2230 .map_err(|e| {
2231 EstimationError::RemlOptimizationFailed(format!(
2232 "spatial adaptive screening eval failed: {e}"
2233 ))
2234 })?;
2235 st.warm_cache = Some(result.warm_start);
2236 Ok(result.objective)
2237 },
2238 );
2239
2240 let outer_result = problem
2241 .run(&mut obj, "exact spatial adaptive regularization")
2242 .map_err(|e| {
2243 EstimationError::InvalidInput(format!(
2244 "exact spatial adaptive outer optimization failed: {e}"
2245 ))
2246 })?;
2247 if !outer_result.converged {
2248 let rel_to_cost_threshold = options.tol * (1.0_f64 + outer_result.final_value.abs());
2265 if let Some(final_grad) = outer_result
2269 .final_grad_norm
2270 .filter(|v| v.is_finite() && *v <= rel_to_cost_threshold)
2271 {
2272 log::info!(
2273 "[spatial-adaptive] outer optimization hit max_iter={} but \
2274 projected gradient norm {:.3e} ≤ τ·(1+|f|) = {:.3e} \
2275 (τ={:.3e}, |f|={:.3e}); accepting iterate under the mgcv-style \
2276 relative-to-cost REML convergence criterion.",
2277 outer_result.iterations,
2278 final_grad,
2279 rel_to_cost_threshold,
2280 options.tol,
2281 outer_result.final_value.abs(),
2282 );
2283 } else {
2284 crate::bail_invalid_estim!(
2285 "exact spatial adaptive outer optimization did not converge after {} iterations (final_objective={:.6e}, final_grad_norm={})",
2286 outer_result.iterations,
2287 outer_result.final_value,
2288 outer_result.final_grad_norm_report(),
2289 );
2290 }
2291 }
2292 let outer_iterations = outer_result.iterations;
2293 let outer_grad_norm: Option<f64> = outer_result.final_grad_norm;
2296 let theta_star = outer_result.rho;
2297 let rho_star = theta_star.slice(s![..rho_dim]).to_owned();
2298 let adaptive_lambda_start = rho_dim;
2299 let adaptive_lambda_end = adaptive_lambda_start + runtime_caches.len() * 3;
2300 let eps_star = [
2301 theta_star[adaptive_lambda_end].exp(),
2302 theta_star[adaptive_lambda_end + 1].exp(),
2303 theta_star[adaptive_lambda_end + 2].exp(),
2304 ];
2305 let adaptive_params = runtime_caches
2306 .iter()
2307 .enumerate()
2308 .map(|(cache_idx, _)| SpatialAdaptiveTermHyperParams {
2309 lambda: [
2310 theta_star[adaptive_lambda_start + cache_idx * 3].exp(),
2311 theta_star[adaptive_lambda_start + cache_idx * 3 + 1].exp(),
2312 theta_star[adaptive_lambda_start + cache_idx * 3 + 2].exp(),
2313 ],
2314 epsilon: eps_star,
2315 })
2316 .collect::<Vec<_>>();
2317 let mut fixed_total = Array2::<f64>::zeros((
2318 baseline.design.design.ncols(),
2319 baseline.design.design.ncols(),
2320 ));
2321 for (idx, penalty) in retained_penalties.iter().enumerate() {
2322 fixed_total.scaled_add(rho_star[idx].exp(), penalty);
2323 }
2324 let final_family =
2325 base_family.with_adaptive_params(adaptive_params.clone(), Arc::new(fixed_total.clone()));
2326 let final_blockspec = ParameterBlockSpec {
2327 name: "eta".to_string(),
2328 design: baseline.design.design.clone(),
2329 offset: offset.to_owned(),
2330 penalties: vec![],
2331 nullspace_dims: vec![],
2332 initial_log_lambdas: Array1::zeros(0),
2333 initial_beta: Some(baseline.fit.beta.clone()),
2334 gauge_priority: 100,
2335 jacobian_callback: None,
2336 stacked_design: None,
2337 stacked_offset: None,
2338 };
2339 let final_fit = fit_custom_family(
2340 &final_family,
2341 &[final_blockspec],
2342 &BlockwiseFitOptions {
2343 inner_max_cycles: options.max_iter,
2344 inner_tol: options.tol,
2345 outer_max_iter: 1,
2346 outer_tol: options.tol,
2347 compute_covariance: true,
2348 ..BlockwiseFitOptions::default()
2349 },
2350 )
2351 .map_err(EstimationError::CustomFamily)?;
2352 let beta = final_fit.block_states[0].beta.clone();
2353 let final_eval = final_family
2354 .exact_evaluation(&beta)
2355 .map_err(EstimationError::InvalidInput)?;
2356 let penalized_hessian = final_eval
2357 .totalobjectivehessian(&final_family.design)
2358 .map_err(EstimationError::InvalidInput)?;
2359 let beta_covariance = final_fit.covariance_conditional.clone();
2360 let beta_standard_errors = beta_covariance
2361 .as_ref()
2362 .map(|cov| Array1::from_iter((0..cov.nrows()).map(|i| cov[[i, i]].max(0.0).sqrt())));
2363
2364 let mut full_lambdas = baseline.fit.lambdas.clone();
2365 for (idx, &global_idx) in retained_global_indices.iter().enumerate() {
2366 full_lambdas[global_idx] = rho_star[idx].exp();
2367 }
2368 for (cache_idx, cache) in runtime_caches.iter().enumerate() {
2369 full_lambdas[cache.mass_penalty_global_idx] = adaptive_params[cache_idx].lambda[0];
2370 full_lambdas[cache.tension_penalty_global_idx] = adaptive_params[cache_idx].lambda[1];
2371 full_lambdas[cache.stiffness_penalty_global_idx] = adaptive_params[cache_idx].lambda[2];
2372 }
2373
2374 let deviance = if family.is_gaussian_identity() {
2375 y.iter()
2376 .zip(final_eval.obs.mu.iter())
2377 .zip(weights.iter())
2378 .map(|((&yy, &mu), &w)| w.max(0.0) * (yy - mu) * (yy - mu))
2379 .sum()
2380 } else {
2381 -2.0 * final_eval.obs.log_likelihood
2382 };
2383 let mut local_penalty_blocks =
2384 Vec::<PenaltySpec>::with_capacity(baseline.design.penalties.len());
2385 for (global_idx, bp) in baseline.design.penalties.iter().enumerate() {
2386 if adaptive_penalty_indices.contains(&global_idx) {
2387 let cache = runtime_caches
2388 .iter()
2389 .find(|cache| {
2390 cache.mass_penalty_global_idx == global_idx
2391 || cache.tension_penalty_global_idx == global_idx
2392 || cache.stiffness_penalty_global_idx == global_idx
2393 })
2394 .ok_or_else(|| {
2395 EstimationError::InvalidInput(format!(
2396 "missing runtime cache for adaptive penalty index {global_idx}"
2397 ))
2398 })?;
2399 let cache_idx = runtime_caches
2400 .iter()
2401 .position(|c| {
2402 c.mass_penalty_global_idx == global_idx
2403 || c.tension_penalty_global_idx == global_idx
2404 || c.stiffness_penalty_global_idx == global_idx
2405 })
2406 .ok_or_else(|| {
2407 EstimationError::InvalidInput(format!(
2408 "missing adaptive cache position for penalty index {global_idx}"
2409 ))
2410 })?;
2411 let state = &final_eval.adaptive_states[cache_idx];
2412 let local = if cache.mass_penalty_global_idx == global_idx {
2413 scalar_operatorhessian(&cache.d0, &state.magnitude.betahessian_diag())
2414 .mapv(|v| adaptive_params[cache_idx].lambda[0] * v)
2415 } else if cache.tension_penalty_global_idx == global_idx {
2416 grouped_operatorhessian(
2417 &cache.d1,
2418 cache.dimension,
2419 &state.gradient.betahessian_blocks(),
2420 )?
2421 .mapv(|v| adaptive_params[cache_idx].lambda[1] * v)
2422 } else {
2423 grouped_operatorhessian(
2424 &cache.d2,
2425 cache.dimension * cache.dimension,
2426 &state.curvature.betahessian_blocks(),
2427 )?
2428 .mapv(|v| adaptive_params[cache_idx].lambda[2] * v)
2429 };
2430 local_penalty_blocks.push(PenaltySpec::Dense(penalty_matrixwith_local_block(
2432 baseline.design.design.ncols(),
2433 cache.coeff_global_range.clone(),
2434 &local,
2435 )));
2436 } else {
2437 local_penalty_blocks.push(PenaltySpec::Dense(
2438 bp.to_global(p_total).mapv(|v| v * full_lambdas[global_idx]),
2439 ));
2440 }
2441 }
2442 let (edf_by_block, penalty_block_trace, edf_total) = if let Some(cov) = beta_covariance.as_ref()
2443 {
2444 exact_bounded_edf(
2445 &local_penalty_blocks,
2446 &Array1::from_elem(local_penalty_blocks.len(), 1.0),
2447 cov,
2448 )?
2449 } else {
2450 (
2451 vec![0.0; local_penalty_blocks.len()],
2452 vec![0.0; local_penalty_blocks.len()],
2453 0.0,
2454 )
2455 };
2456 let stable_penalty_term =
2457 2.0 * final_eval.adaptive_penalty_value + beta.dot(&fixed_total.dot(&beta));
2458 let standard_deviation = if family.is_gaussian_identity() {
2459 let denom = (y.len() as f64 - edf_total).max(1.0);
2460 (deviance / denom).sqrt()
2461 } else {
2462 1.0
2463 };
2464 let maps = compute_spatial_adaptiveweights_for_beta(
2465 &beta,
2466 runtime_caches,
2467 eps_star[0],
2468 eps_star[1],
2469 eps_star[2],
2470 adaptive_opts.weight_floor,
2471 adaptive_opts.weight_ceiling,
2472 beta_covariance.as_ref(),
2476 )?
2477 .into_iter()
2478 .zip(runtime_caches.iter())
2479 .map(|(w, cache)| AdaptiveSpatialMap {
2480 termname: cache.termname.clone(),
2481 feature_cols: cache.feature_cols.clone(),
2482 collocation_points: cache.collocation_points.clone(),
2483 inv_magweight: w.inv_magweight,
2484 invgradweight: w.invgradweight,
2485 inv_lapweight: w.inv_lapweight,
2486 })
2487 .collect::<Vec<_>>();
2488 let fitted_link = if family.is_latent_cloglog() {
2489 FittedLinkState::LatentCLogLog {
2490 state: latent_cloglog_state
2491 .expect("BinomialLatentCLogLog requires an explicit latent-cloglog state"),
2492 }
2493 } else if family.is_binomial_mixture() {
2494 mixture_link_state
2495 .clone()
2496 .map(|state| FittedLinkState::Mixture {
2497 state,
2498 covariance: None,
2499 })
2500 .unwrap_or(FittedLinkState::Standard(None))
2501 } else if family.is_binomial_sas() {
2502 sas_link_state
2503 .map(|state| FittedLinkState::Sas {
2504 state,
2505 covariance: None,
2506 })
2507 .unwrap_or(FittedLinkState::Standard(None))
2508 } else if family.is_binomial_beta_logistic() {
2509 sas_link_state
2510 .map(|state| FittedLinkState::BetaLogistic {
2511 state,
2512 covariance: None,
2513 })
2514 .unwrap_or(FittedLinkState::Standard(None))
2515 } else {
2516 FittedLinkState::Standard(None)
2517 };
2518 let max_abs_eta = final_eval
2519 .obs
2520 .eta
2521 .iter()
2522 .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
2523 let fitted = FittedTermCollection {
2524 fit: {
2525 let log_lambdas = full_lambdas.mapv(|v| v.max(1e-300).ln());
2526 let inf = FitInference {
2527 edf_by_block,
2528 penalty_block_trace,
2529 edf_total,
2530 smoothing_correction: None,
2531 penalized_hessian: penalized_hessian.clone().into(),
2534 working_weights: final_eval.obs.fisherweight.clone(),
2535 working_response: {
2536 let mut out = final_eval.obs.eta.clone();
2537 for i in 0..out.len() {
2538 let wi = final_eval.obs.fisherweight[i].max(1e-12);
2539 out[i] += final_eval.obs.score[i] / wi;
2540 }
2541 out
2542 },
2543 reparam_qs: None,
2544 dispersion: gam_solve::estimate::Dispersion::Known(1.0),
2545 beta_covariance: beta_covariance
2546 .clone()
2547 .map(gam_problem::dispersion_cov::PhiScaledCovariance::from),
2548 beta_standard_errors,
2549 beta_covariance_corrected: None,
2550 beta_standard_errors_corrected: None,
2551 beta_covariance_frequentist: None,
2552 coefficient_influence: None,
2553 weighted_gram: None,
2554 bias_correction_beta: None,
2555 bias_correction_jacobian: None,
2556 };
2557 let geometry = Some(gam_solve::estimate::FitGeometry {
2558 penalized_hessian: penalized_hessian.into(),
2559 working_weights: inf.working_weights.clone(),
2560 working_response: inf.working_response.clone(),
2561 });
2562 let covariance_conditional = beta_covariance;
2563 let pirls_status_val = gam_solve::pirls::PirlsStatus::Converged;
2567 UnifiedFitResult::try_from_parts(UnifiedFitResultParts {
2568 blocks: vec![gam_solve::estimate::FittedBlock {
2569 beta: beta.clone(),
2570 role: gam_problem::BlockRole::Mean,
2571 edf: edf_total,
2572 lambdas: full_lambdas.clone(),
2573 }],
2574 log_lambdas,
2575 lambdas: full_lambdas,
2576 likelihood_scale: family.default_scale_metadata(),
2577 likelihood_family: Some(family),
2578 log_likelihood_normalization:
2579 gam_spec::LogLikelihoodNormalization::UserProvided,
2580 log_likelihood: final_eval.obs.log_likelihood,
2581 deviance,
2582 reml_score: final_fit.penalized_objective,
2583 stable_penalty_term,
2584 penalized_objective: final_fit.penalized_objective,
2585 used_device: false,
2586 outer_iterations,
2587 outer_converged: true,
2589 outer_gradient_norm: outer_grad_norm,
2590 standard_deviation,
2591 covariance_conditional,
2592 covariance_corrected: None,
2593 inference: Some(inf),
2594 fitted_link,
2595 geometry,
2596 block_states: Vec::new(),
2597 pirls_status: pirls_status_val,
2598 max_abs_eta,
2599 constraint_kkt: None,
2600 artifacts: gam_solve::estimate::FitArtifacts {
2601 pirls: None,
2602 ..Default::default()
2603 },
2604 inner_cycles: 0,
2605 })?
2606 },
2607 design: baseline.design,
2608 adaptive_diagnostics: Some(AdaptiveRegularizationDiagnostics {
2609 epsilon_0: eps_star[0],
2610 epsilon_g: eps_star[1],
2611 epsilon_c: eps_star[2],
2612 epsilon_outer_iterations: outer_iterations,
2613 mm_iterations: 0,
2614 converged: true,
2615 maps,
2616 }),
2617 };
2618 enforce_term_constraint_feasibility(&fitted.design, &fitted.fit)?;
2619 Ok(fitted)
2620}
2621
2622fn relax_smoothing_rho_prior(
2654 options: &FitOptions,
2655 design: &TermCollectionDesign,
2656) -> gam_spec::RhoPrior {
2657 use gam_terms::basis::BasisMetadata;
2658 let base = &options.rho_prior;
2659 if matches!(
2662 base,
2663 gam_spec::RhoPrior::Flat | gam_spec::RhoPrior::Independent(_)
2664 ) {
2665 return base.clone();
2666 }
2667 let has_link_aux = options.sas_link.is_some()
2687 || options.optimize_sas
2688 || options.mixture_link.is_some()
2689 || options.optimize_mixture;
2690 let has_moving_kappa = design.smooth.terms.iter().any(|t| {
2691 if let BasisMetadata::Duchon {
2704 length_scale,
2705 aniso_log_scales,
2706 ..
2707 } = &t.metadata
2708 {
2709 return length_scale.is_some() || aniso_log_scales.is_some();
2710 }
2711 matches!(
2712 t.metadata,
2713 BasisMetadata::Matern { .. }
2714 | BasisMetadata::Sphere { .. }
2715 | BasisMetadata::SphereHarmonics { .. }
2716 | BasisMetadata::ConstantCurvature { .. }
2717 | BasisMetadata::MeasureJet { .. }
2718 )
2719 });
2720 let length_safe = !has_link_aux && !has_moving_kappa;
2727 if !length_safe {
2728 return base.clone();
2729 }
2730 let coords = &design.penaltyinfo;
2731 if coords.is_empty() {
2732 return base.clone();
2733 }
2734 let n_obs = design.design.nrows();
2745 let p_total = design.design.ncols();
2746 let underdetermined = n_obs < 2 * p_total;
2777 let relaxable_terms: std::collections::HashSet<&str> = design
2789 .smooth
2790 .terms
2791 .iter()
2792 .filter(|t| {
2793 (matches!(
2794 t.metadata,
2795 BasisMetadata::BSpline1D { .. }
2796 | BasisMetadata::ThinPlate { .. }
2797 | BasisMetadata::TensorBSpline { .. }
2798 )
2799 || matches!(
2811 t.metadata,
2812 BasisMetadata::Duchon {
2813 length_scale: None,
2814 aniso_log_scales: None,
2815 ..
2816 }
2817 ))
2818 && matches!(t.shape, gam_terms::smooth::ShapeConstraint::None)
2832 })
2833 .map(|t| t.name.as_str())
2834 .collect();
2835 let any_relaxed = coords.iter().any(|info| {
2836 info.termname
2837 .as_deref()
2838 .is_some_and(|name| relaxable_terms.contains(name))
2839 });
2840 if !any_relaxed {
2841 return base.clone();
2842 }
2843 let relaxed_prior = if underdetermined {
2848 gam_spec::RhoPrior::Normal {
2849 mean: 0.0,
2850 sd: RELAX_UNDERDETERMINED_RHO_SD,
2851 }
2852 } else {
2853 gam_spec::RhoPrior::Flat
2854 };
2855 let nullspace_select_prior = gam_spec::RhoPrior::PenalizedComplexity {
2882 upper: NULLSPACE_SELECT_PC_UPPER,
2883 tail_prob: NULLSPACE_SELECT_PC_TAIL_PROB,
2884 };
2885 let nullspace_degeneracy_prior = gam_spec::RhoPrior::Normal {
2912 mean: 0.0,
2913 sd: NULLSPACE_WELLDET_DEGENERACY_RHO_SD,
2914 };
2915 let per_coord = coords
2916 .iter()
2917 .map(|info| {
2918 let relax = info
2919 .termname
2920 .as_deref()
2921 .is_some_and(|name| relaxable_terms.contains(name));
2922 if !relax {
2923 return base.clone();
2924 }
2925 let is_nullspace =
2926 matches!(info.penalty.source, PenaltySource::DoublePenaltyNullspace);
2927 if is_nullspace {
2966 if underdetermined {
2967 nullspace_select_prior.clone()
2968 } else {
2969 nullspace_degeneracy_prior.clone()
2970 }
2971 } else {
2972 relaxed_prior.clone()
2973 }
2974 })
2975 .collect::<Vec<_>>();
2976 gam_spec::RhoPrior::Independent(per_coord)
2977}
2978
2979const RELAX_UNDERDETERMINED_RHO_SD: f64 = 15.0;
2992
2993const NULLSPACE_SELECT_PC_UPPER: f64 = 0.05;
3011
3012const NULLSPACE_SELECT_PC_TAIL_PROB: f64 = 0.01;
3022
3023fn adaptive_fit_options_base(options: &FitOptions, design: &TermCollectionDesign) -> FitOptions {
3024 FitOptions {
3025 resource_policy: options.resource_policy.clone(),
3026 latent_cloglog: options.latent_cloglog,
3027 mixture_link: options.mixture_link.clone(),
3028 optimize_mixture: options.optimize_mixture,
3029 sas_link: options.sas_link,
3030 optimize_sas: options.optimize_sas,
3031 compute_inference: options.compute_inference,
3032 skip_rho_posterior_inference: options.skip_rho_posterior_inference,
3033 max_iter: options.max_iter,
3034 tol: options.tol,
3035 nullspace_dims: design.nullspace_dims.clone(),
3036 linear_constraints: design.linear_constraints.clone(),
3037 firth_bias_reduction: options.firth_bias_reduction,
3038 adaptive_regularization: None,
3039 penalty_shrinkage_floor: options.penalty_shrinkage_floor,
3040 rho_prior: options.rho_prior.clone(),
3043 kronecker_penalty_system: design.kronecker_penalty_system(),
3044 kronecker_factored: design
3045 .smooth
3046 .terms
3047 .iter()
3048 .find_map(|t| t.kronecker_factored.clone()),
3049 persist_warm_start_disk: options.persist_warm_start_disk,
3050 }
3051}
3052
3053fn superseded_fit_options(options: &FitOptions) -> FitOptions {
3054 let mut fit_options = options.clone();
3055 fit_options.skip_rho_posterior_inference = true;
3056 fit_options
3057}
3058
3059#[derive(Clone)]
3060struct BoundedLinearTermMeta {
3061 col_idx: usize,
3062 min: f64,
3063 max: f64,
3064 prior: BoundedCoefficientPriorSpec,
3065}
3066
3067struct BoundedEffectiveJacobian {
3091 design: Array2<f64>,
3092 bounded_terms: Vec<BoundedLinearTermMeta>,
3093}
3094
3095impl BlockEffectiveJacobian for BoundedEffectiveJacobian {
3096 fn effective_jacobian_rows(
3097 &self,
3098 state: &FamilyLinearizationState<'_>,
3099 rows: std::ops::Range<usize>,
3100 ) -> Result<Array2<f64>, String> {
3101 let p = self.design.ncols();
3102 let n = self.design.nrows();
3103 let rows = rows.start.min(n)..rows.end.min(n);
3104 if !state.beta.is_empty() {
3105 if state.beta.len() != p {
3106 return Err(format!(
3107 "BoundedEffectiveJacobian::effective_jacobian_at: beta length {} != design \
3108 ncols {p}",
3109 state.beta.len(),
3110 ));
3111 }
3112 if state.beta.iter().any(|v| v.is_nan()) {
3113 return Err(
3114 "BoundedEffectiveJacobian::effective_jacobian_at: beta contains NaN"
3115 .to_string(),
3116 );
3117 }
3118 }
3119 let mut jac = self
3120 .design
3121 .slice(ndarray::s![rows.start..rows.end, ..])
3122 .to_owned();
3123 for term in &self.bounded_terms {
3124 let theta = if state.beta.is_empty() {
3125 0.0
3126 } else {
3127 state.beta[term.col_idx]
3128 };
3129 let (_, _, db_dtheta, _, _) = bounded_latent_derivatives(theta, term.min, term.max);
3130 jac.column_mut(term.col_idx).mapv_inplace(|v| v * db_dtheta);
3131 }
3132 Ok(jac)
3133 }
3134}
3135
3136#[derive(Clone)]
3137struct BoundedLinearFamily {
3138 family: LikelihoodSpec,
3139 latent_cloglog_state: Option<LatentCLogLogState>,
3140 mixture_link_state: Option<MixtureLinkState>,
3141 sas_link_state: Option<SasLinkState>,
3142 y: Array1<f64>,
3143 weights: Array1<f64>,
3144 design: Array2<f64>,
3145 designzeroed: Array2<f64>,
3146 offset: Array1<f64>,
3147 bounded_terms: Vec<BoundedLinearTermMeta>,
3148}
3149
3150#[derive(Clone)]
3151struct StandardFamilyObservationState {
3152 eta: Array1<f64>,
3153 mu: Array1<f64>,
3154 score: Array1<f64>,
3155 fisherweight: Array1<f64>,
3156 neghessian_eta: Array1<f64>,
3157 neghessian_eta_derivative: Array1<f64>,
3158 log_likelihood: f64,
3159}
3160
3161fn bounded_logit(z: f64) -> f64 {
3162 let zc = z.clamp(1e-12, 1.0 - 1e-12);
3163 (zc / (1.0 - zc)).ln()
3164}
3165
3166fn stable_sigmoid(theta: f64) -> f64 {
3167 if theta >= 0.0 {
3168 let exp_neg = (-theta).exp();
3169 1.0 / (1.0 + exp_neg)
3170 } else {
3171 let exp_pos = theta.exp();
3172 exp_pos / (1.0 + exp_pos)
3173 }
3174}
3175
3176fn bounded_latent_to_user(theta: f64, min: f64, max: f64) -> (f64, f64, f64) {
3177 let z = stable_sigmoid(theta);
3178 let width = max - min;
3179 let beta = min + width * z;
3180 let db_dtheta = width * z * (1.0 - z);
3181 (beta, z, db_dtheta)
3182}
3183
3184fn bounded_user_to_latent(beta: f64, min: f64, max: f64) -> f64 {
3195 let width = max - min;
3196 if width <= 0.0 || !width.is_finite() {
3197 return 0.0;
3198 }
3199 let z = (beta - min) / width;
3200 bounded_logit(z)
3201}
3202
3203#[derive(Debug, Clone, Copy)]
3207pub struct BoundedSampleColumn {
3208 pub col_idx: usize,
3210 pub min: f64,
3212 pub max: f64,
3214}
3215
3216pub fn sample_bounded_latent_posterior_internal(
3254 beta_user: &Array1<f64>,
3255 user_hessian: &Array2<f64>,
3256 bounded_columns: &[BoundedSampleColumn],
3257 n_draws: usize,
3258 sqrt_cov_scale: f64,
3259 base_seed: u64,
3260) -> Result<Array2<f64>, EstimationError> {
3261 let p = beta_user.len();
3262 if user_hessian.nrows() != p || user_hessian.ncols() != p {
3263 crate::bail_invalid_estim!(
3264 "bounded posterior sampling dimension mismatch: mode has {p} entries, user Hessian is {}x{}",
3265 user_hessian.nrows(),
3266 user_hessian.ncols()
3267 );
3268 }
3269
3270 let mut theta_mode = beta_user.clone();
3272 let mut jac_diag = Array1::<f64>::ones(p);
3273 for bc in bounded_columns {
3274 if bc.col_idx >= p {
3275 crate::bail_invalid_estim!(
3276 "bounded posterior sampling: bounded column index {} out of range for {p} coefficients",
3277 bc.col_idx
3278 );
3279 }
3280 let theta_i = bounded_user_to_latent(beta_user[bc.col_idx], bc.min, bc.max);
3281 let (_, _, db_dtheta) = bounded_latent_to_user(theta_i, bc.min, bc.max);
3282 theta_mode[bc.col_idx] = theta_i;
3283 jac_diag[bc.col_idx] = db_dtheta.max(1e-12);
3288 }
3289
3290 let mut h_latent = user_hessian.clone();
3293 for i in 0..p {
3294 let ji = jac_diag[i];
3295 if ji != 1.0 {
3296 h_latent.row_mut(i).mapv_inplace(|v| v * ji);
3297 h_latent.column_mut(i).mapv_inplace(|v| v * ji);
3298 }
3299 }
3300
3301 use gam_linalg::faer_ndarray::FaerCholesky as _;
3304 use rand::SeedableRng as _;
3305 let chol = h_latent.cholesky(faer::Side::Lower).map_err(|err| {
3306 EstimationError::InvalidInput(format!(
3307 "bounded posterior sampling: Cholesky of the latent penalized Hessian failed: {err:?}"
3308 ))
3309 })?;
3310 let l = chol.lower_triangular();
3311
3312 let mut draws = Array2::<f64>::zeros((n_draws, p));
3313 let mut eps = Array1::<f64>::zeros(p);
3314 let mut delta = Array1::<f64>::zeros(p);
3315 let mut rng = rand::rngs::StdRng::seed_from_u64(base_seed);
3316 for k in 0..n_draws {
3317 for e in eps.iter_mut() {
3318 *e = standard_normal_draw(&mut rng);
3319 }
3320 solve_lower_transpose_into(&l, &eps, &mut delta);
3321 for i in 0..p {
3322 draws[(k, i)] = theta_mode[i] + sqrt_cov_scale * delta[i];
3325 }
3326 for bc in bounded_columns {
3329 let (beta_draw, _, _) = bounded_latent_to_user(draws[(k, bc.col_idx)], bc.min, bc.max);
3330 draws[(k, bc.col_idx)] = beta_draw;
3331 }
3332 }
3333
3334 Ok(draws)
3335}
3336
3337#[inline]
3340fn standard_normal_draw<R: rand::Rng + ?Sized>(rng: &mut R) -> f64 {
3341 use rand::RngExt as _;
3342 let u1 = rng.random::<f64>().max(1e-16);
3343 let u2 = rng.random::<f64>();
3344 (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos()
3345}
3346
3347fn solve_lower_transpose_into(l: &Array2<f64>, b: &Array1<f64>, out: &mut Array1<f64>) {
3351 let p = l.nrows();
3352 for i in (0..p).rev() {
3353 let mut acc = b[i];
3354 for j in (i + 1)..p {
3355 acc -= l[(j, i)] * out[j];
3356 }
3357 let diag = l[(i, i)];
3358 out[i] = if diag.abs() > 0.0 { acc / diag } else { 0.0 };
3359 }
3360}
3361
3362fn bounded_latent_derivatives(theta: f64, min: f64, max: f64) -> (f64, f64, f64, f64, f64) {
3363 let z = stable_sigmoid(theta);
3364 let width = max - min;
3365 let s = z * (1.0 - z);
3366 let beta = min + width * z;
3367 let db_dtheta = width * s;
3368 let d2b_dtheta2 = width * s * (1.0 - 2.0 * z);
3369 let d3b_dtheta3 = width * s * (1.0 - 6.0 * z + 6.0 * z * z);
3370 (beta, z, db_dtheta, d2b_dtheta2, d3b_dtheta3)
3371}
3372
3373fn bounded_prior_terms(theta: f64, prior: &BoundedCoefficientPriorSpec) -> (f64, f64, f64, f64) {
3374 let (a, b) = match prior {
3375 BoundedCoefficientPriorSpec::None => return (0.0, 0.0, 0.0, 0.0),
3377 BoundedCoefficientPriorSpec::Uniform => (1.0, 1.0),
3380 BoundedCoefficientPriorSpec::Beta { a, b } => (*a, *b),
3381 };
3382 let z = stable_sigmoid(theta).clamp(1e-12, 1.0 - 1e-12);
3383 let logp = a * z.ln() + b * (1.0 - z).ln();
3384 let grad = a - (a + b) * z;
3385 let neghess = (a + b) * z * (1.0 - z);
3386 let neghess_derivative = (a + b) * z * (1.0 - z) * (1.0 - 2.0 * z);
3387 (logp, grad, neghess, neghess_derivative)
3388}
3389
3390#[inline]
3399fn glm_eta_observation_state(
3400 w: f64,
3401 lmu: f64,
3402 lmumu: f64,
3403 lmumumu: f64,
3404 var: f64,
3405 d1: f64,
3406 d2: f64,
3407 d3: f64,
3408 mu_deriv_eps: f64,
3409) -> (f64, f64, f64, f64) {
3410 let score = w * lmu * d1;
3411 let fisherweight = (w * d1 * d1 / var).max(mu_deriv_eps);
3412 let neghessian = -w * (lmumu * d1 * d1 + lmu * d2);
3413 let neghessian_deriv = -w * (lmumumu * d1 * d1 * d1 + 3.0 * lmumu * d1 * d2 + lmu * d3);
3414 (score, fisherweight, neghessian, neghessian_deriv)
3415}
3416
3417fn evaluate_standard_familyobservations(
3418 family: LikelihoodSpec,
3419 latent_cloglog_state: Option<&LatentCLogLogState>,
3420 mixture_link_state: Option<&MixtureLinkState>,
3421 sas_link_state: Option<&SasLinkState>,
3422 y: &Array1<f64>,
3423 weights: &Array1<f64>,
3424 eta: &Array1<f64>,
3425) -> Result<StandardFamilyObservationState, EstimationError> {
3426 const PROB_EPS: f64 = 1e-10;
3427 const MU_DERIV_EPS: f64 = 1e-12;
3428 let n = y.len();
3429 if weights.len() != n || eta.len() != n {
3430 crate::bail_invalid_estim!("bounded family observation size mismatch");
3431 }
3432
3433 let mut mu = Array1::<f64>::zeros(n);
3434 let mut score = Array1::<f64>::zeros(n);
3435 let mut fisherweight = Array1::<f64>::zeros(n);
3436 let mut neghessian_eta = Array1::<f64>::zeros(n);
3437 let mut neghessian_eta_derivative = Array1::<f64>::zeros(n);
3438 let mut log_likelihood = 0.0;
3439
3440 for i in 0..n {
3441 let w = weights[i].max(0.0);
3442 let yi = y[i];
3443 let eta_i = eta[i];
3444 match (&family.response, &family.link) {
3445 (ResponseFamily::Gaussian, _) => {
3446 let resid = yi - eta_i;
3447 mu[i] = eta_i;
3448 score[i] = w * resid;
3449 fisherweight[i] = w.max(MU_DERIV_EPS);
3450 neghessian_eta[i] = w;
3451 neghessian_eta_derivative[i] = 0.0;
3452 log_likelihood += -0.5 * w * resid * resid;
3453 }
3454 (ResponseFamily::Binomial, InverseLink::Standard(StandardLink::Logit)) => {
3455 let jet = logit_inverse_link_jet5(eta_i);
3456 mu[i] = jet.mu;
3457 score[i] = w * (yi - jet.mu);
3458 fisherweight[i] = jet.d1.max(MU_DERIV_EPS);
3459 neghessian_eta[i] = jet.d1;
3460 neghessian_eta_derivative[i] = jet.d2;
3461 let logmu = -gam_linalg::utils::stable_softplus(-eta_i);
3462 let log_one_minusmu = -gam_linalg::utils::stable_softplus(eta_i);
3463 log_likelihood += w * (yi * logmu + (1.0 - yi) * log_one_minusmu);
3464 }
3465 (ResponseFamily::Binomial, _) => {
3466 let inverse_link = if let Some(state) = latent_cloglog_state {
3467 Some(InverseLink::LatentCLogLog(*state))
3468 } else if let Some(state) = mixture_link_state {
3469 Some(InverseLink::Mixture(state.clone()))
3470 } else {
3471 sas_link_state.map(|state| {
3472 if family.is_binomial_beta_logistic() {
3473 InverseLink::BetaLogistic(*state)
3474 } else {
3475 InverseLink::Sas(*state)
3476 }
3477 })
3478 };
3479 let strategy_spec = LikelihoodSpec {
3480 response: family.response.clone(),
3481 link: inverse_link.clone().unwrap_or_else(|| family.link.clone()),
3482 };
3483 let jet = strategy_for_spec(&strategy_spec).inverse_link_jet(eta_i)?;
3484 let mu_i_raw = jet.mu;
3485 let dmu_deta_raw = jet.d1;
3486 let mu_i: f64 = mu_i_raw.clamp(PROB_EPS, 1.0 - PROB_EPS);
3487 let dmu_deta = dmu_deta_raw.max(MU_DERIV_EPS);
3488 let d2mu_deta2 = jet.d2;
3489 let d3mu_deta3 = jet.d3;
3490 let var = (mu_i * (1.0 - mu_i)).max(PROB_EPS);
3491 let lmu = (yi - mu_i) / var;
3492 let lmumu = -(yi / (mu_i * mu_i)) - ((1.0 - yi) / ((1.0 - mu_i) * (1.0 - mu_i)));
3493 let lmumumu = 2.0 * yi / (mu_i * mu_i * mu_i)
3494 - 2.0 * (1.0 - yi) / ((1.0 - mu_i) * (1.0 - mu_i) * (1.0 - mu_i));
3495 mu[i] = mu_i;
3496 score[i] = w * lmu * dmu_deta;
3497 fisherweight[i] = (w * dmu_deta * dmu_deta / var).max(MU_DERIV_EPS);
3498 neghessian_eta[i] = -w * (lmumu * dmu_deta * dmu_deta + lmu * d2mu_deta2);
3499 neghessian_eta_derivative[i] = -w
3500 * (lmumumu * dmu_deta * dmu_deta * dmu_deta
3501 + 3.0 * lmumu * dmu_deta * d2mu_deta2
3502 + lmu * d3mu_deta3);
3503 log_likelihood += w * (yi * mu_i.ln() + (1.0 - yi) * (1.0 - mu_i).ln());
3504 }
3505 (ResponseFamily::Poisson, _) => {
3506 let strategy_spec = LikelihoodSpec {
3509 response: family.response.clone(),
3510 link: family.link.clone(),
3511 };
3512 let jet = strategy_for_spec(&strategy_spec).inverse_link_jet(eta_i)?;
3513 let mu_i = jet.mu.max(PROB_EPS);
3514 let d1 = jet.d1.max(MU_DERIV_EPS);
3515 let var = mu_i;
3516 let lmu = yi / mu_i - 1.0;
3517 let lmumu = -yi / (mu_i * mu_i);
3518 let lmumumu = 2.0 * yi / (mu_i * mu_i * mu_i);
3519 let (s, f, nh, nhd) = glm_eta_observation_state(
3520 w, lmu, lmumu, lmumumu, var, d1, jet.d2, jet.d3, MU_DERIV_EPS,
3521 );
3522 mu[i] = mu_i;
3523 score[i] = s;
3524 fisherweight[i] = f;
3525 neghessian_eta[i] = nh;
3526 neghessian_eta_derivative[i] = nhd;
3527 log_likelihood += w * (yi * mu_i.ln() - mu_i);
3528 }
3529 (ResponseFamily::Tweedie { p }, _) => {
3530 let p = *p;
3535 let strategy_spec = LikelihoodSpec {
3536 response: family.response.clone(),
3537 link: family.link.clone(),
3538 };
3539 let jet = strategy_for_spec(&strategy_spec).inverse_link_jet(eta_i)?;
3540 let mu_i = jet.mu.max(PROB_EPS);
3541 let d1 = jet.d1.max(MU_DERIV_EPS);
3542 let var = mu_i.powf(p);
3543 let resid = yi - mu_i;
3544 let lmu = resid / var;
3545 let lmumu = -mu_i.powf(-p) - p * resid * mu_i.powf(-p - 1.0);
3546 let lmumumu =
3547 2.0 * p * mu_i.powf(-p - 1.0) + p * (p + 1.0) * resid * mu_i.powf(-p - 2.0);
3548 let (s, f, nh, nhd) = glm_eta_observation_state(
3549 w, lmu, lmumu, lmumumu, var, d1, jet.d2, jet.d3, MU_DERIV_EPS,
3550 );
3551 mu[i] = mu_i;
3552 score[i] = s;
3553 fisherweight[i] = f;
3554 neghessian_eta[i] = nh;
3555 neghessian_eta_derivative[i] = nhd;
3556 log_likelihood += w
3558 * (yi * mu_i.powf(1.0 - p) / (1.0 - p) - mu_i.powf(2.0 - p) / (2.0 - p));
3559 }
3560 (ResponseFamily::NegativeBinomial { theta, .. }, _) => {
3561 let theta = (*theta).max(PROB_EPS);
3565 let strategy_spec = LikelihoodSpec {
3566 response: family.response.clone(),
3567 link: family.link.clone(),
3568 };
3569 let jet = strategy_for_spec(&strategy_spec).inverse_link_jet(eta_i)?;
3570 let mu_i = jet.mu.max(PROB_EPS);
3571 let d1 = jet.d1.max(MU_DERIV_EPS);
3572 let mu_plus = mu_i + theta;
3573 let var = mu_i + mu_i * mu_i / theta;
3574 let lmu = yi / mu_i - (yi + theta) / mu_plus;
3575 let lmumu = -yi / (mu_i * mu_i) + (yi + theta) / (mu_plus * mu_plus);
3576 let lmumumu =
3577 2.0 * yi / (mu_i * mu_i * mu_i) - 2.0 * (yi + theta) / (mu_plus * mu_plus * mu_plus);
3578 let (s, f, nh, nhd) = glm_eta_observation_state(
3579 w, lmu, lmumu, lmumumu, var, d1, jet.d2, jet.d3, MU_DERIV_EPS,
3580 );
3581 mu[i] = mu_i;
3582 score[i] = s;
3583 fisherweight[i] = f;
3584 neghessian_eta[i] = nh;
3585 neghessian_eta_derivative[i] = nhd;
3586 log_likelihood += w * (yi * mu_i.ln() - (yi + theta) * mu_plus.ln());
3587 }
3588 (ResponseFamily::Beta { .. }, _) => {
3589 crate::bail_invalid_estim!(
3590 "bounded linear terms are not supported for BetaLogit fits"
3591 );
3592 }
3593 (ResponseFamily::Gamma, _) => {
3594 let strategy_spec = LikelihoodSpec {
3598 response: family.response.clone(),
3599 link: family.link.clone(),
3600 };
3601 let jet = strategy_for_spec(&strategy_spec).inverse_link_jet(eta_i)?;
3602 let mu_i = jet.mu.max(PROB_EPS);
3603 let d1 = jet.d1.max(MU_DERIV_EPS);
3604 let var = mu_i * mu_i;
3605 let lmu = yi / (mu_i * mu_i) - 1.0 / mu_i;
3606 let lmumu = -2.0 * yi / (mu_i * mu_i * mu_i) + 1.0 / (mu_i * mu_i);
3607 let lmumumu =
3608 6.0 * yi / (mu_i * mu_i * mu_i * mu_i) - 2.0 / (mu_i * mu_i * mu_i);
3609 let (s, f, nh, nhd) = glm_eta_observation_state(
3610 w, lmu, lmumu, lmumumu, var, d1, jet.d2, jet.d3, MU_DERIV_EPS,
3611 );
3612 mu[i] = mu_i;
3613 score[i] = s;
3614 fisherweight[i] = f;
3615 neghessian_eta[i] = nh;
3616 neghessian_eta_derivative[i] = nhd;
3617 log_likelihood += w * (-(yi / mu_i) - mu_i.ln());
3618 }
3619 (ResponseFamily::RoystonParmar, _) => {
3620 crate::bail_invalid_estim!(
3621 "bounded linear terms are not supported for survival model fits"
3622 );
3623 }
3624 }
3625 }
3626
3627 Ok(StandardFamilyObservationState {
3628 eta: eta.clone(),
3629 mu,
3630 score,
3631 fisherweight,
3632 neghessian_eta,
3633 neghessian_eta_derivative,
3634 log_likelihood,
3635 })
3636}
3637
3638#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3639enum SpatialAdaptiveHyperKind {
3640 LogLambdaMagnitude,
3641 LogLambdaGradient,
3642 LogLambdaCurvature,
3643 LogEpsilonMagnitude,
3644 LogEpsilonGradient,
3645 LogEpsilonCurvature,
3646}
3647
3648impl SpatialAdaptiveHyperKind {
3649 fn component_index(self) -> usize {
3650 match self {
3651 SpatialAdaptiveHyperKind::LogLambdaMagnitude
3652 | SpatialAdaptiveHyperKind::LogEpsilonMagnitude => 0,
3653 SpatialAdaptiveHyperKind::LogLambdaGradient
3654 | SpatialAdaptiveHyperKind::LogEpsilonGradient => 1,
3655 SpatialAdaptiveHyperKind::LogLambdaCurvature
3656 | SpatialAdaptiveHyperKind::LogEpsilonCurvature => 2,
3657 }
3658 }
3659
3660 fn is_log_lambda(self) -> bool {
3661 matches!(
3662 self,
3663 SpatialAdaptiveHyperKind::LogLambdaMagnitude
3664 | SpatialAdaptiveHyperKind::LogLambdaGradient
3665 | SpatialAdaptiveHyperKind::LogLambdaCurvature
3666 )
3667 }
3668
3669 fn is_log_epsilon(self) -> bool {
3670 matches!(
3671 self,
3672 SpatialAdaptiveHyperKind::LogEpsilonMagnitude
3673 | SpatialAdaptiveHyperKind::LogEpsilonGradient
3674 | SpatialAdaptiveHyperKind::LogEpsilonCurvature
3675 )
3676 }
3677}
3678
3679#[derive(Clone, Copy, Debug)]
3680struct SpatialAdaptiveHyperSpec {
3681 cache_index: usize,
3682 kind: SpatialAdaptiveHyperKind,
3683}
3684
3685#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3686enum SpatialAdaptiveExplicitSecondOrderKind {
3687 StructuralZero,
3688 LocalAlphaAlpha,
3689 LocalAlphaEta,
3690 SharedEtaEta,
3691}
3692
3693#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3698enum AdaptiveComponent {
3699 Magnitude,
3700 Gradient,
3701 Curvature,
3702}
3703
3704impl AdaptiveComponent {
3705 fn from_index(index: usize) -> Result<Self, String> {
3706 match index {
3707 0 => Ok(AdaptiveComponent::Magnitude),
3708 1 => Ok(AdaptiveComponent::Gradient),
3709 2 => Ok(AdaptiveComponent::Curvature),
3710 other => Err(SmoothError::invalid_index(format!(
3711 "invalid adaptive component index {}",
3712 other
3713 ))
3714 .into()),
3715 }
3716 }
3717}
3718
3719#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3724enum HyperDerivativeKind {
3725 Rho,
3727 LogEpsilonFirst,
3729 LogEpsilonSecond,
3731}
3732
3733#[derive(Clone, Copy, Debug, PartialEq, Eq)]
3738enum HyperDriftKind {
3739 Rho,
3740 LogEpsilon,
3741}
3742
3743impl SpatialAdaptiveHyperSpec {
3744 fn component_index(self) -> usize {
3745 self.kind.component_index()
3746 }
3747
3748 fn explicit_second_order_kind(self, other: Self) -> SpatialAdaptiveExplicitSecondOrderKind {
3749 if self.component_index() != other.component_index() {
3750 return SpatialAdaptiveExplicitSecondOrderKind::StructuralZero;
3751 }
3752 match (
3753 self.kind.is_log_lambda(),
3754 other.kind.is_log_lambda(),
3755 self.kind.is_log_epsilon(),
3756 other.kind.is_log_epsilon(),
3757 ) {
3758 (true, true, false, false) if self.cache_index == other.cache_index => {
3759 SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaAlpha
3760 }
3761 (true, false, false, true) | (false, true, true, false) => {
3762 SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaEta
3763 }
3764 (false, false, true, true) => SpatialAdaptiveExplicitSecondOrderKind::SharedEtaEta,
3765 _ => SpatialAdaptiveExplicitSecondOrderKind::StructuralZero,
3766 }
3767 }
3768}
3769
3770#[derive(Clone, Debug)]
3771struct SpatialAdaptiveTermHyperParams {
3772 lambda: [f64; 3],
3773 epsilon: [f64; 3],
3774}
3775
3776#[derive(Clone)]
3777struct SpatialAdaptiveExactEvaluation {
3778 obs: StandardFamilyObservationState,
3779 adaptive_states: Vec<SpatialPenaltyExactState>,
3780 adaptive_penalty_value: f64,
3781 adaptive_penaltygradient: Array1<f64>,
3782 adaptive_penaltyhessian: Array2<f64>,
3783 fixed_quadraticvalue: f64,
3784 fixed_quadraticgradient: Array1<f64>,
3785 fixed_quadratichessian: Array2<f64>,
3786}
3787
3788#[derive(Clone)]
3789struct CachedSpatialAdaptiveExactEvaluation {
3790 beta: Array1<f64>,
3791 eval: Arc<SpatialAdaptiveExactEvaluation>,
3792}
3793
3794impl SpatialAdaptiveExactEvaluation {
3795 fn total_penalty_value(&self) -> f64 {
3796 self.adaptive_penalty_value + self.fixed_quadraticvalue
3797 }
3798
3799 fn total_penaltygradient(&self) -> Array1<f64> {
3800 &self.adaptive_penaltygradient + &self.fixed_quadraticgradient
3801 }
3802
3803 fn total_penaltyhessian(&self) -> Array2<f64> {
3804 &self.adaptive_penaltyhessian + &self.fixed_quadratichessian
3805 }
3806
3807 fn totalobjectivehessian(&self, design: &Array2<f64>) -> Result<Array2<f64>, String> {
3808 let mut out = xt_diag_x_dense(design.view(), self.obs.neghessian_eta.view())?;
3809 out += &self.total_penaltyhessian();
3810 Ok(out)
3811 }
3812}
3813
3814#[derive(Clone)]
3815struct SpatialAdaptiveExactFamily {
3816 family: LikelihoodSpec,
3817 latent_cloglog_state: Option<LatentCLogLogState>,
3818 mixture_link_state: Option<MixtureLinkState>,
3819 sas_link_state: Option<SasLinkState>,
3820 y: Arc<Array1<f64>>,
3821 weights: Arc<Array1<f64>>,
3822 design: Arc<Array2<f64>>,
3823 offset: Arc<Array1<f64>>,
3824 linear_constraints: Option<LinearInequalityConstraints>,
3825 runtime_caches: Arc<Vec<SpatialOperatorRuntimeCache>>,
3826 adaptive_params: Vec<SpatialAdaptiveTermHyperParams>,
3827 fixed_quadratichessian: Arc<Array2<f64>>,
3828 hyperspecs: Arc<Vec<SpatialAdaptiveHyperSpec>>,
3829 exact_eval_cache: Arc<Mutex<Option<CachedSpatialAdaptiveExactEvaluation>>>,
3830}
3831
3832impl SpatialAdaptiveExactFamily {
3833 fn with_adaptive_params(
3834 &self,
3835 adaptive_params: Vec<SpatialAdaptiveTermHyperParams>,
3836 fixed_quadratichessian: Arc<Array2<f64>>,
3837 ) -> Self {
3838 Self {
3839 family: self.family.clone(),
3840 latent_cloglog_state: self.latent_cloglog_state,
3841 mixture_link_state: self.mixture_link_state.clone(),
3842 sas_link_state: self.sas_link_state,
3843 y: self.y.clone(),
3844 weights: self.weights.clone(),
3845 design: self.design.clone(),
3846 offset: self.offset.clone(),
3847 linear_constraints: self.linear_constraints.clone(),
3848 runtime_caches: self.runtime_caches.clone(),
3849 adaptive_params,
3850 fixed_quadratichessian,
3851 hyperspecs: self.hyperspecs.clone(),
3852 exact_eval_cache: Arc::new(Mutex::new(None)),
3853 }
3854 }
3855
3856 fn total_eta(&self, beta: &Array1<f64>) -> Array1<f64> {
3857 gam_linalg::faer_ndarray::fast_av(self.design.as_ref(), beta) + self.offset.as_ref()
3858 }
3859
3860 fn fixed_quadratic_terms(&self, beta: &Array1<f64>) -> (f64, Array1<f64>) {
3861 let grad = self.fixed_quadratichessian.dot(beta);
3862 let value = 0.5 * beta.dot(&grad);
3863 (value, grad)
3864 }
3865
3866 fn adaptive_penalty_value_only(&self, beta: &Array1<f64>) -> Result<f64, String> {
3867 let mut penalty_value = 0.0;
3868 for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
3869 let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
3870 format!(
3871 "missing adaptive parameter block for cache {}",
3872 cache.termname
3873 )
3874 })?;
3875 let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
3876 let state =
3877 SpatialPenaltyExactState::from_beta_local(beta_local, cache, params.epsilon)
3878 .map_err(|e| e.to_string())?;
3879 penalty_value += params.lambda[0] * state.magnitude.penalty_value();
3880 penalty_value += params.lambda[1] * state.gradient.penalty_value();
3881 penalty_value += params.lambda[2] * state.curvature.penalty_value();
3882 }
3883 Ok(penalty_value)
3884 }
3885
3886 fn zero_hyper_parts(&self) -> (Array1<f64>, Array2<f64>) {
3887 let total_dim = self.design.ncols();
3888 (
3889 Array1::<f64>::zeros(total_dim),
3890 Array2::<f64>::zeros((total_dim, total_dim)),
3891 )
3892 }
3893
3894 fn embed_local_hyper_parts(
3895 &self,
3896 coeff_range: &Range<usize>,
3897 local_grad: &Array1<f64>,
3898 local_hess: &Array2<f64>,
3899 ) -> (Array1<f64>, Array2<f64>) {
3900 let (mut beta_mixed, mut betahessian) = self.zero_hyper_parts();
3901 beta_mixed
3902 .slice_mut(s![coeff_range.clone()])
3903 .assign(local_grad);
3904 betahessian
3905 .slice_mut(s![coeff_range.clone(), coeff_range.clone()])
3906 .assign(local_hess);
3907 (beta_mixed, betahessian)
3908 }
3909
3910 fn embed_local_hyper_hessian(
3911 &self,
3912 coeff_range: &Range<usize>,
3913 local_hess: &Array2<f64>,
3914 ) -> Array2<f64> {
3915 let total_dim = self.design.ncols();
3916 let mut out = Array2::<f64>::zeros((total_dim, total_dim));
3917 out.slice_mut(s![coeff_range.clone(), coeff_range.clone()])
3918 .assign(local_hess);
3919 out
3920 }
3921
3922 fn adaptive_block_eval(
3931 &self,
3932 eval: &SpatialAdaptiveExactEvaluation,
3933 cache_idx: usize,
3934 component: AdaptiveComponent,
3935 derivative: HyperDerivativeKind,
3936 ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
3937 let cache = self
3938 .runtime_caches
3939 .get(cache_idx)
3940 .ok_or_else(|| format!("adaptive cache index {} out of bounds", cache_idx))?;
3941 let params = self
3942 .adaptive_params
3943 .get(cache_idx)
3944 .ok_or_else(|| format!("adaptive hyperparameter block {} out of bounds", cache_idx))?;
3945 let state = eval
3946 .adaptive_states
3947 .get(cache_idx)
3948 .ok_or_else(|| format!("adaptive exact state index {} out of bounds", cache_idx))?;
3949
3950 let (objective_local, beta_mixed_local, betahessian_local) = match component {
3951 AdaptiveComponent::Magnitude => {
3952 let lambda = params.lambda[0];
3953 let mag = &state.magnitude;
3954 let (objective, gradient_coeff, hessian_diag) = match derivative {
3955 HyperDerivativeKind::Rho => (
3956 mag.penalty_value(),
3957 mag.betagradient_coeff(),
3958 mag.betahessian_diag(),
3959 ),
3960 HyperDerivativeKind::LogEpsilonFirst => (
3961 mag.log_epsilon_gradient_terms().sum(),
3962 mag.log_epsilon_betagradient_coeff(),
3963 mag.log_epsilon_betahessian_diag(),
3964 ),
3965 HyperDerivativeKind::LogEpsilonSecond => (
3966 mag.log_epsilon_hessian_terms().sum(),
3967 mag.log_epsilon_beta_mixed_second_coeff(),
3968 mag.log_epsilon_betahessian_second_diag(),
3969 ),
3970 };
3971 (
3972 lambda * objective,
3973 lambda * scalar_operatorgradient(&cache.d0, &gradient_coeff),
3974 lambda * scalar_operatorhessian(&cache.d0, &hessian_diag),
3975 )
3976 }
3977 AdaptiveComponent::Gradient => {
3978 let lambda = params.lambda[1];
3979 let grad = &state.gradient;
3980 let (objective, gradient_blocks, hessian_blocks) = match derivative {
3981 HyperDerivativeKind::Rho => (
3982 grad.penalty_value(),
3983 grad.betagradient_blocks(),
3984 grad.betahessian_blocks(),
3985 ),
3986 HyperDerivativeKind::LogEpsilonFirst => (
3987 grad.log_epsilon_gradient_terms().sum(),
3988 grad.log_epsilon_betagradient_blocks(),
3989 grad.log_epsilon_betahessian_blocks(),
3990 ),
3991 HyperDerivativeKind::LogEpsilonSecond => (
3992 grad.log_epsilon_hessian_terms().sum(),
3993 grad.log_epsilon_beta_mixed_second_blocks(),
3994 grad.log_epsilon_betahessian_second_blocks(),
3995 ),
3996 };
3997 (
3998 lambda * objective,
3999 lambda
4000 * grouped_operatorgradient(&cache.d1, cache.dimension, &gradient_blocks)
4001 .map_err(|e| e.to_string())?,
4002 lambda
4003 * grouped_operatorhessian(&cache.d1, cache.dimension, &hessian_blocks)
4004 .map_err(|e| e.to_string())?,
4005 )
4006 }
4007 AdaptiveComponent::Curvature => {
4008 let lambda = params.lambda[2];
4009 let group = cache.dimension * cache.dimension;
4010 let curv = &state.curvature;
4011 let (objective, gradient_blocks, hessian_blocks) = match derivative {
4012 HyperDerivativeKind::Rho => (
4013 curv.penalty_value(),
4014 curv.betagradient_blocks(),
4015 curv.betahessian_blocks(),
4016 ),
4017 HyperDerivativeKind::LogEpsilonFirst => (
4018 curv.log_epsilon_gradient_terms().sum(),
4019 curv.log_epsilon_betagradient_blocks(),
4020 curv.log_epsilon_betahessian_blocks(),
4021 ),
4022 HyperDerivativeKind::LogEpsilonSecond => (
4023 curv.log_epsilon_hessian_terms().sum(),
4024 curv.log_epsilon_beta_mixed_second_blocks(),
4025 curv.log_epsilon_betahessian_second_blocks(),
4026 ),
4027 };
4028 (
4029 lambda * objective,
4030 lambda
4031 * grouped_operatorgradient(&cache.d2, group, &gradient_blocks)
4032 .map_err(|e| e.to_string())?,
4033 lambda
4034 * grouped_operatorhessian(&cache.d2, group, &hessian_blocks)
4035 .map_err(|e| e.to_string())?,
4036 )
4037 }
4038 };
4039
4040 let (beta_mixed, betahessian) = self.embed_local_hyper_parts(
4041 &cache.coeff_global_range,
4042 &beta_mixed_local,
4043 &betahessian_local,
4044 );
4045 Ok((objective_local, beta_mixed, betahessian))
4046 }
4047
4048 fn adaptive_shared_log_epsilon_parts(
4049 &self,
4050 eval: &SpatialAdaptiveExactEvaluation,
4051 component: usize,
4052 ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
4053 self.adaptive_shared_block_eval(eval, component, HyperDerivativeKind::LogEpsilonFirst)
4059 }
4060
4061 fn adaptive_shared_log_epsilon_second_parts(
4062 &self,
4063 eval: &SpatialAdaptiveExactEvaluation,
4064 component: usize,
4065 ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
4066 self.adaptive_shared_block_eval(eval, component, HyperDerivativeKind::LogEpsilonSecond)
4072 }
4073
4074 fn adaptive_shared_block_eval(
4079 &self,
4080 eval: &SpatialAdaptiveExactEvaluation,
4081 component: usize,
4082 derivative: HyperDerivativeKind,
4083 ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
4084 let component = AdaptiveComponent::from_index(component)?;
4085 let (mut score, mut hessian) = self.zero_hyper_parts();
4086 let mut objective = 0.0;
4087 for cache_idx in 0..self.runtime_caches.len() {
4088 let (local_objective, local_score, local_hessian) =
4089 self.adaptive_block_eval(eval, cache_idx, component, derivative)?;
4090 objective += local_objective;
4091 score += &local_score;
4092 hessian += &local_hessian;
4093 }
4094 Ok((objective, score, hessian))
4095 }
4096
4097 fn adaptive_shared_log_epsilon_drift(
4098 &self,
4099 eval: &SpatialAdaptiveExactEvaluation,
4100 component: usize,
4101 direction: &Array1<f64>,
4102 ) -> Result<Array2<f64>, String> {
4103 let component = AdaptiveComponent::from_index(component)?;
4107 let total_dim = self.design.ncols();
4108 let mut total = Array2::<f64>::zeros((total_dim, total_dim));
4109 for cache_idx in 0..self.runtime_caches.len() {
4110 total += &self.adaptive_block_drift_eval(
4111 eval,
4112 cache_idx,
4113 component,
4114 HyperDriftKind::LogEpsilon,
4115 direction,
4116 )?;
4117 }
4118 Ok(total)
4119 }
4120
4121 fn adaptive_explicit_second_order_parts(
4122 &self,
4123 eval: &SpatialAdaptiveExactEvaluation,
4124 left: SpatialAdaptiveHyperSpec,
4125 right: SpatialAdaptiveHyperSpec,
4126 ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
4127 match left.explicit_second_order_kind(right) {
4136 SpatialAdaptiveExplicitSecondOrderKind::StructuralZero => {
4137 let (score, hessian) = self.zero_hyper_parts();
4138 Ok((0.0, score, hessian))
4139 }
4140 SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaAlpha => self.adaptive_block_eval(
4141 eval,
4142 left.cache_index,
4143 AdaptiveComponent::from_index(left.component_index())?,
4144 HyperDerivativeKind::Rho,
4145 ),
4146 SpatialAdaptiveExplicitSecondOrderKind::LocalAlphaEta => {
4147 let local_alpha = if left.kind.is_log_lambda() {
4148 left
4149 } else {
4150 right
4151 };
4152 self.adaptive_block_eval(
4153 eval,
4154 local_alpha.cache_index,
4155 AdaptiveComponent::from_index(local_alpha.component_index())?,
4156 HyperDerivativeKind::LogEpsilonFirst,
4157 )
4158 }
4159 SpatialAdaptiveExplicitSecondOrderKind::SharedEtaEta => {
4160 self.adaptive_shared_log_epsilon_second_parts(eval, left.component_index())
4161 }
4162 }
4163 }
4164
4165 fn adaptive_block_drift_eval(
4173 &self,
4174 eval: &SpatialAdaptiveExactEvaluation,
4175 cache_idx: usize,
4176 component: AdaptiveComponent,
4177 drift: HyperDriftKind,
4178 direction: &Array1<f64>,
4179 ) -> Result<Array2<f64>, String> {
4180 let cache = self
4181 .runtime_caches
4182 .get(cache_idx)
4183 .ok_or_else(|| format!("adaptive cache index {} out of bounds", cache_idx))?;
4184 let params = self
4185 .adaptive_params
4186 .get(cache_idx)
4187 .ok_or_else(|| format!("adaptive hyperparameter block {} out of bounds", cache_idx))?;
4188 let state = eval
4189 .adaptive_states
4190 .get(cache_idx)
4191 .ok_or_else(|| format!("adaptive exact state index {} out of bounds", cache_idx))?;
4192 let direction_local = direction.slice(s![cache.coeff_global_range.clone()]);
4193
4194 let local_hessian = match component {
4195 AdaptiveComponent::Magnitude => {
4196 let d0_u = cache.d0.dot(&direction_local);
4197 let mag = &state.magnitude;
4198 let diag = match drift {
4199 HyperDriftKind::Rho => mag.directionalhessian_diag(&d0_u),
4200 HyperDriftKind::LogEpsilon => {
4201 mag.log_epsilon_betahessian_directional_diag(&d0_u)
4202 }
4203 };
4204 params.lambda[0] * scalar_operatorhessian(&cache.d0, &diag)
4205 }
4206 AdaptiveComponent::Gradient => {
4207 let d1_u = cache.d1.dot(&direction_local);
4208 let direction_blocks = collocationgradient_blocks(&d1_u, cache.dimension)
4209 .map_err(|e| e.to_string())?;
4210 let grad = &state.gradient;
4211 let blocks = match drift {
4212 HyperDriftKind::Rho => grad.directionalhessian_blocks(&direction_blocks),
4213 HyperDriftKind::LogEpsilon => {
4214 grad.log_epsilon_betahessian_directional_blocks(&direction_blocks)
4215 }
4216 };
4217 params.lambda[1]
4218 * grouped_operatorhessian(&cache.d1, cache.dimension, &blocks)
4219 .map_err(|e| e.to_string())?
4220 }
4221 AdaptiveComponent::Curvature => {
4222 let group = cache.dimension * cache.dimension;
4223 let d2_u = cache.d2.dot(&direction_local);
4224 let direction_blocks =
4225 collocationhessian_blocks(&d2_u, cache.dimension).map_err(|e| e.to_string())?;
4226 let curv = &state.curvature;
4227 let blocks = match drift {
4228 HyperDriftKind::Rho => curv.directionalhessian_blocks(&direction_blocks),
4229 HyperDriftKind::LogEpsilon => {
4230 curv.log_epsilon_betahessian_directional_blocks(&direction_blocks)
4231 }
4232 };
4233 params.lambda[2]
4234 * grouped_operatorhessian(&cache.d2, group, &blocks)
4235 .map_err(|e| e.to_string())?
4236 }
4237 };
4238
4239 Ok(self.embed_local_hyper_hessian(&cache.coeff_global_range, &local_hessian))
4240 }
4241
4242 fn adaptive_hyper_parts(
4243 &self,
4244 eval: &SpatialAdaptiveExactEvaluation,
4245 hyper: SpatialAdaptiveHyperSpec,
4246 ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
4247 match hyper.kind {
4248 SpatialAdaptiveHyperKind::LogLambdaMagnitude
4251 | SpatialAdaptiveHyperKind::LogLambdaGradient
4252 | SpatialAdaptiveHyperKind::LogLambdaCurvature => self.adaptive_block_eval(
4253 eval,
4254 hyper.cache_index,
4255 AdaptiveComponent::from_index(hyper.component_index())?,
4256 HyperDerivativeKind::Rho,
4257 ),
4258 SpatialAdaptiveHyperKind::LogEpsilonMagnitude
4260 | SpatialAdaptiveHyperKind::LogEpsilonGradient
4261 | SpatialAdaptiveHyperKind::LogEpsilonCurvature => {
4262 self.adaptive_shared_log_epsilon_parts(eval, hyper.component_index())
4263 }
4264 }
4265 }
4266
4267 fn exact_evaluation_uncached(
4268 &self,
4269 beta: &Array1<f64>,
4270 ) -> Result<SpatialAdaptiveExactEvaluation, String> {
4271 let eta = self.total_eta(beta);
4272 let obs = evaluate_standard_familyobservations(
4273 self.family.clone(),
4274 self.latent_cloglog_state.as_ref(),
4275 self.mixture_link_state.as_ref(),
4276 self.sas_link_state.as_ref(),
4277 &self.y,
4278 &self.weights,
4279 &eta,
4280 )
4281 .map_err(|e| e.to_string())?;
4282 let p = beta.len();
4283 let mut penalty_value = 0.0;
4284 let mut penaltygradient = Array1::<f64>::zeros(p);
4285 let mut penaltyhessian = Array2::<f64>::zeros((p, p));
4286 let mut adaptive_states = Vec::with_capacity(self.runtime_caches.len());
4287
4288 for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
4289 let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
4290 format!(
4291 "missing adaptive parameter block for cache {}",
4292 cache.termname
4293 )
4294 })?;
4295 let beta_local = beta.slice(s![cache.coeff_global_range.clone()]);
4296 let state =
4297 SpatialPenaltyExactState::from_beta_local(beta_local, cache, params.epsilon)
4298 .map_err(|e| e.to_string())?;
4299
4300 let g0 = scalar_operatorgradient(&cache.d0, &state.magnitude.betagradient_coeff());
4301 let gg = grouped_operatorgradient(
4302 &cache.d1,
4303 cache.dimension,
4304 &state.gradient.betagradient_blocks(),
4305 )
4306 .map_err(|e| e.to_string())?;
4307 let gc = grouped_operatorgradient(
4308 &cache.d2,
4309 cache.dimension * cache.dimension,
4310 &state.curvature.betagradient_blocks(),
4311 )
4312 .map_err(|e| e.to_string())?;
4313 let h0 = scalar_operatorhessian(&cache.d0, &state.magnitude.betahessian_diag());
4314 let hg = grouped_operatorhessian(
4315 &cache.d1,
4316 cache.dimension,
4317 &state.gradient.betahessian_blocks(),
4318 )
4319 .map_err(|e| e.to_string())?;
4320 let hc = grouped_operatorhessian(
4321 &cache.d2,
4322 cache.dimension * cache.dimension,
4323 &state.curvature.betahessian_blocks(),
4324 )
4325 .map_err(|e| e.to_string())?;
4326
4327 let lambda0 = params.lambda[0];
4328 let lambdag = params.lambda[1];
4329 let lambdac = params.lambda[2];
4330
4331 penalty_value += lambda0 * state.magnitude.penalty_value();
4332 penalty_value += lambdag * state.gradient.penalty_value();
4333 penalty_value += lambdac * state.curvature.penalty_value();
4334
4335 let range = cache.coeff_global_range.clone();
4336 {
4337 let mut grad_local = penaltygradient.slice_mut(s![range.clone()]);
4338 grad_local += &(g0.mapv(|v| lambda0 * v));
4339 grad_local += &(gg.mapv(|v| lambdag * v));
4340 grad_local += &(gc.mapv(|v| lambdac * v));
4341 }
4342 {
4343 let mut h_local = penaltyhessian.slice_mut(s![range.clone(), range]);
4344 h_local += &h0.mapv(|v| lambda0 * v);
4345 h_local += &hg.mapv(|v| lambdag * v);
4346 h_local += &hc.mapv(|v| lambdac * v);
4347 }
4348
4349 adaptive_states.push(state);
4350 }
4351
4352 let (fixed_quadraticvalue, fixed_quadraticgradient) = self.fixed_quadratic_terms(beta);
4353 Ok(SpatialAdaptiveExactEvaluation {
4354 obs,
4355 adaptive_states,
4356 adaptive_penalty_value: penalty_value,
4357 adaptive_penaltygradient: penaltygradient,
4358 adaptive_penaltyhessian: penaltyhessian,
4359 fixed_quadraticvalue,
4360 fixed_quadraticgradient,
4361 fixed_quadratichessian: self.fixed_quadratichessian.as_ref().clone(),
4362 })
4363 }
4364
4365 fn exact_evaluation(
4366 &self,
4367 beta: &Array1<f64>,
4368 ) -> Result<Arc<SpatialAdaptiveExactEvaluation>, String> {
4369 {
4370 let cache = self
4371 .exact_eval_cache
4372 .lock()
4373 .map_err(|_| "spatial adaptive exact-evaluation cache lock poisoned".to_string())?;
4374 if let Some(cached) = cache.as_ref()
4375 && cached.beta.len() == beta.len()
4376 && cached
4377 .beta
4378 .iter()
4379 .zip(beta.iter())
4380 .all(|(&left, &right)| left == right)
4381 {
4382 return Ok(Arc::clone(&cached.eval));
4383 }
4384 }
4385
4386 let eval = Arc::new(self.exact_evaluation_uncached(beta)?);
4387 let mut cache = self
4388 .exact_eval_cache
4389 .lock()
4390 .map_err(|_| "spatial adaptive exact-evaluation cache lock poisoned".to_string())?;
4391 *cache = Some(CachedSpatialAdaptiveExactEvaluation {
4392 beta: beta.clone(),
4393 eval: Arc::clone(&eval),
4394 });
4395 Ok(eval)
4396 }
4397
4398 fn exacthessian_directional_derivative_from_evaluation(
4399 &self,
4400 beta: &Array1<f64>,
4401 eval: &SpatialAdaptiveExactEvaluation,
4402 direction: &Array1<f64>,
4403 ) -> Result<Array2<f64>, String> {
4404 assert_eq!(
4405 beta.len(),
4406 direction.len(),
4407 "beta/direction length mismatch",
4408 );
4409 let d_eta = gam_linalg::faer_ndarray::fast_av(self.design.as_ref(), direction);
4410 let mut total = xt_diag_x_dense(
4411 self.design.view(),
4412 (&eval.obs.neghessian_eta_derivative * &d_eta).view(),
4413 )?;
4414 for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
4415 let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
4416 format!(
4417 "missing adaptive parameter block for cache {}",
4418 cache.termname
4419 )
4420 })?;
4421 let state = eval
4422 .adaptive_states
4423 .get(cache_idx)
4424 .ok_or_else(|| format!("missing adaptive state for cache {}", cache.termname))?;
4425 let direction_local = direction.slice(s![cache.coeff_global_range.clone()]);
4426 let d0_u = cache.d0.dot(&direction_local);
4427 let d1_u = cache.d1.dot(&direction_local);
4428 let d2_u = cache.d2.dot(&direction_local);
4429 let h0 =
4430 scalar_operatorhessian(&cache.d0, &state.magnitude.directionalhessian_diag(&d0_u))
4431 .mapv(|v| params.lambda[0] * v);
4432 let hg = grouped_operatorhessian(
4433 &cache.d1,
4434 cache.dimension,
4435 &state.gradient.directionalhessian_blocks(
4436 &collocationgradient_blocks(&d1_u, cache.dimension)
4437 .map_err(|e| e.to_string())?,
4438 ),
4439 )
4440 .map_err(|e| e.to_string())?
4441 .mapv(|v| params.lambda[1] * v);
4442 let hc = grouped_operatorhessian(
4443 &cache.d2,
4444 cache.dimension * cache.dimension,
4445 &state.curvature.directionalhessian_blocks(
4446 &collocationhessian_blocks(&d2_u, cache.dimension)
4447 .map_err(|e| e.to_string())?,
4448 ),
4449 )
4450 .map_err(|e| e.to_string())?
4451 .mapv(|v| params.lambda[2] * v);
4452 let range = cache.coeff_global_range.clone();
4453 let mut local = total.slice_mut(s![range.clone(), range]);
4454 local += &h0;
4455 local += &hg;
4456 local += &hc;
4457 }
4458 Ok(total)
4459 }
4460
4461 fn exacthessian_second_directional_derivative_from_evaluation(
4482 &self,
4483 eval: &SpatialAdaptiveExactEvaluation,
4484 direction_u: &Array1<f64>,
4485 direction_v: &Array1<f64>,
4486 ) -> Result<Option<Array2<f64>>, String> {
4487 let p = self.design.ncols();
4488 if eval.obs.neghessian_eta_derivative.iter().any(|&w| w != 0.0) {
4490 return Ok(None);
4491 }
4492 let mut total = Array2::<f64>::zeros((p, p));
4493 for (cache_idx, cache) in self.runtime_caches.iter().enumerate() {
4494 let params = self.adaptive_params.get(cache_idx).ok_or_else(|| {
4495 format!(
4496 "missing adaptive parameter block for cache {}",
4497 cache.termname
4498 )
4499 })?;
4500 let state = eval
4501 .adaptive_states
4502 .get(cache_idx)
4503 .ok_or_else(|| format!("missing adaptive state for cache {}", cache.termname))?;
4504 let u_local = direction_u.slice(s![cache.coeff_global_range.clone()]);
4505 let v_local = direction_v.slice(s![cache.coeff_global_range.clone()]);
4506
4507 let q0_u = cache.d0.dot(&u_local);
4509 let q0_v = cache.d0.dot(&v_local);
4510 let h0 = scalar_operatorhessian(
4511 &cache.d0,
4512 &state.magnitude.second_directionalhessian_diag(&q0_u, &q0_v),
4513 )
4514 .mapv(|x| params.lambda[0] * x);
4515
4516 let a1 = collocationgradient_blocks(&cache.d1.dot(&u_local), cache.dimension)
4518 .map_err(|e| e.to_string())?;
4519 let b1 = collocationgradient_blocks(&cache.d1.dot(&v_local), cache.dimension)
4520 .map_err(|e| e.to_string())?;
4521 let hg = grouped_operatorhessian(
4522 &cache.d1,
4523 cache.dimension,
4524 &state.gradient.second_directionalhessian_blocks(&a1, &b1),
4525 )
4526 .map_err(|e| e.to_string())?
4527 .mapv(|x| params.lambda[1] * x);
4528
4529 let a2 = collocationhessian_blocks(&cache.d2.dot(&u_local), cache.dimension)
4531 .map_err(|e| e.to_string())?;
4532 let b2 = collocationhessian_blocks(&cache.d2.dot(&v_local), cache.dimension)
4533 .map_err(|e| e.to_string())?;
4534 let hc = grouped_operatorhessian(
4535 &cache.d2,
4536 cache.dimension * cache.dimension,
4537 &state.curvature.second_directionalhessian_blocks(&a2, &b2),
4538 )
4539 .map_err(|e| e.to_string())?
4540 .mapv(|x| params.lambda[2] * x);
4541
4542 let range = cache.coeff_global_range.clone();
4543 let mut local = total.slice_mut(s![range.clone(), range]);
4544 local += &h0;
4545 local += &hg;
4546 local += &hc;
4547 }
4548 Ok(Some(total))
4549 }
4550}
4551
4552impl CustomFamily for SpatialAdaptiveExactFamily {
4553 fn joint_jeffreys_term_required(&self) -> bool {
4557 true
4558 }
4559
4560 fn joint_jeffreys_information_with_specs(
4597 &self,
4598 block_states: &[ParameterBlockState],
4599 specs: &[ParameterBlockSpec],
4600 ) -> Result<Option<Array2<f64>>, String> {
4601 let spec = expect_single_blockspec(specs, "spatial adaptive exact family")?;
4602 let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
4603 if spec.design.ncols() != beta.len() {
4604 return Err(SmoothError::dimension_mismatch(format!(
4605 "spatial adaptive Jeffreys information: spec design has {} columns, beta has {}",
4606 spec.design.ncols(),
4607 beta.len()
4608 ))
4609 .into());
4610 }
4611 let eval = self.exact_evaluation(beta)?;
4612 Ok(Some(xt_diag_x_dense(
4613 self.design.view(),
4614 eval.obs.neghessian_eta.view(),
4615 )?))
4616 }
4617
4618 fn joint_jeffreys_information_directional_derivative_with_specs(
4619 &self,
4620 block_states: &[ParameterBlockState],
4621 specs: &[ParameterBlockSpec],
4622 d_beta_flat: &Array1<f64>,
4623 ) -> Result<Option<Array2<f64>>, String> {
4624 let spec = expect_single_blockspec(specs, "spatial adaptive exact family")?;
4630 let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
4631 if spec.design.ncols() != d_beta_flat.len() {
4632 return Err(SmoothError::dimension_mismatch(format!(
4633 "spatial adaptive Jeffreys directional derivative: spec design has {} columns, direction has {}",
4634 spec.design.ncols(),
4635 d_beta_flat.len()
4636 ))
4637 .into());
4638 }
4639 let eval = self.exact_evaluation(beta)?;
4640 let d_eta = gam_linalg::faer_ndarray::fast_av(self.design.as_ref(), d_beta_flat);
4641 Ok(Some(xt_diag_x_dense(
4642 self.design.view(),
4643 (&eval.obs.neghessian_eta_derivative * &d_eta).view(),
4644 )?))
4645 }
4646
4647 fn joint_jeffreys_information_second_directional_derivative_with_specs(
4648 &self,
4649 block_states: &[ParameterBlockState],
4650 specs: &[ParameterBlockSpec],
4651 d_beta_u_flat: &Array1<f64>,
4652 d_betav_flat: &Array1<f64>,
4653 ) -> Result<Option<Array2<f64>>, String> {
4654 let spec = expect_single_blockspec(specs, "spatial adaptive exact family")?;
4661 let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
4662 if spec.design.ncols() != beta.len()
4663 || d_beta_u_flat.len() != beta.len()
4664 || d_betav_flat.len() != beta.len()
4665 {
4666 return Err(SmoothError::dimension_mismatch(format!(
4667 "spatial adaptive Jeffreys second-direction length mismatch: spec cols={}, dirs=({}, {}), expected {}",
4668 spec.design.ncols(),
4669 d_beta_u_flat.len(),
4670 d_betav_flat.len(),
4671 beta.len()
4672 ))
4673 .into());
4674 }
4675 let eval = self.exact_evaluation(beta)?;
4676 if eval.obs.neghessian_eta_derivative.iter().any(|&w| w != 0.0) {
4677 return Ok(None);
4678 }
4679 Ok(Some(Array2::<f64>::zeros((beta.len(), beta.len()))))
4680 }
4681
4682 fn joint_jeffreys_information_matches_observed_hessian(&self) -> bool {
4683 false
4688 }
4689
4690 fn joint_jeffreys_information_depends_on_psi(&self) -> bool {
4691 false
4700 }
4701
4702 fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
4703 let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
4704 let eval = self.exact_evaluation(beta)?;
4705 let mut gradient = fast_atv(&self.design, &eval.obs.score);
4706 gradient -= &eval.total_penaltygradient();
4707 let mut hessian = xt_diag_x_dense(self.design.view(), eval.obs.neghessian_eta.view())?;
4708 hessian += &eval.total_penaltyhessian();
4709 Ok(FamilyEvaluation {
4710 log_likelihood: eval.obs.log_likelihood - eval.total_penalty_value(),
4711 blockworking_sets: vec![BlockWorkingSet::ExactNewton {
4712 gradient,
4713 hessian: SymmetricMatrix::Dense(hessian),
4714 }],
4715 })
4716 }
4717
4718 fn log_likelihood_only(&self, block_states: &[ParameterBlockState]) -> Result<f64, String> {
4719 let state = expect_single_block_state(block_states, "spatial adaptive exact family")?;
4720 let beta = &state.beta;
4721 let obs = evaluate_standard_familyobservations(
4722 self.family.clone(),
4723 self.latent_cloglog_state.as_ref(),
4724 self.mixture_link_state.as_ref(),
4725 self.sas_link_state.as_ref(),
4726 &self.y,
4727 &self.weights,
4728 &state.eta,
4729 )
4730 .map_err(|e| e.to_string())?;
4731 let adaptive_penalty = self.adaptive_penalty_value_only(beta)?;
4732 let (fixed_quadratic, _) = self.fixed_quadratic_terms(beta);
4733 Ok(obs.log_likelihood - adaptive_penalty - fixed_quadratic)
4734 }
4735
4736 fn exact_newton_outerobjective(&self) -> ExactNewtonOuterObjective {
4737 ExactNewtonOuterObjective::StrictPseudoLaplace
4738 }
4739
4740 fn exact_newton_joint_hessian(
4741 &self,
4742 block_states: &[ParameterBlockState],
4743 ) -> Result<Option<Array2<f64>>, String> {
4744 let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
4745 let eval = self.exact_evaluation(beta)?;
4746 Ok(Some(eval.totalobjectivehessian(&self.design)?))
4747 }
4748
4749 fn exact_newton_hessian_directional_derivative(
4750 &self,
4751 block_states: &[ParameterBlockState],
4752 block_idx: usize,
4753 d_beta: &Array1<f64>,
4754 ) -> Result<Option<Array2<f64>>, String> {
4755 expect_block_idx_zero(block_idx, "spatial adaptive exact family", "")?;
4756 self.exact_newton_joint_hessian_directional_derivative(block_states, d_beta)
4757 }
4758
4759 fn exact_newton_joint_hessian_directional_derivative(
4760 &self,
4761 block_states: &[ParameterBlockState],
4762 d_beta_flat: &Array1<f64>,
4763 ) -> Result<Option<Array2<f64>>, String> {
4764 let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
4765 if d_beta_flat.len() != beta.len() {
4766 return Err(SmoothError::dimension_mismatch(format!(
4767 "spatial adaptive exact family direction length mismatch: got {}, expected {}",
4768 d_beta_flat.len(),
4769 beta.len()
4770 ))
4771 .into());
4772 }
4773 let eval = self.exact_evaluation(beta)?;
4774 Ok(Some(
4775 self.exacthessian_directional_derivative_from_evaluation(beta, &eval, d_beta_flat)?,
4776 ))
4777 }
4778
4779 fn exact_newton_joint_hessiansecond_directional_derivative(
4780 &self,
4781 block_states: &[ParameterBlockState],
4782 d_beta_u_flat: &Array1<f64>,
4783 d_betav_flat: &Array1<f64>,
4784 ) -> Result<Option<Array2<f64>>, String> {
4785 let beta = &expect_single_block_state(block_states, "spatial adaptive exact family")?.beta;
4786 if d_beta_u_flat.len() != beta.len() || d_betav_flat.len() != beta.len() {
4787 return Err(SmoothError::dimension_mismatch(format!(
4788 "spatial adaptive exact family second-direction length mismatch: got ({}, {}), expected {}",
4789 d_beta_u_flat.len(),
4790 d_betav_flat.len(),
4791 beta.len()
4792 ))
4793 .into());
4794 }
4795 let eval = self.exact_evaluation(beta)?;
4796 self.exacthessian_second_directional_derivative_from_evaluation(
4797 &eval,
4798 d_beta_u_flat,
4799 d_betav_flat,
4800 )
4801 }
4802
4803 fn block_linear_constraints(
4804 &self,
4805 block_states: &[ParameterBlockState],
4806 block_idx: usize,
4807 block_spec: &ParameterBlockSpec,
4808 ) -> Result<Option<LinearInequalityConstraints>, String> {
4809 assert!(!block_states.is_empty(), "block_states must be non-empty");
4810 assert!(
4811 !block_spec.name.is_empty(),
4812 "block spec name must be non-empty",
4813 );
4814 expect_block_idx_zero(block_idx, "spatial adaptive exact family", "")?;
4815 Ok(self.linear_constraints.clone())
4816 }
4817
4818 fn exact_newton_joint_psi_terms(
4819 &self,
4820 block_states: &[ParameterBlockState],
4821 specs: &[ParameterBlockSpec],
4822 derivative_blocks: &[Vec<CustomFamilyBlockPsiDerivative>],
4823 psi_index: usize,
4824 ) -> Result<Option<ExactNewtonJointPsiTerms>, String> {
4825 if block_states.len() != 1 || specs.len() != 1 || derivative_blocks.len() != 1 {
4826 return Err(SmoothError::dimension_mismatch(format!(
4827 "spatial adaptive exact family expects one block/state/spec/psi payload, got states={} specs={} deriv_blocks={}",
4828 block_states.len(),
4829 specs.len(),
4830 derivative_blocks.len()
4831 ))
4832 .into());
4833 }
4834 derivative_blocks[0]
4835 .get(psi_index)
4836 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
4837 let hyper = self
4838 .hyperspecs
4839 .get(psi_index)
4840 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
4841 let beta = &block_states[0].beta;
4842 let eval = self.exact_evaluation(beta)?;
4843 let (direct, beta_mixed, betahessian_explicit) =
4844 self.adaptive_hyper_parts(&eval, *hyper)?;
4845
4846 Ok(Some(ExactNewtonJointPsiTerms {
4867 objective_psi: direct,
4868 score_psi: beta_mixed,
4869 hessian_psi: betahessian_explicit,
4870 hessian_psi_operator: None,
4871 }))
4872 }
4873
4874 fn exact_newton_joint_psisecond_order_terms(
4875 &self,
4876 block_states: &[ParameterBlockState],
4877 specs: &[ParameterBlockSpec],
4878 derivative_blocks: &[Vec<CustomFamilyBlockPsiDerivative>],
4879 psi_i: usize,
4880 psi_j: usize,
4881 ) -> Result<Option<gam_problem::ExactNewtonJointPsiSecondOrderTerms>, String> {
4882 if block_states.len() != 1 || specs.len() != 1 || derivative_blocks.len() != 1 {
4883 return Err(SmoothError::dimension_mismatch(format!(
4884 "spatial adaptive exact family expects one block/state/spec/psi payload, got states={} specs={} deriv_blocks={}",
4885 block_states.len(),
4886 specs.len(),
4887 derivative_blocks.len()
4888 ))
4889 .into());
4890 }
4891 derivative_blocks[0]
4892 .get(psi_i)
4893 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_i))?;
4894 derivative_blocks[0]
4895 .get(psi_j)
4896 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_j))?;
4897 let hyper_i = self
4898 .hyperspecs
4899 .get(psi_i)
4900 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_i))?;
4901 let hyper_j = self
4902 .hyperspecs
4903 .get(psi_j)
4904 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_j))?;
4905 let beta = &block_states[0].beta;
4906 let eval = self.exact_evaluation(beta)?;
4907 let (objective_psi_psi, score_psi_psi, hessian_psi_psi) =
4908 self.adaptive_explicit_second_order_parts(&eval, *hyper_i, *hyper_j)?;
4909
4910 Ok(Some(
4911 gam_problem::ExactNewtonJointPsiSecondOrderTerms {
4912 objective_psi_psi,
4913 score_psi_psi,
4914 hessian_psi_psi,
4915 hessian_psi_psi_operator: None,
4916 },
4917 ))
4918 }
4919
4920 fn exact_newton_joint_psihessian_directional_derivative(
4921 &self,
4922 block_states: &[ParameterBlockState],
4923 specs: &[ParameterBlockSpec],
4924 derivative_blocks: &[Vec<CustomFamilyBlockPsiDerivative>],
4925 psi_index: usize,
4926 direction: &Array1<f64>,
4927 ) -> Result<Option<Array2<f64>>, String> {
4928 if block_states.len() != 1 || specs.len() != 1 || derivative_blocks.len() != 1 {
4929 return Err(SmoothError::dimension_mismatch(format!(
4930 "spatial adaptive exact family expects one block/state/spec/psi payload, got states={} specs={} deriv_blocks={}",
4931 block_states.len(),
4932 specs.len(),
4933 derivative_blocks.len()
4934 ))
4935 .into());
4936 }
4937 let beta = &block_states[0].beta;
4938 if direction.len() != beta.len() {
4939 return Err(SmoothError::dimension_mismatch(format!(
4940 "spatial adaptive exact family direction length mismatch: got {}, expected {}",
4941 direction.len(),
4942 beta.len()
4943 ))
4944 .into());
4945 }
4946 derivative_blocks[0]
4947 .get(psi_index)
4948 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
4949 let hyper = self
4950 .hyperspecs
4951 .get(psi_index)
4952 .ok_or_else(|| format!("adaptive psi index {} out of bounds", psi_index))?;
4953 let eval = self.exact_evaluation(beta)?;
4954 let drift = match hyper.kind {
4955 SpatialAdaptiveHyperKind::LogLambdaMagnitude
4956 | SpatialAdaptiveHyperKind::LogLambdaGradient
4957 | SpatialAdaptiveHyperKind::LogLambdaCurvature => self.adaptive_block_drift_eval(
4958 &eval,
4959 hyper.cache_index,
4960 AdaptiveComponent::from_index(hyper.kind.component_index())?,
4961 HyperDriftKind::Rho,
4962 direction,
4963 )?,
4964 SpatialAdaptiveHyperKind::LogEpsilonMagnitude
4965 | SpatialAdaptiveHyperKind::LogEpsilonGradient
4966 | SpatialAdaptiveHyperKind::LogEpsilonCurvature => self
4967 .adaptive_shared_log_epsilon_drift(
4968 &eval,
4969 hyper.kind.component_index(),
4970 direction,
4971 )?,
4972 };
4973 Ok(Some(drift))
4974 }
4975}
4976
4977fn expect_single_block_state<'a>(
4978 block_states: &'a [ParameterBlockState],
4979 family_name: &str,
4980) -> Result<&'a ParameterBlockState, String> {
4981 crate::block_layout::block_count::validate_block_count::<SmoothError>(
4982 family_name,
4983 1,
4984 block_states.len(),
4985 )?;
4986 Ok(&block_states[0])
4987}
4988
4989fn expect_single_blockspec<'a>(
4990 specs: &'a [ParameterBlockSpec],
4991 family_name: &str,
4992) -> Result<&'a ParameterBlockSpec, String> {
4993 crate::block_layout::block_count::validate_block_count::<SmoothError>(
4994 family_name,
4995 1,
4996 specs.len(),
4997 )?;
4998 Ok(&specs[0])
4999}
5000
5001fn expect_block_idx_zero(block_idx: usize, family_name: &str, context: &str) -> Result<(), String> {
5002 if block_idx != 0 {
5003 return Err(SmoothError::invalid_index(format!(
5004 "{family_name} expects block_idx 0{context}, got {block_idx}"
5005 ))
5006 .into());
5007 }
5008 Ok::<(), _>(())
5009}
5010
5011impl BoundedLinearFamily {
5012 fn bounded_term_derivative_data(
5013 &self,
5014 latent_beta: &Array1<f64>,
5015 ) -> (
5016 Array1<f64>,
5017 Array1<f64>,
5018 Array1<f64>,
5019 Array1<f64>,
5020 Array1<f64>,
5021 ) {
5022 let p = latent_beta.len();
5023 let mut beta_user = latent_beta.clone();
5024 let mut jac_diag = Array1::<f64>::ones(p);
5025 let mut second_diag = Array1::<f64>::zeros(p);
5026 let mut third_diag = Array1::<f64>::zeros(p);
5027 let mut priorthird = Array1::<f64>::zeros(p);
5028 for term in &self.bounded_terms {
5029 let (beta, _, db_dtheta, d2b_dtheta2, d3b_dtheta3) =
5030 bounded_latent_derivatives(latent_beta[term.col_idx], term.min, term.max);
5031 beta_user[term.col_idx] = beta;
5032 jac_diag[term.col_idx] = db_dtheta;
5033 second_diag[term.col_idx] = d2b_dtheta2;
5034 third_diag[term.col_idx] = d3b_dtheta3;
5035 let (_, _, _, prior_neghess_derivative) =
5036 bounded_prior_terms(latent_beta[term.col_idx], &term.prior);
5037 priorthird[term.col_idx] = prior_neghess_derivative;
5038 }
5039 (beta_user, jac_diag, second_diag, third_diag, priorthird)
5040 }
5041
5042 fn user_beta_and_jacobian(&self, latent_beta: &Array1<f64>) -> (Array1<f64>, Array1<f64>) {
5043 let (beta_user, jac_diag, _, _, _) = self.bounded_term_derivative_data(latent_beta);
5044 (beta_user, jac_diag)
5045 }
5046
5047 fn nonlinear_offset_from_latent(&self, latent_beta: &Array1<f64>) -> Array1<f64> {
5048 let mut offset = self.offset.clone();
5049 for term in &self.bounded_terms {
5050 let (beta, _, _) =
5051 bounded_latent_to_user(latent_beta[term.col_idx], term.min, term.max);
5052 offset.scaled_add(beta, &self.design.column(term.col_idx));
5053 }
5054 offset
5055 }
5056
5057 fn effective_design_for_latent(&self, jac_diag: &Array1<f64>) -> Array2<f64> {
5058 let mut x_eff = self.design.clone();
5059 for term in &self.bounded_terms {
5060 x_eff
5061 .column_mut(term.col_idx)
5062 .mapv_inplace(|v| v * jac_diag[term.col_idx]);
5063 }
5064 x_eff
5065 }
5066
5067 fn exacthessian_andgradient(
5068 &self,
5069 latent_beta: &Array1<f64>,
5070 ) -> Result<
5071 (
5072 StandardFamilyObservationState,
5073 Array2<f64>,
5074 Array1<f64>,
5075 f64,
5076 Array1<f64>,
5077 Array1<f64>,
5078 Array1<f64>,
5079 ),
5080 String,
5081 > {
5082 let (_, jac_diag, second_diag, third_diag, priorthird) =
5083 self.bounded_term_derivative_data(latent_beta);
5084 let x_eff = self.effective_design_for_latent(&jac_diag);
5085 let eta =
5086 self.designzeroed.dot(latent_beta) + self.nonlinear_offset_from_latent(latent_beta);
5087 let obs = evaluate_standard_familyobservations(
5088 self.family.clone(),
5089 self.latent_cloglog_state.as_ref(),
5090 self.mixture_link_state.as_ref(),
5091 self.sas_link_state.as_ref(),
5092 &self.y,
5093 &self.weights,
5094 &eta,
5095 )
5096 .map_err(|e| e.to_string())?;
5097
5098 let mut priorgrad = Array1::<f64>::zeros(latent_beta.len());
5099 let mut prior_neghess = Array2::<f64>::zeros((latent_beta.len(), latent_beta.len()));
5100 let mut prior_loglik = 0.0;
5101 for term in &self.bounded_terms {
5102 let (logp, grad, neghess, _) =
5103 bounded_prior_terms(latent_beta[term.col_idx], &term.prior);
5104 prior_loglik += logp;
5105 priorgrad[term.col_idx] += grad;
5106 prior_neghess[[term.col_idx, term.col_idx]] += neghess;
5107 }
5108
5109 let mut hessian = xt_diag_x_dense(x_eff.view(), obs.neghessian_eta.view())?;
5110 let mut gradient = fast_atv(&x_eff, &obs.score);
5111 for term in &self.bounded_terms {
5112 let score_beta = self.design.column(term.col_idx).dot(&obs.score);
5113 hessian[[term.col_idx, term.col_idx]] -= score_beta * second_diag[term.col_idx];
5114 }
5115 hessian += &prior_neghess;
5116 gradient += &priorgrad;
5117
5118 Ok((
5119 obs,
5120 hessian,
5121 gradient,
5122 prior_loglik,
5123 second_diag,
5124 third_diag,
5125 priorthird,
5126 ))
5127 }
5128
5129 fn evaluation_from_latent(
5130 &self,
5131 latent_beta: &Array1<f64>,
5132 ) -> Result<
5133 (
5134 StandardFamilyObservationState,
5135 Array2<f64>,
5136 Array1<f64>,
5137 f64,
5138 ),
5139 String,
5140 > {
5141 let (obs, hessian, gradient, prior_loglik, _, _, _) =
5142 self.exacthessian_andgradient(latent_beta)?;
5143 Ok((obs, hessian, gradient, prior_loglik))
5144 }
5145}
5146
5147impl CustomFamily for BoundedLinearFamily {
5148 fn joint_jeffreys_term_required(&self) -> bool {
5152 true
5153 }
5154
5155 fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
5156 let latent_beta = &expect_single_block_state(block_states, "bounded linear family")?.beta;
5157 let (obs, hessian, gradient, prior_loglik) = self.evaluation_from_latent(latent_beta)?;
5158 Ok(FamilyEvaluation {
5159 log_likelihood: obs.log_likelihood + prior_loglik,
5160 blockworking_sets: vec![BlockWorkingSet::ExactNewton {
5161 gradient,
5162 hessian: SymmetricMatrix::Dense(hessian),
5163 }],
5164 })
5165 }
5166
5167 fn exact_newton_joint_hessian(
5168 &self,
5169 block_states: &[ParameterBlockState],
5170 ) -> Result<Option<Array2<f64>>, String> {
5171 let latent_beta = &expect_single_block_state(block_states, "bounded linear family")?.beta;
5172 let (_, hessian, _, _) = self.evaluation_from_latent(latent_beta)?;
5173 Ok(Some(hessian))
5174 }
5175
5176 fn exact_newton_hessian_directional_derivative(
5177 &self,
5178 block_states: &[ParameterBlockState],
5179 block_idx: usize,
5180 d_beta: &Array1<f64>,
5181 ) -> Result<Option<Array2<f64>>, String> {
5182 expect_block_idx_zero(block_idx, "bounded linear family", "")?;
5183 self.exact_newton_joint_hessian_directional_derivative(block_states, d_beta)
5184 }
5185
5186 fn exact_newton_joint_hessian_directional_derivative(
5187 &self,
5188 block_states: &[ParameterBlockState],
5189 d_beta_flat: &Array1<f64>,
5190 ) -> Result<Option<Array2<f64>>, String> {
5191 let latent_beta = &expect_single_block_state(block_states, "bounded linear family")?.beta;
5192 if d_beta_flat.len() != latent_beta.len() {
5193 return Err(SmoothError::dimension_mismatch(format!(
5194 "bounded linear family directional derivative length mismatch: got {}, expected {}",
5195 d_beta_flat.len(),
5196 latent_beta.len()
5197 ))
5198 .into());
5199 }
5200
5201 let (obs, _, _, _, second_diag, third_diag, priorthird) =
5202 self.exacthessian_andgradient(latent_beta)?;
5203
5204 let (_, jac_diag, _, _, _) = self.bounded_term_derivative_data(latent_beta);
5205 let x_eff = self.effective_design_for_latent(&jac_diag);
5206 let deta = x_eff.dot(d_beta_flat);
5207 let d_neghess_eta = &obs.neghessian_eta_derivative * &deta;
5208
5209 let mut dx_eff = Array2::<f64>::zeros(x_eff.raw_dim());
5210 for term in &self.bounded_terms {
5211 let scale = second_diag[term.col_idx] * d_beta_flat[term.col_idx];
5212 if scale != 0.0 {
5213 let mut col = dx_eff.column_mut(term.col_idx);
5214 col.assign(&self.design.column(term.col_idx));
5215 col.mapv_inplace(|v| v * scale);
5216 }
5217 }
5218
5219 let mut dhessian = xt_diag_x_dense(x_eff.view(), d_neghess_eta.view())?;
5220 let mut wxdx = Array2::<f64>::zeros((x_eff.ncols(), x_eff.ncols()));
5221 for i in 0..x_eff.nrows() {
5222 let wi = obs.neghessian_eta[i];
5223 if wi == 0.0 {
5224 continue;
5225 }
5226 for a in 0..x_eff.ncols() {
5227 let xa = x_eff[[i, a]];
5228 for b in 0..x_eff.ncols() {
5229 wxdx[[a, b]] += wi * (dx_eff[[i, a]] * x_eff[[i, b]] + xa * dx_eff[[i, b]]);
5230 }
5231 }
5232 }
5233 dhessian += &wxdx;
5234
5235 let d_score = -&obs.neghessian_eta * &deta;
5236 for term in &self.bounded_terms {
5237 let score_beta = self.design.column(term.col_idx).dot(&obs.score);
5238 let d_score_beta = self.design.column(term.col_idx).dot(&d_score);
5239 dhessian[[term.col_idx, term.col_idx]] -= d_score_beta * second_diag[term.col_idx]
5240 + score_beta * third_diag[term.col_idx] * d_beta_flat[term.col_idx];
5241 dhessian[[term.col_idx, term.col_idx]] +=
5242 priorthird[term.col_idx] * d_beta_flat[term.col_idx];
5243 }
5244
5245 Ok(Some(dhessian))
5246 }
5247
5248 fn block_geometry(
5249 &self,
5250 block_states: &[ParameterBlockState],
5251 spec: &ParameterBlockSpec,
5252 ) -> Result<(DesignMatrix, Array1<f64>), String> {
5253 if block_states.is_empty() {
5254 return Ok((
5255 DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
5256 self.designzeroed.clone(),
5257 )),
5258 self.offset.clone(),
5259 ));
5260 }
5261 let offset = self.nonlinear_offset_from_latent(
5262 &expect_single_block_state(block_states, "bounded linear family")?.beta,
5263 );
5264 let x = if spec.design.ncols() == self.designzeroed.ncols() {
5265 self.designzeroed.clone()
5266 } else {
5267 return Err(SmoothError::dimension_mismatch(
5268 "bounded linear family design column mismatch",
5269 )
5270 .into());
5271 };
5272 Ok((
5273 DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(x)),
5274 offset,
5275 ))
5276 }
5277
5278 fn block_geometry_is_dynamic(&self) -> bool {
5279 true
5280 }
5281
5282 fn block_geometry_directional_derivative(
5283 &self,
5284 block_states: &[ParameterBlockState],
5285 block_idx: usize,
5286 spec: &ParameterBlockSpec,
5287 d_beta: &Array1<f64>,
5288 ) -> Result<Option<BlockGeometryDirectionalDerivative>, String> {
5289 expect_block_idx_zero(
5290 block_idx,
5291 "bounded linear family",
5292 " for geometry derivative",
5293 )?;
5294 expect_single_block_state(block_states, "bounded linear family")?;
5295 if d_beta.len() != spec.design.ncols() {
5296 return Err(SmoothError::dimension_mismatch(format!(
5297 "bounded linear family geometry derivative direction mismatch: got {}, expected {}",
5298 d_beta.len(),
5299 spec.design.ncols()
5300 ))
5301 .into());
5302 }
5303 let (_, jac_diag, _, _, _) = self.bounded_term_derivative_data(&block_states[0].beta);
5304 let mut d_offset = Array1::<f64>::zeros(self.offset.len());
5305 let has_drift = self
5306 .bounded_terms
5307 .iter()
5308 .any(|term| jac_diag[term.col_idx] != 0.0 && d_beta[term.col_idx] != 0.0);
5309 if !has_drift {
5310 return Ok(Some(BlockGeometryDirectionalDerivative {
5311 d_design: None,
5312 d_offset,
5313 }));
5314 }
5315 for term in &self.bounded_terms {
5316 let col = term.col_idx;
5317 let drift = jac_diag[col] * d_beta[col];
5318 if drift != 0.0 {
5319 d_offset.scaled_add(drift, &self.design.column(col));
5320 }
5321 }
5322 Ok(Some(BlockGeometryDirectionalDerivative {
5323 d_design: None,
5324 d_offset,
5325 }))
5326 }
5327}
5328
5329#[inline]
5330fn dense_diag_gram_chunkrows(p: usize) -> usize {
5331 const MIN_ROWS: usize = 512;
5332 const MAX_ROWS: usize = 2048;
5333 const TARGET_BYTES: usize = 2 * 1024 * 1024;
5334 let bytes_per_row = p.max(1) * std::mem::size_of::<f64>();
5335 (TARGET_BYTES / bytes_per_row).clamp(MIN_ROWS, MAX_ROWS)
5336}
5337
5338fn xt_diag_x_dense(x: ArrayView2<'_, f64>, w: ArrayView1<'_, f64>) -> Result<Array2<f64>, String> {
5339 if x.nrows() != w.len() {
5340 return Err(SmoothError::dimension_mismatch("xt_diag_x_dense row mismatch").into());
5341 }
5342 let (n, p) = x.dim();
5343 if n == 0 || p == 0 {
5344 return Ok(Array2::<f64>::zeros((p, p)));
5345 }
5346
5347 const STREAMING_BYTES_THRESHOLD: usize = 8 * 1024 * 1024;
5348 let dense_work_bytes = n
5349 .checked_mul(p)
5350 .and_then(|cells| cells.checked_mul(std::mem::size_of::<f64>()))
5351 .unwrap_or(usize::MAX);
5352 if dense_work_bytes <= STREAMING_BYTES_THRESHOLD {
5353 let mut weighted = x.to_owned();
5354 ndarray::Zip::from(weighted.rows_mut())
5355 .and(w)
5356 .par_for_each(|mut row, wi| row *= *wi);
5357 return Ok(fast_atb(&x, &weighted));
5358 }
5359
5360 let chunkrows = dense_diag_gram_chunkrows(p).min(n);
5361 let mut weighted_chunk = Array2::<f64>::zeros((chunkrows, p));
5362 let mut out = Array2::<f64>::zeros((p, p));
5363 for row_start in (0..n).step_by(chunkrows) {
5364 let rows = (n - row_start).min(chunkrows);
5365 let x_chunk = x.slice(s![row_start..row_start + rows, ..]);
5366 {
5367 let mut chunk = weighted_chunk.slice_mut(s![0..rows, ..]);
5368 for local_row in 0..rows {
5369 let scale = w[row_start + local_row];
5370 if scale == 0.0 {
5371 chunk.row_mut(local_row).fill(0.0);
5372 continue;
5373 }
5374 for col in 0..p {
5375 chunk[[local_row, col]] = x_chunk[[local_row, col]] * scale;
5376 }
5377 }
5378 }
5379 out += &fast_atb(&x_chunk, &weighted_chunk.slice(s![0..rows, ..]));
5380 }
5381 Ok(out)
5382}
5383
5384fn trace_of_dense_product(a: &Array2<f64>, b: &Array2<f64>) -> Result<f64, String> {
5385 if a.nrows() != a.ncols() || b.nrows() != b.ncols() || a.nrows() != b.nrows() {
5386 return Err(
5387 SmoothError::dimension_mismatch("trace_of_dense_product dimension mismatch").into(),
5388 );
5389 }
5390 let mut trace = 0.0;
5391 for i in 0..a.nrows() {
5392 for j in 0..a.ncols() {
5393 trace += a[[i, j]] * b[[j, i]];
5394 }
5395 }
5396 Ok(trace)
5397}
5398
5399fn exact_bounded_edf(
5400 penalties: &[PenaltySpec],
5401 lambdas: &Array1<f64>,
5402 latent_cov: &Array2<f64>,
5403) -> Result<(Vec<f64>, Vec<f64>, f64), EstimationError> {
5404 if penalties.len() != lambdas.len() {
5405 crate::bail_invalid_estim!(
5406 "bounded EDF penalty/lambda mismatch: {} penalties vs {} lambdas",
5407 penalties.len(),
5408 lambdas.len()
5409 );
5410 }
5411 if latent_cov.nrows() != latent_cov.ncols() {
5412 crate::bail_invalid_estim!("bounded EDF covariance must be square");
5413 }
5414
5415 let p = latent_cov.nrows();
5416 let mut s_lambda = Array2::<f64>::zeros((p, p));
5417 let mut edf_by_block = Vec::with_capacity(penalties.len());
5418 let mut penalty_block_trace = Vec::with_capacity(penalties.len());
5420 let mut trace_sum = 0.0;
5421
5422 for (k, ps) in penalties.iter().enumerate() {
5423 let lambda_k = lambdas[k];
5424 match ps {
5425 PenaltySpec::Block {
5426 local, col_range, ..
5427 } => {
5428 s_lambda
5429 .slice_mut(ndarray::s![col_range.clone(), col_range.clone()])
5430 .scaled_add(lambda_k, local);
5431 let penalty_rank =
5433 local
5434 .nrows()
5435 .saturating_sub(estimate_penalty_nullity(local).map_err(|e| {
5436 EstimationError::InvalidInput(format!("bounded EDF rank failed: {e}"))
5437 })?);
5438 let cov_block = latent_cov.slice(ndarray::s![col_range.clone(), col_range.clone()]);
5440 let trace_k = lambda_k
5441 * trace_of_dense_product(&cov_block.to_owned(), local)
5442 .map_err(EstimationError::InvalidInput)?;
5443 trace_sum += trace_k;
5444 penalty_block_trace.push(trace_k);
5445 let p_k = penalty_rank as f64;
5446 edf_by_block.push((p_k - trace_k).clamp(0.0, p_k));
5447 }
5448 PenaltySpec::Dense(m) | PenaltySpec::DenseWithMean { matrix: m, .. } => {
5449 s_lambda.scaled_add(lambda_k, m);
5450 let penalty_rank = p.saturating_sub(estimate_penalty_nullity(m).map_err(|e| {
5451 EstimationError::InvalidInput(format!("bounded EDF rank failed: {e}"))
5452 })?);
5453 let trace_k = lambda_k
5454 * trace_of_dense_product(latent_cov, m)
5455 .map_err(EstimationError::InvalidInput)?;
5456 trace_sum += trace_k;
5457 penalty_block_trace.push(trace_k);
5458 let p_k = penalty_rank as f64;
5459 edf_by_block.push((p_k - trace_k).clamp(0.0, p_k));
5460 }
5461 }
5462 }
5463
5464 let nullity_total = estimate_penalty_nullity(&s_lambda)
5465 .map_err(|e| EstimationError::InvalidInput(format!("bounded EDF nullity failed: {e}")))?
5466 as f64;
5467 let edf_total = (p as f64 - trace_sum).clamp(nullity_total, p as f64);
5468 Ok((edf_by_block, penalty_block_trace, edf_total))
5469}
5470
5471fn symmetric_positive_definite_inverse_or_pseudo(
5483 precision: &Array2<f64>,
5484) -> Result<Array2<f64>, EstimationError> {
5485 use gam_linalg::faer_ndarray::FaerEigh;
5486 let p = precision.nrows();
5487 if precision.ncols() != p {
5488 crate::bail_invalid_estim!(
5489 "posterior precision inverse requires a square matrix, got {}x{}",
5490 precision.nrows(),
5491 precision.ncols()
5492 );
5493 }
5494 if p == 0 {
5495 return Ok(Array2::<f64>::zeros((0, 0)));
5496 }
5497 let symmetric = (precision + &precision.t().to_owned()) * 0.5;
5498 let (evals, evecs) = symmetric.eigh(faer::Side::Lower).map_err(|e| {
5499 EstimationError::InvalidInput(format!(
5500 "posterior precision eigendecomposition failed: {e}"
5501 ))
5502 })?;
5503 let max_abs_eval = evals.iter().fold(0.0_f64, |acc, &ev| acc.max(ev.abs()));
5504 let tol =
5505 (10.0 * f64::EPSILON * (p as f64) * (p as f64) * max_abs_eval).max(100.0 * f64::EPSILON);
5506 if let Some(&min_eval) = evals
5507 .iter()
5508 .filter(|&&ev| ev < -tol)
5509 .min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
5510 {
5511 crate::bail_invalid_estim!(
5512 "bounded posterior precision is non-PD at the converged optimum (min eigenvalue \
5513 {min_eval:.6e} < -tol={tol:.6e}); the reported mode is not a strict posterior \
5514 maximum, so a covariance would be meaningless"
5515 );
5516 }
5517 let mut scaled = evecs.clone();
5519 for (j, &ev) in evals.iter().enumerate() {
5520 let inv = if ev > tol { 1.0 / ev } else { 0.0 };
5521 scaled.column_mut(j).mapv_inplace(|v| v * inv);
5522 }
5523 let cov = scaled.dot(&evecs.t());
5524 Ok((&cov + &cov.t().to_owned()) * 0.5)
5525}
5526
5527fn transform_bounded_latent_precision_to_user_internal(
5528 latent_precision: &Array2<f64>,
5529 jac_diag: &Array1<f64>,
5530) -> Result<Array2<f64>, EstimationError> {
5531 let p = latent_precision.nrows();
5532 if latent_precision.ncols() != p || jac_diag.len() != p {
5533 crate::bail_invalid_estim!(
5534 "bounded precision transform dimension mismatch: precision is {}x{}, jacobian has {} entries",
5535 latent_precision.nrows(),
5536 latent_precision.ncols(),
5537 jac_diag.len()
5538 );
5539 }
5540 let mut out = latent_precision.clone();
5541 for i in 0..p {
5542 let scale = jac_diag[i];
5543 if !scale.is_finite() || scale <= 0.0 {
5544 crate::bail_invalid_estim!(
5545 "bounded precision transform requires a positive finite coefficient jacobian; column {i} has {scale}"
5546 );
5547 }
5548 if scale != 1.0 {
5549 out.row_mut(i).mapv_inplace(|v| v / scale);
5550 out.column_mut(i).mapv_inplace(|v| v / scale);
5551 }
5552 }
5553 Ok(out)
5554}
5555
5556fn fit_bounded_term_collection_with_design(
5557 y: ArrayView1<'_, f64>,
5558 weights: ArrayView1<'_, f64>,
5559 offset: ArrayView1<'_, f64>,
5560 spec: &TermCollectionSpec,
5561 design: &TermCollectionDesign,
5562 heuristic_lambdas: Option<&[f64]>,
5563 family: LikelihoodSpec,
5564 options: &FitOptions,
5565) -> Result<FittedTermCollection, EstimationError> {
5566 let conditioning_cols: Vec<usize> = spec
5567 .linear_terms
5568 .iter()
5569 .enumerate()
5570 .filter_map(|(j, linear)| {
5571 (!linear.double_penalty).then_some(design.intercept_range.end + j)
5572 })
5573 .collect();
5574 let conditioning = LinearFitConditioning::from_columns(design, &conditioning_cols);
5575 let dense_design = design.design.to_dense_cow();
5576 let fit_design = conditioning.apply_to_design(&dense_design);
5577 let fit_penalties = conditioning
5578 .transform_blockwise_penalties_to_internal(&design.penalties, design.design.ncols());
5579 if design.linear_constraints.is_some() {
5580 crate::bail_invalid_estim!(
5581 "bounded() terms are not yet compatible with explicit linear constraints"
5582 );
5583 }
5584 let mut bounded_terms = Vec::<BoundedLinearTermMeta>::new();
5585 for (j, term) in spec.linear_terms.iter().enumerate() {
5586 if term.double_penalty
5587 && matches!(
5588 term.coefficient_geometry,
5589 LinearCoefficientGeometry::Bounded { .. }
5590 )
5591 {
5592 crate::bail_invalid_estim!(
5593 "bounded linear term '{}' cannot also use double_penalty",
5594 term.name
5595 );
5596 }
5597 if let LinearCoefficientGeometry::Bounded { min, max, prior } =
5598 term.coefficient_geometry.clone()
5599 {
5600 let col_idx = design.intercept_range.end + j;
5601 let (min_internal, max_internal) = conditioning.internal_bounds_for(col_idx, min, max);
5602 bounded_terms.push(BoundedLinearTermMeta {
5603 col_idx,
5604 min: min_internal,
5605 max: max_internal,
5606 prior,
5607 });
5608 }
5609 }
5610 if bounded_terms.is_empty() {
5611 crate::bail_invalid_estim!("internal bounded fit path called with no bounded terms");
5612 }
5613
5614 let mut designzeroed = fit_design.clone();
5615 let mut initial_beta = Array1::<f64>::zeros(fit_design.ncols());
5616 for term in &bounded_terms {
5617 designzeroed.column_mut(term.col_idx).fill(0.0);
5618 initial_beta[term.col_idx] = bounded_logit(0.5);
5619 }
5620
5621 let initial_log_lambdas = heuristic_lambdas
5622 .map(|vals| Array1::from_vec(vals.to_vec()))
5623 .unwrap_or_else(|| Array1::zeros(fit_penalties.len()));
5624 if initial_log_lambdas.len() != fit_penalties.len() {
5625 crate::bail_invalid_estim!(
5626 "heuristic lambda length mismatch for bounded model: got {}, expected {}",
5627 initial_log_lambdas.len(),
5628 fit_penalties.len()
5629 );
5630 }
5631
5632 let is_beta_logistic = family.is_binomial_beta_logistic();
5633 let family_adapter = BoundedLinearFamily {
5634 family: family.clone(),
5635 latent_cloglog_state: options.latent_cloglog,
5636 mixture_link_state: options
5637 .mixture_link
5638 .clone()
5639 .as_ref()
5640 .map(state_fromspec)
5641 .transpose()
5642 .map_err(EstimationError::InvalidInput)?,
5643 sas_link_state: options
5644 .sas_link
5645 .map(|spec| {
5646 if is_beta_logistic {
5647 state_from_beta_logisticspec(spec)
5648 } else {
5649 state_from_sasspec(spec)
5650 }
5651 })
5652 .transpose()
5653 .map_err(EstimationError::InvalidInput)?,
5654 y: y.to_owned(),
5655 weights: weights.to_owned(),
5656 design: fit_design.clone(),
5657 designzeroed: designzeroed.clone(),
5658 offset: offset.to_owned(),
5659 bounded_terms: bounded_terms.clone(),
5660 };
5661 let blockspec = ParameterBlockSpec {
5662 name: "eta".to_string(),
5663 design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(designzeroed)),
5664 offset: offset.to_owned(),
5665 penalties: fit_penalties
5666 .iter()
5667 .map(|ps| match ps {
5668 PenaltySpec::Block {
5669 local, col_range, ..
5670 } => PenaltyMatrix::Blockwise {
5671 local: local.clone(),
5672 col_range: col_range.clone(),
5673 total_dim: design.design.ncols(),
5674 },
5675 PenaltySpec::Dense(m) | PenaltySpec::DenseWithMean { matrix: m, .. } => {
5676 PenaltyMatrix::Dense(m.clone())
5677 }
5678 })
5679 .collect(),
5680 nullspace_dims: design.nullspace_dims.clone(),
5681 initial_log_lambdas,
5682 initial_beta: Some(initial_beta),
5683 gauge_priority: 100,
5684 jacobian_callback: Some(Arc::new(BoundedEffectiveJacobian {
5690 design: fit_design.clone(),
5691 bounded_terms: bounded_terms.clone(),
5692 })),
5693 stacked_design: None,
5694 stacked_offset: None,
5695 };
5696 let fit = fit_custom_family(
5697 &family_adapter,
5698 &[blockspec],
5699 &BlockwiseFitOptions {
5700 inner_max_cycles: options.max_iter,
5701 inner_tol: options.tol,
5702 outer_max_iter: options.max_iter,
5703 outer_tol: options.tol,
5704 compute_covariance: false,
5714 ..BlockwiseFitOptions::default()
5715 },
5716 )
5717 .map_err(EstimationError::CustomFamily)?;
5718
5719 let latent_beta = fit.block_states[0].beta.clone();
5720 let (beta_user_internal, jac_diag) = family_adapter.user_beta_and_jacobian(&latent_beta);
5721 let beta_user = conditioning.backtransform_beta(&beta_user_internal);
5722
5723 let (eta_state, h_data, _, _) = family_adapter
5724 .evaluation_from_latent(&latent_beta)
5725 .map_err(EstimationError::InvalidInput)?;
5726 let p_fit = fit_design.ncols();
5727 let mut s_lambda_internal = Array2::<f64>::zeros((p_fit, p_fit));
5728 for (k, penalty) in fit_penalties.iter().enumerate() {
5729 match penalty {
5730 PenaltySpec::Block {
5731 local, col_range, ..
5732 } => {
5733 s_lambda_internal
5734 .slice_mut(ndarray::s![col_range.clone(), col_range.clone()])
5735 .scaled_add(fit.lambdas[k], local);
5736 }
5737 PenaltySpec::Dense(m) | PenaltySpec::DenseWithMean { matrix: m, .. } => {
5738 s_lambda_internal.scaled_add(fit.lambdas[k], m);
5739 }
5740 }
5741 }
5742 let mut latent_precision = h_data.clone();
5743 latent_precision += &s_lambda_internal;
5744 let user_precision_internal =
5745 transform_bounded_latent_precision_to_user_internal(&latent_precision, &jac_diag)?;
5746 let penalized_hessian =
5747 conditioning.transform_penalized_hessian_to_original(&user_precision_internal);
5748
5749 let beta_covariance_unscaled = if options.compute_inference {
5777 Some(symmetric_positive_definite_inverse_or_pseudo(
5778 &penalized_hessian,
5779 )?)
5780 } else {
5781 None
5782 };
5783 let latent_cov = if options.compute_inference {
5789 Some(symmetric_positive_definite_inverse_or_pseudo(
5790 &latent_precision,
5791 )?)
5792 } else {
5793 None
5794 };
5795 let s_lambda_original = weighted_blockwise_penalty_sum(
5796 &design.penalties,
5797 fit.lambdas.as_slice().unwrap(),
5798 design.design.ncols(),
5799 );
5800 let penalty_term = beta_user.dot(&s_lambda_original.dot(&beta_user));
5801 let deviance = if family.is_gaussian_identity() {
5802 y.iter()
5803 .zip(eta_state.mu.iter())
5804 .zip(weights.iter())
5805 .map(|((&yy, &mu), &w)| w.max(0.0) * (yy - mu) * (yy - mu))
5806 .sum()
5807 } else {
5808 -2.0 * eta_state.log_likelihood
5809 };
5810 let (edf_by_block, penalty_block_trace, edf_total) = if let Some(cov) = latent_cov.as_ref() {
5811 exact_bounded_edf(&fit_penalties, &fit.lambdas, cov)?
5812 } else {
5813 (
5814 vec![0.0; fit_penalties.len()],
5815 vec![0.0; fit_penalties.len()],
5816 0.0,
5817 )
5818 };
5819
5820 let glm_likelihood = gam_spec::GlmLikelihoodSpec::canonical(family.clone());
5832 let standard_deviation = if family.is_gaussian_identity() {
5833 let denom = if options.compute_inference {
5834 (y.len() as f64 - edf_total).max(1.0)
5835 } else {
5836 (y.len() as f64).max(1.0)
5837 };
5838 (deviance / denom).sqrt()
5839 } else {
5840 1.0
5841 };
5842 let cov_scale = glm_likelihood
5843 .coefficient_covariance_scale(standard_deviation * standard_deviation)
5844 .max(f64::MIN_POSITIVE);
5845 let dispersion = gam_solve::estimate::dispersion_from_likelihood(&glm_likelihood, standard_deviation);
5846 let beta_covariance = beta_covariance_unscaled.map(|mut cov| {
5852 if cov_scale != 1.0 {
5853 cov.mapv_inplace(|v| v * cov_scale);
5854 }
5855 cov
5856 });
5857 let beta_standard_errors = beta_covariance
5858 .as_ref()
5859 .map(|cov| Array1::from_iter((0..cov.nrows()).map(|i| cov[[i, i]].max(0.0).sqrt())));
5860
5861 let geometry = Some(gam_solve::estimate::FitGeometry {
5862 penalized_hessian: penalized_hessian.clone().into(),
5863 working_weights: eta_state.fisherweight.clone(),
5864 working_response: {
5865 let mut working_response = eta_state.eta.clone();
5866 for i in 0..working_response.len() {
5867 let wi = eta_state.fisherweight[i].max(1e-12);
5868 working_response[i] += eta_state.score[i] / wi;
5869 }
5870 working_response
5871 },
5872 });
5873 let max_abs_eta = eta_state
5874 .eta
5875 .iter()
5876 .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
5877 Ok(FittedTermCollection {
5878 fit: {
5879 let log_lambdas = fit.lambdas.mapv(|v| v.max(1e-300).ln());
5880 let inf = FitInference {
5881 edf_by_block,
5882 penalty_block_trace,
5883 edf_total,
5884 smoothing_correction: None,
5885 penalized_hessian: penalized_hessian.clone().into(),
5888 working_weights: eta_state.fisherweight.clone(),
5889 working_response: {
5890 let mut working_response = eta_state.eta.clone();
5891 for i in 0..working_response.len() {
5892 let wi = eta_state.fisherweight[i].max(1e-12);
5893 working_response[i] += eta_state.score[i] / wi;
5894 }
5895 working_response
5896 },
5897 reparam_qs: None,
5898 dispersion,
5899 beta_covariance: beta_covariance
5900 .clone()
5901 .map(gam_problem::dispersion_cov::PhiScaledCovariance::from),
5902 beta_standard_errors,
5903 beta_covariance_corrected: None,
5904 beta_standard_errors_corrected: None,
5905 beta_covariance_frequentist: None,
5906 coefficient_influence: None,
5907 weighted_gram: None,
5908 bias_correction_beta: None,
5909 bias_correction_jacobian: None,
5910 };
5911 let covariance_conditional = beta_covariance;
5912 let pirls_status_val = gam_solve::pirls::PirlsStatus::Converged;
5915 UnifiedFitResult::try_from_parts(UnifiedFitResultParts {
5916 blocks: vec![gam_solve::estimate::FittedBlock {
5917 beta: beta_user.clone(),
5918 role: gam_problem::BlockRole::Mean,
5919 edf: edf_total,
5920 lambdas: fit.lambdas.clone(),
5921 }],
5922 log_lambdas,
5923 lambdas: fit.lambdas,
5924 likelihood_scale: family.default_scale_metadata(),
5925 likelihood_family: Some(family),
5926 log_likelihood_normalization:
5927 gam_spec::LogLikelihoodNormalization::UserProvided,
5928 log_likelihood: eta_state.log_likelihood,
5929 deviance,
5930 reml_score: fit.penalized_objective,
5931 stable_penalty_term: penalty_term,
5932 penalized_objective: fit.penalized_objective,
5933 used_device: false,
5934 outer_iterations: fit.outer_iterations,
5935 outer_converged: true,
5937 outer_gradient_norm: fit.outer_gradient_norm,
5938 standard_deviation,
5939 covariance_conditional,
5940 covariance_corrected: None,
5941 inference: Some(inf),
5942 fitted_link: gam_solve::estimate::FittedLinkState::Standard(None),
5943 geometry,
5944 block_states: Vec::new(),
5945 pirls_status: pirls_status_val,
5946 max_abs_eta,
5947 constraint_kkt: None,
5948 artifacts: gam_solve::estimate::FitArtifacts {
5949 pirls: None,
5950 ..Default::default()
5951 },
5952 inner_cycles: 0,
5953 })?
5954 },
5955 design: design.clone(),
5956 adaptive_diagnostics: None,
5957 })
5958}
5959
5960fn enforce_term_constraint_feasibility(
5961 design: &TermCollectionDesign,
5962 fit: &UnifiedFitResult,
5963) -> Result<(), EstimationError> {
5964 const CONSTRAINT_FEASIBILITY_RAW_TOL: f64 = 1e-7;
5978 let tol = CONSTRAINT_FEASIBILITY_RAW_TOL;
5979 let smooth_start = design
5980 .design
5981 .ncols()
5982 .saturating_sub(design.smooth.total_smooth_cols());
5983 let mut violations: Vec<String> = Vec::new();
5984 for term in &design.smooth.terms {
5985 let gr = (smooth_start + term.coeff_range.start)..(smooth_start + term.coeff_range.end);
5986 let beta_local = fit.beta.slice(s![gr.clone()]).to_owned();
5987 if let Some(lb) = term.lower_bounds_local.as_ref() {
5988 let mut worst = 0.0_f64;
5989 let mut worst_idx = 0usize;
5990 for i in 0..lb.len().min(beta_local.len()) {
5991 if lb[i].is_finite() {
5992 let viol = (lb[i] - beta_local[i]).max(0.0);
5993 if viol > worst {
5994 worst = viol;
5995 worst_idx = i;
5996 }
5997 }
5998 }
5999 if worst > tol {
6000 violations.push(format!(
6001 "term='{}' kind=lower-bound maxviolation={:.3e} coeff_index={}",
6002 term.name, worst, worst_idx
6003 ));
6004 }
6005 }
6006 if let Some(lin) = term.linear_constraints_local.as_ref() {
6007 let mut worst = 0.0_f64;
6008 let mut worstrow = 0usize;
6009 for i in 0..lin.a.nrows() {
6010 let norm = lin.a.row(i).dot(&lin.a.row(i)).sqrt();
6011 let inv = if norm > 0.0 { 1.0 / norm } else { 0.0 };
6012 let s = (lin.a.row(i).dot(&beta_local) - lin.b[i]) * inv;
6013 let viol = (-s).max(0.0);
6014 if viol > worst {
6015 worst = viol;
6016 worstrow = i;
6017 }
6018 }
6019 if worst > tol {
6020 violations.push(format!(
6021 "term='{}' kind=linear-inequality maxviolation={:.3e} row={}",
6022 term.name, worst, worstrow
6023 ));
6024 }
6025 }
6026 }
6027
6028 if !violations.is_empty() {
6029 let mut msg = format!(
6030 "constraint violation after fit ({} violating term constraints): {}",
6031 violations.len(),
6032 violations.join(" | ")
6033 );
6034 if let Some(kkt) = fit.constraint_kkt.as_ref() {
6035 msg.push_str(&format!(
6036 "; KKT[primal={:.3e}, dual={:.3e}, comp={:.3e}, stat={:.3e}]",
6037 kkt.primal_feasibility, kkt.dual_feasibility, kkt.complementarity, kkt.stationarity
6038 ));
6039 }
6040 return Err(EstimationError::ParameterConstraintViolation(msg));
6041 }
6042 Ok(())
6043}
6044
6045fn stratified_spatial_subsample(
6046 data: ArrayView2<'_, f64>,
6047 spec: &TermCollectionSpec,
6048 target_size: usize,
6049) -> Vec<usize> {
6050 use rand::SeedableRng;
6051 use rand::rngs::StdRng;
6052 use rand::seq::SliceRandom;
6053
6054 let n = data.nrows();
6055 if n <= target_size {
6056 return (0..n).collect();
6057 }
6058
6059 let spatial_cols: Option<Vec<usize>> =
6060 spec.smooth_terms.iter().find_map(|term| match &term.basis {
6061 SmoothBasisSpec::ThinPlate { feature_cols, .. }
6062 | SmoothBasisSpec::Matern { feature_cols, .. }
6063 | SmoothBasisSpec::Duchon { feature_cols, .. } => {
6064 if !feature_cols.is_empty() {
6065 Some(feature_cols.clone())
6066 } else {
6067 None
6068 }
6069 }
6070 _ => None,
6071 });
6072
6073 let cols = match spatial_cols {
6074 Some(c) if !c.is_empty() => c,
6075 _ => {
6076 let mut rng = StdRng::seed_from_u64(spatial_subsample_seed(data, &[], target_size));
6077 let mut indices: Vec<usize> = (0..n).collect();
6078 indices.shuffle(&mut rng);
6079 indices.truncate(target_size);
6080 indices.sort_unstable();
6081 return indices;
6082 }
6083 };
6084 let mut rng = StdRng::seed_from_u64(spatial_subsample_seed(data, &cols, target_size));
6085
6086 let d = cols.len();
6087 let mut mins = vec![f64::INFINITY; d];
6088 let mut maxs = vec![f64::NEG_INFINITY; d];
6089 for i in 0..n {
6090 for (ax, &col) in cols.iter().enumerate() {
6091 let v = data[[i, col]];
6092 if v < mins[ax] {
6093 mins[ax] = v;
6094 }
6095 if v > maxs[ax] {
6096 maxs[ax] = v;
6097 }
6098 }
6099 }
6100
6101 const TARGET_POINTS_PER_CELL: usize = 5;
6105 let total_cells_target = (target_size / TARGET_POINTS_PER_CELL).max(1);
6106 let cells_per_axis = ((total_cells_target as f64).powf(1.0 / d as f64)).ceil() as usize;
6107 let cells_per_axis = cells_per_axis.max(1);
6108
6109 let mut cell_members: std::collections::HashMap<Vec<usize>, Vec<usize>> =
6110 std::collections::HashMap::new();
6111 for i in 0..n {
6112 let mut cell_key = Vec::with_capacity(d);
6113 for (ax, &col) in cols.iter().enumerate() {
6114 let range = maxs[ax] - mins[ax];
6115 let cell = if range <= 0.0 {
6116 0
6117 } else {
6118 let frac = (data[[i, col]] - mins[ax]) / range;
6119 (frac * cells_per_axis as f64).floor() as usize
6120 };
6121 cell_key.push(cell.min(cells_per_axis - 1));
6122 }
6123 cell_members.entry(cell_key).or_default().push(i);
6124 }
6125
6126 let mut selected: Vec<usize> = Vec::with_capacity(target_size);
6127 let mut remaining_budget = target_size;
6128 let mut remaining_population = n;
6129
6130 let mut cells: Vec<(Vec<usize>, Vec<usize>)> = cell_members.into_iter().collect();
6131 cells.sort_by(|a, b| a.0.cmp(&b.0));
6132
6133 for (_, members) in &mut cells {
6134 if remaining_budget == 0 {
6135 break;
6136 }
6137 let alloc = ((members.len() as f64 / remaining_population as f64) * remaining_budget as f64)
6138 .round() as usize;
6139 let alloc = alloc.max(1).min(members.len()).min(remaining_budget);
6140 members.shuffle(&mut rng);
6141 selected.extend_from_slice(&members[..alloc]);
6142 remaining_budget = remaining_budget.saturating_sub(alloc);
6143 remaining_population = remaining_population.saturating_sub(members.len());
6144 }
6145
6146 if selected.len() > target_size {
6147 selected.shuffle(&mut rng);
6148 selected.truncate(target_size);
6149 }
6150
6151 selected.sort_unstable();
6152 selected
6153}
6154
6155fn spatial_subsample_seed(
6156 data: ArrayView2<'_, f64>,
6157 spatial_cols: &[usize],
6158 target_size: usize,
6159) -> u64 {
6160 let mut state = 0x5350_4154_4941_4C53_u64;
6161 spatial_seed_mix(&mut state, data.nrows() as u64);
6162 spatial_seed_mix(&mut state, data.ncols() as u64);
6163 spatial_seed_mix(&mut state, target_size as u64);
6164 spatial_seed_mix(&mut state, spatial_cols.len() as u64);
6165 for &col in spatial_cols {
6166 spatial_seed_mix(&mut state, col as u64);
6167 }
6168
6169 if data.nrows() > 0 {
6170 let mid = data.nrows() / 2;
6171 let last = data.nrows() - 1;
6172 for &row in &[0usize, mid, last] {
6173 for &col in spatial_cols {
6174 let value = data[[row, col]];
6175 spatial_seed_mix(&mut state, value.to_bits());
6176 }
6177 }
6178 }
6179 state
6180}
6181
6182#[inline]
6183fn spatial_seed_mix(state: &mut u64, value: u64) {
6184 let mut s = value.wrapping_add(*state);
6187 let z = gam_linalg::utils::splitmix64(&mut s);
6188 *state ^= z;
6189 *state = (*state).rotate_left(27).wrapping_mul(0x3C79_AC49_2BA7_B653);
6190}
6191
6192fn sampled_rows(data: ArrayView2<'_, f64>, indices: &[usize]) -> Array2<f64> {
6193 let mut sampled = Array2::<f64>::zeros((indices.len(), data.ncols()));
6194 for (new_row, &orig_row) in indices.iter().enumerate() {
6195 sampled.row_mut(new_row).assign(&data.row(orig_row));
6196 }
6197 sampled
6198}
6199
6200fn spatial_term_user_centers(term: &SmoothTermSpec) -> Option<ArrayView2<'_, f64>> {
6201 match spatial_term_center_strategy(term) {
6202 Some(CenterStrategy::UserProvided(centers)) => Some(centers.view()),
6203 _ => None,
6204 }
6205}
6206
6207fn finite_centered_axis_contrasts(values: &[f64], expected_dim: usize) -> Option<Vec<f64>> {
6208 if values.len() != expected_dim || expected_dim <= 1 {
6209 return None;
6210 }
6211 if values.iter().any(|value| !value.is_finite()) {
6212 return None;
6213 }
6214 Some(center_aniso_log_scales(values))
6215}
6216
6217fn blended_pilot_axis_contrasts(
6218 pilot_data: ArrayView2<'_, f64>,
6219 term: &SmoothTermSpec,
6220 centers: ArrayView2<'_, f64>,
6221) -> Option<Vec<f64>> {
6222 let d = centers.ncols();
6223 if d <= 1 {
6224 return None;
6225 }
6226 let center_eta = initial_aniso_contrasts(centers);
6227 let data_eta = standardized_spatial_term_data(pilot_data, term)
6228 .ok()
6229 .and_then(|x| finite_centered_axis_contrasts(&initial_aniso_contrasts(x.view()), d));
6230 let center_eta = finite_centered_axis_contrasts(¢er_eta, d)?;
6231 let blended = match data_eta {
6232 Some(data_eta) => center_eta
6233 .iter()
6234 .zip(data_eta.iter())
6235 .map(|(&from_centers, &from_data)| 0.5 * (from_centers + from_data))
6236 .collect::<Vec<_>>(),
6237 None => center_eta,
6238 };
6239 finite_centered_axis_contrasts(&blended, d)
6240}
6241
6242fn apply_pilot_spatial_psi_reseed(
6243 pilot_data: ArrayView2<'_, f64>,
6244 spec: &TermCollectionSpec,
6245 spatial_terms: &[usize],
6246 kappa_options: &SpatialLengthScaleOptimizationOptions,
6247) -> Result<TermCollectionSpec, EstimationError> {
6248 let dims_per_term = spatial_dims_per_term(spec, spatial_terms);
6249 let use_aniso = has_aniso_terms(spec, spatial_terms);
6250 let log_kappa0 = if use_aniso {
6251 SpatialLogKappaCoords::from_length_scales_aniso(spec, spatial_terms, kappa_options)
6252 } else {
6253 SpatialLogKappaCoords::from_length_scales(spec, spatial_terms, kappa_options)
6254 };
6255 let log_kappa0 = log_kappa0.reseed_from_data(pilot_data, spec, spatial_terms, kappa_options);
6256 let log_kappa_lower = if use_aniso {
6257 SpatialLogKappaCoords::lower_bounds_aniso_from_data(
6258 pilot_data,
6259 spec,
6260 spatial_terms,
6261 &dims_per_term,
6262 kappa_options,
6263 )
6264 } else {
6265 SpatialLogKappaCoords::lower_bounds_from_data(
6266 pilot_data,
6267 spec,
6268 spatial_terms,
6269 kappa_options,
6270 )
6271 };
6272 let log_kappa_upper = if use_aniso {
6273 SpatialLogKappaCoords::upper_bounds_aniso_from_data(
6274 pilot_data,
6275 spec,
6276 spatial_terms,
6277 &dims_per_term,
6278 kappa_options,
6279 )
6280 } else {
6281 SpatialLogKappaCoords::upper_bounds_from_data(
6282 pilot_data,
6283 spec,
6284 spatial_terms,
6285 kappa_options,
6286 )
6287 };
6288 log_kappa0
6289 .clamp_to_bounds(&log_kappa_lower, &log_kappa_upper)
6290 .apply_tospec(spec, spatial_terms)
6291}
6292
6293pub(crate) fn apply_spatial_anisotropy_pilot_initializer(
6294 data: ArrayView2<'_, f64>,
6295 spec: &mut TermCollectionSpec,
6296 spatial_terms: &[usize],
6297 target_size: usize,
6298 kappa_options: &SpatialLengthScaleOptimizationOptions,
6299) -> usize {
6300 if target_size == 0 || data.nrows() <= target_size.saturating_mul(2) || spatial_terms.is_empty()
6301 {
6302 return 0;
6303 }
6304 if !has_aniso_terms(spec, spatial_terms) {
6305 return 0;
6306 }
6307 let indices = stratified_spatial_subsample(data, spec, target_size);
6308 let pilot_data = sampled_rows(data, &indices);
6309 let mut working = spec.clone();
6310 let mut updated_terms = 0usize;
6311 const GEOMETRY_UPDATES: usize = 2;
6312
6313 for pass in 0..GEOMETRY_UPDATES {
6314 let planned_terms = match plan_joint_spatial_centers_for_term_blocks(
6315 pilot_data.view(),
6316 &[working.smooth_terms.clone()],
6317 )
6318 .and_then(|mut blocks| {
6319 blocks.pop().ok_or_else(|| {
6320 BasisError::InvalidInput(
6321 "pilot geometry initializer produced no smooth-term block".to_string(),
6322 )
6323 })
6324 }) {
6325 Ok(terms) => terms,
6326 Err(err) => {
6327 log::warn!(
6328 "[spatial-kappa] pilot geometry initializer skipped after center planning failed: {err}"
6329 );
6330 return updated_terms;
6331 }
6332 };
6333
6334 for &term_idx in spatial_terms {
6335 let Some(current_eta) = get_spatial_aniso_log_scales(&working, term_idx) else {
6336 continue;
6337 };
6338 let Some(d) = get_spatial_feature_dim(&working, term_idx) else {
6339 continue;
6340 };
6341 if d <= 1 || current_eta.len() != d {
6342 continue;
6343 }
6344 let Some(planned_term) = planned_terms.get(term_idx) else {
6345 continue;
6346 };
6347 let Some(centers) = spatial_term_user_centers(planned_term) else {
6348 continue;
6349 };
6350 let Some(eta) = blended_pilot_axis_contrasts(pilot_data.view(), planned_term, centers)
6351 else {
6352 continue;
6353 };
6354 if set_spatial_aniso_log_scales(&mut working, term_idx, eta).is_ok() {
6355 updated_terms += usize::from(pass == 0);
6356 }
6357 }
6358
6359 match apply_pilot_spatial_psi_reseed(
6360 pilot_data.view(),
6361 &working,
6362 spatial_terms,
6363 kappa_options,
6364 ) {
6365 Ok(updated) => {
6366 working = updated;
6367 }
6368 Err(err) => {
6369 log::warn!(
6370 "[spatial-kappa] pilot geometry ψ reseed skipped after deterministic initializer error: {err}"
6371 );
6372 break;
6373 }
6374 }
6375 }
6376
6377 if updated_terms > 0 {
6378 log::info!(
6379 "[spatial-kappa] initialized anisotropy from {}-row pilot geometry for {} spatial term(s); proceeding to full-data optimization",
6380 indices.len(),
6381 updated_terms
6382 );
6383 *spec = working;
6384 }
6385 updated_terms
6386}
6387
6388pub(crate) fn spatial_length_scale_term_indices(spec: &TermCollectionSpec) -> Vec<usize> {
6389 spec.smooth_terms
6390 .iter()
6391 .enumerate()
6392 .filter_map(|(idx, _)| spatial_term_supports_hyper_optimization(spec, idx).then_some(idx))
6393 .collect()
6394}
6395
6396fn fit_score(fit: &UnifiedFitResult) -> f64 {
6408 if fit.reml_score.is_finite() {
6409 return fit.reml_score;
6410 }
6411 let score = 0.5 * fit.deviance + 0.5 * fit.stable_penalty_term;
6412 if score.is_finite() {
6413 score
6414 } else {
6415 f64::INFINITY
6416 }
6417}
6418
6419fn is_recoverable_trial_point_error(err: &EstimationError) -> bool {
6441 matches!(err, EstimationError::BasisError(_))
6442 || err.is_inner_solve_retreat()
6443 || is_recoverable_fit_inference_finiteness_error(err)
6444}
6445
6446fn is_recoverable_fit_inference_finiteness_error(err: &EstimationError) -> bool {
6447 let EstimationError::InvalidInput(message) = err else {
6448 return false;
6449 };
6450
6451 message.contains("must be finite")
6452 && [
6453 "fit_result.beta_covariance_frequentist",
6454 "fit_result.coefficient_influence",
6455 "fit_result.weighted_gram",
6456 ]
6457 .iter()
6458 .any(|field| message.contains(field))
6459}
6460
6461#[cfg(test)]
6462mod spatial_trial_recovery_tests {
6463 use super::*;
6464
6465 #[test]
6466 fn nonfinite_frequentist_covariance_is_recoverable_trial_point() {
6467 let err = EstimationError::InvalidInput(
6468 "fit_result.beta_covariance_frequentist[0] must be finite, got NaN".to_string(),
6469 );
6470
6471 assert!(
6472 is_recoverable_trial_point_error(&err),
6473 "singular trial-point curvature should make spatial κ retreat, not abort"
6474 );
6475 }
6476
6477 #[test]
6478 fn arbitrary_invalid_input_remains_fatal_trial_point_error() {
6479 let err = EstimationError::InvalidInput("outer rho bounds are invalid".to_string());
6480
6481 assert!(
6482 !is_recoverable_trial_point_error(&err),
6483 "the spatial κ recovery gate must not mask unrelated invalid inputs"
6484 );
6485 }
6486}
6487
6488fn require_successful_spatial_optimization_result<T>(
6489 initial_score: f64,
6490 result: Result<Option<(T, f64)>, EstimationError>,
6491) -> Result<T, EstimationError> {
6492 match result {
6493 Ok(Some((value, exact_score))) => {
6494 const SCORE_DRIFT_ABS_TOL: f64 = 1e-6;
6503 const SCORE_DRIFT_REL_TOL: f64 = 1e-8;
6504 let tol = SCORE_DRIFT_ABS_TOL.max(initial_score.abs() * SCORE_DRIFT_REL_TOL);
6505 if exact_score <= initial_score + tol {
6506 Ok(value)
6507 } else {
6508 Err(EstimationError::RemlOptimizationFailed(format!(
6509 "spatial kappa optimization made REML score worse ({initial_score:.6e} -> {exact_score:.6e})"
6510 )))
6511 }
6512 }
6513 Ok(None) => Err(EstimationError::RemlOptimizationFailed(
6514 "spatial kappa optimization is unavailable for one or more eligible spatial terms"
6515 .to_string(),
6516 )),
6517 Err(err) => Err(EstimationError::RemlOptimizationFailed(format!(
6518 "spatial kappa optimization failed: {err}"
6519 ))),
6520 }
6521}
6522
6523fn external_opts_for_design(
6524 family: &LikelihoodSpec,
6525 design: &TermCollectionDesign,
6526 options: &FitOptions,
6527) -> ExternalOptimOptions {
6528 ExternalOptimOptions {
6529 family: family.clone(),
6530 latent_cloglog: options.latent_cloglog,
6531 mixture_link: options.mixture_link.clone(),
6532 optimize_mixture: options.optimize_mixture,
6533 sas_link: options.sas_link,
6534 optimize_sas: options.optimize_sas,
6535 compute_inference: options.compute_inference,
6536 skip_rho_posterior_inference: options.skip_rho_posterior_inference,
6537 max_iter: options.max_iter,
6538 tol: options.tol,
6539 nullspace_dims: design.nullspace_dims.clone(),
6540 linear_constraints: design.linear_constraints.clone(),
6541 firth_bias_reduction: Some(options.firth_bias_reduction),
6542 penalty_shrinkage_floor: options.penalty_shrinkage_floor,
6543 rho_prior: options.rho_prior.clone(),
6544 kronecker_penalty_system: design.kronecker_penalty_system(),
6547 kronecker_factored: design
6548 .smooth
6549 .terms
6550 .iter()
6551 .find_map(|t| t.kronecker_factored.clone()),
6552 persist_warm_start_disk: options.persist_warm_start_disk,
6553 }
6554}
6555
6556fn evaluate_joint_reml_outer_eval_at_theta(
6564 evaluator: &mut gam_solve::estimate::ExternalJointHyperEvaluator<'_>,
6565 design: &TermCollectionDesign,
6566 theta: &Array1<f64>,
6567 rho_dim: usize,
6568 hyper_dirs: Vec<gam_solve::estimate::reml::DirectionalHyperParam>,
6569 warm_start_beta: Option<ArrayView1<'_, f64>>,
6570 order: gam_solve::rho_optimizer::OuterEvalOrder,
6571 design_revision: Option<u64>,
6572) -> Result<
6573 (
6574 f64,
6575 Array1<f64>,
6576 gam_problem::HessianValue,
6577 ),
6578 EstimationError,
6579> {
6580 evaluator.evaluate_with_order(
6581 &design.design,
6582 &design.penalties,
6583 &design.nullspace_dims,
6584 design.linear_constraints.clone(),
6585 theta,
6586 rho_dim,
6587 hyper_dirs,
6588 warm_start_beta,
6589 "evaluate_joint_reml_outer_eval_at_theta",
6590 order,
6591 design_revision,
6592 )
6593}
6594
6595fn evaluate_joint_reml_efs_at_theta(
6596 evaluator: &mut gam_solve::estimate::ExternalJointHyperEvaluator<'_>,
6597 design: &TermCollectionDesign,
6598 theta: &Array1<f64>,
6599 rho_dim: usize,
6600 hyper_dirs: Vec<gam_solve::estimate::reml::DirectionalHyperParam>,
6601 warm_start_beta: Option<ArrayView1<'_, f64>>,
6602 design_revision: Option<u64>,
6603) -> Result<gam_problem::EfsEval, EstimationError> {
6604 evaluator.evaluate_efs(
6605 &design.design,
6606 &design.penalties,
6607 &design.nullspace_dims,
6608 design.linear_constraints.clone(),
6609 theta,
6610 rho_dim,
6611 hyper_dirs,
6612 warm_start_beta,
6613 "evaluate_joint_reml_efs_at_theta",
6614 design_revision,
6615 )
6616}
6617
6618fn exact_joint_spatial_outer_hessian_available(
6619 family: &LikelihoodSpec,
6620 design: &TermCollectionDesign,
6621) -> bool {
6622 let family_supported = match &family.response {
6645 ResponseFamily::Gaussian
6646 | ResponseFamily::Binomial
6647 | ResponseFamily::Poisson
6648 | ResponseFamily::Tweedie { .. }
6649 | ResponseFamily::NegativeBinomial { .. }
6650 | ResponseFamily::Beta { .. }
6651 | ResponseFamily::Gamma
6652 | ResponseFamily::RoystonParmar => true,
6653 };
6654 family_supported && design.design.ncols() > 0
6657}
6658
6659fn smooth_term_penalty_index(
6660 spec: &TermCollectionSpec,
6661 design: &TermCollectionDesign,
6662 term_idx: usize,
6663) -> Option<usize> {
6664 if term_idx >= design.smooth.terms.len() || term_idx >= spec.smooth_terms.len() {
6665 return None;
6666 }
6667 if design.smooth.terms[term_idx].penalties_local.is_empty() {
6668 return None;
6669 }
6670 let linear_penalties = spec
6671 .linear_terms
6672 .iter()
6673 .filter(|t| t.double_penalty)
6674 .count()
6675 * 2;
6676 let random_penalties = design
6677 .random_effect_ranges
6678 .iter()
6679 .filter(|(_, range)| !range.is_empty())
6680 .count();
6681 let smooth_offset = linear_penalties + random_penalties;
6682 let local_offset = design
6683 .smooth
6684 .terms
6685 .iter()
6686 .take(term_idx)
6687 .map(|term| term.penalties_local.len())
6688 .sum::<usize>();
6689 Some(smooth_offset + local_offset)
6690}
6691
6692fn try_build_spatial_term_log_kappa_derivativeinfo(
6693 data: ArrayView2<'_, f64>,
6694 resolvedspec: &TermCollectionSpec,
6695 design: &TermCollectionDesign,
6696 term_idx: usize,
6697) -> Result<Option<SpatialPsiDerivative>, EstimationError> {
6698 let Some((
6699 global_range,
6700 total_p,
6701 x_psi_local,
6702 s_psi_local_check,
6703 x_psi_psi_local,
6704 s_psi_psi_local,
6705 s_psi_components_local,
6706 s_psi_psi_components_local,
6707 implicit_operator,
6708 )) = try_build_spatial_term_log_kappa_derivative(data, resolvedspec, design, term_idx)?
6709 else {
6710 return Ok(None);
6711 };
6712 let Some(penalty_start) = smooth_term_penalty_index(resolvedspec, design, term_idx) else {
6713 return Ok(None);
6714 };
6715 if s_psi_components_local.is_empty() || s_psi_psi_components_local.is_empty() {
6716 return Ok(None);
6717 }
6718 if s_psi_components_local.len() != s_psi_psi_components_local.len() {
6719 return Ok(None);
6720 }
6721 let penalty_indices = (0..s_psi_components_local.len())
6722 .map(|j| penalty_start + j)
6723 .collect::<Vec<_>>();
6724 let penalty_index = penalty_indices[0];
6725 if s_psi_local_check.nrows() == 0 || s_psi_psi_local.nrows() == 0 {
6726 return Ok(None);
6727 }
6728 Ok(Some(SpatialPsiDerivative {
6729 penalty_index,
6730 penalty_indices,
6731 global_range,
6732 total_p,
6733 x_psi_local,
6734 s_psi_components_local,
6735 x_psi_psi_local,
6736 s_psi_psi_components_local,
6737 aniso_group_id: None,
6738 aniso_cross_designs: None,
6739 aniso_cross_penalty_provider: None,
6740 implicit_operator,
6741 implicit_axis: 0,
6742 }))
6743}
6744
6745pub(crate) fn try_build_spatial_log_kappa_derivativeinfo_list(
6746 data: ArrayView2<'_, f64>,
6747 resolvedspec: &TermCollectionSpec,
6748 design: &TermCollectionDesign,
6749 spatial_terms: &[usize],
6750) -> Result<Option<Vec<SpatialPsiDerivative>>, EstimationError> {
6751 let mut out = Vec::new();
6752 let mut aniso_gid = 0usize;
6753 for &term_idx in spatial_terms {
6754 if spatial_term_uses_per_axis_psi(resolvedspec, term_idx) {
6755 if let Some(entries) = try_build_spatial_term_log_kappa_aniso_derivativeinfos(
6756 data,
6757 resolvedspec,
6758 design,
6759 term_idx,
6760 aniso_gid,
6761 )? {
6762 aniso_gid += 1;
6763 out.extend(entries);
6764 continue;
6765 } else {
6766 return Ok(None);
6767 }
6768 }
6769 let Some(info) =
6770 try_build_spatial_term_log_kappa_derivativeinfo(data, resolvedspec, design, term_idx)?
6771 else {
6772 return Ok(None);
6773 };
6774 out.push(info);
6775 }
6776 Ok(Some(out))
6777}
6778
6779fn try_build_spatial_term_log_kappa_aniso_derivativeinfos(
6781 data: ArrayView2<'_, f64>,
6782 resolvedspec: &TermCollectionSpec,
6783 design: &TermCollectionDesign,
6784 term_idx: usize,
6785 aniso_group_id: usize,
6786) -> Result<Option<Vec<SpatialPsiDerivative>>, EstimationError> {
6787 let Some(smooth_term) = design.smooth.terms.get(term_idx) else {
6788 return Ok(None);
6789 };
6790 let Some(termspec) = resolvedspec.smooth_terms.get(term_idx) else {
6791 return Ok(None);
6792 };
6793 let mut aniso_result = match &termspec.basis {
6794 SmoothBasisSpec::Sphere { .. } => return Ok(None),
6795 SmoothBasisSpec::Matern {
6796 feature_cols,
6797 spec,
6798 input_scales,
6799 } => {
6800 let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
6801 if let Some(s) = input_scales {
6802 apply_input_standardization(&mut x, s);
6803 }
6804 let mut spec_operator = spec.clone();
6813 spec_operator.double_penalty = false;
6814 build_matern_basis_log_kappa_aniso_derivatives(x.view(), &spec_operator)
6815 .map_err(EstimationError::from)?
6816 }
6817 SmoothBasisSpec::MeasureJet {
6823 feature_cols,
6824 spec,
6825 input_scales,
6826 } => {
6827 let mut x = select_columns(data, feature_cols).map_err(EstimationError::from)?;
6828 if let Some(s) = input_scales {
6829 apply_input_standardization(&mut x, s);
6830 }
6831 build_measure_jet_basis_psi_derivatives(x.view(), spec)
6832 .map_err(EstimationError::from)?
6833 }
6834 _ => return Ok(None),
6835 };
6836 let d = if let Some(ref op) = aniso_result.implicit_operator {
6839 op.n_axes()
6840 } else if !aniso_result.design_first.is_empty() {
6841 aniso_result.design_first.len()
6842 } else {
6843 0
6844 };
6845 if d == 0 {
6846 return Ok(None);
6847 }
6848 let Some(penalty_start) = smooth_term_penalty_index(resolvedspec, design, term_idx) else {
6849 return Ok(None);
6850 };
6851 let p_total = design.design.ncols();
6852 let smooth_start = p_total.saturating_sub(design.smooth.total_smooth_cols());
6853 let global_range = (smooth_start + smooth_term.coeff_range.start)
6854 ..(smooth_start + smooth_term.coeff_range.end);
6855 let num_penalties = aniso_result.penalties_first[0].len();
6856 let penalty_indices: Vec<usize> = (0..num_penalties).map(|j| penalty_start + j).collect();
6857 let penalties_cross_provider = aniso_result.penalties_cross_provider.clone();
6858
6859 let use_implicit_design = aniso_result.design_first.is_empty();
6863 let implicit_op_arc = aniso_result
6864 .implicit_operator
6865 .as_ref()
6866 .map(|op| std::sync::Arc::new(op.clone()));
6867
6868 let mut entries = Vec::with_capacity(d);
6869 for a in 0..d {
6870 let (x_psi_local, x_psi_psi_local) = if use_implicit_design {
6871 (Array2::<f64>::zeros((0, 0)), Array2::<f64>::zeros((0, 0)))
6877 } else {
6878 let x_first = std::mem::take(&mut aniso_result.design_first[a]);
6883 let x_second = std::mem::take(&mut aniso_result.design_second_diag[a]);
6884 if x_first.ncols() != smooth_term.coeff_range.len() {
6885 return Ok(None);
6886 }
6887 (x_first, x_second)
6888 };
6889 let s_psi_components = std::mem::take(&mut aniso_result.penalties_first[a]);
6890 let s_psi_psi_components = std::mem::take(&mut aniso_result.penalties_second_diag[a]);
6891 let cross_designs = if implicit_op_arc.is_some() {
6897 let mut cd = Vec::with_capacity(d - 1);
6898 for b in 0..d {
6899 if b == a {
6900 continue;
6901 }
6902 cd.push((b, Array2::<f64>::zeros((0, 0))));
6903 }
6904 cd
6905 } else if !aniso_result.design_second_cross.is_empty() {
6906 let mut cd = Vec::new();
6907 for (cross_idx, &(pa, pb)) in aniso_result.design_second_cross_pairs.iter().enumerate()
6908 {
6909 if pa == a {
6910 cd.push((pb, aniso_result.design_second_cross[cross_idx].clone()));
6911 } else if pb == a {
6912 cd.push((pa, aniso_result.design_second_cross[cross_idx].clone()));
6913 }
6914 }
6915 cd
6916 } else {
6917 Vec::new()
6918 };
6919 let cross_penalty_provider = if d > 1 {
6920 let penalties_cross_provider = penalties_cross_provider.clone();
6921 Some(std::sync::Arc::new(
6922 move |b_axis: usize| -> Result<Vec<Array2<f64>>, EstimationError> {
6923 if b_axis == a {
6924 return Ok(Vec::new());
6925 }
6926 let (axis_lo, axis_hi) = if a < b_axis { (a, b_axis) } else { (b_axis, a) };
6927 if let Some(provider) = penalties_cross_provider.as_ref() {
6928 provider
6929 .evaluate(axis_lo, axis_hi)
6930 .map_err(EstimationError::from)
6931 } else {
6932 Ok(Vec::new())
6936 }
6937 },
6938 )
6939 as std::sync::Arc<
6940 dyn Fn(usize) -> Result<Vec<Array2<f64>>, EstimationError>
6941 + Send
6942 + Sync
6943 + 'static,
6944 >)
6945 } else {
6946 None
6947 };
6948
6949 entries.push(SpatialPsiDerivative {
6950 penalty_index: penalty_indices[0],
6951 penalty_indices: penalty_indices.clone(),
6952 global_range: global_range.clone(),
6953 total_p: p_total,
6954 x_psi_local,
6955 s_psi_components_local: s_psi_components,
6956 x_psi_psi_local,
6957 s_psi_psi_components_local: s_psi_psi_components,
6958 aniso_group_id: Some(aniso_group_id),
6959 aniso_cross_designs: if cross_designs.is_empty() {
6960 None
6961 } else {
6962 Some(cross_designs)
6963 },
6964 aniso_cross_penalty_provider: cross_penalty_provider,
6965 implicit_operator: implicit_op_arc.clone(),
6966 implicit_axis: a,
6967 });
6968 }
6969 Ok(Some(entries))
6970}
6971
6972#[cfg(test)]
6973mod glm_eta_observation_fd_tests {
6974 use super::*;
6980
6981 fn one_obs(spec: &LikelihoodSpec, y: f64, eta: f64) -> StandardFamilyObservationState {
6982 let yv = Array1::from_vec(vec![y]);
6983 let wv = Array1::from_vec(vec![1.0]);
6984 let ev = Array1::from_vec(vec![eta]);
6985 evaluate_standard_familyobservations(spec.clone(), None, None, None, &yv, &wv, &ev)
6986 .expect("standard family observation state assembles")
6987 }
6988
6989 fn check_fd(label: &str, spec: &LikelihoodSpec, y: f64, eta: f64) {
6990 let h = 1e-5;
6991 let s0 = one_obs(spec, y, eta);
6992 let sp = one_obs(spec, y, eta + h);
6993 let sm = one_obs(spec, y, eta - h);
6994
6995 let score_fd = (sp.log_likelihood - sm.log_likelihood) / (2.0 * h);
6997 let score = s0.score[0];
6998 assert!(
6999 (score - score_fd).abs() <= 1e-4 * (1.0 + score.abs()),
7000 "{label}: score {score} vs FD {score_fd}"
7001 );
7002
7003 let neghess_fd = -(sp.score[0] - sm.score[0]) / (2.0 * h);
7005 let neghess = s0.neghessian_eta[0];
7006 assert!(
7007 (neghess - neghess_fd).abs() <= 1e-3 * (1.0 + neghess.abs()),
7008 "{label}: neghessian_eta {neghess} vs FD {neghess_fd}"
7009 );
7010
7011 let nhd_fd = (sp.neghessian_eta[0] - sm.neghessian_eta[0]) / (2.0 * h);
7013 let nhd = s0.neghessian_eta_derivative[0];
7014 assert!(
7015 (nhd - nhd_fd).abs() <= 1e-2 * (1.0 + nhd.abs()),
7016 "{label}: neghessian_eta_derivative {nhd} vs FD {nhd_fd}"
7017 );
7018 }
7019
7020 #[test]
7021 fn poisson_gamma_nb_tweedie_arms_match_finite_differences_1615_1616() {
7022 let log = InverseLink::Standard(StandardLink::Log);
7023 let poisson = LikelihoodSpec {
7024 response: ResponseFamily::Poisson,
7025 link: log.clone(),
7026 };
7027 check_fd("poisson y=3", &poisson, 3.0, 0.4);
7028 check_fd("poisson y=0", &poisson, 0.0, -0.2);
7029
7030 let gamma = LikelihoodSpec {
7031 response: ResponseFamily::Gamma,
7032 link: log.clone(),
7033 };
7034 check_fd("gamma y=2.5", &gamma, 2.5, 0.3);
7035 check_fd("gamma y=0.7", &gamma, 0.7, -0.1);
7036
7037 let nb = LikelihoodSpec {
7038 response: ResponseFamily::NegativeBinomial {
7039 theta: 1.5,
7040 theta_fixed: true,
7041 },
7042 link: log.clone(),
7043 };
7044 check_fd("negbin y=4", &nb, 4.0, 0.5);
7045 check_fd("negbin y=0", &nb, 0.0, -0.3);
7046
7047 let tweedie = LikelihoodSpec {
7048 response: ResponseFamily::Tweedie { p: 1.5 },
7049 link: log.clone(),
7050 };
7051 check_fd("tweedie y=2", &tweedie, 2.0, 0.25);
7052 check_fd("tweedie y=0.5", &tweedie, 0.5, -0.15);
7053 }
7054}