1use super::weighted_design_products::{mirror_upper_to_lower, xt_diag_x_design, xt_diag_y_design};
7use super::{
10 BlockwiseTermFitResult, GamlssLambdaLayout, LOCATION_SCALE_N_OUTPUTS,
11 LocationScaleFamilyBuilder, build_location_scale_block, fit_location_scale_terms,
12 solve_penalizedweighted_projection, spatial_length_scale_term_indices,
13};
14use crate::block_layout::block_count::validate_block_count;
15use crate::custom_family::{
16 BlockWorkingSet, BlockwiseFitOptions, CustomFamily, CustomFamilyBlockPsiDerivative,
17 FamilyEvaluation, ParameterBlockSpec, ParameterBlockState,
18};
19use crate::gamlss::GamlssError;
20use crate::model_types::UnifiedFitResult;
21use gam_linalg::matrix::LinearOperator;
22use gam_math::jet_scalar::JetScalar;
23use gam_math::nested_dual::JetField;
24use gam_terms::smooth::{
25 SpatialLengthScaleOptimizationOptions, TermCollectionDesign, TermCollectionSpec,
26 get_spatial_length_scale, spatial_term_uses_per_axis_psi,
27};
28use ndarray::{Array1, Array2, s};
29use statrs::function::gamma::ln_gamma;
30
31#[derive(Clone, Copy, Debug, PartialEq)]
81pub enum DispersionFamilyKind {
82 NegativeBinomial,
84 Gamma,
86 Beta,
89 Tweedie { p: f64 },
96}
97
98impl DispersionFamilyKind {
99 pub const fn family_tag(self) -> &'static str {
100 match self {
101 DispersionFamilyKind::NegativeBinomial => FAMILY_NEGBIN_LOCATION_SCALE,
102 DispersionFamilyKind::Gamma => FAMILY_GAMMA_LOCATION_SCALE,
103 DispersionFamilyKind::Beta => FAMILY_BETA_LOCATION_SCALE,
104 DispersionFamilyKind::Tweedie { .. } => FAMILY_TWEEDIE_LOCATION_SCALE,
105 }
106 }
107
108 pub(crate) const fn mean_is_logit(self) -> bool {
110 matches!(self, DispersionFamilyKind::Beta)
111 }
112
113 pub fn base_link(self) -> gam_problem::InverseLink {
118 use gam_problem::{InverseLink, StandardLink};
119 if self.mean_is_logit() {
120 InverseLink::Standard(StandardLink::Logit)
121 } else {
122 InverseLink::Standard(StandardLink::Log)
123 }
124 }
125
126 pub fn likelihood_spec(self) -> gam_problem::LikelihoodSpec {
134 use gam_problem::{InverseLink, LikelihoodSpec, ResponseFamily, StandardLink};
135 let response = match self {
136 DispersionFamilyKind::NegativeBinomial => ResponseFamily::NegativeBinomial {
137 theta: 1.0,
138 theta_fixed: false,
139 },
140 DispersionFamilyKind::Gamma => ResponseFamily::Gamma,
141 DispersionFamilyKind::Beta => ResponseFamily::Beta { phi: 1.0 },
142 DispersionFamilyKind::Tweedie { p } => ResponseFamily::Tweedie { p },
143 };
144 let link = if self.mean_is_logit() {
145 InverseLink::Standard(StandardLink::Logit)
146 } else {
147 InverseLink::Standard(StandardLink::Log)
148 };
149 LikelihoodSpec::new(response, link)
150 }
151}
152
153pub const FAMILY_NEGBIN_LOCATION_SCALE: &str = "negbin-location-scale";
154pub const FAMILY_GAMMA_LOCATION_SCALE: &str = "gamma-location-scale";
155pub const FAMILY_BETA_LOCATION_SCALE: &str = "beta-location-scale";
156pub const FAMILY_TWEEDIE_LOCATION_SCALE: &str = "tweedie-location-scale";
157
158const DISPERSION_PARALLEL_ROW_THRESHOLD: usize = 1024;
164
165pub(super) struct DispersionRowKernel {
167 pub(super) loglik: f64,
168 pub(super) mean_weight: f64,
169 pub(super) mean_response: f64,
170 pub(super) disp_weight: f64,
171 pub(super) disp_response: f64,
172}
173
174#[inline]
175fn dispersion_geometry_error(row: usize, quantity: &'static str, eta: f64, value: f64) -> String {
176 GamlssError::RowGeometryUnrepresentable {
177 row,
178 quantity,
179 eta,
180 value,
181 }
182 .into()
183}
184
185fn validate_dispersion_row_geometry_inputs(
189 kind: DispersionFamilyKind,
190 row: usize,
191 yi: f64,
192 eta_mu: f64,
193 eta_d: f64,
194 prior_weight: f64,
195) -> Result<(), String> {
196 if !eta_mu.is_finite() || !eta_d.is_finite() {
197 return Err(GamlssError::NonFinite {
198 reason: format!(
199 "{} requires finite predictors at row {row}; eta_mu={eta_mu}, eta_d={eta_d}",
200 kind.family_tag()
201 ),
202 }
203 .into());
204 }
205 if !prior_weight.is_finite() || prior_weight < 0.0 {
206 return Err(GamlssError::InvalidInput {
207 reason: format!(
208 "{} requires finite non-negative prior weights; weight[{row}]={prior_weight}",
209 kind.family_tag()
210 ),
211 }
212 .into());
213 }
214 if prior_weight == 0.0 {
215 return Ok(());
216 }
217 let (support_ok, support) = match kind {
218 DispersionFamilyKind::NegativeBinomial => (
219 yi.is_finite() && yi >= 0.0 && yi.fract() == 0.0,
220 "a finite non-negative integer",
221 ),
222 DispersionFamilyKind::Gamma => (yi.is_finite() && yi > 0.0, "finite and > 0"),
223 DispersionFamilyKind::Beta => (
224 yi.is_finite() && yi > 0.0 && yi < 1.0,
225 "finite and strictly inside (0, 1)",
226 ),
227 DispersionFamilyKind::Tweedie { p } => (
228 yi.is_finite() && yi >= 0.0 && p.is_finite() && p > 1.0 && p < 2.0,
229 "finite and >= 0 with power strictly inside (1, 2)",
230 ),
231 };
232 if !support_ok {
233 return Err(GamlssError::InvalidInput {
234 reason: format!(
235 "{} response outside support at row {row}: y={yi} (requires {support})",
236 kind.family_tag()
237 ),
238 }
239 .into());
240 }
241
242 let require_positive = |quantity, eta, value: f64| {
243 if value.is_finite() && value > 0.0 {
244 Ok(())
245 } else {
246 Err(dispersion_geometry_error(row, quantity, eta, value))
247 }
248 };
249 match kind {
250 DispersionFamilyKind::NegativeBinomial => {
251 let mu = eta_mu.exp();
252 let theta = eta_d.exp();
253 require_positive("negative-binomial mean exp(eta_mu)", eta_mu, mu)?;
254 require_positive("negative-binomial precision exp(eta_d)", eta_d, theta)
255 }
256 DispersionFamilyKind::Gamma => {
257 require_positive("Gamma mean exp(eta_mu)", eta_mu, eta_mu.exp())?;
258 require_positive("Gamma precision exp(eta_d)", eta_d, eta_d.exp())
259 }
260 DispersionFamilyKind::Beta => {
261 let mu = gam_linalg::utils::stable_logistic(eta_mu);
262 if !mu.is_finite() || mu <= 0.0 || mu >= 1.0 {
263 return Err(dispersion_geometry_error(
264 row,
265 "Beta mean logistic(eta_mu) in the open unit interval",
266 eta_mu,
267 mu,
268 ));
269 }
270 let phi = eta_d.exp();
271 require_positive("Beta precision exp(eta_d)", eta_d, phi)?;
272 require_positive("Beta first shape mu*phi", eta_mu, mu * phi)?;
273 require_positive("Beta second shape (1-mu)*phi", eta_mu, (1.0 - mu) * phi)
274 }
275 DispersionFamilyKind::Tweedie { .. } => {
276 require_positive("Tweedie mean exp(eta_mu)", eta_mu, eta_mu.exp())?;
277 require_positive("Tweedie dispersion exp(-eta_d)", eta_d, (-eta_d).exp())
278 }
279 }
280}
281
282fn validate_dispersion_row_kernel_output(
283 row: usize,
284 eta_mu: f64,
285 eta_d: f64,
286 prior_weight: f64,
287 output: &DispersionRowKernel,
288) -> Result<(), String> {
289 if prior_weight == 0.0 {
290 return Ok(());
291 }
292 for (quantity, eta, value, strictly_positive) in [
293 (
294 "dispersion-family row log likelihood",
295 eta_mu,
296 output.loglik,
297 false,
298 ),
299 (
300 "dispersion-family mean working weight",
301 eta_mu,
302 output.mean_weight,
303 true,
304 ),
305 (
306 "dispersion-family mean working response",
307 eta_mu,
308 output.mean_response,
309 false,
310 ),
311 (
312 "dispersion-family precision working weight",
313 eta_d,
314 output.disp_weight,
315 true,
316 ),
317 (
318 "dispersion-family precision working response",
319 eta_d,
320 output.disp_response,
321 false,
322 ),
323 ] {
324 if !value.is_finite() || (strictly_positive && value <= 0.0) {
325 return Err(dispersion_geometry_error(row, quantity, eta, value));
326 }
327 }
328 Ok(())
329}
330
331#[cfg(test)]
332mod test_support {
333 use super::*;
334
335 #[inline]
338 pub(super) fn dispersion_nb_nll_generic<S: gam_math::jet_scalar::JetScalar<2>>(
339 yi: f64,
340 mu_value: f64,
341 theta_value: f64,
342 wi: f64,
343 ) -> S {
344 let mu = S::variable(mu_value, 0);
345 let theta = S::variable(theta_value, 1);
346 let tpm = theta.add(&mu);
347 let loglik = theta
350 .add(&S::constant(yi))
351 .ln_gamma()
352 .sub(&theta.ln_gamma())
353 .sub(&S::constant(ln_gamma(yi + 1.0)))
354 .add(&theta.mul(&theta.ln()))
355 .sub(&theta.mul(&tpm.ln()))
356 .add(&mu.ln().scale(yi))
357 .sub(&tpm.ln().scale(yi));
358 loglik.scale(-wi)
359 }
360
361 #[inline]
364 pub(super) fn dispersion_gamma_nll_generic<S: gam_math::jet_scalar::JetScalar<2>>(
365 yi: f64,
366 y_pos: f64,
367 mu_value: f64,
368 nu_value: f64,
369 wi: f64,
370 ) -> S {
371 let mu = S::variable(mu_value, 0);
372 let nu = S::variable(nu_value, 1);
373 let loglik = nu
375 .mul(&nu.ln())
376 .sub(&nu.mul(&mu.ln()))
377 .sub(&nu.ln_gamma())
378 .add(&nu.sub(&S::constant(1.0)).scale(y_pos.ln()))
379 .sub(&nu.mul(&mu.recip().scale(yi)));
380 loglik.scale(-wi)
381 }
382
383 #[inline]
386 pub(super) fn dispersion_beta_nll_generic<S: gam_math::jet_scalar::JetScalar<2>>(
387 yi: f64,
388 mu_value: f64,
389 phi_value: f64,
390 wi: f64,
391 ) -> S {
392 let mu = S::variable(mu_value, 0);
393 let phi = S::variable(phi_value, 1);
394 let one_minus_mu = S::constant(1.0).sub(&mu);
395 let yc = yi;
396 let a = mu.mul(&phi);
397 let b = one_minus_mu.mul(&phi);
398 let loglik = phi
401 .ln_gamma()
402 .sub(&a.ln_gamma())
403 .sub(&b.ln_gamma())
404 .add(&a.sub(&S::constant(1.0)).scale(yc.ln()))
405 .add(&b.sub(&S::constant(1.0)).scale((-yc).ln_1p()));
406 loglik.scale(-wi)
407 }
408
409 #[inline]
418 pub(super) fn dispersion_nb_nll_order2(
419 yi: f64,
420 mu_value: f64,
421 theta_value: f64,
422 wi: f64,
423 ) -> gam_math::jet_scalar::Order2<2> {
424 type O2 = gam_math::jet_scalar::Order2<2>;
425
426 let mu = O2::variable(mu_value, 0);
427 let theta = O2::variable(theta_value, 1);
428 let tpm = theta.add(&mu);
429 let theta_plus_y = theta.add(&O2::constant(yi));
430 let loglik = order2_ln_gamma(&theta_plus_y)
431 .sub(&order2_ln_gamma(&theta))
432 .sub(&O2::constant(ln_gamma(yi + 1.0)))
433 .add(&theta.mul(&theta.ln()))
434 .sub(&theta.mul(&tpm.ln()))
435 .add(&mu.ln().scale(yi))
436 .sub(&tpm.ln().scale(yi));
437 loglik.scale(-wi)
438 }
439
440 #[inline]
445 pub(super) fn dispersion_gamma_nll_order2(
446 yi: f64,
447 y_pos: f64,
448 mu_value: f64,
449 nu_value: f64,
450 wi: f64,
451 ) -> gam_math::jet_scalar::Order2<2> {
452 type O2 = gam_math::jet_scalar::Order2<2>;
453
454 let mu = O2::variable(mu_value, 0);
455 let nu = O2::variable(nu_value, 1);
456 let loglik = nu
457 .mul(&nu.ln())
458 .sub(&nu.mul(&mu.ln()))
459 .sub(&order2_ln_gamma(&nu))
460 .add(&nu.sub(&O2::constant(1.0)).scale(y_pos.ln()))
461 .sub(&nu.mul(&mu.recip().scale(yi)));
462 loglik.scale(-wi)
463 }
464}
465
466#[inline]
469pub(crate) fn dispersion_beta_nll_order2(
470 yi: f64,
471 mu_value: f64,
472 phi_value: f64,
473 wi: f64,
474) -> gam_math::jet_scalar::Order2<2> {
475 type O2 = gam_math::jet_scalar::Order2<2>;
476
477 let mu = O2::variable(mu_value, 0);
478 let phi = O2::variable(phi_value, 1);
479 let one_minus_mu = O2::constant(1.0).sub(&mu);
480 let yc = yi;
481 let a = mu.mul(&phi);
482 let b = one_minus_mu.mul(&phi);
483 let loglik = order2_ln_gamma(&phi)
484 .sub(&order2_ln_gamma(&a))
485 .sub(&order2_ln_gamma(&b))
486 .add(&a.sub(&O2::constant(1.0)).scale(yc.ln()))
487 .add(&b.sub(&O2::constant(1.0)).scale((-yc).ln_1p()));
488 loglik.scale(-wi)
489}
490
491#[inline]
492fn order2_ln_gamma<const K: usize>(
493 x: &gam_math::jet_scalar::Order2<K>,
494) -> gam_math::jet_scalar::Order2<K> {
495 gam_math::jet_scalar::Order2(
496 x.0.compose_unary(gam_math::jet_tower::ln_gamma_derivative_stack_order2(x.0.v)),
497 )
498}
499
500#[inline]
531pub(crate) fn dispersion_gamma_disp_order2(
532 yi: f64,
533 y_pos: f64,
534 mu_value: f64,
535 nu_value: f64,
536 wi: f64,
537) -> gam_math::jet_scalar::Order2<1> {
538 type O1 = gam_math::jet_scalar::Order2<1>;
539
540 let mu = O1::constant(mu_value);
541 let nu = O1::variable(nu_value, 0);
542 let loglik = nu
543 .mul(&nu.ln())
544 .sub(&nu.mul(&mu.ln()))
545 .sub(&order2_ln_gamma(&nu))
546 .add(&nu.sub(&O1::constant(1.0)).scale(y_pos.ln()))
547 .sub(&nu.mul(&mu.recip().scale(yi)));
548 loglik.scale(-wi)
549}
550
551#[inline]
557pub(crate) fn dispersion_tweedie_disp_order2(
558 yi: f64,
559 eta_mu: f64,
560 eta_d: f64,
561 p: f64,
562 wi: f64,
563) -> gam_math::jet_scalar::Order2<1> {
564 type O1 = gam_math::jet_scalar::Order2<1>;
565
566 let one_minus_p = 1.0 - p;
567 let two_minus_p = 2.0 - p;
568 let mu = O1::constant(eta_mu).exp();
569 let phi = O1::variable(eta_d, 0).scale(-1.0).exp();
570 if yi > 0.0 {
571 let dev = mu
572 .powf(two_minus_p)
573 .scale(1.0 / two_minus_p)
574 .sub(&mu.powf(one_minus_p).scale(yi / one_minus_p))
575 .add(&O1::constant(
576 yi.powf(two_minus_p) / (one_minus_p * two_minus_p),
577 ))
578 .scale(2.0);
579 let loglik = dev
580 .mul(&phi.recip().scale(-0.5))
581 .sub(&phi.scale(2.0 * std::f64::consts::PI).ln().scale(0.5))
582 .sub(&O1::constant(0.5 * p * yi.ln()));
583 loglik.scale(-wi)
584 } else {
585 let c = mu.powf(two_minus_p).scale(1.0 / two_minus_p);
586 let loglik = c.mul(&phi.recip()).scale(-1.0);
587 loglik.scale(-wi)
588 }
589}
590
591#[inline]
607fn dispersion_nb_loglik(yi: f64, mu: f64, theta: f64, wi: f64) -> f64 {
608 let log_theta_share = log_positive_share(theta, mu);
609 let log_mu_share = log_positive_share(mu, theta);
610 let s = ln_gamma(theta + yi) - ln_gamma(theta) - ln_gamma(yi + 1.0)
611 + theta * log_theta_share
612 + yi * log_mu_share;
613 -(s * -wi)
614}
615
616#[inline]
619fn log_positive_share(numerator: f64, other: f64) -> f64 {
620 if numerator >= other {
621 -(other / numerator).ln_1p()
622 } else {
623 let ratio = numerator / other;
624 numerator.ln() - other.ln() - ratio.ln_1p()
625 }
626}
627
628#[inline]
629fn positive_share(numerator: f64, other: f64) -> f64 {
630 if numerator >= other {
631 1.0 / (1.0 + other / numerator)
632 } else {
633 let ratio = numerator / other;
634 ratio / (1.0 + ratio)
635 }
636}
637
638#[inline]
643fn nb_log_precision_fisher_jensen(mu: f64, theta: f64) -> f64 {
644 let r = positive_share(theta, mu);
645 let q = positive_share(mu, theta);
646 if theta <= 32.0 {
647 let total = theta + mu;
648 let remainder_theta = gam_math::jet_tower::trigamma(theta) - theta.recip();
649 let remainder_total = gam_math::jet_tower::trigamma(total) - total.recip();
650 return theta * theta * (remainder_theta - remainder_total);
651 }
652 let one_minus_r2 = q * (1.0 + r);
653 let r2 = r * r;
654 let one_minus_r3 = q * (1.0 + r + r2);
655 let r4 = r2 * r2;
656 let one_minus_r5 = q * (1.0 + r + r2 + r2 * r + r4);
657 let r6 = r4 * r2;
658 let one_minus_r7 = q * (1.0 + r + r2 + r2 * r + r4 + r4 * r + r6);
659 let inv = theta.recip();
660 let inv2 = inv * inv;
661 0.5 * one_minus_r2 + (inv / 6.0) * one_minus_r3 - (inv * inv2 / 30.0) * one_minus_r5
662 + (inv * inv2 * inv2 / 42.0) * one_minus_r7
663}
664
665#[inline]
668fn dispersion_gamma_loglik(yi: f64, y_pos: f64, mu: f64, nu: f64, wi: f64) -> f64 {
669 let s = nu * nu.ln() - nu * mu.ln() - ln_gamma(nu) + (nu - 1.0) * y_pos.ln()
673 - nu * ((1.0 / mu) * yi);
674 -(s * -wi)
675}
676
677#[inline]
680fn dispersion_beta_loglik(yi: f64, mu: f64, phi: f64, wi: f64) -> f64 {
681 let one_minus_mu = 1.0 - mu;
682 let yc = yi;
683 let a = mu * phi;
684 let b = one_minus_mu * phi;
685 let s =
686 ln_gamma(phi) - ln_gamma(a) - ln_gamma(b) + (a - 1.0) * yc.ln() + (b - 1.0) * (-yc).ln_1p();
687 -(s * -wi)
688}
689
690#[inline]
693fn dispersion_tweedie_loglik(yi: f64, eta_mu: f64, eta_d: f64, p: f64, wi: f64) -> f64 {
694 let one_minus_p = 1.0 - p;
695 let two_minus_p = 2.0 - p;
696 let mu = eta_mu.exp();
697 let phi = (-eta_d).exp();
698 let s = if yi > 0.0 {
699 let dev = (mu.powf(two_minus_p) * (1.0 / two_minus_p)
700 - mu.powf(one_minus_p) * (yi / one_minus_p)
701 + yi.powf(two_minus_p) / (one_minus_p * two_minus_p))
702 * 2.0;
703 dev * ((1.0 / phi) * -0.5)
704 - (phi * (2.0 * std::f64::consts::PI)).ln() * 0.5
705 - 0.5 * p * yi.ln()
706 } else {
707 let c = mu.powf(two_minus_p) * (1.0 / two_minus_p);
708 (c * (1.0 / phi)) * -1.0
709 };
710 -(s * -wi)
711}
712
713#[inline]
719pub(crate) fn dispersion_row_loglik(
720 kind: DispersionFamilyKind,
721 yi: f64,
722 eta_mu: f64,
723 eta_d: f64,
724 prior_weight: f64,
725) -> f64 {
726 if prior_weight <= 0.0 {
730 return 0.0;
731 }
732 let wi = prior_weight;
733 let em = eta_mu;
734 let ed = eta_d;
735 match kind {
736 DispersionFamilyKind::NegativeBinomial => {
737 let mu = em.exp();
738 let theta = ed.exp();
739 dispersion_nb_loglik(yi, mu, theta, wi)
740 }
741 DispersionFamilyKind::Gamma => {
742 let mu = em.exp();
743 let nu = ed.exp();
744 let y_pos = yi;
745 dispersion_gamma_loglik(yi, y_pos, mu, nu, wi)
746 }
747 DispersionFamilyKind::Beta => {
748 let mu = gam_linalg::utils::stable_logistic(em);
749 let phi = ed.exp();
750 dispersion_beta_loglik(yi, mu, phi, wi)
751 }
752 DispersionFamilyKind::Tweedie { p } => dispersion_tweedie_loglik(yi, em, ed, p, wi),
753 }
754}
755
756pub(crate) fn dispersion_eta_nll_order2(
766 kind: DispersionFamilyKind,
767 yi: f64,
768 em: f64,
769 ed: f64,
770 wi: f64,
771) -> gam_math::jet_scalar::Order2<2> {
772 type O2 = gam_math::jet_scalar::Order2<2>;
773 let eta_mu = O2::variable(em, 0);
774 let eta_d = O2::variable(ed, 1);
775 match kind {
776 DispersionFamilyKind::NegativeBinomial => {
777 let theta = eta_d.exp();
781 let theta_plus_y = theta.add(&O2::constant(yi));
782 let log_total = if em >= ed {
783 eta_mu.add(&eta_d.sub(&eta_mu).exp().add(&O2::constant(1.0)).ln())
784 } else {
785 eta_d.add(&eta_mu.sub(&eta_d).exp().add(&O2::constant(1.0)).ln())
786 };
787 let loglik = order2_ln_gamma(&theta_plus_y)
788 .sub(&order2_ln_gamma(&theta))
789 .sub(&O2::constant(ln_gamma(yi + 1.0)))
790 .add(&theta.mul(&eta_d.sub(&log_total)))
791 .add(&eta_mu.sub(&log_total).scale(yi));
792 loglik.scale(-wi)
793 }
794 DispersionFamilyKind::Gamma => {
795 let mu = eta_mu.exp();
796 let nu = eta_d.exp();
797 let y_pos = yi;
798 let loglik = nu
799 .mul(&nu.ln())
800 .sub(&nu.mul(&mu.ln()))
801 .sub(&order2_ln_gamma(&nu))
802 .add(&nu.sub(&O2::constant(1.0)).scale(y_pos.ln()))
803 .sub(&nu.mul(&mu.recip().scale(yi)));
804 loglik.scale(-wi)
805 }
806 DispersionFamilyKind::Beta => {
807 let mu = eta_mu.scale(-1.0).exp().add(&O2::constant(1.0)).recip();
808 let phi = eta_d.exp();
809 let one_minus_mu = O2::constant(1.0).sub(&mu);
810 let yc = yi;
811 let a = mu.mul(&phi);
812 let b = one_minus_mu.mul(&phi);
813 let loglik = order2_ln_gamma(&phi)
814 .sub(&order2_ln_gamma(&a))
815 .sub(&order2_ln_gamma(&b))
816 .add(&a.sub(&O2::constant(1.0)).scale(yc.ln()))
817 .add(&b.sub(&O2::constant(1.0)).scale((-yc).ln_1p()));
818 loglik.scale(-wi)
819 }
820 DispersionFamilyKind::Tweedie { p } => {
821 let one_minus_p = 1.0 - p;
822 let two_minus_p = 2.0 - p;
823 let mu = eta_mu.exp();
824 let phi = eta_d.scale(-1.0).exp();
825 if yi > 0.0 {
826 let dev = mu
827 .powf(two_minus_p)
828 .scale(1.0 / two_minus_p)
829 .sub(&mu.powf(one_minus_p).scale(yi / one_minus_p))
830 .add(&O2::constant(
831 yi.powf(two_minus_p) / (one_minus_p * two_minus_p),
832 ))
833 .scale(2.0);
834 let loglik = dev
835 .mul(&phi.recip().scale(-0.5))
836 .sub(&phi.scale(2.0 * std::f64::consts::PI).ln().scale(0.5))
837 .sub(&O2::constant(0.5 * p * yi.ln()));
838 loglik.scale(-wi)
839 } else {
840 let c = mu.powf(two_minus_p).scale(1.0 / two_minus_p);
841 let loglik = c.mul(&phi.recip()).scale(-1.0);
842 loglik.scale(-wi)
843 }
844 }
845 }
846}
847
848pub(crate) fn dispersion_row_observed_hessian_weights(
851 kind: DispersionFamilyKind,
852 yi: f64,
853 eta_mu: f64,
854 eta_d: f64,
855 prior_weight: f64,
856) -> (f64, f64, f64) {
857 if prior_weight <= 0.0 {
858 return (0.0, 0.0, 0.0);
859 }
860 let tower = dispersion_eta_nll_order2(kind, yi, eta_mu, eta_d, prior_weight);
861 let h = tower.h();
862 (h[0][0], h[0][1], h[1][1])
863}
864
865type O3 = gam_math::jet_tower::Tower3<2>;
867
868fn o3_exp(x: &O3) -> O3 {
869 x.compose_unary_with(|v| {
870 let e = v.exp();
871 [e, e, e, e]
872 })
873}
874
875fn o3_ln(x: &O3) -> O3 {
876 x.compose_unary_with(|v| [v.ln(), v.recip(), -v.powi(-2), 2.0 * v.powi(-3)])
877}
878
879fn o3_recip(x: &O3) -> O3 {
880 x.compose_unary_with(|v| [v.recip(), -v.powi(-2), 2.0 * v.powi(-3), -6.0 * v.powi(-4)])
881}
882
883fn o3_powf(x: &O3, a: f64) -> O3 {
884 x.compose_unary_with(|v| {
885 [
886 v.powf(a),
887 a * v.powf(a - 1.0),
888 a * (a - 1.0) * v.powf(a - 2.0),
889 a * (a - 1.0) * (a - 2.0) * v.powf(a - 3.0),
890 ]
891 })
892}
893
894fn o3_ln_gamma(x: &O3) -> O3 {
895 x.compose_unary_with(|v| {
896 let stack = gam_math::jet_tower::ln_gamma_derivative_stack(v);
897 [stack[0], stack[1], stack[2], stack[3]]
898 })
899}
900
901pub(crate) fn dispersion_eta_nll_order3(
919 kind: DispersionFamilyKind,
920 yi: f64,
921 em: f64,
922 ed: f64,
923 wi: f64,
924) -> O3 {
925 let eta_mu = O3::variable(em, 0);
926 let eta_d = O3::variable(ed, 1);
927 match kind {
928 DispersionFamilyKind::NegativeBinomial => {
929 let theta = o3_exp(&eta_d);
930 let theta_plus_y = theta.add(&O3::constant(yi));
931 let log_total = if em >= ed {
932 eta_mu.add(&o3_ln(
933 &o3_exp(&eta_d.sub(&eta_mu)).add(&O3::constant(1.0)),
934 ))
935 } else {
936 eta_d.add(&o3_ln(
937 &o3_exp(&eta_mu.sub(&eta_d)).add(&O3::constant(1.0)),
938 ))
939 };
940 let loglik = o3_ln_gamma(&theta_plus_y)
941 .sub(&o3_ln_gamma(&theta))
942 .sub(&O3::constant(ln_gamma(yi + 1.0)))
943 .add(&theta.mul(&eta_d.sub(&log_total)))
944 .add(&eta_mu.sub(&log_total).scale(yi));
945 loglik.scale(-wi)
946 }
947 DispersionFamilyKind::Gamma => {
948 let mu = o3_exp(&eta_mu);
949 let nu = o3_exp(&eta_d);
950 let y_pos = yi;
951 let loglik = nu
952 .mul(&o3_ln(&nu))
953 .sub(&nu.mul(&o3_ln(&mu)))
954 .sub(&o3_ln_gamma(&nu))
955 .add(&nu.sub(&O3::constant(1.0)).scale(y_pos.ln()))
956 .sub(&nu.mul(&o3_recip(&mu).scale(yi)));
957 loglik.scale(-wi)
958 }
959 DispersionFamilyKind::Beta => {
960 let mu = o3_recip(&o3_exp(&eta_mu.scale(-1.0)).add(&O3::constant(1.0)));
961 let phi = o3_exp(&eta_d);
962 let one_minus_mu = O3::constant(1.0).sub(&mu);
963 let yc = yi;
964 let a = mu.mul(&phi);
965 let b = one_minus_mu.mul(&phi);
966 let loglik = o3_ln_gamma(&phi)
967 .sub(&o3_ln_gamma(&a))
968 .sub(&o3_ln_gamma(&b))
969 .add(&a.sub(&O3::constant(1.0)).scale(yc.ln()))
970 .add(&b.sub(&O3::constant(1.0)).scale((-yc).ln_1p()));
971 loglik.scale(-wi)
972 }
973 DispersionFamilyKind::Tweedie { p } => {
974 let one_minus_p = 1.0 - p;
975 let two_minus_p = 2.0 - p;
976 let mu = o3_exp(&eta_mu);
977 let phi = o3_exp(&eta_d.scale(-1.0));
978 if yi > 0.0 {
979 let dev = o3_powf(&mu, two_minus_p)
980 .scale(1.0 / two_minus_p)
981 .sub(&o3_powf(&mu, one_minus_p).scale(yi / one_minus_p))
982 .add(&O3::constant(
983 yi.powf(two_minus_p) / (one_minus_p * two_minus_p),
984 ))
985 .scale(2.0);
986 let loglik = dev
987 .mul(&o3_recip(&phi).scale(-0.5))
988 .sub(&o3_ln(&phi.scale(2.0 * std::f64::consts::PI)).scale(0.5))
989 .sub(&O3::constant(0.5 * p * yi.ln()));
990 loglik.scale(-wi)
991 } else {
992 let c = o3_powf(&mu, two_minus_p).scale(1.0 / two_minus_p);
993 let loglik = c.mul(&o3_recip(&phi)).scale(-1.0);
994 loglik.scale(-wi)
995 }
996 }
997 }
998}
999
1000pub(crate) fn dispersion_row_observed_hessian_directional(
1005 kind: DispersionFamilyKind,
1006 yi: f64,
1007 eta_mu: f64,
1008 eta_d: f64,
1009 prior_weight: f64,
1010 du_mu: f64,
1011 du_d: f64,
1012) -> (f64, f64, f64) {
1013 if prior_weight <= 0.0 {
1014 return (0.0, 0.0, 0.0);
1015 }
1016 let tower = dispersion_eta_nll_order3(kind, yi, eta_mu, eta_d, prior_weight);
1017 let t3 = &tower.t3;
1018 (
1019 t3[0][0][0] * du_mu + t3[0][0][1] * du_d,
1020 t3[0][1][0] * du_mu + t3[0][1][1] * du_d,
1021 t3[1][1][0] * du_mu + t3[1][1][1] * du_d,
1022 )
1023}
1024
1025#[derive(Clone, Copy, Debug, PartialEq)]
1034pub struct DispersionAloRowGeometry {
1035 pub nll_score: [f64; 2],
1036 pub observed_hessian: [[f64; 2]; 2],
1037}
1038
1039pub fn dispersion_alo_row_geometry(
1046 kind: DispersionFamilyKind,
1047 row: usize,
1048 y: f64,
1049 eta_mu: f64,
1050 eta_d: f64,
1051 prior_weight: f64,
1052) -> Result<DispersionAloRowGeometry, String> {
1053 validate_dispersion_row_geometry_inputs(kind, row, y, eta_mu, eta_d, prior_weight)?;
1054 if prior_weight == 0.0 {
1055 return Ok(DispersionAloRowGeometry {
1056 nll_score: [0.0; 2],
1057 observed_hessian: [[0.0; 2]; 2],
1058 });
1059 }
1060 let tower = dispersion_eta_nll_order2(kind, y, eta_mu, eta_d, prior_weight);
1061 let (_, gradient, hessian) = tower.into_channels();
1062 let geometry = DispersionAloRowGeometry {
1063 nll_score: gradient,
1064 observed_hessian: hessian,
1065 };
1066 if geometry
1067 .nll_score
1068 .iter()
1069 .chain(geometry.observed_hessian.iter().flatten())
1070 .any(|value| !value.is_finite())
1071 {
1072 return Err(GamlssError::RowGeometryUnrepresentable {
1073 row,
1074 quantity: "dispersion-family ALO row geometry",
1075 eta: eta_mu,
1076 value: f64::NAN,
1077 }
1078 .into());
1079 }
1080 Ok(geometry)
1081}
1082
1083#[inline]
1084pub(crate) fn tower_score_info<const K: usize>(
1085 tower: &gam_math::jet_scalar::Order2<K>,
1086 idx: usize,
1087 wi: f64,
1088) -> (f64, f64) {
1089 if wi == 0.0 {
1090 (0.0, 0.0)
1091 } else {
1092 (-tower.g()[idx] / wi, tower.h()[idx][idx] / wi)
1093 }
1094}
1095
1096pub(super) fn dispersion_row_kernel(
1101 kind: DispersionFamilyKind,
1102 yi: f64,
1103 eta_mu: f64,
1104 eta_d: f64,
1105 prior_weight: f64,
1106) -> DispersionRowKernel {
1107 let em = eta_mu;
1108 let ed = eta_d;
1109 if prior_weight <= 0.0 {
1113 return DispersionRowKernel {
1114 loglik: 0.0,
1115 mean_weight: 0.0,
1116 mean_response: em,
1117 disp_weight: 0.0,
1118 disp_response: ed,
1119 };
1120 }
1121 let wi = prior_weight;
1122 match kind {
1123 DispersionFamilyKind::NegativeBinomial => {
1124 let mu = em.exp();
1125 let theta = ed.exp(); let loglik = dispersion_nb_loglik(yi, mu, theta, wi);
1127 let mean_eta_information = if mu >= theta {
1128 theta / (1.0 + theta / mu)
1129 } else {
1130 mu / (1.0 + mu / theta)
1131 };
1132 let mean_weight = wi * mean_eta_information;
1133 let mean_response = em + (yi - mu) / mu;
1134 let theta_fraction = if theta >= mu {
1170 (mu / theta - yi / theta) / (1.0 + mu / theta)
1171 } else {
1172 (1.0 - yi / mu) / (1.0 + theta / mu)
1173 };
1174 let score_theta = gam_math::jet_tower::digamma(theta + yi)
1175 - gam_math::jet_tower::digamma(theta)
1176 + log_positive_share(theta, mu)
1177 + theta_fraction;
1178 let score_eta = theta * score_theta;
1179 let eta_information = nb_log_precision_fisher_jensen(mu, theta);
1180 let disp_weight = wi * eta_information;
1181 let disp_response = ed + score_eta / eta_information;
1182 DispersionRowKernel {
1183 loglik,
1184 mean_weight,
1185 mean_response,
1186 disp_weight,
1187 disp_response,
1188 }
1189 }
1190 DispersionFamilyKind::Gamma => {
1191 let mu = em.exp();
1192 let nu = ed.exp(); let tower = dispersion_gamma_disp_order2(yi, yi, mu, nu, wi);
1194 let (s_nu, info_nu_raw) = tower_score_info(&tower, 0, wi);
1195 let loglik = -tower.value();
1196 let mean_weight = wi * nu;
1197 let mean_response = em + (yi - mu) / mu;
1198 let disp_weight = wi * nu * nu * info_nu_raw;
1199 let disp_response = ed + s_nu / (nu * info_nu_raw);
1200 DispersionRowKernel {
1201 loglik,
1202 mean_weight,
1203 mean_response,
1204 disp_weight,
1205 disp_response,
1206 }
1207 }
1208 DispersionFamilyKind::Beta => {
1209 let logit = gam_solve::mixture_link::logit_inverse_link_jet5(em);
1211 let mu = logit.mu;
1212 let phi = ed.exp(); let q = logit.d1;
1214 let tower = dispersion_beta_nll_order2(yi, mu, phi, wi);
1215 let (score_mu, _) = tower_score_info(&tower, 0, wi);
1216 let (s_phi, _) = tower_score_info(&tower, 1, wi);
1217 let loglik = -tower.value();
1218 let a = mu * phi;
1219 let b = (1.0 - mu) * phi;
1220 let tri_a = gam_math::jet_tower::trigamma(a);
1221 let tri_b = gam_math::jet_tower::trigamma(b);
1222 let tri_phi = gam_math::jet_tower::trigamma(phi);
1223 let info_mu = phi * phi * (tri_a + tri_b);
1224 let one_minus_mu = 1.0 - mu;
1225 let info_phi = mu * mu * tri_a + one_minus_mu * one_minus_mu * tri_b - tri_phi;
1226 let mean_weight = wi * q * q * info_mu;
1227 let mean_response = em + score_mu / (q * info_mu);
1228 let disp_weight = wi * phi * phi * info_phi;
1229 let disp_response = ed + s_phi / (phi * info_phi);
1230 DispersionRowKernel {
1231 loglik,
1232 mean_weight,
1233 mean_response,
1234 disp_weight,
1235 disp_response,
1236 }
1237 }
1238 DispersionFamilyKind::Tweedie { p } => {
1239 let mu = em.exp();
1240 let phi = (-ed).exp();
1242 let two_minus_p = 2.0 - p;
1243 let mean_weight = wi * mu.powf(two_minus_p) / phi;
1249 let mean_response = em + (yi - mu) / mu;
1250 let tower = dispersion_tweedie_disp_order2(yi, em, ed, p, wi);
1259 let loglik = -tower.value();
1260 let (s_eta, info_eta_raw) = tower_score_info(&tower, 0, wi);
1264 let curvature_eta = if yi > 0.0 { 0.5 } else { info_eta_raw };
1265 let disp_weight = wi * curvature_eta;
1266 let disp_response = ed + s_eta / curvature_eta;
1267 DispersionRowKernel {
1268 loglik,
1269 mean_weight,
1270 mean_response,
1271 disp_weight,
1272 disp_response,
1273 }
1274 }
1275 }
1276}
1277
1278#[derive(Clone)]
1280pub(crate) struct DispersionGlmLocationScaleFamily {
1281 pub(crate) kind: DispersionFamilyKind,
1282 pub(crate) y: Array1<f64>,
1283 pub(crate) weights: Array1<f64>,
1284}
1285
1286impl DispersionGlmLocationScaleFamily {
1287 pub(crate) const BLOCK_MEAN: usize = 0;
1288 pub(crate) const BLOCK_DISP: usize = 1;
1289}
1290
1291impl CustomFamily for DispersionGlmLocationScaleFamily {
1292 fn joint_jeffreys_term_required(&self) -> bool {
1296 true
1297 }
1298
1299 fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
1300 validate_block_count::<GamlssError>(self.kind.family_tag(), 2, block_states.len())?;
1301 let eta_mu = &block_states[Self::BLOCK_MEAN].eta;
1302 let eta_d = &block_states[Self::BLOCK_DISP].eta;
1303 let n = self.y.len();
1304 if eta_mu.len() != n || eta_d.len() != n || self.weights.len() != n {
1305 return Err(format!(
1306 "{} row-count mismatch: y={n}, eta_mu={}, eta_d={}, weights={}",
1307 self.kind.family_tag(),
1308 eta_mu.len(),
1309 eta_d.len(),
1310 self.weights.len()
1311 ));
1312 }
1313 for i in 0..n {
1314 validate_dispersion_row_geometry_inputs(
1315 self.kind,
1316 i,
1317 self.y[i],
1318 eta_mu[i],
1319 eta_d[i],
1320 self.weights[i],
1321 )?;
1322 }
1323 let kernels: Vec<DispersionRowKernel> =
1334 if rayon::current_thread_index().is_none() && n > DISPERSION_PARALLEL_ROW_THRESHOLD {
1335 use rayon::iter::{IntoParallelIterator, ParallelIterator};
1336 (0..n)
1337 .into_par_iter()
1338 .map(|i| {
1339 dispersion_row_kernel(
1340 self.kind,
1341 self.y[i],
1342 eta_mu[i],
1343 eta_d[i],
1344 self.weights[i],
1345 )
1346 })
1347 .collect()
1348 } else {
1349 (0..n)
1350 .map(|i| {
1351 dispersion_row_kernel(
1352 self.kind,
1353 self.y[i],
1354 eta_mu[i],
1355 eta_d[i],
1356 self.weights[i],
1357 )
1358 })
1359 .collect()
1360 };
1361
1362 let mut log_likelihood = 0.0;
1368 for (i, row) in kernels.iter().enumerate() {
1369 validate_dispersion_row_kernel_output(i, eta_mu[i], eta_d[i], self.weights[i], row)?;
1370 log_likelihood += row.loglik;
1371 if !log_likelihood.is_finite() {
1372 return Err(dispersion_geometry_error(
1373 i,
1374 "dispersion-family cumulative log likelihood",
1375 eta_mu[i],
1376 log_likelihood,
1377 ));
1378 }
1379 }
1380 let mean_weights = Array1::from_iter(kernels.iter().map(|row| row.mean_weight));
1381 let mean_response = Array1::from_iter(kernels.iter().map(|row| row.mean_response));
1382 let disp_weights = Array1::from_iter(kernels.iter().map(|row| row.disp_weight));
1383 let disp_response = Array1::from_iter(kernels.iter().map(|row| row.disp_response));
1384 Ok(FamilyEvaluation {
1385 log_likelihood,
1386 blockworking_sets: vec![
1387 BlockWorkingSet::diagonal_checked(mean_response, mean_weights)?,
1388 BlockWorkingSet::diagonal_checked(disp_response, disp_weights)?,
1389 ],
1390 })
1391 }
1392
1393 fn log_likelihood_only(&self, block_states: &[ParameterBlockState]) -> Result<f64, String> {
1394 validate_block_count::<GamlssError>(self.kind.family_tag(), 2, block_states.len())?;
1395 let eta_mu = &block_states[Self::BLOCK_MEAN].eta;
1396 let eta_d = &block_states[Self::BLOCK_DISP].eta;
1397 let n = self.y.len();
1398 if eta_mu.len() != n || eta_d.len() != n || self.weights.len() != n {
1399 return Err(GamlssError::DimensionMismatch {
1400 reason: format!(
1401 "{} log-likelihood row-count mismatch: y={n}, eta_mu={}, eta_d={}, weights={}",
1402 self.kind.family_tag(),
1403 eta_mu.len(),
1404 eta_d.len(),
1405 self.weights.len()
1406 ),
1407 }
1408 .into());
1409 }
1410 for i in 0..n {
1411 validate_dispersion_row_geometry_inputs(
1412 self.kind,
1413 i,
1414 self.y[i],
1415 eta_mu[i],
1416 eta_d[i],
1417 self.weights[i],
1418 )?;
1419 }
1420 let per_row: Vec<f64> =
1429 if rayon::current_thread_index().is_none() && n > DISPERSION_PARALLEL_ROW_THRESHOLD {
1430 use rayon::iter::{IntoParallelIterator, ParallelIterator};
1431 (0..n)
1432 .into_par_iter()
1433 .map(|i| {
1434 dispersion_row_loglik(
1435 self.kind,
1436 self.y[i],
1437 eta_mu[i],
1438 eta_d[i],
1439 self.weights[i],
1440 )
1441 })
1442 .collect()
1443 } else {
1444 (0..n)
1445 .map(|i| {
1446 dispersion_row_loglik(
1447 self.kind,
1448 self.y[i],
1449 eta_mu[i],
1450 eta_d[i],
1451 self.weights[i],
1452 )
1453 })
1454 .collect()
1455 };
1456 let mut ll = 0.0;
1459 for (i, loglik) in per_row.into_iter().enumerate() {
1460 if !loglik.is_finite() {
1461 return Err(dispersion_geometry_error(
1462 i,
1463 "dispersion-family row log likelihood",
1464 eta_mu[i],
1465 loglik,
1466 ));
1467 }
1468 ll += loglik;
1469 if !ll.is_finite() {
1470 return Err(dispersion_geometry_error(
1471 i,
1472 "dispersion-family cumulative log likelihood",
1473 eta_mu[i],
1474 ll,
1475 ));
1476 }
1477 }
1478 Ok(ll)
1479 }
1480
1481 fn coefficient_hessian_cost(&self, specs: &[ParameterBlockSpec]) -> u64 {
1482 crate::location_scale_engine::location_scale_coefficient_hessian_cost(
1483 self.y.len() as u64,
1484 specs,
1485 )
1486 }
1487
1488 fn exact_newton_joint_hessian_with_specs(
1515 &self,
1516 block_states: &[ParameterBlockState],
1517 specs: &[ParameterBlockSpec],
1518 ) -> Result<Option<Array2<f64>>, String> {
1519 validate_block_count::<GamlssError>(self.kind.family_tag(), 2, block_states.len())?;
1520 if specs.len() != 2 {
1521 return Err(format!(
1522 "{} exact joint Hessian expects 2 specs, got {}",
1523 self.kind.family_tag(),
1524 specs.len()
1525 ));
1526 }
1527 let eta_mu = &block_states[Self::BLOCK_MEAN].eta;
1528 let eta_d = &block_states[Self::BLOCK_DISP].eta;
1529 let n = self.y.len();
1530 if eta_mu.len() != n || eta_d.len() != n || self.weights.len() != n {
1531 return Err(format!(
1532 "{} exact joint Hessian row-count mismatch: y={n}, eta_mu={}, eta_d={}, weights={}",
1533 self.kind.family_tag(),
1534 eta_mu.len(),
1535 eta_d.len(),
1536 self.weights.len()
1537 ));
1538 }
1539 for i in 0..n {
1540 validate_dispersion_row_geometry_inputs(
1541 self.kind,
1542 i,
1543 self.y[i],
1544 eta_mu[i],
1545 eta_d[i],
1546 self.weights[i],
1547 )?;
1548 }
1549
1550 let observed: Vec<(f64, f64, f64)> =
1555 if rayon::current_thread_index().is_none() && n > DISPERSION_PARALLEL_ROW_THRESHOLD {
1556 use rayon::iter::{IntoParallelIterator, ParallelIterator};
1557 (0..n)
1558 .into_par_iter()
1559 .map(|i| {
1560 dispersion_row_observed_hessian_weights(
1561 self.kind,
1562 self.y[i],
1563 eta_mu[i],
1564 eta_d[i],
1565 self.weights[i],
1566 )
1567 })
1568 .collect()
1569 } else {
1570 (0..n)
1571 .map(|i| {
1572 dispersion_row_observed_hessian_weights(
1573 self.kind,
1574 self.y[i],
1575 eta_mu[i],
1576 eta_d[i],
1577 self.weights[i],
1578 )
1579 })
1580 .collect()
1581 };
1582 for (i, &(h_mm, h_md, h_dd)) in observed.iter().enumerate() {
1583 for (quantity, eta, value) in [
1584 ("dispersion-family observed mean curvature", eta_mu[i], h_mm),
1585 (
1586 "dispersion-family observed cross curvature",
1587 eta_mu[i],
1588 h_md,
1589 ),
1590 (
1591 "dispersion-family observed precision curvature",
1592 eta_d[i],
1593 h_dd,
1594 ),
1595 ] {
1596 if !value.is_finite() {
1597 return Err(dispersion_geometry_error(i, quantity, eta, value));
1598 }
1599 }
1600 }
1601 let mean_weights = Array1::from_shape_fn(n, |i| observed[i].0);
1602 let cross_weights = Array1::from_shape_fn(n, |i| observed[i].1);
1603 let disp_weights = Array1::from_shape_fn(n, |i| observed[i].2);
1604 let mean_spec = &specs[Self::BLOCK_MEAN];
1605 let disp_spec = &specs[Self::BLOCK_DISP];
1606 if mean_spec.design.nrows() != n || disp_spec.design.nrows() != n {
1607 return Err(format!(
1608 "{} exact joint Hessian design row mismatch: y={n}, mean rows={}, precision rows={}",
1609 self.kind.family_tag(),
1610 mean_spec.design.nrows(),
1611 disp_spec.design.nrows()
1612 ));
1613 }
1614 let p_mean = mean_spec.design.ncols();
1615 let p_disp = disp_spec.design.ncols();
1616 if block_states[Self::BLOCK_MEAN].beta.len() != p_mean
1617 || block_states[Self::BLOCK_DISP].beta.len() != p_disp
1618 {
1619 return Err(format!(
1620 "{} exact joint Hessian beta/design mismatch: mean beta {} vs cols {}, precision beta {} vs cols {}",
1621 self.kind.family_tag(),
1622 block_states[Self::BLOCK_MEAN].beta.len(),
1623 p_mean,
1624 block_states[Self::BLOCK_DISP].beta.len(),
1625 p_disp
1626 ));
1627 }
1628
1629 let h_mean = xt_diag_x_design(&mean_spec.design, &mean_weights)?;
1630 let h_cross = xt_diag_y_design(&mean_spec.design, &cross_weights, &disp_spec.design)?;
1631 let h_disp = xt_diag_x_design(&disp_spec.design, &disp_weights)?;
1632 let total = p_mean + p_disp;
1633 let mut h = Array2::<f64>::zeros((total, total));
1634 h.slice_mut(s![0..p_mean, 0..p_mean]).assign(&h_mean);
1635 h.slice_mut(s![0..p_mean, p_mean..total]).assign(&h_cross);
1636 h.slice_mut(s![p_mean..total, p_mean..total])
1637 .assign(&h_disp);
1638 mirror_upper_to_lower(&mut h);
1639 Ok(Some(h))
1640 }
1641
1642 fn exact_newton_joint_hessian_directional_derivative_with_specs(
1660 &self,
1661 block_states: &[ParameterBlockState],
1662 specs: &[ParameterBlockSpec],
1663 d_beta_flat: &Array1<f64>,
1664 ) -> Result<Option<Array2<f64>>, String> {
1665 validate_block_count::<GamlssError>(self.kind.family_tag(), 2, block_states.len())?;
1666 if specs.len() != 2 {
1667 return Err(format!(
1668 "{} joint Hessian directional derivative expects 2 specs, got {}",
1669 self.kind.family_tag(),
1670 specs.len()
1671 ));
1672 }
1673 let eta_mu = &block_states[Self::BLOCK_MEAN].eta;
1674 let eta_d = &block_states[Self::BLOCK_DISP].eta;
1675 let n = self.y.len();
1676 if eta_mu.len() != n || eta_d.len() != n || self.weights.len() != n {
1677 return Err(format!(
1678 "{} joint Hessian directional derivative row-count mismatch: y={n}, eta_mu={}, eta_d={}, weights={}",
1679 self.kind.family_tag(),
1680 eta_mu.len(),
1681 eta_d.len(),
1682 self.weights.len()
1683 ));
1684 }
1685 for i in 0..n {
1686 validate_dispersion_row_geometry_inputs(
1687 self.kind,
1688 i,
1689 self.y[i],
1690 eta_mu[i],
1691 eta_d[i],
1692 self.weights[i],
1693 )?;
1694 }
1695 let mean_spec = &specs[Self::BLOCK_MEAN];
1696 let disp_spec = &specs[Self::BLOCK_DISP];
1697 if mean_spec.design.nrows() != n || disp_spec.design.nrows() != n {
1698 return Err(format!(
1699 "{} joint Hessian directional derivative design row mismatch: y={n}, mean rows={}, precision rows={}",
1700 self.kind.family_tag(),
1701 mean_spec.design.nrows(),
1702 disp_spec.design.nrows()
1703 ));
1704 }
1705 let p_mean = mean_spec.design.ncols();
1706 let p_disp = disp_spec.design.ncols();
1707 if d_beta_flat.len() != p_mean + p_disp {
1708 return Err(format!(
1709 "{} joint Hessian directional derivative direction length mismatch: got {}, expected {}",
1710 self.kind.family_tag(),
1711 d_beta_flat.len(),
1712 p_mean + p_disp
1713 ));
1714 }
1715 let u_mu = d_beta_flat.slice(s![0..p_mean]).to_owned();
1716 let u_d = d_beta_flat.slice(s![p_mean..p_mean + p_disp]).to_owned();
1717 let du_mu = mean_spec.design.apply(&u_mu);
1720 let du_d = disp_spec.design.apply(&u_d);
1721 let directional: Vec<(f64, f64, f64)> =
1722 if rayon::current_thread_index().is_none() && n > DISPERSION_PARALLEL_ROW_THRESHOLD {
1723 use rayon::iter::{IntoParallelIterator, ParallelIterator};
1724 (0..n)
1725 .into_par_iter()
1726 .map(|i| {
1727 dispersion_row_observed_hessian_directional(
1728 self.kind,
1729 self.y[i],
1730 eta_mu[i],
1731 eta_d[i],
1732 self.weights[i],
1733 du_mu[i],
1734 du_d[i],
1735 )
1736 })
1737 .collect()
1738 } else {
1739 (0..n)
1740 .map(|i| {
1741 dispersion_row_observed_hessian_directional(
1742 self.kind,
1743 self.y[i],
1744 eta_mu[i],
1745 eta_d[i],
1746 self.weights[i],
1747 du_mu[i],
1748 du_d[i],
1749 )
1750 })
1751 .collect()
1752 };
1753 for (i, &(d_mm, d_md, d_dd)) in directional.iter().enumerate() {
1754 for (quantity, eta, value) in [
1755 (
1756 "dispersion-family directional mean curvature drift",
1757 eta_mu[i],
1758 d_mm,
1759 ),
1760 (
1761 "dispersion-family directional cross curvature drift",
1762 eta_mu[i],
1763 d_md,
1764 ),
1765 (
1766 "dispersion-family directional precision curvature drift",
1767 eta_d[i],
1768 d_dd,
1769 ),
1770 ] {
1771 if !value.is_finite() {
1772 return Err(dispersion_geometry_error(i, quantity, eta, value));
1773 }
1774 }
1775 }
1776 let mean_drift = Array1::from_shape_fn(n, |i| directional[i].0);
1777 let cross_drift = Array1::from_shape_fn(n, |i| directional[i].1);
1778 let disp_drift = Array1::from_shape_fn(n, |i| directional[i].2);
1779 let dh_mean = xt_diag_x_design(&mean_spec.design, &mean_drift)?;
1780 let dh_cross = xt_diag_y_design(&mean_spec.design, &cross_drift, &disp_spec.design)?;
1781 let dh_disp = xt_diag_x_design(&disp_spec.design, &disp_drift)?;
1782 let total = p_mean + p_disp;
1783 let mut dh = Array2::<f64>::zeros((total, total));
1784 dh.slice_mut(s![0..p_mean, 0..p_mean]).assign(&dh_mean);
1785 dh.slice_mut(s![0..p_mean, p_mean..total]).assign(&dh_cross);
1786 dh.slice_mut(s![p_mean..total, p_mean..total])
1787 .assign(&dh_disp);
1788 mirror_upper_to_lower(&mut dh);
1789 Ok(Some(dh))
1790 }
1791
1792 fn likelihood_blocks_uncoupled(&self) -> bool {
1802 false
1803 }
1804
1805 fn has_explicit_joint_hessian(&self) -> bool {
1809 true
1810 }
1811
1812 fn outer_hyper_hessian_dense_available(&self, specs: &[ParameterBlockSpec]) -> bool {
1823 assert!(
1824 crate::custom_family::validate_blockspec_consistency(specs).is_ok(),
1825 "DispersionGlmLocationScale outer hyper-Hessian dense availability: \
1826 inconsistent parameter block specs"
1827 );
1828 specs.len() < 2
1829 }
1830}
1831
1832pub struct DispersionGlmLocationScaleTermSpec {
1836 pub kind: DispersionFamilyKind,
1837 pub y: Array1<f64>,
1838 pub weights: Array1<f64>,
1839 pub meanspec: TermCollectionSpec,
1840 pub log_dispspec: TermCollectionSpec,
1841 pub mean_offset: Array1<f64>,
1842 pub log_disp_offset: Array1<f64>,
1843}
1844
1845pub(crate) struct DispersionGlmLocationScaleTermBuilder {
1846 pub(crate) kind: DispersionFamilyKind,
1847 pub(crate) y: Array1<f64>,
1848 pub(crate) weights: Array1<f64>,
1849 pub(crate) meanspec: TermCollectionSpec,
1850 pub(crate) noisespec: TermCollectionSpec,
1851 pub(crate) mean_offset: Array1<f64>,
1852 pub(crate) noise_offset: Array1<f64>,
1853}
1854
1855pub(crate) fn dispersion_location_scale_warm_start(
1859 kind: DispersionFamilyKind,
1860 y: &Array1<f64>,
1861 weights: &Array1<f64>,
1862 mean_block: &ParameterBlockSpec,
1863 disp_block: &ParameterBlockSpec,
1864 mean_beta_hint: Option<&Array1<f64>>,
1865 disp_beta_hint: Option<&Array1<f64>>,
1866) -> Result<(Array1<f64>, Array1<f64>), String> {
1867 let ridge_floor = 1e-10;
1868 let mean_beta = if let Some(beta) = mean_beta_hint {
1869 beta.clone()
1870 } else {
1871 let target = Array1::from_shape_fn(y.len(), |i| {
1872 if kind.mean_is_logit() {
1873 let yi = y[i].clamp(1e-3, 1.0 - 1e-3);
1874 (yi / (1.0 - yi)).ln()
1875 } else {
1876 (y[i].max(0.0) + 0.1).ln()
1878 }
1879 });
1880 solve_penalizedweighted_projection(
1881 &mean_block.design,
1882 &mean_block.offset,
1883 &target,
1884 weights,
1885 &mean_block.penalties,
1886 &mean_block.initial_log_lambdas,
1887 ridge_floor,
1888 )?
1889 };
1890 let disp_beta = if let Some(beta) = disp_beta_hint {
1891 beta.clone()
1892 } else {
1893 let mean_eta = mean_block.design.apply(&mean_beta) + &mean_block.offset;
1908 let target = Array1::from_shape_fn(y.len(), |i| {
1909 dispersion_moment_log_precision_seed(kind, y[i], mean_eta[i])
1910 });
1911 solve_penalizedweighted_projection(
1912 &disp_block.design,
1913 &disp_block.offset,
1914 &target,
1915 weights,
1916 &disp_block.penalties,
1917 &disp_block.initial_log_lambdas,
1918 ridge_floor,
1919 )?
1920 };
1921 Ok((mean_beta, disp_beta))
1922}
1923
1924#[inline]
1925fn dispersion_moment_log_precision_seed(kind: DispersionFamilyKind, yi: f64, eta_mu: f64) -> f64 {
1926 const LOG_PRECISION_FLOOR: f64 = -10.0;
1927 const LOG_PRECISION_CEILING: f64 = 10.0;
1928 let em = eta_mu;
1929 let raw = match kind {
1930 DispersionFamilyKind::Beta => {
1931 0.0
1940 }
1941 DispersionFamilyKind::Gamma => {
1942 let mu = em.exp().max(1e-12);
1943 let e2 = (yi - mu).powi(2).max(1e-8 * mu * mu);
1944 (mu * mu / e2).max(1e-6).ln()
1945 }
1946 DispersionFamilyKind::NegativeBinomial => {
1947 let mu = em.exp().max(1e-12);
1948 let e2 = (yi - mu).powi(2);
1949 let excess = (e2 - mu).max((mu + 2.0 * mu * mu).sqrt());
1988 (mu * mu / excess).max(1e-6).ln()
1989 }
1990 DispersionFamilyKind::Tweedie { p } => {
1991 let mu = em.exp().max(1e-12);
1992 let e2 = (yi - mu).powi(2).max(1e-8 * mu.powf(p));
1993 (mu.powf(p) / e2).max(1e-6).ln()
1994 }
1995 };
1996 raw.clamp(LOG_PRECISION_FLOOR, LOG_PRECISION_CEILING)
1997}
1998
1999impl LocationScaleFamilyBuilder for DispersionGlmLocationScaleTermBuilder {
2000 type Family = DispersionGlmLocationScaleFamily;
2001
2002 fn meanspec(&self) -> &TermCollectionSpec {
2003 &self.meanspec
2004 }
2005
2006 fn noisespec(&self) -> &TermCollectionSpec {
2007 &self.noisespec
2008 }
2009
2010 fn build_blocks(
2011 &self,
2012 theta: &Array1<f64>,
2013 mean_design: &TermCollectionDesign,
2014 noise_design: &TermCollectionDesign,
2015 mean_beta_hint: Option<Array1<f64>>,
2016 noise_beta_hint: Option<Array1<f64>>,
2017 ) -> Result<Vec<ParameterBlockSpec>, String> {
2018 let layout = GamlssLambdaLayout::two_block(
2019 mean_design.penalties.len(),
2020 self.noise_penalty_count(noise_design),
2021 );
2022 layout.validate_theta_len(theta.len(), "dispersion location-scale")?;
2023
2024 let mean_offset = mean_design
2025 .compose_offset(self.mean_offset.view(), "dispersion location-scale mean")
2026 .map_err(|error| error.to_string())?;
2027 let noise_offset = noise_design
2028 .compose_offset(
2029 self.noise_offset.view(),
2030 "dispersion location-scale log-precision",
2031 )
2032 .map_err(|error| error.to_string())?;
2033 let mut meanspec = build_location_scale_block(
2034 "mu",
2035 mean_design.design.clone(),
2036 mean_offset,
2037 mean_design.penalties_as_penalty_matrix(),
2038 mean_design.nullspace_dims.clone(),
2039 layout.mean_from(theta),
2040 mean_beta_hint,
2041 0,
2042 LOCATION_SCALE_N_OUTPUTS,
2043 "DispersionLocationScale::build_blocks: mu",
2044 )?;
2045
2046 let disp_penalties = noise_design.penalties_as_penalty_matrix();
2058 let disp_nullspace = noise_design.nullspace_dims.clone();
2059 let mut dispspec = build_location_scale_block(
2060 "log_precision",
2061 noise_design.design.clone(),
2062 noise_offset,
2063 disp_penalties,
2064 disp_nullspace,
2065 layout.noise_from(theta),
2066 noise_beta_hint,
2067 1,
2068 LOCATION_SCALE_N_OUTPUTS,
2069 "DispersionLocationScale::build_blocks: log_precision",
2070 )?;
2071
2072 if meanspec.initial_beta.is_none() || dispspec.initial_beta.is_none() {
2073 let (mean_beta0, disp_beta0) = dispersion_location_scale_warm_start(
2074 self.kind,
2075 &self.y,
2076 &self.weights,
2077 &meanspec,
2078 &dispspec,
2079 meanspec.initial_beta.as_ref(),
2080 dispspec.initial_beta.as_ref(),
2081 )?;
2082 if meanspec.initial_beta.is_none() {
2083 meanspec.initial_beta = Some(mean_beta0);
2084 }
2085 if dispspec.initial_beta.is_none() {
2086 dispspec.initial_beta = Some(disp_beta0);
2087 }
2088 }
2089
2090 Ok(vec![meanspec, dispspec])
2091 }
2092
2093 fn build_family(
2094 &self,
2095 mean_design: &TermCollectionDesign,
2096 noise_design: &TermCollectionDesign,
2097 ) -> Self::Family {
2098 assert_eq!(
2105 mean_design.design.nrows(),
2106 self.y.len(),
2107 "DispersionGlmLocationScale::build_family: mean design row count must match y"
2108 );
2109 assert_eq!(
2110 noise_design.design.nrows(),
2111 self.y.len(),
2112 "DispersionGlmLocationScale::build_family: noise design row count must match y"
2113 );
2114 DispersionGlmLocationScaleFamily {
2115 kind: self.kind,
2116 y: self.y.clone(),
2117 weights: self.weights.clone(),
2118 }
2119 }
2120
2121 fn extract_primary_betas(
2122 &self,
2123 fit: &UnifiedFitResult,
2124 ) -> Result<(Array1<f64>, Array1<f64>), String> {
2125 let mean_beta = fit
2126 .block_states
2127 .get(DispersionGlmLocationScaleFamily::BLOCK_MEAN)
2128 .ok_or_else(|| "missing dispersion mean block state".to_string())?
2129 .beta
2130 .clone();
2131 let disp_beta = fit
2132 .block_states
2133 .get(DispersionGlmLocationScaleFamily::BLOCK_DISP)
2134 .ok_or_else(|| "missing dispersion log-precision block state".to_string())?
2135 .beta
2136 .clone();
2137 Ok((mean_beta, disp_beta))
2138 }
2139
2140 fn build_psiderivative_blocks(
2141 &self,
2142 data: ndarray::ArrayView2<'_, f64>,
2143 meanspec: &TermCollectionSpec,
2144 noisespec: &TermCollectionSpec,
2145 mean_design: &TermCollectionDesign,
2146 noise_design: &TermCollectionDesign,
2147 ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String> {
2148 Err(format!(
2155 "dispersion location-scale ({:?}) does not implement analytic spatial \
2156 psi derivatives; the κ/ψ joint optimizer must be explicitly disabled before \
2157 this builder is consulted. Called with data {n_rows}×{n_cols}, mean \
2158 spec (linear={mean_lin}, random={mean_re}, smooth={mean_sm}), noise \
2159 spec (linear={noise_lin}, random={noise_re}, smooth={noise_sm}), \
2160 mean design cols={mean_p}, noise design cols={noise_p}",
2161 self.kind,
2162 n_rows = data.nrows(),
2163 n_cols = data.ncols(),
2164 mean_lin = meanspec.linear_terms.len(),
2165 mean_re = meanspec.random_effect_terms.len(),
2166 mean_sm = meanspec.smooth_terms.len(),
2167 noise_lin = noisespec.linear_terms.len(),
2168 noise_re = noisespec.random_effect_terms.len(),
2169 noise_sm = noisespec.smooth_terms.len(),
2170 mean_p = mean_design.design.ncols(),
2171 noise_p = noise_design.design.ncols(),
2172 ))
2173 }
2174}
2175
2176fn validate_dispersion_family_data(
2187 kind: DispersionFamilyKind,
2188 y: &Array1<f64>,
2189 weights: &Array1<f64>,
2190) -> Result<(), String> {
2191 if y.len() != weights.len() {
2192 return Err(format!(
2193 "{}: response/weights length mismatch: y={}, weights={}",
2194 kind.family_tag(),
2195 y.len(),
2196 weights.len()
2197 ));
2198 }
2199 for (i, &w) in weights.iter().enumerate() {
2200 if !w.is_finite() || w < 0.0 {
2201 return Err(format!(
2202 "{}: prior weights must be finite and non-negative; got weights[{i}] = {w}",
2203 kind.family_tag()
2204 ));
2205 }
2206 }
2207 for (i, &yi) in y.iter().enumerate() {
2208 if weights[i] == 0.0 {
2209 continue;
2210 }
2211 let (ok, requirement) = match kind {
2212 DispersionFamilyKind::NegativeBinomial => (
2213 yi.is_finite() && yi >= 0.0 && yi.fract() == 0.0,
2214 "a finite non-negative integer count",
2215 ),
2216 DispersionFamilyKind::Gamma => (yi.is_finite() && yi > 0.0, "finite and > 0"),
2217 DispersionFamilyKind::Beta => (
2218 yi.is_finite() && yi > 0.0 && yi < 1.0,
2219 "finite and strictly inside (0, 1)",
2220 ),
2221 DispersionFamilyKind::Tweedie { .. } => {
2222 (yi.is_finite() && yi >= 0.0, "finite and >= 0")
2223 }
2224 };
2225 if !ok {
2226 return Err(format!(
2227 "{}: response outside family support at row {i}: y = {yi} (must be {requirement}; \
2228 set the row's prior weight to 0 to exclude it)",
2229 kind.family_tag()
2230 ));
2231 }
2232 }
2233 Ok(())
2234}
2235
2236fn validate_dispersion_spatial_hyperparameter_request(
2248 kind: DispersionFamilyKind,
2249 meanspec: &TermCollectionSpec,
2250 log_dispspec: &TermCollectionSpec,
2251 kappa_options: &SpatialLengthScaleOptimizationOptions,
2252) -> Result<(), GamlssError> {
2253 if !kappa_options.enabled {
2254 return Ok(());
2255 }
2256
2257 let unfrozen_terms = |spec: &TermCollectionSpec| -> Vec<usize> {
2258 spatial_length_scale_term_indices(spec)
2259 .into_iter()
2260 .filter(|&idx| {
2261 let scalar_scale_is_locked = get_spatial_length_scale(spec, idx)
2266 .is_some_and(|scale| scale.is_finite() && scale > 0.0)
2267 && !spatial_term_uses_per_axis_psi(spec, idx);
2268 !scalar_scale_is_locked
2269 })
2270 .collect()
2271 };
2272 let mean_terms = unfrozen_terms(meanspec);
2273 let log_disp_terms = unfrozen_terms(log_dispspec);
2274 if mean_terms.is_empty() && log_disp_terms.is_empty() {
2275 return Ok(());
2276 }
2277 let term_names = |spec: &TermCollectionSpec, indices: &[usize]| -> Vec<String> {
2278 indices
2279 .iter()
2280 .filter_map(|&idx| spec.smooth_terms.get(idx).map(|term| term.name.clone()))
2281 .collect()
2282 };
2283 Err(GamlssError::UnsupportedConfiguration {
2284 reason: format!(
2285 "dispersion location-scale ({kind:?}) cannot optimize spatial hyperparameters: \
2286 exact coupled D_beta H and D_beta H_psi derivatives are unavailable for \
2287 unfrozen spatial terms (mean={:?}, log_precision={:?}). Supply locked spatial \
2288 geometry or explicitly set spatial length-scale optimization enabled=false; the \
2289 fitter will not silently freeze a requested spatial optimization",
2290 term_names(meanspec, &mean_terms),
2291 term_names(log_dispspec, &log_disp_terms),
2292 ),
2293 })
2294}
2295
2296pub fn fit_dispersion_glm_location_scale_terms(
2300 data: ndarray::ArrayView2<'_, f64>,
2301 spec: DispersionGlmLocationScaleTermSpec,
2302 options: &BlockwiseFitOptions,
2303 kappa_options: &SpatialLengthScaleOptimizationOptions,
2304) -> Result<BlockwiseTermFitResult, String> {
2305 if let DispersionFamilyKind::Tweedie { p } = spec.kind {
2306 if !(p.is_finite() && p > 1.0 && p < 2.0) {
2307 return Err(format!(
2308 "Tweedie location-scale requires a variance power strictly in (1, 2); got p={p}"
2309 ));
2310 }
2311 }
2312 validate_dispersion_family_data(spec.kind, &spec.y, &spec.weights)?;
2313 validate_dispersion_spatial_hyperparameter_request(
2314 spec.kind,
2315 &spec.meanspec,
2316 &spec.log_dispspec,
2317 kappa_options,
2318 )?;
2319 let mut options = options.clone();
2337 options.compute_covariance = true;
2338 fit_location_scale_terms(
2339 data,
2340 DispersionGlmLocationScaleTermBuilder {
2341 kind: spec.kind,
2342 y: spec.y,
2343 weights: spec.weights,
2344 meanspec: spec.meanspec,
2345 noisespec: spec.log_dispspec,
2346 mean_offset: spec.mean_offset,
2347 noise_offset: spec.log_disp_offset,
2348 },
2349 &options,
2350 kappa_options,
2351 )
2352}
2353
2354#[cfg(test)]
2355mod tests {
2356 use super::test_support::{dispersion_gamma_nll_order2, dispersion_nb_nll_order2};
2357 use super::*;
2358 use crate::gamlss::test_support::dispersion_tweedie_nll_generic;
2359 use gam_math::nested_dual::JetField;
2360
2361 #[test]
2362 fn saved_alo_gamma_row_geometry_matches_closed_form_and_keeps_meat_distinct() {
2363 let y = 4.0;
2364 let mu: f64 = 2.0;
2365 let nu: f64 = 3.0;
2366 let weight = 1.7;
2367 let geometry = dispersion_alo_row_geometry(
2368 DispersionFamilyKind::Gamma,
2369 0,
2370 y,
2371 mu.ln(),
2372 nu.ln(),
2373 weight,
2374 )
2375 .expect("Gamma row geometry must be representable");
2376
2377 let ratio = y / mu;
2378 let a = gam_math::jet_tower::digamma(nu) - nu.ln() - 1.0 + mu.ln() - y.ln() + ratio;
2379 let expected_score = [weight * nu * (1.0 - ratio), weight * nu * a];
2380 let expected_hessian = [
2381 [weight * nu * ratio, weight * nu * (1.0 - ratio)],
2382 [
2383 weight * nu * (1.0 - ratio),
2384 weight * nu * (a + nu * gam_math::jet_tower::trigamma(nu) - 1.0),
2385 ],
2386 ];
2387 for coordinate in 0..2 {
2388 assert_close(
2389 "Gamma ALO score",
2390 geometry.nll_score[coordinate],
2391 expected_score[coordinate],
2392 2e-12,
2393 );
2394 for other in 0..2 {
2395 assert_close(
2396 "Gamma ALO observed Hessian",
2397 geometry.observed_hessian[coordinate][other],
2398 expected_hessian[coordinate][other],
2399 2e-12,
2400 );
2401 }
2402 }
2403
2404 let score_meat = [
2405 [
2406 expected_score[0] * expected_score[0],
2407 expected_score[0] * expected_score[1],
2408 ],
2409 [
2410 expected_score[1] * expected_score[0],
2411 expected_score[1] * expected_score[1],
2412 ],
2413 ];
2414 assert_ne!(
2415 geometry.observed_hessian, score_meat,
2416 "the deletion curvature must not be replaced by score covariance"
2417 );
2418 }
2419
2420 #[inline]
2427 fn order1_ln_gamma<const K: usize>(
2428 x: &gam_math::jet_scalar::Order1<K>,
2429 ) -> gam_math::jet_scalar::Order1<K> {
2430 x.compose_unary([
2431 ln_gamma(x.v),
2432 gam_math::jet_tower::digamma(x.v),
2433 0.0,
2434 0.0,
2435 0.0,
2436 ])
2437 }
2438
2439 #[inline]
2447 fn dispersion_nb_disp_order1(
2448 yi: f64,
2449 mu_value: f64,
2450 theta_value: f64,
2451 wi: f64,
2452 ) -> gam_math::jet_scalar::Order1<1> {
2453 type O1 = gam_math::jet_scalar::Order1<1>;
2454
2455 let mu = O1::constant(mu_value);
2456 let theta = O1::variable(theta_value, 0);
2457 let tpm = theta.add(&mu);
2458 let theta_plus_y = theta.add(&O1::constant(yi));
2459 let loglik = order1_ln_gamma(&theta_plus_y)
2460 .sub(&order1_ln_gamma(&theta))
2461 .sub(&O1::constant(ln_gamma(yi + 1.0)))
2462 .add(&theta.mul(&theta.ln()))
2463 .sub(&theta.mul(&tpm.ln()))
2464 .add(&mu.ln().scale(yi))
2465 .sub(&tpm.ln().scale(yi));
2466 loglik.scale(-wi)
2467 }
2468
2469 #[inline]
2475 fn dispersion_nb_disp_order2(
2476 yi: f64,
2477 mu_value: f64,
2478 theta_value: f64,
2479 wi: f64,
2480 ) -> gam_math::jet_scalar::Order2<1> {
2481 use gam_math::jet_scalar::JetScalar;
2482 use statrs::function::gamma::ln_gamma;
2483 type O1 = gam_math::jet_scalar::Order2<1>;
2484
2485 let mu = O1::constant(mu_value);
2486 let theta = O1::variable(theta_value, 0);
2487 let tpm = theta.add(&mu);
2488 let theta_plus_y = theta.add(&O1::constant(yi));
2489 let loglik = order2_ln_gamma(&theta_plus_y)
2490 .sub(&order2_ln_gamma(&theta))
2491 .sub(&O1::constant(ln_gamma(yi + 1.0)))
2492 .add(&theta.mul(&theta.ln()))
2493 .sub(&theta.mul(&tpm.ln()))
2494 .add(&mu.ln().scale(yi))
2495 .sub(&tpm.ln().scale(yi));
2496 loglik.scale(-wi)
2497 }
2498
2499 pub(crate) fn beta_fisher_cross_info_mu_phi(mu: f64, phi: f64) -> f64 {
2500 let a = mu * phi;
2501 let b = (1.0 - mu) * phi;
2502 phi * (mu * gam_math::jet_tower::trigamma_derivative_stack(a)[0]
2503 - (1.0 - mu) * gam_math::jet_tower::trigamma_derivative_stack(b)[0])
2504 }
2505
2506 pub(crate) fn assert_close(label: &str, got: f64, want: f64, tol: f64) {
2507 assert!(
2508 (got - want).abs() <= tol,
2509 "{label}: got {got:.12e}, want {want:.12e}, |diff|={:.3e}",
2510 (got - want).abs()
2511 );
2512 }
2513
2514 #[test]
2515 fn spatial_hyperparameter_request_is_a_typed_error_until_explicitly_frozen() {
2516 let locked_meanspec = crate::gamlss::tests::simple_matern_term_collection(&[0, 1], 0.6);
2517 let mut meanspec = locked_meanspec.clone();
2518 let gam_terms::smooth::SmoothBasisSpec::Matern { spec, .. } =
2519 &mut meanspec.smooth_terms[0].basis
2520 else {
2521 panic!("test fixture must contain a Matérn term");
2522 };
2523 spec.aniso_log_scales = Some(vec![0.0, 0.0]);
2524 let log_dispspec = crate::gamlss::tests::empty_term_collection();
2525 let enabled = SpatialLengthScaleOptimizationOptions::default();
2526
2527 let error = validate_dispersion_spatial_hyperparameter_request(
2528 DispersionFamilyKind::Gamma,
2529 &meanspec,
2530 &log_dispspec,
2531 &enabled,
2532 )
2533 .expect_err("enabled dispersion spatial optimization must be rejected");
2534 assert!(matches!(
2535 error,
2536 GamlssError::UnsupportedConfiguration { .. }
2537 ));
2538
2539 let n = 8;
2540 let public_error = match fit_dispersion_glm_location_scale_terms(
2541 Array2::zeros((n, 2)).view(),
2542 DispersionGlmLocationScaleTermSpec {
2543 kind: DispersionFamilyKind::Gamma,
2544 y: Array1::from_elem(n, 1.0),
2545 weights: Array1::from_elem(n, 1.0),
2546 meanspec: meanspec.clone(),
2547 log_dispspec: log_dispspec.clone(),
2548 mean_offset: Array1::zeros(n),
2549 log_disp_offset: Array1::zeros(n),
2550 },
2551 &BlockwiseFitOptions::default(),
2552 &enabled,
2553 ) {
2554 Ok(_) => panic!("public fit must not silently freeze spatial optimization"),
2555 Err(error) => error,
2556 };
2557 assert!(public_error.contains("will not silently freeze"));
2558
2559 validate_dispersion_spatial_hyperparameter_request(
2560 DispersionFamilyKind::Gamma,
2561 &locked_meanspec,
2562 &log_dispspec,
2563 &enabled,
2564 )
2565 .expect("a caller-supplied locked spatial scale is explicit frozen geometry");
2566
2567 let auto_meanspec = crate::gamlss::tests::simple_matern_term_collection(&[0, 1], 0.0);
2568 assert!(matches!(
2569 validate_dispersion_spatial_hyperparameter_request(
2570 DispersionFamilyKind::Gamma,
2571 &auto_meanspec,
2572 &log_dispspec,
2573 &enabled,
2574 ),
2575 Err(GamlssError::UnsupportedConfiguration { .. })
2576 ));
2577
2578 let mut frozen = enabled;
2579 frozen.enabled = false;
2580 validate_dispersion_spatial_hyperparameter_request(
2581 DispersionFamilyKind::Gamma,
2582 &meanspec,
2583 &log_dispspec,
2584 &frozen,
2585 )
2586 .expect("an explicit frozen-geometry request is supported");
2587
2588 validate_dispersion_spatial_hyperparameter_request(
2589 DispersionFamilyKind::Gamma,
2590 &log_dispspec,
2591 &log_dispspec,
2592 &SpatialLengthScaleOptimizationOptions::default(),
2593 )
2594 .expect("enabled spatial optimization is irrelevant without spatial coordinates");
2595 }
2596
2597 #[test]
2598 pub(crate) fn beta_tower_mixed_channel_matches_cross_information_formula() {
2599 let mu = 0.1;
2600 let phi = 10.0;
2601 let a = mu * phi;
2602 let b = (1.0 - mu) * phi;
2603 let digamma_a = gam_math::jet_tower::digamma_derivative_stack(a)[0];
2604 let digamma_b = gam_math::jet_tower::digamma_derivative_stack(b)[0];
2605 let score_neutral_y = 1.0 / (1.0 + (-(digamma_a - digamma_b)).exp());
2606
2607 let tower = dispersion_beta_nll_order2(score_neutral_y, mu, phi, 1.0);
2608 let trigamma_a = std::f64::consts::PI * std::f64::consts::PI / 6.0;
2609 let trigamma_b = gam_math::jet_tower::trigamma_derivative_stack(b)[0];
2610 let analytic = phi * (mu * trigamma_a - (1.0 - mu) * trigamma_b);
2611 let helper = beta_fisher_cross_info_mu_phi(mu, phi);
2612
2613 assert!(
2614 analytic > 0.58,
2615 "audit example should have visibly nonzero cross information, got {analytic}"
2616 );
2617 assert_close("helper cross information", helper, analytic, 1e-12);
2618 assert_close("tower mixed channel", tower.h()[0][1], analytic, 1e-8);
2619
2620 let q = mu * (1.0 - mu);
2623 let em = (mu / (1.0 - mu)).ln();
2624 let ed = phi.ln();
2625 let eta_tower =
2626 dispersion_eta_nll_order2(DispersionFamilyKind::Beta, score_neutral_y, em, ed, 1.0);
2627 assert_close(
2628 "eta-scale observed cross curvature",
2629 eta_tower.h()[0][1],
2630 q * phi * analytic,
2631 1e-8,
2632 );
2633 }
2634
2635 #[test]
2639 pub(crate) fn order2_matches_dense_tower_all_channels() {
2640 use gam_math::jet_scalar::Order2;
2641 use gam_math::jet_tower::Tower4;
2642
2643 fn check_o2_vs_tower4(label: &str, o2: Order2<2>, t4: Tower4<2>) {
2644 let band = |a: f64, b: f64| 1e-9 + 1e-9 * a.abs().max(b.abs());
2645 assert!(
2646 (o2.value() - t4.v).abs() <= band(o2.value(), t4.v),
2647 "{label} value: {} vs {}",
2648 o2.value(),
2649 t4.v
2650 );
2651 for a in 0..2 {
2652 assert!(
2653 (o2.g()[a] - t4.g[a]).abs() <= band(o2.g()[a], t4.g[a]),
2654 "{label} grad[{a}]: {} vs {}",
2655 o2.g()[a],
2656 t4.g[a]
2657 );
2658 for b in 0..2 {
2659 assert!(
2660 (o2.h()[a][b] - t4.h[a][b]).abs() <= band(o2.h()[a][b], t4.h[a][b]),
2661 "{label} hess[{a}][{b}]: {} vs {}",
2662 o2.h()[a][b],
2663 t4.h[a][b]
2664 );
2665 }
2666 }
2667 }
2668
2669 let wi = 1.7_f64;
2670 for &(yi, mu, theta) in &[(0.0, 1.2, 3.0), (4.0, 2.5, 0.7), (10.0, 0.6, 5.0)] {
2672 check_o2_vs_tower4(
2673 "nb",
2674 dispersion_nb_nll_order2(yi, mu, theta, wi),
2675 test_support::dispersion_nb_nll_generic::<Tower4<2>>(yi, mu, theta, wi),
2676 );
2677 }
2678 for &(yi, mu, nu) in &[
2680 (0.5_f64, 1.1_f64, 2.0_f64),
2681 (3.0, 4.0, 0.9),
2682 (1.0, 0.3, 6.0),
2683 ] {
2684 let y_pos = yi.max(1e-300);
2685 check_o2_vs_tower4(
2686 "gamma",
2687 dispersion_gamma_nll_order2(yi, y_pos, mu, nu, wi),
2688 test_support::dispersion_gamma_nll_generic::<Tower4<2>>(yi, y_pos, mu, nu, wi),
2689 );
2690 }
2691 for &(yi, mu, phi) in &[(0.3, 0.4, 5.0), (0.9, 0.6, 12.0), (0.01, 0.2, 3.0)] {
2693 check_o2_vs_tower4(
2694 "beta",
2695 dispersion_beta_nll_order2(yi, mu, phi, wi),
2696 test_support::dispersion_beta_nll_generic::<Tower4<2>>(yi, mu, phi, wi),
2697 );
2698 }
2699 for &(yi, eta_mu, eta_d, p) in &[
2701 (0.0, 0.4, -0.3, 1.5),
2702 (2.5, -0.2, 0.5, 1.3),
2703 (0.0, 1.0, 0.1, 1.7),
2704 (5.0, 0.7, -0.6, 1.6),
2705 ] {
2706 check_o2_vs_tower4(
2707 "tweedie",
2708 dispersion_tweedie_nll_generic::<Order2<2>>(yi, eta_mu, eta_d, p, wi),
2709 dispersion_tweedie_nll_generic::<Tower4<2>>(yi, eta_mu, eta_d, p, wi),
2710 );
2711 }
2712 }
2713
2714 #[test]
2720 pub(crate) fn pruned_disp_towers_bit_identical_to_full_order2() {
2721 use gam_math::jet_scalar::Order2;
2722
2723 let mut state: u64 = 0x9E3779B97F4A7C15;
2725 let mut next = || {
2726 state = state
2727 .wrapping_mul(6364136223846793005)
2728 .wrapping_add(1442695040888963407);
2729 ((state >> 11) as f64) / ((1u64 << 53) as f64)
2730 };
2731 let bits = |x: f64| x.to_bits();
2732
2733 let n_per = 600; for _ in 0..n_per {
2735 let wi = 0.25 + 3.0 * next();
2736 let yi_count = (next() * 12.0).floor();
2737
2738 {
2740 let mu = (0.05 + 4.0 * next()).max(1e-300);
2741 let theta = (0.05 + 6.0 * next()).max(1e-12);
2742 let full = dispersion_nb_nll_order2(yi_count, mu, theta, wi);
2743 let prn = dispersion_nb_disp_order2(yi_count, mu, theta, wi);
2744 assert_eq!(bits(full.value()), bits(prn.value()), "nb value");
2745 assert_eq!(bits(full.g()[1]), bits(prn.g()[0]), "nb grad");
2746 assert_eq!(bits(full.h()[1][1]), bits(prn.h()[0][0]), "nb hess");
2747 let prn1 = dispersion_nb_disp_order1(yi_count, mu, theta, wi);
2752 assert_eq!(bits(prn.value()), bits(prn1.value()), "nb order1 value");
2753 assert_eq!(bits(prn.g()[0]), bits(prn1.g()[0]), "nb order1 grad");
2754 assert_close(
2756 "nb stable value-only",
2757 dispersion_nb_loglik(yi_count, mu, theta, wi),
2758 -prn.value(),
2759 1e-12,
2760 );
2761 }
2762 {
2764 let mu = (0.05 + 4.0 * next()).max(1e-300);
2765 let nu = (0.05 + 6.0 * next()).max(1e-12);
2766 let yi = 0.01 + 8.0 * next();
2767 let y_pos = yi.max(1e-300);
2768 let full = dispersion_gamma_nll_order2(yi, y_pos, mu, nu, wi);
2769 let prn = dispersion_gamma_disp_order2(yi, y_pos, mu, nu, wi);
2770 assert_eq!(bits(full.value()), bits(prn.value()), "gamma value");
2771 assert_eq!(bits(full.g()[1]), bits(prn.g()[0]), "gamma grad");
2772 assert_eq!(bits(full.h()[1][1]), bits(prn.h()[0][0]), "gamma hess");
2773 assert_eq!(
2774 bits(dispersion_gamma_loglik(yi, y_pos, mu, nu, wi)),
2775 bits(-prn.value()),
2776 "gamma value-only"
2777 );
2778 }
2779 {
2781 let mu = (1e-6 + (1.0 - 2e-6) * next()).clamp(1e-12, 1.0 - 1e-12);
2782 let phi = (0.05 + 20.0 * next()).max(1e-12);
2783 let yi = next();
2784 let full = dispersion_beta_nll_order2(yi, mu, phi, wi);
2785 assert_eq!(
2786 bits(dispersion_beta_loglik(yi, mu, phi, wi)),
2787 bits(-full.value()),
2788 "beta value-only"
2789 );
2790 }
2791 for &(yi, eta_mu, eta_d, p) in &[
2793 (
2794 0.0_f64,
2795 -4.0 + 8.0 * next(),
2796 -4.0 + 8.0 * next(),
2797 1.1 + 0.8 * next(),
2798 ),
2799 (
2800 0.01 + 9.0 * next(),
2801 -4.0 + 8.0 * next(),
2802 -4.0 + 8.0 * next(),
2803 1.1 + 0.8 * next(),
2804 ),
2805 (3.0, -8.0, 8.0, 1.5),
2806 ] {
2807 let em = eta_mu;
2808 let ed = eta_d;
2809 let full = dispersion_tweedie_nll_generic::<Order2<2>>(yi, em, ed, p, wi);
2810 let prn = dispersion_tweedie_disp_order2(yi, em, ed, p, wi);
2811 assert_eq!(bits(full.value()), bits(prn.value()), "tweedie value");
2812 assert_eq!(bits(full.g()[1]), bits(prn.g()[0]), "tweedie grad");
2813 assert_eq!(bits(full.h()[1][1]), bits(prn.h()[0][0]), "tweedie hess");
2814 assert_eq!(
2815 bits(dispersion_tweedie_loglik(yi, em, ed, p, wi)),
2816 bits(-prn.value()),
2817 "tweedie value-only"
2818 );
2819 }
2820 }
2821 }
2822
2823 #[test]
2829 pub(crate) fn observed_eta_hessian_matches_gamma_closed_form() {
2830 let (yi, mu, nu): (f64, f64, f64) = (4.0, 2.0, 3.0);
2831 let (h_mm, h_md, h_dd) = dispersion_row_observed_hessian_weights(
2832 DispersionFamilyKind::Gamma,
2833 yi,
2834 mu.ln(),
2835 nu.ln(),
2836 1.0,
2837 );
2838 assert_close("gamma observed d2/d_eta_mu2", h_mm, nu * yi / mu, 1e-10);
2839 assert_close(
2840 "gamma observed cross d2/d_eta_mu d_eta_nu",
2841 h_md,
2842 nu * (1.0 - yi / mu),
2843 1e-10,
2844 );
2845 let nll =
2849 |ed: f64| -dispersion_row_loglik(DispersionFamilyKind::Gamma, yi, mu.ln(), ed, 1.0);
2850 let h = 1e-5;
2851 let ed0 = nu.ln();
2852 let fd = (nll(ed0 + h) - 2.0 * nll(ed0) + nll(ed0 - h)) / (h * h);
2853 assert_close("gamma observed d2/d_eta_nu2 (FD)", h_dd, fd, 1e-4);
2854 }
2855
2856 #[test]
2859 pub(crate) fn observed_eta_hessian_is_exact_beyond_former_clamp() {
2860 let (h_mm, h_md, h_dd) = dispersion_row_observed_hessian_weights(
2861 DispersionFamilyKind::Gamma,
2862 4.0,
2863 35.0,
2864 0.5,
2865 1.0,
2866 );
2867 assert!(h_mm > 0.0);
2868 assert!(h_md.is_finite());
2869 assert!(h_dd != 0.0);
2870 let kernel = dispersion_row_kernel(DispersionFamilyKind::Gamma, 4.0, 35.0, 0.5, 1.0);
2871 assert!(kernel.mean_weight > 0.0);
2872 assert_ne!(kernel.mean_response, 35.0);
2873 assert!(kernel.disp_weight > 0.0);
2874 }
2875
2876 #[test]
2877 fn negative_binomial_balanced_ratios_and_precision_information_keep_tail_geometry() {
2878 let huge = 1.0e200_f64;
2879 let kernel = dispersion_row_kernel(
2880 DispersionFamilyKind::NegativeBinomial,
2881 1.0,
2882 huge.ln(),
2883 huge.ln(),
2884 1.0,
2885 );
2886 assert!(kernel.loglik.is_finite());
2887 assert!(kernel.mean_weight.is_finite());
2888 let precision = huge.ln().exp();
2897 assert!((kernel.mean_weight / precision - 0.5).abs() <= 8.0 * f64::EPSILON);
2898 assert!(kernel.disp_weight.is_finite() && kernel.disp_weight > 0.0);
2899
2900 let eta_info = nb_log_precision_fisher_jensen(1.0, 1.0e17);
2901 assert!(eta_info.is_finite() && eta_info > 0.0);
2902 assert!((eta_info * 1.0e17 - 1.0).abs() < 1.0e-12);
2903
2904 let log_share = log_positive_share((-700.0_f64).exp(), 700.0_f64.exp());
2905 assert!(log_share.is_finite());
2906 assert!((log_share + 1400.0).abs() < 1.0e-12);
2907 }
2908
2909 #[test]
2919 pub(crate) fn parallel_evaluate_matches_serial_reference() {
2920 let n = DISPERSION_PARALLEL_ROW_THRESHOLD * 3 + 7;
2921 let mut state: u64 = 0xD1B5_4A32_D192_ED03;
2923 let mut next = || {
2924 state = state
2925 .wrapping_mul(6364136223846793005)
2926 .wrapping_add(1442695040888963407);
2927 ((state >> 11) as f64) / ((1u64 << 53) as f64)
2928 };
2929
2930 for kind in [
2931 DispersionFamilyKind::NegativeBinomial,
2932 DispersionFamilyKind::Gamma,
2933 DispersionFamilyKind::Beta,
2934 DispersionFamilyKind::Tweedie { p: 1.5 },
2935 ] {
2936 let y = Array1::from_shape_fn(n, |_| match kind {
2937 DispersionFamilyKind::Beta => 1e-3 + (1.0 - 2e-3) * next(),
2938 DispersionFamilyKind::NegativeBinomial => (next() * 12.0).floor(),
2939 _ => 0.05 + 8.0 * next(),
2940 });
2941 let weights = Array1::from_shape_fn(n, |_| 0.25 + 2.0 * next());
2942 let eta_mu = Array1::from_shape_fn(n, |_| -1.0 + 2.0 * next());
2943 let eta_d = Array1::from_shape_fn(n, |_| -1.0 + 2.0 * next());
2944
2945 let family = DispersionGlmLocationScaleFamily {
2946 kind,
2947 y: y.clone(),
2948 weights: weights.clone(),
2949 };
2950 let states = vec![
2951 ParameterBlockState {
2952 beta: Array1::zeros(0),
2953 eta: eta_mu.clone(),
2954 },
2955 ParameterBlockState {
2956 beta: Array1::zeros(0),
2957 eta: eta_d.clone(),
2958 },
2959 ];
2960
2961 let mut ll_ref = 0.0;
2963 let mut mw_ref = Array1::<f64>::zeros(n);
2964 let mut mr_ref = Array1::<f64>::zeros(n);
2965 let mut dw_ref = Array1::<f64>::zeros(n);
2966 let mut dr_ref = Array1::<f64>::zeros(n);
2967 for i in 0..n {
2968 let row = dispersion_row_kernel(kind, y[i], eta_mu[i], eta_d[i], weights[i]);
2969 ll_ref += row.loglik;
2970 mw_ref[i] = row.mean_weight;
2971 mr_ref[i] = row.mean_response;
2972 dw_ref[i] = row.disp_weight;
2973 dr_ref[i] = row.disp_response;
2974 }
2975
2976 let eval = family.evaluate(&states).expect("parallel evaluate");
2977 assert_close(
2978 &format!("{kind:?} evaluate log-likelihood"),
2979 eval.log_likelihood,
2980 ll_ref,
2981 1e-9,
2982 );
2983
2984 let BlockWorkingSet::Diagonal {
2985 working_response: mr,
2986 working_weights: mw,
2987 } = &eval.blockworking_sets[0]
2988 else {
2989 panic!("mean block not diagonal");
2990 };
2991 let BlockWorkingSet::Diagonal {
2992 working_response: dr,
2993 working_weights: dw,
2994 } = &eval.blockworking_sets[1]
2995 else {
2996 panic!("dispersion block not diagonal");
2997 };
2998 for i in 0..n {
2999 assert_close("mean weight", mw[i], mw_ref[i], 1e-9);
3000 assert_close("mean response", mr[i], mr_ref[i], 1e-9);
3001 assert_close("disp weight", dw[i], dw_ref[i], 1e-9);
3002 assert_close("disp response", dr[i], dr_ref[i], 1e-9);
3003 }
3004
3005 let ll_only = family
3008 .log_likelihood_only(&states)
3009 .expect("parallel log_likelihood_only");
3010 assert_close(
3011 &format!("{kind:?} log_likelihood_only"),
3012 ll_only,
3013 ll_ref,
3014 1e-9,
3015 );
3016 }
3017 }
3018}