Skip to content

Add MaxFeatures enum for forest parameter m #531

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
slievens wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Add MaxFeatures enum for forest parameter m #531

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5b552ec
Add MaxFeatures enum for forest parameter `m`
slievens Nov 11, 2026
ab7fc76
Merge branch 'main' into enum-for-m-pr
Mec-iS Nov 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
- `tree`: `with_seed` on `DecisionTreeRegressorParameters` and `DecisionTreeClassifierParameters`.
- `ensemble`: `ExtraTreesRegressorParameters` gets a `bootstrap` field and a `with_bootstrap` method. The default is `false`, so each tree still uses the whole dataset.
- `ensemble`: `ExtraTreesClassifier` added with an API very similar to `ExtraTreesRegressor`
- `ensemble`: `MaxFeatures` enum (`Sqrt`, `Fraction`, `Count`, `All`) sets the number of features to test at each split in the forest estimators.

### Fixed
- `ensemble`: `ExtraTreesRegressor::predict_oob` returns an error when the forest was fitted without bootstrap. Before this change, it returned `NaN` for a float target and panicked for an integer target (#520).
Expand All @@ -31,6 +32,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
- `ensemble`: each forest uses one RNG for bootstrap sampling and for the tree seeds (before: `seed + tree_idx`). Results for a given seed differ from earlier versions.
- **Breaking change:** `ensemble`: `ExtraTreesRegressorParameters` gets a public `bootstrap` field. Struct literals without `..Default::default()` must add `bootstrap: false` to keep the earlier behavior. Serialized parameters without this field still deserialize, with `bootstrap = false` (#520).
- `ensemble`: Extra trees do not sort the feature columns. Trees are not bit-identical to the old ones because of different summation order.
**Breaking change:** `ensemble`: `m` on the four forest parameter structs is now `MaxFeatures` (was `Option<usize>`). `with_m(3)` still works. `fit` returns an error when `m` is not valid for the data. Parameters serialized with an earlier version that contain `m` no longer deserialize.

## [0.7.0] - 2026-10-26
### Added
Expand Down
126 changes: 110 additions & 16 deletions src/ensemble/base_forest.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ use std::fmt::Debug;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};

use crate::ensemble::max_features::MaxFeatures;
use crate::error::{Failed, FailedError};
use crate::linalg::basic::arrays::{Array1, Array2};
use crate::linalg::basic::arrays::{ArrayView1, MutArrayView1};
Expand Down Expand Up @@ -33,8 +34,8 @@ pub struct BaseForestParameters {
/// The number of trees in the forest.
pub n_trees: usize,
#[cfg_attr(feature = "serde", serde(default))]
/// Number of random sample of predictors to use as split candidates.
pub m: Option<usize>,
/// Number of features to test at each split. See [`MaxFeatures`].
pub m: MaxFeatures,
#[cfg_attr(feature = "serde", serde(default))]
/// Whether to keep samples used for tree generation. This is required for OOB prediction.
pub keep_samples: bool,
Expand Down Expand Up @@ -127,9 +128,7 @@ impl<
));
}

let mtry = parameters
.m
.unwrap_or((num_attributes as f64).sqrt().floor() as usize);
let mtry = parameters.m.resolve(num_attributes)?;

let mut rng = get_rng_impl(parameters.seed);
let n_trees = parameters.n_trees;
Expand Down Expand Up @@ -436,7 +435,7 @@ mod tests {
min_samples_leaf: 1,
min_samples_split: 2,
n_trees: 5,
m: None,
m: MaxFeatures::Sqrt,
keep_samples: true,
seed: Some(42),
bootstrap: true,
Expand Down Expand Up @@ -465,7 +464,7 @@ mod tests {
min_samples_leaf: 1,
min_samples_split: 2,
n_trees: 5,
m: None,
m: MaxFeatures::Sqrt,
keep_samples: false,
seed: Some(0),
bootstrap: true,
Expand Down Expand Up @@ -494,7 +493,7 @@ mod tests {
min_samples_leaf: 1,
min_samples_split: 2,
n_trees: 5,
m: None,
m: MaxFeatures::Sqrt,
keep_samples: false,
seed: Some(0),
bootstrap: true,
Expand All @@ -520,7 +519,7 @@ mod tests {
min_samples_leaf: 1,
min_samples_split: 2,
n_trees: 5,
m: None,
m: MaxFeatures::Sqrt,
keep_samples: true,
seed: Some(42),
bootstrap: false,
Expand Down Expand Up @@ -579,7 +578,7 @@ mod tests {
min_samples_leaf: 1,
min_samples_split: 2,
n_trees: 10,
m: None,
m: MaxFeatures::Sqrt,
keep_samples: true, // keep samples used for each tree, so we can check that they are different
seed: Some(42),
bootstrap: true, // Use bootstrapping
Expand Down Expand Up @@ -625,7 +624,7 @@ mod tests {
min_samples_leaf: 1,
min_samples_split: 2,
n_trees: 10,
m: Some(1),
m: MaxFeatures::Count(1),
keep_samples: false,
seed: Some(42),
bootstrap: false,
Expand Down Expand Up @@ -654,7 +653,7 @@ mod tests {
min_samples_leaf: 1,
min_samples_split: 2,
n_trees: 10,
m: None,
m: MaxFeatures::Sqrt,
keep_samples: true,
seed: Some(42),
bootstrap: false, // No bootstrapping
Expand Down Expand Up @@ -697,7 +696,7 @@ mod tests {
min_samples_leaf: 1,
min_samples_split: 2,
n_trees: 500, // Use more trees than before to smooth out randomness
m: None,
m: MaxFeatures::Sqrt,
keep_samples: true, // keep samples used for each tree, so we can check that they are different
seed: Some(42),
bootstrap: true, // Use bootstrapping
Expand Down Expand Up @@ -757,7 +756,7 @@ mod tests {
min_samples_leaf: 1,
min_samples_split: 2,
n_trees: 10,
m: Some(1),
m: MaxFeatures::Count(1),
keep_samples: true,
seed: Some(42),
bootstrap: true,
Expand Down Expand Up @@ -817,7 +816,7 @@ mod tests {
min_samples_leaf: 1,
min_samples_split: 2,
n_trees: 10,
m: Some(1),
m: MaxFeatures::Count(1),
keep_samples: false,
seed: Some(42),
bootstrap,
Expand Down Expand Up @@ -890,7 +889,7 @@ mod tests {
min_samples_leaf: 1,
min_samples_split: 2,
n_trees: 10,
m: Some(1),
m: MaxFeatures::Count(1),
keep_samples: false,
seed,
bootstrap: true,
Expand Down Expand Up @@ -940,4 +939,99 @@ mod tests {
"two fits with seed None give the same forest"
);
}

#[cfg_attr(
all(target_arch = "wasm32", not(target_os = "wasi")),
wasm_bindgen_test::wasm_bindgen_test
)]
#[test]
fn fit_returns_error_for_invalid_max_features() {
let n_rows = 10;
let n_features = 4;
let x: DenseMatrix<f64> = DenseMatrix::from_iterator(
(0..n_features * n_rows).map(|k| ((k * 7919) % 101) as f64),
n_rows,
n_features,
0,
);
let y: Vec<f64> = (0..n_rows).map(|i| i as f64).collect();

let invalid = [
MaxFeatures::Count(0),
MaxFeatures::Count(n_features + 1),
MaxFeatures::Fraction(0.0),
MaxFeatures::Fraction(-0.5),
MaxFeatures::Fraction(1.5),
MaxFeatures::Fraction(f64::NAN),
MaxFeatures::Fraction(f64::INFINITY),
];

for splitter in [Splitter::Best, Splitter::Random] {
for m in invalid {
let params = BaseForestParameters {
max_depth: None,
min_samples_leaf: 1,
min_samples_split: 2,
n_trees: 3,
m,
keep_samples: false,
seed: Some(42),
bootstrap: true,
splitter,
};
let result = BaseForest::fit(&x, &y, None, &MSECriterion, params);
let err = result.err().unwrap_or_else(|| {
panic!("fit must fail for m = {m:?} (splitter: {splitter:?})")
});
assert_eq!(
err.error(),
FailedError::ParametersError,
"m = {m:?} (splitter: {splitter:?})"
);
}
}
}

#[cfg_attr(
all(target_arch = "wasm32", not(target_os = "wasi")),
wasm_bindgen_test::wasm_bindgen_test
)]
#[test]
fn fit_accepts_boundary_max_features() {
let n_rows = 10;
let n_features = 4;
let x: DenseMatrix<f64> = DenseMatrix::from_iterator(
(0..n_features * n_rows).map(|k| ((k * 7919) % 101) as f64),
n_rows,
n_features,
0,
);
let y: Vec<f64> = (0..n_rows).map(|i| i as f64).collect();

let valid = [
MaxFeatures::Sqrt,
MaxFeatures::All,
MaxFeatures::Count(1),
MaxFeatures::Count(n_features),
MaxFeatures::Fraction(1.0),
MaxFeatures::Fraction(0.01),
];

for m in valid {
let params = BaseForestParameters {
max_depth: None,
min_samples_leaf: 1,
min_samples_split: 2,
n_trees: 3,
m,
keep_samples: false,
seed: Some(42),
bootstrap: true,
splitter: Splitter::Best,
};
let forest = BaseForest::fit(&x, &y, None, &MSECriterion, params)
.unwrap_or_else(|e| panic!("fit must succeed for m = {m:?}: {e}"));
assert_eq!(forest.trees.unwrap().len(), 3, "m = {m:?}");
}
}
}
11 changes: 6 additions & 5 deletions src/ensemble/extra_trees_classifier.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,7 @@ use crate::api::{Predictor, SupervisedEstimator};
use crate::ensemble::base_forest::{
BaseForest, BaseForestParameters, predict_mode, predict_mode_oob,
};
use crate::ensemble::max_features::MaxFeatures;
use crate::error::Failed;
use crate::linalg::basic::arrays::{Array1, Array2};
use crate::numbers::basenum::Number;
Expand Down Expand Up @@ -112,8 +113,8 @@ pub struct ExtraTreesClassifierParameters {
/// The number of trees in the forest.
pub n_trees: usize,
#[cfg_attr(feature = "serde", serde(default))]
/// Number of random sample of predictors to use as split candidates.
pub m: Option<usize>,
/// Number of features to test at each split. See [`MaxFeatures`].
pub m: MaxFeatures,
#[cfg_attr(feature = "serde", serde(default))]
/// Whether to keep samples used for tree generation. This is required for OOB prediction.
pub keep_samples: bool,
Expand Down Expand Up @@ -169,8 +170,8 @@ impl ExtraTreesClassifierParameters {
self
}
/// Number of random sample of predictors to use as split candidates.
pub fn with_m(mut self, m: usize) -> Self {
self.m = Some(m);
pub fn with_m(mut self, m: impl Into<MaxFeatures>) -> Self {
self.m = m.into();
self
}
/// Seed used for feature selection and for the random split points of each tree.
Expand Down Expand Up @@ -209,7 +210,7 @@ impl Default for ExtraTreesClassifierParameters {
min_samples_leaf: 1,
min_samples_split: 2,
n_trees: 10,
m: Option::None,
m: MaxFeatures::Sqrt,
keep_samples: false,
seed: Option::None,
bootstrap: false,
Expand Down
13 changes: 7 additions & 6 deletions src/ensemble/extra_trees_regressor.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@ use crate::api::{Predictor, SupervisedEstimator};
use crate::ensemble::base_forest::{
BaseForest, BaseForestParameters, predict_mean, predict_mean_oob,
};
use crate::ensemble::max_features::MaxFeatures;
use crate::error::Failed;
use crate::linalg::basic::arrays::{Array1, Array2};
use crate::numbers::basenum::Number;
Expand All @@ -89,8 +90,8 @@ pub struct ExtraTreesRegressorParameters {
/// The number of trees in the forest.
pub n_trees: usize,
#[cfg_attr(feature = "serde", serde(default))]
/// Number of random sample of predictors to use as split candidates.
pub m: Option<usize>,
/// Number of features to test at each split. See [`MaxFeatures`].
pub m: MaxFeatures,
#[cfg_attr(feature = "serde", serde(default))]
/// Whether to keep samples used for tree generation. This is required for OOB prediction.
pub keep_samples: bool,
Expand Down Expand Up @@ -139,8 +140,8 @@ impl ExtraTreesRegressorParameters {
self
}
/// Number of random sample of predictors to use as split candidates.
pub fn with_m(mut self, m: usize) -> Self {
self.m = Some(m);
pub fn with_m(mut self, m: impl Into<MaxFeatures>) -> Self {
self.m = m.into();
self
}

Expand Down Expand Up @@ -170,7 +171,7 @@ impl Default for ExtraTreesRegressorParameters {
min_samples_leaf: 1,
min_samples_split: 2,
n_trees: 10,
m: Option::None,
m: MaxFeatures::Sqrt,
keep_samples: false,
seed: Option::None,
bootstrap: false,
Expand Down Expand Up @@ -594,7 +595,7 @@ mod tests {
#[test]
fn serde_without_bootstrap_field_uses_default() {
// Parameters serialized before the `bootstrap` field was added.
let json = r#"{"max_depth":null,"min_samples_leaf":1,"min_samples_split":2,"n_trees":10,"m":null,"keep_samples":true,"seed":42}"#;
let json = r#"{"max_depth":null,"min_samples_leaf":1,"min_samples_split":2,"n_trees":10,"keep_samples":true,"seed":42}"#;
let parameters: ExtraTreesRegressorParameters = serde_json::from_str(json).unwrap();
assert!(!parameters.bootstrap);

Expand Down
Loading
Loading