[slievens]

Fixes #526

Checklist

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

Current behaviour

A classification tree stores only the predicted class in each leaf. DecisionTreeClassifier::predict_proba returns 1 for the predicted class and 0 for the other classes, so it gives no information about how certain a prediction is.

New expected behaviour

Classification trees store the weighted class distribution of each leaf. predict_proba returns the weighted class frequencies of the training samples in the leaf (the same as scikit-learn).

Implementation:

• The leaf output of a classification tree is now ClassificationOutput { class_index, slot } (was usize). slot is the row of the leaf in BaseTree::leaf_probabilities, a flat Vec<f32> with n_leaves rows of n_classes values. Regression trees keep an empty vector.
• When a leaf splits, the true child takes the parent's slot and the false child appends a new slot. Thus the slots are always 0..n_leaves, and the vector has no gaps.
• Node<ClassificationOutput> is still 24 bytes (a test checks this). The three classification criteria share one ClassStats::leaf_value helper.
• BaseTree equality also compares the leaf distributions.

Random forest and extra trees store the distributions too, but they do not use them yet. A follow-up PR adds predict_proba to these forests.

Tests

• Slot allocation: n_leaves equals the number of leaf nodes, and the slots are exactly 0..n_leaves for both splitters and for different depths.
• Each distribution sums to 1, its argmax is the class prediction, and it equals the weighted class frequencies of the training rows in the leaf (with sample weights).
• Criterion tests check the stored rows: the append of a new slot, the overwrite of an existing slot, and the uniform distribution for empty stats.
• predict_proba for the Gini, Entropy and ClassificationError criteria, and for a root-only tree.
• Serde round trip keeps the leaf distributions, and equality detects a changed distribution.

Change logs

Changed

• Breaking change: tree, ensemble: the fitted trees store the number of leaves and, for classification, the class distribution of each leaf. Models serialized with an earlier version cannot be deserialized.
• tree: DecisionTreeClassifier::predict_proba returns the weighted class frequencies of the training samples in the leaf.

Note: the follow up issue #532 which implements predict_proba for random forest and extra trees classifier is ready to be merged after this one has landed.

[codecov]

Codecov Report

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

Additional details and impacted files

@@             Coverage Diff             @@
##             main     #533       +/-   ##
===========================================
+ Coverage   43.97%   87.15%   +43.17%     
===========================================
  Files          85      100       +15     
  Lines        7281     9962     +2681     
===========================================
+ Hits         3202     8682     +5480     
+ Misses       4079     1280     -2799     

☔ 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 for this PR, and for the clear description. The design is sound and fits the project's constraints: no unsafe, no FFI, pure Rust, and no allocation at predict time (leaf_distribution_for_row returns a slice into the flat store). Tests are thorough, and CI is green on wasm32, i686, Windows, macOS, MSRV 1.85 and all feature sets.

Verified in the diff

• The slot invariant holds. The early-return path of the split runs before parent_slot is read and before n_leaves is incremented. The true child overwrites the parent's row, and the false child appends exactly one row (start == store.len()). Slots are therefore always 0..n_leaves with no gaps. The tests cover this for both splitters and several depths.
• Node<ClassificationOutput> stays at 24 bytes, and the new size test guards that.
• The serde break is documented in the changelog. n_leaves and leaf_probabilities have no #[serde(default)], so old models fail to load instead of silently predicting wrong. I agree that is the right behaviour.

Requested changes / questions

1. Lint failure. The lint job failed on the latest run, although the checklist says linting was applied. fix(metrics): return 1.0 for zero-entropy labels in HCVScore #534 has the same failure on the same base commit, so it may already fail on main. Please check the log and say in the PR whether this is pre-existing.
2. f32 ties between predict and predict_proba. class_index comes from an f64 argmax of the counts, but the stored distribution is rounded to f32. Two classes with nearly equal weighted counts can round to the same f32, and then the f32 argmax picks the first class. predict_proba's argmax can then differ from predict. This is rare, but class_prediction_equals_argmax asserts the equality and could flake on other data. Either document that the equality holds only up to f32 precision, or use an assertion that tolerates ties.
3. Memory cost for forests. Each tree now stores n_leaves * n_classes * 4 bytes. Deep, unpruned forests with many classes will grow noticeably. For example, 500 trees with 20k leaves each and 100 classes is about 4 GB. Forests store the distributions but don't use them until Add predict_proba to RandomForestClassifier and ExtraTreesClassifier #532 lands. Please add a note on the expected overhead to the PR or the docs. It would also help to either merge Add predict_proba to RandomForestClassifier and ExtraTreesClassifier #532 right after this PR, or say that it should land in the same release.

Minor / nits

• LeafOutput::slot() returns a meaningless 0 for f64 (regressors). A narrower trait or a leaf_probabilities-aware extension for classification only would avoid putting a classification concept on the shared trait. This is not blocking.
• leaf_distribution_for_row derives n_classes as leaf_probabilities.len() / n_leaves. That is correct, but storing n_classes on the tree would make it explicit and cheaper. It is optional.
• The unrelated .map(|e| key(e)) to .map(&key) change in the test helper is fine, but it is outside the scope of this PR. Mention it in the description if it was a clippy fix.
• The PR adds many explanatory comments. Project guidelines ask to keep comments limited, so consider trimming the ones that restate the code (e.g. in the new tests).
• It would be good to add a DecisionTreeClassifier-level test of predict_proba with sample_weight. Today the weighted-frequency check is at the BaseTree level only.

Once item 1 is clarified and 2 is addressed, this looks good to merge from my side.

[slievens]

@Mec-iS

(1) lint failure does not happen at my side and is in parts of the source code unrelated to this PR.
(2) this is good catch. It was fixed and a test in criterion was added to catch it.