regression_diagnostics/regularized/
penalized_glm.rs1use 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#[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 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 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 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 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 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 + δ
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 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 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 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 pub fn lambda(&self) -> f64 {
234 self.lambda
235 }
236
237 pub fn n_observations(&self) -> usize {
239 self.n
240 }
241
242 pub fn n_parameters(&self) -> usize {
244 self.p
245 }
246
247 pub fn has_intercept(&self) -> bool {
249 self.intercept_col.is_some()
250 }
251
252 pub fn iterations(&self) -> usize {
254 self.iterations
255 }
256
257 pub fn design_matrix(&self) -> ArrayView2<'_, f64> {
259 self.x.view()
260 }
261
262 pub fn response(&self) -> ArrayView1<'_, f64> {
264 self.y.view()
265 }
266
267 pub fn coefficients(&self) -> ArrayView1<'_, f64> {
269 self.coefficients.view()
270 }
271
272 pub fn fitted_probabilities(&self) -> ArrayView1<'_, f64> {
274 self.probabilities.view()
275 }
276
277 pub fn covariance(&self) -> ArrayView2<'_, f64> {
279 self.cov.view()
280 }
281
282 pub fn log_likelihood(&self) -> f64 {
284 self.log_likelihood
285 }
286
287 pub fn effective_df(&self) -> f64 {
290 self.effective_df
291 }
292
293 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 pub fn residual_deviance(&self) -> f64 {
300 -2.0 * self.log_likelihood
301 }
302
303 pub fn aic(&self) -> f64 {
305 -2.0 * self.log_likelihood + 2.0 * self.effective_df
306 }
307
308 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}