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, checked_increment, ensure_finite};
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            if !v.is_nan() {
68                ensure_finite("feature", v)?;
69            }
70            out.push(v);
71            out.push(if v.is_nan() { 1.0 } else { 0.0 });
72        }
73        Ok(out)
74    }
75
76    fn update(&mut self, features: &[f64]) -> Result<(), RillError> {
77        self.check_dimension(features)?;
78        for &v in features {
79            if !v.is_nan() {
80                ensure_finite("feature", v)?;
81            }
82        }
83        self.samples_seen = checked_increment(self.samples_seen, "samples_seen")?;
84        Ok(())
85    }
86
87    fn samples_seen(&self) -> u64 {
88        self.samples_seen
89    }
90
91    fn reset(&mut self) {
92        self.samples_seen = 0;
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99
100    #[test]
101    fn no_missing_values() {
102        let mi = MissingIndicator::new(3).unwrap();
103        let out = mi.transform(&[1.0, 2.0, 3.0]).unwrap();
104        assert_eq!(out, vec![1.0, 0.0, 2.0, 0.0, 3.0, 0.0]);
105    }
106
107    #[test]
108    fn all_missing_values() {
109        let mi = MissingIndicator::new(2).unwrap();
110        let out = mi.transform(&[f64::NAN, f64::NAN]).unwrap();
111        assert!(out[0].is_nan());
112        assert_eq!(out[1], 1.0);
113        assert!(out[2].is_nan());
114        assert_eq!(out[3], 1.0);
115    }
116
117    #[test]
118    fn mixed_missing_and_present() {
119        let mi = MissingIndicator::new(3).unwrap();
120        let out = mi.transform(&[1.0, f64::NAN, 3.0]).unwrap();
121        assert_eq!(out[0], 1.0);
122        assert_eq!(out[1], 0.0);
123        assert!(out[2].is_nan());
124        assert_eq!(out[3], 1.0);
125        assert_eq!(out[4], 3.0);
126        assert_eq!(out[5], 0.0);
127    }
128
129    #[test]
130    fn dimension_mismatch_rejected() {
131        let mi = MissingIndicator::new(3).unwrap();
132        assert!(matches!(
133            mi.transform(&[1.0, 2.0]),
134            Err(RillError::DimensionMismatch { .. })
135        ));
136        let mut mi = mi;
137        assert!(matches!(
138            mi.update(&[1.0, 2.0, 3.0, 4.0]),
139            Err(RillError::DimensionMismatch { .. })
140        ));
141    }
142
143    #[test]
144    fn reset_clears_state() {
145        let mut mi = MissingIndicator::new(2).unwrap();
146        mi.update(&[1.0, 2.0]).unwrap();
147        mi.update(&[3.0, 4.0]).unwrap();
148        assert_eq!(mi.samples_seen(), 2);
149        mi.reset();
150        assert_eq!(mi.samples_seen(), 0);
151    }
152
153    #[test]
154    #[cfg(feature = "serde")]
155    fn serde_roundtrip() {
156        let mut mi = MissingIndicator::new(3).unwrap();
157        mi.update(&[1.0, 2.0, 3.0]).unwrap();
158        let json = serde_json::to_string(&mi).unwrap();
159        let restored: MissingIndicator = serde_json::from_str(&json).unwrap();
160        assert_eq!(restored.input_dim(), mi.input_dim());
161        assert_eq!(restored.output_dim(), mi.output_dim());
162        assert_eq!(restored.samples_seen(), mi.samples_seen());
163    }
164}
rill_ml/preprocessing/missing_indicator.rs

rill_ml/preprocessing/
missing_indicator.rs
