use std::collections::BTreeMap;
use prospicio_core::{Error, Result};
use crate::resample::Split;
#[derive(Debug, Clone, PartialEq)]
pub struct TargetEncoding {
pub levels: Vec<String>,
pub ratios: Vec<f64>,
pub k: f64,
}
impl TargetEncoding {
pub fn encode(&self, values: &[String]) -> Vec<f64> {
values
.iter()
.map(|v| match self.levels.binary_search(v) {
Ok(i) => self.ratios[i],
Err(_) => 1.0,
})
.collect()
}
}
pub fn target_encode(
values: &[String],
y: &[f64],
expected: &[f64],
exposure: Option<&[f64]>,
splits: &[Split],
k: f64,
) -> Result<(Vec<f64>, TargetEncoding)> {
let n = values.len();
for (len, name) in [(y.len(), "y"), (expected.len(), "expected")] {
if len != n {
return Err(Error::InvalidParameter {
name,
value: len as f64,
reason: "must have one value per row",
});
}
}
if let Some(e) = exposure {
if e.len() != n {
return Err(Error::InvalidParameter {
name: "exposure",
value: e.len() as f64,
reason: "must have one value per row",
});
}
}
if !(k.is_finite() && k >= 0.0) {
return Err(Error::InvalidParameter {
name: "k",
value: k,
reason: "must be finite and non-negative",
});
}
let mut seen = vec![false; n];
for s in splits {
for &i in &s.test {
if i >= n || seen[i] {
return Err(Error::InvalidParameter {
name: "splits",
value: i as f64,
reason: "every row must be in exactly one test set",
});
}
seen[i] = true;
}
}
if let Some(i) = seen.iter().position(|s| !s) {
return Err(Error::InvalidParameter {
name: "splits",
value: i as f64,
reason: "every row must be in exactly one test set",
});
}
let e = |i: usize| exposure.map_or(1.0, |e| e[i]);
let stats = |rows: &mut dyn Iterator<Item = usize>| {
let mut by: BTreeMap<&str, (f64, f64, f64)> = BTreeMap::new();
for i in rows {
let s = by.entry(values[i].as_str()).or_default();
s.0 += y[i];
s.1 += expected[i];
s.2 += e(i);
}
by
};
let ratio = |(actual, expected, exposure): (f64, f64, f64)| {
let z = if exposure + k > 0.0 {
exposure / (exposure + k)
} else {
0.0
};
if expected == 0.0 {
1.0
} else {
z * actual / expected + (1.0 - z)
}
};
let mut rows = vec![1.0; n];
for s in splits {
let by = stats(&mut s.train.iter().copied());
for &i in &s.test {
rows[i] = by.get(values[i].as_str()).copied().map_or(1.0, ratio);
}
}
let by = stats(&mut (0..n));
let encoding = TargetEncoding {
levels: by.keys().map(|l| l.to_string()).collect(),
ratios: by.values().copied().map(ratio).collect(),
k,
};
Ok((rows, encoding))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::resample::k_fold;
#[test]
fn a_row_never_sees_its_own_response() {
let values: Vec<String> = (0..60).map(|i| format!("l{}", i % 7)).collect();
let y: Vec<f64> = (0..60).map(|i| ((i * 13) % 5) as f64).collect();
let expected = vec![2.0; 60];
let splits = k_fold(60, 5, 3).unwrap();
let (rows, all) = target_encode(&values, &y, &expected, None, &splits, 4.0).unwrap();
let mut y2 = y.clone();
y2[10] += 100.0;
let (rows2, _) = target_encode(&values, &y2, &expected, None, &splits, 4.0).unwrap();
assert_eq!(rows[10], rows2[10]);
assert!(rows2.iter().zip(&rows).any(|(a, b)| a != b));
let (_, plain) = target_encode(&values, &y, &expected, None, &splits, 0.0).unwrap();
let a: f64 = (0..60).filter(|i| i % 7 == 0).map(|i| y[i]).sum();
let n = (0..60).filter(|i| i % 7 == 0).count() as f64;
assert!((plain.encode(&["l0".into()])[0] - a / (2.0 * n)).abs() < 1e-12);
assert_eq!(all.encode(&["new".into()]), [1.0]);
assert!(target_encode(&values, &y, &expected, None, &splits[1..], 4.0).is_err());
}
}