rill_ml/preprocessing/

standard_scaler.rs

//! Online standard scaler.
//!
//! Maintains per-feature Welford variance and mean. Time complexity per
//! update/transform: `O(d)`. Space complexity: `O(d)`.

use crate::error::{RillError, checked_increment, ensure_finite, validate_features};
use crate::traits::Transformer;

/// Configuration for [`StandardScaler`].
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct StandardScalerConfig {
    /// Whether to subtract the running mean. Default: `true`.
    pub with_mean: bool,
    /// Whether to divide by the running standard deviation. Default: `true`.
    pub with_std: bool,
    /// Variance threshold below which the scale is treated as `1.0` to avoid
    /// division by zero. Default: `1e-12`.
    pub epsilon: f64,
}

impl Default for StandardScalerConfig {
    fn default() -> Self {
        Self {
            with_mean: true,
            with_std: true,
            epsilon: 1e-12,
        }
    }
}

/// Online standard scaler that standardizes features to approximately zero
/// mean and unit variance.
///
/// - When `with_mean = false`, the mean subtraction is skipped.
/// - When `with_std = false`, the scaling is skipped.
/// - When a feature has seen zero samples, its mean is `0` and scale is `1`,
///   so the original value is returned unchanged.
/// - When a feature's variance is below `epsilon`, the scale is `1` to avoid
///   NaN or Infinity.
///
/// `transform` does not update state; only `update` does.
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(serde::Serialize))]
pub struct StandardScaler {
    feature_count: usize,
    config: StandardScalerConfig,
    counts: Vec<u64>,
    means: Vec<f64>,
    m2s: Vec<f64>,
}

impl StandardScaler {
    /// Create a new scaler for `feature_count` features with default config.
    pub fn new(feature_count: usize) -> Result<Self, RillError> {
        Self::with_config(feature_count, StandardScalerConfig::default())
    }

    /// Create a new scaler with a custom configuration.
    pub fn with_config(
        feature_count: usize,
        config: StandardScalerConfig,
    ) -> Result<Self, RillError> {
        if feature_count == 0 {
            return Err(RillError::EmptyFeatures);
        }
        ensure_finite("epsilon", config.epsilon)?;
        if config.epsilon < 0.0 {
            return Err(RillError::InvalidParameter {
                name: "epsilon",
                value: config.epsilon,
            });
        }
        Ok(Self {
            feature_count,
            config,
            counts: vec![0; feature_count],
            means: vec![0.0; feature_count],
            m2s: vec![0.0; feature_count],
        })
    }

    /// The per-feature means.
    pub fn means(&self) -> &[f64] {
        &self.means
    }

    /// The per-feature variances (population).
    pub fn variances(&self) -> Vec<f64> {
        self.m2s
            .iter()
            .zip(&self.counts)
            .map(|(&m2, &n)| if n == 0 { 0.0 } else { m2 / n as f64 })
            .collect()
    }

    /// The per-feature standard deviations.
    pub fn std_devs(&self) -> Vec<f64> {
        self.variances().iter().map(|v| v.sqrt()).collect()
    }

    /// The per-feature scales used during transformation.
    pub fn scales(&self) -> Vec<f64> {
        self.variances()
            .iter()
            .map(|&var| {
                if var < self.config.epsilon {
                    1.0
                } else {
                    var.sqrt()
                }
            })
            .collect()
    }

    /// Validate all configuration and persisted-state invariants.
    ///
    /// This is run automatically during deserialization and before operations
    /// that index per-feature state. The hot path (`transform_into` /
    /// `transform`) relies on the invariants established here and by the
    /// private constructor, and only re-checks them under `debug_assert!`.
    pub fn validate(&self) -> Result<(), RillError> {
        if self.feature_count == 0 {
            return Err(RillError::EmptyFeatures);
        }
        ensure_finite("epsilon", self.config.epsilon)?;
        if self.config.epsilon < 0.0 {
            return Err(RillError::InvalidParameter {
                name: "epsilon",
                value: self.config.epsilon,
            });
        }
        if self.counts.len() != self.feature_count
            || self.means.len() != self.feature_count
            || self.m2s.len() != self.feature_count
        {
            return Err(RillError::InvalidState(
                "standard scaler feature_count does not match state lengths".to_owned(),
            ));
        }
        if self.means.iter().any(|value| !value.is_finite())
            || self.m2s.iter().any(|value| !value.is_finite())
        {
            return Err(RillError::InvalidState(
                "standard scaler state must contain only finite values".to_owned(),
            ));
        }
        if self.counts.windows(2).any(|pair| pair[0] != pair[1]) {
            return Err(RillError::InvalidState(
                "standard scaler feature counts must stay synchronized".to_owned(),
            ));
        }
        Ok(())
    }

    /// Transform `features` into the provided `output` buffer, reusing its
    /// allocation instead of allocating a fresh `Vec` on every call.
    ///
    /// This is the hot-path entry point. Compared to [`transform`](Transformer::transform)
    /// it avoids two temporary allocations (the `variances()` and `scales()`
    /// vectors) by fusing the scale computation into the single output loop.
    /// The trust-boundary checks (dimension validation and finite-output
    /// enforcement) are preserved; the internal-state invariants established
    /// by [`validate`](Self::validate) and the private constructor are only
    /// re-checked under `debug_assert!` because they are guaranteed by the
    /// private fields and the validated deserialization path.
    ///
    /// `output` is truncated to the feature count and then filled; callers
    /// that reuse the same buffer across iterations avoid allocation
    /// entirely after the first call.
    pub fn transform_into(&self, features: &[f64], output: &mut Vec<f64>) -> Result<(), RillError> {
        // Trust-boundary: dimension must match. This is the public input
        // boundary and must always be enforced.
        validate_features(self.feature_count, features)?;

        // Internal invariants are guaranteed by the private constructor and
        // the validated Deserialize impl; re-check only in debug builds so
        // release-mode hot paths do not pay for repeated O(d) scans.
        debug_assert!(
            self.counts.len() == self.feature_count
                && self.means.len() == self.feature_count
                && self.m2s.len() == self.feature_count,
            "standard scaler state lengths must match feature_count"
        );
        debug_assert!(
            self.counts.windows(2).all(|pair| pair[0] == pair[1]),
            "standard scaler feature counts must stay synchronized"
        );
        debug_assert!(
            self.means.iter().all(|v| v.is_finite()) && self.m2s.iter().all(|v| v.is_finite()),
            "standard scaler state must contain only finite values"
        );

        output.clear();
        output.reserve(self.feature_count);

        let iter = features
            .iter()
            .zip(&self.counts)
            .zip(&self.means)
            .zip(&self.m2s);
        for (((&x, &n), &mean_storage), &m2) in iter {
            // Population variance = m2 / n; if n == 0 the scale is 1.0 so
            // the original value is returned unchanged.
            let scale = if !self.config.with_std || n == 0 {
                1.0
            } else {
                let variance = m2 / n as f64;
                if variance < self.config.epsilon {
                    1.0
                } else {
                    variance.sqrt()
                }
            };
            let mean = if self.config.with_mean {
                mean_storage
            } else {
                0.0
            };
            let transformed = (x - mean) / scale;
            // Trust-boundary: output must be finite. This catches NaN/Inf
            // introduced by adversarial input even when internal state is
            // already validated.
            ensure_finite("transformed feature", transformed)?;
            output.push(transformed);
        }
        Ok(())
    }
}

impl Transformer for StandardScaler {
    fn input_dim(&self) -> usize {
        self.feature_count
    }

    fn output_dim(&self) -> usize {
        self.feature_count
    }

    fn transform(&self, features: &[f64]) -> Result<Vec<f64>, RillError> {
        // Validate the full state on the public Transformer entry point so
        // a corrupted scaler (e.g. one built via unsafe or a future struct
        // literal) cannot proceed. The hot path `transform_into` relies on
        // the same invariants but only re-checks them under `debug_assert!`.
        self.validate()?;
        let mut output = Vec::with_capacity(self.feature_count);
        self.transform_into(features, &mut output)?;
        Ok(output)
    }

    fn update(&mut self, features: &[f64]) -> Result<(), RillError> {
        self.validate()?;
        validate_features(self.feature_count, features)?;
        let mut next_counts = self.counts.clone();
        let mut next_means = self.means.clone();
        let mut next_m2s = self.m2s.clone();
        for (i, &x) in features.iter().enumerate() {
            let count = checked_increment(self.counts[i], "standard scaler sample")?;
            let delta = x - self.means[i];
            ensure_finite("standard scaler delta", delta)?;
            let mean = self.means[i] + delta / count as f64;
            ensure_finite("standard scaler mean", mean)?;
            let delta2 = x - mean;
            ensure_finite("standard scaler delta", delta2)?;
            let m2 = self.m2s[i] + delta * delta2;
            ensure_finite("standard scaler M2", m2)?;
            next_counts[i] = count;
            next_means[i] = mean;
            next_m2s[i] = m2;
        }
        self.counts = next_counts;
        self.means = next_means;
        self.m2s = next_m2s;
        Ok(())
    }

    fn samples_seen(&self) -> u64 {
        self.counts.iter().copied().max().unwrap_or(0)
    }

    fn reset(&mut self) {
        for c in &mut self.counts {
            *c = 0;
        }
        for m in &mut self.means {
            *m = 0.0;
        }
        for m2 in &mut self.m2s {
            *m2 = 0.0;
        }
    }
}

#[cfg(feature = "serde")]
#[derive(serde::Deserialize)]
struct StandardScalerState {
    feature_count: usize,
    config: StandardScalerConfig,
    counts: Vec<u64>,
    means: Vec<f64>,
    m2s: Vec<f64>,
}

#[cfg(feature = "serde")]
impl<'de> serde::Deserialize<'de> for StandardScaler {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let state = StandardScalerState::deserialize(deserializer)?;
        let scaler = Self {
            feature_count: state.feature_count,
            config: state.config,
            counts: state.counts,
            means: state.means,
            m2s: state.m2s,
        };
        scaler.validate().map_err(serde::de::Error::custom)?;
        Ok(scaler)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn scaler_zero_state_returns_original() {
        let s = StandardScaler::new(3).unwrap();
        let out = s.transform(&[1.0, 2.0, 3.0]).unwrap();
        // count == 0 -> mean=0, scale=1 -> original
        assert!((out[0] - 1.0).abs() < 1e-12);
        assert!((out[1] - 2.0).abs() < 1e-12);
        assert!((out[2] - 3.0).abs() < 1e-12);
    }

    #[test]
    fn scaler_standardizes_after_updates() {
        let mut s = StandardScaler::new(2).unwrap();
        // feature 0: values [1, 3] -> mean 2, var 1, std 1
        // feature 1: values [10, 20] -> mean 15, var 25, std 5
        s.update(&[1.0, 10.0]).unwrap();
        s.update(&[3.0, 20.0]).unwrap();
        let out = s.transform(&[3.0, 20.0]).unwrap();
        // (3-2)/1 = 1, (20-15)/5 = 1
        assert!((out[0] - 1.0).abs() < 1e-9);
        assert!((out[1] - 1.0).abs() < 1e-9);
    }

    #[test]
    fn transform_does_not_update_state() {
        let mut s = StandardScaler::new(1).unwrap();
        s.update(&[10.0]).unwrap();
        let mean_before = s.means()[0];
        let _ = s.transform(&[5.0]).unwrap();
        assert_eq!(s.means()[0], mean_before);
        assert_eq!(s.counts[0], 1);
    }

    #[test]
    fn update_rejects_overflow_without_mutating_state() {
        let mut scaler = StandardScaler::new(1).unwrap();
        scaler.update(&[f64::MAX]).unwrap();
        let before = scaler.clone();
        assert!(scaler.update(&[-f64::MAX]).is_err());
        assert_eq!(scaler.counts, before.counts);
        assert_eq!(scaler.means, before.means);
        assert_eq!(scaler.m2s, before.m2s);
    }

    #[cfg(feature = "serde")]
    #[test]
    fn serde_rejects_malformed_state() {
        let malformed = r#"{
            "feature_count":2,
            "config":{"with_mean":true,"with_std":true,"epsilon":1e-12},
            "counts":[1],
            "means":[0.0],
            "m2s":[0.0]
        }"#;
        assert!(serde_json::from_str::<StandardScaler>(malformed).is_err());
    }

    #[test]
    fn constant_feature_uses_scale_one() {
        let mut s = StandardScaler::new(1).unwrap();
        for _ in 0..10 {
            s.update(&[5.0]).unwrap();
        }
        // var = 0 < epsilon -> scale = 1, mean = 5 -> (5-5)/1 = 0
        let out = s.transform(&[5.0]).unwrap();
        assert!(out[0].abs() < 1e-12);
        assert!(!out[0].is_nan());
    }

    #[test]
    fn with_mean_false_keeps_offset() {
        let mut s = StandardScaler::with_config(
            1,
            StandardScalerConfig {
                with_mean: false,
                with_std: true,
                epsilon: 1e-12,
            },
        )
        .unwrap();
        s.update(&[1.0]).unwrap();
        s.update(&[3.0]).unwrap();
        // mean=2, var=1, std=1, but with_mean=false so x/1 = x
        let out = s.transform(&[3.0]).unwrap();
        assert!((out[0] - 3.0).abs() < 1e-9);
    }

    #[test]
    fn dimension_mismatch_rejected() {
        let mut s = StandardScaler::new(3).unwrap();
        assert!(s.transform(&[1.0, 2.0]).is_err());
        assert!(s.update(&[1.0, 2.0]).is_err());
    }

    #[test]
    fn zero_features_rejected() {
        assert!(matches!(
            StandardScaler::new(0),
            Err(RillError::EmptyFeatures)
        ));
    }

    #[test]
    fn non_finite_rejected() {
        let mut s = StandardScaler::new(2).unwrap();
        assert!(s.update(&[1.0, f64::NAN]).is_err());
    }

    #[test]
    fn reset_clears_state() {
        let mut s = StandardScaler::new(1).unwrap();
        s.update(&[1.0]).unwrap();
        s.update(&[2.0]).unwrap();
        s.reset();
        assert_eq!(s.counts[0], 0);
        assert_eq!(s.means()[0], 0.0);
    }

    #[test]
    fn transform_into_matches_transform_output() {
        let mut scaler = StandardScaler::new(4).unwrap();
        // Feed a few samples so means/variances are non-trivial.
        scaler.update(&[1.0, 10.0, 100.0, 1000.0]).unwrap();
        scaler.update(&[3.0, 20.0, 300.0, 3000.0]).unwrap();
        scaler.update(&[5.0, 30.0, 500.0, 5000.0]).unwrap();
        let features = [2.0, 15.0, 200.0, 2000.0];
        let via_transform = scaler.transform(&features).unwrap();
        let mut via_into = Vec::new();
        scaler.transform_into(&features, &mut via_into).unwrap();
        assert_eq!(via_transform, via_into);
    }

    #[test]
    fn transform_into_reuses_buffer_capacity() {
        let mut scaler = StandardScaler::new(3).unwrap();
        scaler.update(&[1.0, 2.0, 3.0]).unwrap();
        scaler.update(&[4.0, 5.0, 6.0]).unwrap();
        let features = [2.5, 3.5, 4.5];
        let mut buffer = Vec::with_capacity(64);
        // Prime the buffer with sentinel content to prove clear() is called.
        buffer.extend_from_slice(&[-1.0, -2.0, -3.0, -4.0]);
        scaler.transform_into(&features, &mut buffer).unwrap();
        assert_eq!(buffer.len(), 3);
        // Capacity must be preserved (no reallocation).
        assert!(buffer.capacity() >= 64);
        // Content must match the public transform() output.
        assert_eq!(buffer, scaler.transform(&features).unwrap());
    }

    #[test]
    fn transform_into_rejects_dimension_mismatch() {
        let scaler = StandardScaler::new(3).unwrap();
        let mut buffer = Vec::new();
        assert!(scaler.transform_into(&[1.0, 2.0], &mut buffer).is_err());
        // Buffer must remain empty after the dimension error.
        assert!(buffer.is_empty());
    }

    #[test]
    fn transform_into_rejects_non_finite_output() {
        // with_std = false and a non-finite input must still be caught by
        // the finite-output trust-boundary check.
        let scaler = StandardScaler::with_config(
            1,
            StandardScalerConfig {
                with_mean: false,
                with_std: false,
                epsilon: 1e-12,
            },
        )
        .unwrap();
        let mut buffer = Vec::new();
        assert!(scaler.transform_into(&[f64::NAN], &mut buffer).is_err());
    }

    #[test]
    fn transform_into_with_zero_state_returns_original() {
        let scaler = StandardScaler::new(3).unwrap();
        let features = [1.5, 2.5, 3.5];
        let mut buffer = Vec::new();
        scaler.transform_into(&features, &mut buffer).unwrap();
        // count == 0 → mean = 0, scale = 1 → original values.
        assert_eq!(buffer, vec![1.5, 2.5, 3.5]);
    }
}