diff --git a/src/tree/base_tree_regressor.rs b/src/tree/base_tree_regressor.rs index 56059824..b6632851 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,14 @@ 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 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); } @@ -529,48 +539,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 +569,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 +594,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 +619,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 +722,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 +742,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 +765,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 +811,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 +827,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); } @@ -1775,4 +1760,213 @@ 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. + + // Weighted sum of squared errors around the weighted mean + fn weighted_sse(samples: &[(f64, f64)]) -> f64 { + let mass: f64 = samples.iter().map(|&(_, m)| m).sum(); + if mass == 0.0 { + return 0.0; + } + let mean = samples.iter().map(|&(y, m)| m * y).sum::() / mass; + samples + .iter() + .map(|&(y, m)| m * (y - mean) * (y - mean)) + .sum() + } + + let f = split_score_fixture(); + + for max_depth in [Some(2), Some(3), None] { + for min_samples_leaf in [1, 2, 3] { + for splitter in [Splitter::Best, Splitter::Random] { + let parameters = params(max_depth, min_samples_leaf, 2, splitter); + for (kind, tree, mass) in split_score_trees(&f, parameters) { + let case = format!( + "{kind}, max_depth: {max_depth:?}, min_samples_leaf: {min_samples_leaf}, \ + splitter: {splitter:?}" + ); + + let samples: Vec<(f64, f64)> = + f.y.iter().copied().zip(mass.iter().copied()).collect(); + let mut leaves: Vec> = vec![Vec::new(); tree.nodes().len()]; + for (row, &sample) in samples.iter().enumerate() { + leaves[leaf_of(&tree, &f.x, row)].push(sample); + } + let expected = weighted_sse(&samples) + - leaves.iter().map(|leaf| weighted_sse(leaf)).sum::(); + + let total: f64 = tree.nodes().iter().filter_map(|n| n.split_score).sum(); + + assert!(tree.nodes().len() > 1, "tree has no split, {case}"); + assert!( + (total - expected).abs() <= 1e-9 * weighted_sse(&samples), + "sum of split_score: {total}, expected: {expected}, {case}" + ); + } + } + } + } + } + + #[test] + fn zero_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 new file mode 100644 index 00000000..10204fd1 --- /dev/null +++ b/src/tree/criterion.rs @@ -0,0 +1,402 @@ +//! 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) + +/// 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`. + /// 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). + 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; +} + +/// 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 { + 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; + // 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)] + 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(test)] +mod tests { + use super::*; + + const EPS: f64 = 1e-12; + + fn mse_stats(samples: &[(f64, f64)]) -> MeanStats { + let criterion = MSECriterion; + let mut stats = criterion.empty_stats(); + for &(y, w) in samples { + criterion.add(&mut stats, y, w); + } + stats + } + + // Sum of squared errors around the mean + fn sse(ys: &[f64]) -> f64 { + let mean = ys.iter().sum::() / ys.len() as f64; + ys.iter().map(|y| (y - mean) * (y - mean)).sum() + } + + #[test] + fn mse_empty_stats() { + let criterion = MSECriterion; + let empty = criterion.empty_stats(); + assert_eq!(empty.sum, 0.0); + assert_eq!(empty.mass, 0.0); + assert_eq!(criterion.leaf_value(&empty), 0.0); + assert_eq!(criterion.proxy_improvement(&empty, &empty), 0.0); + } + + #[test] + fn mse_leaf_value_is_weighted_mean() { + let criterion = MSECriterion; + let stats = mse_stats(&[(1.0, 1.0), (3.0, 1.0), (5.0, 2.0)]); + assert_eq!(stats.mass, 4.0); + assert_eq!(stats.sum, 14.0); + assert!((criterion.leaf_value(&stats) - 3.5).abs() < EPS); + } + + #[test] + fn mse_reset_clears_stats() { + let criterion = MSECriterion; + let mut stats = mse_stats(&[(1.0, 1.0), (3.0, 2.0)]); + criterion.reset(&mut stats); + assert_eq!(stats.sum, 0.0); + assert_eq!(stats.mass, 0.0); + } + + #[test] + fn mse_difference_gives_complement() { + let criterion = MSECriterion; + let parent = mse_stats(&[(1.0, 1.0), (2.0, 2.0), (10.0, 1.0), (11.0, 3.0)]); + let part = mse_stats(&[(1.0, 1.0), (2.0, 2.0)]); + let expected = mse_stats(&[(10.0, 1.0), (11.0, 3.0)]); + + let mut out = criterion.empty_stats(); + criterion.difference(&parent, &part, &mut out); + assert!((out.sum - expected.sum).abs() < EPS); + assert!((out.mass - expected.mass).abs() < EPS); + } + + #[test] + fn mse_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; + 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); + } + + // 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); + } + } +} 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.