Skip to main content

rill_ml/preprocessing/
missing_indicator.rs

1//! Missing value indicator transformer.
2//!
3//! For each input feature, appends a binary indicator (1.0 if NaN, 0.0
4//! otherwise). Output dimension = 2 * input dimension.
5//!
6//! This transformer accepts NaN values in its input, unlike most other
7//! transformers.
8
9use crate::error::RillError;
10use crate::traits::Transformer;
11
12/// Adds a missing-value indicator for each feature.
13///
14/// `transform([x1, x2, ...])` produces `[x1, is_nan(x1) as f64, x2, is_nan(x2) as f64, ...]`.
15/// The original values are preserved (including NaN), with the indicator
16/// appended after each value.
17#[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    /// Create a new indicator for `feature_count` features.
26    ///
27    /// # Errors
28    /// Returns [`RillError::EmptyFeatures`] if `feature_count` is `0`.
29    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    /// Validate only the dimension, allowing NaN values.
40    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}