rill_ml/preprocessing/
frequency.rs

1//! Online frequency encoder for categorical string features.
2//!
3//! Each category is mapped to its observed frequency `count / total`,
4//! updated incrementally as new samples are seen.
5
6use std::collections::BTreeMap;
7
8use crate::error::RillError;
9
10/// Online frequency encoder for string features.
11///
12/// Maps each category to its observed frequency `count / total`.
13#[derive(Debug, Clone, Default)]
14#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
15pub struct FrequencyEncoder {
16    category_counts: BTreeMap<String, u64>,
17    total: u64,
18    samples_seen: u64,
19}
20
21impl FrequencyEncoder {
22    /// Create a new empty encoder.
23    pub fn new() -> Self {
24        Self {
25            category_counts: BTreeMap::new(),
26            total: 0,
27            samples_seen: 0,
28        }
29    }
30
31    /// The per-category counts.
32    pub fn category_counts(&self) -> &BTreeMap<String, u64> {
33        &self.category_counts
34    }
35
36    /// The total number of category observations seen.
37    pub fn total(&self) -> u64 {
38        self.total
39    }
40
41    /// Return the observed frequency for each input string.
42    ///
43    /// Unknown categories map to `0.0`. Before any sample has been seen,
44    /// all outputs are `0.0`.
45    ///
46    /// # Errors
47    /// Returns [`RillError::EmptyFeatures`] if `features` is empty.
48    pub fn transform_strs(&self, features: &[&str]) -> Result<Vec<f64>, RillError> {
49        if features.is_empty() {
50            return Err(RillError::EmptyFeatures);
51        }
52        if self.total == 0 {
53            return Ok(vec![0.0; features.len()]);
54        }
55        let total = self.total as f64;
56        Ok(features
57            .iter()
58            .map(|&feat| {
59                self.category_counts
60                    .get(feat)
61                    .map(|&c| c as f64 / total)
62                    .unwrap_or(0.0)
63            })
64            .collect())
65    }
66
67    /// Increment counts for each category in `features` and increment
68    /// `samples_seen`.
69    ///
70    /// # Errors
71    /// Returns [`RillError::EmptyFeatures`] if `features` is empty.
72    pub fn update_strs(&mut self, features: &[&str]) -> Result<(), RillError> {
73        if features.is_empty() {
74            return Err(RillError::EmptyFeatures);
75        }
76        for &feat in features {
77            *self.category_counts.entry(feat.to_string()).or_insert(0) += 1;
78        }
79        self.total += features.len() as u64;
80        self.samples_seen += 1;
81        Ok(())
82    }
83
84    /// How many samples have been seen.
85    pub fn samples_seen(&self) -> u64 {
86        self.samples_seen
87    }
88
89    /// Reset the encoder to its initial empty state.
90    pub fn reset(&mut self) {
91        self.category_counts.clear();
92        self.total = 0;
93        self.samples_seen = 0;
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[test]
102    fn frequency_calculation() {
103        let mut enc = FrequencyEncoder::new();
104        // "a" x2, "b" x1 -> total 3
105        enc.update_strs(&["a"]).unwrap();
106        enc.update_strs(&["a", "b"]).unwrap();
107        let out = enc.transform_strs(&["a"]).unwrap();
108        assert!((out[0] - 2.0 / 3.0).abs() < 1e-12);
109        let out = enc.transform_strs(&["b"]).unwrap();
110        assert!((out[0] - 1.0 / 3.0).abs() < 1e-12);
111    }
112
113    #[test]
114    fn unknown_category_returns_zero() {
115        let mut enc = FrequencyEncoder::new();
116        enc.update_strs(&["a"]).unwrap();
117        let out = enc.transform_strs(&["z"]).unwrap();
118        assert_eq!(out, vec![0.0]);
119    }
120
121    #[test]
122    fn multiple_updates_accumulate() {
123        let mut enc = FrequencyEncoder::new();
124        enc.update_strs(&["a", "a"]).unwrap();
125        enc.update_strs(&["a"]).unwrap();
126        // "a" count = 3, total = 3 -> freq = 1.0
127        let out = enc.transform_strs(&["a"]).unwrap();
128        assert!((out[0] - 1.0).abs() < 1e-12);
129    }
130
131    #[test]
132    fn reset_clears_state() {
133        let mut enc = FrequencyEncoder::new();
134        enc.update_strs(&["a", "b"]).unwrap();
135        enc.reset();
136        assert_eq!(enc.total(), 0);
137        assert_eq!(enc.samples_seen(), 0);
138        assert!(enc.category_counts().is_empty());
139    }
140
141    #[test]
142    fn multiple_features_return_one_frequency_each() {
143        let mut enc = FrequencyEncoder::new();
144        // "a" x1, "b" x3 -> total 4
145        enc.update_strs(&["b", "b", "b", "a"]).unwrap();
146        let out = enc.transform_strs(&["a", "b"]).unwrap();
147        assert!((out[0] - 0.25).abs() < 1e-12);
148        assert!((out[1] - 0.75).abs() < 1e-12);
149    }
150
151    #[test]
152    fn total_tracks_observations() {
153        let mut enc = FrequencyEncoder::new();
154        enc.update_strs(&["a", "b"]).unwrap();
155        assert_eq!(enc.total(), 2);
156        enc.update_strs(&["c"]).unwrap();
157        assert_eq!(enc.total(), 3);
158    }
159
160    #[test]
161    #[cfg(feature = "serde")]
162    fn serde_roundtrip() {
163        let mut enc = FrequencyEncoder::new();
164        enc.update_strs(&["a", "b", "a"]).unwrap();
165        let json = serde_json::to_string(&enc).unwrap();
166        let restored: FrequencyEncoder = serde_json::from_str(&json).unwrap();
167        assert_eq!(restored.total(), enc.total());
168        assert_eq!(restored.samples_seen(), enc.samples_seen());
169        assert_eq!(restored.category_counts(), enc.category_counts());
170    }
171}
rill_ml/preprocessing/frequency.rs

rill_ml/preprocessing/
frequency.rs
