zsfm_tabfm/ensemble/
cat_encoder.rs

1//! `CategoricalOrdinalEncoder` (`cat_encoder_mode="appearance"`, the wrapper's default): assigns
2//! integer codes to a categorical column's values in order of first appearance in the training
3//! data; unknown/missing values (at fit or transform time) map to `-1`.
4
5use std::collections::HashMap;
6
7use serde_json::Value;
8
9pub struct CategoricalOrdinalEncoder {
10    index: HashMap<String, i64>,
11}
12
13impl CategoricalOrdinalEncoder {
14    pub fn fit(column: &[Value]) -> Self {
15        let mut index = HashMap::new();
16        let mut next = 0i64;
17        for v in column {
18            if is_missing(v) {
19                continue;
20            }
21            let key = value_key(v);
22            index.entry(key).or_insert_with(|| {
23                let code = next;
24                next += 1;
25                code
26            });
27        }
28        CategoricalOrdinalEncoder { index }
29    }
30
31    pub fn transform(&self, column: &[Value]) -> Vec<f64> {
32        column
33            .iter()
34            .map(|v| {
35                if is_missing(v) {
36                    return -1.0;
37                }
38                *self.index.get(&value_key(v)).unwrap_or(&-1) as f64
39            })
40            .collect()
41    }
42}
43
44/// Classification label encoder (`CategoricalOrdinalEncoder(mode="alphabetical")` as used on
45/// `y`): unique labels sorted alphabetically (by their string form) get codes `0..n_classes-1`.
46pub struct LabelEncoder {
47    classes: Vec<String>, // sorted; index = class code
48}
49
50impl LabelEncoder {
51    pub fn fit(y: &[Value]) -> Self {
52        let mut uniq: Vec<String> = y.iter().filter(|v| !is_missing(v)).map(value_key).collect();
53        uniq.sort();
54        uniq.dedup();
55        LabelEncoder { classes: uniq }
56    }
57
58    pub fn n_classes(&self) -> usize {
59        self.classes.len()
60    }
61
62    pub fn transform(&self, y: &[Value]) -> Vec<f64> {
63        y.iter()
64            .map(|v| {
65                if is_missing(v) {
66                    return -1.0;
67                }
68                let key = value_key(v);
69                self.classes
70                    .iter()
71                    .position(|c| c == &key)
72                    .map(|i| i as f64)
73                    .unwrap_or(-1.0)
74            })
75            .collect()
76    }
77
78    pub fn decode(&self, code: usize) -> &str {
79        &self.classes[code]
80    }
81}
82
83fn is_missing(v: &Value) -> bool {
84    match v {
85        Value::Null => true,
86        Value::String(s) => s.is_empty() || s.eq_ignore_ascii_case("nan"),
87        Value::Number(n) => n.as_f64().map(f64::is_nan).unwrap_or(false),
88        _ => false,
89    }
90}
91
92fn value_key(v: &Value) -> String {
93    match v {
94        Value::String(s) => s.clone(),
95        Value::Number(n) => n.to_string(),
96        Value::Bool(b) => b.to_string(),
97        other => other.to_string(),
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104
105    #[test]
106    fn test_appearance_order() {
107        let col = vec![
108            Value::String("blue".into()),
109            Value::String("red".into()),
110            Value::String("blue".into()),
111            Value::String("green".into()),
112        ];
113        let enc = CategoricalOrdinalEncoder::fit(&col);
114        assert_eq!(enc.transform(&col), vec![0.0, 1.0, 0.0, 2.0]);
115    }
116
117    #[test]
118    fn test_unknown_at_transform_time() {
119        let train = vec![Value::String("a".into()), Value::String("b".into())];
120        let enc = CategoricalOrdinalEncoder::fit(&train);
121        let test = vec![Value::String("a".into()), Value::String("c".into())];
122        assert_eq!(enc.transform(&test), vec![0.0, -1.0]);
123    }
124
125    #[test]
126    fn test_missing_values() {
127        let col = vec![
128            Value::String("a".into()),
129            Value::Null,
130            Value::String("a".into()),
131        ];
132        let enc = CategoricalOrdinalEncoder::fit(&col);
133        assert_eq!(enc.transform(&col), vec![0.0, -1.0, 0.0]);
134    }
135}
zsfm_tabfm/ensemble/cat_encoder.rs

zsfm_tabfm/ensemble/
cat_encoder.rs
