rill_ml/preprocessing/
missing_indicator.rs1use crate::error::RillError;
10use crate::traits::Transformer;
11
12#[derive(Debug, Clone)]
18#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
19pub struct MissingIndicator {
20 feature_count: usize,
21 samples_seen: u64,
22}
23
24impl MissingIndicator {
25 pub fn new(feature_count: usize) -> Result<Self, RillError> {
30 if feature_count == 0 {
31 return Err(RillError::EmptyFeatures);
32 }
33 Ok(Self {
34 feature_count,
35 samples_seen: 0,
36 })
37 }
38
39 fn check_dimension(&self, features: &[f64]) -> Result<(), RillError> {
41 if features.is_empty() {
42 return Err(RillError::EmptyFeatures);
43 }
44 if features.len() != self.feature_count {
45 return Err(RillError::DimensionMismatch {
46 expected: self.feature_count,
47 actual: features.len(),
48 });
49 }
50 Ok(())
51 }
52}
53
54impl Transformer for MissingIndicator {
55 fn input_dim(&self) -> usize {
56 self.feature_count
57 }
58
59 fn output_dim(&self) -> usize {
60 self.feature_count * 2
61 }
62
63 fn transform(&self, features: &[f64]) -> Result<Vec<f64>, RillError> {
64 self.check_dimension(features)?;
65 let mut out = Vec::with_capacity(features.len() * 2);
66 for &v in features {
67 out.push(v);
68 out.push(if v.is_nan() { 1.0 } else { 0.0 });
69 }
70 Ok(out)
71 }
72
73 fn update(&mut self, features: &[f64]) -> Result<(), RillError> {
74 self.check_dimension(features)?;
75 self.samples_seen += 1;
76 Ok(())
77 }
78
79 fn samples_seen(&self) -> u64 {
80 self.samples_seen
81 }
82
83 fn reset(&mut self) {
84 self.samples_seen = 0;
85 }
86}
87
88#[cfg(test)]
89mod tests {
90 use super::*;
91
92 #[test]
93 fn no_missing_values() {
94 let mi = MissingIndicator::new(3).unwrap();
95 let out = mi.transform(&[1.0, 2.0, 3.0]).unwrap();
96 assert_eq!(out, vec![1.0, 0.0, 2.0, 0.0, 3.0, 0.0]);
97 }
98
99 #[test]
100 fn all_missing_values() {
101 let mi = MissingIndicator::new(2).unwrap();
102 let out = mi.transform(&[f64::NAN, f64::NAN]).unwrap();
103 assert!(out[0].is_nan());
104 assert_eq!(out[1], 1.0);
105 assert!(out[2].is_nan());
106 assert_eq!(out[3], 1.0);
107 }
108
109 #[test]
110 fn mixed_missing_and_present() {
111 let mi = MissingIndicator::new(3).unwrap();
112 let out = mi.transform(&[1.0, f64::NAN, 3.0]).unwrap();
113 assert_eq!(out[0], 1.0);
114 assert_eq!(out[1], 0.0);
115 assert!(out[2].is_nan());
116 assert_eq!(out[3], 1.0);
117 assert_eq!(out[4], 3.0);
118 assert_eq!(out[5], 0.0);
119 }
120
121 #[test]
122 fn dimension_mismatch_rejected() {
123 let mi = MissingIndicator::new(3).unwrap();
124 assert!(matches!(
125 mi.transform(&[1.0, 2.0]),
126 Err(RillError::DimensionMismatch { .. })
127 ));
128 let mut mi = mi;
129 assert!(matches!(
130 mi.update(&[1.0, 2.0, 3.0, 4.0]),
131 Err(RillError::DimensionMismatch { .. })
132 ));
133 }
134
135 #[test]
136 fn reset_clears_state() {
137 let mut mi = MissingIndicator::new(2).unwrap();
138 mi.update(&[1.0, 2.0]).unwrap();
139 mi.update(&[3.0, 4.0]).unwrap();
140 assert_eq!(mi.samples_seen(), 2);
141 mi.reset();
142 assert_eq!(mi.samples_seen(), 0);
143 }
144
145 #[test]
146 #[cfg(feature = "serde")]
147 fn serde_roundtrip() {
148 let mut mi = MissingIndicator::new(3).unwrap();
149 mi.update(&[1.0, 2.0, 3.0]).unwrap();
150 let json = serde_json::to_string(&mi).unwrap();
151 let restored: MissingIndicator = serde_json::from_str(&json).unwrap();
152 assert_eq!(restored.input_dim(), mi.input_dim());
153 assert_eq!(restored.output_dim(), mi.output_dim());
154 assert_eq!(restored.samples_seen(), mi.samples_seen());
155 }
156}