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 mod adagrad;
8pub 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))]
16pub enum Optimizer {
17    /// Stochastic gradient descent with optional L2 regularization.
18    Sgd(Sgd),
19    /// AdaGrad with per-parameter squared gradient accumulation.
20    AdaGrad(AdaGrad),
21}
22
23impl Optimizer {
24    /// Create an SGD optimizer for `feature_count` features (plus intercept).
25    pub fn sgd(feature_count: usize, config: SgdConfig) -> Result<Self, RillError> {
26        Ok(Optimizer::Sgd(Sgd::new(feature_count, config)?))
27    }
28
29    /// Create an AdaGrad optimizer for `feature_count` features (plus intercept).
30    pub fn adagrad(feature_count: usize, config: AdaGradConfig) -> Result<Self, RillError> {
31        Ok(Optimizer::AdaGrad(AdaGrad::new(feature_count, config)?))
32    }
33
34    /// The number of parameters this optimizer manages (features + intercept).
35    pub fn param_count(&self) -> usize {
36        match self {
37            Optimizer::Sgd(o) => o.param_count(),
38            Optimizer::AdaGrad(o) => o.param_count(),
39        }
40    }
41
42    /// Number of samples the optimizer has processed.
43    pub fn samples_seen(&self) -> u64 {
44        match self {
45            Optimizer::Sgd(o) => o.samples_seen(),
46            Optimizer::AdaGrad(o) => o.samples_seen(),
47        }
48    }
49
50    /// Apply a gradient step to `weights` (length `feature_count`) and
51    /// `intercept` (single value). The gradient vector passed in must have
52    /// the same length as `weights`; the intercept gradient is passed
53    /// separately.
54    pub fn step(
55        &mut self,
56        weights: &mut [f64],
57        intercept: &mut f64,
58        grad_weights: &[f64],
59        grad_intercept: f64,
60    ) -> Result<(), RillError> {
61        match self {
62            Optimizer::Sgd(o) => o.step(weights, intercept, grad_weights, grad_intercept),
63            Optimizer::AdaGrad(o) => o.step(weights, intercept, grad_weights, grad_intercept),
64        }
65    }
66
67    /// Reset the optimizer to its initial state.
68    pub fn reset(&mut self) {
69        match self {
70            Optimizer::Sgd(o) => o.reset(),
71            Optimizer::AdaGrad(o) => o.reset(),
72        }
73    }
74}
75
76use crate::error::RillError;