Skip to content

Faster rule extraction - #24

Merged
topepo merged 4 commits into
mainfrom
faster-rule-extraction
Oct 25, 2026
Merged

topepo merged 4 commits into
mainfrom
faster-rule-extraction

Conversation

@topepo

@topepo topepo commented Oct 25, 2026 •

Copy link
Copy Markdown
Member

This PR increases the speed of extracting rules. The main cause was that extract_rules.party() looked up each leaf's parent by looping over every node in the tree, which is quadratic or worse per tree. Several other methods had the same per-leaf *_find_parent pattern. Most users extract rules from every tree in an ensemble, so the work was benchmarked on that case.

Changes

Speedups

  • party, cforest, ranger, randomForest, grf: extract_rules.party() now does one top-down walk per tree that carries the split expressions down and emits a rule at each leaf. ranger, randomForest and grf also skipped a row-by-row computation of fitted node ids (nrow × depth per tree) by passing a zero-row copy of the data.
  • rpart: one top-down walk with a precomputed split-row lookup.
  • xgboost, lightgbm, bart, catboost: one model dump, getTrees() call or JSON export now covers all requested trees. Previously it was repeated per tree, which made the total quadratic. Each tree is then converted with a top-down walk.
  • aorsf (ObliqueForest): the slow R6 accessors (get_means(), get_stdev(), get_fctr_info(), get_names_x()) are cached once, and each tree is a top-down walk.
  • C5.0 and Cubist: already fast, unchanged apart from the shared checks below.

Output is identical() to the old code wherever the old code could produce output.

Timings

Method Before After Case
ranger 115 s 1.7 s per 20 trees 2 small trees before, 20 trees on 2,000 rows after
randomForest 422 s 1.7 s per 20 trees 2 small trees before, 20 trees on 2,000 rows after
grf 661 s 0.25 s 10 trees, 1,000 rows
rpart 4.1 s 0.12 s
bart 7.6 s 0.19 s 100 trees × 2 chains
lightgbm 6.8 s 0.25 s 50 trees
xgboost 1.2 s 0.19 s 40 trees
catboost 0.6–1.1 s 0.17 s 100 trees
aorsf 1.5 s 0.39 s 20 trees

The ranger and randomForest "before" and "after" rows are not the same workload, so they understate the improvement.

Test suite timings

File Before After
test-cforest.R skipped (estimated 1–2 hours if enabled) about 2 s
test-party.R skipped (7:43 when enabled) about 1.7 s

Both files were slow because of model fitting, not rule extraction. partykit::cforest() and ctree() were fit on wa_trees, where county has 30 levels (one 3-tree cforest fit took about 170 s, while extraction took 0.01–0.03 s). A new test helper, get_wa_trees_small_data(), keeps 200 rows and the 3 most common counties and labels the rest other. A 3-tree fit now takes 0.02 s. The skip() and the random skip_if(runif(1) <= 0.1) are removed, and tests/testthat/_snaps/cforest.md is new because those snapshots were never recorded.

API and behavior changes

  • tree accepts a vector for ObliqueForest, bart, catboost, lightgbm and xgboost, matching the other ensemble methods.
  • Shared validation: new internal check_tree_type() and check_tree_range() are used by every method with a tree argument (including C5.0 trials, Cubist committees, cforest and active_predictors.cforest()). The range error now states the upper limit, for example "tree must be between 1 and 3, the number of trees in the model."
  • Tree counts: num_ensemble_members() is the single source for the number of trees. Internal lgb_num_trees(), xgb_num_trees() and bart_num_trees() take an already-read dump so the model is not dumped a second time.
  • rlang::check_dots_empty() is now called first in every extract_rules() method, so typos such as trees = 2 error instead of being silently ignored. As a result data is no longer accepted by extract_rules.C5.0().
  • grf: the fallback of 1000 trees is removed, num_ensemble_members.grf() aborts clearly when the number of trees is missing, and matrices without column names no longer error.

Tests

  • New test that every training row matches exactly one party rule and that the rule id equals predict(type = "node").
  • New "handles several trees in one call" tests for lightgbm, xgboost, catboost, ObliqueForest and bart.
  • New tests with snapshots for check_tree_type() and check_tree_range().
  • Snapshots updated for the new error messages. The tree = c(1, 2) error snapshots were removed because that call is now valid. data = was removed from the C5.0 tests.
  • Full suite: 2,545 passing, 0 failing, 1 skipped (aorsf "On Mac").

@topepo
topepo marked this pull request as ready for review October 25, 2026 21:01
@topepo
topepo merged commit d2df22d into main Oct 25, 2026
13 checks passed
@topepo
topepo deleted the faster-rule-extraction branch October 26, 2026 12:38
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant