Skip to main content

rill_ml/optim/
mod.rs

1//! Optimizers for online linear models.
2//!
3//! Optimizers are represented as a concrete enum to avoid trait-object
4//! overhead and simplify serialization. The internal parameter vector has
5//! length `feature_count + 1`, where the last position holds the intercept.
6
7pub(crate) mod adagrad;
8pub(crate) mod sgd;
9
10pub use adagrad::{AdaGrad, AdaGradConfig};
11pub use sgd::{Sgd, SgdConfig};
12
13/// Concrete optimizer enum wrapping all supported optimizers.
14#[derive(Debug, Clone)]
15#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
16#[non_exhaustive]
17pub enum Optimizer {
18    /// Stochastic gradient descent with optional L2 regularization.
19    Sgd(Sgd),
20    /// AdaGrad with per-parameter squared gradient accumulation.
21    AdaGrad(AdaGrad),
22}
23
24impl Optimizer {
25    /// Create an SGD optimizer for `feature_count` features (plus intercept).
26    pub fn sgd(feature_count: usize, config: SgdConfig) -> Result<Self, RillError> {
27        Ok(Optimizer::Sgd(Sgd::new(feature_count, config)?))
28    }
29
30    /// Create an AdaGrad optimizer for `feature_count` features (plus intercept).
31    pub fn adagrad(feature_count: usize, config: AdaGradConfig) -> Result<Self, RillError> {
32        Ok(Optimizer::AdaGrad(AdaGrad::new(feature_count, config)?))
33    }
34
35    /// The number of parameters this optimizer manages (features + intercept).
36    pub fn param_count(&self) -> usize {
37        match self {
38            Optimizer::Sgd(o) => o.param_count(),
39            Optimizer::AdaGrad(o) => o.param_count(),
40        }
41    }
42
43    /// Number of samples the optimizer has processed.
44    pub fn samples_seen(&self) -> u64 {
45        match self {
46            Optimizer::Sgd(o) => o.samples_seen(),
47            Optimizer::AdaGrad(o) => o.samples_seen(),
48        }
49    }
50
51    /// Apply a gradient step to `weights` (length `feature_count`) and
52    /// `intercept` (single value). The gradient vector passed in must have
53    /// the same length as `weights`; the intercept gradient is passed
54    /// separately.
55    pub fn step(
56        &mut self,
57        weights: &mut [f64],
58        intercept: &mut f64,
59        grad_weights: &[f64],
60        grad_intercept: f64,
61    ) -> Result<(), RillError> {
62        match self {
63            Optimizer::Sgd(o) => o.step(weights, intercept, grad_weights, grad_intercept),
64            Optimizer::AdaGrad(o) => o.step(weights, intercept, grad_weights, grad_intercept),
65        }
66    }
67
68    /// Reset the optimizer to its initial state.
69    pub fn reset(&mut self) {
70        match self {
71            Optimizer::Sgd(o) => o.reset(),
72            Optimizer::AdaGrad(o) => o.reset(),
73        }
74    }
75}
76
77#[cfg(feature = "serde")]
78impl crate::persistence::ValidateState for Optimizer {
79    fn validate_state(&self) -> Result<(), RillError> {
80        match self {
81            Optimizer::Sgd(o) => o.validate_state(),
82            Optimizer::AdaGrad(o) => o.validate_state(),
83        }
84    }
85}
86
87use crate::error::RillError;