1use super::*;
6
7#[derive(Clone, Copy)]
8pub(crate) struct GamlssLambdaLayout {
9 pub(crate) k_mean: usize,
10 pub(crate) k_noise: usize,
11 pub(crate) kwiggle: usize,
12}
13
14impl GamlssLambdaLayout {
15 pub(crate) fn two_block(k_mean: usize, k_noise: usize) -> Self {
16 Self {
17 k_mean,
18 k_noise,
19 kwiggle: 0,
20 }
21 }
22
23 pub(crate) fn withwiggle(k_mean: usize, k_noise: usize, kwiggle: usize) -> Self {
24 Self {
25 k_mean,
26 k_noise,
27 kwiggle,
28 }
29 }
30
31 pub(crate) fn total(self) -> usize {
32 self.k_mean + self.k_noise + self.kwiggle
33 }
34
35 pub(crate) fn noise_start(self) -> usize {
36 self.k_mean
37 }
38
39 pub(crate) fn noise_end(self) -> usize {
40 self.k_mean + self.k_noise
41 }
42
43 pub(crate) fn wiggle_start(self) -> usize {
44 self.k_mean + self.k_noise
45 }
46
47 pub(crate) fn wiggle_end(self) -> usize {
48 self.k_mean + self.k_noise + self.kwiggle
49 }
50
51 pub(crate) fn validate_theta_len(self, theta_len: usize, context: &str) -> Result<(), String> {
52 let needed = self.total();
53 if theta_len < needed {
54 return Err(GamlssError::DimensionMismatch {
55 reason: format!(
56 "{context} theta too short: got {}, need at least {}",
57 theta_len, needed
58 ),
59 }
60 .into());
61 }
62 Ok(())
63 }
64
65 pub(crate) fn mean_from(self, theta: &Array1<f64>) -> Array1<f64> {
66 theta.slice(s![0..self.k_mean]).to_owned()
67 }
68
69 pub(crate) fn noise_from(self, theta: &Array1<f64>) -> Array1<f64> {
70 theta
71 .slice(s![self.noise_start()..self.noise_end()])
72 .to_owned()
73 }
74
75 pub(crate) fn wiggle_from(self, theta: &Array1<f64>) -> Array1<f64> {
76 theta
77 .slice(s![self.wiggle_start()..self.wiggle_end()])
78 .to_owned()
79 }
80}
81
82#[derive(Clone, Copy)]
83pub(crate) struct GamlssBetaLayout {
84 pub(crate) pt: usize,
85 pub(crate) pls: usize,
86 pub(crate) pw: usize,
87}
88
89impl GamlssBetaLayout {
90 pub(crate) fn withwiggle(pt: usize, pls: usize, pw: usize) -> Self {
91 Self { pt, pls, pw }
92 }
93
94 pub(crate) fn total(self) -> usize {
95 self.pt + self.pls + self.pw
96 }
97
98 pub(crate) fn split_three(
99 self,
100 flat: &Array1<f64>,
101 context: &str,
102 ) -> Result<(Array1<f64>, Array1<f64>, Array1<f64>), String> {
103 if flat.len() != self.total() {
104 return Err(GamlssError::DimensionMismatch {
105 reason: format!(
106 "{context} length mismatch: got {}, expected {}",
107 flat.len(),
108 self.total()
109 ),
110 }
111 .into());
112 }
113 Ok((
114 flat.slice(s![0..self.pt]).to_owned(),
115 flat.slice(s![self.pt..self.pt + self.pls]).to_owned(),
116 flat.slice(s![self.pt + self.pls..self.total()]).to_owned(),
117 ))
118 }
119}
120
121#[derive(Clone, Debug)]
122pub struct FamilyMetadata {
123 pub name: &'static str,
124 pub parameternames: &'static [&'static str],
125 pub parameter_links: &'static [ParameterLink],
126}
127
128pub(crate) const DEFAULT_GAUGE_PRIORITY: u8 = 100;
129
130pub(crate) const LINK_WIGGLE_GAUGE_PRIORITY: u8 = 80;
131
132pub(crate) fn initial_log_lambdas_orzeros(
133 block: &ParameterBlockInput,
134) -> Result<Array1<f64>, String> {
135 let k = block.penalties.len();
136 let lambdas = block
137 .initial_log_lambdas
138 .clone()
139 .unwrap_or_else(|| Array1::<f64>::zeros(k));
140 if lambdas.len() != k {
141 return Err(GamlssError::DimensionMismatch {
142 reason: format!(
143 "initial_log_lambdas length mismatch: got {}, expected {}",
144 lambdas.len(),
145 k
146 ),
147 }
148 .into());
149 }
150 Ok(lambdas)
151}
152
153pub(crate) fn build_two_block_exact_joint_setup(
154 data: ArrayView2<'_, f64>,
155 meanspec: &TermCollectionSpec,
156 noisespec: &TermCollectionSpec,
157 mean_penalties: usize,
158 noise_penalties: usize,
159 extra_rho0: &[f64],
160 rho0_override: Option<&Array1<f64>>,
161 kappa_options: &SpatialLengthScaleOptimizationOptions,
162) -> ExactJointHyperSetup {
163 let rho_dim = mean_penalties + noise_penalties + extra_rho0.len();
166 let mut rho0vec = Array1::<f64>::zeros(rho_dim);
167 if let Some(rho0) = rho0_override.filter(|rho0| rho0.len() == rho_dim) {
168 rho0vec.assign(rho0);
169 } else {
170 for (i, &rho_init) in extra_rho0.iter().enumerate() {
171 rho0vec[mean_penalties + noise_penalties + i] = rho_init;
172 }
173 }
174
175 build_location_scale_exact_joint_setup(data, &[meanspec, noisespec], rho0vec, kappa_options)
178}
179
180pub(crate) fn gaussian_location_scalewarm_start(
181 y: &Array1<f64>,
182 weights: &Array1<f64>,
183 mu_block: &ParameterBlockSpec,
184 log_sigma_block: &ParameterBlockSpec,
185 ridge_floor: f64,
186 mean_beta_hint: Option<&Array1<f64>>,
187 noise_beta_hint: Option<&Array1<f64>>,
188) -> Result<(Array1<f64>, Array1<f64>, f64), String> {
189 let betamu = if let Some(beta) = mean_beta_hint {
190 beta.clone()
191 } else {
192 solve_penalizedweighted_projection(
193 &mu_block.design,
194 &mu_block.offset,
195 y,
196 weights,
197 &mu_block.penalties,
198 &mu_block.initial_log_lambdas,
199 ridge_floor,
200 )?
201 };
202 let mut mu_hat = mu_block.solver_design().matrixvectormultiply(&betamu);
203 mu_hat += mu_block.solver_offset();
204 let mut weighted_ss = 0.0;
205 let mut weight_sum = 0.0;
206 for i in 0..y.len() {
207 let wi = weights[i].max(0.0);
208 let resid = y[i] - mu_hat[i];
209 weighted_ss += wi * resid * resid;
210 weight_sum += wi;
211 }
212 if !weighted_ss.is_finite() || !weight_sum.is_finite() || weight_sum <= 0.0 {
213 return Err(
214 "gaussian location-scale warm start could not estimate residual scale".to_string(),
215 );
216 }
217 let sigma_hat = (weighted_ss / weight_sum)
222 .sqrt()
223 .max(LOGB_SIGMA_FLOOR * 1.5);
224 let beta_log_sigma = if let Some(beta) = noise_beta_hint {
225 beta.clone()
226 } else {
227 let eta_sigma = (sigma_hat - LOGB_SIGMA_FLOOR).ln();
228 let sigma_target = Array1::from_elem(y.len(), eta_sigma);
229 solve_penalizedweighted_projection(
230 &log_sigma_block.design,
231 &log_sigma_block.offset,
232 &sigma_target,
233 weights,
234 &log_sigma_block.penalties,
235 &log_sigma_block.initial_log_lambdas,
236 ridge_floor,
237 )?
238 };
239 Ok((betamu, beta_log_sigma, sigma_hat))
240}
241
242pub(crate) const LOCATION_SCALE_N_OUTPUTS: usize = 2;
246
247pub(crate) fn build_location_scale_block(
261 name: impl Into<String>,
262 design: DesignMatrix,
263 offset: Array1<f64>,
264 penalties: Vec<PenaltyMatrix>,
265 nullspace_dims: Vec<usize>,
266 initial_log_lambdas: Array1<f64>,
267 initial_beta: Option<Array1<f64>>,
268 own_output: usize,
269 n_family_outputs: usize,
270 caller: &str,
271) -> Result<ParameterBlockSpec, String> {
272 if own_output >= n_family_outputs {
273 return Err(format!(
274 "{caller}: own_output={own_output} >= n_family_outputs={n_family_outputs}"
275 ));
276 }
277 let mut spec = ParameterBlockSpec {
278 name: name.into(),
279 design,
280 offset,
281 penalties,
282 nullspace_dims,
283 initial_log_lambdas,
284 initial_beta,
285 gauge_priority: 100,
286 jacobian_callback: None,
287 stacked_design: None,
288 stacked_offset: None,
289 };
290 let dense = spec.effective_design(caller)?;
291 spec.jacobian_callback = Some(std::sync::Arc::new(AdditiveBlockJacobian {
292 design: dense,
293 own_output,
294 n_family_outputs,
295 }));
296 Ok(spec)
297}
298
299pub(crate) fn build_location_scale_wiggle_block(
305 name: impl Into<String>,
306 design: DesignMatrix,
307 offset: Array1<f64>,
308 penalties: Vec<PenaltyMatrix>,
309 nullspace_dims: Vec<usize>,
310 initial_log_lambdas: Array1<f64>,
311 initial_beta: Option<Array1<f64>>,
312 n_rows: usize,
313) -> Result<ParameterBlockSpec, String> {
314 let p_w = design.ncols();
315 let mut spec = ParameterBlockSpec {
316 name: name.into(),
317 design,
318 offset,
319 penalties,
320 nullspace_dims,
321 initial_log_lambdas,
322 initial_beta,
323 gauge_priority: 100,
324 jacobian_callback: None,
325 stacked_design: None,
326 stacked_offset: None,
327 };
328 spec.jacobian_callback = Some(std::sync::Arc::new(AdditiveBlockJacobian {
329 design: ndarray::Array2::<f64>::zeros((n_rows, p_w)),
330 own_output: 0,
331 n_family_outputs: LOCATION_SCALE_N_OUTPUTS,
332 }));
333 Ok(spec)
334}
335
336pub(crate) fn prepared_gaussian_log_sigma_design(
337 mu_design: &DesignMatrix,
338 log_sigma_design: &DesignMatrix,
339) -> Result<DesignMatrix, String> {
340 if mu_design.nrows() != log_sigma_design.nrows() {
341 return Err(GamlssError::DimensionMismatch {
342 reason: format!(
343 "gaussian log-sigma design row mismatch: mean rows={}, log_sigma rows={}",
344 mu_design.nrows(),
345 log_sigma_design.nrows()
346 ),
347 }
348 .into());
349 }
350 Ok(log_sigma_design.clone())
361}
362
363pub(crate) fn identified_binomial_log_sigma_design(
364 threshold_design: &TermCollectionDesign,
365 log_sigma_design: &TermCollectionDesign,
366 weights: &Array1<f64>,
367) -> Result<DesignMatrix, String> {
368 let non_intercept_start = log_sigma_design
369 .intercept_range
370 .end
371 .min(log_sigma_design.design.ncols());
372 let transform = build_scale_deviation_transform_design(
373 &threshold_design.design,
374 &log_sigma_design.design,
375 weights,
376 non_intercept_start,
377 )?;
378 build_scale_deviation_operator(
379 threshold_design.design.clone(),
380 log_sigma_design.design.clone(),
381 &transform,
382 )
383}
384
385pub(crate) fn identity_penalty(dim: usize) -> Array2<f64> {
386 let mut penalty = Array2::<f64>::zeros((dim, dim));
387 for i in 0..dim {
388 penalty[[i, i]] = 1.0;
389 }
390 penalty
391}
392
393pub(crate) fn penalty_nullspace_projector(penalties: &[PenaltyMatrix], dim: usize) -> Array2<f64> {
416 use gam_linalg::faer_ndarray::FaerEigh;
417 use faer::Side;
418
419 if dim == 0 {
420 return Array2::<f64>::zeros((0, 0));
421 }
422 let mut combined = Array2::<f64>::zeros((dim, dim));
425 for pen in penalties {
426 let dense = pen.to_dense();
427 assert_eq!(
428 dense.nrows(),
429 dim,
430 "scale penalty block dim {} != scale design cols {dim}",
431 dense.nrows()
432 );
433 if dense.nrows() == dim && dense.ncols() == dim {
434 combined += &dense;
435 }
436 }
437 let combined_sym = 0.5 * (&combined + &combined.t());
439 let (eigvals, eigvecs) = match combined_sym.eigh(Side::Lower) {
440 Ok(decomp) => decomp,
441 Err(_) => return identity_penalty(dim),
445 };
446 let max_eig = eigvals.iter().cloned().fold(0.0_f64, f64::max);
452 let tol = (max_eig * 1e-8).max(1e-12);
453 let mut projector = Array2::<f64>::zeros((dim, dim));
454 for (j, &lambda) in eigvals.iter().enumerate() {
455 if lambda <= tol {
456 let v = eigvecs.column(j);
457 for a in 0..dim {
459 let va = v[a];
460 if va == 0.0 {
461 continue;
462 }
463 for b in 0..dim {
464 projector[[a, b]] += va * v[b];
465 }
466 }
467 }
468 }
469 projector
470}
471
472pub(crate) fn append_binomial_log_sigma_shrinkage_penalty_design(
473 design: &mut TermCollectionDesign,
474) {
475 let p = design.design.ncols();
476 design
477 .penalties
478 .push(BlockwisePenalty::new(0..p, identity_penalty(p)));
479 design.nullspace_dims.push(0);
481 design.penaltyinfo.push(PenaltyBlockInfo {
482 global_index: design.penaltyinfo.len(),
483 termname: Some("log_sigma_shrinkage".to_string()),
484 penalty: PenaltyInfo {
485 source: PenaltySource::Other("shrinkage".to_string()),
486 original_index: 0,
487 active: true,
488 effective_rank: p,
489 dropped_reason: None,
490 nullspace_dim_hint: 0,
491 normalization_scale: 1.0,
492 kronecker_factors: None,
493 },
494 });
495}
496
497pub(crate) fn build_gaussian_mean_and_scale_blocks(
504 y: &Array1<f64>,
505 weights: &Array1<f64>,
506 mean_design: &TermCollectionDesign,
507 noise_design: &TermCollectionDesign,
508 mean_offset: &Array1<f64>,
509 noise_offset: &Array1<f64>,
510 mean_log_lambdas: Array1<f64>,
511 noise_log_lambdas: Array1<f64>,
512 mean_beta_hint: Option<Array1<f64>>,
513 noise_beta_hint: Option<Array1<f64>>,
514 context: &str,
515) -> Result<(ParameterBlockSpec, ParameterBlockSpec), String> {
516 let mut meanspec = build_location_scale_block(
517 "mu",
518 mean_design.design.clone(),
519 mean_offset.clone(),
520 mean_design.penalties_as_penalty_matrix(),
521 mean_design.nullspace_dims.clone(),
522 mean_log_lambdas,
523 mean_beta_hint,
524 0,
525 LOCATION_SCALE_N_OUTPUTS,
526 &format!("{context}: mu"),
527 )?;
528 let prepared_noise_design =
529 prepared_gaussian_log_sigma_design(&mean_design.design, &noise_design.design)?;
530 let p_noise = prepared_noise_design.ncols();
531 let mut log_sigma_penalty_matrices = noise_design.penalties_as_penalty_matrix();
532 let shrinkage = penalty_nullspace_projector(&log_sigma_penalty_matrices, p_noise);
538 let shrinkage_rank = (0..p_noise).map(|i| shrinkage[[i, i]]).sum::<f64>().round() as usize;
548 log_sigma_penalty_matrices.push(PenaltyMatrix::Dense(shrinkage));
549 let mut log_sigma_nullspace_dims = noise_design.nullspace_dims.clone();
550 log_sigma_nullspace_dims.push(p_noise.saturating_sub(shrinkage_rank));
553 let mut noisespec = build_location_scale_block(
554 "log_sigma",
555 prepared_noise_design,
556 noise_offset.clone(),
557 log_sigma_penalty_matrices,
558 log_sigma_nullspace_dims,
559 noise_log_lambdas,
560 noise_beta_hint,
561 1,
562 LOCATION_SCALE_N_OUTPUTS,
563 &format!("{context}: log_sigma"),
564 )?;
565 if meanspec.initial_beta.is_none() || noisespec.initial_beta.is_none() {
566 let (betamu0, beta_ls0, _) = gaussian_location_scalewarm_start(
567 y,
568 weights,
569 &meanspec,
570 &noisespec,
571 1e-10,
572 meanspec.initial_beta.as_ref(),
573 noisespec.initial_beta.as_ref(),
574 )?;
575 if meanspec.initial_beta.is_none() {
576 meanspec.initial_beta = Some(betamu0);
577 }
578 if noisespec.initial_beta.is_none() {
579 noisespec.initial_beta = Some(beta_ls0);
580 }
581 }
582 Ok((meanspec, noisespec))
583}
584
585pub(crate) fn build_binomial_threshold_and_scale_blocks(
591 y: &Array1<f64>,
592 weights: &Array1<f64>,
593 link_kind: &InverseLink,
594 mean_design: &TermCollectionDesign,
595 noise_design: &TermCollectionDesign,
596 mean_offset: &Array1<f64>,
597 noise_offset: &Array1<f64>,
598 mean_log_lambdas: Array1<f64>,
599 noise_log_lambdas: Array1<f64>,
600 mean_beta_hint: Option<Array1<f64>>,
601 noise_beta_hint: Option<Array1<f64>>,
602 context: &str,
603) -> Result<(ParameterBlockSpec, ParameterBlockSpec), String> {
604 let identifiednoise_design =
605 identified_binomial_log_sigma_design(mean_design, noise_design, weights)?;
606 let p_noise = identifiednoise_design.ncols();
607 let mut log_sigma_penalty_matrices: Vec<PenaltyMatrix> =
608 noise_design.penalties_as_penalty_matrix();
609 log_sigma_penalty_matrices.push(PenaltyMatrix::Dense(identity_penalty(p_noise)));
610 let mut thresholdspec = build_location_scale_block(
611 "threshold",
612 mean_design.design.clone(),
613 mean_offset.clone(),
614 mean_design.penalties_as_penalty_matrix(),
615 vec![],
616 mean_log_lambdas,
617 mean_beta_hint,
618 0,
619 LOCATION_SCALE_N_OUTPUTS,
620 &format!("{context}: threshold"),
621 )?;
622 let mut log_sigmaspec = build_location_scale_block(
623 "log_sigma",
624 identifiednoise_design,
625 noise_offset.clone(),
626 log_sigma_penalty_matrices,
627 vec![],
628 noise_log_lambdas,
629 noise_beta_hint,
630 1,
631 LOCATION_SCALE_N_OUTPUTS,
632 &format!("{context}: log_sigma"),
633 )?;
634 if thresholdspec.initial_beta.is_none() || log_sigmaspec.initial_beta.is_none() {
635 let (beta_t0, beta_ls0) = binomial_location_scalewarm_start(
636 y,
637 weights,
638 link_kind,
639 &thresholdspec,
640 &log_sigmaspec,
641 thresholdspec.initial_beta.as_ref(),
642 log_sigmaspec.initial_beta.as_ref(),
643 )?;
644 if thresholdspec.initial_beta.is_none() {
645 thresholdspec.initial_beta = Some(beta_t0);
646 }
647 if log_sigmaspec.initial_beta.is_none() {
648 log_sigmaspec.initial_beta = Some(beta_ls0);
649 }
650 }
651 Ok((thresholdspec, log_sigmaspec))
652}
653
654pub(crate) fn wiggle_block_penalty_matrices(
658 wiggle_block: &ParameterBlockInput,
659) -> Vec<PenaltyMatrix> {
660 let p_wiggle = wiggle_block.design.ncols();
661 wiggle_block
662 .penalties
663 .iter()
664 .map(|spec| match spec {
665 crate::model_types::PenaltySpec::Block {
666 local, col_range, ..
667 } => PenaltyMatrix::Blockwise {
668 local: local.clone(),
669 col_range: col_range.clone(),
670 total_dim: p_wiggle,
671 },
672 crate::model_types::PenaltySpec::Dense(m)
673 | crate::model_types::PenaltySpec::DenseWithMean { matrix: m, .. } => {
674 PenaltyMatrix::Dense(m.clone())
675 }
676 })
677 .collect()
678}
679
680pub(crate) fn binomial_location_scale_link_eta_from_probability(
681 link_kind: &InverseLink,
682 probability: f64,
683) -> Result<f64, String> {
684 let target = probability.clamp(1e-6, 1.0 - 1e-6);
685 match link_kind {
686 InverseLink::Standard(StandardLink::Logit) => Ok((target / (1.0 - target)).ln()),
687 InverseLink::Standard(StandardLink::Probit) => standard_normal_quantile(target)
688 .map_err(|err| format!("failed to invert probit warm-start probability: {err}")),
689 InverseLink::Standard(StandardLink::CLogLog) => Ok((-((1.0 - target).ln())).ln()),
690 other => Err(GamlssError::UnsupportedConfiguration { reason: format!(
691 "binomial location-scale warm start requires logit, probit, or cloglog link, got {other:?}"
692 ) }.into()),
693 }
694}
695
696pub(crate) fn weighted_binomial_prevalence(
697 y: &Array1<f64>,
698 weights: &Array1<f64>,
699) -> Result<f64, String> {
700 if y.len() != weights.len() {
701 return Err(GamlssError::DimensionMismatch { reason: format!(
702 "binomial location-scale warm start dimension mismatch: y has length {}, weights have length {}",
703 y.len(),
704 weights.len()
705 ) }.into());
706 }
707 let mut weight_sum = 0.0;
708 let mut success_sum = 0.0;
709 for (&yi, &wi) in y.iter().zip(weights.iter()) {
710 if !yi.is_finite() {
711 return Err(GamlssError::NonFinite {
712 reason: format!(
713 "binomial location-scale warm start encountered non-finite response {yi}"
714 ),
715 }
716 .into());
717 }
718 let weight = floor_positiveweight(wi, MIN_WEIGHT);
719 if weight > 0.0 {
720 weight_sum += weight;
721 success_sum += weight * yi;
722 }
723 }
724 if !weight_sum.is_finite() || weight_sum <= 0.0 {
725 return Err(
726 "binomial location-scale warm start requires positive total weight".to_string(),
727 );
728 }
729 Ok(success_sum / weight_sum)
730}
731
732pub(crate) fn project_constant_eta_into_block(
733 block: &ParameterBlockSpec,
734 weights: &Array1<f64>,
735 eta: f64,
736) -> Result<Array1<f64>, String> {
737 let target_eta = Array1::from_elem(block.design.nrows(), eta);
738 solve_penalizedweighted_projection(
739 &block.design,
740 &block.offset,
741 &target_eta,
742 weights,
743 &block.penalties,
744 &block.initial_log_lambdas,
745 1e-10,
746 )
747}
748
749pub(crate) fn binomial_location_scalewarm_start(
753 y: &Array1<f64>,
754 weights: &Array1<f64>,
755 link_kind: &InverseLink,
756 threshold_block: &ParameterBlockSpec,
757 log_sigma_block: &ParameterBlockSpec,
758 mean_beta_hint: Option<&Array1<f64>>,
759 noise_beta_hint: Option<&Array1<f64>>,
760) -> Result<(Array1<f64>, Array1<f64>), String> {
761 if let (Some(mean_beta), Some(noise_beta)) = (mean_beta_hint, noise_beta_hint) {
762 return Ok((mean_beta.clone(), noise_beta.clone()));
763 }
764
765 let beta_threshold = match mean_beta_hint {
766 Some(beta) => beta.clone(),
767 None => {
768 let prevalence = weighted_binomial_prevalence(y, weights)?;
769 let eta = binomial_location_scale_link_eta_from_probability(link_kind, prevalence)?;
770 project_constant_eta_into_block(threshold_block, weights, eta)?
771 }
772 };
773 let beta_log_sigma = match noise_beta_hint {
774 Some(beta) => beta.clone(),
775 None => project_constant_eta_into_block(log_sigma_block, weights, 0.0)?,
776 };
777 Ok((beta_threshold, beta_log_sigma))
778}
779
780#[derive(Clone)]
781pub(crate) struct BinomialMeanWiggleSpec {
782 pub y: Array1<f64>,
783 pub weights: Array1<f64>,
784 pub link_kind: InverseLink,
785 pub wiggle_knots: Array1<f64>,
786 pub wiggle_degree: usize,
787 pub eta_block: ParameterBlockInput,
788 pub wiggle_block: ParameterBlockInput,
789}
790
791#[derive(Clone)]
792pub struct GaussianLocationScaleTermSpec {
793 pub y: Array1<f64>,
794 pub weights: Array1<f64>,
795 pub meanspec: TermCollectionSpec,
796 pub log_sigmaspec: TermCollectionSpec,
797 pub mean_offset: Array1<f64>,
798 pub log_sigma_offset: Array1<f64>,
799}
800
801#[derive(Clone)]
802pub struct GaussianLocationScaleWiggleTermSpec {
803 pub y: Array1<f64>,
804 pub weights: Array1<f64>,
805 pub meanspec: TermCollectionSpec,
806 pub log_sigmaspec: TermCollectionSpec,
807 pub mean_offset: Array1<f64>,
808 pub log_sigma_offset: Array1<f64>,
809 pub wiggle_knots: Array1<f64>,
810 pub wiggle_degree: usize,
811 pub wiggle_block: ParameterBlockInput,
812}
813
814#[derive(Clone)]
815pub struct BinomialLocationScaleTermSpec {
816 pub y: Array1<f64>,
817 pub weights: Array1<f64>,
818 pub link_kind: InverseLink,
819 pub thresholdspec: TermCollectionSpec,
820 pub log_sigmaspec: TermCollectionSpec,
821 pub threshold_offset: Array1<f64>,
822 pub log_sigma_offset: Array1<f64>,
823}
824
825#[derive(Clone)]
826pub struct BinomialLocationScaleWiggleTermSpec {
827 pub y: Array1<f64>,
828 pub weights: Array1<f64>,
829 pub link_kind: InverseLink,
830 pub thresholdspec: TermCollectionSpec,
831 pub log_sigmaspec: TermCollectionSpec,
832 pub threshold_offset: Array1<f64>,
833 pub log_sigma_offset: Array1<f64>,
834 pub wiggle_knots: Array1<f64>,
835 pub wiggle_degree: usize,
836 pub wiggle_block: ParameterBlockInput,
837}
838
839#[derive(Clone, Debug)]
840pub struct BlockwiseTermFitResult {
841 pub fit: UnifiedFitResult,
842 pub meanspec_resolved: TermCollectionSpec,
843 pub noisespec_resolved: TermCollectionSpec,
844 pub mean_design: TermCollectionDesign,
845 pub noise_design: TermCollectionDesign,
846}
847
848pub(crate) struct BlockwiseTermFitResultParts {
849 pub fit: UnifiedFitResult,
850 pub meanspec_resolved: TermCollectionSpec,
851 pub noisespec_resolved: TermCollectionSpec,
852 pub mean_design: TermCollectionDesign,
853 pub noise_design: TermCollectionDesign,
854}
855
856pub struct BlockwiseTermWiggleFitResult {
857 pub fit: BlockwiseTermFitResult,
858 pub wiggle_knots: Array1<f64>,
859 pub wiggle_degree: usize,
860}
861
862pub struct BinomialMeanWiggleTermFitResult {
863 pub fit: UnifiedFitResult,
864 pub resolvedspec: TermCollectionSpec,
865 pub design: TermCollectionDesign,
866 pub wiggle_knots: Array1<f64>,
867 pub wiggle_degree: usize,
868}
869
870pub(crate) struct BlockwiseTermWiggleFitResultParts {
871 pub fit: BlockwiseTermFitResult,
872 pub wiggle_knots: Array1<f64>,
873 pub wiggle_degree: usize,
874}
875
876pub(crate) fn validate_term_collection_design(
877 label: &str,
878 design: &TermCollectionDesign,
879) -> Result<(), String> {
880 let p = design.design.ncols();
881 let n = design.design.nrows();
882 for rows in exact_design_row_chunks(n, p) {
883 let chunk = design
884 .design
885 .try_row_chunk(rows)
886 .map_err(|e| format!("{label}.design row chunk materialization failed: {e}"))?;
887 validate_all_finite_estimation(&format!("{label}.design"), chunk.iter().copied())
888 .map_err(|e| e.to_string())?;
889 }
890 if design.nullspace_dims.len() != design.penalties.len() {
891 return Err(GamlssError::DimensionMismatch {
892 reason: format!(
893 "{label}.nullspace_dims length mismatch: got {}, expected {}",
894 design.nullspace_dims.len(),
895 design.penalties.len()
896 ),
897 }
898 .into());
899 }
900 if design.penaltyinfo.len() != design.penalties.len() {
901 return Err(GamlssError::DimensionMismatch {
902 reason: format!(
903 "{label}.penaltyinfo length mismatch: got {}, expected {}",
904 design.penaltyinfo.len(),
905 design.penalties.len()
906 ),
907 }
908 .into());
909 }
910 for (idx, bp) in design.penalties.iter().enumerate() {
911 validate_all_finite_estimation(
912 &format!("{label}.penalties[{idx}]"),
913 bp.local.iter().copied(),
914 )
915 .map_err(|e| e.to_string())?;
916 if bp.col_range.end > p {
917 return Err(GamlssError::DimensionMismatch {
918 reason: format!(
919 "{label}.penalties[{idx}] col_range {}..{} exceeds design width {}",
920 bp.col_range.start, bp.col_range.end, p
921 ),
922 }
923 .into());
924 }
925 }
926 if let Some(bounds) = design.coefficient_lower_bounds.as_ref() {
927 if bounds.len() != p {
928 return Err(GamlssError::ConstraintViolation {
929 reason: format!(
930 "{label}.coefficient_lower_bounds length mismatch: got {}, expected {p}",
931 bounds.len()
932 ),
933 }
934 .into());
935 }
936 for (idx, &bound) in bounds.iter().enumerate() {
937 if !(bound.is_finite() || bound == f64::NEG_INFINITY) {
938 return Err(GamlssError::NonFinite { reason: format!(
939 "{label}.coefficient_lower_bounds[{idx}] must be finite or -inf, got {bound}",
940 ) }.into());
941 }
942 }
943 }
944 if let Some(constraints) = design.linear_constraints.as_ref() {
945 validate_all_finite_estimation(
946 &format!("{label}.linear_constraints.a"),
947 constraints.a.iter().copied(),
948 )
949 .map_err(|e| e.to_string())?;
950 validate_all_finite_estimation(
951 &format!("{label}.linear_constraints.b"),
952 constraints.b.iter().copied(),
953 )
954 .map_err(|e| e.to_string())?;
955 if constraints.a.ncols() != p {
956 return Err(GamlssError::DimensionMismatch {
957 reason: format!(
958 "{label}.linear_constraints.a column mismatch: got {}, expected {p}",
959 constraints.a.ncols()
960 ),
961 }
962 .into());
963 }
964 if constraints.a.nrows() != constraints.b.len() {
965 return Err(GamlssError::DimensionMismatch {
966 reason: format!(
967 "{label}.linear_constraints row mismatch: a has {}, b has {}",
968 constraints.a.nrows(),
969 constraints.b.len()
970 ),
971 }
972 .into());
973 }
974 }
975 if design.intercept_range.start > design.intercept_range.end || design.intercept_range.end > p {
976 return Err(GamlssError::ConstraintViolation {
977 reason: format!(
978 "{label}.intercept_range out of bounds: {:?} for {} columns",
979 design.intercept_range, p
980 ),
981 }
982 .into());
983 }
984 Ok(())
985}
986
987impl BlockwiseTermFitResult {
988 pub(crate) fn try_from_parts(parts: BlockwiseTermFitResultParts) -> Result<Self, String> {
989 let BlockwiseTermFitResultParts {
990 fit,
991 meanspec_resolved,
992 noisespec_resolved,
993 mean_design,
994 noise_design,
995 } = parts;
996
997 fit.validate_numeric_finiteness()
998 .map_err(|e| format!("{e}"))?;
999 if fit.block_states.len() < 2 {
1000 return Err(GamlssError::DimensionMismatch {
1001 reason: format!(
1002 "BlockwiseTermFitResult requires at least 2 block states, got {}",
1003 fit.block_states.len()
1004 ),
1005 }
1006 .into());
1007 }
1008 validate_term_collection_design("blockwise_term.mean_design", &mean_design)?;
1009 validate_term_collection_design("blockwise_term.noise_design", &noise_design)?;
1010 if mean_design.design.nrows() != noise_design.design.nrows() {
1011 return Err(GamlssError::DimensionMismatch {
1012 reason: format!(
1013 "BlockwiseTermFitResult row mismatch: mean_design={}, noise_design={}",
1014 mean_design.design.nrows(),
1015 noise_design.design.nrows()
1016 ),
1017 }
1018 .into());
1019 }
1020 if fit.block_states[0].beta.len() != mean_design.design.ncols() {
1021 return Err(GamlssError::DimensionMismatch {
1022 reason: format!(
1023 "BlockwiseTermFitResult mean beta length mismatch: got {}, expected {}",
1024 fit.block_states[0].beta.len(),
1025 mean_design.design.ncols()
1026 ),
1027 }
1028 .into());
1029 }
1030 if fit.block_states[1].beta.len() != noise_design.design.ncols() {
1031 return Err(GamlssError::DimensionMismatch {
1032 reason: format!(
1033 "BlockwiseTermFitResult noise beta length mismatch: got {}, expected {}",
1034 fit.block_states[1].beta.len(),
1035 noise_design.design.ncols()
1036 ),
1037 }
1038 .into());
1039 }
1040 if fit.block_states[0].eta.len() != mean_design.design.nrows() {
1041 return Err(GamlssError::DimensionMismatch {
1042 reason: format!(
1043 "BlockwiseTermFitResult mean eta length mismatch: got {}, expected {}",
1044 fit.block_states[0].eta.len(),
1045 mean_design.design.nrows()
1046 ),
1047 }
1048 .into());
1049 }
1050 if fit.block_states[1].eta.len() != noise_design.design.nrows() {
1051 return Err(GamlssError::DimensionMismatch {
1052 reason: format!(
1053 "BlockwiseTermFitResult noise eta length mismatch: got {}, expected {}",
1054 fit.block_states[1].eta.len(),
1055 noise_design.design.nrows()
1056 ),
1057 }
1058 .into());
1059 }
1060
1061 Ok(Self {
1062 fit,
1063 meanspec_resolved,
1064 noisespec_resolved,
1065 mean_design,
1066 noise_design,
1067 })
1068 }
1069
1070 pub(crate) fn validate_numeric_finiteness(&self) -> Result<(), String> {
1071 Self::try_from_parts(BlockwiseTermFitResultParts {
1072 fit: self.fit.clone(),
1073 meanspec_resolved: self.meanspec_resolved.clone(),
1074 noisespec_resolved: self.noisespec_resolved.clone(),
1075 mean_design: self.mean_design.clone(),
1076 noise_design: self.noise_design.clone(),
1077 })
1078 .map(|_| ())
1079 }
1080}
1081
1082impl BlockwiseTermWiggleFitResult {
1083 pub(crate) fn try_from_parts(parts: BlockwiseTermWiggleFitResultParts) -> Result<Self, String> {
1084 let BlockwiseTermWiggleFitResultParts {
1085 fit,
1086 wiggle_knots,
1087 wiggle_degree,
1088 } = parts;
1089
1090 fit.validate_numeric_finiteness()
1091 .map_err(|e| e.to_string())?;
1092 if fit.fit.block_states.len() < 3 {
1093 return Err(GamlssError::DimensionMismatch {
1094 reason: format!(
1095 "BlockwiseTermWiggleFitResult requires at least 3 block states, got {}",
1096 fit.fit.block_states.len()
1097 ),
1098 }
1099 .into());
1100 }
1101 if wiggle_knots.is_empty() {
1102 return Err(GamlssError::UnsupportedConfiguration {
1103 reason: "BlockwiseTermWiggleFitResult requires non-empty wiggle_knots".to_string(),
1104 }
1105 .into());
1106 }
1107 validate_all_finite_estimation(
1108 "blockwise_term_wiggle.wiggle_knots",
1109 wiggle_knots.iter().copied(),
1110 )
1111 .map_err(|e| e.to_string())?;
1112
1113 Ok(Self {
1114 fit,
1115 wiggle_knots,
1116 wiggle_degree,
1117 })
1118 }
1119}
1120
1121pub struct BinomialLocationScaleFitResult {
1122 pub fit: BlockwiseTermFitResult,
1123 pub wiggle_knots: Option<Array1<f64>>,
1124 pub wiggle_degree: Option<usize>,
1125 pub beta_link_wiggle: Option<Vec<f64>>,
1126}
1127
1128pub struct GaussianLocationScaleFitResult {
1129 pub fit: BlockwiseTermFitResult,
1130 pub wiggle_knots: Option<Array1<f64>>,
1131 pub wiggle_degree: Option<usize>,
1132 pub beta_link_wiggle: Option<Vec<f64>>,
1133 pub response_scale: f64,
1162}
1163
1164pub(crate) fn fit_binomial_mean_wiggle(
1165 spec: BinomialMeanWiggleSpec,
1166 options: &BlockwiseFitOptions,
1167) -> Result<UnifiedFitResult, String> {
1168 let n = spec.y.len();
1169 validate_len_match("weights vs y", n, spec.weights.len())?;
1170 validateweights(&spec.weights, "fit_binomial_mean_wiggle")?;
1171 validate_binomial_response(&spec.y, "fit_binomial_mean_wiggle")?;
1172 validate_blockrows("eta", n, &spec.eta_block)?;
1173 validate_blockrows("wiggle", n, &spec.wiggle_block)?;
1174 if matches!(
1175 spec.link_kind,
1176 InverseLink::Standard(StandardLink::Identity)
1177 ) {
1178 return Err(GamlssError::UnsupportedConfiguration {
1179 reason: "fit_binomial_mean_wiggle does not support identity link".to_string(),
1180 }
1181 .into());
1182 }
1183 gam_terms::inference::formula_dsl::require_binomial_inverse_link_supports_joint_wiggle(
1184 &spec.link_kind,
1185 "fit_binomial_mean_wiggle",
1186 )?;
1187 if spec.wiggle_degree < 2 {
1188 return Err(GamlssError::ConstraintViolation {
1189 reason: format!(
1190 "fit_binomial_mean_wiggle: wiggle_degree must be >= 2, got {}",
1191 spec.wiggle_degree
1192 ),
1193 }
1194 .into());
1195 }
1196 let minimum_knots = minimum_monotone_wiggle_knot_count(spec.wiggle_degree)?;
1197 if spec.wiggle_knots.len() < minimum_knots {
1198 return Err(GamlssError::DimensionMismatch { reason: format!(
1199 "fit_binomial_mean_wiggle: wiggle_knots length {} is too short for degree {} (need at least {})",
1200 spec.wiggle_knots.len(),
1201 spec.wiggle_degree,
1202 minimum_knots
1203 ) }.into());
1204 }
1205
1206 let family = BinomialMeanWiggleFamily {
1207 y: spec.y,
1208 weights: spec.weights,
1209 link_kind: spec.link_kind,
1210 wiggle_knots: spec.wiggle_knots,
1211 wiggle_degree: spec.wiggle_degree,
1212 policy: gam_runtime::resource::ResourcePolicy::default_library(),
1213 };
1214 let blocks = vec![
1215 spec.eta_block
1227 .intospec_with_gauge_priority("eta", LINK_WIGGLE_GAUGE_PRIORITY)?,
1228 spec.wiggle_block.intospec("wiggle")?,
1229 ];
1230 fit_custom_family(&family, &blocks, options).map_err(|e| e.to_string())
1231}
1232
1233pub(crate) trait LocationScaleFamilyBuilder {
1234 type Family: CustomFamily + Clone + Send + Sync + 'static;
1235
1236 fn meanspec(&self) -> &TermCollectionSpec;
1237 fn noisespec(&self) -> &TermCollectionSpec;
1238
1239 fn build_blocks(
1240 &self,
1241 theta: &Array1<f64>,
1242 mean_design: &TermCollectionDesign,
1243 noise_design: &TermCollectionDesign,
1244 mean_beta_hint: Option<Array1<f64>>,
1245 noise_beta_hint: Option<Array1<f64>>,
1246 ) -> Result<Vec<ParameterBlockSpec>, String>;
1247
1248 fn build_family(
1249 &self,
1250 mean_design: &TermCollectionDesign,
1251 noise_design: &TermCollectionDesign,
1252 ) -> Self::Family;
1253
1254 fn extract_primary_betas(
1255 &self,
1256 fit: &UnifiedFitResult,
1257 ) -> Result<(Array1<f64>, Array1<f64>), String>;
1258
1259 fn mean_penalty_count(&self, mean_design: &TermCollectionDesign) -> usize {
1260 mean_design.penalties.len()
1261 }
1262
1263 fn noise_penalty_count(&self, noise_design: &TermCollectionDesign) -> usize {
1264 noise_design.penalties.len()
1265 }
1266
1267 fn exact_spatial_joint_supported(&self) -> bool {
1268 false
1269 }
1270
1271 fn require_exact_spatial_joint(&self) -> bool {
1272 false
1273 }
1274
1275 fn exact_spatial_seed_risk_profile(&self) -> crate::seeding::SeedRiskProfile {
1276 crate::seeding::SeedRiskProfile::GeneralizedLinear
1277 }
1278
1279 fn extra_rho0(&self) -> Result<Array1<f64>, String> {
1280 Ok(Array1::zeros(0))
1281 }
1282
1283 fn build_psiderivative_blocks(
1284 &self,
1285 arr: ndarray::ArrayView2<'_, f64>,
1286 term_spec: &TermCollectionSpec,
1287 term_spec2: &TermCollectionSpec,
1288 term_design: &TermCollectionDesign,
1289 term_design2: &TermCollectionDesign,
1290 ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String>;
1291}
1292
1293pub(crate) fn fit_location_scale_terms<B: LocationScaleFamilyBuilder>(
1294 data: ndarray::ArrayView2<'_, f64>,
1295 builder: B,
1296 options: &BlockwiseFitOptions,
1297 kappa_options: &SpatialLengthScaleOptimizationOptions,
1298) -> Result<BlockwiseTermFitResult, String> {
1299 let mut mean_beta_hint: Option<Array1<f64>> = None;
1305 let mut noise_beta_hint: Option<Array1<f64>> = None;
1306 let extra_rho0 = builder.extra_rho0()?;
1307
1308 let mean_boot_design =
1309 build_term_collection_design(data, builder.meanspec()).map_err(|e| e.to_string())?;
1310 let noise_boot_design =
1311 build_term_collection_design(data, builder.noisespec()).map_err(|e| e.to_string())?;
1312 let mean_bootspec = freeze_term_collection_from_design(builder.meanspec(), &mean_boot_design)
1313 .map_err(|e| e.to_string())?;
1314 let noise_bootspec =
1315 freeze_term_collection_from_design(builder.noisespec(), &noise_boot_design)
1316 .map_err(|e| e.to_string())?;
1317
1318 let require_exact_spatial_joint = builder.require_exact_spatial_joint();
1319 let analytic_joint_derivatives_check = if builder.exact_spatial_joint_supported() {
1320 builder
1321 .build_psiderivative_blocks(
1322 data,
1323 &mean_bootspec,
1324 &noise_bootspec,
1325 &mean_boot_design,
1326 &noise_boot_design,
1327 )
1328 .map(|_| ())
1329 } else {
1330 Err(
1331 "analytic spatial psi derivatives are unavailable for this location-scale family"
1332 .to_string(),
1333 )
1334 };
1335 let analytic_joint_derivatives_available = analytic_joint_derivatives_check.is_ok();
1336 if require_exact_spatial_joint {
1337 analytic_joint_derivatives_check.map_err(|err| {
1338 format!("exact two-block spatial path requires analytic psi derivatives: {err}")
1339 })?;
1340 }
1341 let mean_penalty_count = builder.mean_penalty_count(&mean_boot_design);
1342 let noise_penalty_count = builder.noise_penalty_count(&noise_boot_design);
1343
1344 let mut effective_kappa_options = kappa_options.clone();
1355 if effective_kappa_options.enabled
1356 && gam_terms::smooth::all_spatial_terms_kappa_fixed(&mean_bootspec)
1357 && gam_terms::smooth::all_spatial_terms_kappa_fixed(&noise_bootspec)
1358 {
1359 log::info!(
1360 "[GAMLSS spatial] disabling κ/ψ optimization: every spatial term in \
1361 both blocks has an explicit length_scale and no anisotropy; \
1362 user-supplied kernel scale is fixed"
1363 );
1364 effective_kappa_options.enabled = false;
1365 }
1366 let kappa_options: &SpatialLengthScaleOptimizationOptions = &effective_kappa_options;
1367
1368 macro_rules! run_exact_joint_spatial {
1372 () => {{
1373 let joint_setup = build_two_block_exact_joint_setup(
1374 data,
1375 builder.meanspec(),
1376 builder.noisespec(),
1377 mean_penalty_count,
1378 noise_penalty_count,
1379 extra_rho0.as_slice().unwrap_or(&[]),
1380 None,
1381 kappa_options,
1382 );
1383 let mean_terms = spatial_length_scale_term_indices(builder.meanspec());
1384 let noise_terms = spatial_length_scale_term_indices(builder.noisespec());
1385 let mean_beta_hint_cell = std::cell::RefCell::new(mean_beta_hint.clone());
1386 let noise_beta_hint_cell = std::cell::RefCell::new(noise_beta_hint.clone());
1387 let hyper_warm_start_cell =
1388 std::cell::RefCell::new(None::<CustomFamilyWarmStart>);
1389 let gamlss_disable_fixed_point = true;
1399 let outer_policy = {
1400 let theta_seed = joint_setup.theta0();
1412 let rho_dim = joint_setup.rho_dim();
1413 let psi_dim = theta_seed.len() - rho_dim;
1414 let rho_seed = theta_seed.slice(s![..rho_dim]).to_owned();
1415 let policy_blocks_res = builder.build_blocks(
1416 &rho_seed,
1417 &mean_boot_design,
1418 &noise_boot_design,
1419 mean_beta_hint_cell.borrow().clone(),
1420 noise_beta_hint_cell.borrow().clone(),
1421 );
1422 let mut policy = match policy_blocks_res {
1423 Ok(policy_blocks) => {
1424 let policy_family =
1425 builder.build_family(&mean_boot_design, &noise_boot_design);
1426 crate::custom_family::CustomFamily::outer_derivative_policy(
1427 &policy_family,
1428 &policy_blocks,
1429 psi_dim,
1430 options,
1431 )
1432 }
1433 Err(err) => {
1434 log::warn!(
1442 "[GAMLSS spatial] failed to realize policy blocks at seed rho ({err}); \
1443 routing outer optimizer through gradient-only BFGS"
1444 );
1445 let capability = if analytic_joint_derivatives_available {
1446 crate::custom_family::ExactOuterDerivativeOrder::Second
1447 } else {
1448 crate::custom_family::ExactOuterDerivativeOrder::First
1449 };
1450 crate::custom_family::OuterDerivativePolicy {
1451 capability,
1452 predicted_gradient_work: u128::MAX,
1453 predicted_hessian_work: u128::MAX,
1454 subsample_capable: false,
1459 }
1460 }
1461 };
1462 if !analytic_joint_derivatives_available {
1463 policy.capability =
1467 crate::custom_family::ExactOuterDerivativeOrder::First;
1468 }
1469 policy
1470 };
1471 optimize_spatial_length_scale_exact_joint(
1472 data,
1473 &[builder.meanspec().clone(), builder.noisespec().clone()],
1474 &[mean_terms, noise_terms],
1475 kappa_options,
1476 &joint_setup,
1477 builder.exact_spatial_seed_risk_profile(),
1478 analytic_joint_derivatives_available,
1479 analytic_joint_derivatives_available,
1480 gamlss_disable_fixed_point,
1481 None,
1482 outer_policy,
1483 |theta, specs: &[TermCollectionSpec], designs: &[TermCollectionDesign]| {
1484 assert_eq!(
1485 specs.len(),
1486 2,
1487 "joint spatial closure expects exactly two block specs (mean, noise); got {}",
1488 specs.len(),
1489 );
1490 assert_eq!(
1491 designs.len(),
1492 2,
1493 "joint spatial closure expects exactly two block designs (mean, noise); got {}",
1494 designs.len(),
1495 );
1496 let rho = theta.slice(s![..joint_setup.rho_dim()]).to_owned();
1497 let fit = {
1498 let blocks = builder.build_blocks(
1499 &rho,
1500 &designs[0],
1501 &designs[1],
1502 mean_beta_hint_cell.borrow().clone(),
1503 noise_beta_hint_cell.borrow().clone(),
1504 )?;
1505 if mean_beta_hint_cell.borrow().is_none()
1506 && let Some(beta) = blocks.first().and_then(|block| block.initial_beta.clone())
1507 {
1508 *mean_beta_hint_cell.borrow_mut() = Some(beta);
1509 }
1510 if noise_beta_hint_cell.borrow().is_none()
1511 && let Some(beta) =
1512 blocks.get(1).and_then(|block| block.initial_beta.clone())
1513 {
1514 *noise_beta_hint_cell.borrow_mut() = Some(beta);
1515 }
1516 let family = builder.build_family(&designs[0], &designs[1]);
1517 if joint_setup.log_kappa_dim() > 0 && kappa_options.enabled {
1539 let warm_start = hyper_warm_start_cell.borrow().clone();
1540 fit_custom_family_fixed_log_lambdas(
1541 &family,
1542 &blocks,
1543 options,
1544 warm_start.as_ref(),
1545 0,
1546 None,
1547 true,
1548 )?
1549 } else {
1550 fit_custom_family(&family, &blocks, options)?
1551 }
1552 };
1553 let (mean_beta, noise_beta) = builder.extract_primary_betas(&fit)?;
1554 mean_beta_hint = Some(mean_beta);
1555 noise_beta_hint = Some(noise_beta);
1556 *mean_beta_hint_cell.borrow_mut() = mean_beta_hint.clone();
1557 *noise_beta_hint_cell.borrow_mut() = noise_beta_hint.clone();
1558 Ok(fit)
1559 },
1560 |theta,
1561 specs: &[TermCollectionSpec],
1562 designs: &[TermCollectionDesign],
1563 eval_mode,
1564 row_set: &crate::row_kernel::RowSet| {
1565 use gam_problem::EvalMode;
1566 if !analytic_joint_derivatives_available {
1567 return Err(
1568 "analytic spatial psi derivatives are unavailable for this exact two-block path"
1569 .to_string(),
1570 );
1571 }
1572 let rho = theta.slice(s![..joint_setup.rho_dim()]).to_owned();
1573 let blocks = builder.build_blocks(
1574 &rho,
1575 &designs[0],
1576 &designs[1],
1577 mean_beta_hint_cell.borrow().clone(),
1578 noise_beta_hint_cell.borrow().clone(),
1579 )?;
1580 if mean_beta_hint_cell.borrow().is_none()
1581 && let Some(beta) = blocks.first().and_then(|block| block.initial_beta.clone())
1582 {
1583 *mean_beta_hint_cell.borrow_mut() = Some(beta);
1584 }
1585 if noise_beta_hint_cell.borrow().is_none()
1586 && let Some(beta) = blocks.get(1).and_then(|block| block.initial_beta.clone())
1587 {
1588 *noise_beta_hint_cell.borrow_mut() = Some(beta);
1589 }
1590 let family = builder.build_family(&designs[0], &designs[1]);
1591 let psiderivative_blocks = builder.build_psiderivative_blocks(
1592 data,
1593 &specs[0],
1594 &specs[1],
1595 &designs[0],
1596 &designs[1],
1597 )?;
1598 let warm_start = hyper_warm_start_cell.borrow().clone();
1599 let eval_options = match row_set {
1606 crate::row_kernel::RowSet::All => {
1607 std::borrow::Cow::Borrowed(options)
1608 }
1609 crate::row_kernel::RowSet::Subsample {
1610 rows,
1611 n_full,
1612 } => {
1613 let subsample = crate::outer_subsample::
1614 OuterScoreSubsample::from_weighted_rows(
1615 (**rows).clone(),
1616 *n_full,
1617 *n_full as u64,
1618 );
1619 let mut cloned = options.clone();
1620 cloned.outer_score_subsample =
1621 Some(std::sync::Arc::new(subsample));
1622 std::borrow::Cow::Owned(cloned)
1623 }
1624 };
1625 let eval = evaluate_custom_family_joint_hyper(
1626 &family,
1627 &blocks,
1628 eval_options.as_ref(),
1629 &rho,
1630 &psiderivative_blocks,
1631 warm_start.as_ref(),
1632 eval_mode,
1633 )?;
1634 *hyper_warm_start_cell.borrow_mut() = Some(eval.warm_start.clone());
1635 if !eval.inner_converged {
1636 return Err(
1637 "exact two-block spatial inner solve did not converge".to_string(),
1638 );
1639 }
1640 if matches!(eval_mode, EvalMode::ValueGradientHessian)
1641 && !eval.outer_hessian.is_analytic()
1642 {
1643 return Err(
1644 "exact two-block spatial objective requires a full joint [rho, psi] hessian"
1645 .to_string(),
1646 );
1647 }
1648 Ok((eval.objective, eval.gradient, eval.outer_hessian))
1649 },
1650 |theta, specs: &[TermCollectionSpec], designs: &[TermCollectionDesign]| {
1651 if !analytic_joint_derivatives_available {
1652 return Err(
1653 "analytic spatial psi derivatives are unavailable for this exact two-block path"
1654 .to_string(),
1655 );
1656 }
1657 let rho = theta.slice(s![..joint_setup.rho_dim()]).to_owned();
1658 let blocks = builder.build_blocks(
1659 &rho,
1660 &designs[0],
1661 &designs[1],
1662 mean_beta_hint_cell.borrow().clone(),
1663 noise_beta_hint_cell.borrow().clone(),
1664 )?;
1665 if mean_beta_hint_cell.borrow().is_none()
1666 && let Some(beta) = blocks.first().and_then(|block| block.initial_beta.clone())
1667 {
1668 *mean_beta_hint_cell.borrow_mut() = Some(beta);
1669 }
1670 if noise_beta_hint_cell.borrow().is_none()
1671 && let Some(beta) = blocks.get(1).and_then(|block| block.initial_beta.clone())
1672 {
1673 *noise_beta_hint_cell.borrow_mut() = Some(beta);
1674 }
1675 let family = builder.build_family(&designs[0], &designs[1]);
1676 let psiderivative_blocks = builder.build_psiderivative_blocks(
1677 data,
1678 &specs[0],
1679 &specs[1],
1680 &designs[0],
1681 &designs[1],
1682 )?;
1683 let warm_start = hyper_warm_start_cell.borrow().clone();
1684 let eval = evaluate_custom_family_joint_hyper_efs(
1685 &family,
1686 &blocks,
1687 options,
1688 &rho,
1689 &psiderivative_blocks,
1690 warm_start.as_ref(),
1691 )?;
1692 *hyper_warm_start_cell.borrow_mut() = Some(eval.warm_start.clone());
1693 if !eval.inner_converged {
1694 return Err(
1695 "exact two-block spatial EFS inner solve did not converge".to_string(),
1696 );
1697 }
1698 Ok(eval.efs_eval)
1699 },
1700 |_beta: &Array1<f64>| Ok(gam_solve::rho_optimizer::SeedOutcome::NoSlot),
1701 )
1702 }};
1703 }
1704
1705 let mut solved = run_exact_joint_spatial!()
1706 .map_err(|err| format!("exact two-block spatial optimization failed: {err}"))?;
1707
1708 let expected_noise_penalty_count = builder.noise_penalty_count(&solved.designs[1]);
1709 let actual_noise_penalty_count = solved.designs[1].penalties.len();
1710 if expected_noise_penalty_count > actual_noise_penalty_count {
1711 if expected_noise_penalty_count != actual_noise_penalty_count + 1 {
1712 return Err(GamlssError::UnsupportedConfiguration {
1713 reason: format!(
1714 "location-scale result noise design expected {} penalties after augmentation, got {} before augmentation",
1715 expected_noise_penalty_count, actual_noise_penalty_count
1716 ),
1717 }
1718 .into());
1719 }
1720 append_binomial_log_sigma_shrinkage_penalty_design(&mut solved.designs[1]);
1721 }
1722
1723 BlockwiseTermFitResult::try_from_parts(BlockwiseTermFitResultParts {
1724 fit: solved.fit,
1725 meanspec_resolved: solved.resolved_specs.remove(0),
1726 noisespec_resolved: solved.resolved_specs.remove(0),
1727 mean_design: solved.designs.remove(0),
1728 noise_design: solved.designs.remove(0),
1729 })
1730}
1731
1732pub(crate) struct GaussianLocationScaleTermBuilder {
1733 pub(crate) y: Array1<f64>,
1734 pub(crate) weights: Array1<f64>,
1735 pub(crate) meanspec: TermCollectionSpec,
1736 pub(crate) noisespec: TermCollectionSpec,
1737 pub(crate) mean_offset: Array1<f64>,
1738 pub(crate) noise_offset: Array1<f64>,
1739}
1740
1741impl LocationScaleFamilyBuilder for GaussianLocationScaleTermBuilder {
1742 type Family = GaussianLocationScaleFamily;
1743
1744 fn meanspec(&self) -> &TermCollectionSpec {
1745 &self.meanspec
1746 }
1747
1748 fn noisespec(&self) -> &TermCollectionSpec {
1749 &self.noisespec
1750 }
1751
1752 fn noise_penalty_count(&self, noise_design: &TermCollectionDesign) -> usize {
1753 noise_design.penalties.len() + 1
1762 }
1763
1764 fn exact_spatial_joint_supported(&self) -> bool {
1765 true
1766 }
1767
1768 fn exact_spatial_seed_risk_profile(&self) -> crate::seeding::SeedRiskProfile {
1769 crate::seeding::SeedRiskProfile::GaussianLocationScale
1770 }
1771
1772 fn build_blocks(
1773 &self,
1774 theta: &Array1<f64>,
1775 mean_design: &TermCollectionDesign,
1776 noise_design: &TermCollectionDesign,
1777 mean_beta_hint: Option<Array1<f64>>,
1778 noise_beta_hint: Option<Array1<f64>>,
1779 ) -> Result<Vec<ParameterBlockSpec>, String> {
1780 let layout = GamlssLambdaLayout::two_block(
1781 mean_design.penalties.len(),
1782 self.noise_penalty_count(noise_design),
1783 );
1784 layout.validate_theta_len(theta.len(), "gaussian location-scale")?;
1785 let (meanspec, noisespec) = build_gaussian_mean_and_scale_blocks(
1786 &self.y,
1787 &self.weights,
1788 mean_design,
1789 noise_design,
1790 &self.mean_offset,
1791 &self.noise_offset,
1792 layout.mean_from(theta),
1793 layout.noise_from(theta),
1794 mean_beta_hint,
1795 noise_beta_hint,
1796 "GaussianLocationScale::build_blocks",
1797 )?;
1798 Ok(vec![meanspec, noisespec])
1799 }
1800
1801 fn build_family(
1802 &self,
1803 mean_design: &TermCollectionDesign,
1804 noise_design: &TermCollectionDesign,
1805 ) -> Self::Family {
1806 let preparednoise_design =
1807 prepared_gaussian_log_sigma_design(&mean_design.design, &noise_design.design)
1808 .expect("prepared Gaussian log-sigma design should match block construction");
1809 GaussianLocationScaleFamily {
1810 y: self.y.clone(),
1811 weights: self.weights.clone(),
1812 mu_design: Some(mean_design.design.clone()),
1813 log_sigma_design: Some(preparednoise_design),
1814 policy: gam_runtime::resource::ResourcePolicy::default_library(),
1815 cached_row_scalars: std::sync::RwLock::new(None),
1816 }
1817 }
1818
1819 fn extract_primary_betas(
1820 &self,
1821 fit: &UnifiedFitResult,
1822 ) -> Result<(Array1<f64>, Array1<f64>), String> {
1823 let mean_beta = fit
1824 .block_states
1825 .get(GaussianLocationScaleFamily::BLOCK_MU)
1826 .ok_or_else(|| "missing Gaussian mu block state".to_string())?
1827 .beta
1828 .clone();
1829 let noise_beta = fit
1830 .block_states
1831 .get(GaussianLocationScaleFamily::BLOCK_LOG_SIGMA)
1832 .ok_or_else(|| "missing Gaussian log_sigma block state".to_string())?
1833 .beta
1834 .clone();
1835 Ok((mean_beta, noise_beta))
1836 }
1837
1838 fn build_psiderivative_blocks(
1839 &self,
1840 data: ndarray::ArrayView2<'_, f64>,
1841 meanspec_resolved: &TermCollectionSpec,
1842 noisespec_resolved: &TermCollectionSpec,
1843 mean_design: &TermCollectionDesign,
1844 noise_design: &TermCollectionDesign,
1845 ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String> {
1846 let mean_derivs =
1847 build_block_spatial_psi_derivatives(data, meanspec_resolved, mean_design)?
1848 .ok_or_else(|| "missing Gaussian mean spatial psi derivatives".to_string())?;
1849 let noise_derivs =
1850 build_block_spatial_psi_derivatives(data, noisespec_resolved, noise_design)?
1851 .ok_or_else(|| "missing Gaussian log-sigma spatial psi derivatives".to_string())?;
1852 Ok(vec![mean_derivs, noise_derivs])
1853 }
1854}
1855
1856pub(crate) struct GaussianLocationScaleWiggleTermBuilder {
1857 pub(crate) y: Array1<f64>,
1858 pub(crate) weights: Array1<f64>,
1859 pub(crate) meanspec: TermCollectionSpec,
1860 pub(crate) noisespec: TermCollectionSpec,
1861 pub(crate) mean_offset: Array1<f64>,
1862 pub(crate) noise_offset: Array1<f64>,
1863 pub(crate) wiggle_knots: Array1<f64>,
1864 pub(crate) wiggle_degree: usize,
1865 pub(crate) wiggle_block: ParameterBlockInput,
1866}
1867
1868impl LocationScaleFamilyBuilder for GaussianLocationScaleWiggleTermBuilder {
1869 type Family = GaussianLocationScaleWiggleFamily;
1870
1871 fn meanspec(&self) -> &TermCollectionSpec {
1872 &self.meanspec
1873 }
1874
1875 fn noisespec(&self) -> &TermCollectionSpec {
1876 &self.noisespec
1877 }
1878
1879 fn noise_penalty_count(&self, noise_design: &TermCollectionDesign) -> usize {
1880 noise_design.penalties.len() + 1
1883 }
1884
1885 fn exact_spatial_joint_supported(&self) -> bool {
1886 true
1887 }
1888
1889 fn exact_spatial_seed_risk_profile(&self) -> crate::seeding::SeedRiskProfile {
1890 crate::seeding::SeedRiskProfile::GaussianLocationScale
1891 }
1892
1893 fn require_exact_spatial_joint(&self) -> bool {
1894 true
1895 }
1896
1897 fn extra_rho0(&self) -> Result<Array1<f64>, String> {
1898 initial_log_lambdas_orzeros(&self.wiggle_block)
1899 }
1900
1901 fn build_blocks(
1902 &self,
1903 theta: &Array1<f64>,
1904 mean_design: &TermCollectionDesign,
1905 noise_design: &TermCollectionDesign,
1906 mean_beta_hint: Option<Array1<f64>>,
1907 noise_beta_hint: Option<Array1<f64>>,
1908 ) -> Result<Vec<ParameterBlockSpec>, String> {
1909 let layout = GamlssLambdaLayout::withwiggle(
1910 mean_design.penalties.len(),
1911 self.noise_penalty_count(noise_design),
1912 self.wiggle_block.penalties.len(),
1913 );
1914 layout.validate_theta_len(theta.len(), "gaussian location-scale wiggle")?;
1915 let (mut meanspec, mut noisespec) = build_gaussian_mean_and_scale_blocks(
1916 &self.y,
1917 &self.weights,
1918 mean_design,
1919 noise_design,
1920 &self.mean_offset,
1921 &self.noise_offset,
1922 layout.mean_from(theta),
1923 layout.noise_from(theta),
1924 mean_beta_hint,
1925 noise_beta_hint,
1926 "GaussianLocationScaleWiggle::build_blocks",
1927 )?;
1928 meanspec.gauge_priority = LINK_WIGGLE_GAUGE_PRIORITY;
1934 noisespec.gauge_priority = LINK_WIGGLE_GAUGE_PRIORITY;
1935 let n_rows = meanspec.design.nrows();
1936 let wigglespec = build_location_scale_wiggle_block(
1937 "wiggle",
1938 self.wiggle_block.design.clone(),
1939 self.wiggle_block.offset.clone(),
1940 wiggle_block_penalty_matrices(&self.wiggle_block),
1941 self.wiggle_block.nullspace_dims.clone(),
1942 layout.wiggle_from(theta),
1943 self.wiggle_block.initial_beta.clone(),
1944 n_rows,
1945 )?;
1946 Ok(vec![meanspec, noisespec, wigglespec])
1947 }
1948
1949 fn build_family(
1950 &self,
1951 mean_design: &TermCollectionDesign,
1952 noise_design: &TermCollectionDesign,
1953 ) -> Self::Family {
1954 let preparednoise_design =
1955 prepared_gaussian_log_sigma_design(&mean_design.design, &noise_design.design).expect(
1956 "prepared Gaussian log-sigma design should match wiggle block construction",
1957 );
1958 GaussianLocationScaleWiggleFamily {
1959 y: self.y.clone(),
1960 weights: self.weights.clone(),
1961 mu_design: Some(mean_design.design.clone()),
1962 log_sigma_design: Some(preparednoise_design),
1963 wiggle_knots: self.wiggle_knots.clone(),
1964 wiggle_degree: self.wiggle_degree,
1965 policy: gam_runtime::resource::ResourcePolicy::default_library(),
1966 cached_row_scalars: std::sync::RwLock::new(None),
1967 }
1968 }
1969
1970 fn extract_primary_betas(
1971 &self,
1972 fit: &UnifiedFitResult,
1973 ) -> Result<(Array1<f64>, Array1<f64>), String> {
1974 let mean_beta = fit
1975 .block_states
1976 .get(GaussianLocationScaleWiggleFamily::BLOCK_MU)
1977 .ok_or_else(|| "missing Gaussian wiggle mu block state".to_string())?
1978 .beta
1979 .clone();
1980 let noise_beta = fit
1981 .block_states
1982 .get(GaussianLocationScaleWiggleFamily::BLOCK_LOG_SIGMA)
1983 .ok_or_else(|| "missing Gaussian wiggle log_sigma block state".to_string())?
1984 .beta
1985 .clone();
1986 Ok((mean_beta, noise_beta))
1987 }
1988
1989 fn build_psiderivative_blocks(
1990 &self,
1991 data: ndarray::ArrayView2<'_, f64>,
1992 meanspec_resolved: &TermCollectionSpec,
1993 noisespec_resolved: &TermCollectionSpec,
1994 mean_design: &TermCollectionDesign,
1995 noise_design: &TermCollectionDesign,
1996 ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String> {
1997 let mean_derivs =
1998 build_block_spatial_psi_derivatives(data, meanspec_resolved, mean_design)?.ok_or_else(
1999 || "missing Gaussian wiggle mean spatial psi derivatives".to_string(),
2000 )?;
2001 let noise_derivs =
2002 build_block_spatial_psi_derivatives(data, noisespec_resolved, noise_design)?
2003 .ok_or_else(|| {
2004 "missing Gaussian wiggle log-sigma spatial psi derivatives".to_string()
2005 })?;
2006 Ok(vec![mean_derivs, noise_derivs, Vec::new()])
2007 }
2008}
2009
2010pub(crate) struct BinomialLocationScaleTermBuilder {
2011 pub(crate) y: Array1<f64>,
2012 pub(crate) weights: Array1<f64>,
2013 pub(crate) link_kind: InverseLink,
2014 pub(crate) meanspec: TermCollectionSpec,
2015 pub(crate) noisespec: TermCollectionSpec,
2016 pub(crate) mean_offset: Array1<f64>,
2017 pub(crate) noise_offset: Array1<f64>,
2018}
2019
2020impl LocationScaleFamilyBuilder for BinomialLocationScaleTermBuilder {
2021 type Family = BinomialLocationScaleFamily;
2022
2023 fn meanspec(&self) -> &TermCollectionSpec {
2024 &self.meanspec
2025 }
2026
2027 fn noisespec(&self) -> &TermCollectionSpec {
2028 &self.noisespec
2029 }
2030
2031 fn exact_spatial_joint_supported(&self) -> bool {
2032 true
2033 }
2034
2035 fn require_exact_spatial_joint(&self) -> bool {
2036 true
2037 }
2038
2039 fn noise_penalty_count(&self, noise_design: &TermCollectionDesign) -> usize {
2040 noise_design.penalties.len() + 1
2041 }
2042
2043 fn build_blocks(
2044 &self,
2045 theta: &Array1<f64>,
2046 mean_design: &TermCollectionDesign,
2047 noise_design: &TermCollectionDesign,
2048 mean_beta_hint: Option<Array1<f64>>,
2049 noise_beta_hint: Option<Array1<f64>>,
2050 ) -> Result<Vec<ParameterBlockSpec>, String> {
2051 let layout = GamlssLambdaLayout::two_block(
2052 mean_design.penalties.len(),
2053 self.noise_penalty_count(noise_design),
2054 );
2055 layout.validate_theta_len(theta.len(), "binomial location-scale")?;
2056 let (thresholdspec, log_sigmaspec) = build_binomial_threshold_and_scale_blocks(
2057 &self.y,
2058 &self.weights,
2059 &self.link_kind,
2060 mean_design,
2061 noise_design,
2062 &self.mean_offset,
2063 &self.noise_offset,
2064 layout.mean_from(theta),
2065 layout.noise_from(theta),
2066 mean_beta_hint,
2067 noise_beta_hint,
2068 "BinomialLocationScale::build_blocks",
2069 )?;
2070 Ok(vec![thresholdspec, log_sigmaspec])
2071 }
2072
2073 fn build_family(
2074 &self,
2075 mean_design: &TermCollectionDesign,
2076 noise_design: &TermCollectionDesign,
2077 ) -> Self::Family {
2078 let identifiednoise_design =
2079 identified_binomial_log_sigma_design(mean_design, noise_design, &self.weights)
2080 .expect("identified binomial log-sigma design");
2081 BinomialLocationScaleFamily {
2082 y: self.y.clone(),
2083 weights: self.weights.clone(),
2084 link_kind: self.link_kind.clone(),
2085 threshold_design: Some(mean_design.design.clone()),
2086 log_sigma_design: Some(identifiednoise_design),
2087 policy: gam_runtime::resource::ResourcePolicy::default_library(),
2088 }
2089 }
2090
2091 fn extract_primary_betas(
2092 &self,
2093 fit: &UnifiedFitResult,
2094 ) -> Result<(Array1<f64>, Array1<f64>), String> {
2095 let mean_beta = fit
2096 .block_states
2097 .get(BinomialLocationScaleFamily::BLOCK_T)
2098 .ok_or_else(|| "missing Binomial threshold block state".to_string())?
2099 .beta
2100 .clone();
2101 let noise_beta = fit
2102 .block_states
2103 .get(BinomialLocationScaleFamily::BLOCK_LOG_SIGMA)
2104 .ok_or_else(|| "missing Binomial log_sigma block state".to_string())?
2105 .beta
2106 .clone();
2107 Ok((mean_beta, noise_beta))
2108 }
2109
2110 fn build_psiderivative_blocks(
2111 &self,
2112 data: ndarray::ArrayView2<'_, f64>,
2113 meanspec_resolved: &TermCollectionSpec,
2114 noisespec_resolved: &TermCollectionSpec,
2115 mean_design: &TermCollectionDesign,
2116 noise_design: &TermCollectionDesign,
2117 ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String> {
2118 let mean_derivs =
2119 build_block_spatial_psi_derivatives(data, meanspec_resolved, mean_design)?
2120 .ok_or_else(|| "missing threshold spatial psi derivatives".to_string())?;
2121 let noise_derivs =
2122 build_block_spatial_psi_derivatives(data, noisespec_resolved, noise_design)?
2123 .ok_or_else(|| "missing log_sigma spatial psi derivatives".to_string())?;
2124 Ok(vec![mean_derivs, noise_derivs])
2125 }
2126}
2127
2128pub(crate) struct BinomialLocationScaleWiggleTermBuilder {
2129 pub(crate) y: Array1<f64>,
2130 pub(crate) weights: Array1<f64>,
2131 pub(crate) link_kind: InverseLink,
2132 pub(crate) meanspec: TermCollectionSpec,
2133 pub(crate) noisespec: TermCollectionSpec,
2134 pub(crate) mean_offset: Array1<f64>,
2135 pub(crate) noise_offset: Array1<f64>,
2136 pub(crate) wiggle_knots: Array1<f64>,
2137 pub(crate) wiggle_degree: usize,
2138 pub(crate) wiggle_block: ParameterBlockInput,
2139}
2140
2141impl LocationScaleFamilyBuilder for BinomialLocationScaleWiggleTermBuilder {
2142 type Family = BinomialLocationScaleWiggleFamily;
2143
2144 fn meanspec(&self) -> &TermCollectionSpec {
2145 &self.meanspec
2146 }
2147
2148 fn noisespec(&self) -> &TermCollectionSpec {
2149 &self.noisespec
2150 }
2151
2152 fn exact_spatial_joint_supported(&self) -> bool {
2153 true
2154 }
2155
2156 fn require_exact_spatial_joint(&self) -> bool {
2157 true
2158 }
2159
2160 fn extra_rho0(&self) -> Result<Array1<f64>, String> {
2161 initial_log_lambdas_orzeros(&self.wiggle_block)
2162 }
2163
2164 fn noise_penalty_count(&self, noise_design: &TermCollectionDesign) -> usize {
2165 noise_design.penalties.len() + 1
2166 }
2167
2168 fn build_blocks(
2169 &self,
2170 theta: &Array1<f64>,
2171 mean_design: &TermCollectionDesign,
2172 noise_design: &TermCollectionDesign,
2173 mean_beta_hint: Option<Array1<f64>>,
2174 noise_beta_hint: Option<Array1<f64>>,
2175 ) -> Result<Vec<ParameterBlockSpec>, String> {
2176 let layout = GamlssLambdaLayout::withwiggle(
2177 mean_design.penalties.len(),
2178 self.noise_penalty_count(noise_design),
2179 self.wiggle_block.penalties.len(),
2180 );
2181 layout.validate_theta_len(theta.len(), "wiggle location-scale")?;
2182 let (mut thresholdspec, mut log_sigmaspec) = build_binomial_threshold_and_scale_blocks(
2183 &self.y,
2184 &self.weights,
2185 &self.link_kind,
2186 mean_design,
2187 noise_design,
2188 &self.mean_offset,
2189 &self.noise_offset,
2190 layout.mean_from(theta),
2191 layout.noise_from(theta),
2192 mean_beta_hint,
2193 noise_beta_hint,
2194 "BinomialLocationScaleWiggle::build_blocks",
2195 )?;
2196 thresholdspec.gauge_priority = LINK_WIGGLE_GAUGE_PRIORITY;
2208 log_sigmaspec.gauge_priority = LINK_WIGGLE_GAUGE_PRIORITY;
2209 let n_rows = thresholdspec.design.nrows();
2210 let wigglespec = build_location_scale_wiggle_block(
2211 "wiggle",
2212 self.wiggle_block.design.clone(),
2213 self.wiggle_block.offset.clone(),
2214 wiggle_block_penalty_matrices(&self.wiggle_block),
2215 vec![],
2216 layout.wiggle_from(theta),
2217 self.wiggle_block.initial_beta.clone(),
2218 n_rows,
2219 )?;
2220 Ok(vec![thresholdspec, log_sigmaspec, wigglespec])
2221 }
2222
2223 fn build_family(
2224 &self,
2225 mean_design: &TermCollectionDesign,
2226 noise_design: &TermCollectionDesign,
2227 ) -> Self::Family {
2228 let identifiednoise_design =
2229 identified_binomial_log_sigma_design(mean_design, noise_design, &self.weights)
2230 .expect("identified binomial log-sigma design should match block construction");
2231 BinomialLocationScaleWiggleFamily {
2232 y: self.y.clone(),
2233 weights: self.weights.clone(),
2234 link_kind: self.link_kind.clone(),
2235 threshold_design: Some(mean_design.design.clone()),
2236 log_sigma_design: Some(identifiednoise_design),
2237 wiggle_knots: self.wiggle_knots.clone(),
2238 wiggle_degree: self.wiggle_degree,
2239 policy: gam_runtime::resource::ResourcePolicy::default_library(),
2240 }
2241 }
2242
2243 fn extract_primary_betas(
2244 &self,
2245 fit: &UnifiedFitResult,
2246 ) -> Result<(Array1<f64>, Array1<f64>), String> {
2247 let mean_beta = fit
2248 .block_states
2249 .get(BinomialLocationScaleWiggleFamily::BLOCK_T)
2250 .ok_or_else(|| "missing Binomial wiggle threshold block state".to_string())?
2251 .beta
2252 .clone();
2253 let noise_beta = fit
2254 .block_states
2255 .get(BinomialLocationScaleWiggleFamily::BLOCK_LOG_SIGMA)
2256 .ok_or_else(|| "missing Binomial wiggle log_sigma block state".to_string())?
2257 .beta
2258 .clone();
2259 Ok((mean_beta, noise_beta))
2260 }
2261
2262 fn build_psiderivative_blocks(
2263 &self,
2264 data: ndarray::ArrayView2<'_, f64>,
2265 meanspec_resolved: &TermCollectionSpec,
2266 noisespec_resolved: &TermCollectionSpec,
2267 mean_design: &TermCollectionDesign,
2268 noise_design: &TermCollectionDesign,
2269 ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String> {
2270 let mean_derivs =
2271 build_block_spatial_psi_derivatives(data, meanspec_resolved, mean_design)?
2272 .ok_or_else(|| "missing threshold spatial psi derivatives".to_string())?;
2273 let noise_derivs =
2274 build_block_spatial_psi_derivatives(data, noisespec_resolved, noise_design)?
2275 .ok_or_else(|| "missing log_sigma spatial psi derivatives".to_string())?;
2276 Ok(vec![mean_derivs, noise_derivs, Vec::new()])
2285 }
2286}
2287
2288pub(crate) fn fit_gaussian_location_scale_terms(
2289 data: ndarray::ArrayView2<'_, f64>,
2290 spec: GaussianLocationScaleTermSpec,
2291 options: &BlockwiseFitOptions,
2292 kappa_options: &SpatialLengthScaleOptimizationOptions,
2293) -> Result<BlockwiseTermFitResult, String> {
2294 validate_gaussian_location_scale_termspec(data, &spec, "fit_gaussian_location_scale_terms")?;
2295 fit_location_scale_terms(
2296 data,
2297 GaussianLocationScaleTermBuilder {
2298 y: spec.y,
2299 weights: spec.weights,
2300 meanspec: spec.meanspec,
2301 noisespec: spec.log_sigmaspec,
2302 mean_offset: spec.mean_offset,
2303 noise_offset: spec.log_sigma_offset,
2304 },
2305 options,
2306 kappa_options,
2307 )
2308}
2309
2310pub(crate) fn fit_gaussian_location_scalewiggle_terms(
2311 data: ndarray::ArrayView2<'_, f64>,
2312 spec: GaussianLocationScaleWiggleTermSpec,
2313 options: &BlockwiseFitOptions,
2314 kappa_options: &SpatialLengthScaleOptimizationOptions,
2315) -> Result<BlockwiseTermFitResult, String> {
2316 validate_gaussian_location_scalewiggle_termspec(
2317 data,
2318 &spec,
2319 "fit_gaussian_location_scalewiggle_terms",
2320 )?;
2321 fit_location_scale_terms(
2322 data,
2323 GaussianLocationScaleWiggleTermBuilder {
2324 y: spec.y,
2325 weights: spec.weights,
2326 meanspec: spec.meanspec,
2327 noisespec: spec.log_sigmaspec,
2328 mean_offset: spec.mean_offset,
2329 noise_offset: spec.log_sigma_offset,
2330 wiggle_knots: spec.wiggle_knots,
2331 wiggle_degree: spec.wiggle_degree,
2332 wiggle_block: spec.wiggle_block,
2333 },
2334 options,
2335 kappa_options,
2336 )
2337}
2338
2339pub(crate) fn select_gaussian_location_scale_link_wiggle_basis_from_pilot(
2340 pilot: &BlockwiseTermFitResult,
2341 wiggle_cfg: &WiggleBlockConfig,
2342 wiggle_penalty_orders: &[usize],
2343) -> Result<SelectedWiggleBasis, String> {
2344 let q_seed = pilot
2345 .fit
2346 .block_states
2347 .first()
2348 .ok_or_else(|| "pilot Gaussian wiggle fit is missing mean block".to_string())?
2349 .eta
2350 .view();
2351 select_wiggle_basis_from_seed(q_seed, wiggle_cfg, wiggle_penalty_orders)
2352}
2353
2354pub(crate) fn fit_gaussian_location_scale_terms_with_selected_wiggle(
2355 data: ndarray::ArrayView2<'_, f64>,
2356 spec: GaussianLocationScaleTermSpec,
2357 selected_wiggle_basis: SelectedWiggleBasis,
2358 options: &BlockwiseFitOptions,
2359 kappa_options: &SpatialLengthScaleOptimizationOptions,
2360) -> Result<BlockwiseTermWiggleFitResult, String> {
2361 let SelectedWiggleBasis {
2362 knots: wiggle_knots,
2363 degree: wiggle_degree,
2364 block: wiggle_block,
2365 ..
2366 } = selected_wiggle_basis;
2367 let solved = fit_gaussian_location_scalewiggle_terms(
2368 data,
2369 GaussianLocationScaleWiggleTermSpec {
2370 y: spec.y,
2371 weights: spec.weights,
2372 meanspec: spec.meanspec,
2373 log_sigmaspec: spec.log_sigmaspec,
2374 mean_offset: spec.mean_offset,
2375 log_sigma_offset: spec.log_sigma_offset,
2376 wiggle_knots: wiggle_knots.clone(),
2377 wiggle_degree,
2378 wiggle_block,
2379 },
2380 options,
2381 kappa_options,
2382 )?;
2383
2384 BlockwiseTermWiggleFitResult::try_from_parts(BlockwiseTermWiggleFitResultParts {
2385 fit: solved,
2386 wiggle_knots,
2387 wiggle_degree,
2388 })
2389}
2390
2391pub(crate) fn fit_binomial_location_scale_terms(
2392 data: ndarray::ArrayView2<'_, f64>,
2393 spec: BinomialLocationScaleTermSpec,
2394 options: &BlockwiseFitOptions,
2395 kappa_options: &SpatialLengthScaleOptimizationOptions,
2396) -> Result<BlockwiseTermFitResult, String> {
2397 validate_binomial_location_scale_termspec(data, &spec, "fit_binomial_location_scale_terms")?;
2398 fit_location_scale_terms(
2399 data,
2400 BinomialLocationScaleTermBuilder {
2401 y: spec.y,
2402 weights: spec.weights,
2403 link_kind: spec.link_kind,
2404 meanspec: spec.thresholdspec,
2405 noisespec: spec.log_sigmaspec,
2406 mean_offset: spec.threshold_offset,
2407 noise_offset: spec.log_sigma_offset,
2408 },
2409 options,
2410 kappa_options,
2411 )
2412}
2413
2414pub(crate) fn fit_binomial_location_scalewiggle_terms(
2415 data: ndarray::ArrayView2<'_, f64>,
2416 spec: BinomialLocationScaleWiggleTermSpec,
2417 options: &BlockwiseFitOptions,
2418 kappa_options: &SpatialLengthScaleOptimizationOptions,
2419) -> Result<BlockwiseTermFitResult, String> {
2420 validate_binomial_location_scalewiggle_termspec(
2421 data,
2422 &spec,
2423 "fit_binomial_location_scalewiggle_terms",
2424 )?;
2425 fit_location_scale_terms(
2426 data,
2427 BinomialLocationScaleWiggleTermBuilder {
2428 y: spec.y,
2429 weights: spec.weights,
2430 link_kind: spec.link_kind,
2431 meanspec: spec.thresholdspec,
2432 noisespec: spec.log_sigmaspec,
2433 mean_offset: spec.threshold_offset,
2434 noise_offset: spec.log_sigma_offset,
2435 wiggle_knots: spec.wiggle_knots,
2436 wiggle_degree: spec.wiggle_degree,
2437 wiggle_block: spec.wiggle_block,
2438 },
2439 options,
2440 kappa_options,
2441 )
2442}
2443
2444pub(crate) fn select_binomial_location_scale_link_wiggle_basis_from_pilot(
2445 pilot: &BlockwiseTermFitResult,
2446 wiggle_cfg: &WiggleBlockConfig,
2447 wiggle_penalty_orders: &[usize],
2448) -> Result<SelectedWiggleBasis, String> {
2449 let eta_t = pilot
2450 .fit
2451 .block_states
2452 .first()
2453 .ok_or_else(|| "pilot fit is missing threshold block".to_string())?
2454 .eta
2455 .view();
2456 let eta_ls = pilot
2457 .fit
2458 .block_states
2459 .get(1)
2460 .ok_or_else(|| "pilot fit is missing log_sigma block".to_string())?
2461 .eta
2462 .view();
2463 let sigma = eta_ls.mapv(safe_exp);
2464 let q_seed = Array1::from_iter(eta_t.iter().zip(sigma.iter()).map(|(&t, &s)| -t / s));
2465 select_wiggle_basis_from_seed(q_seed.view(), wiggle_cfg, wiggle_penalty_orders)
2466}
2467
2468pub(crate) fn fit_binomial_location_scale_terms_with_selected_wiggle(
2469 data: ndarray::ArrayView2<'_, f64>,
2470 spec: BinomialLocationScaleTermSpec,
2471 selected_wiggle_basis: SelectedWiggleBasis,
2472 options: &BlockwiseFitOptions,
2473 kappa_options: &SpatialLengthScaleOptimizationOptions,
2474) -> Result<BlockwiseTermWiggleFitResult, String> {
2475 let SelectedWiggleBasis {
2476 knots: wiggle_knots,
2477 degree: wiggle_degree,
2478 block: wiggle_block,
2479 ..
2480 } = selected_wiggle_basis;
2481 let solved = fit_binomial_location_scalewiggle_terms(
2482 data,
2483 BinomialLocationScaleWiggleTermSpec {
2484 y: spec.y,
2485 weights: spec.weights,
2486 link_kind: spec.link_kind,
2487 thresholdspec: spec.thresholdspec,
2488 log_sigmaspec: spec.log_sigmaspec,
2489 threshold_offset: spec.threshold_offset,
2490 log_sigma_offset: spec.log_sigma_offset,
2491 wiggle_knots: wiggle_knots.clone(),
2492 wiggle_degree,
2493 wiggle_block,
2494 },
2495 options,
2496 kappa_options,
2497 )?;
2498
2499 BlockwiseTermWiggleFitResult::try_from_parts(BlockwiseTermWiggleFitResultParts {
2500 fit: solved,
2501 wiggle_knots,
2502 wiggle_degree,
2503 })
2504}
2505
2506pub(crate) fn select_binomial_mean_link_wiggle_basis_from_pilot(
2507 pilot_design: &TermCollectionDesign,
2508 pilot_fit: &UnifiedFitResult,
2509 wiggle_cfg: &WiggleBlockConfig,
2510 wiggle_penalty_orders: &[usize],
2511) -> Result<SelectedWiggleBasis, String> {
2512 let q_seed = pilot_design.design.dot(&pilot_fit.beta);
2513 select_wiggle_basis_from_seed(q_seed.view(), wiggle_cfg, wiggle_penalty_orders)
2514}
2515
2516pub(crate) fn fit_binomial_mean_wiggle_terms_with_selected_basis(
2517 data: ndarray::ArrayView2<'_, f64>,
2518 pilot_spec: &TermCollectionSpec,
2519 pilot_design: &TermCollectionDesign,
2520 pilot_fit: &UnifiedFitResult,
2521 y: &Array1<f64>,
2522 weights: &Array1<f64>,
2523 link_kind: InverseLink,
2524 selected_wiggle_basis: SelectedWiggleBasis,
2525 options: &BlockwiseFitOptions,
2526 kappa_options: &SpatialLengthScaleOptimizationOptions,
2527) -> Result<BinomialMeanWiggleTermFitResult, String> {
2528 const RHO_BOUND: f64 = 12.0;
2529
2530 validate_term_weights(
2531 data,
2532 y.len(),
2533 weights,
2534 "fit_binomial_mean_wiggle_terms_with_selected_basis",
2535 )?;
2536 validate_binomial_response(y, "fit_binomial_mean_wiggle_terms_with_selected_basis")?;
2537
2538 let SelectedWiggleBasis {
2544 knots: wiggle_knots,
2545 degree: wiggle_degree,
2546 block: wiggle_block,
2547 ..
2548 } = selected_wiggle_basis;
2549
2550 let spatial_terms = spatial_length_scale_term_indices(pilot_spec);
2551 if spatial_terms.is_empty() {
2552 let fit = fit_binomial_mean_wiggle(
2553 BinomialMeanWiggleSpec {
2554 y: y.clone(),
2555 weights: weights.clone(),
2556 link_kind,
2557 wiggle_knots: wiggle_knots.clone(),
2558 wiggle_degree,
2559 eta_block: ParameterBlockInput {
2560 design: pilot_design.design.clone(),
2561 offset: Array1::zeros(y.len()),
2562 penalties: pilot_design
2563 .penalties
2564 .iter()
2565 .map(crate::model_types::PenaltySpec::from_blockwise_ref)
2566 .collect(),
2567 nullspace_dims: vec![],
2568 initial_log_lambdas: Some(
2569 pilot_fit
2570 .lambdas
2571 .mapv(|v| v.max(WARMSTART_LOG_LAMBDA_FLOOR).ln()),
2572 ),
2573 initial_beta: Some(pilot_fit.beta.clone()),
2574 },
2575 wiggle_block,
2576 },
2577 options,
2578 )?;
2579 return Ok(BinomialMeanWiggleTermFitResult {
2580 fit,
2581 resolvedspec: pilot_spec.clone(),
2582 design: pilot_design.clone(),
2583 wiggle_knots,
2584 wiggle_degree,
2585 });
2586 }
2587
2588 let dims_per_term = spatial_dims_per_term(pilot_spec, &spatial_terms);
2589 let log_kappa0 =
2590 SpatialLogKappaCoords::from_length_scales_aniso(pilot_spec, &spatial_terms, kappa_options)
2591 .reseed_from_data(data, pilot_spec, &spatial_terms, kappa_options);
2592 let log_kappa_lower = SpatialLogKappaCoords::lower_bounds_aniso_from_data(
2593 data,
2594 pilot_spec,
2595 &spatial_terms,
2596 &dims_per_term,
2597 kappa_options,
2598 );
2599 let log_kappa_upper = SpatialLogKappaCoords::upper_bounds_aniso_from_data(
2600 data,
2601 pilot_spec,
2602 &spatial_terms,
2603 &dims_per_term,
2604 kappa_options,
2605 );
2606 let log_kappa0 = log_kappa0.clamp_to_bounds(&log_kappa_lower, &log_kappa_upper);
2608
2609 let eta_penalty_count = pilot_design.penalties.len();
2610 let wiggle_penalty_count = initial_log_lambdas_orzeros(&wiggle_block)?.len();
2611 let rho_dim = eta_penalty_count + wiggle_penalty_count;
2612 let baseline_resolvedspec = log_kappa0
2613 .apply_tospec(pilot_spec, &spatial_terms)
2614 .map_err(|e| e.to_string())?;
2615 let baseline_design =
2616 build_term_collection_design(data, &baseline_resolvedspec).map_err(|e| e.to_string())?;
2617 let baseline_fit = fit_binomial_mean_wiggle(
2618 BinomialMeanWiggleSpec {
2619 y: y.clone(),
2620 weights: weights.clone(),
2621 link_kind: link_kind.clone(),
2622 wiggle_knots: wiggle_knots.clone(),
2623 wiggle_degree,
2624 eta_block: ParameterBlockInput {
2625 design: baseline_design.design.clone(),
2626 offset: Array1::zeros(y.len()),
2627 penalties: baseline_design
2628 .penalties
2629 .iter()
2630 .map(crate::model_types::PenaltySpec::from_blockwise_ref)
2631 .collect(),
2632 nullspace_dims: vec![],
2633 initial_log_lambdas: Some(
2634 pilot_fit
2635 .lambdas
2636 .mapv(|v| v.max(WARMSTART_LOG_LAMBDA_FLOOR).ln()),
2637 ),
2638 initial_beta: Some(pilot_fit.beta.clone()),
2639 },
2640 wiggle_block: wiggle_block.clone(),
2641 },
2642 options,
2643 )?;
2644 let baseline_log_lambdas = baseline_fit
2645 .lambdas
2646 .mapv(|v| v.max(WARMSTART_LOG_LAMBDA_FLOOR).ln());
2647 if baseline_log_lambdas.len() != rho_dim {
2648 return Err(GamlssError::DimensionMismatch {
2649 reason: format!(
2650 "baseline binomial mean-wiggle fit returned {} log-lambdas, expected {rho_dim}",
2651 baseline_log_lambdas.len()
2652 ),
2653 }
2654 .into());
2655 }
2656 let baseline_eta_beta = baseline_fit
2657 .block_states
2658 .get(BinomialMeanWiggleFamily::BLOCK_ETA)
2659 .ok_or_else(|| "baseline binomial mean-wiggle fit missing eta block".to_string())?
2660 .beta
2661 .clone();
2662 let baseline_wiggle_beta = Some(
2663 baseline_fit
2664 .block_states
2665 .get(BinomialMeanWiggleFamily::BLOCK_WIGGLE)
2666 .ok_or_else(|| "baseline binomial mean-wiggle fit missing wiggle block".to_string())?
2667 .beta
2668 .clone(),
2669 );
2670 let theta_dim = rho_dim + log_kappa0.len();
2671 let mut theta0 = Array1::<f64>::zeros(theta_dim);
2672 theta0
2673 .slice_mut(s![0..rho_dim])
2674 .assign(&baseline_log_lambdas);
2675 theta0
2676 .slice_mut(s![rho_dim..theta_dim])
2677 .assign(log_kappa0.as_array());
2678
2679 let mut lower = Array1::<f64>::from_elem(theta_dim, -RHO_BOUND);
2680 let mut upper = Array1::<f64>::from_elem(theta_dim, RHO_BOUND);
2681 lower
2682 .slice_mut(s![rho_dim..theta_dim])
2683 .assign(log_kappa_lower.as_array());
2684 upper
2685 .slice_mut(s![rho_dim..theta_dim])
2686 .assign(log_kappa_upper.as_array());
2687
2688 let pilot_spec_cloned = pilot_spec.clone();
2689 let pilot_beta = baseline_eta_beta;
2690 let wiggle_design = wiggle_block.design.clone();
2691 let wiggle_offset = wiggle_block.offset.clone();
2692 let wiggle_penalties = wiggle_block.penalties.clone();
2693 let wiggle_initial_beta = baseline_wiggle_beta;
2694 let wiggle_knots_cloned = wiggle_knots.clone();
2695 let y_cloned = y.clone();
2696 let weights_cloned = weights.clone();
2697 let link_kind_cloned = link_kind.clone();
2698 let outer_family = BinomialMeanWiggleFamily {
2699 y: y_cloned.clone(),
2700 weights: weights_cloned.clone(),
2701 link_kind: link_kind_cloned.clone(),
2702 wiggle_knots: wiggle_knots_cloned.clone(),
2703 wiggle_degree,
2704 policy: gam_runtime::resource::ResourcePolicy::default_library(),
2705 };
2706 let screening_cap = Arc::new(AtomicUsize::new(0));
2707 let mut outer_options = options.clone();
2708 outer_options.screening_max_inner_iterations = Some(Arc::clone(&screening_cap));
2709 struct MeanWiggleOuterState {
2710 pub(crate) warm_cache: Option<crate::custom_family::CustomFamilyWarmStart>,
2711 pub(crate) last_eval: Option<(
2712 Array1<f64>,
2713 f64,
2714 Array1<f64>,
2715 gam_problem::HessianResult,
2716 crate::custom_family::CustomFamilyWarmStart,
2717 )>,
2718 }
2719
2720 let build_realized_blocks = |theta: &Array1<f64>| -> Result<
2721 (
2722 TermCollectionSpec,
2723 TermCollectionDesign,
2724 Vec<ParameterBlockSpec>,
2725 Vec<CustomFamilyBlockPsiDerivative>,
2726 ),
2727 String,
2728 > {
2729 let log_kappa =
2730 SpatialLogKappaCoords::from_theta_tail_with_dims(theta, rho_dim, dims_per_term.clone());
2731 let resolvedspec = log_kappa
2732 .apply_tospec(&pilot_spec_cloned, &spatial_terms)
2733 .map_err(|e| e.to_string())?;
2734 let design =
2735 build_term_collection_design(data, &resolvedspec).map_err(|e| e.to_string())?;
2736 let eta_derivs = build_block_spatial_psi_derivatives(data, &resolvedspec, &design)?
2737 .ok_or_else(|| {
2738 "missing eta spatial psi derivatives for binomial mean wiggle".to_string()
2739 })?;
2740 let blocks = vec![
2741 ParameterBlockSpec {
2742 name: "eta".to_string(),
2743 design: design.design.clone(),
2744 offset: Array1::zeros(y_cloned.len()),
2745 penalties: design.penalties_as_penalty_matrix(),
2746 nullspace_dims: vec![],
2747 initial_log_lambdas: theta.slice(s![0..eta_penalty_count]).to_owned(),
2748 initial_beta: Some(pilot_beta.clone()),
2749 gauge_priority: LINK_WIGGLE_GAUGE_PRIORITY,
2753 jacobian_callback: None,
2754 stacked_design: None,
2755 stacked_offset: None,
2756 },
2757 ParameterBlockSpec {
2758 name: "wiggle".to_string(),
2759 design: wiggle_design.clone(),
2760 offset: wiggle_offset.clone(),
2761 penalties: {
2762 let p_wiggle = wiggle_design.ncols();
2763 wiggle_penalties
2764 .iter()
2765 .map(|spec| match spec {
2766 crate::model_types::PenaltySpec::Block {
2767 local, col_range, ..
2768 } => PenaltyMatrix::Blockwise {
2769 local: local.clone(),
2770 col_range: col_range.clone(),
2771 total_dim: p_wiggle,
2772 },
2773 crate::model_types::PenaltySpec::Dense(m)
2774 | crate::model_types::PenaltySpec::DenseWithMean {
2775 matrix: m, ..
2776 } => PenaltyMatrix::Dense(m.clone()),
2777 })
2778 .collect()
2779 },
2780 nullspace_dims: vec![],
2781 initial_log_lambdas: theta.slice(s![eta_penalty_count..rho_dim]).to_owned(),
2782 initial_beta: wiggle_initial_beta.clone(),
2783 gauge_priority: DEFAULT_GAUGE_PRIORITY,
2784 jacobian_callback: None,
2785 stacked_design: None,
2786 stacked_offset: None,
2787 },
2788 ];
2789 Ok((resolvedspec, design, blocks, eta_derivs))
2790 };
2791
2792 let build_eval = |theta: &Array1<f64>,
2793 warm_cache: Option<&crate::custom_family::CustomFamilyWarmStart>,
2794 need_hessian: bool|
2795 -> Result<
2796 (
2797 crate::custom_family::CustomFamilyJointHyperResult,
2798 TermCollectionSpec,
2799 TermCollectionDesign,
2800 ),
2801 String,
2802 > {
2803 let (resolvedspec, design, blocks, eta_derivs) = build_realized_blocks(theta)?;
2804 let eval = evaluate_custom_family_joint_hyper(
2805 &outer_family,
2806 &blocks,
2807 &outer_options,
2808 &theta.slice(s![0..rho_dim]).to_owned(),
2809 &[eta_derivs, Vec::new()],
2810 warm_cache,
2811 if need_hessian {
2812 gam_problem::EvalMode::ValueGradientHessian
2813 } else {
2814 gam_problem::EvalMode::ValueAndGradient
2815 },
2816 )?;
2817 Ok((eval, resolvedspec, design))
2818 };
2819
2820 let build_efs = |theta: &Array1<f64>,
2821 warm_cache: Option<&crate::custom_family::CustomFamilyWarmStart>|
2822 -> Result<crate::custom_family::CustomFamilyJointHyperEfsResult, String> {
2823 let (_, _, blocks, eta_derivs) = build_realized_blocks(theta)?;
2824 evaluate_custom_family_joint_hyper_efs(
2825 &outer_family,
2826 &blocks,
2827 &outer_options,
2828 &theta.slice(s![0..rho_dim]).to_owned(),
2829 &[eta_derivs, Vec::new()],
2830 warm_cache,
2831 )
2832 .map_err(|e| e.to_string())
2833 };
2834
2835 use crate::model_types::EstimationError;
2836 use gam_solve::rho_optimizer::OuterEvalOrder;
2837 use gam_problem::{DeclaredHessianForm, Derivative, OuterEval};
2838
2839 let analytic_outer_hessian_available = true;
2850 let mut seed_heuristic = theta0.to_vec();
2851 for value in &mut seed_heuristic[..rho_dim] {
2852 *value = value.exp();
2853 }
2854 let problem = gam_solve::rho_optimizer::OuterProblem::new(theta_dim)
2855 .with_gradient(Derivative::Analytic)
2856 .with_hessian(if analytic_outer_hessian_available {
2857 DeclaredHessianForm::Either
2858 } else {
2859 DeclaredHessianForm::Unavailable
2860 })
2861 .with_psi_dim(theta_dim - rho_dim)
2862 .with_tolerance(options.outer_tol)
2863 .with_max_iter(options.outer_max_iter)
2864 .with_bounds(lower.clone(), upper.clone())
2865 .with_initial_rho(theta0.clone())
2866 .with_seed_config(crate::seeding::SeedConfig {
2867 max_seeds: 4,
2868 seed_budget: 2,
2869 risk_profile: crate::seeding::SeedRiskProfile::GeneralizedLinear,
2870 num_auxiliary_trailing: theta_dim - rho_dim,
2871 ..Default::default()
2872 })
2873 .with_screening_cap(Arc::clone(&screening_cap))
2874 .with_rho_bound(12.0)
2875 .with_heuristic_lambdas(seed_heuristic);
2876
2877 let eval_outer = |state: &mut MeanWiggleOuterState,
2878 theta: &Array1<f64>,
2879 order: OuterEvalOrder|
2880 -> Result<OuterEval, EstimationError> {
2881 if let Some((cached_theta, cached_cost, cached_grad, cached_hess, cached_warm)) =
2882 &state.last_eval
2883 && cached_theta == theta
2884 && (!matches!(order, OuterEvalOrder::ValueGradientHessian)
2885 || matches!(
2886 cached_hess,
2887 gam_problem::HessianResult::Analytic(_)
2888 | gam_problem::HessianResult::Operator(_)
2889 ))
2890 {
2891 state.warm_cache = Some(cached_warm.clone());
2892 return Ok(OuterEval {
2893 cost: *cached_cost,
2894 gradient: cached_grad.clone(),
2895 hessian: cached_hess.clone(),
2896 inner_beta_hint: None,
2897 });
2898 }
2899 let need_hessian = matches!(order, OuterEvalOrder::ValueGradientHessian)
2900 && analytic_outer_hessian_available;
2901 let (eval, _, _) = build_eval(theta, state.warm_cache.as_ref(), need_hessian)
2902 .map_err(EstimationError::InvalidInput)?;
2903 if !eval.inner_converged {
2904 state.warm_cache = Some(eval.warm_start);
2905 crate::bail_invalid_estim!(
2906 "binomial mean-wiggle exact spatial inner solve did not converge"
2907 );
2908 }
2909 let hessian_result = eval.outer_hessian.clone();
2910 state.last_eval = Some((
2911 theta.clone(),
2912 eval.objective,
2913 eval.gradient.clone(),
2914 eval.outer_hessian.clone(),
2915 eval.warm_start.clone(),
2916 ));
2917 state.warm_cache = Some(eval.warm_start);
2918 Ok(OuterEval {
2919 cost: eval.objective,
2920 gradient: eval.gradient,
2921 hessian: hessian_result,
2922 inner_beta_hint: None,
2923 })
2924 };
2925
2926 let mut obj = problem.build_objective_with_screening_proxy(
2927 MeanWiggleOuterState {
2928 warm_cache: None,
2929 last_eval: None,
2930 },
2931 |state: &mut MeanWiggleOuterState, theta: &Array1<f64>| {
2932 if let Some((cached_theta, cached_cost, _, _, cached_warm)) = &state.last_eval
2933 && cached_theta == theta
2934 {
2935 state.warm_cache = Some(cached_warm.clone());
2936 return Ok(*cached_cost);
2937 }
2938 let (eval, _, _) = build_eval(theta, state.warm_cache.as_ref(), false)
2939 .map_err(EstimationError::InvalidInput)?;
2940 if !eval.inner_converged {
2941 state.warm_cache = Some(eval.warm_start);
2942 crate::bail_invalid_estim!(
2943 "binomial mean-wiggle exact spatial cost inner solve did not converge"
2944 .to_string(),
2945 );
2946 }
2947 state.warm_cache = Some(eval.warm_start);
2948 Ok(eval.objective)
2949 },
2950 |state: &mut MeanWiggleOuterState, theta: &Array1<f64>| {
2951 eval_outer(
2952 state,
2953 theta,
2954 if analytic_outer_hessian_available {
2955 OuterEvalOrder::ValueGradientHessian
2956 } else {
2957 OuterEvalOrder::ValueAndGradient
2958 },
2959 )
2960 },
2961 |state: &mut MeanWiggleOuterState, theta: &Array1<f64>, order: OuterEvalOrder| {
2962 eval_outer(state, theta, order)
2963 },
2964 Some(|state: &mut MeanWiggleOuterState| {
2965 state.warm_cache = None;
2966 state.last_eval = None;
2967 }),
2968 Some(|state: &mut MeanWiggleOuterState, theta: &Array1<f64>| {
2969 let eval = build_efs(theta, state.warm_cache.as_ref())
2970 .map_err(EstimationError::InvalidInput)?;
2971 if !eval.inner_converged {
2972 state.warm_cache = Some(eval.warm_start);
2973 crate::bail_invalid_estim!(
2974 "binomial mean-wiggle exact spatial EFS inner solve did not converge"
2975 .to_string(),
2976 );
2977 }
2978 state.warm_cache = Some(eval.warm_start);
2979 Ok(eval.efs_eval)
2980 }),
2981 |state: &mut MeanWiggleOuterState, theta: &Array1<f64>| {
2991 if let Some((cached_theta, cached_cost, _, _, cached_warm)) = &state.last_eval
2992 && cached_theta == theta
2993 {
2994 state.warm_cache = Some(cached_warm.clone());
2995 return Ok(*cached_cost);
2996 }
2997 let (eval, _, _) = build_eval(theta, state.warm_cache.as_ref(), false)
2998 .map_err(EstimationError::InvalidInput)?;
2999 state.warm_cache = Some(eval.warm_start);
3000 Ok(eval.objective)
3001 },
3002 );
3003
3004 let outer = problem
3005 .run(&mut obj, "binomial mean wiggle exact spatial hyper")
3006 .map_err(|e| e.to_string())?;
3007 if !outer.converged {
3008 return Err(GamlssError::NumericalFailure { reason: format!(
3009 "binomial mean wiggle exact spatial hyper did not converge after {} iterations (final_objective={:.6e}, final_grad_norm={})",
3010 outer.iterations,
3011 outer.final_value,
3012 outer.final_grad_norm_report(),
3013 ) }.into());
3014 }
3015 let theta_star = outer.rho;
3016
3017 let log_kappa =
3018 SpatialLogKappaCoords::from_theta_tail_with_dims(&theta_star, rho_dim, dims_per_term);
3019 let resolvedspec = log_kappa
3020 .apply_tospec(&pilot_spec_cloned, &spatial_terms)
3021 .map_err(|e| e.to_string())?;
3022 let design = build_term_collection_design(data, &resolvedspec).map_err(|e| e.to_string())?;
3023 let resolvedspec =
3024 freeze_term_collection_from_design(&resolvedspec, &design).map_err(|e| e.to_string())?;
3025 let fit = fit_binomial_mean_wiggle(
3026 BinomialMeanWiggleSpec {
3027 y: y_cloned,
3028 weights: weights_cloned,
3029 link_kind: link_kind_cloned,
3030 wiggle_knots: wiggle_knots.clone(),
3031 wiggle_degree,
3032 eta_block: ParameterBlockInput {
3033 design: design.design.clone(),
3034 offset: Array1::zeros(y.len()),
3035 penalties: design
3036 .penalties
3037 .iter()
3038 .map(crate::model_types::PenaltySpec::from_blockwise_ref)
3039 .collect(),
3040 nullspace_dims: vec![],
3041 initial_log_lambdas: Some(theta_star.slice(s![0..eta_penalty_count]).to_owned()),
3042 initial_beta: Some(pilot_beta),
3043 },
3044 wiggle_block: ParameterBlockInput {
3045 design: wiggle_design,
3046 offset: wiggle_offset,
3047 penalties: wiggle_penalties,
3048 nullspace_dims: vec![],
3049 initial_log_lambdas: Some(
3050 theta_star.slice(s![eta_penalty_count..rho_dim]).to_owned(),
3051 ),
3052 initial_beta: wiggle_initial_beta,
3053 },
3054 },
3055 options,
3056 )?;
3057
3058 Ok(BinomialMeanWiggleTermFitResult {
3059 fit,
3060 resolvedspec,
3061 design,
3062 wiggle_knots,
3063 wiggle_degree,
3064 })
3065}