rill_ml/preprocessing/
forward_fill.rs

1//! Forward-fill imputer for missing data.
2//!
3//! Replaces `NaN` values with the last observed non-NaN value for
4//! each feature. If no value has been seen yet, `NaN` is replaced
5//! with `0.0`.
6
7use crate::error::{RillError, ensure_finite};
8use crate::traits::Transformer;
9
10/// Replaces `NaN` values with the last observed non-NaN value.
11///
12/// This transformer accepts `NaN` in its input. When no value has
13/// been seen for a feature, `0.0` is used as the fill value.
14#[derive(Debug, Clone)]
15#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
16pub struct ForwardFill {
17    feature_count: usize,
18    last_values: Vec<Option<f64>>,
19    samples_seen: u64,
20}
21
22impl ForwardFill {
23    /// Create a new forward-fill imputer for `feature_count` features.
24    ///
25    /// # Errors
26    /// Returns [`RillError::EmptyFeatures`] if `feature_count` is `0`.
27    pub fn new(feature_count: usize) -> Result<Self, RillError> {
28        if feature_count == 0 {
29            return Err(RillError::EmptyFeatures);
30        }
31        Ok(Self {
32            feature_count,
33            last_values: vec![None; feature_count],
34            samples_seen: 0,
35        })
36    }
37
38    /// The per-feature last observed non-NaN values (`None` if unseen).
39    pub fn last_values(&self) -> &[Option<f64>] {
40        &self.last_values
41    }
42
43    /// Validate only the dimension, allowing NaN values.
44    fn check_dimension(&self, features: &[f64]) -> Result<(), RillError> {
45        if features.is_empty() {
46            return Err(RillError::EmptyFeatures);
47        }
48        if features.len() != self.feature_count {
49            return Err(RillError::DimensionMismatch {
50                expected: self.feature_count,
51                actual: features.len(),
52            });
53        }
54        Ok(())
55    }
56}
57
58impl Transformer for ForwardFill {
59    fn input_dim(&self) -> usize {
60        self.feature_count
61    }
62
63    fn output_dim(&self) -> usize {
64        self.feature_count
65    }
66
67    fn transform(&self, features: &[f64]) -> Result<Vec<f64>, RillError> {
68        self.check_dimension(features)?;
69        Ok(features
70            .iter()
71            .enumerate()
72            .map(|(i, &x)| {
73                if x.is_nan() {
74                    self.last_values[i].unwrap_or(0.0)
75                } else {
76                    x
77                }
78            })
79            .collect())
80    }
81
82    fn update(&mut self, features: &[f64]) -> Result<(), RillError> {
83        self.check_dimension(features)?;
84        for (i, &x) in features.iter().enumerate() {
85            if x.is_nan() {
86                continue;
87            }
88            ensure_finite("feature", x)?;
89            self.last_values[i] = Some(x);
90        }
91        self.samples_seen += 1;
92        Ok(())
93    }
94
95    fn samples_seen(&self) -> u64 {
96        self.samples_seen
97    }
98
99    fn reset(&mut self) {
100        for v in &mut self.last_values {
101            *v = None;
102        }
103        self.samples_seen = 0;
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110
111    #[test]
112    fn nan_replaced_with_last_value() {
113        let mut ff = ForwardFill::new(2).unwrap();
114        ff.update(&[1.0, 2.0]).unwrap();
115        ff.update(&[3.0, 4.0]).unwrap();
116        let out = ff.transform(&[f64::NAN, f64::NAN]).unwrap();
117        assert_eq!(out, vec![3.0, 4.0]);
118    }
119
120    #[test]
121    fn nan_replaced_with_zero_when_no_data() {
122        let ff = ForwardFill::new(2).unwrap();
123        let out = ff.transform(&[f64::NAN, f64::NAN]).unwrap();
124        assert_eq!(out, vec![0.0, 0.0]);
125    }
126
127    #[test]
128    fn non_nan_passed_through() {
129        let mut ff = ForwardFill::new(2).unwrap();
130        ff.update(&[5.0, 6.0]).unwrap();
131        let out = ff.transform(&[1.5, -2.0]).unwrap();
132        assert_eq!(out, vec![1.5, -2.0]);
133    }
134
135    #[test]
136    fn last_value_updates_correctly() {
137        let mut ff = ForwardFill::new(2).unwrap();
138        ff.update(&[1.0, 10.0]).unwrap();
139        assert_eq!(ff.last_values()[0], Some(1.0));
140        assert_eq!(ff.last_values()[1], Some(10.0));
141        ff.update(&[2.0, f64::NAN]).unwrap();
142        assert_eq!(ff.last_values()[0], Some(2.0));
143        assert_eq!(ff.last_values()[1], Some(10.0));
144    }
145
146    #[test]
147    fn nan_skipped_in_update() {
148        let mut ff = ForwardFill::new(2).unwrap();
149        ff.update(&[f64::NAN, 5.0]).unwrap();
150        assert_eq!(ff.last_values()[0], None);
151        assert_eq!(ff.last_values()[1], Some(5.0));
152        ff.update(&[7.0, f64::NAN]).unwrap();
153        assert_eq!(ff.last_values()[0], Some(7.0));
154        assert_eq!(ff.last_values()[1], Some(5.0));
155    }
156
157    #[test]
158    fn dimension_mismatch_rejected() {
159        let ff = ForwardFill::new(3).unwrap();
160        assert!(matches!(
161            ff.transform(&[1.0, 2.0]),
162            Err(RillError::DimensionMismatch { .. })
163        ));
164        let mut ff = ff;
165        assert!(matches!(
166            ff.update(&[1.0, 2.0, 3.0, 4.0]),
167            Err(RillError::DimensionMismatch { .. })
168        ));
169    }
170
171    #[test]
172    fn reset_clears_state() {
173        let mut ff = ForwardFill::new(2).unwrap();
174        ff.update(&[1.0, 2.0]).unwrap();
175        ff.update(&[3.0, 4.0]).unwrap();
176        assert_eq!(ff.samples_seen(), 2);
177        ff.reset();
178        assert_eq!(ff.samples_seen(), 0);
179        assert_eq!(ff.last_values()[0], None);
180        assert_eq!(ff.last_values()[1], None);
181    }
182
183    #[test]
184    #[cfg(feature = "serde")]
185    fn serde_roundtrip() {
186        let mut ff = ForwardFill::new(2).unwrap();
187        ff.update(&[1.0, f64::NAN]).unwrap();
188        ff.update(&[3.0, 5.0]).unwrap();
189        let json = serde_json::to_string(&ff).unwrap();
190        let restored: ForwardFill = serde_json::from_str(&json).unwrap();
191        assert_eq!(restored.input_dim(), ff.input_dim());
192        assert_eq!(restored.output_dim(), ff.output_dim());
193        assert_eq!(restored.samples_seen(), ff.samples_seen());
194        assert_eq!(restored.last_values(), ff.last_values());
195    }
196}
rill_ml/preprocessing/forward_fill.rs

rill_ml/preprocessing/
forward_fill.rs
