rill_ml/metrics/
regression.rs

1//! Regression metrics: MAE, MSE, RMSE, R².
2
3use crate::error::{
4    RillError, checked_finite_add, checked_increment, ensure_finite, ensure_finite_target,
5};
6use crate::traits::Metric;
7
8/// Mean Absolute Error.
9#[derive(Debug, Clone, Default)]
10#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
11pub struct Mae {
12    sum_abs_error: f64,
13    count: u64,
14}
15
16impl Mae {
17    /// Create a new MAE accumulator.
18    pub const fn new() -> Self {
19        Self {
20            sum_abs_error: 0.0,
21            count: 0,
22        }
23    }
24}
25
26impl Metric for Mae {
27    type Truth = f64;
28    type Prediction = f64;
29
30    fn update(&mut self, truth: f64, prediction: f64) -> Result<(), RillError> {
31        ensure_finite_target(truth)?;
32        ensure_finite("prediction", prediction)?;
33        let error = truth - prediction;
34        ensure_finite("absolute error", error)?;
35        let next_sum = checked_finite_add(self.sum_abs_error, error.abs(), "MAE sum")?;
36        let next_count = checked_increment(self.count, "MAE sample")?;
37        self.sum_abs_error = next_sum;
38        self.count = next_count;
39        Ok(())
40    }
41
42    fn value(&self) -> Option<f64> {
43        if self.count == 0 {
44            None
45        } else {
46            Some(self.sum_abs_error / self.count as f64)
47        }
48    }
49
50    fn samples_seen(&self) -> u64 {
51        self.count
52    }
53
54    fn reset(&mut self) {
55        self.sum_abs_error = 0.0;
56        self.count = 0;
57    }
58}
59
60/// Mean Squared Error.
61#[derive(Debug, Clone, Default)]
62#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
63pub struct Mse {
64    sum_sq_error: f64,
65    count: u64,
66}
67
68impl Mse {
69    /// Create a new MSE accumulator.
70    pub const fn new() -> Self {
71        Self {
72            sum_sq_error: 0.0,
73            count: 0,
74        }
75    }
76}
77
78impl Metric for Mse {
79    type Truth = f64;
80    type Prediction = f64;
81
82    fn update(&mut self, truth: f64, prediction: f64) -> Result<(), RillError> {
83        ensure_finite_target(truth)?;
84        ensure_finite("prediction", prediction)?;
85        let err = truth - prediction;
86        ensure_finite("squared error input", err)?;
87        let squared_error = err * err;
88        ensure_finite("squared error", squared_error)?;
89        let next_sum = checked_finite_add(self.sum_sq_error, squared_error, "MSE sum")?;
90        let next_count = checked_increment(self.count, "MSE sample")?;
91        self.sum_sq_error = next_sum;
92        self.count = next_count;
93        Ok(())
94    }
95
96    fn value(&self) -> Option<f64> {
97        if self.count == 0 {
98            None
99        } else {
100            Some(self.sum_sq_error / self.count as f64)
101        }
102    }
103
104    fn samples_seen(&self) -> u64 {
105        self.count
106    }
107
108    fn reset(&mut self) {
109        self.sum_sq_error = 0.0;
110        self.count = 0;
111    }
112}
113
114/// Root Mean Squared Error.
115#[derive(Debug, Clone, Default)]
116#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
117pub struct Rmse {
118    mse: Mse,
119}
120
121impl Rmse {
122    /// Create a new RMSE accumulator.
123    pub const fn new() -> Self {
124        Self { mse: Mse::new() }
125    }
126}
127
128impl Metric for Rmse {
129    type Truth = f64;
130    type Prediction = f64;
131
132    fn update(&mut self, truth: f64, prediction: f64) -> Result<(), RillError> {
133        self.mse.update(truth, prediction)
134    }
135
136    fn value(&self) -> Option<f64> {
137        self.mse.value().map(|v| v.sqrt())
138    }
139
140    fn samples_seen(&self) -> u64 {
141        self.mse.samples_seen()
142    }
143
144    fn reset(&mut self) {
145        self.mse.reset();
146    }
147}
148
149/// R² (coefficient of determination).
150///
151/// Computed online as `1 - SS_res / SS_tot`, where `SS_tot` uses the running
152/// mean of the truth. Returns `None` when fewer than 2 samples have been seen
153/// or when `SS_tot` is zero (constant truth).
154#[derive(Debug, Clone, Default)]
155#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
156pub struct R2 {
157    ss_res: f64,
158    sum_truth: f64,
159    sum_truth_sq: f64,
160    count: u64,
161}
162
163impl R2 {
164    /// Create a new R² accumulator.
165    pub const fn new() -> Self {
166        Self {
167            ss_res: 0.0,
168            sum_truth: 0.0,
169            sum_truth_sq: 0.0,
170            count: 0,
171        }
172    }
173}
174
175impl Metric for R2 {
176    type Truth = f64;
177    type Prediction = f64;
178
179    fn update(&mut self, truth: f64, prediction: f64) -> Result<(), RillError> {
180        ensure_finite_target(truth)?;
181        ensure_finite("prediction", prediction)?;
182        let err = truth - prediction;
183        ensure_finite("R2 error", err)?;
184        let squared_error = err * err;
185        let squared_truth = truth * truth;
186        ensure_finite("R2 squared error", squared_error)?;
187        ensure_finite("R2 squared truth", squared_truth)?;
188        let next_ss_res = checked_finite_add(self.ss_res, squared_error, "R2 residual sum")?;
189        let next_sum_truth = checked_finite_add(self.sum_truth, truth, "R2 truth sum")?;
190        let next_sum_truth_sq =
191            checked_finite_add(self.sum_truth_sq, squared_truth, "R2 squared truth sum")?;
192        let next_count = checked_increment(self.count, "R2 sample")?;
193
194        self.ss_res = next_ss_res;
195        self.sum_truth = next_sum_truth;
196        self.sum_truth_sq = next_sum_truth_sq;
197        self.count = next_count;
198        Ok(())
199    }
200
201    fn value(&self) -> Option<f64> {
202        if self.count < 2 {
203            return None;
204        }
205        let n = self.count as f64;
206        let mean = self.sum_truth / n;
207        let ss_tot = self.sum_truth_sq - n * mean * mean;
208        if ss_tot.abs() < f64::EPSILON {
209            return None;
210        }
211        Some(1.0 - self.ss_res / ss_tot)
212    }
213
214    fn samples_seen(&self) -> u64 {
215        self.count
216    }
217
218    fn reset(&mut self) {
219        self.ss_res = 0.0;
220        self.sum_truth = 0.0;
221        self.sum_truth_sq = 0.0;
222        self.count = 0;
223    }
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229
230    #[test]
231    fn mae_basic() {
232        let mut m = Mae::new();
233        m.update(3.0, 5.0).unwrap(); // err=2
234        m.update(5.0, 4.0).unwrap(); // err=1
235        assert!((m.value().unwrap() - 1.5).abs() < 1e-12);
236    }
237
238    #[test]
239    fn mse_basic() {
240        let mut m = Mse::new();
241        m.update(3.0, 5.0).unwrap(); // err=2, sq=4
242        m.update(5.0, 4.0).unwrap(); // err=1, sq=1
243        assert!((m.value().unwrap() - 2.5).abs() < 1e-12);
244    }
245
246    #[test]
247    fn metrics_reject_overflow_without_mutating_state() {
248        let mut mae = Mae::new();
249        let mut mse = Mse::new();
250        let mut r2 = R2::new();
251
252        assert!(mae.update(f64::MAX, -f64::MAX).is_err());
253        assert!(mse.update(f64::MAX, 0.0).is_err());
254        assert!(r2.update(f64::MAX, 0.0).is_err());
255
256        assert_eq!(mae.samples_seen(), 0);
257        assert_eq!(mse.samples_seen(), 0);
258        assert_eq!(r2.samples_seen(), 0);
259    }
260
261    #[test]
262    fn rmse_basic() {
263        let mut m = Rmse::new();
264        m.update(3.0, 5.0).unwrap();
265        m.update(5.0, 4.0).unwrap();
266        assert!((m.value().unwrap() - 2.5_f64.sqrt()).abs() < 1e-12);
267    }
268
269    #[test]
270    fn r2_perfect_prediction_is_one() {
271        let mut m = R2::new();
272        m.update(1.0, 1.0).unwrap();
273        m.update(2.0, 2.0).unwrap();
274        m.update(3.0, 3.0).unwrap();
275        assert!((m.value().unwrap() - 1.0).abs() < 1e-9);
276    }
277
278    #[test]
279    fn r2_mean_prediction_is_zero() {
280        let mut m = R2::new();
281        // predict the mean every time
282        m.update(1.0, 2.0).unwrap();
283        m.update(3.0, 2.0).unwrap();
284        // mean=2, ss_res = 1+1=2, ss_tot = 1+1=2 -> R2=0
285        assert!((m.value().unwrap()).abs() < 1e-9);
286    }
287
288    #[test]
289    fn r2_insufficient_data_returns_none() {
290        let mut m = R2::new();
291        m.update(1.0, 1.0).unwrap();
292        assert!(m.value().is_none());
293    }
294
295    #[test]
296    fn r2_constant_truth_returns_none() {
297        let mut m = R2::new();
298        m.update(5.0, 3.0).unwrap();
299        m.update(5.0, 4.0).unwrap();
300        assert!(m.value().is_none());
301    }
302
303    #[test]
304    fn non_finite_rejected() {
305        let mut m = Mae::new();
306        assert!(m.update(f64::NAN, 1.0).is_err());
307        assert!(m.update(1.0, f64::INFINITY).is_err());
308    }
309
310    #[test]
311    fn empty_metric_returns_none() {
312        assert!(Mae::new().value().is_none());
313        assert!(Mse::new().value().is_none());
314        assert!(Rmse::new().value().is_none());
315        assert!(R2::new().value().is_none());
316    }
317}
rill_ml/metrics/regression.rs

rill_ml/metrics/
regression.rs
