Skip to main content

regression_diagnostics/regularized/
penalized_glm.rs

1use ndarray::{Array1, Array2, ArrayView1, ArrayView2};
2
3use crate::error::{RegressionError, Result};
4use crate::linalg::dmatrix_from_rows;
5
6const PROB_EPS: f64 = 1e-12;
7
8/// A fitted **ridge-penalized logistic regression** — a penalized GLM — and its
9/// shrinkage-aware diagnostics.
10///
11/// Maximizes the L2-penalized log-likelihood
12///
13/// `ℓ(β) − ½λ‖β_pen‖²`
14///
15/// (a detected intercept column is left unpenalized), fit by **penalized IRLS**:
16/// each Newton step solves `(XᵀWX + λP) Δ = Xᵀ(y − p) − λPβ` with
17/// `W = diag(pᵢ(1 − pᵢ))` and `P` the penalty selector. The penalty is the GLM
18/// analogue of ridge — it tames separation and multicollinearity in logistic
19/// regression, at the cost of biased-but-lower-variance coefficients.
20///
21/// # The diagnostic that changes: effective degrees of freedom
22///
23/// As in ridge OLS, the penalty means the model no longer spends `p` degrees of
24/// freedom. The **effective df** is the trace of the penalized hat matrix,
25///
26/// `df = tr[ (XᵀWX + λP)⁻¹ XᵀWX ]`,
27///
28/// which falls from `p` toward the unpenalized/intercept count as `λ` grows. It
29/// replaces the raw parameter count in [`aic`](Self::aic) / [`bic`](Self::bic),
30/// and the coefficient covariance is the **sandwich**
31/// `(XᵀWX + λP)⁻¹ (XᵀWX) (XᵀWX + λP)⁻¹`, not the naïve inverse information — so
32/// the reported standard errors account for the shrinkage.
33#[derive(Debug, Clone)]
34pub struct PenalizedLogisticFit {
35    x: Array2<f64>,
36    y: Array1<f64>,
37    lambda: f64,
38    coefficients: Array1<f64>,
39    probabilities: Array1<f64>,
40    cov: Array2<f64>,
41    log_likelihood: f64,
42    effective_df: f64,
43    intercept_col: Option<usize>,
44    iterations: usize,
45    n: usize,
46    p: usize,
47}
48
49impl PenalizedLogisticFit {
50    /// Fit ridge-penalized logistic regression of binary `y` on `X` with penalty
51    /// `lambda ≥ 0` (default: up to 100 IRLS iterations, tolerance `1e-10`).
52    ///
53    /// At `lambda = 0` this reproduces the ordinary
54    /// [`LogisticFit`](crate::logistic::LogisticFit).
55    ///
56    /// # Errors
57    ///
58    /// * [`RegressionError::EmptyInput`] / [`RegressionError::ShapeMismatch`].
59    /// * [`RegressionError::InvalidParameter`] if `lambda < 0`.
60    /// * [`RegressionError::InvalidResponse`] if `y` is not `0/1` or is one class.
61    /// * [`RegressionError::RankDeficient`] if the penalized information is
62    ///   singular.
63    /// * [`RegressionError::NotConverged`] if IRLS fails to converge.
64    pub fn new(x: Array2<f64>, y: Array1<f64>, lambda: f64) -> Result<Self> {
65        Self::with_options(x, y, lambda, 100, 1e-10)
66    }
67
68    /// Like [`PenalizedLogisticFit::new`] with an explicit iteration cap and
69    /// tolerance.
70    pub fn with_options(
71        x: Array2<f64>,
72        y: Array1<f64>,
73        lambda: f64,
74        max_iter: usize,
75        tol: f64,
76    ) -> Result<Self> {
77        let n = x.nrows();
78        let p = x.ncols();
79        if n == 0 || p == 0 {
80            return Err(RegressionError::EmptyInput { what: "X" });
81        }
82        if y.len() != n {
83            return Err(RegressionError::ShapeMismatch {
84                what: "y length vs X rows",
85                expected: n,
86                got: y.len(),
87            });
88        }
89        if lambda < 0.0 || lambda.is_nan() {
90            return Err(RegressionError::InvalidParameter {
91                msg: format!("penalty lambda must be >= 0, got {lambda}"),
92            });
93        }
94        let (mut saw0, mut saw1) = (false, false);
95        for &v in y.iter() {
96            if v == 0.0 {
97                saw0 = true;
98            } else if v == 1.0 {
99                saw1 = true;
100            } else {
101                return Err(RegressionError::InvalidResponse {
102                    msg: format!("response must be 0 or 1, found {v}"),
103                });
104            }
105        }
106        if !(saw0 && saw1) {
107            return Err(RegressionError::InvalidResponse {
108                msg: "response is entirely one class".into(),
109            });
110        }
111
112        let intercept_col = detect_constant_column(&x);
113        // Penalty selector: 1 for penalized columns, 0 for the intercept.
114        let pen: Vec<f64> = (0..p)
115            .map(|j| if Some(j) == intercept_col { 0.0 } else { 1.0 })
116            .collect();
117
118        let mut beta = Array1::<f64>::zeros(p);
119        let mut probabilities = Array1::<f64>::zeros(n);
120        let mut weights = Array1::<f64>::zeros(n);
121        let mut xtwx_pen_inv = Array2::<f64>::zeros((p, p));
122        let mut iterations = 0usize;
123        let mut converged = false;
124
125        while iterations < max_iter {
126            iterations += 1;
127            let eta = x.dot(&beta);
128            for i in 0..n {
129                let pi = sigmoid(eta[i]).clamp(PROB_EPS, 1.0 - PROB_EPS);
130                probabilities[i] = pi;
131                weights[i] = pi * (1.0 - pi);
132            }
133
134            // Penalized gradient g = Xᵀ(y − p) − λ P β.
135            let resid = &y - &probabilities;
136            let mut grad = x.t().dot(&resid);
137            for j in 0..p {
138                grad[j] -= lambda * pen[j] * beta[j];
139            }
140            // Penalized information A = XᵀWX + λ P.
141            let mut a = Array2::<f64>::zeros((p, p));
142            for r in 0..p {
143                for c in r..p {
144                    let mut s = 0.0;
145                    for i in 0..n {
146                        s += x[(i, r)] * weights[i] * x[(i, c)];
147                    }
148                    a[(r, c)] = s;
149                    a[(c, r)] = s;
150                }
151            }
152            for j in 0..p {
153                a[(j, j)] += lambda * pen[j];
154            }
155
156            let a_dm = dmatrix_from_rows(p, p, a.as_standard_layout().as_slice().unwrap());
157            let inv = a_dm.try_inverse().ok_or(RegressionError::RankDeficient)?;
158            let inv_arr = Array2::from_shape_fn((p, p), |(i, j)| inv[(i, j)]);
159
160            let delta = inv_arr.dot(&grad);
161            beta = &beta + &delta;
162            xtwx_pen_inv = inv_arr;
163
164            let step = delta.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
165            if !beta.iter().all(|v| v.is_finite()) {
166                return Err(RegressionError::NotConverged {
167                    iterations,
168                    msg: "coefficients diverging".into(),
169                });
170            }
171            if step < tol {
172                converged = true;
173                break;
174            }
175        }
176        if !converged {
177            return Err(RegressionError::NotConverged {
178                iterations,
179                msg: "penalized IRLS did not reach tolerance".into(),
180            });
181        }
182
183        // Final probabilities, weights, and the unpenalized XᵀWX.
184        let eta = x.dot(&beta);
185        for i in 0..n {
186            let pi = sigmoid(eta[i]).clamp(PROB_EPS, 1.0 - PROB_EPS);
187            probabilities[i] = pi;
188            weights[i] = pi * (1.0 - pi);
189        }
190        let mut xtwx = Array2::<f64>::zeros((p, p));
191        for r in 0..p {
192            for c in r..p {
193                let mut s = 0.0;
194                for i in 0..n {
195                    s += x[(i, r)] * weights[i] * x[(i, c)];
196                }
197                xtwx[(r, c)] = s;
198                xtwx[(c, r)] = s;
199            }
200        }
201        // effective df = tr[(XᵀWX + λP)⁻¹ XᵀWX].
202        let effective_df = (0..p)
203            .map(|i| (0..p).map(|kk| xtwx_pen_inv[(i, kk)] * xtwx[(kk, i)]).sum::<f64>())
204            .sum();
205        // Sandwich covariance A⁻¹ (XᵀWX) A⁻¹.
206        let mid = xtwx.dot(&xtwx_pen_inv);
207        let cov = xtwx_pen_inv.dot(&mid);
208
209        let log_likelihood = (0..n)
210            .map(|i| {
211                let pi = probabilities[i];
212                y[i] * pi.ln() + (1.0 - y[i]) * (1.0 - pi).ln()
213            })
214            .sum();
215
216        Ok(Self {
217            x,
218            y,
219            lambda,
220            coefficients: beta,
221            probabilities,
222            cov,
223            log_likelihood,
224            effective_df,
225            intercept_col,
226            iterations,
227            n,
228            p,
229        })
230    }
231
232    /// The penalty `λ` this model was fit with.
233    pub fn lambda(&self) -> f64 {
234        self.lambda
235    }
236
237    /// Number of observations.
238    pub fn n_observations(&self) -> usize {
239        self.n
240    }
241
242    /// Number of coefficients (design columns, intercept included).
243    pub fn n_parameters(&self) -> usize {
244        self.p
245    }
246
247    /// Whether an unpenalized intercept is present.
248    pub fn has_intercept(&self) -> bool {
249        self.intercept_col.is_some()
250    }
251
252    /// IRLS iterations taken to converge.
253    pub fn iterations(&self) -> usize {
254        self.iterations
255    }
256
257    /// The design matrix as fitted.
258    pub fn design_matrix(&self) -> ArrayView2<'_, f64> {
259        self.x.view()
260    }
261
262    /// The binary response.
263    pub fn response(&self) -> ArrayView1<'_, f64> {
264        self.y.view()
265    }
266
267    /// Penalized coefficients (log-odds scale), aligned to the design columns.
268    pub fn coefficients(&self) -> ArrayView1<'_, f64> {
269        self.coefficients.view()
270    }
271
272    /// Fitted probabilities `pᵢ = P(yᵢ = 1)`.
273    pub fn fitted_probabilities(&self) -> ArrayView1<'_, f64> {
274        self.probabilities.view()
275    }
276
277    /// Sandwich coefficient covariance `(XᵀWX + λP)⁻¹ (XᵀWX) (XᵀWX + λP)⁻¹`.
278    pub fn covariance(&self) -> ArrayView2<'_, f64> {
279        self.cov.view()
280    }
281
282    /// Unpenalized log-likelihood at the penalized estimate.
283    pub fn log_likelihood(&self) -> f64 {
284        self.log_likelihood
285    }
286
287    /// Effective degrees of freedom `tr[(XᵀWX + λP)⁻¹ XᵀWX]` — falls from `p`
288    /// toward the intercept count as `λ` grows.
289    pub fn effective_df(&self) -> f64 {
290        self.effective_df
291    }
292
293    /// Coefficient standard errors from the sandwich covariance.
294    pub fn coefficient_standard_errors(&self) -> Array1<f64> {
295        Array1::from_shape_fn(self.p, |j| self.cov[(j, j)].max(0.0).sqrt())
296    }
297
298    /// Residual deviance `−2ℓ` (unpenalized log-likelihood).
299    pub fn residual_deviance(&self) -> f64 {
300        -2.0 * self.log_likelihood
301    }
302
303    /// AIC using the effective degrees of freedom, `−2ℓ + 2·df`.
304    pub fn aic(&self) -> f64 {
305        -2.0 * self.log_likelihood + 2.0 * self.effective_df
306    }
307
308    /// BIC using the effective degrees of freedom, `−2ℓ + ln(n)·df`.
309    pub fn bic(&self) -> f64 {
310        -2.0 * self.log_likelihood + (self.n as f64).ln() * self.effective_df
311    }
312}
313
314fn sigmoid(z: f64) -> f64 {
315    if z >= 0.0 {
316        1.0 / (1.0 + (-z).exp())
317    } else {
318        let e = z.exp();
319        e / (1.0 + e)
320    }
321}
322
323fn detect_constant_column(x: &Array2<f64>) -> Option<usize> {
324    for (j, col) in x.columns().into_iter().enumerate() {
325        let first = col[0];
326        let scale = first.abs().max(1.0);
327        if col.iter().all(|&v| (v - first).abs() <= 1e-12 * scale) {
328            return Some(j);
329        }
330    }
331    None
332}