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 gam_problem::validate_log_strengths(lambdas.iter().copied())
151 .map_err(|error| format!("initial_log_lambdas: {error}"))?;
152 Ok(lambdas)
153}
154
155fn fitted_log_lambdas(lambdas: &Array1<f64>, context: &str) -> Result<Array1<f64>, String> {
156 lambdas
157 .iter()
158 .copied()
159 .enumerate()
160 .map(|(coordinate, value)| {
161 gam_problem::checked_log_strength(value)
162 .map_err(|error| format!("{context} coordinate {coordinate}: {error}"))
163 })
164 .collect::<Result<Vec<_>, _>>()
165 .map(Array1::from_vec)
166}
167
168pub(crate) fn build_two_block_exact_joint_setup(
169 data: ArrayView2<'_, f64>,
170 meanspec: &TermCollectionSpec,
171 noisespec: &TermCollectionSpec,
172 mean_penalties: usize,
173 noise_penalties: usize,
174 extra_rho0: &[f64],
175 rho0_override: Option<&Array1<f64>>,
176 kappa_options: &SpatialLengthScaleOptimizationOptions,
177) -> Result<ExactJointHyperSetup, gam_terms::basis::BasisError> {
178 let rho_dim = mean_penalties + noise_penalties + extra_rho0.len();
181 let mut rho0vec = Array1::<f64>::zeros(rho_dim);
182 if let Some(rho0) = rho0_override.filter(|rho0| rho0.len() == rho_dim) {
183 rho0vec.assign(rho0);
184 } else {
185 for (i, &rho_init) in extra_rho0.iter().enumerate() {
186 rho0vec[mean_penalties + noise_penalties + i] = rho_init;
187 }
188 }
189
190 build_location_scale_exact_joint_setup(data, &[meanspec, noisespec], rho0vec, kappa_options)
193}
194
195pub(crate) fn gaussian_location_scalewarm_start(
196 y: &Array1<f64>,
197 weights: &Array1<f64>,
198 mu_block: &ParameterBlockSpec,
199 log_sigma_block: &ParameterBlockSpec,
200 ridge_floor: f64,
201 mean_beta_hint: Option<&Array1<f64>>,
202 noise_beta_hint: Option<&Array1<f64>>,
203) -> Result<(Array1<f64>, Array1<f64>, f64), String> {
204 let betamu = if let Some(beta) = mean_beta_hint {
205 beta.clone()
206 } else {
207 solve_penalizedweighted_projection(
208 &mu_block.design,
209 &mu_block.offset,
210 y,
211 weights,
212 &mu_block.penalties,
213 &mu_block.initial_log_lambdas,
214 ridge_floor,
215 )?
216 };
217 let mut mu_hat = mu_block.solver_design().matrixvectormultiply(&betamu);
218 mu_hat += mu_block.solver_offset();
219 let mut weighted_ss = 0.0;
220 let mut weight_sum = 0.0;
221 for i in 0..y.len() {
222 let wi = weights[i].max(0.0);
223 let resid = y[i] - mu_hat[i];
224 weighted_ss += wi * resid * resid;
225 weight_sum += wi;
226 }
227 if !weighted_ss.is_finite() || !weight_sum.is_finite() || weight_sum <= 0.0 {
228 return Err(
229 "gaussian location-scale warm start could not estimate residual scale".to_string(),
230 );
231 }
232 let sigma_hat = (weighted_ss / weight_sum)
237 .sqrt()
238 .max(LOGB_SIGMA_FLOOR * 1.5);
239 let beta_log_sigma = if let Some(beta) = noise_beta_hint {
240 beta.clone()
241 } else {
242 let eta_sigma = (sigma_hat - LOGB_SIGMA_FLOOR).ln();
243 let sigma_target = Array1::from_elem(y.len(), eta_sigma);
244 solve_penalizedweighted_projection(
245 &log_sigma_block.design,
246 &log_sigma_block.offset,
247 &sigma_target,
248 weights,
249 &log_sigma_block.penalties,
250 &log_sigma_block.initial_log_lambdas,
251 ridge_floor,
252 )?
253 };
254 Ok((betamu, beta_log_sigma, sigma_hat))
255}
256
257pub(crate) const LOCATION_SCALE_N_OUTPUTS: usize = 2;
261
262pub(crate) fn build_location_scale_block(
276 name: impl Into<String>,
277 design: DesignMatrix,
278 offset: Array1<f64>,
279 penalties: Vec<PenaltyMatrix>,
280 nullspace_dims: Vec<usize>,
281 initial_log_lambdas: Array1<f64>,
282 initial_beta: Option<Array1<f64>>,
283 own_output: usize,
284 n_family_outputs: usize,
285 caller: &str,
286) -> Result<ParameterBlockSpec, String> {
287 if own_output >= n_family_outputs {
288 return Err(format!(
289 "{caller}: own_output={own_output} >= n_family_outputs={n_family_outputs}"
290 ));
291 }
292 let mut spec = ParameterBlockSpec {
293 name: name.into(),
294 design,
295 offset,
296 penalties,
297 nullspace_dims,
298 initial_log_lambdas,
299 initial_beta,
300 gauge_priority: 100,
301 jacobian_callback: None,
302 stacked_design: None,
303 stacked_offset: None,
304 };
305 let dense = spec.effective_design(caller)?;
306 spec.jacobian_callback = Some(std::sync::Arc::new(AdditiveBlockJacobian {
307 design: dense,
308 own_output,
309 n_family_outputs,
310 }));
311 Ok(spec)
312}
313
314pub(crate) fn build_location_scale_wiggle_block(
320 name: impl Into<String>,
321 design: DesignMatrix,
322 offset: Array1<f64>,
323 penalties: Vec<PenaltyMatrix>,
324 nullspace_dims: Vec<usize>,
325 initial_log_lambdas: Array1<f64>,
326 initial_beta: Option<Array1<f64>>,
327 n_rows: usize,
328) -> Result<ParameterBlockSpec, String> {
329 let p_w = design.ncols();
330 let mut spec = ParameterBlockSpec {
331 name: name.into(),
332 design,
333 offset,
334 penalties,
335 nullspace_dims,
336 initial_log_lambdas,
337 initial_beta,
338 gauge_priority: 100,
339 jacobian_callback: None,
340 stacked_design: None,
341 stacked_offset: None,
342 };
343 spec.jacobian_callback = Some(std::sync::Arc::new(AdditiveBlockJacobian {
344 design: ndarray::Array2::<f64>::zeros((n_rows, p_w)),
345 own_output: 0,
346 n_family_outputs: LOCATION_SCALE_N_OUTPUTS,
347 }));
348 Ok(spec)
349}
350
351pub(crate) fn prepared_gaussian_log_sigma_design(
352 mu_design: &DesignMatrix,
353 log_sigma_design: &DesignMatrix,
354) -> Result<DesignMatrix, String> {
355 if mu_design.nrows() != log_sigma_design.nrows() {
356 return Err(GamlssError::DimensionMismatch {
357 reason: format!(
358 "gaussian log-sigma design row mismatch: mean rows={}, log_sigma rows={}",
359 mu_design.nrows(),
360 log_sigma_design.nrows()
361 ),
362 }
363 .into());
364 }
365 Ok(log_sigma_design.clone())
376}
377
378pub(crate) fn identified_binomial_log_sigma_design(
379 threshold_design: &TermCollectionDesign,
380 log_sigma_design: &TermCollectionDesign,
381 weights: &Array1<f64>,
382) -> Result<DesignMatrix, String> {
383 let non_intercept_start = log_sigma_design
384 .intercept_range
385 .end
386 .min(log_sigma_design.design.ncols());
387 let transform = build_scale_deviation_transform_design(
388 &threshold_design.design,
389 &log_sigma_design.design,
390 weights,
391 non_intercept_start,
392 )?;
393 build_scale_deviation_operator(
394 threshold_design.design.clone(),
395 log_sigma_design.design.clone(),
396 &transform,
397 )
398}
399
400pub(crate) fn identity_penalty(dim: usize) -> Array2<f64> {
401 let mut penalty = Array2::<f64>::zeros((dim, dim));
402 for i in 0..dim {
403 penalty[[i, i]] = 1.0;
404 }
405 penalty
406}
407
408pub(crate) fn append_binomial_log_sigma_shrinkage_penalty_design(
409 design: &mut TermCollectionDesign,
410) {
411 let p = design.design.ncols();
412 design
413 .penalties
414 .push(BlockwisePenalty::new(0..p, identity_penalty(p)));
415 design.nullspace_dims.push(0);
417 design.penaltyinfo.push(PenaltyBlockInfo {
418 global_index: design.penaltyinfo.len(),
419 termname: Some("log_sigma_shrinkage".to_string()),
420 penalty: ActivePenaltyInfo {
421 source: PenaltySource::Other("shrinkage".to_string()),
422 original_index: 0,
423 effective_rank: p,
424 normalization_scale: 1.0,
425 kronecker_factors: None,
426 },
427 });
428}
429
430pub(crate) fn build_gaussian_mean_and_scale_blocks(
437 y: &Array1<f64>,
438 weights: &Array1<f64>,
439 mean_design: &TermCollectionDesign,
440 noise_design: &TermCollectionDesign,
441 mean_offset: &Array1<f64>,
442 noise_offset: &Array1<f64>,
443 mean_log_lambdas: Array1<f64>,
444 noise_log_lambdas: Array1<f64>,
445 mean_beta_hint: Option<Array1<f64>>,
446 noise_beta_hint: Option<Array1<f64>>,
447 context: &str,
448) -> Result<(ParameterBlockSpec, ParameterBlockSpec), String> {
449 let mean_offset = mean_design
450 .compose_offset(mean_offset.view(), &format!("{context}: mu"))
451 .map_err(|error| error.to_string())?;
452 let noise_offset = noise_design
453 .compose_offset(noise_offset.view(), &format!("{context}: log_sigma"))
454 .map_err(|error| error.to_string())?;
455 let mut meanspec = build_location_scale_block(
456 "mu",
457 mean_design.design.clone(),
458 mean_offset,
459 mean_design.penalties_as_penalty_matrix(),
460 mean_design.nullspace_dims.clone(),
461 mean_log_lambdas,
462 mean_beta_hint,
463 0,
464 LOCATION_SCALE_N_OUTPUTS,
465 &format!("{context}: mu"),
466 )?;
467 let prepared_noise_design =
468 prepared_gaussian_log_sigma_design(&mean_design.design, &noise_design.design)?;
469 let mut noisespec = build_location_scale_block(
477 "log_sigma",
478 prepared_noise_design,
479 noise_offset,
480 noise_design.penalties_as_penalty_matrix(),
481 noise_design.nullspace_dims.clone(),
482 noise_log_lambdas,
483 noise_beta_hint,
484 1,
485 LOCATION_SCALE_N_OUTPUTS,
486 &format!("{context}: log_sigma"),
487 )?;
488 if meanspec.initial_beta.is_none() || noisespec.initial_beta.is_none() {
489 let (betamu0, beta_ls0, _) = gaussian_location_scalewarm_start(
490 y,
491 weights,
492 &meanspec,
493 &noisespec,
494 1e-10,
495 meanspec.initial_beta.as_ref(),
496 noisespec.initial_beta.as_ref(),
497 )?;
498 if meanspec.initial_beta.is_none() {
499 meanspec.initial_beta = Some(betamu0);
500 }
501 if noisespec.initial_beta.is_none() {
502 noisespec.initial_beta = Some(beta_ls0);
503 }
504 }
505 Ok((meanspec, noisespec))
506}
507
508pub(crate) fn build_binomial_threshold_and_scale_blocks(
514 y: &Array1<f64>,
515 weights: &Array1<f64>,
516 link_kind: &InverseLink,
517 mean_design: &TermCollectionDesign,
518 noise_design: &TermCollectionDesign,
519 mean_offset: &Array1<f64>,
520 noise_offset: &Array1<f64>,
521 mean_log_lambdas: Array1<f64>,
522 noise_log_lambdas: Array1<f64>,
523 mean_beta_hint: Option<Array1<f64>>,
524 noise_beta_hint: Option<Array1<f64>>,
525 context: &str,
526) -> Result<(ParameterBlockSpec, ParameterBlockSpec), String> {
527 let mean_offset = mean_design
528 .compose_offset(mean_offset.view(), &format!("{context}: threshold"))
529 .map_err(|error| error.to_string())?;
530 let noise_offset = noise_design
531 .compose_offset(noise_offset.view(), &format!("{context}: log_sigma"))
532 .map_err(|error| error.to_string())?;
533 let identifiednoise_design =
534 identified_binomial_log_sigma_design(mean_design, noise_design, weights)?;
535 let p_noise = identifiednoise_design.ncols();
536 let mut log_sigma_penalty_matrices: Vec<PenaltyMatrix> =
537 noise_design.penalties_as_penalty_matrix();
538 log_sigma_penalty_matrices.push(PenaltyMatrix::Dense(identity_penalty(p_noise)));
539 let mut thresholdspec = build_location_scale_block(
540 "threshold",
541 mean_design.design.clone(),
542 mean_offset,
543 mean_design.penalties_as_penalty_matrix(),
544 vec![],
545 mean_log_lambdas,
546 mean_beta_hint,
547 0,
548 LOCATION_SCALE_N_OUTPUTS,
549 &format!("{context}: threshold"),
550 )?;
551 let mut log_sigmaspec = build_location_scale_block(
552 "log_sigma",
553 identifiednoise_design,
554 noise_offset,
555 log_sigma_penalty_matrices,
556 vec![],
557 noise_log_lambdas,
558 noise_beta_hint,
559 1,
560 LOCATION_SCALE_N_OUTPUTS,
561 &format!("{context}: log_sigma"),
562 )?;
563 if thresholdspec.initial_beta.is_none() || log_sigmaspec.initial_beta.is_none() {
564 let (beta_t0, beta_ls0) = binomial_location_scalewarm_start(
565 y,
566 weights,
567 link_kind,
568 &thresholdspec,
569 &log_sigmaspec,
570 thresholdspec.initial_beta.as_ref(),
571 log_sigmaspec.initial_beta.as_ref(),
572 )?;
573 if thresholdspec.initial_beta.is_none() {
574 thresholdspec.initial_beta = Some(beta_t0);
575 }
576 if log_sigmaspec.initial_beta.is_none() {
577 log_sigmaspec.initial_beta = Some(beta_ls0);
578 }
579 }
580 Ok((thresholdspec, log_sigmaspec))
581}
582
583pub(crate) fn wiggle_block_penalty_matrices(
587 wiggle_block: &ParameterBlockInput,
588) -> Vec<PenaltyMatrix> {
589 let p_wiggle = wiggle_block.design.ncols();
590 wiggle_block
591 .penalties
592 .iter()
593 .map(|spec| match spec {
594 crate::model_types::PenaltySpec::Block {
595 local, col_range, ..
596 } => PenaltyMatrix::Blockwise {
597 local: local.clone(),
598 col_range: col_range.clone(),
599 total_dim: p_wiggle,
600 },
601 crate::model_types::PenaltySpec::Dense(m)
602 | crate::model_types::PenaltySpec::DenseWithMean { matrix: m, .. } => {
603 PenaltyMatrix::Dense(m.clone())
604 }
605 })
606 .collect()
607}
608
609pub(crate) fn binomial_location_scale_link_eta_from_probability(
610 link_kind: &InverseLink,
611 probability: f64,
612) -> Result<f64, String> {
613 let target = probability.clamp(1e-6, 1.0 - 1e-6);
614 match link_kind {
615 InverseLink::Standard(StandardLink::Logit) => Ok((target / (1.0 - target)).ln()),
616 InverseLink::Standard(StandardLink::Probit) => standard_normal_quantile(target)
617 .map_err(|err| format!("failed to invert probit warm-start probability: {err}")),
618 InverseLink::Standard(StandardLink::CLogLog) => Ok((-((1.0 - target).ln())).ln()),
619 other => Err(GamlssError::UnsupportedConfiguration { reason: format!(
620 "binomial location-scale warm start requires logit, probit, or cloglog link, got {other:?}"
621 ) }.into()),
622 }
623}
624
625pub(crate) fn weighted_binomial_prevalence(
626 y: &Array1<f64>,
627 weights: &Array1<f64>,
628) -> Result<f64, String> {
629 if y.len() != weights.len() {
630 return Err(GamlssError::DimensionMismatch { reason: format!(
631 "binomial location-scale warm start dimension mismatch: y has length {}, weights have length {}",
632 y.len(),
633 weights.len()
634 ) }.into());
635 }
636 let mut weight_sum = 0.0;
637 let mut success_sum = 0.0;
638 for (&yi, &wi) in y.iter().zip(weights.iter()) {
639 if !yi.is_finite() {
640 return Err(GamlssError::NonFinite {
641 reason: format!(
642 "binomial location-scale warm start encountered non-finite response {yi}"
643 ),
644 }
645 .into());
646 }
647 if !wi.is_finite() || wi < 0.0 {
648 return Err(GamlssError::InvalidInput {
649 reason: format!(
650 "binomial location-scale warm start requires finite non-negative weights; weight={wi}"
651 ),
652 }
653 .into());
654 }
655 if wi > 0.0 {
656 weight_sum += wi;
657 success_sum += wi * yi;
658 }
659 }
660 if !weight_sum.is_finite() || weight_sum <= 0.0 {
661 return Err(
662 "binomial location-scale warm start requires positive total weight".to_string(),
663 );
664 }
665 Ok(success_sum / weight_sum)
666}
667
668pub(crate) fn project_constant_eta_into_block(
669 block: &ParameterBlockSpec,
670 weights: &Array1<f64>,
671 eta: f64,
672) -> Result<Array1<f64>, String> {
673 let target_eta = Array1::from_elem(block.design.nrows(), eta);
674 solve_penalizedweighted_projection(
675 &block.design,
676 &block.offset,
677 &target_eta,
678 weights,
679 &block.penalties,
680 &block.initial_log_lambdas,
681 1e-10,
682 )
683}
684
685pub(crate) fn binomial_location_scalewarm_start(
689 y: &Array1<f64>,
690 weights: &Array1<f64>,
691 link_kind: &InverseLink,
692 threshold_block: &ParameterBlockSpec,
693 log_sigma_block: &ParameterBlockSpec,
694 mean_beta_hint: Option<&Array1<f64>>,
695 noise_beta_hint: Option<&Array1<f64>>,
696) -> Result<(Array1<f64>, Array1<f64>), String> {
697 if let (Some(mean_beta), Some(noise_beta)) = (mean_beta_hint, noise_beta_hint) {
698 return Ok((mean_beta.clone(), noise_beta.clone()));
699 }
700
701 let beta_threshold = match mean_beta_hint {
702 Some(beta) => beta.clone(),
703 None => {
704 let prevalence = weighted_binomial_prevalence(y, weights)?;
705 let eta = binomial_location_scale_link_eta_from_probability(link_kind, prevalence)?;
706 project_constant_eta_into_block(threshold_block, weights, eta)?
707 }
708 };
709 let beta_log_sigma = match noise_beta_hint {
710 Some(beta) => beta.clone(),
711 None => project_constant_eta_into_block(log_sigma_block, weights, 0.0)?,
712 };
713 Ok((beta_threshold, beta_log_sigma))
714}
715
716#[derive(Clone)]
717pub(crate) struct BinomialMeanWiggleSpec {
718 pub y: Array1<f64>,
719 pub weights: Array1<f64>,
720 pub link_kind: InverseLink,
721 pub wiggle_knots: Array1<f64>,
722 pub wiggle_degree: usize,
723 pub eta_block: ParameterBlockInput,
724 pub wiggle_block: ParameterBlockInput,
725}
726
727#[derive(Clone)]
728pub struct GaussianLocationScaleTermSpec {
729 pub y: Array1<f64>,
730 pub weights: Array1<f64>,
731 pub meanspec: TermCollectionSpec,
732 pub log_sigmaspec: TermCollectionSpec,
733 pub mean_offset: Array1<f64>,
734 pub log_sigma_offset: Array1<f64>,
735}
736
737#[derive(Clone)]
738pub struct GaussianLocationScaleWiggleTermSpec {
739 pub y: Array1<f64>,
740 pub weights: Array1<f64>,
741 pub meanspec: TermCollectionSpec,
742 pub log_sigmaspec: TermCollectionSpec,
743 pub mean_offset: Array1<f64>,
744 pub log_sigma_offset: Array1<f64>,
745 pub wiggle_knots: Array1<f64>,
746 pub wiggle_degree: usize,
747 pub wiggle_block: ParameterBlockInput,
748}
749
750#[derive(Clone)]
751pub struct BinomialLocationScaleTermSpec {
752 pub y: Array1<f64>,
753 pub weights: Array1<f64>,
754 pub link_kind: InverseLink,
755 pub thresholdspec: TermCollectionSpec,
756 pub log_sigmaspec: TermCollectionSpec,
757 pub threshold_offset: Array1<f64>,
758 pub log_sigma_offset: Array1<f64>,
759}
760
761#[derive(Clone)]
762pub struct BinomialLocationScaleWiggleTermSpec {
763 pub y: Array1<f64>,
764 pub weights: Array1<f64>,
765 pub link_kind: InverseLink,
766 pub thresholdspec: TermCollectionSpec,
767 pub log_sigmaspec: TermCollectionSpec,
768 pub threshold_offset: Array1<f64>,
769 pub log_sigma_offset: Array1<f64>,
770 pub wiggle_knots: Array1<f64>,
771 pub wiggle_degree: usize,
772 pub wiggle_block: ParameterBlockInput,
773}
774
775#[derive(Clone, Debug)]
776pub struct BlockwiseTermFitResult {
777 pub fit: UnifiedFitResult,
778 pub meanspec_resolved: TermCollectionSpec,
779 pub noisespec_resolved: TermCollectionSpec,
780 pub mean_design: TermCollectionDesign,
781 pub noise_design: TermCollectionDesign,
782}
783
784pub(crate) struct BlockwiseTermFitResultParts {
785 pub fit: UnifiedFitResult,
786 pub meanspec_resolved: TermCollectionSpec,
787 pub noisespec_resolved: TermCollectionSpec,
788 pub mean_design: TermCollectionDesign,
789 pub noise_design: TermCollectionDesign,
790}
791
792pub struct BlockwiseTermWiggleFitResult {
793 pub fit: BlockwiseTermFitResult,
794 pub wiggle_knots: Array1<f64>,
795 pub wiggle_degree: usize,
796}
797
798pub struct BinomialMeanWiggleTermFitResult {
799 pub fit: UnifiedFitResult,
800 pub resolvedspec: TermCollectionSpec,
801 pub design: TermCollectionDesign,
802 pub wiggle_knots: Array1<f64>,
803 pub wiggle_degree: usize,
804 pub saved_warp_beta: Option<Vec<f64>>,
809 pub saved_index_shift: Option<Vec<f64>>,
814}
815
816pub(crate) struct BlockwiseTermWiggleFitResultParts {
817 pub fit: BlockwiseTermFitResult,
818 pub wiggle_knots: Array1<f64>,
819 pub wiggle_degree: usize,
820}
821
822pub(crate) fn validate_term_collection_design(
823 label: &str,
824 design: &TermCollectionDesign,
825) -> Result<(), String> {
826 let p = design.design.ncols();
827 let n = design.design.nrows();
828 for rows in exact_design_row_chunks(n, p) {
829 let chunk = design
830 .design
831 .try_row_chunk(rows)
832 .map_err(|e| format!("{label}.design row chunk materialization failed: {e}"))?;
833 validate_all_finite_estimation(&format!("{label}.design"), chunk.iter().copied())
834 .map_err(|e| e.to_string())?;
835 }
836 if design.nullspace_dims.len() != design.penalties.len() {
837 return Err(GamlssError::DimensionMismatch {
838 reason: format!(
839 "{label}.nullspace_dims length mismatch: got {}, expected {}",
840 design.nullspace_dims.len(),
841 design.penalties.len()
842 ),
843 }
844 .into());
845 }
846 if design.penaltyinfo.len() != design.penalties.len() {
847 return Err(GamlssError::DimensionMismatch {
848 reason: format!(
849 "{label}.penaltyinfo length mismatch: got {}, expected {}",
850 design.penaltyinfo.len(),
851 design.penalties.len()
852 ),
853 }
854 .into());
855 }
856 for (idx, bp) in design.penalties.iter().enumerate() {
857 validate_all_finite_estimation(
858 &format!("{label}.penalties[{idx}]"),
859 bp.local.iter().copied(),
860 )
861 .map_err(|e| e.to_string())?;
862 if bp.col_range.end > p {
863 return Err(GamlssError::DimensionMismatch {
864 reason: format!(
865 "{label}.penalties[{idx}] col_range {}..{} exceeds design width {}",
866 bp.col_range.start, bp.col_range.end, p
867 ),
868 }
869 .into());
870 }
871 }
872 if let Some(bounds) = design.coefficient_lower_bounds.as_ref() {
873 if bounds.len() != p {
874 return Err(GamlssError::ConstraintViolation {
875 reason: format!(
876 "{label}.coefficient_lower_bounds length mismatch: got {}, expected {p}",
877 bounds.len()
878 ),
879 }
880 .into());
881 }
882 for (idx, &bound) in bounds.iter().enumerate() {
883 if !(bound.is_finite() || bound == f64::NEG_INFINITY) {
884 return Err(GamlssError::NonFinite { reason: format!(
885 "{label}.coefficient_lower_bounds[{idx}] must be finite or -inf, got {bound}",
886 ) }.into());
887 }
888 }
889 }
890 if let Some(constraints) = design.linear_constraints.as_ref() {
891 validate_all_finite_estimation(
892 &format!("{label}.linear_constraints.a"),
893 constraints.a.iter().copied(),
894 )
895 .map_err(|e| e.to_string())?;
896 validate_all_finite_estimation(
897 &format!("{label}.linear_constraints.b"),
898 constraints.b.iter().copied(),
899 )
900 .map_err(|e| e.to_string())?;
901 if constraints.a.ncols() != p {
902 return Err(GamlssError::DimensionMismatch {
903 reason: format!(
904 "{label}.linear_constraints.a column mismatch: got {}, expected {p}",
905 constraints.a.ncols()
906 ),
907 }
908 .into());
909 }
910 if constraints.a.nrows() != constraints.b.len() {
911 return Err(GamlssError::DimensionMismatch {
912 reason: format!(
913 "{label}.linear_constraints row mismatch: a has {}, b has {}",
914 constraints.a.nrows(),
915 constraints.b.len()
916 ),
917 }
918 .into());
919 }
920 }
921 if design.intercept_range.start > design.intercept_range.end || design.intercept_range.end > p {
922 return Err(GamlssError::ConstraintViolation {
923 reason: format!(
924 "{label}.intercept_range out of bounds: {:?} for {} columns",
925 design.intercept_range, p
926 ),
927 }
928 .into());
929 }
930 Ok(())
931}
932
933impl BlockwiseTermFitResult {
934 pub(crate) fn try_from_parts(parts: BlockwiseTermFitResultParts) -> Result<Self, String> {
935 let BlockwiseTermFitResultParts {
936 fit,
937 meanspec_resolved,
938 noisespec_resolved,
939 mean_design,
940 noise_design,
941 } = parts;
942
943 fit.validate_numeric_finiteness()
944 .map_err(|e| format!("{e}"))?;
945 if fit.block_states.len() < 2 {
946 return Err(GamlssError::DimensionMismatch {
947 reason: format!(
948 "BlockwiseTermFitResult requires at least 2 block states, got {}",
949 fit.block_states.len()
950 ),
951 }
952 .into());
953 }
954 validate_term_collection_design("blockwise_term.mean_design", &mean_design)?;
955 validate_term_collection_design("blockwise_term.noise_design", &noise_design)?;
956 if mean_design.design.nrows() != noise_design.design.nrows() {
957 return Err(GamlssError::DimensionMismatch {
958 reason: format!(
959 "BlockwiseTermFitResult row mismatch: mean_design={}, noise_design={}",
960 mean_design.design.nrows(),
961 noise_design.design.nrows()
962 ),
963 }
964 .into());
965 }
966 if fit.block_states[0].beta.len() != mean_design.design.ncols() {
967 return Err(GamlssError::DimensionMismatch {
968 reason: format!(
969 "BlockwiseTermFitResult mean beta length mismatch: got {}, expected {}",
970 fit.block_states[0].beta.len(),
971 mean_design.design.ncols()
972 ),
973 }
974 .into());
975 }
976 if fit.block_states[1].beta.len() != noise_design.design.ncols() {
977 return Err(GamlssError::DimensionMismatch {
978 reason: format!(
979 "BlockwiseTermFitResult noise beta length mismatch: got {}, expected {}",
980 fit.block_states[1].beta.len(),
981 noise_design.design.ncols()
982 ),
983 }
984 .into());
985 }
986 if fit.block_states[0].eta.len() != mean_design.design.nrows() {
987 return Err(GamlssError::DimensionMismatch {
988 reason: format!(
989 "BlockwiseTermFitResult mean eta length mismatch: got {}, expected {}",
990 fit.block_states[0].eta.len(),
991 mean_design.design.nrows()
992 ),
993 }
994 .into());
995 }
996 if fit.block_states[1].eta.len() != noise_design.design.nrows() {
997 return Err(GamlssError::DimensionMismatch {
998 reason: format!(
999 "BlockwiseTermFitResult noise eta length mismatch: got {}, expected {}",
1000 fit.block_states[1].eta.len(),
1001 noise_design.design.nrows()
1002 ),
1003 }
1004 .into());
1005 }
1006
1007 Ok(Self {
1008 fit,
1009 meanspec_resolved,
1010 noisespec_resolved,
1011 mean_design,
1012 noise_design,
1013 })
1014 }
1015
1016 pub(crate) fn validate_numeric_finiteness(&self) -> Result<(), String> {
1017 Self::try_from_parts(BlockwiseTermFitResultParts {
1018 fit: self.fit.clone(),
1019 meanspec_resolved: self.meanspec_resolved.clone(),
1020 noisespec_resolved: self.noisespec_resolved.clone(),
1021 mean_design: self.mean_design.clone(),
1022 noise_design: self.noise_design.clone(),
1023 })
1024 .map(|_| ())
1025 }
1026}
1027
1028impl BlockwiseTermWiggleFitResult {
1029 pub(crate) fn try_from_parts(parts: BlockwiseTermWiggleFitResultParts) -> Result<Self, String> {
1030 let BlockwiseTermWiggleFitResultParts {
1031 fit,
1032 wiggle_knots,
1033 wiggle_degree,
1034 } = parts;
1035
1036 fit.validate_numeric_finiteness()
1037 .map_err(|e| e.to_string())?;
1038 if fit.fit.block_states.len() < 3 {
1039 return Err(GamlssError::DimensionMismatch {
1040 reason: format!(
1041 "BlockwiseTermWiggleFitResult requires at least 3 block states, got {}",
1042 fit.fit.block_states.len()
1043 ),
1044 }
1045 .into());
1046 }
1047 if wiggle_knots.is_empty() {
1048 return Err(GamlssError::UnsupportedConfiguration {
1049 reason: "BlockwiseTermWiggleFitResult requires non-empty wiggle_knots".to_string(),
1050 }
1051 .into());
1052 }
1053 validate_all_finite_estimation(
1054 "blockwise_term_wiggle.wiggle_knots",
1055 wiggle_knots.iter().copied(),
1056 )
1057 .map_err(|e| e.to_string())?;
1058
1059 Ok(Self {
1060 fit,
1061 wiggle_knots,
1062 wiggle_degree,
1063 })
1064 }
1065}
1066
1067pub struct BinomialLocationScaleFitResult {
1068 pub fit: BlockwiseTermFitResult,
1069 pub wiggle_knots: Option<Array1<f64>>,
1070 pub wiggle_degree: Option<usize>,
1071 pub beta_link_wiggle: Option<Vec<f64>>,
1072}
1073
1074pub struct GaussianLocationScaleFitResult {
1075 pub fit: BlockwiseTermFitResult,
1076 pub wiggle_knots: Option<Array1<f64>>,
1077 pub wiggle_degree: Option<usize>,
1078 pub beta_link_wiggle: Option<Vec<f64>>,
1079 pub response_scale: f64,
1108}
1109
1110fn binomial_mean_wiggle_saved_frame_gauge(
1127 alias: &Array2<f64>,
1128 mean_width: usize,
1129 wiggle_width: usize,
1130) -> Result<gam_problem::Gauge, String> {
1131 if alias.dim() != (mean_width, wiggle_width) {
1132 return Err(format!(
1133 "binomial mean-wiggle de-alias map is {}x{}, expected {mean_width}x{wiggle_width}",
1134 alias.nrows(),
1135 alias.ncols(),
1136 ));
1137 }
1138 let total_width = mean_width
1139 .checked_add(wiggle_width)
1140 .ok_or_else(|| "binomial mean-wiggle coefficient dimension overflows usize".to_string())?;
1141 let mut transform = Array2::<f64>::eye(total_width);
1142 for row in 0..mean_width {
1143 for column in 0..wiggle_width {
1144 transform[[row, mean_width + column]] = -alias[[row, column]];
1145 }
1146 }
1147 let gauge = gam_problem::Gauge::from_t(
1148 transform,
1149 &[mean_width, wiggle_width],
1150 &[mean_width, wiggle_width],
1151 );
1152 gauge.validate().map_err(|reason| {
1153 format!("binomial mean-wiggle saved coefficient gauge is invalid: {reason}")
1154 })?;
1155 Ok(gauge)
1156}
1157
1158fn binomial_mean_wiggle_saved_geometry(
1159 geometry: &gam_solve::model_types::FitGeometry,
1160 saved_frame: &gam_problem::Gauge,
1161) -> Result<gam_solve::model_types::FitGeometry, String> {
1162 let mut saved_geometry = geometry.clone();
1163 saved_geometry.coefficient_gauge = geometry
1164 .coefficient_gauge
1165 .left_compose(saved_frame)
1166 .map_err(|reason| {
1167 format!(
1168 "binomial mean-wiggle active geometry cannot compose with its exact saved-result gauge: {reason}"
1169 )
1170 })?;
1171 Ok(saved_geometry)
1172}
1173
1174fn binomial_mean_wiggle_saved_covariance(
1175 covariance: &Array2<f64>,
1176 saved_frame: &gam_problem::Gauge,
1177 label: &str,
1178) -> Result<Array2<f64>, String> {
1179 let expected = saved_frame.reduced_total();
1180 if covariance.dim() != (expected, expected) {
1181 return Err(format!(
1182 "binomial mean-wiggle {label} is {}x{}; exact saved-result gauge requires {expected}x{expected} solver-frame coordinates",
1183 covariance.nrows(),
1184 covariance.ncols(),
1185 ));
1186 }
1187 if let Some(((row, column), value)) = covariance
1188 .indexed_iter()
1189 .find(|(_, value)| !value.is_finite())
1190 {
1191 return Err(format!(
1192 "binomial mean-wiggle {label} is non-finite at ({row}, {column}): {value}"
1193 ));
1194 }
1195 let saved = saved_frame.lift_covariance(covariance);
1196 if let Some(((row, column), value)) = saved.indexed_iter().find(|(_, value)| !value.is_finite())
1197 {
1198 return Err(format!(
1199 "binomial mean-wiggle saved-frame {label} is non-finite at ({row}, {column}): {value}"
1200 ));
1201 }
1202 Ok(saved)
1203}
1204
1205fn finalize_binomial_mean_wiggle_saved_frame(
1214 fit: &mut UnifiedFitResult,
1215 alias: &Array2<f64>,
1216 mean_design: &Array2<f64>,
1217 mean_offset: &Array1<f64>,
1218) -> Result<(), String> {
1219 use gam_problem::BlockRole;
1220
1221 if fit.blocks.len() != 2
1222 || fit.blocks[0].role != BlockRole::Mean
1223 || fit.blocks[1].role != BlockRole::LinkWiggle
1224 {
1225 return Err(format!(
1226 "binomial mean-wiggle saved-frame finalization requires fitted blocks [Mean, LinkWiggle], got {:?}",
1227 fit.blocks
1228 .iter()
1229 .map(|block| block.role)
1230 .collect::<Vec<_>>()
1231 ));
1232 }
1233 if fit.block_states.len() != 2 {
1234 return Err(format!(
1235 "binomial mean-wiggle saved-frame finalization requires two fitted block states, got {}",
1236 fit.block_states.len(),
1237 ));
1238 }
1239 if mean_offset.len() != mean_design.nrows() {
1240 return Err(format!(
1241 "binomial mean-wiggle mean offset has {} rows, expected {}",
1242 mean_offset.len(),
1243 mean_design.nrows(),
1244 ));
1245 }
1246
1247 let mean_width = fit.blocks[0].beta.len();
1248 let wiggle_width = fit.blocks[1].beta.len();
1249 if mean_design.ncols() != mean_width {
1250 return Err(format!(
1251 "binomial mean-wiggle saved mean design has {} columns, expected fitted width {mean_width}",
1252 mean_design.ncols(),
1253 ));
1254 }
1255 for block_index in 0..2 {
1256 if fit.block_states[block_index].beta != fit.blocks[block_index].beta {
1257 return Err(format!(
1258 "binomial mean-wiggle fitted block {block_index} and block-state coefficients disagree before saved-frame finalization"
1259 ));
1260 }
1261 }
1262 let total_width = mean_width
1263 .checked_add(wiggle_width)
1264 .ok_or_else(|| "binomial mean-wiggle coefficient dimension overflows usize".to_string())?;
1265 if fit.beta.len() != total_width {
1266 return Err(format!(
1267 "binomial mean-wiggle flat coefficient vector has width {}, expected {total_width}",
1268 fit.beta.len(),
1269 ));
1270 }
1271 if fit.beta.slice(s![0..mean_width]) != fit.blocks[0].beta
1272 || fit.beta.slice(s![mean_width..total_width]) != fit.blocks[1].beta
1273 {
1274 return Err(
1275 "binomial mean-wiggle flat and block coefficient vectors disagree before saved-frame finalization"
1276 .to_string(),
1277 );
1278 }
1279
1280 let saved_frame = binomial_mean_wiggle_saved_frame_gauge(alias, mean_width, wiggle_width)?;
1281 let saved_blocks =
1282 saved_frame.lift_block_betas(&[fit.blocks[0].beta.clone(), fit.blocks[1].beta.clone()]);
1283 let saved_mean_eta = mean_design.dot(&saved_blocks[0]) + mean_offset;
1284 let mut saved_beta = Array1::<f64>::zeros(total_width);
1285 saved_beta
1286 .slice_mut(s![0..mean_width])
1287 .assign(&saved_blocks[0]);
1288 saved_beta
1289 .slice_mut(s![mean_width..total_width])
1290 .assign(&saved_blocks[1]);
1291
1292 let saved_conditional = fit
1293 .covariance_conditional
1294 .as_ref()
1295 .map(|covariance| {
1296 binomial_mean_wiggle_saved_covariance(
1297 covariance,
1298 &saved_frame,
1299 "conditional covariance",
1300 )
1301 })
1302 .transpose()?;
1303 let saved_corrected = fit
1304 .covariance_corrected
1305 .as_ref()
1306 .map(|covariance| {
1307 binomial_mean_wiggle_saved_covariance(covariance, &saved_frame, "corrected covariance")
1308 })
1309 .transpose()?;
1310 let saved_geometry = binomial_mean_wiggle_saved_geometry(
1311 fit.geometry.as_ref().ok_or_else(|| {
1312 "binomial mean-wiggle fit is missing its exact active geometry".to_string()
1313 })?,
1314 &saved_frame,
1315 )?;
1316
1317 let mut saved_inference = fit.inference.clone();
1318 if let Some(inference) = saved_inference.as_mut() {
1319 if inference.beta_covariance.is_none() && inference.beta_standard_errors.is_some() {
1320 return Err(
1321 "binomial mean-wiggle inference has conditional standard errors without their covariance"
1322 .to_string(),
1323 );
1324 }
1325 if inference.beta_covariance_corrected.is_none()
1326 && inference.beta_standard_errors_corrected.is_some()
1327 {
1328 return Err(
1329 "binomial mean-wiggle inference has corrected standard errors without their covariance"
1330 .to_string(),
1331 );
1332 }
1333 if let Some(covariance) = inference.beta_covariance.take() {
1334 let covariance = binomial_mean_wiggle_saved_covariance(
1335 covariance.as_array(),
1336 &saved_frame,
1337 "inference conditional covariance",
1338 )?;
1339 if inference.beta_standard_errors.is_some() {
1340 inference.beta_standard_errors = Some(
1341 gam_problem::se_from_covariance(&covariance).map_err(|reason| {
1342 format!(
1343 "binomial mean-wiggle saved conditional standard errors are invalid: {reason}"
1344 )
1345 })?,
1346 );
1347 }
1348 inference.beta_covariance = Some(covariance.into());
1349 }
1350 if let Some(covariance) = inference.beta_covariance_corrected.take() {
1351 let covariance = binomial_mean_wiggle_saved_covariance(
1352 &covariance,
1353 &saved_frame,
1354 "inference corrected covariance",
1355 )?;
1356 if inference.beta_standard_errors_corrected.is_some() {
1357 inference.beta_standard_errors_corrected = Some(
1358 gam_problem::se_from_covariance(&covariance).map_err(|reason| {
1359 format!(
1360 "binomial mean-wiggle saved corrected standard errors are invalid: {reason}"
1361 )
1362 })?,
1363 );
1364 }
1365 inference.beta_covariance_corrected = Some(covariance);
1366 }
1367 if let Some(covariance) = inference.beta_covariance_frequentist.take() {
1368 inference.beta_covariance_frequentist = Some(binomial_mean_wiggle_saved_covariance(
1369 &covariance,
1370 &saved_frame,
1371 "frequentist covariance",
1372 )?);
1373 }
1374 if let Some(correction) = inference.smoothing_correction.take() {
1375 inference.smoothing_correction = Some(binomial_mean_wiggle_saved_covariance(
1376 &correction,
1377 &saved_frame,
1378 "smoothing covariance correction",
1379 )?);
1380 }
1381 }
1382
1383 fit.blocks[0].beta = saved_blocks[0].clone();
1384 fit.blocks[1].beta = saved_blocks[1].clone();
1385 fit.block_states[0].beta = saved_blocks[0].clone();
1386 fit.block_states[0].eta = saved_mean_eta;
1387 fit.block_states[1].beta = saved_blocks[1].clone();
1388 fit.beta = saved_beta;
1389 fit.covariance_conditional = saved_conditional;
1390 fit.covariance_corrected = saved_corrected;
1391 fit.geometry = Some(saved_geometry);
1392 fit.inference = saved_inference;
1393 Ok(())
1394}
1395
1396#[cfg(test)]
1397mod binomial_mean_wiggle_saved_frame_tests {
1398 use super::*;
1399 use ndarray::array;
1400
1401 #[test]
1402 fn cross_block_dealias_composes_non_square_geometry_and_pushes_covariance() {
1403 let alias = array![[2.0], [-0.5]];
1404 let saved_frame = binomial_mean_wiggle_saved_frame_gauge(&alias, 2, 1)
1405 .expect("valid cross-block de-alias map");
1406
1407 let active_to_solver = gam_problem::Gauge::from_t(
1410 array![[1.0, 0.0], [0.0, 0.0], [0.0, 1.0]],
1411 &[2, 1],
1412 &[1, 1],
1413 );
1414 let active_hessian = array![[7.0, 1.5], [1.5, 4.0]];
1415 let geometry = gam_solve::model_types::FitGeometry {
1416 coefficient_gauge: active_to_solver,
1417 penalized_hessian: active_hessian.clone().into(),
1418 working: None,
1419 };
1420 let saved_geometry = binomial_mean_wiggle_saved_geometry(&geometry, &saved_frame)
1421 .expect("non-square active geometry composes through saved frame");
1422
1423 assert_eq!(
1424 saved_geometry.coefficient_gauge.t_full,
1425 array![[1.0, -2.0], [0.0, 0.5], [0.0, 1.0]],
1426 );
1427 assert_eq!(
1428 saved_geometry.penalized_hessian.as_array(),
1429 &active_hessian,
1430 "precision stays in the canonical active frame",
1431 );
1432
1433 let solver_covariance = Array2::<f64>::eye(3);
1434 let saved_covariance = binomial_mean_wiggle_saved_covariance(
1435 &solver_covariance,
1436 &saved_frame,
1437 "test covariance",
1438 )
1439 .expect("covariance pushes into saved frame");
1440 assert_eq!(
1441 saved_covariance,
1442 array![[5.0, -1.0, -2.0], [-1.0, 1.25, 0.5], [-2.0, 0.5, 1.0]],
1443 "the -A cross block must alter both Mean variance and Mean/Wiggle covariance",
1444 );
1445 }
1446}
1447
1448pub(crate) fn fit_binomial_mean_wiggle(
1452 spec: BinomialMeanWiggleSpec,
1453 options: &BlockwiseFitOptions,
1454) -> Result<(UnifiedFitResult, Option<Vec<f64>>, Option<Vec<f64>>), String> {
1455 let n = spec.y.len();
1456 validate_len_match("weights vs y", n, spec.weights.len())?;
1457 validateweights(&spec.weights, "fit_binomial_mean_wiggle")?;
1458 validate_binomial_response(&spec.y, "fit_binomial_mean_wiggle")?;
1459 validate_blockrows("eta", n, &spec.eta_block)?;
1460 validate_blockrows("wiggle", n, &spec.wiggle_block)?;
1461 if matches!(
1462 spec.link_kind,
1463 InverseLink::Standard(StandardLink::Identity)
1464 ) {
1465 return Err(GamlssError::UnsupportedConfiguration {
1466 reason: "fit_binomial_mean_wiggle does not support identity link".to_string(),
1467 }
1468 .into());
1469 }
1470 gam_terms::inference::formula_dsl::require_binomial_inverse_link_supports_joint_wiggle(
1471 &spec.link_kind,
1472 "fit_binomial_mean_wiggle",
1473 )?;
1474 if spec.wiggle_degree < 2 {
1475 return Err(GamlssError::ConstraintViolation {
1476 reason: format!(
1477 "fit_binomial_mean_wiggle: wiggle_degree must be >= 2, got {}",
1478 spec.wiggle_degree
1479 ),
1480 }
1481 .into());
1482 }
1483 let minimum_knots = minimum_monotone_wiggle_knot_count(spec.wiggle_degree)?;
1484 if spec.wiggle_knots.len() < minimum_knots {
1485 return Err(GamlssError::DimensionMismatch { reason: format!(
1486 "fit_binomial_mean_wiggle: wiggle_knots length {} is too short for degree {} (need at least {})",
1487 spec.wiggle_knots.len(),
1488 spec.wiggle_degree,
1489 minimum_knots
1490 ) }.into());
1491 }
1492
1493 let x_dense: Array2<f64> = spec.eta_block.design.to_dense();
1502 let (pilot_beta, pilot_eta): (Array1<f64>, Array1<f64>) = {
1503 let pilot_beta = spec.eta_block.initial_beta.clone().ok_or_else(|| {
1504 "fit_binomial_mean_wiggle: eta block carries no pilot β to seed the \
1505 frozen-basis warp index"
1506 .to_string()
1507 })?;
1508 if x_dense.ncols() != pilot_beta.len() {
1509 return Err(GamlssError::DimensionMismatch {
1510 reason: format!(
1511 "fit_binomial_mean_wiggle: eta design has {} columns but pilot β has {} \
1512 coefficients",
1513 x_dense.ncols(),
1514 pilot_beta.len()
1515 ),
1516 }
1517 .into());
1518 }
1519 let mut eta = x_dense.dot(&pilot_beta);
1520 eta += &spec.eta_block.offset;
1521 (pilot_beta, eta)
1522 };
1523
1524 let wiggle_penalties_full = spec.wiggle_block.penalties.clone();
1528 let wiggle_nullspace_dims = spec.wiggle_block.nullspace_dims.clone();
1529 if !wiggle_nullspace_dims.is_empty()
1530 && wiggle_nullspace_dims.len() != wiggle_penalties_full.len()
1531 {
1532 return Err(GamlssError::DimensionMismatch {
1533 reason: format!(
1534 "fit_binomial_mean_wiggle: wiggle block has {} penalties but {} nullspace dimensions",
1535 wiggle_penalties_full.len(),
1536 wiggle_nullspace_dims.len()
1537 ),
1538 }
1539 .into());
1540 }
1541 let wiggle_log_lambdas = spec.wiggle_block.initial_log_lambdas.clone();
1542 let wiggle_beta_initial = spec.wiggle_block.initial_beta.clone();
1543 let eta_block_input = spec.eta_block.clone();
1544
1545 let family = BinomialMeanWiggleFamily {
1546 y: spec.y,
1547 weights: spec.weights,
1548 link_kind: spec.link_kind,
1549 wiggle_knots: spec.wiggle_knots,
1550 wiggle_degree: spec.wiggle_degree,
1551 policy: gam_runtime::resource::ResourcePolicy::default_library(),
1552 frozen_warp_design: None,
1553 };
1554
1555 let build_dealiased = |frozen: &Array1<f64>,
1573 beta_hint: Option<&Array1<f64>>,
1574 log_lambda_hint: Option<&Array1<f64>>|
1575 -> Result<
1576 (
1577 ParameterBlockInput,
1578 Array2<f64>,
1579 std::sync::Arc<Array2<f64>>,
1580 ),
1581 String,
1582 > {
1583 use faer::Side;
1584 use gam_linalg::faer_ndarray::FaerEigh;
1585
1586 let b_full = family.wiggle_design(frozen.view())?;
1587 let xtx = x_dense.t().dot(&x_dense);
1588 let xtb = x_dense.t().dot(&b_full);
1589 let (evals, evecs) = xtx
1590 .eigh(Side::Lower)
1591 .map_err(|e| format!("frozen-basis warp de-aliasing mean QR failed: {e}"))?;
1592 let max_eval = evals.iter().map(|v| v.abs()).fold(0.0_f64, f64::max);
1593 let cutoff = 1.0e3 * f64::EPSILON * (xtx.nrows().max(1) as f64) * max_eval.max(1.0);
1594 let mut alias = Array2::<f64>::zeros((x_dense.ncols(), b_full.ncols()));
1595 for k in 0..evals.len() {
1596 let lam = evals[k];
1597 if !lam.is_finite() || lam.abs() <= cutoff {
1598 continue;
1599 }
1600 let uk = evecs.column(k);
1601 let uk_xtb = uk.t().dot(&xtb);
1602 for i in 0..alias.nrows() {
1603 for j in 0..alias.ncols() {
1604 alias[[i, j]] += uk[i] * uk_xtb[j] / lam;
1605 }
1606 }
1607 }
1608 let bda = &b_full - &x_dense.dot(&alias);
1609 let max_b = b_full.iter().map(|v| v.abs()).fold(0.0_f64, f64::max);
1610 let max_resid = bda.iter().map(|v| v.abs()).fold(0.0_f64, f64::max);
1611 let resid_tol =
1612 1.0e3 * f64::EPSILON * (bda.nrows().max(bda.ncols()).max(1) as f64) * max_b.max(1.0);
1613 if max_resid <= resid_tol {
1614 return Err("frozen-basis warp de-aliasing left no identifiable warp \
1615 direction (the mean block already spans the warp in \
1616 observation space)"
1617 .to_string());
1618 }
1619 let penalties: Vec<crate::model_types::PenaltySpec> = wiggle_penalties_full
1620 .iter()
1621 .map(|p| {
1622 let s = penalty_spec_to_dense(p, b_full.ncols())?;
1623 Ok(crate::model_types::PenaltySpec::Dense(s))
1624 })
1625 .collect::<Result<_, String>>()?;
1626 let q = bda.ncols();
1627 let initial_beta = match beta_hint {
1628 Some(beta) if beta.len() == q => Some(beta.clone()),
1629 Some(beta) => {
1630 return Err(format!(
1631 "frozen-basis warp warm start has {} coefficients but the realized basis has {q}",
1632 beta.len()
1633 ));
1634 }
1635 None => Some(Array1::zeros(q)),
1636 };
1637 let block = ParameterBlockInput {
1638 design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(bda.clone())),
1639 offset: Array1::zeros(frozen.len()),
1640 penalties,
1641 nullspace_dims: wiggle_nullspace_dims.clone(),
1642 initial_log_lambdas: log_lambda_hint
1643 .cloned()
1644 .or_else(|| wiggle_log_lambdas.clone()),
1645 initial_beta,
1646 };
1647 Ok((block, alias, std::sync::Arc::new(bda)))
1648 };
1649
1650 if options.outer_max_iter == 0 || !options.outer_tol.is_finite() || options.outer_tol <= 0.0 {
1657 return Err(GamlssError::InvalidInput {
1658 reason: format!(
1659 "fit_binomial_mean_wiggle requires positive outer convergence policy; outer_max_iter={}, outer_tol={}",
1660 options.outer_max_iter, options.outer_tol
1661 ),
1662 }
1663 .into());
1664 }
1665 let mut frozen_source_beta = pilot_beta;
1666 let mut frozen_eta = pilot_eta;
1667 let mut eta_block_warm = eta_block_input.clone();
1668 let mut wiggle_beta_warm = wiggle_beta_initial;
1669 let mut wiggle_log_lambda_warm = wiggle_log_lambdas.clone();
1670 let mut converged: Option<(UnifiedFitResult, Array2<f64>, Array1<f64>)> = None;
1671 let mut last_delta = f64::INFINITY;
1672 let mut last_scale = 1.0_f64;
1673 for _outer in 0..options.outer_max_iter {
1674 let (wiggle_block, alias, bda) = build_dealiased(
1675 &frozen_eta,
1676 wiggle_beta_warm.as_ref(),
1677 wiggle_log_lambda_warm.as_ref(),
1678 )?;
1679 let eta_penalty_count = eta_block_warm.penalties.len();
1680 let wiggle_penalty_count = wiggle_block.penalties.len();
1681 let blocks = vec![
1682 eta_block_warm.clone().intospec("eta")?,
1683 wiggle_block.intospec("wiggle")?,
1684 ];
1685 let mut fam = family.clone();
1686 fam.frozen_warp_design = Some(bda);
1687 let fit = fit_custom_family(&fam, &blocks, options).map_err(|e| e.to_string())?;
1688 let mean_state = fit
1689 .block_states
1690 .get(BinomialMeanWiggleFamily::BLOCK_ETA)
1691 .ok_or_else(|| {
1692 "fit_binomial_mean_wiggle: frozen-basis refit did not expose a fitted eta block"
1693 .to_string()
1694 })?;
1695 if mean_state.eta.len() != frozen_eta.len()
1696 || mean_state.beta.len() != frozen_source_beta.len()
1697 {
1698 return Err(GamlssError::DimensionMismatch {
1699 reason: "fit_binomial_mean_wiggle: frozen-basis refit returned an incompatible eta block"
1700 .to_string(),
1701 }
1702 .into());
1703 }
1704 let new_eta = mean_state.eta.clone();
1705 let new_source_beta = mean_state.beta.clone();
1706 let new_wiggle_beta = fit
1707 .block_states
1708 .get(BinomialMeanWiggleFamily::BLOCK_WIGGLE)
1709 .map(|state| state.beta.clone())
1710 .ok_or_else(|| {
1711 "fit_binomial_mean_wiggle: frozen-basis refit did not expose a fitted wiggle block"
1712 .to_string()
1713 })?;
1714 last_scale = frozen_eta
1715 .iter()
1716 .chain(new_eta.iter())
1717 .map(|value| value.abs())
1718 .fold(1.0_f64, f64::max);
1719 last_delta = new_eta
1720 .iter()
1721 .zip(frozen_eta.iter())
1722 .map(|(a, b)| (a - b).abs())
1723 .fold(0.0_f64, f64::max);
1724 if last_delta <= options.outer_tol * last_scale {
1725 converged = Some((fit, alias, frozen_source_beta));
1726 break;
1727 }
1728
1729 let expected_log_lambdas = eta_penalty_count + wiggle_penalty_count;
1730 if fit.log_lambdas.len() != expected_log_lambdas {
1731 return Err(GamlssError::DimensionMismatch {
1732 reason: format!(
1733 "fit_binomial_mean_wiggle: refit returned {} log-lambdas for {expected_log_lambdas} penalties",
1734 fit.log_lambdas.len()
1735 ),
1736 }
1737 .into());
1738 }
1739 eta_block_warm.initial_beta = Some(new_source_beta.clone());
1740 eta_block_warm.initial_log_lambdas =
1741 Some(fit.log_lambdas.slice(s![0..eta_penalty_count]).to_owned());
1742 wiggle_beta_warm = Some(new_wiggle_beta);
1743 wiggle_log_lambda_warm = Some(
1744 fit.log_lambdas
1745 .slice(s![eta_penalty_count..expected_log_lambdas])
1746 .to_owned(),
1747 );
1748 frozen_source_beta = new_source_beta;
1749 frozen_eta = new_eta;
1750 }
1751 let (mut fit, last_alias, frozen_source_beta) = converged.ok_or_else(|| {
1752 GamlssError::NumericalFailure {
1753 reason: format!(
1754 "fit_binomial_mean_wiggle frozen-index fixed point did not converge in {} outer iterations: delta={last_delta:.3e}, scale={last_scale:.3e}, tolerance={:.3e}",
1755 options.outer_max_iter,
1756 options.outer_tol * last_scale,
1757 ),
1758 }
1759 .to_string()
1760 })?;
1761 let saved_warp_beta = fit
1777 .block_states
1778 .get(BinomialMeanWiggleFamily::BLOCK_WIGGLE)
1779 .map(|state| state.beta.to_vec())
1780 .ok_or_else(|| {
1781 "fit_binomial_mean_wiggle: converged fit is missing its LinkWiggle block state"
1782 .to_string()
1783 })?;
1784 validate_monotone_wiggle_beta_nonnegative(
1785 &saved_warp_beta,
1786 "fit_binomial_mean_wiggle saved warp",
1787 )?;
1788 finalize_binomial_mean_wiggle_saved_frame(
1789 &mut fit,
1790 &last_alias,
1791 &x_dense,
1792 &eta_block_input.offset,
1793 )?;
1794 let saved_mean_state = fit
1799 .block_states
1800 .get(BinomialMeanWiggleFamily::BLOCK_ETA)
1801 .ok_or_else(|| {
1802 "fit_binomial_mean_wiggle: finalized fit is missing its Mean block state".to_string()
1803 })?;
1804 if frozen_source_beta.len() != saved_mean_state.beta.len() {
1805 return Err(format!(
1806 "fit_binomial_mean_wiggle: frozen-index source has {} coefficients, but saved Mean block has {}",
1807 frozen_source_beta.len(),
1808 saved_mean_state.beta.len(),
1809 ));
1810 }
1811 let saved_index_shift = Some((&frozen_source_beta - &saved_mean_state.beta).to_vec());
1812 Ok((fit, Some(saved_warp_beta), saved_index_shift))
1813}
1814
1815fn penalty_spec_to_dense(
1819 spec: &crate::model_types::PenaltySpec,
1820 p: usize,
1821) -> Result<Array2<f64>, String> {
1822 use crate::model_types::PenaltySpec;
1823 match spec {
1824 PenaltySpec::Dense(m) | PenaltySpec::DenseWithMean { matrix: m, .. } => {
1825 if m.nrows() != p || m.ncols() != p {
1826 return Err(format!(
1827 "frozen-basis warp penalty must be {p}x{p}, got {}x{}",
1828 m.nrows(),
1829 m.ncols()
1830 ));
1831 }
1832 Ok(m.clone())
1833 }
1834 PenaltySpec::Block {
1835 local, col_range, ..
1836 } => {
1837 let mut full = Array2::<f64>::zeros((p, p));
1838 if col_range.end > p || local.nrows() != col_range.len() {
1839 return Err("frozen-basis warp penalty block range out of bounds".to_string());
1840 }
1841 full.slice_mut(s![col_range.clone(), col_range.clone()])
1842 .assign(local);
1843 Ok(full)
1844 }
1845 }
1846}
1847
1848pub(crate) trait LocationScaleFamilyBuilder {
1849 type Family: CustomFamily + Clone + Send + Sync + 'static;
1850
1851 fn meanspec(&self) -> &TermCollectionSpec;
1852 fn noisespec(&self) -> &TermCollectionSpec;
1853
1854 fn build_blocks(
1855 &self,
1856 theta: &Array1<f64>,
1857 mean_design: &TermCollectionDesign,
1858 noise_design: &TermCollectionDesign,
1859 mean_beta_hint: Option<Array1<f64>>,
1860 noise_beta_hint: Option<Array1<f64>>,
1861 ) -> Result<Vec<ParameterBlockSpec>, String>;
1862
1863 fn build_family(
1864 &self,
1865 mean_design: &TermCollectionDesign,
1866 noise_design: &TermCollectionDesign,
1867 ) -> Self::Family;
1868
1869 fn extract_primary_betas(
1870 &self,
1871 fit: &UnifiedFitResult,
1872 ) -> Result<(Array1<f64>, Array1<f64>), String>;
1873
1874 fn mean_penalty_count(&self, mean_design: &TermCollectionDesign) -> usize {
1875 mean_design.penalties.len()
1876 }
1877
1878 fn noise_penalty_count(&self, noise_design: &TermCollectionDesign) -> usize {
1879 noise_design.penalties.len()
1880 }
1881
1882 fn exact_spatial_joint_supported(&self) -> bool {
1883 false
1884 }
1885
1886 fn require_exact_spatial_joint(&self) -> bool {
1887 false
1888 }
1889
1890 fn exact_spatial_seed_risk_profile(&self) -> crate::seeding::SeedRiskProfile {
1891 crate::seeding::SeedRiskProfile::GeneralizedLinear
1892 }
1893
1894 fn extra_rho0(&self) -> Result<Array1<f64>, String> {
1895 Ok(Array1::zeros(0))
1896 }
1897
1898 fn build_psiderivative_blocks(
1899 &self,
1900 arr: ndarray::ArrayView2<'_, f64>,
1901 term_spec: &TermCollectionSpec,
1902 term_spec2: &TermCollectionSpec,
1903 term_design: &TermCollectionDesign,
1904 term_design2: &TermCollectionDesign,
1905 ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String>;
1906}
1907
1908pub(crate) fn fit_location_scale_terms<B: LocationScaleFamilyBuilder>(
1909 data: ndarray::ArrayView2<'_, f64>,
1910 builder: B,
1911 options: &BlockwiseFitOptions,
1912 kappa_options: &SpatialLengthScaleOptimizationOptions,
1913) -> Result<BlockwiseTermFitResult, String> {
1914 let mut mean_beta_hint: Option<Array1<f64>> = None;
1920 let mut noise_beta_hint: Option<Array1<f64>> = None;
1921 let extra_rho0 = builder.extra_rho0()?;
1922
1923 let mean_boot_design =
1924 build_term_collection_design(data, builder.meanspec()).map_err(|e| e.to_string())?;
1925 let noise_boot_design =
1926 build_term_collection_design(data, builder.noisespec()).map_err(|e| e.to_string())?;
1927 let mean_bootspec = freeze_term_collection_from_design(builder.meanspec(), &mean_boot_design)
1928 .map_err(|e| e.to_string())?;
1929 let noise_bootspec =
1930 freeze_term_collection_from_design(builder.noisespec(), &noise_boot_design)
1931 .map_err(|e| e.to_string())?;
1932
1933 let require_exact_spatial_joint = builder.require_exact_spatial_joint();
1934 let analytic_joint_derivatives_check = if builder.exact_spatial_joint_supported() {
1935 builder
1936 .build_psiderivative_blocks(
1937 data,
1938 &mean_bootspec,
1939 &noise_bootspec,
1940 &mean_boot_design,
1941 &noise_boot_design,
1942 )
1943 .map(|_| ())
1944 } else {
1945 Err(
1946 "analytic spatial psi derivatives are unavailable for this location-scale family"
1947 .to_string(),
1948 )
1949 };
1950 let analytic_joint_derivatives_available = analytic_joint_derivatives_check.is_ok();
1951 if require_exact_spatial_joint {
1952 analytic_joint_derivatives_check.map_err(|err| {
1953 format!("exact two-block spatial path requires analytic psi derivatives: {err}")
1954 })?;
1955 }
1956 let mean_penalty_count = builder.mean_penalty_count(&mean_boot_design);
1957 let noise_penalty_count = builder.noise_penalty_count(&noise_boot_design);
1958
1959 let mut effective_kappa_options = kappa_options.clone();
1970 if effective_kappa_options.enabled
1971 && gam_terms::smooth::all_spatial_terms_kappa_fixed(&mean_bootspec)
1972 && gam_terms::smooth::all_spatial_terms_kappa_fixed(&noise_bootspec)
1973 {
1974 log::info!(
1975 "[GAMLSS spatial] disabling κ/ψ optimization: every spatial term in \
1976 both blocks has an explicit length_scale and no anisotropy; \
1977 user-supplied kernel scale is fixed"
1978 );
1979 effective_kappa_options.enabled = false;
1980 }
1981 let kappa_options: &SpatialLengthScaleOptimizationOptions = &effective_kappa_options;
1982
1983 macro_rules! run_exact_joint_spatial {
1987 () => {{
1988 let joint_setup = build_two_block_exact_joint_setup(
1989 data,
1990 builder.meanspec(),
1991 builder.noisespec(),
1992 mean_penalty_count,
1993 noise_penalty_count,
1994 extra_rho0.as_slice().unwrap_or(&[]),
1995 None,
1996 kappa_options,
1997 )
1998 .map_err(|error| error.to_string())?;
1999 let mean_terms = spatial_length_scale_term_indices(builder.meanspec());
2000 let noise_terms = spatial_length_scale_term_indices(builder.noisespec());
2001 let mean_beta_hint_cell = std::cell::RefCell::new(mean_beta_hint.clone());
2002 let noise_beta_hint_cell = std::cell::RefCell::new(noise_beta_hint.clone());
2003 let hyper_warm_start_cell =
2004 std::cell::RefCell::new(None::<CustomFamilyWarmStart>);
2005 let gamlss_disable_fixed_point = true;
2015 let outer_policy = {
2016 let theta_seed = joint_setup.theta0();
2028 let rho_dim = joint_setup.rho_dim();
2029 let psi_dim = theta_seed.len() - rho_dim;
2030 let rho_seed = theta_seed.slice(s![..rho_dim]).to_owned();
2031 let policy_blocks_res = builder.build_blocks(
2032 &rho_seed,
2033 &mean_boot_design,
2034 &noise_boot_design,
2035 mean_beta_hint_cell.borrow().clone(),
2036 noise_beta_hint_cell.borrow().clone(),
2037 );
2038 let mut policy = match policy_blocks_res {
2039 Ok(policy_blocks) => {
2040 let policy_family =
2041 builder.build_family(&mean_boot_design, &noise_boot_design);
2042 crate::custom_family::CustomFamily::outer_derivative_policy(
2043 &policy_family,
2044 &policy_blocks,
2045 psi_dim,
2046 options,
2047 )
2048 }
2049 Err(err) => {
2050 log::warn!(
2058 "[GAMLSS spatial] failed to realize policy blocks at seed rho ({err}); \
2059 routing outer optimizer through gradient-only BFGS"
2060 );
2061 let capability = if analytic_joint_derivatives_available {
2062 crate::custom_family::ExactOuterDerivativeOrder::Second
2063 } else {
2064 crate::custom_family::ExactOuterDerivativeOrder::First
2065 };
2066 crate::custom_family::OuterDerivativePolicy {
2067 capability,
2068 predicted_gradient_work: u128::MAX,
2069 predicted_hessian_work: u128::MAX,
2070 subsample_capable: false,
2075 }
2076 }
2077 };
2078 if !analytic_joint_derivatives_available {
2079 policy.capability =
2083 crate::custom_family::ExactOuterDerivativeOrder::First;
2084 }
2085 policy
2086 };
2087 optimize_spatial_length_scale_exact_joint(
2088 data,
2089 &[builder.meanspec().clone(), builder.noisespec().clone()],
2090 &[mean_terms, noise_terms],
2091 kappa_options,
2092 &joint_setup,
2093 builder.exact_spatial_seed_risk_profile(),
2094 analytic_joint_derivatives_available,
2095 analytic_joint_derivatives_available,
2096 gamlss_disable_fixed_point,
2097 None,
2098 outer_policy,
2099 |theta,
2100 specs: &[TermCollectionSpec],
2101 designs: &[TermCollectionDesign],
2102 provenance| {
2103 assert_eq!(
2104 specs.len(),
2105 2,
2106 "joint spatial closure expects exactly two block specs (mean, noise); got {}",
2107 specs.len(),
2108 );
2109 assert_eq!(
2110 designs.len(),
2111 2,
2112 "joint spatial closure expects exactly two block designs (mean, noise); got {}",
2113 designs.len(),
2114 );
2115 let rho = theta.slice(s![..joint_setup.rho_dim()]).to_owned();
2116 let fit = {
2117 let blocks = builder.build_blocks(
2118 &rho,
2119 &designs[0],
2120 &designs[1],
2121 mean_beta_hint_cell.borrow().clone(),
2122 noise_beta_hint_cell.borrow().clone(),
2123 )?;
2124 if mean_beta_hint_cell.borrow().is_none()
2125 && let Some(beta) = blocks.first().and_then(|block| block.initial_beta.clone())
2126 {
2127 *mean_beta_hint_cell.borrow_mut() = Some(beta);
2128 }
2129 if noise_beta_hint_cell.borrow().is_none()
2130 && let Some(beta) =
2131 blocks.get(1).and_then(|block| block.initial_beta.clone())
2132 {
2133 *noise_beta_hint_cell.borrow_mut() = Some(beta);
2134 }
2135 let family = builder.build_family(&designs[0], &designs[1]);
2136 if joint_setup.log_kappa_dim() > 0 && kappa_options.enabled {
2158 let (certified_outer, mode) = match provenance {
2159 SpatialFitProvenance::Certified { outer, mode } => (outer, mode),
2160 SpatialFitProvenance::NoOuterOptimization => {
2161 return Err(
2162 "active GAMLSS spatial optimization returned no certified outer provenance"
2163 .to_string(),
2164 );
2165 }
2166 };
2167 let exact_options =
2168 crate::outer_subsample::exact_outer_options_for_row_set(
2169 options,
2170 &crate::row_kernel::RowSet::All,
2171 );
2172 fit_custom_family_fixed_log_lambdas_from_owned_mode(
2173 &family,
2174 &blocks,
2175 &exact_options,
2176 mode,
2177 theta,
2178 certified_outer,
2179 )?
2180 } else {
2181 fit_custom_family(&family, &blocks, options)?
2182 }
2183 };
2184 let (mean_beta, noise_beta) = builder.extract_primary_betas(&fit)?;
2185 mean_beta_hint = Some(mean_beta);
2186 noise_beta_hint = Some(noise_beta);
2187 *mean_beta_hint_cell.borrow_mut() = mean_beta_hint.clone();
2188 *noise_beta_hint_cell.borrow_mut() = noise_beta_hint.clone();
2189 Ok(fit)
2190 },
2191 |theta,
2192 specs: &[TermCollectionSpec],
2193 designs: &[TermCollectionDesign],
2194 eval_mode,
2195 row_set: &crate::row_kernel::RowSet| {
2196 use gam_problem::EvalMode;
2197 if !analytic_joint_derivatives_available {
2198 return Err(
2199 "analytic spatial psi derivatives are unavailable for this exact two-block path"
2200 .to_string(),
2201 );
2202 }
2203 let rho = theta.slice(s![..joint_setup.rho_dim()]).to_owned();
2204 let blocks = builder.build_blocks(
2205 &rho,
2206 &designs[0],
2207 &designs[1],
2208 mean_beta_hint_cell.borrow().clone(),
2209 noise_beta_hint_cell.borrow().clone(),
2210 )?;
2211 if mean_beta_hint_cell.borrow().is_none()
2212 && let Some(beta) = blocks.first().and_then(|block| block.initial_beta.clone())
2213 {
2214 *mean_beta_hint_cell.borrow_mut() = Some(beta);
2215 }
2216 if noise_beta_hint_cell.borrow().is_none()
2217 && let Some(beta) = blocks.get(1).and_then(|block| block.initial_beta.clone())
2218 {
2219 *noise_beta_hint_cell.borrow_mut() = Some(beta);
2220 }
2221 let family = builder.build_family(&designs[0], &designs[1]);
2222 let psiderivative_blocks = builder.build_psiderivative_blocks(
2223 data,
2224 &specs[0],
2225 &specs[1],
2226 &designs[0],
2227 &designs[1],
2228 )?;
2229 let hyper_layout = crate::custom_family::CustomFamilyHyperLayout::new(
2230 psiderivative_blocks,
2231 Vec::new(),
2232 theta.slice(s![joint_setup.rho_dim()..]).to_owned(),
2233 )?;
2234 let warm_start = hyper_warm_start_cell.borrow().clone();
2235 let eval_options =
2242 crate::outer_subsample::exact_outer_options_for_row_set(options, row_set);
2243 let owned = evaluate_custom_family_joint_hyper_owned(
2244 &family,
2245 &blocks,
2246 &eval_options,
2247 &rho,
2248 &hyper_layout,
2249 warm_start.as_ref(),
2250 eval_mode,
2251 )?;
2252 *hyper_warm_start_cell.borrow_mut() = Some(owned.result.warm_start.clone());
2253 if !owned.result.inner_converged {
2254 return Err(
2255 "exact two-block spatial inner solve did not converge".to_string(),
2256 );
2257 }
2258 if matches!(eval_mode, EvalMode::ValueGradientHessian)
2259 && !owned.result.outer_hessian.is_analytic()
2260 {
2261 return Err(
2262 "exact two-block spatial objective requires a full joint [rho, psi] hessian"
2263 .to_string(),
2264 );
2265 }
2266 Ok(ExactJointEvaluation {
2267 objective: owned.result.objective,
2268 gradient: owned.result.gradient,
2269 hessian: owned.result.outer_hessian,
2270 mode: owned.mode,
2271 })
2272 },
2273 |theta,
2274 specs: &[TermCollectionSpec],
2275 designs: &[TermCollectionDesign],
2276 row_set: &crate::row_kernel::RowSet| {
2277 if !analytic_joint_derivatives_available {
2278 return Err(
2279 "analytic spatial psi derivatives are unavailable for this exact two-block path"
2280 .to_string(),
2281 );
2282 }
2283 let rho = theta.slice(s![..joint_setup.rho_dim()]).to_owned();
2284 let blocks = builder.build_blocks(
2285 &rho,
2286 &designs[0],
2287 &designs[1],
2288 mean_beta_hint_cell.borrow().clone(),
2289 noise_beta_hint_cell.borrow().clone(),
2290 )?;
2291 if mean_beta_hint_cell.borrow().is_none()
2292 && let Some(beta) = blocks.first().and_then(|block| block.initial_beta.clone())
2293 {
2294 *mean_beta_hint_cell.borrow_mut() = Some(beta);
2295 }
2296 if noise_beta_hint_cell.borrow().is_none()
2297 && let Some(beta) = blocks.get(1).and_then(|block| block.initial_beta.clone())
2298 {
2299 *noise_beta_hint_cell.borrow_mut() = Some(beta);
2300 }
2301 let family = builder.build_family(&designs[0], &designs[1]);
2302 let psiderivative_blocks = builder.build_psiderivative_blocks(
2303 data,
2304 &specs[0],
2305 &specs[1],
2306 &designs[0],
2307 &designs[1],
2308 )?;
2309 let hyper_layout = crate::custom_family::CustomFamilyHyperLayout::new(
2310 psiderivative_blocks,
2311 Vec::new(),
2312 theta.slice(s![joint_setup.rho_dim()..]).to_owned(),
2313 )?;
2314 let warm_start = hyper_warm_start_cell.borrow().clone();
2315 let eval_options =
2316 crate::outer_subsample::exact_outer_options_for_row_set(options, row_set);
2317 let owned = evaluate_custom_family_joint_hyper_efs_owned(
2318 &family,
2319 &blocks,
2320 &eval_options,
2321 &rho,
2322 &hyper_layout,
2323 warm_start.as_ref(),
2324 )?;
2325 *hyper_warm_start_cell.borrow_mut() = Some(owned.result.warm_start.clone());
2326 if !owned.result.inner_converged {
2327 return Err(
2328 "exact two-block spatial EFS inner solve did not converge".to_string(),
2329 );
2330 }
2331 Ok(ExactJointEfsEvaluation {
2332 evaluation: owned.result.efs_eval,
2333 mode: owned.mode,
2334 })
2335 },
2336 |_beta: &Array1<f64>| Ok(gam_solve::rho_optimizer::SeedOutcome::NoSlot),
2337 )
2338 }};
2339 }
2340
2341 let mut solved = run_exact_joint_spatial!()
2342 .map_err(|err| format!("exact two-block spatial optimization failed: {err}"))?;
2343
2344 let expected_noise_penalty_count = builder.noise_penalty_count(&solved.designs[1]);
2345 let actual_noise_penalty_count = solved.designs[1].penalties.len();
2346 if expected_noise_penalty_count > actual_noise_penalty_count {
2347 if expected_noise_penalty_count != actual_noise_penalty_count + 1 {
2348 return Err(GamlssError::UnsupportedConfiguration {
2349 reason: format!(
2350 "location-scale result noise design expected {} penalties after augmentation, got {} before augmentation",
2351 expected_noise_penalty_count, actual_noise_penalty_count
2352 ),
2353 }
2354 .into());
2355 }
2356 append_binomial_log_sigma_shrinkage_penalty_design(&mut solved.designs[1]);
2357 }
2358
2359 BlockwiseTermFitResult::try_from_parts(BlockwiseTermFitResultParts {
2360 fit: solved.fit,
2361 meanspec_resolved: solved.resolved_specs.remove(0),
2362 noisespec_resolved: solved.resolved_specs.remove(0),
2363 mean_design: solved.designs.remove(0),
2364 noise_design: solved.designs.remove(0),
2365 })
2366}
2367
2368pub(crate) struct GaussianLocationScaleTermBuilder {
2369 pub(crate) y: Array1<f64>,
2370 pub(crate) weights: Array1<f64>,
2371 pub(crate) meanspec: TermCollectionSpec,
2372 pub(crate) noisespec: TermCollectionSpec,
2373 pub(crate) mean_offset: Array1<f64>,
2374 pub(crate) noise_offset: Array1<f64>,
2375}
2376
2377impl LocationScaleFamilyBuilder for GaussianLocationScaleTermBuilder {
2378 type Family = GaussianLocationScaleFamily;
2379
2380 fn meanspec(&self) -> &TermCollectionSpec {
2381 &self.meanspec
2382 }
2383
2384 fn noisespec(&self) -> &TermCollectionSpec {
2385 &self.noisespec
2386 }
2387
2388 fn exact_spatial_joint_supported(&self) -> bool {
2389 true
2390 }
2391
2392 fn exact_spatial_seed_risk_profile(&self) -> crate::seeding::SeedRiskProfile {
2393 crate::seeding::SeedRiskProfile::GaussianLocationScale
2394 }
2395
2396 fn build_blocks(
2397 &self,
2398 theta: &Array1<f64>,
2399 mean_design: &TermCollectionDesign,
2400 noise_design: &TermCollectionDesign,
2401 mean_beta_hint: Option<Array1<f64>>,
2402 noise_beta_hint: Option<Array1<f64>>,
2403 ) -> Result<Vec<ParameterBlockSpec>, String> {
2404 let layout = GamlssLambdaLayout::two_block(
2405 mean_design.penalties.len(),
2406 self.noise_penalty_count(noise_design),
2407 );
2408 layout.validate_theta_len(theta.len(), "gaussian location-scale")?;
2409 let (meanspec, noisespec) = build_gaussian_mean_and_scale_blocks(
2410 &self.y,
2411 &self.weights,
2412 mean_design,
2413 noise_design,
2414 &self.mean_offset,
2415 &self.noise_offset,
2416 layout.mean_from(theta),
2417 layout.noise_from(theta),
2418 mean_beta_hint,
2419 noise_beta_hint,
2420 "GaussianLocationScale::build_blocks",
2421 )?;
2422 Ok(vec![meanspec, noisespec])
2423 }
2424
2425 fn build_family(
2426 &self,
2427 mean_design: &TermCollectionDesign,
2428 noise_design: &TermCollectionDesign,
2429 ) -> Self::Family {
2430 let preparednoise_design =
2431 prepared_gaussian_log_sigma_design(&mean_design.design, &noise_design.design)
2432 .expect("prepared Gaussian log-sigma design should match block construction");
2433 GaussianLocationScaleFamily {
2434 y: self.y.clone(),
2435 weights: self.weights.clone(),
2436 mu_design: Some(mean_design.design.clone()),
2437 log_sigma_design: Some(preparednoise_design),
2438 policy: gam_runtime::resource::ResourcePolicy::default_library(),
2439 cached_row_scalars: std::sync::RwLock::new(None),
2440 }
2441 }
2442
2443 fn extract_primary_betas(
2444 &self,
2445 fit: &UnifiedFitResult,
2446 ) -> Result<(Array1<f64>, Array1<f64>), String> {
2447 let mean_beta = fit
2448 .block_states
2449 .get(GaussianLocationScaleFamily::BLOCK_MU)
2450 .ok_or_else(|| "missing Gaussian mu block state".to_string())?
2451 .beta
2452 .clone();
2453 let noise_beta = fit
2454 .block_states
2455 .get(GaussianLocationScaleFamily::BLOCK_LOG_SIGMA)
2456 .ok_or_else(|| "missing Gaussian log_sigma block state".to_string())?
2457 .beta
2458 .clone();
2459 Ok((mean_beta, noise_beta))
2460 }
2461
2462 fn build_psiderivative_blocks(
2463 &self,
2464 data: ndarray::ArrayView2<'_, f64>,
2465 meanspec_resolved: &TermCollectionSpec,
2466 noisespec_resolved: &TermCollectionSpec,
2467 mean_design: &TermCollectionDesign,
2468 noise_design: &TermCollectionDesign,
2469 ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String> {
2470 let mean_derivs =
2471 build_block_spatial_psi_derivatives(data, meanspec_resolved, mean_design)?
2472 .ok_or_else(|| "missing Gaussian mean spatial psi derivatives".to_string())?;
2473 let noise_derivs =
2474 build_block_spatial_psi_derivatives(data, noisespec_resolved, noise_design)?
2475 .ok_or_else(|| "missing Gaussian log-sigma spatial psi derivatives".to_string())?;
2476 Ok(vec![mean_derivs, noise_derivs])
2477 }
2478}
2479
2480pub(crate) struct GaussianLocationScaleWiggleTermBuilder {
2481 pub(crate) y: Array1<f64>,
2482 pub(crate) weights: Array1<f64>,
2483 pub(crate) meanspec: TermCollectionSpec,
2484 pub(crate) noisespec: TermCollectionSpec,
2485 pub(crate) mean_offset: Array1<f64>,
2486 pub(crate) noise_offset: Array1<f64>,
2487 pub(crate) wiggle_knots: Array1<f64>,
2488 pub(crate) wiggle_degree: usize,
2489 pub(crate) wiggle_block: ParameterBlockInput,
2490}
2491
2492impl LocationScaleFamilyBuilder for GaussianLocationScaleWiggleTermBuilder {
2493 type Family = GaussianLocationScaleWiggleFamily;
2494
2495 fn meanspec(&self) -> &TermCollectionSpec {
2496 &self.meanspec
2497 }
2498
2499 fn noisespec(&self) -> &TermCollectionSpec {
2500 &self.noisespec
2501 }
2502
2503 fn exact_spatial_joint_supported(&self) -> bool {
2504 true
2505 }
2506
2507 fn exact_spatial_seed_risk_profile(&self) -> crate::seeding::SeedRiskProfile {
2508 crate::seeding::SeedRiskProfile::GaussianLocationScale
2509 }
2510
2511 fn require_exact_spatial_joint(&self) -> bool {
2512 true
2513 }
2514
2515 fn extra_rho0(&self) -> Result<Array1<f64>, String> {
2516 initial_log_lambdas_orzeros(&self.wiggle_block)
2517 }
2518
2519 fn build_blocks(
2520 &self,
2521 theta: &Array1<f64>,
2522 mean_design: &TermCollectionDesign,
2523 noise_design: &TermCollectionDesign,
2524 mean_beta_hint: Option<Array1<f64>>,
2525 noise_beta_hint: Option<Array1<f64>>,
2526 ) -> Result<Vec<ParameterBlockSpec>, String> {
2527 let layout = GamlssLambdaLayout::withwiggle(
2528 mean_design.penalties.len(),
2529 self.noise_penalty_count(noise_design),
2530 self.wiggle_block.penalties.len(),
2531 );
2532 layout.validate_theta_len(theta.len(), "gaussian location-scale wiggle")?;
2533 let (mut meanspec, mut noisespec) = build_gaussian_mean_and_scale_blocks(
2534 &self.y,
2535 &self.weights,
2536 mean_design,
2537 noise_design,
2538 &self.mean_offset,
2539 &self.noise_offset,
2540 layout.mean_from(theta),
2541 layout.noise_from(theta),
2542 mean_beta_hint,
2543 noise_beta_hint,
2544 "GaussianLocationScaleWiggle::build_blocks",
2545 )?;
2546 meanspec.gauge_priority = LINK_WIGGLE_GAUGE_PRIORITY;
2552 noisespec.gauge_priority = LINK_WIGGLE_GAUGE_PRIORITY;
2553 let n_rows = meanspec.design.nrows();
2554 let wigglespec = build_location_scale_wiggle_block(
2555 "wiggle",
2556 self.wiggle_block.design.clone(),
2557 self.wiggle_block.offset.clone(),
2558 wiggle_block_penalty_matrices(&self.wiggle_block),
2559 self.wiggle_block.nullspace_dims.clone(),
2560 layout.wiggle_from(theta),
2561 self.wiggle_block.initial_beta.clone(),
2562 n_rows,
2563 )?;
2564 Ok(vec![meanspec, noisespec, wigglespec])
2565 }
2566
2567 fn build_family(
2568 &self,
2569 mean_design: &TermCollectionDesign,
2570 noise_design: &TermCollectionDesign,
2571 ) -> Self::Family {
2572 let preparednoise_design =
2573 prepared_gaussian_log_sigma_design(&mean_design.design, &noise_design.design).expect(
2574 "prepared Gaussian log-sigma design should match wiggle block construction",
2575 );
2576 GaussianLocationScaleWiggleFamily {
2577 y: self.y.clone(),
2578 weights: self.weights.clone(),
2579 mu_design: Some(mean_design.design.clone()),
2580 log_sigma_design: Some(preparednoise_design),
2581 wiggle_knots: self.wiggle_knots.clone(),
2582 wiggle_degree: self.wiggle_degree,
2583 policy: gam_runtime::resource::ResourcePolicy::default_library(),
2584 cached_row_scalars: std::sync::RwLock::new(None),
2585 }
2586 }
2587
2588 fn extract_primary_betas(
2589 &self,
2590 fit: &UnifiedFitResult,
2591 ) -> Result<(Array1<f64>, Array1<f64>), String> {
2592 let mean_beta = fit
2593 .block_states
2594 .get(GaussianLocationScaleWiggleFamily::BLOCK_MU)
2595 .ok_or_else(|| "missing Gaussian wiggle mu block state".to_string())?
2596 .beta
2597 .clone();
2598 let noise_beta = fit
2599 .block_states
2600 .get(GaussianLocationScaleWiggleFamily::BLOCK_LOG_SIGMA)
2601 .ok_or_else(|| "missing Gaussian wiggle log_sigma block state".to_string())?
2602 .beta
2603 .clone();
2604 Ok((mean_beta, noise_beta))
2605 }
2606
2607 fn build_psiderivative_blocks(
2608 &self,
2609 data: ndarray::ArrayView2<'_, f64>,
2610 meanspec_resolved: &TermCollectionSpec,
2611 noisespec_resolved: &TermCollectionSpec,
2612 mean_design: &TermCollectionDesign,
2613 noise_design: &TermCollectionDesign,
2614 ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String> {
2615 let mean_derivs =
2616 build_block_spatial_psi_derivatives(data, meanspec_resolved, mean_design)?.ok_or_else(
2617 || "missing Gaussian wiggle mean spatial psi derivatives".to_string(),
2618 )?;
2619 let noise_derivs =
2620 build_block_spatial_psi_derivatives(data, noisespec_resolved, noise_design)?
2621 .ok_or_else(|| {
2622 "missing Gaussian wiggle log-sigma spatial psi derivatives".to_string()
2623 })?;
2624 Ok(vec![mean_derivs, noise_derivs, Vec::new()])
2625 }
2626}
2627
2628pub(crate) struct BinomialLocationScaleTermBuilder {
2629 pub(crate) y: Array1<f64>,
2630 pub(crate) weights: Array1<f64>,
2631 pub(crate) link_kind: InverseLink,
2632 pub(crate) meanspec: TermCollectionSpec,
2633 pub(crate) noisespec: TermCollectionSpec,
2634 pub(crate) mean_offset: Array1<f64>,
2635 pub(crate) noise_offset: Array1<f64>,
2636}
2637
2638impl LocationScaleFamilyBuilder for BinomialLocationScaleTermBuilder {
2639 type Family = BinomialLocationScaleFamily;
2640
2641 fn meanspec(&self) -> &TermCollectionSpec {
2642 &self.meanspec
2643 }
2644
2645 fn noisespec(&self) -> &TermCollectionSpec {
2646 &self.noisespec
2647 }
2648
2649 fn exact_spatial_joint_supported(&self) -> bool {
2650 true
2651 }
2652
2653 fn require_exact_spatial_joint(&self) -> bool {
2654 true
2655 }
2656
2657 fn noise_penalty_count(&self, noise_design: &TermCollectionDesign) -> usize {
2658 noise_design.penalties.len() + 1
2659 }
2660
2661 fn build_blocks(
2662 &self,
2663 theta: &Array1<f64>,
2664 mean_design: &TermCollectionDesign,
2665 noise_design: &TermCollectionDesign,
2666 mean_beta_hint: Option<Array1<f64>>,
2667 noise_beta_hint: Option<Array1<f64>>,
2668 ) -> Result<Vec<ParameterBlockSpec>, String> {
2669 let layout = GamlssLambdaLayout::two_block(
2670 mean_design.penalties.len(),
2671 self.noise_penalty_count(noise_design),
2672 );
2673 layout.validate_theta_len(theta.len(), "binomial location-scale")?;
2674 let (thresholdspec, log_sigmaspec) = build_binomial_threshold_and_scale_blocks(
2675 &self.y,
2676 &self.weights,
2677 &self.link_kind,
2678 mean_design,
2679 noise_design,
2680 &self.mean_offset,
2681 &self.noise_offset,
2682 layout.mean_from(theta),
2683 layout.noise_from(theta),
2684 mean_beta_hint,
2685 noise_beta_hint,
2686 "BinomialLocationScale::build_blocks",
2687 )?;
2688 Ok(vec![thresholdspec, log_sigmaspec])
2689 }
2690
2691 fn build_family(
2692 &self,
2693 mean_design: &TermCollectionDesign,
2694 noise_design: &TermCollectionDesign,
2695 ) -> Self::Family {
2696 let identifiednoise_design =
2697 identified_binomial_log_sigma_design(mean_design, noise_design, &self.weights)
2698 .expect("identified binomial log-sigma design");
2699 BinomialLocationScaleFamily {
2700 y: self.y.clone(),
2701 weights: self.weights.clone(),
2702 link_kind: self.link_kind.clone(),
2703 threshold_design: Some(mean_design.design.clone()),
2704 log_sigma_design: Some(identifiednoise_design),
2705 policy: gam_runtime::resource::ResourcePolicy::default_library(),
2706 }
2707 }
2708
2709 fn extract_primary_betas(
2710 &self,
2711 fit: &UnifiedFitResult,
2712 ) -> Result<(Array1<f64>, Array1<f64>), String> {
2713 let mean_beta = fit
2714 .block_states
2715 .get(BinomialLocationScaleFamily::BLOCK_T)
2716 .ok_or_else(|| "missing Binomial threshold block state".to_string())?
2717 .beta
2718 .clone();
2719 let noise_beta = fit
2720 .block_states
2721 .get(BinomialLocationScaleFamily::BLOCK_LOG_SIGMA)
2722 .ok_or_else(|| "missing Binomial log_sigma block state".to_string())?
2723 .beta
2724 .clone();
2725 Ok((mean_beta, noise_beta))
2726 }
2727
2728 fn build_psiderivative_blocks(
2729 &self,
2730 data: ndarray::ArrayView2<'_, f64>,
2731 meanspec_resolved: &TermCollectionSpec,
2732 noisespec_resolved: &TermCollectionSpec,
2733 mean_design: &TermCollectionDesign,
2734 noise_design: &TermCollectionDesign,
2735 ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String> {
2736 let mean_derivs =
2737 build_block_spatial_psi_derivatives(data, meanspec_resolved, mean_design)?
2738 .ok_or_else(|| "missing threshold spatial psi derivatives".to_string())?;
2739 let noise_derivs =
2740 build_block_spatial_psi_derivatives(data, noisespec_resolved, noise_design)?
2741 .ok_or_else(|| "missing log_sigma spatial psi derivatives".to_string())?;
2742 Ok(vec![mean_derivs, noise_derivs])
2743 }
2744}
2745
2746pub(crate) struct BinomialLocationScaleWiggleTermBuilder {
2747 pub(crate) y: Array1<f64>,
2748 pub(crate) weights: Array1<f64>,
2749 pub(crate) link_kind: InverseLink,
2750 pub(crate) meanspec: TermCollectionSpec,
2751 pub(crate) noisespec: TermCollectionSpec,
2752 pub(crate) mean_offset: Array1<f64>,
2753 pub(crate) noise_offset: Array1<f64>,
2754 pub(crate) wiggle_knots: Array1<f64>,
2755 pub(crate) wiggle_degree: usize,
2756 pub(crate) wiggle_block: ParameterBlockInput,
2757}
2758
2759impl LocationScaleFamilyBuilder for BinomialLocationScaleWiggleTermBuilder {
2760 type Family = BinomialLocationScaleWiggleFamily;
2761
2762 fn meanspec(&self) -> &TermCollectionSpec {
2763 &self.meanspec
2764 }
2765
2766 fn noisespec(&self) -> &TermCollectionSpec {
2767 &self.noisespec
2768 }
2769
2770 fn exact_spatial_joint_supported(&self) -> bool {
2771 true
2772 }
2773
2774 fn require_exact_spatial_joint(&self) -> bool {
2775 true
2776 }
2777
2778 fn extra_rho0(&self) -> Result<Array1<f64>, String> {
2779 initial_log_lambdas_orzeros(&self.wiggle_block)
2780 }
2781
2782 fn noise_penalty_count(&self, noise_design: &TermCollectionDesign) -> usize {
2783 noise_design.penalties.len() + 1
2784 }
2785
2786 fn build_blocks(
2787 &self,
2788 theta: &Array1<f64>,
2789 mean_design: &TermCollectionDesign,
2790 noise_design: &TermCollectionDesign,
2791 mean_beta_hint: Option<Array1<f64>>,
2792 noise_beta_hint: Option<Array1<f64>>,
2793 ) -> Result<Vec<ParameterBlockSpec>, String> {
2794 let layout = GamlssLambdaLayout::withwiggle(
2795 mean_design.penalties.len(),
2796 self.noise_penalty_count(noise_design),
2797 self.wiggle_block.penalties.len(),
2798 );
2799 layout.validate_theta_len(theta.len(), "wiggle location-scale")?;
2800 let (mut thresholdspec, mut log_sigmaspec) = build_binomial_threshold_and_scale_blocks(
2801 &self.y,
2802 &self.weights,
2803 &self.link_kind,
2804 mean_design,
2805 noise_design,
2806 &self.mean_offset,
2807 &self.noise_offset,
2808 layout.mean_from(theta),
2809 layout.noise_from(theta),
2810 mean_beta_hint,
2811 noise_beta_hint,
2812 "BinomialLocationScaleWiggle::build_blocks",
2813 )?;
2814 thresholdspec.gauge_priority = LINK_WIGGLE_GAUGE_PRIORITY;
2826 log_sigmaspec.gauge_priority = LINK_WIGGLE_GAUGE_PRIORITY;
2827 let n_rows = thresholdspec.design.nrows();
2828 let wigglespec = build_location_scale_wiggle_block(
2829 "wiggle",
2830 self.wiggle_block.design.clone(),
2831 self.wiggle_block.offset.clone(),
2832 wiggle_block_penalty_matrices(&self.wiggle_block),
2833 vec![],
2834 layout.wiggle_from(theta),
2835 self.wiggle_block.initial_beta.clone(),
2836 n_rows,
2837 )?;
2838 Ok(vec![thresholdspec, log_sigmaspec, wigglespec])
2839 }
2840
2841 fn build_family(
2842 &self,
2843 mean_design: &TermCollectionDesign,
2844 noise_design: &TermCollectionDesign,
2845 ) -> Self::Family {
2846 let identifiednoise_design =
2847 identified_binomial_log_sigma_design(mean_design, noise_design, &self.weights)
2848 .expect("identified binomial log-sigma design should match block construction");
2849 BinomialLocationScaleWiggleFamily {
2850 y: self.y.clone(),
2851 weights: self.weights.clone(),
2852 link_kind: self.link_kind.clone(),
2853 threshold_design: Some(mean_design.design.clone()),
2854 log_sigma_design: Some(identifiednoise_design),
2855 wiggle_knots: self.wiggle_knots.clone(),
2856 wiggle_degree: self.wiggle_degree,
2857 policy: gam_runtime::resource::ResourcePolicy::default_library(),
2858 }
2859 }
2860
2861 fn extract_primary_betas(
2862 &self,
2863 fit: &UnifiedFitResult,
2864 ) -> Result<(Array1<f64>, Array1<f64>), String> {
2865 let mean_beta = fit
2866 .block_states
2867 .get(BinomialLocationScaleWiggleFamily::BLOCK_T)
2868 .ok_or_else(|| "missing Binomial wiggle threshold block state".to_string())?
2869 .beta
2870 .clone();
2871 let noise_beta = fit
2872 .block_states
2873 .get(BinomialLocationScaleWiggleFamily::BLOCK_LOG_SIGMA)
2874 .ok_or_else(|| "missing Binomial wiggle log_sigma block state".to_string())?
2875 .beta
2876 .clone();
2877 Ok((mean_beta, noise_beta))
2878 }
2879
2880 fn build_psiderivative_blocks(
2881 &self,
2882 data: ndarray::ArrayView2<'_, f64>,
2883 meanspec_resolved: &TermCollectionSpec,
2884 noisespec_resolved: &TermCollectionSpec,
2885 mean_design: &TermCollectionDesign,
2886 noise_design: &TermCollectionDesign,
2887 ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String> {
2888 let mean_derivs =
2889 build_block_spatial_psi_derivatives(data, meanspec_resolved, mean_design)?
2890 .ok_or_else(|| "missing threshold spatial psi derivatives".to_string())?;
2891 let noise_derivs =
2892 build_block_spatial_psi_derivatives(data, noisespec_resolved, noise_design)?
2893 .ok_or_else(|| "missing log_sigma spatial psi derivatives".to_string())?;
2894 Ok(vec![mean_derivs, noise_derivs, Vec::new()])
2903 }
2904}
2905
2906pub(crate) fn fit_gaussian_location_scale_terms(
2907 data: ndarray::ArrayView2<'_, f64>,
2908 spec: GaussianLocationScaleTermSpec,
2909 options: &BlockwiseFitOptions,
2910 kappa_options: &SpatialLengthScaleOptimizationOptions,
2911) -> Result<BlockwiseTermFitResult, String> {
2912 validate_gaussian_location_scale_termspec(data, &spec, "fit_gaussian_location_scale_terms")?;
2913 fit_location_scale_terms(
2914 data,
2915 GaussianLocationScaleTermBuilder {
2916 y: spec.y,
2917 weights: spec.weights,
2918 meanspec: spec.meanspec,
2919 noisespec: spec.log_sigmaspec,
2920 mean_offset: spec.mean_offset,
2921 noise_offset: spec.log_sigma_offset,
2922 },
2923 options,
2924 kappa_options,
2925 )
2926}
2927
2928pub(crate) fn fit_gaussian_location_scalewiggle_terms(
2929 data: ndarray::ArrayView2<'_, f64>,
2930 spec: GaussianLocationScaleWiggleTermSpec,
2931 options: &BlockwiseFitOptions,
2932 kappa_options: &SpatialLengthScaleOptimizationOptions,
2933) -> Result<BlockwiseTermFitResult, String> {
2934 validate_gaussian_location_scalewiggle_termspec(
2935 data,
2936 &spec,
2937 "fit_gaussian_location_scalewiggle_terms",
2938 )?;
2939 fit_location_scale_terms(
2940 data,
2941 GaussianLocationScaleWiggleTermBuilder {
2942 y: spec.y,
2943 weights: spec.weights,
2944 meanspec: spec.meanspec,
2945 noisespec: spec.log_sigmaspec,
2946 mean_offset: spec.mean_offset,
2947 noise_offset: spec.log_sigma_offset,
2948 wiggle_knots: spec.wiggle_knots,
2949 wiggle_degree: spec.wiggle_degree,
2950 wiggle_block: spec.wiggle_block,
2951 },
2952 options,
2953 kappa_options,
2954 )
2955}
2956
2957pub(crate) fn select_gaussian_location_scale_link_wiggle_basis_from_pilot(
2958 pilot: &BlockwiseTermFitResult,
2959 wiggle_cfg: &WiggleBlockConfig,
2960 wiggle_penalty_orders: &[usize],
2961) -> Result<SelectedWiggleBasis, String> {
2962 let q_seed = pilot
2963 .fit
2964 .block_states
2965 .first()
2966 .ok_or_else(|| "pilot Gaussian wiggle fit is missing mean block".to_string())?
2967 .eta
2968 .view();
2969 select_wiggle_basis_from_seed(q_seed, wiggle_cfg, wiggle_penalty_orders)
2970}
2971
2972pub(crate) fn fit_gaussian_location_scale_terms_with_selected_wiggle(
2973 data: ndarray::ArrayView2<'_, f64>,
2974 spec: GaussianLocationScaleTermSpec,
2975 selected_wiggle_basis: SelectedWiggleBasis,
2976 options: &BlockwiseFitOptions,
2977 kappa_options: &SpatialLengthScaleOptimizationOptions,
2978) -> Result<BlockwiseTermWiggleFitResult, String> {
2979 let SelectedWiggleBasis {
2980 knots: wiggle_knots,
2981 degree: wiggle_degree,
2982 block: wiggle_block,
2983 ..
2984 } = selected_wiggle_basis;
2985 let solved = fit_gaussian_location_scalewiggle_terms(
2986 data,
2987 GaussianLocationScaleWiggleTermSpec {
2988 y: spec.y,
2989 weights: spec.weights,
2990 meanspec: spec.meanspec,
2991 log_sigmaspec: spec.log_sigmaspec,
2992 mean_offset: spec.mean_offset,
2993 log_sigma_offset: spec.log_sigma_offset,
2994 wiggle_knots: wiggle_knots.clone(),
2995 wiggle_degree,
2996 wiggle_block,
2997 },
2998 options,
2999 kappa_options,
3000 )?;
3001
3002 BlockwiseTermWiggleFitResult::try_from_parts(BlockwiseTermWiggleFitResultParts {
3003 fit: solved,
3004 wiggle_knots,
3005 wiggle_degree,
3006 })
3007}
3008
3009pub(crate) fn fit_binomial_location_scale_terms(
3010 data: ndarray::ArrayView2<'_, f64>,
3011 spec: BinomialLocationScaleTermSpec,
3012 options: &BlockwiseFitOptions,
3013 kappa_options: &SpatialLengthScaleOptimizationOptions,
3014) -> Result<BlockwiseTermFitResult, String> {
3015 validate_binomial_location_scale_termspec(data, &spec, "fit_binomial_location_scale_terms")?;
3016 fit_location_scale_terms(
3017 data,
3018 BinomialLocationScaleTermBuilder {
3019 y: spec.y,
3020 weights: spec.weights,
3021 link_kind: spec.link_kind,
3022 meanspec: spec.thresholdspec,
3023 noisespec: spec.log_sigmaspec,
3024 mean_offset: spec.threshold_offset,
3025 noise_offset: spec.log_sigma_offset,
3026 },
3027 options,
3028 kappa_options,
3029 )
3030}
3031
3032pub(crate) fn fit_binomial_location_scalewiggle_terms(
3033 data: ndarray::ArrayView2<'_, f64>,
3034 spec: BinomialLocationScaleWiggleTermSpec,
3035 options: &BlockwiseFitOptions,
3036 kappa_options: &SpatialLengthScaleOptimizationOptions,
3037) -> Result<BlockwiseTermFitResult, String> {
3038 validate_binomial_location_scalewiggle_termspec(
3039 data,
3040 &spec,
3041 "fit_binomial_location_scalewiggle_terms",
3042 )?;
3043 fit_location_scale_terms(
3044 data,
3045 BinomialLocationScaleWiggleTermBuilder {
3046 y: spec.y,
3047 weights: spec.weights,
3048 link_kind: spec.link_kind,
3049 meanspec: spec.thresholdspec,
3050 noisespec: spec.log_sigmaspec,
3051 mean_offset: spec.threshold_offset,
3052 noise_offset: spec.log_sigma_offset,
3053 wiggle_knots: spec.wiggle_knots,
3054 wiggle_degree: spec.wiggle_degree,
3055 wiggle_block: spec.wiggle_block,
3056 },
3057 options,
3058 kappa_options,
3059 )
3060}
3061
3062pub(crate) fn select_binomial_location_scale_link_wiggle_basis_from_pilot(
3063 pilot: &BlockwiseTermFitResult,
3064 wiggle_cfg: &WiggleBlockConfig,
3065 wiggle_penalty_orders: &[usize],
3066) -> Result<SelectedWiggleBasis, String> {
3067 let eta_t = pilot
3068 .fit
3069 .block_states
3070 .first()
3071 .ok_or_else(|| "pilot fit is missing threshold block".to_string())?
3072 .eta
3073 .view();
3074 let eta_ls = pilot
3075 .fit
3076 .block_states
3077 .get(1)
3078 .ok_or_else(|| "pilot fit is missing log_sigma block".to_string())?
3079 .eta
3080 .view();
3081 let sigma = eta_ls.mapv(safe_exp);
3082 let q_seed = Array1::from_iter(eta_t.iter().zip(sigma.iter()).map(|(&t, &s)| -t / s));
3083 select_wiggle_basis_from_seed(q_seed.view(), wiggle_cfg, wiggle_penalty_orders)
3084}
3085
3086pub(crate) fn fit_binomial_location_scale_terms_with_selected_wiggle(
3087 data: ndarray::ArrayView2<'_, f64>,
3088 spec: BinomialLocationScaleTermSpec,
3089 selected_wiggle_basis: SelectedWiggleBasis,
3090 options: &BlockwiseFitOptions,
3091 kappa_options: &SpatialLengthScaleOptimizationOptions,
3092) -> Result<BlockwiseTermWiggleFitResult, String> {
3093 let SelectedWiggleBasis {
3094 knots: wiggle_knots,
3095 degree: wiggle_degree,
3096 block: wiggle_block,
3097 ..
3098 } = selected_wiggle_basis;
3099 let solved = fit_binomial_location_scalewiggle_terms(
3100 data,
3101 BinomialLocationScaleWiggleTermSpec {
3102 y: spec.y,
3103 weights: spec.weights,
3104 link_kind: spec.link_kind,
3105 thresholdspec: spec.thresholdspec,
3106 log_sigmaspec: spec.log_sigmaspec,
3107 threshold_offset: spec.threshold_offset,
3108 log_sigma_offset: spec.log_sigma_offset,
3109 wiggle_knots: wiggle_knots.clone(),
3110 wiggle_degree,
3111 wiggle_block,
3112 },
3113 options,
3114 kappa_options,
3115 )?;
3116
3117 BlockwiseTermWiggleFitResult::try_from_parts(BlockwiseTermWiggleFitResultParts {
3118 fit: solved,
3119 wiggle_knots,
3120 wiggle_degree,
3121 })
3122}
3123
3124pub(crate) fn select_binomial_mean_link_wiggle_basis_from_pilot(
3125 pilot_design: &TermCollectionDesign,
3126 pilot_fit: &UnifiedFitResult,
3127 wiggle_cfg: &WiggleBlockConfig,
3128 wiggle_penalty_orders: &[usize],
3129) -> Result<SelectedWiggleBasis, String> {
3130 let q_seed = pilot_design
3131 .apply(pilot_fit.beta.view())
3132 .map_err(|error| error.to_string())?;
3133 select_wiggle_basis_from_seed(q_seed.view(), wiggle_cfg, wiggle_penalty_orders)
3134}
3135
3136pub(crate) fn fit_binomial_mean_wiggle_terms_with_selected_basis(
3137 data: ndarray::ArrayView2<'_, f64>,
3138 pilot_spec: &TermCollectionSpec,
3139 pilot_design: &TermCollectionDesign,
3140 pilot_fit: &UnifiedFitResult,
3141 y: &Array1<f64>,
3142 weights: &Array1<f64>,
3143 link_kind: InverseLink,
3144 selected_wiggle_basis: SelectedWiggleBasis,
3145 options: &BlockwiseFitOptions,
3146 kappa_options: &SpatialLengthScaleOptimizationOptions,
3147) -> Result<BinomialMeanWiggleTermFitResult, String> {
3148 const RHO_BOUND: f64 = 12.0;
3149
3150 validate_term_weights(
3151 data,
3152 y.len(),
3153 weights,
3154 "fit_binomial_mean_wiggle_terms_with_selected_basis",
3155 )?;
3156 validate_binomial_response(y, "fit_binomial_mean_wiggle_terms_with_selected_basis")?;
3157
3158 let SelectedWiggleBasis {
3164 knots: wiggle_knots,
3165 degree: wiggle_degree,
3166 block: wiggle_block,
3167 ..
3168 } = selected_wiggle_basis;
3169
3170 let spatial_terms = spatial_length_scale_term_indices(pilot_spec);
3171 if spatial_terms.is_empty() {
3172 let (fit, saved_warp_beta, saved_index_shift) = fit_binomial_mean_wiggle(
3173 BinomialMeanWiggleSpec {
3174 y: y.clone(),
3175 weights: weights.clone(),
3176 link_kind,
3177 wiggle_knots: wiggle_knots.clone(),
3178 wiggle_degree,
3179 eta_block: ParameterBlockInput {
3180 design: pilot_design.design.clone(),
3181 offset: pilot_design.affine_offset.clone(),
3182 penalties: pilot_design
3183 .penalties
3184 .iter()
3185 .map(crate::model_types::PenaltySpec::from_blockwise_ref)
3186 .collect(),
3187 nullspace_dims: vec![],
3188 initial_log_lambdas: Some(fitted_log_lambdas(
3189 &pilot_fit.lambdas,
3190 "binomial mean-wiggle pilot lambda",
3191 )?),
3192 initial_beta: Some(pilot_fit.beta.clone()),
3193 },
3194 wiggle_block,
3195 },
3196 options,
3197 )?;
3198 return Ok(BinomialMeanWiggleTermFitResult {
3199 fit,
3200 resolvedspec: pilot_spec.clone(),
3201 design: pilot_design.clone(),
3202 wiggle_knots,
3203 wiggle_degree,
3204 saved_warp_beta,
3205 saved_index_shift,
3206 });
3207 }
3208
3209 let dims_per_term = spatial_dims_per_term(pilot_spec, &spatial_terms);
3210 let log_kappa0 =
3211 SpatialLogKappaCoords::from_length_scales_aniso(pilot_spec, &spatial_terms, kappa_options)
3212 .reseed_from_data(data, pilot_spec, &spatial_terms, kappa_options)
3213 .map_err(|error| error.to_string())?;
3214 let log_kappa_lower = SpatialLogKappaCoords::lower_bounds_aniso_from_data(
3215 data,
3216 pilot_spec,
3217 &spatial_terms,
3218 &dims_per_term,
3219 kappa_options,
3220 )
3221 .map_err(|error| error.to_string())?;
3222 let log_kappa_upper = SpatialLogKappaCoords::upper_bounds_aniso_from_data(
3223 data,
3224 pilot_spec,
3225 &spatial_terms,
3226 &dims_per_term,
3227 kappa_options,
3228 )
3229 .map_err(|error| error.to_string())?;
3230 let log_kappa0 = log_kappa0.clamp_to_bounds(&log_kappa_lower, &log_kappa_upper);
3232
3233 let eta_penalty_count = pilot_design.penalties.len();
3234 let wiggle_penalty_count = initial_log_lambdas_orzeros(&wiggle_block)?.len();
3235 let rho_dim = eta_penalty_count + wiggle_penalty_count;
3236 let baseline_resolvedspec = log_kappa0
3237 .apply_tospec(pilot_spec, &spatial_terms)
3238 .map_err(|e| e.to_string())?;
3239 let baseline_design =
3240 build_term_collection_design(data, &baseline_resolvedspec).map_err(|e| e.to_string())?;
3241 let baseline_fit = fit_binomial_mean_wiggle(
3242 BinomialMeanWiggleSpec {
3243 y: y.clone(),
3244 weights: weights.clone(),
3245 link_kind: link_kind.clone(),
3246 wiggle_knots: wiggle_knots.clone(),
3247 wiggle_degree,
3248 eta_block: ParameterBlockInput {
3249 design: baseline_design.design.clone(),
3250 offset: baseline_design.affine_offset.clone(),
3251 penalties: baseline_design
3252 .penalties
3253 .iter()
3254 .map(crate::model_types::PenaltySpec::from_blockwise_ref)
3255 .collect(),
3256 nullspace_dims: vec![],
3257 initial_log_lambdas: Some(fitted_log_lambdas(
3258 &pilot_fit.lambdas,
3259 "binomial mean-wiggle pilot lambda",
3260 )?),
3261 initial_beta: Some(pilot_fit.beta.clone()),
3262 },
3263 wiggle_block: wiggle_block.clone(),
3264 },
3265 options,
3266 )?
3267 .0;
3268 let baseline_log_lambdas = fitted_log_lambdas(
3269 &baseline_fit.lambdas,
3270 "binomial mean-wiggle baseline lambda",
3271 )?;
3272 if baseline_log_lambdas.len() != rho_dim {
3273 return Err(GamlssError::DimensionMismatch {
3274 reason: format!(
3275 "baseline binomial mean-wiggle fit returned {} log-lambdas, expected {rho_dim}",
3276 baseline_log_lambdas.len()
3277 ),
3278 }
3279 .into());
3280 }
3281 let baseline_eta_beta = baseline_fit
3282 .block_states
3283 .get(BinomialMeanWiggleFamily::BLOCK_ETA)
3284 .ok_or_else(|| "baseline binomial mean-wiggle fit missing eta block".to_string())?
3285 .beta
3286 .clone();
3287 let baseline_wiggle_beta = Some(
3288 baseline_fit
3289 .block_states
3290 .get(BinomialMeanWiggleFamily::BLOCK_WIGGLE)
3291 .ok_or_else(|| "baseline binomial mean-wiggle fit missing wiggle block".to_string())?
3292 .beta
3293 .clone(),
3294 );
3295 let theta_dim = rho_dim + log_kappa0.len();
3296 let mut theta0 = Array1::<f64>::zeros(theta_dim);
3297 theta0
3298 .slice_mut(s![0..rho_dim])
3299 .assign(&baseline_log_lambdas);
3300 theta0
3301 .slice_mut(s![rho_dim..theta_dim])
3302 .assign(log_kappa0.as_array());
3303
3304 let mut lower = Array1::<f64>::from_elem(theta_dim, -RHO_BOUND);
3305 let mut upper = Array1::<f64>::from_elem(theta_dim, RHO_BOUND);
3306 lower
3307 .slice_mut(s![rho_dim..theta_dim])
3308 .assign(log_kappa_lower.as_array());
3309 upper
3310 .slice_mut(s![rho_dim..theta_dim])
3311 .assign(log_kappa_upper.as_array());
3312
3313 let pilot_spec_cloned = pilot_spec.clone();
3314 let pilot_beta = baseline_eta_beta;
3315 let wiggle_design = wiggle_block.design.clone();
3316 let wiggle_offset = wiggle_block.offset.clone();
3317 let wiggle_penalties = wiggle_block.penalties.clone();
3318 let wiggle_initial_beta = baseline_wiggle_beta;
3319 let wiggle_knots_cloned = wiggle_knots.clone();
3320 let y_cloned = y.clone();
3321 let weights_cloned = weights.clone();
3322 let link_kind_cloned = link_kind.clone();
3323 let outer_family = BinomialMeanWiggleFamily {
3324 y: y_cloned.clone(),
3325 weights: weights_cloned.clone(),
3326 link_kind: link_kind_cloned.clone(),
3327 wiggle_knots: wiggle_knots_cloned.clone(),
3328 wiggle_degree,
3329 policy: gam_runtime::resource::ResourcePolicy::default_library(),
3330 frozen_warp_design: None,
3333 };
3334 let screening_cap = Arc::new(AtomicUsize::new(0));
3335 let mut outer_options = options.clone();
3336 outer_options.screening_max_inner_iterations = Some(Arc::clone(&screening_cap));
3337 struct MeanWiggleOuterState {
3338 pub(crate) warm_cache: Option<crate::custom_family::CustomFamilyWarmStart>,
3339 pub(crate) last_eval: Option<(
3340 Array1<f64>,
3341 f64,
3342 Array1<f64>,
3343 gam_problem::HessianValue,
3344 crate::custom_family::CustomFamilyWarmStart,
3345 )>,
3346 }
3347
3348 let build_realized_blocks = |theta: &Array1<f64>| -> Result<
3349 (
3350 TermCollectionSpec,
3351 TermCollectionDesign,
3352 Vec<ParameterBlockSpec>,
3353 Vec<CustomFamilyBlockPsiDerivative>,
3354 ),
3355 String,
3356 > {
3357 let log_kappa =
3358 SpatialLogKappaCoords::from_theta_tail_with_dims(theta, rho_dim, dims_per_term.clone());
3359 let resolvedspec = log_kappa
3360 .apply_tospec(&pilot_spec_cloned, &spatial_terms)
3361 .map_err(|e| e.to_string())?;
3362 let design =
3363 build_term_collection_design(data, &resolvedspec).map_err(|e| e.to_string())?;
3364 let eta_derivs = build_block_spatial_psi_derivatives(data, &resolvedspec, &design)?
3365 .ok_or_else(|| {
3366 "missing eta spatial psi derivatives for binomial mean wiggle".to_string()
3367 })?;
3368 let blocks = vec![
3369 ParameterBlockSpec {
3370 name: "eta".to_string(),
3371 design: design.design.clone(),
3372 offset: design.affine_offset.clone(),
3373 penalties: design.penalties_as_penalty_matrix(),
3374 nullspace_dims: vec![],
3375 initial_log_lambdas: theta.slice(s![0..eta_penalty_count]).to_owned(),
3376 initial_beta: Some(pilot_beta.clone()),
3377 gauge_priority: LINK_WIGGLE_GAUGE_PRIORITY,
3381 jacobian_callback: None,
3382 stacked_design: None,
3383 stacked_offset: None,
3384 },
3385 ParameterBlockSpec {
3386 name: "wiggle".to_string(),
3387 design: wiggle_design.clone(),
3388 offset: wiggle_offset.clone(),
3389 penalties: {
3390 let p_wiggle = wiggle_design.ncols();
3391 wiggle_penalties
3392 .iter()
3393 .map(|spec| match spec {
3394 crate::model_types::PenaltySpec::Block {
3395 local, col_range, ..
3396 } => PenaltyMatrix::Blockwise {
3397 local: local.clone(),
3398 col_range: col_range.clone(),
3399 total_dim: p_wiggle,
3400 },
3401 crate::model_types::PenaltySpec::Dense(m)
3402 | crate::model_types::PenaltySpec::DenseWithMean {
3403 matrix: m, ..
3404 } => PenaltyMatrix::Dense(m.clone()),
3405 })
3406 .collect()
3407 },
3408 nullspace_dims: vec![],
3409 initial_log_lambdas: theta.slice(s![eta_penalty_count..rho_dim]).to_owned(),
3410 initial_beta: wiggle_initial_beta.clone(),
3411 gauge_priority: DEFAULT_GAUGE_PRIORITY,
3412 jacobian_callback: None,
3413 stacked_design: None,
3414 stacked_offset: None,
3415 },
3416 ];
3417 Ok((resolvedspec, design, blocks, eta_derivs))
3418 };
3419
3420 let build_eval = |theta: &Array1<f64>,
3421 warm_cache: Option<&crate::custom_family::CustomFamilyWarmStart>,
3422 need_hessian: bool|
3423 -> Result<
3424 (
3425 crate::custom_family::CustomFamilyJointHyperResult,
3426 TermCollectionSpec,
3427 TermCollectionDesign,
3428 ),
3429 String,
3430 > {
3431 let (resolvedspec, design, blocks, eta_derivs) = build_realized_blocks(theta)?;
3432 let hyper_layout = crate::custom_family::CustomFamilyHyperLayout::new(
3433 vec![eta_derivs, Vec::new()],
3434 Vec::new(),
3435 theta.slice(s![rho_dim..]).to_owned(),
3436 )?;
3437 let eval = evaluate_custom_family_joint_hyper(
3438 &outer_family,
3439 &blocks,
3440 &outer_options,
3441 &theta.slice(s![0..rho_dim]).to_owned(),
3442 &hyper_layout,
3443 warm_cache,
3444 if need_hessian {
3445 gam_problem::EvalMode::ValueGradientHessian
3446 } else {
3447 gam_problem::EvalMode::ValueAndGradient
3448 },
3449 )?;
3450 Ok((eval, resolvedspec, design))
3451 };
3452
3453 let build_efs = |theta: &Array1<f64>,
3454 warm_cache: Option<&crate::custom_family::CustomFamilyWarmStart>|
3455 -> Result<crate::custom_family::CustomFamilyJointHyperEfsResult, String> {
3456 let (_, _, blocks, eta_derivs) = build_realized_blocks(theta)?;
3457 let hyper_layout = crate::custom_family::CustomFamilyHyperLayout::new(
3458 vec![eta_derivs, Vec::new()],
3459 Vec::new(),
3460 theta.slice(s![rho_dim..]).to_owned(),
3461 )?;
3462 evaluate_custom_family_joint_hyper_efs(
3463 &outer_family,
3464 &blocks,
3465 &outer_options,
3466 &theta.slice(s![0..rho_dim]).to_owned(),
3467 &hyper_layout,
3468 warm_cache,
3469 )
3470 .map_err(|e| e.to_string())
3471 };
3472
3473 use crate::model_types::EstimationError;
3474 use gam_problem::{DeclaredHessianForm, Derivative, OuterEval};
3475 use gam_solve::rho_optimizer::OuterEvalOrder;
3476
3477 let analytic_outer_hessian_available = true;
3488 let mut seed_heuristic = theta0.to_vec();
3489 for value in &mut seed_heuristic[..rho_dim] {
3490 *value = value.exp();
3491 }
3492 let problem = gam_solve::rho_optimizer::OuterProblem::new(theta_dim)
3493 .with_gradient(Derivative::Analytic)
3494 .with_hessian(if analytic_outer_hessian_available {
3495 DeclaredHessianForm::Either
3496 } else {
3497 DeclaredHessianForm::Unavailable
3498 })
3499 .with_psi_dim(theta_dim - rho_dim)
3500 .with_tolerance(options.outer_tol)
3501 .with_max_iter(options.outer_max_iter)
3502 .with_bounds(lower.clone(), upper.clone())
3503 .with_initial_rho(theta0.clone())
3504 .with_seed_config(crate::seeding::SeedConfig {
3505 max_seeds: 4,
3506 seed_budget: 2,
3507 risk_profile: crate::seeding::SeedRiskProfile::GeneralizedLinear,
3508 num_auxiliary_trailing: theta_dim - rho_dim,
3509 ..Default::default()
3510 })
3511 .with_screening_cap(Arc::clone(&screening_cap))
3512 .with_rho_bound(12.0)
3513 .with_heuristic_lambdas(seed_heuristic);
3514
3515 let eval_outer = |state: &mut MeanWiggleOuterState,
3516 theta: &Array1<f64>,
3517 order: OuterEvalOrder|
3518 -> Result<OuterEval, EstimationError> {
3519 if let Some((cached_theta, cached_cost, cached_grad, cached_hess, cached_warm)) =
3520 &state.last_eval
3521 && cached_theta == theta
3522 && (!matches!(order, OuterEvalOrder::ValueGradientHessian)
3523 || matches!(
3524 cached_hess,
3525 gam_problem::HessianValue::Dense(_) | gam_problem::HessianValue::Operator(_)
3526 ))
3527 {
3528 state.warm_cache = Some(cached_warm.clone());
3529 return Ok(OuterEval {
3530 cost: *cached_cost,
3531 gradient: cached_grad.clone(),
3532 hessian: cached_hess.clone(),
3533 inner_beta_hint: None,
3534 });
3535 }
3536 let need_hessian = matches!(order, OuterEvalOrder::ValueGradientHessian)
3537 && analytic_outer_hessian_available;
3538 let (eval, _, _) = build_eval(theta, state.warm_cache.as_ref(), need_hessian)
3539 .map_err(EstimationError::InvalidInput)?;
3540 if !eval.inner_converged {
3541 state.warm_cache = Some(eval.warm_start);
3542 crate::bail_invalid_estim!(
3543 "binomial mean-wiggle exact spatial inner solve did not converge"
3544 );
3545 }
3546 let hessian_result = eval.outer_hessian.clone();
3547 state.last_eval = Some((
3548 theta.clone(),
3549 eval.objective,
3550 eval.gradient.clone(),
3551 eval.outer_hessian.clone(),
3552 eval.warm_start.clone(),
3553 ));
3554 state.warm_cache = Some(eval.warm_start);
3555 Ok(OuterEval {
3556 cost: eval.objective,
3557 gradient: eval.gradient,
3558 hessian: hessian_result,
3559 inner_beta_hint: None,
3560 })
3561 };
3562
3563 let mut obj = problem.build_objective_with_screening_proxy(
3564 MeanWiggleOuterState {
3565 warm_cache: None,
3566 last_eval: None,
3567 },
3568 |state: &mut MeanWiggleOuterState, theta: &Array1<f64>| {
3569 if let Some((cached_theta, cached_cost, _, _, cached_warm)) = &state.last_eval
3570 && cached_theta == theta
3571 {
3572 state.warm_cache = Some(cached_warm.clone());
3573 return Ok(*cached_cost);
3574 }
3575 let (eval, _, _) = build_eval(theta, state.warm_cache.as_ref(), false)
3576 .map_err(EstimationError::InvalidInput)?;
3577 if !eval.inner_converged {
3578 state.warm_cache = Some(eval.warm_start);
3579 crate::bail_invalid_estim!(
3580 "binomial mean-wiggle exact spatial cost inner solve did not converge"
3581 .to_string(),
3582 );
3583 }
3584 state.warm_cache = Some(eval.warm_start);
3585 Ok(eval.objective)
3586 },
3587 |state: &mut MeanWiggleOuterState, theta: &Array1<f64>| {
3588 eval_outer(
3589 state,
3590 theta,
3591 if analytic_outer_hessian_available {
3592 OuterEvalOrder::ValueGradientHessian
3593 } else {
3594 OuterEvalOrder::ValueAndGradient
3595 },
3596 )
3597 },
3598 |state: &mut MeanWiggleOuterState, theta: &Array1<f64>, order: OuterEvalOrder| {
3599 eval_outer(state, theta, order)
3600 },
3601 Some(|state: &mut MeanWiggleOuterState| {
3602 state.warm_cache = None;
3603 state.last_eval = None;
3604 }),
3605 Some(|state: &mut MeanWiggleOuterState, theta: &Array1<f64>| {
3606 let eval = build_efs(theta, state.warm_cache.as_ref())
3607 .map_err(EstimationError::InvalidInput)?;
3608 if !eval.inner_converged {
3609 state.warm_cache = Some(eval.warm_start);
3610 crate::bail_invalid_estim!(
3611 "binomial mean-wiggle exact spatial EFS inner solve did not converge"
3612 .to_string(),
3613 );
3614 }
3615 state.warm_cache = Some(eval.warm_start);
3616 Ok(eval.efs_eval)
3617 }),
3618 |state: &mut MeanWiggleOuterState, theta: &Array1<f64>| {
3628 if let Some((cached_theta, cached_cost, _, _, cached_warm)) = &state.last_eval
3629 && cached_theta == theta
3630 {
3631 state.warm_cache = Some(cached_warm.clone());
3632 return Ok(*cached_cost);
3633 }
3634 let (eval, _, _) = build_eval(theta, state.warm_cache.as_ref(), false)
3635 .map_err(EstimationError::InvalidInput)?;
3636 state.warm_cache = Some(eval.warm_start);
3637 Ok(eval.objective)
3638 },
3639 );
3640
3641 let outer = problem
3642 .run(&mut obj, "binomial mean wiggle exact spatial hyper")
3643 .map_err(|e| e.to_string())?;
3644 if !outer.converged {
3645 return Err(GamlssError::NumericalFailure { reason: format!(
3646 "binomial mean wiggle exact spatial hyper did not converge after {} iterations (final_objective={:.6e}, final_grad_norm={})",
3647 outer.iterations,
3648 outer.final_value,
3649 outer.final_grad_norm_report(),
3650 ) }.into());
3651 }
3652 let theta_star = outer.rho;
3653
3654 let log_kappa =
3655 SpatialLogKappaCoords::from_theta_tail_with_dims(&theta_star, rho_dim, dims_per_term);
3656 let resolvedspec = log_kappa
3657 .apply_tospec(&pilot_spec_cloned, &spatial_terms)
3658 .map_err(|e| e.to_string())?;
3659 let design = build_term_collection_design(data, &resolvedspec).map_err(|e| e.to_string())?;
3660 let resolvedspec =
3661 freeze_term_collection_from_design(&resolvedspec, &design).map_err(|e| e.to_string())?;
3662 let fit = fit_binomial_mean_wiggle(
3663 BinomialMeanWiggleSpec {
3664 y: y_cloned,
3665 weights: weights_cloned,
3666 link_kind: link_kind_cloned,
3667 wiggle_knots: wiggle_knots.clone(),
3668 wiggle_degree,
3669 eta_block: ParameterBlockInput {
3670 design: design.design.clone(),
3671 offset: design.affine_offset.clone(),
3672 penalties: design
3673 .penalties
3674 .iter()
3675 .map(crate::model_types::PenaltySpec::from_blockwise_ref)
3676 .collect(),
3677 nullspace_dims: vec![],
3678 initial_log_lambdas: Some(theta_star.slice(s![0..eta_penalty_count]).to_owned()),
3679 initial_beta: Some(pilot_beta),
3680 },
3681 wiggle_block: ParameterBlockInput {
3682 design: wiggle_design,
3683 offset: wiggle_offset,
3684 penalties: wiggle_penalties,
3685 nullspace_dims: vec![],
3686 initial_log_lambdas: Some(
3687 theta_star.slice(s![eta_penalty_count..rho_dim]).to_owned(),
3688 ),
3689 initial_beta: wiggle_initial_beta,
3690 },
3691 },
3692 options,
3693 )?;
3694 let (fit, saved_warp_beta, saved_index_shift) = fit;
3695
3696 Ok(BinomialMeanWiggleTermFitResult {
3697 fit,
3698 resolvedspec,
3699 design,
3700 wiggle_knots,
3701 wiggle_degree,
3702 saved_warp_beta,
3703 saved_index_shift,
3704 })
3705}