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/4] 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/4] 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/4] 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/4] 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;