Skip to main content

gam_models/
penalized_vector_glm.rs

1//! Generic penalized vector-response GLM Newton solver (fixed λ).
2//!
3//! This is the shared scaffold extracted from
4//! [`crate::multinomial::fit_penalized_multinomial`] (dense softmax
5//! Fisher block) and
6//! [`crate::binomial_multi::fit_penalized_binomial_multi`]
7//! (row-diagonal independent-binomial Fisher block). Both families fit a
8//! penalized vector-response GLM with a shared design `X ∈ ℝ^{N×P}` and a
9//! shared penalty `S ∈ ℝ^{P×P}` replicated per output, differing **only** in
10//! the per-row Fisher-block algebra and the likelihood/residual. Everything
11//! else — input validation, penalized objective / gradient / Hessian assembly,
12//! damped Newton with backtracking, convergence certification, and the final
13//! penalized-objective / deviance tally — is written once here.
14//!
15//! # Fit problem
16//!
17//! With `β = [β_0; β_1; …; β_{M-1}]` stacked in output-major order
18//! (`β_a ∈ ℝ^P` is the coefficient block for output `a`), minimise the
19//! penalized negative log-likelihood
20//!
21//! ```text
22//!   F(β) = − log L(β) + ½ Σ_{a=0}^{M-1} λ_a · β_aᵀ S β_a
23//! ```
24//!
25//! where `log L` and its η-derivatives are supplied by the family's
26//! [`VectorLikelihood`] adapter and `λ_a` is a per-output smoothing parameter
27//! scaling the shared penalty `S`. The active linear predictor is
28//! `η_{n,a} = (X β_a)_n`, shape `(N, M)`.
29//!
30//! # Newton step
31//!
32//! Each iteration assembles the coupled penalized Hessian and gradient in
33//! output-major coefficient ordering `flat[a·P + i] = β[i, a]` (matching
34//! [`gam_solve::pirls::dense_block_xtwx`]):
35//!
36//! ```text
37//!   H[a·P + i, b·P + j] = Σ_n W_{n,a,b} · X[n,i] · X[n,j]   (+ δ_{ab} λ_a S[i,j])
38//!   g[a·P + i]          = Σ_n r_{n,a} · X[n,i]              (+ λ_a (S β_a)[i])
39//! ```
40//!
41//! with the per-row Fisher block `W_{n,·,·} = −∂² log L / ∂η ∂η` (the family's
42//! [`VectorLikelihood::hess_block`], or a caller override) and the residual
43//! `r_{n,a} = −∂ log L / ∂η_a` (`−`[`VectorLikelihood::grad_eta`]). The step
44//! `δ = − H^{-1} g` is solved through faer's symmetric-PD-with-fallback
45//! factorisation under an adaptive Levenberg–Marquardt ridge: when a
46//! rank-deficient block (collinear / quasi-separated columns under a small
47//! per-output λ) makes the Bunch–Kaufman fallback back-substitute through
48//! near-zero pivots into a non-finite δ, a diagonal ridge `τ·I` — scaled by the
49//! Hessian's largest diagonal so it is curvature-scale invariant — is added and
50//! the system re-solved, escalating τ geometrically until δ is finite. The
51//! step is then accepted by a backtracking line search on `F` (full step first,
52//! halve up to 8 times). Because the line search validates against the
53//! *unridged* objective `F`, the ridge never biases the converged β̂ (at the
54//! optimum the gradient vanishes and δ → 0 for any τ). Convergence requires
55//! both the relative coefficient step `‖δ‖ / (1 + ‖β‖) ≤ tol` and an exact
56//! curvature-scaled first-order score certificate recomputed at the accepted
57//! final iterate.
58//!
59//! # Fisher-block override
60//!
61//! When `fisher_w_override` is `Some`, each Newton step uses the supplied
62//! per-row `(N, M, M)` curvature block in place of the analytic
63//! [`VectorLikelihood::hess_block`]; the gradient/residual path stays analytic
64//! (issue #349). The two families differ in what they accept off the diagonal:
65//! multinomial admits a full dense block, while independent-binomial columns
66//! only consume the per-output diagonal (a non-zero cross term cannot be
67//! represented by the separable columns). That family-specific precondition is
68//! enforced by the adapter before it constructs the override view; the engine
69//! consumes whatever block it is given.
70
71use crate::model_types::EstimationError;
72use crate::vector_response::VectorLikelihood;
73use faer::Side;
74use gam_linalg::faer_ndarray::{FaerArrayView, array2_to_matmut, factorize_symmetricwith_fallback};
75use gam_problem::{
76    FixedLambdaCheckpoint, FixedLambdaResidualKind, FixedLambdaSolverStage, FixedLambdaStallReason,
77    FixedLambdaStationarityEvidence,
78};
79use gam_solve::pirls::dense_block_xtwx;
80use ndarray::{Array1, Array2, ArrayView1, ArrayView2, ArrayView3};
81use opt::{BacktrackConfig, RidgeSchedule, backtracking_line_search, escalate_ridge};
82
83/// Base Levenberg–Marquardt ridge as a fraction of the penalized Hessian's
84/// largest diagonal entry (so it is invariant to the problem's overall
85/// curvature scale). At ~1e-10 of the dominant curvature it is negligible
86/// relative to identified-direction curvature — it never biases the identified
87/// optimum (at β̂ the unridged gradient still vanishes there) — yet large
88/// enough to lift an exactly rank-deficient null direction off zero so the
89/// Bunch–Kaufman fallback yields a finite, descent Newton step (gam#856).
90const BASE_RIDGE_FRACTION_OF_MAX_DIAG: f64 = 1.0e-10;
91
92/// Geometric ridge-escalation budget for a single Newton step. 30 doublings
93/// span ~9 orders of magnitude over the base ridge, which covers any
94/// conditioning a finite-curvature softmax/binomial block can present.
95const MAX_RIDGE_ESCALATIONS: usize = 30;
96
97/// Backtracking budget for the damped-Newton line search: full step first, then
98/// halve up to this many times if the penalized objective fails to decrease.
99const MAX_BACKTRACKS: usize = 8;
100
101/// Per-step line-search contraction factor (halving).
102const LINE_SEARCH_SHRINK: f64 = 0.5;
103
104/// Slack on the "objective decreased" acceptance test, absorbing floating-point
105/// round-off so a step that is flat to machine precision is not rejected.
106const OBJECTIVE_DECREASE_SLACK: f64 = 1.0e-12;
107
108/// First-order optimality gate (gam#856) as a fraction of `1 + max_diag`: the
109/// unridged penalized gradient norm must fall below this curvature-scaled
110/// threshold before convergence is declared, certifying stationarity on the
111/// identified subspace rather than a premature step-norm stall.
112const OPTIMALITY_GRAD_FRACTION: f64 = 1.0e-6;
113
114/// Class-space metric of the replicated smoothing penalty (#1587).
115///
116/// * `Diagonal` — the historical `diag_a(λ_a) ⊗ S`: each active output's
117///   coefficient block is penalised independently. Correct for genuinely
118///   independent outputs (independent-binomial columns), but for a *softmax*
119///   multinomial it penalises the reference-anchored log-odds contrasts
120///   `η_a = log(p_a/p_ref)`, so the fit is NOT invariant to the arbitrary
121///   reference-class choice (#1587).
122/// * `Centered` — the reference-symmetric `λ · ((I_{M} − J_{M}/K) ⊗ S)` with a
123///   single shared `λ` (= `lambdas[0]`; the caller must pass uniform `lambdas`)
124///   and `K = M + 1`. This is exactly the symmetric CLR penalty
125///   `Σ_{k=0}^{K-1} β̃_kᵀ S β̃_k` (with `Σ_k β̃_k = 0`) written in the active-class
126///   (ALR) gauge — invariant to which class is the baseline (the multinomial
127///   analogue of #1549's `G^{1/2}` Aitchison whitening). Couples the class
128///   blocks via the `−(λ/K)·S` off-diagonals; the engine already factors a
129///   class-coupled Hessian (the softmax Fisher block is dense), so this is a
130///   penalty-assembly change only.
131#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
132pub enum ClassPenaltyMetric {
133    /// Independent per-output penalty `diag_a(λ_a) ⊗ S` (historical default).
134    #[default]
135    Diagonal,
136    /// Reference-symmetric centered penalty `λ·((I − J/K) ⊗ S)`, `K = M + 1`.
137    Centered,
138    /// Permutation-equivariant heterogeneous per-CLASS penalty (#2344, the
139    /// fixed-λ twin of the REML equivariant carrier from `1326d0794`): the
140    /// caller supplies `K = M + 1` lambdas — one per class, REFERENCE CLASS
141    /// INCLUDED — and the quadratic is `Σ_c λ_c · γ_cᵀ S γ_c` on the CENTERED
142    /// class functions `γ_c = β_c − β̄` (with `β_ref ≡ 0`, `β̄ = Σ_b β_b / K`).
143    /// In the active-class (ALR) gauge that is `A(λ) ⊗ S` with the M×M class
144    /// metric `A[a,b] = λ_a δ_ab − (λ_a + λ_b)/K + (Σ_c λ_c)/K²` — closed
145    /// under class relabeling (the (γ_c, λ_c) pairs permute together), unlike
146    /// `Diagonal`'s ALR-anchored family, and collapsing exactly to `Centered`
147    /// when every λ_c is equal. The engine reads `M = lambdas.len() − 1` for
148    /// this variant.
149    EquivariantPerClass,
150}
151
152impl ClassPenaltyMetric {
153    /// Number of ACTIVE outputs `M` implied by a lambda vector under this
154    /// metric: `Diagonal`/`Centered` carry one λ per active output;
155    /// `EquivariantPerClass` carries one λ per CLASS (`K = M + 1`, reference
156    /// included).
157    pub fn active_outputs(self, lambdas_len: usize) -> usize {
158        match self {
159            ClassPenaltyMetric::Diagonal | ClassPenaltyMetric::Centered => lambdas_len,
160            ClassPenaltyMetric::EquivariantPerClass => lambdas_len.saturating_sub(1),
161        }
162    }
163}
164
165/// The M×M equivariant class metric `A[a,b] = Σ_c λ_c·(δ_ca − 1/K)(δ_cb − 1/K)`
166/// over the active (ALR) coordinates, `c` ranging over ALL `K = M + 1` classes
167/// (the reference contributes through its centering row `−𝟙/K`). Expanded:
168/// `A[a,b] = λ_a δ_ab − (λ_a + λ_b)/K + (Σ_c λ_c)/K²`. PSD by construction
169/// (a nonnegative sum of rank-1 outer products).
170pub(crate) fn equivariant_class_metric(lambdas: ArrayView1<'_, f64>, m: usize) -> Array2<f64> {
171    let k = (m + 1) as f64;
172    let total: f64 = lambdas.iter().sum();
173    let mut a_mat = Array2::<f64>::zeros((m, m));
174    for a in 0..m {
175        for b in 0..m {
176            let mut value = -(lambdas[a] + lambdas[b]) / k + total / (k * k);
177            if a == b {
178                value += lambdas[a];
179            }
180            a_mat[[a, b]] = value;
181        }
182    }
183    a_mat
184}
185
186/// Inputs to [`fit_penalized_vector_glm`].
187///
188/// `M` (the number of active outputs / linear-predictor columns) is derived
189/// from `lambdas.len()` under the selected [`ClassPenaltyMetric`]
190/// (`Diagonal`/`Centered`: `M = lambdas.len()`; `EquivariantPerClass`:
191/// `M = lambdas.len() − 1`, one λ per CLASS, reference included); the engine
192/// validates it against the design and override shapes. The response `y` is passed verbatim to the [`VectorLikelihood`]
193/// adapter, which owns its own `(N, ·)` shape contract (binomial columns use
194/// `K = M`; multinomial one-hot uses `K = M + 1`), so the engine does not
195/// constrain its column count beyond `y.nrows() == N`.
196pub struct PenalizedVectorGlmInputs<'a> {
197    /// Design matrix `X ∈ ℝ^{N×P}` (one row per observation, shared across
198    /// every output column).
199    pub design: ArrayView2<'a, f64>,
200    /// Response `Y ∈ ℝ^{N×·}`, interpreted by the [`VectorLikelihood`].
201    pub y: ArrayView2<'a, f64>,
202    /// Shared smoothing penalty `S ∈ ℝ^{P×P}` (symmetric, PSD).
203    pub penalty: ArrayView2<'a, f64>,
204    /// Per-output smoothing parameter `λ_a`, length `M`.
205    pub lambdas: ArrayView1<'a, f64>,
206    /// Optional per-row Fisher-block override, shape `(N, M, M)`. When `Some`,
207    /// it replaces the analytic [`VectorLikelihood::hess_block`] as the Newton
208    /// curvature; the gradient/residual path stays analytic (issue #349). The
209    /// adapter is responsible for any family-specific structural precondition
210    /// on the block (e.g. zero off-diagonals for independent columns).
211    pub fisher_w_override: Option<ArrayView3<'a, f64>>,
212    /// Number of Newton iterations available to this invocation. On resume,
213    /// this is an additional budget beyond the checkpoint's completed count.
214    pub max_iter: usize,
215    /// Relative-step convergence tolerance.
216    pub tol: f64,
217    /// Class-space metric of the replicated penalty (#1587). `Diagonal`
218    /// preserves the historical independent-per-output penalty; `Centered`
219    /// selects the reference-symmetric softmax penalty (requires uniform
220    /// `lambdas`). See [`ClassPenaltyMetric`].
221    pub class_penalty_metric: ClassPenaltyMetric,
222    /// Optional checkpoint from the SAME design/response/penalty/weight
223    /// problem. Coefficients are sufficient to resume because η, the score,
224    /// Hessian, and objective are deterministically rebuilt before the first
225    /// additional Newton step.
226    pub resume_from: Option<VectorGlmResume<'a>>,
227}
228
229/// Borrowed fixed-λ vector-GLM checkpoint used to continue a stalled solve.
230#[derive(Debug, Clone, Copy)]
231pub struct VectorGlmResume<'a> {
232    pub coefficients: ArrayView2<'a, f64>,
233    pub completed_iterations: usize,
234}
235
236/// Outputs of a CONVERGED [`fit_penalized_vector_glm`] solve.
237///
238/// SPEC: a fit object only ever comes from a converged optimization. This
239/// struct is constructed exclusively on the [`VectorGlmSolve::Converged`] arm,
240/// so every consumer holding one holds a certified stationary point; there is
241/// no `converged` flag to check. A budget-exhausted solve surfaces instead as
242/// [`VectorGlmSolve::Stalled`], which carries the abandoned iterate as
243/// checkpoint evidence but deliberately has NO Laplace covariance — posterior
244/// uncertainty evaluated at a non-stationary iterate is not a posterior.
245pub struct PenalizedVectorGlmOutputs {
246    /// Coefficient matrix, shape `(P, M)` (column `a` is `β_a`).
247    pub coefficients: Array2<f64>,
248    /// Final active linear predictor `η = X β̂`, shape `(N, M)`. The adapter
249    /// turns this into fitted probabilities via its own inverse link.
250    pub eta: Array2<f64>,
251    /// Number of Newton iterations executed (including the final step that
252    /// satisfied the tolerance).
253    pub iterations: usize,
254    /// Unpenalized log-likelihood `log L(β̂)`.
255    pub log_likelihood: f64,
256    /// Penalty term `½ Σ_a λ_a · β̂_aᵀ S β̂_a` at the returned `β̂`.
257    pub penalty_term: f64,
258    /// Joint Laplace posterior coefficient covariance `H⁻¹` at the converged
259    /// `β̂`, shape `(P·M)×(P·M)` (#1101). `H = block(XᵀWX) + diag_a(λ_a)⊗S` is
260    /// the penalized Hessian the Newton loop already assembles and factors at
261    /// every step, discarding the factor; here it is re-assembled once at the
262    /// mode and inverted (solve against the identity through the same symmetric
263    /// factorization used for the Newton step). Block-ordered to match the
264    /// stacked coefficient vector `θ[a·P + i] = β̂[i, a]`, i.e.
265    /// `β = [β_0; …; β_{M-1}]`. This is the covariance the predict / inference
266    /// surface uses for posterior-mean probabilities and prediction intervals.
267    pub coefficient_covariance: Array2<f64>,
268}
269
270/// Checkpoint evidence for a Newton solve that stopped without certification.
271///
272/// This is NOT a fit: it exists so family adapters can inspect the abandoned
273/// iterate (e.g. the multinomial separation fingerprint `|η| ≥ 25` that routes
274/// to the Firth/Jeffreys proper-prior refit) and so the typed non-convergence
275/// error can carry honest evidence — the iteration count and the penalized
276/// objective at the last iterate. It carries no covariance and no fitted
277/// probabilities on purpose: nothing downstream may dress it up as a result.
278pub struct VectorGlmStall {
279    /// Why the convergence certificate was not reached.
280    pub reason: VectorGlmStallReason,
281    /// Coefficient checkpoint at the last accepted iterate, shape `(P, M)`.
282    pub coefficients: Array2<f64>,
283    /// Linear predictor `η = X β` at the abandoned iterate, shape `(N, M)`.
284    pub eta: Array2<f64>,
285    /// Newton iterations executed before the stall was diagnosed.
286    pub iterations: usize,
287    /// Unpenalized log-likelihood at the abandoned iterate.
288    pub log_likelihood: f64,
289    /// Penalty term at the abandoned iterate.
290    pub penalty_term: f64,
291    /// Norm of the exact penalized score at the checkpoint.
292    pub gradient_norm: f64,
293    /// Curvature-scaled score bound required by the stationarity certificate.
294    pub gradient_bound: f64,
295}
296
297impl VectorGlmStall {
298    /// Convert this solver checkpoint into the canonical typed fixed-lambda
299    /// non-convergence error. Family adapters supply only the objective stage
300    /// and a human-readable entry-point name; the evidence and resumable
301    /// coefficient state come from the solver that produced the stall.
302    pub fn into_nonconvergence_error(
303        self,
304        stage: FixedLambdaSolverStage,
305        context: impl Into<String>,
306    ) -> Result<EstimationError, EstimationError> {
307        let rows = self.coefficients.nrows();
308        let cols = self.coefficients.ncols();
309        let checkpoint = FixedLambdaCheckpoint::new(
310            stage,
311            self.coefficients.iter().copied().collect(),
312            rows,
313            cols,
314            self.iterations,
315        )
316        .map_err(|reason| {
317            EstimationError::InvalidInput(format!(
318                "fixed-lambda vector-GLM produced an invalid internal checkpoint: {reason}"
319            ))
320        })?;
321        let reason = match self.reason {
322            VectorGlmStallReason::IterationBudgetExhausted => {
323                FixedLambdaStallReason::IterationBudgetExhausted
324            }
325            VectorGlmStallReason::LineSearchExhausted => {
326                FixedLambdaStallReason::LineSearchExhausted
327            }
328            VectorGlmStallReason::PostStepCertificateFailed => {
329                FixedLambdaStallReason::StationarityCertificateFailed
330            }
331        };
332        Ok(EstimationError::FixedLambdaNewtonDidNotConverge {
333            context: context.into(),
334            reason,
335            objective_value: -self.log_likelihood + self.penalty_term,
336            stationarity: FixedLambdaStationarityEvidence {
337                kind: FixedLambdaResidualKind::PenalizedGradientNorm,
338                residual: self.gradient_norm,
339                bound: self.gradient_bound,
340            },
341            checkpoint,
342        })
343    }
344}
345
346/// Exhaustive reason a fixed-λ vector solve produced checkpoint evidence
347/// instead of a converged result.
348#[derive(Debug, Clone, Copy, PartialEq, Eq)]
349pub enum VectorGlmStallReason {
350    /// The caller's iteration budget ended before both certificates passed.
351    IterationBudgetExhausted,
352    /// No backtracked candidate satisfied the objective-descent certificate.
353    LineSearchExhausted,
354    /// The small-step gate passed, but the exact score at the accepted iterate
355    /// exceeded its curvature-scaled stationarity bound.
356    PostStepCertificateFailed,
357}
358
359/// Two-outcome result of the fixed-λ vector-GLM Newton solve. Hard input /
360/// linear-algebra failures remain `Err`; any terminal state without a
361/// stationarity certificate is a first-class `Stalled` outcome so adapters must
362/// decide explicitly (typed error, or the multinomial separation → Firth
363/// escalation) instead of ever forwarding a non-converged iterate as a fit.
364pub enum VectorGlmSolve {
365    /// Certified stationary point (step-norm AND first-order optimality gates
366    /// passed), with the Laplace covariance computed at the mode.
367    Converged(PenalizedVectorGlmOutputs),
368    /// Solver stopped without a convergence certificate.
369    Stalled(VectorGlmStall),
370}
371
372/// Add `A(λ) ⊗ S` — the equivariant per-class metric's coupled blocks
373/// (#2344, see [`equivariant_class_metric`]) — onto the penalized Hessian.
374/// Shared by the in-loop and final-iterate Hessian assemblies so both see the
375/// identical algebra.
376fn add_equivariant_penalty_blocks(
377    hessian: &mut Array2<f64>,
378    penalty: ArrayView2<'_, f64>,
379    lambdas: ArrayView1<'_, f64>,
380    p: usize,
381    m: usize,
382) {
383    if m == 0 {
384        return;
385    }
386    let a_mat = equivariant_class_metric(lambdas, m);
387    for a in 0..m {
388        for b in 0..m {
389            let coef = a_mat[[a, b]];
390            if coef == 0.0 {
391                continue;
392            }
393            let (ba, bb) = (a * p, b * p);
394            for i in 0..p {
395                for j in 0..p {
396                    hessian[[ba + i, bb + j]] += coef * penalty[[i, j]];
397                }
398            }
399        }
400    }
401}
402
403/// Quadratic form `½ β_aᵀ S β_a` accumulated across outputs with per-output
404/// weight `λ_a`. Shared by the objective evaluator and the final tally.
405fn weighted_penalty_sum(
406    beta: &Array2<f64>,
407    penalty: ArrayView2<'_, f64>,
408    lambdas: ArrayView1<'_, f64>,
409    metric: ClassPenaltyMetric,
410) -> f64 {
411    let (p, m) = beta.dim();
412    match metric {
413        ClassPenaltyMetric::Diagonal => {
414            let mut pen = 0.0_f64;
415            for a in 0..m {
416                let la = lambdas[a];
417                if la == 0.0 {
418                    continue;
419                }
420                let beta_col = beta.column(a);
421                let mut quad = 0.0_f64;
422                for i in 0..p {
423                    let mut s_beta_i = 0.0_f64;
424                    for j in 0..p {
425                        s_beta_i += penalty[[i, j]] * beta_col[j];
426                    }
427                    quad += beta_col[i] * s_beta_i;
428                }
429                pen += 0.5 * la * quad;
430            }
431            pen
432        }
433        // Centered (#1587): ½·λ·[ Σ_a β_aᵀSβ_a − (1/K)·gᵀSg ], g = Σ_a β_a,
434        // K = M + 1. Equals the symmetric CLR penalty Σ_k β̃_kᵀSβ̃_k (Σβ̃=0) in
435        // the active-class gauge — reference-invariant. Shared λ = lambdas[0].
436        ClassPenaltyMetric::Centered => {
437            if m == 0 {
438                return 0.0;
439            }
440            let lam = lambdas[0];
441            if lam == 0.0 {
442                return 0.0;
443            }
444            let k = (m + 1) as f64;
445            // g = Σ_a β_a (the active-class coefficient sum, a p-vector).
446            let mut g = vec![0.0_f64; p];
447            for a in 0..m {
448                let col = beta.column(a);
449                for i in 0..p {
450                    g[i] += col[i];
451                }
452            }
453            // Σ_a β_aᵀSβ_a.
454            let mut sum_quad = 0.0_f64;
455            for a in 0..m {
456                let col = beta.column(a);
457                for i in 0..p {
458                    let mut s_beta_i = 0.0_f64;
459                    for j in 0..p {
460                        s_beta_i += penalty[[i, j]] * col[j];
461                    }
462                    sum_quad += col[i] * s_beta_i;
463                }
464            }
465            // gᵀSg.
466            let mut g_quad = 0.0_f64;
467            for i in 0..p {
468                let mut s_g_i = 0.0_f64;
469                for j in 0..p {
470                    s_g_i += penalty[[i, j]] * g[j];
471                }
472                g_quad += g[i] * s_g_i;
473            }
474            0.5 * lam * (sum_quad - g_quad / k)
475        }
476        // EquivariantPerClass (#2344): ½·Σ_{a,b} A[a,b]·β_aᵀSβ_b with the
477        // heterogeneous per-class metric A(λ) (see
478        // [`equivariant_class_metric`]); equal λ collapses to `Centered`.
479        ClassPenaltyMetric::EquivariantPerClass => {
480            if m == 0 {
481                return 0.0;
482            }
483            let a_mat = equivariant_class_metric(lambdas, m);
484            let mut s_beta = Array2::<f64>::zeros((p, m));
485            for b in 0..m {
486                let col = beta.column(b);
487                for i in 0..p {
488                    let mut acc = 0.0_f64;
489                    for j in 0..p {
490                        acc += penalty[[i, j]] * col[j];
491                    }
492                    s_beta[[i, b]] = acc;
493                }
494            }
495            let mut pen = 0.0_f64;
496            for a in 0..m {
497                let col = beta.column(a);
498                for b in 0..m {
499                    let coef = a_mat[[a, b]];
500                    if coef == 0.0 {
501                        continue;
502                    }
503                    let mut cross = 0.0_f64;
504                    for i in 0..p {
505                        cross += col[i] * s_beta[[i, b]];
506                    }
507                    pen += 0.5 * coef * cross;
508                }
509            }
510            pen
511        }
512    }
513}
514
515/// Fill the gradient of the penalized negative log-likelihood in the engine's
516/// class-major coefficient order. `residual = -∂ log L / ∂η`; the penalty
517/// contribution uses the same class-space metric as the objective and Hessian.
518/// Keeping this algebra in one production helper lets the loop and the final
519/// convergence certificate evaluate exactly the same score at different
520/// iterates.
521fn fill_penalized_gradient(
522    design: ArrayView2<'_, f64>,
523    residual: ArrayView2<'_, f64>,
524    beta: &Array2<f64>,
525    penalty: ArrayView2<'_, f64>,
526    lambdas: ArrayView1<'_, f64>,
527    metric: ClassPenaltyMetric,
528    out: &mut Array1<f64>,
529) {
530    let (p, m) = beta.dim();
531    for a in 0..m {
532        for i in 0..p {
533            let mut acc = 0.0_f64;
534            for row in 0..design.nrows() {
535                acc += design[[row, i]] * residual[[row, a]];
536            }
537            out[a * p + i] = acc;
538        }
539    }
540    match metric {
541        ClassPenaltyMetric::Diagonal => {
542            for a in 0..m {
543                let la = lambdas[a];
544                if la == 0.0 {
545                    continue;
546                }
547                let beta_col = beta.column(a);
548                for i in 0..p {
549                    let mut s_beta_i = 0.0_f64;
550                    for j in 0..p {
551                        s_beta_i += penalty[[i, j]] * beta_col[j];
552                    }
553                    out[a * p + i] += la * s_beta_i;
554                }
555            }
556        }
557        ClassPenaltyMetric::Centered if m > 0 && lambdas[0] != 0.0 => {
558            let lam = lambdas[0];
559            let inv_k = 1.0 / ((m + 1) as f64);
560            let mut beta_bar = vec![0.0_f64; p];
561            for a in 0..m {
562                let col = beta.column(a);
563                for i in 0..p {
564                    beta_bar[i] += col[i];
565                }
566            }
567            for value in &mut beta_bar {
568                *value *= inv_k;
569            }
570            for a in 0..m {
571                let beta_col = beta.column(a);
572                for i in 0..p {
573                    let mut s_centered_i = 0.0_f64;
574                    for j in 0..p {
575                        s_centered_i += penalty[[i, j]] * (beta_col[j] - beta_bar[j]);
576                    }
577                    out[a * p + i] += lam * s_centered_i;
578                }
579            }
580        }
581        ClassPenaltyMetric::Centered => {}
582        // EquivariantPerClass (#2344): out_a += Σ_b A[a,b]·S·β_b — the exact
583        // gradient of the ½·Σ A[a,b]·β_aᵀSβ_b objective arm (A symmetric).
584        ClassPenaltyMetric::EquivariantPerClass if m > 0 => {
585            let a_mat = equivariant_class_metric(lambdas, m);
586            let mut s_beta = Array2::<f64>::zeros((p, m));
587            for b in 0..m {
588                let col = beta.column(b);
589                for i in 0..p {
590                    let mut acc = 0.0_f64;
591                    for j in 0..p {
592                        acc += penalty[[i, j]] * col[j];
593                    }
594                    s_beta[[i, b]] = acc;
595                }
596            }
597            for a in 0..m {
598                for i in 0..p {
599                    let mut acc = 0.0_f64;
600                    for b in 0..m {
601                        acc += a_mat[[a, b]] * s_beta[[i, b]];
602                    }
603                    out[a * p + i] += acc;
604                }
605            }
606        }
607        ClassPenaltyMetric::EquivariantPerClass => {}
608    }
609}
610
611/// Invert the symmetric penalized Hessian `H` to the joint Laplace covariance
612/// `Σ = H⁻¹` by solving `H·Σ = I` through the shared symmetric factorization
613/// (#1101). `dim` is the flat block dimension `P·M`; `context` prefixes any
614/// diagnostic. A curvature-scaled Tikhonov ridge `τ·I` — floored at
615/// [`BASE_RIDGE_FRACTION_OF_MAX_DIAG`]·max_diag and escalated geometrically up
616/// to [`MAX_RIDGE_ESCALATIONS`] times — is added ONLY when the raw factor/solve
617/// is non-finite (a rank-deficient null direction), exactly mirroring the
618/// Newton step's ridge so the covariance is always finite; at full rank the
619/// ridge is never engaged and `Σ` is the exact `H⁻¹`. The returned matrix is
620/// symmetrized `(Σ + Σᵀ)/2` to null round-off asymmetry from the back-solve.
621fn invert_symmetric_penalized_hessian(
622    hessian: &Array2<f64>,
623    dim: usize,
624    context: &str,
625) -> Result<Array2<f64>, EstimationError> {
626    let max_diag = (0..dim).fold(0.0_f64, |acc, idx| acc.max(hessian[[idx, idx]].abs()));
627    let base_ridge = if max_diag.is_finite() && max_diag > 0.0 {
628        max_diag * BASE_RIDGE_FRACTION_OF_MAX_DIAG
629    } else {
630        BASE_RIDGE_FRACTION_OF_MAX_DIAG
631    };
632    // `last_failure` distinguishes the two exhaustion modes so their distinct
633    // terminal errors survive the migration: `Some((ridge, err))` when the
634    // final attempt died in the factorization, `None` when it factored but the
635    // back-solve stayed non-finite.
636    let mut last_failure: Option<(f64, String)> = None;
637    let mut try_ridge = |ridge: f64| -> Option<Array2<f64>> {
638        let mut ridged = hessian.clone();
639        if ridge > 0.0 {
640            for idx in 0..dim {
641                ridged[[idx, idx]] += ridge;
642            }
643        }
644        let factor = match factorize_symmetricwith_fallback(
645            FaerArrayView::new(&ridged).as_ref(),
646            Side::Lower,
647        ) {
648            Ok(factor) => factor,
649            Err(err) => {
650                last_failure = Some((ridge, err.to_string()));
651                return None;
652            }
653        };
654        // Solve H·Σ = I: identity RHS, back-solved in place to yield Σ = H⁻¹.
655        let mut rhs = Array2::<f64>::eye(dim);
656        {
657            let rhs_view = array2_to_matmut(&mut rhs);
658            factor.solve_in_place(rhs_view);
659        }
660        if !rhs.iter().all(|v| v.is_finite()) {
661            last_failure = None;
662            return None;
663        }
664        // Symmetrize to remove round-off asymmetry from the back-solve.
665        let mut cov = Array2::<f64>::zeros((dim, dim));
666        for i in 0..dim {
667            for j in 0..dim {
668                cov[[i, j]] = 0.5 * (rhs[[i, j]] + rhs[[j, i]]);
669            }
670        }
671        Some(cov)
672    };
673    // Bare (unridged) attempt first — at full rank the ridge is never engaged —
674    // then the geometric escalation from `base_ridge` with the doubling growth
675    // this site has always used.
676    if let Some(cov) = try_ridge(0.0) {
677        return Ok(cov);
678    }
679    match escalate_ridge(
680        RidgeSchedule {
681            initial: base_ridge,
682            growth: 2.0,
683            max_escalations: MAX_RIDGE_ESCALATIONS,
684        },
685        &mut try_ridge,
686    ) {
687        Ok(success) => Ok(success.value),
688        Err(_) => match last_failure {
689            Some((ridge, err)) => Err(EstimationError::InvalidInput(format!(
690                "{context}: covariance factorization failed even with ridge \
691                 {ridge:.3e}: {err}"
692            ))),
693            None => Err(EstimationError::InvalidInput(format!(
694                "{context}: covariance solve remained non-finite after {} ridge escalations \
695                 (max_diag={max_diag:.3e})",
696                MAX_RIDGE_ESCALATIONS,
697            ))),
698        },
699    }
700}
701
702/// Fit a penalized vector-response GLM at fixed `λ` via damped Newton.
703///
704/// The `likelihood` adapter supplies the per-row Fisher block, the residual
705/// gradient, and the log-likelihood; the engine owns the entire optimisation
706/// scaffold. See the module docs for the optimisation problem, the
707/// output-major coefficient ordering, and the convergence semantics.
708///
709/// `context` is woven into every diagnostic message so each family keeps its
710/// own error prefix (e.g. `"fit_penalized_multinomial"`).
711pub fn fit_penalized_vector_glm<L: VectorLikelihood>(
712    inputs: PenalizedVectorGlmInputs<'_>,
713    likelihood: &L,
714    context: &str,
715) -> Result<VectorGlmSolve, EstimationError> {
716    let PenalizedVectorGlmInputs {
717        design,
718        y,
719        penalty,
720        lambdas,
721        fisher_w_override,
722        max_iter,
723        tol,
724        class_penalty_metric,
725        resume_from,
726    } = inputs;
727
728    // ────────────────────────────── shape checks ──────────────────────────
729    let n_obs = design.nrows();
730    let p = design.ncols();
731    if n_obs == 0 || p == 0 {
732        crate::bail_invalid_estim!("{context}: design must be nonempty (got {n_obs}x{p})");
733    }
734    let m = class_penalty_metric.active_outputs(lambdas.len());
735    if m == 0 {
736        crate::bail_invalid_estim!("{context}: need at least one active output (got M=0)");
737    }
738    if y.nrows() != n_obs {
739        crate::bail_invalid_estim!("{context}: y rows {} ≠ design rows {n_obs}", y.nrows());
740    }
741    if penalty.dim() != (p, p) {
742        crate::bail_invalid_estim!(
743            "{context}: penalty shape {:?} ≠ (P, P) = ({p}, {p})",
744            penalty.dim()
745        );
746    }
747    for (i, &v) in lambdas.iter().enumerate() {
748        if !(v.is_finite() && v >= 0.0) {
749            crate::bail_invalid_estim!("{context}: lambdas[{i}] must be finite and ≥ 0 (got {v})");
750        }
751    }
752    if let Some(fw) = fisher_w_override.as_ref() {
753        if fw.dim() != (n_obs, m, m) {
754            crate::bail_invalid_estim!(
755                "{context}: fisher_w_override shape {:?} ≠ (N, M, M) = ({n_obs}, {m}, {m})",
756                fw.dim()
757            );
758        }
759    }
760    for ((i, j), &v) in design.indexed_iter() {
761        if !v.is_finite() {
762            crate::bail_invalid_estim!("{context}: design[{i},{j}] must be finite (got {v})");
763        }
764    }
765
766    // ────────────────────────── Newton iteration ──────────────────────────
767    // β stored as (P, M) column-major-per-output; flat index uses output-major
768    // ordering `flat[a · P + i] = β[i, a]` to align with `dense_block_xtwx`.
769    let (mut beta, completed_iterations) = match resume_from {
770        Some(resume) => {
771            if resume.coefficients.dim() != (p, m) {
772                crate::bail_invalid_estim!(
773                    "{context}: resume checkpoint coefficient shape {:?} ≠ (P, M) = ({p}, {m})",
774                    resume.coefficients.dim()
775                );
776            }
777            for ((i, a), &value) in resume.coefficients.indexed_iter() {
778                if !value.is_finite() {
779                    crate::bail_invalid_estim!(
780                        "{context}: resume checkpoint coefficient[{i},{a}] must be finite (got {value})"
781                    );
782                }
783            }
784            (resume.coefficients.to_owned(), resume.completed_iterations)
785        }
786        None => (Array2::<f64>::zeros((p, m)), 0),
787    };
788    let mut eta = Array2::<f64>::zeros((n_obs, m));
789    // Reused η scratch for the line-search objective probes (see
790    // `evaluate_objective`): overwritten in full on every call, so it carries
791    // no state between calls and hoisting it out of the backtracking loop is a
792    // pure heap-allocation removal with no effect on the computed objective.
793    let mut eta_objective_scratch = Array2::<f64>::zeros((n_obs, m));
794    let beta_flat_dim = p * m;
795    // Reused penalized-gradient buffer: each Newton iteration writes every entry
796    // `grad_flat[a·p + i] = Xᵀr` (direct assignment over all a∈0..m, i∈0..p)
797    // before adding the penalty term and before any read, so it carries no state
798    // across iterations and hoisting it out of the Newton loop is a pure
799    // heap-allocation removal with no effect on the computed gradient.
800    let mut grad_flat = Array1::<f64>::zeros(beta_flat_dim);
801
802    let mut iterations = completed_iterations;
803    let mut small_step_reached = false;
804    let mut stall_reason = VectorGlmStallReason::IterationBudgetExhausted;
805    let mut last_objective = f64::INFINITY;
806
807    // η = X · β for the current β, reused by the analytic Fisher / gradient.
808    let recompute_eta = |beta: &Array2<f64>, eta: &mut Array2<f64>| {
809        for a in 0..m {
810            let beta_col = beta.column(a);
811            for row in 0..n_obs {
812                let mut eta_val = 0.0_f64;
813                for i in 0..p {
814                    eta_val += design[[row, i]] * beta_col[i];
815                }
816                eta[[row, a]] = eta_val;
817            }
818        }
819    };
820
821    // Penalized objective F(β) = − log L(X β) + ½ Σ_a λ_a β_aᵀ S β_a.
822    // The caller supplies a reused `(n_obs, m)` scratch for η = X·β so the
823    // backtracking line search (which calls this up to `MAX_BACKTRACKS + 1`
824    // times per Newton iteration) does not heap-allocate a fresh η buffer on
825    // every probe. The scratch is overwritten in full by `recompute_eta` before
826    // it is read, so reusing it is bit-for-bit identical to the prior
827    // allocate-fresh body: `recompute_eta` runs the SAME `Σ_i design·β` loop in
828    // the SAME order this closure used inline.
829    let evaluate_objective =
830        |beta_trial: &Array2<f64>, eta_scratch: &mut Array2<f64>| -> Result<f64, EstimationError> {
831            recompute_eta(beta_trial, eta_scratch);
832            let ll = likelihood.log_lik(eta_scratch.view(), y)?;
833            let pen = weighted_penalty_sum(beta_trial, penalty, lambdas, class_penalty_metric);
834            Ok(-ll + pen)
835        };
836
837    for iter in 0..max_iter {
838        iterations = completed_iterations.checked_add(iter + 1).ok_or_else(|| {
839            EstimationError::InvalidInput(format!(
840                "{context}: resume checkpoint iteration count overflowed usize"
841            ))
842        })?;
843
844        recompute_eta(&beta, &mut eta);
845
846        // Per-row dense Fisher block W_{n,a,b} = −∂² log L / ∂η_a ∂η_b: either
847        // the caller-supplied curvature override (issue #349 escape-hatch —
848        // curvature only) or the analytic [`VectorLikelihood::hess_block`]. The
849        // residual r_{n,a} = −∂ log L / ∂η_a stays analytic in both cases.
850        let analytic_fisher = match fisher_w_override.as_ref() {
851            Some(_) => None,
852            None => Some(likelihood.hess_block(eta.view(), y)?),
853        };
854        let fisher_blocks = match fisher_w_override.as_ref() {
855            Some(fw) => *fw,
856            None => analytic_fisher
857                .as_ref()
858                .expect("analytic Fisher computed when no override")
859                .view(),
860        };
861        let residual = likelihood.grad_eta(eta.view(), y)?.mapv(|v| -v);
862
863        // Penalized Hessian: H = block(XᵀWX) + diag_a(λ_a S).
864        let mut hessian = dense_block_xtwx(design, fisher_blocks, None)?;
865        if hessian.nrows() != beta_flat_dim || hessian.ncols() != beta_flat_dim {
866            crate::bail_invalid_estim!(
867                "{context}: assembled Hessian shape {:?} ≠ ({beta_flat_dim}, {beta_flat_dim})",
868                hessian.dim()
869            );
870        }
871        match class_penalty_metric {
872            ClassPenaltyMetric::Diagonal => {
873                for a in 0..m {
874                    let la = lambdas[a];
875                    if la == 0.0 {
876                        continue;
877                    }
878                    let base = a * p;
879                    for i in 0..p {
880                        for j in 0..p {
881                            hessian[[base + i, base + j]] += la * penalty[[i, j]];
882                        }
883                    }
884                }
885            }
886            // Centered (#1587): H_{ab} += λ·(δ_ab − 1/K)·S, K = M+1, shared
887            // λ = lambdas[0] — couples every class pair via the −(λ/K)·S
888            // off-diagonals. Reference-invariant softmax penalty.
889            ClassPenaltyMetric::Centered if m > 0 && lambdas[0] != 0.0 => {
890                let lam = lambdas[0];
891                let inv_k = 1.0 / ((m + 1) as f64);
892                for a in 0..m {
893                    for b in 0..m {
894                        let coef = lam * (if a == b { 1.0 } else { 0.0 } - inv_k);
895                        let (ba, bb) = (a * p, b * p);
896                        for i in 0..p {
897                            for j in 0..p {
898                                hessian[[ba + i, bb + j]] += coef * penalty[[i, j]];
899                            }
900                        }
901                    }
902                }
903            }
904            ClassPenaltyMetric::Centered => {}
905            // EquivariantPerClass (#2344): H += A(λ) ⊗ S, the coupled
906            // heterogeneous per-class blocks.
907            ClassPenaltyMetric::EquivariantPerClass => {
908                add_equivariant_penalty_blocks(&mut hessian, penalty, lambdas, p, m);
909            }
910        }
911
912        fill_penalized_gradient(
913            design,
914            residual.view(),
915            &beta,
916            penalty,
917            lambdas,
918            class_penalty_metric,
919            &mut grad_flat,
920        );
921
922        // δ = − H^{-1} · grad, solved through an adaptive Levenberg–Marquardt
923        // ridge. The penalized Hessian `H = block(XᵀWX) + diag_a(λ_a S)` can be
924        // rank-deficient — a multinomial class block with quasi-separated /
925        // collinear columns and a small per-class λ leaves `XᵀW_aX + λ_a S`
926        // singular. faer's symmetric fallback chain ends at Bunch–Kaufman
927        // (LBLᵀ), which factorizes indefinite/singular matrices "successfully"
928        // and then back-substitutes through near-zero pivots, yielding a
929        // non-finite δ. Rather than aborting the whole fit on one bad block, we
930        // add a small ridge `τ·I` (Levenberg style) to the diagonal and
931        // re-factorize, escalating τ geometrically until the step is finite.
932        //
933        // The base ridge is scaled by the Hessian's largest diagonal entry so
934        // it is invariant to the problem's overall curvature scale: a tiny
935        // nudge relative to the dominant curvature, large enough to lift the
936        // null directions off zero. A finite δ from the ridged system is a
937        // descent direction for the *unridged* penalized objective `F`
938        // (ridging only shrinks the step toward the gradient direction), and
939        // the backtracking line search below validates it against `F` itself,
940        // so the ridge never biases the converged β̂ — at the optimum the
941        // gradient vanishes and the step → 0 regardless of τ.
942        let max_diag =
943            (0..beta_flat_dim).fold(0.0_f64, |acc, idx| acc.max(hessian[[idx, idx]].abs()));
944        // The ridge floors at `base_ridge` (not 0) for every solve. An exactly
945        // rank-deficient block (e.g. duplicate / collinear design columns under
946        // a near-zero λ) leaves `H = block(XᵀWX) + diag_a(λ_a S)` singular along
947        // a null direction. faer's Bunch–Kaufman fallback factorizes a singular
948        // matrix "successfully" and back-substitutes through the zero pivot to a
949        // *finite but arbitrary* component in the null space, so the resulting
950        // Newton direction is not a descent direction in the identified
951        // subspace — the line search then shrinks α toward 0 and the step-norm
952        // test declares a false convergence at a point where the unridged
953        // penalized gradient on identified directions is still large (gam#856).
954        // A minimal Tikhonov ridge `base_ridge·I` resolves the null direction to
955        // its minimum-norm representative, giving a true descent direction.
956        let base_ridge = if max_diag.is_finite() && max_diag > 0.0 {
957            max_diag * BASE_RIDGE_FRACTION_OF_MAX_DIAG
958        } else {
959            BASE_RIDGE_FRACTION_OF_MAX_DIAG
960        };
961        // A genuine factorization failure (not just a singular pivot) is
962        // remembered so exhaustion can surface its distinct terminal error;
963        // singular pivots back-substituted to ±inf/NaN just escalate.
964        let mut last_factor_err: Option<(f64, String)> = None;
965        let delta = match escalate_ridge(
966            RidgeSchedule {
967                initial: base_ridge,
968                growth: 2.0,
969                max_escalations: MAX_RIDGE_ESCALATIONS + 1,
970            },
971            |ridge| {
972                let mut ridged = hessian.clone();
973                for idx in 0..beta_flat_dim {
974                    ridged[[idx, idx]] += ridge;
975                }
976                let factor = match factorize_symmetricwith_fallback(
977                    FaerArrayView::new(&ridged).as_ref(),
978                    Side::Lower,
979                ) {
980                    Ok(factor) => factor,
981                    Err(err) => {
982                        last_factor_err = Some((ridge, err.to_string()));
983                        return None;
984                    }
985                };
986                last_factor_err = None;
987                let mut rhs = Array2::<f64>::zeros((beta_flat_dim, 1));
988                for i in 0..beta_flat_dim {
989                    rhs[[i, 0]] = -grad_flat[i];
990                }
991                {
992                    let rhs_view = array2_to_matmut(&mut rhs);
993                    factor.solve_in_place(rhs_view);
994                }
995                (0..beta_flat_dim)
996                    .all(|i| rhs[[i, 0]].is_finite())
997                    .then(|| Array1::from_iter((0..beta_flat_dim).map(|i| rhs[[i, 0]])))
998            },
999        ) {
1000            Ok(success) => success.value,
1001            Err(exhausted) => {
1002                if let Some((ridge, err)) = last_factor_err {
1003                    return Err(EstimationError::InvalidInput(format!(
1004                        "{context}: Hessian factorization failed at iter {iter} \
1005                         even with ridge {ridge:.3e}: {err}"
1006                    )));
1007                }
1008                return Err(EstimationError::InvalidInput(format!(
1009                    "{context}: Newton step remained non-finite at iter {iter} after {} ridge \
1010                     escalations up to {:.3e}; the penalized Hessian is pathologically \
1011                     rank-deficient (grad_norm={:.3e}, max_diag={max_diag:.3e})",
1012                    MAX_RIDGE_ESCALATIONS,
1013                    exhausted.next_ridge,
1014                    grad_flat.iter().map(|v| v * v).sum::<f64>().sqrt(),
1015                )));
1016            }
1017        };
1018
1019        // Damped acceptance: full step first, halve up to `MAX_BACKTRACKS` times
1020        // if the penalized negative log-likelihood fails to decrease. The first
1021        // iteration seeds `last_objective` from the initial β.
1022        let proposed_beta = |alpha: f64| -> Array2<f64> {
1023            let mut out = beta.clone();
1024            for a in 0..m {
1025                for i in 0..p {
1026                    out[[i, a]] += alpha * delta[a * p + i];
1027                }
1028            }
1029            out
1030        };
1031        if iter == 0 {
1032            last_objective = evaluate_objective(&beta, &mut eta_objective_scratch)?;
1033            if !last_objective.is_finite() {
1034                crate::bail_invalid_estim!("{context}: non-finite objective at β = 0");
1035            }
1036        }
1037        let accepted = backtracking_line_search::<_, EstimationError>(
1038            BacktrackConfig {
1039                contraction: LINE_SEARCH_SHRINK,
1040                max_steps: MAX_BACKTRACKS + 1,
1041                ..BacktrackConfig::default()
1042            },
1043            |alpha| {
1044                let candidate = proposed_beta(alpha);
1045                let objective = evaluate_objective(&candidate, &mut eta_objective_scratch)?;
1046                Ok(Some((objective, candidate)))
1047            },
1048            |_alpha, f| f.is_finite() && f <= last_objective + OBJECTIVE_DECREASE_SLACK,
1049        )?;
1050        let Some(accepted) = accepted else {
1051            // Every candidate failed the descent certificate. Keep the last
1052            // ACCEPTED iterate as checkpoint evidence; a rejected trial can
1053            // never become a result merely because the line-search budget was
1054            // exhausted.
1055            stall_reason = VectorGlmStallReason::LineSearchExhausted;
1056            break;
1057        };
1058        let accepted_beta = accepted.payload;
1059        let new_objective = accepted.value;
1060
1061        let mut step_norm_sq = 0.0_f64;
1062        let mut beta_norm_sq = 0.0_f64;
1063        for a in 0..m {
1064            for i in 0..p {
1065                let d = accepted_beta[[i, a]] - beta[[i, a]];
1066                step_norm_sq += d * d;
1067                let v = accepted_beta[[i, a]];
1068                beta_norm_sq += v * v;
1069            }
1070        }
1071
1072        beta = accepted_beta;
1073        last_objective = new_objective;
1074
1075        let step_norm = step_norm_sq.sqrt();
1076        let beta_norm = beta_norm_sq.sqrt();
1077        // First-order optimality gate (gam#856): the step-norm test alone can
1078        // fire prematurely when a backtracking line search has shrunk α on a
1079        // poor direction, leaving a point that is NOT stationary. `grad_flat`
1080        // is the unridged penalized gradient ∇F(β) at the pre-step β; with a
1081        // small step it is ≈ ∇F at the accepted β. Its norm reflects only
1082        // identified directions (it is exactly zero along an unidentified null
1083        // direction such as a duplicate-column e₁−e₂ split), so requiring it to
1084        // be small certifies first-order optimality on the identified subspace
1085        // without penalizing legitimate non-identifiability. Scale the gate by
1086        // the data magnitude so it is invariant to problem scale.
1087        let grad_norm = grad_flat.iter().map(|v| v * v).sum::<f64>().sqrt();
1088        // Curvature-scaled optimality threshold: `max_diag` is the dominant
1089        // penalized-Hessian diagonal entry, so `OPTIMALITY_GRAD_FRACTION·max_diag`
1090        // is a tiny gradient relative to the problem's curvature scale and is
1091        // reached by a few quadratically-converging Newton steps on this smooth,
1092        // bounded softmax/binomial likelihood.
1093        let grad_optimal = grad_norm <= OPTIMALITY_GRAD_FRACTION * (1.0 + max_diag);
1094        if step_norm <= tol * (1.0 + beta_norm) && grad_optimal {
1095            small_step_reached = true;
1096            break;
1097        }
1098    }
1099
1100    // ──────────────────────────── post-process ────────────────────────────
1101    recompute_eta(&beta, &mut eta);
1102    let log_likelihood = likelihood.log_lik(eta.view(), y)?;
1103    let penalty_term = weighted_penalty_sum(&beta, penalty, lambdas, class_penalty_metric);
1104
1105    // Re-assemble the final penalized Hessian before certification. This is not
1106    // posterior work: its diagonal supplies the same curvature scale used by
1107    // the loop's first-order gate. Covariance inversion remains below the gate
1108    // and is therefore impossible for an uncertified iterate.
1109    //
1110    // Joint Laplace covariance `H⁻¹` at the converged mode (#1101). Re-assemble
1111    // the penalized Hessian `H = block(XᵀWX) + penalty` at β̂ — the SAME algebra
1112    // the Newton loop runs each iteration — and invert it by solving `H·Σ = I`
1113    // through the shared symmetric factorization. The Newton loop discarded its
1114    // per-step factor; this recomputes the factor once at the mode where the
1115    // curvature is the correct posterior precision. A tiny curvature-scaled
1116    // ridge is added only when the raw factorization / solve is non-finite
1117    // (rank-deficient null direction), mirroring the Newton step's ridge logic,
1118    // so the covariance is always finite; at full rank the ridge is never used.
1119    let analytic_fisher_final = match fisher_w_override.as_ref() {
1120        Some(_) => None,
1121        None => Some(likelihood.hess_block(eta.view(), y)?),
1122    };
1123    let fisher_blocks_final = match fisher_w_override.as_ref() {
1124        Some(fw) => *fw,
1125        None => analytic_fisher_final
1126            .as_ref()
1127            .expect("analytic Fisher computed when no override")
1128            .view(),
1129    };
1130    let mut hessian_final = dense_block_xtwx(design, fisher_blocks_final, None)?;
1131    match class_penalty_metric {
1132        ClassPenaltyMetric::Diagonal => {
1133            for a in 0..m {
1134                let la = lambdas[a];
1135                if la == 0.0 {
1136                    continue;
1137                }
1138                let base = a * p;
1139                for i in 0..p {
1140                    for j in 0..p {
1141                        hessian_final[[base + i, base + j]] += la * penalty[[i, j]];
1142                    }
1143                }
1144            }
1145        }
1146        ClassPenaltyMetric::Centered if m > 0 && lambdas[0] != 0.0 => {
1147            let lam = lambdas[0];
1148            let inv_k = 1.0 / ((m + 1) as f64);
1149            for a in 0..m {
1150                for b in 0..m {
1151                    let coef = lam * (if a == b { 1.0 } else { 0.0 } - inv_k);
1152                    let (ba, bb) = (a * p, b * p);
1153                    for i in 0..p {
1154                        for j in 0..p {
1155                            hessian_final[[ba + i, bb + j]] += coef * penalty[[i, j]];
1156                        }
1157                    }
1158                }
1159            }
1160        }
1161        ClassPenaltyMetric::Centered => {}
1162        // EquivariantPerClass (#2344): H += A(λ) ⊗ S, the coupled
1163        // heterogeneous per-class blocks.
1164        ClassPenaltyMetric::EquivariantPerClass => {
1165            add_equivariant_penalty_blocks(&mut hessian_final, penalty, lambdas, p, m);
1166        }
1167    }
1168
1169    // Re-evaluate the exact penalized score AT the accepted final iterate. The
1170    // loop's inexpensive gate uses the pre-step score (valid to first order
1171    // when the accepted step is tiny); this second evaluation closes the only
1172    // gap through which heavy backtracking could otherwise certify a point
1173    // whose post-step score is still material.
1174    let final_residual = likelihood.grad_eta(eta.view(), y)?.mapv(|value| -value);
1175    fill_penalized_gradient(
1176        design,
1177        final_residual.view(),
1178        &beta,
1179        penalty,
1180        lambdas,
1181        class_penalty_metric,
1182        &mut grad_flat,
1183    );
1184    let final_grad_norm = grad_flat
1185        .iter()
1186        .map(|value| value * value)
1187        .sum::<f64>()
1188        .sqrt();
1189    let final_max_diag =
1190        (0..beta_flat_dim).fold(0.0_f64, |acc, i| acc.max(hessian_final[[i, i]].abs()));
1191    let final_grad_optimal = final_grad_norm <= OPTIMALITY_GRAD_FRACTION * (1.0 + final_max_diag);
1192    if !(small_step_reached && final_grad_optimal) {
1193        if small_step_reached {
1194            stall_reason = VectorGlmStallReason::PostStepCertificateFailed;
1195        }
1196        // Budget exhausted (or the post-step score failed certification). Hand
1197        // back checkpoint evidence — never a covariance or fitted probabilities.
1198        // The adapter decides between a typed non-convergence error and the
1199        // multinomial separation → Firth/Jeffreys escalation.
1200        return Ok(VectorGlmSolve::Stalled(VectorGlmStall {
1201            reason: stall_reason,
1202            coefficients: beta,
1203            eta,
1204            iterations,
1205            log_likelihood,
1206            penalty_term,
1207            gradient_norm: final_grad_norm,
1208            gradient_bound: OPTIMALITY_GRAD_FRACTION * (1.0 + final_max_diag),
1209        }));
1210    }
1211
1212    let coefficient_covariance =
1213        invert_symmetric_penalized_hessian(&hessian_final, beta_flat_dim, context)?;
1214
1215    Ok(VectorGlmSolve::Converged(PenalizedVectorGlmOutputs {
1216        coefficients: beta,
1217        eta,
1218        iterations,
1219        log_likelihood,
1220        penalty_term,
1221        coefficient_covariance,
1222    }))
1223}
1224
1225#[cfg(test)]
1226mod parity_tests {
1227    //! Parity tests for the shared scaffold across both Fisher-block families
1228    //! (issue #409). The engine is exercised through the two public adapters —
1229    //! [`crate::binomial_multi::fit_penalized_binomial_multi`]
1230    //! (row-diagonal block) and
1231    //! [`crate::multinomial::fit_penalized_multinomial`] (dense
1232    //! softmax block) — and we assert, with un-weakened bounds, that:
1233    //!
1234    //!   1. each fit hits the first-order optimality condition `∇F(β̂) = 0`,
1235    //!      verified by a central finite difference of the penalized objective
1236    //!      (the engine never sees this gradient, so this is an independent
1237    //!      check that the shared Newton scaffold converged correctly);
1238    //!   2. the reported fitted probabilities are consistent with `β̂` and the
1239    //!      reported deviance equals `−2 · log L(β̂)`;
1240    //!   3. for the binomial family, the `K`-column joint solve reproduces a
1241    //!      from-scratch single-column penalized logistic Newton solve column
1242    //!      for column (the row-diagonal block must decouple exactly).
1243
1244    use super::{ClassPenaltyMetric, weighted_penalty_sum};
1245    use crate::binomial_multi::{BinomialMultiFitInputs, fit_penalized_binomial_multi};
1246    use crate::multinomial::{MultinomialFitInputs, fit_penalized_multinomial};
1247    use gam_test_support::fd_checker::numerical_gradient_central_diff;
1248    use ndarray::{Array1, Array2};
1249
1250    /// #1587: the `Centered` class-penalty metric is invariant to the arbitrary
1251    /// reference-class choice. Penalizing the `K−1` ALR contrasts under ANY of
1252    /// the `K` baselines yields the same value (the symmetric CLR penalty
1253    /// `Σ_k β̃_kᵀSβ̃_k`), whereas the historical `Diagonal` metric does not — that
1254    /// non-invariance is exactly the #1587 defect. Pure-algebra check on the
1255    /// penalty form (no fit), so it pins the engine foundation the production
1256    /// wiring (REML per-term λ re-key) will build on.
1257    #[test]
1258    fn centered_penalty_is_reference_class_invariant_1587() {
1259        // K = 3 classes, p = 2 coefficients; symmetric PSD penalty S.
1260        let s = ndarray::array![[2.0_f64, 0.5], [0.5, 1.0]];
1261        // A CLR (sum-to-zero) coefficient set: β̃_0 + β̃_1 + β̃_2 = 0.
1262        let bt = [[1.0_f64, 0.5], [-0.3, 0.2], [-0.7, -0.7]];
1263        for j in 0..2 {
1264            let colsum: f64 = (0..3).map(|k| bt[k][j]).sum();
1265            assert!(colsum.abs() < 1e-12, "test CLR set must sum to zero");
1266        }
1267        // Direct symmetric penalty Σ_k β̃_kᵀ S β̃_k.
1268        let mut symmetric = 0.0_f64;
1269        for k in 0..3 {
1270            for i in 0..2 {
1271                for j in 0..2 {
1272                    symmetric += bt[k][i] * s[[i, j]] * bt[k][j];
1273                }
1274            }
1275        }
1276        let lambdas = Array1::from(vec![1.0_f64, 1.0]);
1277        let mut centered_vals = Vec::new();
1278        let mut diagonal_vals = Vec::new();
1279        // For each reference class r, the two ALR contrasts are β̃_a − β̃_r (a≠r).
1280        for r in 0..3 {
1281            let others: Vec<usize> = (0..3).filter(|&k| k != r).collect();
1282            let mut beta = Array2::<f64>::zeros((2, 2));
1283            for (a, &o) in others.iter().enumerate() {
1284                for i in 0..2 {
1285                    beta[[i, a]] = bt[o][i] - bt[r][i];
1286                }
1287            }
1288            let c = weighted_penalty_sum(
1289                &beta,
1290                s.view(),
1291                lambdas.view(),
1292                ClassPenaltyMetric::Centered,
1293            );
1294            let d = weighted_penalty_sum(
1295                &beta,
1296                s.view(),
1297                lambdas.view(),
1298                ClassPenaltyMetric::Diagonal,
1299            );
1300            assert!(
1301                (c - 0.5 * symmetric).abs() < 1e-12,
1302                "ref {r}: Centered penalty {c} must equal ½·symmetric {}",
1303                0.5 * symmetric
1304            );
1305            centered_vals.push(c);
1306            diagonal_vals.push(d);
1307        }
1308        let cspread = centered_vals.iter().cloned().fold(f64::MIN, f64::max)
1309            - centered_vals.iter().cloned().fold(f64::MAX, f64::min);
1310        assert!(
1311            cspread < 1e-12,
1312            "Centered must be reference-invariant; got {centered_vals:?}"
1313        );
1314        let dspread = diagonal_vals.iter().cloned().fold(f64::MIN, f64::max)
1315            - diagonal_vals.iter().cloned().fold(f64::MAX, f64::min);
1316        assert!(
1317            dspread > 1e-6,
1318            "Diagonal is the non-invariant #1587 path; references must disagree, got {diagonal_vals:?}"
1319        );
1320    }
1321
1322    fn sigmoid(eta: f64) -> f64 {
1323        if eta >= 0.0 {
1324            1.0 / (1.0 + (-eta).exp())
1325        } else {
1326            let e = eta.exp();
1327            e / (1.0 + e)
1328        }
1329    }
1330
1331    /// Softmax with implicit reference column (η_ref = 0) over `M` active η.
1332    fn softmax_ref(eta_active: &[f64]) -> Vec<f64> {
1333        let m = eta_active.len();
1334        let mut out = vec![0.0_f64; m + 1];
1335        let mut max_eta = 0.0_f64;
1336        for &v in eta_active {
1337            if v > max_eta {
1338                max_eta = v;
1339            }
1340        }
1341        let baseline = (-max_eta).exp();
1342        let mut denom = baseline;
1343        for (idx, &v) in eta_active.iter().enumerate() {
1344            let e = (v - max_eta).exp();
1345            out[idx] = e;
1346            denom += e;
1347        }
1348        for v in out.iter_mut().take(m) {
1349            *v /= denom;
1350        }
1351        out[m] = baseline / denom;
1352        out
1353    }
1354
1355    /// Penalized negative log-likelihood for the independent-binomial family at
1356    /// a candidate coefficient matrix `β ∈ ℝ^{P×K}`, computed directly from the
1357    /// definition (no engine internals).
1358    fn binomial_objective(
1359        design: &Array2<f64>,
1360        y: &Array2<f64>,
1361        penalty: &Array2<f64>,
1362        lambdas: &Array1<f64>,
1363        beta: &Array2<f64>,
1364    ) -> f64 {
1365        let (n, p) = design.dim();
1366        let k = y.ncols();
1367        let mut ll = 0.0_f64;
1368        for row in 0..n {
1369            for a in 0..k {
1370                let mut eta = 0.0_f64;
1371                for i in 0..p {
1372                    eta += design[[row, i]] * beta[[i, a]];
1373                }
1374                let mu = sigmoid(eta).clamp(1.0e-12, 1.0 - 1.0e-12);
1375                let yv = y[[row, a]];
1376                ll += yv * mu.ln() + (1.0 - yv) * (1.0 - mu).ln();
1377            }
1378        }
1379        let mut pen = 0.0_f64;
1380        for a in 0..k {
1381            let la = lambdas[a];
1382            for i in 0..p {
1383                let mut sbi = 0.0_f64;
1384                for j in 0..p {
1385                    sbi += penalty[[i, j]] * beta[[j, a]];
1386                }
1387                pen += 0.5 * la * beta[[i, a]] * sbi;
1388            }
1389        }
1390        -ll + pen
1391    }
1392
1393    /// Penalized negative log-likelihood for the multinomial family at a
1394    /// candidate active-class coefficient matrix `β ∈ ℝ^{P×(K-1)}`.
1395    fn multinomial_objective(
1396        design: &Array2<f64>,
1397        y_one_hot: &Array2<f64>,
1398        penalty: &Array2<f64>,
1399        lambdas: &Array1<f64>,
1400        beta: &Array2<f64>,
1401    ) -> f64 {
1402        let (n, p) = design.dim();
1403        let k = y_one_hot.ncols();
1404        let m = k - 1;
1405        let mut ll = 0.0_f64;
1406        let mut eta_active = vec![0.0_f64; m];
1407        for row in 0..n {
1408            for a in 0..m {
1409                let mut eta = 0.0_f64;
1410                for i in 0..p {
1411                    eta += design[[row, i]] * beta[[i, a]];
1412                }
1413                eta_active[a] = eta;
1414            }
1415            let probs = softmax_ref(&eta_active);
1416            for c in 0..k {
1417                let yc = y_one_hot[[row, c]];
1418                if yc != 0.0 {
1419                    ll += yc * probs[c].max(1.0e-300).ln();
1420                }
1421            }
1422        }
1423        // #2344 equivariant per-class penalty, RE-DERIVED independently of the
1424        // engine's metric assembly (Σ_c λ_c·γ_cᵀSγ_c on the centered class
1425        // functions, γ_c = β_c − β̄ with β_ref ≡ 0) so the FD parity witness
1426        // still checks the production algebra against a second formulation.
1427        let kf = k as f64;
1428        let mut pen = 0.0_f64;
1429        let mut beta_bar = vec![0.0_f64; p];
1430        for a in 0..m {
1431            for i in 0..p {
1432                beta_bar[i] += beta[[i, a]] / kf;
1433            }
1434        }
1435        for c in 0..k {
1436            let lc = lambdas[c];
1437            if lc == 0.0 {
1438                continue;
1439            }
1440            // γ_c[i] = β_c[i] − β̄[i]; the reference class has β_ref ≡ 0.
1441            let gamma_i = |i: usize| -> f64 {
1442                if c < m {
1443                    beta[[i, c]] - beta_bar[i]
1444                } else {
1445                    -beta_bar[i]
1446                }
1447            };
1448            for i in 0..p {
1449                let mut s_gamma_i = 0.0_f64;
1450                for j in 0..p {
1451                    s_gamma_i += penalty[[i, j]] * gamma_i(j);
1452                }
1453                pen += 0.5 * lc * gamma_i(i) * s_gamma_i;
1454            }
1455        }
1456        -ll + pen
1457    }
1458
1459    /// Max-norm of the central finite-difference gradient of an objective over
1460    /// every entry of a `(P, C)` coefficient matrix. The optimum must drive every
1461    /// component to ~0; we assert the max |component| against an un-weakened
1462    /// bound. This is a thin matrix reshape over the canonical scalar-objective
1463    /// FD helper — the finite-difference math itself lives in
1464    /// [`gam_test_support::fd_checker::numerical_gradient_central_diff`].
1465    fn fd_grad<F: Fn(&Array2<f64>) -> f64>(beta: &Array2<f64>, f: F) -> f64 {
1466        let (p, c) = beta.dim();
1467        let flat = Array1::from_iter(beta.iter().copied());
1468        let grad = numerical_gradient_central_diff(
1469            |x: &Array1<f64>| {
1470                let m = Array2::from_shape_vec((p, c), x.to_vec())
1471                    .expect("row-major reshape of coefficient vector");
1472                f(&m)
1473            },
1474            &flat,
1475            1.0e-6,
1476        );
1477        grad.iter().fold(0.0_f64, |acc, &g| acc.max(g.abs()))
1478    }
1479
1480    fn binomial_fixture() -> (Array2<f64>, Array2<f64>, Array2<f64>, Array1<f64>) {
1481        let n = 40;
1482        let p = 3;
1483        let k = 3;
1484        let design = Array2::<f64>::from_shape_fn((n, p), |(i, j)| match j {
1485            0 => 1.0,
1486            1 => ((i + 1) as f64 * 0.37).sin(),
1487            _ => ((i + 1) as f64 * 0.11).cos(),
1488        });
1489        let y = Array2::<f64>::from_shape_fn((n, k), |(i, a)| {
1490            // Deterministic but non-degenerate {0,1} labels per column.
1491            if ((i * 7 + a * 13 + 3) % 5) < 3 {
1492                1.0
1493            } else {
1494                0.0
1495            }
1496        });
1497        let penalty = Array2::<f64>::eye(p);
1498        let lambdas = Array1::from(vec![0.3_f64, 1.2, 2.5]);
1499        (design, y, penalty, lambdas)
1500    }
1501
1502    fn multinomial_fixture() -> (Array2<f64>, Array2<f64>, Array2<f64>, Array1<f64>) {
1503        let n = 45;
1504        let p = 3;
1505        let k = 4;
1506        let design = Array2::<f64>::from_shape_fn((n, p), |(i, j)| match j {
1507            0 => 1.0,
1508            1 => ((i + 2) as f64 * 0.29).sin(),
1509            _ => ((i + 2) as f64 * 0.17).cos(),
1510        });
1511        let mut y = Array2::<f64>::zeros((n, k));
1512        for i in 0..n {
1513            y[[i, (i * 3 + 1) % k]] = 1.0;
1514        }
1515        let penalty = Array2::<f64>::eye(p);
1516        // #2344: K per-class lambdas (reference class included), heterogeneous
1517        // so the equivariant metric's off-diagonal coupling is exercised.
1518        let lambdas = Array1::from(vec![0.5_f64, 1.0, 2.0, 0.8]);
1519        (design, y, penalty, lambdas)
1520    }
1521
1522    #[test]
1523    fn binomial_engine_hits_optimum_and_is_self_consistent() {
1524        let (design, y, penalty, lambdas) = binomial_fixture();
1525        let fit = fit_penalized_binomial_multi(BinomialMultiFitInputs {
1526            design: design.view(),
1527            y: y.view(),
1528            penalty: penalty.view(),
1529            lambdas: lambdas.view(),
1530            row_weights: None,
1531            fisher_w_override: None,
1532            max_iter: 100,
1533            tol: 1.0e-12,
1534        })
1535        .expect("binomial fit must succeed");
1536        // First-order optimality: ∇F(β̂) = 0 (engine never used this gradient).
1537        let g = fd_grad(&fit.coefficients, |b| {
1538            binomial_objective(&design, &y, &penalty, &lambdas, b)
1539        });
1540        assert!(
1541            g < 1.0e-6,
1542            "binomial penalized gradient at β̂ must vanish (max |∂F| = {g})"
1543        );
1544
1545        // Fitted probabilities reproduce σ(X β̂) and deviance = −2 log L.
1546        let (n, p) = design.dim();
1547        let k = y.ncols();
1548        let mut log_lik = 0.0_f64;
1549        for row in 0..n {
1550            for a in 0..k {
1551                let mut eta = 0.0_f64;
1552                for i in 0..p {
1553                    eta += design[[row, i]] * fit.coefficients[[i, a]];
1554                }
1555                let mu = sigmoid(eta);
1556                assert!(
1557                    (fit.fitted_probabilities[[row, a]] - mu).abs() < 1.0e-10,
1558                    "fitted probability must equal σ(X β̂)"
1559                );
1560                let muc = mu.clamp(1.0e-12, 1.0 - 1.0e-12);
1561                let yv = y[[row, a]];
1562                log_lik += yv * muc.ln() + (1.0 - yv) * (1.0 - muc).ln();
1563            }
1564        }
1565        assert!(
1566            (fit.deviance - (-2.0 * log_lik)).abs() < 1.0e-9,
1567            "deviance must equal −2 log L"
1568        );
1569    }
1570
1571    #[test]
1572    fn binomial_joint_solve_decouples_into_single_column_solves() {
1573        // Parity: the row-diagonal Fisher block means the K-column joint solve
1574        // must reproduce, column for column, an independent single-column
1575        // penalized logistic Newton solve. This is the defining property the
1576        // shared engine preserves for the independent-binomial family.
1577        let (design, y, penalty, lambdas) = binomial_fixture();
1578        let joint = fit_penalized_binomial_multi(BinomialMultiFitInputs {
1579            design: design.view(),
1580            y: y.view(),
1581            penalty: penalty.view(),
1582            lambdas: lambdas.view(),
1583            row_weights: None,
1584            fisher_w_override: None,
1585            max_iter: 100,
1586            tol: 1.0e-12,
1587        })
1588        .expect("joint fit must succeed");
1589
1590        let k = y.ncols();
1591        for a in 0..k {
1592            // Single-column problem: one binomial response, one λ.
1593            let y_col = y.column(a).to_owned().insert_axis(ndarray::Axis(1));
1594            let lam = Array1::from(vec![lambdas[a]]);
1595            let single = fit_penalized_binomial_multi(BinomialMultiFitInputs {
1596                design: design.view(),
1597                y: y_col.view(),
1598                penalty: penalty.view(),
1599                lambdas: lam.view(),
1600                row_weights: None,
1601                fisher_w_override: None,
1602                max_iter: 100,
1603                tol: 1.0e-12,
1604            })
1605            .expect("single-column fit must succeed");
1606            for i in 0..design.ncols() {
1607                let dj = joint.coefficients[[i, a]];
1608                let ds = single.coefficients[[i, 0]];
1609                assert!(
1610                    (dj - ds).abs() < 1.0e-8,
1611                    "joint column {a} coef {i} ({dj}) must match single-column solve ({ds})"
1612                );
1613            }
1614        }
1615    }
1616
1617    #[test]
1618    fn multinomial_engine_hits_optimum_and_is_self_consistent() {
1619        let (design, y, penalty, lambdas) = multinomial_fixture();
1620        let fit = fit_penalized_multinomial(MultinomialFitInputs {
1621            design: design.view(),
1622            y_one_hot: y.view(),
1623            penalty: penalty.view(),
1624            lambdas: lambdas.view(),
1625            row_weights: None,
1626            fisher_w_override: None,
1627            max_iter: 100,
1628            tol: 1.0e-12,
1629            resume_from: None,
1630        })
1631        .expect("multinomial fit must succeed");
1632        // First-order optimality: ∇F(β̂) = 0.
1633        let g = fd_grad(&fit.coefficients_active, |b| {
1634            multinomial_objective(&design, &y, &penalty, &lambdas, b)
1635        });
1636        assert!(
1637            g < 1.0e-6,
1638            "multinomial penalized gradient at β̂ must vanish (max |∂F| = {g})"
1639        );
1640
1641        // Fitted probabilities are a valid simplex per row and reproduce the
1642        // softmax of X β̂; deviance = −2 log L.
1643        let (n, p) = design.dim();
1644        let k = y.ncols();
1645        let m = k - 1;
1646        let mut log_lik = 0.0_f64;
1647        let mut eta_active = vec![0.0_f64; m];
1648        for row in 0..n {
1649            for a in 0..m {
1650                let mut eta = 0.0_f64;
1651                for i in 0..p {
1652                    eta += design[[row, i]] * fit.coefficients_active[[i, a]];
1653                }
1654                eta_active[a] = eta;
1655            }
1656            let probs = softmax_ref(&eta_active);
1657            let mut row_sum = 0.0_f64;
1658            for c in 0..k {
1659                assert!(
1660                    (fit.fitted_probabilities[[row, c]] - probs[c]).abs() < 1.0e-10,
1661                    "fitted probability must equal softmax(X β̂)"
1662                );
1663                row_sum += fit.fitted_probabilities[[row, c]];
1664                let yc = y[[row, c]];
1665                if yc != 0.0 {
1666                    log_lik += yc * probs[c].max(1.0e-300).ln();
1667                }
1668            }
1669            assert!(
1670                (row_sum - 1.0).abs() < 1.0e-10,
1671                "fitted probabilities must sum to 1 per row"
1672            );
1673        }
1674        assert!(
1675            (fit.deviance - (-2.0 * log_lik)).abs() < 1.0e-9,
1676            "deviance must equal −2 log L"
1677        );
1678    }
1679
1680    #[test]
1681    fn multinomial_rank_deficient_block_recovers_via_ridge_not_crash() {
1682        // Issue #557: a rank-deficient class block under a tiny per-class λ used
1683        // to make faer's Bunch–Kaufman fallback back-substitute through near-zero
1684        // pivots into a non-finite Newton step δ, and the solver aborted with
1685        // "Newton step is non-finite". The adaptive Levenberg–Marquardt ridge
1686        // must instead lift the null direction off zero, keep δ finite, and let
1687        // the backtracking line search converge to the penalized optimum.
1688        //
1689        // Construct an exactly rank-deficient design: column 2 is a perfect
1690        // duplicate of column 1, so XᵀWX is singular along (e₁ − e₂) for every
1691        // class, and we drive the corresponding λ to a tiny value so the penalty
1692        // cannot regularize that null direction. A non-robust solver crashes
1693        // here; the ridge path must produce a finite, self-consistent fit.
1694        let n = 50;
1695        let p = 4;
1696        let k = 4;
1697        let design = Array2::<f64>::from_shape_fn((n, p), |(i, j)| match j {
1698            0 => 1.0,
1699            1 => ((i + 1) as f64 * 0.23).sin(),
1700            2 => ((i + 1) as f64 * 0.23).sin(), // exact duplicate of column 1
1701            _ => ((i + 1) as f64 * 0.19).cos(),
1702        });
1703        let mut y = Array2::<f64>::zeros((n, k));
1704        for i in 0..n {
1705            y[[i, (i * 5 + 2) % k]] = 1.0;
1706        }
1707        // Penalty touches only the smooth-ish columns 1..p; columns 0/1/2 share
1708        // the collinearity, and a near-zero λ leaves the (e₁ − e₂) null direction
1709        // unregularized — exactly the rank-deficient regime that triggered #557.
1710        let mut penalty = Array2::<f64>::zeros((p, p));
1711        penalty[[3, 3]] = 1.0;
1712        // #2344: K per-class lambdas (reference class included).
1713        let lambdas = Array1::from(vec![1.0e-10_f64, 1.0e-10, 1.0e-10, 1.0e-10]);
1714
1715        let fit = fit_penalized_multinomial(MultinomialFitInputs {
1716            design: design.view(),
1717            y_one_hot: y.view(),
1718            penalty: penalty.view(),
1719            lambdas: lambdas.view(),
1720            row_weights: None,
1721            fisher_w_override: None,
1722            max_iter: 200,
1723            tol: 1.0e-10,
1724            resume_from: None,
1725        })
1726        .expect("rank-deficient multinomial fit must NOT crash (#557): the ridge path recovers it");
1727
1728        // Every coefficient and fitted probability must be finite (no inf/NaN
1729        // leaked from the near-singular solve).
1730        for &c in fit.coefficients_active.iter() {
1731            assert!(c.is_finite(), "coefficient must be finite, got {c}");
1732        }
1733        for &pr in fit.fitted_probabilities.iter() {
1734            assert!(
1735                pr.is_finite() && (-1.0e-9..=1.0 + 1.0e-9).contains(&pr),
1736                "fitted probability must be a finite simplex entry, got {pr}"
1737            );
1738        }
1739        // Rows must remain on the simplex.
1740        let (nn, kk) = fit.fitted_probabilities.dim();
1741        for row in 0..nn {
1742            let s: f64 = (0..kk).map(|c| fit.fitted_probabilities[[row, c]]).sum();
1743            assert!(
1744                (s - 1.0).abs() < 1.0e-9,
1745                "row {row} probabilities must sum to 1, got {s}"
1746            );
1747        }
1748
1749        // The recovered fit must satisfy first-order optimality of the penalized
1750        // objective along every NON-NULL coordinate. The (e₁ − e₂) null
1751        // direction is unidentified (the ridge picks the minimum-norm split
1752        // between the duplicate columns), so the gradient is exactly zero along
1753        // every identified direction; a central finite difference of F over the
1754        // full coefficient matrix is dominated by the identified part and must be
1755        // small. We assert the penalized objective gradient is near-zero — the
1756        // ridge biases the step but never the optimum (at β̂ the unridged
1757        // gradient vanishes for any τ).
1758        let g = fd_grad(&fit.coefficients_active, |b| {
1759            multinomial_objective(&design, &y, &penalty, &lambdas, b)
1760        });
1761        assert!(
1762            g < 1.0e-4,
1763            "penalized objective gradient at the ridge-recovered β̂ must (near-)vanish \
1764             along identified directions (max |∂F| = {g})"
1765        );
1766    }
1767}