1use super::*;
6
7#[derive(Debug)]
13pub enum GamlssError {
14 DimensionMismatch { reason: String },
17 InvalidInput { reason: String },
21 NonFinite { reason: String },
24 UnsupportedConfiguration { reason: String },
28 ConstraintViolation { reason: String },
31 NumericalFailure { reason: String },
35}
36
37impl_reason_error_boilerplate! {
38 GamlssError {
39 DimensionMismatch,
40 InvalidInput,
41 NonFinite,
42 UnsupportedConfiguration,
43 ConstraintViolation,
44 NumericalFailure,
45 }
46}
47
48impl From<crate::block_layout::block_count::BlockCountMismatch> for GamlssError {
49 fn from(err: crate::block_layout::block_count::BlockCountMismatch) -> GamlssError {
50 GamlssError::DimensionMismatch {
51 reason: err.message(),
52 }
53 }
54}
55
56pub(crate) const MIN_PROB: f64 = 1e-10;
69
70pub(crate) const MIN_DERIV: f64 = 1e-8;
71
72use gam_problem::MIN_WEIGHT;
89
90pub(crate) const ETA_HARD_CLAMP: f64 = 30.0;
95
96#[inline]
103pub(crate) fn saturated_exp_eta(eta: f64) -> f64 {
104 eta.clamp(-ETA_HARD_CLAMP, ETA_HARD_CLAMP)
105 .exp()
106 .max(MIN_WEIGHT)
107}
108
109pub(crate) const WARMSTART_LOG_LAMBDA_FLOOR: f64 = 1e-12;
118
119pub(crate) const EXACT_DENSE_BLOCK_BUDGET_BYTES: usize = 512 * 1024 * 1024;
120
121pub(crate) const EXACT_DENSE_TOTAL_BUDGET_BYTES: usize = 2 * 1024 * 1024 * 1024;
122
123pub(crate) const GAMLSS_ROWWISE_PAR_MIN_N: usize = 4096;
124
125pub(crate) const GAMLSS_PROJECTED_TRACE_TARGET_BYTES: usize = 32 * 1024 * 1024;
126
127pub(crate) const GAMLSS_PROJECTED_TRACE_MIN_CHUNK_ROWS: usize = 64;
128
129pub(crate) const GAMLSS_PROJECTED_TRACE_MAX_CHUNK_ROWS: usize = 8192;
130
131pub(crate) fn gamlss_projected_trace_chunk_rows(
132 rank: usize,
133 projected_channel_count: usize,
134 gram_column_count: usize,
135) -> usize {
136 let per_row_values = rank
137 .saturating_mul(projected_channel_count.max(1))
138 .saturating_add(gram_column_count.max(1))
139 .max(1);
140 let per_row_bytes = per_row_values.saturating_mul(std::mem::size_of::<f64>());
141 let rows = GAMLSS_PROJECTED_TRACE_TARGET_BYTES / per_row_bytes.max(1);
142 rows.clamp(
143 GAMLSS_PROJECTED_TRACE_MIN_CHUNK_ROWS,
144 GAMLSS_PROJECTED_TRACE_MAX_CHUNK_ROWS,
145 )
146}
147
148pub(crate) fn gamlss_rowwise_map<F>(n: usize, f: F) -> Array1<f64>
149where
150 F: Fn(usize) -> f64 + Sync,
151{
152 if n >= GAMLSS_ROWWISE_PAR_MIN_N {
153 Array1::from((0..n).into_par_iter().map(&f).collect::<Vec<f64>>())
154 } else {
155 Array1::from_iter((0..n).map(f))
156 }
157}
158
159pub(crate) fn gamlss_rowwise_map_result<F>(n: usize, f: F) -> Result<Array1<f64>, String>
160where
161 F: Fn(usize) -> Result<f64, String> + Sync,
162{
163 if n >= GAMLSS_ROWWISE_PAR_MIN_N {
164 let values: Result<Vec<f64>, String> = (0..n).into_par_iter().map(&f).collect();
165 Ok(Array1::from(values?))
166 } else {
167 let mut out = Array1::<f64>::zeros(n);
168 for i in 0..n {
169 out[i] = f(i)?;
170 }
171 Ok(out)
172 }
173}
174
175pub(crate) enum DenseOrOperator<'a> {
176 Borrowed(&'a Array2<f64>),
177 Owned(Array2<f64>),
178 Operator(DesignMatrix),
179}
180
181impl DenseOrOperator<'_> {
182 pub(crate) fn nrows(&self) -> usize {
183 match self {
184 Self::Borrowed(dense) => dense.nrows(),
185 Self::Owned(dense) => dense.nrows(),
186 Self::Operator(design) => design.nrows(),
187 }
188 }
189
190 pub(crate) fn ncols(&self) -> usize {
191 match self {
192 Self::Borrowed(dense) => dense.ncols(),
193 Self::Owned(dense) => dense.ncols(),
194 Self::Operator(design) => design.ncols(),
195 }
196 }
197
198 pub(crate) fn row_chunk(&self, rows: std::ops::Range<usize>) -> Result<Array2<f64>, String> {
199 match self {
200 Self::Borrowed(dense) => Ok(dense.slice(s![rows, ..]).to_owned()),
201 Self::Owned(dense) => Ok(dense.slice(s![rows, ..]).to_owned()),
202 Self::Operator(design) => design.try_row_chunk(rows).map_err(|e| e.to_string()),
203 }
204 }
205
206 pub(crate) fn dot(&self, beta: ArrayView1<'_, f64>) -> Array1<f64> {
207 let n = self.nrows();
208 let p = self.ncols();
209 assert_eq!(beta.len(), p);
210 match self {
211 Self::Borrowed(dense) => fast_av(*dense, &beta),
212 Self::Owned(dense) => fast_av(dense, &beta),
213 Self::Operator(design) => {
214 let mut out = Array1::<f64>::zeros(n);
215 for rows in exact_design_row_chunks(n, p) {
216 let chunk = design
217 .try_row_chunk(rows.clone())
218 .expect("gamlss DesignSlot::dot: design row chunk materialization failed");
219 out.slice_mut(s![rows]).assign(&fast_av(&chunk, &beta));
220 }
221 out
222 }
223 }
224 }
225}
226
227pub(crate) fn dense_block_from_spec<'a>(
234 spec: &'a ParameterBlockSpec,
235 material_policy: &gam_runtime::resource::MaterializationPolicy,
236 materialization_label: &str,
237) -> Result<Cow<'a, Array2<f64>>, String> {
238 match spec.design.as_dense_ref() {
239 Some(d) => Ok(Cow::Borrowed(d)),
240 None => Ok(Cow::Owned(
241 spec.design
242 .try_to_dense_with_policy(material_policy, "gamlss dense_block_from_spec")
243 .map_err(|e| format!("{materialization_label}: {e}"))?
244 .as_ref()
245 .clone(),
246 )),
247 }
248}
249
250pub(crate) fn dense_locscale_block_designs_fromspecs<'a>(
257 specs: &'a [ParameterBlockSpec],
258 expected_count: usize,
259 family_name: &str,
260 short_family_name: &str,
261 primary_block_idx: usize,
262 log_sigma_block_idx: usize,
263 primary_label: &str,
264 material_policy: &gam_runtime::resource::MaterializationPolicy,
265) -> Result<(Cow<'a, Array2<f64>>, Cow<'a, Array2<f64>>), String> {
266 if specs.len() != expected_count {
267 return Err(GamlssError::DimensionMismatch {
268 reason: format!(
269 "{family_name} expects {expected_count} specs, got {}",
270 specs.len()
271 ),
272 }
273 .into());
274 }
275 let primary = dense_block_from_spec(
276 &specs[primary_block_idx],
277 material_policy,
278 &format!("{short_family_name} dense_block_designs_fromspecs {primary_label}"),
279 )?;
280 let log_sigma = dense_block_from_spec(
281 &specs[log_sigma_block_idx],
282 material_policy,
283 &format!("{short_family_name} dense_block_designs_fromspecs log_sigma"),
284 )?;
285 Ok((primary, log_sigma))
286}
287
288pub(crate) fn dense_locscale_block_designs_cached<'a>(
295 primary_design: Option<&'a DesignMatrix>,
296 log_sigma_design: Option<&'a DesignMatrix>,
297 family_name: &str,
298 short_family_name: &str,
299 primary_label: &str,
300 material_policy: &gam_runtime::resource::MaterializationPolicy,
301) -> Result<(Cow<'a, Array2<f64>>, Cow<'a, Array2<f64>>), String> {
302 let primary_design = primary_design
303 .ok_or_else(|| format!("{family_name} exact path is missing {primary_label} design"))?;
304 let log_sigma_design = log_sigma_design
305 .ok_or_else(|| format!("{family_name} exact path is missing log-sigma design"))?;
306 let primary = match primary_design.as_dense_ref() {
307 Some(d) => Cow::Borrowed(d),
308 None => Cow::Owned(
309 primary_design
310 .try_to_dense_with_policy(material_policy, "gamlss dense_locscale_block_designs")
311 .map_err(|e| {
312 format!("{short_family_name} dense_block_designs {primary_label}: {e}")
313 })?
314 .as_ref()
315 .clone(),
316 ),
317 };
318 let log_sigma = match log_sigma_design.as_dense_ref() {
319 Some(d) => Cow::Borrowed(d),
320 None => Cow::Owned(
321 log_sigma_design
322 .try_to_dense_with_policy(material_policy, "gamlss dense_locscale_block_designs")
323 .map_err(|e| format!("{short_family_name} dense_block_designs log_sigma: {e}"))?
324 .as_ref()
325 .clone(),
326 ),
327 };
328 Ok((primary, log_sigma))
329}
330
331pub(crate) struct LocScalePsiDirectionParts {
336 pub(crate) block_idx: usize,
337 pub(crate) local_idx: usize,
338 pub(crate) primary_psi: PsiDesignMap,
339 pub(crate) log_sigma_psi: PsiDesignMap,
340 pub(crate) primary_z: Array1<f64>,
341 pub(crate) log_sigma_z: Array1<f64>,
342}
343
344pub(crate) fn locscale_joint_psi_direction_parts(
354 block_states: &[ParameterBlockState],
355 derivative_blocks: &[Vec<crate::custom_family::CustomFamilyBlockPsiDerivative>],
356 psi_index: usize,
357 n: usize,
358 p_primary: usize,
359 p_log_sigma: usize,
360 primary_block_idx: usize,
361 log_sigma_block_idx: usize,
362 expected_blocks: usize,
363 family_name: &str,
364 primary_label: &str,
365 policy: &gam_runtime::resource::ResourcePolicy,
366) -> Result<Option<LocScalePsiDirectionParts>, String> {
367 validate_block_count::<GamlssError>(family_name, expected_blocks, block_states.len())?;
368 if derivative_blocks.len() != expected_blocks {
369 return Err(GamlssError::DimensionMismatch {
370 reason: format!(
371 "{family_name} joint psi direction expects {expected_blocks} derivative block lists, got {}",
372 derivative_blocks.len()
373 ),
374 }
375 .into());
376 }
377 let beta_primary = &block_states[primary_block_idx].beta;
378 let beta_log_sigma = &block_states[log_sigma_block_idx].beta;
379
380 let mut global = 0usize;
381 for (block_idx, block_derivs) in derivative_blocks.iter().enumerate() {
382 for (local_idx, deriv) in block_derivs.iter().enumerate() {
383 if global == psi_index {
384 let primary_psi;
385 let log_sigma_psi;
386 let primary_z;
387 let log_sigma_z;
388 if block_idx == primary_block_idx {
389 primary_psi = resolve_custom_family_x_psi_map(
390 deriv,
391 n,
392 p_primary,
393 0..n,
394 &format!("{family_name} {primary_label}"),
395 policy,
396 )?;
397 primary_z = primary_psi
398 .forward_mul(beta_primary.view())
399 .map_err(|e| format!("{family_name} {primary_label} forward_mul: {e}"))?;
400 log_sigma_psi = PsiDesignMap::Zero {
401 nrows: n,
402 ncols: p_log_sigma,
403 };
404 log_sigma_z = Array1::<f64>::zeros(n);
405 } else if block_idx == log_sigma_block_idx {
406 log_sigma_psi = resolve_custom_family_x_psi_map(
407 deriv,
408 n,
409 p_log_sigma,
410 0..n,
411 &format!("{family_name} log-sigma"),
412 policy,
413 )?;
414 log_sigma_z = log_sigma_psi
415 .forward_mul(beta_log_sigma.view())
416 .map_err(|e| format!("{family_name} log-sigma forward_mul: {e}"))?;
417 primary_psi = PsiDesignMap::Zero {
418 nrows: n,
419 ncols: p_primary,
420 };
421 primary_z = Array1::<f64>::zeros(n);
422 } else {
423 return Ok(None);
424 }
425 return Ok(Some(LocScalePsiDirectionParts {
426 block_idx,
427 local_idx,
428 primary_psi,
429 log_sigma_psi,
430 primary_z,
431 log_sigma_z,
432 }));
433 }
434 global += 1;
435 }
436 }
437 Ok(None)
438}
439
440pub(crate) struct LocScalePsiDriftConfig<'a> {
445 pub(crate) n: usize,
446 pub(crate) p_primary: usize,
447 pub(crate) p_log_sigma: usize,
448 pub(crate) primary_block_idx: usize,
449 pub(crate) log_sigma_block_idx: usize,
450 pub(crate) family_name: &'a str,
451 pub(crate) primary_label: &'a str,
452 pub(crate) policy: &'a gam_runtime::resource::ResourcePolicy,
453}
454
455pub(crate) fn locscale_joint_psisecond_design_drifts(
456 block_states: &[ParameterBlockState],
457 derivative_blocks: &[Vec<crate::custom_family::CustomFamilyBlockPsiDerivative>],
458 psi_a: &LocationScaleJointPsiDirection,
459 psi_b: &LocationScaleJointPsiDirection,
460 cfg: LocScalePsiDriftConfig<'_>,
461) -> Result<LocationScaleJointPsiSecondDrifts, String> {
462 let beta_primary = &block_states[cfg.primary_block_idx].beta;
463 let beta_log_sigma = &block_states[cfg.log_sigma_block_idx].beta;
464 let mut primary_ab_action = None;
465 let mut log_sigma_ab_action = None;
466 let mut primary_ab = None;
467 let mut log_sigma_ab = None;
468
469 if psi_a.block_idx == psi_b.block_idx {
473 let deriv = &derivative_blocks[psi_a.block_idx][psi_a.local_idx];
474 let deriv_b = &derivative_blocks[psi_b.block_idx][psi_b.local_idx];
475 if psi_a.block_idx == cfg.primary_block_idx {
476 let (action, matrix) = psi_psi_map_to_drift_slots(
477 deriv,
478 deriv_b,
479 psi_b.local_idx,
480 cfg.n,
481 cfg.p_primary,
482 &format!("{} {}", cfg.family_name, cfg.primary_label),
483 cfg.policy,
484 )?;
485 primary_ab_action = action;
486 primary_ab = matrix;
487 } else if psi_a.block_idx == cfg.log_sigma_block_idx {
488 let (action, matrix) = psi_psi_map_to_drift_slots(
489 deriv,
490 deriv_b,
491 psi_b.local_idx,
492 cfg.n,
493 cfg.p_log_sigma,
494 &format!("{} log-sigma", cfg.family_name),
495 cfg.policy,
496 )?;
497 log_sigma_ab_action = action;
498 log_sigma_ab = matrix;
499 }
500 }
501
502 let z_primary_ab = second_psi_linear_map(
503 primary_ab_action.as_ref(),
504 primary_ab.as_ref(),
505 cfg.n,
506 cfg.p_primary,
507 )
508 .forward_mul(beta_primary.view());
509 let z_ls_ab = second_psi_linear_map(
510 log_sigma_ab_action.as_ref(),
511 log_sigma_ab.as_ref(),
512 cfg.n,
513 cfg.p_log_sigma,
514 )
515 .forward_mul(beta_log_sigma.view());
516
517 Ok(LocationScaleJointPsiSecondDrifts {
518 x_primary_ab_action: primary_ab_action,
519 x_ls_ab_action: log_sigma_ab_action,
520 x_primary_ab: primary_ab,
521 x_ls_ab: log_sigma_ab,
522 z_primary_ab,
523 z_ls_ab,
524 })
525}
526
527pub(crate) fn psi_psi_map_to_drift_slots(
528 deriv: &crate::custom_family::CustomFamilyBlockPsiDerivative,
529 deriv_b: &crate::custom_family::CustomFamilyBlockPsiDerivative,
530 local_idx_b: usize,
531 n: usize,
532 p: usize,
533 label: &str,
534 policy: &gam_runtime::resource::ResourcePolicy,
535) -> Result<
536 (
537 Option<crate::custom_family::CustomFamilyPsiSecondDesignAction>,
538 Option<Array2<f64>>,
539 ),
540 String,
541> {
542 match resolve_custom_family_x_psi_psi_map(
543 deriv,
544 deriv_b,
545 local_idx_b,
546 n,
547 p,
548 0..n,
549 label,
550 policy,
551 )? {
552 crate::custom_family::PsiDesignMap::Second { action } => Ok((Some(action), None)),
553 crate::custom_family::PsiDesignMap::Dense { matrix } => Ok((None, Some((*matrix).clone()))),
554 crate::custom_family::PsiDesignMap::Zero { .. } => Ok((None, None)),
555 crate::custom_family::PsiDesignMap::First { .. } => {
556 Err(GamlssError::UnsupportedConfiguration {
557 reason: format!("{label}: unexpected First variant from _psi_psi_map"),
558 }
559 .into())
560 }
561 }
562}
563
564pub(crate) fn dense_block_or_operator<'a>(
565 design: &'a DesignMatrix,
566 n: usize,
567 p: usize,
568 budget_bytes: usize,
569 policy: &gam_runtime::resource::ResourcePolicy,
570) -> DenseOrOperator<'a> {
571 if let Some(dense) = design.as_dense_ref() {
572 return DenseOrOperator::Borrowed(dense);
573 }
574
575 let dense_bytes = 8usize.saturating_mul(n).saturating_mul(p);
576 if dense_bytes <= budget_bytes
577 && let Ok(arc) = design
578 .try_to_dense_with_policy(&policy.material_policy(), "gamlss dense_block_or_operator")
579 {
580 return DenseOrOperator::Owned(arc.as_ref().clone());
581 }
582
583 DenseOrOperator::Operator(design.clone())
584}
585
586pub(crate) fn dense_blocks_planned_budget(blocks: &[&DesignMatrix]) -> Vec<usize> {
587 let mut planned = vec![0; blocks.len()];
588 let mut total = 0usize;
589 for (idx, design) in blocks.iter().enumerate() {
590 if design.as_dense_ref().is_some() {
591 continue;
592 }
593 let bytes = 8usize
594 .saturating_mul(design.nrows())
595 .saturating_mul(design.ncols());
596 if bytes <= EXACT_DENSE_BLOCK_BUDGET_BYTES
597 && total.saturating_add(bytes) <= EXACT_DENSE_TOTAL_BUDGET_BYTES
598 {
599 planned[idx] = bytes;
600 total += bytes;
601 }
602 }
603 planned
604}
605
606pub(crate) fn exact_design_row_chunks(
607 n: usize,
608 p: usize,
609) -> impl Iterator<Item = std::ops::Range<usize>> {
610 const TARGET_BYTES: usize = 8 * 1024 * 1024;
611 const MIN_ROWS: usize = 512;
612 const MAX_ROWS: usize = 131_072;
613 let rows = (TARGET_BYTES / (p.max(1) * 8))
614 .clamp(MIN_ROWS, MAX_ROWS)
615 .min(n.max(1));
616 (0..n)
617 .step_by(rows)
618 .map(move |start| start..(start + rows).min(n))
619}
620
621pub(crate) fn design_weighted_column_squares(
622 design: &DesignMatrix,
623 weights: &Array1<f64>,
624) -> Result<Array1<f64>, String> {
625 let n = design.nrows();
626 let p = design.ncols();
627 if weights.len() != n {
628 return Err(GamlssError::DimensionMismatch {
629 reason: format!(
630 "design weighted column squares dimension mismatch: weights={}, rows={}",
631 weights.len(),
632 n
633 ),
634 }
635 .into());
636 }
637 let mut out = Array1::<f64>::zeros(p);
638 for rows in exact_design_row_chunks(n, p) {
639 let chunk = design.try_row_chunk(rows.clone()).map_err(|e| {
640 format!("design weighted column squares row chunk materialization failed: {e}")
641 })?;
642 for (local_i, row) in chunk.outer_iter().enumerate() {
643 let w = weights[rows.start + local_i];
644 if w == 0.0 {
645 continue;
646 }
647 for j in 0..p {
648 let x = row[j];
649 out[j] += w * x * x;
650 }
651 }
652 }
653 Ok(out)
654}
655
656#[inline]
657pub(crate) fn floor_positiveweight(rawweight: f64, minweight: f64) -> f64 {
658 if !rawweight.is_finite() || rawweight <= 0.0 {
659 0.0
660 } else {
661 rawweight.max(minweight)
662 }
663}
664
665#[inline]
666pub(crate) fn logb_dlog_sigma_deta(sigma: f64, d_sigma_deta: f64) -> f64 {
667 if d_sigma_deta.is_infinite() {
668 1.0
669 } else {
670 let value = d_sigma_deta / sigma;
671 if value.is_finite() {
672 value.clamp(0.0, 1.0)
673 } else {
674 0.0
675 }
676 }
677}
678
679#[inline]
680pub(crate) fn gaussian_log_sigma_irlsinfo_directional_derivative(
681 weight: f64,
682 sigma: f64,
683 d_sigma_deta: f64,
684 d_eta: f64,
685) -> f64 {
686 if weight == 0.0 || d_eta == 0.0 || !sigma.is_finite() || sigma <= 0.0 {
687 return 0.0;
688 }
689 let g = logb_dlog_sigma_deta(sigma, d_sigma_deta);
695 if !g.is_finite() || !(0.0..1.0).contains(&g) {
696 return 0.0;
697 }
698 let rawinfo = 2.0 * weight * g * g;
699 if !rawinfo.is_finite() || rawinfo <= MIN_WEIGHT {
700 return 0.0;
701 }
702 let dg_deta = g * (1.0 - g);
703 let dw = 4.0 * weight * g * dg_deta * d_eta;
704 if dw.is_finite() { dw } else { 0.0 }
705}
706
707#[derive(Clone, Copy)]
708pub(crate) struct GaussianDiagonalRowKernel {
709 pub(crate) log_likelihood: f64,
710 pub(crate) location_working_weight: f64,
711 pub(crate) location_working_shift: f64,
712 pub(crate) log_sigma_working_weight: f64,
713 pub(crate) log_sigma_working_response: f64,
714}
715
716#[inline]
717pub(crate) fn gaussian_diagonal_row_kernel(
718 y: f64,
719 location_eta: f64,
720 eta_log_sigma: f64,
721 obs_weight: f64,
722 ln2pi: f64,
723) -> GaussianDiagonalRowKernel {
724 if obs_weight == 0.0 {
725 return GaussianDiagonalRowKernel {
726 log_likelihood: 0.0,
727 location_working_weight: 0.0,
728 location_working_shift: 0.0,
729 log_sigma_working_weight: 0.0,
730 log_sigma_working_response: eta_log_sigma,
731 };
732 }
733
734 let SigmaJet1 { sigma, d1 } = logb_sigma_jet1_scalar(eta_log_sigma);
742 let inv_s2 = (sigma * sigma).recip();
743 let residual = y - location_eta;
744 let location_working_weight = floor_positiveweight(obs_weight * inv_s2, MIN_WEIGHT);
745 let dlog_sigma_deta = logb_dlog_sigma_deta(sigma, d1);
752 let log_sigma_working_weight = floor_positiveweight(
753 2.0 * obs_weight * dlog_sigma_deta * dlog_sigma_deta,
754 MIN_WEIGHT,
755 );
756 let log_sigma_score = obs_weight * (residual * residual * inv_s2 - 1.0) * dlog_sigma_deta;
757 let log_sigma_working_response = if log_sigma_working_weight == 0.0 {
758 eta_log_sigma
759 } else {
760 eta_log_sigma + log_sigma_score / log_sigma_working_weight
761 };
762
763 GaussianDiagonalRowKernel {
764 log_likelihood: obs_weight
765 * (-0.5 * (residual * residual * inv_s2 + ln2pi + 2.0 * sigma.ln())),
766 location_working_weight,
767 location_working_shift: residual,
768 log_sigma_working_weight,
769 log_sigma_working_response,
770 }
771}
772
773#[derive(Clone, Copy, Debug, PartialEq, Eq)]
775pub enum ParameterLink {
776 Identity,
777 Log,
778 Logit,
779 Probit,
780 InverseLink,
781 Wiggle,
783}