From d3dae39f35883b4273effd8ee808424132232b16 Mon Sep 17 00:00:00 2001 From: Johan Larsson Date: Thu, 8 Oct 2026 22:43:19 +0200 Subject: [PATCH 01/18] feat: add lazymatrix feature --- .github/workflows/ci.yml | 4 +- CHANGELOG.md | 5 + Cargo.toml | 4 + README.md | 23 + src/linear/elastic_net.rs | 323 +++++++++++- src/linear/lasso.rs | 290 ++++++++++- src/linear/lasso_optimizer.rs | 75 ++- src/linear/lazy_normalization.rs | 830 +++++++++++++++++++++++++++++++ src/linear/mod.rs | 2 + 9 files changed, 1527 insertions(+), 29 deletions(-) create mode 100644 src/linear/lazy_normalization.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7715aed7..01bcc12f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -57,6 +57,8 @@ jobs: - "--features serde" - "--features datasets" - "--features ndarray-bindings" + - "--features lazy-normalization" + - "--features lazy-normalization,ndarray-bindings" - "" steps: - uses: actions/checkout@v4 @@ -73,7 +75,7 @@ jobs: - name: Stable Build run: cargo build --no-default-features ${{ matrix.features }} - name: Tests - if: matrix.features == '--features ndarray-bindings' + if: matrix.features == '--features ndarray-bindings' || contains(matrix.features, 'lazy-normalization') run: cargo test --no-default-features ${{ matrix.features }} msrv: diff --git a/CHANGELOG.md b/CHANGELOG.md index 2c0ff5c5..e32e47a4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,11 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +### Added +- Optional `lazy-normalization` support for Lasso and Elastic Net. It applies normalization to a borrowed design and keeps Elastic Net's identity-block augmentation implicit, while retaining the dense Gram workspace. Enable the feature and set `with_lazy_normalization(true)` to select this path. The default remains eager. The new parameter field is available only with the feature and defaults to false during deserialization (#479). + +### Changed +- **Breaking change:** Enabling `lazy-normalization` adds a public `lazy_normalization` field to `LassoParameters`, `LassoSearchParameters`, `ElasticNetParameters`, and `ElasticNetSearchParameters`. Struct literals without `..Default::default()` must add `lazy_normalization: false` to retain the eager path, or `true` to select the lazy path. This also applies when another dependency enables the feature through Cargo feature unification (#479). ## [0.6.16] ### Fixed diff --git a/Cargo.toml b/Cargo.toml index 4ec3a4b5..3e8ab12e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,6 +24,7 @@ exclude = [ [dependencies] approx = "0.5.1" cfg-if = "1.0.0" +lazymatrix = { version = "0.6.0", optional = true, default-features = false } ndarray = { version = "0.17", optional = true } num-traits = "0.2.12" num = "0.4" @@ -46,6 +47,9 @@ serde = ["dep:serde", "dep:typetag"] # Optional bindings for the ndarray crate (DenseMatrix <-> ndarray). ndarray-bindings = ["dep:ndarray"] +# Apply Lasso and Elastic Net normalization without copying the design matrix. +lazy-normalization = ["dep:lazymatrix"] + # Sample datasets (iris, boston, digits, ...) and dataset generators. # Enables std_rand and serde because datasets ship serialized bundles. datasets = ["std_rand", "serde"] diff --git a/README.md b/README.md index b219b2f9..dac4ac5c 100644 --- a/README.md +++ b/README.md @@ -100,6 +100,29 @@ smartcore organizes algorithms into clear modules with consistent traits: Recent refactors emphasize reusable components in trees/forests and expanded multiclass SVM capabilities. XGBoost-style regression and single-linkage clustering have been added. See CHANGELOG for API changes and migration notes. +## Optional lazy normalization + +The optional `lazy-normalization` feature lets Lasso and Elastic Net borrow the +design matrix during normalization. Elastic Net also applies its identity-block +augmentation without storing the augmented design. Enable it with `features = +["lazy-normalization"]`, then select it with +`LassoParameters::default().with_lazy_normalization(true)` or +`ElasticNetParameters::default().with_lazy_normalization(true)`. The default +remains eager, even when the feature is enabled. The `normalize` option still +controls whether columns are centered and scaled. Parameter searches use the +fixed `lazy_normalization` setting. It works with `DenseMatrix` and +`ndarray-bindings`. + +Lazy normalization reduces memory use by computing normalized values when the +solver needs them. This adds repeated arithmetic and can increase fit time for +both estimators, including Elastic Net on large inputs. The solver still stores +a dense Gram matrix with one row and column per feature. + +Enabling the feature adds a public `lazy_normalization` field to both estimators' +parameter and search parameter structs. Code that constructs these structs +without `..Default::default()` must add the field. Missing fields default to +`false` during deserialization. + ## Data access and readers - CSV readers: Read matrices from CSV with configurable delimiter and header rows, with helpful error messages and testing utilities (including non-IO reader abstractions). diff --git a/src/linear/elastic_net.rs b/src/linear/elastic_net.rs index 30544bcd..f1045420 100644 --- a/src/linear/elastic_net.rs +++ b/src/linear/elastic_net.rs @@ -68,12 +68,20 @@ use crate::numbers::floatnum::FloatNumber; use crate::numbers::realnum::RealNumber; use crate::linear::lasso_optimizer::InteriorPointOptimizer; +#[cfg(feature = "lazy-normalization")] +use crate::linear::lazy_normalization::LazyDesign; /// Elastic net parameters #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[derive(Debug, Clone)] #[must_use] pub struct ElasticNetParameters { + #[cfg(feature = "lazy-normalization")] + #[cfg_attr(feature = "serde", serde(default))] + /// Borrow the design matrix instead of storing its normalized values. + /// This reduces memory use but can increase fit time. The default is false. + /// This option does not change whether normalization is applied. + pub lazy_normalization: bool, #[cfg_attr(feature = "serde", serde(default))] /// Regularization parameter. pub alpha: f64, @@ -104,6 +112,24 @@ pub struct ElasticNet, Y } impl ElasticNetParameters { + /// Select lazy matrix operations. Requires the `lazy-normalization` feature. + /// This reduces memory use but can increase fit time. + /// + /// ``` + /// # #[cfg(feature = "lazy-normalization")] + /// # { + /// use smartcore::linear::elastic_net::ElasticNetParameters; + /// + /// let parameters = ElasticNetParameters::default().with_lazy_normalization(true); + /// assert!(parameters.lazy_normalization); + /// # } + /// ``` + #[cfg(feature = "lazy-normalization")] + pub fn with_lazy_normalization(mut self, lazy_normalization: bool) -> Self { + self.lazy_normalization = lazy_normalization; + self + } + /// Regularization parameter. pub fn with_alpha(mut self, alpha: f64) -> Self { self.alpha = alpha; @@ -136,6 +162,8 @@ impl ElasticNetParameters { impl Default for ElasticNetParameters { fn default() -> Self { ElasticNetParameters { + #[cfg(feature = "lazy-normalization")] + lazy_normalization: false, alpha: 1.0, l1_ratio: 0.5, normalize: true, @@ -150,6 +178,13 @@ impl Default for ElasticNetParameters { #[derive(Debug, Clone)] #[must_use] pub struct ElasticNetSearchParameters { + #[cfg(feature = "lazy-normalization")] + #[cfg_attr(feature = "serde", serde(default))] + /// Borrow the design matrix instead of storing its normalized values. + /// This reduces memory use but can increase fit time. The default is false. + /// This setting applies to all parameter combinations in the search. + /// This option does not change whether normalization is applied. + pub lazy_normalization: bool, #[cfg_attr(feature = "serde", serde(default))] /// Regularization parameter. pub alpha: Vec, @@ -209,6 +244,8 @@ impl Iterator for ElasticNetSearchParametersIterator { } let next = ElasticNetParameters { + #[cfg(feature = "lazy-normalization")] + lazy_normalization: self.lasso_regression_search_parameters.lazy_normalization, alpha: self.lasso_regression_search_parameters.alpha[self.current_alpha], l1_ratio: self.lasso_regression_search_parameters.alpha[self.current_l1_ratio], normalize: self.lasso_regression_search_parameters.normalize[self.current_normalize], @@ -257,6 +294,8 @@ impl Default for ElasticNetSearchParameters { let default_params = ElasticNetParameters::default(); ElasticNetSearchParameters { + #[cfg(feature = "lazy-normalization")] + lazy_normalization: false, alpha: vec![default_params.alpha], l1_ratio: vec![default_params.l1_ratio], normalize: vec![default_params.normalize], @@ -320,7 +359,7 @@ impl, Y: Array1> y: &Y, parameters: ElasticNetParameters, ) -> Result, Failed> { - let (n, p) = x.shape(); + let (n, _) = x.shape(); if y.shape() != n { return Err(Failed::fit("Number of rows in X should = len(y)")); @@ -334,6 +373,32 @@ impl, Y: Array1> let y_mean = TX::from_f64(y.mean_by()).unwrap(); + #[cfg(feature = "lazy-normalization")] + let (w, b) = if parameters.lazy_normalization { + Self::fit_lazy_coefficients(x, y, ¶meters, l1_reg, l2_reg, y_mean)? + } else { + Self::fit_eager_coefficients(x, y, ¶meters, l1_reg, l2_reg, y_mean)? + }; + #[cfg(not(feature = "lazy-normalization"))] + let (w, b) = Self::fit_eager_coefficients(x, y, ¶meters, l1_reg, l2_reg, y_mean)?; + + Ok(ElasticNet { + intercept: Some(b), + coefficients: Some(w), + _phantom_ty: PhantomData, + _phantom_y: PhantomData, + }) + } + + fn fit_eager_coefficients( + x: &X, + y: &Y, + parameters: &ElasticNetParameters, + l1_reg: TX, + l2_reg: TX, + y_mean: TX, + ) -> Result<(X, TX), Failed> { + let (_, p) = x.shape(); let (w, b) = if parameters.normalize { let (scaled_x, col_mean, col_std) = Self::rescale_x(x)?; @@ -383,13 +448,73 @@ impl, Y: Array1> (X::from_column(&w), y_mean) }; + Ok((w, b)) + } - Ok(ElasticNet { - intercept: Some(b), - coefficients: Some(w), - _phantom_ty: PhantomData, - _phantom_y: PhantomData, - }) + #[cfg(feature = "lazy-normalization")] + fn fit_lazy_coefficients( + x: &X, + y: &Y, + parameters: &ElasticNetParameters, + l1_reg: TX, + l2_reg: TX, + y_mean: TX, + ) -> Result<(X, TX), Failed> { + let (n, p) = x.shape(); + let (w, b) = if parameters.normalize { + let x = { + let (col_mean, col_std) = Self::normalization_stats(x)?; + LazyDesign::new(x, Some(col_mean), Some(col_std)).with_l2(l2_reg) + }; + let (col_mean, col_std, gamma) = (x.centers(), x.scales(), x.gamma()); + let y = Self::augment_y(y, n + p); + + let mut optimizer = InteriorPointOptimizer::new_lazy(&x); + + let mut w = optimizer.optimize_lazy( + &x, + &y, + l1_reg * gamma, + parameters.max_iter, + TX::from_f64(parameters.tol).unwrap(), + true, + )?; + + for i in 0..p { + w.set(i, gamma * *w.get(i) / col_std[i]); + } + + let mut b = TX::zero(); + + for i in 0..p { + b += *w.get(i) * col_mean[i]; + } + + b = y_mean - b; + + (X::from_column(&w), b) + } else { + let x = LazyDesign::new(x, None, None).with_l2(l2_reg); + let (y, gamma) = (Self::augment_y(y, n + p), x.gamma()); + + let mut optimizer = InteriorPointOptimizer::new_lazy(&x); + + let mut w = optimizer.optimize_lazy( + &x, + &y, + l1_reg * gamma, + parameters.max_iter, + TX::from_f64(parameters.tol).unwrap(), + true, + )?; + + for i in 0..p { + w.set(i, gamma * *w.get(i)); + } + + (X::from_column(&w), y_mean) + }; + Ok((w, b)) } /// Predict target values from `x` @@ -422,7 +547,7 @@ impl, Y: Array1> self.intercept.as_ref().unwrap() } - fn rescale_x(x: &X) -> Result<(X, Vec, Vec), Failed> { + fn normalization_stats(x: &X) -> Result<(Vec, Vec), Failed> { let col_mean: Vec = x .mean_by(0) .iter() @@ -440,21 +565,31 @@ impl, Y: Array1> } } + Ok((col_mean, col_std)) + } + + fn rescale_x(x: &X) -> Result<(X, Vec, Vec), Failed> { + let (col_mean, col_std) = Self::normalization_stats(x)?; let mut scaled_x = x.clone(); scaled_x.scale_mut(&col_mean, &col_std, 0); Ok((scaled_x, col_mean, col_std)) } + fn augment_y(y: &Y, length: usize) -> Vec { + let mut result = Vec::::zeros(length); + for (dest, &value) in result.iter_mut().zip(y.iterator(0)) { + *dest = TX::from(value).unwrap(); + } + result + } + fn augment_x_and_y(x: &X, y: &Y, l2_reg: TX) -> (X, Vec, TX) { let (n, p) = x.shape(); let gamma = TX::one() / (TX::one() + l2_reg).sqrt(); let padding = gamma * l2_reg.sqrt(); - let mut y2 = Vec::::zeros(n + p); - for i in 0..y.shape() { - y2.set(i, TX::from(*y.get(i)).unwrap()); - } + let y2 = Self::augment_y(y, n + p); let mut x2 = X::zeros(n + p, p); @@ -476,9 +611,167 @@ mod tests { use crate::linalg::basic::matrix::DenseMatrix; use crate::metrics::mean_absolute_error; + #[cfg(feature = "lazy-normalization")] + fn compare_eager_fit>(tolerance: f64) { + for lazy_normalization in [false, true] { + use crate::linalg::basic::arrays::ArrayView1; + for (n, p) in [(8, 2), (3, 5)] { + let values: Vec = (0..n * p) + .map(|i| T::from_f64(((i * 7 + i / 3) % 13) as f64 / 3.0).unwrap()) + .collect(); + let x = X::from_iterator(values.iter().copied(), n, p, 0); + let y: Vec = (0..n) + .map(|i| { + T::from_f64(3.0).unwrap() + T::from_f64(2.0).unwrap() * values[i * p] + - values[i * p + 1] + }) + .collect(); + for normalize in [false, true] { + // The existing interior-point solver is unstable at a zero L1 penalty. + for ratio in [0.05, 0.5, 1.0] { + let parameters = ElasticNetParameters::default() + .with_alpha(0.04) + .with_lazy_normalization(lazy_normalization) + .with_l1_ratio(ratio) + .with_normalize(normalize); + let l1 = T::from_f64(0.04 * ratio * n as f64).unwrap(); + let l2 = T::from_f64(0.04 * (1.0 - ratio) * n as f64).unwrap(); + let (eager, centers, scales) = if normalize { + ElasticNet::>::rescale_x(&x).unwrap() + } else { + (x.clone(), vec![T::zero(); p], vec![T::one(); p]) + }; + let (eager, augmented_y, gamma) = + ElasticNet::>::augment_x_and_y(&eager, &y, l2); + let mut optimizer = InteriorPointOptimizer::new(&eager, p); + let expected = optimizer.optimize( + &eager, + &augmented_y, + l1 * gamma, + parameters.max_iter, + T::from_f64(parameters.tol).unwrap(), + true, + ); + let model = ElasticNet::fit(&x, &y, parameters); + let (mut expected, model) = match (expected, model) { + (Ok(expected), Ok(model)) => (expected, model), + (Err(expected), Err(actual)) => { + assert_eq!(actual.to_string(), expected.to_string()); + continue; + } + (expected, actual) => panic!( + "different fit outcomes: eager={expected:?}, lazy={actual:?}; n={n}, p={p}, ratio={ratio}, normalize={normalize}, scalar={}", + std::any::type_name::() + ), + }; + for (w, scale) in expected.iter_mut().zip(scales) { + *w = gamma * *w / scale; + } + for (&a, &b) in model.coefficients().iterator(0).zip(&expected) { + assert!( + (a - b).abs().to_f64().unwrap() < tolerance, + "{a} != {b}; n={n}, p={p}, ratio={ratio}" + ); + } + let intercept = T::from_f64(y.mean_by()).unwrap() + - expected.iter().zip(centers).map(|(&w, c)| w * c).sum::(); + assert!( + (*model.intercept() - intercept).abs().to_f64().unwrap() + < 10.0 * tolerance + ); + let prediction = model.predict(&x).unwrap(); + for (i, &actual) in prediction.iter().enumerate() { + let expected = + (0..p).map(|j| *x.get((i, j)) * expected[j]).sum::() + intercept; + assert!((actual - expected).abs().to_f64().unwrap() < 10.0 * tolerance); + } + assert!(x.iterator(0).zip(&values).all(|(a, b)| a == b)); + } + } + } + } + } + + #[cfg(feature = "lazy-normalization")] + #[cfg_attr( + all(target_arch = "wasm32", not(target_os = "wasi")), + wasm_bindgen_test::wasm_bindgen_test + )] + #[test] + fn lazy_fit_matches_eager_for_normalization_and_penalty_options() { + compare_eager_fit::>(1e-6); + compare_eager_fit::>(1e-2); + #[cfg(feature = "ndarray-bindings")] + { + compare_eager_fit::>(1e-6); + compare_eager_fit::>(1e-2); + } + } + + #[cfg_attr( + all(target_arch = "wasm32", not(target_os = "wasi")), + wasm_bindgen_test::wasm_bindgen_test + )] + #[test] + fn normalization_rejects_constant_and_near_constant_columns() { + for step in [0.0, f64::EPSILON / 16.0] { + let x = DenseMatrix::from_2d_array(&[&[0.0, 1.0], &[step, 2.0], &[2.0 * step, 4.0]]) + .unwrap(); + let error = ElasticNet::fit(&x, &vec![1., 2., 3.], ElasticNetParameters::default()) + .unwrap_err(); + assert!( + error + .to_string() + .contains("Cannot rescale constant column 0") + ); + #[cfg(feature = "lazy-normalization")] + { + let lazy_error = ElasticNet::fit( + &x, + &vec![1., 2., 3.], + ElasticNetParameters::default().with_lazy_normalization(true), + ) + .unwrap_err(); + assert_eq!(lazy_error, error); + } + } + } + + #[cfg(feature = "lazy-normalization")] + #[test] + fn lazy_selection_defaults_and_search() { + assert!(!ElasticNetParameters::default().lazy_normalization); + assert!(!ElasticNetSearchParameters::default().lazy_normalization); + for selected in [false, true] { + let search = ElasticNetSearchParameters { + lazy_normalization: selected, + alpha: vec![0.1, 0.2], + ..Default::default() + }; + let parameters: Vec<_> = search.into_iter().collect(); + assert_eq!(parameters.len(), 2); + assert!(parameters.iter().all(|p| p.lazy_normalization == selected)); + } + } + + #[cfg(all(feature = "lazy-normalization", feature = "serde"))] + #[test] + fn missing_lazy_selection_deserializes_as_eager() { + let parameters: ElasticNetParameters = serde_json::from_str("{}").unwrap(); + let search: ElasticNetSearchParameters = serde_json::from_str("{}").unwrap(); + assert!(!parameters.lazy_normalization); + assert!(!search.lazy_normalization); + let parameters = ElasticNetParameters::default().with_lazy_normalization(true); + let json = serde_json::to_string(¶meters).unwrap(); + let restored: ElasticNetParameters = serde_json::from_str(&json).unwrap(); + assert!(restored.lazy_normalization); + } + #[test] fn search_parameters() { let parameters = ElasticNetSearchParameters { + #[cfg(feature = "lazy-normalization")] + lazy_normalization: false, alpha: vec![0., 1.], max_iter: vec![10, 100], ..Default::default() @@ -534,6 +827,8 @@ mod tests { &x, &y, ElasticNetParameters { + #[cfg(feature = "lazy-normalization")] + lazy_normalization: false, alpha: 1.0, l1_ratio: 0.5, normalize: false, @@ -586,6 +881,8 @@ mod tests { &x, &y, ElasticNetParameters { + #[cfg(feature = "lazy-normalization")] + lazy_normalization: false, alpha: 1.0, l1_ratio: 1.0, normalize: true, @@ -599,6 +896,8 @@ mod tests { &x, &y, ElasticNetParameters { + #[cfg(feature = "lazy-normalization")] + lazy_normalization: false, alpha: 1.0, l1_ratio: 0.0, normalize: true, diff --git a/src/linear/lasso.rs b/src/linear/lasso.rs index 45493ea3..a57810eb 100644 --- a/src/linear/lasso.rs +++ b/src/linear/lasso.rs @@ -32,6 +32,8 @@ use crate::api::{Predictor, SupervisedEstimator}; use crate::error::Failed; use crate::linalg::basic::arrays::{Array1, Array2, ArrayView1}; use crate::linear::lasso_optimizer::InteriorPointOptimizer; +#[cfg(feature = "lazy-normalization")] +use crate::linear::lazy_normalization::LazyDesign; use crate::numbers::basenum::Number; use crate::numbers::floatnum::FloatNumber; use crate::numbers::realnum::RealNumber; @@ -41,6 +43,12 @@ use crate::numbers::realnum::RealNumber; #[derive(Debug, Clone)] #[must_use] pub struct LassoParameters { + #[cfg(feature = "lazy-normalization")] + #[cfg_attr(feature = "serde", serde(default))] + /// Borrow the design matrix instead of storing its normalized values. + /// This reduces memory use but can increase fit time. The default is false. + /// This option does not change whether normalization is applied. + pub lazy_normalization: bool, #[cfg_attr(feature = "serde", serde(default))] /// Controls the strength of the penalty to the loss function. pub alpha: f64, @@ -70,6 +78,24 @@ pub struct Lasso, Y: Arr } impl LassoParameters { + /// Select lazy matrix operations. Requires the `lazy-normalization` feature. + /// This reduces memory use but can increase fit time. + /// + /// ``` + /// # #[cfg(feature = "lazy-normalization")] + /// # { + /// use smartcore::linear::lasso::LassoParameters; + /// + /// let parameters = LassoParameters::default().with_lazy_normalization(true); + /// assert!(parameters.lazy_normalization); + /// # } + /// ``` + #[cfg(feature = "lazy-normalization")] + pub fn with_lazy_normalization(mut self, lazy_normalization: bool) -> Self { + self.lazy_normalization = lazy_normalization; + self + } + /// Regularization parameter. pub fn with_alpha(mut self, alpha: f64) -> Self { self.alpha = alpha; @@ -101,6 +127,8 @@ impl LassoParameters { impl Default for LassoParameters { fn default() -> Self { LassoParameters { + #[cfg(feature = "lazy-normalization")] + lazy_normalization: false, alpha: 1f64, normalize: true, tol: 1e-4, @@ -154,6 +182,13 @@ impl, Y: Array1> Pre #[derive(Debug, Clone)] #[must_use] pub struct LassoSearchParameters { + #[cfg(feature = "lazy-normalization")] + #[cfg_attr(feature = "serde", serde(default))] + /// Borrow the design matrix instead of storing its normalized values. + /// This reduces memory use but can increase fit time. The default is false. + /// This setting applies to all parameter combinations in the search. + /// This option does not change whether normalization is applied. + pub lazy_normalization: bool, #[cfg_attr(feature = "serde", serde(default))] /// Controls the strength of the penalty to the loss function. pub alpha: Vec, @@ -212,6 +247,8 @@ impl Iterator for LassoSearchParametersIterator { } let next = LassoParameters { + #[cfg(feature = "lazy-normalization")] + lazy_normalization: self.lasso_search_parameters.lazy_normalization, alpha: self.lasso_search_parameters.alpha[self.current_alpha], normalize: self.lasso_search_parameters.normalize[self.current_normalize], tol: self.lasso_search_parameters.tol[self.current_tol], @@ -257,6 +294,8 @@ impl Default for LassoSearchParameters { let default_params = LassoParameters::default(); LassoSearchParameters { + #[cfg(feature = "lazy-normalization")] + lazy_normalization: false, alpha: vec![default_params.alpha], normalize: vec![default_params.normalize], tol: vec![default_params.tol], @@ -300,6 +339,30 @@ impl, Y: Array1> Las let l1_reg = TX::from_f64(parameters.alpha * n as f64).unwrap(); + #[cfg(feature = "lazy-normalization")] + let (w, b) = if parameters.lazy_normalization { + Self::fit_lazy_coefficients(x, &y, ¶meters, l1_reg)? + } else { + Self::fit_eager_coefficients(x, &y, ¶meters, l1_reg)? + }; + #[cfg(not(feature = "lazy-normalization"))] + let (w, b) = Self::fit_eager_coefficients(x, &y, ¶meters, l1_reg)?; + + Ok(Lasso { + intercept: b, + coefficients: Some(w), + _phantom_ty: PhantomData, + _phantom_y: PhantomData, + }) + } + + fn fit_eager_coefficients( + x: &X, + y: &Vec, + parameters: &LassoParameters, + l1_reg: TX, + ) -> Result<(X, Option), Failed> { + let (_, p) = x.shape(); let (w, b) = if parameters.normalize { let (scaled_x, col_mean, col_std) = Self::rescale_x(x)?; @@ -307,7 +370,7 @@ impl, Y: Array1> Las let mut w = optimizer.optimize( &scaled_x, - &y, + y, l1_reg, parameters.max_iter, TX::from_f64(parameters.tol).unwrap(), @@ -334,7 +397,7 @@ impl, Y: Array1> Las let w = optimizer.optimize( x, - &y, + y, l1_reg, parameters.max_iter, TX::from_f64(parameters.tol).unwrap(), @@ -350,13 +413,72 @@ impl, Y: Array1> Las }, ) }; + Ok((w, b)) + } - Ok(Lasso { - intercept: b, - coefficients: Some(w), - _phantom_ty: PhantomData, - _phantom_y: PhantomData, - }) + #[cfg(feature = "lazy-normalization")] + fn fit_lazy_coefficients( + x: &X, + y: &Vec, + parameters: &LassoParameters, + l1_reg: TX, + ) -> Result<(X, Option), Failed> { + let (_, p) = x.shape(); + let (w, b) = if parameters.normalize { + let scaled_x = { + let (col_mean, col_std) = Self::normalization_stats(x)?; + LazyDesign::new(x, Some(col_mean), Some(col_std)) + }; + let (col_mean, col_std) = (scaled_x.centers(), scaled_x.scales()); + + let mut optimizer = InteriorPointOptimizer::new_lazy(&scaled_x); + + let mut w = optimizer.optimize_lazy( + &scaled_x, + y, + l1_reg, + parameters.max_iter, + TX::from_f64(parameters.tol).unwrap(), + parameters.fit_intercept, + )?; + + for (j, col_std_j) in col_std.iter().enumerate().take(p) { + w[j] /= *col_std_j; + } + + let b = if parameters.fit_intercept { + let mut xw_mean = TX::zero(); + for (i, col_mean_i) in col_mean.iter().enumerate().take(p) { + xw_mean += w[i] * *col_mean_i; + } + + Some(TX::from_f64(y.mean_by()).unwrap() - xw_mean) + } else { + None + }; + (X::from_column(&w), b) + } else { + let mut optimizer = InteriorPointOptimizer::new(x, p); + + let w = optimizer.optimize( + x, + y, + l1_reg, + parameters.max_iter, + TX::from_f64(parameters.tol).unwrap(), + parameters.fit_intercept, + )?; + + ( + X::from_column(&w), + if parameters.fit_intercept { + Some(TX::from_f64(y.mean_by()).unwrap()) + } else { + None + }, + ) + }; + Ok((w, b)) } /// Predict target values from `x` @@ -389,7 +511,7 @@ impl, Y: Array1> Las self.intercept.as_ref().unwrap() } - fn rescale_x(x: &X) -> Result<(X, Vec, Vec), Failed> { + fn normalization_stats(x: &X) -> Result<(Vec, Vec), Failed> { let col_mean: Vec = x .mean_by(0) .iter() @@ -407,6 +529,11 @@ impl, Y: Array1> Las } } + Ok((col_mean, col_std)) + } + + fn rescale_x(x: &X) -> Result<(X, Vec, Vec), Failed> { + let (col_mean, col_std) = Self::normalization_stats(x)?; let mut scaled_x = x.clone(); scaled_x.scale_mut(&col_mean, &col_std, 0); Ok((scaled_x, col_mean, col_std)) @@ -420,9 +547,150 @@ mod tests { use crate::linalg::basic::matrix::DenseMatrix; use crate::metrics::mean_absolute_error; + #[cfg(feature = "lazy-normalization")] + fn compare_eager_fit>(tolerance: f64) { + for lazy_normalization in [false, true] { + let values = [ + 1., 2., 2., 1., 3., 5., 4., 3., 5., 8., 6., 4., 7., 6., 8., 7., + ]; + let x = X::from_iterator(values.into_iter().map(|v| T::from_f64(v).unwrap()), 8, 2, 0); + let y: Vec = (0..8) + .map(|i| T::from_f64(3.0 + 2.0 * values[2 * i] - 0.4 * values[2 * i + 1]).unwrap()) + .collect(); + for normalize in [false, true] { + for fit_intercept in [false, true] { + let parameters = LassoParameters::default() + .with_alpha(0.03) + .with_lazy_normalization(lazy_normalization) + .with_normalize(normalize) + .with_fit_intercept(fit_intercept); + let (eager, centers, scales) = if normalize { + Lasso::>::rescale_x(&x).unwrap() + } else { + (x.clone(), vec![T::zero(); 2], vec![T::one(); 2]) + }; + let mut optimizer = InteriorPointOptimizer::new(&eager, 2); + let mut expected = optimizer + .optimize( + &eager, + &y, + T::from_f64(0.03 * 8.0).unwrap(), + parameters.max_iter, + T::from_f64(parameters.tol).unwrap(), + fit_intercept, + ) + .unwrap(); + for (w, scale) in expected.iter_mut().zip(scales) { + *w /= scale; + } + let model = Lasso::fit(&x, &y, parameters).unwrap(); + for (&a, &b) in model.coefficients().iterator(0).zip(&expected) { + assert!((a - b).abs().to_f64().unwrap() < tolerance, "{a} != {b}"); + } + if fit_intercept { + let intercept = T::from_f64(y.mean_by()).unwrap() + - expected.iter().zip(centers).map(|(&w, c)| w * c).sum::(); + assert!( + (*model.intercept() - intercept).abs().to_f64().unwrap() < tolerance + ); + let prediction = model.predict(&x).unwrap(); + for (i, &actual) in prediction.iter().enumerate() { + let expected = + (0..2).map(|j| *x.get((i, j)) * expected[j]).sum::() + intercept; + assert!((actual - expected).abs().to_f64().unwrap() < 10.0 * tolerance); + } + } else { + assert!(model.intercept.is_none()); + } + assert!( + x.iterator(0) + .zip(values) + .all(|(&actual, original)| actual == T::from_f64(original).unwrap()) + ); + } + } + } + } + + #[cfg(feature = "lazy-normalization")] + #[cfg_attr( + all(target_arch = "wasm32", not(target_os = "wasi")), + wasm_bindgen_test::wasm_bindgen_test + )] + #[test] + fn lazy_fit_matches_eager_for_normalization_and_intercept_options() { + compare_eager_fit::>(1e-8); + compare_eager_fit::>(5e-3); + #[cfg(feature = "ndarray-bindings")] + { + compare_eager_fit::>(1e-8); + compare_eager_fit::>(5e-3); + } + } + + #[cfg_attr( + all(target_arch = "wasm32", not(target_os = "wasi")), + wasm_bindgen_test::wasm_bindgen_test + )] + #[test] + fn normalization_rejects_constant_and_near_constant_columns() { + for step in [0.0, f64::EPSILON / 16.0] { + let x = DenseMatrix::from_2d_array(&[&[0.0, 1.0], &[step, 2.0], &[2.0 * step, 4.0]]) + .unwrap(); + let error = Lasso::fit(&x, &vec![1., 2., 3.], LassoParameters::default()).unwrap_err(); + assert!( + error + .to_string() + .contains("Cannot rescale constant column 0") + ); + #[cfg(feature = "lazy-normalization")] + { + let lazy_error = Lasso::fit( + &x, + &vec![1., 2., 3.], + LassoParameters::default().with_lazy_normalization(true), + ) + .unwrap_err(); + assert_eq!(lazy_error, error); + } + } + } + + #[cfg(feature = "lazy-normalization")] + #[test] + fn lazy_selection_defaults_and_search() { + assert!(!LassoParameters::default().lazy_normalization); + assert!(!LassoSearchParameters::default().lazy_normalization); + for selected in [false, true] { + let search = LassoSearchParameters { + lazy_normalization: selected, + alpha: vec![0.1, 0.2], + ..Default::default() + }; + let parameters: Vec<_> = search.into_iter().collect(); + assert_eq!(parameters.len(), 2); + assert!(parameters.iter().all(|p| p.lazy_normalization == selected)); + } + } + + #[cfg(all(feature = "lazy-normalization", feature = "serde"))] + #[test] + fn missing_lazy_selection_deserializes_as_eager() { + let parameters: LassoParameters = serde_json::from_str("{}").unwrap(); + let search: LassoSearchParameters = serde_json::from_str("{}").unwrap(); + assert!(!parameters.lazy_normalization); + assert!(!search.lazy_normalization); + let parameters = LassoParameters::default().with_lazy_normalization(true); + let json = serde_json::to_string(¶meters).unwrap(); + let restored: LassoParameters = serde_json::from_str(&json).unwrap(); + assert!(restored.lazy_normalization); + } + #[test] fn search_parameters() { let parameters = LassoSearchParameters { + #[cfg(feature = "lazy-normalization")] + lazy_normalization: false, alpha: vec![0., 1.], max_iter: vec![10, 100], fit_intercept: vec![false, true], @@ -493,6 +761,8 @@ mod tests { &x, &y, LassoParameters { + #[cfg(feature = "lazy-normalization")] + lazy_normalization: false, alpha: 0.1, normalize: false, tol: 1e-4, @@ -547,6 +817,8 @@ mod tests { &x, &y, LassoParameters { + #[cfg(feature = "lazy-normalization")] + lazy_normalization: false, alpha: 0.1, normalize: false, tol: 1e-8, diff --git a/src/linear/lasso_optimizer.rs b/src/linear/lasso_optimizer.rs index 5969347a..f3675d08 100644 --- a/src/linear/lasso_optimizer.rs +++ b/src/linear/lasso_optimizer.rs @@ -16,6 +16,9 @@ use crate::linalg::basic::arrays::{Array1, Array2, ArrayView1, MutArray, MutArra use crate::linear::bg_solver::BiconjugateGradientSolver; use crate::numbers::floatnum::FloatNumber; +#[cfg(feature = "lazy-normalization")] +use crate::linear::lazy_normalization::LazyDesign; + /// Interior Point Optimizer pub struct InteriorPointOptimizer> { ata: X, @@ -28,8 +31,12 @@ pub struct InteriorPointOptimizer> { impl> InteriorPointOptimizer { /// Initialize a new Interior Point Optimizer pub fn new(a: &X, n: usize) -> InteriorPointOptimizer { + Self::from_gram(a.ab(true, a, false), n) + } + + fn from_gram(ata: X, n: usize) -> Self { InteriorPointOptimizer { - ata: a.ab(true, a, false), + ata, d1: vec![T::zero(); n], d2: vec![T::zero(); n], prb: vec![T::zero(); n], @@ -37,6 +44,33 @@ impl> InteriorPointOptimizer { } } + #[cfg(feature = "lazy-normalization")] + pub(super) fn new_lazy(design: &LazyDesign<'_, T, X>) -> Self { + Self::from_gram(design.gram(), design.shape().1) + } + + #[cfg(feature = "lazy-normalization")] + pub(super) fn optimize_lazy( + &mut self, + design: &LazyDesign<'_, T, X>, + y: &Vec, + lambda: T, + max_iter: usize, + tol: T, + fit_intercept: bool, + ) -> Result, Failed> { + self.optimize_with_products( + design.shape(), + y, + lambda, + max_iter, + tol, + fit_intercept, + |v| design.matvec(v), + |v| design.transpose_matvec(v), + ) + } + /// Run the optimization pub fn optimize( &mut self, @@ -47,7 +81,34 @@ impl> InteriorPointOptimizer { tol: T, fit_intercept: bool, ) -> Result, Failed> { - let (n, p) = x.shape(); + self.optimize_with_products( + x.shape(), + y, + lambda, + max_iter, + tol, + fit_intercept, + |v| v.xa(true, x), + |v| v.xa(false, x), + ) + } + + fn optimize_with_products( + &mut self, + shape: (usize, usize), + y: &Vec, + lambda: T, + max_iter: usize, + tol: T, + fit_intercept: bool, + matvec: F, + transpose_matvec: G, + ) -> Result, Failed> + where + F: Fn(&Vec) -> Vec, + G: Fn(&Vec) -> Vec, + { + let (n, p) = shape; let p_f64 = T::from_usize(p).unwrap(); let lambda = lambda.max(T::epsilon()); @@ -96,7 +157,7 @@ impl> InteriorPointOptimizer { let lambda_f64 = lambda.to_f64().unwrap(); for ntiter in 0..max_iter { - let mut z = w.xa(true, x); + let mut z = matvec(&w); for i in 0..n { z[i] -= y[i]; @@ -104,7 +165,7 @@ impl> InteriorPointOptimizer { } // CALCULATE DUALITY GAP - let xnu = nu.xa(false, x); + let xnu = transpose_matvec(&nu); let max_xnu = xnu.norm(f64::INFINITY); if max_xnu > lambda_f64 { let lnu = T::from_f64(lambda_f64 / max_xnu).unwrap(); @@ -136,7 +197,7 @@ impl> InteriorPointOptimizer { self.d2[i] = (q1i * q1i - q2i * q2i) / t; } - let mut gradphi = z.xa(false, x); + let mut gradphi = transpose_matvec(&z); for i in 0..p { let g1 = T::two() * gradphi[i] - (q1[i] - q2[i]) / t; @@ -157,7 +218,7 @@ impl> InteriorPointOptimizer { pcgtol *= min_pcgtol; } - let error = self.solve_mut(x, &grad, &mut dxu, pcgtol, pcgmaxi)?; + let error = self.solve_mut(&self.ata, &grad, &mut dxu, pcgtol, pcgmaxi)?; if error > pcgtol { pitr = pcgmaxi; } @@ -184,7 +245,7 @@ impl> InteriorPointOptimizer { .fold(T::neg_infinity(), |max, v| v.max(max)) < T::zero() { - let mut newz = neww.xa(true, x); + let mut newz = matvec(&neww); for i in 0..n { newz[i] -= y[i]; } diff --git a/src/linear/lazy_normalization.rs b/src/linear/lazy_normalization.rs new file mode 100644 index 00000000..783f7446 --- /dev/null +++ b/src/linear/lazy_normalization.rs @@ -0,0 +1,830 @@ +//! Borrowed normalized designs for the interior-point method of Kim et al. +//! +//! See . + +use std::convert::Infallible; +use std::marker::PhantomData; + +use lazymatrix::{ + DotSlice, ElemDivAssign, LazyMatrix, MatTransposeVecInto, MatVecInto, MatrixErrorType, + MatrixShape, MatrixWrite, Scalar, ScaledSubSlice, SubScalarAssign, SumEntries, VectorView, + WeightedGramInto, WeightedGramKernel, +}; + +use crate::linalg::basic::arrays::{Array1, Array2}; +use crate::numbers::floatnum::FloatNumber; + +/// Keep the design borrowed while coefficient matrices retain the caller's storage type. +pub(super) struct LazyDesign<'a, T: FloatNumber, X: Array2> { + normalized: LazyMatrix, T>, + augmentation: Option<(T, T)>, +} + +impl<'a, T: FloatNumber, X: Array2> LazyDesign<'a, T, X> { + pub(super) fn new(x: &'a X, centers: Option>, scales: Option>) -> Self { + Self { + normalized: LazyMatrix::from_parts(MatrixAdapter(x, PhantomData), centers, scales), + augmentation: None, + } + } + + pub(super) fn with_l2(mut self, penalty: T) -> Self { + let gamma = T::one() / (T::one() + penalty).sqrt(); + self.augmentation = Some((gamma, gamma * penalty.sqrt())); + self + } + + pub(super) fn shape(&self) -> (usize, usize) { + let n = self.normalized.nrows(); + let p = self.normalized.ncols(); + (n + if self.augmentation.is_some() { p } else { 0 }, p) + } + + pub(super) fn centers(&self) -> &[T] { + self.normalized.centers().unwrap() + } + + pub(super) fn scales(&self) -> &[T] { + self.normalized.scales().unwrap() + } + + pub(super) fn gamma(&self) -> T { + self.augmentation.map_or(T::one(), |(gamma, _)| gamma) + } + + pub(super) fn matvec(&self, x: &Vec) -> Vec { + let mut out = Vec::with_capacity(self.shape().0); + out.resize(self.normalized.nrows(), T::zero()); + self.normalized.matvec_into(x, &mut out).unwrap(); + if let Some((gamma, padding)) = self.augmentation { + for value in &mut out { + *value *= gamma; + } + out.extend(x.iter().map(|&value| padding * value)); + } + out + } + + pub(super) fn transpose_matvec(&self, x: &[T]) -> Vec { + check_vector_length(x.len(), self.shape().0); + let n = self.normalized.nrows(); + let mut out = Vec::zeros(self.normalized.ncols()); + self.normalized + .mat_transpose_vec_into(&VectorRef(&x[..n]), &mut out) + .unwrap(); + if let Some((gamma, padding)) = self.augmentation { + for (value, &tail) in out.iter_mut().zip(&x[n..]) { + *value = gamma * *value + padding * tail; + } + } + out + } + + pub(super) fn gram(&self) -> X { + let p = self.normalized.ncols(); + let mut out = X::zeros(p, p); + self.normalized + .weighted_gram_into( + &UnitWeights(self.normalized.nrows()), + &mut MatrixOutput(&mut out, PhantomData), + ) + .unwrap(); + if let Some((gamma, padding)) = self.augmentation { + out.mul_scalar_mut(gamma * gamma); + for j in 0..p { + out.add_element_mut((j, j), padding * padding); + } + } + out + } +} + +struct MatrixAdapter<'a, T: FloatNumber, X: Array2>(&'a X, PhantomData); + +impl> MatrixShape for MatrixAdapter<'_, T, X> { + fn nrows(&self) -> usize { + self.0.shape().0 + } + fn ncols(&self) -> usize { + self.0.shape().1 + } +} + +impl> MatrixErrorType for MatrixAdapter<'_, T, X> { + type Error = Infallible; +} + +impl> MatVecInto> for MatrixAdapter<'_, T, X> { + fn matvec_into(&self, x: &Vec, out: &mut Vec) -> Result<(), Infallible> { + self.matvec_normalized_into::(x, None, None, out) + } + + fn matvec_normalized_into( + &self, + x: &Vec, + centers: Option<&[F]>, + scales: Option<&[F]>, + out: &mut Vec, + ) -> Result<(), Infallible> + where + Vec: ElemDivAssign + DotSlice + SubScalarAssign, + { + check_vector_length(x.len(), self.ncols()); + check_vector_length(out.len(), self.nrows()); + out.fill(T::zero()); + if self.nrows() == 0 { + return Ok(()); + } + // Column traversal reuses each view while preserving each row's summation order. + for (j, &v) in x.iter().enumerate() { + let center = centers.map_or(T::zero(), |c| T::from(c[j]).unwrap()); + let scale = scales.map_or(T::one(), |s| T::from(s[j]).unwrap()); + // Keep division when a reciprocal would overflow or lose normal precision. + let reciprocal = T::one() / scale; + let use_reciprocal = scale.is_normal() && reciprocal.is_normal(); + let column = self.0.get_col(j); + for (&raw, result) in column.iterator(0).zip(out.iter_mut()) { + // Center before accumulation to retain small variations at large offsets. + let normalized = if use_reciprocal { + (raw - center) * reciprocal + } else { + (raw - center) / scale + }; + *result += normalized * v; + } + } + Ok(()) + } +} + +struct VectorRef<'a, T>(&'a [T]); + +impl SumEntries for VectorRef<'_, T> { + fn sum_entries(&self) -> T { + self.0.iter().copied().sum() + } +} + +impl<'v, T: FloatNumber, X: Array2> MatTransposeVecInto, Vec> + for MatrixAdapter<'_, T, X> +{ + fn mat_transpose_vec_into( + &self, + x: &VectorRef<'v, T>, + out: &mut Vec, + ) -> Result<(), Infallible> { + self.mat_transpose_vec_normalized_into::(x, None, None, out) + } + + fn mat_transpose_vec_normalized_into( + &self, + x: &VectorRef<'v, T>, + centers: Option<&[F]>, + scales: Option<&[F]>, + out: &mut Vec, + ) -> Result<(), Infallible> + where + VectorRef<'v, T>: SumEntries, + Vec: ScaledSubSlice + ElemDivAssign, + { + check_vector_length(x.0.len(), self.nrows()); + check_vector_length(out.len(), self.ncols()); + if self.nrows() == 0 { + out.fill(T::zero()); + return Ok(()); + } + for (j, result) in out.iter_mut().enumerate() { + let center = centers.map_or(T::zero(), |c| T::from(c[j]).unwrap()); + let scale = scales.map_or(T::one(), |s| T::from(s[j]).unwrap()); + // Keep division when a reciprocal would overflow or lose normal precision. + let reciprocal = T::one() / scale; + let use_reciprocal = scale.is_normal() && reciprocal.is_normal(); + let column = self.0.get_col(j); + *result = column + .iterator(0) + .zip(x.0) + .fold(T::zero(), |sum, (&raw, &v)| { + let normalized = if use_reciprocal { + (raw - center) * reciprocal + } else { + (raw - center) / scale + }; + sum + normalized * v + }); + } + Ok(()) + } +} + +impl> WeightedGramKernel for MatrixAdapter<'_, T, X> { + fn weighted_gram_normalized_into( + &self, + weights: &W, + centers: Option<&[T]>, + scales: Option<&[T]>, + out: &mut O, + ) -> Result<(), Infallible> + where + W: VectorView + ?Sized, + O: MatrixWrite + ?Sized, + { + // Panels reuse normalized values across column pairs without storing the design. + const ROWS: usize = 256; + const COLS: usize = 32; + let (n, p) = self.0.shape(); + let rows = n.min(ROWS); + let cols = p.min(COLS); + let mut left = Vec::zeros(rows * cols); + let mut right = Vec::zeros(rows * cols); + let mut sums = Vec::zeros(cols * cols); + for j in (0..p).step_by(COLS) { + let nj = (p - j).min(COLS); + for k in (0..=j).step_by(COLS) { + let nk = (p - k).min(COLS); + sums.fill(T::zero()); + for row in (0..n).step_by(ROWS) { + let nr = (n - row).min(ROWS); + self.normalized_panel(row, nr, j, nj, centers, scales, &mut left); + if j == k { + right[..nr * nk].copy_from_slice(&left[..nr * nj]); + } else { + self.normalized_panel(row, nr, k, nk, centers, scales, &mut right); + } + for column in left[..nr * nj].chunks_exact_mut(nr) { + for (i, value) in column.iter_mut().enumerate() { + *value = weights.get(row + i) * *value; + } + } + for a in 0..nj { + let left_column = &left[a * nr..(a + 1) * nr]; + let width = if j == k { a + 1 } else { nk }; + for b in 0..width { + let right_column = &right[b * nr..(b + 1) * nr]; + let sum = &mut sums[a * cols + b]; + // Carry the accumulator across panels to retain row order. + for (&x, &y) in left_column.iter().zip(right_column) { + *sum += x * y; + } + } + } + } + for a in 0..nj { + let width = if j == k { a + 1 } else { nk }; + for b in 0..width { + let value = sums[a * cols + b]; + out.set(j + a, k + b, value); + if j + a != k + b { + out.set(k + b, j + a, value); + } + } + } + } + } + Ok(()) + } +} + +impl> MatrixAdapter<'_, T, X> { + fn normalized_panel( + &self, + row: usize, + nrows: usize, + column: usize, + ncols: usize, + centers: Option<&[T]>, + scales: Option<&[T]>, + out: &mut [T], + ) { + for j in 0..ncols { + let center = centers.map_or(T::zero(), |c| c[column + j]); + let scale = scales.map_or(T::one(), |s| s[column + j]); + let view = self.0.slice(row..row + nrows, column + j..column + j + 1); + for (value, &raw) in out[j * nrows..(j + 1) * nrows] + .iter_mut() + .zip(view.iterator(0)) + { + *value = (raw - center) / scale; + } + } + } +} + +struct UnitWeights(usize); + +impl VectorView for UnitWeights { + fn len(&self) -> usize { + self.0 + } + fn get(&self, index: usize) -> T { + if index >= self.0 { + std::panic::panic_any("weight index is out of bounds"); + } + T::one() + } +} + +fn check_vector_length(actual: usize, expected: usize) { + if actual != expected { + std::panic::panic_any("matrix and vector dimensions do not agree"); + } +} + +struct MatrixOutput<'a, T: FloatNumber, X: Array2>(&'a mut X, PhantomData); + +impl> MatrixShape for MatrixOutput<'_, T, X> { + fn nrows(&self) -> usize { + self.0.shape().0 + } + fn ncols(&self) -> usize { + self.0.shape().1 + } +} + +impl> MatrixWrite for MatrixOutput<'_, T, X> { + fn set(&mut self, row: usize, column: usize, value: T) { + self.0.set((row, column), value); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::linalg::basic::arrays::{Array, Array1, Array2, MutArray, MutArrayView2}; + use crate::linalg::basic::matrix::DenseMatrix; + use crate::linear::lasso_optimizer::InteriorPointOptimizer; + use crate::numbers::floatnum::FloatNumber; + + fn check_products>(x: X, tolerance: f64) { + let value = |v| T::from_f64(v).unwrap(); + let original = x.clone(); + for centered in [false, true] { + for scaled in [false, true] { + let centers = centered.then(|| vec![value(1.0), value(2.0)]); + let scales = scaled.then(|| vec![value(2.0), value(4.0)]); + let mut eager = x.clone(); + eager.scale_mut( + centers.as_deref().unwrap_or(&[T::zero(); 2]), + scales.as_deref().unwrap_or(&[T::one(); 2]), + 0, + ); + let design = LazyDesign::new(&x, centers, scales); + let coefficients = vec![value(2.0), value(-1.0)]; + let response = vec![value(1.0), value(-2.0), value(3.0)]; + assert_close( + &design.matvec(&coefficients), + &coefficients.xa(true, &eager), + tolerance, + ); + assert_close( + &design.transpose_matvec(&response), + &response.xa(false, &eager), + tolerance, + ); + let gram = design.gram(); + assert_close( + &gram.iterator(0).copied().collect::>(), + &eager + .ab(true, &eager, false) + .iterator(0) + .copied() + .collect::>(), + tolerance, + ); + let left = design + .matvec(&coefficients) + .iter() + .zip(&response) + .map(|(&a, &b)| a * b) + .sum::(); + let right = coefficients + .iter() + .zip(design.transpose_matvec(&response)) + .map(|(&a, b)| a * b) + .sum::(); + assert!((left - right).abs().to_f64().unwrap() <= tolerance); + } + } + assert!(x.iterator(0).zip(original.iterator(0)).all(|(a, b)| a == b)); + } + + fn assert_close(actual: &[T], expected: &[T], tolerance: f64) { + assert_eq!(actual.len(), expected.len()); + for (&a, &b) in actual.iter().zip(expected) { + assert!((a - b).abs().to_f64().unwrap() <= tolerance, "{a} != {b}"); + } + } + + fn assert_relative_close(actual: &[T], expected: &[T], tolerance: f64) { + assert_eq!(actual.len(), expected.len()); + for (&a, &b) in actual.iter().zip(expected) { + let scale = a.abs().max(b.abs()).max(T::one()).to_f64().unwrap(); + assert!( + (a - b).abs().to_f64().unwrap() <= tolerance * scale, + "{a} != {b}" + ); + } + } + + #[cfg_attr( + all(target_arch = "wasm32", not(target_os = "wasi")), + wasm_bindgen_test::wasm_bindgen_test + )] + #[test] + fn adapter_products_overwrite_output_without_normalization() { + for (n, p) in [(3, 2), (0, 2), (3, 0)] { + let x = DenseMatrix::from_iterator((0..n * p).map(|i| i as f64 + 1.0), n, p, 0); + let adapter = MatrixAdapter(&x, PhantomData); + let coefficients = vec![2.0; p]; + let response = vec![-1.0; n]; + let mut product = vec![f64::NAN; n]; + adapter.matvec_into(&coefficients, &mut product).unwrap(); + assert_eq!(product, coefficients.xa(true, &x)); + + let mut transpose_product = vec![f64::NAN; p]; + adapter + .mat_transpose_vec_into(&VectorRef(&response), &mut transpose_product) + .unwrap(); + assert_eq!(transpose_product, response.xa(false, &x)); + assert_eq!(VectorRef(&response).sum_entries(), -(n as f64)); + + let design = LazyDesign::new(&x, None, None); + let mut gram = DenseMatrix::fill(p, p, f64::NAN); + design + .normalized + .weighted_gram_into(&UnitWeights(n), &mut MatrixOutput(&mut gram, PhantomData)) + .unwrap(); + assert_eq!(gram, x.ab(true, &x, false)); + } + } + + #[cfg(not(target_arch = "wasm32"))] + #[test] + fn adapter_products_reject_incorrect_vector_lengths() { + use std::panic::{AssertUnwindSafe, catch_unwind}; + + let x = DenseMatrix::from_2d_array(&[&[1., 2.], &[3., 4.], &[5., 6.]]).unwrap(); + let adapter = MatrixAdapter(&x, PhantomData); + for (input, output) in [(1, 3), (3, 3), (2, 2), (2, 4)] { + let mut out = vec![0.; output]; + let error = catch_unwind(AssertUnwindSafe(|| { + adapter.matvec_into(&vec![1.; input], &mut out).unwrap(); + })) + .unwrap_err(); + assert_eq!( + error.downcast_ref::<&str>(), + Some(&"matrix and vector dimensions do not agree") + ); + } + for (input, output) in [(2, 2), (4, 2), (3, 1), (3, 3)] { + let mut out = vec![0.; output]; + let error = catch_unwind(AssertUnwindSafe(|| { + adapter + .mat_transpose_vec_into(&VectorRef(&vec![1.; input]), &mut out) + .unwrap(); + })) + .unwrap_err(); + assert_eq!( + error.downcast_ref::<&str>(), + Some(&"matrix and vector dimensions do not agree") + ); + } + for augmented in [false, true] { + let design = LazyDesign::new(&x, None, None); + let design = if augmented { + design.with_l2(0.7) + } else { + design + }; + for length in [design.shape().0 - 1, design.shape().0 + 1] { + let error = catch_unwind(AssertUnwindSafe(|| { + design.transpose_matvec(&vec![1.; length]); + })) + .unwrap_err(); + assert_eq!( + error.downcast_ref::<&str>(), + Some(&"matrix and vector dimensions do not agree") + ); + } + } + } + + #[cfg(not(target_arch = "wasm32"))] + #[test] + #[should_panic(expected = "weight index is out of bounds")] + fn unit_weights_reject_out_of_bounds_index() { + let _: f64 = UnitWeights(3).get(3); + } + + fn check_extreme_product_scales() { + for scale in [ + T::min_positive_value() * T::epsilon(), + T::min_positive_value(), + ::max_value(), + -::max_value(), + T::infinity(), + T::neg_infinity(), + T::nan(), + ] { + let raw = if scale.is_finite() { scale } else { T::one() }; + let x = DenseMatrix::new(2, 1, vec![raw, -raw], true).unwrap(); + let design = LazyDesign::new(&x, None, Some(vec![scale])); + let expected = vec![raw / scale, -raw / scale]; + let actual = design.matvec(&vec![T::one()]); + for (&a, &b) in actual.iter().zip(&expected) { + if b.is_nan() { + assert!(a.is_nan()); + } else { + assert_eq!(a, b); + } + } + let actual = design.transpose_matvec(&[T::one(), -T::one()]); + let expected = expected[0] - expected[1]; + if expected.is_nan() { + assert!(actual[0].is_nan()); + } else { + assert_eq!(actual[0], expected); + } + } + } + + #[test] + fn products_preserve_extreme_scales() { + check_extreme_product_scales::(); + check_extreme_product_scales::(); + } + + #[cfg_attr( + all(target_arch = "wasm32", not(target_os = "wasi")), + wasm_bindgen_test::wasm_bindgen_test + )] + #[test] + fn products_match_eager_for_both_dense_layouts_and_precisions() { + for column_major in [false, true] { + let data = if column_major { + vec![1., 2., 4., 4., 6., 10.] + } else { + vec![1., 4., 2., 6., 4., 10.] + }; + check_products( + DenseMatrix::new(3, 2, data.clone(), column_major).unwrap(), + 1e-12, + ); + check_products( + DenseMatrix::new(3, 2, data.iter().map(|&v| v as f32).collect(), column_major) + .unwrap(), + 1e-5, + ); + } + } + + #[cfg(feature = "ndarray-bindings")] + #[cfg_attr( + all(target_arch = "wasm32", not(target_os = "wasi")), + wasm_bindgen_test::wasm_bindgen_test + )] + #[test] + fn products_match_eager_for_ndarray() { + use ndarray::ShapeBuilder; + for column_major in [false, true] { + let shape = (3, 2).set_f(column_major); + let data = if column_major { + vec![1., 2., 4., 4., 6., 10.] + } else { + vec![1., 4., 2., 6., 4., 10.] + }; + check_products( + ndarray::Array2::from_shape_vec(shape, data.clone()).unwrap(), + 1e-12, + ); + check_products( + ndarray::Array2::from_shape_vec(shape, data.iter().map(|&v| v as f32).collect()) + .unwrap(), + 1e-5, + ); + } + } + + #[cfg_attr( + all(target_arch = "wasm32", not(target_os = "wasi")), + wasm_bindgen_test::wasm_bindgen_test + )] + #[test] + fn centered_products_preserve_small_variation_at_large_offsets() { + let x = DenseMatrix::from_2d_array(&[ + &[1e12 + 1., 1e12 + 4.], + &[1e12 + 2., 1e12 + 6.], + &[1e12 + 4., 1e12 + 10.], + ]) + .unwrap(); + let design = LazyDesign::new(&x, Some(vec![1e12 + 1., 1e12 + 2.]), Some(vec![2., 4.])); + assert_eq!(design.matvec(&vec![2., -1.]), vec![-0.5, 0., 1.]); + assert_eq!(design.transpose_matvec(&vec![1., -2., 3.]), vec![3.5, 4.5]); + let nonbinary = LazyDesign::new(&x, Some(vec![1e12 + 1., 1e12 + 2.]), Some(vec![3., 7.])); + let mut eager = x.clone(); + eager.scale_mut(&[1e12 + 1., 1e12 + 2.], &[3., 7.], 0); + assert_relative_close( + &nonbinary.matvec(&vec![2., -1.]), + &vec![2., -1.].xa(true, &eager), + 8.0 * f64::EPSILON, + ); + assert_relative_close( + &nonbinary.transpose_matvec(&[1., -2., 3.]), + &vec![1., -2., 3.].xa(false, &eager), + 8.0 * f64::EPSILON, + ); + let gram = design.gram(); + assert_eq!(*gram.get((0, 0)), 2.5); + assert_eq!(*gram.get((0, 1)), 3.5); + assert_eq!(*gram.get((1, 1)), 5.25); + } + + fn check_panel_gram>() { + for (n, p) in [ + (0, 3), + (3, 0), + (1, 1), + (255, 31), + (256, 32), + (257, 33), + (513, 65), + ] { + for axis in [0, 1] { + let x = X::from_iterator( + (0..n * p).map(|index| { + let (i, j) = if axis == 0 { + (index / p, index % p) + } else { + (index % n, index / n) + }; + T::from_f64(100_000.0 + 2.0 * j as f64 + ((i * 7 + j) % 23) as f64 / 8.0) + .unwrap() + }), + n, + p, + axis, + ); + let weights: Vec = (0..n) + .map(|i| T::from_f64([1.0, 0.0, -0.5, 2.0][i % 4]).unwrap()) + .collect(); + for centered in [false, true] { + for scaled in [false, true] { + let centers: Vec = (0..p) + .map(|j| { + T::from_f64(if centered { + 100_000.0 + 2.0 * j as f64 + } else { + 0.0 + }) + .unwrap() + }) + .collect(); + let scales: Vec = (0..p) + .map(|j| { + T::from_f64(if scaled { [3.0, 0.5, -2.0][j % 3] } else { 1.0 }) + .unwrap() + }) + .collect(); + let mut eager = x.clone(); + eager.scale_mut(¢ers, &scales, 0); + let design = LazyDesign::new( + &x, + centered.then_some(centers), + scaled.then_some(scales), + ); + let coefficients: Vec = (0..p) + .map(|j| T::from_f64((j % 5) as f64 - 2.0).unwrap()) + .collect(); + let tolerance = 128.0 * T::epsilon().to_f64().unwrap(); + assert_relative_close( + &design.matvec(&coefficients), + &coefficients.xa(true, &eager), + tolerance, + ); + assert_relative_close( + &design.transpose_matvec(&weights), + &weights.xa(false, &eager), + tolerance, + ); + let mut actual = X::fill(p, p, T::nan()); + design + .normalized + .weighted_gram_into( + &weights, + &mut MatrixOutput(&mut actual, PhantomData), + ) + .unwrap(); + for j in 0..p { + for k in 0..=j { + let expected = (0..n).fold(T::zero(), |sum, i| { + sum + weights[i] * *eager.get((i, j)) * *eager.get((i, k)) + }); + assert_eq!( + *actual.get((j, k)), + expected, + "shape=({n},{p}), axis={axis}, pair=({j},{k})" + ); + assert_eq!(*actual.get((k, j)), expected); + } + } + } + } + } + } + } + + #[cfg_attr( + all(target_arch = "wasm32", not(target_os = "wasi")), + wasm_bindgen_test::wasm_bindgen_test + )] + #[test] + fn panel_gram_preserves_eager_accumulation_across_boundaries() { + check_panel_gram::>(); + check_panel_gram::>(); + #[cfg(feature = "ndarray-bindings")] + { + check_panel_gram::>(); + check_panel_gram::>(); + } + } + + #[cfg_attr( + all(target_arch = "wasm32", not(target_os = "wasi")), + wasm_bindgen_test::wasm_bindgen_test + )] + #[test] + fn panel_gram_preserves_nonfinite_products() { + let x = DenseMatrix::from_2d_array(&[ + &[0.0, f64::INFINITY, 2.0], + &[1.0, 3.0, f64::NAN], + &[-2.0, 0.0, 4.0], + ]) + .unwrap(); + let weights = vec![0.0, 1.0, f64::INFINITY]; + let design = LazyDesign::new(&x, None, Some(vec![1.0, -2.0, f64::INFINITY])); + let mut eager = x.clone(); + eager.scale_mut(&[0.0; 3], &[1.0, -2.0, f64::INFINITY], 0); + let mut actual = DenseMatrix::zeros(3, 3); + design + .normalized + .weighted_gram_into(&weights, &mut MatrixOutput(&mut actual, PhantomData)) + .unwrap(); + for j in 0..3 { + for k in 0..=j { + let expected = (0..3).fold(0.0, |sum, i| { + sum + weights[i] * *eager.get((i, j)) * *eager.get((i, k)) + }); + for index in [(j, k), (k, j)] { + if expected.is_nan() { + assert!(actual.get(index).is_nan()); + } else { + assert_eq!(*actual.get(index), expected); + } + } + } + } + } + + #[cfg_attr( + all(target_arch = "wasm32", not(target_os = "wasi")), + wasm_bindgen_test::wasm_bindgen_test + )] + #[test] + fn implicit_augmentation_matches_materialized_products_and_fit() { + let x = DenseMatrix::from_2d_array(&[&[1., 4.], &[2., 6.], &[4., 10.]]).unwrap(); + for penalty in [0.0_f64, 0.7, 100.0] { + let gamma = 1.0 / (1.0 + penalty).sqrt(); + let padding = gamma * penalty.sqrt(); + let design = + LazyDesign::new(&x, Some(vec![1., 2.]), Some(vec![2., 4.])).with_l2(penalty); + let mut eager = DenseMatrix::zeros(5, 2); + for j in 0..2 { + for i in 0..3 { + eager.set((i, j), gamma * (x.get((i, j)) - [1., 2.][j]) / [2., 4.][j]); + } + eager.set((3 + j, j), padding); + } + let w = vec![2., -1.]; + let y = vec![7., 9., 13., 0., 0.]; + assert_close(&design.matvec(&w), &w.xa(true, &eager), 1e-12); + assert_close(&design.transpose_matvec(&y), &y.xa(false, &eager), 1e-12); + assert_close( + &design.gram().iterator(0).copied().collect::>(), + &eager + .ab(true, &eager, false) + .iterator(0) + .copied() + .collect::>(), + 1e-12, + ); + let mut lazy_opt = InteriorPointOptimizer::new_lazy(&design); + let mut eager_opt = InteriorPointOptimizer::new(&eager, 2); + let actual = lazy_opt + .optimize_lazy(&design, &y, 0.2 * gamma, 1000, 1e-8, true) + .unwrap(); + let expected = eager_opt + .optimize(&eager, &y, 0.2 * gamma, 1000, 1e-8, true) + .unwrap(); + assert_close(&actual, &expected, 1e-6); + } + } +} diff --git a/src/linear/mod.rs b/src/linear/mod.rs index fb3bf8a4..d3f4f129 100644 --- a/src/linear/mod.rs +++ b/src/linear/mod.rs @@ -24,6 +24,8 @@ pub mod bg_solver; pub mod elastic_net; pub mod lasso; pub mod lasso_optimizer; +#[cfg(feature = "lazy-normalization")] +mod lazy_normalization; pub mod linear_regression; pub mod logistic_regression; pub mod ridge_regression; 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 02/18] 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 03/18] 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 04/18] 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 05/18] 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 06/18] 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 07/18] 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 08/18] 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 09/18] 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] + ); } } From 0ff587f5ba3b66f3d8fd0453ba3890bfb6bff9a7 Mon Sep 17 00:00:00 2001 From: slievens <5492873+slievens@users.noreply.github.com> Date: Thu, 22 Oct 2026 20:14:48 +0200 Subject: [PATCH 10/18] Make the base tree generic over a split criterion --- src/ensemble/base_forest_regressor.rs | 7 +- src/tree/base_tree_regressor.rs | 633 ++++++++++++++++++-------- src/tree/criterion.rs | 50 +- src/tree/decision_tree_regressor.rs | 6 +- 4 files changed, 494 insertions(+), 202 deletions(-) diff --git a/src/ensemble/base_forest_regressor.rs b/src/ensemble/base_forest_regressor.rs index cea579fe..a30865ae 100644 --- a/src/ensemble/base_forest_regressor.rs +++ b/src/ensemble/base_forest_regressor.rs @@ -12,6 +12,7 @@ use crate::numbers::floatnum::FloatNumber; use crate::rand_custom::get_rng_impl; use crate::tree::base_tree_regressor::{BaseTreeRegressor, BaseTreeRegressorParameters, Splitter}; +use crate::tree::criterion::MSECriterion; #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[derive(Debug, Clone)] @@ -70,7 +71,7 @@ pub struct BaseForestRegressor< X: Array2, Y: Array1, > { - trees: Option>>, + trees: Option>>, num_features: Option, samples: Option>>, } @@ -106,7 +107,8 @@ impl, Y: Array1 let mut rng = get_rng_impl(Some(parameters.seed)); let n_trees = parameters.n_trees; - let mut trees: Vec> = Vec::with_capacity(n_trees); + let mut trees: Vec> = + Vec::with_capacity(n_trees); let mut maybe_all_samples: Option>> = Option::None; if parameters.keep_samples { @@ -164,6 +166,7 @@ impl, Y: Array1 samples.clone(), mtry, &order, + &MSECriterion, params, )?; trees.push(tree); diff --git a/src/tree/base_tree_regressor.rs b/src/tree/base_tree_regressor.rs index 0720631d..2a067b85 100644 --- a/src/tree/base_tree_regressor.rs +++ b/src/tree/base_tree_regressor.rs @@ -11,7 +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}; +use crate::tree::criterion::Criterion; #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[derive(Debug, Clone, Copy, Default)] @@ -45,9 +45,22 @@ pub struct BaseTreeRegressorParameters { /// Regression base_tree #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +#[cfg_attr( + feature = "serde", + serde(bound( + serialize = "C::Output: Serialize", + deserialize = "C::Output: Deserialize<'de>" + )) +)] #[derive(Debug)] -pub struct BaseTreeRegressor, Y: Array1> { - nodes: Vec, +pub struct BaseTreeRegressor< + TX: Number + PartialOrd, + TY: Number, + X: Array2, + Y: Array1, + C: Criterion, +> { + nodes: Vec>, parameters: Option, depth: u16, root_mass: f64, @@ -57,11 +70,11 @@ pub struct BaseTreeRegressor, _phantom_y: PhantomData, } -impl, Y: Array1> - BaseTreeRegressor +impl, Y: Array1, C: Criterion> + BaseTreeRegressor { /// Get nodes, return a shared reference - fn nodes(&self) -> &Vec { + fn nodes(&self) -> &Vec> { self.nodes.as_ref() } /// Get parameters, return a shared reference @@ -72,8 +85,8 @@ impl, Y: Array1> #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[derive(Debug, Clone)] -struct Node { - output: f64, +struct Node { + output: Out, split_feature: usize, split_value: Option, split_score: Option, @@ -81,8 +94,8 @@ struct Node { false_child: Option, } -impl Node { - fn new(output: f64) -> Self { +impl Node { + fn new(output: Out) -> Self { Node { output, split_feature: 0, @@ -94,9 +107,9 @@ impl Node { } } -impl PartialEq for Node { - fn eq(&self, other: &Self) -> bool { - (self.output - other.output).abs() < f64::EPSILON +impl Node { + fn approximate_eq>(&self, other: &Self) -> bool { + C::output_approximate_eq(&self.output, &other.output) && self.split_feature == other.split_feature && match (self.split_value, other.split_value) { (Some(a), Some(b)) => (a - b).abs() < f64::EPSILON, @@ -111,8 +124,8 @@ impl PartialEq for Node { } } -impl, Y: Array1> PartialEq - for BaseTreeRegressor +impl, Y: Array1, C: Criterion> PartialEq + for BaseTreeRegressor { fn eq(&self, other: &Self) -> bool { if self.depth != other.depth || self.nodes.len() != other.nodes.len() { @@ -121,25 +134,32 @@ impl, Y: Array1> PartialE self.nodes .iter() .zip(other.nodes().iter()) - .all(|(a, b)| a == b) + .all(|(a, b)| a.approximate_eq::(b)) } } } -struct NodeVisitor<'a, TX: Number + PartialOrd, TY: Number, X: Array2, Y: Array1> { +struct NodeVisitor< + 'a, + TX: Number + PartialOrd, + TY: Number, + X: Array2, + Y: Array1, + C: Criterion, +> { x: &'a X, y: &'a Y, node: usize, start_idx: usize, // start index for the elements in `sorted_by_feature` of SplitWorkspace end_idx: usize, // end index (exclusive) in the same vector(s) level: u16, - stats: MeanStats, // statistics of the samples in this node + stats: C::Stats, // statistics of the samples in this node _phantom_tx: PhantomData, _phantom_ty: PhantomData, } -impl<'a, TX: Number + PartialOrd, TY: Number, X: Array2, Y: Array1> - NodeVisitor<'a, TX, TY, X, Y> +impl<'a, TX: Number + PartialOrd, TY: Number, X: Array2, Y: Array1, C: Criterion> + NodeVisitor<'a, TX, TY, X, Y, C> { fn new( node_id: usize, @@ -148,7 +168,7 @@ impl<'a, TX: Number + PartialOrd, TY: Number, X: Array2, Y: Array1> x: &'a X, y: &'a Y, level: u16, - stats: MeanStats, + stats: C::Stats, ) -> Self { NodeVisitor { x, @@ -282,7 +302,7 @@ impl NodeElement for UnitElement { } // Checks whether the example indicated by node_element is a "true child" for this node -fn is_true_sample(node_element: &E, x: &X, node: &Node) -> bool +fn is_true_sample(node_element: &E, x: &X, node: &Node) -> bool where TX: Number + PartialOrd, X: Array2, @@ -320,44 +340,49 @@ where w } -struct SplitWorkspace { +struct SplitWorkspace { in_true_branch: Vec, // indexed by row index in X partition_buffer: Vec, sorted_by_feature: Vec>, variables: Vec, // feature indices which will be shuffled when mtry < n_features - true_stats: MeanStats, // scratch stats for split functions to use - false_stats: MeanStats, - parent_stats: MeanStats, + true_stats: C::Stats, // scratch stats for split functions to use + false_stats: C::Stats, + parent_stats: C::Stats, } -impl SplitWorkspace { +impl SplitWorkspace +where + C: Criterion, +{ fn new( in_true_branch: Vec, partition_buffer: Vec, sorted_by_feature: Vec>, n_features: usize, + criterion: &C, ) -> Self { Self { in_true_branch, partition_buffer, sorted_by_feature, variables: (0..n_features).collect::>(), - true_stats: MSECriterion.empty_stats(), - false_stats: MSECriterion.empty_stats(), - parent_stats: MSECriterion.empty_stats(), + true_stats: criterion.empty_stats(), + false_stats: criterion.empty_stats(), + parent_stats: criterion.empty_stats(), } } } -impl, Y: Array1> - BaseTreeRegressor +impl, Y: Array1, C: Criterion> + BaseTreeRegressor { pub(crate) fn fit_inner( x: &X, y: &Y, sample_weights: Option<&[f64]>, + criterion: &C, parameters: BaseTreeRegressorParameters, - ) -> Result, Failed> { + ) -> Result, Failed> { let (x_nrows, num_attributes) = x.shape(); if x_nrows != y.shape() { return Err(Failed::fit("Size of x should equal size of y")); @@ -385,6 +410,7 @@ impl, Y: Array1> samples, num_attributes, &order, + criterion, parameters, ) } @@ -396,8 +422,9 @@ impl, Y: Array1> samples: Vec, mtry: usize, order: &[Vec], + criterion: &C, parameters: BaseTreeRegressorParameters, - ) -> Result, Failed> { + ) -> Result, Failed> { check_row_count(y.shape())?; match sample_weights { Some(_) => Self::grow::( @@ -407,14 +434,17 @@ impl, Y: Array1> samples, mtry, order, + criterion, parameters, ), None if samples.iter().all(|&s| s <= 1) => { // Note: if any sample has count zero, it will be filtered out, so we should // use the smallest possible NodeElement - Self::grow::(x, y, None, samples, mtry, order, parameters) + Self::grow::(x, y, None, samples, mtry, order, criterion, parameters) } - None => Self::grow::(x, y, None, samples, mtry, order, parameters), + None => Self::grow::( + x, y, None, samples, mtry, order, criterion, parameters, + ), } } @@ -425,14 +455,15 @@ impl, Y: Array1> samples: Vec, mtry: usize, order: &[Vec], + criterion: &C, parameters: BaseTreeRegressorParameters, - ) -> Result, Failed> { - let mut nodes: Vec = Vec::new(); + ) -> Result, Failed> { + let mut nodes: Vec> = Vec::new(); let mut rng = get_rng_impl(parameters.seed); let mut n = 0usize; let mut root_mass = 0f64; - let mut root_stats = MSECriterion.empty_stats(); + let mut root_stats = criterion.empty_stats(); // A row with mass 0 (bootstrap count 0 or sample weight 0) has no effect on // the statistics. Give it count 0, so that it is left out of the node elements and @@ -450,12 +481,12 @@ impl, Y: Array1> for (i, &c) in counts.iter().enumerate() { // A row with mass 0 adds 0 to the statistics: no branch is necessary let mass = mass_of(i, &samples, sample_weights); - MSECriterion.add(&mut root_stats, y.get(i).to_f64().unwrap(), mass); + criterion.add(&mut root_stats, criterion.target(*y.get(i)), mass); root_mass += mass; n += c as usize; } - let root = Node::new(MSECriterion.leaf_value(&root_stats)); + let root = Node::new(criterion.leaf_value(&root_stats)); nodes.push(root); // number of distinct rows of x in this tree with a positive mass @@ -483,6 +514,7 @@ impl, Y: Array1> vec![E::default(); end_idx], sorted_by_feature, x.shape().1, + criterion, ); let mut base_tree = BaseTreeRegressor { @@ -496,14 +528,21 @@ impl, Y: Array1> _phantom_y: PhantomData, }; - let mut visitor = NodeVisitor::::new(0, 0, end_idx, x, y, 1, root_stats); + let mut visitor = NodeVisitor::::new(0, 0, end_idx, x, y, 1, root_stats); - let mut visitor_stack: Vec> = Vec::new(); + let mut visitor_stack: Vec> = Vec::new(); if base_tree.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) + && base_tree.find_best_cutoff( + &mut visitor, + mtry, + &mut rng, + &mut workspace, + n, + criterion, + ) { visitor_stack.push(visitor); } @@ -511,33 +550,26 @@ impl, Y: Array1> let max_depth = base_tree.parameters().max_depth.unwrap_or(u16::MAX); while let Some(node) = visitor_stack.pop() { if node.level < max_depth { - base_tree.split(node, mtry, &mut visitor_stack, &mut rng, &mut workspace); + base_tree.split( + node, + mtry, + &mut visitor_stack, + &mut rng, + &mut workspace, + criterion, + ); } } Ok(base_tree) } - /// Predict regression value for `x`. - /// * `x` - _KxM_ data where _K_ is number of observations and _M_ is number of features. - pub fn predict(&self, x: &X) -> Result { - let mut result = Y::zeros(x.shape().0); - - let (n, _) = x.shape(); - - for i in 0..n { - result.set(i, self.predict_for_row(x, i)); - } - - Ok(result) - } - - pub(crate) fn predict_for_row(&self, x: &X, row: usize) -> TY { + pub(crate) fn leaf_output_for_row(&self, x: &X, row: usize) -> &C::Output { let mut node_id = 0; loop { let node = &self.nodes[node_id]; let Some(true_child) = node.true_child else { - return TY::from_f64(node.output).unwrap(); + return &node.output; }; let false_child = node.false_child.unwrap(); node_id = if x.get((row, node.split_feature)).to_f64().unwrap() @@ -552,11 +584,12 @@ impl, Y: Array1> fn find_best_cutoff( &mut self, - visitor: &mut NodeVisitor<'_, TX, TY, X, Y>, + visitor: &mut NodeVisitor<'_, TX, TY, X, Y, C>, mtry: usize, rng: &mut impl rand::Rng, - workspace: &mut SplitWorkspace, + workspace: &mut SplitWorkspace, n: usize, + criterion: &C, ) -> bool { workspace.variables.shuffle(rng); @@ -566,10 +599,10 @@ impl, Y: Array1> let variable = workspace.variables[k]; match splitter { Splitter::Random => { - self.find_random_split(visitor, n, variable, rng, workspace); + self.find_random_split(visitor, n, variable, rng, workspace, criterion); } Splitter::Best => { - self.find_best_split(visitor, n, variable, workspace); + self.find_best_split(visitor, n, variable, workspace, criterion); } } } @@ -579,11 +612,12 @@ impl, Y: Array1> fn find_random_split( &mut self, - visitor: &NodeVisitor<'_, TX, TY, X, Y>, + visitor: &NodeVisitor<'_, TX, TY, X, Y, C>, n: usize, j: usize, rng: &mut impl rand::Rng, - workspace: &mut SplitWorkspace, + workspace: &mut SplitWorkspace, + criterion: &C, ) { if visitor.start_idx == visitor.end_idx { return; @@ -606,13 +640,13 @@ impl, Y: Array1> let mut true_count = 0; let true_stats = &mut workspace.true_stats; - MSECriterion.reset(true_stats); + criterion.reset(true_stats); for elem in &workspace.sorted_by_feature[j][visitor.start_idx..visitor.end_idx] { if visitor.x.get((elem.row(), j)).to_f64().unwrap() <= split_value { - MSECriterion.add( + criterion.add( true_stats, - visitor.y.get(elem.row()).to_f64().unwrap(), + criterion.target(*visitor.y.get(elem.row())), elem.mass(), ); true_count += elem.count(); @@ -631,8 +665,8 @@ impl, Y: Array1> // 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); + criterion.difference(&visitor.stats, true_stats, false_stats); + let proxy = criterion.proxy_improvement(true_stats, false_stats); if self.nodes[visitor.node].split_score.is_none() || proxy > self.nodes[visitor.node].split_score.unwrap() @@ -645,14 +679,15 @@ impl, Y: Array1> fn find_best_split( &mut self, - visitor: &NodeVisitor<'_, TX, TY, X, Y>, + visitor: &NodeVisitor<'_, TX, TY, X, Y, C>, n: usize, j: usize, - workspace: &mut SplitWorkspace, + workspace: &mut SplitWorkspace, + criterion: &C, ) { let mut true_count = 0; let true_stats = &mut workspace.true_stats; - MSECriterion.reset(true_stats); + criterion.reset(true_stats); // no need to reset false_stats, will be overwritten by 'difference' let false_stats = &mut workspace.false_stats; // set the "previous" element equal to the first one, so we don't have to deal with Option @@ -668,9 +703,9 @@ impl, Y: Array1> if x_ij == prevx { true_count += elem.count(); - MSECriterion.add( + criterion.add( true_stats, - visitor.y.get(elem.row()).to_f64().unwrap(), + criterion.target(*visitor.y.get(elem.row())), elem.mass(), ); continue; @@ -681,16 +716,16 @@ impl, Y: Array1> if true_count < min_samples_leaf || false_count < min_samples_leaf { prevx = x_ij; true_count += elem.count(); - MSECriterion.add( + criterion.add( true_stats, - visitor.y.get(elem.row()).to_f64().unwrap(), + criterion.target(*visitor.y.get(elem.row())), elem.mass(), ); continue; } - MSECriterion.difference(&visitor.stats, true_stats, false_stats); - let proxy = MSECriterion.proxy_improvement(true_stats, false_stats); + criterion.difference(&visitor.stats, true_stats, false_stats); + let proxy = criterion.proxy_improvement(true_stats, false_stats); if self.nodes[visitor.node].split_score.is_none() || proxy > self.nodes[visitor.node].split_score.unwrap() @@ -710,9 +745,9 @@ impl, Y: Array1> prevx = x_ij; true_count += elem.count(); - MSECriterion.add( + criterion.add( true_stats, - visitor.y.get(elem.row()).to_f64().unwrap(), + criterion.target(*visitor.y.get(elem.row())), elem.mass(), ); } @@ -726,14 +761,15 @@ impl, Y: Array1> /// keep the node as a leaf. fn split<'a, E: NodeElement>( &mut self, - visitor: NodeVisitor<'a, TX, TY, X, Y>, + visitor: NodeVisitor<'a, TX, TY, X, Y, C>, mtry: usize, - visitor_stack: &mut Vec>, + visitor_stack: &mut Vec>, rng: &mut impl rand::Rng, - workspace: &mut SplitWorkspace, + workspace: &mut SplitWorkspace, + criterion: &C, ) -> bool { let parent_stats = &mut workspace.parent_stats; - MSECriterion.reset(parent_stats); + criterion.reset(parent_stats); let this_node = &self.nodes[visitor.node]; let mut true_count = 0usize; @@ -751,26 +787,26 @@ impl, Y: Array1> } in_true_branch[e.row()] = true; true_count += e.count(); - MSECriterion.add( + criterion.add( parent_stats, - visitor.y.get(e.row()).to_f64().unwrap(), + criterion.target(*visitor.y.get(e.row())), e.mass(), ); n_true += 1; } let true_stats = parent_stats.clone(); - let mut false_stats = MSECriterion.empty_stats(); + let mut false_stats = criterion.empty_stats(); for e in &col[n_true..] { in_true_branch[e.row()] = false; false_count += e.count(); - MSECriterion.add( + criterion.add( parent_stats, - visitor.y.get(e.row()).to_f64().unwrap(), + criterion.target(*visitor.y.get(e.row())), e.mass(), ); - MSECriterion.add( + criterion.add( &mut false_stats, - visitor.y.get(e.row()).to_f64().unwrap(), + criterion.target(*visitor.y.get(e.row())), e.mass(), ); } @@ -789,7 +825,7 @@ impl, Y: Array1> // Store the real (mass-weighted) impurity improvement in split_score, not the proxy. // false_stats is a direct sum, not parent - true: no cancellation for a small false side self.nodes[visitor.node].split_score = - Some(MSECriterion.impurity_improvement(parent_stats, &true_stats, &false_stats)); + Some(criterion.impurity_improvement(parent_stats, &true_stats, &false_stats)); // Add the child nodes to the tree let split_idx = visitor.start_idx + n_true; @@ -797,10 +833,10 @@ impl, Y: Array1> let true_child_idx = self.nodes.len(); self.nodes - .push(Node::new(MSECriterion.leaf_value(&true_stats))); + .push(Node::new(criterion.leaf_value(&true_stats))); let false_child_idx = self.nodes.len(); self.nodes - .push(Node::new(MSECriterion.leaf_value(&false_stats))); + .push(Node::new(criterion.leaf_value(&false_stats))); self.nodes[visitor.node].true_child = Some(true_child_idx); self.nodes[visitor.node].false_child = Some(false_child_idx); @@ -825,7 +861,7 @@ impl, Y: Array1> ); } - let mut true_visitor = NodeVisitor::::new( + let mut true_visitor = NodeVisitor::::new( true_child_idx, visitor.start_idx, split_idx, @@ -836,12 +872,19 @@ impl, Y: Array1> ); if true_can_split - && self.find_best_cutoff(&mut true_visitor, mtry, rng, workspace, true_count) + && self.find_best_cutoff( + &mut true_visitor, + mtry, + rng, + workspace, + true_count, + criterion, + ) { visitor_stack.push(true_visitor); } - let mut false_visitor = NodeVisitor::::new( + let mut false_visitor = NodeVisitor::::new( false_child_idx, split_idx, visitor.end_idx, @@ -852,7 +895,14 @@ impl, Y: Array1> ); if false_can_split - && self.find_best_cutoff(&mut false_visitor, mtry, rng, workspace, false_count) + && self.find_best_cutoff( + &mut false_visitor, + mtry, + rng, + workspace, + false_count, + criterion, + ) { visitor_stack.push(false_visitor); } @@ -890,12 +940,41 @@ impl, Y: Array1> } } +// Regression: the leaf holds the prediction itself. +impl BaseTreeRegressor +where + TX: Number + PartialOrd, + TY: Number, + X: Array2, + Y: Array1, + C: Criterion, +{ + pub(crate) fn predict_for_row(&self, x: &X, row: usize) -> TY { + TY::from_f64(*self.leaf_output_for_row(x, row)).unwrap() + } + + /// Predict regression value for `x`. + /// * `x` - _KxM_ data where _K_ is number of observations and _M_ is number of features. + pub fn predict(&self, x: &X) -> Result { + let mut result = Y::zeros(x.shape().0); + + let (n, _) = x.shape(); + + for i in 0..n { + result.set(i, self.predict_for_row(x, i)); + } + + Ok(result) + } +} + #[cfg(test)] mod tests { use super::*; use crate::linalg::basic::arrays::Array; use crate::linalg::basic::matrix::DenseMatrix; use crate::metrics::mean_absolute_error; + use crate::tree::criterion::MSECriterion; #[test] #[cfg(target_pointer_width = "64")] @@ -922,6 +1001,7 @@ mod tests { &empty, &y, None, + &MSECriterion, BaseTreeRegressorParameters { max_depth: None, min_samples_leaf: 1, @@ -946,6 +1026,7 @@ mod tests { &no_features, &y, None, + &MSECriterion, BaseTreeRegressorParameters { max_depth: None, min_samples_leaf: 1, @@ -972,6 +1053,7 @@ mod tests { &x, &y, Some(&sample_weights), + &MSECriterion, BaseTreeRegressorParameters { max_depth: Some(0), min_samples_leaf: 1, @@ -1006,12 +1088,17 @@ mod tests { }; let tree_no_weights = - BaseTreeRegressor::fit_inner(&x_rand, &y_rand, None, parameters.clone()) + BaseTreeRegressor::fit_inner(&x_rand, &y_rand, None, &MSECriterion, parameters.clone()) .expect("Fit should work"); let uniform_weights = vec![1.0f64; 17]; - let tree_with_weights = - BaseTreeRegressor::fit_inner(&x_rand, &y_rand, Some(&uniform_weights), parameters) - .expect("Fit should work"); + let tree_with_weights = BaseTreeRegressor::fit_inner( + &x_rand, + &y_rand, + Some(&uniform_weights), + &MSECriterion, + parameters, + ) + .expect("Fit should work"); let y_pred_no_weights = tree_no_weights .predict(&x_rand) @@ -1057,12 +1144,22 @@ mod tests { splitter: Splitter::Best, }; - let tree_weighted = - BaseTreeRegressor::fit_inner(&x, &y, Some(&sample_weights), weighted_parameters) - .expect("Fit should work"); - let tree_repeated = - BaseTreeRegressor::fit_inner(&x_repeated, &y_repeated, None, repeated_parameters) - .expect("Fit should work"); + let tree_weighted = BaseTreeRegressor::fit_inner( + &x, + &y, + Some(&sample_weights), + &MSECriterion, + weighted_parameters, + ) + .expect("Fit should work"); + let tree_repeated = BaseTreeRegressor::fit_inner( + &x_repeated, + &y_repeated, + None, + &MSECriterion, + repeated_parameters, + ) + .expect("Fit should work"); // Predict on the same data let y_pred_weighted = tree_weighted.predict(&x).expect("Predict should work"); let y_pred_repeated = tree_repeated.predict(&x).expect("Predict should work"); @@ -1091,7 +1188,8 @@ mod tests { splitter: Splitter::Best, }; - let tree = BaseTreeRegressor::fit_inner(&x, &y, None, parameters).expect("Fit should work"); + let tree = BaseTreeRegressor::fit_inner(&x, &y, None, &MSECriterion, parameters) + .expect("Fit should work"); let y_expected = vec![1.0, 2.0, 6.5, 6.5, 11.50, 11.50]; let y_hat = tree.predict(&x).expect("Predict should work"); assert_eq!(tree.nodes().len(), 7); @@ -1120,8 +1218,9 @@ mod tests { }; let sample_weights = [1.0f64, 1.0, 1.0, 1.0, 10.0, 10.0]; - let tree = BaseTreeRegressor::fit_inner(&x, &y, Some(&sample_weights), parameters) - .expect("Fit should work"); + let tree = + BaseTreeRegressor::fit_inner(&x, &y, Some(&sample_weights), &MSECriterion, parameters) + .expect("Fit should work"); let x_test: DenseMatrix = DenseMatrix::from_iterator((1..=12).map(|i| i as f64), 12, 1, 0); let y_expected = vec![ @@ -1147,7 +1246,8 @@ mod tests { splitter: Splitter::Best, }; - let tree = BaseTreeRegressor::fit_inner(&x, &y, None, parameters).expect("Fit should work"); + let tree = BaseTreeRegressor::fit_inner(&x, &y, None, &MSECriterion, parameters) + .expect("Fit should work"); assert_eq!(tree.nodes().len(), 3); assert_eq!(tree.depth, 2); @@ -1160,7 +1260,8 @@ mod tests { splitter: Splitter::Best, }; - let tree = BaseTreeRegressor::fit_inner(&x, &y, None, parameters).expect("Fit should work"); + let tree = BaseTreeRegressor::fit_inner(&x, &y, None, &MSECriterion, parameters) + .expect("Fit should work"); assert_eq!(tree.nodes().len(), 1); assert_eq!(tree.depth, 0); } @@ -1181,8 +1282,9 @@ mod tests { splitter: Splitter::Best, }; - let tree = BaseTreeRegressor::fit_inner(&x, &y, Some(&sample_weights), parameters) - .expect("Fit should work"); + let tree = + BaseTreeRegressor::fit_inner(&x, &y, Some(&sample_weights), &MSECriterion, parameters) + .expect("Fit should work"); assert!(tree.nodes().iter().all(|node| node.output.is_finite())); assert!( @@ -1216,8 +1318,9 @@ mod tests { splitter: Splitter::Best, }; - let tree = BaseTreeRegressor::fit_inner(&x, &y, Some(&sample_weights), parameters) - .expect("Fit should work"); + let tree = + BaseTreeRegressor::fit_inner(&x, &y, Some(&sample_weights), &MSECriterion, parameters) + .expect("Fit should work"); assert!(tree.nodes().iter().all(|node| node.output.is_finite())); assert!( @@ -1247,8 +1350,14 @@ mod tests { splitter: Splitter::Best, }; - let tree = BaseTreeRegressor::fit_inner(&x, &y, Some(&sample_weights), parameters.clone()) - .expect("Fit should work"); + let tree = BaseTreeRegressor::fit_inner( + &x, + &y, + Some(&sample_weights), + &MSECriterion, + parameters.clone(), + ) + .expect("Fit should work"); assert_eq!(tree.nodes().len(), 3); assert_eq!(tree.depth, 2); @@ -1272,7 +1381,7 @@ mod tests { .unwrap(); let y_repeated = vec![0.0f64, 0.0, 0.0, 10.0, 20.0, 30.0, 30.0, 30.0]; let tree_repeated = - BaseTreeRegressor::fit_inner(&x_repeated, &y_repeated, None, parameters) + BaseTreeRegressor::fit_inner(&x_repeated, &y_repeated, None, &MSECriterion, parameters) .expect("Fit should work"); let y_hat_repeated = tree_repeated.predict(&x).expect("Predict should work"); assert!(mean_absolute_error(&y_hat, &y_hat_repeated) < 1e-9); @@ -1300,9 +1409,14 @@ mod tests { DenseMatrix::from_2d_vec(&vec![vec![TWO_POW_53 + 2], vec![TWO_POW_53 + 4]]).unwrap(); let y = vec![0.0f64, 10.0]; - let tree = - BaseTreeRegressor::fit_inner(&x, &y, None, split_value_parameters(Splitter::Best)) - .expect("Fit should work"); + let tree = BaseTreeRegressor::fit_inner( + &x, + &y, + None, + &MSECriterion, + split_value_parameters(Splitter::Best), + ) + .expect("Fit should work"); let threshold = tree.nodes()[0].split_value.expect("Root should split"); assert!((TWO_POW_53 + 2) as f64 <= threshold); @@ -1319,9 +1433,14 @@ mod tests { let x = DenseMatrix::from_2d_vec(&vec![vec![TWO_POW_53], vec![TWO_POW_53 + 1]]).unwrap(); let y = vec![0.0f64, 10.0]; - let tree = - BaseTreeRegressor::fit_inner(&x, &y, None, split_value_parameters(Splitter::Best)) - .expect("Fit should work"); + let tree = BaseTreeRegressor::fit_inner( + &x, + &y, + None, + &MSECriterion, + split_value_parameters(Splitter::Best), + ) + .expect("Fit should work"); let y_hat = tree.predict(&x).expect("Predict should work"); assert_eq!(y_hat, vec![5.0, 5.0]); @@ -1333,9 +1452,14 @@ mod tests { let x = DenseMatrix::from_2d_vec(&vec![vec![TWO_POW_53], vec![TWO_POW_53 + 1]]).unwrap(); let y = vec![0.0f64, 10.0]; - let tree = - BaseTreeRegressor::fit_inner(&x, &y, None, split_value_parameters(Splitter::Random)) - .expect("Fit should work"); + let tree = BaseTreeRegressor::fit_inner( + &x, + &y, + None, + &MSECriterion, + split_value_parameters(Splitter::Random), + ) + .expect("Fit should work"); let y_hat = tree.predict(&x).expect("Predict should work"); assert_eq!(y_hat, vec![5.0, 5.0]); @@ -1353,9 +1477,14 @@ mod tests { let x = DenseMatrix::from_2d_vec(&vec![vec![a], vec![b]]).unwrap(); let y = vec![0.0f64, 10.0]; - let tree = - BaseTreeRegressor::fit_inner(&x, &y, None, split_value_parameters(Splitter::Best)) - .expect("Fit should work"); + let tree = BaseTreeRegressor::fit_inner( + &x, + &y, + None, + &MSECriterion, + split_value_parameters(Splitter::Best), + ) + .expect("Fit should work"); let threshold = tree.nodes()[0].split_value.expect("Root should split"); assert!(a <= threshold); @@ -1372,7 +1501,7 @@ mod tests { assert_eq!(std::mem::size_of::(), 16); } - type TestTree = BaseTreeRegressor, Vec>; + type TestTree = BaseTreeRegressor, Vec, MSECriterion>; fn params( max_depth: Option, @@ -1492,29 +1621,45 @@ mod tests { (None, 15, 4), ]; for (max_depth, n_nodes, depth) in cases { - let tree: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, None, params(max_depth, 1, 2, Splitter::Best)) - .expect("Fit should work"); + let tree: TestTree = BaseTreeRegressor::fit_inner( + &x, + &y, + None, + &MSECriterion, + params(max_depth, 1, 2, Splitter::Best), + ) + .expect("Fit should work"); assert_eq!(tree.nodes().len(), n_nodes, "max_depth: {max_depth:?}"); assert_eq!(tree.depth, depth, "max_depth: {max_depth:?}"); } // With max_depth = 2 the root split separates the two halves. - let tree: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, None, params(Some(2), 1, 2, Splitter::Best)) - .expect("Fit should work"); + let tree: TestTree = BaseTreeRegressor::fit_inner( + &x, + &y, + None, + &MSECriterion, + params(Some(2), 1, 2, Splitter::Best), + ) + .expect("Fit should work"); assert_eq!(tree.nodes()[0].split_value, Some(4.5)); let y_hat = tree.predict(&x).expect("Predict should work"); assert_eq!(y_hat, vec![5.5, 5.5, 5.5, 5.5, 105.5, 105.5, 105.5, 105.5]); // No limit is the same as the largest limit, and the leaves memorize y. - let tree_none: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, None, params(None, 1, 2, Splitter::Best)) - .expect("Fit should work"); + let tree_none: TestTree = BaseTreeRegressor::fit_inner( + &x, + &y, + None, + &MSECriterion, + params(None, 1, 2, Splitter::Best), + ) + .expect("Fit should work"); let tree_max: TestTree = BaseTreeRegressor::fit_inner( &x, &y, None, + &MSECriterion, params(Some(u16::MAX), 1, 2, Splitter::Best), ) .expect("Fit should work"); @@ -1526,6 +1671,7 @@ mod tests { &x, &y, None, + &MSECriterion, params(Some(max_depth), 1, 2, Splitter::Random), ) .expect("Fit should work"); @@ -1547,9 +1693,14 @@ mod tests { .unwrap(); let y = vec![0.0f64, 1.0, 2.0, 100.0, 101.0]; - let tree: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, None, params(None, 1, 3, Splitter::Best)) - .expect("Fit should work"); + let tree: TestTree = BaseTreeRegressor::fit_inner( + &x, + &y, + None, + &MSECriterion, + params(None, 1, 3, Splitter::Best), + ) + .expect("Fit should work"); assert_eq!(tree.nodes()[0].split_value, Some(3.5)); let true_child = tree.nodes()[0].true_child.unwrap(); let false_child = tree.nodes()[0].false_child.unwrap(); @@ -1560,9 +1711,14 @@ mod tests { // If min_samples_split is larger than the number of samples, the root is a leaf. for splitter in [Splitter::Best, Splitter::Random] { - let tree: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, None, params(None, 1, 6, splitter)) - .expect("Fit should work"); + let tree: TestTree = BaseTreeRegressor::fit_inner( + &x, + &y, + None, + &MSECriterion, + params(None, 1, 6, splitter), + ) + .expect("Fit should work"); assert_eq!(tree.nodes().len(), 1, "splitter: {splitter:?}"); } } @@ -1574,24 +1730,39 @@ mod tests { let y = vec![0.0f64, 0.0, 0.0, 100.0]; // With no limit, the best split isolates the last sample. - let tree: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, None, params(None, 1, 2, Splitter::Best)) - .expect("Fit should work"); + let tree: TestTree = BaseTreeRegressor::fit_inner( + &x, + &y, + None, + &MSECriterion, + params(None, 1, 2, Splitter::Best), + ) + .expect("Fit should work"); assert_eq!(tree.nodes()[0].split_value, Some(3.5)); // The best split is not permitted. The split with exactly 2 samples on each side is used. - let tree: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, None, params(None, 2, 2, Splitter::Best)) - .expect("Fit should work"); + let tree: TestTree = BaseTreeRegressor::fit_inner( + &x, + &y, + None, + &MSECriterion, + params(None, 2, 2, Splitter::Best), + ) + .expect("Fit should work"); assert_eq!(tree.nodes()[0].split_value, Some(2.5)); assert_eq!(tree.nodes().len(), 3); assert_eq!(leaf_counts(&tree, &x, &[1; 4]), vec![2, 2]); // No split can give 3 samples on each side. for splitter in [Splitter::Best, Splitter::Random] { - let tree: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, None, params(None, 3, 2, splitter)) - .expect("Fit should work"); + let tree: TestTree = BaseTreeRegressor::fit_inner( + &x, + &y, + None, + &MSECriterion, + params(None, 3, 2, splitter), + ) + .expect("Fit should work"); assert_eq!(tree.nodes().len(), 1, "splitter: {splitter:?}"); assert_eq!(tree.nodes()[0].output, 25.0); } @@ -1633,6 +1804,7 @@ mod tests { samples.clone(), 3, &order, + &MSECriterion, parameters, ) .expect("Fit should work"); @@ -1665,9 +1837,14 @@ mod tests { let y = vec![1.0f64, 2.0, 3.0, 6.0]; for splitter in [Splitter::Best, Splitter::Random] { - let tree: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, None, params(None, 1, 2, splitter)) - .expect("Fit should work"); + let tree: TestTree = BaseTreeRegressor::fit_inner( + &x, + &y, + None, + &MSECriterion, + params(None, 1, 2, splitter), + ) + .expect("Fit should work"); assert_eq!(tree.nodes().len(), 1, "splitter: {splitter:?}"); assert_eq!(tree.depth, 0, "splitter: {splitter:?}"); let y_hat = tree.predict(&x).expect("Predict should work"); @@ -1688,9 +1865,14 @@ mod tests { let y = vec![0.0f64, 0.0, 10.0, 10.0]; for splitter in [Splitter::Best, Splitter::Random] { - let tree: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, None, params(None, 1, 2, splitter)) - .expect("Fit should work"); + let tree: TestTree = BaseTreeRegressor::fit_inner( + &x, + &y, + None, + &MSECriterion, + params(None, 1, 2, splitter), + ) + .expect("Fit should work"); assert!(tree.nodes().len() > 1, "splitter: {splitter:?}"); assert!( tree.nodes() @@ -1701,9 +1883,14 @@ mod tests { ); } - let tree: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, None, params(None, 1, 2, Splitter::Best)) - .expect("Fit should work"); + let tree: TestTree = BaseTreeRegressor::fit_inner( + &x, + &y, + None, + &MSECriterion, + params(None, 1, 2, Splitter::Best), + ) + .expect("Fit should work"); assert_eq!(tree.nodes()[0].split_value, Some(2.5)); } @@ -1726,11 +1913,16 @@ mod tests { seed: Some(seed), ..params(None, 1, 2, Splitter::Best) }; - let tree_single: TestTree = - BaseTreeRegressor::fit_inner(&x_single, &y, None, parameters.clone()) - .expect("Fit should work"); + let tree_single: TestTree = BaseTreeRegressor::fit_inner( + &x_single, + &y, + None, + &MSECriterion, + parameters.clone(), + ) + .expect("Fit should work"); let tree_double: TestTree = - BaseTreeRegressor::fit_inner(&x_double, &y, None, parameters) + BaseTreeRegressor::fit_inner(&x_double, &y, None, &MSECriterion, parameters) .expect("Fit should work"); assert!(tree_double.nodes().len() > 1); @@ -1774,12 +1966,18 @@ mod tests { counts.clone(), 2, &order, + &MSECriterion, parameters.clone(), ) .expect("Fit should work"); - let tree_rep: TestTree = - BaseTreeRegressor::fit_inner(&x_rep, &y_rep, None, parameters) - .expect("Fit should work"); + let tree_rep: TestTree = BaseTreeRegressor::fit_inner( + &x_rep, + &y_rep, + None, + &MSECriterion, + parameters, + ) + .expect("Fit should work"); assert!(tree.nodes().len() > 1); assert_same_tree(&tree, &tree_rep); @@ -1813,13 +2011,19 @@ mod tests { counts.clone(), 2, &order, + &MSECriterion, parameters.clone(), ) .expect("Fit should work"); // in this tree, there are no counts but the rows are repeated physically - let tree_rep: TestTree = - BaseTreeRegressor::fit_inner(&x_rep, &y_rep, Some(&w_rep), parameters) - .expect("Fit should work"); + let tree_rep: TestTree = BaseTreeRegressor::fit_inner( + &x_rep, + &y_rep, + Some(&w_rep), + &MSECriterion, + parameters, + ) + .expect("Fit should work"); assert!(tree.nodes().len() > 1); assert_same_tree(&tree, &tree_rep); @@ -1877,13 +2081,20 @@ mod tests { vec![ ( "unit", - BaseTreeRegressor::fit_inner(&f.x, &f.y, None, parameters.clone()).unwrap(), + BaseTreeRegressor::fit_inner(&f.x, &f.y, None, &MSECriterion, parameters.clone()) + .unwrap(), vec![1.0; n], ), ( "weighted", - BaseTreeRegressor::fit_inner(&f.x, &f.y, Some(&f.weights), parameters.clone()) - .unwrap(), + BaseTreeRegressor::fit_inner( + &f.x, + &f.y, + Some(&f.weights), + &MSECriterion, + parameters.clone(), + ) + .unwrap(), f.weights.clone(), ), ( @@ -1895,6 +2106,7 @@ mod tests { f.counts.clone(), 2, &f.order, + &MSECriterion, parameters, ) .unwrap(), @@ -2015,12 +2227,22 @@ mod tests { 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 tree: TestTree = BaseTreeRegressor::fit_inner( + &x, + &y, + Some(&weights), + &MSECriterion, + parameters.clone(), + ) + .expect("Fit should work"); + let tree_kept: TestTree = BaseTreeRegressor::fit_inner( + &x_kept, + &y_kept, + Some(&weights_kept), + &MSECriterion, + parameters, + ) + .expect("Fit should work"); assert!( tree == tree_kept, @@ -2047,8 +2269,14 @@ mod tests { 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") + BaseTreeRegressor::fit_inner( + &x, + &y, + None, + &MSECriterion, + params(Some(4), 1, 2, splitter), + ) + .expect("Fit should work") }; let reference = fit(0.0); for offset in [1e6, 1e8, 1e9] { @@ -2088,9 +2316,14 @@ mod tests { ]) .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"); + let tree: TestTree = BaseTreeRegressor::fit_inner( + &x, + &y, + None, + &MSECriterion, + params(Some(2), 1, 2, Splitter::Best), + ) + .expect("Fit should work"); assert_eq!(tree.nodes().len(), 3); // Split at 1.5: SSE(root) = 17, SSE(true) = SSE(false) = 0.5, root mass = 4. @@ -2121,6 +2354,7 @@ mod tests { &x, &y, sample_weights, + &MSECriterion, params(None, 1, 2, splitter), ) .expect("Fit should work"); @@ -2152,11 +2386,16 @@ mod tests { for splitter in [Splitter::Best, Splitter::Random] { for max_depth in [Some(2), None] { let parameters = params(max_depth, 1, 2, splitter); - let tree: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, Some(&weights), parameters.clone()) - .expect("Fit should work"); + let tree: TestTree = BaseTreeRegressor::fit_inner( + &x, + &y, + Some(&weights), + &MSECriterion, + parameters.clone(), + ) + .expect("Fit should work"); let tree_scaled: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, Some(&scaled), parameters) + BaseTreeRegressor::fit_inner(&x, &y, Some(&scaled), &MSECriterion, parameters) .expect("Fit should work"); assert_same_tree(&tree, &tree_scaled); diff --git a/src/tree/criterion.rs b/src/tree/criterion.rs index bfd918e6..d75fa67a 100644 --- a/src/tree/criterion.rs +++ b/src/tree/criterion.rs @@ -7,6 +7,9 @@ //! * ["Classification and regression trees", Breiman L, Friedman J H, Olshen R A, and Stone C J, 1984](https://www.taylorfrancis.com/books/mono/10.1201/9781315139470/classification-regression-trees-leo-breiman) //! * scikit-learn, `sklearn/tree/_criterion.pyx` (proxy improvement) +use crate::numbers::basenum::Number; +use std::fmt::Debug; + /// Impurity measure that a tree uses to select splits. pub(crate) trait Criterion { /// Sufficient statistics for a set of samples. @@ -14,7 +17,7 @@ pub(crate) trait Criterion { /// Target value of one sample. type Target: Copy; /// Value of a leaf. - type Output: Clone; + type Output: Clone + Debug; /// Statistics for an empty set. Use only this function to make new statistics. fn empty_stats(&self) -> Self::Stats; @@ -43,6 +46,12 @@ pub(crate) trait Criterion { /// Value of a leaf that contains the samples of `stats`. fn leaf_value(&self, stats: &Self::Stats) -> Self::Output; + + /// Are the outputs approximately equal? + fn output_approximate_eq(out: &Self::Output, other_out: &Self::Output) -> bool; + + /// Convert a target value of the training data to the target of this criterion. + fn target(&self, y: T) -> Self::Target; } /// Statistics for the weighted mean of the targets. @@ -53,6 +62,7 @@ pub(crate) struct MeanStats { } /// Mean squared error criterion for regression. The leaf value is the weighted mean. +#[derive(Debug)] pub(crate) struct MSECriterion; impl Criterion for MSECriterion { @@ -113,6 +123,14 @@ impl Criterion for MSECriterion { 0.0 } } + + fn output_approximate_eq(out: &Self::Output, other_out: &Self::Output) -> bool { + (out - other_out).abs() < f64::EPSILON + } + + fn target(&self, y: T) -> Self::Target { + y.to_f64().unwrap() + } } #[cfg(test)] @@ -370,4 +388,34 @@ mod tests { ); } } + + #[test] + fn mse_target_converts_numbers_to_f64() { + let criterion = MSECriterion; + assert_eq!(criterion.target(3.25_f64), 3.25); + assert_eq!(criterion.target(-1.5_f32), -1.5); + assert_eq!(criterion.target(-7_i32), -7.0); + assert_eq!(criterion.target(200_u8), 200.0); + assert_eq!(criterion.target(42_usize), 42.0); + // Largest integer that f64 represents exactly + assert_eq!( + criterion.target(9_007_199_254_740_992_i64), + 9_007_199_254_740_992.0 + ); + } + + #[test] + fn mse_output_approximate_eq() { + assert!(MSECriterion::output_approximate_eq(&1.5, &1.5)); + assert!(MSECriterion::output_approximate_eq(&0.0, &-0.0)); + // Difference smaller than EPSILON + assert!(MSECriterion::output_approximate_eq( + &0.0, + &(f64::EPSILON / 2.0) + )); + // Difference of exactly EPSILON is not "approximately equal" (strict comparison) + assert!(!MSECriterion::output_approximate_eq(&0.0, &f64::EPSILON)); + assert!(!MSECriterion::output_approximate_eq(&1.0, &1.1)); + assert!(!MSECriterion::output_approximate_eq(&f64::NAN, &f64::NAN)); + } } diff --git a/src/tree/decision_tree_regressor.rs b/src/tree/decision_tree_regressor.rs index 46126cc1..9fa5abb8 100644 --- a/src/tree/decision_tree_regressor.rs +++ b/src/tree/decision_tree_regressor.rs @@ -69,6 +69,7 @@ use crate::api::{Predictor, SupervisedEstimator}; use crate::error::Failed; use crate::linalg::basic::arrays::{Array1, Array2}; use crate::numbers::basenum::Number; +use crate::tree::criterion::MSECriterion; /// Validates the sample weights fn validate_sample_weights(sample_weights: &[f64], n_rows: usize) -> Result<(), Failed> { @@ -112,7 +113,7 @@ pub struct DecisionTreeRegressorParameters { #[derive(Debug)] pub struct DecisionTreeRegressor, Y: Array1> { - tree_regressor: Option>, + tree_regressor: Option>, num_features: Option, } @@ -351,7 +352,8 @@ impl, Y: Array1> seed: parameters.seed, splitter: Splitter::Best, }; - let tree = BaseTreeRegressor::fit_inner(x, y, sample_weights, tree_parameters)?; + let tree = + BaseTreeRegressor::fit_inner(x, y, sample_weights, &MSECriterion, tree_parameters)?; Ok(Self { tree_regressor: Some(tree), num_features: Some(x.shape().1), From 4cef0b0f8e5bc4adbd1bc3708b858866fc90a650 Mon Sep 17 00:00:00 2001 From: slievens <5492873+slievens@users.noreply.github.com> Date: Thu, 22 Oct 2026 20:31:30 +0200 Subject: [PATCH 11/18] Rename BaseTreeRegressor to BaseTree --- src/ensemble/base_forest_regressor.rs | 25 +- src/ensemble/extra_trees_regressor.rs | 2 +- src/ensemble/random_forest_regressor.rs | 2 +- .../{base_tree_regressor.rs => base_tree.rs} | 290 +++++++----------- src/tree/decision_tree_regressor.rs | 9 +- src/tree/mod.rs | 2 +- 6 files changed, 134 insertions(+), 196 deletions(-) rename src/tree/{base_tree_regressor.rs => base_tree.rs} (91%) diff --git a/src/ensemble/base_forest_regressor.rs b/src/ensemble/base_forest_regressor.rs index a30865ae..c4a49986 100644 --- a/src/ensemble/base_forest_regressor.rs +++ b/src/ensemble/base_forest_regressor.rs @@ -11,7 +11,7 @@ use crate::numbers::basenum::Number; use crate::numbers::floatnum::FloatNumber; use crate::rand_custom::get_rng_impl; -use crate::tree::base_tree_regressor::{BaseTreeRegressor, BaseTreeRegressorParameters, Splitter}; +use crate::tree::base_tree::{BaseTree, BaseTreeParameters, Splitter}; use crate::tree::criterion::MSECriterion; #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] @@ -71,7 +71,7 @@ pub struct BaseForestRegressor< X: Array2, Y: Array1, > { - trees: Option>>, + trees: Option>>, num_features: Option, samples: Option>>, } @@ -107,8 +107,7 @@ impl, Y: Array1 let mut rng = get_rng_impl(Some(parameters.seed)); let n_trees = parameters.n_trees; - let mut trees: Vec> = - Vec::with_capacity(n_trees); + let mut trees: Vec> = Vec::with_capacity(n_trees); let mut maybe_all_samples: Option>> = Option::None; if parameters.keep_samples { @@ -146,7 +145,7 @@ impl, Y: Array1 all_samples.push(samples.iter().map(|x| *x != 0).collect()) } - let params = BaseTreeRegressorParameters { + let params = BaseTreeParameters { max_depth: parameters.max_depth, min_samples_leaf: parameters.min_samples_leaf, min_samples_split: parameters.min_samples_split, @@ -159,7 +158,7 @@ impl, Y: Array1 } else { sample_weights }; - let tree = BaseTreeRegressor::fit_weak_learner( + let tree = BaseTree::fit_weak_learner( x, y, sample_weights_for_base_tree, @@ -330,7 +329,7 @@ mod tests { keep_samples: true, seed: 42, bootstrap: true, - splitter: crate::tree::base_tree_regressor::Splitter::Best, + splitter: crate::tree::base_tree::Splitter::Best, }; let regressor = BaseForestRegressor::fit(&x, &y, None, params).unwrap(); assert_eq!(regressor.trees.unwrap().len(), 5); @@ -358,7 +357,7 @@ mod tests { keep_samples: false, seed: 0, bootstrap: true, - splitter: crate::tree::base_tree_regressor::Splitter::Best, + splitter: crate::tree::base_tree::Splitter::Best, }, ); assert!(result.is_err()); @@ -386,7 +385,7 @@ mod tests { keep_samples: false, seed: 0, bootstrap: true, - splitter: crate::tree::base_tree_regressor::Splitter::Best, + splitter: crate::tree::base_tree::Splitter::Best, }, ); assert!(result.is_err()); @@ -412,7 +411,7 @@ mod tests { keep_samples: true, seed: 42, bootstrap: false, - splitter: crate::tree::base_tree_regressor::Splitter::Best, + splitter: crate::tree::base_tree::Splitter::Best, }; let forest = BaseForestRegressor::fit(&x, &y, None, forest_parameters.clone()) @@ -466,7 +465,7 @@ mod tests { keep_samples: true, // keep samples used for each tree, so we can check that they are different seed: 42, bootstrap: true, // Use bootstrapping - splitter: crate::tree::base_tree_regressor::Splitter::Best, + splitter: crate::tree::base_tree::Splitter::Best, }; let regressor = BaseForestRegressor::fit(&x, &y, weights, params).unwrap(); let samples = regressor.samples.unwrap(); @@ -541,7 +540,7 @@ mod tests { keep_samples: true, seed: 42, bootstrap: false, // No bootstrapping - splitter: crate::tree::base_tree_regressor::Splitter::Best, + splitter: crate::tree::base_tree::Splitter::Best, }; let forest = BaseForestRegressor::fit(&x, &y, Some(&sample_weights), parameters.clone()) @@ -577,7 +576,7 @@ mod tests { keep_samples: true, // keep samples used for each tree, so we can check that they are different seed: 42, bootstrap: true, // Use bootstrapping - splitter: crate::tree::base_tree_regressor::Splitter::Best, + splitter: crate::tree::base_tree::Splitter::Best, }; let forest = BaseForestRegressor::fit(&x, &y, Some(&sample_weights), parameters.clone()) diff --git a/src/ensemble/extra_trees_regressor.rs b/src/ensemble/extra_trees_regressor.rs index 67a420fe..6ac6b849 100644 --- a/src/ensemble/extra_trees_regressor.rs +++ b/src/ensemble/extra_trees_regressor.rs @@ -64,7 +64,7 @@ use crate::error::Failed; use crate::linalg::basic::arrays::{Array1, Array2}; use crate::numbers::basenum::Number; use crate::numbers::floatnum::FloatNumber; -use crate::tree::base_tree_regressor::Splitter; +use crate::tree::base_tree::Splitter; /// Validates the sample weights fn validate_sample_weights(sample_weights: &[f64], n_rows: usize) -> Result<(), Failed> { diff --git a/src/ensemble/random_forest_regressor.rs b/src/ensemble/random_forest_regressor.rs index b2978e07..3b9a04e4 100644 --- a/src/ensemble/random_forest_regressor.rs +++ b/src/ensemble/random_forest_regressor.rs @@ -55,7 +55,7 @@ use crate::error::Failed; use crate::linalg::basic::arrays::{Array1, Array2}; use crate::numbers::basenum::Number; use crate::numbers::floatnum::FloatNumber; -use crate::tree::base_tree_regressor::Splitter; +use crate::tree::base_tree::Splitter; /// Validates the sample weights fn validate_sample_weights(sample_weights: &[f64], n_rows: usize) -> Result<(), Failed> { diff --git a/src/tree/base_tree_regressor.rs b/src/tree/base_tree.rs similarity index 91% rename from src/tree/base_tree_regressor.rs rename to src/tree/base_tree.rs index 2a067b85..6f0fb07d 100644 --- a/src/tree/base_tree_regressor.rs +++ b/src/tree/base_tree.rs @@ -23,11 +23,11 @@ pub enum Splitter { #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[derive(Debug, Clone)] -/// Parameters of Regression base_tree +/// Parameters of base tree #[must_use] -pub struct BaseTreeRegressorParameters { +pub struct BaseTreeParameters { #[cfg_attr(feature = "serde", serde(default))] - /// The maximum depth of the base_tree. + /// The maximum depth of the tree pub max_depth: Option, #[cfg_attr(feature = "serde", serde(default))] /// The minimum number of samples required to be at a leaf node. @@ -43,7 +43,7 @@ pub struct BaseTreeRegressorParameters { pub splitter: Splitter, } -/// Regression base_tree +/// Base Decision Tree #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[cfg_attr( feature = "serde", @@ -53,15 +53,10 @@ pub struct BaseTreeRegressorParameters { )) )] #[derive(Debug)] -pub struct BaseTreeRegressor< - TX: Number + PartialOrd, - TY: Number, - X: Array2, - Y: Array1, - C: Criterion, -> { +pub struct BaseTree, Y: Array1, C: Criterion> +{ nodes: Vec>, - parameters: Option, + parameters: Option, depth: u16, root_mass: f64, _phantom_tx: PhantomData, @@ -71,14 +66,14 @@ pub struct BaseTreeRegressor< } impl, Y: Array1, C: Criterion> - BaseTreeRegressor + BaseTree { /// Get nodes, return a shared reference fn nodes(&self) -> &Vec> { self.nodes.as_ref() } /// Get parameters, return a shared reference - fn parameters(&self) -> &BaseTreeRegressorParameters { + fn parameters(&self) -> &BaseTreeParameters { self.parameters.as_ref().unwrap() } } @@ -125,7 +120,7 @@ impl Node { } impl, Y: Array1, C: Criterion> PartialEq - for BaseTreeRegressor + for BaseTree { fn eq(&self, other: &Self) -> bool { if self.depth != other.depth || self.nodes.len() != other.nodes.len() { @@ -374,15 +369,15 @@ where } impl, Y: Array1, C: Criterion> - BaseTreeRegressor + BaseTree { pub(crate) fn fit_inner( x: &X, y: &Y, sample_weights: Option<&[f64]>, criterion: &C, - parameters: BaseTreeRegressorParameters, - ) -> Result, Failed> { + parameters: BaseTreeParameters, + ) -> Result, Failed> { let (x_nrows, num_attributes) = x.shape(); if x_nrows != y.shape() { return Err(Failed::fit("Size of x should equal size of y")); @@ -403,7 +398,7 @@ impl, Y: Array1, C: Crite } let samples = vec![1; x_nrows]; - BaseTreeRegressor::fit_weak_learner( + BaseTree::fit_weak_learner( x, y, sample_weights, @@ -423,8 +418,8 @@ impl, Y: Array1, C: Crite mtry: usize, order: &[Vec], criterion: &C, - parameters: BaseTreeRegressorParameters, - ) -> Result, Failed> { + parameters: BaseTreeParameters, + ) -> Result, Failed> { check_row_count(y.shape())?; match sample_weights { Some(_) => Self::grow::( @@ -456,8 +451,8 @@ impl, Y: Array1, C: Crite mtry: usize, order: &[Vec], criterion: &C, - parameters: BaseTreeRegressorParameters, - ) -> Result, Failed> { + parameters: BaseTreeParameters, + ) -> Result, Failed> { let mut nodes: Vec> = Vec::new(); let mut rng = get_rng_impl(parameters.seed); @@ -517,7 +512,7 @@ impl, Y: Array1, C: Crite criterion, ); - let mut base_tree = BaseTreeRegressor { + let mut base_tree = BaseTree { nodes, parameters: Some(parameters), depth: 0u16, @@ -941,7 +936,7 @@ impl, Y: Array1, C: Crite } // Regression: the leaf holds the prediction itself. -impl BaseTreeRegressor +impl BaseTree where TX: Number + PartialOrd, TY: Number, @@ -997,12 +992,12 @@ mod tests { assert_eq!(empty.shape(), (0, 2)); let y: Vec = vec![]; - let result = BaseTreeRegressor::fit_inner( + let result = BaseTree::fit_inner( &empty, &y, None, &MSECriterion, - BaseTreeRegressorParameters { + BaseTreeParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 2, @@ -1022,12 +1017,12 @@ mod tests { assert_eq!(no_features.shape(), (2, 0)); let y = vec![1.0_f64, 2.0]; - let result = BaseTreeRegressor::fit_inner( + let result = BaseTree::fit_inner( &no_features, &y, None, &MSECriterion, - BaseTreeRegressorParameters { + BaseTreeParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 2, @@ -1049,12 +1044,12 @@ mod tests { let y = vec![1.0, 2.0, 3.0]; let sample_weights = vec![5.0, 6.0, 7.0]; - let result = BaseTreeRegressor::fit_inner( + let result = BaseTree::fit_inner( &x, &y, Some(&sample_weights), &MSECriterion, - BaseTreeRegressorParameters { + BaseTreeParameters { max_depth: Some(0), min_samples_leaf: 1, min_samples_split: 2, @@ -1079,7 +1074,7 @@ mod tests { // Test that using no weights is equivalent to using uniform weights let x_rand: DenseMatrix = DenseMatrix::::rand(17, 5); let y_rand: Vec = (0..17).collect::>().map(|y| *y as f64); - let parameters = BaseTreeRegressorParameters { + let parameters = BaseTreeParameters { max_depth: Some(5), min_samples_leaf: 1, min_samples_split: 1, @@ -1088,10 +1083,10 @@ mod tests { }; let tree_no_weights = - BaseTreeRegressor::fit_inner(&x_rand, &y_rand, None, &MSECriterion, parameters.clone()) + BaseTree::fit_inner(&x_rand, &y_rand, None, &MSECriterion, parameters.clone()) .expect("Fit should work"); let uniform_weights = vec![1.0f64; 17]; - let tree_with_weights = BaseTreeRegressor::fit_inner( + let tree_with_weights = BaseTree::fit_inner( &x_rand, &y_rand, Some(&uniform_weights), @@ -1128,7 +1123,7 @@ mod tests { .unwrap(); let y_repeated = vec![4.0, 5.0, 5.0, 6.0, 6.0, 6.0]; - let weighted_parameters = BaseTreeRegressorParameters { + let weighted_parameters = BaseTreeParameters { max_depth: Some(1), // tree should not be able to fully separate all the examples min_samples_leaf: 1, min_samples_split: 1, @@ -1136,7 +1131,7 @@ mod tests { splitter: Splitter::Best, }; - let repeated_parameters = BaseTreeRegressorParameters { + let repeated_parameters = BaseTreeParameters { max_depth: Some(1), min_samples_leaf: 1, min_samples_split: 1, @@ -1144,7 +1139,7 @@ mod tests { splitter: Splitter::Best, }; - let tree_weighted = BaseTreeRegressor::fit_inner( + let tree_weighted = BaseTree::fit_inner( &x, &y, Some(&sample_weights), @@ -1152,7 +1147,7 @@ mod tests { weighted_parameters, ) .expect("Fit should work"); - let tree_repeated = BaseTreeRegressor::fit_inner( + let tree_repeated = BaseTree::fit_inner( &x_repeated, &y_repeated, None, @@ -1180,7 +1175,7 @@ mod tests { .unwrap(); let y = vec![1.0f64, 2.0, 6.0, 7.0, 11., 12.]; - let parameters = BaseTreeRegressorParameters { + let parameters = BaseTreeParameters { max_depth: Some(3), min_samples_leaf: 1, min_samples_split: 2, @@ -1188,8 +1183,8 @@ mod tests { splitter: Splitter::Best, }; - let tree = BaseTreeRegressor::fit_inner(&x, &y, None, &MSECriterion, parameters) - .expect("Fit should work"); + let tree = + BaseTree::fit_inner(&x, &y, None, &MSECriterion, parameters).expect("Fit should work"); let y_expected = vec![1.0, 2.0, 6.5, 6.5, 11.50, 11.50]; let y_hat = tree.predict(&x).expect("Predict should work"); assert_eq!(tree.nodes().len(), 7); @@ -1209,7 +1204,7 @@ mod tests { ]) .unwrap(); let y = vec![1.0f64, 2.0, 6.0, 7.0, 11., 12.]; - let parameters = BaseTreeRegressorParameters { + let parameters = BaseTreeParameters { max_depth: Some(3), min_samples_leaf: 1, min_samples_split: 2, @@ -1218,9 +1213,8 @@ mod tests { }; let sample_weights = [1.0f64, 1.0, 1.0, 1.0, 10.0, 10.0]; - let tree = - BaseTreeRegressor::fit_inner(&x, &y, Some(&sample_weights), &MSECriterion, parameters) - .expect("Fit should work"); + let tree = BaseTree::fit_inner(&x, &y, Some(&sample_weights), &MSECriterion, parameters) + .expect("Fit should work"); let x_test: DenseMatrix = DenseMatrix::from_iterator((1..=12).map(|i| i as f64), 12, 1, 0); let y_expected = vec![ @@ -1238,7 +1232,7 @@ mod tests { let x = DenseMatrix::from_2d_vec(&vec![vec![1.0_f64], vec![2.0], vec![3.0]]).unwrap(); let y = vec![1.0f64, 2.0, 6.0]; - let parameters = BaseTreeRegressorParameters { + let parameters = BaseTreeParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 3, @@ -1246,13 +1240,13 @@ mod tests { splitter: Splitter::Best, }; - let tree = BaseTreeRegressor::fit_inner(&x, &y, None, &MSECriterion, parameters) - .expect("Fit should work"); + let tree = + BaseTree::fit_inner(&x, &y, None, &MSECriterion, parameters).expect("Fit should work"); assert_eq!(tree.nodes().len(), 3); assert_eq!(tree.depth, 2); // A node with fewer than `min_samples_split` samples must stay a leaf. - let parameters = BaseTreeRegressorParameters { + let parameters = BaseTreeParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 4, @@ -1260,8 +1254,8 @@ mod tests { splitter: Splitter::Best, }; - let tree = BaseTreeRegressor::fit_inner(&x, &y, None, &MSECriterion, parameters) - .expect("Fit should work"); + let tree = + BaseTree::fit_inner(&x, &y, None, &MSECriterion, parameters).expect("Fit should work"); assert_eq!(tree.nodes().len(), 1); assert_eq!(tree.depth, 0); } @@ -1274,7 +1268,7 @@ mod tests { let y = vec![1.0f64, 2.0, 3.0]; let sample_weights = [0.0f64, 1.0, 1.0]; - let parameters = BaseTreeRegressorParameters { + let parameters = BaseTreeParameters { max_depth: Some(2), min_samples_leaf: 1, min_samples_split: 2, @@ -1282,9 +1276,8 @@ mod tests { splitter: Splitter::Best, }; - let tree = - BaseTreeRegressor::fit_inner(&x, &y, Some(&sample_weights), &MSECriterion, parameters) - .expect("Fit should work"); + let tree = BaseTree::fit_inner(&x, &y, Some(&sample_weights), &MSECriterion, parameters) + .expect("Fit should work"); assert!(tree.nodes().iter().all(|node| node.output.is_finite())); assert!( @@ -1310,7 +1303,7 @@ mod tests { let y = vec![1.0f64, 2.0, 3.0]; let sample_weights = [1.0f64, 1.0, 0.0]; - let parameters = BaseTreeRegressorParameters { + let parameters = BaseTreeParameters { max_depth: Some(2), min_samples_leaf: 1, min_samples_split: 2, @@ -1318,9 +1311,8 @@ mod tests { splitter: Splitter::Best, }; - let tree = - BaseTreeRegressor::fit_inner(&x, &y, Some(&sample_weights), &MSECriterion, parameters) - .expect("Fit should work"); + let tree = BaseTree::fit_inner(&x, &y, Some(&sample_weights), &MSECriterion, parameters) + .expect("Fit should work"); assert!(tree.nodes().iter().all(|node| node.output.is_finite())); assert!( @@ -1342,7 +1334,7 @@ mod tests { let y = vec![0.0f64, 10.0, 20.0, 30.0]; let sample_weights = [3.0f64, 1.0, 1.0, 3.0]; - let parameters = BaseTreeRegressorParameters { + let parameters = BaseTreeParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 2, @@ -1350,7 +1342,7 @@ mod tests { splitter: Splitter::Best, }; - let tree = BaseTreeRegressor::fit_inner( + let tree = BaseTree::fit_inner( &x, &y, Some(&sample_weights), @@ -1381,7 +1373,7 @@ mod tests { .unwrap(); let y_repeated = vec![0.0f64, 0.0, 0.0, 10.0, 20.0, 30.0, 30.0, 30.0]; let tree_repeated = - BaseTreeRegressor::fit_inner(&x_repeated, &y_repeated, None, &MSECriterion, parameters) + BaseTree::fit_inner(&x_repeated, &y_repeated, None, &MSECriterion, parameters) .expect("Fit should work"); let y_hat_repeated = tree_repeated.predict(&x).expect("Predict should work"); assert!(mean_absolute_error(&y_hat, &y_hat_repeated) < 1e-9); @@ -1391,8 +1383,8 @@ mod tests { // but partitioning and predict compare them as f64 against an f64 threshold. const TWO_POW_53: i64 = 1 << 53; - fn split_value_parameters(splitter: Splitter) -> BaseTreeRegressorParameters { - BaseTreeRegressorParameters { + fn split_value_parameters(splitter: Splitter) -> BaseTreeParameters { + BaseTreeParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 2, @@ -1409,7 +1401,7 @@ mod tests { DenseMatrix::from_2d_vec(&vec![vec![TWO_POW_53 + 2], vec![TWO_POW_53 + 4]]).unwrap(); let y = vec![0.0f64, 10.0]; - let tree = BaseTreeRegressor::fit_inner( + let tree = BaseTree::fit_inner( &x, &y, None, @@ -1433,7 +1425,7 @@ mod tests { let x = DenseMatrix::from_2d_vec(&vec![vec![TWO_POW_53], vec![TWO_POW_53 + 1]]).unwrap(); let y = vec![0.0f64, 10.0]; - let tree = BaseTreeRegressor::fit_inner( + let tree = BaseTree::fit_inner( &x, &y, None, @@ -1452,7 +1444,7 @@ mod tests { let x = DenseMatrix::from_2d_vec(&vec![vec![TWO_POW_53], vec![TWO_POW_53 + 1]]).unwrap(); let y = vec![0.0f64, 10.0]; - let tree = BaseTreeRegressor::fit_inner( + let tree = BaseTree::fit_inner( &x, &y, None, @@ -1477,7 +1469,7 @@ mod tests { let x = DenseMatrix::from_2d_vec(&vec![vec![a], vec![b]]).unwrap(); let y = vec![0.0f64, 10.0]; - let tree = BaseTreeRegressor::fit_inner( + let tree = BaseTree::fit_inner( &x, &y, None, @@ -1501,15 +1493,15 @@ mod tests { assert_eq!(std::mem::size_of::(), 16); } - type TestTree = BaseTreeRegressor, Vec, MSECriterion>; + type TestTree = BaseTree, Vec, MSECriterion>; fn params( max_depth: Option, min_samples_leaf: usize, min_samples_split: usize, splitter: Splitter, - ) -> BaseTreeRegressorParameters { - BaseTreeRegressorParameters { + ) -> BaseTreeParameters { + BaseTreeParameters { max_depth, min_samples_leaf, min_samples_split, @@ -1621,7 +1613,7 @@ mod tests { (None, 15, 4), ]; for (max_depth, n_nodes, depth) in cases { - let tree: TestTree = BaseTreeRegressor::fit_inner( + let tree: TestTree = BaseTree::fit_inner( &x, &y, None, @@ -1634,7 +1626,7 @@ mod tests { } // With max_depth = 2 the root split separates the two halves. - let tree: TestTree = BaseTreeRegressor::fit_inner( + let tree: TestTree = BaseTree::fit_inner( &x, &y, None, @@ -1647,7 +1639,7 @@ mod tests { assert_eq!(y_hat, vec![5.5, 5.5, 5.5, 5.5, 105.5, 105.5, 105.5, 105.5]); // No limit is the same as the largest limit, and the leaves memorize y. - let tree_none: TestTree = BaseTreeRegressor::fit_inner( + let tree_none: TestTree = BaseTree::fit_inner( &x, &y, None, @@ -1655,7 +1647,7 @@ mod tests { params(None, 1, 2, Splitter::Best), ) .expect("Fit should work"); - let tree_max: TestTree = BaseTreeRegressor::fit_inner( + let tree_max: TestTree = BaseTree::fit_inner( &x, &y, None, @@ -1667,7 +1659,7 @@ mod tests { assert_eq!(tree_none.predict(&x).expect("Predict should work"), y); for max_depth in 1..=4 { - let tree: TestTree = BaseTreeRegressor::fit_inner( + let tree: TestTree = BaseTree::fit_inner( &x, &y, None, @@ -1693,7 +1685,7 @@ mod tests { .unwrap(); let y = vec![0.0f64, 1.0, 2.0, 100.0, 101.0]; - let tree: TestTree = BaseTreeRegressor::fit_inner( + let tree: TestTree = BaseTree::fit_inner( &x, &y, None, @@ -1711,14 +1703,9 @@ mod tests { // If min_samples_split is larger than the number of samples, the root is a leaf. for splitter in [Splitter::Best, Splitter::Random] { - let tree: TestTree = BaseTreeRegressor::fit_inner( - &x, - &y, - None, - &MSECriterion, - params(None, 1, 6, splitter), - ) - .expect("Fit should work"); + let tree: TestTree = + BaseTree::fit_inner(&x, &y, None, &MSECriterion, params(None, 1, 6, splitter)) + .expect("Fit should work"); assert_eq!(tree.nodes().len(), 1, "splitter: {splitter:?}"); } } @@ -1730,7 +1717,7 @@ mod tests { let y = vec![0.0f64, 0.0, 0.0, 100.0]; // With no limit, the best split isolates the last sample. - let tree: TestTree = BaseTreeRegressor::fit_inner( + let tree: TestTree = BaseTree::fit_inner( &x, &y, None, @@ -1741,7 +1728,7 @@ mod tests { assert_eq!(tree.nodes()[0].split_value, Some(3.5)); // The best split is not permitted. The split with exactly 2 samples on each side is used. - let tree: TestTree = BaseTreeRegressor::fit_inner( + let tree: TestTree = BaseTree::fit_inner( &x, &y, None, @@ -1755,14 +1742,9 @@ mod tests { // No split can give 3 samples on each side. for splitter in [Splitter::Best, Splitter::Random] { - let tree: TestTree = BaseTreeRegressor::fit_inner( - &x, - &y, - None, - &MSECriterion, - params(None, 3, 2, splitter), - ) - .expect("Fit should work"); + let tree: TestTree = + BaseTree::fit_inner(&x, &y, None, &MSECriterion, params(None, 3, 2, splitter)) + .expect("Fit should work"); assert_eq!(tree.nodes().len(), 1, "splitter: {splitter:?}"); assert_eq!(tree.nodes()[0].output, 25.0); } @@ -1793,11 +1775,11 @@ mod tests { for splitter in [Splitter::Best, Splitter::Random] { for min_samples_leaf in [1, 3, 5] { for seed in 0..5 { - let parameters = BaseTreeRegressorParameters { + let parameters = BaseTreeParameters { seed: Some(seed), ..params(None, min_samples_leaf, 2, splitter) }; - let tree: TestTree = BaseTreeRegressor::fit_weak_learner( + let tree: TestTree = BaseTree::fit_weak_learner( &x, &y, sample_weights, @@ -1837,14 +1819,9 @@ mod tests { let y = vec![1.0f64, 2.0, 3.0, 6.0]; for splitter in [Splitter::Best, Splitter::Random] { - let tree: TestTree = BaseTreeRegressor::fit_inner( - &x, - &y, - None, - &MSECriterion, - params(None, 1, 2, splitter), - ) - .expect("Fit should work"); + let tree: TestTree = + BaseTree::fit_inner(&x, &y, None, &MSECriterion, params(None, 1, 2, splitter)) + .expect("Fit should work"); assert_eq!(tree.nodes().len(), 1, "splitter: {splitter:?}"); assert_eq!(tree.depth, 0, "splitter: {splitter:?}"); let y_hat = tree.predict(&x).expect("Predict should work"); @@ -1865,14 +1842,9 @@ mod tests { let y = vec![0.0f64, 0.0, 10.0, 10.0]; for splitter in [Splitter::Best, Splitter::Random] { - let tree: TestTree = BaseTreeRegressor::fit_inner( - &x, - &y, - None, - &MSECriterion, - params(None, 1, 2, splitter), - ) - .expect("Fit should work"); + let tree: TestTree = + BaseTree::fit_inner(&x, &y, None, &MSECriterion, params(None, 1, 2, splitter)) + .expect("Fit should work"); assert!(tree.nodes().len() > 1, "splitter: {splitter:?}"); assert!( tree.nodes() @@ -1883,7 +1855,7 @@ mod tests { ); } - let tree: TestTree = BaseTreeRegressor::fit_inner( + let tree: TestTree = BaseTree::fit_inner( &x, &y, None, @@ -1909,20 +1881,15 @@ mod tests { let mut used_columns = [false; 2]; for seed in 0..10 { - let parameters = BaseTreeRegressorParameters { + let parameters = BaseTreeParameters { seed: Some(seed), ..params(None, 1, 2, Splitter::Best) }; - let tree_single: TestTree = BaseTreeRegressor::fit_inner( - &x_single, - &y, - None, - &MSECriterion, - parameters.clone(), - ) - .expect("Fit should work"); + let tree_single: TestTree = + BaseTree::fit_inner(&x_single, &y, None, &MSECriterion, parameters.clone()) + .expect("Fit should work"); let tree_double: TestTree = - BaseTreeRegressor::fit_inner(&x_double, &y, None, &MSECriterion, parameters) + BaseTree::fit_inner(&x_double, &y, None, &MSECriterion, parameters) .expect("Fit should work"); assert!(tree_double.nodes().len() > 1); @@ -1959,7 +1926,7 @@ mod tests { for splitter in [Splitter::Best, Splitter::Random] { for max_depth in [Some(2), None] { let parameters = params(max_depth, 1, 2, splitter); - let tree: TestTree = BaseTreeRegressor::fit_weak_learner( + let tree: TestTree = BaseTree::fit_weak_learner( &x, &y, None, @@ -1970,14 +1937,9 @@ mod tests { parameters.clone(), ) .expect("Fit should work"); - let tree_rep: TestTree = BaseTreeRegressor::fit_inner( - &x_rep, - &y_rep, - None, - &MSECriterion, - parameters, - ) - .expect("Fit should work"); + let tree_rep: TestTree = + BaseTree::fit_inner(&x_rep, &y_rep, None, &MSECriterion, parameters) + .expect("Fit should work"); assert!(tree.nodes().len() > 1); assert_same_tree(&tree, &tree_rep); @@ -2004,7 +1966,7 @@ mod tests { for splitter in [Splitter::Best, Splitter::Random] { for max_depth in [Some(2), None] { let parameters = params(max_depth, 1, 2, splitter); - let tree: TestTree = BaseTreeRegressor::fit_weak_learner( + let tree: TestTree = BaseTree::fit_weak_learner( &x, &y, Some(&weights), @@ -2016,14 +1978,9 @@ mod tests { ) .expect("Fit should work"); // in this tree, there are no counts but the rows are repeated physically - let tree_rep: TestTree = BaseTreeRegressor::fit_inner( - &x_rep, - &y_rep, - Some(&w_rep), - &MSECriterion, - parameters, - ) - .expect("Fit should work"); + let tree_rep: TestTree = + BaseTree::fit_inner(&x_rep, &y_rep, Some(&w_rep), &MSECriterion, parameters) + .expect("Fit should work"); assert!(tree.nodes().len() > 1); assert_same_tree(&tree, &tree_rep); @@ -2075,19 +2032,18 @@ mod tests { /// Fit one tree for each element type. Each tree comes with the mass of each row. fn split_score_trees( f: &SplitScoreFixture, - parameters: BaseTreeRegressorParameters, + parameters: BaseTreeParameters, ) -> Vec<(&'static str, TestTree, Vec)> { let n = f.y.len(); vec![ ( "unit", - BaseTreeRegressor::fit_inner(&f.x, &f.y, None, &MSECriterion, parameters.clone()) - .unwrap(), + BaseTree::fit_inner(&f.x, &f.y, None, &MSECriterion, parameters.clone()).unwrap(), vec![1.0; n], ), ( "weighted", - BaseTreeRegressor::fit_inner( + BaseTree::fit_inner( &f.x, &f.y, Some(&f.weights), @@ -2099,7 +2055,7 @@ mod tests { ), ( "counted", - BaseTreeRegressor::fit_weak_learner( + BaseTree::fit_weak_learner( &f.x, &f.y, None, @@ -2227,15 +2183,10 @@ mod tests { for splitter in [Splitter::Best, Splitter::Random] { let parameters = params(None, 1, 2, splitter); - let tree: TestTree = BaseTreeRegressor::fit_inner( - &x, - &y, - Some(&weights), - &MSECriterion, - parameters.clone(), - ) - .expect("Fit should work"); - let tree_kept: TestTree = BaseTreeRegressor::fit_inner( + let tree: TestTree = + BaseTree::fit_inner(&x, &y, Some(&weights), &MSECriterion, parameters.clone()) + .expect("Fit should work"); + let tree_kept: TestTree = BaseTree::fit_inner( &x_kept, &y_kept, Some(&weights_kept), @@ -2269,14 +2220,8 @@ mod tests { 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, - &MSECriterion, - params(Some(4), 1, 2, splitter), - ) - .expect("Fit should work") + BaseTree::fit_inner(&x, &y, None, &MSECriterion, params(Some(4), 1, 2, splitter)) + .expect("Fit should work") }; let reference = fit(0.0); for offset in [1e6, 1e8, 1e9] { @@ -2316,7 +2261,7 @@ mod tests { ]) .unwrap(); let y = vec![1.0, 2.0, 5.0, 6.0]; - let tree: TestTree = BaseTreeRegressor::fit_inner( + let tree: TestTree = BaseTree::fit_inner( &x, &y, None, @@ -2350,7 +2295,7 @@ mod tests { for splitter in [Splitter::Best, Splitter::Random] { for sample_weights in [None, Some(weights.as_slice())] { - let tree: TestTree = BaseTreeRegressor::fit_inner( + let tree: TestTree = BaseTree::fit_inner( &x, &y, sample_weights, @@ -2386,16 +2331,11 @@ mod tests { for splitter in [Splitter::Best, Splitter::Random] { for max_depth in [Some(2), None] { let parameters = params(max_depth, 1, 2, splitter); - let tree: TestTree = BaseTreeRegressor::fit_inner( - &x, - &y, - Some(&weights), - &MSECriterion, - parameters.clone(), - ) - .expect("Fit should work"); + let tree: TestTree = + BaseTree::fit_inner(&x, &y, Some(&weights), &MSECriterion, parameters.clone()) + .expect("Fit should work"); let tree_scaled: TestTree = - BaseTreeRegressor::fit_inner(&x, &y, Some(&scaled), &MSECriterion, parameters) + BaseTree::fit_inner(&x, &y, Some(&scaled), &MSECriterion, parameters) .expect("Fit should work"); assert_same_tree(&tree, &tree_scaled); diff --git a/src/tree/decision_tree_regressor.rs b/src/tree/decision_tree_regressor.rs index 9fa5abb8..41a7c014 100644 --- a/src/tree/decision_tree_regressor.rs +++ b/src/tree/decision_tree_regressor.rs @@ -64,7 +64,7 @@ use std::fmt::Debug; #[cfg(feature = "serde")] use serde::{Deserialize, Serialize}; -use super::base_tree_regressor::{BaseTreeRegressor, BaseTreeRegressorParameters, Splitter}; +use super::base_tree::{BaseTree, BaseTreeParameters, Splitter}; use crate::api::{Predictor, SupervisedEstimator}; use crate::error::Failed; use crate::linalg::basic::arrays::{Array1, Array2}; @@ -113,7 +113,7 @@ pub struct DecisionTreeRegressorParameters { #[derive(Debug)] pub struct DecisionTreeRegressor, Y: Array1> { - tree_regressor: Option>, + tree_regressor: Option>, num_features: Option, } @@ -345,15 +345,14 @@ impl, Y: Array1> sample_weights: Option<&[f64]>, parameters: DecisionTreeRegressorParameters, ) -> Result, Failed> { - let tree_parameters = BaseTreeRegressorParameters { + let tree_parameters = BaseTreeParameters { max_depth: parameters.max_depth, min_samples_leaf: parameters.min_samples_leaf, min_samples_split: parameters.min_samples_split, seed: parameters.seed, splitter: Splitter::Best, }; - let tree = - BaseTreeRegressor::fit_inner(x, y, sample_weights, &MSECriterion, tree_parameters)?; + let tree = BaseTree::fit_inner(x, y, sample_weights, &MSECriterion, tree_parameters)?; Ok(Self { tree_regressor: Some(tree), num_features: Some(x.shape().1), diff --git a/src/tree/mod.rs b/src/tree/mod.rs index 601d67cd..84a72d73 100644 --- a/src/tree/mod.rs +++ b/src/tree/mod.rs @@ -19,7 +19,7 @@ //! //! -pub(crate) mod base_tree_regressor; +pub(crate) mod base_tree; pub(crate) mod criterion; /// Classification tree for dependent variables that take a finite number of unordered values. From f2f29e623f665a82d25d0f7274ec8a6c8afe176d Mon Sep 17 00:00:00 2001 From: slievens <5492873+slievens@users.noreply.github.com> Date: Thu, 22 Oct 2026 21:03:56 +0200 Subject: [PATCH 12/18] Move regression predict out of BaseTree BaseTree only finds the leaf. DecisionTreeRegressor and BaseForestRegressor convert the leaf output to TY. --- src/ensemble/base_forest_regressor.rs | 8 +-- src/tree/base_tree.rs | 92 +++++++++++---------------- src/tree/decision_tree_regressor.rs | 12 +++- 3 files changed, 52 insertions(+), 60 deletions(-) diff --git a/src/ensemble/base_forest_regressor.rs b/src/ensemble/base_forest_regressor.rs index c4a49986..1b24febc 100644 --- a/src/ensemble/base_forest_regressor.rs +++ b/src/ensemble/base_forest_regressor.rs @@ -181,10 +181,10 @@ impl, Y: Array1 /// Predict class for `x` /// * `x` - _KxM_ data where _K_ is number of observations and _M_ is number of features. pub fn predict(&self, x: &X) -> Result { - let mut result = Y::zeros(x.shape().0); - let (n, _) = x.shape(); + let mut result = Y::zeros(n); + for i in 0..n { result.set(i, self.predict_for_row(x, i)); } @@ -198,7 +198,7 @@ impl, Y: Array1 let mut result = TY::zero(); for tree in self.trees.as_ref().unwrap().iter() { - result += tree.predict_for_row(x, row); + result += TY::from_f64(*tree.leaf_output_for_row(x, row)).unwrap(); } result / TY::from_usize(n_trees).unwrap() @@ -246,7 +246,7 @@ impl, Y: Array1 .zip(self.samples.as_ref().unwrap()) { if !samples[row] { - result += tree.predict_for_row(x, row); + result += TY::from_f64(*tree.leaf_output_for_row(x, row)).unwrap(); n_trees += 1; } } diff --git a/src/tree/base_tree.rs b/src/tree/base_tree.rs index 6f0fb07d..52bc1ada 100644 --- a/src/tree/base_tree.rs +++ b/src/tree/base_tree.rs @@ -935,34 +935,6 @@ impl, Y: Array1, C: Crite } } -// Regression: the leaf holds the prediction itself. -impl BaseTree -where - TX: Number + PartialOrd, - TY: Number, - X: Array2, - Y: Array1, - C: Criterion, -{ - pub(crate) fn predict_for_row(&self, x: &X, row: usize) -> TY { - TY::from_f64(*self.leaf_output_for_row(x, row)).unwrap() - } - - /// Predict regression value for `x`. - /// * `x` - _KxM_ data where _K_ is number of observations and _M_ is number of features. - pub fn predict(&self, x: &X) -> Result { - let mut result = Y::zeros(x.shape().0); - - let (n, _) = x.shape(); - - for i in 0..n { - result.set(i, self.predict_for_row(x, i)); - } - - Ok(result) - } -} - #[cfg(test)] mod tests { use super::*; @@ -1066,7 +1038,7 @@ mod tests { .sum::() / sample_weights.iter().sum::(); - assert!((tree.predict_for_row(&x, 0) - expected).abs() < 1e-9); + assert!((*tree.leaf_output_for_row(&x, 0) - expected).abs() < 1e-9); } #[test] @@ -1095,12 +1067,8 @@ mod tests { ) .expect("Fit should work"); - let y_pred_no_weights = tree_no_weights - .predict(&x_rand) - .expect("Predict should work"); - let y_pred_with_weights = tree_with_weights - .predict(&x_rand) - .expect("Predict should work"); + let y_pred_no_weights = predict(&tree_no_weights, &x_rand); + let y_pred_with_weights = predict(&tree_with_weights, &x_rand); assert!(mean_absolute_error(&y_pred_no_weights, &y_pred_with_weights) < 1e-9); } @@ -1156,8 +1124,8 @@ mod tests { ) .expect("Fit should work"); // Predict on the same data - let y_pred_weighted = tree_weighted.predict(&x).expect("Predict should work"); - let y_pred_repeated = tree_repeated.predict(&x).expect("Predict should work"); + let y_pred_weighted = predict(&tree_weighted, &x); + let y_pred_repeated = predict(&tree_repeated, &x); assert!(mean_absolute_error(&y_pred_weighted, &y_pred_repeated) < 1e-9); } @@ -1186,7 +1154,7 @@ mod tests { let tree = BaseTree::fit_inner(&x, &y, None, &MSECriterion, parameters).expect("Fit should work"); let y_expected = vec![1.0, 2.0, 6.5, 6.5, 11.50, 11.50]; - let y_hat = tree.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); assert_eq!(tree.nodes().len(), 7); assert_eq!(tree.depth, 3); assert!(mean_absolute_error(&y_expected, &y_hat) < 1e-9); @@ -1220,7 +1188,7 @@ mod tests { let y_expected = vec![ 1.50, 1.50, 6.50, 6.50, 11.0, 12.0, 12.0, 12.0, 12.0, 12.0, 12.0, 12.0, ]; - let y_hat = tree.predict(&x_test).expect("Predict should work"); + let y_hat = predict(&tree, &x_test); assert_eq!(tree.nodes().len(), 7); assert_eq!(tree.depth, 3); assert!(mean_absolute_error(&y_expected, &y_hat) < 1e-9); @@ -1286,7 +1254,7 @@ mod tests { .all(|node| node.split_score.is_none_or(f64::is_finite)) ); - let y_hat = tree.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); assert!(y_hat.iter().all(|v| v.is_finite())); // The best split separates x=3 from the other rows. @@ -1321,7 +1289,7 @@ mod tests { .all(|node| node.split_score.is_none_or(f64::is_finite)) ); - let y_hat = tree.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); assert!(y_hat.iter().all(|v| v.is_finite())); } @@ -1355,7 +1323,7 @@ mod tests { assert_eq!(tree.depth, 2); assert!((tree.nodes()[0].split_value.unwrap() - 1.5).abs() < 1e-9); // Split should be at 1.5 - let y_hat = tree.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); let y_expected = vec![2.5, 2.5, 27.5, 27.5]; // Expected values are weighted means assert!(mean_absolute_error(&y_expected, &y_hat) < 1e-9); @@ -1375,7 +1343,7 @@ mod tests { let tree_repeated = BaseTree::fit_inner(&x_repeated, &y_repeated, None, &MSECriterion, parameters) .expect("Fit should work"); - let y_hat_repeated = tree_repeated.predict(&x).expect("Predict should work"); + let y_hat_repeated = predict(&tree_repeated, &x); assert!(mean_absolute_error(&y_hat, &y_hat_repeated) < 1e-9); } @@ -1414,7 +1382,7 @@ mod tests { assert!((TWO_POW_53 + 2) as f64 <= threshold); assert!(((TWO_POW_53 + 4) as f64) > threshold); - let y_hat = tree.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); assert_eq!(y_hat, vec![0.0, 10.0]); } @@ -1434,7 +1402,7 @@ mod tests { ) .expect("Fit should work"); - let y_hat = tree.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); assert_eq!(y_hat, vec![5.0, 5.0]); } @@ -1453,7 +1421,7 @@ mod tests { ) .expect("Fit should work"); - let y_hat = tree.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); assert_eq!(y_hat, vec![5.0, 5.0]); } @@ -1482,7 +1450,7 @@ mod tests { assert!(a <= threshold); assert!(b > threshold); - let y_hat = tree.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); assert_eq!(y_hat, vec![0.0, 10.0]); } @@ -1495,6 +1463,20 @@ mod tests { type TestTree = BaseTree, Vec, MSECriterion>; + /// Leaf output of `tree` for each row of `x`. + fn predict(tree: &BaseTree, x: &X) -> Vec + where + TX: Number + PartialOrd, + TY: Number, + X: Array2, + Y: Array1, + C: Criterion, + { + (0..x.shape().0) + .map(|i| tree.leaf_output_for_row(x, i).clone()) + .collect() + } + fn params( max_depth: Option, min_samples_leaf: usize, @@ -1520,7 +1502,7 @@ mod tests { .collect() } - /// Index of the leaf that gets row `row` of `x`. Uses the same rule as `predict_for_row`. + /// Index of the leaf that gets row `row` of `x`. Uses the same rule as `leaf_output_for_row`. fn leaf_of(tree: &TestTree, x: &DenseMatrix, row: usize) -> usize { let mut node_id = 0; while let Some(true_child) = tree.nodes()[node_id].true_child { @@ -1635,7 +1617,7 @@ mod tests { ) .expect("Fit should work"); assert_eq!(tree.nodes()[0].split_value, Some(4.5)); - let y_hat = tree.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); assert_eq!(y_hat, vec![5.5, 5.5, 5.5, 5.5, 105.5, 105.5, 105.5, 105.5]); // No limit is the same as the largest limit, and the leaves memorize y. @@ -1656,7 +1638,7 @@ mod tests { ) .expect("Fit should work"); assert_same_tree(&tree_none, &tree_max); - assert_eq!(tree_none.predict(&x).expect("Predict should work"), y); + assert_eq!(predict(&tree_none, &x), y); for max_depth in 1..=4 { let tree: TestTree = BaseTree::fit_inner( @@ -1824,7 +1806,7 @@ mod tests { .expect("Fit should work"); assert_eq!(tree.nodes().len(), 1, "splitter: {splitter:?}"); assert_eq!(tree.depth, 0, "splitter: {splitter:?}"); - let y_hat = tree.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); assert_eq!(y_hat, vec![3.0; 4], "splitter: {splitter:?}"); } } @@ -1943,8 +1925,8 @@ mod tests { assert!(tree.nodes().len() > 1); assert_same_tree(&tree, &tree_rep); - let y_hat = tree.predict(&x).expect("Predict should work"); - let y_hat_rep = tree_rep.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); + let y_hat_rep = predict(&tree_rep, &x); assert!(mean_absolute_error(&y_hat, &y_hat_rep) < 1e-9); } } @@ -1984,8 +1966,8 @@ mod tests { assert!(tree.nodes().len() > 1); assert_same_tree(&tree, &tree_rep); - let y_hat = tree.predict(&x).expect("Predict should work"); - let y_hat_rep = tree_rep.predict(&x).expect("Predict should work"); + let y_hat = predict(&tree, &x); + let y_hat_rep = predict(&tree_rep, &x); assert!(mean_absolute_error(&y_hat, &y_hat_rep) < 1e-9); } } diff --git a/src/tree/decision_tree_regressor.rs b/src/tree/decision_tree_regressor.rs index 41a7c014..93290923 100644 --- a/src/tree/decision_tree_regressor.rs +++ b/src/tree/decision_tree_regressor.rs @@ -363,7 +363,17 @@ impl, Y: Array1> /// * `x` - _KxM_ data where _K_ is number of observations and _M_ is number of features. pub fn predict(&self, x: &X) -> Result { match &self.tree_regressor { - Some(tree) => tree.predict(x), + Some(tree) => { + let mut result = Y::zeros(x.shape().0); + + let (n, _) = x.shape(); + + for i in 0..n { + result.set(i, TY::from_f64(*tree.leaf_output_for_row(x, i)).unwrap()); + } + + Ok(result) + } None => Err(Failed::predict( "'fit' should be called before calling 'predict'", )), From 553d3bf64d0c929db49035b7eec0e19af96746b0 Mon Sep 17 00:00:00 2001 From: slievens <5492873+slievens@users.noreply.github.com> Date: Thu, 22 Oct 2026 21:40:56 +0200 Subject: [PATCH 13/18] Make BaseForestRegressor generic over the split criterion The forest passes a Criterion to its trees. The predict methods are restricted to Output = f64. The RF and ET regressors use MSECriterion. --- src/ensemble/base_forest_regressor.rs | 207 ++++++++++++++---------- src/ensemble/extra_trees_regressor.rs | 6 +- src/ensemble/random_forest_regressor.rs | 6 +- 3 files changed, 134 insertions(+), 85 deletions(-) diff --git a/src/ensemble/base_forest_regressor.rs b/src/ensemble/base_forest_regressor.rs index 1b24febc..bbf57591 100644 --- a/src/ensemble/base_forest_regressor.rs +++ b/src/ensemble/base_forest_regressor.rs @@ -12,7 +12,7 @@ use crate::numbers::floatnum::FloatNumber; use crate::rand_custom::get_rng_impl; use crate::tree::base_tree::{BaseTree, BaseTreeParameters, Splitter}; -use crate::tree::criterion::MSECriterion; +use crate::tree::criterion::Criterion; #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[derive(Debug, Clone)] @@ -47,8 +47,8 @@ pub struct BaseForestRegressorParameters { pub splitter: Splitter, } -impl, Y: Array1> PartialEq - for BaseForestRegressor +impl, Y: Array1, C: Criterion> + PartialEq for BaseForestRegressor { fn eq(&self, other: &Self) -> bool { if self.trees.as_ref().unwrap().len() != other.trees.as_ref().unwrap().len() { @@ -64,20 +64,28 @@ impl, Y: Array1 /// Forest Regressor #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +#[cfg_attr( + feature = "serde", + serde(bound( + serialize = "C::Output: Serialize", + deserialize = "C::Output: Deserialize<'de>" + )) +)] #[derive(Debug)] pub struct BaseForestRegressor< TX: Number + FloatNumber + PartialOrd, TY: Number, X: Array2, Y: Array1, + C: Criterion, > { - trees: Option>>, + trees: Option>>, num_features: Option, samples: Option>>, } -impl, Y: Array1> - BaseForestRegressor +impl, Y: Array1, C: Criterion> + BaseForestRegressor { /// Build a forest of trees from the training set. /// * `x` - _NxM_ matrix with _N_ observations and _M_ features in each observation. @@ -87,8 +95,9 @@ impl, Y: Array1 x: &X, y: &Y, sample_weights: Option<&[f64]>, + criterion: &C, parameters: BaseForestRegressorParameters, - ) -> Result, Failed> { + ) -> Result, Failed> { let (n_rows, num_attributes) = x.shape(); if n_rows != y.shape() { @@ -107,7 +116,7 @@ impl, Y: Array1 let mut rng = get_rng_impl(Some(parameters.seed)); let n_trees = parameters.n_trees; - let mut trees: Vec> = Vec::with_capacity(n_trees); + let mut trees: Vec> = Vec::with_capacity(n_trees); let mut maybe_all_samples: Option>> = Option::None; if parameters.keep_samples { @@ -133,7 +142,7 @@ impl, Y: Array1 for tree_idx in 0..parameters.n_trees { if parameters.bootstrap { - samples = BaseForestRegressor::::sample_with_replacement( + samples = BaseForestRegressor::::sample_with_replacement( n_rows, &mut rng, dist.as_ref(), @@ -165,7 +174,7 @@ impl, Y: Array1 samples.clone(), mtry, &order, - &MSECriterion, + criterion, params, )?; trees.push(tree); @@ -178,6 +187,68 @@ impl, Y: Array1 }) } + fn sample_with_replacement( + nrows: usize, + rng: &mut impl rand::Rng, + distribution: Option<&rand::distr::weighted::WeightedIndex>, + ) -> Vec { + let mut samples = vec![0; nrows]; + for _ in 0..nrows { + let xi = match distribution { + Some(dist) => rng.sample(dist), + None => rng.random_range(0..nrows), + }; + samples[xi] += 1; + } + + samples + } + + /// 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'", + )), + } + } +} + +impl< + TX: Number + FloatNumber + PartialOrd, + TY: Number, + X: Array2, + Y: Array1, + C: Criterion, +> BaseForestRegressor +{ /// Predict class for `x` /// * `x` - _KxM_ data where _K_ is number of observations and _M_ is number of features. pub fn predict(&self, x: &X) -> Result { @@ -254,59 +325,6 @@ impl, Y: Array1 // TODO: What to do if there are no oob trees? result / TY::from(n_trees).unwrap() } - - fn sample_with_replacement( - nrows: usize, - rng: &mut impl rand::Rng, - distribution: Option<&rand::distr::weighted::WeightedIndex>, - ) -> Vec { - let mut samples = vec![0; nrows]; - for _ in 0..nrows { - let xi = match distribution { - Some(dist) => rng.sample(dist), - None => rng.random_range(0..nrows), - }; - samples[xi] += 1; - } - - samples - } - - /// 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)] @@ -315,6 +333,7 @@ mod tests { use super::*; use crate::linalg::basic::arrays::Array; use crate::linalg::basic::matrix::DenseMatrix; + use crate::tree::criterion::MSECriterion; #[test] fn test_base_forest_regressor_keep_samples() { @@ -331,7 +350,7 @@ mod tests { bootstrap: true, splitter: crate::tree::base_tree::Splitter::Best, }; - let regressor = BaseForestRegressor::fit(&x, &y, None, params).unwrap(); + let regressor = BaseForestRegressor::fit(&x, &y, None, &MSECriterion, params).unwrap(); assert_eq!(regressor.trees.unwrap().len(), 5); assert!(regressor.samples.is_some()); } @@ -348,6 +367,7 @@ mod tests { &empty, &y, None, + &MSECriterion, BaseForestRegressorParameters { max_depth: None, min_samples_leaf: 1, @@ -376,6 +396,7 @@ mod tests { &no_features, &y, None, + &MSECriterion, BaseForestRegressorParameters { max_depth: None, min_samples_leaf: 1, @@ -414,8 +435,9 @@ mod tests { splitter: crate::tree::base_tree::Splitter::Best, }; - let forest = BaseForestRegressor::fit(&x, &y, None, forest_parameters.clone()) - .expect("Fit should work"); + let forest = + BaseForestRegressor::fit(&x, &y, None, &MSECriterion, forest_parameters.clone()) + .expect("Fit should work"); let y_hat = forest.predict(&x).expect("Predict should work"); assert!((y_hat.iter().sum::() - y.iter().sum::()).abs() < 1e-9); @@ -424,9 +446,14 @@ mod tests { // Positive weights in [0.5, 2.0) let sample_weights: Vec = (0..1000).map(|_| rng.random_range(0.5..2.0)).collect(); - let forest = - BaseForestRegressor::fit(&x, &y, Some(&sample_weights), forest_parameters.clone()) - .expect("Fit should work"); + let forest = BaseForestRegressor::fit( + &x, + &y, + Some(&sample_weights), + &MSECriterion, + forest_parameters.clone(), + ) + .expect("Fit should work"); let y_hat = forest.predict(&x).expect("Predict should work"); let s: f64 = sample_weights.iter().sum(); @@ -467,7 +494,8 @@ mod tests { bootstrap: true, // Use bootstrapping splitter: crate::tree::base_tree::Splitter::Best, }; - let regressor = BaseForestRegressor::fit(&x, &y, weights, params).unwrap(); + let regressor = + BaseForestRegressor::fit(&x, &y, weights, &MSECriterion, params).unwrap(); let samples = regressor.samples.unwrap(); for (t, in_bag) in samples.iter().enumerate() { @@ -513,7 +541,7 @@ mod tests { bootstrap: false, splitter: splitter.clone(), }; - let forest = BaseForestRegressor::fit(&x, &y, None, params).unwrap(); + let forest = BaseForestRegressor::fit(&x, &y, None, &MSECriterion, params).unwrap(); let trees = forest.trees.unwrap(); assert!( trees.iter().any(|tree| tree != &trees[0]), @@ -543,8 +571,14 @@ mod tests { splitter: crate::tree::base_tree::Splitter::Best, }; - let forest = BaseForestRegressor::fit(&x, &y, Some(&sample_weights), parameters.clone()) - .expect("Fit should work"); + let forest = BaseForestRegressor::fit( + &x, + &y, + Some(&sample_weights), + &MSECriterion, + parameters.clone(), + ) + .expect("Fit should work"); let y_hat = forest.predict(&x).expect("Predict should work"); // weighted mean is (10.0 * 0 + 90*10) / 100 = 9 @@ -556,7 +590,8 @@ mod tests { } // Without weights, the predicted value should be close to 5 - let forest = BaseForestRegressor::fit(&x, &y, None, parameters).expect("Fit should work"); + let forest = BaseForestRegressor::fit(&x, &y, None, &MSECriterion, parameters) + .expect("Fit should work"); let y_hat = forest.predict(&x).expect("Predict should work"); for p in y_hat.iter() { @@ -579,8 +614,14 @@ mod tests { splitter: crate::tree::base_tree::Splitter::Best, }; - let forest = BaseForestRegressor::fit(&x, &y, Some(&sample_weights), parameters.clone()) - .expect("Fit should work"); + let forest = BaseForestRegressor::fit( + &x, + &y, + Some(&sample_weights), + &MSECriterion, + parameters.clone(), + ) + .expect("Fit should work"); let y_hat = forest.predict(&x).expect("Predict should work"); // with bootstrapping the weight should be close to 9, but not extremely close @@ -592,7 +633,8 @@ mod tests { } // Without weights, the predicted value should be reasonably close to 5 - let forest = BaseForestRegressor::fit(&x, &y, None, parameters).expect("Fit should work"); + let forest = BaseForestRegressor::fit(&x, &y, None, &MSECriterion, parameters) + .expect("Fit should work"); let y_hat = forest.predict(&x).expect("Predict should work"); // The bootstrapping makes that we will not be very close to 5 @@ -632,10 +674,12 @@ mod tests { splitter: splitter.clone(), }; - let forest_a = BaseForestRegressor::fit(&x, &y, weights, params.clone()) - .expect("Fit should work"); - let forest_b = BaseForestRegressor::fit(&x, &y, weights, params.clone()) - .expect("Fit should work"); + let forest_a = + BaseForestRegressor::fit(&x, &y, weights, &MSECriterion, params.clone()) + .expect("Fit should work"); + let forest_b = + BaseForestRegressor::fit(&x, &y, weights, &MSECriterion, params.clone()) + .expect("Fit should work"); assert_eq!( forest_a, forest_b, @@ -656,6 +700,7 @@ mod tests { &x, &y, weights, + &MSECriterion, BaseForestRegressorParameters { seed: 43, ..params }, ) .expect("Fit should work"); diff --git a/src/ensemble/extra_trees_regressor.rs b/src/ensemble/extra_trees_regressor.rs index 6ac6b849..74fdbc9a 100644 --- a/src/ensemble/extra_trees_regressor.rs +++ b/src/ensemble/extra_trees_regressor.rs @@ -65,6 +65,7 @@ use crate::linalg::basic::arrays::{Array1, Array2}; use crate::numbers::basenum::Number; use crate::numbers::floatnum::FloatNumber; use crate::tree::base_tree::Splitter; +use crate::tree::criterion::MSECriterion; /// Validates the sample weights fn validate_sample_weights(sample_weights: &[f64], n_rows: usize) -> Result<(), Failed> { @@ -122,7 +123,7 @@ pub struct ExtraTreesRegressor< X: Array2, Y: Array1, > { - forest_regressor: Option>, + forest_regressor: Option>, } impl ExtraTreesRegressorParameters { @@ -245,7 +246,8 @@ impl, Y: Array1 bootstrap: false, splitter: Splitter::Random, }; - let forest_regressor = BaseForestRegressor::fit(x, y, sample_weights, regressor_params)?; + let forest_regressor = + BaseForestRegressor::fit(x, y, sample_weights, &MSECriterion, regressor_params)?; Ok(ExtraTreesRegressor { forest_regressor: Some(forest_regressor), diff --git a/src/ensemble/random_forest_regressor.rs b/src/ensemble/random_forest_regressor.rs index 3b9a04e4..014f1918 100644 --- a/src/ensemble/random_forest_regressor.rs +++ b/src/ensemble/random_forest_regressor.rs @@ -56,6 +56,7 @@ use crate::linalg::basic::arrays::{Array1, Array2}; use crate::numbers::basenum::Number; use crate::numbers::floatnum::FloatNumber; use crate::tree::base_tree::Splitter; +use crate::tree::criterion::MSECriterion; /// Validates the sample weights fn validate_sample_weights(sample_weights: &[f64], n_rows: usize) -> Result<(), Failed> { @@ -113,7 +114,7 @@ pub struct RandomForestRegressor< X: Array2, Y: Array1, > { - forest_regressor: Option>, + forest_regressor: Option>, } impl RandomForestRegressorParameters { @@ -438,7 +439,8 @@ impl, Y: Array1 bootstrap: true, splitter: Splitter::Best, }; - let forest_regressor = BaseForestRegressor::fit(x, y, sample_weights, regressor_params)?; + let forest_regressor = + BaseForestRegressor::fit(x, y, sample_weights, &MSECriterion, regressor_params)?; Ok(RandomForestRegressor { forest_regressor: Some(forest_regressor), From b8184a0ace874d7d6744e4d963916ab6e5104910 Mon Sep 17 00:00:00 2001 From: slievens <5492873+slievens@users.noreply.github.com> Date: Thu, 22 Oct 2026 21:51:21 +0200 Subject: [PATCH 14/18] Rename BaseForestRegressor to BaseForest --- ...ase_forest_regressor.rs => base_forest.rs} | 90 +++++++++---------- src/ensemble/extra_trees_regressor.rs | 8 +- src/ensemble/mod.rs | 2 +- src/ensemble/random_forest_regressor.rs | 8 +- 4 files changed, 52 insertions(+), 56 deletions(-) rename src/ensemble/{base_forest_regressor.rs => base_forest.rs} (90%) diff --git a/src/ensemble/base_forest_regressor.rs b/src/ensemble/base_forest.rs similarity index 90% rename from src/ensemble/base_forest_regressor.rs rename to src/ensemble/base_forest.rs index bbf57591..c75d6520 100644 --- a/src/ensemble/base_forest_regressor.rs +++ b/src/ensemble/base_forest.rs @@ -16,10 +16,10 @@ use crate::tree::criterion::Criterion; #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[derive(Debug, Clone)] -/// Parameters of the Forest Regressor +/// Parameters of a base forest. /// Some parameters here are passed directly into base estimator. #[must_use] -pub struct BaseForestRegressorParameters { +pub struct BaseForestParameters { #[cfg_attr(feature = "serde", serde(default))] /// Tree max depth. See [Decision Tree Regressor](../../tree/decision_tree_regressor/index.html) pub max_depth: Option, @@ -48,7 +48,7 @@ pub struct BaseForestRegressorParameters { } impl, Y: Array1, C: Criterion> - PartialEq for BaseForestRegressor + PartialEq for BaseForest { fn eq(&self, other: &Self) -> bool { if self.trees.as_ref().unwrap().len() != other.trees.as_ref().unwrap().len() { @@ -62,7 +62,7 @@ impl, Y: Array1 } } -/// Forest Regressor +/// Base forest of decision trees. The criterion `C` selects the splits and gives the leaf values. #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[cfg_attr( feature = "serde", @@ -72,7 +72,7 @@ impl, Y: Array1 )) )] #[derive(Debug)] -pub struct BaseForestRegressor< +pub struct BaseForest< TX: Number + FloatNumber + PartialOrd, TY: Number, X: Array2, @@ -85,7 +85,7 @@ pub struct BaseForestRegressor< } impl, Y: Array1, C: Criterion> - BaseForestRegressor + BaseForest { /// Build a forest of trees from the training set. /// * `x` - _NxM_ matrix with _N_ observations and _M_ features in each observation. @@ -96,8 +96,8 @@ impl, Y: Array1 y: &Y, sample_weights: Option<&[f64]>, criterion: &C, - parameters: BaseForestRegressorParameters, - ) -> Result, Failed> { + parameters: BaseForestParameters, + ) -> Result, Failed> { let (n_rows, num_attributes) = x.shape(); if n_rows != y.shape() { @@ -142,7 +142,7 @@ impl, Y: Array1 for tree_idx in 0..parameters.n_trees { if parameters.bootstrap { - samples = BaseForestRegressor::::sample_with_replacement( + samples = BaseForest::::sample_with_replacement( n_rows, &mut rng, dist.as_ref(), @@ -180,7 +180,7 @@ impl, Y: Array1 trees.push(tree); } - Ok(BaseForestRegressor { + Ok(BaseForest { trees: Some(trees), num_features: Some(x.shape().1), samples: maybe_all_samples, @@ -247,9 +247,9 @@ impl< X: Array2, Y: Array1, C: Criterion, -> BaseForestRegressor +> BaseForest { - /// Predict class for `x` + /// Predict values for `x` /// * `x` - _KxM_ data where _K_ is number of observations and _M_ is number of features. pub fn predict(&self, x: &X) -> Result { let (n, _) = x.shape(); @@ -275,7 +275,7 @@ impl< result / TY::from_usize(n_trees).unwrap() } - /// Predict OOB classes for `x`. `x` is expected to be equal to the dataset used in training. + /// Predict OOB values for `x`. `x` is expected to be equal to the dataset used in training. pub fn predict_oob(&self, x: &X) -> Result { let (n, _) = x.shape(); @@ -336,10 +336,10 @@ mod tests { use crate::tree::criterion::MSECriterion; #[test] - fn test_base_forest_regressor_keep_samples() { + fn test_base_forest_keep_samples() { let x = DenseMatrix::from_2d_array(&[&[1.0, 2.0], &[3.0, 4.0], &[5.0, 6.0]]).unwrap(); let y = vec![1.0, 2.0, 3.0]; - let params = BaseForestRegressorParameters { + let params = BaseForestParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 2, @@ -350,7 +350,7 @@ mod tests { bootstrap: true, splitter: crate::tree::base_tree::Splitter::Best, }; - let regressor = BaseForestRegressor::fit(&x, &y, None, &MSECriterion, params).unwrap(); + let regressor = BaseForest::fit(&x, &y, None, &MSECriterion, params).unwrap(); assert_eq!(regressor.trees.unwrap().len(), 5); assert!(regressor.samples.is_some()); } @@ -363,12 +363,12 @@ mod tests { assert_eq!(empty.shape(), (0, 2)); let y: Vec = vec![]; - let result = BaseForestRegressor::fit( + let result = BaseForest::fit( &empty, &y, None, &MSECriterion, - BaseForestRegressorParameters { + BaseForestParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 2, @@ -392,12 +392,12 @@ mod tests { assert_eq!(no_features.shape(), (2, 0)); let y: Vec = vec![1.0, 2.0]; - let result = BaseForestRegressor::fit( + let result = BaseForest::fit( &no_features, &y, None, &MSECriterion, - BaseForestRegressorParameters { + BaseForestParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 2, @@ -423,7 +423,7 @@ mod tests { let model_parameters = (0..=9).map(|x| x as f64).collect::>(); let y: Vec = model_parameters.xa(true, &x); - let forest_parameters = BaseForestRegressorParameters { + let forest_parameters = BaseForestParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 2, @@ -435,9 +435,8 @@ mod tests { splitter: crate::tree::base_tree::Splitter::Best, }; - let forest = - BaseForestRegressor::fit(&x, &y, None, &MSECriterion, forest_parameters.clone()) - .expect("Fit should work"); + let forest = BaseForest::fit(&x, &y, None, &MSECriterion, forest_parameters.clone()) + .expect("Fit should work"); let y_hat = forest.predict(&x).expect("Predict should work"); assert!((y_hat.iter().sum::() - y.iter().sum::()).abs() < 1e-9); @@ -446,7 +445,7 @@ mod tests { // Positive weights in [0.5, 2.0) let sample_weights: Vec = (0..1000).map(|_| rng.random_range(0.5..2.0)).collect(); - let forest = BaseForestRegressor::fit( + let forest = BaseForest::fit( &x, &y, Some(&sample_weights), @@ -483,7 +482,7 @@ mod tests { let sample_weights: Vec = (0..n_rows).map(|i| 1.0 + (i % 4) as f64).collect(); for weights in [None, Some(sample_weights.as_slice())] { - let params = BaseForestRegressorParameters { + let params = BaseForestParameters { max_depth: Some(1), min_samples_leaf: 1, min_samples_split: 2, @@ -494,8 +493,7 @@ mod tests { bootstrap: true, // Use bootstrapping splitter: crate::tree::base_tree::Splitter::Best, }; - let regressor = - BaseForestRegressor::fit(&x, &y, weights, &MSECriterion, params).unwrap(); + let regressor = BaseForest::fit(&x, &y, weights, &MSECriterion, params).unwrap(); let samples = regressor.samples.unwrap(); for (t, in_bag) in samples.iter().enumerate() { @@ -530,7 +528,7 @@ mod tests { let y: Vec = (0..n_rows).map(|i| ((i * 31) % 17) as f64).collect(); for splitter in [Splitter::Best, Splitter::Random] { - let params = BaseForestRegressorParameters { + let params = BaseForestParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 2, @@ -541,7 +539,7 @@ mod tests { bootstrap: false, splitter: splitter.clone(), }; - let forest = BaseForestRegressor::fit(&x, &y, None, &MSECriterion, params).unwrap(); + let forest = BaseForest::fit(&x, &y, None, &MSECriterion, params).unwrap(); let trees = forest.trees.unwrap(); assert!( trees.iter().any(|tree| tree != &trees[0]), @@ -559,7 +557,7 @@ mod tests { // Rows with y = 10 have weight 9, rows with y = 0 have weight 1. let sample_weights: Vec = (0..20).map(|i| if i < 10 { 1.0 } else { 9.0 }).collect(); - let parameters = BaseForestRegressorParameters { + let parameters = BaseForestParameters { max_depth: Some(0), min_samples_leaf: 1, min_samples_split: 2, @@ -571,7 +569,7 @@ mod tests { splitter: crate::tree::base_tree::Splitter::Best, }; - let forest = BaseForestRegressor::fit( + let forest = BaseForest::fit( &x, &y, Some(&sample_weights), @@ -590,8 +588,8 @@ mod tests { } // Without weights, the predicted value should be close to 5 - let forest = BaseForestRegressor::fit(&x, &y, None, &MSECriterion, parameters) - .expect("Fit should work"); + let forest = + BaseForest::fit(&x, &y, None, &MSECriterion, parameters).expect("Fit should work"); let y_hat = forest.predict(&x).expect("Predict should work"); for p in y_hat.iter() { @@ -602,7 +600,7 @@ mod tests { } // Use bootstrapping - let parameters = BaseForestRegressorParameters { + let parameters = BaseForestParameters { max_depth: Some(0), // Match sibling test on RandomForestRegressor min_samples_leaf: 1, min_samples_split: 2, @@ -614,7 +612,7 @@ mod tests { splitter: crate::tree::base_tree::Splitter::Best, }; - let forest = BaseForestRegressor::fit( + let forest = BaseForest::fit( &x, &y, Some(&sample_weights), @@ -633,8 +631,8 @@ mod tests { } // Without weights, the predicted value should be reasonably close to 5 - let forest = BaseForestRegressor::fit(&x, &y, None, &MSECriterion, parameters) - .expect("Fit should work"); + let forest = + BaseForest::fit(&x, &y, None, &MSECriterion, parameters).expect("Fit should work"); let y_hat = forest.predict(&x).expect("Predict should work"); // The bootstrapping makes that we will not be very close to 5 @@ -662,7 +660,7 @@ mod tests { for splitter in [Splitter::Best, Splitter::Random] { for weights in [None, Some(sample_weights.as_slice())] { - let params = BaseForestRegressorParameters { + let params = BaseForestParameters { max_depth: None, min_samples_leaf: 1, min_samples_split: 2, @@ -674,12 +672,10 @@ mod tests { splitter: splitter.clone(), }; - let forest_a = - BaseForestRegressor::fit(&x, &y, weights, &MSECriterion, params.clone()) - .expect("Fit should work"); - let forest_b = - BaseForestRegressor::fit(&x, &y, weights, &MSECriterion, params.clone()) - .expect("Fit should work"); + let forest_a = BaseForest::fit(&x, &y, weights, &MSECriterion, params.clone()) + .expect("Fit should work"); + let forest_b = BaseForest::fit(&x, &y, weights, &MSECriterion, params.clone()) + .expect("Fit should work"); assert_eq!( forest_a, forest_b, @@ -696,12 +692,12 @@ mod tests { ); // A different seed must give a different forest, else the check above is trivial - let forest_c = BaseForestRegressor::fit( + let forest_c = BaseForest::fit( &x, &y, weights, &MSECriterion, - BaseForestRegressorParameters { seed: 43, ..params }, + BaseForestParameters { seed: 43, ..params }, ) .expect("Fit should work"); assert_ne!( diff --git a/src/ensemble/extra_trees_regressor.rs b/src/ensemble/extra_trees_regressor.rs index 74fdbc9a..092d95e5 100644 --- a/src/ensemble/extra_trees_regressor.rs +++ b/src/ensemble/extra_trees_regressor.rs @@ -59,7 +59,7 @@ use std::fmt::Debug; use serde::{Deserialize, Serialize}; use crate::api::{Predictor, SupervisedEstimator}; -use crate::ensemble::base_forest_regressor::{BaseForestRegressor, BaseForestRegressorParameters}; +use crate::ensemble::base_forest::{BaseForest, BaseForestParameters}; use crate::error::Failed; use crate::linalg::basic::arrays::{Array1, Array2}; use crate::numbers::basenum::Number; @@ -123,7 +123,7 @@ pub struct ExtraTreesRegressor< X: Array2, Y: Array1, > { - forest_regressor: Option>, + forest_regressor: Option>, } impl ExtraTreesRegressorParameters { @@ -235,7 +235,7 @@ impl, Y: Array1 sample_weights: Option<&[f64]>, parameters: ExtraTreesRegressorParameters, ) -> Result, Failed> { - let regressor_params = BaseForestRegressorParameters { + let regressor_params = BaseForestParameters { max_depth: parameters.max_depth, min_samples_leaf: parameters.min_samples_leaf, min_samples_split: parameters.min_samples_split, @@ -247,7 +247,7 @@ impl, Y: Array1 splitter: Splitter::Random, }; let forest_regressor = - BaseForestRegressor::fit(x, y, sample_weights, &MSECriterion, regressor_params)?; + BaseForest::fit(x, y, sample_weights, &MSECriterion, regressor_params)?; Ok(ExtraTreesRegressor { forest_regressor: Some(forest_regressor), diff --git a/src/ensemble/mod.rs b/src/ensemble/mod.rs index 4f5eefc5..e56b9256 100644 --- a/src/ensemble/mod.rs +++ b/src/ensemble/mod.rs @@ -16,7 +16,7 @@ //! //! * ["An Introduction to Statistical Learning", James G., Witten D., Hastie T., Tibshirani R., 8.2 Bagging, Random Forests, Boosting](http://faculty.marshall.usc.edu/gareth-james/ISL/) -mod base_forest_regressor; +mod base_forest; pub mod extra_trees_regressor; /// Random forest classifier pub mod random_forest_classifier; diff --git a/src/ensemble/random_forest_regressor.rs b/src/ensemble/random_forest_regressor.rs index 014f1918..ff5e27c8 100644 --- a/src/ensemble/random_forest_regressor.rs +++ b/src/ensemble/random_forest_regressor.rs @@ -50,7 +50,7 @@ use std::fmt::Debug; use serde::{Deserialize, Serialize}; use crate::api::{Predictor, SupervisedEstimator}; -use crate::ensemble::base_forest_regressor::{BaseForestRegressor, BaseForestRegressorParameters}; +use crate::ensemble::base_forest::{BaseForest, BaseForestParameters}; use crate::error::Failed; use crate::linalg::basic::arrays::{Array1, Array2}; use crate::numbers::basenum::Number; @@ -114,7 +114,7 @@ pub struct RandomForestRegressor< X: Array2, Y: Array1, > { - forest_regressor: Option>, + forest_regressor: Option>, } impl RandomForestRegressorParameters { @@ -428,7 +428,7 @@ impl, Y: Array1 sample_weights: Option<&[f64]>, parameters: RandomForestRegressorParameters, ) -> Result, Failed> { - let regressor_params = BaseForestRegressorParameters { + let regressor_params = BaseForestParameters { max_depth: parameters.max_depth, min_samples_leaf: parameters.min_samples_leaf, min_samples_split: parameters.min_samples_split, @@ -440,7 +440,7 @@ impl, Y: Array1 splitter: Splitter::Best, }; let forest_regressor = - BaseForestRegressor::fit(x, y, sample_weights, &MSECriterion, regressor_params)?; + BaseForest::fit(x, y, sample_weights, &MSECriterion, regressor_params)?; Ok(RandomForestRegressor { forest_regressor: Some(forest_regressor), From 808caada7cc728fdc612a9b7a5a6dee2e09a1c78 Mon Sep 17 00:00:00 2001 From: slievens <5492873+slievens@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:47:49 +0200 Subject: [PATCH 15/18] Move regression predict out of BaseForest BaseForest only gives access to its trees and samples. The mean of the tree predictions is in predict_mean and predict_mean_oob. --- src/ensemble/base_forest.rs | 140 ++++++++++++------------ src/ensemble/extra_trees_regressor.rs | 12 +- src/ensemble/random_forest_regressor.rs | 12 +- 3 files changed, 84 insertions(+), 80 deletions(-) diff --git a/src/ensemble/base_forest.rs b/src/ensemble/base_forest.rs index c75d6520..1b6e760b 100644 --- a/src/ensemble/base_forest.rs +++ b/src/ensemble/base_forest.rs @@ -87,6 +87,16 @@ pub struct BaseForest< impl, Y: Array1, C: Criterion> BaseForest { + /// The trees of the forest. Call only after `fit`. + pub(crate) fn trees(&self) -> &[BaseTree] { + self.trees.as_ref().unwrap() + } + + /// For each tree, the rows it used in training. `None` if `keep_samples` was false. + pub(crate) fn samples(&self) -> Option<&[Vec]> { + self.samples.as_deref() + } + /// Build a forest of trees from the training set. /// * `x` - _NxM_ matrix with _N_ observations and _M_ features in each observation. /// * `y` - the target class values @@ -241,90 +251,80 @@ impl, Y: Array1 } } -impl< - TX: Number + FloatNumber + PartialOrd, +/// Mean of the tree predictions for each row of `x`. +pub(crate) fn predict_mean( + forest: &BaseForest, + x: &X, +) -> Result +where + TX: Number + FloatNumber, TY: Number, X: Array2, Y: Array1, C: Criterion, -> BaseForest { - /// Predict values for `x` - /// * `x` - _KxM_ data where _K_ is number of observations and _M_ is number of features. - pub fn predict(&self, x: &X) -> Result { - let (n, _) = x.shape(); - - let mut result = Y::zeros(n); - - for i in 0..n { - result.set(i, self.predict_for_row(x, i)); + let (n, _) = x.shape(); + let mut result = Y::zeros(n); + let num_trees = TY::from_usize(forest.trees().len()).unwrap(); + + for row in 0..n { + let mut sum = TY::zero(); + for tree in forest.trees() { + sum += TY::from_f64(*tree.leaf_output_for_row(x, row)).unwrap(); } - - Ok(result) - } - - fn predict_for_row(&self, x: &X, row: usize) -> TY { - let n_trees = self.trees.as_ref().unwrap().len(); - - let mut result = TY::zero(); - - for tree in self.trees.as_ref().unwrap().iter() { - result += TY::from_f64(*tree.leaf_output_for_row(x, row)).unwrap(); - } - - result / TY::from_usize(n_trees).unwrap() + result.set(row, sum / num_trees); } + Ok(result) +} - /// Predict OOB values for `x`. `x` is expected to be equal to the dataset used in training. - pub fn predict_oob(&self, x: &X) -> Result { - let (n, _) = x.shape(); - - let samples = match &self.samples { - Some(s) => s, - None => { - return Err(Failed::because( - FailedError::PredictFailed, - "Need samples=true for OOB predictions.", - )); - } - }; +/// Mean of the out-of-bag tree predictions for each row of `x`. +/// `x` is expected to be equal to the dataset used in training. +pub(crate) fn predict_mean_oob( + forest: &BaseForest, + x: &X, +) -> Result +where + TX: Number + FloatNumber, + TY: Number, + X: Array2, + Y: Array1, + C: Criterion, +{ + let (n, _) = x.shape(); - if samples[0].len() != n { + let samples = match forest.samples() { + Some(s) => s, + None => { return Err(Failed::because( FailedError::PredictFailed, - "Prediction matrix must match matrix used in training for OOB predictions.", + "Need samples=true for OOB predictions.", )); } + }; - let mut result = Y::zeros(n); - - for i in 0..n { - result.set(i, self.predict_for_row_oob(x, i)); - } - - Ok(result) + if samples[0].len() != n { + return Err(Failed::because( + FailedError::PredictFailed, + "Prediction matrix must match matrix used in training for OOB predictions.", + )); } - fn predict_for_row_oob(&self, x: &X, row: usize) -> TY { - let mut n_trees = 0; - let mut result = TY::zero(); + let mut result = Y::zeros(n); - for (tree, samples) in self - .trees - .as_ref() - .unwrap() - .iter() - .zip(self.samples.as_ref().unwrap()) - { - if !samples[row] { - result += TY::from_f64(*tree.leaf_output_for_row(x, row)).unwrap(); + for row in 0..n { + let mut sum = TY::zero(); + let mut n_trees = 0usize; + for (tree, in_bag) in forest.trees().iter().zip(samples) { + if !in_bag[row] { + sum += TY::from_f64(*tree.leaf_output_for_row(x, row)).unwrap(); n_trees += 1; } } - // TODO: What to do if there are no oob trees? - result / TY::from(n_trees).unwrap() + result.set(row, sum / TY::from_usize(n_trees).unwrap()); } + + Ok(result) } #[cfg(test)] @@ -437,7 +437,7 @@ mod tests { let forest = BaseForest::fit(&x, &y, None, &MSECriterion, forest_parameters.clone()) .expect("Fit should work"); - let y_hat = forest.predict(&x).expect("Predict should work"); + let y_hat = predict_mean(&forest, &x).expect("Predict should work"); assert!((y_hat.iter().sum::() - y.iter().sum::()).abs() < 1e-9); // Seeded RNG: the test gives the same result on each run @@ -453,7 +453,7 @@ mod tests { forest_parameters.clone(), ) .expect("Fit should work"); - let y_hat = forest.predict(&x).expect("Predict should work"); + let y_hat = predict_mean(&forest, &x).expect("Predict should work"); let s: f64 = sample_weights.iter().sum(); let normalized_weights = sample_weights.iter().map(|w| *w / s).collect::>(); @@ -577,7 +577,7 @@ mod tests { parameters.clone(), ) .expect("Fit should work"); - let y_hat = forest.predict(&x).expect("Predict should work"); + let y_hat = predict_mean(&forest, &x).expect("Predict should work"); // weighted mean is (10.0 * 0 + 90*10) / 100 = 9 for p in y_hat.iter() { @@ -590,7 +590,7 @@ mod tests { // Without weights, the predicted value should be close to 5 let forest = BaseForest::fit(&x, &y, None, &MSECriterion, parameters).expect("Fit should work"); - let y_hat = forest.predict(&x).expect("Predict should work"); + let y_hat = predict_mean(&forest, &x).expect("Predict should work"); for p in y_hat.iter() { assert!( @@ -620,7 +620,7 @@ mod tests { parameters.clone(), ) .expect("Fit should work"); - let y_hat = forest.predict(&x).expect("Predict should work"); + let y_hat = predict_mean(&forest, &x).expect("Predict should work"); // with bootstrapping the weight should be close to 9, but not extremely close for p in y_hat.iter() { @@ -633,7 +633,7 @@ mod tests { // Without weights, the predicted value should be reasonably close to 5 let forest = BaseForest::fit(&x, &y, None, &MSECriterion, parameters).expect("Fit should work"); - let y_hat = forest.predict(&x).expect("Predict should work"); + let y_hat = predict_mean(&forest, &x).expect("Predict should work"); // The bootstrapping makes that we will not be very close to 5 for p in y_hat.iter() { @@ -686,8 +686,8 @@ mod tests { "bootstrap samples differ (splitter: {splitter:?}, weights: {weights:?})" ); assert_eq!( - forest_a.predict(&x).unwrap(), - forest_b.predict(&x).unwrap(), + predict_mean(&forest_a, &x).unwrap(), + predict_mean(&forest_b, &x).unwrap(), "predictions differ (splitter: {splitter:?}, weights: {weights:?})" ); diff --git a/src/ensemble/extra_trees_regressor.rs b/src/ensemble/extra_trees_regressor.rs index 092d95e5..2ab46a9e 100644 --- a/src/ensemble/extra_trees_regressor.rs +++ b/src/ensemble/extra_trees_regressor.rs @@ -59,7 +59,9 @@ use std::fmt::Debug; use serde::{Deserialize, Serialize}; use crate::api::{Predictor, SupervisedEstimator}; -use crate::ensemble::base_forest::{BaseForest, BaseForestParameters}; +use crate::ensemble::base_forest::{ + BaseForest, BaseForestParameters, predict_mean, predict_mean_oob, +}; use crate::error::Failed; use crate::linalg::basic::arrays::{Array1, Array2}; use crate::numbers::basenum::Number; @@ -254,21 +256,21 @@ impl, Y: Array1 }) } - /// Predict class for `x` + /// Predict values for `x` /// * `x` - _KxM_ data where _K_ is number of observations and _M_ is number of features. pub fn predict(&self, x: &X) -> Result { match &self.forest_regressor { - Some(forest) => forest.predict(x), + Some(forest) => predict_mean(forest, x), None => Err(Failed::predict( "'fit' should be called before calling 'predict'", )), } } - /// Predict OOB classes for `x`. `x` is expected to be equal to the dataset used in training. + /// Predict OOB values for `x`. `x` is expected to be equal to the dataset used in training. pub fn predict_oob(&self, x: &X) -> Result { match &self.forest_regressor { - Some(forest) => forest.predict_oob(x), + Some(forest) => predict_mean_oob(forest, x), None => Err(Failed::predict( "'fit' should be called before calling 'predict'", )), diff --git a/src/ensemble/random_forest_regressor.rs b/src/ensemble/random_forest_regressor.rs index ff5e27c8..f03faa9f 100644 --- a/src/ensemble/random_forest_regressor.rs +++ b/src/ensemble/random_forest_regressor.rs @@ -50,7 +50,9 @@ use std::fmt::Debug; use serde::{Deserialize, Serialize}; use crate::api::{Predictor, SupervisedEstimator}; -use crate::ensemble::base_forest::{BaseForest, BaseForestParameters}; +use crate::ensemble::base_forest::{ + BaseForest, BaseForestParameters, predict_mean, predict_mean_oob, +}; use crate::error::Failed; use crate::linalg::basic::arrays::{Array1, Array2}; use crate::numbers::basenum::Number; @@ -447,21 +449,21 @@ impl, Y: Array1 }) } - /// Predict class for `x` + /// Predict value for `x` /// * `x` - _KxM_ data where _K_ is number of observations and _M_ is number of features. pub fn predict(&self, x: &X) -> Result { match &self.forest_regressor { - Some(forest) => forest.predict(x), + Some(forest) => predict_mean(forest, x), None => Err(Failed::predict( "'fit' should be called before calling 'predict'", )), } } - /// Predict OOB classes for `x`. `x` is expected to be equal to the dataset used in training. + /// Predict OOB values for `x`. `x` is expected to be equal to the dataset used in training. pub fn predict_oob(&self, x: &X) -> Result { match &self.forest_regressor { - Some(forest) => forest.predict_oob(x), + Some(forest) => predict_mean_oob(forest, x), None => Err(Failed::predict( "'fit' should be called before calling 'predict'", )), From 7a8237c06af319703e05b870d73567a4f92711d3 Mon Sep 17 00:00:00 2001 From: slievens <5492873+slievens@users.noreply.github.com> Date: Fri, 23 Oct 2026 21:00:28 +0200 Subject: [PATCH 16/18] Make BaseTree and BaseForest generic over the leaf output BaseTree and BaseForest do not store the criterion. Only the type of the leaf value is kept in the tree. Make the types generic over this output type and move the Criterion type parameter to the fit functions. Move output_approximate_eq from Criterion to a new LeafOutput trait on the output type. --- src/ensemble/base_forest.rs | 51 ++++++++-------- src/ensemble/extra_trees_regressor.rs | 2 +- src/ensemble/random_forest_regressor.rs | 2 +- src/tree/base_tree.rs | 78 +++++++++++++------------ src/tree/criterion.rs | 35 +++++------ src/tree/decision_tree_regressor.rs | 2 +- 6 files changed, 90 insertions(+), 80 deletions(-) diff --git a/src/ensemble/base_forest.rs b/src/ensemble/base_forest.rs index 1b6e760b..905433c0 100644 --- a/src/ensemble/base_forest.rs +++ b/src/ensemble/base_forest.rs @@ -12,7 +12,7 @@ use crate::numbers::floatnum::FloatNumber; use crate::rand_custom::get_rng_impl; use crate::tree::base_tree::{BaseTree, BaseTreeParameters, Splitter}; -use crate::tree::criterion::Criterion; +use crate::tree::criterion::{Criterion, LeafOutput}; #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[derive(Debug, Clone)] @@ -47,8 +47,13 @@ pub struct BaseForestParameters { pub splitter: Splitter, } -impl, Y: Array1, C: Criterion> - PartialEq for BaseForest +impl< + TX: Number + FloatNumber + PartialOrd, + TY: Number, + X: Array2, + Y: Array1, + Out: LeafOutput, +> PartialEq for BaseForest { fn eq(&self, other: &Self) -> bool { if self.trees.as_ref().unwrap().len() != other.trees.as_ref().unwrap().len() { @@ -62,14 +67,11 @@ impl, Y: Array1 } } -/// Base forest of decision trees. The criterion `C` selects the splits and gives the leaf values. +/// Base forest of decision trees. Out is the type of the leaf values. The criterion is passed to fit. #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[cfg_attr( feature = "serde", - serde(bound( - serialize = "C::Output: Serialize", - deserialize = "C::Output: Deserialize<'de>" - )) + serde(bound(serialize = "Out: Serialize", deserialize = "Out: Deserialize<'de>")) )] #[derive(Debug)] pub struct BaseForest< @@ -77,18 +79,23 @@ pub struct BaseForest< TY: Number, X: Array2, Y: Array1, - C: Criterion, + Out: LeafOutput, > { - trees: Option>>, + trees: Option>>, num_features: Option, samples: Option>>, } -impl, Y: Array1, C: Criterion> - BaseForest +impl< + TX: Number + FloatNumber + PartialOrd, + TY: Number, + X: Array2, + Y: Array1, + Out: LeafOutput, +> BaseForest { /// The trees of the forest. Call only after `fit`. - pub(crate) fn trees(&self) -> &[BaseTree] { + pub(crate) fn trees(&self) -> &[BaseTree] { self.trees.as_ref().unwrap() } @@ -101,13 +108,13 @@ impl, Y: Array1 /// * `x` - _NxM_ matrix with _N_ observations and _M_ features in each observation. /// * `y` - the target class values /// * `sample_weights`: optional sample_weights to use during fitting - pub fn fit( + pub fn fit>( x: &X, y: &Y, sample_weights: Option<&[f64]>, criterion: &C, parameters: BaseForestParameters, - ) -> Result, Failed> { + ) -> Result, Failed> { let (n_rows, num_attributes) = x.shape(); if n_rows != y.shape() { @@ -126,7 +133,7 @@ impl, Y: Array1 let mut rng = get_rng_impl(Some(parameters.seed)); let n_trees = parameters.n_trees; - let mut trees: Vec> = Vec::with_capacity(n_trees); + let mut trees: Vec> = Vec::with_capacity(n_trees); let mut maybe_all_samples: Option>> = Option::None; if parameters.keep_samples { @@ -152,7 +159,7 @@ impl, Y: Array1 for tree_idx in 0..parameters.n_trees { if parameters.bootstrap { - samples = BaseForest::::sample_with_replacement( + samples = BaseForest::::sample_with_replacement( n_rows, &mut rng, dist.as_ref(), @@ -252,8 +259,8 @@ impl, Y: Array1 } /// Mean of the tree predictions for each row of `x`. -pub(crate) fn predict_mean( - forest: &BaseForest, +pub(crate) fn predict_mean( + forest: &BaseForest, x: &X, ) -> Result where @@ -261,7 +268,6 @@ where TY: Number, X: Array2, Y: Array1, - C: Criterion, { let (n, _) = x.shape(); let mut result = Y::zeros(n); @@ -279,8 +285,8 @@ where /// Mean of the out-of-bag tree predictions for each row of `x`. /// `x` is expected to be equal to the dataset used in training. -pub(crate) fn predict_mean_oob( - forest: &BaseForest, +pub(crate) fn predict_mean_oob( + forest: &BaseForest, x: &X, ) -> Result where @@ -288,7 +294,6 @@ where TY: Number, X: Array2, Y: Array1, - C: Criterion, { let (n, _) = x.shape(); diff --git a/src/ensemble/extra_trees_regressor.rs b/src/ensemble/extra_trees_regressor.rs index 2ab46a9e..10dc42f5 100644 --- a/src/ensemble/extra_trees_regressor.rs +++ b/src/ensemble/extra_trees_regressor.rs @@ -125,7 +125,7 @@ pub struct ExtraTreesRegressor< X: Array2, Y: Array1, > { - forest_regressor: Option>, + forest_regressor: Option>, } impl ExtraTreesRegressorParameters { diff --git a/src/ensemble/random_forest_regressor.rs b/src/ensemble/random_forest_regressor.rs index f03faa9f..917d292a 100644 --- a/src/ensemble/random_forest_regressor.rs +++ b/src/ensemble/random_forest_regressor.rs @@ -116,7 +116,7 @@ pub struct RandomForestRegressor< X: Array2, Y: Array1, > { - forest_regressor: Option>, + forest_regressor: Option>, } impl RandomForestRegressorParameters { diff --git a/src/tree/base_tree.rs b/src/tree/base_tree.rs index 52bc1ada..6594338c 100644 --- a/src/tree/base_tree.rs +++ b/src/tree/base_tree.rs @@ -11,7 +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; +use crate::tree::criterion::{Criterion, LeafOutput}; #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[derive(Debug, Clone, Copy, Default)] @@ -47,15 +47,17 @@ pub struct BaseTreeParameters { #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[cfg_attr( feature = "serde", - serde(bound( - serialize = "C::Output: Serialize", - deserialize = "C::Output: Deserialize<'de>" - )) + serde(bound(serialize = "Out: Serialize", deserialize = "Out: Deserialize<'de>")) )] #[derive(Debug)] -pub struct BaseTree, Y: Array1, C: Criterion> -{ - nodes: Vec>, +pub struct BaseTree< + TX: Number + PartialOrd, + TY: Number, + X: Array2, + Y: Array1, + Out: LeafOutput, +> { + nodes: Vec>, parameters: Option, depth: u16, root_mass: f64, @@ -65,11 +67,11 @@ pub struct BaseTree, Y: Array _phantom_y: PhantomData, } -impl, Y: Array1, C: Criterion> - BaseTree +impl, Y: Array1, Out: LeafOutput> + BaseTree { /// Get nodes, return a shared reference - fn nodes(&self) -> &Vec> { + fn nodes(&self) -> &Vec> { self.nodes.as_ref() } /// Get parameters, return a shared reference @@ -102,9 +104,9 @@ impl Node { } } -impl Node { - fn approximate_eq>(&self, other: &Self) -> bool { - C::output_approximate_eq(&self.output, &other.output) +impl Node { + fn approximate_eq(&self, other: &Self) -> bool { + Out::approximate_eq(&self.output, &other.output) && self.split_feature == other.split_feature && match (self.split_value, other.split_value) { (Some(a), Some(b)) => (a - b).abs() < f64::EPSILON, @@ -119,8 +121,8 @@ impl Node { } } -impl, Y: Array1, C: Criterion> PartialEq - for BaseTree +impl, Y: Array1, Out: LeafOutput> PartialEq + for BaseTree { fn eq(&self, other: &Self) -> bool { if self.depth != other.depth || self.nodes.len() != other.nodes.len() { @@ -129,7 +131,7 @@ impl, Y: Array1, C: Crite self.nodes .iter() .zip(other.nodes().iter()) - .all(|(a, b)| a.approximate_eq::(b)) + .all(|(a, b)| a.approximate_eq(b)) } } } @@ -368,16 +370,16 @@ where } } -impl, Y: Array1, C: Criterion> - BaseTree +impl, Y: Array1, Out: LeafOutput> + BaseTree { - pub(crate) fn fit_inner( + pub(crate) fn fit_inner>( x: &X, y: &Y, sample_weights: Option<&[f64]>, criterion: &C, parameters: BaseTreeParameters, - ) -> Result, Failed> { + ) -> Result, Failed> { let (x_nrows, num_attributes) = x.shape(); if x_nrows != y.shape() { return Err(Failed::fit("Size of x should equal size of y")); @@ -410,7 +412,7 @@ impl, Y: Array1, C: Crite ) } - pub(crate) fn fit_weak_learner( + pub(crate) fn fit_weak_learner>( x: &X, y: &Y, sample_weights: Option<&[f64]>, @@ -419,10 +421,10 @@ impl, Y: Array1, C: Crite order: &[Vec], criterion: &C, parameters: BaseTreeParameters, - ) -> Result, Failed> { + ) -> Result, Failed> { check_row_count(y.shape())?; match sample_weights { - Some(_) => Self::grow::( + Some(_) => Self::grow::( x, y, sample_weights, @@ -435,15 +437,17 @@ impl, Y: Array1, C: Crite None if samples.iter().all(|&s| s <= 1) => { // Note: if any sample has count zero, it will be filtered out, so we should // use the smallest possible NodeElement - Self::grow::(x, y, None, samples, mtry, order, criterion, parameters) + Self::grow::( + x, y, None, samples, mtry, order, criterion, parameters, + ) } - None => Self::grow::( + None => Self::grow::( x, y, None, samples, mtry, order, criterion, parameters, ), } } - fn grow( + fn grow>( x: &X, y: &Y, sample_weights: Option<&[f64]>, @@ -452,8 +456,8 @@ impl, Y: Array1, C: Crite order: &[Vec], criterion: &C, parameters: BaseTreeParameters, - ) -> Result, Failed> { - let mut nodes: Vec> = Vec::new(); + ) -> Result, Failed> { + let mut nodes: Vec> = Vec::new(); let mut rng = get_rng_impl(parameters.seed); let mut n = 0usize; @@ -559,7 +563,7 @@ impl, Y: Array1, C: Crite Ok(base_tree) } - pub(crate) fn leaf_output_for_row(&self, x: &X, row: usize) -> &C::Output { + pub(crate) fn leaf_output_for_row(&self, x: &X, row: usize) -> &Out { let mut node_id = 0; loop { let node = &self.nodes[node_id]; @@ -577,7 +581,7 @@ impl, Y: Array1, C: Crite } } - fn find_best_cutoff( + fn find_best_cutoff>( &mut self, visitor: &mut NodeVisitor<'_, TX, TY, X, Y, C>, mtry: usize, @@ -605,7 +609,7 @@ impl, Y: Array1, C: Crite self.nodes[visitor.node].split_score.is_some() } - fn find_random_split( + fn find_random_split>( &mut self, visitor: &NodeVisitor<'_, TX, TY, X, Y, C>, n: usize, @@ -672,7 +676,7 @@ impl, Y: Array1, C: Crite } } - fn find_best_split( + fn find_best_split>( &mut self, visitor: &NodeVisitor<'_, TX, TY, X, Y, C>, n: usize, @@ -754,7 +758,7 @@ impl, Y: Array1, C: Crite /// split, push it on `visitor_stack`. /// If a branch has fewer than `min_samples_leaf` samples, clear the split and /// keep the node as a leaf. - fn split<'a, E: NodeElement>( + fn split<'a, E: NodeElement, C: Criterion>( &mut self, visitor: NodeVisitor<'a, TX, TY, X, Y, C>, mtry: usize, @@ -1461,16 +1465,16 @@ mod tests { assert_eq!(std::mem::size_of::(), 16); } - type TestTree = BaseTree, Vec, MSECriterion>; + type TestTree = BaseTree, Vec, f64>; /// Leaf output of `tree` for each row of `x`. - fn predict(tree: &BaseTree, x: &X) -> Vec + fn predict(tree: &BaseTree, x: &X) -> Vec where TX: Number + PartialOrd, TY: Number, X: Array2, Y: Array1, - C: Criterion, + Out: LeafOutput, { (0..x.shape().0) .map(|i| tree.leaf_output_for_row(x, i).clone()) diff --git a/src/tree/criterion.rs b/src/tree/criterion.rs index d75fa67a..aa6f3f50 100644 --- a/src/tree/criterion.rs +++ b/src/tree/criterion.rs @@ -47,13 +47,21 @@ pub(crate) trait Criterion { /// Value of a leaf that contains the samples of `stats`. fn leaf_value(&self, stats: &Self::Stats) -> Self::Output; - /// Are the outputs approximately equal? - fn output_approximate_eq(out: &Self::Output, other_out: &Self::Output) -> bool; - /// Convert a target value of the training data to the target of this criterion. fn target(&self, y: T) -> Self::Target; } +/// Value in a leaf of a tree. +pub(crate) trait LeafOutput: Clone + Debug { + fn approximate_eq(&self, other: &Self) -> bool; +} + +impl LeafOutput for f64 { + fn approximate_eq(&self, other: &f64) -> bool { + (self - other).abs() < f64::EPSILON + } +} + /// Statistics for the weighted mean of the targets. #[derive(Clone, Debug, Default)] pub(crate) struct MeanStats { @@ -124,10 +132,6 @@ impl Criterion for MSECriterion { } } - fn output_approximate_eq(out: &Self::Output, other_out: &Self::Output) -> bool { - (out - other_out).abs() < f64::EPSILON - } - fn target(&self, y: T) -> Self::Target { y.to_f64().unwrap() } @@ -405,17 +409,14 @@ mod tests { } #[test] - fn mse_output_approximate_eq() { - assert!(MSECriterion::output_approximate_eq(&1.5, &1.5)); - assert!(MSECriterion::output_approximate_eq(&0.0, &-0.0)); + fn f64_approximate_eq() { + assert!(1.5.approximate_eq(&1.5)); + assert!(0.0.approximate_eq(&-0.0)); // Difference smaller than EPSILON - assert!(MSECriterion::output_approximate_eq( - &0.0, - &(f64::EPSILON / 2.0) - )); + assert!(0.0.approximate_eq(&(f64::EPSILON / 2.0))); // Difference of exactly EPSILON is not "approximately equal" (strict comparison) - assert!(!MSECriterion::output_approximate_eq(&0.0, &f64::EPSILON)); - assert!(!MSECriterion::output_approximate_eq(&1.0, &1.1)); - assert!(!MSECriterion::output_approximate_eq(&f64::NAN, &f64::NAN)); + assert!(!0.0.approximate_eq(&f64::EPSILON)); + assert!(!1.0.approximate_eq(&1.1)); + assert!(!f64::NAN.approximate_eq(&f64::NAN)); } } diff --git a/src/tree/decision_tree_regressor.rs b/src/tree/decision_tree_regressor.rs index 93290923..de0c207c 100644 --- a/src/tree/decision_tree_regressor.rs +++ b/src/tree/decision_tree_regressor.rs @@ -113,7 +113,7 @@ pub struct DecisionTreeRegressorParameters { #[derive(Debug)] pub struct DecisionTreeRegressor, Y: Array1> { - tree_regressor: Option>, + tree_regressor: Option>, num_features: Option, } From a67a93f19becf042531355dc12820af9ea05bb46 Mon Sep 17 00:00:00 2001 From: "Lorenzo (Mec-iS)" Date: Sat, 24 Oct 2026 06:20:41 +0100 Subject: [PATCH 17/18] chore: require ndarray 0.17.2 --- Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 9710bda4..b9eb009b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,7 +26,7 @@ exclude = [ approx = "0.5.1" cfg-if = "1.0.0" lazymatrix = { version = "0.6.0", optional = true, default-features = false } -ndarray = { version = "0.17", optional = true, default-features = false } +ndarray = { version = "0.17.2", optional = true, default-features = false } num-traits = "0.2.12" num = "0.4" rand = { version = "0.10.1", default-features = false, features = ["alloc"] } From 3cd696e50e6055021d7e4338d210f7aa903d4b71 Mon Sep 17 00:00:00 2001 From: "Lorenzo (Mec-iS)" Date: Sat, 24 Oct 2026 06:23:47 +0100 Subject: [PATCH 18/18] chore(changelog): consolidate candidate changes under [v0.7] section --- CHANGELOG.md | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2459db79..db9f89f8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,21 +5,21 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] -### Added -- A `Dependency audit` CI workflow runs `cargo deny` and `cargo audit` on each push and pull request, and every week. It fails when a dependency has a known vulnerability, is yanked, uses a license that is not allowed, or comes from an unknown source. The `deny.toml` file holds the rules. -<<<<<<< HEAD +## [v0.7] ### Added +- A `Dependency audit` CI workflow runs `cargo deny` and `cargo audit` on each push and pull request, and every week. It fails when a dependency has a known vulnerability, is yanked, uses a license that is not allowed, or comes from an unknown source. The `deny.toml` file holds the rules. - Optional `lazy-normalization` support for Lasso and Elastic Net. It applies normalization to a borrowed design and keeps Elastic Net's identity-block augmentation implicit, while retaining the dense Gram workspace. Enable the feature and set `with_lazy_normalization(true)` to select this path. The default remains eager. The new parameter field is available only with the feature and defaults to false during deserialization (#479). +- `tree`, `ensemble`: `DecisionTreeRegressor`, `RandomForestRegressor` and `ExtraTreesRegressor` get a `compute_feature_importances` method. ### Changed - `Cargo.toml`: optional `ndarray` is now `default-features = false`, and the default feature set is `std` (`ndarray?/std`). Ordinary builds still enable ndarray's standard-library support when `ndarray-bindings` is on, while `std` alone does not activate ndarray. The `std` feature only controls ndarray's standard-library support; smartcore itself is not `no_std` (#421). - **Migration:** builds that use `--no-default-features --features ndarray-bindings` (or `default-features = false` with `ndarray-bindings`) no longer get ndarray's `std` feature. Add `std` to the feature list to keep the previous behavior. - **Breaking change:** Enabling `lazy-normalization` adds a public `lazy_normalization` field to `LassoParameters`, `LassoSearchParameters`, `ElasticNetParameters`, and `ElasticNetSearchParameters`. Struct literals without `..Default::default()` must add `lazy_normalization: false` to retain the eager path, or `true` to select the lazy path. This also applies when another dependency enables the feature through Cargo feature unification (#479). +- **Breaking change:** 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. - `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. +- `Cargo.toml`: `ndarray` requires version `0.17.2` or later. ## [0.7.0] ### Added