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 append_binomial_log_sigma_shrinkage_penalty_design(
394 design: &mut TermCollectionDesign,
395) {
396 let p = design.design.ncols();
397 design
398 .penalties
399 .push(BlockwisePenalty::new(0..p, identity_penalty(p)));
400 design.nullspace_dims.push(0);
402 design.penaltyinfo.push(PenaltyBlockInfo {
403 global_index: design.penaltyinfo.len(),
404 termname: Some("log_sigma_shrinkage".to_string()),
405 penalty: PenaltyInfo {
406 source: PenaltySource::Other("shrinkage".to_string()),
407 original_index: 0,
408 active: true,
409 effective_rank: p,
410 dropped_reason: None,
411 nullspace_dim_hint: 0,
412 normalization_scale: 1.0,
413 kronecker_factors: None,
414 },
415 });
416}
417
418pub(crate) fn build_gaussian_mean_and_scale_blocks(
425 y: &Array1<f64>,
426 weights: &Array1<f64>,
427 mean_design: &TermCollectionDesign,
428 noise_design: &TermCollectionDesign,
429 mean_offset: &Array1<f64>,
430 noise_offset: &Array1<f64>,
431 mean_log_lambdas: Array1<f64>,
432 noise_log_lambdas: Array1<f64>,
433 mean_beta_hint: Option<Array1<f64>>,
434 noise_beta_hint: Option<Array1<f64>>,
435 context: &str,
436) -> Result<(ParameterBlockSpec, ParameterBlockSpec), String> {
437 let mut meanspec = build_location_scale_block(
438 "mu",
439 mean_design.design.clone(),
440 mean_offset.clone(),
441 mean_design.penalties_as_penalty_matrix(),
442 mean_design.nullspace_dims.clone(),
443 mean_log_lambdas,
444 mean_beta_hint,
445 0,
446 LOCATION_SCALE_N_OUTPUTS,
447 &format!("{context}: mu"),
448 )?;
449 let prepared_noise_design =
450 prepared_gaussian_log_sigma_design(&mean_design.design, &noise_design.design)?;
451 let mut noisespec = build_location_scale_block(
459 "log_sigma",
460 prepared_noise_design,
461 noise_offset.clone(),
462 noise_design.penalties_as_penalty_matrix(),
463 noise_design.nullspace_dims.clone(),
464 noise_log_lambdas,
465 noise_beta_hint,
466 1,
467 LOCATION_SCALE_N_OUTPUTS,
468 &format!("{context}: log_sigma"),
469 )?;
470 if meanspec.initial_beta.is_none() || noisespec.initial_beta.is_none() {
471 let (betamu0, beta_ls0, _) = gaussian_location_scalewarm_start(
472 y,
473 weights,
474 &meanspec,
475 &noisespec,
476 1e-10,
477 meanspec.initial_beta.as_ref(),
478 noisespec.initial_beta.as_ref(),
479 )?;
480 if meanspec.initial_beta.is_none() {
481 meanspec.initial_beta = Some(betamu0);
482 }
483 if noisespec.initial_beta.is_none() {
484 noisespec.initial_beta = Some(beta_ls0);
485 }
486 }
487 Ok((meanspec, noisespec))
488}
489
490pub(crate) fn build_binomial_threshold_and_scale_blocks(
496 y: &Array1<f64>,
497 weights: &Array1<f64>,
498 link_kind: &InverseLink,
499 mean_design: &TermCollectionDesign,
500 noise_design: &TermCollectionDesign,
501 mean_offset: &Array1<f64>,
502 noise_offset: &Array1<f64>,
503 mean_log_lambdas: Array1<f64>,
504 noise_log_lambdas: Array1<f64>,
505 mean_beta_hint: Option<Array1<f64>>,
506 noise_beta_hint: Option<Array1<f64>>,
507 context: &str,
508) -> Result<(ParameterBlockSpec, ParameterBlockSpec), String> {
509 let identifiednoise_design =
510 identified_binomial_log_sigma_design(mean_design, noise_design, weights)?;
511 let p_noise = identifiednoise_design.ncols();
512 let mut log_sigma_penalty_matrices: Vec<PenaltyMatrix> =
513 noise_design.penalties_as_penalty_matrix();
514 log_sigma_penalty_matrices.push(PenaltyMatrix::Dense(identity_penalty(p_noise)));
515 let mut thresholdspec = build_location_scale_block(
516 "threshold",
517 mean_design.design.clone(),
518 mean_offset.clone(),
519 mean_design.penalties_as_penalty_matrix(),
520 vec![],
521 mean_log_lambdas,
522 mean_beta_hint,
523 0,
524 LOCATION_SCALE_N_OUTPUTS,
525 &format!("{context}: threshold"),
526 )?;
527 let mut log_sigmaspec = build_location_scale_block(
528 "log_sigma",
529 identifiednoise_design,
530 noise_offset.clone(),
531 log_sigma_penalty_matrices,
532 vec![],
533 noise_log_lambdas,
534 noise_beta_hint,
535 1,
536 LOCATION_SCALE_N_OUTPUTS,
537 &format!("{context}: log_sigma"),
538 )?;
539 if thresholdspec.initial_beta.is_none() || log_sigmaspec.initial_beta.is_none() {
540 let (beta_t0, beta_ls0) = binomial_location_scalewarm_start(
541 y,
542 weights,
543 link_kind,
544 &thresholdspec,
545 &log_sigmaspec,
546 thresholdspec.initial_beta.as_ref(),
547 log_sigmaspec.initial_beta.as_ref(),
548 )?;
549 if thresholdspec.initial_beta.is_none() {
550 thresholdspec.initial_beta = Some(beta_t0);
551 }
552 if log_sigmaspec.initial_beta.is_none() {
553 log_sigmaspec.initial_beta = Some(beta_ls0);
554 }
555 }
556 Ok((thresholdspec, log_sigmaspec))
557}
558
559pub(crate) fn wiggle_block_penalty_matrices(
563 wiggle_block: &ParameterBlockInput,
564) -> Vec<PenaltyMatrix> {
565 let p_wiggle = wiggle_block.design.ncols();
566 wiggle_block
567 .penalties
568 .iter()
569 .map(|spec| match spec {
570 crate::model_types::PenaltySpec::Block {
571 local, col_range, ..
572 } => PenaltyMatrix::Blockwise {
573 local: local.clone(),
574 col_range: col_range.clone(),
575 total_dim: p_wiggle,
576 },
577 crate::model_types::PenaltySpec::Dense(m)
578 | crate::model_types::PenaltySpec::DenseWithMean { matrix: m, .. } => {
579 PenaltyMatrix::Dense(m.clone())
580 }
581 })
582 .collect()
583}
584
585pub(crate) fn binomial_location_scale_link_eta_from_probability(
586 link_kind: &InverseLink,
587 probability: f64,
588) -> Result<f64, String> {
589 let target = probability.clamp(1e-6, 1.0 - 1e-6);
590 match link_kind {
591 InverseLink::Standard(StandardLink::Logit) => Ok((target / (1.0 - target)).ln()),
592 InverseLink::Standard(StandardLink::Probit) => standard_normal_quantile(target)
593 .map_err(|err| format!("failed to invert probit warm-start probability: {err}")),
594 InverseLink::Standard(StandardLink::CLogLog) => Ok((-((1.0 - target).ln())).ln()),
595 other => Err(GamlssError::UnsupportedConfiguration { reason: format!(
596 "binomial location-scale warm start requires logit, probit, or cloglog link, got {other:?}"
597 ) }.into()),
598 }
599}
600
601pub(crate) fn weighted_binomial_prevalence(
602 y: &Array1<f64>,
603 weights: &Array1<f64>,
604) -> Result<f64, String> {
605 if y.len() != weights.len() {
606 return Err(GamlssError::DimensionMismatch { reason: format!(
607 "binomial location-scale warm start dimension mismatch: y has length {}, weights have length {}",
608 y.len(),
609 weights.len()
610 ) }.into());
611 }
612 let mut weight_sum = 0.0;
613 let mut success_sum = 0.0;
614 for (&yi, &wi) in y.iter().zip(weights.iter()) {
615 if !yi.is_finite() {
616 return Err(GamlssError::NonFinite {
617 reason: format!(
618 "binomial location-scale warm start encountered non-finite response {yi}"
619 ),
620 }
621 .into());
622 }
623 let weight = floor_positiveweight(wi, MIN_WEIGHT);
624 if weight > 0.0 {
625 weight_sum += weight;
626 success_sum += weight * yi;
627 }
628 }
629 if !weight_sum.is_finite() || weight_sum <= 0.0 {
630 return Err(
631 "binomial location-scale warm start requires positive total weight".to_string(),
632 );
633 }
634 Ok(success_sum / weight_sum)
635}
636
637pub(crate) fn project_constant_eta_into_block(
638 block: &ParameterBlockSpec,
639 weights: &Array1<f64>,
640 eta: f64,
641) -> Result<Array1<f64>, String> {
642 let target_eta = Array1::from_elem(block.design.nrows(), eta);
643 solve_penalizedweighted_projection(
644 &block.design,
645 &block.offset,
646 &target_eta,
647 weights,
648 &block.penalties,
649 &block.initial_log_lambdas,
650 1e-10,
651 )
652}
653
654pub(crate) fn binomial_location_scalewarm_start(
658 y: &Array1<f64>,
659 weights: &Array1<f64>,
660 link_kind: &InverseLink,
661 threshold_block: &ParameterBlockSpec,
662 log_sigma_block: &ParameterBlockSpec,
663 mean_beta_hint: Option<&Array1<f64>>,
664 noise_beta_hint: Option<&Array1<f64>>,
665) -> Result<(Array1<f64>, Array1<f64>), String> {
666 if let (Some(mean_beta), Some(noise_beta)) = (mean_beta_hint, noise_beta_hint) {
667 return Ok((mean_beta.clone(), noise_beta.clone()));
668 }
669
670 let beta_threshold = match mean_beta_hint {
671 Some(beta) => beta.clone(),
672 None => {
673 let prevalence = weighted_binomial_prevalence(y, weights)?;
674 let eta = binomial_location_scale_link_eta_from_probability(link_kind, prevalence)?;
675 project_constant_eta_into_block(threshold_block, weights, eta)?
676 }
677 };
678 let beta_log_sigma = match noise_beta_hint {
679 Some(beta) => beta.clone(),
680 None => project_constant_eta_into_block(log_sigma_block, weights, 0.0)?,
681 };
682 Ok((beta_threshold, beta_log_sigma))
683}
684
685#[derive(Clone)]
686pub(crate) struct BinomialMeanWiggleSpec {
687 pub y: Array1<f64>,
688 pub weights: Array1<f64>,
689 pub link_kind: InverseLink,
690 pub wiggle_knots: Array1<f64>,
691 pub wiggle_degree: usize,
692 pub eta_block: ParameterBlockInput,
693 pub wiggle_block: ParameterBlockInput,
694}
695
696#[derive(Clone)]
697pub struct GaussianLocationScaleTermSpec {
698 pub y: Array1<f64>,
699 pub weights: Array1<f64>,
700 pub meanspec: TermCollectionSpec,
701 pub log_sigmaspec: TermCollectionSpec,
702 pub mean_offset: Array1<f64>,
703 pub log_sigma_offset: Array1<f64>,
704}
705
706#[derive(Clone)]
707pub struct GaussianLocationScaleWiggleTermSpec {
708 pub y: Array1<f64>,
709 pub weights: Array1<f64>,
710 pub meanspec: TermCollectionSpec,
711 pub log_sigmaspec: TermCollectionSpec,
712 pub mean_offset: Array1<f64>,
713 pub log_sigma_offset: Array1<f64>,
714 pub wiggle_knots: Array1<f64>,
715 pub wiggle_degree: usize,
716 pub wiggle_block: ParameterBlockInput,
717}
718
719#[derive(Clone)]
720pub struct BinomialLocationScaleTermSpec {
721 pub y: Array1<f64>,
722 pub weights: Array1<f64>,
723 pub link_kind: InverseLink,
724 pub thresholdspec: TermCollectionSpec,
725 pub log_sigmaspec: TermCollectionSpec,
726 pub threshold_offset: Array1<f64>,
727 pub log_sigma_offset: Array1<f64>,
728}
729
730#[derive(Clone)]
731pub struct BinomialLocationScaleWiggleTermSpec {
732 pub y: Array1<f64>,
733 pub weights: Array1<f64>,
734 pub link_kind: InverseLink,
735 pub thresholdspec: TermCollectionSpec,
736 pub log_sigmaspec: TermCollectionSpec,
737 pub threshold_offset: Array1<f64>,
738 pub log_sigma_offset: Array1<f64>,
739 pub wiggle_knots: Array1<f64>,
740 pub wiggle_degree: usize,
741 pub wiggle_block: ParameterBlockInput,
742}
743
744#[derive(Clone, Debug)]
745pub struct BlockwiseTermFitResult {
746 pub fit: UnifiedFitResult,
747 pub meanspec_resolved: TermCollectionSpec,
748 pub noisespec_resolved: TermCollectionSpec,
749 pub mean_design: TermCollectionDesign,
750 pub noise_design: TermCollectionDesign,
751}
752
753pub(crate) struct BlockwiseTermFitResultParts {
754 pub fit: UnifiedFitResult,
755 pub meanspec_resolved: TermCollectionSpec,
756 pub noisespec_resolved: TermCollectionSpec,
757 pub mean_design: TermCollectionDesign,
758 pub noise_design: TermCollectionDesign,
759}
760
761pub struct BlockwiseTermWiggleFitResult {
762 pub fit: BlockwiseTermFitResult,
763 pub wiggle_knots: Array1<f64>,
764 pub wiggle_degree: usize,
765}
766
767pub struct BinomialMeanWiggleTermFitResult {
768 pub fit: UnifiedFitResult,
769 pub resolvedspec: TermCollectionSpec,
770 pub design: TermCollectionDesign,
771 pub wiggle_knots: Array1<f64>,
772 pub wiggle_degree: usize,
773 pub saved_warp_beta: Option<Vec<f64>>,
778 pub saved_index_shift: Option<Vec<f64>>,
783}
784
785pub(crate) struct BlockwiseTermWiggleFitResultParts {
786 pub fit: BlockwiseTermFitResult,
787 pub wiggle_knots: Array1<f64>,
788 pub wiggle_degree: usize,
789}
790
791pub(crate) fn validate_term_collection_design(
792 label: &str,
793 design: &TermCollectionDesign,
794) -> Result<(), String> {
795 let p = design.design.ncols();
796 let n = design.design.nrows();
797 for rows in exact_design_row_chunks(n, p) {
798 let chunk = design
799 .design
800 .try_row_chunk(rows)
801 .map_err(|e| format!("{label}.design row chunk materialization failed: {e}"))?;
802 validate_all_finite_estimation(&format!("{label}.design"), chunk.iter().copied())
803 .map_err(|e| e.to_string())?;
804 }
805 if design.nullspace_dims.len() != design.penalties.len() {
806 return Err(GamlssError::DimensionMismatch {
807 reason: format!(
808 "{label}.nullspace_dims length mismatch: got {}, expected {}",
809 design.nullspace_dims.len(),
810 design.penalties.len()
811 ),
812 }
813 .into());
814 }
815 if design.penaltyinfo.len() != design.penalties.len() {
816 return Err(GamlssError::DimensionMismatch {
817 reason: format!(
818 "{label}.penaltyinfo length mismatch: got {}, expected {}",
819 design.penaltyinfo.len(),
820 design.penalties.len()
821 ),
822 }
823 .into());
824 }
825 for (idx, bp) in design.penalties.iter().enumerate() {
826 validate_all_finite_estimation(
827 &format!("{label}.penalties[{idx}]"),
828 bp.local.iter().copied(),
829 )
830 .map_err(|e| e.to_string())?;
831 if bp.col_range.end > p {
832 return Err(GamlssError::DimensionMismatch {
833 reason: format!(
834 "{label}.penalties[{idx}] col_range {}..{} exceeds design width {}",
835 bp.col_range.start, bp.col_range.end, p
836 ),
837 }
838 .into());
839 }
840 }
841 if let Some(bounds) = design.coefficient_lower_bounds.as_ref() {
842 if bounds.len() != p {
843 return Err(GamlssError::ConstraintViolation {
844 reason: format!(
845 "{label}.coefficient_lower_bounds length mismatch: got {}, expected {p}",
846 bounds.len()
847 ),
848 }
849 .into());
850 }
851 for (idx, &bound) in bounds.iter().enumerate() {
852 if !(bound.is_finite() || bound == f64::NEG_INFINITY) {
853 return Err(GamlssError::NonFinite { reason: format!(
854 "{label}.coefficient_lower_bounds[{idx}] must be finite or -inf, got {bound}",
855 ) }.into());
856 }
857 }
858 }
859 if let Some(constraints) = design.linear_constraints.as_ref() {
860 validate_all_finite_estimation(
861 &format!("{label}.linear_constraints.a"),
862 constraints.a.iter().copied(),
863 )
864 .map_err(|e| e.to_string())?;
865 validate_all_finite_estimation(
866 &format!("{label}.linear_constraints.b"),
867 constraints.b.iter().copied(),
868 )
869 .map_err(|e| e.to_string())?;
870 if constraints.a.ncols() != p {
871 return Err(GamlssError::DimensionMismatch {
872 reason: format!(
873 "{label}.linear_constraints.a column mismatch: got {}, expected {p}",
874 constraints.a.ncols()
875 ),
876 }
877 .into());
878 }
879 if constraints.a.nrows() != constraints.b.len() {
880 return Err(GamlssError::DimensionMismatch {
881 reason: format!(
882 "{label}.linear_constraints row mismatch: a has {}, b has {}",
883 constraints.a.nrows(),
884 constraints.b.len()
885 ),
886 }
887 .into());
888 }
889 }
890 if design.intercept_range.start > design.intercept_range.end || design.intercept_range.end > p {
891 return Err(GamlssError::ConstraintViolation {
892 reason: format!(
893 "{label}.intercept_range out of bounds: {:?} for {} columns",
894 design.intercept_range, p
895 ),
896 }
897 .into());
898 }
899 Ok(())
900}
901
902impl BlockwiseTermFitResult {
903 pub(crate) fn try_from_parts(parts: BlockwiseTermFitResultParts) -> Result<Self, String> {
904 let BlockwiseTermFitResultParts {
905 fit,
906 meanspec_resolved,
907 noisespec_resolved,
908 mean_design,
909 noise_design,
910 } = parts;
911
912 fit.validate_numeric_finiteness()
913 .map_err(|e| format!("{e}"))?;
914 if fit.block_states.len() < 2 {
915 return Err(GamlssError::DimensionMismatch {
916 reason: format!(
917 "BlockwiseTermFitResult requires at least 2 block states, got {}",
918 fit.block_states.len()
919 ),
920 }
921 .into());
922 }
923 validate_term_collection_design("blockwise_term.mean_design", &mean_design)?;
924 validate_term_collection_design("blockwise_term.noise_design", &noise_design)?;
925 if mean_design.design.nrows() != noise_design.design.nrows() {
926 return Err(GamlssError::DimensionMismatch {
927 reason: format!(
928 "BlockwiseTermFitResult row mismatch: mean_design={}, noise_design={}",
929 mean_design.design.nrows(),
930 noise_design.design.nrows()
931 ),
932 }
933 .into());
934 }
935 if fit.block_states[0].beta.len() != mean_design.design.ncols() {
936 return Err(GamlssError::DimensionMismatch {
937 reason: format!(
938 "BlockwiseTermFitResult mean beta length mismatch: got {}, expected {}",
939 fit.block_states[0].beta.len(),
940 mean_design.design.ncols()
941 ),
942 }
943 .into());
944 }
945 if fit.block_states[1].beta.len() != noise_design.design.ncols() {
946 return Err(GamlssError::DimensionMismatch {
947 reason: format!(
948 "BlockwiseTermFitResult noise beta length mismatch: got {}, expected {}",
949 fit.block_states[1].beta.len(),
950 noise_design.design.ncols()
951 ),
952 }
953 .into());
954 }
955 if fit.block_states[0].eta.len() != mean_design.design.nrows() {
956 return Err(GamlssError::DimensionMismatch {
957 reason: format!(
958 "BlockwiseTermFitResult mean eta length mismatch: got {}, expected {}",
959 fit.block_states[0].eta.len(),
960 mean_design.design.nrows()
961 ),
962 }
963 .into());
964 }
965 if fit.block_states[1].eta.len() != noise_design.design.nrows() {
966 return Err(GamlssError::DimensionMismatch {
967 reason: format!(
968 "BlockwiseTermFitResult noise eta length mismatch: got {}, expected {}",
969 fit.block_states[1].eta.len(),
970 noise_design.design.nrows()
971 ),
972 }
973 .into());
974 }
975
976 Ok(Self {
977 fit,
978 meanspec_resolved,
979 noisespec_resolved,
980 mean_design,
981 noise_design,
982 })
983 }
984
985 pub(crate) fn validate_numeric_finiteness(&self) -> Result<(), String> {
986 Self::try_from_parts(BlockwiseTermFitResultParts {
987 fit: self.fit.clone(),
988 meanspec_resolved: self.meanspec_resolved.clone(),
989 noisespec_resolved: self.noisespec_resolved.clone(),
990 mean_design: self.mean_design.clone(),
991 noise_design: self.noise_design.clone(),
992 })
993 .map(|_| ())
994 }
995}
996
997impl BlockwiseTermWiggleFitResult {
998 pub(crate) fn try_from_parts(parts: BlockwiseTermWiggleFitResultParts) -> Result<Self, String> {
999 let BlockwiseTermWiggleFitResultParts {
1000 fit,
1001 wiggle_knots,
1002 wiggle_degree,
1003 } = parts;
1004
1005 fit.validate_numeric_finiteness()
1006 .map_err(|e| e.to_string())?;
1007 if fit.fit.block_states.len() < 3 {
1008 return Err(GamlssError::DimensionMismatch {
1009 reason: format!(
1010 "BlockwiseTermWiggleFitResult requires at least 3 block states, got {}",
1011 fit.fit.block_states.len()
1012 ),
1013 }
1014 .into());
1015 }
1016 if wiggle_knots.is_empty() {
1017 return Err(GamlssError::UnsupportedConfiguration {
1018 reason: "BlockwiseTermWiggleFitResult requires non-empty wiggle_knots".to_string(),
1019 }
1020 .into());
1021 }
1022 validate_all_finite_estimation(
1023 "blockwise_term_wiggle.wiggle_knots",
1024 wiggle_knots.iter().copied(),
1025 )
1026 .map_err(|e| e.to_string())?;
1027
1028 Ok(Self {
1029 fit,
1030 wiggle_knots,
1031 wiggle_degree,
1032 })
1033 }
1034}
1035
1036pub struct BinomialLocationScaleFitResult {
1037 pub fit: BlockwiseTermFitResult,
1038 pub wiggle_knots: Option<Array1<f64>>,
1039 pub wiggle_degree: Option<usize>,
1040 pub beta_link_wiggle: Option<Vec<f64>>,
1041}
1042
1043pub struct GaussianLocationScaleFitResult {
1044 pub fit: BlockwiseTermFitResult,
1045 pub wiggle_knots: Option<Array1<f64>>,
1046 pub wiggle_degree: Option<usize>,
1047 pub beta_link_wiggle: Option<Vec<f64>>,
1048 pub response_scale: f64,
1077}
1078
1079pub(crate) fn fit_binomial_mean_wiggle(
1083 spec: BinomialMeanWiggleSpec,
1084 options: &BlockwiseFitOptions,
1085) -> Result<(UnifiedFitResult, Option<Vec<f64>>, Option<Vec<f64>>), String> {
1086 let n = spec.y.len();
1087 validate_len_match("weights vs y", n, spec.weights.len())?;
1088 validateweights(&spec.weights, "fit_binomial_mean_wiggle")?;
1089 validate_binomial_response(&spec.y, "fit_binomial_mean_wiggle")?;
1090 validate_blockrows("eta", n, &spec.eta_block)?;
1091 validate_blockrows("wiggle", n, &spec.wiggle_block)?;
1092 if matches!(
1093 spec.link_kind,
1094 InverseLink::Standard(StandardLink::Identity)
1095 ) {
1096 return Err(GamlssError::UnsupportedConfiguration {
1097 reason: "fit_binomial_mean_wiggle does not support identity link".to_string(),
1098 }
1099 .into());
1100 }
1101 gam_terms::inference::formula_dsl::require_binomial_inverse_link_supports_joint_wiggle(
1102 &spec.link_kind,
1103 "fit_binomial_mean_wiggle",
1104 )?;
1105 if spec.wiggle_degree < 2 {
1106 return Err(GamlssError::ConstraintViolation {
1107 reason: format!(
1108 "fit_binomial_mean_wiggle: wiggle_degree must be >= 2, got {}",
1109 spec.wiggle_degree
1110 ),
1111 }
1112 .into());
1113 }
1114 let minimum_knots = minimum_monotone_wiggle_knot_count(spec.wiggle_degree)?;
1115 if spec.wiggle_knots.len() < minimum_knots {
1116 return Err(GamlssError::DimensionMismatch { reason: format!(
1117 "fit_binomial_mean_wiggle: wiggle_knots length {} is too short for degree {} (need at least {})",
1118 spec.wiggle_knots.len(),
1119 spec.wiggle_degree,
1120 minimum_knots
1121 ) }.into());
1122 }
1123
1124 let x_dense: Array2<f64> = spec.eta_block.design.to_dense();
1133 let (pilot_beta, pilot_eta): (Array1<f64>, Array1<f64>) = {
1134 let pilot_beta = spec.eta_block.initial_beta.clone().ok_or_else(|| {
1135 "fit_binomial_mean_wiggle: eta block carries no pilot β to seed the \
1136 frozen-basis warp index"
1137 .to_string()
1138 })?;
1139 if x_dense.ncols() != pilot_beta.len() {
1140 return Err(GamlssError::DimensionMismatch {
1141 reason: format!(
1142 "fit_binomial_mean_wiggle: eta design has {} columns but pilot β has {} \
1143 coefficients",
1144 x_dense.ncols(),
1145 pilot_beta.len()
1146 ),
1147 }
1148 .into());
1149 }
1150 let mut eta = x_dense.dot(&pilot_beta);
1151 eta += &spec.eta_block.offset;
1152 (pilot_beta, eta)
1153 };
1154
1155 let wiggle_penalties_full = spec.wiggle_block.penalties.clone();
1159 let wiggle_nullspace_dims = spec.wiggle_block.nullspace_dims.clone();
1160 if !wiggle_nullspace_dims.is_empty()
1161 && wiggle_nullspace_dims.len() != wiggle_penalties_full.len()
1162 {
1163 return Err(GamlssError::DimensionMismatch {
1164 reason: format!(
1165 "fit_binomial_mean_wiggle: wiggle block has {} penalties but {} nullspace dimensions",
1166 wiggle_penalties_full.len(),
1167 wiggle_nullspace_dims.len()
1168 ),
1169 }
1170 .into());
1171 }
1172 let wiggle_log_lambdas = spec.wiggle_block.initial_log_lambdas.clone();
1173 let wiggle_beta_initial = spec.wiggle_block.initial_beta.clone();
1174 let eta_block_input = spec.eta_block.clone();
1175
1176 let family = BinomialMeanWiggleFamily {
1177 y: spec.y,
1178 weights: spec.weights,
1179 link_kind: spec.link_kind,
1180 wiggle_knots: spec.wiggle_knots,
1181 wiggle_degree: spec.wiggle_degree,
1182 policy: gam_runtime::resource::ResourcePolicy::default_library(),
1183 frozen_warp_design: None,
1184 };
1185
1186 let build_dealiased = |frozen: &Array1<f64>,
1204 beta_hint: Option<&Array1<f64>>,
1205 log_lambda_hint: Option<&Array1<f64>>|
1206 -> Result<
1207 (
1208 ParameterBlockInput,
1209 Array2<f64>,
1210 std::sync::Arc<Array2<f64>>,
1211 ),
1212 String,
1213 > {
1214 use faer::Side;
1215 use gam_linalg::faer_ndarray::FaerEigh;
1216
1217 let b_full = family.wiggle_design(frozen.view())?;
1218 let xtx = x_dense.t().dot(&x_dense);
1219 let xtb = x_dense.t().dot(&b_full);
1220 let (evals, evecs) = xtx
1221 .eigh(Side::Lower)
1222 .map_err(|e| format!("frozen-basis warp de-aliasing mean QR failed: {e}"))?;
1223 let max_eval = evals.iter().map(|v| v.abs()).fold(0.0_f64, f64::max);
1224 let cutoff = 1.0e3 * f64::EPSILON * (xtx.nrows().max(1) as f64) * max_eval.max(1.0);
1225 let mut alias = Array2::<f64>::zeros((x_dense.ncols(), b_full.ncols()));
1226 for k in 0..evals.len() {
1227 let lam = evals[k];
1228 if !lam.is_finite() || lam.abs() <= cutoff {
1229 continue;
1230 }
1231 let uk = evecs.column(k);
1232 let uk_xtb = uk.t().dot(&xtb);
1233 for i in 0..alias.nrows() {
1234 for j in 0..alias.ncols() {
1235 alias[[i, j]] += uk[i] * uk_xtb[j] / lam;
1236 }
1237 }
1238 }
1239 let bda = &b_full - &x_dense.dot(&alias);
1240 let max_b = b_full.iter().map(|v| v.abs()).fold(0.0_f64, f64::max);
1241 let max_resid = bda.iter().map(|v| v.abs()).fold(0.0_f64, f64::max);
1242 let resid_tol =
1243 1.0e3 * f64::EPSILON * (bda.nrows().max(bda.ncols()).max(1) as f64) * max_b.max(1.0);
1244 if max_resid <= resid_tol {
1245 return Err("frozen-basis warp de-aliasing left no identifiable warp \
1246 direction (the mean block already spans the warp in \
1247 observation space)"
1248 .to_string());
1249 }
1250 let penalties: Vec<crate::model_types::PenaltySpec> = wiggle_penalties_full
1251 .iter()
1252 .map(|p| {
1253 let s = penalty_spec_to_dense(p, b_full.ncols())?;
1254 Ok(crate::model_types::PenaltySpec::Dense(s))
1255 })
1256 .collect::<Result<_, String>>()?;
1257 let q = bda.ncols();
1258 let initial_beta = match beta_hint {
1259 Some(beta) if beta.len() == q => Some(beta.clone()),
1260 Some(beta) => {
1261 return Err(format!(
1262 "frozen-basis warp warm start has {} coefficients but the realized basis has {q}",
1263 beta.len()
1264 ));
1265 }
1266 None => Some(Array1::zeros(q)),
1267 };
1268 let block = ParameterBlockInput {
1269 design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(bda.clone())),
1270 offset: Array1::zeros(frozen.len()),
1271 penalties,
1272 nullspace_dims: wiggle_nullspace_dims.clone(),
1273 initial_log_lambdas: log_lambda_hint
1274 .cloned()
1275 .or_else(|| wiggle_log_lambdas.clone()),
1276 initial_beta,
1277 };
1278 Ok((block, alias, std::sync::Arc::new(bda)))
1279 };
1280
1281 if options.outer_max_iter == 0 || !options.outer_tol.is_finite() || options.outer_tol <= 0.0 {
1288 return Err(GamlssError::InvalidInput {
1289 reason: format!(
1290 "fit_binomial_mean_wiggle requires positive outer convergence policy; outer_max_iter={}, outer_tol={}",
1291 options.outer_max_iter, options.outer_tol
1292 ),
1293 }
1294 .into());
1295 }
1296 let mut frozen_source_beta = pilot_beta;
1297 let mut frozen_eta = pilot_eta;
1298 let mut eta_block_warm = eta_block_input.clone();
1299 let mut wiggle_beta_warm = wiggle_beta_initial;
1300 let mut wiggle_log_lambda_warm = wiggle_log_lambdas.clone();
1301 let mut converged: Option<(UnifiedFitResult, Array2<f64>, Array1<f64>)> = None;
1302 let mut last_delta = f64::INFINITY;
1303 let mut last_scale = 1.0_f64;
1304 for _outer in 0..options.outer_max_iter {
1305 let (wiggle_block, alias, bda) = build_dealiased(
1306 &frozen_eta,
1307 wiggle_beta_warm.as_ref(),
1308 wiggle_log_lambda_warm.as_ref(),
1309 )?;
1310 let eta_penalty_count = eta_block_warm.penalties.len();
1311 let wiggle_penalty_count = wiggle_block.penalties.len();
1312 let blocks = vec![
1313 eta_block_warm.clone().intospec("eta")?,
1314 wiggle_block.intospec("wiggle")?,
1315 ];
1316 let mut fam = family.clone();
1317 fam.frozen_warp_design = Some(bda);
1318 let fit = fit_custom_family(&fam, &blocks, options).map_err(|e| e.to_string())?;
1319 let mean_state = fit
1320 .block_states
1321 .get(BinomialMeanWiggleFamily::BLOCK_ETA)
1322 .ok_or_else(|| {
1323 "fit_binomial_mean_wiggle: frozen-basis refit did not expose a fitted eta block"
1324 .to_string()
1325 })?;
1326 if mean_state.eta.len() != frozen_eta.len()
1327 || mean_state.beta.len() != frozen_source_beta.len()
1328 {
1329 return Err(GamlssError::DimensionMismatch {
1330 reason: "fit_binomial_mean_wiggle: frozen-basis refit returned an incompatible eta block"
1331 .to_string(),
1332 }
1333 .into());
1334 }
1335 let new_eta = mean_state.eta.clone();
1336 let new_source_beta = mean_state.beta.clone();
1337 let new_wiggle_beta = fit
1338 .block_states
1339 .get(BinomialMeanWiggleFamily::BLOCK_WIGGLE)
1340 .map(|state| state.beta.clone())
1341 .ok_or_else(|| {
1342 "fit_binomial_mean_wiggle: frozen-basis refit did not expose a fitted wiggle block"
1343 .to_string()
1344 })?;
1345 last_scale = frozen_eta
1346 .iter()
1347 .chain(new_eta.iter())
1348 .map(|value| value.abs())
1349 .fold(1.0_f64, f64::max);
1350 last_delta = new_eta
1351 .iter()
1352 .zip(frozen_eta.iter())
1353 .map(|(a, b)| (a - b).abs())
1354 .fold(0.0_f64, f64::max);
1355 if last_delta <= options.outer_tol * last_scale {
1356 converged = Some((fit, alias, frozen_source_beta));
1357 break;
1358 }
1359
1360 let expected_log_lambdas = eta_penalty_count + wiggle_penalty_count;
1361 if fit.log_lambdas.len() != expected_log_lambdas {
1362 return Err(GamlssError::DimensionMismatch {
1363 reason: format!(
1364 "fit_binomial_mean_wiggle: refit returned {} log-lambdas for {expected_log_lambdas} penalties",
1365 fit.log_lambdas.len()
1366 ),
1367 }
1368 .into());
1369 }
1370 eta_block_warm.initial_beta = Some(new_source_beta.clone());
1371 eta_block_warm.initial_log_lambdas =
1372 Some(fit.log_lambdas.slice(s![0..eta_penalty_count]).to_owned());
1373 wiggle_beta_warm = Some(new_wiggle_beta);
1374 wiggle_log_lambda_warm = Some(
1375 fit.log_lambdas
1376 .slice(s![eta_penalty_count..expected_log_lambdas])
1377 .to_owned(),
1378 );
1379 frozen_source_beta = new_source_beta;
1380 frozen_eta = new_eta;
1381 }
1382 let (mut fit, last_alias, frozen_source_beta) = converged.ok_or_else(|| {
1383 GamlssError::NumericalFailure {
1384 reason: format!(
1385 "fit_binomial_mean_wiggle frozen-index fixed point did not converge in {} outer iterations: delta={last_delta:.3e}, scale={last_scale:.3e}, tolerance={:.3e}",
1386 options.outer_max_iter,
1387 options.outer_tol * last_scale,
1388 ),
1389 }
1390 .to_string()
1391 })?;
1392 let frozen_source_beta = Some(frozen_source_beta);
1402 let saved_warp_beta = fit
1409 .block_states
1410 .get(BinomialMeanWiggleFamily::BLOCK_WIGGLE)
1411 .map(|state| state.beta.to_vec());
1412 if let Some(beta_w) = saved_warp_beta.as_ref() {
1413 validate_monotone_wiggle_beta_nonnegative(beta_w, "fit_binomial_mean_wiggle saved warp")?;
1414 }
1415 if let Some(beta_w) = saved_warp_beta.as_ref() {
1416 let alias = &last_alias;
1417 let beta_w = Array1::from_vec(beta_w.clone());
1418 if alias.ncols() == beta_w.len() && alias.nrows() == eta_block_input.design.ncols() {
1419 let shift = alias.dot(&beta_w);
1420 if let Some(block) = fit.blocks.get_mut(BinomialMeanWiggleFamily::BLOCK_ETA) {
1421 if block.beta.len() == shift.len() {
1422 block.beta -= &shift;
1423 }
1424 }
1425 if let Some(state) = fit
1426 .block_states
1427 .get_mut(BinomialMeanWiggleFamily::BLOCK_ETA)
1428 {
1429 if state.beta.len() == shift.len() {
1430 state.beta -= &shift;
1431 state.eta = x_dense.dot(&state.beta) + &eta_block_input.offset;
1432 }
1433 }
1434 if fit.beta.len() >= shift.len() {
1435 for i in 0..shift.len() {
1436 fit.beta[i] -= shift[i];
1437 }
1438 }
1439 }
1440 }
1441 let saved_index_shift: Option<Vec<f64>> = match (
1446 saved_warp_beta.as_ref(),
1447 frozen_source_beta.as_ref(),
1448 fit.block_states.get(BinomialMeanWiggleFamily::BLOCK_ETA),
1449 ) {
1450 (Some(_), Some(source), Some(state)) if source.len() == state.beta.len() => {
1451 Some((source - &state.beta).to_vec())
1452 }
1453 _ => None,
1454 };
1455 Ok((fit, saved_warp_beta, saved_index_shift))
1456}
1457
1458fn penalty_spec_to_dense(
1462 spec: &crate::model_types::PenaltySpec,
1463 p: usize,
1464) -> Result<Array2<f64>, String> {
1465 use crate::model_types::PenaltySpec;
1466 match spec {
1467 PenaltySpec::Dense(m) | PenaltySpec::DenseWithMean { matrix: m, .. } => {
1468 if m.nrows() != p || m.ncols() != p {
1469 return Err(format!(
1470 "frozen-basis warp penalty must be {p}x{p}, got {}x{}",
1471 m.nrows(),
1472 m.ncols()
1473 ));
1474 }
1475 Ok(m.clone())
1476 }
1477 PenaltySpec::Block {
1478 local, col_range, ..
1479 } => {
1480 let mut full = Array2::<f64>::zeros((p, p));
1481 if col_range.end > p || local.nrows() != col_range.len() {
1482 return Err("frozen-basis warp penalty block range out of bounds".to_string());
1483 }
1484 full.slice_mut(s![col_range.clone(), col_range.clone()])
1485 .assign(local);
1486 Ok(full)
1487 }
1488 }
1489}
1490
1491pub(crate) trait LocationScaleFamilyBuilder {
1492 type Family: CustomFamily + Clone + Send + Sync + 'static;
1493
1494 fn meanspec(&self) -> &TermCollectionSpec;
1495 fn noisespec(&self) -> &TermCollectionSpec;
1496
1497 fn build_blocks(
1498 &self,
1499 theta: &Array1<f64>,
1500 mean_design: &TermCollectionDesign,
1501 noise_design: &TermCollectionDesign,
1502 mean_beta_hint: Option<Array1<f64>>,
1503 noise_beta_hint: Option<Array1<f64>>,
1504 ) -> Result<Vec<ParameterBlockSpec>, String>;
1505
1506 fn build_family(
1507 &self,
1508 mean_design: &TermCollectionDesign,
1509 noise_design: &TermCollectionDesign,
1510 ) -> Self::Family;
1511
1512 fn extract_primary_betas(
1513 &self,
1514 fit: &UnifiedFitResult,
1515 ) -> Result<(Array1<f64>, Array1<f64>), String>;
1516
1517 fn mean_penalty_count(&self, mean_design: &TermCollectionDesign) -> usize {
1518 mean_design.penalties.len()
1519 }
1520
1521 fn noise_penalty_count(&self, noise_design: &TermCollectionDesign) -> usize {
1522 noise_design.penalties.len()
1523 }
1524
1525 fn exact_spatial_joint_supported(&self) -> bool {
1526 false
1527 }
1528
1529 fn require_exact_spatial_joint(&self) -> bool {
1530 false
1531 }
1532
1533 fn exact_spatial_seed_risk_profile(&self) -> crate::seeding::SeedRiskProfile {
1534 crate::seeding::SeedRiskProfile::GeneralizedLinear
1535 }
1536
1537 fn extra_rho0(&self) -> Result<Array1<f64>, String> {
1538 Ok(Array1::zeros(0))
1539 }
1540
1541 fn build_psiderivative_blocks(
1542 &self,
1543 arr: ndarray::ArrayView2<'_, f64>,
1544 term_spec: &TermCollectionSpec,
1545 term_spec2: &TermCollectionSpec,
1546 term_design: &TermCollectionDesign,
1547 term_design2: &TermCollectionDesign,
1548 ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String>;
1549}
1550
1551pub(crate) fn fit_location_scale_terms<B: LocationScaleFamilyBuilder>(
1552 data: ndarray::ArrayView2<'_, f64>,
1553 builder: B,
1554 options: &BlockwiseFitOptions,
1555 kappa_options: &SpatialLengthScaleOptimizationOptions,
1556) -> Result<BlockwiseTermFitResult, String> {
1557 let mut mean_beta_hint: Option<Array1<f64>> = None;
1563 let mut noise_beta_hint: Option<Array1<f64>> = None;
1564 let extra_rho0 = builder.extra_rho0()?;
1565
1566 let mean_boot_design =
1567 build_term_collection_design(data, builder.meanspec()).map_err(|e| e.to_string())?;
1568 let noise_boot_design =
1569 build_term_collection_design(data, builder.noisespec()).map_err(|e| e.to_string())?;
1570 let mean_bootspec = freeze_term_collection_from_design(builder.meanspec(), &mean_boot_design)
1571 .map_err(|e| e.to_string())?;
1572 let noise_bootspec =
1573 freeze_term_collection_from_design(builder.noisespec(), &noise_boot_design)
1574 .map_err(|e| e.to_string())?;
1575
1576 let require_exact_spatial_joint = builder.require_exact_spatial_joint();
1577 let analytic_joint_derivatives_check = if builder.exact_spatial_joint_supported() {
1578 builder
1579 .build_psiderivative_blocks(
1580 data,
1581 &mean_bootspec,
1582 &noise_bootspec,
1583 &mean_boot_design,
1584 &noise_boot_design,
1585 )
1586 .map(|_| ())
1587 } else {
1588 Err(
1589 "analytic spatial psi derivatives are unavailable for this location-scale family"
1590 .to_string(),
1591 )
1592 };
1593 let analytic_joint_derivatives_available = analytic_joint_derivatives_check.is_ok();
1594 if require_exact_spatial_joint {
1595 analytic_joint_derivatives_check.map_err(|err| {
1596 format!("exact two-block spatial path requires analytic psi derivatives: {err}")
1597 })?;
1598 }
1599 let mean_penalty_count = builder.mean_penalty_count(&mean_boot_design);
1600 let noise_penalty_count = builder.noise_penalty_count(&noise_boot_design);
1601
1602 let mut effective_kappa_options = kappa_options.clone();
1613 if effective_kappa_options.enabled
1614 && gam_terms::smooth::all_spatial_terms_kappa_fixed(&mean_bootspec)
1615 && gam_terms::smooth::all_spatial_terms_kappa_fixed(&noise_bootspec)
1616 {
1617 log::info!(
1618 "[GAMLSS spatial] disabling κ/ψ optimization: every spatial term in \
1619 both blocks has an explicit length_scale and no anisotropy; \
1620 user-supplied kernel scale is fixed"
1621 );
1622 effective_kappa_options.enabled = false;
1623 }
1624 let kappa_options: &SpatialLengthScaleOptimizationOptions = &effective_kappa_options;
1625
1626 macro_rules! run_exact_joint_spatial {
1630 () => {{
1631 let joint_setup = build_two_block_exact_joint_setup(
1632 data,
1633 builder.meanspec(),
1634 builder.noisespec(),
1635 mean_penalty_count,
1636 noise_penalty_count,
1637 extra_rho0.as_slice().unwrap_or(&[]),
1638 None,
1639 kappa_options,
1640 );
1641 let mean_terms = spatial_length_scale_term_indices(builder.meanspec());
1642 let noise_terms = spatial_length_scale_term_indices(builder.noisespec());
1643 let mean_beta_hint_cell = std::cell::RefCell::new(mean_beta_hint.clone());
1644 let noise_beta_hint_cell = std::cell::RefCell::new(noise_beta_hint.clone());
1645 let hyper_warm_start_cell =
1646 std::cell::RefCell::new(None::<CustomFamilyWarmStart>);
1647 let gamlss_disable_fixed_point = true;
1657 let outer_policy = {
1658 let theta_seed = joint_setup.theta0();
1670 let rho_dim = joint_setup.rho_dim();
1671 let psi_dim = theta_seed.len() - rho_dim;
1672 let rho_seed = theta_seed.slice(s![..rho_dim]).to_owned();
1673 let policy_blocks_res = builder.build_blocks(
1674 &rho_seed,
1675 &mean_boot_design,
1676 &noise_boot_design,
1677 mean_beta_hint_cell.borrow().clone(),
1678 noise_beta_hint_cell.borrow().clone(),
1679 );
1680 let mut policy = match policy_blocks_res {
1681 Ok(policy_blocks) => {
1682 let policy_family =
1683 builder.build_family(&mean_boot_design, &noise_boot_design);
1684 crate::custom_family::CustomFamily::outer_derivative_policy(
1685 &policy_family,
1686 &policy_blocks,
1687 psi_dim,
1688 options,
1689 )
1690 }
1691 Err(err) => {
1692 log::warn!(
1700 "[GAMLSS spatial] failed to realize policy blocks at seed rho ({err}); \
1701 routing outer optimizer through gradient-only BFGS"
1702 );
1703 let capability = if analytic_joint_derivatives_available {
1704 crate::custom_family::ExactOuterDerivativeOrder::Second
1705 } else {
1706 crate::custom_family::ExactOuterDerivativeOrder::First
1707 };
1708 crate::custom_family::OuterDerivativePolicy {
1709 capability,
1710 predicted_gradient_work: u128::MAX,
1711 predicted_hessian_work: u128::MAX,
1712 subsample_capable: false,
1717 }
1718 }
1719 };
1720 if !analytic_joint_derivatives_available {
1721 policy.capability =
1725 crate::custom_family::ExactOuterDerivativeOrder::First;
1726 }
1727 policy
1728 };
1729 optimize_spatial_length_scale_exact_joint(
1730 data,
1731 &[builder.meanspec().clone(), builder.noisespec().clone()],
1732 &[mean_terms, noise_terms],
1733 kappa_options,
1734 &joint_setup,
1735 builder.exact_spatial_seed_risk_profile(),
1736 analytic_joint_derivatives_available,
1737 analytic_joint_derivatives_available,
1738 gamlss_disable_fixed_point,
1739 None,
1740 outer_policy,
1741 |theta, specs: &[TermCollectionSpec], designs: &[TermCollectionDesign]| {
1742 assert_eq!(
1743 specs.len(),
1744 2,
1745 "joint spatial closure expects exactly two block specs (mean, noise); got {}",
1746 specs.len(),
1747 );
1748 assert_eq!(
1749 designs.len(),
1750 2,
1751 "joint spatial closure expects exactly two block designs (mean, noise); got {}",
1752 designs.len(),
1753 );
1754 let rho = theta.slice(s![..joint_setup.rho_dim()]).to_owned();
1755 let fit = {
1756 let blocks = builder.build_blocks(
1757 &rho,
1758 &designs[0],
1759 &designs[1],
1760 mean_beta_hint_cell.borrow().clone(),
1761 noise_beta_hint_cell.borrow().clone(),
1762 )?;
1763 if mean_beta_hint_cell.borrow().is_none()
1764 && let Some(beta) = blocks.first().and_then(|block| block.initial_beta.clone())
1765 {
1766 *mean_beta_hint_cell.borrow_mut() = Some(beta);
1767 }
1768 if noise_beta_hint_cell.borrow().is_none()
1769 && let Some(beta) =
1770 blocks.get(1).and_then(|block| block.initial_beta.clone())
1771 {
1772 *noise_beta_hint_cell.borrow_mut() = Some(beta);
1773 }
1774 let family = builder.build_family(&designs[0], &designs[1]);
1775 if joint_setup.log_kappa_dim() > 0 && kappa_options.enabled {
1797 let warm_start = hyper_warm_start_cell.borrow().clone();
1798 fit_custom_family_fixed_log_lambdas(
1799 &family,
1800 &blocks,
1801 options,
1802 warm_start.as_ref(),
1803 0,
1804 None,
1805 true,
1806 )?
1807 } else {
1808 fit_custom_family(&family, &blocks, options)?
1809 }
1810 };
1811 let (mean_beta, noise_beta) = builder.extract_primary_betas(&fit)?;
1812 mean_beta_hint = Some(mean_beta);
1813 noise_beta_hint = Some(noise_beta);
1814 *mean_beta_hint_cell.borrow_mut() = mean_beta_hint.clone();
1815 *noise_beta_hint_cell.borrow_mut() = noise_beta_hint.clone();
1816 Ok(fit)
1817 },
1818 |theta,
1819 specs: &[TermCollectionSpec],
1820 designs: &[TermCollectionDesign],
1821 eval_mode,
1822 row_set: &crate::row_kernel::RowSet| {
1823 use gam_problem::EvalMode;
1824 if !analytic_joint_derivatives_available {
1825 return Err(
1826 "analytic spatial psi derivatives are unavailable for this exact two-block path"
1827 .to_string(),
1828 );
1829 }
1830 let rho = theta.slice(s![..joint_setup.rho_dim()]).to_owned();
1831 let blocks = builder.build_blocks(
1832 &rho,
1833 &designs[0],
1834 &designs[1],
1835 mean_beta_hint_cell.borrow().clone(),
1836 noise_beta_hint_cell.borrow().clone(),
1837 )?;
1838 if mean_beta_hint_cell.borrow().is_none()
1839 && let Some(beta) = blocks.first().and_then(|block| block.initial_beta.clone())
1840 {
1841 *mean_beta_hint_cell.borrow_mut() = Some(beta);
1842 }
1843 if noise_beta_hint_cell.borrow().is_none()
1844 && let Some(beta) = blocks.get(1).and_then(|block| block.initial_beta.clone())
1845 {
1846 *noise_beta_hint_cell.borrow_mut() = Some(beta);
1847 }
1848 let family = builder.build_family(&designs[0], &designs[1]);
1849 let psiderivative_blocks = if matches!(eval_mode, EvalMode::ValueOnly) {
1850 (0..specs.len()).map(|_| Vec::new()).collect()
1861 } else {
1862 builder.build_psiderivative_blocks(
1863 data,
1864 &specs[0],
1865 &specs[1],
1866 &designs[0],
1867 &designs[1],
1868 )?
1869 };
1870 let warm_start = hyper_warm_start_cell.borrow().clone();
1871 let eval_options = match row_set {
1878 crate::row_kernel::RowSet::All => {
1879 std::borrow::Cow::Borrowed(options)
1880 }
1881 crate::row_kernel::RowSet::Subsample {
1882 rows,
1883 n_full,
1884 } => {
1885 let subsample = crate::outer_subsample::
1886 OuterScoreSubsample::from_weighted_rows(
1887 (**rows).clone(),
1888 *n_full,
1889 *n_full as u64,
1890 );
1891 let mut cloned = options.clone();
1892 cloned.outer_score_subsample =
1893 Some(std::sync::Arc::new(subsample));
1894 std::borrow::Cow::Owned(cloned)
1895 }
1896 };
1897 let eval = evaluate_custom_family_joint_hyper(
1898 &family,
1899 &blocks,
1900 eval_options.as_ref(),
1901 &rho,
1902 &psiderivative_blocks,
1903 warm_start.as_ref(),
1904 eval_mode,
1905 )?;
1906 *hyper_warm_start_cell.borrow_mut() = Some(eval.warm_start.clone());
1907 if !eval.inner_converged {
1908 return Err(
1909 "exact two-block spatial inner solve did not converge".to_string(),
1910 );
1911 }
1912 if matches!(eval_mode, EvalMode::ValueGradientHessian)
1913 && !eval.outer_hessian.is_analytic()
1914 {
1915 return Err(
1916 "exact two-block spatial objective requires a full joint [rho, psi] hessian"
1917 .to_string(),
1918 );
1919 }
1920 Ok((eval.objective, eval.gradient, eval.outer_hessian))
1921 },
1922 |theta, specs: &[TermCollectionSpec], designs: &[TermCollectionDesign]| {
1923 if !analytic_joint_derivatives_available {
1924 return Err(
1925 "analytic spatial psi derivatives are unavailable for this exact two-block path"
1926 .to_string(),
1927 );
1928 }
1929 let rho = theta.slice(s![..joint_setup.rho_dim()]).to_owned();
1930 let blocks = builder.build_blocks(
1931 &rho,
1932 &designs[0],
1933 &designs[1],
1934 mean_beta_hint_cell.borrow().clone(),
1935 noise_beta_hint_cell.borrow().clone(),
1936 )?;
1937 if mean_beta_hint_cell.borrow().is_none()
1938 && let Some(beta) = blocks.first().and_then(|block| block.initial_beta.clone())
1939 {
1940 *mean_beta_hint_cell.borrow_mut() = Some(beta);
1941 }
1942 if noise_beta_hint_cell.borrow().is_none()
1943 && let Some(beta) = blocks.get(1).and_then(|block| block.initial_beta.clone())
1944 {
1945 *noise_beta_hint_cell.borrow_mut() = Some(beta);
1946 }
1947 let family = builder.build_family(&designs[0], &designs[1]);
1948 let psiderivative_blocks = builder.build_psiderivative_blocks(
1949 data,
1950 &specs[0],
1951 &specs[1],
1952 &designs[0],
1953 &designs[1],
1954 )?;
1955 let warm_start = hyper_warm_start_cell.borrow().clone();
1956 let eval = evaluate_custom_family_joint_hyper_efs(
1957 &family,
1958 &blocks,
1959 options,
1960 &rho,
1961 &psiderivative_blocks,
1962 warm_start.as_ref(),
1963 )?;
1964 *hyper_warm_start_cell.borrow_mut() = Some(eval.warm_start.clone());
1965 if !eval.inner_converged {
1966 return Err(
1967 "exact two-block spatial EFS inner solve did not converge".to_string(),
1968 );
1969 }
1970 Ok(eval.efs_eval)
1971 },
1972 |_beta: &Array1<f64>| Ok(gam_solve::rho_optimizer::SeedOutcome::NoSlot),
1973 )
1974 }};
1975 }
1976
1977 let mut solved = run_exact_joint_spatial!()
1978 .map_err(|err| format!("exact two-block spatial optimization failed: {err}"))?;
1979
1980 let expected_noise_penalty_count = builder.noise_penalty_count(&solved.designs[1]);
1981 let actual_noise_penalty_count = solved.designs[1].penalties.len();
1982 if expected_noise_penalty_count > actual_noise_penalty_count {
1983 if expected_noise_penalty_count != actual_noise_penalty_count + 1 {
1984 return Err(GamlssError::UnsupportedConfiguration {
1985 reason: format!(
1986 "location-scale result noise design expected {} penalties after augmentation, got {} before augmentation",
1987 expected_noise_penalty_count, actual_noise_penalty_count
1988 ),
1989 }
1990 .into());
1991 }
1992 append_binomial_log_sigma_shrinkage_penalty_design(&mut solved.designs[1]);
1993 }
1994
1995 BlockwiseTermFitResult::try_from_parts(BlockwiseTermFitResultParts {
1996 fit: solved.fit,
1997 meanspec_resolved: solved.resolved_specs.remove(0),
1998 noisespec_resolved: solved.resolved_specs.remove(0),
1999 mean_design: solved.designs.remove(0),
2000 noise_design: solved.designs.remove(0),
2001 })
2002}
2003
2004pub(crate) struct GaussianLocationScaleTermBuilder {
2005 pub(crate) y: Array1<f64>,
2006 pub(crate) weights: Array1<f64>,
2007 pub(crate) meanspec: TermCollectionSpec,
2008 pub(crate) noisespec: TermCollectionSpec,
2009 pub(crate) mean_offset: Array1<f64>,
2010 pub(crate) noise_offset: Array1<f64>,
2011}
2012
2013impl LocationScaleFamilyBuilder for GaussianLocationScaleTermBuilder {
2014 type Family = GaussianLocationScaleFamily;
2015
2016 fn meanspec(&self) -> &TermCollectionSpec {
2017 &self.meanspec
2018 }
2019
2020 fn noisespec(&self) -> &TermCollectionSpec {
2021 &self.noisespec
2022 }
2023
2024 fn exact_spatial_joint_supported(&self) -> bool {
2025 true
2026 }
2027
2028 fn exact_spatial_seed_risk_profile(&self) -> crate::seeding::SeedRiskProfile {
2029 crate::seeding::SeedRiskProfile::GaussianLocationScale
2030 }
2031
2032 fn build_blocks(
2033 &self,
2034 theta: &Array1<f64>,
2035 mean_design: &TermCollectionDesign,
2036 noise_design: &TermCollectionDesign,
2037 mean_beta_hint: Option<Array1<f64>>,
2038 noise_beta_hint: Option<Array1<f64>>,
2039 ) -> Result<Vec<ParameterBlockSpec>, String> {
2040 let layout = GamlssLambdaLayout::two_block(
2041 mean_design.penalties.len(),
2042 self.noise_penalty_count(noise_design),
2043 );
2044 layout.validate_theta_len(theta.len(), "gaussian location-scale")?;
2045 let (meanspec, noisespec) = build_gaussian_mean_and_scale_blocks(
2046 &self.y,
2047 &self.weights,
2048 mean_design,
2049 noise_design,
2050 &self.mean_offset,
2051 &self.noise_offset,
2052 layout.mean_from(theta),
2053 layout.noise_from(theta),
2054 mean_beta_hint,
2055 noise_beta_hint,
2056 "GaussianLocationScale::build_blocks",
2057 )?;
2058 Ok(vec![meanspec, noisespec])
2059 }
2060
2061 fn build_family(
2062 &self,
2063 mean_design: &TermCollectionDesign,
2064 noise_design: &TermCollectionDesign,
2065 ) -> Self::Family {
2066 let preparednoise_design =
2067 prepared_gaussian_log_sigma_design(&mean_design.design, &noise_design.design)
2068 .expect("prepared Gaussian log-sigma design should match block construction");
2069 GaussianLocationScaleFamily {
2070 y: self.y.clone(),
2071 weights: self.weights.clone(),
2072 mu_design: Some(mean_design.design.clone()),
2073 log_sigma_design: Some(preparednoise_design),
2074 policy: gam_runtime::resource::ResourcePolicy::default_library(),
2075 cached_row_scalars: std::sync::RwLock::new(None),
2076 }
2077 }
2078
2079 fn extract_primary_betas(
2080 &self,
2081 fit: &UnifiedFitResult,
2082 ) -> Result<(Array1<f64>, Array1<f64>), String> {
2083 let mean_beta = fit
2084 .block_states
2085 .get(GaussianLocationScaleFamily::BLOCK_MU)
2086 .ok_or_else(|| "missing Gaussian mu block state".to_string())?
2087 .beta
2088 .clone();
2089 let noise_beta = fit
2090 .block_states
2091 .get(GaussianLocationScaleFamily::BLOCK_LOG_SIGMA)
2092 .ok_or_else(|| "missing Gaussian log_sigma block state".to_string())?
2093 .beta
2094 .clone();
2095 Ok((mean_beta, noise_beta))
2096 }
2097
2098 fn build_psiderivative_blocks(
2099 &self,
2100 data: ndarray::ArrayView2<'_, f64>,
2101 meanspec_resolved: &TermCollectionSpec,
2102 noisespec_resolved: &TermCollectionSpec,
2103 mean_design: &TermCollectionDesign,
2104 noise_design: &TermCollectionDesign,
2105 ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String> {
2106 let mean_derivs =
2107 build_block_spatial_psi_derivatives(data, meanspec_resolved, mean_design)?
2108 .ok_or_else(|| "missing Gaussian mean spatial psi derivatives".to_string())?;
2109 let noise_derivs =
2110 build_block_spatial_psi_derivatives(data, noisespec_resolved, noise_design)?
2111 .ok_or_else(|| "missing Gaussian log-sigma spatial psi derivatives".to_string())?;
2112 Ok(vec![mean_derivs, noise_derivs])
2113 }
2114}
2115
2116pub(crate) struct GaussianLocationScaleWiggleTermBuilder {
2117 pub(crate) y: Array1<f64>,
2118 pub(crate) weights: Array1<f64>,
2119 pub(crate) meanspec: TermCollectionSpec,
2120 pub(crate) noisespec: TermCollectionSpec,
2121 pub(crate) mean_offset: Array1<f64>,
2122 pub(crate) noise_offset: Array1<f64>,
2123 pub(crate) wiggle_knots: Array1<f64>,
2124 pub(crate) wiggle_degree: usize,
2125 pub(crate) wiggle_block: ParameterBlockInput,
2126}
2127
2128impl LocationScaleFamilyBuilder for GaussianLocationScaleWiggleTermBuilder {
2129 type Family = GaussianLocationScaleWiggleFamily;
2130
2131 fn meanspec(&self) -> &TermCollectionSpec {
2132 &self.meanspec
2133 }
2134
2135 fn noisespec(&self) -> &TermCollectionSpec {
2136 &self.noisespec
2137 }
2138
2139 fn exact_spatial_joint_supported(&self) -> bool {
2140 true
2141 }
2142
2143 fn exact_spatial_seed_risk_profile(&self) -> crate::seeding::SeedRiskProfile {
2144 crate::seeding::SeedRiskProfile::GaussianLocationScale
2145 }
2146
2147 fn require_exact_spatial_joint(&self) -> bool {
2148 true
2149 }
2150
2151 fn extra_rho0(&self) -> Result<Array1<f64>, String> {
2152 initial_log_lambdas_orzeros(&self.wiggle_block)
2153 }
2154
2155 fn build_blocks(
2156 &self,
2157 theta: &Array1<f64>,
2158 mean_design: &TermCollectionDesign,
2159 noise_design: &TermCollectionDesign,
2160 mean_beta_hint: Option<Array1<f64>>,
2161 noise_beta_hint: Option<Array1<f64>>,
2162 ) -> Result<Vec<ParameterBlockSpec>, String> {
2163 let layout = GamlssLambdaLayout::withwiggle(
2164 mean_design.penalties.len(),
2165 self.noise_penalty_count(noise_design),
2166 self.wiggle_block.penalties.len(),
2167 );
2168 layout.validate_theta_len(theta.len(), "gaussian location-scale wiggle")?;
2169 let (mut meanspec, mut noisespec) = build_gaussian_mean_and_scale_blocks(
2170 &self.y,
2171 &self.weights,
2172 mean_design,
2173 noise_design,
2174 &self.mean_offset,
2175 &self.noise_offset,
2176 layout.mean_from(theta),
2177 layout.noise_from(theta),
2178 mean_beta_hint,
2179 noise_beta_hint,
2180 "GaussianLocationScaleWiggle::build_blocks",
2181 )?;
2182 meanspec.gauge_priority = LINK_WIGGLE_GAUGE_PRIORITY;
2188 noisespec.gauge_priority = LINK_WIGGLE_GAUGE_PRIORITY;
2189 let n_rows = meanspec.design.nrows();
2190 let wigglespec = build_location_scale_wiggle_block(
2191 "wiggle",
2192 self.wiggle_block.design.clone(),
2193 self.wiggle_block.offset.clone(),
2194 wiggle_block_penalty_matrices(&self.wiggle_block),
2195 self.wiggle_block.nullspace_dims.clone(),
2196 layout.wiggle_from(theta),
2197 self.wiggle_block.initial_beta.clone(),
2198 n_rows,
2199 )?;
2200 Ok(vec![meanspec, noisespec, wigglespec])
2201 }
2202
2203 fn build_family(
2204 &self,
2205 mean_design: &TermCollectionDesign,
2206 noise_design: &TermCollectionDesign,
2207 ) -> Self::Family {
2208 let preparednoise_design =
2209 prepared_gaussian_log_sigma_design(&mean_design.design, &noise_design.design).expect(
2210 "prepared Gaussian log-sigma design should match wiggle block construction",
2211 );
2212 GaussianLocationScaleWiggleFamily {
2213 y: self.y.clone(),
2214 weights: self.weights.clone(),
2215 mu_design: Some(mean_design.design.clone()),
2216 log_sigma_design: Some(preparednoise_design),
2217 wiggle_knots: self.wiggle_knots.clone(),
2218 wiggle_degree: self.wiggle_degree,
2219 policy: gam_runtime::resource::ResourcePolicy::default_library(),
2220 cached_row_scalars: std::sync::RwLock::new(None),
2221 }
2222 }
2223
2224 fn extract_primary_betas(
2225 &self,
2226 fit: &UnifiedFitResult,
2227 ) -> Result<(Array1<f64>, Array1<f64>), String> {
2228 let mean_beta = fit
2229 .block_states
2230 .get(GaussianLocationScaleWiggleFamily::BLOCK_MU)
2231 .ok_or_else(|| "missing Gaussian wiggle mu block state".to_string())?
2232 .beta
2233 .clone();
2234 let noise_beta = fit
2235 .block_states
2236 .get(GaussianLocationScaleWiggleFamily::BLOCK_LOG_SIGMA)
2237 .ok_or_else(|| "missing Gaussian wiggle log_sigma block state".to_string())?
2238 .beta
2239 .clone();
2240 Ok((mean_beta, noise_beta))
2241 }
2242
2243 fn build_psiderivative_blocks(
2244 &self,
2245 data: ndarray::ArrayView2<'_, f64>,
2246 meanspec_resolved: &TermCollectionSpec,
2247 noisespec_resolved: &TermCollectionSpec,
2248 mean_design: &TermCollectionDesign,
2249 noise_design: &TermCollectionDesign,
2250 ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String> {
2251 let mean_derivs =
2252 build_block_spatial_psi_derivatives(data, meanspec_resolved, mean_design)?.ok_or_else(
2253 || "missing Gaussian wiggle mean spatial psi derivatives".to_string(),
2254 )?;
2255 let noise_derivs =
2256 build_block_spatial_psi_derivatives(data, noisespec_resolved, noise_design)?
2257 .ok_or_else(|| {
2258 "missing Gaussian wiggle log-sigma spatial psi derivatives".to_string()
2259 })?;
2260 Ok(vec![mean_derivs, noise_derivs, Vec::new()])
2261 }
2262}
2263
2264pub(crate) struct BinomialLocationScaleTermBuilder {
2265 pub(crate) y: Array1<f64>,
2266 pub(crate) weights: Array1<f64>,
2267 pub(crate) link_kind: InverseLink,
2268 pub(crate) meanspec: TermCollectionSpec,
2269 pub(crate) noisespec: TermCollectionSpec,
2270 pub(crate) mean_offset: Array1<f64>,
2271 pub(crate) noise_offset: Array1<f64>,
2272}
2273
2274impl LocationScaleFamilyBuilder for BinomialLocationScaleTermBuilder {
2275 type Family = BinomialLocationScaleFamily;
2276
2277 fn meanspec(&self) -> &TermCollectionSpec {
2278 &self.meanspec
2279 }
2280
2281 fn noisespec(&self) -> &TermCollectionSpec {
2282 &self.noisespec
2283 }
2284
2285 fn exact_spatial_joint_supported(&self) -> bool {
2286 true
2287 }
2288
2289 fn require_exact_spatial_joint(&self) -> bool {
2290 true
2291 }
2292
2293 fn noise_penalty_count(&self, noise_design: &TermCollectionDesign) -> usize {
2294 noise_design.penalties.len() + 1
2295 }
2296
2297 fn build_blocks(
2298 &self,
2299 theta: &Array1<f64>,
2300 mean_design: &TermCollectionDesign,
2301 noise_design: &TermCollectionDesign,
2302 mean_beta_hint: Option<Array1<f64>>,
2303 noise_beta_hint: Option<Array1<f64>>,
2304 ) -> Result<Vec<ParameterBlockSpec>, String> {
2305 let layout = GamlssLambdaLayout::two_block(
2306 mean_design.penalties.len(),
2307 self.noise_penalty_count(noise_design),
2308 );
2309 layout.validate_theta_len(theta.len(), "binomial location-scale")?;
2310 let (thresholdspec, log_sigmaspec) = build_binomial_threshold_and_scale_blocks(
2311 &self.y,
2312 &self.weights,
2313 &self.link_kind,
2314 mean_design,
2315 noise_design,
2316 &self.mean_offset,
2317 &self.noise_offset,
2318 layout.mean_from(theta),
2319 layout.noise_from(theta),
2320 mean_beta_hint,
2321 noise_beta_hint,
2322 "BinomialLocationScale::build_blocks",
2323 )?;
2324 Ok(vec![thresholdspec, log_sigmaspec])
2325 }
2326
2327 fn build_family(
2328 &self,
2329 mean_design: &TermCollectionDesign,
2330 noise_design: &TermCollectionDesign,
2331 ) -> Self::Family {
2332 let identifiednoise_design =
2333 identified_binomial_log_sigma_design(mean_design, noise_design, &self.weights)
2334 .expect("identified binomial log-sigma design");
2335 BinomialLocationScaleFamily {
2336 y: self.y.clone(),
2337 weights: self.weights.clone(),
2338 link_kind: self.link_kind.clone(),
2339 threshold_design: Some(mean_design.design.clone()),
2340 log_sigma_design: Some(identifiednoise_design),
2341 policy: gam_runtime::resource::ResourcePolicy::default_library(),
2342 }
2343 }
2344
2345 fn extract_primary_betas(
2346 &self,
2347 fit: &UnifiedFitResult,
2348 ) -> Result<(Array1<f64>, Array1<f64>), String> {
2349 let mean_beta = fit
2350 .block_states
2351 .get(BinomialLocationScaleFamily::BLOCK_T)
2352 .ok_or_else(|| "missing Binomial threshold block state".to_string())?
2353 .beta
2354 .clone();
2355 let noise_beta = fit
2356 .block_states
2357 .get(BinomialLocationScaleFamily::BLOCK_LOG_SIGMA)
2358 .ok_or_else(|| "missing Binomial log_sigma block state".to_string())?
2359 .beta
2360 .clone();
2361 Ok((mean_beta, noise_beta))
2362 }
2363
2364 fn build_psiderivative_blocks(
2365 &self,
2366 data: ndarray::ArrayView2<'_, f64>,
2367 meanspec_resolved: &TermCollectionSpec,
2368 noisespec_resolved: &TermCollectionSpec,
2369 mean_design: &TermCollectionDesign,
2370 noise_design: &TermCollectionDesign,
2371 ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String> {
2372 let mean_derivs =
2373 build_block_spatial_psi_derivatives(data, meanspec_resolved, mean_design)?
2374 .ok_or_else(|| "missing threshold spatial psi derivatives".to_string())?;
2375 let noise_derivs =
2376 build_block_spatial_psi_derivatives(data, noisespec_resolved, noise_design)?
2377 .ok_or_else(|| "missing log_sigma spatial psi derivatives".to_string())?;
2378 Ok(vec![mean_derivs, noise_derivs])
2379 }
2380}
2381
2382pub(crate) struct BinomialLocationScaleWiggleTermBuilder {
2383 pub(crate) y: Array1<f64>,
2384 pub(crate) weights: Array1<f64>,
2385 pub(crate) link_kind: InverseLink,
2386 pub(crate) meanspec: TermCollectionSpec,
2387 pub(crate) noisespec: TermCollectionSpec,
2388 pub(crate) mean_offset: Array1<f64>,
2389 pub(crate) noise_offset: Array1<f64>,
2390 pub(crate) wiggle_knots: Array1<f64>,
2391 pub(crate) wiggle_degree: usize,
2392 pub(crate) wiggle_block: ParameterBlockInput,
2393}
2394
2395impl LocationScaleFamilyBuilder for BinomialLocationScaleWiggleTermBuilder {
2396 type Family = BinomialLocationScaleWiggleFamily;
2397
2398 fn meanspec(&self) -> &TermCollectionSpec {
2399 &self.meanspec
2400 }
2401
2402 fn noisespec(&self) -> &TermCollectionSpec {
2403 &self.noisespec
2404 }
2405
2406 fn exact_spatial_joint_supported(&self) -> bool {
2407 true
2408 }
2409
2410 fn require_exact_spatial_joint(&self) -> bool {
2411 true
2412 }
2413
2414 fn extra_rho0(&self) -> Result<Array1<f64>, String> {
2415 initial_log_lambdas_orzeros(&self.wiggle_block)
2416 }
2417
2418 fn noise_penalty_count(&self, noise_design: &TermCollectionDesign) -> usize {
2419 noise_design.penalties.len() + 1
2420 }
2421
2422 fn build_blocks(
2423 &self,
2424 theta: &Array1<f64>,
2425 mean_design: &TermCollectionDesign,
2426 noise_design: &TermCollectionDesign,
2427 mean_beta_hint: Option<Array1<f64>>,
2428 noise_beta_hint: Option<Array1<f64>>,
2429 ) -> Result<Vec<ParameterBlockSpec>, String> {
2430 let layout = GamlssLambdaLayout::withwiggle(
2431 mean_design.penalties.len(),
2432 self.noise_penalty_count(noise_design),
2433 self.wiggle_block.penalties.len(),
2434 );
2435 layout.validate_theta_len(theta.len(), "wiggle location-scale")?;
2436 let (mut thresholdspec, mut log_sigmaspec) = build_binomial_threshold_and_scale_blocks(
2437 &self.y,
2438 &self.weights,
2439 &self.link_kind,
2440 mean_design,
2441 noise_design,
2442 &self.mean_offset,
2443 &self.noise_offset,
2444 layout.mean_from(theta),
2445 layout.noise_from(theta),
2446 mean_beta_hint,
2447 noise_beta_hint,
2448 "BinomialLocationScaleWiggle::build_blocks",
2449 )?;
2450 thresholdspec.gauge_priority = LINK_WIGGLE_GAUGE_PRIORITY;
2462 log_sigmaspec.gauge_priority = LINK_WIGGLE_GAUGE_PRIORITY;
2463 let n_rows = thresholdspec.design.nrows();
2464 let wigglespec = build_location_scale_wiggle_block(
2465 "wiggle",
2466 self.wiggle_block.design.clone(),
2467 self.wiggle_block.offset.clone(),
2468 wiggle_block_penalty_matrices(&self.wiggle_block),
2469 vec![],
2470 layout.wiggle_from(theta),
2471 self.wiggle_block.initial_beta.clone(),
2472 n_rows,
2473 )?;
2474 Ok(vec![thresholdspec, log_sigmaspec, wigglespec])
2475 }
2476
2477 fn build_family(
2478 &self,
2479 mean_design: &TermCollectionDesign,
2480 noise_design: &TermCollectionDesign,
2481 ) -> Self::Family {
2482 let identifiednoise_design =
2483 identified_binomial_log_sigma_design(mean_design, noise_design, &self.weights)
2484 .expect("identified binomial log-sigma design should match block construction");
2485 BinomialLocationScaleWiggleFamily {
2486 y: self.y.clone(),
2487 weights: self.weights.clone(),
2488 link_kind: self.link_kind.clone(),
2489 threshold_design: Some(mean_design.design.clone()),
2490 log_sigma_design: Some(identifiednoise_design),
2491 wiggle_knots: self.wiggle_knots.clone(),
2492 wiggle_degree: self.wiggle_degree,
2493 policy: gam_runtime::resource::ResourcePolicy::default_library(),
2494 }
2495 }
2496
2497 fn extract_primary_betas(
2498 &self,
2499 fit: &UnifiedFitResult,
2500 ) -> Result<(Array1<f64>, Array1<f64>), String> {
2501 let mean_beta = fit
2502 .block_states
2503 .get(BinomialLocationScaleWiggleFamily::BLOCK_T)
2504 .ok_or_else(|| "missing Binomial wiggle threshold block state".to_string())?
2505 .beta
2506 .clone();
2507 let noise_beta = fit
2508 .block_states
2509 .get(BinomialLocationScaleWiggleFamily::BLOCK_LOG_SIGMA)
2510 .ok_or_else(|| "missing Binomial wiggle log_sigma block state".to_string())?
2511 .beta
2512 .clone();
2513 Ok((mean_beta, noise_beta))
2514 }
2515
2516 fn build_psiderivative_blocks(
2517 &self,
2518 data: ndarray::ArrayView2<'_, f64>,
2519 meanspec_resolved: &TermCollectionSpec,
2520 noisespec_resolved: &TermCollectionSpec,
2521 mean_design: &TermCollectionDesign,
2522 noise_design: &TermCollectionDesign,
2523 ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String> {
2524 let mean_derivs =
2525 build_block_spatial_psi_derivatives(data, meanspec_resolved, mean_design)?
2526 .ok_or_else(|| "missing threshold spatial psi derivatives".to_string())?;
2527 let noise_derivs =
2528 build_block_spatial_psi_derivatives(data, noisespec_resolved, noise_design)?
2529 .ok_or_else(|| "missing log_sigma spatial psi derivatives".to_string())?;
2530 Ok(vec![mean_derivs, noise_derivs, Vec::new()])
2539 }
2540}
2541
2542pub(crate) fn fit_gaussian_location_scale_terms(
2543 data: ndarray::ArrayView2<'_, f64>,
2544 spec: GaussianLocationScaleTermSpec,
2545 options: &BlockwiseFitOptions,
2546 kappa_options: &SpatialLengthScaleOptimizationOptions,
2547) -> Result<BlockwiseTermFitResult, String> {
2548 validate_gaussian_location_scale_termspec(data, &spec, "fit_gaussian_location_scale_terms")?;
2549 fit_location_scale_terms(
2550 data,
2551 GaussianLocationScaleTermBuilder {
2552 y: spec.y,
2553 weights: spec.weights,
2554 meanspec: spec.meanspec,
2555 noisespec: spec.log_sigmaspec,
2556 mean_offset: spec.mean_offset,
2557 noise_offset: spec.log_sigma_offset,
2558 },
2559 options,
2560 kappa_options,
2561 )
2562}
2563
2564pub(crate) fn fit_gaussian_location_scalewiggle_terms(
2565 data: ndarray::ArrayView2<'_, f64>,
2566 spec: GaussianLocationScaleWiggleTermSpec,
2567 options: &BlockwiseFitOptions,
2568 kappa_options: &SpatialLengthScaleOptimizationOptions,
2569) -> Result<BlockwiseTermFitResult, String> {
2570 validate_gaussian_location_scalewiggle_termspec(
2571 data,
2572 &spec,
2573 "fit_gaussian_location_scalewiggle_terms",
2574 )?;
2575 fit_location_scale_terms(
2576 data,
2577 GaussianLocationScaleWiggleTermBuilder {
2578 y: spec.y,
2579 weights: spec.weights,
2580 meanspec: spec.meanspec,
2581 noisespec: spec.log_sigmaspec,
2582 mean_offset: spec.mean_offset,
2583 noise_offset: spec.log_sigma_offset,
2584 wiggle_knots: spec.wiggle_knots,
2585 wiggle_degree: spec.wiggle_degree,
2586 wiggle_block: spec.wiggle_block,
2587 },
2588 options,
2589 kappa_options,
2590 )
2591}
2592
2593pub(crate) fn select_gaussian_location_scale_link_wiggle_basis_from_pilot(
2594 pilot: &BlockwiseTermFitResult,
2595 wiggle_cfg: &WiggleBlockConfig,
2596 wiggle_penalty_orders: &[usize],
2597) -> Result<SelectedWiggleBasis, String> {
2598 let q_seed = pilot
2599 .fit
2600 .block_states
2601 .first()
2602 .ok_or_else(|| "pilot Gaussian wiggle fit is missing mean block".to_string())?
2603 .eta
2604 .view();
2605 select_wiggle_basis_from_seed(q_seed, wiggle_cfg, wiggle_penalty_orders)
2606}
2607
2608pub(crate) fn fit_gaussian_location_scale_terms_with_selected_wiggle(
2609 data: ndarray::ArrayView2<'_, f64>,
2610 spec: GaussianLocationScaleTermSpec,
2611 selected_wiggle_basis: SelectedWiggleBasis,
2612 options: &BlockwiseFitOptions,
2613 kappa_options: &SpatialLengthScaleOptimizationOptions,
2614) -> Result<BlockwiseTermWiggleFitResult, String> {
2615 let SelectedWiggleBasis {
2616 knots: wiggle_knots,
2617 degree: wiggle_degree,
2618 block: wiggle_block,
2619 ..
2620 } = selected_wiggle_basis;
2621 let solved = fit_gaussian_location_scalewiggle_terms(
2622 data,
2623 GaussianLocationScaleWiggleTermSpec {
2624 y: spec.y,
2625 weights: spec.weights,
2626 meanspec: spec.meanspec,
2627 log_sigmaspec: spec.log_sigmaspec,
2628 mean_offset: spec.mean_offset,
2629 log_sigma_offset: spec.log_sigma_offset,
2630 wiggle_knots: wiggle_knots.clone(),
2631 wiggle_degree,
2632 wiggle_block,
2633 },
2634 options,
2635 kappa_options,
2636 )?;
2637
2638 BlockwiseTermWiggleFitResult::try_from_parts(BlockwiseTermWiggleFitResultParts {
2639 fit: solved,
2640 wiggle_knots,
2641 wiggle_degree,
2642 })
2643}
2644
2645pub(crate) fn fit_binomial_location_scale_terms(
2646 data: ndarray::ArrayView2<'_, f64>,
2647 spec: BinomialLocationScaleTermSpec,
2648 options: &BlockwiseFitOptions,
2649 kappa_options: &SpatialLengthScaleOptimizationOptions,
2650) -> Result<BlockwiseTermFitResult, String> {
2651 validate_binomial_location_scale_termspec(data, &spec, "fit_binomial_location_scale_terms")?;
2652 fit_location_scale_terms(
2653 data,
2654 BinomialLocationScaleTermBuilder {
2655 y: spec.y,
2656 weights: spec.weights,
2657 link_kind: spec.link_kind,
2658 meanspec: spec.thresholdspec,
2659 noisespec: spec.log_sigmaspec,
2660 mean_offset: spec.threshold_offset,
2661 noise_offset: spec.log_sigma_offset,
2662 },
2663 options,
2664 kappa_options,
2665 )
2666}
2667
2668pub(crate) fn fit_binomial_location_scalewiggle_terms(
2669 data: ndarray::ArrayView2<'_, f64>,
2670 spec: BinomialLocationScaleWiggleTermSpec,
2671 options: &BlockwiseFitOptions,
2672 kappa_options: &SpatialLengthScaleOptimizationOptions,
2673) -> Result<BlockwiseTermFitResult, String> {
2674 validate_binomial_location_scalewiggle_termspec(
2675 data,
2676 &spec,
2677 "fit_binomial_location_scalewiggle_terms",
2678 )?;
2679 fit_location_scale_terms(
2680 data,
2681 BinomialLocationScaleWiggleTermBuilder {
2682 y: spec.y,
2683 weights: spec.weights,
2684 link_kind: spec.link_kind,
2685 meanspec: spec.thresholdspec,
2686 noisespec: spec.log_sigmaspec,
2687 mean_offset: spec.threshold_offset,
2688 noise_offset: spec.log_sigma_offset,
2689 wiggle_knots: spec.wiggle_knots,
2690 wiggle_degree: spec.wiggle_degree,
2691 wiggle_block: spec.wiggle_block,
2692 },
2693 options,
2694 kappa_options,
2695 )
2696}
2697
2698pub(crate) fn select_binomial_location_scale_link_wiggle_basis_from_pilot(
2699 pilot: &BlockwiseTermFitResult,
2700 wiggle_cfg: &WiggleBlockConfig,
2701 wiggle_penalty_orders: &[usize],
2702) -> Result<SelectedWiggleBasis, String> {
2703 let eta_t = pilot
2704 .fit
2705 .block_states
2706 .first()
2707 .ok_or_else(|| "pilot fit is missing threshold block".to_string())?
2708 .eta
2709 .view();
2710 let eta_ls = pilot
2711 .fit
2712 .block_states
2713 .get(1)
2714 .ok_or_else(|| "pilot fit is missing log_sigma block".to_string())?
2715 .eta
2716 .view();
2717 let sigma = eta_ls.mapv(safe_exp);
2718 let q_seed = Array1::from_iter(eta_t.iter().zip(sigma.iter()).map(|(&t, &s)| -t / s));
2719 select_wiggle_basis_from_seed(q_seed.view(), wiggle_cfg, wiggle_penalty_orders)
2720}
2721
2722pub(crate) fn fit_binomial_location_scale_terms_with_selected_wiggle(
2723 data: ndarray::ArrayView2<'_, f64>,
2724 spec: BinomialLocationScaleTermSpec,
2725 selected_wiggle_basis: SelectedWiggleBasis,
2726 options: &BlockwiseFitOptions,
2727 kappa_options: &SpatialLengthScaleOptimizationOptions,
2728) -> Result<BlockwiseTermWiggleFitResult, String> {
2729 let SelectedWiggleBasis {
2730 knots: wiggle_knots,
2731 degree: wiggle_degree,
2732 block: wiggle_block,
2733 ..
2734 } = selected_wiggle_basis;
2735 let solved = fit_binomial_location_scalewiggle_terms(
2736 data,
2737 BinomialLocationScaleWiggleTermSpec {
2738 y: spec.y,
2739 weights: spec.weights,
2740 link_kind: spec.link_kind,
2741 thresholdspec: spec.thresholdspec,
2742 log_sigmaspec: spec.log_sigmaspec,
2743 threshold_offset: spec.threshold_offset,
2744 log_sigma_offset: spec.log_sigma_offset,
2745 wiggle_knots: wiggle_knots.clone(),
2746 wiggle_degree,
2747 wiggle_block,
2748 },
2749 options,
2750 kappa_options,
2751 )?;
2752
2753 BlockwiseTermWiggleFitResult::try_from_parts(BlockwiseTermWiggleFitResultParts {
2754 fit: solved,
2755 wiggle_knots,
2756 wiggle_degree,
2757 })
2758}
2759
2760pub(crate) fn select_binomial_mean_link_wiggle_basis_from_pilot(
2761 pilot_design: &TermCollectionDesign,
2762 pilot_fit: &UnifiedFitResult,
2763 wiggle_cfg: &WiggleBlockConfig,
2764 wiggle_penalty_orders: &[usize],
2765) -> Result<SelectedWiggleBasis, String> {
2766 let q_seed = pilot_design.design.dot(&pilot_fit.beta);
2767 select_wiggle_basis_from_seed(q_seed.view(), wiggle_cfg, wiggle_penalty_orders)
2768}
2769
2770pub(crate) fn fit_binomial_mean_wiggle_terms_with_selected_basis(
2771 data: ndarray::ArrayView2<'_, f64>,
2772 pilot_spec: &TermCollectionSpec,
2773 pilot_design: &TermCollectionDesign,
2774 pilot_fit: &UnifiedFitResult,
2775 y: &Array1<f64>,
2776 weights: &Array1<f64>,
2777 link_kind: InverseLink,
2778 selected_wiggle_basis: SelectedWiggleBasis,
2779 options: &BlockwiseFitOptions,
2780 kappa_options: &SpatialLengthScaleOptimizationOptions,
2781) -> Result<BinomialMeanWiggleTermFitResult, String> {
2782 const RHO_BOUND: f64 = 12.0;
2783
2784 validate_term_weights(
2785 data,
2786 y.len(),
2787 weights,
2788 "fit_binomial_mean_wiggle_terms_with_selected_basis",
2789 )?;
2790 validate_binomial_response(y, "fit_binomial_mean_wiggle_terms_with_selected_basis")?;
2791
2792 let SelectedWiggleBasis {
2798 knots: wiggle_knots,
2799 degree: wiggle_degree,
2800 block: wiggle_block,
2801 ..
2802 } = selected_wiggle_basis;
2803
2804 let spatial_terms = spatial_length_scale_term_indices(pilot_spec);
2805 if spatial_terms.is_empty() {
2806 let (fit, saved_warp_beta, saved_index_shift) = fit_binomial_mean_wiggle(
2807 BinomialMeanWiggleSpec {
2808 y: y.clone(),
2809 weights: weights.clone(),
2810 link_kind,
2811 wiggle_knots: wiggle_knots.clone(),
2812 wiggle_degree,
2813 eta_block: ParameterBlockInput {
2814 design: pilot_design.design.clone(),
2815 offset: Array1::zeros(y.len()),
2816 penalties: pilot_design
2817 .penalties
2818 .iter()
2819 .map(crate::model_types::PenaltySpec::from_blockwise_ref)
2820 .collect(),
2821 nullspace_dims: vec![],
2822 initial_log_lambdas: Some(
2823 pilot_fit
2824 .lambdas
2825 .mapv(|v| v.max(WARMSTART_LOG_LAMBDA_FLOOR).ln()),
2826 ),
2827 initial_beta: Some(pilot_fit.beta.clone()),
2828 },
2829 wiggle_block,
2830 },
2831 options,
2832 )?;
2833 return Ok(BinomialMeanWiggleTermFitResult {
2834 fit,
2835 resolvedspec: pilot_spec.clone(),
2836 design: pilot_design.clone(),
2837 wiggle_knots,
2838 wiggle_degree,
2839 saved_warp_beta,
2840 saved_index_shift,
2841 });
2842 }
2843
2844 let dims_per_term = spatial_dims_per_term(pilot_spec, &spatial_terms);
2845 let log_kappa0 =
2846 SpatialLogKappaCoords::from_length_scales_aniso(pilot_spec, &spatial_terms, kappa_options)
2847 .reseed_from_data(data, pilot_spec, &spatial_terms, kappa_options);
2848 let log_kappa_lower = SpatialLogKappaCoords::lower_bounds_aniso_from_data(
2849 data,
2850 pilot_spec,
2851 &spatial_terms,
2852 &dims_per_term,
2853 kappa_options,
2854 );
2855 let log_kappa_upper = SpatialLogKappaCoords::upper_bounds_aniso_from_data(
2856 data,
2857 pilot_spec,
2858 &spatial_terms,
2859 &dims_per_term,
2860 kappa_options,
2861 );
2862 let log_kappa0 = log_kappa0.clamp_to_bounds(&log_kappa_lower, &log_kappa_upper);
2864
2865 let eta_penalty_count = pilot_design.penalties.len();
2866 let wiggle_penalty_count = initial_log_lambdas_orzeros(&wiggle_block)?.len();
2867 let rho_dim = eta_penalty_count + wiggle_penalty_count;
2868 let baseline_resolvedspec = log_kappa0
2869 .apply_tospec(pilot_spec, &spatial_terms)
2870 .map_err(|e| e.to_string())?;
2871 let baseline_design =
2872 build_term_collection_design(data, &baseline_resolvedspec).map_err(|e| e.to_string())?;
2873 let baseline_fit = fit_binomial_mean_wiggle(
2874 BinomialMeanWiggleSpec {
2875 y: y.clone(),
2876 weights: weights.clone(),
2877 link_kind: link_kind.clone(),
2878 wiggle_knots: wiggle_knots.clone(),
2879 wiggle_degree,
2880 eta_block: ParameterBlockInput {
2881 design: baseline_design.design.clone(),
2882 offset: Array1::zeros(y.len()),
2883 penalties: baseline_design
2884 .penalties
2885 .iter()
2886 .map(crate::model_types::PenaltySpec::from_blockwise_ref)
2887 .collect(),
2888 nullspace_dims: vec![],
2889 initial_log_lambdas: Some(
2890 pilot_fit
2891 .lambdas
2892 .mapv(|v| v.max(WARMSTART_LOG_LAMBDA_FLOOR).ln()),
2893 ),
2894 initial_beta: Some(pilot_fit.beta.clone()),
2895 },
2896 wiggle_block: wiggle_block.clone(),
2897 },
2898 options,
2899 )?
2900 .0;
2901 let baseline_log_lambdas = baseline_fit
2902 .lambdas
2903 .mapv(|v| v.max(WARMSTART_LOG_LAMBDA_FLOOR).ln());
2904 if baseline_log_lambdas.len() != rho_dim {
2905 return Err(GamlssError::DimensionMismatch {
2906 reason: format!(
2907 "baseline binomial mean-wiggle fit returned {} log-lambdas, expected {rho_dim}",
2908 baseline_log_lambdas.len()
2909 ),
2910 }
2911 .into());
2912 }
2913 let baseline_eta_beta = baseline_fit
2914 .block_states
2915 .get(BinomialMeanWiggleFamily::BLOCK_ETA)
2916 .ok_or_else(|| "baseline binomial mean-wiggle fit missing eta block".to_string())?
2917 .beta
2918 .clone();
2919 let baseline_wiggle_beta = Some(
2920 baseline_fit
2921 .block_states
2922 .get(BinomialMeanWiggleFamily::BLOCK_WIGGLE)
2923 .ok_or_else(|| "baseline binomial mean-wiggle fit missing wiggle block".to_string())?
2924 .beta
2925 .clone(),
2926 );
2927 let theta_dim = rho_dim + log_kappa0.len();
2928 let mut theta0 = Array1::<f64>::zeros(theta_dim);
2929 theta0
2930 .slice_mut(s![0..rho_dim])
2931 .assign(&baseline_log_lambdas);
2932 theta0
2933 .slice_mut(s![rho_dim..theta_dim])
2934 .assign(log_kappa0.as_array());
2935
2936 let mut lower = Array1::<f64>::from_elem(theta_dim, -RHO_BOUND);
2937 let mut upper = Array1::<f64>::from_elem(theta_dim, RHO_BOUND);
2938 lower
2939 .slice_mut(s![rho_dim..theta_dim])
2940 .assign(log_kappa_lower.as_array());
2941 upper
2942 .slice_mut(s![rho_dim..theta_dim])
2943 .assign(log_kappa_upper.as_array());
2944
2945 let pilot_spec_cloned = pilot_spec.clone();
2946 let pilot_beta = baseline_eta_beta;
2947 let wiggle_design = wiggle_block.design.clone();
2948 let wiggle_offset = wiggle_block.offset.clone();
2949 let wiggle_penalties = wiggle_block.penalties.clone();
2950 let wiggle_initial_beta = baseline_wiggle_beta;
2951 let wiggle_knots_cloned = wiggle_knots.clone();
2952 let y_cloned = y.clone();
2953 let weights_cloned = weights.clone();
2954 let link_kind_cloned = link_kind.clone();
2955 let outer_family = BinomialMeanWiggleFamily {
2956 y: y_cloned.clone(),
2957 weights: weights_cloned.clone(),
2958 link_kind: link_kind_cloned.clone(),
2959 wiggle_knots: wiggle_knots_cloned.clone(),
2960 wiggle_degree,
2961 policy: gam_runtime::resource::ResourcePolicy::default_library(),
2962 frozen_warp_design: None,
2965 };
2966 let screening_cap = Arc::new(AtomicUsize::new(0));
2967 let mut outer_options = options.clone();
2968 outer_options.screening_max_inner_iterations = Some(Arc::clone(&screening_cap));
2969 struct MeanWiggleOuterState {
2970 pub(crate) warm_cache: Option<crate::custom_family::CustomFamilyWarmStart>,
2971 pub(crate) last_eval: Option<(
2972 Array1<f64>,
2973 f64,
2974 Array1<f64>,
2975 gam_problem::HessianValue,
2976 crate::custom_family::CustomFamilyWarmStart,
2977 )>,
2978 }
2979
2980 let build_realized_blocks = |theta: &Array1<f64>| -> Result<
2981 (
2982 TermCollectionSpec,
2983 TermCollectionDesign,
2984 Vec<ParameterBlockSpec>,
2985 Vec<CustomFamilyBlockPsiDerivative>,
2986 ),
2987 String,
2988 > {
2989 let log_kappa =
2990 SpatialLogKappaCoords::from_theta_tail_with_dims(theta, rho_dim, dims_per_term.clone());
2991 let resolvedspec = log_kappa
2992 .apply_tospec(&pilot_spec_cloned, &spatial_terms)
2993 .map_err(|e| e.to_string())?;
2994 let design =
2995 build_term_collection_design(data, &resolvedspec).map_err(|e| e.to_string())?;
2996 let eta_derivs = build_block_spatial_psi_derivatives(data, &resolvedspec, &design)?
2997 .ok_or_else(|| {
2998 "missing eta spatial psi derivatives for binomial mean wiggle".to_string()
2999 })?;
3000 let blocks = vec![
3001 ParameterBlockSpec {
3002 name: "eta".to_string(),
3003 design: design.design.clone(),
3004 offset: Array1::zeros(y_cloned.len()),
3005 penalties: design.penalties_as_penalty_matrix(),
3006 nullspace_dims: vec![],
3007 initial_log_lambdas: theta.slice(s![0..eta_penalty_count]).to_owned(),
3008 initial_beta: Some(pilot_beta.clone()),
3009 gauge_priority: LINK_WIGGLE_GAUGE_PRIORITY,
3013 jacobian_callback: None,
3014 stacked_design: None,
3015 stacked_offset: None,
3016 },
3017 ParameterBlockSpec {
3018 name: "wiggle".to_string(),
3019 design: wiggle_design.clone(),
3020 offset: wiggle_offset.clone(),
3021 penalties: {
3022 let p_wiggle = wiggle_design.ncols();
3023 wiggle_penalties
3024 .iter()
3025 .map(|spec| match spec {
3026 crate::model_types::PenaltySpec::Block {
3027 local, col_range, ..
3028 } => PenaltyMatrix::Blockwise {
3029 local: local.clone(),
3030 col_range: col_range.clone(),
3031 total_dim: p_wiggle,
3032 },
3033 crate::model_types::PenaltySpec::Dense(m)
3034 | crate::model_types::PenaltySpec::DenseWithMean {
3035 matrix: m, ..
3036 } => PenaltyMatrix::Dense(m.clone()),
3037 })
3038 .collect()
3039 },
3040 nullspace_dims: vec![],
3041 initial_log_lambdas: theta.slice(s![eta_penalty_count..rho_dim]).to_owned(),
3042 initial_beta: wiggle_initial_beta.clone(),
3043 gauge_priority: DEFAULT_GAUGE_PRIORITY,
3044 jacobian_callback: None,
3045 stacked_design: None,
3046 stacked_offset: None,
3047 },
3048 ];
3049 Ok((resolvedspec, design, blocks, eta_derivs))
3050 };
3051
3052 let build_eval = |theta: &Array1<f64>,
3053 warm_cache: Option<&crate::custom_family::CustomFamilyWarmStart>,
3054 need_hessian: bool|
3055 -> Result<
3056 (
3057 crate::custom_family::CustomFamilyJointHyperResult,
3058 TermCollectionSpec,
3059 TermCollectionDesign,
3060 ),
3061 String,
3062 > {
3063 let (resolvedspec, design, blocks, eta_derivs) = build_realized_blocks(theta)?;
3064 let eval = evaluate_custom_family_joint_hyper(
3065 &outer_family,
3066 &blocks,
3067 &outer_options,
3068 &theta.slice(s![0..rho_dim]).to_owned(),
3069 &[eta_derivs, Vec::new()],
3070 warm_cache,
3071 if need_hessian {
3072 gam_problem::EvalMode::ValueGradientHessian
3073 } else {
3074 gam_problem::EvalMode::ValueAndGradient
3075 },
3076 )?;
3077 Ok((eval, resolvedspec, design))
3078 };
3079
3080 let build_efs = |theta: &Array1<f64>,
3081 warm_cache: Option<&crate::custom_family::CustomFamilyWarmStart>|
3082 -> Result<crate::custom_family::CustomFamilyJointHyperEfsResult, String> {
3083 let (_, _, blocks, eta_derivs) = build_realized_blocks(theta)?;
3084 evaluate_custom_family_joint_hyper_efs(
3085 &outer_family,
3086 &blocks,
3087 &outer_options,
3088 &theta.slice(s![0..rho_dim]).to_owned(),
3089 &[eta_derivs, Vec::new()],
3090 warm_cache,
3091 )
3092 .map_err(|e| e.to_string())
3093 };
3094
3095 use crate::model_types::EstimationError;
3096 use gam_problem::{DeclaredHessianForm, Derivative, OuterEval};
3097 use gam_solve::rho_optimizer::OuterEvalOrder;
3098
3099 let analytic_outer_hessian_available = true;
3110 let mut seed_heuristic = theta0.to_vec();
3111 for value in &mut seed_heuristic[..rho_dim] {
3112 *value = value.exp();
3113 }
3114 let problem = gam_solve::rho_optimizer::OuterProblem::new(theta_dim)
3115 .with_gradient(Derivative::Analytic)
3116 .with_hessian(if analytic_outer_hessian_available {
3117 DeclaredHessianForm::Either
3118 } else {
3119 DeclaredHessianForm::Unavailable
3120 })
3121 .with_psi_dim(theta_dim - rho_dim)
3122 .with_tolerance(options.outer_tol)
3123 .with_max_iter(options.outer_max_iter)
3124 .with_bounds(lower.clone(), upper.clone())
3125 .with_initial_rho(theta0.clone())
3126 .with_seed_config(crate::seeding::SeedConfig {
3127 max_seeds: 4,
3128 seed_budget: 2,
3129 risk_profile: crate::seeding::SeedRiskProfile::GeneralizedLinear,
3130 num_auxiliary_trailing: theta_dim - rho_dim,
3131 ..Default::default()
3132 })
3133 .with_screening_cap(Arc::clone(&screening_cap))
3134 .with_rho_bound(12.0)
3135 .with_heuristic_lambdas(seed_heuristic);
3136
3137 let eval_outer = |state: &mut MeanWiggleOuterState,
3138 theta: &Array1<f64>,
3139 order: OuterEvalOrder|
3140 -> Result<OuterEval, EstimationError> {
3141 if let Some((cached_theta, cached_cost, cached_grad, cached_hess, cached_warm)) =
3142 &state.last_eval
3143 && cached_theta == theta
3144 && (!matches!(order, OuterEvalOrder::ValueGradientHessian)
3145 || matches!(
3146 cached_hess,
3147 gam_problem::HessianValue::Dense(_) | gam_problem::HessianValue::Operator(_)
3148 ))
3149 {
3150 state.warm_cache = Some(cached_warm.clone());
3151 return Ok(OuterEval {
3152 cost: *cached_cost,
3153 gradient: cached_grad.clone(),
3154 hessian: cached_hess.clone(),
3155 inner_beta_hint: None,
3156 });
3157 }
3158 let need_hessian = matches!(order, OuterEvalOrder::ValueGradientHessian)
3159 && analytic_outer_hessian_available;
3160 let (eval, _, _) = build_eval(theta, state.warm_cache.as_ref(), need_hessian)
3161 .map_err(EstimationError::InvalidInput)?;
3162 if !eval.inner_converged {
3163 state.warm_cache = Some(eval.warm_start);
3164 crate::bail_invalid_estim!(
3165 "binomial mean-wiggle exact spatial inner solve did not converge"
3166 );
3167 }
3168 let hessian_result = eval.outer_hessian.clone();
3169 state.last_eval = Some((
3170 theta.clone(),
3171 eval.objective,
3172 eval.gradient.clone(),
3173 eval.outer_hessian.clone(),
3174 eval.warm_start.clone(),
3175 ));
3176 state.warm_cache = Some(eval.warm_start);
3177 Ok(OuterEval {
3178 cost: eval.objective,
3179 gradient: eval.gradient,
3180 hessian: hessian_result,
3181 inner_beta_hint: None,
3182 })
3183 };
3184
3185 let mut obj = problem.build_objective_with_screening_proxy(
3186 MeanWiggleOuterState {
3187 warm_cache: None,
3188 last_eval: None,
3189 },
3190 |state: &mut MeanWiggleOuterState, theta: &Array1<f64>| {
3191 if let Some((cached_theta, cached_cost, _, _, cached_warm)) = &state.last_eval
3192 && cached_theta == theta
3193 {
3194 state.warm_cache = Some(cached_warm.clone());
3195 return Ok(*cached_cost);
3196 }
3197 let (eval, _, _) = build_eval(theta, state.warm_cache.as_ref(), false)
3198 .map_err(EstimationError::InvalidInput)?;
3199 if !eval.inner_converged {
3200 state.warm_cache = Some(eval.warm_start);
3201 crate::bail_invalid_estim!(
3202 "binomial mean-wiggle exact spatial cost inner solve did not converge"
3203 .to_string(),
3204 );
3205 }
3206 state.warm_cache = Some(eval.warm_start);
3207 Ok(eval.objective)
3208 },
3209 |state: &mut MeanWiggleOuterState, theta: &Array1<f64>| {
3210 eval_outer(
3211 state,
3212 theta,
3213 if analytic_outer_hessian_available {
3214 OuterEvalOrder::ValueGradientHessian
3215 } else {
3216 OuterEvalOrder::ValueAndGradient
3217 },
3218 )
3219 },
3220 |state: &mut MeanWiggleOuterState, theta: &Array1<f64>, order: OuterEvalOrder| {
3221 eval_outer(state, theta, order)
3222 },
3223 Some(|state: &mut MeanWiggleOuterState| {
3224 state.warm_cache = None;
3225 state.last_eval = None;
3226 }),
3227 Some(|state: &mut MeanWiggleOuterState, theta: &Array1<f64>| {
3228 let eval = build_efs(theta, state.warm_cache.as_ref())
3229 .map_err(EstimationError::InvalidInput)?;
3230 if !eval.inner_converged {
3231 state.warm_cache = Some(eval.warm_start);
3232 crate::bail_invalid_estim!(
3233 "binomial mean-wiggle exact spatial EFS inner solve did not converge"
3234 .to_string(),
3235 );
3236 }
3237 state.warm_cache = Some(eval.warm_start);
3238 Ok(eval.efs_eval)
3239 }),
3240 |state: &mut MeanWiggleOuterState, theta: &Array1<f64>| {
3250 if let Some((cached_theta, cached_cost, _, _, cached_warm)) = &state.last_eval
3251 && cached_theta == theta
3252 {
3253 state.warm_cache = Some(cached_warm.clone());
3254 return Ok(*cached_cost);
3255 }
3256 let (eval, _, _) = build_eval(theta, state.warm_cache.as_ref(), false)
3257 .map_err(EstimationError::InvalidInput)?;
3258 state.warm_cache = Some(eval.warm_start);
3259 Ok(eval.objective)
3260 },
3261 );
3262
3263 let outer = problem
3264 .run(&mut obj, "binomial mean wiggle exact spatial hyper")
3265 .map_err(|e| e.to_string())?;
3266 if !outer.converged {
3267 return Err(GamlssError::NumericalFailure { reason: format!(
3268 "binomial mean wiggle exact spatial hyper did not converge after {} iterations (final_objective={:.6e}, final_grad_norm={})",
3269 outer.iterations,
3270 outer.final_value,
3271 outer.final_grad_norm_report(),
3272 ) }.into());
3273 }
3274 let theta_star = outer.rho;
3275
3276 let log_kappa =
3277 SpatialLogKappaCoords::from_theta_tail_with_dims(&theta_star, rho_dim, dims_per_term);
3278 let resolvedspec = log_kappa
3279 .apply_tospec(&pilot_spec_cloned, &spatial_terms)
3280 .map_err(|e| e.to_string())?;
3281 let design = build_term_collection_design(data, &resolvedspec).map_err(|e| e.to_string())?;
3282 let resolvedspec =
3283 freeze_term_collection_from_design(&resolvedspec, &design).map_err(|e| e.to_string())?;
3284 let fit = fit_binomial_mean_wiggle(
3285 BinomialMeanWiggleSpec {
3286 y: y_cloned,
3287 weights: weights_cloned,
3288 link_kind: link_kind_cloned,
3289 wiggle_knots: wiggle_knots.clone(),
3290 wiggle_degree,
3291 eta_block: ParameterBlockInput {
3292 design: design.design.clone(),
3293 offset: Array1::zeros(y.len()),
3294 penalties: design
3295 .penalties
3296 .iter()
3297 .map(crate::model_types::PenaltySpec::from_blockwise_ref)
3298 .collect(),
3299 nullspace_dims: vec![],
3300 initial_log_lambdas: Some(theta_star.slice(s![0..eta_penalty_count]).to_owned()),
3301 initial_beta: Some(pilot_beta),
3302 },
3303 wiggle_block: ParameterBlockInput {
3304 design: wiggle_design,
3305 offset: wiggle_offset,
3306 penalties: wiggle_penalties,
3307 nullspace_dims: vec![],
3308 initial_log_lambdas: Some(
3309 theta_star.slice(s![eta_penalty_count..rho_dim]).to_owned(),
3310 ),
3311 initial_beta: wiggle_initial_beta,
3312 },
3313 },
3314 options,
3315 )?;
3316 let (fit, saved_warp_beta, saved_index_shift) = fit;
3317
3318 Ok(BinomialMeanWiggleTermFitResult {
3319 fit,
3320 resolvedspec,
3321 design,
3322 wiggle_knots,
3323 wiggle_degree,
3324 saved_warp_beta,
3325 saved_index_shift,
3326 })
3327}