From 39f64c109137f8ccee6baaa8a060f7c4c35cd6b2 Mon Sep 17 00:00:00 2001 From: slievens <5492873+slievens@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:40:28 +0200 Subject: [PATCH 1/8] Use MSE criterion in base tree regressor This is the first step in making base tree "regressor" be used by all the tree classifiers. Next step is to make the implementation generic over the used Criterion with callers deciding it. Add a Criterion trait with MSE and Gini implementations in src/tree/criterion.rs. Use MeanStats in BaseTreeRegressor in place of explicit sum/mass bookkeeping. - Store node stats in NodeVisitor; remove true/false_child_output - Keep scratch stats in SplitWorkspace to avoid new allocations - Rank splits by proxy improvement; store the real impurity improvement in split_score - Pass the sample count to find_best_cutoff and check min_samples_split at the caller --- src/tree/base_tree_regressor.rs | 214 +++++----- src/tree/criterion.rs | 667 ++++++++++++++++++++++++++++++++ src/tree/mod.rs | 2 + 3 files changed, 768 insertions(+), 115 deletions(-) create mode 100644 src/tree/criterion.rs diff --git a/src/tree/base_tree_regressor.rs b/src/tree/base_tree_regressor.rs index 56059824..385a6c30 100644 --- a/src/tree/base_tree_regressor.rs +++ b/src/tree/base_tree_regressor.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, MSECriterion, MeanStats}; #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[derive(Debug, Clone, Copy, Default)] @@ -130,9 +131,8 @@ struct NodeVisitor<'a, TX: Number + PartialOrd, TY: Number, X: Array2, Y: Ar 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: MeanStats, // statistics of the samples in this node _phantom_tx: PhantomData, _phantom_ty: PhantomData, } @@ -147,6 +147,7 @@ impl<'a, TX: Number + PartialOrd, TY: Number, X: Array2, Y: Array1> x: &'a X, y: &'a Y, level: u16, + stats: MeanStats, ) -> Self { NodeVisitor { x, @@ -154,9 +155,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, } @@ -324,6 +324,8 @@ struct SplitWorkspace { partition_buffer: Vec, sorted_by_feature: Vec>, variables: Vec, // feature indices which will be shuffled when mtry < n_features + true_stats: MeanStats, // scratch stats for split functions to use + false_stats: MeanStats, } impl SplitWorkspace { @@ -338,6 +340,8 @@ impl SplitWorkspace { partition_buffer, sorted_by_feature, variables: (0..n_features).collect::>(), + true_stats: MSECriterion.empty_stats(), + false_stats: MSECriterion.empty_stats(), } } } @@ -425,16 +429,19 @@ impl, Y: Array1> 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_stats = MSECriterion.empty_stats(); 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(); + MSECriterion.add( + &mut root_stats, + y.get(i).to_f64().unwrap(), + mass_of(i, &samples, sample_weights), + ); + n += samples[i]; } - let root = Node::new(sum / mass); + let root = Node::new(MSECriterion.leaf_value(&root_stats)); nodes.push(root); let counts: Vec = samples.iter().map(|&s| s as u32).collect(); @@ -475,11 +482,13 @@ impl, Y: Array1> _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(); - if base_tree.find_best_cutoff(&mut visitor, mtry, mass, &mut rng, &mut workspace) { + if n >= base_tree.parameters().min_samples_split + && base_tree.find_best_cutoff(&mut visitor, mtry, &mut rng, &mut workspace, n) + { visitor_stack.push(visitor); } @@ -529,48 +538,27 @@ impl, Y: Array1> &mut self, visitor: &mut NodeVisitor<'_, TX, TY, X, Y>, mtry: usize, - mass: f64, rng: &mut impl rand::Rng, workspace: &mut SplitWorkspace, + n: usize, ) -> bool { let (_, n_attr) = visitor.x.shape(); - 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; - // Note: in sklearn, the attributes are always considered in a random order if mtry < n_attr { 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); } Splitter::Best => { - self.find_best_split(visitor, n, mass, sum, parent_gain, *variable, workspace); + self.find_best_split(visitor, n, variable, workspace); } } } @@ -580,14 +568,11 @@ impl, Y: Array1> fn find_random_split( &mut self, - visitor: &mut NodeVisitor<'_, TX, TY, X, Y>, + visitor: &NodeVisitor<'_, TX, TY, X, Y>, n: usize, - mass: f64, - sum: f64, - parent_gain: f64, j: usize, rng: &mut impl rand::Rng, - workspace: &SplitWorkspace, + workspace: &mut SplitWorkspace, ) { if visitor.start_idx == visitor.end_idx { return; @@ -608,14 +593,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; + MSECriterion.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(); + MSECriterion.add( + true_stats, + visitor.y.get(elem.row()).to_f64().unwrap(), + elem.mass(), + ); true_count += elem.count(); - true_mass += elem.mass(); } else { break; } @@ -629,110 +618,92 @@ 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; + MSECriterion.difference(&visitor.stats, true_stats, false_stats); + let proxy = MSECriterion.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( &mut self, - visitor: &mut NodeVisitor<'_, TX, TY, X, Y>, + visitor: &NodeVisitor<'_, TX, TY, X, Y>, n: usize, - mass: f64, - sum: f64, - parent_gain: f64, j: usize, - workspace: &SplitWorkspace, + workspace: &mut SplitWorkspace, ) { - 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; + MSECriterion.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(); + MSECriterion.add( + true_stats, + visitor.y.get(elem.row()).to_f64().unwrap(), + 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(); + MSECriterion.add( + true_stats, + visitor.y.get(elem.row()).to_f64().unwrap(), + 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; + MSECriterion.difference(&visitor.stats, true_stats, false_stats); + let proxy = MSECriterion.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(); + MSECriterion.add( + true_stats, + visitor.y.get(elem.row()).to_f64().unwrap(), + elem.mass(), + ); } } @@ -750,12 +721,13 @@ impl, Y: Array1> rng: &mut impl rand::Rng, workspace: &mut SplitWorkspace, ) -> bool { + let parent_stats = &visitor.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; + let mut true_stats = MSECriterion.empty_stats(); + let mut false_stats = MSECriterion.empty_stats(); // for each row_index, does it belong in the true branch or not? let in_true_branch = &mut workspace.in_true_branch; @@ -769,13 +741,16 @@ impl, Y: Array1> } in_true_branch[e.row()] = true; true_count += e.count(); - true_mass += e.mass(); + MSECriterion.add( + &mut true_stats, + visitor.y.get(e.row()).to_f64().unwrap(), + e.mass(), + ); n_true += 1; } for e in &col[n_true..] { in_true_branch[e.row()] = false; false_count += e.count(); - false_mass += e.mass(); } // Stop early if it is clear that there will be too few examples in the leaf @@ -789,14 +764,21 @@ impl, Y: Array1> return false; } + // Store the real (mass-weighted) impurity improvement in split_score, not the proxy. + MSECriterion.difference(parent_stats, &true_stats, &mut false_stats); + self.nodes[visitor.node].split_score = + Some(MSECriterion.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(MSECriterion.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(MSECriterion.leaf_value(&false_stats))); self.nodes[visitor.node].true_child = Some(true_child_idx); self.nodes[visitor.node].false_child = Some(false_child_idx); @@ -828,10 +810,11 @@ impl, Y: Array1> 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) { visitor_stack.push(true_visitor); } @@ -843,10 +826,11 @@ impl, Y: Array1> 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) { visitor_stack.push(false_visitor); } diff --git a/src/tree/criterion.rs b/src/tree/criterion.rs new file mode 100644 index 00000000..27ab6fed --- /dev/null +++ b/src/tree/criterion.rs @@ -0,0 +1,667 @@ +//! 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::linalg::basic::arrays::ArrayView1; + +/// 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; + + /// 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; +} + +/// Statistics for the weighted mean of the targets. +#[derive(Clone, Copy, 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. +pub(crate) struct MSECriterion; + +impl MSECriterion { + // One term of the proxy improvement: sum^2 / mass + fn term(s: &MeanStats) -> f64 { + if s.mass > 0.0 { + s.sum * s.sum / s.mass + } else { + 0.0 + } + } +} + +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(always)] + fn add(&self, stats: &mut Self::Stats, target: Self::Target, mass: f64) { + stats.mass += mass; + stats.sum += target * mass; + } + + #[inline(always)] + 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(always)] + fn proxy_improvement(&self, true_side: &Self::Stats, false_side: &Self::Stats) -> f64 { + Self::term(true_side) + Self::term(false_side) + } + + fn impurity_improvement( + &self, + parent: &Self::Stats, + true_side: &Self::Stats, + false_side: &Self::Stats, + ) -> f64 { + self.proxy_improvement(true_side, false_side) - Self::term(parent) + } + + fn leaf_value(&self, stats: &Self::Stats) -> Self::Output { + if stats.mass > 0.0 { + stats.sum / stats.mass + } else { + 0.0 + } + } +} + +#[cfg_attr( + not(test), + expect(dead_code, reason = "not yet used by the classifier") +)] +/// Gini impurity criterion for classification. The leaf value is the class with the largest mass. +pub(crate) struct GiniCriterion { + num_classes: usize, +} + +#[cfg_attr( + not(test), + expect(dead_code, reason = "not yet used by the classifier") +)] +#[derive(Clone, Debug, Default)] +/// Weighted count for each class. +pub(crate) struct GiniStats { + class_counts: Vec, + mass: f64, +} + +impl GiniCriterion { + // One term of the proxy improvement: sum_k count_k^2 / mass + #[cfg_attr( + not(test), + expect(dead_code, reason = "not yet used by the classifier") + )] + fn term(stats: &GiniStats) -> f64 { + if stats.mass > 0.0 { + stats.class_counts.iter().map(|c| c * c).sum::() / stats.mass + } else { + 0.0 + } + } +} + +impl Criterion for GiniCriterion { + type Stats = GiniStats; + type Target = usize; + type Output = usize; + + fn empty_stats(&self) -> Self::Stats { + GiniStats { + class_counts: vec![0.0; self.num_classes], + mass: 0.0, + } + } + + fn reset(&self, stats: &mut Self::Stats) { + stats.class_counts.fill(0.0); + stats.mass = 0.0; + } + + #[inline(always)] + fn add(&self, stats: &mut Self::Stats, target: Self::Target, mass: f64) { + debug_assert!(target < self.num_classes); + stats.class_counts[target] += mass; + stats.mass += mass; + } + + #[inline(always)] + fn difference(&self, parent: &Self::Stats, part: &Self::Stats, out: &mut Self::Stats) { + debug_assert_eq!(parent.class_counts.len(), part.class_counts.len()); + debug_assert_eq!(parent.class_counts.len(), out.class_counts.len()); + for ((diff, p), c) in out + .class_counts + .iter_mut() + .zip(parent.class_counts.iter()) + .zip(part.class_counts.iter()) + { + *diff = p - c; + } + out.mass = parent.mass - part.mass; + } + + #[inline(always)] + fn proxy_improvement(&self, true_side: &Self::Stats, false_side: &Self::Stats) -> f64 { + Self::term(true_side) + Self::term(false_side) + } + + fn impurity_improvement( + &self, + parent: &Self::Stats, + true_side: &Self::Stats, + false_side: &Self::Stats, + ) -> f64 { + self.proxy_improvement(true_side, false_side) - Self::term(parent) + } + + fn leaf_value(&self, stats: &Self::Stats) -> Self::Output { + stats.class_counts.argmax() + } +} + +// TODO: add the entropy and classification-error criteria. + +#[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 + } + + fn gini_stats(num_classes: usize, samples: &[(usize, f64)]) -> GiniStats { + let criterion = GiniCriterion { num_classes }; + let mut stats = criterion.empty_stats(); + for &(class, w) in samples { + criterion.add(&mut stats, class, w); + } + stats + } + + // Weighted Gini impurity: \sum_c n_c * (1 - \sum_k p_k^2) + fn weighted_gini(sides: &[&GiniStats]) -> f64 { + sides + .iter() + .filter(|s| s.mass > 0.0) + .map(|s| { + let sum_p2: f64 = s + .class_counts + .iter() + .map(|c| (c / s.mass) * (c / s.mass)) + .sum(); + s.mass * (1.0 - sum_p2) + }) + .sum() + } + + // 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_matches_sse_reduction() { + 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 - 225.0).abs() < EPS); + + // proxy - S^2 / N is the reduction of the sum of squared errors + let reduction = proxy - parent.sum * parent.sum / parent.mass; + let expected = sse(&ys) - sse(&ys[..2]) - sse(&ys[2..]); + assert!((reduction - expected).abs() < EPS); + assert!((reduction - 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); + } + + #[test] + fn gini_empty_stats() { + let criterion = GiniCriterion { num_classes: 3 }; + let empty = criterion.empty_stats(); + assert_eq!(empty.class_counts, vec![0.0; 3]); + assert_eq!(empty.mass, 0.0); + assert_eq!(criterion.leaf_value(&empty), 0); + assert_eq!(criterion.proxy_improvement(&empty, &empty), 0.0); + } + + #[test] + fn gini_add_with_weights() { + let stats = gini_stats(3, &[(2, 2.5), (0, 1.0)]); + assert_eq!(stats.class_counts, vec![1.0, 0.0, 2.5]); + assert_eq!(stats.mass, 3.5); + } + + #[test] + fn gini_term_known_answer() { + let criterion = GiniCriterion { num_classes: 2 }; + let stats = gini_stats(2, &[(0, 7.0), (1, 3.0)]); + let empty = criterion.empty_stats(); + // (7^2 + 3^2) / 10 + assert!((criterion.proxy_improvement(&stats, &empty) - 5.8).abs() < EPS); + assert!((criterion.proxy_improvement(&empty, &stats) - 5.8).abs() < EPS); + } + + #[test] + fn gini_proxy_equals_n_minus_weighted_gini() { + let criterion = GiniCriterion { num_classes: 2 }; + let left = gini_stats(2, &[(0, 3.0), (1, 1.0)]); + let right = gini_stats(2, &[(0, 1.0), (1, 3.0)]); + let n = left.mass + right.mass; + + let impurity = weighted_gini(&[&left, &right]); + assert!((impurity - 3.0).abs() < EPS); + + let proxy = criterion.proxy_improvement(&left, &right); + assert!((proxy - (n - impurity)).abs() < EPS); + assert!((proxy - 5.0).abs() < EPS); + } + + #[test] + fn gini_difference_gives_complement() { + let criterion = GiniCriterion { num_classes: 2 }; + let parent = gini_stats(2, &[(0, 4.0), (1, 4.0)]); + let part = gini_stats(2, &[(0, 3.0), (1, 1.0)]); + + let mut out = criterion.empty_stats(); + criterion.difference(&parent, &part, &mut out); + assert_eq!(out.class_counts, vec![1.0, 3.0]); + assert_eq!(out.mass, 4.0); + } + + #[test] + fn gini_reset_clears_stats() { + let criterion = GiniCriterion { num_classes: 3 }; + let mut stats = gini_stats(3, &[(0, 1.0), (2, 2.0)]); + criterion.reset(&mut stats); + assert_eq!(stats.class_counts, vec![0.0; 3]); + assert_eq!(stats.mass, 0.0); + } + + #[test] + fn gini_pure_split_beats_mixed() { + let criterion = GiniCriterion { num_classes: 2 }; + let pure = + criterion.proxy_improvement(&gini_stats(2, &[(0, 4.0)]), &gini_stats(2, &[(1, 4.0)])); + let mixed = criterion.proxy_improvement( + &gini_stats(2, &[(0, 3.0), (1, 1.0)]), + &gini_stats(2, &[(0, 1.0), (1, 3.0)]), + ); + assert!((pure - 8.0).abs() < EPS); + assert!(pure > mixed); + } + + #[test] + fn gini_leaf_value() { + let criterion = GiniCriterion { num_classes: 3 }; + let stats = gini_stats(3, &[(0, 1.0), (1, 5.0), (2, 2.0)]); + assert_eq!(criterion.leaf_value(&stats), 1); + + // A tie gives the first class, same as `which_max` + let tie = gini_stats(3, &[(0, 3.0), (1, 3.0)]); + assert_eq!(criterion.leaf_value(&tie), 0); + } + + // 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_improvement_differs_from_proxy_by_a_constant() { + // 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(); + let mut offsets = Vec::new(); + 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); + assert!(improvement > -1e-9, "split {k}: {improvement}"); + offsets.push(criterion.proxy_improvement(&left, &right) - improvement); + } + for offset in &offsets { + assert!((offset - offsets[0]).abs() < 1e-9); + } + } + + #[test] + fn gini_improvement_known_answer() { + let criterion = GiniCriterion { num_classes: 2 }; + let parent = gini_stats(2, &[(0, 4.0), (1, 4.0)]); + let left = gini_stats(2, &[(0, 3.0), (1, 1.0)]); + let right = gini_stats(2, &[(0, 1.0), (1, 3.0)]); + + // 8 * gini(parent) = 4, weighted gini of children = 3 + let improvement = criterion.impurity_improvement(&parent, &left, &right); + assert!((improvement - 1.0).abs() < EPS); + } + + #[test] + fn gini_improvement_matches_weighted_gini_reduction() { + let criterion = GiniCriterion { num_classes: 3 }; + let left_samples = [(0, 1.5), (1, 0.25), (2, 2.0), (0, 0.75)]; + let right_samples = [(1, 3.0), (2, 0.5), (0, 1.0)]; + let all_samples: Vec<(usize, f64)> = left_samples + .iter() + .chain(right_samples.iter()) + .copied() + .collect(); + + let parent = gini_stats(3, &all_samples); + let left = gini_stats(3, &left_samples); + let right = gini_stats(3, &right_samples); + + let improvement = criterion.impurity_improvement(&parent, &left, &right); + let expected = weighted_gini(&[&parent]) - weighted_gini(&[&left, &right]); + assert!((improvement - expected).abs() < EPS); + } + + #[test] + fn gini_improvement_pure_children_removes_all_impurity() { + let criterion = GiniCriterion { num_classes: 3 }; + let parent = gini_stats(3, &[(0, 2.0), (1, 1.0), (2, 3.0)]); + let left = gini_stats(3, &[(0, 2.0), (1, 1.0)]); + let right = gini_stats(3, &[(2, 3.0)]); + + let improvement = criterion.impurity_improvement(&parent, &left, &right); + let expected = weighted_gini(&[&parent]) - weighted_gini(&[&left]); + assert!((improvement - expected).abs() < EPS); + + // A split into two pure children removes all impurity of the parent + let parent = gini_stats(2, &[(0, 4.0), (1, 4.0)]); + let improvement = criterion.impurity_improvement( + &parent, + &gini_stats(2, &[(0, 4.0)]), + &gini_stats(2, &[(1, 4.0)]), + ); + assert!((improvement - 4.0).abs() < EPS); + } + + #[test] + fn gini_improvement_is_zero_without_information() { + let criterion = GiniCriterion { num_classes: 2 }; + // Both children have the same class proportions as the parent + let parent = gini_stats(2, &[(0, 6.0), (1, 3.0)]); + let left = gini_stats(2, &[(0, 2.0), (1, 1.0)]); + let right = gini_stats(2, &[(0, 4.0), (1, 2.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 gini_improvement_differs_from_proxy_by_a_constant() { + // For all splits of one parent, the ranking must be the same as for the proxy + let criterion = GiniCriterion { num_classes: 3 }; + let samples = [(0, 1.0), (2, 2.0), (1, 0.5), (0, 1.0), (2, 1.5), (1, 3.0)]; + let parent = gini_stats(3, &samples); + let mut right = criterion.empty_stats(); + let mut offsets = Vec::new(); + for k in 0..=samples.len() { + let left = gini_stats(3, &samples[..k]); + criterion.difference(&parent, &left, &mut right); + let improvement = criterion.impurity_improvement(&parent, &left, &right); + assert!(improvement > -1e-9, "split {k}: {improvement}"); + offsets.push(criterion.proxy_improvement(&left, &right) - improvement); + } + for offset in &offsets { + assert!((offset - offsets[0]).abs() < 1e-9); + } + } +} diff --git a/src/tree/mod.rs b/src/tree/mod.rs index 82937a5b..601d67cd 100644 --- a/src/tree/mod.rs +++ b/src/tree/mod.rs @@ -20,6 +20,8 @@ //! pub(crate) mod base_tree_regressor; +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. From ce1ee4897ed9c36d3c1b3b5cb7215c454047aada Mon Sep 17 00:00:00 2001 From: slievens <5492873+slievens@users.noreply.github.com> Date: Thu, 15 Oct 2026 19:04:21 +0200 Subject: [PATCH 2/8] No leaf has a split_score. Make sure no leaf has a split_core. Add tests to check this. Also check that no "proxy" values are leftover by verifying that the sum of all split_scores adds up to the total impurity decrease. --- src/tree/base_tree_regressor.rs | 177 +++++++++++++++++++++++++++++++- 1 file changed, 176 insertions(+), 1 deletion(-) diff --git a/src/tree/base_tree_regressor.rs b/src/tree/base_tree_regressor.rs index 385a6c30..feec097d 100644 --- a/src/tree/base_tree_regressor.rs +++ b/src/tree/base_tree_regressor.rs @@ -486,7 +486,8 @@ impl, Y: Array1> let mut visitor_stack: Vec> = Vec::new(); - if n >= base_tree.parameters().min_samples_split + if base_tree.parameters().max_depth.unwrap_or(u16::MAX) > 1 + && n >= base_tree.parameters().min_samples_split && base_tree.find_best_cutoff(&mut visitor, mtry, &mut rng, &mut workspace, n) { visitor_stack.push(visitor); @@ -1759,4 +1760,178 @@ mod tests { } } } + + /// 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: BaseTreeRegressorParameters, + ) -> Vec<(&'static str, TestTree, Vec)> { + let n = f.y.len(); + vec![ + ( + "unit", + BaseTreeRegressor::fit_inner(&f.x, &f.y, None, parameters.clone()).unwrap(), + vec![1.0; n], + ), + ( + "weighted", + BaseTreeRegressor::fit_inner(&f.x, &f.y, Some(&f.weights), parameters.clone()) + .unwrap(), + f.weights.clone(), + ), + ( + "counted", + BaseTreeRegressor::fit_weak_learner( + &f.x, + &f.y, + None, + f.counts.clone(), + 2, + &f.order, + 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. + + // Index of the leaf that contains `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 { + let node = &tree.nodes()[node_id]; + node_id = if *x.get((row, node.split_feature)) <= node.split_value.unwrap() { + true_child + } else { + node.false_child.unwrap() + }; + } + node_id + } + + // 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}" + ); + } + } + } + } + } } From 6505833c8236e51efb897222ec37c439c5660389 Mon Sep 17 00:00:00 2001 From: slievens <5492873+slievens@users.noreply.github.com> Date: Thu, 15 Oct 2026 19:09:09 +0200 Subject: [PATCH 3/8] Remove gini and tests --- src/tree/criterion.rs | 307 ------------------------------------------ 1 file changed, 307 deletions(-) diff --git a/src/tree/criterion.rs b/src/tree/criterion.rs index 27ab6fed..f5f0d62a 100644 --- a/src/tree/criterion.rs +++ b/src/tree/criterion.rs @@ -7,8 +7,6 @@ //! * ["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::linalg::basic::arrays::ArrayView1; - /// Impurity measure that a tree uses to select splits. pub(crate) trait Criterion { /// Sufficient statistics for a set of samples. @@ -120,101 +118,6 @@ impl Criterion for MSECriterion { } } -#[cfg_attr( - not(test), - expect(dead_code, reason = "not yet used by the classifier") -)] -/// Gini impurity criterion for classification. The leaf value is the class with the largest mass. -pub(crate) struct GiniCriterion { - num_classes: usize, -} - -#[cfg_attr( - not(test), - expect(dead_code, reason = "not yet used by the classifier") -)] -#[derive(Clone, Debug, Default)] -/// Weighted count for each class. -pub(crate) struct GiniStats { - class_counts: Vec, - mass: f64, -} - -impl GiniCriterion { - // One term of the proxy improvement: sum_k count_k^2 / mass - #[cfg_attr( - not(test), - expect(dead_code, reason = "not yet used by the classifier") - )] - fn term(stats: &GiniStats) -> f64 { - if stats.mass > 0.0 { - stats.class_counts.iter().map(|c| c * c).sum::() / stats.mass - } else { - 0.0 - } - } -} - -impl Criterion for GiniCriterion { - type Stats = GiniStats; - type Target = usize; - type Output = usize; - - fn empty_stats(&self) -> Self::Stats { - GiniStats { - class_counts: vec![0.0; self.num_classes], - mass: 0.0, - } - } - - fn reset(&self, stats: &mut Self::Stats) { - stats.class_counts.fill(0.0); - stats.mass = 0.0; - } - - #[inline(always)] - fn add(&self, stats: &mut Self::Stats, target: Self::Target, mass: f64) { - debug_assert!(target < self.num_classes); - stats.class_counts[target] += mass; - stats.mass += mass; - } - - #[inline(always)] - fn difference(&self, parent: &Self::Stats, part: &Self::Stats, out: &mut Self::Stats) { - debug_assert_eq!(parent.class_counts.len(), part.class_counts.len()); - debug_assert_eq!(parent.class_counts.len(), out.class_counts.len()); - for ((diff, p), c) in out - .class_counts - .iter_mut() - .zip(parent.class_counts.iter()) - .zip(part.class_counts.iter()) - { - *diff = p - c; - } - out.mass = parent.mass - part.mass; - } - - #[inline(always)] - fn proxy_improvement(&self, true_side: &Self::Stats, false_side: &Self::Stats) -> f64 { - Self::term(true_side) + Self::term(false_side) - } - - fn impurity_improvement( - &self, - parent: &Self::Stats, - true_side: &Self::Stats, - false_side: &Self::Stats, - ) -> f64 { - self.proxy_improvement(true_side, false_side) - Self::term(parent) - } - - fn leaf_value(&self, stats: &Self::Stats) -> Self::Output { - stats.class_counts.argmax() - } -} - -// TODO: add the entropy and classification-error criteria. - #[cfg(test)] mod tests { use super::*; @@ -230,31 +133,6 @@ mod tests { stats } - fn gini_stats(num_classes: usize, samples: &[(usize, f64)]) -> GiniStats { - let criterion = GiniCriterion { num_classes }; - let mut stats = criterion.empty_stats(); - for &(class, w) in samples { - criterion.add(&mut stats, class, w); - } - stats - } - - // Weighted Gini impurity: \sum_c n_c * (1 - \sum_k p_k^2) - fn weighted_gini(sides: &[&GiniStats]) -> f64 { - sides - .iter() - .filter(|s| s.mass > 0.0) - .map(|s| { - let sum_p2: f64 = s - .class_counts - .iter() - .map(|c| (c / s.mass) * (c / s.mass)) - .sum(); - s.mass * (1.0 - sum_p2) - }) - .sum() - } - // Sum of squared errors around the mean fn sse(ys: &[f64]) -> f64 { let mean = ys.iter().sum::() / ys.len() as f64; @@ -334,93 +212,6 @@ mod tests { assert!(good > bad); } - #[test] - fn gini_empty_stats() { - let criterion = GiniCriterion { num_classes: 3 }; - let empty = criterion.empty_stats(); - assert_eq!(empty.class_counts, vec![0.0; 3]); - assert_eq!(empty.mass, 0.0); - assert_eq!(criterion.leaf_value(&empty), 0); - assert_eq!(criterion.proxy_improvement(&empty, &empty), 0.0); - } - - #[test] - fn gini_add_with_weights() { - let stats = gini_stats(3, &[(2, 2.5), (0, 1.0)]); - assert_eq!(stats.class_counts, vec![1.0, 0.0, 2.5]); - assert_eq!(stats.mass, 3.5); - } - - #[test] - fn gini_term_known_answer() { - let criterion = GiniCriterion { num_classes: 2 }; - let stats = gini_stats(2, &[(0, 7.0), (1, 3.0)]); - let empty = criterion.empty_stats(); - // (7^2 + 3^2) / 10 - assert!((criterion.proxy_improvement(&stats, &empty) - 5.8).abs() < EPS); - assert!((criterion.proxy_improvement(&empty, &stats) - 5.8).abs() < EPS); - } - - #[test] - fn gini_proxy_equals_n_minus_weighted_gini() { - let criterion = GiniCriterion { num_classes: 2 }; - let left = gini_stats(2, &[(0, 3.0), (1, 1.0)]); - let right = gini_stats(2, &[(0, 1.0), (1, 3.0)]); - let n = left.mass + right.mass; - - let impurity = weighted_gini(&[&left, &right]); - assert!((impurity - 3.0).abs() < EPS); - - let proxy = criterion.proxy_improvement(&left, &right); - assert!((proxy - (n - impurity)).abs() < EPS); - assert!((proxy - 5.0).abs() < EPS); - } - - #[test] - fn gini_difference_gives_complement() { - let criterion = GiniCriterion { num_classes: 2 }; - let parent = gini_stats(2, &[(0, 4.0), (1, 4.0)]); - let part = gini_stats(2, &[(0, 3.0), (1, 1.0)]); - - let mut out = criterion.empty_stats(); - criterion.difference(&parent, &part, &mut out); - assert_eq!(out.class_counts, vec![1.0, 3.0]); - assert_eq!(out.mass, 4.0); - } - - #[test] - fn gini_reset_clears_stats() { - let criterion = GiniCriterion { num_classes: 3 }; - let mut stats = gini_stats(3, &[(0, 1.0), (2, 2.0)]); - criterion.reset(&mut stats); - assert_eq!(stats.class_counts, vec![0.0; 3]); - assert_eq!(stats.mass, 0.0); - } - - #[test] - fn gini_pure_split_beats_mixed() { - let criterion = GiniCriterion { num_classes: 2 }; - let pure = - criterion.proxy_improvement(&gini_stats(2, &[(0, 4.0)]), &gini_stats(2, &[(1, 4.0)])); - let mixed = criterion.proxy_improvement( - &gini_stats(2, &[(0, 3.0), (1, 1.0)]), - &gini_stats(2, &[(0, 1.0), (1, 3.0)]), - ); - assert!((pure - 8.0).abs() < EPS); - assert!(pure > mixed); - } - - #[test] - fn gini_leaf_value() { - let criterion = GiniCriterion { num_classes: 3 }; - let stats = gini_stats(3, &[(0, 1.0), (1, 5.0), (2, 2.0)]); - assert_eq!(criterion.leaf_value(&stats), 1); - - // A tie gives the first class, same as `which_max` - let tie = gini_stats(3, &[(0, 3.0), (1, 3.0)]); - assert_eq!(criterion.leaf_value(&tie), 0); - } - // 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(); @@ -566,102 +357,4 @@ mod tests { assert!((offset - offsets[0]).abs() < 1e-9); } } - - #[test] - fn gini_improvement_known_answer() { - let criterion = GiniCriterion { num_classes: 2 }; - let parent = gini_stats(2, &[(0, 4.0), (1, 4.0)]); - let left = gini_stats(2, &[(0, 3.0), (1, 1.0)]); - let right = gini_stats(2, &[(0, 1.0), (1, 3.0)]); - - // 8 * gini(parent) = 4, weighted gini of children = 3 - let improvement = criterion.impurity_improvement(&parent, &left, &right); - assert!((improvement - 1.0).abs() < EPS); - } - - #[test] - fn gini_improvement_matches_weighted_gini_reduction() { - let criterion = GiniCriterion { num_classes: 3 }; - let left_samples = [(0, 1.5), (1, 0.25), (2, 2.0), (0, 0.75)]; - let right_samples = [(1, 3.0), (2, 0.5), (0, 1.0)]; - let all_samples: Vec<(usize, f64)> = left_samples - .iter() - .chain(right_samples.iter()) - .copied() - .collect(); - - let parent = gini_stats(3, &all_samples); - let left = gini_stats(3, &left_samples); - let right = gini_stats(3, &right_samples); - - let improvement = criterion.impurity_improvement(&parent, &left, &right); - let expected = weighted_gini(&[&parent]) - weighted_gini(&[&left, &right]); - assert!((improvement - expected).abs() < EPS); - } - - #[test] - fn gini_improvement_pure_children_removes_all_impurity() { - let criterion = GiniCriterion { num_classes: 3 }; - let parent = gini_stats(3, &[(0, 2.0), (1, 1.0), (2, 3.0)]); - let left = gini_stats(3, &[(0, 2.0), (1, 1.0)]); - let right = gini_stats(3, &[(2, 3.0)]); - - let improvement = criterion.impurity_improvement(&parent, &left, &right); - let expected = weighted_gini(&[&parent]) - weighted_gini(&[&left]); - assert!((improvement - expected).abs() < EPS); - - // A split into two pure children removes all impurity of the parent - let parent = gini_stats(2, &[(0, 4.0), (1, 4.0)]); - let improvement = criterion.impurity_improvement( - &parent, - &gini_stats(2, &[(0, 4.0)]), - &gini_stats(2, &[(1, 4.0)]), - ); - assert!((improvement - 4.0).abs() < EPS); - } - - #[test] - fn gini_improvement_is_zero_without_information() { - let criterion = GiniCriterion { num_classes: 2 }; - // Both children have the same class proportions as the parent - let parent = gini_stats(2, &[(0, 6.0), (1, 3.0)]); - let left = gini_stats(2, &[(0, 2.0), (1, 1.0)]); - let right = gini_stats(2, &[(0, 4.0), (1, 2.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 gini_improvement_differs_from_proxy_by_a_constant() { - // For all splits of one parent, the ranking must be the same as for the proxy - let criterion = GiniCriterion { num_classes: 3 }; - let samples = [(0, 1.0), (2, 2.0), (1, 0.5), (0, 1.0), (2, 1.5), (1, 3.0)]; - let parent = gini_stats(3, &samples); - let mut right = criterion.empty_stats(); - let mut offsets = Vec::new(); - for k in 0..=samples.len() { - let left = gini_stats(3, &samples[..k]); - criterion.difference(&parent, &left, &mut right); - let improvement = criterion.impurity_improvement(&parent, &left, &right); - assert!(improvement > -1e-9, "split {k}: {improvement}"); - offsets.push(criterion.proxy_improvement(&left, &right) - improvement); - } - for offset in &offsets { - assert!((offset - offsets[0]).abs() < 1e-9); - } - } } From d4e67bc743657b0f1843349abe5c0820ff7708a3 Mon Sep 17 00:00:00 2001 From: slievens <5492873+slievens@users.noreply.github.com> Date: Thu, 15 Oct 2026 19:40:17 +0200 Subject: [PATCH 4/8] Clamp mass residue to zero --- src/tree/base_tree_regressor.rs | 63 +++++++++++++++++++++++++-------- src/tree/criterion.rs | 42 ++++++++++++++++++++++ 2 files changed, 91 insertions(+), 14 deletions(-) diff --git a/src/tree/base_tree_regressor.rs b/src/tree/base_tree_regressor.rs index feec097d..b6632851 100644 --- a/src/tree/base_tree_regressor.rs +++ b/src/tree/base_tree_regressor.rs @@ -1874,20 +1874,6 @@ mod tests { // weighted SSE(root) - sum of weighted SSE(leaves). // The depths 0 and 1 are covered by split_score_is_some_only_for_internal_nodes. - // Index of the leaf that contains `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 { - let node = &tree.nodes()[node_id]; - node_id = if *x.get((row, node.split_feature)) <= node.split_value.unwrap() { - true_child - } else { - node.false_child.unwrap() - }; - } - node_id - } - // 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(); @@ -1934,4 +1920,53 @@ mod tests { } } } + + #[test] + fn zero_mass_leaf_with_fractional_weights() { + // The root adds the weights in row order, but split() adds them in the order of x. + // With fractional weights, the two sums can differ by one ulp. Then `parent - true` + // gives a tiny residue mass for a side that only has zero-weight rows. + // A leaf with no mass must get the value 0.0, and not residue_sum / residue_mass. + let x = DenseMatrix::from_2d_vec(&vec![vec![3.0_f64], vec![1.0], vec![2.0], vec![4.0]]) + .unwrap(); + let y = vec![1.0f64, 2.0, 3.0, 10.0]; + let weights = vec![0.1f64, 0.2, 0.3, 0.0]; + + let tree: TestTree = BaseTreeRegressor::fit_inner( + &x, + &y, + Some(&weights), + params(None, 1, 2, Splitter::Best), + ) + .expect("Fit should work"); + + let n_nodes = tree.nodes().len(); + let mut leaf_mass = vec![0f64; n_nodes]; + let mut leaf_sum = vec![0f64; n_nodes]; + for row in 0..y.len() { + let leaf = leaf_of(&tree, &x, row); + leaf_mass[leaf] += weights[row]; + leaf_sum[leaf] += weights[row] * y[row]; + } + + let mut n_zero_mass_leaves = 0; + for (i, node) in tree.nodes().iter().enumerate() { + if node.true_child.is_some() { + continue; + } + if leaf_mass[i] == 0.0 { + n_zero_mass_leaves += 1; + assert_eq!(node.output, 0.0, "leaf {i} has no mass"); + } else { + let mean = leaf_sum[i] / leaf_mass[i]; + assert!( + (node.output - mean).abs() <= 1e-9 * mean.abs(), + "leaf {i}: output {}, weighted mean {mean}", + node.output + ); + } + } + // The test must contain the case that it checks + assert_eq!(n_zero_mass_leaves, 1); + } } diff --git a/src/tree/criterion.rs b/src/tree/criterion.rs index f5f0d62a..10204fd1 100644 --- a/src/tree/criterion.rs +++ b/src/tree/criterion.rs @@ -26,6 +26,7 @@ pub(crate) trait Criterion { fn add(&self, stats: &mut Self::Stats, target: Self::Target, mass: f64); /// Set `out = parent - part`. `part` must be a subset of `parent`. + /// If the mass of `out` is only a rounding residue, `out` is set to the empty set. 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). @@ -45,6 +46,9 @@ pub(crate) trait Criterion { fn leaf_value(&self, stats: &Self::Stats) -> Self::Output; } +/// Relative mass below which `difference` gives the empty set. +const MASS_EPS: f64 = 1e-12; + /// Statistics for the weighted mean of the targets. #[derive(Clone, Copy, Debug, Default)] pub(crate) struct MeanStats { @@ -93,6 +97,11 @@ impl Criterion for MSECriterion { 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; + // Remove a rounding residue, see MASS_EPS. This also catches a negative residue. + if out.mass <= MASS_EPS * parent.mass { + out.mass = 0.0; + out.sum = 0.0; + } } #[inline(always)] @@ -180,6 +189,39 @@ mod tests { assert!((out.mass - expected.mass).abs() < EPS); } + #[test] + fn mse_difference_removes_rounding_residue() { + let criterion = MSECriterion; + // The same samples, added in a different order: the sums differ by one ulp. + let parent = mse_stats(&[(1.0, 0.1), (2.0, 0.2), (3.0, 0.3)]); + let part = mse_stats(&[(2.0, 0.2), (3.0, 0.3), (1.0, 0.1)]); + assert_ne!(parent.mass, part.mass, "the test needs a residue"); + + let mut out = criterion.empty_stats(); + criterion.difference(&parent, &part, &mut out); + assert_eq!(out.mass, 0.0); + assert_eq!(out.sum, 0.0); + assert_eq!(criterion.leaf_value(&out), 0.0); + + // Also when part is the larger sum (negative residue) + criterion.difference(&part, &parent, &mut out); + assert_eq!(out.mass, 0.0); + assert_eq!(out.sum, 0.0); + } + + #[test] + fn mse_difference_keeps_small_real_mass() { + // A real side with a small, but not tiny, mass must not be removed. + let criterion = MSECriterion; + let parent = mse_stats(&[(1.0, 1.0), (5.0, 1e-6)]); + let part = mse_stats(&[(1.0, 1.0)]); + + let mut out = criterion.empty_stats(); + criterion.difference(&parent, &part, &mut out); + assert!(out.mass > 0.0); + assert!((criterion.leaf_value(&out) - 5.0).abs() < 1e-6); + } + #[test] fn mse_proxy_matches_sse_reduction() { let criterion = MSECriterion; From e7a8008558c68a1c04aac3ff9b1cb0dd7cee9c0d Mon Sep 17 00:00:00 2001 From: slievens <5492873+slievens@users.noreply.github.com> Date: Sun, 18 Oct 2026 19:28:26 +0200 Subject: [PATCH 5/8] Use a stable MSE improvement and drop zero-mass rows MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Compute the MSE improvement as (S_T·M_F − S_F·M_T)² / (M_T·M_F·M_P). This prevents cancellation for targets with a large offset. - Remove rows with mass 0 in grow(). Zero-weight rows no longer count toward min_samples_leaf/min_samples_split and do not move thresholds. - Remove the MASS_EPS clamp from difference(). - In split(), add the false side directly instead of parent − true. - Add tests for target-offset invariance and zero-weight rows. --- src/tree/base_tree_regressor.rs | 156 ++++++++++++++++++++------------ src/tree/criterion.rs | 113 +++++++++-------------- 2 files changed, 142 insertions(+), 127 deletions(-) diff --git a/src/tree/base_tree_regressor.rs b/src/tree/base_tree_regressor.rs index 97bebc54..7f4f4a88 100644 --- a/src/tree/base_tree_regressor.rs +++ b/src/tree/base_tree_regressor.rs @@ -326,6 +326,7 @@ struct SplitWorkspace { variables: Vec, // feature indices which will be shuffled when mtry < n_features true_stats: MeanStats, // scratch stats for split functions to use false_stats: MeanStats, + parent_stats: MeanStats, } impl SplitWorkspace { @@ -342,6 +343,7 @@ impl SplitWorkspace { variables: (0..n_features).collect::>(), true_stats: MSECriterion.empty_stats(), false_stats: MSECriterion.empty_stats(), + parent_stats: MSECriterion.empty_stats(), } } } @@ -432,21 +434,25 @@ impl, Y: Array1> let mut n = 0usize; let mut root_stats = MSECriterion.empty_stats(); + // A row with mass 0 (bootstrap count 0 or sample weight 0) has no effect on + // the statistics. Leave it out, so that every side of a split with elements + // has a positive mass. Rows with mass zero are not counted towards min_samples_leaf etc. + let masses: Vec = (0..n_rows) + .map(|i| mass_of(i, &samples, sample_weights)) + .collect(); + for i in 0..n_rows { - MSECriterion.add( - &mut root_stats, - y.get(i).to_f64().unwrap(), - mass_of(i, &samples, sample_weights), - ); - n += samples[i]; + if masses[i] > 0.0 { + MSECriterion.add(&mut root_stats, y.get(i).to_f64().unwrap(), masses[i]); + n += samples[i]; + } } let root = Node::new(MSECriterion.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 - let n_kept = counts.iter().filter(|&&c| c > 0).count(); + // number of distinct rows of x in this tree with a postive mass + let n_kept = masses.iter().filter(|&&m| m > 0.0).count(); let sorted_by_feature: Vec> = order .iter() @@ -454,9 +460,8 @@ impl, Y: Array1> let mut out = vec![E::default(); n_kept + 1]; // preallocate, one additional place let mut w = 0usize; // index to write to for &i in col_order { - let c = counts[i]; - out[w] = E::new(i, c as usize, mass_of(i, &samples, sample_weights)); - w += (c > 0) as usize; // update w in a branchless way + out[w] = E::new(i, samples[i], masses[i]); + w += (masses[i] > 0.0) as usize; // update w in a branchless way } out.truncate(n_kept); out @@ -487,6 +492,7 @@ impl, Y: Array1> let mut visitor_stack: Vec> = Vec::new(); 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) { @@ -717,13 +723,12 @@ impl, Y: Array1> rng: &mut impl rand::Rng, workspace: &mut SplitWorkspace, ) -> bool { - let parent_stats = &visitor.stats; + let parent_stats = &mut workspace.parent_stats; + MSECriterion.reset(parent_stats); let this_node = &self.nodes[visitor.node]; let mut true_count = 0usize; let mut false_count = 0usize; - let mut true_stats = MSECriterion.empty_stats(); - let mut false_stats = MSECriterion.empty_stats(); // for each row_index, does it belong in the true branch or not? let in_true_branch = &mut workspace.in_true_branch; @@ -738,15 +743,27 @@ impl, Y: Array1> in_true_branch[e.row()] = true; true_count += e.count(); MSECriterion.add( - &mut true_stats, + parent_stats, visitor.y.get(e.row()).to_f64().unwrap(), e.mass(), ); n_true += 1; } + let true_stats = parent_stats.clone(); + let mut false_stats = MSECriterion.empty_stats(); for e in &col[n_true..] { in_true_branch[e.row()] = false; false_count += e.count(); + MSECriterion.add( + parent_stats, + visitor.y.get(e.row()).to_f64().unwrap(), + e.mass(), + ); + MSECriterion.add( + &mut false_stats, + visitor.y.get(e.row()).to_f64().unwrap(), + e.mass(), + ); } // Stop early if it is clear that there will be too few examples in the leaf @@ -761,7 +778,7 @@ impl, Y: Array1> } // Store the real (mass-weighted) impurity improvement in split_score, not the proxy. - MSECriterion.difference(parent_stats, &true_stats, &mut false_stats); + // false_stats is a direct sum, not parent - true: no cancellation for a small false side self.nodes[visitor.node].split_score = Some(MSECriterion.impurity_improvement(parent_stats, &true_stats, &false_stats)); @@ -1936,51 +1953,78 @@ mod tests { } #[test] - fn zero_mass_leaf_with_fractional_weights() { - // The root adds the weights in row order, but split() adds them in the order of x. - // With fractional weights, the two sums can differ by one ulp. Then `parent - true` - // gives a tiny residue mass for a side that only has zero-weight rows. - // A leaf with no mass must get the value 0.0, and not residue_sum / residue_mass. - let x = DenseMatrix::from_2d_vec(&vec![vec![3.0_f64], vec![1.0], vec![2.0], vec![4.0]]) - .unwrap(); - let y = vec![1.0f64, 2.0, 3.0, 10.0]; - let weights = vec![0.1f64, 0.2, 0.3, 0.0]; + 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 tree: TestTree = BaseTreeRegressor::fit_inner( - &x, - &y, - Some(&weights), - params(None, 1, 2, Splitter::Best), - ) - .expect("Fit should work"); + 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 = + BaseTreeRegressor::fit_inner(&x, &y, Some(&weights), parameters.clone()) + .expect("Fit should work"); + let tree_kept: TestTree = + BaseTreeRegressor::fit_inner(&x_kept, &y_kept, Some(&weights_kept), parameters) + .expect("Fit should work"); - let n_nodes = tree.nodes().len(); - let mut leaf_mass = vec![0f64; n_nodes]; - let mut leaf_sum = vec![0f64; n_nodes]; - for row in 0..y.len() { - let leaf = leaf_of(&tree, &x, row); - leaf_mass[leaf] += weights[row]; - leaf_sum[leaf] += weights[row] * y[row]; + assert!( + tree == tree_kept, + "{splitter:?}: zero-weight row changed the tree" + ); } + } - let mut n_zero_mass_leaves = 0; - for (i, node) in tree.nodes().iter().enumerate() { - if node.true_child.is_some() { - continue; - } - if leaf_mass[i] == 0.0 { - n_zero_mass_leaves += 1; - assert_eq!(node.output, 0.0, "leaf {i} has no mass"); - } else { - let mean = leaf_sum[i] / leaf_mass[i]; - assert!( - (node.output - mean).abs() <= 1e-9 * mean.abs(), - "leaf {i}: output {}, weighted mean {mean}", - node.output + #[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(); + BaseTreeRegressor::fit_inner(&x, &y, None, 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}"); + } } } - // The test must contain the case that it checks - assert_eq!(n_zero_mass_leaves, 1); } } diff --git a/src/tree/criterion.rs b/src/tree/criterion.rs index 10204fd1..bfd918e6 100644 --- a/src/tree/criterion.rs +++ b/src/tree/criterion.rs @@ -26,7 +26,6 @@ pub(crate) trait Criterion { fn add(&self, stats: &mut Self::Stats, target: Self::Target, mass: f64); /// Set `out = parent - part`. `part` must be a subset of `parent`. - /// If the mass of `out` is only a rounding residue, `out` is set to the empty set. 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). @@ -46,11 +45,8 @@ pub(crate) trait Criterion { fn leaf_value(&self, stats: &Self::Stats) -> Self::Output; } -/// Relative mass below which `difference` gives the empty set. -const MASS_EPS: f64 = 1e-12; - /// Statistics for the weighted mean of the targets. -#[derive(Clone, Copy, Debug, Default)] +#[derive(Clone, Debug, Default)] pub(crate) struct MeanStats { sum: f64, // \sum w_i y_i mass: f64, // \sum w_i @@ -59,17 +55,6 @@ pub(crate) struct MeanStats { /// Mean squared error criterion for regression. The leaf value is the weighted mean. pub(crate) struct MSECriterion; -impl MSECriterion { - // One term of the proxy improvement: sum^2 / mass - fn term(s: &MeanStats) -> f64 { - if s.mass > 0.0 { - s.sum * s.sum / s.mass - } else { - 0.0 - } - } -} - impl Criterion for MSECriterion { type Stats = MeanStats; type Target = f64; @@ -87,26 +72,25 @@ impl Criterion for MSECriterion { stats.mass = 0.0; } - #[inline(always)] + #[inline] fn add(&self, stats: &mut Self::Stats, target: Self::Target, mass: f64) { stats.mass += mass; stats.sum += target * mass; } - #[inline(always)] + #[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; - // Remove a rounding residue, see MASS_EPS. This also catches a negative residue. - if out.mass <= MASS_EPS * parent.mass { - out.mass = 0.0; - out.sum = 0.0; - } } - #[inline(always)] + #[inline] fn proxy_improvement(&self, true_side: &Self::Stats, false_side: &Self::Stats) -> f64 { - Self::term(true_side) + Self::term(false_side) + 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( @@ -115,7 +99,11 @@ impl Criterion for MSECriterion { true_side: &Self::Stats, false_side: &Self::Stats, ) -> f64 { - self.proxy_improvement(true_side, false_side) - Self::term(parent) + 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 { @@ -190,40 +178,7 @@ mod tests { } #[test] - fn mse_difference_removes_rounding_residue() { - let criterion = MSECriterion; - // The same samples, added in a different order: the sums differ by one ulp. - let parent = mse_stats(&[(1.0, 0.1), (2.0, 0.2), (3.0, 0.3)]); - let part = mse_stats(&[(2.0, 0.2), (3.0, 0.3), (1.0, 0.1)]); - assert_ne!(parent.mass, part.mass, "the test needs a residue"); - - let mut out = criterion.empty_stats(); - criterion.difference(&parent, &part, &mut out); - assert_eq!(out.mass, 0.0); - assert_eq!(out.sum, 0.0); - assert_eq!(criterion.leaf_value(&out), 0.0); - - // Also when part is the larger sum (negative residue) - criterion.difference(&part, &parent, &mut out); - assert_eq!(out.mass, 0.0); - assert_eq!(out.sum, 0.0); - } - - #[test] - fn mse_difference_keeps_small_real_mass() { - // A real side with a small, but not tiny, mass must not be removed. - let criterion = MSECriterion; - let parent = mse_stats(&[(1.0, 1.0), (5.0, 1e-6)]); - let part = mse_stats(&[(1.0, 1.0)]); - - let mut out = criterion.empty_stats(); - criterion.difference(&parent, &part, &mut out); - assert!(out.mass > 0.0); - assert!((criterion.leaf_value(&out) - 5.0).abs() < 1e-6); - } - - #[test] - fn mse_proxy_matches_sse_reduction() { + 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)]); @@ -231,13 +186,11 @@ mod tests { let right = mse_stats(&[(10.0, 1.0), (11.0, 1.0)]); let proxy = criterion.proxy_improvement(&left, &right); - assert!((proxy - 225.0).abs() < EPS); + assert!((proxy - 324.0).abs() < EPS); - // proxy - S^2 / N is the reduction of the sum of squared errors - let reduction = proxy - parent.sum * parent.sum / parent.mass; let expected = sse(&ys) - sse(&ys[..2]) - sse(&ys[2..]); - assert!((reduction - expected).abs() < EPS); - assert!((reduction - 81.0).abs() < EPS); + assert!((proxy / parent.mass - expected).abs() < EPS); + assert!((proxy / parent.mass - 81.0).abs() < EPS); } #[test] @@ -381,22 +334,40 @@ mod tests { } #[test] - fn mse_improvement_differs_from_proxy_by_a_constant() { + 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(); - let mut offsets = Vec::new(); 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); - assert!(improvement > -1e-9, "split {k}: {improvement}"); - offsets.push(criterion.proxy_improvement(&left, &right) - improvement); + 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}" + ); } - for offset in &offsets { - assert!((offset - offsets[0]).abs() < 1e-9); + } + + #[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" + ); } } } From 501e5d54849ecc4dbcfbcddff50893269e8f07de Mon Sep 17 00:00:00 2001 From: slievens <5492873+slievens@users.noreply.github.com> Date: Sun, 18 Oct 2026 20:47:53 +0200 Subject: [PATCH 6/8] Refactor grow --- src/tree/base_tree_regressor.rs | 40 ++++++++++++++++++++------------- 1 file changed, 24 insertions(+), 16 deletions(-) diff --git a/src/tree/base_tree_regressor.rs b/src/tree/base_tree_regressor.rs index 7f4f4a88..92bd7a36 100644 --- a/src/tree/base_tree_regressor.rs +++ b/src/tree/base_tree_regressor.rs @@ -426,8 +426,6 @@ impl, Y: Array1> order: &[Vec], parameters: BaseTreeRegressorParameters, ) -> Result, Failed> { - let n_rows = y.shape(); - let mut nodes: Vec = Vec::new(); let mut rng = get_rng_impl(parameters.seed); @@ -435,24 +433,33 @@ impl, Y: Array1> let mut root_stats = MSECriterion.empty_stats(); // A row with mass 0 (bootstrap count 0 or sample weight 0) has no effect on - // the statistics. Leave it out, so that every side of a split with elements - // has a positive mass. Rows with mass zero are not counted towards min_samples_leaf etc. - let masses: Vec = (0..n_rows) - .map(|i| mass_of(i, &samples, sample_weights)) - .collect(); + // 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 { - if masses[i] > 0.0 { - MSECriterion.add(&mut root_stats, y.get(i).to_f64().unwrap(), masses[i]); - n += samples[i]; - } + for (i, &c) in counts.iter().enumerate() { + // A row with mass 0 adds 0 to the statistics: no branch is necessary + MSECriterion.add( + &mut root_stats, + y.get(i).to_f64().unwrap(), + mass_of(i, &samples, sample_weights), + ); + n += c as usize; } let root = Node::new(MSECriterion.leaf_value(&root_stats)); nodes.push(root); - // number of distinct rows of x in this tree with a postive mass - let n_kept = masses.iter().filter(|&&m| m > 0.0).count(); + // 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 .iter() @@ -460,8 +467,9 @@ impl, Y: Array1> let mut out = vec![E::default(); n_kept + 1]; // preallocate, one additional place let mut w = 0usize; // index to write to for &i in col_order { - out[w] = E::new(i, samples[i], masses[i]); - w += (masses[i] > 0.0) as usize; // update w in a branchless way + let c = counts[i]; + out[w] = E::new(i, c as usize, mass_of(i, &samples, sample_weights)); + w += (c > 0) as usize; // update w in a branchless way } out.truncate(n_kept); out From c7eb8a382b6127ded6116681fc0797809cef893a Mon Sep 17 00:00:00 2001 From: slievens <5492873+slievens@users.noreply.github.com> Date: Sun, 18 Oct 2026 22:00:07 +0200 Subject: [PATCH 7/8] update changelog --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9d0dc5d7..2be393d5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,8 @@ 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. ## [0.6.16] ### Fixed From c34b214207f93af1e05cd9555e040fbe7fe56a75 Mon Sep 17 00:00:00 2001 From: slievens <5492873+slievens@users.noreply.github.com> Date: Tue, 20 Oct 2026 21:52:25 +0200 Subject: [PATCH 8/8] Add feature importances to the regression trees and forests Add compute_feature_importances (mean decrease in impurity, as in scikit-learn) to DecisionTreeRegressor, RandomForestRegressor and ExtraTreesRegressor. The forests normalize the importances of each tree, take the mean and normalize again. The fitted models store new fields (root_mass, num_features), so models serialized with an earlier version do not deserialize. --- CHANGELOG.md | 2 + src/ensemble/base_forest_regressor.rs | 38 ++++++ src/ensemble/extra_trees_regressor.rs | 165 ++++++++++++++++++++++++ src/ensemble/random_forest_regressor.rs | 163 +++++++++++++++++++++++ src/tree/base_tree_regressor.rs | 145 ++++++++++++++++++++- src/tree/decision_tree_regressor.rs | 64 +++++++++ 6 files changed, 572 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2be393d5..38c888e2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - 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. ## [0.6.16] ### Fixed diff --git a/src/ensemble/base_forest_regressor.rs b/src/ensemble/base_forest_regressor.rs index b48f8dad..cea579fe 100644 --- a/src/ensemble/base_forest_regressor.rs +++ b/src/ensemble/base_forest_regressor.rs @@ -71,6 +71,7 @@ pub struct BaseForestRegressor< Y: Array1, > { trees: Option>>, + num_features: Option, samples: Option>>, } @@ -170,6 +171,7 @@ impl, Y: Array1 Ok(BaseForestRegressor { trees: Some(trees), + num_features: Some(x.shape().1), samples: maybe_all_samples, }) } @@ -267,6 +269,42 @@ impl, Y: Array1 samples } + + /// 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; + } + } + + // compute the mean of each feature importance + let num_trees = trees.len() as f64; + for i in importances.iter_mut() { + *i /= num_trees; + } + + // 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'", + )), + } + } } #[cfg(test)] diff --git a/src/ensemble/extra_trees_regressor.rs b/src/ensemble/extra_trees_regressor.rs index a5a06f77..67a420fe 100644 --- a/src/ensemble/extra_trees_regressor.rs +++ b/src/ensemble/extra_trees_regressor.rs @@ -272,6 +272,36 @@ impl, Y: Array1 )), } } + + /// 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 +734,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/random_forest_regressor.rs b/src/ensemble/random_forest_regressor.rs index 364806e4..b2978e07 100644 --- a/src/ensemble/random_forest_regressor.rs +++ b/src/ensemble/random_forest_regressor.rs @@ -465,6 +465,42 @@ impl, Y: Array1 )), } } + + /// 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 +1135,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_regressor.rs index 92bd7a36..0720631d 100644 --- a/src/tree/base_tree_regressor.rs +++ b/src/tree/base_tree_regressor.rs @@ -50,6 +50,7 @@ pub struct BaseTreeRegressor, nodes: Vec, parameters: Option, depth: u16, + root_mass: f64, _phantom_tx: PhantomData, _phantom_ty: PhantomData, _phantom_x: PhantomData, @@ -430,6 +431,7 @@ impl, Y: Array1> let mut rng = get_rng_impl(parameters.seed); let mut n = 0usize; + let mut root_mass = 0f64; let mut root_stats = MSECriterion.empty_stats(); // A row with mass 0 (bootstrap count 0 or sample weight 0) has no effect on @@ -447,11 +449,9 @@ impl, Y: Array1> for (i, &c) in counts.iter().enumerate() { // A row with mass 0 adds 0 to the statistics: no branch is necessary - MSECriterion.add( - &mut root_stats, - y.get(i).to_f64().unwrap(), - mass_of(i, &samples, sample_weights), - ); + let mass = mass_of(i, &samples, sample_weights); + MSECriterion.add(&mut root_stats, y.get(i).to_f64().unwrap(), mass); + root_mass += mass; n += c as usize; } @@ -489,6 +489,7 @@ impl, Y: Array1> nodes, parameters: Some(parameters), depth: 0u16, + root_mass, _phantom_tx: PhantomData, _phantom_ty: PhantomData, _phantom_x: PhantomData, @@ -858,6 +859,35 @@ impl, Y: Array1> 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)] @@ -2035,4 +2065,109 @@ mod tests { } } } + + /// 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 = + BaseTreeRegressor::fit_inner(&x, &y, None, 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 = BaseTreeRegressor::fit_inner( + &x, + &y, + sample_weights, + 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 = + BaseTreeRegressor::fit_inner(&x, &y, Some(&weights), parameters.clone()) + .expect("Fit should work"); + let tree_scaled: TestTree = + BaseTreeRegressor::fit_inner(&x, &y, Some(&scaled), 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/decision_tree_regressor.rs b/src/tree/decision_tree_regressor.rs index 12ab72b6..46126cc1 100644 --- a/src/tree/decision_tree_regressor.rs +++ b/src/tree/decision_tree_regressor.rs @@ -113,6 +113,7 @@ pub struct DecisionTreeRegressorParameters { pub struct DecisionTreeRegressor, Y: Array1> { tree_regressor: Option>, + num_features: Option, } impl DecisionTreeRegressorParameters { @@ -292,6 +293,7 @@ impl, Y: Array1> fn new() -> Self { Self { tree_regressor: None, + num_features: None, } } @@ -352,6 +354,7 @@ impl, Y: Array1> let tree = BaseTreeRegressor::fit_inner(x, y, sample_weights, tree_parameters)?; Ok(Self { tree_regressor: Some(tree), + num_features: Some(x.shape().1), }) } @@ -365,6 +368,37 @@ impl, Y: Array1> )), } } + + /// 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 +721,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] + ); } }