1use crate::error::{RillError, checked_increment, ensure_finite, validate_features};
7#[cfg(feature = "serde")]
8use crate::persistence::ValidateState;
9use crate::traits::Transformer;
10
11#[derive(Debug, Clone)]
13#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
14#[non_exhaustive]
15pub struct StandardScalerConfig {
16 pub with_mean: bool,
18 pub with_std: bool,
20 pub epsilon: f64,
23}
24
25impl Default for StandardScalerConfig {
26 fn default() -> Self {
27 Self {
28 with_mean: true,
29 with_std: true,
30 epsilon: 1e-12,
31 }
32 }
33}
34
35#[derive(Debug, Clone)]
47#[cfg_attr(feature = "serde", derive(serde::Serialize))]
48pub struct StandardScaler {
49 feature_count: usize,
50 config: StandardScalerConfig,
51 counts: Vec<u64>,
52 means: Vec<f64>,
53 m2s: Vec<f64>,
54}
55
56impl StandardScaler {
57 pub fn new(feature_count: usize) -> Result<Self, RillError> {
59 Self::with_config(feature_count, StandardScalerConfig::default())
60 }
61
62 pub fn with_config(
64 feature_count: usize,
65 config: StandardScalerConfig,
66 ) -> Result<Self, RillError> {
67 if feature_count == 0 {
68 return Err(RillError::EmptyFeatures);
69 }
70 ensure_finite("epsilon", config.epsilon)?;
71 if config.epsilon < 0.0 {
72 return Err(RillError::InvalidParameter {
73 name: "epsilon",
74 value: config.epsilon,
75 });
76 }
77 Ok(Self {
78 feature_count,
79 config,
80 counts: vec![0; feature_count],
81 means: vec![0.0; feature_count],
82 m2s: vec![0.0; feature_count],
83 })
84 }
85
86 pub fn means(&self) -> &[f64] {
88 &self.means
89 }
90
91 pub fn variances(&self) -> Vec<f64> {
93 self.m2s
94 .iter()
95 .zip(&self.counts)
96 .map(|(&m2, &n)| if n == 0 { 0.0 } else { m2 / n as f64 })
97 .collect()
98 }
99
100 pub fn std_devs(&self) -> Vec<f64> {
102 self.variances().iter().map(|v| v.sqrt()).collect()
103 }
104
105 pub fn scales(&self) -> Vec<f64> {
107 self.variances()
108 .iter()
109 .map(|&var| {
110 if var < self.config.epsilon {
111 1.0
112 } else {
113 var.sqrt()
114 }
115 })
116 .collect()
117 }
118
119 pub fn validate(&self) -> Result<(), RillError> {
126 if self.feature_count == 0 {
127 return Err(RillError::EmptyFeatures);
128 }
129 ensure_finite("epsilon", self.config.epsilon)?;
130 if self.config.epsilon < 0.0 {
131 return Err(RillError::InvalidParameter {
132 name: "epsilon",
133 value: self.config.epsilon,
134 });
135 }
136 if self.counts.len() != self.feature_count
137 || self.means.len() != self.feature_count
138 || self.m2s.len() != self.feature_count
139 {
140 return Err(RillError::InvalidState(
141 "standard scaler feature_count does not match state lengths".to_owned(),
142 ));
143 }
144 if self.means.iter().any(|value| !value.is_finite())
145 || self.m2s.iter().any(|value| !value.is_finite())
146 {
147 return Err(RillError::InvalidState(
148 "standard scaler state must contain only finite values".to_owned(),
149 ));
150 }
151 if self.m2s.iter().any(|value| *value < 0.0) {
157 return Err(RillError::InvalidState(
158 "standard scaler m2s must be non-negative".to_owned(),
159 ));
160 }
161 if self.counts.windows(2).any(|pair| pair[0] != pair[1]) {
162 return Err(RillError::InvalidState(
163 "standard scaler feature counts must stay synchronized".to_owned(),
164 ));
165 }
166 let n = self.counts.first().copied().unwrap_or(0);
169 if n == 0 {
175 for (i, &mean) in self.means.iter().enumerate() {
176 if mean != 0.0 {
177 return Err(RillError::InvalidState(format!(
178 "standard scaler means[{i}] must be 0 when count == 0, got {mean}"
179 )));
180 }
181 }
182 for (i, &m2) in self.m2s.iter().enumerate() {
183 if m2 != 0.0 {
184 return Err(RillError::InvalidState(format!(
185 "standard scaler m2s[{i}] must be 0 when count == 0, got {m2}"
186 )));
187 }
188 }
189 }
190 if n == 1 {
196 for (i, &m2) in self.m2s.iter().enumerate() {
197 if m2 != 0.0 {
198 return Err(RillError::InvalidState(format!(
199 "standard scaler m2s[{i}] must be 0 when count == 1, got {m2}"
200 )));
201 }
202 }
203 }
204 Ok(())
205 }
206
207 pub fn transform_into(&self, features: &[f64], output: &mut Vec<f64>) -> Result<(), RillError> {
223 validate_features(self.feature_count, features)?;
226
227 debug_assert!(
231 self.counts.len() == self.feature_count
232 && self.means.len() == self.feature_count
233 && self.m2s.len() == self.feature_count,
234 "standard scaler state lengths must match feature_count"
235 );
236 debug_assert!(
237 self.counts.windows(2).all(|pair| pair[0] == pair[1]),
238 "standard scaler feature counts must stay synchronized"
239 );
240 debug_assert!(
241 self.means.iter().all(|v| v.is_finite()) && self.m2s.iter().all(|v| v.is_finite()),
242 "standard scaler state must contain only finite values"
243 );
244
245 output.clear();
246 output.reserve(self.feature_count);
247
248 let iter = features
249 .iter()
250 .zip(&self.counts)
251 .zip(&self.means)
252 .zip(&self.m2s);
253 for (((&x, &n), &mean_storage), &m2) in iter {
254 let scale = if !self.config.with_std || n == 0 {
257 1.0
258 } else {
259 let variance = m2 / n as f64;
260 if variance < self.config.epsilon {
261 1.0
262 } else {
263 variance.sqrt()
264 }
265 };
266 let mean = if self.config.with_mean {
267 mean_storage
268 } else {
269 0.0
270 };
271 let transformed = (x - mean) / scale;
272 ensure_finite("transformed feature", transformed)?;
276 output.push(transformed);
277 }
278 Ok(())
279 }
280}
281
282impl Transformer for StandardScaler {
283 fn input_dim(&self) -> usize {
284 self.feature_count
285 }
286
287 fn output_dim(&self) -> usize {
288 self.feature_count
289 }
290
291 fn transform(&self, features: &[f64]) -> Result<Vec<f64>, RillError> {
292 let mut output = Vec::with_capacity(self.feature_count);
299 self.transform_into(features, &mut output)?;
300 Ok(output)
301 }
302
303 fn update(&mut self, features: &[f64]) -> Result<(), RillError> {
304 self.validate()?;
305 validate_features(self.feature_count, features)?;
306 let mut next_counts = self.counts.clone();
307 let mut next_means = self.means.clone();
308 let mut next_m2s = self.m2s.clone();
309 for (i, &x) in features.iter().enumerate() {
310 let count = checked_increment(self.counts[i], "standard scaler sample")?;
311 let delta = x - self.means[i];
312 ensure_finite("standard scaler delta", delta)?;
313 let mean = self.means[i] + delta / count as f64;
314 ensure_finite("standard scaler mean", mean)?;
315 let delta2 = x - mean;
316 ensure_finite("standard scaler delta", delta2)?;
317 let m2 = self.m2s[i] + delta * delta2;
318 ensure_finite("standard scaler M2", m2)?;
319 next_counts[i] = count;
320 next_means[i] = mean;
321 next_m2s[i] = m2;
322 }
323 self.counts = next_counts;
324 self.means = next_means;
325 self.m2s = next_m2s;
326 Ok(())
327 }
328
329 fn samples_seen(&self) -> u64 {
330 self.counts.iter().copied().max().unwrap_or(0)
331 }
332
333 fn reset(&mut self) {
334 self.counts.fill(0);
335 self.means.fill(0.0);
336 self.m2s.fill(0.0);
337 }
338}
339
340#[cfg(feature = "serde")]
341#[derive(serde::Deserialize)]
342struct StandardScalerState {
343 feature_count: usize,
344 config: StandardScalerConfig,
345 counts: Vec<u64>,
346 means: Vec<f64>,
347 m2s: Vec<f64>,
348}
349
350#[cfg(feature = "serde")]
351impl<'de> serde::Deserialize<'de> for StandardScaler {
352 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
353 where
354 D: serde::Deserializer<'de>,
355 {
356 let state = StandardScalerState::deserialize(deserializer)?;
357 let scaler = Self {
358 feature_count: state.feature_count,
359 config: state.config,
360 counts: state.counts,
361 means: state.means,
362 m2s: state.m2s,
363 };
364 scaler.validate().map_err(serde::de::Error::custom)?;
365 Ok(scaler)
366 }
367}
368
369#[cfg(feature = "serde")]
370impl ValidateState for StandardScaler {
371 fn validate_state(&self) -> Result<(), RillError> {
372 StandardScaler::validate(self)
373 }
374}
375
376#[cfg(test)]
377mod tests {
378 use super::*;
379
380 #[test]
381 fn scaler_zero_state_returns_original() {
382 let s = StandardScaler::new(3).unwrap();
383 let out = s.transform(&[1.0, 2.0, 3.0]).unwrap();
384 assert!((out[0] - 1.0).abs() < 1e-12);
386 assert!((out[1] - 2.0).abs() < 1e-12);
387 assert!((out[2] - 3.0).abs() < 1e-12);
388 }
389
390 #[test]
391 fn scaler_standardizes_after_updates() {
392 let mut s = StandardScaler::new(2).unwrap();
393 s.update(&[1.0, 10.0]).unwrap();
396 s.update(&[3.0, 20.0]).unwrap();
397 let out = s.transform(&[3.0, 20.0]).unwrap();
398 assert!((out[0] - 1.0).abs() < 1e-9);
400 assert!((out[1] - 1.0).abs() < 1e-9);
401 }
402
403 #[test]
404 fn transform_does_not_update_state() {
405 let mut s = StandardScaler::new(1).unwrap();
406 s.update(&[10.0]).unwrap();
407 let mean_before = s.means()[0];
408 let _ = s.transform(&[5.0]).unwrap();
409 assert_eq!(s.means()[0], mean_before);
410 assert_eq!(s.counts[0], 1);
411 }
412
413 #[test]
414 fn update_rejects_overflow_without_mutating_state() {
415 let mut scaler = StandardScaler::new(1).unwrap();
416 scaler.update(&[f64::MAX]).unwrap();
417 let before = scaler.clone();
418 assert!(scaler.update(&[-f64::MAX]).is_err());
419 assert_eq!(scaler.counts, before.counts);
420 assert_eq!(scaler.means, before.means);
421 assert_eq!(scaler.m2s, before.m2s);
422 }
423
424 #[cfg(feature = "serde")]
425 #[test]
426 fn serde_rejects_malformed_state() {
427 let malformed = r#"{
428 "feature_count":2,
429 "config":{"with_mean":true,"with_std":true,"epsilon":1e-12},
430 "counts":[1],
431 "means":[0.0],
432 "m2s":[0.0]
433 }"#;
434 assert!(serde_json::from_str::<StandardScaler>(malformed).is_err());
435 }
436
437 #[cfg(feature = "serde")]
438 #[test]
439 fn serde_rejects_negative_m2() {
440 let malformed = r#"{
445 "feature_count":1,
446 "config":{"with_mean":true,"with_std":true,"epsilon":1e-12},
447 "counts":[2],
448 "means":[0.0],
449 "m2s":[-1.0]
450 }"#;
451 assert!(serde_json::from_str::<StandardScaler>(malformed).is_err());
452 }
453
454 #[cfg(feature = "serde")]
455 #[test]
456 fn scaler_serde_rejects_zero_count_nonzero_mean() {
457 let malformed = r#"{
460 "feature_count":1,
461 "config":{"with_mean":true,"with_std":true,"epsilon":1e-12},
462 "counts":[0],
463 "means":[10.0],
464 "m2s":[0.0]
465 }"#;
466 assert!(
467 serde_json::from_str::<StandardScaler>(malformed).is_err(),
468 "count=0 with non-zero mean must be rejected"
469 );
470 }
471
472 #[cfg(feature = "serde")]
473 #[test]
474 fn scaler_serde_rejects_zero_count_nonzero_m2() {
475 let malformed = r#"{
478 "feature_count":1,
479 "config":{"with_mean":true,"with_std":true,"epsilon":1e-12},
480 "counts":[0],
481 "means":[0.0],
482 "m2s":[1.0]
483 }"#;
484 assert!(
485 serde_json::from_str::<StandardScaler>(malformed).is_err(),
486 "count=0 with non-zero m2 must be rejected"
487 );
488 }
489
490 #[cfg(feature = "serde")]
491 #[test]
492 fn scaler_serde_rejects_one_count_nonzero_m2() {
493 let malformed = r#"{
497 "feature_count":1,
498 "config":{"with_mean":true,"with_std":true,"epsilon":1e-12},
499 "counts":[1],
500 "means":[3.0],
501 "m2s":[0.25]
502 }"#;
503 assert!(
504 serde_json::from_str::<StandardScaler>(malformed).is_err(),
505 "count=1 with non-zero m2 must be rejected"
506 );
507 }
508
509 #[cfg(feature = "serde")]
510 #[test]
511 fn scaler_serde_accepts_one_count_finite_mean() {
512 let json = r#"{
516 "feature_count":2,
517 "config":{"with_mean":true,"with_std":true,"epsilon":1e-12},
518 "counts":[1,1],
519 "means":[3.0,-7.5],
520 "m2s":[0.0,0.0]
521 }"#;
522 let scaler: StandardScaler =
523 serde_json::from_str(json).expect("count=1 with finite mean and m2=0 must be accepted");
524 let re = serde_json::to_string(&scaler).unwrap();
526 let _: StandardScaler = serde_json::from_str(&re).unwrap();
527 assert_eq!(scaler.counts, vec![1, 1]);
528 assert_eq!(scaler.means, vec![3.0, -7.5]);
529 assert_eq!(scaler.m2s, vec![0.0, 0.0]);
530 }
531
532 #[cfg(feature = "serde")]
533 #[test]
534 fn scaler_serde_roundtrip_preserves_state() {
535 let mut scaler = StandardScaler::new(2).unwrap();
538 scaler.update(&[1.0, 10.0]).unwrap();
539 scaler.update(&[3.0, 20.0]).unwrap();
540 scaler.update(&[5.0, 30.0]).unwrap();
541 let json = serde_json::to_string(&scaler).unwrap();
542 let restored: StandardScaler = serde_json::from_str(&json).unwrap();
543 assert_eq!(restored.counts, scaler.counts);
544 assert_eq!(restored.means, scaler.means);
545 assert_eq!(restored.m2s, scaler.m2s);
546 let features = [2.0, 15.0];
548 assert_eq!(
549 scaler.transform(&features).unwrap(),
550 restored.transform(&features).unwrap()
551 );
552 }
553
554 #[test]
555 fn transform_hot_path_does_not_scan_invariants() {
556 let mut scaler = StandardScaler::new(3).unwrap();
564 scaler.update(&[1.0, 2.0, 3.0]).unwrap();
565 scaler.update(&[4.0, 5.0, 6.0]).unwrap();
566 let before = scaler.clone();
567 let out = scaler.transform(&[2.5, 3.5, 4.5]).unwrap();
568 assert_eq!(scaler.counts, before.counts);
570 assert_eq!(scaler.means, before.means);
571 assert_eq!(scaler.m2s, before.m2s);
572 assert_eq!(out.len(), 3);
574 }
575
576 #[test]
577 fn constant_feature_uses_scale_one() {
578 let mut s = StandardScaler::new(1).unwrap();
579 for _ in 0..10 {
580 s.update(&[5.0]).unwrap();
581 }
582 let out = s.transform(&[5.0]).unwrap();
584 assert!(out[0].abs() < 1e-12);
585 assert!(!out[0].is_nan());
586 }
587
588 #[test]
589 fn with_mean_false_keeps_offset() {
590 let mut s = StandardScaler::with_config(
591 1,
592 StandardScalerConfig {
593 with_mean: false,
594 with_std: true,
595 epsilon: 1e-12,
596 },
597 )
598 .unwrap();
599 s.update(&[1.0]).unwrap();
600 s.update(&[3.0]).unwrap();
601 let out = s.transform(&[3.0]).unwrap();
603 assert!((out[0] - 3.0).abs() < 1e-9);
604 }
605
606 #[test]
607 fn dimension_mismatch_rejected() {
608 let mut s = StandardScaler::new(3).unwrap();
609 assert!(s.transform(&[1.0, 2.0]).is_err());
610 assert!(s.update(&[1.0, 2.0]).is_err());
611 }
612
613 #[test]
614 fn zero_features_rejected() {
615 assert!(matches!(
616 StandardScaler::new(0),
617 Err(RillError::EmptyFeatures)
618 ));
619 }
620
621 #[test]
622 fn non_finite_rejected() {
623 let mut s = StandardScaler::new(2).unwrap();
624 assert!(s.update(&[1.0, f64::NAN]).is_err());
625 }
626
627 #[test]
628 fn reset_clears_state() {
629 let mut s = StandardScaler::new(1).unwrap();
630 s.update(&[1.0]).unwrap();
631 s.update(&[2.0]).unwrap();
632 s.reset();
633 assert_eq!(s.counts[0], 0);
634 assert_eq!(s.means()[0], 0.0);
635 }
636
637 #[test]
638 fn transform_into_matches_transform_output() {
639 let mut scaler = StandardScaler::new(4).unwrap();
640 scaler.update(&[1.0, 10.0, 100.0, 1000.0]).unwrap();
642 scaler.update(&[3.0, 20.0, 300.0, 3000.0]).unwrap();
643 scaler.update(&[5.0, 30.0, 500.0, 5000.0]).unwrap();
644 let features = [2.0, 15.0, 200.0, 2000.0];
645 let via_transform = scaler.transform(&features).unwrap();
646 let mut via_into = Vec::new();
647 scaler.transform_into(&features, &mut via_into).unwrap();
648 assert_eq!(via_transform, via_into);
649 }
650
651 #[test]
652 fn transform_into_reuses_buffer_capacity() {
653 let mut scaler = StandardScaler::new(3).unwrap();
654 scaler.update(&[1.0, 2.0, 3.0]).unwrap();
655 scaler.update(&[4.0, 5.0, 6.0]).unwrap();
656 let features = [2.5, 3.5, 4.5];
657 let mut buffer = Vec::with_capacity(64);
658 buffer.extend_from_slice(&[-1.0, -2.0, -3.0, -4.0]);
660 scaler.transform_into(&features, &mut buffer).unwrap();
661 assert_eq!(buffer.len(), 3);
662 assert!(buffer.capacity() >= 64);
664 assert_eq!(buffer, scaler.transform(&features).unwrap());
666 }
667
668 #[test]
669 fn transform_into_rejects_dimension_mismatch() {
670 let scaler = StandardScaler::new(3).unwrap();
671 let mut buffer = Vec::new();
672 assert!(scaler.transform_into(&[1.0, 2.0], &mut buffer).is_err());
673 assert!(buffer.is_empty());
675 }
676
677 #[test]
678 fn transform_into_rejects_non_finite_output() {
679 let scaler = StandardScaler::with_config(
682 1,
683 StandardScalerConfig {
684 with_mean: false,
685 with_std: false,
686 epsilon: 1e-12,
687 },
688 )
689 .unwrap();
690 let mut buffer = Vec::new();
691 assert!(scaler.transform_into(&[f64::NAN], &mut buffer).is_err());
692 }
693
694 #[test]
695 fn transform_into_with_zero_state_returns_original() {
696 let scaler = StandardScaler::new(3).unwrap();
697 let features = [1.5, 2.5, 3.5];
698 let mut buffer = Vec::new();
699 scaler.transform_into(&features, &mut buffer).unwrap();
700 assert_eq!(buffer, vec![1.5, 2.5, 3.5]);
702 }
703}