⚡💥 Bankai Random Forest

Bankai is a high-performance random forest for Python, powered by a native Rust core based on a maintained fork of XRF, the engine behind FRU. It combines parallel tree building and efficient permutation importance with optional LightGBM-style histograms for continuous features.

Preliminary results from an earlier run of the Colab benchmark reported Bankai's fit-plus-OOB-permutation workload 21× faster than scikit-learn fit plus validation permutation importance, and Bankai histogram fit 1.8× faster than LightGBM RF. These figures have not been remeasured with the notebook's current 1,000-row setup and are workload- and hardware-specific; see the notebook for current settings and comparison caveats.

Bankai follows the familiar scikit-learn estimator API, and full compatibility is a project goal. The project is in alpha, so APIs and serialized models may change.

Install

Python 3.11 or newer is required. Install the package from PyPI:

python -m pip install bankai-random-forest

Quick start

from sklearn.datasets import load_breast_cancer
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split

from bankai_random_forest import BankaiRandomForestClassifier

X, y = load_breast_cancer(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

model = BankaiRandomForestClassifier(
    n_estimators=200,
    max_bins=63,
    binning_strategy="sampled_select",
    importance_type="permutation",
    n_jobs=-1,
    random_state=42,
)
model.fit(X_train, y_train)

predictions = model.predict(X_test)
probabilities = model.predict_proba(X_test)
print(f"Accuracy: {accuracy_score(y_test, predictions):.3f}")
print("Probabilities:", probabilities[:3])
print("Feature importances:", model.feature_importances_)

max_bins enables histogram-based split search in this example. Set max_bins=None to use exact split search. Histogram cuts can change the fitted trees and the accuracy/speed trade-off; evaluate both on your data when that matters.

Training options

Bankai follows the familiar RandomForestClassifier estimator pattern and supports common options such as n_estimators, criterion, max_depth, min_samples_split, min_samples_leaf, max_features, max_leaf_nodes, bootstrap, max_samples, class_weight, random_state, n_jobs, oob_score, and ccp_alpha.

Sampled binning strategies use a deterministic subset of rows to construct feature cuts. Selection strategies avoid fully sorting each feature, while repeated feature values can produce cuts that differ from the sort strategies. For reproducible model comparisons, keep the strategy, sample size, seed, and thread count fixed.

API and supported inputs

• Dense NumPy arrays and SciPy CSR/CSC matrices are supported. Numeric features are expected; encode categorical columns in a scikit-learn transformer such as OneHotEncoder.
• NaN feature values are supported in dense and sparse inputs and are routed by the fitted trees. Infinite values are rejected.
• Binary, multiclass, multioutput, and multilabel classification targets are supported. Bankai trains a separate native forest for each output, so its multioutput trees and probabilities need not match scikit-learn's shared-tree implementation.
• feature_importances_ defaults to gain-based importance; its values are not guaranteed to match scikit-learn's impurity importance. Bankai also provides split-count and out-of-bag permutation importance.
• warm_start=True does not append only new trees: a later fit rebuilds the forest. For multioutput models, estimators_, apply, decision_path, and TreeSHAP inspection expose the first output forest.

Tested with scikit-learn 1.9.1 (>=1.9.1,<1.10).

Feature explanations

SHAP is optional and is not a runtime dependency. For TreeSHAP, pass a fitted single-output classifier directly to TreeExplainer:

import shap

explainer = shap.TreeExplainer(model)
explanation = explainer(X_test[:10])

See the SHAP compatibility notes for output shapes, multioutput limitations, and model-agnostic alternatives.

Performance

Fit time depends on the data, tree settings, hardware, and importance calculation. The benchmark report includes stage breakdowns, later Bankai measurements, and comparison caveats; results are specific to the recorded workload and are not a performance guarantee.

Run the binary classification fit benchmark in Google Colab or view the notebook in this repository. It compares scikit-learn RF, LightGBM RF, and Bankai on one shared dataset with n_jobs=-1; the first cell installs the packages. The scikit-learn comparison times fit plus validation permutation importance against Bankai fit with OOB permutation importance. The LightGBM comparison remains fit-only with gain on both models.

Save and restore models

Fitted models can be saved with joblib:

import joblib

joblib.dump(model, "classifier.joblib")
restored = joblib.load("classifier.joblib")
predictions = restored.predict(X_test)

Bankai's serialized state includes training arrays and rebuilds the native forest when loaded. Loading therefore has reconstruction cost and model files can be large. Keep Bankai and Python versions consistent when restoring a model, and only load joblib files from trusted sources. See the serialization notes.

Development

Building from source requires Python 3.11 or newer, Rust, and uv:

uv sync --locked --no-install-project
uv run maturin develop --release --locked
uv run pytest -q
uv run cargo test --workspace --locked

Install optional benchmark dependencies with uv sync --locked --group benchmark. The roadmap tracks project status and planned work.

License

Bankai is licensed under GPL-3.0-or-later. The maintained XRF fork is also GPL-3.0-or-later; see NOTICE for upstream attribution.