#[derive(Debug, Clone, PartialEq)]
pub(crate) struct ClassHistogram {
pub counts: Vec<f64>,
}
#[derive(Debug, Clone, PartialEq)]
pub(crate) enum ValueSummary {
Classification(ClassHistogram),
Regression {
sum: f64,
sum_sq: f64,
n: f64,
},
}
impl ValueSummary {
pub fn argmax_class(&self) -> Result<usize> {
match self {
ValueSummary::Classification(h) => {
let mut best = 0usize;
let mut best_count = -1.0f64;
for (k, &c) in h.counts.iter().enumerate() {
if c > best_count {
best = k;
best_count = c;
}
}
if best_count <= 0.0 {
return Err(DatarustError::InvalidInput(
"cannot predict from an empty node".into(),
));
}
Ok(best)
}
ValueSummary::Regression { .. } => Err(DatarustError::InvalidInput(
"classification prediction on a regression tree".into(),
)),
}
}
pub fn regression_mean(&self) -> Result<f64> {
match self {
ValueSummary::Regression { sum, n, .. } => {
if *n <= 0.0 {
return Err(DatarustError::InvalidInput(
"cannot predict from an empty node".into(),
));
}
Ok(sum / *n)
}
ValueSummary::Classification(_) => Err(DatarustError::InvalidInput(
"regression prediction on a classification tree".into(),
)),
}
}
pub fn class_probabilities(&self) -> Result<Vec<f64>> {
match self {
ValueSummary::Classification(h) => {
let n: f64 = h.counts.iter().sum();
if n <= 0.0 {
return Err(DatarustError::InvalidInput(
"cannot produce probabilities from an empty node".into(),
));
}
Ok(h.counts.iter().map(|&c| c / n).collect())
}
ValueSummary::Regression { .. } => Err(DatarustError::InvalidInput(
"predict_proba on a regression tree".into(),
)),
}
}
}
use crate::error::{DatarustError, Result};