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(),
            ));
        }
        // ``m2s`` is the running sum of squared deviations from the mean
        // (Welford M2). It is mathematically non-negative; a negative value
        // indicates corruption or a maliciously crafted serde payload. The
        // finite check above already rules out NaN/Infinity, so here we only
        // need to reject strictly negative values.
        if self.m2s.iter().any(|value| *value < 0.0) {
            return Err(RillError::InvalidState(
                "standard scaler m2s must be non-negative".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(),
            ));
        }
        // counts are synchronized (verified above), so the first entry
        // represents every feature's sample count.
        let n = self.counts.first().copied().unwrap_or(0);
        // count == 0: no samples seen → mean and M2 must be exactly 0 for
        // every feature. The normal paths (new() / reset()) guarantee this;
        // a non-zero value indicates a corrupted or malicious payload. The
        // public docs promise mean=0 and scale=1 in this state, so accepting
        // a non-zero mean would silently break transform() output.
        if n == 0 {
            for (i, &mean) in self.means.iter().enumerate() {
                if mean != 0.0 {
                    return Err(RillError::InvalidState(format!(
                        "standard scaler means[{i}] must be 0 when count == 0, got {mean}"
                    )));
                }
            }
            for (i, &m2) in self.m2s.iter().enumerate() {
                if m2 != 0.0 {
                    return Err(RillError::InvalidState(format!(
                        "standard scaler m2s[{i}] must be 0 when count == 0, got {m2}"
                    )));
                }
            }
        }
        // count == 1: after a single Welford update, delta2 = x - mean = 0,
        // so m2_delta = delta * delta2 = 0 and M2 stays exactly 0. The
        // normal update path guarantees exact 0, so no floating-point
        // tolerance is introduced here. The single training sample is
        // unknown, so mean is NOT constrained to a specific value.
        if n == 1 {
            for (i, &m2) in self.m2s.iter().enumerate() {
                if m2 != 0.0 {
                    return Err(RillError::InvalidState(format!(
                        "standard scaler m2s[{i}] must be 0 when count == 1, got {m2}"
                    )));
                }
            }
        }
        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> {
        // Hot path: rely on the invariants established by the private
        // constructor and the validated Deserialize impl (private fields
        // cannot be set to an inconsistent state from outside the crate).
        // Only input dimension and output finiteness are checked on the
        // release path; full state scanning is left to `validate()` and
        // `debug_assert!` inside `transform_into`.
        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());
    }

    #[cfg(feature = "serde")]
    #[test]
    fn serde_rejects_negative_m2() {
        // Regression: a malicious or corrupted state with a negative Welford
        // M2 must be rejected. ``m2`` is a sum of squared deviations and is
        // mathematically non-negative; accepting a negative value would let
        // an attacker poison the scaler with imaginary variances.
        let malformed = r#"{
            "feature_count":1,
            "config":{"with_mean":true,"with_std":true,"epsilon":1e-12},
            "counts":[2],
            "means":[0.0],
            "m2s":[-1.0]
        }"#;
        assert!(serde_json::from_str::<StandardScaler>(malformed).is_err());
    }

    #[cfg(feature = "serde")]
    #[test]
    fn scaler_serde_rejects_zero_count_nonzero_mean() {
        // count == 0 promises mean=0, scale=1 (input returned unchanged).
        // A non-zero mean would silently break transform() output.
        let malformed = r#"{
            "feature_count":1,
            "config":{"with_mean":true,"with_std":true,"epsilon":1e-12},
            "counts":[0],
            "means":[10.0],
            "m2s":[0.0]
        }"#;
        assert!(
            serde_json::from_str::<StandardScaler>(malformed).is_err(),
            "count=0 with non-zero mean must be rejected"
        );
    }

    #[cfg(feature = "serde")]
    #[test]
    fn scaler_serde_rejects_zero_count_nonzero_m2() {
        // count == 0 implies no samples → M2 (sum of squared deviations)
        // must be exactly 0. A non-zero M2 is a corrupted/malicious state.
        let malformed = r#"{
            "feature_count":1,
            "config":{"with_mean":true,"with_std":true,"epsilon":1e-12},
            "counts":[0],
            "means":[0.0],
            "m2s":[1.0]
        }"#;
        assert!(
            serde_json::from_str::<StandardScaler>(malformed).is_err(),
            "count=0 with non-zero m2 must be rejected"
        );
    }

    #[cfg(feature = "serde")]
    #[test]
    fn scaler_serde_rejects_one_count_nonzero_m2() {
        // count == 1: after a single Welford update, delta2 = x - mean = 0,
        // so M2 stays exactly 0. The training sample itself is unknown, so
        // mean is NOT constrained — only M2 must be 0.
        let malformed = r#"{
            "feature_count":1,
            "config":{"with_mean":true,"with_std":true,"epsilon":1e-12},
            "counts":[1],
            "means":[3.0],
            "m2s":[0.25]
        }"#;
        assert!(
            serde_json::from_str::<StandardScaler>(malformed).is_err(),
            "count=1 with non-zero m2 must be rejected"
        );
    }

    #[cfg(feature = "serde")]
    #[test]
    fn scaler_serde_accepts_one_count_finite_mean() {
        // count == 1 with an arbitrary finite mean and M2 == 0 is the
        // legitimate post-single-update state (mean == the single training
        // sample). It must round-trip successfully.
        let json = r#"{
            "feature_count":2,
            "config":{"with_mean":true,"with_std":true,"epsilon":1e-12},
            "counts":[1,1],
            "means":[3.0,-7.5],
            "m2s":[0.0,0.0]
        }"#;
        let scaler: StandardScaler =
            serde_json::from_str(json).expect("count=1 with finite mean and m2=0 must be accepted");
        // Round-trip back to JSON and re-parse to confirm stability.
        let re = serde_json::to_string(&scaler).unwrap();
        let _: StandardScaler = serde_json::from_str(&re).unwrap();
        assert_eq!(scaler.counts, vec![1, 1]);
        assert_eq!(scaler.means, vec![3.0, -7.5]);
        assert_eq!(scaler.m2s, vec![0.0, 0.0]);
    }

    #[cfg(feature = "serde")]
    #[test]
    fn scaler_serde_roundtrip_preserves_state() {
        // A scaler trained on a few samples must survive a serialize →
        // deserialize round-trip with all invariants intact.
        let mut scaler = StandardScaler::new(2).unwrap();
        scaler.update(&[1.0, 10.0]).unwrap();
        scaler.update(&[3.0, 20.0]).unwrap();
        scaler.update(&[5.0, 30.0]).unwrap();
        let json = serde_json::to_string(&scaler).unwrap();
        let restored: StandardScaler = serde_json::from_str(&json).unwrap();
        assert_eq!(restored.counts, scaler.counts);
        assert_eq!(restored.means, scaler.means);
        assert_eq!(restored.m2s, scaler.m2s);
        // transform output must match as well.
        let features = [2.0, 15.0];
        assert_eq!(
            scaler.transform(&features).unwrap(),
            restored.transform(&features).unwrap()
        );
    }

    #[test]
    fn transform_hot_path_does_not_scan_invariants() {
        // The release transform() hot path must remain a single loop and
        // must NOT re-run the O(d) invariant scan from validate(). This
        // test constructs a valid scaler and confirms transform() succeeds
        // without invoking validate() (which would walk counts/means/m2s).
        // We can't directly count scans, but we can confirm a scaler that
        // has valid invariants transforms correctly and that the hot path
        // doesn't change state.
        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 before = scaler.clone();
        let out = scaler.transform(&[2.5, 3.5, 4.5]).unwrap();
        // State must be unchanged.
        assert_eq!(scaler.counts, before.counts);
        assert_eq!(scaler.means, before.means);
        assert_eq!(scaler.m2s, before.m2s);
        // Output length matches feature count (single loop).
        assert_eq!(out.len(), 3);
    }

    #[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]);
    }
}