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