Skip to main content

rill_ml/loss/
mod.rs

1//! Loss functions for online models.
2//!
3//! Losses are represented as concrete enums, not trait objects, to keep
4//! serialization and state management simple.
5
6pub mod huber;
7pub mod log_loss;
8pub mod squared;
9
10pub use huber::HuberLoss;
11pub use log_loss::BinaryLogLoss;
12pub use squared::SquaredError;
13
14/// Regression loss variants.
15///
16/// Used by [`LinearRegression`](crate::models::LinearRegression) to select
17/// the loss function applied to each update.
18#[derive(Debug, Clone, Default)]
19#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
20pub enum RegressionLoss {
21    /// `0.5 * (y - y_hat)^2`
22    #[default]
23    SquaredError,
24    /// Huber loss, robust to outliers.
25    Huber(HuberLoss),
26}
27
28impl RegressionLoss {
29    /// Compute the loss value given a prediction and target.
30    pub fn loss(&self, prediction: f64, target: f64) -> f64 {
31        match self {
32            RegressionLoss::SquaredError => SquaredError::loss(prediction, target),
33            RegressionLoss::Huber(h) => h.loss(prediction, target),
34        }
35    }
36
37    /// Compute the derivative of the loss with respect to the prediction.
38    pub fn gradient(&self, prediction: f64, target: f64) -> f64 {
39        match self {
40            RegressionLoss::SquaredError => SquaredError::gradient(prediction, target),
41            RegressionLoss::Huber(h) => h.gradient(prediction, target),
42        }
43    }
44}