Skip to main content

rill_ml/preprocessing/
constant_imputer.rs

1//! Constant value imputer for missing data.
2//!
3//! Replaces `NaN` values with a fixed fill value. Non-NaN values are
4//! passed through unchanged. This imputer has no learnable state.
5
6use crate::error::{RillError, ensure_finite};
7use crate::traits::Transformer;
8
9/// Configuration for [`ConstantImputer`].
10#[derive(Debug, Clone)]
11#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
12pub struct ConstantImputerConfig {
13    /// The value to replace NaN with.
14    pub fill_value: f64,
15}
16
17impl Default for ConstantImputerConfig {
18    fn default() -> Self {
19        Self { fill_value: 0.0 }
20    }
21}
22
23/// Replaces `NaN` values with a constant.
24///
25/// This transformer accepts `NaN` in its input, unlike most other
26/// transformers. The fill value must be finite.
27#[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    /// Create a new imputer for `feature_count` features with default config.
37    ///
38    /// # Errors
39    /// Returns [`RillError::EmptyFeatures`] if `feature_count` is `0`.
40    pub fn new(feature_count: usize) -> Result<Self, RillError> {
41        Self::with_config(feature_count, ConstantImputerConfig::default())
42    }
43
44    /// Create a new imputer with a custom configuration.
45    ///
46    /// # Errors
47    /// Returns [`RillError::EmptyFeatures`] if `feature_count` is `0`.
48    /// Returns [`RillError::NonFiniteValue`] if `fill_value` is not finite.
49    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    /// The fill value used to replace NaN.
65    pub fn fill_value(&self) -> f64 {
66        self.config.fill_value
67    }
68
69    /// Validate only the dimension, allowing NaN values.
70    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}