[slievens]

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

Fixes #491

Checklist

• [ x] My branch is up-to-date with main branch.
• [ x] Everything works and tested on latest stable Rust.
• [ x] Coverage and Linting have been applied

Current behaviour

BaseTreeRegressor has hardcoded MSE stats.

New expected behaviour

BaseTreeRegressor uses a MSECriterion.

Change logs

Added

• src/tree/criterion.rs File with the Criterion trait and implementations for MSE and Gini.

• src/tree/base_tree_regressor.rs Uses MSECriterion.
• src/tree/mod.rs Declares criterion as pub(mod)

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 85.62%. Comparing base (9eaae9e) to head (c7eb8a3).
⚠️ Report is 202 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff             @@
##             main     #492       +/-   ##
===========================================
+ Coverage   43.97%   85.62%   +41.64%     
===========================================
  Files          85       97       +12     
  Lines        7281     9396     +2115     
===========================================
+ Hits         3202     8045     +4843     
+ Misses       4079     1351     -2728     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Mec-iS]

Thanks @slievens. Centralising the sufficient statistics in one place is a good direction, and recomputing child leaf values from true_stats/false_stats in split_node is cleaner than carrying true_child_output/false_child_output through the visitor. Notes from reading the diff:

Process

• GiniCriterion/GiniStats (plus roughly 250 lines of tests) are unused in this PR and are kept alive with #[cfg_attr(not(test), expect(dead_code))]. The title and description only cover the regressor. Please move the Gini code to the follow-up PR that wires it into the classifier.
• Project conventions are to limit comments and not to add TODOs. // TODO: add the entropy and classification-error criteria. should become an issue. The long module-level docs and several inline comments could be trimmed.
• #[expect(...)] needs Rust 1.81+. Please confirm it fits the crate's MSRV. It goes away if the Gini code is dropped.

Possible behaviour changes in base_tree_regressor.rs

1. find_best_split now seeds prevx from sorted_by_feature[j][visitor.start_idx]. The old Option handled an empty range, but this indexes unconditionally. It panics if start_idx == end_idx (or start_idx == len). find_random_split guards this case explicitly. Please add the same early return, or show that the callers guarantee a non-empty range.
2. The n < min_samples_split check moved out of find_best_cutoff. The root now checks it explicitly, but the children rely on true_can_split/false_can_split. Please confirm those conditions include min_samples_split, and add a test where min_samples_split is the binding constraint on a child.
3. split_score meaning changes. The search functions now store the proxy (sum²/mass terms, without the parent term), and split_node overwrites it with impurity_improvement only when the split is accepted. If a node is evaluated but not split (for example min_samples_leaf rejects it in split_node), it keeps the proxy value. Please check whether split_score is read anywhere else, such as feature importance, and either always store the real improvement or use a separate field for the proxy.
4. difference() computes parent - part for both sum and mass. With fractional sample weights, the residue can leave mass at a tiny positive value, so sum²/mass is evaluated on noise. The old code had the same subtraction, so this is not a regression. A small relative epsilon on the mass > 0.0 guards, or a test with fractional weights and a near-empty side, would document the behaviour.

Tests and evidence

• All new tests are unit tests of the criterion. Please add a regression test that fits the pre-PR and post-PR regressor on fixed seeds (both Splitter::Best and Splitter::Random, with and without sample_weight, and with max_features < n_features). It should assert identical split_feature, split_value and predictions. This is what shows the refactor is behaviour-preserving.
• Please include a before/after fit-time benchmark. The change adds scratch stats to SplitWorkspace and calls through #[inline(always)] methods in the hot loop, so it is probably neutral, but the numbers would confirm it.

Design

• The Criterion trait is generic, but the regressor hard-codes MSECriterion at every call site. Until a second regressor criterion exists, the trait is speculative. Either make BaseTreeRegressor generic over C: Criterion, or keep the trait minimal and add the generic parameter when it is needed.

Happy to re-review after the retarget and the empty-range fix.

[Mec-iS]

Thanks @slievens, moving the regressor onto a Criterion abstraction is a good direction, and the incremental MeanStats + difference approach (as in sklearn) keeps the maths equivalent to the old gain. I reviewed the diff and have the following points.

Possible behaviour changes (please verify)

1. min_samples_split was removed from find_best_cutoff and is now checked only for the root (n >= min_samples_split in fit). The child calls (find_best_cutoff(&mut true_visitor, ..., true_count) and the false one) are guarded only by true_can_split / false_can_split. Unless those flags already include min_samples_split, children smaller than min_samples_split will now be split further. Please confirm, or restore the check inside find_best_cutoff.
2. In find_best_split, workspace.sorted_by_feature[j][visitor.start_idx] is indexed before the loop. find_random_split guards start_idx == end_idx, but this function does not, so an empty range would panic. The old Option logic did not have this problem. Add an early return, or explain why the range can never be empty.
3. false_stats is computed by parent - true at every level, so floating-point residue accumulates down the false path. A child that should have zero mass can end up with mass of about 1e-16 and a garbage leaf_value, because the mass > 0.0 guard passes. The old code had the same weakness, but it is now centralised, so a tolerance or a direct recompute is cheap to add.
4. split_score now holds the proxy during the search and is overwritten with impurity_improvement in split. If any other code reads split_score before split succeeds (for example feature importance on a node whose split was rejected), it will see the proxy value. Please check the consumers.

Tests

• All the new tests are unit tests of criterion.rs. There is no tree-level regression test showing that fitted trees (split features, thresholds, leaf values, feature_importances) are unchanged against the previous implementation. A golden test on a fixed seed with sample weights, Splitter::Random and Splitter::Best, and min_samples_leaf / min_samples_split > 1 would cover the points above.
• Please add a benchmark or a timing note for fit before and after, since add / difference now run in the hot loop.

Scope and project conventions

• GiniCriterion, GiniStats and about 15 Gini tests are not used by this PR (hence the expect(dead_code) attributes). This is dead code in the PR, and it adds roughly 450 lines to a PR titled "Use MSE criterion in base tree regressor". Please move them to the follow-up PR that wires up the classifier.
• #[expect(...)] needs Rust 1.81 or later. Please check it against the crate's MSRV and the CI toolchain.
• // TODO: add the entropy and classification-error criteria. goes against the project rule of opening an issue instead of leaving TODOs. Please remove it and open an issue.
• The target branch is main. The project convention is development, so please retarget.
• GiniStats holds a Vec, so empty_stats() allocates per node. That is fine for now, but it is worth noting before the classifier uses it, given the zero-copy goals.
• SplitWorkspace hardcodes MeanStats, so the trait is not yet usable generically. Please say whether making the workspace generic over Criterion::Stats is planned for the classifier PR.

Nits

• The for k in 0..mtry.min(workspace.variables.len()) change is needed to avoid the borrow conflict. A short comment would help, or keep the .take(mtry) form if it compiles.
• The module doc cites sklearn's _criterion.pyx. Please add the proxy-improvement derivation ((\sum_{side} S^2/W)) briefly to the Criterion docs.

Happy to re-review once 1 and 2 are clarified.

[slievens]

1. min_samples_split for children

Confirmed: the child flags include it. In split():

true_can_split = child_level < max_depth && true_count >= min_samples_split

(same for false), and find_best_cutoff is only called when the flag is true. true_count/false_count are sample counts

min_samples_split_one_child_only covers the case where min_samples_split is the binding constraint: the false child (2 samples, min_samples_split = 3) stays a leaf and the true child (exactly 3) splits. Removing false_count >= min_split makes it fail.

2. Empty range in find_best_split

The seed read only panics if start_idx == len. With start_idx == end_idx < len the index is in bounds and the loop does not run.

start_idx == len cannot happen:

• A true child starts at its parent's (non-empty) start.
• A false child is never empty, because the threshold is always below the first false value.
• The root is empty only if all sample counts are 0, which no caller produces (fit_inner uses all ones, and the bootstrap draws n ≥ 1 samples).

find_random_split needs its guard for a different reason: it reads end_idx - 1 and calls random_range.

3. Residue in parent - true

Agreed, and confirmed with a test: zero_mass_leaf_with_fractional_weights (weights 0.1/0.2/0.3/0.0) produced a leaf value of 1.0 from residue / residue.

• MSECriterion::difference now treats a side as empty when its mass is ≤ 1e-12 × parent mass. This also catches negative residue.
• Unit tests cover the residue case and a small real mass (1e-6), which is kept.

4. split_score consumers

split_score is currently only read by the search, Node::eq, and tests. A split that min_samples_leaf rejects resets it to None.

The only leak was the root with max_depth <= 1, which was searched but never split. The root is no longer searched in that case.

Two new tests enforce the invariant:

• split_score_is_some_only_for_internal_nodes (Some exactly for internal nodes)
• split_scores_add_up_to_total_impurity_decrease (the stored values are real improvements, not the proxy)

The stored values can be used to compute feature importances which will be done in a next PR.

5. speed

It seems the new implementation is slightly slower than the existing one. I think it is worth it given that we will be able to reuse one common implementation for classifiers and regressors.

I will look into making the overall implementation faster still.

6. Removal of currently unused code

All Gini-related code which would be used on a next PR (but not on this one) is removed.

[Mec-iS]

Thanks @slievens. I read the current diff (head d4e67bc) together with the replies. The Criterion / MeanStats split is clean, and the invariant tests (split_score_is_some_only_for_internal_nodes, split_scores_add_up_to_total_impurity_decrease, zero_mass_leaf_with_fractional_weights) are good additions. Points below are mostly new or still open.

Still open from the previous reviews

1. No equivalence test against the pre-PR implementation. The new tree-level tests check invariants, but not that the refactor preserves behaviour. Before this merges I would like a golden test, with expected values captured from main, over Splitter::Best and Splitter::Random, with and without sample_weight, with max_features < n_features, and with min_samples_leaf / min_samples_split > 1. It should assert split_feature, split_value, output and predictions. The old gain was true_mass*true_mean^2 + false_mass*false_mean^2 - parent_gain, which equals the new impurity_improvement. split_score can therefore be compared to the old value within a tolerance. That is a strong check, and it also covers Node::eq.
2. Speed: "slightly slower" needs numbers. Please post a before/after fit time (for example the existing benches or a small Instant harness, with n and p stated). Also try plain #[inline] instead of #[inline(always)] on add / difference / proxy_improvement. With a unit struct and a single impl, the compiler will normally inline these anyway, and always can hurt other call sites later.
3. PR is still based on main; the convention is development. Mergeable state is blocked.
4. The PR description still lists Gini in the change log and says "pub(mod)", but the code has pub(crate) mod criterion; and no Gini. Please update the description.

New points

1. root_stats is summed in row order, but the true_stats in split_node are summed in sorted-feature order, and false_stats = parent - true. The MASS_EPS clamp fixes the symptom (the 0.1/0.2/0.3 case). However, the error still accumulates down the false path. The clamp is also relative to the parent mass, so a real side with weights below 1e-12 * parent mass is now treated as empty. That is a small behaviour change, so it is worth one line in the docs. split_node already loops over col[n_true..] to count false_count. Adding MSECriterion.add(&mut false_stats, ...) there costs almost nothing and gives exact child stats for the leaf value and the child visitor. The difference call would then stay only in the search loops, where the sklearn-style incremental approach makes sense.
2. split_score_is_some_only_for_internal_nodes and split_scores_add_up_to_total_impurity_decrease assert tree.nodes().len() > 1 for every combination, including Splitter::Random with min_samples_leaf = 3 on 10 rows. That depends on the seed and on the rand version. If the default seed is fixed, it is fine. If not, skip or relax the assertion for Random so a rand upgrade does not make CI fail.
3. Root guard: max_depth.unwrap_or(u16::MAX) > 1 is new in fit. Please confirm with the golden test (point 1 above) that trees with max_depth of Some(0) / Some(1) have the same node count as on main. If the old code produced a different structure there, that is a behaviour change to mention in the changelog.
4. find_best_split reads sorted_by_feature[j][visitor.start_idx] without a guard. Your argument that the range is never empty is sound for current callers. A debug_assert!(visitor.start_idx < visitor.end_idx) would document the invariant at no runtime cost, and protects the reader from the asymmetry with find_random_split.
5. find_random_split breaks out of the loop at the first value above split_value. This relies on the column being sorted ascending. That was also true before, but the comment above the loop would help.
6. Minor: Output and Target associated types are not used by any generic code yet (every call site names MSECriterion). I agree with the earlier "speculative trait" comment. It is acceptable as the first step if the next PR makes BaseTreeRegressor / SplitWorkspace generic over C: Criterion. Please state that plan in the description.

After the golden test and the benchmark are in, I am happy to approve.

[slievens]

tree: Criterion trait with numerically stable MSE split improvement

Fixes #491

Checklist

• My branch is up-to-date with main branch.
• Everything works and tested on latest stable Rust.
• [ x] Coverage and Linting have been applied

Current behaviour

BaseTreeRegressor keeps the split statistics (sum, mass) as separate variables in the tree code. The split ranking uses S_T²/M_T + S_F²/M_F.

When the targets have a large offset compared to their spread (for example y ≈ 1e9 ± 1), this formula loses all significant digits. As a result:

• The tree selects splits at random, and split_score can be 0 or negative.
• Rows with sample weight 0 stay in the nodes. They count toward min_samples_leaf/min_samples_split, they can move a split threshold, and they can make a leaf with zero mass.

New expected behaviour

• A new Criterion trait (src/tree/criterion.rs) keeps the split statistics. MSECriterion is its first implementation. This is the first step toward using one tree implementation for regression and classification.

• The MSE improvement uses the stable form (S_T·M_F − S_F·M_T)² / (M_T·M_F·M_P). The selected splits do not change when a constant is added to y.

• split_score holds the real (mass-weighted) impurity decrease, not the proxy. The sum of the split_score values of a tree is the total decrease of the impurity.

• Rows with mass 0 (bootstrap count 0 or weight 0) are removed when the tree grows.

• Performance:

  Benchmark	Result
  compare_rf_bench.sh (own test-suite) (1000–10000 rows, 50–200 features, 100–500 trees)	Same speed as main (geometric mean speedup: 1.02)
  criterion random_forest_regressor bench (very small trees)	fit: 2–9% slower
  predict	No change

• Fitted trees can differ from main, because exactly tied splits can be broken in a different way. On the bench data, the change in test MSE has no direction (mean −0.2%).

Change logs

Added

• tree: Criterion trait with MSECriterion.
• Tests: invariance to a target offset, zero-weight rows, split_score sums.

Changed

• tree: numerically stable MSE split improvement.
• tree: rows with sample weight 0 no longer count toward min_samples_leaf/min_samples_split, and they do not move split thresholds.

[Mec-iS]

This looks good. This is a relevant change in API for developers so my idea is to keep this for version 0.7.

Let me know. I will deal into preparing the release in the next weeks, together with the other features already in the queue.

[Mec-iS]

Also we have to be extra careful because these kind of spikes can happen:

[slievens]

Tree-based models: planned work

I would like to continue working on the tree-based classifiers. Things I still would like to do:

• Add feature importances to the regressors.
• Move the classifiers onto the common base_tree.
  • ⚠️ This will break the serialization format of those trees.
• Make building the trees more efficient.
• Look at the Node struct and see if it can be made more efficient.
  • ⚠️ This will also break the serialization format of regression trees.
• Add an ExtraTreesClassifier.
• Try to remove some inefficiencies from the predict methods.
• Try to compute the OOB scores without having to store the samples.
• Add "poisson" as a criterion for regression trees.

Question

If you agree that these changes are worthwhile, where should I implement them? If different branches are used, we will get a lot of merge conflicts.

[Mec-iS]

you can open different sequential PRs and specify "Depends on #..."