rill_ml/preprocessing/
constant_imputer.rs1use crate::error::{RillError, ensure_finite};
7use crate::traits::Transformer;
8
9#[derive(Debug, Clone)]
11#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
12pub struct ConstantImputerConfig {
13 pub fill_value: f64,
15}
16
17impl Default for ConstantImputerConfig {
18 fn default() -> Self {
19 Self { fill_value: 0.0 }
20 }
21}
22
23#[derive(Debug, Clone)]
28#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
29pub struct ConstantImputer {
30 feature_count: usize,
31 config: ConstantImputerConfig,
32 samples_seen: u64,
33}
34
35impl ConstantImputer {
36 pub fn new(feature_count: usize) -> Result<Self, RillError> {
41 Self::with_config(feature_count, ConstantImputerConfig::default())
42 }
43
44 pub fn with_config(
50 feature_count: usize,
51 config: ConstantImputerConfig,
52 ) -> Result<Self, RillError> {
53 if feature_count == 0 {
54 return Err(RillError::EmptyFeatures);
55 }
56 ensure_finite("fill_value", config.fill_value)?;
57 Ok(Self {
58 feature_count,
59 config,
60 samples_seen: 0,
61 })
62 }
63
64 pub fn fill_value(&self) -> f64 {
66 self.config.fill_value
67 }
68
69 fn check_dimension(&self, features: &[f64]) -> Result<(), RillError> {
71 if features.is_empty() {
72 return Err(RillError::EmptyFeatures);
73 }
74 if features.len() != self.feature_count {
75 return Err(RillError::DimensionMismatch {
76 expected: self.feature_count,
77 actual: features.len(),
78 });
79 }
80 Ok(())
81 }
82}
83
84impl Transformer for ConstantImputer {
85 fn input_dim(&self) -> usize {
86 self.feature_count
87 }
88
89 fn output_dim(&self) -> usize {
90 self.feature_count
91 }
92
93 fn transform(&self, features: &[f64]) -> Result<Vec<f64>, RillError> {
94 self.check_dimension(features)?;
95 Ok(features
96 .iter()
97 .map(|&x| {
98 if x.is_nan() {
99 self.config.fill_value
100 } else {
101 x
102 }
103 })
104 .collect())
105 }
106
107 fn update(&mut self, features: &[f64]) -> Result<(), RillError> {
108 self.check_dimension(features)?;
109 self.samples_seen += 1;
110 Ok(())
111 }
112
113 fn samples_seen(&self) -> u64 {
114 self.samples_seen
115 }
116
117 fn reset(&mut self) {
118 self.samples_seen = 0;
119 }
120}
121
122#[cfg(test)]
123mod tests {
124 use super::*;
125
126 #[test]
127 fn nan_replaced_with_fill_value() {
128 let imp = ConstantImputer::new(3).unwrap();
129 let out = imp.transform(&[1.0, f64::NAN, 3.0]).unwrap();
130 assert_eq!(out, vec![1.0, 0.0, 3.0]);
131 }
132
133 #[test]
134 fn non_nan_passed_through() {
135 let imp = ConstantImputer::new(3).unwrap();
136 let out = imp.transform(&[1.5, -2.0, 3.0]).unwrap();
137 assert_eq!(out, vec![1.5, -2.0, 3.0]);
138 }
139
140 #[test]
141 fn custom_fill_value() {
142 let imp =
143 ConstantImputer::with_config(2, ConstantImputerConfig { fill_value: -1.0 }).unwrap();
144 let out = imp.transform(&[f64::NAN, 5.0]).unwrap();
145 assert_eq!(out, vec![-1.0, 5.0]);
146 }
147
148 #[test]
149 fn dimension_mismatch_rejected() {
150 let imp = ConstantImputer::new(3).unwrap();
151 assert!(matches!(
152 imp.transform(&[1.0, 2.0]),
153 Err(RillError::DimensionMismatch { .. })
154 ));
155 let mut imp = imp;
156 assert!(matches!(
157 imp.update(&[1.0, 2.0, 3.0, 4.0]),
158 Err(RillError::DimensionMismatch { .. })
159 ));
160 }
161
162 #[test]
163 fn reset_clears_state() {
164 let mut imp = ConstantImputer::new(2).unwrap();
165 imp.update(&[1.0, f64::NAN]).unwrap();
166 imp.update(&[f64::NAN, 2.0]).unwrap();
167 assert_eq!(imp.samples_seen(), 2);
168 imp.reset();
169 assert_eq!(imp.samples_seen(), 0);
170 }
171
172 #[test]
173 #[cfg(feature = "serde")]
174 fn serde_roundtrip() {
175 let mut imp =
176 ConstantImputer::with_config(3, ConstantImputerConfig { fill_value: 7.0 }).unwrap();
177 imp.update(&[1.0, f64::NAN, 3.0]).unwrap();
178 let json = serde_json::to_string(&imp).unwrap();
179 let restored: ConstantImputer = serde_json::from_str(&json).unwrap();
180 assert_eq!(restored.input_dim(), imp.input_dim());
181 assert_eq!(restored.output_dim(), imp.output_dim());
182 assert_eq!(restored.samples_seen(), imp.samples_seen());
183 assert_eq!(restored.fill_value(), imp.fill_value());
184 }
185}