diff --git a/CHANGELOG.md b/CHANGELOG.md index 75dba18c..b2b80cb2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] - In the `...Regressor`s the features are always shuffled even when `mtry` equals the number of attributes. This changes the constructed trees. +- `tree`: `BaseTreeRegressor` (used by `DecisionTreeRegressor`, `RandomForestRegressor` and `ExtraTreesRegressor`) uses a new internal `Criterion` trait with `MSECriterion`. The MSE split improvement uses a numerically stable formula. The selected splits do not change when a constant is added to y. Exactly tied splits can be broken in a different way, so fitted trees can differ slightly from earlier versions. +- `tree`: in the `...Regressor`s, rows with sample weight 0 are removed before the tree grows. They no longer count toward `min_samples_leaf` and `min_samples_split`, and they no longer move a split threshold. A tree fitted with a zero-weight row is the same as a tree fitted without that row. +- `tree`, `ensemble`: `DecisionTreeRegressor`, `RandomForestRegressor` and `ExtraTreesRegressor` get a `compute_feature_importances` method. +- `tree`, `ensemble`: the fitted `DecisionTreeRegressor`, `RandomForestRegressor` and `ExtraTreesRegressor` store new fields that the feature importances use. Models serialized with an earlier version cannot be deserialized. Fit these models again. ### Added - `make_blobs_with_seed`, `make_circles_with_seed` and `make_moons_with_seed` take a `seed: Option` for the generator RNG. With `Some(seed)` the output is reproducible with all features. diff --git a/src/ensemble/base_forest_regressor.rs b/src/ensemble/base_forest.rs similarity index 69% rename from src/ensemble/base_forest_regressor.rs rename to src/ensemble/base_forest.rs index b48f8dad..905433c0 100644 --- a/src/ensemble/base_forest_regressor.rs +++ b/src/ensemble/base_forest.rs @@ -11,14 +11,15 @@ use crate::numbers::basenum::Number; use crate::numbers::floatnum::FloatNumber; use crate::rand_custom::get_rng_impl; -use crate::tree::base_tree_regressor::{BaseTreeRegressor, BaseTreeRegressorParameters, Splitter}; +use crate::tree::base_tree::{BaseTree, BaseTreeParameters, Splitter}; +use crate::tree::criterion::{Criterion, LeafOutput}; #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[derive(Debug, Clone)] -/// Parameters of the Forest Regressor +/// Parameters of a base forest. /// Some parameters here are passed directly into base estimator. #[must_use] -pub struct BaseForestRegressorParameters { +pub struct BaseForestParameters { #[cfg_attr(feature = "serde", serde(default))] /// Tree max depth. See [Decision Tree Regressor](../../tree/decision_tree_regressor/index.html) pub max_depth: Option, @@ -46,8 +47,13 @@ pub struct BaseForestRegressorParameters { pub splitter: Splitter, } -impl, Y: Array1> PartialEq - for BaseForestRegressor +impl< + TX: Number + FloatNumber + PartialOrd, + TY: Number, + X: Array2, + Y: Array1, + Out: LeafOutput, +> PartialEq for BaseForest { fn eq(&self, other: &Self) -> bool { if self.trees.as_ref().unwrap().len() != other.trees.as_ref().unwrap().len() { @@ -61,32 +67,54 @@ impl, Y: Array1 } } -/// Forest Regressor +/// Base forest of decision trees. Out is the type of the leaf values. The criterion is passed to fit. #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +#[cfg_attr( + feature = "serde", + serde(bound(serialize = "Out: Serialize", deserialize = "Out: Deserialize<'de>")) +)] #[derive(Debug)] -pub struct BaseForestRegressor< +pub struct BaseForest< TX: Number + FloatNumber + PartialOrd, TY: Number, X: Array2, Y: Array1, + Out: LeafOutput, > { - trees: Option>>, + trees: Option>>, + num_features: Option, samples: Option>>, } -impl, Y: Array1> - BaseForestRegressor +impl< + TX: Number + FloatNumber + PartialOrd, + TY: Number, + X: Array2, + Y: Array1, + Out: LeafOutput, +> BaseForest { + /// The trees of the forest. Call only after `fit`. + pub(crate) fn trees(&self) -> &[BaseTree] { + self.trees.as_ref().unwrap() + } + + /// For each tree, the rows it used in training. `None` if `keep_samples` was false. + pub(crate) fn samples(&self) -> Option<&[Vec]> { + self.samples.as_deref() + } + /// Build a forest of trees from the training set. /// * `x` - _NxM_ matrix with _N_ observations and _M_ features in each observation. /// * `y` - the target class values /// * `sample_weights`: optional sample_weights to use during fitting - pub fn fit( + pub fn fit>( x: &X, y: &Y, sample_weights: Option<&[f64]>, - parameters: BaseForestRegressorParameters, - ) -> Result, Failed> { + criterion: &C, + parameters: BaseForestParameters, + ) -> Result, Failed> { let (n_rows, num_attributes) = x.shape(); if n_rows != y.shape() { @@ -105,7 +133,7 @@ impl, Y: Array1 let mut rng = get_rng_impl(Some(parameters.seed)); let n_trees = parameters.n_trees; - let mut trees: Vec> = Vec::with_capacity(n_trees); + let mut trees: Vec> = Vec::with_capacity(n_trees); let mut maybe_all_samples: Option>> = Option::None; if parameters.keep_samples { @@ -131,7 +159,7 @@ impl, Y: Array1 for tree_idx in 0..parameters.n_trees { if parameters.bootstrap { - samples = BaseForestRegressor::::sample_with_replacement( + samples = BaseForest::::sample_with_replacement( n_rows, &mut rng, dist.as_ref(), @@ -143,7 +171,7 @@ impl, Y: Array1 all_samples.push(samples.iter().map(|x| *x != 0).collect()) } - let params = BaseTreeRegressorParameters { + let params = BaseTreeParameters { max_depth: parameters.max_depth, min_samples_leaf: parameters.min_samples_leaf, min_samples_split: parameters.min_samples_split, @@ -156,117 +184,152 @@ impl, Y: Array1 } else { sample_weights }; - let tree = BaseTreeRegressor::fit_weak_learner( + let tree = BaseTree::fit_weak_learner( x, y, sample_weights_for_base_tree, samples.clone(), mtry, &order, + criterion, params, )?; trees.push(tree); } - Ok(BaseForestRegressor { + Ok(BaseForest { trees: Some(trees), + num_features: Some(x.shape().1), samples: maybe_all_samples, }) } - /// Predict class for `x` - /// * `x` - _KxM_ data where _K_ is number of observations and _M_ is number of features. - pub fn predict(&self, x: &X) -> Result { - let mut result = Y::zeros(x.shape().0); - - let (n, _) = x.shape(); - - for i in 0..n { - result.set(i, self.predict_for_row(x, i)); + fn sample_with_replacement( + nrows: usize, + rng: &mut impl rand::Rng, + distribution: Option<&rand::distr::weighted::WeightedIndex>, + ) -> Vec { + let mut samples = vec![0; nrows]; + for _ in 0..nrows { + let xi = match distribution { + Some(dist) => rng.sample(dist), + None => rng.random_range(0..nrows), + }; + samples[xi] += 1; } - Ok(result) + samples } - fn predict_for_row(&self, x: &X, row: usize) -> TY { - let n_trees = self.trees.as_ref().unwrap().len(); + /// Mean of the normalized tree importances, normalized again. + pub fn compute_feature_importances(&self) -> Result, Failed> { + match (&self.trees, self.num_features) { + (Some(trees), Some(num_features)) => { + let mut importances = vec![0.0f64; num_features]; + let mut scratch = vec![0.0f64; num_features]; + + for tree in trees { + tree.compute_feature_importances(true, &mut scratch); + for (i, s) in importances.iter_mut().zip(scratch.iter()) { + *i += s; + } + } - let mut result = TY::zero(); + // compute the mean of each feature importance + let num_trees = trees.len() as f64; + for i in importances.iter_mut() { + *i /= num_trees; + } - for tree in self.trees.as_ref().unwrap().iter() { - result += tree.predict_for_row(x, row); + // normalize once more + let sum: f64 = importances.iter().sum(); + if sum > 0.0 { + for i in importances.iter_mut() { + *i /= sum; + } + } + + Ok(importances) + } + _ => Err(Failed::predict( + "'fit' should be called before calling 'compute_feature_importances'", + )), } + } +} - result / TY::from_usize(n_trees).unwrap() +/// Mean of the tree predictions for each row of `x`. +pub(crate) fn predict_mean( + forest: &BaseForest, + x: &X, +) -> Result +where + TX: Number + FloatNumber, + TY: Number, + X: Array2, + Y: Array1, +{ + let (n, _) = x.shape(); + let mut result = Y::zeros(n); + let num_trees = TY::from_usize(forest.trees().len()).unwrap(); + + for row in 0..n { + let mut sum = TY::zero(); + for tree in forest.trees() { + sum += TY::from_f64(*tree.leaf_output_for_row(x, row)).unwrap(); + } + result.set(row, sum / num_trees); } + Ok(result) +} - /// Predict OOB classes for `x`. `x` is expected to be equal to the dataset used in training. - pub fn predict_oob(&self, x: &X) -> Result { - let (n, _) = x.shape(); - - let samples = match &self.samples { - Some(s) => s, - None => { - return Err(Failed::because( - FailedError::PredictFailed, - "Need samples=true for OOB predictions.", - )); - } - }; +/// Mean of the out-of-bag tree predictions for each row of `x`. +/// `x` is expected to be equal to the dataset used in training. +pub(crate) fn predict_mean_oob( + forest: &BaseForest, + x: &X, +) -> Result +where + TX: Number + FloatNumber, + TY: Number, + X: Array2, + Y: Array1, +{ + let (n, _) = x.shape(); - if samples[0].len() != n { + let samples = match forest.samples() { + Some(s) => s, + None => { return Err(Failed::because( FailedError::PredictFailed, - "Prediction matrix must match matrix used in training for OOB predictions.", + "Need samples=true for OOB predictions.", )); } + }; - let mut result = Y::zeros(n); - - for i in 0..n { - result.set(i, self.predict_for_row_oob(x, i)); - } - - Ok(result) + if samples[0].len() != n { + return Err(Failed::because( + FailedError::PredictFailed, + "Prediction matrix must match matrix used in training for OOB predictions.", + )); } - fn predict_for_row_oob(&self, x: &X, row: usize) -> TY { - let mut n_trees = 0; - let mut result = TY::zero(); + let mut result = Y::zeros(n); - for (tree, samples) in self - .trees - .as_ref() - .unwrap() - .iter() - .zip(self.samples.as_ref().unwrap()) - { - if !samples[row] { - result += tree.predict_for_row(x, row); + for row in 0..n { + let mut sum = TY::zero(); + let mut n_trees = 0usize; + for (tree, in_bag) in forest.trees().iter().zip(samples) { + if !in_bag[row] { + sum += TY::from_f64(*tree.leaf_output_for_row(x, row)).unwrap(); n_trees += 1; } } - // TODO: What to do if there are no oob trees? - result / TY::from(n_trees).unwrap() + result.set(row, sum / TY::from_usize(n_trees).unwrap()); } - fn sample_with_replacement( - nrows: usize, - rng: &mut impl rand::Rng, - distribution: Option<&rand::distr::weighted::WeightedIndex>, - ) -> Vec { - let mut samples = vec![0; nrows]; - for _ in 0..nrows { - let xi = match distribution { - Some(dist) => rng.sample(dist), - None => rng.random_range(0..nrows), - }; - samples[xi] += 1; - } - - samples - } + Ok(result) } #[cfg(test)] @@ -275,12 +338,13 @@ mod tests { use super::*; use crate::linalg::basic::arrays::Array; use crate::linalg::basic::matrix::DenseMatrix; + use crate::tree::criterion::MSECriterion; #[test] - fn test_base_forest_regressor_keep_samples() { + fn test_base_forest_keep_samples() { let x = DenseMatrix::from_2d_array(&[&[1.0, 2.0], &[3.0, 4.0], &[5.0, 6.0]]).unwrap(); let y = vec![1.0, 2.0, 3.0]; - let params = BaseForestRegressorParameters { + let params = BaseForestParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 2, @@ -289,9 +353,9 @@ mod tests { keep_samples: true, seed: 42, bootstrap: true, - splitter: crate::tree::base_tree_regressor::Splitter::Best, + splitter: crate::tree::base_tree::Splitter::Best, }; - let regressor = BaseForestRegressor::fit(&x, &y, None, params).unwrap(); + let regressor = BaseForest::fit(&x, &y, None, &MSECriterion, params).unwrap(); assert_eq!(regressor.trees.unwrap().len(), 5); assert!(regressor.samples.is_some()); } @@ -304,11 +368,12 @@ mod tests { assert_eq!(empty.shape(), (0, 2)); let y: Vec = vec![]; - let result = BaseForestRegressor::fit( + let result = BaseForest::fit( &empty, &y, None, - BaseForestRegressorParameters { + &MSECriterion, + BaseForestParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 2, @@ -317,7 +382,7 @@ mod tests { keep_samples: false, seed: 0, bootstrap: true, - splitter: crate::tree::base_tree_regressor::Splitter::Best, + splitter: crate::tree::base_tree::Splitter::Best, }, ); assert!(result.is_err()); @@ -332,11 +397,12 @@ mod tests { assert_eq!(no_features.shape(), (2, 0)); let y: Vec = vec![1.0, 2.0]; - let result = BaseForestRegressor::fit( + let result = BaseForest::fit( &no_features, &y, None, - BaseForestRegressorParameters { + &MSECriterion, + BaseForestParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 2, @@ -345,7 +411,7 @@ mod tests { keep_samples: false, seed: 0, bootstrap: true, - splitter: crate::tree::base_tree_regressor::Splitter::Best, + splitter: crate::tree::base_tree::Splitter::Best, }, ); assert!(result.is_err()); @@ -362,7 +428,7 @@ mod tests { let model_parameters = (0..=9).map(|x| x as f64).collect::>(); let y: Vec = model_parameters.xa(true, &x); - let forest_parameters = BaseForestRegressorParameters { + let forest_parameters = BaseForestParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 2, @@ -371,12 +437,12 @@ mod tests { keep_samples: true, seed: 42, bootstrap: false, - splitter: crate::tree::base_tree_regressor::Splitter::Best, + splitter: crate::tree::base_tree::Splitter::Best, }; - let forest = BaseForestRegressor::fit(&x, &y, None, forest_parameters.clone()) + let forest = BaseForest::fit(&x, &y, None, &MSECriterion, forest_parameters.clone()) .expect("Fit should work"); - let y_hat = forest.predict(&x).expect("Predict should work"); + let y_hat = predict_mean(&forest, &x).expect("Predict should work"); assert!((y_hat.iter().sum::() - y.iter().sum::()).abs() < 1e-9); // Seeded RNG: the test gives the same result on each run @@ -384,10 +450,15 @@ mod tests { // Positive weights in [0.5, 2.0) let sample_weights: Vec = (0..1000).map(|_| rng.random_range(0.5..2.0)).collect(); - let forest = - BaseForestRegressor::fit(&x, &y, Some(&sample_weights), forest_parameters.clone()) - .expect("Fit should work"); - let y_hat = forest.predict(&x).expect("Predict should work"); + let forest = BaseForest::fit( + &x, + &y, + Some(&sample_weights), + &MSECriterion, + forest_parameters.clone(), + ) + .expect("Fit should work"); + let y_hat = predict_mean(&forest, &x).expect("Predict should work"); let s: f64 = sample_weights.iter().sum(); let normalized_weights = sample_weights.iter().map(|w| *w / s).collect::>(); @@ -416,7 +487,7 @@ mod tests { let sample_weights: Vec = (0..n_rows).map(|i| 1.0 + (i % 4) as f64).collect(); for weights in [None, Some(sample_weights.as_slice())] { - let params = BaseForestRegressorParameters { + let params = BaseForestParameters { max_depth: Some(1), min_samples_leaf: 1, min_samples_split: 2, @@ -425,9 +496,9 @@ mod tests { keep_samples: true, // keep samples used for each tree, so we can check that they are different seed: 42, bootstrap: true, // Use bootstrapping - splitter: crate::tree::base_tree_regressor::Splitter::Best, + splitter: crate::tree::base_tree::Splitter::Best, }; - let regressor = BaseForestRegressor::fit(&x, &y, weights, params).unwrap(); + let regressor = BaseForest::fit(&x, &y, weights, &MSECriterion, params).unwrap(); let samples = regressor.samples.unwrap(); for (t, in_bag) in samples.iter().enumerate() { @@ -462,7 +533,7 @@ mod tests { let y: Vec = (0..n_rows).map(|i| ((i * 31) % 17) as f64).collect(); for splitter in [Splitter::Best, Splitter::Random] { - let params = BaseForestRegressorParameters { + let params = BaseForestParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 2, @@ -473,7 +544,7 @@ mod tests { bootstrap: false, splitter: splitter.clone(), }; - let forest = BaseForestRegressor::fit(&x, &y, None, params).unwrap(); + let forest = BaseForest::fit(&x, &y, None, &MSECriterion, params).unwrap(); let trees = forest.trees.unwrap(); assert!( trees.iter().any(|tree| tree != &trees[0]), @@ -491,7 +562,7 @@ mod tests { // Rows with y = 10 have weight 9, rows with y = 0 have weight 1. let sample_weights: Vec = (0..20).map(|i| if i < 10 { 1.0 } else { 9.0 }).collect(); - let parameters = BaseForestRegressorParameters { + let parameters = BaseForestParameters { max_depth: Some(0), min_samples_leaf: 1, min_samples_split: 2, @@ -500,12 +571,18 @@ mod tests { keep_samples: true, seed: 42, bootstrap: false, // No bootstrapping - splitter: crate::tree::base_tree_regressor::Splitter::Best, + splitter: crate::tree::base_tree::Splitter::Best, }; - let forest = BaseForestRegressor::fit(&x, &y, Some(&sample_weights), parameters.clone()) - .expect("Fit should work"); - let y_hat = forest.predict(&x).expect("Predict should work"); + let forest = BaseForest::fit( + &x, + &y, + Some(&sample_weights), + &MSECriterion, + parameters.clone(), + ) + .expect("Fit should work"); + let y_hat = predict_mean(&forest, &x).expect("Predict should work"); // weighted mean is (10.0 * 0 + 90*10) / 100 = 9 for p in y_hat.iter() { @@ -516,8 +593,9 @@ mod tests { } // Without weights, the predicted value should be close to 5 - let forest = BaseForestRegressor::fit(&x, &y, None, parameters).expect("Fit should work"); - let y_hat = forest.predict(&x).expect("Predict should work"); + let forest = + BaseForest::fit(&x, &y, None, &MSECriterion, parameters).expect("Fit should work"); + let y_hat = predict_mean(&forest, &x).expect("Predict should work"); for p in y_hat.iter() { assert!( @@ -527,7 +605,7 @@ mod tests { } // Use bootstrapping - let parameters = BaseForestRegressorParameters { + let parameters = BaseForestParameters { max_depth: Some(0), // Match sibling test on RandomForestRegressor min_samples_leaf: 1, min_samples_split: 2, @@ -536,12 +614,18 @@ mod tests { keep_samples: true, // keep samples used for each tree, so we can check that they are different seed: 42, bootstrap: true, // Use bootstrapping - splitter: crate::tree::base_tree_regressor::Splitter::Best, + splitter: crate::tree::base_tree::Splitter::Best, }; - let forest = BaseForestRegressor::fit(&x, &y, Some(&sample_weights), parameters.clone()) - .expect("Fit should work"); - let y_hat = forest.predict(&x).expect("Predict should work"); + let forest = BaseForest::fit( + &x, + &y, + Some(&sample_weights), + &MSECriterion, + parameters.clone(), + ) + .expect("Fit should work"); + let y_hat = predict_mean(&forest, &x).expect("Predict should work"); // with bootstrapping the weight should be close to 9, but not extremely close for p in y_hat.iter() { @@ -552,8 +636,9 @@ mod tests { } // Without weights, the predicted value should be reasonably close to 5 - let forest = BaseForestRegressor::fit(&x, &y, None, parameters).expect("Fit should work"); - let y_hat = forest.predict(&x).expect("Predict should work"); + let forest = + BaseForest::fit(&x, &y, None, &MSECriterion, parameters).expect("Fit should work"); + let y_hat = predict_mean(&forest, &x).expect("Predict should work"); // The bootstrapping makes that we will not be very close to 5 for p in y_hat.iter() { @@ -580,7 +665,7 @@ mod tests { for splitter in [Splitter::Best, Splitter::Random] { for weights in [None, Some(sample_weights.as_slice())] { - let params = BaseForestRegressorParameters { + let params = BaseForestParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 2, @@ -592,9 +677,9 @@ mod tests { splitter: splitter.clone(), }; - let forest_a = BaseForestRegressor::fit(&x, &y, weights, params.clone()) + let forest_a = BaseForest::fit(&x, &y, weights, &MSECriterion, params.clone()) .expect("Fit should work"); - let forest_b = BaseForestRegressor::fit(&x, &y, weights, params.clone()) + let forest_b = BaseForest::fit(&x, &y, weights, &MSECriterion, params.clone()) .expect("Fit should work"); assert_eq!( @@ -606,17 +691,18 @@ mod tests { "bootstrap samples differ (splitter: {splitter:?}, weights: {weights:?})" ); assert_eq!( - forest_a.predict(&x).unwrap(), - forest_b.predict(&x).unwrap(), + predict_mean(&forest_a, &x).unwrap(), + predict_mean(&forest_b, &x).unwrap(), "predictions differ (splitter: {splitter:?}, weights: {weights:?})" ); // A different seed must give a different forest, else the check above is trivial - let forest_c = BaseForestRegressor::fit( + let forest_c = BaseForest::fit( &x, &y, weights, - BaseForestRegressorParameters { seed: 43, ..params }, + &MSECriterion, + BaseForestParameters { seed: 43, ..params }, ) .expect("Fit should work"); assert_ne!( diff --git a/src/ensemble/extra_trees_regressor.rs b/src/ensemble/extra_trees_regressor.rs index a5a06f77..10dc42f5 100644 --- a/src/ensemble/extra_trees_regressor.rs +++ b/src/ensemble/extra_trees_regressor.rs @@ -59,12 +59,15 @@ use std::fmt::Debug; use serde::{Deserialize, Serialize}; use crate::api::{Predictor, SupervisedEstimator}; -use crate::ensemble::base_forest_regressor::{BaseForestRegressor, BaseForestRegressorParameters}; +use crate::ensemble::base_forest::{ + BaseForest, BaseForestParameters, predict_mean, predict_mean_oob, +}; use crate::error::Failed; use crate::linalg::basic::arrays::{Array1, Array2}; use crate::numbers::basenum::Number; use crate::numbers::floatnum::FloatNumber; -use crate::tree::base_tree_regressor::Splitter; +use crate::tree::base_tree::Splitter; +use crate::tree::criterion::MSECriterion; /// Validates the sample weights fn validate_sample_weights(sample_weights: &[f64], n_rows: usize) -> Result<(), Failed> { @@ -122,7 +125,7 @@ pub struct ExtraTreesRegressor< X: Array2, Y: Array1, > { - forest_regressor: Option>, + forest_regressor: Option>, } impl ExtraTreesRegressorParameters { @@ -234,7 +237,7 @@ impl, Y: Array1 sample_weights: Option<&[f64]>, parameters: ExtraTreesRegressorParameters, ) -> Result, Failed> { - let regressor_params = BaseForestRegressorParameters { + let regressor_params = BaseForestParameters { max_depth: parameters.max_depth, min_samples_leaf: parameters.min_samples_leaf, min_samples_split: parameters.min_samples_split, @@ -245,33 +248,64 @@ impl, Y: Array1 bootstrap: false, splitter: Splitter::Random, }; - let forest_regressor = BaseForestRegressor::fit(x, y, sample_weights, regressor_params)?; + let forest_regressor = + BaseForest::fit(x, y, sample_weights, &MSECriterion, regressor_params)?; Ok(ExtraTreesRegressor { forest_regressor: Some(forest_regressor), }) } - /// Predict class for `x` + /// Predict values for `x` /// * `x` - _KxM_ data where _K_ is number of observations and _M_ is number of features. pub fn predict(&self, x: &X) -> Result { match &self.forest_regressor { - Some(forest) => forest.predict(x), + Some(forest) => predict_mean(forest, x), None => Err(Failed::predict( "'fit' should be called before calling 'predict'", )), } } - /// Predict OOB classes for `x`. `x` is expected to be equal to the dataset used in training. + /// Predict OOB values for `x`. `x` is expected to be equal to the dataset used in training. pub fn predict_oob(&self, x: &X) -> Result { match &self.forest_regressor { - Some(forest) => forest.predict_oob(x), + Some(forest) => predict_mean_oob(forest, x), None => Err(Failed::predict( "'fit' should be called before calling 'predict'", )), } } + + /// Compute the feature importances for this forest. + /// * returns: a vector where index `j` gives the feature importances for feature `j` + /// + /// ``` + /// use smartcore::ensemble::extra_trees_regressor::*; + /// use smartcore::linalg::basic::matrix::DenseMatrix; + /// + /// // Column 1 is constant, so it cannot split the data. + /// let rows: Vec> = (0..20).map(|i| vec![i as f64, 7.0]).collect(); + /// let x = DenseMatrix::from_2d_vec(&rows).unwrap(); + /// let y: Vec = (0..20).map(|i| (i * i) as f64).collect(); + /// + /// let parameters = ExtraTreesRegressorParameters::default() + /// .with_n_trees(10) + /// .with_seed(42); + /// let forest = ExtraTreesRegressor::fit(&x, &y, parameters).unwrap(); + /// + /// // All splits use column 0, so it gets all the importance. + /// // The importances of a forest are normalized: they add up to 1. + /// assert_eq!(forest.compute_feature_importances().unwrap(), vec![1.0, 0.0]); + /// ``` + pub fn compute_feature_importances(&self) -> Result, Failed> { + match &self.forest_regressor { + Some(forest) => forest.compute_feature_importances(), + None => Err(Failed::predict( + "'fit' should be called before calling 'compute_feature_importances'", + )), + } + } } #[cfg(test)] @@ -704,4 +738,139 @@ mod tests { ); } } + + /// Data where only column 0 carries information: `y` is a step function of column 0. + /// Columns 1 and 2 are permutations of the row index, so they do not separate the two steps. + fn step_data() -> (DenseMatrix, Vec) { + let n = 50; + let rows: Vec> = (0..n) + .map(|i| vec![i as f64, ((i * 37) % n) as f64, ((i * 17) % n) as f64]) + .collect(); + let y = (0..n).map(|i| if i < n / 2 { 0.0 } else { 10.0 }).collect(); + (DenseMatrix::from_2d_vec(&rows).unwrap(), y) + } + + #[test] + fn feature_importances_unfitted_returns_error() { + let et = + , Vec> as SupervisedEstimator< + DenseMatrix, + Vec, + ExtraTreesRegressorParameters, + >>::new(); + let result = et.compute_feature_importances(); + assert!(result.is_err()); + assert_eq!(result.unwrap_err().error(), FailedError::PredictFailed); + } + + #[test] + fn feature_importances_sum_to_one() { + let (x, y) = step_data(); + for m in [1, 2, 3] { + let et = ExtraTreesRegressor::fit( + &x, + &y, + ExtraTreesRegressorParameters::default() + .with_n_trees(50) + .with_m(m) + .with_seed(42), + ) + .unwrap(); + let importances = et.compute_feature_importances().unwrap(); + assert_eq!(importances.len(), 3); + assert!(importances.iter().all(|&v| v >= 0.0 && v.is_finite())); + assert!((importances.iter().sum::() - 1.0).abs() < 1e-12); + } + } + + #[test] + fn feature_importances_constant_column_gets_zero() { + // The random splitter skips a constant column, so every split is on column 0, + // whatever the random thresholds are. + let x = DenseMatrix::from_2d_vec(&(0..20).map(|i| vec![i as f64, 7.0]).collect::>()) + .unwrap(); + let y: Vec = (0..20).map(|i| (i * i) as f64).collect(); + let et = ExtraTreesRegressor::fit( + &x, + &y, + ExtraTreesRegressorParameters::default() + .with_n_trees(10) + .with_seed(42), + ) + .unwrap(); + assert_eq!(et.compute_feature_importances().unwrap(), vec![1.0, 0.0]); + } + + #[test] + fn feature_importances_find_informative_feature() { + let (x, y) = step_data(); + for m in [1, 3] { + let et = ExtraTreesRegressor::fit( + &x, + &y, + ExtraTreesRegressorParameters::default() + .with_n_trees(100) + .with_m(m) + .with_seed(42), + ) + .unwrap(); + let importances = et.compute_feature_importances().unwrap(); + assert!(importances[0] > importances[1], "m = {m}: {importances:?}"); + assert!(importances[0] > importances[2], "m = {m}: {importances:?}"); + } + } + + #[test] + fn feature_importances_do_not_depend_on_weight_scale() { + // There is no bootstrap, so the weights go directly to the trees. + // Multiplying all weights by 4 (exact in f64) must give the same importances. + let (x, y) = step_data(); + let weights: Vec = (0..y.len()).map(|i| 0.5 + (i % 4) as f64 * 0.25).collect(); + let scaled: Vec = weights.iter().map(|w| 4.0 * w).collect(); + let parameters = ExtraTreesRegressorParameters::default() + .with_n_trees(20) + .with_m(2) + .with_seed(42); + + let a = ExtraTreesRegressor::fit_with_weights(&x, &y, &weights, parameters.clone()) + .unwrap() + .compute_feature_importances() + .unwrap(); + let b = ExtraTreesRegressor::fit_with_weights(&x, &y, &scaled, parameters) + .unwrap() + .compute_feature_importances() + .unwrap(); + for (ai, bi) in a.iter().zip(b.iter()) { + assert!((ai - bi).abs() < 1e-12); + } + } + + #[test] + fn feature_importances_are_zero_without_improvement() { + let (x, y) = step_data(); + + // max_depth 1: every tree is only a root. + let et = ExtraTreesRegressor::fit( + &x, + &y, + ExtraTreesRegressorParameters::default() + .with_n_trees(5) + .with_max_depth(1) + .with_seed(42), + ) + .unwrap(); + assert_eq!(et.compute_feature_importances().unwrap(), vec![0.0; 3]); + + // Constant y: trees can split, but no split decreases the impurity. + let y_const = vec![3.0; y.len()]; + let et = ExtraTreesRegressor::fit( + &x, + &y_const, + ExtraTreesRegressorParameters::default() + .with_n_trees(5) + .with_seed(42), + ) + .unwrap(); + assert_eq!(et.compute_feature_importances().unwrap(), vec![0.0; 3]); + } } diff --git a/src/ensemble/mod.rs b/src/ensemble/mod.rs index 4f5eefc5..e56b9256 100644 --- a/src/ensemble/mod.rs +++ b/src/ensemble/mod.rs @@ -16,7 +16,7 @@ //! //! * ["An Introduction to Statistical Learning", James G., Witten D., Hastie T., Tibshirani R., 8.2 Bagging, Random Forests, Boosting](http://faculty.marshall.usc.edu/gareth-james/ISL/) -mod base_forest_regressor; +mod base_forest; pub mod extra_trees_regressor; /// Random forest classifier pub mod random_forest_classifier; diff --git a/src/ensemble/random_forest_regressor.rs b/src/ensemble/random_forest_regressor.rs index 364806e4..917d292a 100644 --- a/src/ensemble/random_forest_regressor.rs +++ b/src/ensemble/random_forest_regressor.rs @@ -50,12 +50,15 @@ use std::fmt::Debug; use serde::{Deserialize, Serialize}; use crate::api::{Predictor, SupervisedEstimator}; -use crate::ensemble::base_forest_regressor::{BaseForestRegressor, BaseForestRegressorParameters}; +use crate::ensemble::base_forest::{ + BaseForest, BaseForestParameters, predict_mean, predict_mean_oob, +}; use crate::error::Failed; use crate::linalg::basic::arrays::{Array1, Array2}; use crate::numbers::basenum::Number; use crate::numbers::floatnum::FloatNumber; -use crate::tree::base_tree_regressor::Splitter; +use crate::tree::base_tree::Splitter; +use crate::tree::criterion::MSECriterion; /// Validates the sample weights fn validate_sample_weights(sample_weights: &[f64], n_rows: usize) -> Result<(), Failed> { @@ -113,7 +116,7 @@ pub struct RandomForestRegressor< X: Array2, Y: Array1, > { - forest_regressor: Option>, + forest_regressor: Option>, } impl RandomForestRegressorParameters { @@ -427,7 +430,7 @@ impl, Y: Array1 sample_weights: Option<&[f64]>, parameters: RandomForestRegressorParameters, ) -> Result, Failed> { - let regressor_params = BaseForestRegressorParameters { + let regressor_params = BaseForestParameters { max_depth: parameters.max_depth, min_samples_leaf: parameters.min_samples_leaf, min_samples_split: parameters.min_samples_split, @@ -438,33 +441,70 @@ impl, Y: Array1 bootstrap: true, splitter: Splitter::Best, }; - let forest_regressor = BaseForestRegressor::fit(x, y, sample_weights, regressor_params)?; + let forest_regressor = + BaseForest::fit(x, y, sample_weights, &MSECriterion, regressor_params)?; Ok(RandomForestRegressor { forest_regressor: Some(forest_regressor), }) } - /// Predict class for `x` + /// Predict value for `x` /// * `x` - _KxM_ data where _K_ is number of observations and _M_ is number of features. pub fn predict(&self, x: &X) -> Result { match &self.forest_regressor { - Some(forest) => forest.predict(x), + Some(forest) => predict_mean(forest, x), None => Err(Failed::predict( "'fit' should be called before calling 'predict'", )), } } - /// Predict OOB classes for `x`. `x` is expected to be equal to the dataset used in training. + /// Predict OOB values for `x`. `x` is expected to be equal to the dataset used in training. pub fn predict_oob(&self, x: &X) -> Result { match &self.forest_regressor { - Some(forest) => forest.predict_oob(x), + Some(forest) => predict_mean_oob(forest, x), None => Err(Failed::predict( "'fit' should be called before calling 'predict'", )), } } + + /// Compute the feature importances for this forest. + /// * returns: a vector where index `j` gives the feature importances for feature `j` + /// + /// ``` + /// use smartcore::ensemble::random_forest_regressor::*; + /// use smartcore::linalg::basic::matrix::DenseMatrix; + /// + /// // Only column 0 carries information: y is a step function of column 0. + /// // Column 1 is a permutation of the row index, so it is noise. + /// let n = 50; + /// let rows: Vec> = (0..n) + /// .map(|i| vec![i as f64, ((i * 37) % n) as f64]) + /// .collect(); + /// let x = DenseMatrix::from_2d_vec(&rows).unwrap(); + /// let y: Vec = (0..n).map(|i| if i < n / 2 { 0.0 } else { 10.0 }).collect(); + /// + /// let parameters = RandomForestRegressorParameters::default() + /// .with_n_trees(20) + /// .with_m(2) + /// .with_seed(42); + /// let forest = RandomForestRegressor::fit(&x, &y, parameters).unwrap(); + /// + /// let importances = forest.compute_feature_importances().unwrap(); + /// // The importances of a forest are normalized: they add up to 1. + /// assert!((importances.iter().sum::() - 1.0).abs() < 1e-12); + /// assert!(importances[0] > importances[1]); + /// ``` + pub fn compute_feature_importances(&self) -> Result, Failed> { + match &self.forest_regressor { + Some(forest) => forest.compute_feature_importances(), + None => Err(Failed::predict( + "'fit' should be called before calling 'compute_feature_importances'", + )), + } + } } #[cfg(test)] @@ -1099,4 +1139,131 @@ mod tests { ); } } + + /// Data where only column 0 carries information: `y` is a step function of column 0. + /// Columns 1 and 2 are permutations of the row index, so they do not separate the two steps. + fn step_data() -> (DenseMatrix, Vec) { + let n = 50; + let rows: Vec> = (0..n) + .map(|i| vec![i as f64, ((i * 37) % n) as f64, ((i * 17) % n) as f64]) + .collect(); + let y = (0..n).map(|i| if i < n / 2 { 0.0 } else { 10.0 }).collect(); + (DenseMatrix::from_2d_vec(&rows).unwrap(), y) + } + + #[test] + fn feature_importances_unfitted_returns_error() { + let rf = + , Vec> as SupervisedEstimator< + DenseMatrix, + Vec, + RandomForestRegressorParameters, + >>::new(); + let result = rf.compute_feature_importances(); + assert!(result.is_err()); + assert_eq!(result.unwrap_err().error(), FailedError::PredictFailed); + } + + #[test] + fn feature_importances_sum_to_one() { + let (x, y) = step_data(); + for m in [1, 2, 3] { + let rf = RandomForestRegressor::fit( + &x, + &y, + RandomForestRegressorParameters::default() + .with_n_trees(50) + .with_m(m) + .with_seed(42), + ) + .unwrap(); + let importances = rf.compute_feature_importances().unwrap(); + assert_eq!(importances.len(), 3); + assert!(importances.iter().all(|&v| v >= 0.0 && v.is_finite())); + assert!((importances.iter().sum::() - 1.0).abs() < 1e-12); + } + } + + #[test] + fn feature_importances_find_informative_feature() { + let (x, y) = step_data(); + + // All features are tried at each split. The root split on column 0 is perfect, so the + // children are pure, and all the impurity decrease goes to column 0. + let rf = RandomForestRegressor::fit( + &x, + &y, + RandomForestRegressorParameters::default() + .with_n_trees(20) + .with_m(3) + .with_seed(42), + ) + .unwrap(); + let importances = rf.compute_feature_importances().unwrap(); + assert!((importances[0] - 1.0).abs() < 1e-12); + assert!(importances[1] < 1e-12); + assert!(importances[2] < 1e-12); + + // One random feature per split: noise columns get some importance, + // but column 0 must still be the most important. + let rf = RandomForestRegressor::fit( + &x, + &y, + RandomForestRegressorParameters::default() + .with_n_trees(100) + .with_m(1) + .with_seed(42), + ) + .unwrap(); + let importances = rf.compute_feature_importances().unwrap(); + assert!(importances[0] > importances[1]); + assert!(importances[0] > importances[2]); + } + + #[test] + fn feature_importances_are_deterministic_for_a_seed() { + let (x, y) = step_data(); + let parameters = RandomForestRegressorParameters::default() + .with_n_trees(20) + .with_m(1) + .with_seed(7); + let a = RandomForestRegressor::fit(&x, &y, parameters.clone()) + .unwrap() + .compute_feature_importances() + .unwrap(); + let b = RandomForestRegressor::fit(&x, &y, parameters) + .unwrap() + .compute_feature_importances() + .unwrap(); + assert_eq!(a, b); + } + + #[test] + fn feature_importances_are_zero_without_improvement() { + let (x, y) = step_data(); + + // max_depth 1: every tree is only a root. + let rf = RandomForestRegressor::fit( + &x, + &y, + RandomForestRegressorParameters::default() + .with_n_trees(5) + .with_max_depth(1) + .with_seed(42), + ) + .unwrap(); + assert_eq!(rf.compute_feature_importances().unwrap(), vec![0.0; 3]); + + // Constant y: trees can split, but no split decreases the impurity. + let y_const = vec![3.0; y.len()]; + let rf = RandomForestRegressor::fit( + &x, + &y_const, + RandomForestRegressorParameters::default() + .with_n_trees(5) + .with_seed(42), + ) + .unwrap(); + assert_eq!(rf.compute_feature_importances().unwrap(), vec![0.0; 3]); + } } diff --git a/src/tree/base_tree_regressor.rs b/src/tree/base_tree.rs similarity index 59% rename from src/tree/base_tree_regressor.rs rename to src/tree/base_tree.rs index 7e129cfe..6594338c 100644 --- a/src/tree/base_tree_regressor.rs +++ b/src/tree/base_tree.rs @@ -11,6 +11,7 @@ use crate::error::{Failed, FailedError}; use crate::linalg::basic::arrays::{Array1, Array2, MutArrayView1}; use crate::numbers::basenum::Number; use crate::rand_custom::get_rng_impl; +use crate::tree::criterion::{Criterion, LeafOutput}; #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[derive(Debug, Clone, Copy, Default)] @@ -22,11 +23,11 @@ pub enum Splitter { #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[derive(Debug, Clone)] -/// Parameters of Regression base_tree +/// Parameters of base tree #[must_use] -pub struct BaseTreeRegressorParameters { +pub struct BaseTreeParameters { #[cfg_attr(feature = "serde", serde(default))] - /// The maximum depth of the base_tree. + /// The maximum depth of the tree pub max_depth: Option, #[cfg_attr(feature = "serde", serde(default))] /// The minimum number of samples required to be at a leaf node. @@ -42,36 +43,47 @@ pub struct BaseTreeRegressorParameters { pub splitter: Splitter, } -/// Regression base_tree +/// Base Decision Tree #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +#[cfg_attr( + feature = "serde", + serde(bound(serialize = "Out: Serialize", deserialize = "Out: Deserialize<'de>")) +)] #[derive(Debug)] -pub struct BaseTreeRegressor, Y: Array1> { - nodes: Vec, - parameters: Option, +pub struct BaseTree< + TX: Number + PartialOrd, + TY: Number, + X: Array2, + Y: Array1, + Out: LeafOutput, +> { + nodes: Vec>, + parameters: Option, depth: u16, + root_mass: f64, _phantom_tx: PhantomData, _phantom_ty: PhantomData, _phantom_x: PhantomData, _phantom_y: PhantomData, } -impl, Y: Array1> - BaseTreeRegressor +impl, Y: Array1, Out: LeafOutput> + BaseTree { /// Get nodes, return a shared reference - fn nodes(&self) -> &Vec { + fn nodes(&self) -> &Vec> { self.nodes.as_ref() } /// Get parameters, return a shared reference - fn parameters(&self) -> &BaseTreeRegressorParameters { + fn parameters(&self) -> &BaseTreeParameters { self.parameters.as_ref().unwrap() } } #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[derive(Debug, Clone)] -struct Node { - output: f64, +struct Node { + output: Out, split_feature: usize, split_value: Option, split_score: Option, @@ -79,8 +91,8 @@ struct Node { false_child: Option, } -impl Node { - fn new(output: f64) -> Self { +impl Node { + fn new(output: Out) -> Self { Node { output, split_feature: 0, @@ -92,9 +104,9 @@ impl Node { } } -impl PartialEq for Node { - fn eq(&self, other: &Self) -> bool { - (self.output - other.output).abs() < f64::EPSILON +impl Node { + fn approximate_eq(&self, other: &Self) -> bool { + Out::approximate_eq(&self.output, &other.output) && self.split_feature == other.split_feature && match (self.split_value, other.split_value) { (Some(a), Some(b)) => (a - b).abs() < f64::EPSILON, @@ -109,8 +121,8 @@ impl PartialEq for Node { } } -impl, Y: Array1> PartialEq - for BaseTreeRegressor +impl, Y: Array1, Out: LeafOutput> PartialEq + for BaseTree { fn eq(&self, other: &Self) -> bool { if self.depth != other.depth || self.nodes.len() != other.nodes.len() { @@ -119,26 +131,32 @@ impl, Y: Array1> PartialE self.nodes .iter() .zip(other.nodes().iter()) - .all(|(a, b)| a == b) + .all(|(a, b)| a.approximate_eq(b)) } } } -struct NodeVisitor<'a, TX: Number + PartialOrd, TY: Number, X: Array2, Y: Array1> { +struct NodeVisitor< + 'a, + TX: Number + PartialOrd, + TY: Number, + X: Array2, + Y: Array1, + C: Criterion, +> { x: &'a X, y: &'a Y, node: usize, start_idx: usize, // start index for the elements in `sorted_by_feature` of SplitWorkspace end_idx: usize, // end index (exclusive) in the same vector(s) - true_child_output: f64, - false_child_output: f64, level: u16, + stats: C::Stats, // statistics of the samples in this node _phantom_tx: PhantomData, _phantom_ty: PhantomData, } -impl<'a, TX: Number + PartialOrd, TY: Number, X: Array2, Y: Array1> - NodeVisitor<'a, TX, TY, X, Y> +impl<'a, TX: Number + PartialOrd, TY: Number, X: Array2, Y: Array1, C: Criterion> + NodeVisitor<'a, TX, TY, X, Y, C> { fn new( node_id: usize, @@ -147,6 +165,7 @@ impl<'a, TX: Number + PartialOrd, TY: Number, X: Array2, Y: Array1> x: &'a X, y: &'a Y, level: u16, + stats: C::Stats, ) -> Self { NodeVisitor { x, @@ -154,9 +173,8 @@ impl<'a, TX: Number + PartialOrd, TY: Number, X: Array2, Y: Array1> node: node_id, start_idx, end_idx, - true_child_output: 0f64, - false_child_output: 0f64, level, + stats, _phantom_tx: PhantomData, _phantom_ty: PhantomData, } @@ -281,7 +299,7 @@ impl NodeElement for UnitElement { } // Checks whether the example indicated by node_element is a "true child" for this node -fn is_true_sample(node_element: &E, x: &X, node: &Node) -> bool +fn is_true_sample(node_element: &E, x: &X, node: &Node) -> bool where TX: Number + PartialOrd, X: Array2, @@ -319,38 +337,49 @@ where w } -struct SplitWorkspace { +struct SplitWorkspace { in_true_branch: Vec, // indexed by row index in X partition_buffer: Vec, sorted_by_feature: Vec>, variables: Vec, // feature indices which will be shuffled when mtry < n_features + true_stats: C::Stats, // scratch stats for split functions to use + false_stats: C::Stats, + parent_stats: C::Stats, } -impl SplitWorkspace { +impl SplitWorkspace +where + C: Criterion, +{ fn new( in_true_branch: Vec, partition_buffer: Vec, sorted_by_feature: Vec>, n_features: usize, + criterion: &C, ) -> Self { Self { in_true_branch, partition_buffer, sorted_by_feature, variables: (0..n_features).collect::>(), + true_stats: criterion.empty_stats(), + false_stats: criterion.empty_stats(), + parent_stats: criterion.empty_stats(), } } } -impl, Y: Array1> - BaseTreeRegressor +impl, Y: Array1, Out: LeafOutput> + BaseTree { - pub(crate) fn fit_inner( + pub(crate) fn fit_inner>( x: &X, y: &Y, sample_weights: Option<&[f64]>, - parameters: BaseTreeRegressorParameters, - ) -> Result, Failed> { + criterion: &C, + parameters: BaseTreeParameters, + ) -> Result, Failed> { let (x_nrows, num_attributes) = x.shape(); if x_nrows != y.shape() { return Err(Failed::fit("Size of x should equal size of y")); @@ -371,74 +400,95 @@ impl, Y: Array1> } let samples = vec![1; x_nrows]; - BaseTreeRegressor::fit_weak_learner( + BaseTree::fit_weak_learner( x, y, sample_weights, samples, num_attributes, &order, + criterion, parameters, ) } - pub(crate) fn fit_weak_learner( + pub(crate) fn fit_weak_learner>( x: &X, y: &Y, sample_weights: Option<&[f64]>, samples: Vec, mtry: usize, order: &[Vec], - parameters: BaseTreeRegressorParameters, - ) -> Result, Failed> { + criterion: &C, + parameters: BaseTreeParameters, + ) -> Result, Failed> { check_row_count(y.shape())?; match sample_weights { - Some(_) => Self::grow::( + Some(_) => Self::grow::( x, y, sample_weights, samples, mtry, order, + criterion, parameters, ), None if samples.iter().all(|&s| s <= 1) => { // Note: if any sample has count zero, it will be filtered out, so we should // use the smallest possible NodeElement - Self::grow::(x, y, None, samples, mtry, order, parameters) + Self::grow::( + x, y, None, samples, mtry, order, criterion, parameters, + ) } - None => Self::grow::(x, y, None, samples, mtry, order, parameters), + None => Self::grow::( + x, y, None, samples, mtry, order, criterion, parameters, + ), } } - fn grow( + fn grow>( x: &X, y: &Y, sample_weights: Option<&[f64]>, samples: Vec, mtry: usize, order: &[Vec], - parameters: BaseTreeRegressorParameters, - ) -> Result, Failed> { - let n_rows = y.shape(); - - let mut nodes: Vec = Vec::new(); + criterion: &C, + parameters: BaseTreeParameters, + ) -> Result, Failed> { + let mut nodes: Vec> = Vec::new(); let mut rng = get_rng_impl(parameters.seed); - let mut sum = 0f64; - let mut mass = 0f64; + let mut n = 0usize; + let mut root_mass = 0f64; + let mut root_stats = criterion.empty_stats(); + + // A row with mass 0 (bootstrap count 0 or sample weight 0) has no effect on + // the statistics. Give it count 0, so that it is left out of the node elements and + // is not counted towards min_samples_leaf etc. + let counts: Vec = match sample_weights { + // Without weights, the mass is the count + None => samples.iter().map(|&s| s as u32).collect(), + Some(weights) => samples + .iter() + .zip(weights) + .map(|(&s, &w)| if w > 0.0 { s as u32 } else { 0 }) + .collect(), + }; - for i in 0..n_rows { - let mass_i = mass_of(i, &samples, sample_weights); - mass += mass_i; - sum += mass_i * y.get(i).to_f64().unwrap(); + for (i, &c) in counts.iter().enumerate() { + // A row with mass 0 adds 0 to the statistics: no branch is necessary + let mass = mass_of(i, &samples, sample_weights); + criterion.add(&mut root_stats, criterion.target(*y.get(i)), mass); + root_mass += mass; + n += c as usize; } - let root = Node::new(sum / mass); + let root = Node::new(criterion.leaf_value(&root_stats)); nodes.push(root); - let counts: Vec = samples.iter().map(|&s| s as u32).collect(); - // number of distinct rows of x in this tree + // number of distinct rows of x in this tree with a positive mass let n_kept = counts.iter().filter(|&&c| c > 0).count(); let sorted_by_feature: Vec> = order @@ -463,56 +513,62 @@ impl, Y: Array1> vec![E::default(); end_idx], sorted_by_feature, x.shape().1, + criterion, ); - let mut base_tree = BaseTreeRegressor { + let mut base_tree = BaseTree { nodes, parameters: Some(parameters), depth: 0u16, + root_mass, _phantom_tx: PhantomData, _phantom_ty: PhantomData, _phantom_x: PhantomData, _phantom_y: PhantomData, }; - let mut visitor = NodeVisitor::::new(0, 0, end_idx, x, y, 1); + let mut visitor = NodeVisitor::::new(0, 0, end_idx, x, y, 1, root_stats); - let mut visitor_stack: Vec> = Vec::new(); + let mut visitor_stack: Vec> = Vec::new(); - if base_tree.find_best_cutoff(&mut visitor, mtry, mass, &mut rng, &mut workspace) { + if base_tree.parameters().max_depth.unwrap_or(u16::MAX) > 1 + && n_kept > 0 + && n >= base_tree.parameters().min_samples_split + && base_tree.find_best_cutoff( + &mut visitor, + mtry, + &mut rng, + &mut workspace, + n, + criterion, + ) + { visitor_stack.push(visitor); } let max_depth = base_tree.parameters().max_depth.unwrap_or(u16::MAX); while let Some(node) = visitor_stack.pop() { if node.level < max_depth { - base_tree.split(node, mtry, &mut visitor_stack, &mut rng, &mut workspace); + base_tree.split( + node, + mtry, + &mut visitor_stack, + &mut rng, + &mut workspace, + criterion, + ); } } Ok(base_tree) } - /// Predict regression value for `x`. - /// * `x` - _KxM_ data where _K_ is number of observations and _M_ is number of features. - pub fn predict(&self, x: &X) -> Result { - let mut result = Y::zeros(x.shape().0); - - let (n, _) = x.shape(); - - for i in 0..n { - result.set(i, self.predict_for_row(x, i)); - } - - Ok(result) - } - - pub(crate) fn predict_for_row(&self, x: &X, row: usize) -> TY { + pub(crate) fn leaf_output_for_row(&self, x: &X, row: usize) -> &Out { let mut node_id = 0; loop { let node = &self.nodes[node_id]; let Some(true_child) = node.true_child else { - return TY::from_f64(node.output).unwrap(); + return &node.output; }; let false_child = node.false_child.unwrap(); node_id = if x.get((row, node.split_feature)).to_f64().unwrap() @@ -525,47 +581,27 @@ impl, Y: Array1> } } - fn find_best_cutoff( + fn find_best_cutoff>( &mut self, - visitor: &mut NodeVisitor<'_, TX, TY, X, Y>, + visitor: &mut NodeVisitor<'_, TX, TY, X, Y, C>, mtry: usize, - mass: f64, rng: &mut impl rand::Rng, - workspace: &mut SplitWorkspace, + workspace: &mut SplitWorkspace, + n: usize, + criterion: &C, ) -> bool { - let n: usize = workspace.sorted_by_feature[0][visitor.start_idx..visitor.end_idx] - .iter() - .map(|elem| elem.count()) - .sum(); - - if n < self.parameters().min_samples_split { - return false; - } - - let sum = self.nodes[visitor.node].output * mass; - workspace.variables.shuffle(rng); - let parent_gain = mass * self.nodes[visitor.node].output * self.nodes[visitor.node].output; - let splitter = self.parameters().splitter; - for variable in workspace.variables.iter().take(mtry) { + for k in 0..mtry.min(workspace.variables.len()) { + let variable = workspace.variables[k]; match splitter { Splitter::Random => { - self.find_random_split( - visitor, - n, - mass, - sum, - parent_gain, - *variable, - rng, - workspace, - ); + self.find_random_split(visitor, n, variable, rng, workspace, criterion); } Splitter::Best => { - self.find_best_split(visitor, n, mass, sum, parent_gain, *variable, workspace); + self.find_best_split(visitor, n, variable, workspace, criterion); } } } @@ -573,16 +609,14 @@ impl, Y: Array1> self.nodes[visitor.node].split_score.is_some() } - fn find_random_split( + fn find_random_split>( &mut self, - visitor: &mut NodeVisitor<'_, TX, TY, X, Y>, + visitor: &NodeVisitor<'_, TX, TY, X, Y, C>, n: usize, - mass: f64, - sum: f64, - parent_gain: f64, j: usize, rng: &mut impl rand::Rng, - workspace: &SplitWorkspace, + workspace: &mut SplitWorkspace, + criterion: &C, ) { if visitor.start_idx == visitor.end_idx { return; @@ -603,14 +637,18 @@ impl, Y: Array1> let split_value = rng.random_range(min_val..max_val); - let mut true_sum = 0f64; - let mut true_mass = 0f64; let mut true_count = 0; + let true_stats = &mut workspace.true_stats; + criterion.reset(true_stats); + for elem in &workspace.sorted_by_feature[j][visitor.start_idx..visitor.end_idx] { if visitor.x.get((elem.row(), j)).to_f64().unwrap() <= split_value { - true_sum += elem.mass() * visitor.y.get(elem.row()).to_f64().unwrap(); + criterion.add( + true_stats, + criterion.target(*visitor.y.get(elem.row())), + elem.mass(), + ); true_count += elem.count(); - true_mass += elem.mass(); } else { break; } @@ -624,110 +662,93 @@ impl, Y: Array1> return; } - let true_mean = if true_mass > 0f64 { - true_sum / true_mass - } else { - 0.0 - }; - let false_mass = mass - true_mass; - let false_mean = if false_mass > 0f64 { - (sum - true_sum) / false_mass - } else { - 0.0 - }; - let gain = (true_mass * true_mean * true_mean + false_mass * false_mean * false_mean) - - parent_gain; + // no need to reset false_stats, will be overwritten by 'difference' + let false_stats = &mut workspace.false_stats; + criterion.difference(&visitor.stats, true_stats, false_stats); + let proxy = criterion.proxy_improvement(true_stats, false_stats); if self.nodes[visitor.node].split_score.is_none() - || gain > self.nodes[visitor.node].split_score.unwrap() + || proxy > self.nodes[visitor.node].split_score.unwrap() { self.nodes[visitor.node].split_feature = j; self.nodes[visitor.node].split_value = Some(split_value); - self.nodes[visitor.node].split_score = Some(gain); - visitor.true_child_output = true_mean; - visitor.false_child_output = false_mean; + self.nodes[visitor.node].split_score = Some(proxy); } } - fn find_best_split( + fn find_best_split>( &mut self, - visitor: &mut NodeVisitor<'_, TX, TY, X, Y>, + visitor: &NodeVisitor<'_, TX, TY, X, Y, C>, n: usize, - mass: f64, - sum: f64, - parent_gain: f64, j: usize, - workspace: &SplitWorkspace, + workspace: &mut SplitWorkspace, + criterion: &C, ) { - let mut true_sum = 0f64; let mut true_count = 0; - let mut true_mass = 0f64; - let mut prevx = Option::None; + let true_stats = &mut workspace.true_stats; + criterion.reset(true_stats); + // no need to reset false_stats, will be overwritten by 'difference' + let false_stats = &mut workspace.false_stats; + // set the "previous" element equal to the first one, so we don't have to deal with Option + let first_row_idx = workspace.sorted_by_feature[j][visitor.start_idx].row(); + let mut prevx = visitor.x.get((first_row_idx, j)).to_f64().unwrap(); + + let min_samples_leaf = self.parameters().min_samples_leaf; for elem in &workspace.sorted_by_feature[j][visitor.start_idx..visitor.end_idx] { // Comparison for equality is done directly on f64 to avoid later problems // with values being unequal in TX but equal in f64. let x_ij = visitor.x.get((elem.row(), j)).to_f64().unwrap(); - if prevx.is_none() || x_ij == prevx.unwrap() { - prevx = Some(x_ij); + if x_ij == prevx { true_count += elem.count(); - true_mass += elem.mass(); - true_sum += elem.mass() * visitor.y.get(elem.row()).to_f64().unwrap(); + criterion.add( + true_stats, + criterion.target(*visitor.y.get(elem.row())), + elem.mass(), + ); continue; } let false_count = n - true_count; - if true_count < self.parameters().min_samples_leaf - || false_count < self.parameters().min_samples_leaf - { - prevx = Some(x_ij); + if true_count < min_samples_leaf || false_count < min_samples_leaf { + prevx = x_ij; true_count += elem.count(); - true_mass += elem.mass(); - true_sum += elem.mass() * visitor.y.get(elem.row()).to_f64().unwrap(); + criterion.add( + true_stats, + criterion.target(*visitor.y.get(elem.row())), + elem.mass(), + ); continue; } - let true_mean = if true_mass > 0.0 { - true_sum / true_mass - } else { - 0.0 - }; - let false_mass = mass - true_mass; - let false_mean = if false_mass > 0.0 { - (sum - true_sum) / false_mass - } else { - 0.0 - }; - - let gain = (true_mass * true_mean * true_mean + false_mass * false_mean * false_mean) - - parent_gain; + criterion.difference(&visitor.stats, true_stats, false_stats); + let proxy = criterion.proxy_improvement(true_stats, false_stats); if self.nodes[visitor.node].split_score.is_none() - || gain > self.nodes[visitor.node].split_score.unwrap() + || proxy > self.nodes[visitor.node].split_score.unwrap() { - let prev = prevx.unwrap(); // compute the midpoint in a way that does not overflow when x_ij and prevx have the same sign - let mid = prev + (x_ij - prev) / 2.0; + let mid = prevx + (x_ij - prevx) / 2.0; // The midpoint should put prevx in the true branch, and x_ij in the false branch. // However, mid could round up to x_ij, in which case we would also assign x_ij to the true branch. // In this case, we set the midpoint equal to prevx, which is always smaller // than x_ij. - let threshold = if mid < x_ij { mid } else { prev }; + let threshold = if mid < x_ij { mid } else { prevx }; self.nodes[visitor.node].split_feature = j; self.nodes[visitor.node].split_value = Option::Some(threshold); - self.nodes[visitor.node].split_score = Option::Some(gain); - - visitor.true_child_output = true_mean; - visitor.false_child_output = false_mean; + self.nodes[visitor.node].split_score = Option::Some(proxy); } - prevx = Some(x_ij); - true_sum += elem.mass() * visitor.y.get(elem.row()).to_f64().unwrap(); + prevx = x_ij; true_count += elem.count(); - true_mass += elem.mass(); + criterion.add( + true_stats, + criterion.target(*visitor.y.get(elem.row())), + elem.mass(), + ); } } @@ -737,20 +758,21 @@ impl, Y: Array1> /// split, push it on `visitor_stack`. /// If a branch has fewer than `min_samples_leaf` samples, clear the split and /// keep the node as a leaf. - fn split<'a, E: NodeElement>( + fn split<'a, E: NodeElement, C: Criterion>( &mut self, - visitor: NodeVisitor<'a, TX, TY, X, Y>, + visitor: NodeVisitor<'a, TX, TY, X, Y, C>, mtry: usize, - visitor_stack: &mut Vec>, + visitor_stack: &mut Vec>, rng: &mut impl rand::Rng, - workspace: &mut SplitWorkspace, + workspace: &mut SplitWorkspace, + criterion: &C, ) -> bool { + let parent_stats = &mut workspace.parent_stats; + criterion.reset(parent_stats); let this_node = &self.nodes[visitor.node]; let mut true_count = 0usize; - let mut true_mass = 0f64; let mut false_count = 0usize; - let mut false_mass = 0f64; // for each row_index, does it belong in the true branch or not? let in_true_branch = &mut workspace.in_true_branch; @@ -764,13 +786,28 @@ impl, Y: Array1> } in_true_branch[e.row()] = true; true_count += e.count(); - true_mass += e.mass(); + criterion.add( + parent_stats, + criterion.target(*visitor.y.get(e.row())), + e.mass(), + ); n_true += 1; } + let true_stats = parent_stats.clone(); + let mut false_stats = criterion.empty_stats(); for e in &col[n_true..] { in_true_branch[e.row()] = false; false_count += e.count(); - false_mass += e.mass(); + criterion.add( + parent_stats, + criterion.target(*visitor.y.get(e.row())), + e.mass(), + ); + criterion.add( + &mut false_stats, + criterion.target(*visitor.y.get(e.row())), + e.mass(), + ); } // Stop early if it is clear that there will be too few examples in the leaf @@ -784,14 +821,21 @@ impl, Y: Array1> return false; } + // Store the real (mass-weighted) impurity improvement in split_score, not the proxy. + // false_stats is a direct sum, not parent - true: no cancellation for a small false side + self.nodes[visitor.node].split_score = + Some(criterion.impurity_improvement(parent_stats, &true_stats, &false_stats)); + // Add the child nodes to the tree let split_idx = visitor.start_idx + n_true; let true_child_idx = self.nodes.len(); - self.nodes.push(Node::new(visitor.true_child_output)); + self.nodes + .push(Node::new(criterion.leaf_value(&true_stats))); let false_child_idx = self.nodes.len(); - self.nodes.push(Node::new(visitor.false_child_output)); + self.nodes + .push(Node::new(criterion.leaf_value(&false_stats))); self.nodes[visitor.node].true_child = Some(true_child_idx); self.nodes[visitor.node].false_child = Some(false_child_idx); @@ -816,38 +860,83 @@ impl, Y: Array1> ); } - let mut true_visitor = NodeVisitor::::new( + let mut true_visitor = NodeVisitor::::new( true_child_idx, visitor.start_idx, split_idx, visitor.x, visitor.y, child_level, + true_stats, ); if true_can_split - && self.find_best_cutoff(&mut true_visitor, mtry, true_mass, rng, workspace) + && self.find_best_cutoff( + &mut true_visitor, + mtry, + rng, + workspace, + true_count, + criterion, + ) { visitor_stack.push(true_visitor); } - let mut false_visitor = NodeVisitor::::new( + let mut false_visitor = NodeVisitor::::new( false_child_idx, split_idx, visitor.end_idx, visitor.x, visitor.y, child_level, + false_stats, ); if false_can_split - && self.find_best_cutoff(&mut false_visitor, mtry, false_mass, rng, workspace) + && self.find_best_cutoff( + &mut false_visitor, + mtry, + rng, + workspace, + false_count, + criterion, + ) { visitor_stack.push(false_visitor); } true } + + /// Compute the feature importances and store the result in `importances`. + /// * `normalize`: when `false` feature importances are the impurity improvements, when `true` they + /// are normalized to sum to one. + /// * `importances` must have length (at least) equal to the number of features. + /// The result is stored in this slice. + pub fn compute_feature_importances(&self, normalize: bool, importances: &mut [f64]) { + importances.fill(0.0); + for node in self.nodes().iter() { + if let Some(split_score) = node.split_score { + importances[node.split_feature] += split_score; + } + } + + if normalize { + let sum: f64 = importances.iter().sum(); + if sum > 0.0 { + for importance in importances.iter_mut() { + *importance /= sum; + } + } + } else { + if self.root_mass > 0.0 { + for importance in importances.iter_mut() { + *importance /= self.root_mass; + } + } + } + } } #[cfg(test)] @@ -856,6 +945,7 @@ mod tests { use crate::linalg::basic::arrays::Array; use crate::linalg::basic::matrix::DenseMatrix; use crate::metrics::mean_absolute_error; + use crate::tree::criterion::MSECriterion; #[test] #[cfg(target_pointer_width = "64")] @@ -878,11 +968,12 @@ mod tests { assert_eq!(empty.shape(), (0, 2)); let y: Vec = vec![]; - let result = BaseTreeRegressor::fit_inner( + let result = BaseTree::fit_inner( &empty, &y, None, - BaseTreeRegressorParameters { + &MSECriterion, + BaseTreeParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 2, @@ -902,11 +993,12 @@ mod tests { assert_eq!(no_features.shape(), (2, 0)); let y = vec![1.0_f64, 2.0]; - let result = BaseTreeRegressor::fit_inner( + let result = BaseTree::fit_inner( &no_features, &y, None, - BaseTreeRegressorParameters { + &MSECriterion, + BaseTreeParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 2, @@ -928,11 +1020,12 @@ mod tests { let y = vec![1.0, 2.0, 3.0]; let sample_weights = vec![5.0, 6.0, 7.0]; - let result = BaseTreeRegressor::fit_inner( + let result = BaseTree::fit_inner( &x, &y, Some(&sample_weights), - BaseTreeRegressorParameters { + &MSECriterion, + BaseTreeParameters { max_depth: Some(0), min_samples_leaf: 1, min_samples_split: 2, @@ -949,7 +1042,7 @@ mod tests { .sum::() / sample_weights.iter().sum::(); - assert!((tree.predict_for_row(&x, 0) - expected).abs() < 1e-9); + assert!((*tree.leaf_output_for_row(&x, 0) - expected).abs() < 1e-9); } #[test] @@ -957,7 +1050,7 @@ mod tests { // Test that using no weights is equivalent to using uniform weights let x_rand: DenseMatrix = DenseMatrix::::rand(17, 5); let y_rand: Vec = (0..17).collect::>().map(|y| *y as f64); - let parameters = BaseTreeRegressorParameters { + let parameters = BaseTreeParameters { max_depth: Some(5), min_samples_leaf: 1, min_samples_split: 1, @@ -966,19 +1059,20 @@ mod tests { }; let tree_no_weights = - BaseTreeRegressor::fit_inner(&x_rand, &y_rand, None, parameters.clone()) + BaseTree::fit_inner(&x_rand, &y_rand, None, &MSECriterion, parameters.clone()) .expect("Fit should work"); let uniform_weights = vec![1.0f64; 17]; - let tree_with_weights = - BaseTreeRegressor::fit_inner(&x_rand, &y_rand, Some(&uniform_weights), parameters) - .expect("Fit should work"); + let tree_with_weights = BaseTree::fit_inner( + &x_rand, + &y_rand, + Some(&uniform_weights), + &MSECriterion, + parameters, + ) + .expect("Fit should work"); - let y_pred_no_weights = tree_no_weights - .predict(&x_rand) - .expect("Predict should work"); - let y_pred_with_weights = tree_with_weights - .predict(&x_rand) - .expect("Predict should work"); + let y_pred_no_weights = predict(&tree_no_weights, &x_rand); + let y_pred_with_weights = predict(&tree_with_weights, &x_rand); assert!(mean_absolute_error(&y_pred_no_weights, &y_pred_with_weights) < 1e-9); } @@ -1001,7 +1095,7 @@ mod tests { .unwrap(); let y_repeated = vec![4.0, 5.0, 5.0, 6.0, 6.0, 6.0]; - let weighted_parameters = BaseTreeRegressorParameters { + let weighted_parameters = BaseTreeParameters { max_depth: Some(1), // tree should not be able to fully separate all the examples min_samples_leaf: 1, min_samples_split: 1, @@ -1009,7 +1103,7 @@ mod tests { splitter: Splitter::Best, }; - let repeated_parameters = BaseTreeRegressorParameters { + let repeated_parameters = BaseTreeParameters { max_depth: Some(1), min_samples_leaf: 1, min_samples_split: 1, @@ -1017,15 +1111,25 @@ mod tests { splitter: Splitter::Best, }; - let tree_weighted = - BaseTreeRegressor::fit_inner(&x, &y, Some(&sample_weights), weighted_parameters) - .expect("Fit should work"); - let tree_repeated = - BaseTreeRegressor::fit_inner(&x_repeated, &y_repeated, None, repeated_parameters) - .expect("Fit should work"); + let tree_weighted = BaseTree::fit_inner( + &x, + &y, + Some(&sample_weights), + &MSECriterion, + weighted_parameters, + ) + .expect("Fit should work"); + let tree_repeated = BaseTree::fit_inner( + &x_repeated, + &y_repeated, + None, + &MSECriterion, + repeated_parameters, + ) + .expect("Fit should work"); // Predict on the same data - let y_pred_weighted = tree_weighted.predict(&x).expect("Predict should work"); - let y_pred_repeated = tree_repeated.predict(&x).expect("Predict should work"); + let y_pred_weighted = predict(&tree_weighted, &x); + let y_pred_repeated = predict(&tree_repeated, &x); assert!(mean_absolute_error(&y_pred_weighted, &y_pred_repeated) < 1e-9); } @@ -1043,7 +1147,7 @@ mod tests { .unwrap(); let y = vec![1.0f64, 2.0, 6.0, 7.0, 11., 12.]; - let parameters = BaseTreeRegressorParameters { + let parameters = BaseTreeParameters { max_depth: Some(3), min_samples_leaf: 1, min_samples_split: 2, @@ -1051,9 +1155,10 @@ mod tests { splitter: Splitter::Best, }; - let tree = BaseTreeRegressor::fit_inner(&x, &y, None, parameters).expect("Fit should work"); + let tree = + BaseTree::fit_inner(&x, &y, None, &MSECriterion, parameters).expect("Fit should work"); let y_expected = vec![1.0, 2.0, 6.5, 6.5, 11.50, 11.50]; - let y_hat = tree.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); assert_eq!(tree.nodes().len(), 7); assert_eq!(tree.depth, 3); assert!(mean_absolute_error(&y_expected, &y_hat) < 1e-9); @@ -1071,7 +1176,7 @@ mod tests { ]) .unwrap(); let y = vec![1.0f64, 2.0, 6.0, 7.0, 11., 12.]; - let parameters = BaseTreeRegressorParameters { + let parameters = BaseTreeParameters { max_depth: Some(3), min_samples_leaf: 1, min_samples_split: 2, @@ -1080,14 +1185,14 @@ mod tests { }; let sample_weights = [1.0f64, 1.0, 1.0, 1.0, 10.0, 10.0]; - let tree = BaseTreeRegressor::fit_inner(&x, &y, Some(&sample_weights), parameters) + let tree = BaseTree::fit_inner(&x, &y, Some(&sample_weights), &MSECriterion, parameters) .expect("Fit should work"); let x_test: DenseMatrix = DenseMatrix::from_iterator((1..=12).map(|i| i as f64), 12, 1, 0); let y_expected = vec![ 1.50, 1.50, 6.50, 6.50, 11.0, 12.0, 12.0, 12.0, 12.0, 12.0, 12.0, 12.0, ]; - let y_hat = tree.predict(&x_test).expect("Predict should work"); + let y_hat = predict(&tree, &x_test); assert_eq!(tree.nodes().len(), 7); assert_eq!(tree.depth, 3); assert!(mean_absolute_error(&y_expected, &y_hat) < 1e-9); @@ -1099,7 +1204,7 @@ mod tests { let x = DenseMatrix::from_2d_vec(&vec![vec![1.0_f64], vec![2.0], vec![3.0]]).unwrap(); let y = vec![1.0f64, 2.0, 6.0]; - let parameters = BaseTreeRegressorParameters { + let parameters = BaseTreeParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 3, @@ -1107,12 +1212,13 @@ mod tests { splitter: Splitter::Best, }; - let tree = BaseTreeRegressor::fit_inner(&x, &y, None, parameters).expect("Fit should work"); + let tree = + BaseTree::fit_inner(&x, &y, None, &MSECriterion, parameters).expect("Fit should work"); assert_eq!(tree.nodes().len(), 3); assert_eq!(tree.depth, 2); // A node with fewer than `min_samples_split` samples must stay a leaf. - let parameters = BaseTreeRegressorParameters { + let parameters = BaseTreeParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 4, @@ -1120,7 +1226,8 @@ mod tests { splitter: Splitter::Best, }; - let tree = BaseTreeRegressor::fit_inner(&x, &y, None, parameters).expect("Fit should work"); + let tree = + BaseTree::fit_inner(&x, &y, None, &MSECriterion, parameters).expect("Fit should work"); assert_eq!(tree.nodes().len(), 1); assert_eq!(tree.depth, 0); } @@ -1133,7 +1240,7 @@ mod tests { let y = vec![1.0f64, 2.0, 3.0]; let sample_weights = [0.0f64, 1.0, 1.0]; - let parameters = BaseTreeRegressorParameters { + let parameters = BaseTreeParameters { max_depth: Some(2), min_samples_leaf: 1, min_samples_split: 2, @@ -1141,7 +1248,7 @@ mod tests { splitter: Splitter::Best, }; - let tree = BaseTreeRegressor::fit_inner(&x, &y, Some(&sample_weights), parameters) + let tree = BaseTree::fit_inner(&x, &y, Some(&sample_weights), &MSECriterion, parameters) .expect("Fit should work"); assert!(tree.nodes().iter().all(|node| node.output.is_finite())); @@ -1151,7 +1258,7 @@ mod tests { .all(|node| node.split_score.is_none_or(f64::is_finite)) ); - let y_hat = tree.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); assert!(y_hat.iter().all(|v| v.is_finite())); // The best split separates x=3 from the other rows. @@ -1168,7 +1275,7 @@ mod tests { let y = vec![1.0f64, 2.0, 3.0]; let sample_weights = [1.0f64, 1.0, 0.0]; - let parameters = BaseTreeRegressorParameters { + let parameters = BaseTreeParameters { max_depth: Some(2), min_samples_leaf: 1, min_samples_split: 2, @@ -1176,7 +1283,7 @@ mod tests { splitter: Splitter::Best, }; - let tree = BaseTreeRegressor::fit_inner(&x, &y, Some(&sample_weights), parameters) + let tree = BaseTree::fit_inner(&x, &y, Some(&sample_weights), &MSECriterion, parameters) .expect("Fit should work"); assert!(tree.nodes().iter().all(|node| node.output.is_finite())); @@ -1186,7 +1293,7 @@ mod tests { .all(|node| node.split_score.is_none_or(f64::is_finite)) ); - let y_hat = tree.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); assert!(y_hat.iter().all(|v| v.is_finite())); } @@ -1199,7 +1306,7 @@ mod tests { let y = vec![0.0f64, 10.0, 20.0, 30.0]; let sample_weights = [3.0f64, 1.0, 1.0, 3.0]; - let parameters = BaseTreeRegressorParameters { + let parameters = BaseTreeParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 2, @@ -1207,14 +1314,20 @@ mod tests { splitter: Splitter::Best, }; - let tree = BaseTreeRegressor::fit_inner(&x, &y, Some(&sample_weights), parameters.clone()) - .expect("Fit should work"); + let tree = BaseTree::fit_inner( + &x, + &y, + Some(&sample_weights), + &MSECriterion, + parameters.clone(), + ) + .expect("Fit should work"); assert_eq!(tree.nodes().len(), 3); assert_eq!(tree.depth, 2); assert!((tree.nodes()[0].split_value.unwrap() - 1.5).abs() < 1e-9); // Split should be at 1.5 - let y_hat = tree.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); let y_expected = vec![2.5, 2.5, 27.5, 27.5]; // Expected values are weighted means assert!(mean_absolute_error(&y_expected, &y_hat) < 1e-9); @@ -1232,9 +1345,9 @@ mod tests { .unwrap(); let y_repeated = vec![0.0f64, 0.0, 0.0, 10.0, 20.0, 30.0, 30.0, 30.0]; let tree_repeated = - BaseTreeRegressor::fit_inner(&x_repeated, &y_repeated, None, parameters) + BaseTree::fit_inner(&x_repeated, &y_repeated, None, &MSECriterion, parameters) .expect("Fit should work"); - let y_hat_repeated = tree_repeated.predict(&x).expect("Predict should work"); + let y_hat_repeated = predict(&tree_repeated, &x); assert!(mean_absolute_error(&y_hat, &y_hat_repeated) < 1e-9); } @@ -1242,8 +1355,8 @@ mod tests { // but partitioning and predict compare them as f64 against an f64 threshold. const TWO_POW_53: i64 = 1 << 53; - fn split_value_parameters(splitter: Splitter) -> BaseTreeRegressorParameters { - BaseTreeRegressorParameters { + fn split_value_parameters(splitter: Splitter) -> BaseTreeParameters { + BaseTreeParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 2, @@ -1260,15 +1373,20 @@ mod tests { DenseMatrix::from_2d_vec(&vec![vec![TWO_POW_53 + 2], vec![TWO_POW_53 + 4]]).unwrap(); let y = vec![0.0f64, 10.0]; - let tree = - BaseTreeRegressor::fit_inner(&x, &y, None, split_value_parameters(Splitter::Best)) - .expect("Fit should work"); + let tree = BaseTree::fit_inner( + &x, + &y, + None, + &MSECriterion, + split_value_parameters(Splitter::Best), + ) + .expect("Fit should work"); let threshold = tree.nodes()[0].split_value.expect("Root should split"); assert!((TWO_POW_53 + 2) as f64 <= threshold); assert!(((TWO_POW_53 + 4) as f64) > threshold); - let y_hat = tree.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); assert_eq!(y_hat, vec![0.0, 10.0]); } @@ -1279,11 +1397,16 @@ mod tests { let x = DenseMatrix::from_2d_vec(&vec![vec![TWO_POW_53], vec![TWO_POW_53 + 1]]).unwrap(); let y = vec![0.0f64, 10.0]; - let tree = - BaseTreeRegressor::fit_inner(&x, &y, None, split_value_parameters(Splitter::Best)) - .expect("Fit should work"); + let tree = BaseTree::fit_inner( + &x, + &y, + None, + &MSECriterion, + split_value_parameters(Splitter::Best), + ) + .expect("Fit should work"); - let y_hat = tree.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); assert_eq!(y_hat, vec![5.0, 5.0]); } @@ -1293,11 +1416,16 @@ mod tests { let x = DenseMatrix::from_2d_vec(&vec![vec![TWO_POW_53], vec![TWO_POW_53 + 1]]).unwrap(); let y = vec![0.0f64, 10.0]; - let tree = - BaseTreeRegressor::fit_inner(&x, &y, None, split_value_parameters(Splitter::Random)) - .expect("Fit should work"); + let tree = BaseTree::fit_inner( + &x, + &y, + None, + &MSECriterion, + split_value_parameters(Splitter::Random), + ) + .expect("Fit should work"); - let y_hat = tree.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); assert_eq!(y_hat, vec![5.0, 5.0]); } @@ -1313,15 +1441,20 @@ mod tests { let x = DenseMatrix::from_2d_vec(&vec![vec![a], vec![b]]).unwrap(); let y = vec![0.0f64, 10.0]; - let tree = - BaseTreeRegressor::fit_inner(&x, &y, None, split_value_parameters(Splitter::Best)) - .expect("Fit should work"); + let tree = BaseTree::fit_inner( + &x, + &y, + None, + &MSECriterion, + split_value_parameters(Splitter::Best), + ) + .expect("Fit should work"); let threshold = tree.nodes()[0].split_value.expect("Root should split"); assert!(a <= threshold); assert!(b > threshold); - let y_hat = tree.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); assert_eq!(y_hat, vec![0.0, 10.0]); } @@ -1332,15 +1465,29 @@ mod tests { assert_eq!(std::mem::size_of::(), 16); } - type TestTree = BaseTreeRegressor, Vec>; + type TestTree = BaseTree, Vec, f64>; + + /// Leaf output of `tree` for each row of `x`. + fn predict(tree: &BaseTree, x: &X) -> Vec + where + TX: Number + PartialOrd, + TY: Number, + X: Array2, + Y: Array1, + Out: LeafOutput, + { + (0..x.shape().0) + .map(|i| tree.leaf_output_for_row(x, i).clone()) + .collect() + } fn params( max_depth: Option, min_samples_leaf: usize, min_samples_split: usize, splitter: Splitter, - ) -> BaseTreeRegressorParameters { - BaseTreeRegressorParameters { + ) -> BaseTreeParameters { + BaseTreeParameters { max_depth, min_samples_leaf, min_samples_split, @@ -1359,7 +1506,7 @@ mod tests { .collect() } - /// Index of the leaf that gets row `row` of `x`. Uses the same rule as `predict_for_row`. + /// Index of the leaf that gets row `row` of `x`. Uses the same rule as `leaf_output_for_row`. fn leaf_of(tree: &TestTree, x: &DenseMatrix, row: usize) -> usize { let mut node_id = 0; while let Some(true_child) = tree.nodes()[node_id].true_child { @@ -1452,40 +1599,57 @@ mod tests { (None, 15, 4), ]; for (max_depth, n_nodes, depth) in cases { - let tree: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, None, params(max_depth, 1, 2, Splitter::Best)) - .expect("Fit should work"); + let tree: TestTree = BaseTree::fit_inner( + &x, + &y, + None, + &MSECriterion, + params(max_depth, 1, 2, Splitter::Best), + ) + .expect("Fit should work"); assert_eq!(tree.nodes().len(), n_nodes, "max_depth: {max_depth:?}"); assert_eq!(tree.depth, depth, "max_depth: {max_depth:?}"); } // With max_depth = 2 the root split separates the two halves. - let tree: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, None, params(Some(2), 1, 2, Splitter::Best)) - .expect("Fit should work"); + let tree: TestTree = BaseTree::fit_inner( + &x, + &y, + None, + &MSECriterion, + params(Some(2), 1, 2, Splitter::Best), + ) + .expect("Fit should work"); assert_eq!(tree.nodes()[0].split_value, Some(4.5)); - let y_hat = tree.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); assert_eq!(y_hat, vec![5.5, 5.5, 5.5, 5.5, 105.5, 105.5, 105.5, 105.5]); // No limit is the same as the largest limit, and the leaves memorize y. - let tree_none: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, None, params(None, 1, 2, Splitter::Best)) - .expect("Fit should work"); - let tree_max: TestTree = BaseTreeRegressor::fit_inner( + let tree_none: TestTree = BaseTree::fit_inner( &x, &y, None, + &MSECriterion, + params(None, 1, 2, Splitter::Best), + ) + .expect("Fit should work"); + let tree_max: TestTree = BaseTree::fit_inner( + &x, + &y, + None, + &MSECriterion, params(Some(u16::MAX), 1, 2, Splitter::Best), ) .expect("Fit should work"); assert_same_tree(&tree_none, &tree_max); - assert_eq!(tree_none.predict(&x).expect("Predict should work"), y); + assert_eq!(predict(&tree_none, &x), y); for max_depth in 1..=4 { - let tree: TestTree = BaseTreeRegressor::fit_inner( + let tree: TestTree = BaseTree::fit_inner( &x, &y, None, + &MSECriterion, params(Some(max_depth), 1, 2, Splitter::Random), ) .expect("Fit should work"); @@ -1507,9 +1671,14 @@ mod tests { .unwrap(); let y = vec![0.0f64, 1.0, 2.0, 100.0, 101.0]; - let tree: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, None, params(None, 1, 3, Splitter::Best)) - .expect("Fit should work"); + let tree: TestTree = BaseTree::fit_inner( + &x, + &y, + None, + &MSECriterion, + params(None, 1, 3, Splitter::Best), + ) + .expect("Fit should work"); assert_eq!(tree.nodes()[0].split_value, Some(3.5)); let true_child = tree.nodes()[0].true_child.unwrap(); let false_child = tree.nodes()[0].false_child.unwrap(); @@ -1521,7 +1690,7 @@ mod tests { // If min_samples_split is larger than the number of samples, the root is a leaf. for splitter in [Splitter::Best, Splitter::Random] { let tree: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, None, params(None, 1, 6, splitter)) + BaseTree::fit_inner(&x, &y, None, &MSECriterion, params(None, 1, 6, splitter)) .expect("Fit should work"); assert_eq!(tree.nodes().len(), 1, "splitter: {splitter:?}"); } @@ -1534,15 +1703,25 @@ mod tests { let y = vec![0.0f64, 0.0, 0.0, 100.0]; // With no limit, the best split isolates the last sample. - let tree: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, None, params(None, 1, 2, Splitter::Best)) - .expect("Fit should work"); + let tree: TestTree = BaseTree::fit_inner( + &x, + &y, + None, + &MSECriterion, + params(None, 1, 2, Splitter::Best), + ) + .expect("Fit should work"); assert_eq!(tree.nodes()[0].split_value, Some(3.5)); // The best split is not permitted. The split with exactly 2 samples on each side is used. - let tree: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, None, params(None, 2, 2, Splitter::Best)) - .expect("Fit should work"); + let tree: TestTree = BaseTree::fit_inner( + &x, + &y, + None, + &MSECriterion, + params(None, 2, 2, Splitter::Best), + ) + .expect("Fit should work"); assert_eq!(tree.nodes()[0].split_value, Some(2.5)); assert_eq!(tree.nodes().len(), 3); assert_eq!(leaf_counts(&tree, &x, &[1; 4]), vec![2, 2]); @@ -1550,7 +1729,7 @@ mod tests { // No split can give 3 samples on each side. for splitter in [Splitter::Best, Splitter::Random] { let tree: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, None, params(None, 3, 2, splitter)) + BaseTree::fit_inner(&x, &y, None, &MSECriterion, params(None, 3, 2, splitter)) .expect("Fit should work"); assert_eq!(tree.nodes().len(), 1, "splitter: {splitter:?}"); assert_eq!(tree.nodes()[0].output, 25.0); @@ -1582,17 +1761,18 @@ mod tests { for splitter in [Splitter::Best, Splitter::Random] { for min_samples_leaf in [1, 3, 5] { for seed in 0..5 { - let parameters = BaseTreeRegressorParameters { + let parameters = BaseTreeParameters { seed: Some(seed), ..params(None, min_samples_leaf, 2, splitter) }; - let tree: TestTree = BaseTreeRegressor::fit_weak_learner( + let tree: TestTree = BaseTree::fit_weak_learner( &x, &y, sample_weights, samples.clone(), 3, &order, + &MSECriterion, parameters, ) .expect("Fit should work"); @@ -1626,11 +1806,11 @@ mod tests { for splitter in [Splitter::Best, Splitter::Random] { let tree: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, None, params(None, 1, 2, splitter)) + BaseTree::fit_inner(&x, &y, None, &MSECriterion, params(None, 1, 2, splitter)) .expect("Fit should work"); assert_eq!(tree.nodes().len(), 1, "splitter: {splitter:?}"); assert_eq!(tree.depth, 0, "splitter: {splitter:?}"); - let y_hat = tree.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); assert_eq!(y_hat, vec![3.0; 4], "splitter: {splitter:?}"); } } @@ -1649,7 +1829,7 @@ mod tests { for splitter in [Splitter::Best, Splitter::Random] { let tree: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, None, params(None, 1, 2, splitter)) + BaseTree::fit_inner(&x, &y, None, &MSECriterion, params(None, 1, 2, splitter)) .expect("Fit should work"); assert!(tree.nodes().len() > 1, "splitter: {splitter:?}"); assert!( @@ -1661,9 +1841,14 @@ mod tests { ); } - let tree: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, None, params(None, 1, 2, Splitter::Best)) - .expect("Fit should work"); + let tree: TestTree = BaseTree::fit_inner( + &x, + &y, + None, + &MSECriterion, + params(None, 1, 2, Splitter::Best), + ) + .expect("Fit should work"); assert_eq!(tree.nodes()[0].split_value, Some(2.5)); } @@ -1682,15 +1867,15 @@ mod tests { let mut used_columns = [false; 2]; for seed in 0..10 { - let parameters = BaseTreeRegressorParameters { + let parameters = BaseTreeParameters { seed: Some(seed), ..params(None, 1, 2, Splitter::Best) }; let tree_single: TestTree = - BaseTreeRegressor::fit_inner(&x_single, &y, None, parameters.clone()) + BaseTree::fit_inner(&x_single, &y, None, &MSECriterion, parameters.clone()) .expect("Fit should work"); let tree_double: TestTree = - BaseTreeRegressor::fit_inner(&x_double, &y, None, parameters) + BaseTree::fit_inner(&x_double, &y, None, &MSECriterion, parameters) .expect("Fit should work"); assert!(tree_double.nodes().len() > 1); @@ -1727,24 +1912,25 @@ mod tests { for splitter in [Splitter::Best, Splitter::Random] { for max_depth in [Some(2), None] { let parameters = params(max_depth, 1, 2, splitter); - let tree: TestTree = BaseTreeRegressor::fit_weak_learner( + let tree: TestTree = BaseTree::fit_weak_learner( &x, &y, None, counts.clone(), 2, &order, + &MSECriterion, parameters.clone(), ) .expect("Fit should work"); let tree_rep: TestTree = - BaseTreeRegressor::fit_inner(&x_rep, &y_rep, None, parameters) + BaseTree::fit_inner(&x_rep, &y_rep, None, &MSECriterion, parameters) .expect("Fit should work"); assert!(tree.nodes().len() > 1); assert_same_tree(&tree, &tree_rep); - let y_hat = tree.predict(&x).expect("Predict should work"); - let y_hat_rep = tree_rep.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); + let y_hat_rep = predict(&tree_rep, &x); assert!(mean_absolute_error(&y_hat, &y_hat_rep) < 1e-9); } } @@ -1766,27 +1952,387 @@ mod tests { for splitter in [Splitter::Best, Splitter::Random] { for max_depth in [Some(2), None] { let parameters = params(max_depth, 1, 2, splitter); - let tree: TestTree = BaseTreeRegressor::fit_weak_learner( + let tree: TestTree = BaseTree::fit_weak_learner( &x, &y, Some(&weights), counts.clone(), 2, &order, + &MSECriterion, parameters.clone(), ) .expect("Fit should work"); // in this tree, there are no counts but the rows are repeated physically let tree_rep: TestTree = - BaseTreeRegressor::fit_inner(&x_rep, &y_rep, Some(&w_rep), parameters) + BaseTree::fit_inner(&x_rep, &y_rep, Some(&w_rep), &MSECriterion, parameters) .expect("Fit should work"); assert!(tree.nodes().len() > 1); assert_same_tree(&tree, &tree_rep); - let y_hat = tree.predict(&x).expect("Predict should work"); - let y_hat_rep = tree_rep.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); + let y_hat_rep = predict(&tree_rep, &x); assert!(mean_absolute_error(&y_hat, &y_hat_rep) < 1e-9); } } } + + /// Data for the split_score tests: 2 features with ties, weights, and bootstrap counts. + struct SplitScoreFixture { + x: DenseMatrix, + y: Vec, + weights: Vec, + counts: Vec, + order: Vec>, + } + + fn split_score_fixture() -> SplitScoreFixture { + let x = DenseMatrix::from_2d_vec(&vec![ + vec![1.0_f64, 7.0], + vec![2.0, 3.0], + vec![2.0, 5.0], + vec![3.0, 1.0], + vec![4.0, 8.0], + vec![5.0, 2.0], + vec![6.0, 6.0], + vec![6.0, 4.0], + vec![7.0, 9.0], + vec![8.0, 0.0], + ]) + .unwrap(); + let order: Vec> = (0..2) + .map(|j| { + let mut col: Vec = x.get_col(j).iterator(0).copied().collect(); + col.argsort_mut() + }) + .collect(); + SplitScoreFixture { + x, + y: vec![1.0, 3.0, 2.0, 8.0, 5.0, 9.0, 4.0, 7.0, 6.0, 10.0], + weights: vec![1.0, 2.0, 0.5, 1.0, 3.0, 1.0, 1.0, 0.25, 2.0, 1.0], + counts: vec![1, 2, 0, 1, 3, 1, 0, 2, 1, 1], + order, + } + } + + /// Fit one tree for each element type. Each tree comes with the mass of each row. + fn split_score_trees( + f: &SplitScoreFixture, + parameters: BaseTreeParameters, + ) -> Vec<(&'static str, TestTree, Vec)> { + let n = f.y.len(); + vec![ + ( + "unit", + BaseTree::fit_inner(&f.x, &f.y, None, &MSECriterion, parameters.clone()).unwrap(), + vec![1.0; n], + ), + ( + "weighted", + BaseTree::fit_inner( + &f.x, + &f.y, + Some(&f.weights), + &MSECriterion, + parameters.clone(), + ) + .unwrap(), + f.weights.clone(), + ), + ( + "counted", + BaseTree::fit_weak_learner( + &f.x, + &f.y, + None, + f.counts.clone(), + 2, + &f.order, + &MSECriterion, + parameters, + ) + .unwrap(), + f.counts.iter().map(|&c| c as f64).collect(), + ), + ] + } + + #[test] + fn split_score_is_some_only_for_internal_nodes() { + // A leaf has no split, so its split_score must be None. + // An internal node has a split, so its split_score must be Some. + let f = split_score_fixture(); + + for max_depth in [Some(0), Some(1), Some(2), Some(3), None] { + for min_samples_leaf in [1, 2, 3] { + for splitter in [Splitter::Best, Splitter::Random] { + let parameters = params(max_depth, min_samples_leaf, 2, splitter); + for (kind, tree, _) in split_score_trees(&f, parameters) { + for (i, node) in tree.nodes().iter().enumerate() { + let case = format!( + "{kind}, max_depth: {max_depth:?}, min_samples_leaf: {min_samples_leaf}, \ + splitter: {splitter:?}, node: {i}" + ); + assert_eq!( + node.true_child.is_some(), + node.false_child.is_some(), + "{case}" + ); + if node.true_child.is_some() { + assert!(node.split_score.is_some(), "internal node, {case}"); + } else { + assert!(node.split_score.is_none(), "leaf, {case}"); + } + } + } + } + } + } + } + + #[test] + fn split_scores_add_up_to_total_impurity_decrease() { + // Each split_score must be the real impurity improvement, not the proxy. + // The improvements of all splits then add up to + // weighted SSE(root) - sum of weighted SSE(leaves). + // The depths 0 and 1 are covered by split_score_is_some_only_for_internal_nodes. + + // Weighted sum of squared errors around the weighted mean + fn weighted_sse(samples: &[(f64, f64)]) -> f64 { + let mass: f64 = samples.iter().map(|&(_, m)| m).sum(); + if mass == 0.0 { + return 0.0; + } + let mean = samples.iter().map(|&(y, m)| m * y).sum::() / mass; + samples + .iter() + .map(|&(y, m)| m * (y - mean) * (y - mean)) + .sum() + } + + let f = split_score_fixture(); + + for max_depth in [Some(2), Some(3), None] { + for min_samples_leaf in [1, 2, 3] { + for splitter in [Splitter::Best, Splitter::Random] { + let parameters = params(max_depth, min_samples_leaf, 2, splitter); + for (kind, tree, mass) in split_score_trees(&f, parameters) { + let case = format!( + "{kind}, max_depth: {max_depth:?}, min_samples_leaf: {min_samples_leaf}, \ + splitter: {splitter:?}" + ); + + let samples: Vec<(f64, f64)> = + f.y.iter().copied().zip(mass.iter().copied()).collect(); + let mut leaves: Vec> = vec![Vec::new(); tree.nodes().len()]; + for (row, &sample) in samples.iter().enumerate() { + leaves[leaf_of(&tree, &f.x, row)].push(sample); + } + let expected = weighted_sse(&samples) + - leaves.iter().map(|leaf| weighted_sse(leaf)).sum::(); + + let total: f64 = tree.nodes().iter().filter_map(|n| n.split_score).sum(); + + assert!(tree.nodes().len() > 1, "tree has no split, {case}"); + assert!( + (total - expected).abs() <= 1e-9 * weighted_sse(&samples), + "sum of split_score: {total}, expected: {expected}, {case}" + ); + } + } + } + } + } + + #[test] + fn zero_weight_rows_do_not_change_the_tree() { + // A row with weight 0 has mass 0, so grow() leaves it out. The tree must be the + // same as a tree fitted without that row: no zero-mass leaf, and the x value of + // the row (2.5, between the other values) must not move a split threshold. + // The fractional weights make the sums depend on the order of the additions. + let x = DenseMatrix::from_2d_vec(&vec![ + vec![3.0_f64], + vec![1.0], + vec![2.5], + vec![2.0], + vec![4.0], + ]) + .unwrap(); + let y = vec![1.0f64, 2.0, 10.0, 3.0, 7.0]; + let weights = vec![0.1f64, 0.2, 0.0, 0.3, 0.4]; + + let x_kept = + DenseMatrix::from_2d_vec(&vec![vec![3.0_f64], vec![1.0], vec![2.0], vec![4.0]]) + .unwrap(); + let y_kept = vec![1.0f64, 2.0, 3.0, 7.0]; + let weights_kept = vec![0.1f64, 0.2, 0.3, 0.4]; + + for splitter in [Splitter::Best, Splitter::Random] { + let parameters = params(None, 1, 2, splitter); + let tree: TestTree = + BaseTree::fit_inner(&x, &y, Some(&weights), &MSECriterion, parameters.clone()) + .expect("Fit should work"); + let tree_kept: TestTree = BaseTree::fit_inner( + &x_kept, + &y_kept, + Some(&weights_kept), + &MSECriterion, + parameters, + ) + .expect("Fit should work"); + + assert!( + tree == tree_kept, + "{splitter:?}: zero-weight row changed the tree" + ); + } + } + + #[test] + fn split_is_invariant_to_target_offset() { + // Adding a constant to y must not change the splits. With a fixed seed, both fits + // see the same candidate thresholds, so only the ranking of the candidates can differ. + let n = 40; + let x = DenseMatrix::from_2d_vec( + &(0..n) + .map(|i| vec![i as f64, ((i * 7) % n) as f64, ((i * 13) % n) as f64]) + .collect::>(), + ) + .unwrap(); + let pattern: Vec = (0..n) + .map(|i| (i / 10) as f64 + 0.1 * (i % 3) as f64) + .collect(); + + for splitter in [Splitter::Best, Splitter::Random] { + let fit = |offset: f64| -> TestTree { + let y: Vec = pattern.iter().map(|v| offset + v).collect(); + BaseTree::fit_inner(&x, &y, None, &MSECriterion, params(Some(4), 1, 2, splitter)) + .expect("Fit should work") + }; + let reference = fit(0.0); + for offset in [1e6, 1e8, 1e9] { + let shifted = fit(offset); + assert_eq!( + reference.nodes().len(), + shifted.nodes().len(), + "{splitter:?} {offset}" + ); + for (a, b) in reference.nodes().iter().zip(shifted.nodes()) { + assert_eq!(a.split_feature, b.split_feature, "{splitter:?} {offset}"); + assert_eq!(a.split_value, b.split_value, "{splitter:?} {offset}"); + } + } + } + } + + /// Weighted variance of `y`: sum w_i (y_i - mean)^2 / sum w_i. + fn weighted_variance(y: &[f64], w: &[f64]) -> f64 { + let mass: f64 = w.iter().sum(); + let mean = y.iter().zip(w).map(|(yi, wi)| yi * wi).sum::() / mass; + y.iter() + .zip(w) + .map(|(yi, wi)| wi * (yi - mean).powi(2)) + .sum::() + / mass + } + + #[test] + fn feature_importances_single_split_known_answer() { + // Column 1 is constant, so only column 0 can be used for a split. + let x = DenseMatrix::from_2d_vec(&vec![ + vec![0.0_f64, 7.0], + vec![1.0, 7.0], + vec![2.0, 7.0], + vec![3.0, 7.0], + ]) + .unwrap(); + let y = vec![1.0, 2.0, 5.0, 6.0]; + let tree: TestTree = BaseTree::fit_inner( + &x, + &y, + None, + &MSECriterion, + params(Some(2), 1, 2, Splitter::Best), + ) + .expect("Fit should work"); + assert_eq!(tree.nodes().len(), 3); + + // Split at 1.5: SSE(root) = 17, SSE(true) = SSE(false) = 0.5, root mass = 4. + // Importance = (17 - 0.5 - 0.5) / 4 = 4. + let mut importances = vec![0.0; 2]; + tree.compute_feature_importances(false, &mut importances); + assert_eq!(importances.len(), 2); + assert!((importances[0] - 4.0).abs() < 1e-12); + assert_eq!(importances[1], 0.0); + + let mut normalized = vec![0.0; 2]; + tree.compute_feature_importances(true, &mut normalized); + assert!((normalized[0] - 1.0).abs() < 1e-12); + assert_eq!(normalized[1], 0.0); + } + + #[test] + fn feature_importances_full_depth_sum_to_weighted_variance() { + // All rows are different, so at full depth every leaf is pure. The sum of the + // unnormalized importances is then the weighted MSE of the root (the weighted variance of y). + let (x_rows, y) = element_path_data(); + let x = DenseMatrix::from_2d_vec(&x_rows).unwrap(); + let weights = vec![0.5, 1.5, 2.0, 0.25, 1.0, 3.0, 0.75, 1.25]; + + for splitter in [Splitter::Best, Splitter::Random] { + for sample_weights in [None, Some(weights.as_slice())] { + let tree: TestTree = BaseTree::fit_inner( + &x, + &y, + sample_weights, + &MSECriterion, + params(None, 1, 2, splitter), + ) + .expect("Fit should work"); + let w = sample_weights.map_or(vec![1.0; y.len()], |w| w.to_vec()); + + let mut importances = vec![0.0f64; 2]; + tree.compute_feature_importances(false, &mut importances); + assert!(importances.iter().all(|&v| v >= 0.0)); + let sum: f64 = importances.iter().sum(); + assert!((sum - weighted_variance(&y, &w)).abs() < 1e-9); + + let mut normalized = vec![0.0f64; 2]; + tree.compute_feature_importances(true, &mut normalized); + assert!((normalized.iter().sum::() - 1.0).abs() < 1e-12); + } + } + } + + #[test] + fn feature_importances_do_not_depend_on_weight_scale() { + // Multiply all weights by 4 (exact in f64). The tree and the importances must not change. + let (x_rows, y) = element_path_data(); + let x = DenseMatrix::from_2d_vec(&x_rows).unwrap(); + let weights = vec![0.5, 1.5, 2.0, 0.25, 1.0, 3.0, 0.75, 1.25]; + let scaled: Vec = weights.iter().map(|w| 4.0 * w).collect(); + let mut fi_tree = vec![0f64; x.shape().1]; + let mut fi_scaled_tree = vec![0f64; x.shape().1]; + + for splitter in [Splitter::Best, Splitter::Random] { + for max_depth in [Some(2), None] { + let parameters = params(max_depth, 1, 2, splitter); + let tree: TestTree = + BaseTree::fit_inner(&x, &y, Some(&weights), &MSECriterion, parameters.clone()) + .expect("Fit should work"); + let tree_scaled: TestTree = + BaseTree::fit_inner(&x, &y, Some(&scaled), &MSECriterion, parameters) + .expect("Fit should work"); + assert_same_tree(&tree, &tree_scaled); + + for normalize in [false, true] { + tree.compute_feature_importances(normalize, &mut fi_tree); + tree_scaled.compute_feature_importances(normalize, &mut fi_scaled_tree); + for (ai, bi) in fi_tree.iter().zip(fi_scaled_tree.iter()) { + assert!((ai - bi).abs() < 1e-9); + } + } + } + } + } } diff --git a/src/tree/criterion.rs b/src/tree/criterion.rs new file mode 100644 index 00000000..aa6f3f50 --- /dev/null +++ b/src/tree/criterion.rs @@ -0,0 +1,422 @@ +//! Split criteria for decision trees. +//! +//! A criterion keeps sufficient statistics for a set of samples. With these statistics, +//! the tree finds the best split and the value of a leaf. +//! +//! References: +//! * ["Classification and regression trees", Breiman L, Friedman J H, Olshen R A, and Stone C J, 1984](https://www.taylorfrancis.com/books/mono/10.1201/9781315139470/classification-regression-trees-leo-breiman) +//! * scikit-learn, `sklearn/tree/_criterion.pyx` (proxy improvement) + +use crate::numbers::basenum::Number; +use std::fmt::Debug; + +/// Impurity measure that a tree uses to select splits. +pub(crate) trait Criterion { + /// Sufficient statistics for a set of samples. + type Stats: Clone; + /// Target value of one sample. + type Target: Copy; + /// Value of a leaf. + type Output: Clone + Debug; + + /// Statistics for an empty set. Use only this function to make new statistics. + fn empty_stats(&self) -> Self::Stats; + + /// Set `stats` to the empty set. Keeps the allocation. + fn reset(&self, stats: &mut Self::Stats); + + /// Add one sample with weight `mass` to `stats`. + fn add(&self, stats: &mut Self::Stats, target: Self::Target, mass: f64); + + /// Set `out = parent - part`. `part` must be a subset of `parent`. + fn difference(&self, parent: &Self::Stats, part: &Self::Stats, out: &mut Self::Stats); + + /// Larger is better. Only comparable between splits of the same parent (proxy, as in sklearn). + fn proxy_improvement(&self, true_side: &Self::Stats, false_side: &Self::Stats) -> f64; + + /// Weighted decrease in impurity from this split: + /// `mass(parent) * imp(parent) - mass(true) * imp(true) - mass(false) * imp(false)`. + /// Not normalized by the total mass. Compute it one time for each split. + fn impurity_improvement( + &self, + parent: &Self::Stats, + true_side: &Self::Stats, + false_side: &Self::Stats, + ) -> f64; + + /// Value of a leaf that contains the samples of `stats`. + fn leaf_value(&self, stats: &Self::Stats) -> Self::Output; + + /// Convert a target value of the training data to the target of this criterion. + fn target(&self, y: T) -> Self::Target; +} + +/// Value in a leaf of a tree. +pub(crate) trait LeafOutput: Clone + Debug { + fn approximate_eq(&self, other: &Self) -> bool; +} + +impl LeafOutput for f64 { + fn approximate_eq(&self, other: &f64) -> bool { + (self - other).abs() < f64::EPSILON + } +} + +/// Statistics for the weighted mean of the targets. +#[derive(Clone, Debug, Default)] +pub(crate) struct MeanStats { + sum: f64, // \sum w_i y_i + mass: f64, // \sum w_i +} + +/// Mean squared error criterion for regression. The leaf value is the weighted mean. +#[derive(Debug)] +pub(crate) struct MSECriterion; + +impl Criterion for MSECriterion { + type Stats = MeanStats; + type Target = f64; + type Output = f64; + + fn empty_stats(&self) -> Self::Stats { + MeanStats { + sum: 0.0, + mass: 0.0, + } + } + + fn reset(&self, stats: &mut Self::Stats) { + stats.sum = 0.0; + stats.mass = 0.0; + } + + #[inline] + fn add(&self, stats: &mut Self::Stats, target: Self::Target, mass: f64) { + stats.mass += mass; + stats.sum += target * mass; + } + + #[inline] + fn difference(&self, parent: &Self::Stats, part: &Self::Stats, out: &mut Self::Stats) { + out.mass = parent.mass - part.mass; + out.sum = parent.sum - part.sum; + } + + #[inline] + fn proxy_improvement(&self, true_side: &Self::Stats, false_side: &Self::Stats) -> f64 { + if true_side.mass <= 0.0 || false_side.mass <= 0.0 { + return 0.0; + } + let d = true_side.sum * false_side.mass - false_side.sum * true_side.mass; + (d * d) / (true_side.mass * false_side.mass) + } + + fn impurity_improvement( + &self, + parent: &Self::Stats, + true_side: &Self::Stats, + false_side: &Self::Stats, + ) -> f64 { + if true_side.mass <= 0.0 || false_side.mass <= 0.0 { + return 0.0; + } + let d = true_side.sum * false_side.mass - false_side.sum * true_side.mass; + (d * d) / (true_side.mass * false_side.mass * parent.mass) + } + + fn leaf_value(&self, stats: &Self::Stats) -> Self::Output { + if stats.mass > 0.0 { + stats.sum / stats.mass + } else { + 0.0 + } + } + + fn target(&self, y: T) -> Self::Target { + y.to_f64().unwrap() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const EPS: f64 = 1e-12; + + fn mse_stats(samples: &[(f64, f64)]) -> MeanStats { + let criterion = MSECriterion; + let mut stats = criterion.empty_stats(); + for &(y, w) in samples { + criterion.add(&mut stats, y, w); + } + stats + } + + // Sum of squared errors around the mean + fn sse(ys: &[f64]) -> f64 { + let mean = ys.iter().sum::() / ys.len() as f64; + ys.iter().map(|y| (y - mean) * (y - mean)).sum() + } + + #[test] + fn mse_empty_stats() { + let criterion = MSECriterion; + let empty = criterion.empty_stats(); + assert_eq!(empty.sum, 0.0); + assert_eq!(empty.mass, 0.0); + assert_eq!(criterion.leaf_value(&empty), 0.0); + assert_eq!(criterion.proxy_improvement(&empty, &empty), 0.0); + } + + #[test] + fn mse_leaf_value_is_weighted_mean() { + let criterion = MSECriterion; + let stats = mse_stats(&[(1.0, 1.0), (3.0, 1.0), (5.0, 2.0)]); + assert_eq!(stats.mass, 4.0); + assert_eq!(stats.sum, 14.0); + assert!((criterion.leaf_value(&stats) - 3.5).abs() < EPS); + } + + #[test] + fn mse_reset_clears_stats() { + let criterion = MSECriterion; + let mut stats = mse_stats(&[(1.0, 1.0), (3.0, 2.0)]); + criterion.reset(&mut stats); + assert_eq!(stats.sum, 0.0); + assert_eq!(stats.mass, 0.0); + } + + #[test] + fn mse_difference_gives_complement() { + let criterion = MSECriterion; + let parent = mse_stats(&[(1.0, 1.0), (2.0, 2.0), (10.0, 1.0), (11.0, 3.0)]); + let part = mse_stats(&[(1.0, 1.0), (2.0, 2.0)]); + let expected = mse_stats(&[(10.0, 1.0), (11.0, 3.0)]); + + let mut out = criterion.empty_stats(); + criterion.difference(&parent, &part, &mut out); + assert!((out.sum - expected.sum).abs() < EPS); + assert!((out.mass - expected.mass).abs() < EPS); + } + + #[test] + fn mse_proxy_is_sse_reduction_times_parent_mass() { + let criterion = MSECriterion; + let ys = [1.0, 2.0, 10.0, 11.0]; + let parent = mse_stats(&[(1.0, 1.0), (2.0, 1.0), (10.0, 1.0), (11.0, 1.0)]); + let left = mse_stats(&[(1.0, 1.0), (2.0, 1.0)]); + let right = mse_stats(&[(10.0, 1.0), (11.0, 1.0)]); + + let proxy = criterion.proxy_improvement(&left, &right); + assert!((proxy - 324.0).abs() < EPS); + + let expected = sse(&ys) - sse(&ys[..2]) - sse(&ys[2..]); + assert!((proxy / parent.mass - expected).abs() < EPS); + assert!((proxy / parent.mass - 81.0).abs() < EPS); + } + + #[test] + fn mse_proxy_ranks_better_split_higher() { + let criterion = MSECriterion; + let good = criterion.proxy_improvement( + &mse_stats(&[(1.0, 1.0), (2.0, 1.0)]), + &mse_stats(&[(10.0, 1.0), (11.0, 1.0)]), + ); + let bad = criterion.proxy_improvement( + &mse_stats(&[(1.0, 1.0)]), + &mse_stats(&[(2.0, 1.0), (10.0, 1.0), (11.0, 1.0)]), + ); + assert!(good > bad); + } + + // Weighted sum of squared errors around the weighted mean + fn weighted_sse(samples: &[(f64, f64)]) -> f64 { + let mass: f64 = samples.iter().map(|&(_, w)| w).sum(); + if mass == 0.0 { + return 0.0; + } + let mean = samples.iter().map(|&(y, w)| w * y).sum::() / mass; + samples + .iter() + .map(|&(y, w)| w * (y - mean) * (y - mean)) + .sum() + } + + #[test] + fn mse_improvement_known_answer() { + let criterion = MSECriterion; + let parent = mse_stats(&[(1.0, 1.0), (2.0, 1.0), (10.0, 1.0), (11.0, 1.0)]); + let left = mse_stats(&[(1.0, 1.0), (2.0, 1.0)]); + let right = mse_stats(&[(10.0, 1.0), (11.0, 1.0)]); + + // SSE(parent) = 82, SSE(left) = SSE(right) = 0.5 + let improvement = criterion.impurity_improvement(&parent, &left, &right); + assert!((improvement - 81.0).abs() < EPS); + } + + #[test] + fn mse_improvement_matches_weighted_sse_reduction() { + let criterion = MSECriterion; + let left_samples = [(1.5, 0.5), (2.0, 2.0), (4.0, 1.25)]; + let right_samples = [(7.0, 3.0), (9.5, 0.75)]; + let all_samples: Vec<(f64, f64)> = left_samples + .iter() + .chain(right_samples.iter()) + .copied() + .collect(); + + let improvement = criterion.impurity_improvement( + &mse_stats(&all_samples), + &mse_stats(&left_samples), + &mse_stats(&right_samples), + ); + let expected = + weighted_sse(&all_samples) - weighted_sse(&left_samples) - weighted_sse(&right_samples); + assert!((improvement - expected).abs() < 1e-9); + } + + #[test] + fn mse_improvement_weight_equals_repeated_rows() { + let criterion = MSECriterion; + let weighted = criterion.impurity_improvement( + &mse_stats(&[(1.0, 2.0), (3.0, 1.0), (8.0, 3.0)]), + &mse_stats(&[(1.0, 2.0), (3.0, 1.0)]), + &mse_stats(&[(8.0, 3.0)]), + ); + let repeated = criterion.impurity_improvement( + &mse_stats(&[ + (1.0, 1.0), + (1.0, 1.0), + (3.0, 1.0), + (8.0, 1.0), + (8.0, 1.0), + (8.0, 1.0), + ]), + &mse_stats(&[(1.0, 1.0), (1.0, 1.0), (3.0, 1.0)]), + &mse_stats(&[(8.0, 1.0), (8.0, 1.0), (8.0, 1.0)]), + ); + assert!((weighted - repeated).abs() < EPS); + } + + #[test] + fn mse_improvement_pure_children_removes_all_error() { + let criterion = MSECriterion; + let samples = [(2.0, 1.0), (2.0, 3.0), (6.0, 2.0)]; + let improvement = criterion.impurity_improvement( + &mse_stats(&samples), + &mse_stats(&samples[..2]), + &mse_stats(&samples[2..]), + ); + assert!((improvement - weighted_sse(&samples)).abs() < EPS); + } + + #[test] + fn mse_improvement_is_zero_without_information() { + let criterion = MSECriterion; + // Both children have the same mean as the parent + let parent = mse_stats(&[(1.0, 1.0), (3.0, 1.0), (1.0, 1.0), (3.0, 1.0)]); + let left = mse_stats(&[(1.0, 1.0), (3.0, 1.0)]); + let right = mse_stats(&[(1.0, 1.0), (3.0, 1.0)]); + assert!(criterion.impurity_improvement(&parent, &left, &right).abs() < EPS); + + // An empty side does not change the impurity + let empty = criterion.empty_stats(); + assert!( + criterion + .impurity_improvement(&parent, &parent, &empty) + .abs() + < EPS + ); + assert!( + criterion + .impurity_improvement(&parent, &empty, &parent) + .abs() + < EPS + ); + } + + #[test] + fn mse_improvement_is_additive_over_the_tree() { + // The sum of the improvements of all splits is SSE(root) - sum of SSE(leaves) + let criterion = MSECriterion; + let samples = [(1.0, 1.0), (2.0, 2.0), (5.0, 1.0), (9.0, 0.5), (12.0, 1.5)]; + let root = mse_stats(&samples); + let left = mse_stats(&samples[..3]); + let right = mse_stats(&samples[3..]); + let left_left = mse_stats(&samples[..2]); + let left_right = mse_stats(&samples[2..3]); + + let total = criterion.impurity_improvement(&root, &left, &right) + + criterion.impurity_improvement(&left, &left_left, &left_right); + let expected = weighted_sse(&samples) + - weighted_sse(&samples[..2]) + - weighted_sse(&samples[2..3]) + - weighted_sse(&samples[3..]); + assert!((total - expected).abs() < 1e-9); + } + + #[test] + fn mse_proxy_is_improvement_times_parent_mass() { + // For all splits of one parent, the ranking must be the same as for the proxy + let criterion = MSECriterion; + let samples = [(1.0, 1.0), (4.0, 2.0), (2.0, 0.5), (8.0, 1.0), (3.0, 1.5)]; + let parent = mse_stats(&samples); + let mut right = criterion.empty_stats(); + for k in 0..=samples.len() { + let left = mse_stats(&samples[..k]); + criterion.difference(&parent, &left, &mut right); + let improvement = criterion.impurity_improvement(&parent, &left, &right); + let proxy = criterion.proxy_improvement(&left, &right); + assert!(improvement >= 0.0, "split {k}: {improvement}"); + assert!( + (proxy - improvement * parent.mass).abs() < 1e-9, + "split {k}: {improvement} {proxy}" + ); + } + } + + #[test] + fn mse_improvement_is_invariant_to_target_offset() { + // The SSE does not change when all targets get the same offset. + // A large offset must not change the improvement (cancellation in sum^2 / mass). + let criterion = MSECriterion; + for offset in [0.0, 1e6, 1e8, 1e9, 1e10] { + let y = |v: f64| (offset + v, 1.0); + let parent = mse_stats(&[y(1.0), y(2.0), y(10.0), y(11.0)]); + let left = mse_stats(&[y(1.0), y(2.0)]); + let right = mse_stats(&[y(10.0), y(11.0)]); + let improvement = criterion.impurity_improvement(&parent, &left, &right); + assert!( + (improvement - 81.0).abs() < 1e-6, + "offset {offset}: improvement {improvement}, expected 81" + ); + } + } + + #[test] + fn mse_target_converts_numbers_to_f64() { + let criterion = MSECriterion; + assert_eq!(criterion.target(3.25_f64), 3.25); + assert_eq!(criterion.target(-1.5_f32), -1.5); + assert_eq!(criterion.target(-7_i32), -7.0); + assert_eq!(criterion.target(200_u8), 200.0); + assert_eq!(criterion.target(42_usize), 42.0); + // Largest integer that f64 represents exactly + assert_eq!( + criterion.target(9_007_199_254_740_992_i64), + 9_007_199_254_740_992.0 + ); + } + + #[test] + fn f64_approximate_eq() { + assert!(1.5.approximate_eq(&1.5)); + assert!(0.0.approximate_eq(&-0.0)); + // Difference smaller than EPSILON + assert!(0.0.approximate_eq(&(f64::EPSILON / 2.0))); + // Difference of exactly EPSILON is not "approximately equal" (strict comparison) + assert!(!0.0.approximate_eq(&f64::EPSILON)); + assert!(!1.0.approximate_eq(&1.1)); + assert!(!f64::NAN.approximate_eq(&f64::NAN)); + } +} diff --git a/src/tree/decision_tree_regressor.rs b/src/tree/decision_tree_regressor.rs index 12ab72b6..de0c207c 100644 --- a/src/tree/decision_tree_regressor.rs +++ b/src/tree/decision_tree_regressor.rs @@ -64,11 +64,12 @@ use std::fmt::Debug; #[cfg(feature = "serde")] use serde::{Deserialize, Serialize}; -use super::base_tree_regressor::{BaseTreeRegressor, BaseTreeRegressorParameters, Splitter}; +use super::base_tree::{BaseTree, BaseTreeParameters, Splitter}; use crate::api::{Predictor, SupervisedEstimator}; use crate::error::Failed; use crate::linalg::basic::arrays::{Array1, Array2}; use crate::numbers::basenum::Number; +use crate::tree::criterion::MSECriterion; /// Validates the sample weights fn validate_sample_weights(sample_weights: &[f64], n_rows: usize) -> Result<(), Failed> { @@ -112,7 +113,8 @@ pub struct DecisionTreeRegressorParameters { #[derive(Debug)] pub struct DecisionTreeRegressor, Y: Array1> { - tree_regressor: Option>, + tree_regressor: Option>, + num_features: Option, } impl DecisionTreeRegressorParameters { @@ -292,6 +294,7 @@ impl, Y: Array1> fn new() -> Self { Self { tree_regressor: None, + num_features: None, } } @@ -342,16 +345,17 @@ impl, Y: Array1> sample_weights: Option<&[f64]>, parameters: DecisionTreeRegressorParameters, ) -> Result, Failed> { - let tree_parameters = BaseTreeRegressorParameters { + let tree_parameters = BaseTreeParameters { max_depth: parameters.max_depth, min_samples_leaf: parameters.min_samples_leaf, min_samples_split: parameters.min_samples_split, seed: parameters.seed, splitter: Splitter::Best, }; - let tree = BaseTreeRegressor::fit_inner(x, y, sample_weights, tree_parameters)?; + let tree = BaseTree::fit_inner(x, y, sample_weights, &MSECriterion, tree_parameters)?; Ok(Self { tree_regressor: Some(tree), + num_features: Some(x.shape().1), }) } @@ -359,12 +363,53 @@ impl, Y: Array1> /// * `x` - _KxM_ data where _K_ is number of observations and _M_ is number of features. pub fn predict(&self, x: &X) -> Result { match &self.tree_regressor { - Some(tree) => tree.predict(x), + Some(tree) => { + let mut result = Y::zeros(x.shape().0); + + let (n, _) = x.shape(); + + for i in 0..n { + result.set(i, TY::from_f64(*tree.leaf_output_for_row(x, i)).unwrap()); + } + + Ok(result) + } None => Err(Failed::predict( "'fit' should be called before calling 'predict'", )), } } + + /// Compute the feature importances for this tree. + /// * `normalize` - if `true` then the importances are normalized to one. + /// * returns: a vector where index `j` gives the feature importances for feature `j` + /// + /// ``` + /// use smartcore::linalg::basic::matrix::DenseMatrix; + /// use smartcore::tree::decision_tree_regressor::*; + /// + /// // Column 1 is constant, so only column 0 can split the data. + /// let x = DenseMatrix::from_2d_array(&[&[0., 7.], &[1., 7.], &[2., 7.], &[3., 7.]]).unwrap(); + /// let y: Vec = vec![1.0, 2.0, 5.0, 6.0]; + /// let parameters = DecisionTreeRegressorParameters::default().with_max_depth(2); + /// let tree = DecisionTreeRegressor::fit(&x, &y, parameters).unwrap(); + /// + /// // The split decreases the mean squared error from 4.25 to 0.25. + /// assert_eq!(tree.compute_feature_importances(false).unwrap(), vec![4.0, 0.0]); + /// assert_eq!(tree.compute_feature_importances(true).unwrap(), vec![1.0, 0.0]); + /// ``` + pub fn compute_feature_importances(&self, normalize: bool) -> Result, Failed> { + match (&self.tree_regressor, self.num_features) { + (Some(tree), Some(num_features)) => { + let mut importances = vec![0.0f64; num_features]; + tree.compute_feature_importances(normalize, &mut importances); + Ok(importances) + } + _ => Err(Failed::predict( + "'fit' should be called before calling 'compute_feature_importances'", + )), + } + } } #[cfg(test)] @@ -687,6 +732,36 @@ mod tests { let deserialized_tree: DecisionTreeRegressor, Vec> = postcard::from_bytes(&postcard::to_allocvec(&tree).unwrap()).unwrap(); + // Compare tree structure assert_eq!(tree, deserialized_tree); + + // Indirectly check that `root_mass` and `num_features` are correctly restored. + for normalize in [false, true] { + assert_eq!( + tree.compute_feature_importances(normalize).unwrap(), + deserialized_tree + .compute_feature_importances(normalize) + .unwrap() + ); + } + } + + #[test] + fn feature_importances_single_split() { + // Column 1 is constant: only column 0 can split. + // Split at 1.5: SSE(root) = 17, SSE(children) = 0.5 + 0.5, root mass = 4. + let x = DenseMatrix::from_2d_array(&[&[0., 7.], &[1., 7.], &[2., 7.], &[3., 7.]]).unwrap(); + let y: Vec = vec![1.0, 2.0, 5.0, 6.0]; + let parameters = DecisionTreeRegressorParameters::default().with_max_depth(2); + let tree = DecisionTreeRegressor::fit(&x, &y, parameters).expect("Fit should work"); + + assert_eq!( + tree.compute_feature_importances(false).unwrap(), + vec![4.0, 0.0] + ); + assert_eq!( + tree.compute_feature_importances(true).unwrap(), + vec![1.0, 0.0] + ); } } diff --git a/src/tree/mod.rs b/src/tree/mod.rs index 82937a5b..84a72d73 100644 --- a/src/tree/mod.rs +++ b/src/tree/mod.rs @@ -19,7 +19,9 @@ //! //! -pub(crate) mod base_tree_regressor; +pub(crate) mod base_tree; +pub(crate) mod criterion; + /// Classification tree for dependent variables that take a finite number of unordered values. pub mod decision_tree_classifier; /// Regression tree for for dependent variables that take continuous or ordered discrete values.