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 for c in &mut self.counts {
335 *c = 0;
336 }
337 for m in &mut self.means {
338 *m = 0.0;
339 }
340 for m2 in &mut self.m2s {
341 *m2 = 0.0;
342 }
343 }
344}
345
346#[cfg(feature = "serde")]
347#[derive(serde::Deserialize)]
348struct StandardScalerState {
349 feature_count: usize,
350 config: StandardScalerConfig,
351 counts: Vec<u64>,
352 means: Vec<f64>,
353 m2s: Vec<f64>,
354}
355
356#[cfg(feature = "serde")]
357impl<'de> serde::Deserialize<'de> for StandardScaler {
358 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
359 where
360 D: serde::Deserializer<'de>,
361 {
362 let state = StandardScalerState::deserialize(deserializer)?;
363 let scaler = Self {
364 feature_count: state.feature_count,
365 config: state.config,
366 counts: state.counts,
367 means: state.means,
368 m2s: state.m2s,
369 };
370 scaler.validate().map_err(serde::de::Error::custom)?;
371 Ok(scaler)
372 }
373}
374
375#[cfg(feature = "serde")]
376impl ValidateState for StandardScaler {
377 fn validate_state(&self) -> Result<(), RillError> {
378 StandardScaler::validate(self)
379 }
380}
381
382#[cfg(test)]
383mod tests {
384 use super::*;
385
386 #[test]
387 fn scaler_zero_state_returns_original() {
388 let s = StandardScaler::new(3).unwrap();
389 let out = s.transform(&[1.0, 2.0, 3.0]).unwrap();
390 assert!((out[0] - 1.0).abs() < 1e-12);
392 assert!((out[1] - 2.0).abs() < 1e-12);
393 assert!((out[2] - 3.0).abs() < 1e-12);
394 }
395
396 #[test]
397 fn scaler_standardizes_after_updates() {
398 let mut s = StandardScaler::new(2).unwrap();
399 s.update(&[1.0, 10.0]).unwrap();
402 s.update(&[3.0, 20.0]).unwrap();
403 let out = s.transform(&[3.0, 20.0]).unwrap();
404 assert!((out[0] - 1.0).abs() < 1e-9);
406 assert!((out[1] - 1.0).abs() < 1e-9);
407 }
408
409 #[test]
410 fn transform_does_not_update_state() {
411 let mut s = StandardScaler::new(1).unwrap();
412 s.update(&[10.0]).unwrap();
413 let mean_before = s.means()[0];
414 let _ = s.transform(&[5.0]).unwrap();
415 assert_eq!(s.means()[0], mean_before);
416 assert_eq!(s.counts[0], 1);
417 }
418
419 #[test]
420 fn update_rejects_overflow_without_mutating_state() {
421 let mut scaler = StandardScaler::new(1).unwrap();
422 scaler.update(&[f64::MAX]).unwrap();
423 let before = scaler.clone();
424 assert!(scaler.update(&[-f64::MAX]).is_err());
425 assert_eq!(scaler.counts, before.counts);
426 assert_eq!(scaler.means, before.means);
427 assert_eq!(scaler.m2s, before.m2s);
428 }
429
430 #[cfg(feature = "serde")]
431 #[test]
432 fn serde_rejects_malformed_state() {
433 let malformed = r#"{
434 "feature_count":2,
435 "config":{"with_mean":true,"with_std":true,"epsilon":1e-12},
436 "counts":[1],
437 "means":[0.0],
438 "m2s":[0.0]
439 }"#;
440 assert!(serde_json::from_str::<StandardScaler>(malformed).is_err());
441 }
442
443 #[cfg(feature = "serde")]
444 #[test]
445 fn serde_rejects_negative_m2() {
446 let malformed = r#"{
451 "feature_count":1,
452 "config":{"with_mean":true,"with_std":true,"epsilon":1e-12},
453 "counts":[2],
454 "means":[0.0],
455 "m2s":[-1.0]
456 }"#;
457 assert!(serde_json::from_str::<StandardScaler>(malformed).is_err());
458 }
459
460 #[cfg(feature = "serde")]
461 #[test]
462 fn scaler_serde_rejects_zero_count_nonzero_mean() {
463 let malformed = r#"{
466 "feature_count":1,
467 "config":{"with_mean":true,"with_std":true,"epsilon":1e-12},
468 "counts":[0],
469 "means":[10.0],
470 "m2s":[0.0]
471 }"#;
472 assert!(
473 serde_json::from_str::<StandardScaler>(malformed).is_err(),
474 "count=0 with non-zero mean must be rejected"
475 );
476 }
477
478 #[cfg(feature = "serde")]
479 #[test]
480 fn scaler_serde_rejects_zero_count_nonzero_m2() {
481 let malformed = r#"{
484 "feature_count":1,
485 "config":{"with_mean":true,"with_std":true,"epsilon":1e-12},
486 "counts":[0],
487 "means":[0.0],
488 "m2s":[1.0]
489 }"#;
490 assert!(
491 serde_json::from_str::<StandardScaler>(malformed).is_err(),
492 "count=0 with non-zero m2 must be rejected"
493 );
494 }
495
496 #[cfg(feature = "serde")]
497 #[test]
498 fn scaler_serde_rejects_one_count_nonzero_m2() {
499 let malformed = r#"{
503 "feature_count":1,
504 "config":{"with_mean":true,"with_std":true,"epsilon":1e-12},
505 "counts":[1],
506 "means":[3.0],
507 "m2s":[0.25]
508 }"#;
509 assert!(
510 serde_json::from_str::<StandardScaler>(malformed).is_err(),
511 "count=1 with non-zero m2 must be rejected"
512 );
513 }
514
515 #[cfg(feature = "serde")]
516 #[test]
517 fn scaler_serde_accepts_one_count_finite_mean() {
518 let json = r#"{
522 "feature_count":2,
523 "config":{"with_mean":true,"with_std":true,"epsilon":1e-12},
524 "counts":[1,1],
525 "means":[3.0,-7.5],
526 "m2s":[0.0,0.0]
527 }"#;
528 let scaler: StandardScaler =
529 serde_json::from_str(json).expect("count=1 with finite mean and m2=0 must be accepted");
530 let re = serde_json::to_string(&scaler).unwrap();
532 let _: StandardScaler = serde_json::from_str(&re).unwrap();
533 assert_eq!(scaler.counts, vec![1, 1]);
534 assert_eq!(scaler.means, vec![3.0, -7.5]);
535 assert_eq!(scaler.m2s, vec![0.0, 0.0]);
536 }
537
538 #[cfg(feature = "serde")]
539 #[test]
540 fn scaler_serde_roundtrip_preserves_state() {
541 let mut scaler = StandardScaler::new(2).unwrap();
544 scaler.update(&[1.0, 10.0]).unwrap();
545 scaler.update(&[3.0, 20.0]).unwrap();
546 scaler.update(&[5.0, 30.0]).unwrap();
547 let json = serde_json::to_string(&scaler).unwrap();
548 let restored: StandardScaler = serde_json::from_str(&json).unwrap();
549 assert_eq!(restored.counts, scaler.counts);
550 assert_eq!(restored.means, scaler.means);
551 assert_eq!(restored.m2s, scaler.m2s);
552 let features = [2.0, 15.0];
554 assert_eq!(
555 scaler.transform(&features).unwrap(),
556 restored.transform(&features).unwrap()
557 );
558 }
559
560 #[test]
561 fn transform_hot_path_does_not_scan_invariants() {
562 let mut scaler = StandardScaler::new(3).unwrap();
570 scaler.update(&[1.0, 2.0, 3.0]).unwrap();
571 scaler.update(&[4.0, 5.0, 6.0]).unwrap();
572 let before = scaler.clone();
573 let out = scaler.transform(&[2.5, 3.5, 4.5]).unwrap();
574 assert_eq!(scaler.counts, before.counts);
576 assert_eq!(scaler.means, before.means);
577 assert_eq!(scaler.m2s, before.m2s);
578 assert_eq!(out.len(), 3);
580 }
581
582 #[test]
583 fn constant_feature_uses_scale_one() {
584 let mut s = StandardScaler::new(1).unwrap();
585 for _ in 0..10 {
586 s.update(&[5.0]).unwrap();
587 }
588 let out = s.transform(&[5.0]).unwrap();
590 assert!(out[0].abs() < 1e-12);
591 assert!(!out[0].is_nan());
592 }
593
594 #[test]
595 fn with_mean_false_keeps_offset() {
596 let mut s = StandardScaler::with_config(
597 1,
598 StandardScalerConfig {
599 with_mean: false,
600 with_std: true,
601 epsilon: 1e-12,
602 },
603 )
604 .unwrap();
605 s.update(&[1.0]).unwrap();
606 s.update(&[3.0]).unwrap();
607 let out = s.transform(&[3.0]).unwrap();
609 assert!((out[0] - 3.0).abs() < 1e-9);
610 }
611
612 #[test]
613 fn dimension_mismatch_rejected() {
614 let mut s = StandardScaler::new(3).unwrap();
615 assert!(s.transform(&[1.0, 2.0]).is_err());
616 assert!(s.update(&[1.0, 2.0]).is_err());
617 }
618
619 #[test]
620 fn zero_features_rejected() {
621 assert!(matches!(
622 StandardScaler::new(0),
623 Err(RillError::EmptyFeatures)
624 ));
625 }
626
627 #[test]
628 fn non_finite_rejected() {
629 let mut s = StandardScaler::new(2).unwrap();
630 assert!(s.update(&[1.0, f64::NAN]).is_err());
631 }
632
633 #[test]
634 fn reset_clears_state() {
635 let mut s = StandardScaler::new(1).unwrap();
636 s.update(&[1.0]).unwrap();
637 s.update(&[2.0]).unwrap();
638 s.reset();
639 assert_eq!(s.counts[0], 0);
640 assert_eq!(s.means()[0], 0.0);
641 }
642
643 #[test]
644 fn transform_into_matches_transform_output() {
645 let mut scaler = StandardScaler::new(4).unwrap();
646 scaler.update(&[1.0, 10.0, 100.0, 1000.0]).unwrap();
648 scaler.update(&[3.0, 20.0, 300.0, 3000.0]).unwrap();
649 scaler.update(&[5.0, 30.0, 500.0, 5000.0]).unwrap();
650 let features = [2.0, 15.0, 200.0, 2000.0];
651 let via_transform = scaler.transform(&features).unwrap();
652 let mut via_into = Vec::new();
653 scaler.transform_into(&features, &mut via_into).unwrap();
654 assert_eq!(via_transform, via_into);
655 }
656
657 #[test]
658 fn transform_into_reuses_buffer_capacity() {
659 let mut scaler = StandardScaler::new(3).unwrap();
660 scaler.update(&[1.0, 2.0, 3.0]).unwrap();
661 scaler.update(&[4.0, 5.0, 6.0]).unwrap();
662 let features = [2.5, 3.5, 4.5];
663 let mut buffer = Vec::with_capacity(64);
664 buffer.extend_from_slice(&[-1.0, -2.0, -3.0, -4.0]);
666 scaler.transform_into(&features, &mut buffer).unwrap();
667 assert_eq!(buffer.len(), 3);
668 assert!(buffer.capacity() >= 64);
670 assert_eq!(buffer, scaler.transform(&features).unwrap());
672 }
673
674 #[test]
675 fn transform_into_rejects_dimension_mismatch() {
676 let scaler = StandardScaler::new(3).unwrap();
677 let mut buffer = Vec::new();
678 assert!(scaler.transform_into(&[1.0, 2.0], &mut buffer).is_err());
679 assert!(buffer.is_empty());
681 }
682
683 #[test]
684 fn transform_into_rejects_non_finite_output() {
685 let scaler = StandardScaler::with_config(
688 1,
689 StandardScalerConfig {
690 with_mean: false,
691 with_std: false,
692 epsilon: 1e-12,
693 },
694 )
695 .unwrap();
696 let mut buffer = Vec::new();
697 assert!(scaler.transform_into(&[f64::NAN], &mut buffer).is_err());
698 }
699
700 #[test]
701 fn transform_into_with_zero_state_returns_original() {
702 let scaler = StandardScaler::new(3).unwrap();
703 let features = [1.5, 2.5, 3.5];
704 let mut buffer = Vec::new();
705 scaler.transform_into(&features, &mut buffer).unwrap();
706 assert_eq!(buffer, vec![1.5, 2.5, 3.5]);
708 }
709}