From 54a266b630e1ccd2a86e1ba935952825b62191f3 Mon Sep 17 00:00:00 2001
From: Palash Tyagi <23239946+Magnus167@users.noreply.github.com>
Date: Sun, 6 Jul 2025 18:52:49 +0100
Subject: [PATCH] Add unit tests for logistic regression fit and predict
 methods

---
 src/compute/models/logreg.rs | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/src/compute/models/logreg.rs b/src/compute/models/logreg.rs
index 1bdaa86..473d5e9 100644
--- a/src/compute/models/logreg.rs
+++ b/src/compute/models/logreg.rs
@@ -34,3 +34,22 @@ impl LogReg {
         self.predict_proba(x).map(|p| if p >= 0.5 { 1.0 } else { 0.0 })
     }
 }
+
+
+mod tests {
+    use super::LogReg;
+    use crate::matrix::Matrix;
+
+    #[test]
+    fn test_logreg_fit_predict() {
+        let x = Matrix::from_vec(vec![1.0, 2.0, 3.0, 4.0], 4, 1);
+        let y = Matrix::from_vec(vec![0.0, 0.0, 1.0, 1.0], 4, 1);
+        let mut model = LogReg::new(1);
+        model.fit(&x, &y, 0.01, 10000);
+        let preds = model.predict(&x);
+        assert_eq!(preds[(0, 0)], 0.0);
+        assert_eq!(preds[(1, 0)], 0.0);
+        assert_eq!(preds[(2, 0)], 1.0);
+        assert_eq!(preds[(3, 0)], 1.0);
+    }
+}
\ No newline at end of file