Skip to main content

gam_models/gamlss/
errors.rs

1// Real concern-organized submodule of the gamlss family stack.
2// Cross-module items are re-exported flat through the parent (`gamlss.rs`),
3// so `use super::*;` makes the sibling-concern symbols this module references
4// resolve through the parent namespace.
5use super::*;
6
7/// Typed errors surfaced from this module's helpers and family
8/// implementations. The `Display` impl writes the carried `reason` verbatim,
9/// so callers that historically returned `Result<_, String>` keep their
10/// user-visible text byte-for-byte identical after coercion via the
11/// `From<GamlssError> for String` impl below.
12#[derive(Debug)]
13pub enum GamlssError {
14    /// Shape, length, row, or column mismatches between matrices,
15    /// vectors, specs, or block configurations.
16    DimensionMismatch { reason: String },
17    /// Generic input validation that doesn't fit a more specific
18    /// variant (e.g. positivity-of-response checks, shape parameter
19    /// must be finite > 0).
20    InvalidInput { reason: String },
21    /// Non-finite values discovered in inputs, coefficients, seeds,
22    /// or intermediate quantities required to remain finite.
23    NonFinite { reason: String },
24    /// A model configuration or feature combination is not supported
25    /// by the requested family / link / engine (e.g. identity link on
26    /// a binomial mean-wiggle family, unexpected design-map variant).
27    UnsupportedConfiguration { reason: String },
28    /// Bound, range, monotonicity, or sign constraints violated by
29    /// supplied parameters or coefficients.
30    ConstraintViolation { reason: String },
31    /// Numerical failures during inner solves, integration, or
32    /// optimization (invalid probabilities, non-finite log-likelihood,
33    /// invalid λ, divergence).
34    NumericalFailure { reason: String },
35}
36
37impl_reason_error_boilerplate! {
38    GamlssError {
39        DimensionMismatch,
40        InvalidInput,
41        NonFinite,
42        UnsupportedConfiguration,
43        ConstraintViolation,
44        NumericalFailure,
45    }
46}
47
48impl From<crate::block_layout::block_count::BlockCountMismatch> for GamlssError {
49    fn from(err: crate::block_layout::block_count::BlockCountMismatch) -> GamlssError {
50        GamlssError::DimensionMismatch {
51            reason: err.message(),
52        }
53    }
54}
55
56/// Numerical floor on μ ∈ (0, 1) used only for downstream `1/μ` and
57/// `1/(1-μ)` divisions and for `μ.ln()` / `(1-μ).ln()` in the generic
58/// composed-link binomial log-likelihood (where the logit-stable
59/// `log_expit` form is unavailable because `q` is the composed link
60/// argument, not the raw logit η). Pure numerical safety, NOT a model
61/// assumption — when the optimizer pushes μ to the floor it indicates a
62/// separated/saturated fit which is detected and surfaced upstream
63/// (`detect_logit_instability`, `Unstable` PIRLS status). For
64/// composed-link μ, derivatives `dμ/dq` etc. are NOT zeroed when the
65/// floor is hit; they carry the legitimate gradient signal of the
66/// outer link and zeroing them would create a phantom flat region that
67/// the optimizer would converge to as a stationary point.
68pub(crate) const MIN_PROB: f64 = 1e-10;
69
70pub(crate) const MIN_DERIV: f64 = 1e-8;
71
72/// Lower clamp on POSITIVE working weights `w_i = (dμ/dη)² / V(μ_i)`
73/// to keep `Xᵀ W X` numerically representable. Strictly numerical:
74/// `w` enters subsequent dense matrix products and a true zero (which
75/// happens when `dμ/dη = 0` at saturation, e.g. logistic μ → 0 with
76/// `dμ/dη = μ(1-μ)`) is harmless but a denormal `w` propagates as
77/// inf/NaN through `XᵀWX` because `w * (x_i x_j)` underflows
78/// non-uniformly. `floor_positiveweight` returns 0 for non-finite or
79/// non-positive inputs (so saturation correctly drops the row from
80/// the inner Newton system); the floor only fires for *strictly
81/// positive* tiny weights. The 1e-12 magnitude is chosen so that
82/// `1e-12 · max|x|² · n` stays comfortably above `f64::MIN_POSITIVE`
83/// at large scale.
84///
85/// This is the canonical positive-weight floor (`1e-12`); the value is owned by
86/// [`gam_problem::MIN_WEIGHT`] so every floored family shares one definition
87/// rather than re-declaring it per module.
88use gam_problem::MIN_WEIGHT;
89
90/// Hard symmetric clamp on η used by the Poisson / Gaussian / Gamma working-
91/// model log-likelihood loops to keep `exp(η)` and `log(σ)` finite under the
92/// IRLS step. Hoisted out of each loop so all three families share the same
93/// numerical regime.
94pub(crate) const ETA_HARD_CLAMP: f64 = 30.0;
95
96/// Saturated `exp(η)` used by every log-link mean reconstruction in this
97/// module: clamp η into `[−ETA_HARD_CLAMP, ETA_HARD_CLAMP]` so `exp` stays
98/// finite, then floor at `MIN_WEIGHT` so downstream divisions never see
99/// exact zero. Centralising the formula here means a tolerance change
100/// propagates to all three families (Poisson / Gaussian / Gamma) without
101/// risk of one path drifting.
102#[inline]
103pub(crate) fn saturated_exp_eta(eta: f64) -> f64 {
104    eta.clamp(-ETA_HARD_CLAMP, ETA_HARD_CLAMP)
105        .exp()
106        .max(MIN_WEIGHT)
107}
108
109/// Floor applied to a fitted smoothing parameter λ before `ln(λ)` is taken to
110/// seed an outer-loop `initial_log_lambdas` warm start. A pilot fit can return
111/// λ underflowed to exactly 0 for a deselected (effectively unpenalized) term;
112/// `ln(0) = -inf` would poison the seed, so we floor at the smallest λ that is
113/// still numerically distinguishable from zero in the log-domain rather than a
114/// modelling-meaningful value. `ln(1e-12) ≈ -27.6` sits well below any λ the
115/// outer optimizer would select, so a genuinely tiny pilot λ still seeds the
116/// search near its lower edge.
117pub(crate) const WARMSTART_LOG_LAMBDA_FLOOR: f64 = 1e-12;
118
119pub(crate) const EXACT_DENSE_BLOCK_BUDGET_BYTES: usize = 512 * 1024 * 1024;
120
121pub(crate) const EXACT_DENSE_TOTAL_BUDGET_BYTES: usize = 2 * 1024 * 1024 * 1024;
122
123pub(crate) const GAMLSS_ROWWISE_PAR_MIN_N: usize = 4096;
124
125pub(crate) const GAMLSS_PROJECTED_TRACE_TARGET_BYTES: usize = 32 * 1024 * 1024;
126
127pub(crate) const GAMLSS_PROJECTED_TRACE_MIN_CHUNK_ROWS: usize = 64;
128
129pub(crate) const GAMLSS_PROJECTED_TRACE_MAX_CHUNK_ROWS: usize = 8192;
130
131pub(crate) fn gamlss_projected_trace_chunk_rows(
132    rank: usize,
133    projected_channel_count: usize,
134    gram_column_count: usize,
135) -> usize {
136    let per_row_values = rank
137        .saturating_mul(projected_channel_count.max(1))
138        .saturating_add(gram_column_count.max(1))
139        .max(1);
140    let per_row_bytes = per_row_values.saturating_mul(std::mem::size_of::<f64>());
141    let rows = GAMLSS_PROJECTED_TRACE_TARGET_BYTES / per_row_bytes.max(1);
142    rows.clamp(
143        GAMLSS_PROJECTED_TRACE_MIN_CHUNK_ROWS,
144        GAMLSS_PROJECTED_TRACE_MAX_CHUNK_ROWS,
145    )
146}
147
148pub(crate) fn gamlss_rowwise_map<F>(n: usize, f: F) -> Array1<f64>
149where
150    F: Fn(usize) -> f64 + Sync,
151{
152    if n >= GAMLSS_ROWWISE_PAR_MIN_N {
153        Array1::from((0..n).into_par_iter().map(&f).collect::<Vec<f64>>())
154    } else {
155        Array1::from_iter((0..n).map(f))
156    }
157}
158
159pub(crate) fn gamlss_rowwise_map_result<F>(n: usize, f: F) -> Result<Array1<f64>, String>
160where
161    F: Fn(usize) -> Result<f64, String> + Sync,
162{
163    if n >= GAMLSS_ROWWISE_PAR_MIN_N {
164        let values: Result<Vec<f64>, String> = (0..n).into_par_iter().map(&f).collect();
165        Ok(Array1::from(values?))
166    } else {
167        let mut out = Array1::<f64>::zeros(n);
168        for i in 0..n {
169            out[i] = f(i)?;
170        }
171        Ok(out)
172    }
173}
174
175pub(crate) enum DenseOrOperator<'a> {
176    Borrowed(&'a Array2<f64>),
177    Owned(Array2<f64>),
178    Operator(DesignMatrix),
179}
180
181impl DenseOrOperator<'_> {
182    pub(crate) fn nrows(&self) -> usize {
183        match self {
184            Self::Borrowed(dense) => dense.nrows(),
185            Self::Owned(dense) => dense.nrows(),
186            Self::Operator(design) => design.nrows(),
187        }
188    }
189
190    pub(crate) fn ncols(&self) -> usize {
191        match self {
192            Self::Borrowed(dense) => dense.ncols(),
193            Self::Owned(dense) => dense.ncols(),
194            Self::Operator(design) => design.ncols(),
195        }
196    }
197
198    pub(crate) fn row_chunk(&self, rows: std::ops::Range<usize>) -> Result<Array2<f64>, String> {
199        match self {
200            Self::Borrowed(dense) => Ok(dense.slice(s![rows, ..]).to_owned()),
201            Self::Owned(dense) => Ok(dense.slice(s![rows, ..]).to_owned()),
202            Self::Operator(design) => design.try_row_chunk(rows).map_err(|e| e.to_string()),
203        }
204    }
205
206    pub(crate) fn dot(&self, beta: ArrayView1<'_, f64>) -> Array1<f64> {
207        let n = self.nrows();
208        let p = self.ncols();
209        assert_eq!(beta.len(), p);
210        match self {
211            Self::Borrowed(dense) => fast_av(*dense, &beta),
212            Self::Owned(dense) => fast_av(dense, &beta),
213            Self::Operator(design) => {
214                let mut out = Array1::<f64>::zeros(n);
215                for rows in exact_design_row_chunks(n, p) {
216                    let chunk = design
217                        .try_row_chunk(rows.clone())
218                        .expect("gamlss DesignSlot::dot: design row chunk materialization failed");
219                    out.slice_mut(s![rows]).assign(&fast_av(&chunk, &beta));
220                }
221                out
222            }
223        }
224    }
225}
226
227/// Resolve a single dense block design from a `ParameterBlockSpec`, falling
228/// back to materializing the sparse representation through the policy when
229/// the dense form isn't already cached. Returns `Cow::Borrowed` whenever the
230/// spec already holds a dense array; `Cow::Owned` only after a forced
231/// materialization. The `materialization_label` string is woven into the
232/// materializer's error so callers can pin which block failed.
233pub(crate) fn dense_block_from_spec<'a>(
234    spec: &'a ParameterBlockSpec,
235    material_policy: &gam_runtime::resource::MaterializationPolicy,
236    materialization_label: &str,
237) -> Result<Cow<'a, Array2<f64>>, String> {
238    match spec.design.as_dense_ref() {
239        Some(d) => Ok(Cow::Borrowed(d)),
240        None => Ok(Cow::Owned(
241            spec.design
242                .try_to_dense_with_policy(material_policy, "gamlss dense_block_from_spec")
243                .map_err(|e| format!("{materialization_label}: {e}"))?
244                .as_ref()
245                .clone(),
246        )),
247    }
248}
249
250/// Resolve the (primary, log-σ) pair of dense block designs that every
251/// LocationScale family's spec-aware exact path needs. The primary block is
252/// the family-specific "mean" axis (μ for Gaussian, latent t for Binomial);
253/// the `short_family_name` ("GaussianLocationScale", "BinomialLocationScale",
254/// or their Wiggle siblings) and `primary_label` ("mu" / "threshold") are
255/// woven into the per-block materialization label for diagnostics.
256pub(crate) fn dense_locscale_block_designs_fromspecs<'a>(
257    specs: &'a [ParameterBlockSpec],
258    expected_count: usize,
259    family_name: &str,
260    short_family_name: &str,
261    primary_block_idx: usize,
262    log_sigma_block_idx: usize,
263    primary_label: &str,
264    material_policy: &gam_runtime::resource::MaterializationPolicy,
265) -> Result<(Cow<'a, Array2<f64>>, Cow<'a, Array2<f64>>), String> {
266    if specs.len() != expected_count {
267        return Err(GamlssError::DimensionMismatch {
268            reason: format!(
269                "{family_name} expects {expected_count} specs, got {}",
270                specs.len()
271            ),
272        }
273        .into());
274    }
275    let primary = dense_block_from_spec(
276        &specs[primary_block_idx],
277        material_policy,
278        &format!("{short_family_name} dense_block_designs_fromspecs {primary_label}"),
279    )?;
280    let log_sigma = dense_block_from_spec(
281        &specs[log_sigma_block_idx],
282        material_policy,
283        &format!("{short_family_name} dense_block_designs_fromspecs log_sigma"),
284    )?;
285    Ok((primary, log_sigma))
286}
287
288/// Materialize a single location-scale family's two cached block designs
289/// (`primary` = mu/threshold, plus `log_sigma`) into dense matrices, borrowing
290/// when the design is already dense and owning a policy-materialized copy
291/// otherwise. Every non-wiggle and wiggle location-scale family's
292/// `dense_block_designs` method is identical bar the accessed field and the
293/// diagnostic labels, so both bits are passed in.
294pub(crate) fn dense_locscale_block_designs_cached<'a>(
295    primary_design: Option<&'a DesignMatrix>,
296    log_sigma_design: Option<&'a DesignMatrix>,
297    family_name: &str,
298    short_family_name: &str,
299    primary_label: &str,
300    material_policy: &gam_runtime::resource::MaterializationPolicy,
301) -> Result<(Cow<'a, Array2<f64>>, Cow<'a, Array2<f64>>), String> {
302    let primary_design = primary_design
303        .ok_or_else(|| format!("{family_name} exact path is missing {primary_label} design"))?;
304    let log_sigma_design = log_sigma_design
305        .ok_or_else(|| format!("{family_name} exact path is missing log-sigma design"))?;
306    let primary = match primary_design.as_dense_ref() {
307        Some(d) => Cow::Borrowed(d),
308        None => Cow::Owned(
309            primary_design
310                .try_to_dense_with_policy(material_policy, "gamlss dense_locscale_block_designs")
311                .map_err(|e| {
312                    format!("{short_family_name} dense_block_designs {primary_label}: {e}")
313                })?
314                .as_ref()
315                .clone(),
316        ),
317    };
318    let log_sigma = match log_sigma_design.as_dense_ref() {
319        Some(d) => Cow::Borrowed(d),
320        None => Cow::Owned(
321            log_sigma_design
322                .try_to_dense_with_policy(material_policy, "gamlss dense_locscale_block_designs")
323                .map_err(|e| format!("{short_family_name} dense_block_designs log_sigma: {e}"))?
324                .as_ref()
325                .clone(),
326        ),
327    };
328    Ok((primary, log_sigma))
329}
330
331/// One resolved ψ-direction for a two-axis (primary + log-σ) location-scale
332/// family. Holds the neutral pieces shared by every such family's
333/// `exact_newton_joint_psi_direction`; each family wraps these into its own
334/// named struct (mu/threshold field renames only).
335pub(crate) struct LocScalePsiDirectionParts {
336    pub(crate) block_idx: usize,
337    pub(crate) local_idx: usize,
338    pub(crate) primary_psi: PsiDesignMap,
339    pub(crate) log_sigma_psi: PsiDesignMap,
340    pub(crate) primary_z: Array1<f64>,
341    pub(crate) log_sigma_z: Array1<f64>,
342}
343
344/// Shared body of every two-axis location-scale family's
345/// `exact_newton_joint_psi_direction`. Walks the flat ψ-derivative list,
346/// resolves the ψ-design map for the selected block (primary = block 0, log-σ
347/// = block 1; the off-axis map is the matching `Zero`), and applies each
348/// block's β via `forward_mul`. The wiggle block (and any other index) yields
349/// `None`, matching the per-family methods. The only per-family variation —
350/// the column counts, the two block betas, the block-list length (2 or 3) and
351/// the diagnostic label prefix — is passed in; the math is identical across
352/// Gaussian/Binomial × wiggle/non-wiggle.
353pub(crate) fn locscale_joint_psi_direction_parts(
354    block_states: &[ParameterBlockState],
355    derivative_blocks: &[Vec<crate::custom_family::CustomFamilyBlockPsiDerivative>],
356    psi_index: usize,
357    n: usize,
358    p_primary: usize,
359    p_log_sigma: usize,
360    primary_block_idx: usize,
361    log_sigma_block_idx: usize,
362    expected_blocks: usize,
363    family_name: &str,
364    primary_label: &str,
365    policy: &gam_runtime::resource::ResourcePolicy,
366) -> Result<Option<LocScalePsiDirectionParts>, String> {
367    validate_block_count::<GamlssError>(family_name, expected_blocks, block_states.len())?;
368    if derivative_blocks.len() != expected_blocks {
369        return Err(GamlssError::DimensionMismatch {
370            reason: format!(
371                "{family_name} joint psi direction expects {expected_blocks} derivative block lists, got {}",
372                derivative_blocks.len()
373            ),
374        }
375        .into());
376    }
377    let beta_primary = &block_states[primary_block_idx].beta;
378    let beta_log_sigma = &block_states[log_sigma_block_idx].beta;
379
380    let mut global = 0usize;
381    for (block_idx, block_derivs) in derivative_blocks.iter().enumerate() {
382        for (local_idx, deriv) in block_derivs.iter().enumerate() {
383            if global == psi_index {
384                let primary_psi;
385                let log_sigma_psi;
386                let primary_z;
387                let log_sigma_z;
388                if block_idx == primary_block_idx {
389                    primary_psi = resolve_custom_family_x_psi_map(
390                        deriv,
391                        n,
392                        p_primary,
393                        0..n,
394                        &format!("{family_name} {primary_label}"),
395                        policy,
396                    )?;
397                    primary_z = primary_psi
398                        .forward_mul(beta_primary.view())
399                        .map_err(|e| format!("{family_name} {primary_label} forward_mul: {e}"))?;
400                    log_sigma_psi = PsiDesignMap::Zero {
401                        nrows: n,
402                        ncols: p_log_sigma,
403                    };
404                    log_sigma_z = Array1::<f64>::zeros(n);
405                } else if block_idx == log_sigma_block_idx {
406                    log_sigma_psi = resolve_custom_family_x_psi_map(
407                        deriv,
408                        n,
409                        p_log_sigma,
410                        0..n,
411                        &format!("{family_name} log-sigma"),
412                        policy,
413                    )?;
414                    log_sigma_z = log_sigma_psi
415                        .forward_mul(beta_log_sigma.view())
416                        .map_err(|e| format!("{family_name} log-sigma forward_mul: {e}"))?;
417                    primary_psi = PsiDesignMap::Zero {
418                        nrows: n,
419                        ncols: p_primary,
420                    };
421                    primary_z = Array1::<f64>::zeros(n);
422                } else {
423                    return Ok(None);
424                }
425                return Ok(Some(LocScalePsiDirectionParts {
426                    block_idx,
427                    local_idx,
428                    primary_psi,
429                    log_sigma_psi,
430                    primary_z,
431                    log_sigma_z,
432                }));
433            }
434            global += 1;
435        }
436    }
437    Ok(None)
438}
439
440/// Shared second-derivative design drift assembly for two-axis location-scale
441/// joint-ψ paths. The family-specific methods differ only by block constants,
442/// labels, and field names; the ψψ map lookup and `X_{ab} β` action are the
443/// same for Gaussian/Binomial and wiggle/non-wiggle variants.
444pub(crate) struct LocScalePsiDriftConfig<'a> {
445    pub(crate) n: usize,
446    pub(crate) p_primary: usize,
447    pub(crate) p_log_sigma: usize,
448    pub(crate) primary_block_idx: usize,
449    pub(crate) log_sigma_block_idx: usize,
450    pub(crate) family_name: &'a str,
451    pub(crate) primary_label: &'a str,
452    pub(crate) policy: &'a gam_runtime::resource::ResourcePolicy,
453}
454
455pub(crate) fn locscale_joint_psisecond_design_drifts(
456    block_states: &[ParameterBlockState],
457    derivative_blocks: &[Vec<crate::custom_family::CustomFamilyBlockPsiDerivative>],
458    psi_a: &LocationScaleJointPsiDirection,
459    psi_b: &LocationScaleJointPsiDirection,
460    cfg: LocScalePsiDriftConfig<'_>,
461) -> Result<LocationScaleJointPsiSecondDrifts, String> {
462    let beta_primary = &block_states[cfg.primary_block_idx].beta;
463    let beta_log_sigma = &block_states[cfg.log_sigma_block_idx].beta;
464    let mut primary_ab_action = None;
465    let mut log_sigma_ab_action = None;
466    let mut primary_ab = None;
467    let mut log_sigma_ab = None;
468
469    // Smooth ψ second derivatives are block-local. Cross-block ψ_a/ψ_b
470    // design second derivatives are therefore zero unless the derivative
471    // payload itself supplies them for the same moving block.
472    if psi_a.block_idx == psi_b.block_idx {
473        let deriv = &derivative_blocks[psi_a.block_idx][psi_a.local_idx];
474        let deriv_b = &derivative_blocks[psi_b.block_idx][psi_b.local_idx];
475        if psi_a.block_idx == cfg.primary_block_idx {
476            let (action, matrix) = psi_psi_map_to_drift_slots(
477                deriv,
478                deriv_b,
479                psi_b.local_idx,
480                cfg.n,
481                cfg.p_primary,
482                &format!("{} {}", cfg.family_name, cfg.primary_label),
483                cfg.policy,
484            )?;
485            primary_ab_action = action;
486            primary_ab = matrix;
487        } else if psi_a.block_idx == cfg.log_sigma_block_idx {
488            let (action, matrix) = psi_psi_map_to_drift_slots(
489                deriv,
490                deriv_b,
491                psi_b.local_idx,
492                cfg.n,
493                cfg.p_log_sigma,
494                &format!("{} log-sigma", cfg.family_name),
495                cfg.policy,
496            )?;
497            log_sigma_ab_action = action;
498            log_sigma_ab = matrix;
499        }
500    }
501
502    let z_primary_ab = second_psi_linear_map(
503        primary_ab_action.as_ref(),
504        primary_ab.as_ref(),
505        cfg.n,
506        cfg.p_primary,
507    )
508    .forward_mul(beta_primary.view());
509    let z_ls_ab = second_psi_linear_map(
510        log_sigma_ab_action.as_ref(),
511        log_sigma_ab.as_ref(),
512        cfg.n,
513        cfg.p_log_sigma,
514    )
515    .forward_mul(beta_log_sigma.view());
516
517    Ok(LocationScaleJointPsiSecondDrifts {
518        x_primary_ab_action: primary_ab_action,
519        x_ls_ab_action: log_sigma_ab_action,
520        x_primary_ab: primary_ab,
521        x_ls_ab: log_sigma_ab,
522        z_primary_ab,
523        z_ls_ab,
524    })
525}
526
527pub(crate) fn psi_psi_map_to_drift_slots(
528    deriv: &crate::custom_family::CustomFamilyBlockPsiDerivative,
529    deriv_b: &crate::custom_family::CustomFamilyBlockPsiDerivative,
530    local_idx_b: usize,
531    n: usize,
532    p: usize,
533    label: &str,
534    policy: &gam_runtime::resource::ResourcePolicy,
535) -> Result<
536    (
537        Option<crate::custom_family::CustomFamilyPsiSecondDesignAction>,
538        Option<Array2<f64>>,
539    ),
540    String,
541> {
542    match resolve_custom_family_x_psi_psi_map(
543        deriv,
544        deriv_b,
545        local_idx_b,
546        n,
547        p,
548        0..n,
549        label,
550        policy,
551    )? {
552        crate::custom_family::PsiDesignMap::Second { action } => Ok((Some(action), None)),
553        crate::custom_family::PsiDesignMap::Dense { matrix } => Ok((None, Some((*matrix).clone()))),
554        crate::custom_family::PsiDesignMap::Zero { .. } => Ok((None, None)),
555        crate::custom_family::PsiDesignMap::First { .. } => {
556            Err(GamlssError::UnsupportedConfiguration {
557                reason: format!("{label}: unexpected First variant from _psi_psi_map"),
558            }
559            .into())
560        }
561    }
562}
563
564pub(crate) fn dense_block_or_operator<'a>(
565    design: &'a DesignMatrix,
566    n: usize,
567    p: usize,
568    budget_bytes: usize,
569    policy: &gam_runtime::resource::ResourcePolicy,
570) -> DenseOrOperator<'a> {
571    if let Some(dense) = design.as_dense_ref() {
572        return DenseOrOperator::Borrowed(dense);
573    }
574
575    let dense_bytes = 8usize.saturating_mul(n).saturating_mul(p);
576    if dense_bytes <= budget_bytes
577        && let Ok(arc) = design
578            .try_to_dense_with_policy(&policy.material_policy(), "gamlss dense_block_or_operator")
579    {
580        return DenseOrOperator::Owned(arc.as_ref().clone());
581    }
582
583    DenseOrOperator::Operator(design.clone())
584}
585
586pub(crate) fn dense_blocks_planned_budget(blocks: &[&DesignMatrix]) -> Vec<usize> {
587    let mut planned = vec![0; blocks.len()];
588    let mut total = 0usize;
589    for (idx, design) in blocks.iter().enumerate() {
590        if design.as_dense_ref().is_some() {
591            continue;
592        }
593        let bytes = 8usize
594            .saturating_mul(design.nrows())
595            .saturating_mul(design.ncols());
596        if bytes <= EXACT_DENSE_BLOCK_BUDGET_BYTES
597            && total.saturating_add(bytes) <= EXACT_DENSE_TOTAL_BUDGET_BYTES
598        {
599            planned[idx] = bytes;
600            total += bytes;
601        }
602    }
603    planned
604}
605
606pub(crate) fn exact_design_row_chunks(
607    n: usize,
608    p: usize,
609) -> impl Iterator<Item = std::ops::Range<usize>> {
610    const TARGET_BYTES: usize = 8 * 1024 * 1024;
611    const MIN_ROWS: usize = 512;
612    const MAX_ROWS: usize = 131_072;
613    let rows = (TARGET_BYTES / (p.max(1) * 8))
614        .clamp(MIN_ROWS, MAX_ROWS)
615        .min(n.max(1));
616    (0..n)
617        .step_by(rows)
618        .map(move |start| start..(start + rows).min(n))
619}
620
621pub(crate) fn design_weighted_column_squares(
622    design: &DesignMatrix,
623    weights: &Array1<f64>,
624) -> Result<Array1<f64>, String> {
625    let n = design.nrows();
626    let p = design.ncols();
627    if weights.len() != n {
628        return Err(GamlssError::DimensionMismatch {
629            reason: format!(
630                "design weighted column squares dimension mismatch: weights={}, rows={}",
631                weights.len(),
632                n
633            ),
634        }
635        .into());
636    }
637    let mut out = Array1::<f64>::zeros(p);
638    for rows in exact_design_row_chunks(n, p) {
639        let chunk = design.try_row_chunk(rows.clone()).map_err(|e| {
640            format!("design weighted column squares row chunk materialization failed: {e}")
641        })?;
642        for (local_i, row) in chunk.outer_iter().enumerate() {
643            let w = weights[rows.start + local_i];
644            if w == 0.0 {
645                continue;
646            }
647            for j in 0..p {
648                let x = row[j];
649                out[j] += w * x * x;
650            }
651        }
652    }
653    Ok(out)
654}
655
656#[inline]
657pub(crate) fn floor_positiveweight(rawweight: f64, minweight: f64) -> f64 {
658    if !rawweight.is_finite() || rawweight <= 0.0 {
659        0.0
660    } else {
661        rawweight.max(minweight)
662    }
663}
664
665#[inline]
666pub(crate) fn logb_dlog_sigma_deta(sigma: f64, d_sigma_deta: f64) -> f64 {
667    if d_sigma_deta.is_infinite() {
668        1.0
669    } else {
670        let value = d_sigma_deta / sigma;
671        if value.is_finite() {
672            value.clamp(0.0, 1.0)
673        } else {
674            0.0
675        }
676    }
677}
678
679#[inline]
680pub(crate) fn gaussian_log_sigma_irlsinfo_directional_derivative(
681    weight: f64,
682    sigma: f64,
683    d_sigma_deta: f64,
684    d_eta: f64,
685) -> f64 {
686    if weight == 0.0 || d_eta == 0.0 || !sigma.is_finite() || sigma <= 0.0 {
687        return 0.0;
688    }
689    // Logb form mirrors gaussian_jointrow_scalars: κ = exp(η)/(b + exp(η)) ∈ [0, 1)
690    // and dκ/dη = κ(1−κ). Use dσ/dη over σ directly so the η → −∞ tail
691    // preserves subnormal information instead of cancelling in `1 − b/σ`;
692    // the helper handles the η → +∞ inf/inf case by returning the analytic
693    // limit 1.
694    let g = logb_dlog_sigma_deta(sigma, d_sigma_deta);
695    if !g.is_finite() || !(0.0..1.0).contains(&g) {
696        return 0.0;
697    }
698    let rawinfo = 2.0 * weight * g * g;
699    if !rawinfo.is_finite() || rawinfo <= MIN_WEIGHT {
700        return 0.0;
701    }
702    let dg_deta = g * (1.0 - g);
703    let dw = 4.0 * weight * g * dg_deta * d_eta;
704    if dw.is_finite() { dw } else { 0.0 }
705}
706
707#[derive(Clone, Copy)]
708pub(crate) struct GaussianDiagonalRowKernel {
709    pub(crate) log_likelihood: f64,
710    pub(crate) location_working_weight: f64,
711    pub(crate) location_working_shift: f64,
712    pub(crate) log_sigma_working_weight: f64,
713    pub(crate) log_sigma_working_response: f64,
714}
715
716#[inline]
717pub(crate) fn gaussian_diagonal_row_kernel(
718    y: f64,
719    location_eta: f64,
720    eta_log_sigma: f64,
721    obs_weight: f64,
722    ln2pi: f64,
723) -> GaussianDiagonalRowKernel {
724    if obs_weight == 0.0 {
725        return GaussianDiagonalRowKernel {
726            log_likelihood: 0.0,
727            location_working_weight: 0.0,
728            location_working_shift: 0.0,
729            log_sigma_working_weight: 0.0,
730            log_sigma_working_response: eta_log_sigma,
731        };
732    }
733
734    // logb noise link σ = b + exp(η) bounds σ ≥ b > 0 by construction, so the
735    // Gaussian location-scale objective ½Σ(y−μ)²/σ² + Σlog σ is bounded below
736    // for any finite data. Its working weight 1/σ² is bounded by 1/b², so
737    // H_μμ has bounded condition number — no after-the-fact floor or cap is
738    // needed (the previous (1e-12, 1e24) clamp was a numerical bandaid for the
739    // pure-exp link's σ→0 singularity and is structurally unnecessary here).
740    // ApproxKind: Exact — working weight analytically bounded in (0, 1/b²].
741    let SigmaJet1 { sigma, d1 } = logb_sigma_jet1_scalar(eta_log_sigma);
742    let inv_s2 = (sigma * sigma).recip();
743    let residual = y - location_eta;
744    let location_working_weight = floor_positiveweight(obs_weight * inv_s2, MIN_WEIGHT);
745    // dlog σ/dη = (∂σ/∂η)/σ = exp(η)/(b + exp(η)) ∈ [0, 1).
746    // Use dσ/dη over σ directly so the η→−∞ tail preserves subnormal
747    // derivative information instead of cancelling in `1 − b/σ`; the helper
748    // returns the analytic limit 1 for the η→+∞ inf/inf case.
749    // Fisher info per obs = 2·(dσ/dη)²/σ² = 2·dlog_sigma_deta², matching the
750    // formula for the pure-exp link (where dlog_sigma_deta ≡ 1).
751    let dlog_sigma_deta = logb_dlog_sigma_deta(sigma, d1);
752    let log_sigma_working_weight = floor_positiveweight(
753        2.0 * obs_weight * dlog_sigma_deta * dlog_sigma_deta,
754        MIN_WEIGHT,
755    );
756    let log_sigma_score = obs_weight * (residual * residual * inv_s2 - 1.0) * dlog_sigma_deta;
757    let log_sigma_working_response = if log_sigma_working_weight == 0.0 {
758        eta_log_sigma
759    } else {
760        eta_log_sigma + log_sigma_score / log_sigma_working_weight
761    };
762
763    GaussianDiagonalRowKernel {
764        log_likelihood: obs_weight
765            * (-0.5 * (residual * residual * inv_s2 + ln2pi + 2.0 * sigma.ln())),
766        location_working_weight,
767        location_working_shift: residual,
768        log_sigma_working_weight,
769        log_sigma_working_response,
770    }
771}
772
773/// Link identifiers for distribution parameters in multi-parameter GAMLSS families.
774#[derive(Clone, Copy, Debug, PartialEq, Eq)]
775pub enum ParameterLink {
776    Identity,
777    Log,
778    Logit,
779    Probit,
780    InverseLink,
781    /// Learnable smooth departure from a known base link.
782    Wiggle,
783}