Skip to main content

gam_solve/inference/
alo.rs

1use crate::estimate::EstimationError;
2use crate::estimate::{FitGeometry, UnifiedFitResult};
3use crate::pirls;
4use faer::Mat as FaerMat;
5use faer::linalg::matmul::matmul;
6use faer::prelude::ReborrowMut;
7use faer::{Accum, Par};
8use gam_linalg::faer_ndarray::{FaerArrayView, FaerCholesky};
9use gam_linalg::matrix::{PsdWeightsView, SignedWeightsView};
10use gam_linalg::utils::StableSolver;
11use gam_problem::LinkFunction;
12use ndarray::{Array1, Array2, ArrayView1, ShapeBuilder, s};
13use std::fmt;
14
15/// Typed error variants for the ALO (approximate leave-one-out) diagnostics
16/// module.
17///
18/// Public entry points continue to return `Result<_, EstimationError>`; this
19/// enum is materialized at leaf sites and converted at the boundary via
20/// `From<AloError> for EstimationError` so error text remains byte-identical
21/// to the previous `EstimationError::InvalidInput(format!(...))` /
22/// `ModelIsIllConditioned { ... }` output.
23#[derive(Debug, Clone)]
24pub enum AloError {
25    /// Caller-supplied configuration is structurally invalid: dimension
26    /// mismatch, non-finite inputs that are not weights/response, missing
27    /// PIRLS / geometry artifacts, or out-of-range scalar parameters.
28    InvalidInput { reason: String },
29    /// IRLS weights or working response contain a non-finite entry, or the
30    /// working response itself is invalid.
31    WeightInvalid { reason: String },
32    /// The dense design matrix required for ALO could not be materialized
33    /// from the underlying PIRLS artifact (e.g. sparse-only export).
34    DesignDegenerate { reason: String },
35    /// The penalized Hessian factorization failed, or downstream diagnostics
36    /// produced NaN values that indicate the influence matrix is unusable.
37    InfluenceMatrixFailed { condition_number: f64 },
38    /// Per-observation ALO computation produced a non-finite value (variance,
39    /// denominator, or corrected η̃) at convergence.
40    LooComputationFailed { reason: String },
41}
42
43impl fmt::Display for AloError {
44    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
45        match self {
46            AloError::InvalidInput { reason }
47            | AloError::WeightInvalid { reason }
48            | AloError::DesignDegenerate { reason }
49            | AloError::LooComputationFailed { reason } => f.write_str(reason),
50            AloError::InfluenceMatrixFailed { condition_number } => {
51                write!(
52                    f,
53                    "ALO influence matrix failed (condition number {condition_number:.3e})"
54                )
55            }
56        }
57    }
58}
59
60impl std::error::Error for AloError {}
61
62impl From<AloError> for EstimationError {
63    fn from(err: AloError) -> EstimationError {
64        match err {
65            AloError::InvalidInput { reason }
66            | AloError::WeightInvalid { reason }
67            | AloError::DesignDegenerate { reason }
68            | AloError::LooComputationFailed { reason } => EstimationError::InvalidInput(reason),
69            AloError::InfluenceMatrixFailed { condition_number } => {
70                EstimationError::ModelIsIllConditioned { condition_number }
71            }
72        }
73    }
74}
75
76impl From<AloError> for String {
77    fn from(err: AloError) -> String {
78        err.to_string()
79    }
80}
81
82/// Approximate leave-one-out diagnostics derived from a fitted model.
83#[derive(Debug, Clone)]
84pub struct AloDiagnostics {
85    pub eta_tilde: Array1<f64>,
86    /// Bayesian/conditional standard error on eta:
87    /// sqrt(phi * x_i^T H^{-1} x_i).
88    pub se_bayes: Array1<f64>,
89    /// Frequentist sandwich-style standard error on eta:
90    /// sqrt(phi * x_i^T H^{-1} X^T W X H^{-1} x_i).
91    pub se_sandwich: Array1<f64>,
92    pub pred_identity: Array1<f64>,
93    pub leverage: Array1<f64>,
94    pub fisherweights: Array1<f64>,
95}
96
97#[inline]
98fn alo_eta_updatewith_offset(
99    eta_hat: f64,
100    z: f64,
101    offset: f64,
102    x_hinv_x: f64,
103    score_weight: f64,
104    denom: f64,
105) -> f64 {
106    // PIRLS working-response algebra is centered on offset, so the scalar
107    // score uses (eta - offset) - (z - offset).
108    let eta_centered = eta_hat - offset;
109    let z_centered = z - offset;
110    let score = score_weight * (eta_centered - z_centered);
111    offset + eta_centered + x_hinv_x * score / denom
112}
113
114/// Per-row score and curvature of the penalized NLL contribution as functions
115/// of the row's linear predictor `eta`.
116///
117/// Returns `(ℓ_i'(eta), ℓ_i''(eta))` where `ℓ_i` is the (dispersion-scaled)
118/// negative log-likelihood of observation `i` viewed as a univariate function
119/// of `eta_i = x_i^T β`. This is the local family geometry that the ALO
120/// frozen-curvature fixed point [`alo_eta_exact_frozen_curvature`] iterates to
121/// convergence; supplying it upgrades the single-Newton-step ALO correction to
122/// the exact leave-`i`-out predictor under a frozen penalized Hessian.
123pub type AloScalarScoreCurvature<'a> = dyn Fn(usize, f64) -> (f64, f64) + Sync + 'a;
124
125/// Maximum scalar Newton iterations for the exact frozen-curvature ALO fixed
126/// point. The map `r(η) = η − η̂ − a_ii ℓ_i'(η)` is one-dimensional and
127/// strongly contractive for the well-leveraged majority of points, so this
128/// caps the rare high-leverage / near-separation rows where convergence is
129/// slow without ever exceeding O(1) work per observation.
130const ALO_EXACT_SCALAR_MAX_ITERS: usize = 64;
131
132/// Absolute convergence tolerance on the scalar residual `r(η)` for the exact
133/// frozen-curvature ALO fixed point. Well below the `1e-2` predictive bar the
134/// LOO comparison asserts, so the refinement is not the limiting error term.
135const ALO_EXACT_SCALAR_TOL: f64 = 1e-12;
136
137/// Solve the frozen-curvature ALO leave-`i`-out fixed point exactly.
138///
139/// The leave-`i`-out optimum differs from the full fit only through the removed
140/// observation, whose gradient/Hessian depend on `β` solely via the scalar
141/// `η_i = x_i^T β`. Freezing the penalized Hessian `H` at its converged value
142/// reduces the exact leave-`i`-out condition to the scalar equation
143///
144///   η = η̂_i + a_ii · ℓ_i'(η),     a_ii = x_i^T H^{-1} x_i,
145///
146/// where `ℓ_i'(η)` is the row's NLL score (so that `∇F = ℓ_i'(η_i) x_i` at the
147/// leave-`i`-out point). The single-Newton-step ALO is exactly the first
148/// iterate of Newton's method on `r(η) = η − η̂_i − a_ii ℓ_i'(η)` started at
149/// `η̂_i`; iterating to convergence captures the change in the held-out point's
150/// likelihood curvature (the dominant first-order error on small-`n`, curved
151/// likelihoods such as binomial logistic regression near separation).
152///
153/// `score_curvature(eta)` returns `(ℓ_i'(eta), ℓ_i''(eta))`. The returned value
154/// is the corrected linear predictor `η̃_i`. Failure to reach the residual
155/// tolerance is reported to the caller; no one-step approximation is substituted
156/// for a failed exact solve.
157#[derive(Debug, Clone, Copy, PartialEq)]
158enum AloExactScalarError {
159    NonFiniteScoreCurvature {
160        eta: f64,
161        ell_prime: f64,
162        ell_double: f64,
163    },
164    DegenerateJacobian {
165        eta: f64,
166        jacobian: f64,
167    },
168    NonFiniteStep {
169        eta: f64,
170        residual: f64,
171        jacobian: f64,
172        next: f64,
173    },
174    MaxIterations {
175        iterations: usize,
176        residual: f64,
177        eta: f64,
178    },
179}
180
181impl fmt::Display for AloExactScalarError {
182    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
183        match *self {
184            AloExactScalarError::NonFiniteScoreCurvature {
185                eta,
186                ell_prime,
187                ell_double,
188            } => write!(
189                f,
190                "non-finite score/curvature at eta={eta:.6e}: ell_prime={ell_prime:.6e}, ell_double={ell_double:.6e}"
191            ),
192            AloExactScalarError::DegenerateJacobian { eta, jacobian } => write!(
193                f,
194                "degenerate Newton Jacobian at eta={eta:.6e}: jacobian={jacobian:.6e}, min={ALO_DENOMINATOR_MIN:.1e}"
195            ),
196            AloExactScalarError::NonFiniteStep {
197                eta,
198                residual,
199                jacobian,
200                next,
201            } => write!(
202                f,
203                "non-finite Newton step from eta={eta:.6e}: residual={residual:.6e}, jacobian={jacobian:.6e}, next={next:.6e}"
204            ),
205            AloExactScalarError::MaxIterations {
206                iterations,
207                residual,
208                eta,
209            } => write!(
210                f,
211                "did not converge within {iterations} iterations: residual={residual:.6e}, eta={eta:.6e}, tol={ALO_EXACT_SCALAR_TOL:.1e}"
212            ),
213        }
214    }
215}
216
217/// Maximum number of step halvings in the backtracking line search that
218/// globalizes the scalar Newton iteration. `2^{-40}` shrinks a unit step well
219/// below `ALO_EXACT_SCALAR_TOL` relative to any η of practical magnitude, so a
220/// row that cannot make progress within this budget is genuinely stalled rather
221/// than merely under-damped.
222const ALO_EXACT_SCALAR_BACKTRACKS: usize = 40;
223
224#[inline]
225fn alo_eta_exact_frozen_curvature(
226    eta_hat: f64,
227    a_ii: f64,
228    score_curvature: &dyn Fn(f64) -> (f64, f64),
229) -> Result<f64, AloExactScalarError> {
230    // Residual of the leave-i-out fixed point η = η̂ + a_ii ℓ'(η):
231    //   r(η) = η − η̂ − a_ii ℓ'(η),     r'(η) = 1 − a_ii ℓ''(η) = jac.
232    // For an exponential-family NLL score ℓ'(η) = c_i(μ(η) − y) on a non-linear
233    // (e.g. log) link the curvature ℓ''(η) = c_i μ'(η) grows without bound, so
234    // r(η) is concave with an interior maximum where the weighted leverage
235    // a_ii ℓ'' passes 1 (jac = 0): the leave-i-out root that limits to η̂ as
236    // a_ii → 0 sits on the jac > 0 branch anchored at η̂, while beyond the
237    // maximum r turns over and diverges as μ(η) explodes.
238    //
239    // Two safeguards make the scalar solve globally convergent to that root:
240    //
241    //   1. Anchor the iteration at η̂ itself, not at the classical one-step ALO
242    //      predictor. At η̂ the weighted leverage a_ii ℓ''(η̂) < 1, so jac ≈ 1
243    //      and we start strictly inside the correct basin; the brute-force
244    //      n-fold reference solves the identical fixed point anchored at η̂.
245    //      Seeding at the one-step predictor instead can land a high-leverage
246    //      row *past* the interior maximum on the runaway branch, from which no
247    //      Newton iteration returns (Poisson/log row 198: η ≈ 6.3, r ≈ −577).
248    //
249    //   2. Backtrack on the merit ½r(η)². The Newton direction d = −r/jac
250    //      satisfies (½r²)'·d = r·jac·(−r/jac) = −r² < 0 for any finite nonzero
251    //      jac, so halving the step until |r| strictly decreases never leaves
252    //      the basin even if a full step would overshoot the maximum.
253    let residual_and_jac = |eta: f64| -> Result<(f64, f64), AloExactScalarError> {
254        let (ell_prime, ell_double) = score_curvature(eta);
255        if !ell_prime.is_finite() || !ell_double.is_finite() {
256            return Err(AloExactScalarError::NonFiniteScoreCurvature {
257                eta,
258                ell_prime,
259                ell_double,
260            });
261        }
262        Ok((eta - eta_hat - a_ii * ell_prime, 1.0 - a_ii * ell_double))
263    };
264
265    let mut eta = eta_hat;
266    let (mut residual, mut jac) = residual_and_jac(eta)?;
267    for _ in 0..ALO_EXACT_SCALAR_MAX_ITERS {
268        if residual.abs() <= ALO_EXACT_SCALAR_TOL {
269            return Ok(eta);
270        }
271        if jac.abs() <= ALO_DENOMINATOR_MIN || !jac.is_finite() {
272            return Err(AloExactScalarError::DegenerateJacobian { eta, jacobian: jac });
273        }
274        let step = residual / jac;
275        if !step.is_finite() {
276            return Err(AloExactScalarError::NonFiniteStep {
277                eta,
278                residual,
279                jacobian: jac,
280                next: eta - step,
281            });
282        }
283        // Backtracking line search: take the longest damped Newton step
284        // 2^{-k} that strictly reduces the merit |r|. A non-finite trial
285        // (score/curvature evaluated in the runaway branch) is treated as no
286        // improvement and rejected, so the search retreats toward η̂.
287        let mut t = 1.0;
288        let mut advanced = false;
289        for _ in 0..ALO_EXACT_SCALAR_BACKTRACKS {
290            let trial = eta - t * step;
291            if let Ok((r_trial, j_trial)) = residual_and_jac(trial) {
292                if r_trial.abs() < residual.abs() {
293                    eta = trial;
294                    residual = r_trial;
295                    jac = j_trial;
296                    advanced = true;
297                    break;
298                }
299            }
300            t *= 0.5;
301        }
302        if !advanced {
303            break;
304        }
305    }
306    Err(AloExactScalarError::MaxIterations {
307        iterations: ALO_EXACT_SCALAR_MAX_ITERS,
308        residual,
309        eta,
310    })
311}
312
313#[inline]
314fn bayesvar_eta(phi: f64, x_hinv_x: f64) -> f64 {
315    phi * x_hinv_x
316}
317
318#[inline]
319fn sandwichvar_eta_from_meat(phi: f64, meat_quad: f64) -> f64 {
320    phi * meat_quad
321}
322
323#[inline]
324fn variance_negative_tolerance(scale: f64) -> f64 {
325    // Tight relative tolerance for cancellation from x'H^{-1}x - ||E t||^2 - ridge||t||^2.
326    1e-12 * scale.abs().max(1.0)
327}
328
329const LEVERAGE_HIGH_THRESHOLD: f64 = 0.99;
330const LEVERAGE_VERY_HIGH_THRESHOLD: f64 = 0.999;
331const LEVERAGE_RATE_THRESHOLDS: [f64; 3] = [0.90, 0.95, 0.99];
332const LEVERAGE_PERCENTILES: [f64; 3] = [0.50, 0.95, 0.99];
333const ALO_DENOMINATOR_MIN: f64 = 1e-12;
334const MULTIBLOCK_ALO_MEMORY_BUDGET_BYTES: usize = 256 * 1024 * 1024;
335
336/// Number of observation columns solved per blocked right-hand-side batch in the
337/// scalar-leverage path. Sizes the reusable `(p, .)` and `(e_rank, .)` scratch
338/// buffers so the dense multi-RHS solve stays BLAS-3 (good cache reuse) without
339/// materializing all `n` columns at once. The final batch is the remainder.
340const ALO_RHS_BLOCK_COLS: usize = 8192;
341
342/// Relative tolerance for accepting the input penalised Hessian `H` as
343/// symmetric. We require `|H_ij − H_ji| ≤ HESSIAN_SYMMETRY_REL_TOL ·
344/// max(|H_ij|, |H_ji|, 1)`. `1e-8` matches the loosest tolerance any
345/// upstream symmetrisation pass leaves on the matrix and is tight enough
346/// that a genuinely asymmetric Hessian (a real bug) is caught.
347const HESSIAN_SYMMETRY_REL_TOL: f64 = 1e-8;
348
349/// Diagonal ridge added to the local block precision when its LU pivot is
350/// below [`LU_PIVOT_SINGULAR_TOL`]. Matches the legacy `eps = 1e-6`
351/// regularisation in the prior `det_small < 1e-12` branch — bumping the
352/// determinant of `I − W A` (or `I − A W`) safely off zero without
353/// perturbing well-conditioned blocks.
354const ALO_LOCAL_BLOCK_RIDGE: f64 = 1e-6;
355
356/// Pivot magnitude below which [`lu_factor_in_place`] reports the block
357/// `I − W A` as singular and triggers the ridge-regularised refactor.
358/// Equivalent to the original `det_small < 1e-12` test on the unfactored
359/// determinant.
360const LU_PIVOT_SINGULAR_TOL: f64 = 1e-12;
361
362#[inline]
363fn percentile_index(sample_size: usize, quantile: f64) -> usize {
364    if sample_size <= 1 {
365        return 0;
366    }
367    let max_index = sample_size - 1;
368    ((quantile * max_index as f64).round() as usize).min(max_index)
369}
370
371#[inline]
372fn percentile_from_sorted(sorted: &[f64], quantile: f64) -> f64 {
373    if sorted.is_empty() {
374        0.0
375    } else {
376        sorted[percentile_index(sorted.len(), quantile)]
377    }
378}
379
380#[inline]
381fn multiblock_col_offsets(block_designs: &[Array2<f64>]) -> Vec<usize> {
382    let mut offsets = Vec::with_capacity(block_designs.len());
383    let mut off = 0usize;
384    for design in block_designs {
385        offsets.push(off);
386        off += design.ncols();
387    }
388    offsets
389}
390
391#[inline]
392fn multiblock_alo_parallel_leverage_chunk_size(
393    p_tot: usize,
394    n_blocks: usize,
395    n_obs: usize,
396    max_workers: usize,
397) -> usize {
398    if p_tot == 0 || n_blocks == 0 || n_obs == 0 {
399        return 1;
400    }
401
402    // Each parallel leverage chunk owns q_storage for all block RHS products
403    // (B * p_tot * chunk_len) plus one transposed design chunk across all
404    // blocks (p_tot * chunk_len).  Divide the global scratch budget by the
405    // maximum number of chunks Rayon can execute concurrently so total live
406    // per-chunk scratch remains bounded.
407    let workers = max_workers.max(1);
408    let per_worker_budget = (MULTIBLOCK_ALO_MEMORY_BUDGET_BYTES / workers).max(1);
409    let elem_count_per_obs = p_tot.saturating_mul(n_blocks.saturating_add(1)).max(1);
410    let bytes_per_obs = elem_count_per_obs
411        .saturating_mul(std::mem::size_of::<f64>())
412        .max(1);
413    let budget_obs = (per_worker_budget / bytes_per_obs).max(1);
414    budget_obs.min(n_obs)
415}
416
417fn compute_alo_diagnostics_from_pirls_impl(
418    base: &pirls::PirlsResult,
419    y: ArrayView1<f64>,
420    link: LinkFunction,
421) -> Result<AloDiagnostics, EstimationError> {
422    compute_alo_diagnostics_from_pirls_inner(base, y, link).map_err(EstimationError::from)
423}
424
425/// True when the fitted GLM uses a *curved* canonical link, so that the row NLL
426/// score and curvature satisfy `ℓ_i'(η) = c_i(μ(η)−y_i)` and `ℓ_i''(η) = c_i μ'(η)`
427/// with a single per-row scale `c_i = (prior weight)/φ`. This is the exact
428/// condition under which the frozen-curvature ALO scalar fixed point matches
429/// the leave-`i`-out refit; only these families enable the exact refinement.
430///
431/// Gaussian identity is canonical too, but its per-row curvature is *constant*
432/// (`μ'(η) ≡ 1`), so the classical Sherman–Morrison one-step ALO is already the
433/// exact frozen-Hessian leave-`i`-out solution. Routing it through the scalar
434/// Newton closure would only add an O(n) nonlinear solve to diagnostics and
435/// quality sweeps without changing the answer, so it is excluded here and falls
436/// back to the (exact, for this family) one-step formula.
437fn alo_link_needs_exact_curvature_refinement(likelihood: &gam_problem::GlmLikelihoodSpec) -> bool {
438    use gam_problem::ResponseFamily;
439    matches!(
440        (&likelihood.spec.response, likelihood.link_function()),
441        (ResponseFamily::Binomial, LinkFunction::Logit)
442            | (ResponseFamily::Poisson, LinkFunction::Log)
443    )
444}
445
446fn compute_alo_diagnostics_from_pirls_inner(
447    base: &pirls::PirlsResult,
448    y: ArrayView1<f64>,
449    link: LinkFunction,
450) -> Result<AloDiagnostics, AloError> {
451    let x_dense_arc = base
452        .x_transformed
453        .try_to_dense_arc("ALO diagnostics require dense transformed design")
454        .map_err(|reason| AloError::DesignDegenerate { reason })?;
455    let x_dense = x_dense_arc.as_ref();
456    let n = x_dense.nrows();
457
458    // Compute dispersion parameter.
459    let phi = match link {
460        LinkFunction::Log => 1.0,
461        LinkFunction::Logit
462        | LinkFunction::Probit
463        | LinkFunction::CLogLog
464        | LinkFunction::LogLog
465        | LinkFunction::Cauchit
466        | LinkFunction::Sas
467        | LinkFunction::BetaLogistic => 1.0,
468        LinkFunction::Identity => {
469            use rayon::iter::{IntoParallelIterator, ParallelIterator};
470            let rss: f64 = (0..n)
471                .into_par_iter()
472                .map(|i| {
473                    let r = y[i] - base.finalmu[i];
474                    base.finalweights[i] * r * r
475                })
476                .sum();
477            // Effective sample size for dispersion (#584): a zero prior weight
478            // makes w_i·r_i² = 0, so the row is already excluded from the RSS
479            // numerator and must be excluded from the denominator too. Count only
480            // positive-weight rows, exactly as the main optimizer path does
481            // (optimizer.rs ~1567); using the raw row count over a zero-excluding
482            // numerator biases φ̂ low and shrinks every ALO SE.
483            let n_pos = (0..n).filter(|&i| base.finalweights[i] > 0.0).count();
484            let dof = (n_pos as f64) - base.edf;
485            let denom = dof.max(1.0);
486            rss / denom
487        }
488    };
489
490    let e = &base.reparam_result.e_transformed;
491    let ridge = base.ridge_passport.laplacehessianridge().max(0.0);
492
493    // ALO needs the exact penalized Hessian materialized densely for chunked
494    // column solves via StableSolver.  The PIRLS export path validates the
495    // matrix instead of falling back to a numerical Hessian approximation.
496    let h_dense_for_alo = base
497        .dense_stabilizedhessian_transformed(
498            "ALO diagnostics require exact dense stabilized penalized Hessian",
499        )
500        .map_err(|e| match e {
501            EstimationError::InvalidInput(reason) => AloError::InvalidInput { reason },
502            other => AloError::InvalidInput {
503                reason: format!("{other:?}"),
504            },
505        })?;
506
507    // Exact frozen-curvature ALO refinement for canonical-link GLMs.
508    //
509    // For a canonical link the row NLL score and curvature are
510    //   ℓ_i'(η)  = c_i · (μ(η) − y_i),     ℓ_i''(η) = c_i · μ'(η),
511    // with c_i = (prior weight)/φ recovered from the converged geometry as
512    // c_i = W_H[i] / μ'(η̂_i) (since W_H[i] = c_i μ'(η̂_i) at convergence).
513    // Supplying this evaluator lets `compute_alo_from_input_inner` solve the
514    // leave-i-out scalar fixed point η = η̂_i + a_ii ℓ_i'(η) exactly instead of
515    // taking a single Newton step, removing the first-order linearization error
516    // that dominates on small-n, strongly curved likelihoods (binomial logit).
517    //
518    // Restricted to canonical links because only there does the observed
519    // curvature carried by the frozen Hessian (W_H) coincide with c_i μ'(η) for
520    // every trial η; non-canonical links retain the classical one-step ALO.
521    // Per-row scale c_i = W_H[i]/μ'(η̂_i). Rows whose μ'(η̂_i) is negligible
522    // (saturated / near-separation) get c_i = NaN, which makes the exact solver
523    // reject that row explicitly rather than substituting the classical one-step
524    // ALO.
525    let canonical_scale: Option<Array1<f64>> =
526        if alo_link_needs_exact_curvature_refinement(&base.likelihood) {
527            let mut c = Array1::<f64>::zeros(n);
528            for i in 0..n {
529                let dmu = base.solve_dmu_deta[i];
530                let w_h = base.finalweights[i];
531                c[i] = if dmu.abs() <= ALO_DENOMINATOR_MIN || !dmu.is_finite() || !w_h.is_finite() {
532                    f64::NAN
533                } else {
534                    w_h / dmu
535                };
536            }
537            Some(c)
538        } else {
539            None
540        };
541
542    let inv_link_for_closure = base.likelihood.spec.link.clone();
543    let score_curvature_closure = canonical_scale.as_ref().map(|scale| {
544        move |i: usize, eta: f64| -> (f64, f64) {
545            let (mu, dmu) = crate::mixture_link::inverse_link_mu_d1_for_inverse_link(
546                &inv_link_for_closure,
547                eta,
548            )
549            .unwrap_or((f64::NAN, f64::NAN));
550            let c_i = scale[i];
551            (c_i * (mu - y[i]), c_i * dmu)
552        }
553    });
554    let score_curvature_ref: Option<&AloScalarScoreCurvature> = score_curvature_closure
555        .as_ref()
556        .map(|f| f as &AloScalarScoreCurvature);
557
558    // Build model-agnostic AloInput from PIRLS geometry, then delegate.
559    // #1868: the PIRLS row fields are now shared `ArcArray1`; `AloInput` borrows
560    // `&Array1`, so materialise owned copies for this cold post-fit inference
561    // path (ALO runs once after the fit, not per κ trial).
562    let alo_working_response = base.solveworking_response.to_owned();
563    let alo_final_eta = base.final_eta.to_owned();
564    let alo_final_offset = base.final_offset.to_owned();
565    let input = AloInput {
566        design: x_dense,
567        penalized_hessian: &h_dense_for_alo,
568        hessian_weights: base.final_weights_signed(),
569        score_weights: base.solve_weights_psd(),
570        working_response: &alo_working_response,
571        eta: &alo_final_eta,
572        offset: &alo_final_offset,
573        link,
574        phi,
575        penalty_root: if e.nrows() > 0 { Some(e) } else { None },
576        ridge,
577        score_curvature: score_curvature_ref,
578    };
579
580    let result = compute_alo_from_input_inner(&input)?;
581
582    // PIRLS-specific post-hoc leverage diagnostics logging.
583    log_leverage_diagnostics(&result.leverage, phi);
584
585    // Final NaN guard with detailed error reporting.
586    let has_nan_pred = result.eta_tilde.iter().any(|&x| x.is_nan());
587    let has_nan_se_bayes = result.se_bayes.iter().any(|&x| x.is_nan());
588    let has_nan_se_sandwich = result.se_sandwich.iter().any(|&x| x.is_nan());
589    let has_nan_leverage = result.leverage.iter().any(|&x| x.is_nan());
590
591    if has_nan_pred || has_nan_se_bayes || has_nan_se_sandwich || has_nan_leverage {
592        log::error!("[GAM ALO] NaN values found in ALO diagnostics:");
593        log::error!(
594            "[GAM ALO] eta_tilde: {} NaN values",
595            result.eta_tilde.iter().filter(|&&x| x.is_nan()).count()
596        );
597        log::error!(
598            "[GAM ALO] se_bayes: {} NaN values",
599            result.se_bayes.iter().filter(|&&x| x.is_nan()).count()
600        );
601        log::error!(
602            "[GAM ALO] se_sandwich: {} NaN values",
603            result.se_sandwich.iter().filter(|&&x| x.is_nan()).count()
604        );
605        log::error!(
606            "[GAM ALO] leverage: {} NaN values",
607            result.leverage.iter().filter(|&&x| x.is_nan()).count()
608        );
609        return Err(AloError::InfluenceMatrixFailed {
610            condition_number: f64::INFINITY,
611        });
612    }
613
614    Ok(result)
615}
616
617/// Log detailed leverage percentile diagnostics for a completed ALO computation.
618fn log_leverage_diagnostics(leverage: &Array1<f64>, phi: f64) {
619    let n = leverage.len();
620    if n == 0 {
621        return;
622    }
623
624    let mut invalid_count = 0usize;
625    let mut high_leverage_count = 0usize;
626    let mut threshold_counts = [0usize; LEVERAGE_RATE_THRESHOLDS.len()];
627    let mut finite_leverage = Vec::with_capacity(n);
628
629    for (obs, &ai) in leverage.iter().enumerate() {
630        if ai.is_finite() {
631            finite_leverage.push(ai);
632        }
633
634        if !(0.0..=1.0).contains(&ai) || !ai.is_finite() {
635            invalid_count += 1;
636            log::warn!("[GAM ALO] invalid leverage at i={}, a_ii={:.6e}", obs, ai);
637        } else if ai > LEVERAGE_HIGH_THRESHOLD {
638            high_leverage_count += 1;
639            if ai > LEVERAGE_VERY_HIGH_THRESHOLD {
640                log::warn!("[GAM ALO] very high leverage at i={}, a_ii={:.6e}", obs, ai);
641            }
642        }
643
644        for (idx, threshold) in LEVERAGE_RATE_THRESHOLDS.iter().enumerate() {
645            if ai > *threshold {
646                threshold_counts[idx] += 1;
647            }
648        }
649    }
650
651    if invalid_count > 0 || high_leverage_count > 0 {
652        log::warn!(
653            "[GAM ALO] leverage diagnostics: {} invalid values, {} high values (>0.99)",
654            invalid_count,
655            high_leverage_count
656        );
657    }
658
659    finite_leverage.sort_by(f64::total_cmp);
660
661    let finite_n = finite_leverage.len();
662    let a_mean = if finite_n > 0 {
663        finite_leverage.iter().copied().sum::<f64>() / finite_n as f64
664    } else {
665        0.0
666    };
667    let a_median = percentile_from_sorted(&finite_leverage, LEVERAGE_PERCENTILES[0]);
668    let a_p95 = percentile_from_sorted(&finite_leverage, LEVERAGE_PERCENTILES[1]);
669    let a_p99 = percentile_from_sorted(&finite_leverage, LEVERAGE_PERCENTILES[2]);
670    let a_max = finite_leverage.last().copied().unwrap_or(0.0);
671
672    // Routine per-ALO leverage summary: a diagnostic snapshot, not an
673    // anomaly. Emitted at `info!` so it is visible when the host raises
674    // verbosity (CLI `-v`; `gamfit.set_log_level("info")`) but silent at the
675    // default `Warn` level (genuine anomalies — invalid / very
676    // high leverage — are logged at `warn!` above and stay visible). This
677    // line fires once per ALO computation, which recurs across the outer
678    // smoothing loop, so at `warn!` it was a dominant source of stderr noise
679    // on perfectly healthy fits (#1689).
680    log::info!(
681        "[GAM ALO] leverage: n={}, mean={:.3e}, median={:.3e}, p95={:.3e}, p99={:.3e}, max={:.3e}",
682        n,
683        a_mean,
684        a_median,
685        a_p95,
686        a_p99,
687        a_max
688    );
689    log::info!(
690        "[GAM ALO] high-leverage: a>0.90: {:.2}%, a>0.95: {:.2}%, a>0.99: {:.2}%, dispersion phi={:.3e}",
691        100.0 * (threshold_counts[0] as f64) / n as f64,
692        100.0 * (threshold_counts[1] as f64) / n as f64,
693        100.0 * (threshold_counts[2] as f64) / n as f64,
694        phi
695    );
696}
697
698/// Model-agnostic input for ALO diagnostics.
699///
700/// Any model with a design matrix, penalized Hessian, and IRLS geometry can
701/// compute ALO leverages and leave-one-out predictions. This decouples ALO
702/// from the single-block PIRLS solver and enables diagnostics for GAMLSS,
703/// survival, and joint models.
704pub struct AloInput<'a> {
705    /// Dense design matrix X (n × p).
706    pub design: &'a Array2<f64>,
707    /// Penalized Hessian H = X'WX + S(λ) at convergence (p × p).
708    pub penalized_hessian: &'a Array2<f64>,
709    /// Hessian-side IRLS weights W_H at convergence (n). Sign-honest: for
710    /// non-canonical links the observed-information diagonal can have negative
711    /// entries, so the typed [`SignedWeightsView`] is the contract here. PSD
712    /// callers needing to promote (e.g. the canonical-link case where the
713    /// caller has discharged W_H ≥ 0 algebraically) can route through
714    /// `SignedWeightsView::as_psd()` at the consumer.
715    pub hessian_weights: SignedWeightsView<'a>,
716    /// Score-side IRLS weights W_S paired with `working_response` (n).
717    /// PSD-by-construction: the score-side Fisher weights `h'²/(φ V(μ)) ≥ 0`.
718    pub score_weights: PsdWeightsView<'a>,
719    /// IRLS working response at convergence (n).
720    pub working_response: &'a Array1<f64>,
721    /// Fitted linear predictor η̂ (n).
722    pub eta: &'a Array1<f64>,
723    /// Offset vector (n). Pass zeros if no offset.
724    pub offset: &'a Array1<f64>,
725    /// Link function (for phi determination).
726    pub link: LinkFunction,
727    /// Dispersion parameter φ. For non-Gaussian families this is 1.0.
728    pub phi: f64,
729    /// Optional penalty square root E with E^T E = S(λ) (rank × p) for sandwich SE.
730    /// When `None`, sandwich SE is set equal to Bayesian SE.
731    pub penalty_root: Option<&'a Array2<f64>>,
732    /// Ridge added to the Hessian for logdet surface.
733    pub ridge: f64,
734    /// Optional per-row score/curvature evaluator `(i, η) → (ℓ_i'(η), ℓ_i''(η))`.
735    ///
736    /// When supplied, the leave-`i`-out predictor is obtained by solving the
737    /// frozen-curvature scalar fixed point `η = η̂_i + a_ii ℓ_i'(η)` to
738    /// convergence (see [`alo_eta_exact_frozen_curvature`]) instead of taking a
739    /// single Newton step. This eliminates the first-order linearization error
740    /// that the one-step ALO incurs on small-`n`, strongly curved likelihoods
741    /// (e.g. binomial logistic regression). Non-convergence or invalid scalar
742    /// Newton geometry is returned as an ALO error. When `None`, the classical
743    /// single-Newton-step ALO formula is used. The evaluator must be consistent
744    /// with `hessian_weights` at convergence: `ℓ_i''(η̂_i) = W_H[i]` and
745    /// `ℓ_i'(η̂_i) = W_S[i]·((η̂_i−o_i) − (z_i−o_i))`.
746    pub score_curvature: Option<&'a AloScalarScoreCurvature<'a>>,
747}
748
749impl<'a> AloInput<'a> {
750    /// Build an `AloInput` from `FitGeometry` and associated vectors.
751    pub fn from_geometry(
752        geom: &'a FitGeometry,
753        design: &'a Array2<f64>,
754        eta: &'a Array1<f64>,
755        offset: &'a Array1<f64>,
756        link: LinkFunction,
757        phi: f64,
758    ) -> Self {
759        // FitGeometry stores one working-weight vector, so this constructor is
760        // exact only when the score- and Hessian-side IRLS weights coincide
761        // (canonical-link case where Fisher == Observed). In that path the
762        // diagonal is the Fisher weight `h'²/(φ V(μ)) ≥ 0`, so the PSD
763        // obligation is discharged algebraically without a runtime scan;
764        // `as_signed()` re-views the same buffer for the Hessian-side slot.
765        let psd_w = PsdWeightsView::from_view_unchecked(geom.working_weights.view());
766        Self {
767            design,
768            penalized_hessian: &geom.penalized_hessian,
769            hessian_weights: psd_w.as_signed(),
770            score_weights: psd_w,
771            working_response: &geom.working_response,
772            eta,
773            offset,
774            link,
775            phi,
776            penalty_root: None,
777            ridge: 0.0,
778            score_curvature: None,
779        }
780    }
781
782    /// Build an `AloInput` from a `FitGeometry`'s penalized Hessian plus
783    /// externally supplied working weights / working response.
784    ///
785    /// The row-sized IRLS working vectors are *derived* quantities: at
786    /// convergence they are deterministic functions of the linear predictor
787    /// `η̂ = Xβ̂`, the response `y`, and the family (`w_i = h'(η̂_i)²/(φ V(μ̂_i))·
788    /// prior_i`, `z_i = η̂_i + (y_i−μ̂_i)/h'(η̂_i)`). A size-compacted saved model
789    /// keeps the p×p `penalized_hessian` (n-independent) but drops those n-sized
790    /// vectors; a post-fit consumer such as `gam diagnose` reconstructs them from
791    /// the saved `β` by replaying the same PIRLS working-state update the fit
792    /// used, then feeds them here. This preserves the size win of dropping the
793    /// working vectors from persistence while still serving the exact geometry
794    /// ALO path (no refit, exact saved Hessian).
795    ///
796    /// Same canonical (Fisher == Observed) contract as [`from_geometry`]: the
797    /// supplied `working_weights` are the score-side Fisher weights and are
798    /// re-viewed for the Hessian-side slot via `as_signed()`.
799    ///
800    /// [`from_geometry`]: AloInput::from_geometry
801    pub fn from_geometry_with_working_state(
802        geom: &'a FitGeometry,
803        design: &'a Array2<f64>,
804        eta: &'a Array1<f64>,
805        offset: &'a Array1<f64>,
806        link: LinkFunction,
807        phi: f64,
808        working_weights: &'a Array1<f64>,
809        working_response: &'a Array1<f64>,
810    ) -> Self {
811        let psd_w = PsdWeightsView::from_view_unchecked(working_weights.view());
812        Self {
813            design,
814            penalized_hessian: &geom.penalized_hessian,
815            hessian_weights: psd_w.as_signed(),
816            score_weights: psd_w,
817            working_response,
818            eta,
819            offset,
820            link,
821            phi,
822            penalty_root: None,
823            ridge: 0.0,
824            score_curvature: None,
825        }
826    }
827}
828
829/// Compute ALO diagnostics from model-agnostic inputs.
830///
831/// This is the generalized entry point that works for any model type.
832/// For standard single-block GAMs, prefer `compute_alo_diagnostics_from_fit`
833/// which automatically extracts the PIRLS geometry (including sandwich SE).
834pub fn compute_alo_from_input(input: &AloInput) -> Result<AloDiagnostics, EstimationError> {
835    compute_alo_from_input_inner(input).map_err(EstimationError::from)
836}
837
838fn compute_alo_from_input_inner(input: &AloInput) -> Result<AloDiagnostics, AloError> {
839    let x_dense = input.design;
840    let n = x_dense.nrows();
841    let p = x_dense.ncols();
842    // Bind the underlying ArrayView1 once so the loop body can index and
843    // borrow as before; the sign-character contract lives in the
844    // `AloInput` field types, not in this local binding.
845    let w_h = input.hessian_weights.view();
846    let w_s = input.score_weights.view();
847
848    validate_alo_solve_setup(input, n, p)?;
849
850    let factor = StableSolver::new("alo penalized hessian")
851        .factorize(input.penalized_hessian)
852        .map_err(|_| AloError::InfluenceMatrixFailed {
853            condition_number: f64::INFINITY,
854        })?;
855
856    let xt = x_dense.t();
857    let phi = input.phi;
858
859    let mut aii = Array1::<f64>::zeros(n);
860    let mut x_hinv_x_diag = Array1::<f64>::zeros(n);
861    let mut se_bayes = Array1::<f64>::zeros(n);
862    let mut se_sandwich = Array1::<f64>::zeros(n);
863
864    let block_cols = ALO_RHS_BLOCK_COLS;
865    // Allocate the RHS scratch in column-major (Fortran) order so its column
866    // slices are contiguous and align with faer's column-major solve output.
867    // This removes redundant `xrow = x_dense.row(obs)` indirection inside the
868    // per-observation loop: rhs_chunk_buf already holds X^T at the right cols.
869    let mut rhs_chunk_buf = Array2::<f64>::zeros((p, block_cols).f());
870    // Reusable faer column-major buffer for X*S, where S = H^{-1}X_i for the
871    // current RHS chunk.  The sandwich SE must use the same frozen-curvature
872    // meat as the exact LOO reference, `X' W X`, directly; reconstructing it as
873    // `H - S_penalty - ridge*I` is brittle because the exported stabilized
874    // Hessian may include curvature/stabilization details that are not exactly
875    // represented by the penalty root plus public ridge scalar.
876    let mut xs_chunk_storage = FaerMat::<f64>::zeros(n, block_cols);
877    let x_dense_view = FaerArrayView::new(x_dense);
878
879    for chunk_start in (0..n).step_by(block_cols) {
880        let chunk_end = (chunk_start + block_cols).min(n);
881        let width = chunk_end - chunk_start;
882
883        rhs_chunk_buf
884            .slice_mut(s![.., ..width])
885            .assign(&xt.slice(s![.., chunk_start..chunk_end]));
886
887        let rhs_chunkview = rhs_chunk_buf.slice(s![.., ..width]);
888        let rhs_chunk = FaerArrayView::new(&rhs_chunkview);
889        // s_chunk is owned column-major faer storage; its column slices are
890        // contiguous and can be read directly via `col_as_slice` — no need to
891        // materialize a parallel ndarray copy.
892        let s_chunk = factor.solve(rhs_chunk.as_ref());
893
894        let mut xs_target = xs_chunk_storage.as_mut().subcols_mut(0, width);
895        matmul(
896            xs_target.rb_mut(),
897            Accum::Replace,
898            x_dense_view.as_ref(),
899            s_chunk.as_ref(),
900            1.0,
901            Par::Seq,
902        );
903
904        let rhs_view = rhs_chunk_buf.slice(s![.., ..width]);
905
906        for local_col in 0..width {
907            let obs = chunk_start + local_col;
908            // rhs is column-major Fortran ndarray; faer Mat columns are
909            // contiguous by construction. Both accesses borrow the existing
910            // storage directly — no per-column copy.
911            let rhs_col = rhs_view.column(local_col);
912            let rhs_slice = rhs_col.as_slice().expect("column-major col contiguous");
913            let s_slice = s_chunk.col_as_slice(local_col);
914
915            let mut x_hinv_x = 0.0f64;
916            // Fused dot product over the current solve column.
917            for k in 0..p {
918                let sval = s_slice[k];
919                let xval = rhs_slice[k];
920                x_hinv_x = sval.mul_add(xval, x_hinv_x);
921            }
922            let ai = w_h[obs].max(0.0) * x_hinv_x;
923            aii[obs] = ai;
924            x_hinv_x_diag[obs] = x_hinv_x;
925
926            let var_bayes = bayesvar_eta(phi, x_hinv_x);
927            let xs_slice = xs_chunk_storage.col_as_slice(local_col);
928            let mut meat_quad = 0.0f64;
929            for row in 0..n {
930                let xs = xs_slice[row];
931                // Sandwich meat is the SCORE covariance Xᵀ diag(W_S) X (Fisher,
932                // PSD by construction), not the observed-information Hessian
933                // weight W_H: the estimator is Var = H⁻¹·Cov(score)·H⁻¹ with the
934                // bread H = Xᵀ W_H X + S. For non-canonical links W_H ≠ W_S (and
935                // W_H can be negative), so using W_H here gives a wrong — even
936                // negative — sandwich SE. See `AloInput::score_weights`.
937                meat_quad += w_s[row] * xs * xs;
938            }
939            let var_sandwich = sandwichvar_eta_from_meat(phi, meat_quad);
940
941            if !var_bayes.is_finite() || !var_sandwich.is_finite() {
942                return Err(AloError::LooComputationFailed {
943                    reason: format!(
944                        "ALO variance is not finite at row {obs}: bayes={var_bayes:.6e}, sandwich={var_sandwich:.6e}"
945                    ),
946                });
947            }
948            let bayes_tol = variance_negative_tolerance(phi * x_hinv_x.abs());
949            if var_bayes < -bayes_tol {
950                return Err(AloError::LooComputationFailed {
951                    reason: format!(
952                        "ALO Bayesian variance is materially negative at row {obs}: var={var_bayes:.6e}, tol={bayes_tol:.6e}"
953                    ),
954                });
955            }
956            let sandwich_scale = phi * meat_quad.abs().max(x_hinv_x.abs());
957            let sandwich_tol = variance_negative_tolerance(sandwich_scale);
958            if var_sandwich < -sandwich_tol {
959                return Err(AloError::LooComputationFailed {
960                    reason: format!(
961                        "ALO sandwich variance is materially negative at row {obs}: var={var_sandwich:.6e}, tol={sandwich_tol:.6e}"
962                    ),
963                });
964            }
965
966            se_bayes[obs] = var_bayes.max(0.0).sqrt();
967            se_sandwich[obs] = var_sandwich.max(0.0).sqrt();
968        }
969    }
970
971    let eta_hat = input.eta;
972    let z = input.working_response;
973    let offset = input.offset;
974
975    use rayon::prelude::*;
976    let eta_tilde_vec: Vec<f64> = (0..n)
977        .into_par_iter()
978        .map(|i| {
979            let denom_raw = 1.0 - aii[i];
980            if denom_raw <= ALO_DENOMINATOR_MIN || !denom_raw.is_finite() {
981                return Err(AloError::LooComputationFailed {
982                    reason: format!(
983                        "ALO denominator is too small at row {i}: a_ii={:.6e}, 1-a_ii={:.6e}, min={:.1e}",
984                        aii[i], denom_raw, ALO_DENOMINATOR_MIN
985                    ),
986                });
987            }
988            let one_step = alo_eta_updatewith_offset(
989                eta_hat[i],
990                z[i],
991                offset[i],
992                x_hinv_x_diag[i],
993                w_s[i],
994                denom_raw,
995            );
996            // When the family score/curvature evaluator is supplied, solve the
997            // exact frozen-curvature leave-i-out fixed point (anchored at η̂_i,
998            // the basin that limits to the in-sample fit) instead of taking the
999            // single Newton step. a_ii here is the unweighted influence
1000            // x_i^T H^{-1} x_i (= x_hinv_x_diag[i]); the per-row curvature
1001            // W_H[i] = ℓ_i''(η̂_i) is folded into the scalar fixed point via
1002            // score_curvature. Non-canonical links fall back to `one_step`.
1003            let v = if let Some(score_curvature) = input.score_curvature {
1004                alo_eta_exact_frozen_curvature(
1005                    eta_hat[i],
1006                    x_hinv_x_diag[i],
1007                    &|eta| score_curvature(i, eta),
1008                )
1009                .map_err(|err| AloError::LooComputationFailed {
1010                    reason: format!(
1011                        "ALO exact frozen-curvature solve failed at row {i}: {err}"
1012                    ),
1013                })?
1014            } else {
1015                one_step
1016            };
1017            if !v.is_finite() {
1018                return Err(AloError::LooComputationFailed {
1019                    reason: format!("ALO eta_tilde is not finite at row {i}: eta_tilde={v}"),
1020                });
1021            }
1022            Ok(v)
1023        })
1024        .collect::<Result<_, _>>()?;
1025    let eta_tilde = Array1::from(eta_tilde_vec);
1026
1027    Ok(AloDiagnostics {
1028        eta_tilde,
1029        se_bayes,
1030        se_sandwich,
1031        pred_identity: eta_hat.clone(),
1032        leverage: aii,
1033        fisherweights: w_h.to_owned(),
1034    })
1035}
1036
1037fn validate_alo_solve_setup(input: &AloInput, n: usize, p: usize) -> Result<(), AloError> {
1038    let h = input.penalized_hessian;
1039    if h.nrows() != p || h.ncols() != p {
1040        return Err(AloError::InvalidInput {
1041            reason: format!(
1042                "ALO diagnostics require a dense exact penalized Hessian with shape {p}x{p}; got {}x{}",
1043                h.nrows(),
1044                h.ncols()
1045            ),
1046        });
1047    }
1048    if h.iter().any(|v| !v.is_finite()) {
1049        return Err(AloError::InvalidInput {
1050            reason: "ALO diagnostics require a finite dense exact penalized Hessian".to_string(),
1051        });
1052    }
1053    for i in 0..p {
1054        for j in 0..i {
1055            let a = h[[i, j]];
1056            let b = h[[j, i]];
1057            let scale = a.abs().max(b.abs()).max(1.0);
1058            if (a - b).abs() > HESSIAN_SYMMETRY_REL_TOL * scale {
1059                return Err(AloError::InvalidInput {
1060                    reason: format!(
1061                        "ALO diagnostics require a symmetric dense exact penalized Hessian; entries ({i},{j}) and ({j},{i}) differ by {:.3e}",
1062                        (a - b).abs()
1063                    ),
1064                });
1065            }
1066        }
1067    }
1068
1069    let vector_lengths = [
1070        ("hessian_weights", input.hessian_weights.len()),
1071        ("score_weights", input.score_weights.len()),
1072        ("working_response", input.working_response.len()),
1073        ("eta", input.eta.len()),
1074        ("offset", input.offset.len()),
1075    ];
1076    for (name, len) in vector_lengths {
1077        if len != n {
1078            return Err(AloError::InvalidInput {
1079                reason: format!("ALO diagnostics require {name} length {n}; got {len}"),
1080            });
1081        }
1082    }
1083    if input.hessian_weights.view().iter().any(|v| !v.is_finite()) {
1084        return Err(AloError::WeightInvalid {
1085            reason: "ALO diagnostics require finite Hessian-side weights".to_string(),
1086        });
1087    }
1088    if input.score_weights.view().iter().any(|v| !v.is_finite()) {
1089        return Err(AloError::WeightInvalid {
1090            reason: "ALO diagnostics require finite score-side weights".to_string(),
1091        });
1092    }
1093    if input.working_response.iter().any(|v| !v.is_finite()) {
1094        return Err(AloError::WeightInvalid {
1095            reason: "ALO diagnostics require finite working responses".to_string(),
1096        });
1097    }
1098    if input.eta.iter().any(|v| !v.is_finite()) || input.offset.iter().any(|v| !v.is_finite()) {
1099        return Err(AloError::InvalidInput {
1100            reason: "ALO diagnostics require finite linear predictors and offsets".to_string(),
1101        });
1102    }
1103    if !input.phi.is_finite() || input.phi <= 0.0 {
1104        return Err(AloError::InvalidInput {
1105            reason: format!(
1106                "ALO diagnostics require positive finite dispersion phi; got {}",
1107                input.phi
1108            ),
1109        });
1110    }
1111    if !input.ridge.is_finite() || input.ridge < 0.0 {
1112        return Err(AloError::InvalidInput {
1113            reason: format!(
1114                "ALO diagnostics require a finite non-negative Hessian ridge; got {}",
1115                input.ridge
1116            ),
1117        });
1118    }
1119    if let Some(e) = input.penalty_root {
1120        if e.ncols() != p {
1121            return Err(AloError::InvalidInput {
1122                reason: format!(
1123                    "ALO diagnostics require penalty root to have {p} columns; got {}",
1124                    e.ncols()
1125                ),
1126            });
1127        }
1128        if e.iter().any(|v| !v.is_finite()) {
1129            return Err(AloError::InvalidInput {
1130                reason: "ALO diagnostics require finite penalty-root entries".to_string(),
1131            });
1132        }
1133    }
1134    Ok(())
1135}
1136
1137/// Compute ALO diagnostics (eta_tilde, SE, leverage) from a fitted GAM result.
1138pub fn compute_alo_diagnostics_from_fit(
1139    fit: &UnifiedFitResult,
1140    y: ArrayView1<f64>,
1141    link: LinkFunction,
1142) -> Result<AloDiagnostics, EstimationError> {
1143    let pirls = fit
1144        .artifacts
1145        .pirls
1146        .as_ref()
1147        .ok_or_else(|| AloError::InvalidInput {
1148            reason:
1149                "ALO diagnostics require a PIRLS-backed fit; this fit does not expose PIRLS geometry"
1150                    .to_string(),
1151        })
1152        .map_err(EstimationError::from)?;
1153    compute_alo_diagnostics_from_pirls_impl(pirls, y, link)
1154}
1155
1156/// Compute ALO diagnostics from a `UnifiedFitResult`.
1157///
1158/// Extracts `FitGeometry` from `unified.geometry`, builds an `AloInput`
1159/// via `from_geometry`, and delegates to `compute_alo_from_input`.
1160/// This avoids requiring a full `UnifiedFitResult` with PIRLS artifacts.
1161pub fn compute_alo_diagnostics_from_unified(
1162    unified: &UnifiedFitResult,
1163    design: &Array2<f64>,
1164    eta: &Array1<f64>,
1165    offset: &Array1<f64>,
1166    link: LinkFunction,
1167    phi: f64,
1168) -> Result<AloDiagnostics, EstimationError> {
1169    let geom = unified
1170        .geometry
1171        .as_ref()
1172        .ok_or_else(|| AloError::InvalidInput {
1173            reason: "UnifiedFitResult does not contain working-set geometry; \
1174             ALO diagnostics require geometry at convergence"
1175                .to_string(),
1176        })
1177        .map_err(EstimationError::from)?;
1178    let input = AloInput::from_geometry(geom, design, eta, offset, link, phi);
1179    compute_alo_from_input(&input)
1180}
1181
1182/// Compute ALO diagnostics from a PIRLS result for lower-level callers.
1183pub fn compute_alo_diagnostics_from_pirls(
1184    base: &pirls::PirlsResult,
1185    y: ArrayView1<f64>,
1186    link: LinkFunction,
1187) -> Result<AloDiagnostics, EstimationError> {
1188    compute_alo_diagnostics_from_pirls_impl(base, y, link)
1189}
1190
1191/// Exact (one-step) case-deletion influence from a converged PIRLS fit, via
1192/// the one `FitSensitivity` operator (#935).
1193///
1194/// This is the diagnostic the sensitivity operator's `case_deletion` channel
1195/// was built to expose but had no production entry point for: per-observation
1196/// dfbetas `β̂ − β̂₍ᵢ₎`, hat-value leverage `h_ii = w_i x_iᵀ H⁻¹ x_i`, and
1197/// Cook's distance. It is the same factored inverse the REML gradient (IFT),
1198/// ALO, and the Riesz debias already contract — built once at the optimum,
1199/// asked in the leave-one-out direction — so no call site can disagree about
1200/// which `H⁻¹` is meant (the bug class #935 dismantles).
1201///
1202/// The penalized Hessian, design, working weights `w_i = W_H[i]` and working
1203/// residual `z_i − η̂_i` are read straight from the converged geometry — the
1204/// same PIRLS state [`compute_alo_diagnostics_from_pirls`] consumes — so the
1205/// IRLS reduction `scale = w_i r_i / (1 − h_ii)` is exact for the Gaussian
1206/// identity link and the one-step Newton deletion for canonical-link GLMs.
1207/// Returns `None` (rather than emitting `∞`) for any observation whose
1208/// leverage is one, or if the dense Hessian / design is unavailable.
1209pub fn compute_case_deletion_from_pirls(
1210    base: &pirls::PirlsResult,
1211    y: ArrayView1<f64>,
1212    link: LinkFunction,
1213) -> Result<Option<crate::sensitivity::CaseDeletionInfluence>, EstimationError> {
1214    let x_dense_arc = base
1215        .x_transformed
1216        .try_to_dense_arc("case-deletion diagnostics require dense transformed design")
1217        .map_err(|reason| EstimationError::InvalidInput(reason))?;
1218    let x_dense = x_dense_arc.as_ref();
1219    let n = x_dense.nrows();
1220    let p = x_dense.ncols();
1221    if n == 0 || p == 0 {
1222        return Ok(None);
1223    }
1224
1225    // Dispersion φ matches the ALO entry point: estimated RSS/(n−edf) for the
1226    // Gaussian identity link, fixed at 1 for the single-parameter families.
1227    let phi = match link {
1228        LinkFunction::Identity => {
1229            use rayon::iter::{IntoParallelIterator, ParallelIterator};
1230            let rss: f64 = (0..n)
1231                .into_par_iter()
1232                .map(|i| {
1233                    let r = y[i] - base.finalmu[i];
1234                    base.finalweights[i] * r * r
1235                })
1236                .sum();
1237            let dof = (n as f64) - base.edf;
1238            rss / dof.max(1.0)
1239        }
1240        _ => 1.0,
1241    };
1242    if !(phi.is_finite() && phi > 0.0) {
1243        return Ok(None);
1244    }
1245
1246    // The same dense stabilized penalized Hessian ALO materializes; the one
1247    // factored inverse every sensitivity channel shares.
1248    let h_dense = base
1249        .dense_stabilizedhessian_transformed(
1250            "case-deletion diagnostics require exact dense stabilized penalized Hessian",
1251        )
1252        .map_err(|e| match e {
1253            EstimationError::InvalidInput(reason) => EstimationError::InvalidInput(reason),
1254            other => EstimationError::InvalidInput(format!("{other:?}")),
1255        })?;
1256
1257    let factor = match h_dense.cholesky(faer::Side::Lower) {
1258        Ok(f) => f,
1259        // A non-SPD stabilized Hessian means the optimum is rank-deficient in a
1260        // way the dense Cholesky case-deletion path cannot invert; decline
1261        // rather than fabricate an influence diagnostic.
1262        Err(_) => return Ok(None),
1263    };
1264
1265    // Working weights and working residual straight from the IRLS reduction:
1266    // w_i = W_H[i] and r_i = z_i − η̂_i, so w_i r_i is the working score the
1267    // closed-form deletion `scale = w_i r_i / (1 − h_ii)` consumes.
1268    let working_weights = base.finalweights.clone();
1269    let working_residual = &base.solveworking_response - &base.final_eta;
1270
1271    let sensitivity = crate::sensitivity::FitSensitivity::from_faer_cholesky(&factor, p);
1272    Ok(sensitivity.case_deletion(
1273        x_dense,
1274        working_weights.view(),
1275        working_residual.view(),
1276        phi,
1277    ))
1278}
1279
1280// Multi-block ALO for multi-predictor models (GAMLSS, survival, joint)
1281
1282/// Diagnostics returned by multi-block ALO.
1283#[derive(Debug, Clone)]
1284pub struct MultiBlockAloDiagnostics {
1285    /// Corrected linear predictors η̃^{(-i)} for each observation.
1286    /// Outer length = n_obs, inner length = n_blocks (B).
1287    pub eta_tilde: Vec<Array1<f64>>,
1288    /// Per-observation leverage tr(H_ii) where H_ii is the B×B hat-matrix block.
1289    pub leverage: Array1<f64>,
1290    /// Per-observation ALO variance diagonals: for each observation i,
1291    /// Var(Δη_i) ≈ A_i (I - W_i A_i)⁻¹ W_i (I - A_i W_i)⁻¹ A_iᵀ.
1292    /// Outer length = n_obs, inner length = n_blocks (B) containing the
1293    /// diagonal entries of the variance matrix.
1294    pub alo_variance: Vec<Array1<f64>>,
1295    /// Cook-type ALO influence: D_i = Δη_iᵀ W_i Δη_i.
1296    /// Length = n_obs.
1297    pub cook_distance: Array1<f64>,
1298}
1299
1300/// Model-agnostic input for multi-predictor ALO diagnostics.
1301///
1302/// Generalises [`AloInput`] to models with B > 1 linear predictors per
1303/// observation (e.g. location-scale GAMLSS with B=2, or survival models
1304/// with time-dependent predictors).
1305///
1306/// # Mathematical setup
1307///
1308/// For observation i the per-observation Jacobian is a B × p_tot block matrix
1309/// X_i whose b-th row is the i-th row of `block_designs[b]`.  The joint
1310/// hat-matrix block is
1311///
1312///   H_ii = X_i H⁻¹ X_iᵀ W_i     (B × B)
1313///
1314/// where H = Σ_i X_iᵀ W_i X_i + S is the total penalized Hessian and W_i
1315/// is the B × B per-observation weight matrix (negative Hessian of the
1316/// log-likelihood w.r.t. the B predictors at observation i).
1317///
1318/// The ALO leave-one-out correction is
1319///
1320///   Δη_i^ALO = A_i (I_B − W_i A_i)⁻¹ s_i
1321///
1322/// where A_i = X_i H⁻¹ X_iᵀ (the B×B per-observation influence matrix),
1323/// W_i is the B×B per-observation NLL Hessian, and
1324/// s_i = ∇_{η_i} NLL_i(η̂_i) is the B-dimensional score vector.
1325/// This is algebraically equivalent to (I_B − H_ii)⁻¹ H_ii W_i⁻¹ s_i
1326/// but does NOT require W_i⁻¹, which is critical when W_i is singular
1327/// (e.g. at boundary observations in survival models).
1328/// For B = 1 this reduces to the classical scalar ALO formula.
1329pub struct MultiBlockAloInput<'a> {
1330    /// Number of observations.
1331    pub n_obs: usize,
1332    /// Number of predictors per observation (B).
1333    pub n_blocks: usize,
1334    /// B design matrices, each n_obs × p_b.  The total parameter count is
1335    /// p_tot = Σ_b p_b.
1336    pub block_designs: &'a [Array2<f64>],
1337    /// Inverse of the penalized Hessian, H⁻¹ (p_tot × p_tot).
1338    pub penalized_hessian_inv: &'a Array2<f64>,
1339    /// Per-observation weight matrices W_i (B × B).  Length = n_obs.
1340    pub block_weights: Vec<Array2<f64>>,
1341    /// Per-observation score vectors s_i = ∇_{η_i} NLL_i.  Length = n_obs,
1342    /// each entry is B-dimensional.
1343    pub scores: Vec<Array1<f64>>,
1344    /// Fitted linear predictor vectors η̂_i.  Length = n_obs, each entry is
1345    /// B-dimensional.
1346    pub eta_hat: Vec<Array1<f64>>,
1347}
1348
1349/// Compute multi-block ALO diagnostics: corrected η̃ and leverages.
1350///
1351/// # Optimisation note
1352///
1353/// The dominant cost is forming X_i H⁻¹ X_iᵀ for every observation.
1354/// Rather than forming the B × p_tot row-block X_i and multiplying naïvely,
1355/// we precompute for each block b the matrix
1356///
1357///   Q_b = H⁻¹ X_bᵀ      (p_tot × n)
1358///
1359/// Then the (a, b) entry of the B × B matrix X_i H⁻¹ X_iᵀ is simply
1360///
1361///   (X_i H⁻¹ X_iᵀ)_{a,b} = x_{a,i}ᵀ Q_b[:,i]
1362///                           = Σ_k  X_a[i,k] · Q_b[k,i]
1363///
1364/// where x_{a,i} is the i-th row of block-design a.  This turns the per-
1365/// observation work from O(B · p_tot²) into O(B² · p_tot), and the
1366/// precomputation is O(B · p_tot² · n) total via a single blocked solve.
1367pub fn compute_multiblock_alo(
1368    input: &MultiBlockAloInput,
1369) -> Result<MultiBlockAloDiagnostics, EstimationError> {
1370    compute_multiblock_alo_inner(input).map_err(EstimationError::from)
1371}
1372
1373fn compute_multiblock_alo_inner(
1374    input: &MultiBlockAloInput,
1375) -> Result<MultiBlockAloDiagnostics, AloError> {
1376    use rayon::prelude::*;
1377
1378    let n = input.n_obs;
1379    let b = input.n_blocks;
1380    let p_tot = input.penalized_hessian_inv.nrows();
1381
1382    // --- Validate dimensions ---
1383    if input.block_designs.len() != b {
1384        return Err(AloError::InvalidInput {
1385            reason: format!(
1386                "MultiBlockAloInput: expected {} block designs, got {}",
1387                b,
1388                input.block_designs.len()
1389            ),
1390        });
1391    }
1392
1393    // Verify total column count matches p_tot.
1394    let col_sum: usize = input.block_designs.iter().map(|d| d.ncols()).sum();
1395    if col_sum != p_tot {
1396        return Err(AloError::InvalidInput {
1397            reason: format!(
1398                "MultiBlockAloInput: total design columns ({}) != penalized_hessian_inv size ({})",
1399                col_sum, p_tot
1400            ),
1401        });
1402    }
1403
1404    let col_offsets = multiblock_col_offsets(input.block_designs);
1405    let (chunk_size, max_concurrent_chunks) = multiblock_alo_parallel_plan(p_tot, b, n);
1406    let chunk_starts: Vec<usize> = (0..n).step_by(chunk_size).collect();
1407
1408    // Each Rayon worker owns its small B×B/B-vector scratch buffers via
1409    // `map_init`, avoiding cross-thread mutation and avoiding per-observation
1410    // allocations.  The much larger Q panels are bounded by the parallel chunk
1411    // size and by wave-level concurrency, so at most roughly one global memory
1412    // budget worth of p_total × chunk_len panels can be live across workers.
1413    let mut chunk_results: Vec<Result<MultiBlockAloChunkDiagnostics, AloError>> =
1414        Vec::with_capacity(chunk_starts.len());
1415    for chunk_wave in chunk_starts.chunks(max_concurrent_chunks) {
1416        let mut wave_results: Vec<Result<MultiBlockAloChunkDiagnostics, AloError>> = chunk_wave
1417            .par_iter()
1418            .map_init(
1419                || MultiBlockAloScratch::new(b),
1420                |scratch, &chunk_start| {
1421                    let chunk_end = (chunk_start + chunk_size).min(n);
1422                    compute_multiblock_alo_chunk(
1423                        input,
1424                        &col_offsets,
1425                        chunk_start,
1426                        chunk_end,
1427                        scratch,
1428                    )
1429                },
1430            )
1431            .collect();
1432        chunk_results.append(&mut wave_results);
1433    }
1434
1435    let mut eta_tilde = Vec::with_capacity(n);
1436    let mut leverage = Array1::<f64>::zeros(n);
1437    let mut alo_variance = Vec::with_capacity(n);
1438    let mut cook_distance = Array1::<f64>::zeros(n);
1439
1440    let mut chunks = Vec::with_capacity(chunk_results.len());
1441    for result in chunk_results {
1442        chunks.push(result?);
1443    }
1444    chunks.sort_unstable_by_key(|chunk| chunk.chunk_start);
1445
1446    for chunk in chunks {
1447        let chunk_start = chunk.chunk_start;
1448        eta_tilde.extend(chunk.eta_tilde);
1449        alo_variance.extend(chunk.alo_variance);
1450        for (local_i, lev) in chunk.leverage.into_iter().enumerate() {
1451            leverage[chunk_start + local_i] = lev;
1452        }
1453        for (local_i, cook) in chunk.cook_distance.into_iter().enumerate() {
1454            cook_distance[chunk_start + local_i] = cook;
1455        }
1456    }
1457
1458    Ok(MultiBlockAloDiagnostics {
1459        eta_tilde,
1460        leverage,
1461        alo_variance,
1462        cook_distance,
1463    })
1464}
1465
1466#[inline]
1467fn multiblock_alo_parallel_plan(p_tot: usize, n_blocks: usize, n_obs: usize) -> (usize, usize) {
1468    if p_tot == 0 || n_blocks == 0 || n_obs == 0 {
1469        return (1, 1);
1470    }
1471    let bytes_per_obs = (p_tot * n_blocks * std::mem::size_of::<f64>()).max(1);
1472    let workers = rayon::current_num_threads().max(1);
1473    let max_concurrent_chunks = (MULTIBLOCK_ALO_MEMORY_BUDGET_BYTES / bytes_per_obs)
1474        .max(1)
1475        .min(workers);
1476    let per_worker_budget =
1477        (MULTIBLOCK_ALO_MEMORY_BUDGET_BYTES / max_concurrent_chunks).max(bytes_per_obs);
1478    let budget_obs = (per_worker_budget / bytes_per_obs).max(1);
1479    (budget_obs.min(n_obs), max_concurrent_chunks)
1480}
1481
1482struct MultiBlockAloScratch {
1483    a_i: Vec<f64>,
1484    wa: Vec<f64>,
1485    aw: Vec<f64>,
1486    imwa: Vec<f64>,
1487    imaw: Vec<f64>,
1488    perm_imwa: Vec<usize>,
1489    perm_imaw: Vec<usize>,
1490    delta_eta: Vec<f64>,
1491    rhs_buf: Vec<f64>,
1492    w_u: Vec<f64>,
1493    var_diag_buf: Vec<f64>,
1494    w_flat: Vec<f64>,
1495    lu_scratch: Vec<f64>,
1496}
1497
1498impl MultiBlockAloScratch {
1499    fn new(b: usize) -> Self {
1500        let bb_sz = b * b;
1501        Self {
1502            a_i: vec![0.0f64; bb_sz],
1503            wa: vec![0.0f64; bb_sz],
1504            aw: vec![0.0f64; bb_sz],
1505            imwa: vec![0.0f64; bb_sz],
1506            imaw: vec![0.0f64; bb_sz],
1507            perm_imwa: vec![0usize; b],
1508            perm_imaw: vec![0usize; b],
1509            delta_eta: vec![0.0f64; b],
1510            rhs_buf: vec![0.0f64; b],
1511            w_u: vec![0.0f64; b],
1512            var_diag_buf: vec![0.0f64; b],
1513            w_flat: vec![0.0f64; bb_sz],
1514            lu_scratch: vec![0.0f64; b],
1515        }
1516    }
1517}
1518
1519struct MultiBlockAloChunkDiagnostics {
1520    chunk_start: usize,
1521    eta_tilde: Vec<Array1<f64>>,
1522    leverage: Vec<f64>,
1523    alo_variance: Vec<Array1<f64>>,
1524    cook_distance: Vec<f64>,
1525}
1526
1527fn compute_multiblock_alo_chunk(
1528    input: &MultiBlockAloInput,
1529    col_offsets: &[usize],
1530    chunk_start: usize,
1531    chunk_end: usize,
1532    scratch: &mut MultiBlockAloScratch,
1533) -> Result<MultiBlockAloChunkDiagnostics, AloError> {
1534    let b = input.n_blocks;
1535    let chunk_len = chunk_end - chunk_start;
1536
1537    let mut q_blocks = Vec::with_capacity(b);
1538    for blk in 0..b {
1539        let x_chunk_t = input.block_designs[blk]
1540            .slice(s![chunk_start..chunk_end, ..])
1541            .t()
1542            .to_owned();
1543        let off_b = col_offsets[blk];
1544        let h_slice = input
1545            .penalized_hessian_inv
1546            .slice(s![.., off_b..off_b + x_chunk_t.nrows()])
1547            .to_owned();
1548        q_blocks.push(h_slice.dot(&x_chunk_t));
1549    }
1550
1551    let mut eta_tilde = Vec::with_capacity(chunk_len);
1552    let mut leverage = vec![0.0f64; chunk_len];
1553    let mut alo_variance = Vec::with_capacity(chunk_len);
1554    let mut cook_distance = vec![0.0f64; chunk_len];
1555
1556    for local_i in 0..chunk_len {
1557        let i = chunk_start + local_i;
1558        let w_i = &input.block_weights[i];
1559
1560        // Flatten W_i once per observation (row-major).
1561        for r in 0..b {
1562            for c in 0..b {
1563                scratch.w_flat[r * b + c] = w_i[(r, c)];
1564            }
1565        }
1566
1567        // --- Assemble A_i = X_i H⁻¹ X_iᵀ  (B × B), row-major flat. ---
1568        for a in 0..b {
1569            let x_a = &input.block_designs[a];
1570            let p_a = x_a.ncols();
1571            let off_a = col_offsets[a];
1572            let xa_row = x_a.row(i);
1573            for bb in 0..b {
1574                let q_bb = &q_blocks[bb];
1575                let mut dot = 0.0f64;
1576                for k in 0..p_a {
1577                    dot += xa_row[k] * q_bb[(off_a + k, local_i)];
1578                }
1579                scratch.a_i[a * b + bb] = dot;
1580            }
1581        }
1582
1583        // WA = W_i · A_i (row-major).
1584        mat_mul_flat(&scratch.w_flat, &scratch.a_i, &mut scratch.wa, b);
1585        // AW = A_i · W_i (row-major).
1586        mat_mul_flat(&scratch.a_i, &scratch.w_flat, &mut scratch.aw, b);
1587
1588        // Trace of H_ii = A_i W_i (= AW): leverage[i].
1589        // (Original code wrote H_ii = A · W — the same operator we already have in `aw`.)
1590        let mut tr = 0.0f64;
1591        for d in 0..b {
1592            tr += scratch.aw[d * b + d];
1593        }
1594        leverage[local_i] = tr;
1595
1596        // Build (I - W A) and (I - A W) into imwa/imaw.
1597        for r in 0..b {
1598            for c in 0..b {
1599                let idx = r * b + c;
1600                let id = if r == c { 1.0 } else { 0.0 };
1601                scratch.imwa[idx] = id - scratch.wa[idx];
1602                scratch.imaw[idx] = id - scratch.aw[idx];
1603            }
1604        }
1605
1606        // Factor in place with partial pivoting; ridge on the diagonal if singular.
1607        // Equivalence with original: original computed det via det_small, regularized
1608        // by adding eps=1e-6 to the diagonal when |det| < 1e-12, then re-factored on
1609        // the regularized matrix. Here we factor directly; if any pivot is below the
1610        // singular threshold we add the ridge once and re-factor — same numerical path.
1611        if !lu_factor_in_place(&mut scratch.imwa, &mut scratch.perm_imwa, b) {
1612            for r in 0..b {
1613                for c in 0..b {
1614                    let idx = r * b + c;
1615                    let id = if r == c { 1.0 } else { 0.0 };
1616                    scratch.imwa[idx] = id - scratch.wa[idx];
1617                }
1618            }
1619            for d in 0..b {
1620                scratch.imwa[d * b + d] += ALO_LOCAL_BLOCK_RIDGE;
1621            }
1622            let refactored = lu_factor_in_place(&mut scratch.imwa, &mut scratch.perm_imwa, b);
1623            assert!(
1624                refactored,
1625                "ALO local block remained singular after ridge regularization"
1626            );
1627        }
1628        if !lu_factor_in_place(&mut scratch.imaw, &mut scratch.perm_imaw, b) {
1629            for r in 0..b {
1630                for c in 0..b {
1631                    let idx = r * b + c;
1632                    let id = if r == c { 1.0 } else { 0.0 };
1633                    scratch.imaw[idx] = id - scratch.aw[idx];
1634                }
1635            }
1636            for d in 0..b {
1637                scratch.imaw[d * b + d] += ALO_LOCAL_BLOCK_RIDGE;
1638            }
1639            let refactored = lu_factor_in_place(&mut scratch.imaw, &mut scratch.perm_imaw, b);
1640            assert!(
1641                refactored,
1642                "ALO local variance block remained singular after ridge regularization"
1643            );
1644        }
1645
1646        // v_i = (I - W A)⁻¹ s_i  -- solve into rhs_buf.
1647        let s_i = &input.scores[i];
1648        for k in 0..b {
1649            scratch.rhs_buf[k] = s_i[k];
1650        }
1651        lu_solve_in_place(
1652            &scratch.imwa,
1653            &scratch.perm_imwa,
1654            &mut scratch.rhs_buf,
1655            &mut scratch.lu_scratch,
1656            b,
1657        );
1658        // delta_eta = A_i · v_i
1659        for r in 0..b {
1660            let mut acc = 0.0f64;
1661            let row_off = r * b;
1662            for k in 0..b {
1663                acc += scratch.a_i[row_off + k] * scratch.rhs_buf[k];
1664            }
1665            scratch.delta_eta[r] = acc;
1666        }
1667
1668        let eta_i = &input.eta_hat[i];
1669        let mut corrected = Array1::<f64>::zeros(b);
1670        for d in 0..b {
1671            corrected[d] = eta_i[d] + scratch.delta_eta[d];
1672        }
1673        eta_tilde.push(corrected);
1674
1675        // Cook's distance: δη^T W δη.
1676        let mut cook = 0.0f64;
1677        for r in 0..b {
1678            let mut w_delta_r = 0.0f64;
1679            let row_off = r * b;
1680            for k in 0..b {
1681                w_delta_r += scratch.w_flat[row_off + k] * scratch.delta_eta[k];
1682            }
1683            cook += scratch.delta_eta[r] * w_delta_r;
1684        }
1685        cook_distance[local_i] = cook;
1686
1687        // var_diag[d] = a_d^T (I-WA)⁻¹ W (I-AW)⁻¹ a_d
1688        // where a_d is the d-th row of A_i.
1689        // Reuses already-factored imwa and imaw (one LU factorization each, reused
1690        // across all B right-hand sides — major saving over the original which redid
1691        // both LU decompositions B times per observation).
1692        for d in 0..b {
1693            let row_off = d * b;
1694            // u_d = (I - A W)⁻¹ a_d
1695            for k in 0..b {
1696                scratch.rhs_buf[k] = scratch.a_i[row_off + k];
1697            }
1698            lu_solve_in_place(
1699                &scratch.imaw,
1700                &scratch.perm_imaw,
1701                &mut scratch.rhs_buf,
1702                &mut scratch.lu_scratch,
1703                b,
1704            );
1705            // w_u = W u_d
1706            for r in 0..b {
1707                let mut acc = 0.0f64;
1708                let wr = r * b;
1709                for k in 0..b {
1710                    acc += scratch.w_flat[wr + k] * scratch.rhs_buf[k];
1711                }
1712                scratch.w_u[r] = acc;
1713            }
1714            // t_d = (I - W A)⁻¹ w_u  (back-solve in place using w_u as RHS).
1715            lu_solve_in_place(
1716                &scratch.imwa,
1717                &scratch.perm_imwa,
1718                &mut scratch.w_u,
1719                &mut scratch.lu_scratch,
1720                b,
1721            );
1722            // v_dd = a_d^T t_d
1723            let mut v_dd = 0.0f64;
1724            for k in 0..b {
1725                v_dd += scratch.a_i[row_off + k] * scratch.w_u[k];
1726            }
1727            scratch.var_diag_buf[d] = v_dd.max(0.0);
1728        }
1729        let mut var_diag = Array1::<f64>::zeros(b);
1730        for d in 0..b {
1731            var_diag[d] = scratch.var_diag_buf[d];
1732        }
1733        alo_variance.push(var_diag);
1734    }
1735
1736    Ok(MultiBlockAloChunkDiagnostics {
1737        chunk_start,
1738        eta_tilde,
1739        leverage,
1740        alo_variance,
1741        cook_distance,
1742    })
1743}
1744
1745/// B × B row-major matmul: out = a · b.
1746#[inline]
1747fn mat_mul_flat(a: &[f64], b_mat: &[f64], out: &mut [f64], b: usize) {
1748    for r in 0..b {
1749        let ar = r * b;
1750        let or = r * b;
1751        for c in 0..b {
1752            let mut acc = 0.0f64;
1753            for k in 0..b {
1754                acc += a[ar + k] * b_mat[k * b + c];
1755            }
1756            out[or + c] = acc;
1757        }
1758    }
1759}
1760
1761/// LU-decompose a B × B row-major matrix in place with partial pivoting and
1762/// physical row swaps. Returns false if any pivot |a_kk| < 1e-12 (singular).
1763/// On success, `m` holds L (strict lower, unit diag implicit) and U (upper, diag
1764/// included); `perm[k]` records the original-row index that ended up in physical
1765/// row k after pivoting. Pivot threshold matches the original `det_small < 1e-12`
1766/// path so the regularization branch fires under equivalent conditions.
1767fn lu_factor_in_place(m: &mut [f64], perm: &mut [usize], b: usize) -> bool {
1768    for i in 0..b {
1769        perm[i] = i;
1770    }
1771    for col in 0..b {
1772        // Partial pivot on column `col` over physical rows `[col..b]`.
1773        let mut max_val = m[col * b + col].abs();
1774        let mut max_idx = col;
1775        for row in (col + 1)..b {
1776            let v = m[row * b + col].abs();
1777            if v > max_val {
1778                max_val = v;
1779                max_idx = row;
1780            }
1781        }
1782        if max_val < LU_PIVOT_SINGULAR_TOL {
1783            return false;
1784        }
1785        if max_idx != col {
1786            // Physically swap rows `col` and `max_idx` (full row, all columns).
1787            for k in 0..b {
1788                m.swap(col * b + k, max_idx * b + k);
1789            }
1790            perm.swap(col, max_idx);
1791        }
1792        let pivot = m[col * b + col];
1793        for row in (col + 1)..b {
1794            let factor = m[row * b + col] / pivot;
1795            m[row * b + col] = factor; // store L below diag
1796            for k in (col + 1)..b {
1797                let upd = factor * m[col * b + k];
1798                m[row * b + k] -= upd;
1799            }
1800        }
1801    }
1802    true
1803}
1804
1805/// Solve L U x = P rhs using a previously factored matrix (LU in `m`, perm).
1806/// Writes the solution back into `rhs`. `scratch` must have length ≥ b.
1807fn lu_solve_in_place(m: &[f64], perm: &[usize], rhs: &mut [f64], scratch: &mut [f64], b: usize) {
1808    // Forward substitution Ly = P rhs (L is unit-diag, strict lower of m).
1809    let y = &mut scratch[..b];
1810    for row in 0..b {
1811        let mut s = rhs[perm[row]];
1812        for k in 0..row {
1813            s -= m[row * b + k] * y[k];
1814        }
1815        y[row] = s;
1816    }
1817    // Back substitution U x = y.  Write into rhs[].
1818    for row in (0..b).rev() {
1819        let mut s = y[row];
1820        for k in (row + 1)..b {
1821            s -= m[row * b + k] * rhs[k];
1822        }
1823        rhs[row] = s / m[row * b + row];
1824    }
1825}
1826
1827/// Compute only per-observation leverages tr(H_ii) for multi-predictor models.
1828///
1829/// This is cheaper than the full ALO correction when only EDF or leverage
1830/// diagnostics are needed (no scores or W⁻¹ computation required).
1831///
1832/// Returns an n-length array of leverages.  The total model EDF is the sum
1833/// of all leverages.
1834pub fn compute_multiblock_alo_leverages(
1835    n_obs: usize,
1836    n_blocks: usize,
1837    block_designs: &[Array2<f64>],
1838    penalized_hessian_inv: &Array2<f64>,
1839    block_weights: &[Array2<f64>],
1840) -> Result<Array1<f64>, EstimationError> {
1841    use rayon::prelude::*;
1842
1843    let n = n_obs;
1844    let b = n_blocks;
1845    let p_tot = penalized_hessian_inv.nrows();
1846
1847    let col_offsets = multiblock_col_offsets(block_designs);
1848    let max_workers = rayon::current_num_threads();
1849    let chunk_size = multiblock_alo_parallel_leverage_chunk_size(p_tot, b, n, max_workers);
1850
1851    let mut leverage = Array1::<f64>::zeros(n);
1852
1853    // Per-block H_inv stripe scratch (p_tot × p_blk) is read-only once built
1854    // and shared by the parallel chunks.  Only per-chunk q/XT/B×B scratch is
1855    // replicated across Rayon workers.
1856    let block_widths: Vec<usize> = block_designs.iter().map(|d| d.ncols()).collect();
1857    let mut h_stripes: Vec<FaerMat<f64>> = block_widths
1858        .iter()
1859        .map(|&p_blk| FaerMat::<f64>::zeros(p_tot, p_blk))
1860        .collect();
1861    // Populate the H_inv stripes once: each block reads a constant column slab
1862    // out of `penalized_hessian_inv` and copies it into a column-major faer Mat.
1863    for blk in 0..b {
1864        let off_b = col_offsets[blk];
1865        let p_blk = block_widths[blk];
1866        let stripe = &mut h_stripes[blk];
1867        for c in 0..p_blk {
1868            for r in 0..p_tot {
1869                stripe[(r, c)] = penalized_hessian_inv[(r, off_b + c)];
1870            }
1871        }
1872    }
1873
1874    leverage
1875        .as_slice_mut()
1876        .expect("newly allocated Array1 is contiguous")
1877        .par_chunks_mut(chunk_size)
1878        .enumerate()
1879        .for_each(|(chunk_idx, leverage_chunk)| {
1880            let chunk_start = chunk_idx * chunk_size;
1881            let chunk_len = leverage_chunk.len();
1882            let chunk_end = chunk_start + chunk_len;
1883
1884            // Chunk-local scratch: B×B flat row-major buffers for A_i, W_i
1885            // and AW = A·W.  Each worker writes only its `leverage_chunk`, so
1886            // output writes are disjoint and require no synchronization.
1887            let bb_sz = b * b;
1888            let mut a_i = vec![0.0f64; bb_sz];
1889            let mut aw = vec![0.0f64; bb_sz];
1890            let mut w_flat = vec![0.0f64; bb_sz];
1891
1892            // Column-major faer storage for q_blocks: q_k has shape
1893            // (p_tot, chunk_len) with contiguous columns, so
1894            // `col_as_slice(local_i)` is a direct stripe.
1895            let mut q_storage: Vec<FaerMat<f64>> = block_widths
1896                .iter()
1897                .map(|_| FaerMat::<f64>::zeros(p_tot, chunk_len))
1898                .collect();
1899
1900            // Per-block X^T scratch in column-major faer storage
1901            // (p_blk × chunk_len), owned by this chunk to keep the matmul input
1902            // contiguous without sharing mutable scratch across threads.
1903            let mut xt_storage: Vec<FaerMat<f64>> = block_widths
1904                .iter()
1905                .map(|&p_blk| FaerMat::<f64>::zeros(p_blk, chunk_len))
1906                .collect();
1907
1908            // Build q_blocks[blk] = H_inv[:, off..off+p_blk] · X_blk[chunk, :]^T
1909            // entirely in column-major faer storage so subsequent column reads
1910            // are contiguous f64 stripes — replaces the per-chunk `to_owned()`
1911            // ndarray slicing + row-major `dot()` from the original.
1912            for blk in 0..b {
1913                let p_blk = block_widths[blk];
1914
1915                let x_chunk = block_designs[blk].slice(s![chunk_start..chunk_end, ..]);
1916                let xt = &mut xt_storage[blk];
1917                for local_i in 0..chunk_len {
1918                    let row = x_chunk.row(local_i);
1919                    for j in 0..p_blk {
1920                        xt[(j, local_i)] = row[j];
1921                    }
1922                }
1923
1924                matmul(
1925                    q_storage[blk].as_mut(),
1926                    Accum::Replace,
1927                    h_stripes[blk].as_ref(),
1928                    xt_storage[blk].as_ref(),
1929                    1.0,
1930                    Par::Seq,
1931                );
1932            }
1933
1934            for local_i in 0..chunk_len {
1935                let i = chunk_start + local_i;
1936                let w_i = &block_weights[i];
1937
1938                // Flatten W_i once per observation (row-major).
1939                for r in 0..b {
1940                    for c in 0..b {
1941                        w_flat[r * b + c] = w_i[(r, c)];
1942                    }
1943                }
1944
1945                // Assemble A_i[a, k] = X_a[i, :] · q_k[off_a:off_a+p_a, local_i].
1946                // For each k, read its column once (contiguous f64 stripe), then
1947                // for each a take the matching offset slab.
1948                for r in 0..bb_sz {
1949                    a_i[r] = 0.0;
1950                }
1951                for k in 0..b {
1952                    let q_k = &q_storage[k];
1953                    let q_col = q_k.col_as_slice(local_i);
1954                    for a in 0..b {
1955                        let p_a = block_widths[a];
1956                        let off_a = col_offsets[a];
1957                        let xa_row = block_designs[a].row(i);
1958                        let mut dot = 0.0f64;
1959                        for j in 0..p_a {
1960                            dot = xa_row[j].mul_add(q_col[off_a + j], dot);
1961                        }
1962                        a_i[a * b + k] = dot;
1963                    }
1964                }
1965
1966                // AW = A_i · W_i (B×B), then leverage = trace(AW) = sum_{a,k} A[a,k]·W[k,a].
1967                mat_mul_flat(&a_i, &w_flat, &mut aw, b);
1968                let mut tr = 0.0f64;
1969                for d in 0..b {
1970                    tr += aw[d * b + d];
1971                }
1972                leverage_chunk[local_i] = tr;
1973            }
1974        });
1975
1976    Ok(leverage)
1977}
1978
1979// (Allocation-free, factor-once-reuse-many B×B LU helpers live next to the
1980// multi-block ALO callsite — see `lu_factor_in_place` and `lu_solve_in_place`.)
1981
1982#[cfg(test)]
1983mod tests {
1984    use super::{
1985        ALO_EXACT_SCALAR_MAX_ITERS, AloExactScalarError, AloInput, alo_eta_exact_frozen_curvature,
1986        alo_eta_updatewith_offset, bayesvar_eta, compute_alo_from_input_inner,
1987        percentile_from_sorted, percentile_index, sandwichvar_eta_from_meat,
1988    };
1989    use gam_linalg::matrix::{PsdWeightsView, SignedWeightsView};
1990    use gam_problem::LinkFunction;
1991
1992    #[test]
1993    fn alo_offset_update_matches_centered_algebra() {
1994        let eta_hat = 11.0;
1995        let z = 13.0;
1996        let offset = 10.0;
1997        let x_hinv_x = 0.2;
1998        let hessian_weight = 1.0;
1999        let score_weight = 1.0;
2000        // centered: eta~=off + ((eta-off)-a(z-off))/(1-a) when W_S = W_H.
2001        let leverage = hessian_weight * x_hinv_x;
2002        let expected = offset + ((eta_hat - offset) - leverage * (z - offset)) / (1.0 - leverage);
2003        let got =
2004            alo_eta_updatewith_offset(eta_hat, z, offset, x_hinv_x, score_weight, 1.0 - leverage);
2005        assert!((got - expected).abs() < 1e-12);
2006    }
2007
2008    #[test]
2009    fn alo_offset_update_reduces_to_classicwhen_offsetzero() {
2010        let eta_hat = 1.25;
2011        let z = -0.5;
2012        let x_hinv_x = 0.35;
2013        let hessian_weight = 1.0;
2014        let score_weight = 1.0;
2015        let leverage = hessian_weight * x_hinv_x;
2016        let expected = (eta_hat - leverage * z) / (1.0 - leverage);
2017        let got =
2018            alo_eta_updatewith_offset(eta_hat, z, 0.0, x_hinv_x, score_weight, 1.0 - leverage);
2019        assert!((got - expected).abs() < 1e-12);
2020    }
2021
2022    #[test]
2023    fn alo_offset_update_uses_distinct_score_and_hessian_weights() {
2024        let eta_hat = 1.7;
2025        let z = 0.4;
2026        let offset = -0.2;
2027        let x_hinv_x = 0.15;
2028        let hessian_weight = 3.0;
2029        let score_weight = 5.0;
2030        let expected = offset
2031            + (eta_hat - offset)
2032            + x_hinv_x * score_weight * ((eta_hat - offset) - (z - offset))
2033                / (1.0 - hessian_weight * x_hinv_x);
2034        let got = alo_eta_updatewith_offset(
2035            eta_hat,
2036            z,
2037            offset,
2038            x_hinv_x,
2039            score_weight,
2040            1.0 - hessian_weight * x_hinv_x,
2041        );
2042        assert!((got - expected).abs() < 1e-12);
2043    }
2044
2045    #[test]
2046    fn alo_offset_update_handles_zero_hessian_weight() {
2047        let eta_hat = 0.8;
2048        let z = -0.3;
2049        let offset = 0.1;
2050        let x_hinv_x = 0.4;
2051        let hessian_weight = 0.0;
2052        let score_weight = 2.5;
2053        let expected = offset
2054            + (eta_hat - offset)
2055            + x_hinv_x * score_weight * ((eta_hat - offset) - (z - offset));
2056        let got = alo_eta_updatewith_offset(
2057            eta_hat,
2058            z,
2059            offset,
2060            x_hinv_x,
2061            score_weight,
2062            1.0 - hessian_weight * x_hinv_x,
2063        );
2064        assert!((got - expected).abs() < 1e-12);
2065    }
2066
2067    #[test]
2068    fn alo_exact_frozen_curvature_converges_to_fixed_point() {
2069        let eta_hat = 1.0;
2070        let a_ii = 0.4;
2071        let got = alo_eta_exact_frozen_curvature(eta_hat, a_ii, &|eta| (0.5 * (eta - 2.0), 0.5))
2072            .expect("linear scalar fixed point should converge in one Newton step");
2073        assert!((got - 0.75).abs() < 1e-12);
2074    }
2075
2076    #[test]
2077    fn alo_exact_frozen_curvature_reports_nonconvergence() {
2078        let err = alo_eta_exact_frozen_curvature(0.0, 1.0, &|eta| (eta + 1.0, 0.0))
2079            .expect_err("constant residual should exhaust the scalar iteration budget");
2080        let AloExactScalarError::MaxIterations { iterations, .. } = err else {
2081            panic!("constant residual must report MaxIterations, got {err:?}");
2082        };
2083        assert_eq!(
2084            iterations, ALO_EXACT_SCALAR_MAX_ITERS,
2085            "non-convergence must report the full scalar iteration budget"
2086        );
2087    }
2088
2089    #[test]
2090    fn alo_input_reports_exact_scalar_nonconvergence_with_row_context() {
2091        let design = Array2::from_elem((1, 1), 1.0);
2092        let penalized_hessian = Array2::from_elem((1, 1), 1.0);
2093        let hessian_weights = Array1::from_vec(vec![0.0]);
2094        let score_weights = Array1::from_vec(vec![0.0]);
2095        let working_response = Array1::from_vec(vec![0.0]);
2096        let eta = Array1::from_vec(vec![0.0]);
2097        let offset = Array1::from_vec(vec![0.0]);
2098        let score_curvature = |_: usize, eta: f64| (eta + 1.0, 0.0);
2099        let input = AloInput {
2100            design: &design,
2101            penalized_hessian: &penalized_hessian,
2102            hessian_weights: SignedWeightsView::from_array(&hessian_weights),
2103            score_weights: PsdWeightsView::try_from_array(&score_weights).expect("psd weights"),
2104            working_response: &working_response,
2105            eta: &eta,
2106            offset: &offset,
2107            link: LinkFunction::Logit,
2108            phi: 1.0,
2109            penalty_root: None,
2110            ridge: 0.0,
2111            score_curvature: Some(&score_curvature),
2112        };
2113
2114        let err =
2115            compute_alo_from_input_inner(&input).expect_err("non-converged exact ALO must error");
2116        let msg = err.to_string();
2117        assert!(
2118            msg.contains("ALO exact frozen-curvature solve failed at row 0"),
2119            "missing row context in exact ALO error: {msg}"
2120        );
2121        assert!(
2122            msg.contains("did not converge within"),
2123            "missing non-convergence cause in exact ALO error: {msg}"
2124        );
2125    }
2126
2127    #[test]
2128    fn gaussian_unpenalized_direct_sandwich_equals_bayes() {
2129        // In a Gaussian linear model with H = X'WX, direct meat
2130        // x_i'H^{-1}X'WXH^{-1}x_i equals x_i'H^{-1}x_i.
2131        let phi = 2.5;
2132        let x_hinv_x = 0.3;
2133        let vb = bayesvar_eta(phi, x_hinv_x);
2134        let vs = sandwichvar_eta_from_meat(phi, x_hinv_x);
2135        assert!((vb - vs).abs() < 1e-12);
2136    }
2137
2138    #[test]
2139    fn sandwich_from_direct_meat_scales_by_phi() {
2140        let phi = 1.7;
2141        let meat_quad = 0.358;
2142        let got = sandwichvar_eta_from_meat(phi, meat_quad);
2143        let expected = phi * meat_quad;
2144        assert!((got - expected).abs() < 1e-12);
2145    }
2146
2147    #[test]
2148    fn sandwich_meat_uses_score_weights_not_hessian_weights_noncanonical() {
2149        // Regression for the sandwich-SE "meat" weight bug: the meat must be the
2150        // SCORE covariance Xᵀ diag(W_S) X (Fisher, PSD), NOT the observed-info
2151        // Hessian weight W_H (signed). This fixture mimics a non-canonical link
2152        // (W_H ≠ W_S) with mixed-sign observed curvature.
2153        //
2154        // Single column (p = 1) makes H a scalar, so the sandwich variance is
2155        // closed form: with H = Σ W_H·x² + s0 (> 0 after the penalty), the meat
2156        // for obs is x_obs²·H⁻²·Σ_row W_S·x_row², and se = sqrt(φ·meat).
2157        let x = Array2::from_shape_vec((5, 1), vec![1.0, 2.0, 1.0, 2.0, 1.0]).unwrap();
2158        // Mixed-sign observed-information weights; the negative rows carry the
2159        // larger design values so Σ W_H·x² is NEGATIVE (see assert below).
2160        let w_h_vec = Array1::from_vec(vec![1.0, -1.0, 1.0, -1.0, 0.5]);
2161        // Score/Fisher weights are strictly positive (PSD by construction).
2162        let w_s_vec = Array1::from_vec(vec![1.0, 0.8, 1.2, 0.6, 0.9]);
2163        let phi = 1.3;
2164
2165        let n = x.nrows();
2166        let sum_wh_x2: f64 = (0..n).map(|i| w_h_vec[i] * x[[i, 0]] * x[[i, 0]]).sum();
2167        let sum_ws_x2: f64 = (0..n).map(|i| w_s_vec[i] * x[[i, 0]] * x[[i, 0]]).sum();
2168        // The whole point: Σ W_H·x² < 0 < Σ W_S·x². With W_H the meat is negative
2169        // and the "materially negative sandwich variance" guard would trip
2170        // (spurious LooComputationFailed); with W_S it is a valid PSD meat.
2171        assert!(sum_wh_x2 < 0.0, "fixture must exercise a negative W_H meat");
2172        assert!(sum_ws_x2 > 0.0);
2173
2174        // Penalize enough that the penalized Hessian is PD despite Σ W_H·x² < 0.
2175        let s0 = 8.0_f64;
2176        let h = s0 + sum_wh_x2; // = 2.5
2177        assert!(h > 0.0, "penalized Hessian must stay PD");
2178        let penalized_hessian = Array2::from_elem((1, 1), h);
2179
2180        // Pre-fix arithmetic check: the OLD W_H meat would be materially negative
2181        // for the larger-x rows, so the old code returned LooComputationFailed.
2182        let old_meat_obs1 = x[[1, 0]] * x[[1, 0]] / (h * h) * sum_wh_x2;
2183        assert!(
2184            phi * old_meat_obs1 < -super::variance_negative_tolerance(phi * old_meat_obs1.abs()),
2185            "the pre-fix W_H meat must be materially negative (guard would trip)"
2186        );
2187
2188        let working_response = Array1::from_vec(vec![0.3, -0.2, 0.5, 0.1, -0.4]);
2189        let eta = Array1::from_vec(vec![0.2, 0.1, 0.4, -0.1, 0.05]);
2190        let offset = Array1::zeros(n);
2191        let input = AloInput {
2192            design: &x,
2193            penalized_hessian: &penalized_hessian,
2194            hessian_weights: SignedWeightsView::from_array(&w_h_vec),
2195            score_weights: PsdWeightsView::try_from_array(&w_s_vec).expect("psd weights"),
2196            working_response: &working_response,
2197            eta: &eta,
2198            offset: &offset,
2199            link: LinkFunction::Probit,
2200            phi,
2201            penalty_root: None,
2202            ridge: 0.0,
2203            score_curvature: None,
2204        };
2205
2206        // The fix must let this succeed (no spurious negative-meat failure)...
2207        let diag = compute_alo_from_input_inner(&input)
2208            .expect("fixed sandwich meat (W_S) must not trip the negative-variance guard");
2209
2210        // ...and match the closed-form W_S reference for every row.
2211        for obs in 0..n {
2212            let expected =
2213                (phi * x[[obs, 0]] * x[[obs, 0]] / (h * h) * sum_ws_x2).sqrt();
2214            assert!(
2215                (diag.se_sandwich[obs] - expected).abs() <= 1e-10 * expected.max(1.0),
2216                "row {obs}: se_sandwich={} expected={expected}",
2217                diag.se_sandwich[obs]
2218            );
2219        }
2220    }
2221
2222    #[test]
2223    fn percentile_index_matches_expected_rounding() {
2224        assert_eq!(percentile_index(0, 0.95), 0);
2225        assert_eq!(percentile_index(1, 0.95), 0);
2226        assert_eq!(percentile_index(10, 0.50), 5);
2227        assert_eq!(percentile_index(10, 0.95), 9);
2228    }
2229
2230    #[test]
2231    fn percentile_from_sorted_returns_order_statistic() {
2232        let values = [1.0, 2.0, 3.0, 4.0, 5.0];
2233        assert_eq!(percentile_from_sorted(&values, 0.50), 3.0);
2234        assert_eq!(percentile_from_sorted(&values, 0.95), 5.0);
2235        assert_eq!(percentile_from_sorted(&[], 0.95), 0.0);
2236    }
2237
2238    // --- Multi-block ALO tests ---
2239
2240    use super::{MultiBlockAloInput, compute_multiblock_alo, compute_multiblock_alo_leverages};
2241    use ndarray::{Array1, Array2};
2242
2243    #[test]
2244    fn multiblock_b1_matches_scalar_leverage() {
2245        // With B=1 the multi-block formula should reduce to the scalar case.
2246        // H_ii = x_i^T H^{-1} x_i * w_i  (scalar).
2247        let n = 3;
2248        let p = 2;
2249        let x = Array2::from_shape_vec((n, p), vec![1.0, 0.5, 0.8, -0.3, 0.2, 1.1]).unwrap();
2250        // H = X'WX + I (simple regularisation).
2251        let w = [1.0, 2.0, 0.5];
2252        let mut h = Array2::<f64>::eye(p);
2253        for i in 0..n {
2254            for r in 0..p {
2255                for c in 0..p {
2256                    h[(r, c)] += w[i] * x[(i, r)] * x[(i, c)];
2257                }
2258            }
2259        }
2260        // Invert H (2x2).
2261        let det = h[(0, 0)] * h[(1, 1)] - h[(0, 1)] * h[(1, 0)];
2262        let mut h_inv = Array2::<f64>::zeros((p, p));
2263        h_inv[(0, 0)] = h[(1, 1)] / det;
2264        h_inv[(1, 1)] = h[(0, 0)] / det;
2265        h_inv[(0, 1)] = -h[(0, 1)] / det;
2266        h_inv[(1, 0)] = -h[(1, 0)] / det;
2267
2268        // Scalar leverages: a_ii = w_i * x_i^T H^{-1} x_i
2269        let mut scalar_lev = vec![0.0f64; n];
2270        for i in 0..n {
2271            let mut xhx = 0.0;
2272            for r in 0..p {
2273                for c in 0..p {
2274                    xhx += x[(i, r)] * h_inv[(r, c)] * x[(i, c)];
2275                }
2276            }
2277            scalar_lev[i] = w[i] * xhx;
2278        }
2279
2280        // Multi-block with B=1.
2281        let block_designs = vec![x.clone()];
2282        let block_weights: Vec<Array2<f64>> =
2283            w.iter().map(|&wi| Array2::from_elem((1, 1), wi)).collect();
2284        let scores: Vec<Array1<f64>> = (0..n).map(|_| Array1::from_vec(vec![0.1])).collect();
2285        let eta_hat: Vec<Array1<f64>> = (0..n).map(|i| Array1::from_vec(vec![i as f64])).collect();
2286
2287        let input = MultiBlockAloInput {
2288            n_obs: n,
2289            n_blocks: 1,
2290            block_designs: &block_designs,
2291            penalized_hessian_inv: &h_inv,
2292            block_weights,
2293            scores,
2294            eta_hat,
2295        };
2296
2297        let result = compute_multiblock_alo(&input).unwrap();
2298        for i in 0..n {
2299            assert!(
2300                (result.leverage[i] - scalar_lev[i]).abs() < 1e-10,
2301                "leverage mismatch at i={}: got {}, expected {}",
2302                i,
2303                result.leverage[i],
2304                scalar_lev[i]
2305            );
2306        }
2307    }
2308
2309    #[test]
2310    fn multiblock_leverage_only_matches_full() {
2311        // Verify that compute_multiblock_alo_leverages returns the same
2312        // leverages as compute_multiblock_alo.
2313        let n = 4;
2314        let p1 = 2;
2315        let p2 = 3;
2316        let x1 = Array2::from_shape_fn((n, p1), |(i, j)| (i + j + 1) as f64 * 0.3);
2317        let x2 = Array2::from_shape_fn((n, p2), |(i, j)| (i * 2 + j) as f64 * 0.2 - 0.1);
2318        let p_tot = p1 + p2;
2319        let h_inv = Array2::<f64>::eye(p_tot); // Simple identity for test.
2320        let block_weights: Vec<Array2<f64>> = (0..n)
2321            .map(|i| {
2322                let v = (i + 1) as f64;
2323                Array2::from_shape_vec((2, 2), vec![v, 0.1, 0.1, v * 0.5]).unwrap()
2324            })
2325            .collect();
2326        let scores: Vec<Array1<f64>> = (0..n).map(|_| Array1::from_vec(vec![0.0, 0.0])).collect();
2327        let eta_hat: Vec<Array1<f64>> = (0..n).map(|_| Array1::from_vec(vec![0.0, 0.0])).collect();
2328        let block_designs = vec![x1.clone(), x2.clone()];
2329
2330        let input = MultiBlockAloInput {
2331            n_obs: n,
2332            n_blocks: 2,
2333            block_designs: &block_designs,
2334            penalized_hessian_inv: &h_inv,
2335            block_weights: block_weights.clone(),
2336            scores,
2337            eta_hat,
2338        };
2339        let full = compute_multiblock_alo(&input).unwrap();
2340        let lev_only =
2341            compute_multiblock_alo_leverages(n, 2, &block_designs, &h_inv, &block_weights).unwrap();
2342
2343        for i in 0..n {
2344            assert!(
2345                (full.leverage[i] - lev_only[i]).abs() < 1e-12,
2346                "leverage mismatch at i={}: full={}, lev_only={}",
2347                i,
2348                full.leverage[i],
2349                lev_only[i]
2350            );
2351        }
2352    }
2353
2354    #[test]
2355    fn multiblock_singular_weight_still_corrects() {
2356        // When W_i = 0 (singular), the W_i⁻¹-free formula still works:
2357        // (I - W_i A_i)⁻¹ = I, so Δη = A_i s_i.
2358        // A_i = x H⁻¹ xᵀ = 1.0² + 0.5² = 1.25 (scalar, B=1).
2359        let n = 1;
2360        let p = 2;
2361        let x = Array2::from_shape_vec((1, p), vec![1.0, 0.5]).unwrap();
2362        let h_inv = Array2::eye(p);
2363        let block_designs = vec![x.clone()];
2364        let block_weights = vec![Array2::from_elem((1, 1), 0.0)]; // singular
2365        let scores = vec![Array1::from_vec(vec![1.0])];
2366        let eta_hat = vec![Array1::from_vec(vec![std::f64::consts::PI])];
2367
2368        let input = MultiBlockAloInput {
2369            n_obs: n,
2370            n_blocks: 1,
2371            block_designs: &block_designs,
2372            penalized_hessian_inv: &h_inv,
2373            block_weights,
2374            scores,
2375            eta_hat,
2376        };
2377        let result = compute_multiblock_alo(&input).unwrap();
2378        // Δη = A_i * s_i = 1.25 * 1.0 = 1.25
2379        let expected = std::f64::consts::PI + 1.25;
2380        assert!(
2381            (result.eta_tilde[0][0] - expected).abs() < 1e-12,
2382            "expected {}, got {}",
2383            expected,
2384            result.eta_tilde[0][0]
2385        );
2386        // Cook's distance should be 0 since W_i = 0.
2387        assert!(result.cook_distance[0].abs() < 1e-14);
2388        // ALO variance should be 0 since W_i = 0.
2389        assert!(result.alo_variance[0][0].abs() < 1e-14);
2390    }
2391
2392    #[test]
2393    fn multiblock_cook_and_variance_basic() {
2394        // B=1 with known values: verify Cook's distance and variance.
2395        let n = 1;
2396        let x = Array2::from_elem((1, 1), 1.0);
2397        // H⁻¹ = [[0.5]]
2398        let h_inv = Array2::from_elem((1, 1), 0.5);
2399        let block_designs = vec![x.clone()];
2400        let w_val = 2.0;
2401        let s_val = 0.4;
2402        let block_weights = vec![Array2::from_elem((1, 1), w_val)];
2403        let scores = vec![Array1::from_vec(vec![s_val])];
2404        let eta_hat = vec![Array1::from_vec(vec![1.0])];
2405
2406        let input = MultiBlockAloInput {
2407            n_obs: n,
2408            n_blocks: 1,
2409            block_designs: &block_designs,
2410            penalized_hessian_inv: &h_inv,
2411            block_weights,
2412            scores,
2413            eta_hat,
2414        };
2415        let result = compute_multiblock_alo(&input).unwrap();
2416
2417        // A_i = x H⁻¹ xᵀ = 1 * 0.5 * 1 = 0.5
2418        // (I - W A)⁻¹ = 1 / (1 - 2.0 * 0.5) = 1/0 => regularised
2419        // Actually 1 - w*a = 1 - 1.0 = 0.0, so det < 1e-12 => regularised with eps=1e-6
2420        // (I - W A + eps) = 1e-6, so v = s / 1e-6 = 4e5
2421        // delta_eta = A * v = 0.5 * 4e5 = 2e5
2422        // This is the regularised case; just check it doesn't panic and returns finite values.
2423        assert!(result.eta_tilde[0][0].is_finite());
2424        assert!(result.cook_distance[0].is_finite());
2425        assert!(result.alo_variance[0][0].is_finite());
2426    }
2427}