rustframe/src/compute/models/logreg.rs

use crate::matrix::{Matrix, SeriesOps};
use crate::compute::activations::sigmoid;

pub struct LogReg {
    w: Matrix<f64>,
    b: f64,
}

impl LogReg {
    pub fn new(n_features: usize) -> Self {
        Self {
            w: Matrix::zeros(n_features, 1),
            b: 0.0,
        }
    }

    pub fn predict_proba(&self, x: &Matrix<f64>) -> Matrix<f64> {
        sigmoid(&(x.dot(&self.w) + self.b))          // σ(Xw + b)
    }

    pub fn fit(&mut self, x: &Matrix<f64>, y: &Matrix<f64>, lr: f64, epochs: usize) {
        let m = x.rows() as f64;
        for _ in 0..epochs {
            let p  = self.predict_proba(x);          // shape (m,1)
            let err = &p - y;                        // derivative of BCE wrt pre-sigmoid
            let grad_w = x.transpose().dot(&err) / m;
            let grad_b = err.sum_vertical().iter().sum::<f64>() / m;
            self.w = &self.w - &(grad_w * lr);
            self.b -= lr * grad_b;
        }
    }

    pub fn predict(&self, x: &Matrix<f64>) -> Matrix<f64> {
        self.predict_proba(x).map(|p| if p >= 0.5 { 1.0 } else { 0.0 })
    }
}