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}