use datarust::ensemble::{RandomForestClassifier, RandomForestRegressor};
use datarust::traits::{PredictProba, Predictor};
use datarust::{Matrix, MaxFeatures, ParamValue, Params};
fn fixture() -> Matrix {
Matrix::new(vec![
vec![0.0, 0.0, 1.0],
vec![0.0, 1.0, 0.0],
vec![1.0, 0.0, 0.0],
vec![1.0, 1.0, 1.0],
vec![2.0, 0.0, 1.0],
vec![2.0, 1.0, 0.0],
vec![3.0, 0.0, 0.0],
vec![3.0, 1.0, 1.0],
])
.unwrap()
}
#[test]
fn classifier_matches_sklearn_reference_predictions_probabilities_and_importances() {
let x = fixture();
let y = vec![0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0];
let mut model = RandomForestClassifier::new()
.with_n_estimators(5)
.with_max_depth(Some(3))
.with_max_features(MaxFeatures::Count(2))
.with_random_state(13);
model.fit(&x, &y).unwrap();
assert_eq!(model.predict(&x).unwrap(), y);
assert_eq!(
model.predict_proba(&x).unwrap().as_slice(),
&[0.8, 0.2, 0.2, 0.8, 0.0, 1.0, 1.0, 0.0, 1.0, 0.0, 0.2, 0.8, 0.4, 0.6, 1.0, 0.0,]
);
let expected_importances: [f64; 3] =
[0.13142857142857145, 0.04571428571428573, 0.8228571428571428];
for (actual, expected) in model
.feature_importances()
.unwrap()
.iter()
.zip(expected_importances)
{
assert_eq!(actual.to_bits(), expected.to_bits());
}
}
#[test]
fn regressor_matches_sklearn_reference_predictions_importances_and_bootstrap_indices() {
let x = fixture();
let y = vec![0.5, 2.0, 1.0, 2.5, 1.5, 3.0, 2.0, 3.5];
let mut model = RandomForestRegressor::new()
.with_n_estimators(5)
.with_max_depth(Some(3))
.with_max_features(MaxFeatures::Count(2))
.with_random_state(13);
model.fit(&x, &y).unwrap();
let expected_predictions: [f64; 8] = [0.9, 2.2, 1.0, 2.4, 1.4, 2.8, 1.5, 3.5];
for (actual, expected) in model.predict(&x).unwrap().iter().zip(expected_predictions) {
assert_eq!(actual.to_bits(), expected.to_bits());
}
let expected_importances: [f64; 3] =
[0.24508660036871355, 0.6780638747828597, 0.07684952484842672];
for (actual, expected) in model
.feature_importances()
.unwrap()
.iter()
.zip(expected_importances)
{
assert_eq!(actual.to_bits(), expected.to_bits());
}
assert_eq!(
model.estimators_samples(),
vec![
vec![7, 4, 4, 3, 1, 2, 7, 6],
vec![3, 6, 3, 2, 1, 7, 5, 1],
vec![3, 5, 7, 2, 2, 4, 5, 0],
vec![0, 1, 7, 5, 3, 0, 4, 2],
vec![1, 7, 2, 4, 5, 4, 5, 1],
]
);
}
#[test]
fn bootstrap_false_uses_each_training_row_once_and_params_round_trip() {
let x = fixture();
let y = vec![0.5, 2.0, 1.0, 2.5, 1.5, 3.0, 2.0, 3.5];
let mut model = RandomForestRegressor::new()
.with_n_estimators(3)
.with_bootstrap(false)
.with_max_features(MaxFeatures::Count(2));
model.fit(&x, &y).unwrap();
assert!(model
.estimators_samples()
.iter()
.all(|samples| samples == &(0..x.nrows()).collect::<Vec<_>>()));
for params in model.get_params() {
model.set_params(params.0, params.1).unwrap();
}
assert!(model
.get_params()
.contains(&("max_features", ParamValue::Int(4))));
}
#[test]
fn classifier_accepts_one_class_and_returns_one_probability_column() {
let x = fixture();
let y = vec![7.0; 8];
let mut model = RandomForestClassifier::new()
.with_n_estimators(3)
.with_random_state(5);
model.fit(&x, &y).unwrap();
assert_eq!(model.classes(), Some(&[7.0][..]));
let proba = model.predict_proba(&x).unwrap();
assert_eq!(proba.ncols(), 1);
assert!(proba.as_slice().iter().all(|&value| value == 1.0));
assert_eq!(model.predict(&x).unwrap(), y);
}
#[test]
fn classifier_is_deterministic_and_preserves_external_class_labels() {
let x = fixture();
let y = vec![10.0, 20.0, 20.0, 10.0, 10.0, 20.0, 20.0, 10.0];
let build = || {
let mut model = RandomForestClassifier::new()
.with_n_estimators(7)
.with_max_features(MaxFeatures::Count(1))
.with_random_state(99);
model.fit(&x, &y).unwrap();
model
};
let a = build();
let b = build();
assert_eq!(a.classes(), Some(&[10.0, 20.0][..]));
assert_eq!(a.predict(&x).unwrap(), b.predict(&x).unwrap());
assert_eq!(a.predict_proba(&x).unwrap(), b.predict_proba(&x).unwrap());
}
#[test]
fn classifier_rejects_regression_criterion_at_fit() {
use datarust::tree::Criterion;
use datarust::DatarustError;
let x = fixture();
let y = vec![0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0];
let mut model = RandomForestClassifier::new()
.with_n_estimators(3)
.with_criterion(Criterion::SquaredError);
let err = model.fit(&x, &y).unwrap_err();
assert!(matches!(err, DatarustError::InvalidInput(_)), "got {err:?}");
assert_eq!(model.n_estimators(), 0);
}