Skip to main content

gam_models/gamlss/
dispersion_family.rs

1//! #913: dispersion-channel GAMLSS location-scale families.
2//!
3//! Extracted from `gamlss.rs` (issue #780); this module now owns the
4//! dispersion-channel joint-curvature corrections.
5
6use super::weighted_design_products::{mirror_upper_to_lower, xt_diag_x_design, xt_diag_y_design};
7// Concrete `Order2` algebra methods live on the shared `JetField` supertrait;
8// keep it in scope alongside `JetScalar` so the row programs resolve them.
9use 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// ============================================================================
32// #913: dispersion-channel GAMLSS location-scale families.
33//
34// `noise_formula` (a second linear predictor on the dispersion channel) was
35// wired only for Gaussian/Binomial location-scale and the survival families.
36// The genuine-dispersion mean families — NegativeBinomial, Gamma, Beta and
37// Tweedie — were mean-only with a single scalar dispersion. This module adds a
38// SINGLE generic two-block family that routes all four through the existing
39// blockwise REML engine and the shared `LocationScaleFamilyBuilder` /
40// `fit_location_scale_terms` plumbing, so the κ-coordinate assembly, warm
41// start, shrinkage-penalised scale block and result extraction are reused
42// verbatim. A family is added by supplying only its per-row log-likelihood and
43// the (mean, log-precision) working sets — everything else is shared.
44//
45// Block layout: block 0 = mean predictor (η_μ, log link for NB/Gamma/Tweedie,
46// logit for Beta); block 1 = log-precision predictor (η_d). The dispersion
47// channel models log(precision) uniformly — `θ` for NegativeBinomial, the
48// shape `ν` for Gamma, `φ` for Beta, and `1/φ` for Tweedie — so a larger η_d
49// always means *less* dispersion, matching the Gaussian/Binomial convention
50// where η_logσ smaller ⇒ tighter. With no `noise_formula` the log-precision
51// block is a single intercept and the fit reduces to the scalar-dispersion
52// model.
53//
54// NB2 with `(μ, θ)` and the exponential-dispersion members here with
55// `(μ, φ)` are Fisher-orthogonal in their standard mean/dispersion
56// parameterizations: Gamma uses shape `ν = 1/φ`, and Tweedie models
57// `log(1/φ)`, so those precision-channel transforms preserve zero expected
58// mean/dispersion cross information. Beta is the exception in this module's
59// mean/precision parameterization. For `Beta(μφ, (1−μ)φ)`,
60//
61//   I_{μ,φ} = φ · (μ ψ'(μφ) − (1−μ) ψ'((1−μ)φ)),
62//
63// so in predictor coordinates `(η_μ = logit μ, η_φ = log φ)` the Fisher cross
64// block is
65//
66//   I_{η_μ,η_φ} = μ(1−μ) φ² · (μ ψ'(μφ) − (1−μ) ψ'((1−μ)φ)),
67//
68// which is generically nonzero. Block-cyclic Fisher-scoring IRLS is still a
69// valid block coordinate solve for the point estimate, but joint-curvature
70// consumers (`log|H|`, coefficient covariance, posterior draws) must receive
71// Beta's off-diagonal coefficient block. Smoothing-parameter selection still
72// runs through the engine's first-order (gradient-only) outer path: the family
73// declines the dense outer Hessian capability because its working weights
74// couple the two blocks (`W_μ` depends on the precision and vice-versa), which
75// the block-local diagonal-drift hook cannot represent exactly.
76// ============================================================================
77
78/// The genuine-dispersion mean family whose precision (overdispersion) channel
79/// can carry a second `noise_formula` linear predictor (issue #913).
80#[derive(Clone, Copy, Debug, PartialEq)]
81pub enum DispersionFamilyKind {
82    /// NB2: `Var = μ + μ²/θ`; the precision channel models `log θ`.
83    NegativeBinomial,
84    /// Gamma with `Var = μ²/ν`; the precision channel models `log ν` (shape).
85    Gamma,
86    /// Beta(μφ, (1−μ)φ) with a logit mean link; the precision channel models
87    /// `log φ`.
88    Beta,
89    /// Tweedie compound Poisson–Gamma with `Var = φ μ^p`, fixed power `p`; the
90    /// precision channel models `log(1/φ)`. The per-row density uses the
91    /// saddlepoint (Nelder–Pregibon) approximation for `y > 0` and the exact
92    /// point mass at `y = 0`; this is the standard tractable Tweedie ML
93    /// surface (an exact-series φ-derivative is the remaining hard sub-item of
94    /// #913).
95    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    /// The mean link is logit for Beta (a probability mean) and log otherwise.
109    pub(crate) const fn mean_is_logit(self) -> bool {
110        matches!(self, DispersionFamilyKind::Beta)
111    }
112
113    /// The mean inverse link this dispersion family fits on: log for
114    /// NegativeBinomial / Gamma / Tweedie, logit for Beta. Single source of
115    /// truth shared by the CLI and FFI save paths so the persisted
116    /// `base_link` never diverges from the fitted channel.
117    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    /// The family's canonical `LikelihoodSpec` (mean response × mean link).
127    /// The overdispersion parameter is estimated by the log-precision channel,
128    /// so the response-family placeholder parameters (`phi`, `theta`) mirror
129    /// the `resolve_family` defaults
130    /// and are not consumed as fixed values at predict time. This is the single
131    /// source of truth for the persisted location-scale likelihood so the CLI
132    /// and FFI save paths cannot diverge.
133    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
158/// Row count above which the per-row dispersion-kernel map fans out across
159/// rayon workers (only when not already running on a worker, to avoid nested
160/// oversubscription). Below it the serial map beats the fork/join overhead.
161/// Mirrors the row-chunk guard in
162/// [`row_coeff_operator`](super::gaussian::row_coeff_operator).
163const DISPERSION_PARALLEL_ROW_THRESHOLD: usize = 1024;
164
165/// Per-row working quantities for both channels at the current `(η_μ, η_d)`.
166pub(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
185/// Certify the exact open parameter domain used by the row towers.  The domain
186/// is defined by representability of the linked distribution parameters, not
187/// by an arbitrary predictor box.
188fn 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    /// Test-oracle NB2 row NLL over a generic [`JetScalar<2>`], seeded on the
336    /// natural parameters `(μ, θ)`.
337    #[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        // (theta + yi).ln_gamma() - theta.ln_gamma() - ln_gamma(yi+1)
348        //   + theta*theta.ln() - theta*tpm.ln() + mu.ln()*yi - tpm.ln()*yi
349        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    /// Test-oracle Gamma row NLL over a generic [`JetScalar<2>`], seeded on
362    /// `(μ, ν)`.
363    #[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        // nu*nu.ln() - nu*mu.ln() - nu.ln_gamma() + (nu-1)*y_pos.ln() - nu*(mu.recip()*yi)
374        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    /// Test-oracle Beta row NLL over a generic [`JetScalar<2>`], seeded on
384    /// `(μ, φ)`.
385    #[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        // phi.ln_gamma() - a.ln_gamma() - b.ln_gamma()
399        //   + (a-1)*yc.ln() + (b-1)*(1-yc).ln()
400        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    /// #1591 jet-prune oracle: full `Order2<2>` (value/grad/Hessian) NB2 row NLL.
410    ///
411    /// Production no longer consumes the mean (`μ`-axis) derivative channels of
412    /// this tower — the NB mean block is Fisher-orthogonal and hand-written
413    /// exactly in [`dispersion_row_kernel`] — so the hot path uses the pruned
414    /// single-axis [`dispersion_nb_disp_order2`] instead. This `K=2` form
415    /// survives only as the dense-`Tower4<2>` oracle pin
416    /// (`order2_matches_dense_tower_all_channels`).
417    #[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    /// #1591 jet-prune oracle: full `Order2<2>` Gamma row NLL. As with NB, the
441    /// mean axis is unused in production (hand-written, Fisher-orthogonal); the
442    /// hot path uses the single-axis [`dispersion_gamma_disp_order2`]. Kept only
443    /// as the dense-tower oracle pin.
444    #[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/// Production `Order2<2>` Beta row NLL (value/grad/Hessian hot path; the cross
467/// channel `h()[0][1]` feeds the Beta observed cross weight).
468#[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// ============================================================================
501// #1591 jet-prune: single-axis (`K=1`) dispersion-channel towers.
502//
503// For NegativeBinomial / Gamma / Tweedie the production row kernel consumes ONLY
504// the dispersion-axis derivatives (`g[disp]`, `h[disp][disp]`) and the value;
505// the mean block is Fisher-orthogonal and assembled in closed form. Seeding the
506// mean as a CONSTANT and the dispersion parameter as the SOLE jet variable
507// therefore yields a tower whose `(value, g[0], h[0][0])` are `to_bits`-
508// identical to the consumed `(value, g[1], h[1][1])` of the old `Order2<2>`
509// tower — the mean seed only ever populated the now-discarded `g[mean]` /
510// `h[mean][·]` channels (Leibniz/Faà-di-Bruno never read the dispersion-axis
511// channels off the mean seed). Collapsing `K=2 → K=1` quarters the Hessian
512// tensor (1 entry vs 4) and halves the gradient, with no change to any consumed
513// float bit. The `ln_gamma` derivative stacks are unchanged (the irreducible
514// transcendental cost), so this trims the rational composition, not the special
515// functions.
516// ============================================================================
517
518// `dispersion_nb_disp_order1` / `dispersion_nb_disp_order2` — the pruned
519// `Order1<1>` / `Order2<1>` NB2 dispersion oracle pins used only by
520// `prune_towers_*` — now live in the `#[cfg(test)] mod tests` below, next to
521// their sole consumer (they are genuinely test-support, not production —
522// the real NB2 dispersion row kernel below computes score/curvature in
523// closed form via `digamma`/`nb_log_precision_fisher_jensen`, never through
524// either jet tower — so they belong in a `#[cfg(test)]` mod rather than
525// carrying `allow(dead_code)`).
526
527/// Pruned single-axis Gamma dispersion tower: `ν` is the sole jet variable
528/// (axis 0), `μ` a constant. Consumed channels match
529/// `dispersion_gamma_nll_order2` index-1 bit-for-bit.
530#[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/// Pruned single-axis Tweedie dispersion tower seeded on the predictor `η_d`
552/// (axis 0), with `η_μ` a constant (so `μ = exp(η_μ)` carries no jet). The
553/// `φ = exp(−η_d)` chain and its nonlinear `∂²φ/∂η_d²` curvature are carried
554/// exactly as in `dispersion_tweedie_nll_generic`; `value`/`g[0]`/`h[0][0]`
555/// match that program's `value`/`g[1]`/`h[1][1]` bit-for-bit.
556#[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// ============================================================================
592// #1591 jet-prune: value-only (`K=0`) row negative-log-likelihood.
593//
594// `log_likelihood_only` reads ONLY `row.loglik = -tower.value()`; the full row
595// kernel it used to call evaluated every dispersion tower's gradient AND Hessian
596// — including the digamma/trigamma derivative stacks — purely to discard them.
597// These functions evaluate the SAME value-channel program in plain `f64`, so
598// they are `to_bits`-identical to `-tower.value()` (the jet value channel is the
599// naive scalar evaluation: `mul.v = a.v*b.v`, `compose.v = stack[0]`), while
600// touching only `ln_gamma` (stack slot 0) and never the digamma/trigamma slots.
601// On a per-row loglik that is the dominant transcendental saving.
602// ============================================================================
603
604/// NB2 row log-likelihood, evaluated through stable log shares so `mu+theta`
605/// is never formed and may mathematically exceed `f64::MAX`.
606#[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/// `log(numerator / (numerator + other))` without forming the potentially
617/// overflowing sum or subtracting nearly equal logarithms.
618#[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/// Jensen NB precision Fisher information already transformed to log-precision
639/// coordinates, `theta^2 I_theta`.  For large theta, expand
640/// `trigamma(x)-1/x` after the transformation so the representable O(1)
641/// result is never obtained by subtracting underflowed O(theta^-2) terms.
642#[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/// Gamma row log-likelihood, plain `f64`, bit-identical to
666/// `-dispersion_gamma_disp_order2(..).value()`.
667#[inline]
668fn dispersion_gamma_loglik(yi: f64, y_pos: f64, mu: f64, nu: f64, wi: f64) -> f64 {
669    // NB: the jet forms `μ.recip().scale(yi)` = `(1/μ)·yᵢ` (reciprocal then
670    // multiply), NOT `yᵢ/μ` (single divide) — these differ in the last bit, so
671    // the value path must reproduce the reciprocal-then-multiply exactly.
672    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/// Beta row log-likelihood, plain `f64`, bit-identical to
678/// `-dispersion_beta_nll_order2(..).value()`.
679#[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/// Tweedie row log-likelihood, plain `f64`, bit-identical to
691/// `-dispersion_tweedie_disp_order2(..).value()` (both density branches).
692#[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/// Value-only row negative log-likelihood for one observation — the pruned hot
714/// path for [`CustomFamily::log_likelihood_only`]. Mirrors the exact-link
715/// preamble of [`dispersion_row_kernel`] exactly, then evaluates ONLY the value
716/// channel (no gradient/Hessian, no digamma/trigamma). Returns `row.loglik`
717/// `to_bits`-identically.
718#[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    // Zero-weight rows are excluded from the likelihood entirely (and exempt
727    // from the boundary support validation), so their row term must be an
728    // exact 0 rather than `0 · (±inf)` = NaN.
729    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
756/// Observed η-space row NLL tower: both exact predictors are jet variables
757/// (`η_μ` axis 0, `η_d` axis 1) and the full mean-link / precision-link
758/// chains are carried by the jet algebra, so `h()` is the exact per-row
759/// OBSERVED Hessian in `(η_μ, η_d)` — including the inverse-link
760/// second-derivative terms and the mean/dispersion cross curvature that the
761/// expected (Fisher) working weights do not represent. Example: Gamma with
762/// log links at `y = 4, μ = 2, ν = 3` has exact per-row `∂²NLL/∂η_μ² =
763/// νy/μ = 6` and `∂²NLL/∂η_μ∂η_ν = ν(1 − y/μ) = −3`, where the Fisher
764/// working weights give `ν = 3` and `0`.
765pub(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            // The NB log-likelihood below is written directly in the linear
778            // predictors (log-scale) via `log_total`, so the mean `exp(eta_mu)`
779            // is never materialized here (unlike the Gamma arm).
780            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
848/// Per-row observed `(∂²NLL/∂η_μ², ∂²NLL/∂η_μ∂η_d, ∂²NLL/∂η_d²)` weights for
849/// the exact joint Hessian at the supplied predictors.
850pub(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
865/// Order-3 alias for the two-predictor η-space NLL tower.
866type 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
901/// Observed η-space row NLL tower to THIRD order: the order-3 sibling of
902/// [`dispersion_eta_nll_order2`], written with the identical expression
903/// structure per family so the `v`/`g`/`h` channels agree with the order-2
904/// tower and `t3` is the exact per-row third-derivative tensor
905/// `∂³NLL/∂η_a∂η_b∂η_c`.
906///
907/// This tensor is what the β-directional derivative of the observed joint
908/// Hessian contracts row-wise — the object the Jeffreys/Firth gradient
909/// (`joint_jeffreys_term`'s `Hdot[e_k]`) and the outer mode-response
910/// correction (`D_β H_L[u]`) both need. Before it existed the family declined
911/// the directional-derivative hook, and `joint_jeffreys_term` degraded to
912/// `(Φ, 0, 0)`: the Firth value entered the inner merit while its gradient
913/// was silently zero, desynchronizing the inner joint-Newton objective from
914/// its KKT residual whenever the conditioning gate armed (the flat-residual
915/// stall on Beta/NB/Tweedie dispersion location-scale fits — #1561's
916/// `quality_vs_gamlss_beta_dispersion_location_scale_1060` null-model
917/// collapse).
918pub(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
1000/// Per-row directional derivative of the observed η-space Hessian channels
1001/// `(∂²NLL/∂η_μ², ∂²NLL/∂η_μ∂η_d, ∂²NLL/∂η_d²)` along the per-row η-motion
1002/// `(du_mu, du_d)` — the row-wise contraction of the exact third-derivative
1003/// tensor from [`dispersion_eta_nll_order3`].
1004pub(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/// Exact row-local geometry consumed by saved-model case deletion.
1026///
1027/// The score is the gradient of the weighted negative log-likelihood in the
1028/// affine coordinates `(eta_mu, eta_d)`.  `observed_hessian` is its observed
1029/// Hessian, not a Fisher working-weight surrogate and not the outer product of
1030/// the score.  Keeping those two objects separate is essential for ALO: the
1031/// observed Hessian controls the deletion denominator, while the score outer
1032/// product controls the sandwich variance.
1033#[derive(Clone, Copy, Debug, PartialEq)]
1034pub struct DispersionAloRowGeometry {
1035    pub nll_score: [f64; 2],
1036    pub observed_hessian: [[f64; 2]; 2],
1037}
1038
1039/// Replay the exact fitted row likelihood in its two affine predictor
1040/// coordinates for saved-model ALO.
1041///
1042/// This is intentionally a thin public boundary over the same order-two jet
1043/// program used by the fitter, so diagnostics cannot drift onto a second,
1044/// hand-maintained approximation of the dispersion likelihood.
1045pub 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
1096/// Evaluate the row log-likelihood and the (mean, log-precision) Fisher-scoring
1097/// working sets for one observation. `eta_mu`/`eta_d` already include any
1098/// per-channel offset (they are the block predictors). `prior_weight` is the
1099/// observation's prior weight.
1100pub(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    // Zero-weight rows are excluded from the likelihood (and exempt from the
1110    // boundary support validation): return exact zeros rather than letting
1111    // `0 · (±inf)` poison the objective sum.
1112    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(); // precision (size)
1126            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            // Dispersion (log-θ) IRLS curvature: use the EXPECTED (Fisher)
1135            // information in θ, not the per-row OBSERVED Hessian channel
1136            // (`_info_theta_observed`). The NB2 log-likelihood is strongly
1137            // non-quadratic in θ: `−∂²ℓ/∂θ²` carries the row-specific term
1138            // `ψ′(θ+y)` and goes NEGATIVE for every row whose count sits below
1139            // its current fitted precision (overestimated size / underestimated
1140            // overdispersion). Far from the optimum a majority of rows can be
1141            // negative, so the assembled block curvature `Xᵀdiag(w)X` loses
1142            // positive-definiteness; replacing each negative row by an
1143            // arbitrary epsilon then divides the exact score by
1144            // ~0 in the working response, producing O(1e12) IRLS targets that
1145            // make the dispersion block step explode and the inner block-cyclic
1146            // solve stall (never reaching KKT within the cycle budget — the
1147            // `nb` location-scale `IntegrationError`, gam#1606). The mean block
1148            // already uses its closed-form expected info `θ/(μ(θ+μ))`; the
1149            // dispersion block must do the same.
1150            //
1151            // The Fisher information in θ has the closed form
1152            //   I(θ) = ψ′(θ) − E[ψ′(θ+Y)] − 1/θ + 1/(θ+μ),
1153            // whose only costly piece is the per-row infinite expectation
1154            // `E[ψ′(θ+Y)]`. Replacing it with the Jensen plug-in `ψ′(θ+μ)`
1155            // (valid because ψ′ is convex, so this is a tight lower bound on the
1156            // expectation) gives a per-row, sum-free, STRICTLY POSITIVE
1157            // curvature
1158            //   I_θ ≈ ψ′(θ) − ψ′(θ+μ) − 1/θ + 1/(θ+μ) > 0  for all (μ,θ),
1159            // since ψ′ is strictly decreasing. The working RESPONSE still
1160            // carries the EXACT score `s_theta` (= ∂ℓ/∂θ from the tower), so the
1161            // penalized stationary point (score = 0) is byte-unchanged — this is
1162            // Fisher scoring, which only re-conditions the inner solve and never
1163            // shifts the optimum. The observed channel `_info_theta_observed` is no
1164            // longer consumed for the weight.
1165            // #1591-follow-up: scalar `trigamma` (== `trigamma_derivative_stack
1166            // (·)[0]` bit-for-bit) evaluates ONLY ψ′; the old `[0]`-index form
1167            // built the full order-1..5 polygamma stack and discarded four of
1168            // five per call (8 wasted polygamma evaluations per NB2 row).
1169            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(); // precision = shape ν
1193            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            // logit mean link.
1210            let logit = gam_solve::mixture_link::logit_inverse_link_jet5(em);
1211            let mu = logit.mu;
1212            let phi = ed.exp(); // precision
1213            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            // Precision channel models log(1/φ) ⇒ φ = exp(−η_d).
1241            let phi = (-ed).exp();
1242            let two_minus_p = 2.0 - p;
1243            // Mean channel: the quasi-score `(y−μ)/μ` and Fisher weight
1244            // `μ^{2−p}/φ` are simple closed forms (and the mean block is
1245            // Fisher-orthogonal to the dispersion block in this
1246            // parameterization), so they stay hand-written exactly as the
1247            // NB/Gamma mean arms do.
1248            let mean_weight = wi * mu.powf(two_minus_p) / phi;
1249            let mean_response = em + (yi - mu) / mu;
1250            // Dispersion channel: the η_d-space score and OBSERVED information
1251            // come straight off the single-expression tower seeded on `η_d`
1252            // (#932), so the saddlepoint/point-mass branch split, the
1253            // `φ = exp(−η_d)` chain and its nonlinear `∂²φ/∂η_d²` curvature
1254            // correction are all mechanically carried — no per-branch
1255            // `s_phi`/`s_eta`/`curvature_eta` hand calculus. #1591: only the
1256            // η_d axis is consumed, so the tower is the pruned single-axis
1257            // `Order2<1>` (`η_μ` enters as a constant).
1258            let tower = dispersion_tweedie_disp_order2(yi, em, ed, p, wi);
1259            let loglik = -tower.value();
1260            // η_d-space score and observed information off the tower, via the
1261            // same helper the NB/Gamma/Beta arms use (returns `(0, 0)` when the
1262            // prior weight is zero, so the row stays excluded below).
1263            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/// Two-block GAMLSS family for the genuine-dispersion mean families (#913).
1279#[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    // Preserve the pre-gam#1395 behavior: the trait default flipped to OFF (the
1293    // flat-prior exact-Newton objective carries no Jeffreys term), so families
1294    // that historically armed the term by default opt back in explicitly.
1295    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        // `dispersion_row_kernel` is a pure, row-independent map — each row reads
1324        // only `y[i]`/`eta_mu[i]`/`eta_d[i]`/`weights[i]` and writes nothing
1325        // shared — and it is transcendental-heavy (per-row digamma/trigamma
1326        // derivative stacks), so the per-row evaluation is embarrassingly
1327        // row-parallel. Materialize the per-row kernels (in parallel for large
1328        // `n` when not already on a rayon worker; mirrors the
1329        // `row_coeff_operator` guard), then reduce SERIALLY in index order so
1330        // the log-likelihood sum is bit-identical to the old serial loop — no
1331        // float reassociation. The reduction touches no transcendentals, so the
1332        // parallel kernel map captures essentially all the savings.
1333        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        // The objective is the honest sum: with support/weight validation at
1363        // the public boundary and zero-weight rows short-circuited in the
1364        // kernel, a non-finite row term means the likelihood genuinely
1365        // diverges at this (β_μ, β_d) — silently dropping such rows would
1366        // evaluate a different dataset's objective.
1367        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        // #1591 prune: the objective needs only the row log-likelihood, so each
1421        // row evaluates the value channel alone (`to_bits`-identical to
1422        // `dispersion_row_kernel(..).loglik`), skipping every gradient/Hessian
1423        // and digamma/trigamma derivative-stack evaluation. That value-only map
1424        // is still a pure, row-independent per-row `ln_gamma` evaluation, so it
1425        // is row-parallel; fan it out (large `n`, off a rayon worker) into a
1426        // per-row buffer, then sum SERIALLY in index order to keep the objective
1427        // bit-identical to the serial loop (no float reassociation).
1428        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        // Honest sum — see `evaluate`: non-finite row terms signal genuine
1457        // divergence and must reach the caller, not be silently dropped.
1458        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    /// Exact joint coefficient-space Hessian `H_L = -∇²log L` in flattened
1489    /// `[mean | log-precision]` block order.
1490    ///
1491    /// All four members assemble `Xᵀ diag(W) X` blocks from the per-row
1492    /// OBSERVED η-space second derivatives
1493    /// (`dispersion_row_observed_hessian_weights`): the full mean-link and
1494    /// precision-link chains, the inverse-link second-derivative terms, and
1495    /// the mean/dispersion cross curvature are all carried exactly by the
1496    /// `Order2<2>` jet tower. This is deliberately NOT the Fisher-scoring
1497    /// working-weight matrix that `evaluate` returns for the inner IRLS —
1498    /// expected information is a legitimate inner-solve preconditioner (the
1499    /// working response keeps the exact score, so the optimum is unchanged),
1500    /// but LAML/REML log-determinants, Jeffreys corrections, EDF, and the
1501    /// joint posterior covariance all require the observed Hessian. The
1502    /// Fisher-orthogonal members (NB2 / Gamma / Tweedie) have EXPECTED cross
1503    /// information zero, yet their per-row observed cross curvature is
1504    /// nonzero (Gamma at `y=4, μ=2, ν=3`: `∂²NLL/∂η_μ∂η_ν = −3`), so the
1505    /// assembled `H_L` is genuinely coupled for every member.
1506    ///
1507    /// Returning this dense `H_L` — rather than `None` — is what lets the
1508    /// multi-block outer-REML path (`build_joint_hessian_closures` →
1509    /// `joint_outer_evaluate`) and the joint posterior covariance
1510    /// (`compute_joint_covariance`) run for these families instead of failing
1511    /// the "multi-block families must provide a joint outer path" gate and
1512    /// silently escalating to a degraded ρ-seed fit with no covariance/EDF
1513    /// (gam#1119).
1514    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        // Per-row observed `(∂²/∂η_μ², ∂²/∂η_μ∂η_d, ∂²/∂η_d²)` weights — one
1551        // full `Order2<2>` η-space tower per row. Row-independent, so fan it
1552        // out for large `n` (off a rayon worker) into a per-row buffer —
1553        // index-ordered, no reduction, so byte-identical to the serial map.
1554        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    /// Exact β-directional derivative of the observed joint Hessian,
1643    /// `D_β H_L[u]`, assembled row-wise from the third-order η-space tower
1644    /// ([`dispersion_eta_nll_order3`]): with per-row η-motion
1645    /// `du_μ = X_μ u_μ`, `du_d = X_d u_d`, each Hessian channel drifts by the
1646    /// exact tensor contraction `dW_ab = Σ_c (∂³NLL/∂η_a∂η_b∂η_c) du_c`, and
1647    /// the blocks are the same `Xᵀ diag(dW) X` grams the Hessian itself uses.
1648    ///
1649    /// Supplying this hook (instead of the previous silent `None`) is
1650    /// load-bearing twice over:
1651    ///  * the inner Firth/Jeffreys term `joint_jeffreys_term` builds its
1652    ///    `∇Φ`/`H_Φ` from `Hdot[e_k]`; with `None` it degrades to `(Φ, 0, 0)`,
1653    ///    so the inner merit contains a β-dependent `−Φ` the KKT gradient
1654    ///    cannot see — the objective↔gradient desync behind the flat-residual
1655    ///    inner stall (and, post rail-face certification, the λ=∞ null-model
1656    ///    collapse) on Beta/NB/Tweedie dispersion location-scale fits (#1561);
1657    ///  * the outer profiled-Laplace mode-response correction
1658    ///    (`dot H_k = A_k + D_β H_L[u_k]`) consumes the same object.
1659    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        // η-motion of the direction: the offset is β-independent, so
1718        // `dη_b = X_b u_b` exactly.
1719        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    /// The joint likelihood Hessian is NOT block-diagonal for any member:
1793    /// even the Fisher-orthogonal parameterizations — NB2 `(μ, θ)`, Gamma
1794    /// shape `ν`, Tweedie `log(1/φ)` — have zero EXPECTED cross information
1795    /// but nonzero per-row OBSERVED cross curvature `∂²NLL/∂η_μ∂η_d`
1796    /// (Gamma at `y=4, μ=2, ν=3` has `ν(1−y/μ) = −3`). The former
1797    /// `uncoupled = true` shortcut for these members made the outer calculus
1798    /// consume a block-diagonal matrix as if it were the exact Hessian.
1799    /// The explicit-joint-Hessian marker below is what routes the outer
1800    /// dispatch to the trusted coupled override instead (gam#1119).
1801    fn likelihood_blocks_uncoupled(&self) -> bool {
1802        false
1803    }
1804
1805    /// `exact_newton_joint_hessian_with_specs` above returns the true coupled
1806    /// observed joint Hessian for every member, so mark it explicit for the
1807    /// outer-REML trust dispatch.
1808    fn has_explicit_joint_hessian(&self) -> bool {
1809        true
1810    }
1811
1812    /// The mean and precision working weights couple across both blocks, which
1813    /// the block-local diagonal drift hook cannot represent, so decline the
1814    /// dense outer Hessian capability whenever the actual two-block (or
1815    /// larger) geometry is in play; a degenerate single-block probe — there
1816    /// is no cross-block coupling to reject — keeps the trait default's
1817    /// availability verdict.
1818    ///
1819    /// The override still validates the block-spec slice it is handed (the
1820    /// same consistency check the trait default's assertion bottoms out in)
1821    /// so a malformed probe is reported here rather than downstream.
1822    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
1832/// Term spec consumed by [`fit_dispersion_glm_location_scale_terms`]; mirrors
1833/// [`GaussianLocationScaleTermSpec`](super::GaussianLocationScaleTermSpec) with
1834/// the dispersion channel in place of the Gaussian log-σ channel.
1835pub 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
1855/// Warm start for a dispersion location-scale fit: project a link-transformed
1856/// response onto the mean block and seed the log-precision block at a constant
1857/// (precision ≈ 1) baseline. The block-cyclic IRLS then refines both jointly.
1858pub(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                // log mean link; the +0.1 keeps zero counts finite.
1877                (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        // Seed the precision block from a smoothed method-of-moments surface
1894        // rather than the old flat η_d=0 constant.  A single observation cannot
1895        // identify its own variance, but for the Fisher-orthogonal dispersion
1896        // members the residual-squared moment contains the correct first-order
1897        // signal:
1898        //
1899        //   Gamma:   Var(Y)=μ²/ν              ⇒ log ν     ≈ log(μ²/e²)
1900        //   NB2:     Var(Y)=μ+μ²/θ            ⇒ log θ     ≈ log(μ²/(e²-μ))
1901        //   Tweedie: Var(Y)=φ μ^p, η_d=log1/φ ⇒ η_d       ≈ log(μ^p/e²)
1902        //
1903        // The targets are deliberately conservative (finite residual floor,
1904        // precision cap, and no fixture-specific constants): they only give the
1905        // block-cyclic likelihood solve a correctly-signed non-flat starting
1906        // surface, while the final estimate is still the penalized joint MLE.
1907        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            // Beta's mean and precision scores are not Fisher-orthogonal in
1932            // the (logit μ, log φ) parameterization.  Per-row residual moments
1933            // therefore make a poor block-cyclic seed: an outlying y near 0/1
1934            // can imply a near-zero φ and pull the coupled mean block onto the
1935            // boundary before the joint likelihood has had a chance to settle.
1936            // Keep the neutral precision seed for this one coupled member; the
1937            // exact Beta cross-Hessian below still drives the joint solve and
1938            // covariance with the coherent two-block likelihood geometry.
1939            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            // `theta_hat_row = mu^2/(e^2 - mu)` inverts a statistic whose
1950            // sampling distribution STRADDLES ZERO, so the denominator has a
1951            // pole inside the range the data can produce. The signal is
1952            // `E[e^2 - mu] = mu^2/theta`; the noise on the same quantity is
1953            // `sd[e^2] = sqrt(mu + 2 mu^2)` (Poisson-limit fourth central
1954            // moment `mu + 3 mu^2`, less `(mu + mu^2)^2`'s leading `mu^2`),
1955            // i.e. `~ sqrt(2) * mu`. Signal-to-noise is `~ (mu/theta)/sqrt(2)`,
1956            // well under one for every `theta >~ mu`, so a large fraction of
1957            // rows come out with a NEGATIVE excess carrying no information
1958            // about theta at all.
1959            //
1960            // The old relative guard `1e-6 * (mu + mu^2)` is not a floor on
1961            // anything measurable: it sits six orders of magnitude BELOW that
1962            // noise, so a negative row lands on it and seeds
1963            // `log theta = ln(1e6 * mu/(1 + mu)) -> ~13.8`, which
1964            // `LOG_PRECISION_CEILING` then clamps to +10 -- `theta = 22026`,
1965            // numerically Poisson -- for every such row. Measured: 49% of rows
1966            // on the #1119 NB fixture and 61% on the generate fixture seeded at
1967            // that cap. (The fraction is regime-dependent, not a fixed half:
1968            // `e^2 <= mu` is `|y - mu| <= sqrt(mu)`, whose probability falls as
1969            // the overdispersion grows. Both measurements are large.)
1970            //
1971            // The Gamma/Beta/Tweedie siblings cannot reach this state, which is
1972            // why the NB arms are the only red ones in two different files
1973            // whose siblings are all green: Gamma floors `e^2` itself at
1974            // `1e-8 mu^2`, so saturating needs `|y - mu| <= 1e-4 mu`; Tweedie
1975            // likewise floors `e^2`, and has the `y = 0` atom besides; Beta
1976            // deliberately seeds a flat 0.0. Only NB floors a SIGNED
1977            // difference, and only a signed difference has a pole.
1978            //
1979            // So floor the denominator at the statistic's own sampling
1980            // resolution. The resulting cap `theta_seed <= mu^2/sqrt(mu + 2
1981            // mu^2) ~ mu/sqrt(2)` errs toward MORE overdispersion -- the
1982            // direction in which the log-theta Fisher information is largest
1983            // (see `nb_log_precision_fisher_jensen`) -- so the block-cyclic
1984            // solve still gets a usable gradient and can climb back out. Do not
1985            // "simplify" this to a relative floor: the quantity being floored
1986            // is not small, it is NOISE, and a relative floor cannot know that.
1987            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        // SPEC-5: the log-precision block is penalized by its formula-native
2047        // function-space penalties only (a smooth term carries its own
2048        // REML-selected function-metric null-space shrinkage when
2049        // `double_penalty=true`, the default). The previous full-span
2050        // `identity_penalty` ridge was a basis-dependent coefficient-space prior
2051        // that double-penalized the range space and shrank the fitted dispersion
2052        // surface toward its coordinate origin — the exact over-shrinkage the
2053        // Gaussian location-scale path dropped in `de5599435` (#1561). Mirroring
2054        // that path here removes the extra REML smoothing coordinate the ridge
2055        // introduced, so the coupled inner solve no longer optimizes the
2056        // dispersion smoothing against a phantom full-span penalty.
2057        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        // The family stores y/weights/kind directly and does not need the
2099        // designs at construction time, but the row geometry of the offered
2100        // designs is the only cross-check that ties this family back to the
2101        // builder's data — assert it before handing the family to the engine
2102        // so a misaligned design surfaces here rather than downstream in the
2103        // inner solver.
2104        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        // The dispersion location-scale families do not expose the complete
2149        // coupled higher-order calculus needed for analytic spatial psi
2150        // derivatives. The public fit boundary rejects enabled κ/ψ requests;
2151        // if a future caller bypasses that boundary, return a real diagnostic
2152        // rather than a sentinel. Include the exact data/design shape so the
2153        // invalid call is diagnosable from the error string alone.
2154        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
2176/// Validate family support and prior weights at the public boundary.
2177///
2178/// The row kernels evaluate the requested likelihood verbatim; they do not
2179/// winsorize out-of-range Beta responses, floor nonpositive Gamma responses,
2180/// zero negative Tweedie responses, accept noninteger negative-binomial
2181/// counts, or clamp negative weights — all of those silently fit a DIFFERENT
2182/// dataset than the one supplied. Invalid rows must therefore be rejected
2183/// here. Rows with an exactly-zero prior weight are exempt from the response
2184/// support check (they are excluded from the likelihood entirely), which is
2185/// the supported way to carry deliberately masked observations.
2186fn 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
2236/// Reject a spatial-hyperparameter request that this coupled family cannot
2237/// differentiate exactly.
2238///
2239/// The shared spatial bridge can provide exact design/penalty jets
2240/// (`X_psi`, `S_psi`, and their second derivatives), but the dispersion
2241/// likelihood's observed two-block Hessian also moves through both fitted
2242/// predictors. An exact profiled LAML gradient therefore additionally needs
2243/// the coupled `D_beta H` and `D_beta H_psi` contractions. This family does not
2244/// expose those higher-order row jets yet. Silently switching `enabled` off
2245/// changes the requested model; exposing the existing typed configuration
2246/// error makes fixed geometry an explicit caller choice instead.
2247fn 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                // On the incoming (pre-build) spec, `0.0` is the Matérn
2262                // auto-initialization sentinel, not a user-locked scale.
2263                // A positive scalar scale freezes only an isotropic axis;
2264                // per-axis psi coordinates remain an optimization request.
2265                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
2296/// Fit a dispersion-channel GAMLSS location-scale model (#913). All four
2297/// genuine-dispersion mean families share this single entry; the per-family
2298/// likelihood lives in `dispersion_row_kernel`.
2299pub 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    // A dispersion location-scale model is an inherently *predictable* model:
2320    // posterior-mean prediction (the response-scale predict path the CLI/FFI
2321    // drive) needs the joint `(β_μ, β_d)` posterior covariance, and so does the
2322    // reported total EDF / coefficient SEs. The block-diagonal joint Hessian is
2323    // always assembled here (`exact_newton_joint_hessian_with_specs` →
2324    // `compute_joint_covariance`, which for this family's `RidgedQuadraticReml`
2325    // outer objective uses the never-erroring SPD-retry → positive-part
2326    // pseudo-inverse), so we can — and must — request the covariance
2327    // unconditionally rather than leaving `covariance_conditional = None`
2328    // whenever the outer optimizer happens to *converge* (the only family-
2329    // independent reason NB sometimes populated covariance was that it escalated
2330    // into the never-fail posterior-sampling rung, while a cleanly-converged
2331    // Gamma/Tweedie fit took the `!options.compute_covariance ⇒ None` early
2332    // return and stranded its covariance/EDF — gam#1119). Forcing the flag here
2333    // makes all four genuine-dispersion mean families assemble the joint
2334    // covariance + EDF deterministically, exactly as a predictable model
2335    // requires.
2336    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    /// Order-≤1 `ln Γ` compose: only the value (`d[0] = lnΓ`) and first-derivative
2421    /// (`d[1] = ψ`) stack slots are consumed by an [`Order1`] jet, so we evaluate
2422    /// ONLY `lnΓ` and `ψ` — never `ψ′` (trigamma). This is the value/gradient
2423    /// twin of [`order2_ln_gamma`]; its `(value, g)` channels are bit-identical to
2424    /// that function's order-≤1 channels because [`Order1`] runs the same Leibniz /
2425    /// Faà-di-Bruno value+gradient float ops as [`Order2`] (doc on `Order1`).
2426    #[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    /// Value+gradient-only NB2 dispersion tower: `θ` is the sole jet variable
2440    /// (axis 0), `μ` a constant. `value`/`g[0]` reproduce the consumed
2441    /// `value`/`g[0]` of [`dispersion_nb_disp_order2`] bit-for-bit, but as an
2442    /// [`Order1`] jet it never evaluates the trigamma (`ψ′`) that the discarded
2443    /// observed-Hessian channel would need. `h[0][0]` is pure discarded work —
2444    /// dropping it here removes two `ψ′` evaluations (at `θ+y` and `θ`) per
2445    /// evaluation on top of the tensor-shrink.
2446    #[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    /// Pruned single-axis NB2 dispersion tower: `θ` is the sole jet variable
2470    /// (axis 0), `μ` a constant. `value`/`g[0]`/`h[0][0]` reproduce the consumed
2471    /// `value`/`g[1]`/`h[1][1]` of `dispersion_nb_nll_order2` bit-for-bit. The
2472    /// `Order2` oracle pin for `prune_towers_match_dense_all_channels`, and the
2473    /// `Order1` oracle target for `dispersion_nb_disp_order1` above.
2474    #[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        // η-space chain: ∂²NLL/∂η_μ∂η_d = q·φ·f_μφ with q = dμ/dη_μ (the
2621        // cross entry carries no ∂²μ/∂η² term because q is η_d-free).
2622        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    /// #932 oracle: the production `Order2<2>` evaluation of each dispersion
2636    /// row NLL must reproduce, channel-for-channel (value/grad/Hessian), the
2637    /// dense `Tower4<2>` evaluation of the same row expression.
2638    #[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        // NB2: (μ, θ).
2671        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        // Gamma: (μ, ν).
2679        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        // Beta: (μ, φ).
2692        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        // Tweedie: (η_μ, η_d), both density branches.
2700        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    /// #1591 prune oracle: the pruned single-axis (`K=1`) dispersion towers
2715    /// reproduce, `to_bits`-exactly, the CONSUMED channels (`value`, dispersion-
2716    /// axis `g`/`h`) of the full `Order2<2>` towers — across ≥2000 randomized
2717    /// rows per family (both Tweedie density branches). This is the bit-identity guarantee that the K-prune changes no
2718    /// observable float.
2719    #[test]
2720    pub(crate) fn pruned_disp_towers_bit_identical_to_full_order2() {
2721        use gam_math::jet_scalar::Order2;
2722
2723        // Deterministic LCG so the sweep is reproducible without an rng dep.
2724        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; // 600 rows × 4 families (Tweedie ×2 branches) > 2000.
2734        for _ in 0..n_per {
2735            let wi = 0.25 + 3.0 * next();
2736            let yi_count = (next() * 12.0).floor();
2737
2738            // NB: full O2<2> seeds (μ, θ); pruned seeds θ only.
2739            {
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                // Value+gradient-only production tower (`Order1`, trigamma-free):
2748                // its consumed `value`/`g[0]` must match the `Order2` form
2749                // bit-for-bit (the observed Hessian it drops is unused by the NB2
2750                // Fisher-scoring row kernel).
2751                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                // value-only path == -tower.value(), bit-for-bit.
2755                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            // Gamma: seeds (μ, ν) / ν.
2763            {
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            // Beta value-only path vs full K=2 tower value.
2780            {
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            // Tweedie: seeds (η_μ, η_d) / η_d, both density branches.
2792            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    /// Audit finding 34 pin: the exact joint Hessian consumes OBSERVED
2824    /// per-row η-space curvature, not expected (Fisher) information. Gamma
2825    /// with log links at `y = 4, μ = 2, ν = 3` has closed-form per-row NLL
2826    /// second derivatives `∂²/∂η_μ² = νy/μ = 6` and `∂²/∂η_μ∂η_ν =
2827    /// ν(1 − y/μ) = −3`; the Fisher weights are `ν = 3` and `0`.
2828    #[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        // ∂²NLL/∂η_ν² = ν²(ψ′(ν) − 1/ν) + [ν(lnμ − lnν − 1 + ψ(ν) − ln y + y/μ)]·(−1)…
2846        // pin against a central finite difference of the value channel instead
2847        // of a second hand derivation.
2848        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    /// Finite predictors beyond the former arbitrary clamp remain on the exact
2857    /// likelihood surface and retain their score and curvature.
2858    #[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        // The kernel takes LOG-space inputs and re-exponentiates: its internal
2889        // precision is `huge.ln().exp()`, which differs from `huge` by the
2890        // exp∘ln round-trip (~|ln huge|·EPSILON ≈ 460·EPSILON relative at
2891        // η≈460), so dividing by the pre-log `huge` injects ~50·EPSILON of pure
2892        // round-trip noise. Reference the precision the kernel actually sees:
2893        // with μ==θ, `mean_weight = θ/(1+θ/μ) = θ/2` is exact (a bit-for-bit
2894        // exponent decrement), so `mean_weight / precision == 0.5` exactly and
2895        // the tight tolerance stays a genuine balanced-tail geometry guard.
2896        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    /// Speed-path guard (#932): `evaluate` / `log_likelihood_only` materialize
2910    /// the row-kernel map in parallel for large `n`, then reduce SERIALLY in
2911    /// index order. This pins the parallel output (log-likelihood + both
2912    /// blocks' working response/weight vectors) to a hand-rolled serial
2913    /// reference so CI catches any reassociation or row-misindex regression.
2914    /// `n` sits well above `DISPERSION_PARALLEL_ROW_THRESHOLD`, and the test
2915    /// runs on the main thread (not a rayon worker), so the parallel branch is
2916    /// the one exercised. Because the reduction order is preserved the match is
2917    /// in fact bit-exact; the `1e-9` band is the contract floor.
2918    #[test]
2919    pub(crate) fn parallel_evaluate_matches_serial_reference() {
2920        let n = DISPERSION_PARALLEL_ROW_THRESHOLD * 3 + 7;
2921        // Deterministic LCG row data (no rng dependency).
2922        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            // Serial reference, computed exactly as the pre-parallel loop did.
2962            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            // `log_likelihood_only` takes the same parallel-then-serial-sum
3006            // path; its value-only kernel is bit-identical to evaluate's loglik.
3007            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}