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, checked_increment};
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            let count = self.category_counts.entry(feat.to_string()).or_insert(0);
78            *count = checked_increment(*count, "category_count")?;
79        }
80        self.total = self
81            .total
82            .checked_add(features.len() as u64)
83            .ok_or_else(|| RillError::InvalidState("total counter overflow".to_string()))?;
84        self.samples_seen = checked_increment(self.samples_seen, "samples_seen")?;
85        Ok(())
86    }
87
88    /// How many samples have been seen.
89    pub fn samples_seen(&self) -> u64 {
90        self.samples_seen
91    }
92
93    /// Reset the encoder to its initial empty state.
94    pub fn reset(&mut self) {
95        self.category_counts.clear();
96        self.total = 0;
97        self.samples_seen = 0;
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104
105    #[test]
106    fn frequency_calculation() {
107        let mut enc = FrequencyEncoder::new();
108        // "a" x2, "b" x1 -> total 3
109        enc.update_strs(&["a"]).unwrap();
110        enc.update_strs(&["a", "b"]).unwrap();
111        let out = enc.transform_strs(&["a"]).unwrap();
112        assert!((out[0] - 2.0 / 3.0).abs() < 1e-12);
113        let out = enc.transform_strs(&["b"]).unwrap();
114        assert!((out[0] - 1.0 / 3.0).abs() < 1e-12);
115    }
116
117    #[test]
118    fn unknown_category_returns_zero() {
119        let mut enc = FrequencyEncoder::new();
120        enc.update_strs(&["a"]).unwrap();
121        let out = enc.transform_strs(&["z"]).unwrap();
122        assert_eq!(out, vec![0.0]);
123    }
124
125    #[test]
126    fn multiple_updates_accumulate() {
127        let mut enc = FrequencyEncoder::new();
128        enc.update_strs(&["a", "a"]).unwrap();
129        enc.update_strs(&["a"]).unwrap();
130        // "a" count = 3, total = 3 -> freq = 1.0
131        let out = enc.transform_strs(&["a"]).unwrap();
132        assert!((out[0] - 1.0).abs() < 1e-12);
133    }
134
135    #[test]
136    fn reset_clears_state() {
137        let mut enc = FrequencyEncoder::new();
138        enc.update_strs(&["a", "b"]).unwrap();
139        enc.reset();
140        assert_eq!(enc.total(), 0);
141        assert_eq!(enc.samples_seen(), 0);
142        assert!(enc.category_counts().is_empty());
143    }
144
145    #[test]
146    fn multiple_features_return_one_frequency_each() {
147        let mut enc = FrequencyEncoder::new();
148        // "a" x1, "b" x3 -> total 4
149        enc.update_strs(&["b", "b", "b", "a"]).unwrap();
150        let out = enc.transform_strs(&["a", "b"]).unwrap();
151        assert!((out[0] - 0.25).abs() < 1e-12);
152        assert!((out[1] - 0.75).abs() < 1e-12);
153    }
154
155    #[test]
156    fn total_tracks_observations() {
157        let mut enc = FrequencyEncoder::new();
158        enc.update_strs(&["a", "b"]).unwrap();
159        assert_eq!(enc.total(), 2);
160        enc.update_strs(&["c"]).unwrap();
161        assert_eq!(enc.total(), 3);
162    }
163
164    #[test]
165    #[cfg(feature = "serde")]
166    fn serde_roundtrip() {
167        let mut enc = FrequencyEncoder::new();
168        enc.update_strs(&["a", "b", "a"]).unwrap();
169        let json = serde_json::to_string(&enc).unwrap();
170        let restored: FrequencyEncoder = serde_json::from_str(&json).unwrap();
171        assert_eq!(restored.total(), enc.total());
172        assert_eq!(restored.samples_seen(), enc.samples_seen());
173        assert_eq!(restored.category_counts(), enc.category_counts());
174    }
175}
rill_ml/preprocessing/frequency.rs

rill_ml/preprocessing/
frequency.rs
