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(crate) mod huber;
7pub(crate) mod log_loss;
8pub(crate) 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))]
20#[non_exhaustive]
21pub enum RegressionLoss {
22    /// `0.5 * (y - y_hat)^2`
23    #[default]
24    SquaredError,
25    /// Huber loss, robust to outliers.
26    Huber(HuberLoss),
27}
28
29impl RegressionLoss {
30    /// Compute the loss value given a prediction and target.
31    pub fn loss(&self, prediction: f64, target: f64) -> f64 {
32        match self {
33            RegressionLoss::SquaredError => SquaredError::loss(prediction, target),
34            RegressionLoss::Huber(h) => h.loss(prediction, target),
35        }
36    }
37
38    /// Compute the derivative of the loss with respect to the prediction.
39    pub fn gradient(&self, prediction: f64, target: f64) -> f64 {
40        match self {
41            RegressionLoss::SquaredError => SquaredError::gradient(prediction, target),
42            RegressionLoss::Huber(h) => h.gradient(prediction, target),
43        }
44    }
45}