regression_diagnostics/logistic/
fit.rs

1use ndarray::{Array1, Array2, ArrayView1, ArrayView2};
2use statrs::distribution::{ContinuousCDF, Normal};
3
4use crate::error::{RegressionError, Result};
5use crate::linalg::dmatrix_from_rows;
6
7/// Probability floor/ceiling used to keep IRLS weights away from exactly zero.
8const PROB_EPS: f64 = 1e-12;
9
10/// A fitted binary logistic-regression model.
11///
12/// Fit by **iteratively reweighted least squares** (equivalently, Fisher
13/// scoring / Newton–Raphson on the log-likelihood). At convergence the score
14/// equations `Xᵀ(y − p) = 0` hold, and the inverse Fisher information
15/// `(XᵀWX)⁻¹` — with `W = diag(pᵢ(1−pᵢ))` — gives the coefficient covariance the
16/// Wald statistics are built from.
17///
18/// See the [module docs](crate::logistic) for the response convention and the
19/// separation caveat.
20#[derive(Debug, Clone)]
21pub struct LogisticFit {
22    x: Array2<f64>,
23    y: Array1<f64>,
24    coefficients: Array1<f64>,
25    /// Fitted probabilities `pᵢ`.
26    probabilities: Array1<f64>,
27    /// IRLS weights `wᵢ = pᵢ(1 − pᵢ)`.
28    weights: Array1<f64>,
29    /// Coefficient covariance `(XᵀWX)⁻¹` at the MLE.
30    cov: Array2<f64>,
31    log_likelihood: f64,
32    intercept_col: Option<usize>,
33    iterations: usize,
34    n: usize,
35    p: usize,
36}
37
38impl LogisticFit {
39    /// Fit logistic regression of binary `y` on `X` (default: up to 100 IRLS
40    /// iterations, tolerance `1e-10` on the coefficient step).
41    ///
42    /// # Errors
43    ///
44    /// * [`RegressionError::EmptyInput`] / [`RegressionError::ShapeMismatch`].
45    /// * [`RegressionError::InvalidResponse`] if `y` is not `0/1` or is all one
46    ///   class.
47    /// * [`RegressionError::RankDeficient`] if the (weighted) design is singular.
48    /// * [`RegressionError::NotConverged`] if IRLS fails to converge (typically
49    ///   separation).
50    pub fn new(x: Array2<f64>, y: Array1<f64>) -> Result<Self> {
51        Self::with_options(x, y, 100, 1e-10)
52    }
53
54    /// Like [`LogisticFit::new`] with an explicit iteration cap and tolerance.
55    pub fn with_options(x: Array2<f64>, y: Array1<f64>, max_iter: usize, tol: f64) -> Result<Self> {
56        let n = x.nrows();
57        let p = x.ncols();
58        if n == 0 || p == 0 {
59            return Err(RegressionError::EmptyInput { what: "X" });
60        }
61        if y.len() != n {
62            return Err(RegressionError::ShapeMismatch {
63                what: "y length vs X rows",
64                expected: n,
65                got: y.len(),
66            });
67        }
68        // Validate binary response with both classes present.
69        let mut saw0 = false;
70        let mut saw1 = false;
71        for &v in y.iter() {
72            if v == 0.0 {
73                saw0 = true;
74            } else if v == 1.0 {
75                saw1 = true;
76            } else {
77                return Err(RegressionError::InvalidResponse {
78                    msg: format!("response must be 0 or 1, found {v}"),
79                });
80            }
81        }
82        if !(saw0 && saw1) {
83            return Err(RegressionError::InvalidResponse {
84                msg: "response is entirely one class; the fit is not identifiable".into(),
85            });
86        }
87
88        let intercept_col = detect_constant_column(&x);
89
90        let mut beta = Array1::<f64>::zeros(p);
91        let mut probabilities = Array1::<f64>::zeros(n);
92        let mut weights = Array1::<f64>::zeros(n);
93        let mut cov = Array2::<f64>::zeros((p, p));
94        let mut iterations = 0usize;
95        let mut converged = false;
96
97        while iterations < max_iter {
98            iterations += 1;
99
100            // η = Xβ, p = σ(η), w = p(1−p), clamped away from 0/1.
101            let eta = x.dot(&beta);
102            for i in 0..n {
103                let pi = sigmoid(eta[i]).clamp(PROB_EPS, 1.0 - PROB_EPS);
104                probabilities[i] = pi;
105                weights[i] = pi * (1.0 - pi);
106            }
107
108            // Gradient Xᵀ(y − p) and Fisher information XᵀWX.
109            let resid = &y - &probabilities;
110            let grad = x.t().dot(&resid); // length p
111            let mut xtwx = Array2::<f64>::zeros((p, p));
112            for a in 0..p {
113                for b in a..p {
114                    let mut s = 0.0;
115                    for i in 0..n {
116                        s += x[(i, a)] * weights[i] * x[(i, b)];
117                    }
118                    xtwx[(a, b)] = s;
119                    xtwx[(b, a)] = s;
120                }
121            }
122
123            let xtwx_dm = dmatrix_from_rows(p, p, xtwx.as_standard_layout().as_slice().unwrap());
124            let inv = xtwx_dm
125                .try_inverse()
126                .ok_or(RegressionError::RankDeficient)?;
127            let inv_arr = Array2::from_shape_fn((p, p), |(i, j)| inv[(i, j)]);
128
129            // Newton step Δ = (XᵀWX)⁻¹ Xᵀ(y − p).
130            let delta = inv_arr.dot(&grad);
131            beta = &beta + &delta;
132            cov = inv_arr;
133
134            let step = delta.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
135            if !beta.iter().all(|v| v.is_finite()) || beta.iter().any(|v| v.abs() > 1e8) {
136                return Err(RegressionError::NotConverged {
137                    iterations,
138                    msg: "coefficients diverging (likely perfect separation)".into(),
139                });
140            }
141            if step < tol {
142                converged = true;
143                break;
144            }
145        }
146
147        if !converged {
148            return Err(RegressionError::NotConverged {
149                iterations,
150                msg: "IRLS did not reach tolerance (possible quasi-separation)".into(),
151            });
152        }
153
154        // Final probabilities and log-likelihood at the converged β.
155        let eta = x.dot(&beta);
156        for i in 0..n {
157            let pi = sigmoid(eta[i]).clamp(PROB_EPS, 1.0 - PROB_EPS);
158            probabilities[i] = pi;
159            weights[i] = pi * (1.0 - pi);
160        }
161        let log_likelihood = (0..n)
162            .map(|i| {
163                let pi = probabilities[i];
164                y[i] * pi.ln() + (1.0 - y[i]) * (1.0 - pi).ln()
165            })
166            .sum();
167
168        Ok(Self {
169            x,
170            y,
171            coefficients: beta,
172            probabilities,
173            weights,
174            cov,
175            log_likelihood,
176            intercept_col,
177            iterations,
178            n,
179            p,
180        })
181    }
182
183    /// Number of observations.
184    pub fn n_observations(&self) -> usize {
185        self.n
186    }
187
188    /// Number of coefficients (design columns, intercept included).
189    pub fn n_parameters(&self) -> usize {
190        self.p
191    }
192
193    /// Whether an intercept (constant) column is present.
194    pub fn has_intercept(&self) -> bool {
195        self.intercept_col.is_some()
196    }
197
198    /// IRLS iterations taken to converge.
199    pub fn iterations(&self) -> usize {
200        self.iterations
201    }
202
203    /// The design matrix as fitted.
204    pub fn design_matrix(&self) -> ArrayView2<'_, f64> {
205        self.x.view()
206    }
207
208    /// The binary response.
209    pub fn response(&self) -> ArrayView1<'_, f64> {
210        self.y.view()
211    }
212
213    /// Estimated coefficients (log-odds scale), aligned to the design columns.
214    pub fn coefficients(&self) -> ArrayView1<'_, f64> {
215        self.coefficients.view()
216    }
217
218    /// Fitted probabilities `pᵢ = P(yᵢ = 1)`.
219    pub fn fitted_probabilities(&self) -> ArrayView1<'_, f64> {
220        self.probabilities.view()
221    }
222
223    /// IRLS weights `wᵢ = pᵢ(1 − pᵢ)` at the MLE.
224    pub fn weights(&self) -> ArrayView1<'_, f64> {
225        self.weights.view()
226    }
227
228    /// Coefficient covariance matrix `(XᵀWX)⁻¹`.
229    pub fn covariance(&self) -> ArrayView2<'_, f64> {
230        self.cov.view()
231    }
232
233    /// Maximized log-likelihood.
234    pub fn log_likelihood(&self) -> f64 {
235        self.log_likelihood
236    }
237
238    /// Coefficient standard errors `√diag((XᵀWX)⁻¹)`.
239    pub fn coefficient_standard_errors(&self) -> Array1<f64> {
240        Array1::from_shape_fn(self.p, |j| self.cov[(j, j)].max(0.0).sqrt())
241    }
242
243    /// Wald `z`-statistics `βⱼ / seⱼ`.
244    pub fn z_values(&self) -> Array1<f64> {
245        let se = self.coefficient_standard_errors();
246        Array1::from_shape_fn(self.p, |j| {
247            if se[j] > 0.0 {
248                self.coefficients[j] / se[j]
249            } else {
250                f64::NAN
251            }
252        })
253    }
254
255    /// Two-sided Wald p-values from the standard normal.
256    pub fn p_values(&self) -> Array1<f64> {
257        let z = self.z_values();
258        let normal = Normal::new(0.0, 1.0).expect("standard normal");
259        Array1::from_shape_fn(self.p, |j| {
260            if z[j].is_finite() {
261                2.0 * (1.0 - normal.cdf(z[j].abs()))
262            } else {
263                f64::NAN
264            }
265        })
266    }
267
268    /// Predicted probabilities for a new design matrix `x` (same column layout
269    /// as the training design).
270    pub fn predict_proba(&self, x: ArrayView2<'_, f64>) -> Array1<f64> {
271        x.dot(&self.coefficients).mapv(sigmoid)
272    }
273}
274
275fn sigmoid(z: f64) -> f64 {
276    if z >= 0.0 {
277        1.0 / (1.0 + (-z).exp())
278    } else {
279        let e = z.exp();
280        e / (1.0 + e)
281    }
282}
283
284/// Detect the first constant column (treated as the intercept).
285fn detect_constant_column(x: &Array2<f64>) -> Option<usize> {
286    for (j, col) in x.columns().into_iter().enumerate() {
287        let first = col[0];
288        let scale = first.abs().max(1.0);
289        if col.iter().all(|&v| (v - first).abs() <= 1e-12 * scale) {
290            return Some(j);
291        }
292    }
293    None
294}
regression_diagnostics/logistic/fit.rs

regression_diagnostics/logistic/
fit.rs
