Scikit-Learn Estimator Wrappers for BART¶

stochtree.BARTModel is fundamentally a Bayesian interface in which users specify a prior, provide data, sample from the posterior, and manage / inspect the resulting posterior samples.

However, the basic BART model

$$y_i \sim \mathcal{N}\left(f(X_i), \sigma^2\right)$$

involves samples of a nonparametric function $f$ which estimates the expected value of $y$ given $X$. Averaging over these draws, the posterior mean $\bar{f}$ alone may satisfy some supervised learning use cases. In order to serve this use case straightforwardly, we offer scikit-learn-compatible estimator wrappers around BARTModel which implement the familiar API of sklearn models.

For continuous outcomes, the stochtree.StochTreeBARTRegressor class provides fit, predict and score methods.

For binary outcomes (deployed via probit BART), the stochtree.StochTreeBARTBinaryClassifier class provides fit, predict, predict_proba, decision_function, and score methods.

Users can fit multi-class classifiers by wrapping a OneVsRestClassifier around StochTreeBARTBinaryClassifier.

We begin by loading necessary libraries

In [1]:

Copied!





import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.datasets import load_wine, load_breast_cancer
from sklearn.model_selection import GridSearchCV
from sklearn.multiclass import OneVsRestClassifier
from stochtree import (
  StochTreeBARTRegressor, 
  StochTreeBARTBinaryClassifier, 
)
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.datasets import load_wine, load_breast_cancer
from sklearn.model_selection import GridSearchCV
from sklearn.multiclass import OneVsRestClassifier
from stochtree import (
  StochTreeBARTRegressor, 
  StochTreeBARTBinaryClassifier, 
)

Next, we seed a random number generator

In [2]:

Copied!

random_seed = 1234
rng = np.random.default_rng(random_seed)
random_seed = 1234
rng = np.random.default_rng(random_seed)

BART Regression via `sklearn` Estimator¶

We simulate some simple regression data to demonstrate the continuous outcome use case

In [3]:

Copied!





n = 100
p = 10
X = rng.normal(size=(n, p))
y = X[:, 0] * 3 + rng.normal(size=n)
n = 100
p = 10
X = rng.normal(size=(n, p))
y = X[:, 0] * 3 + rng.normal(size=n)

We fit a BART regression model by initializing a StochTreeBARTRegressor and calling its fit() method.

Since stochtree.BARTModel is configured primarily through parameter dictionaries, any downstream parameters that we wish to set are passed through as parameter dictionaries. In this case, we only specify the random seed.

In [4]:

Copied!

reg = StochTreeBARTRegressor(general_params={"random_seed": random_seed})
reg.fit(X, y)
reg = StochTreeBARTRegressor(general_params={"random_seed": random_seed})
reg.fit(X, y)

Out[4]:

StochTreeBARTRegressor(general_params={'random_seed': 1234})

In a Jupyter environment, please rerun this cell to show the HTML representation or trust the notebook.
On GitHub, the HTML representation is unable to render, please try loading this page with nbviewer.org.

Now, we can predict from this model and compare the (posterior mean) predictions to the true outcome

In [5]:

Copied!





pred = reg.predict(X)
plt.scatter(pred, y)
plt.xlabel("Predicted")
plt.ylabel("Actual")
plt.show()
pred = reg.predict(X)
plt.scatter(pred, y)
plt.xlabel("Predicted")
plt.ylabel("Actual")
plt.show()

No description has been provided for this image

We can also test the determinism of the model by running it again with the same seed and comparing predictions to the first model

In [6]:

Copied!





reg2 = StochTreeBARTRegressor(general_params={"random_seed": random_seed})
reg2.fit(X, y)
pred2 = reg2.predict(X)
plt.scatter(pred, pred2)
plt.xlabel("First model")
plt.ylabel("Second model")
plt.show()
reg2 = StochTreeBARTRegressor(general_params={"random_seed": random_seed})
reg2.fit(X, y)
pred2 = reg2.predict(X)
plt.scatter(pred, pred2)
plt.xlabel("First model")
plt.ylabel("Second model")
plt.show()

Cross-Validating a BART Model¶

While the default hyperparameters of stochtree.BARTModel are designed to work well "out of the box," we can use posterior mean prediction error to cross-validate the model's parameters.

Below we use grid search to consider the effect of several BART parameters:

Number of GFR iterations (num_gfr)
Number of MCMC iterations (num_mcmc)
num_trees, alpha and beta for the mean forest

In [7]:

Copied!





param_grid = {
    "num_gfr": [10, 40],
    "num_mcmc": [0, 1000],
    "mean_forest_params": [
        {"num_trees": 50, "alpha": 0.95, "beta": 2.0},
        {"num_trees": 100, "alpha": 0.90, "beta": 1.5},
        {"num_trees": 200, "alpha": 0.85, "beta": 1.0},
    ],
}
grid_search = GridSearchCV(
    estimator=StochTreeBARTRegressor(),
    param_grid=param_grid,
    cv=5,
    scoring="r2",
    n_jobs=-1,
)
grid_search.fit(X, y)
param_grid = {
    "num_gfr": [10, 40],
    "num_mcmc": [0, 1000],
    "mean_forest_params": [
        {"num_trees": 50, "alpha": 0.95, "beta": 2.0},
        {"num_trees": 100, "alpha": 0.90, "beta": 1.5},
        {"num_trees": 200, "alpha": 0.85, "beta": 1.0},
    ],
}
grid_search = GridSearchCV(
    estimator=StochTreeBARTRegressor(),
    param_grid=param_grid,
    cv=5,
    scoring="r2",
    n_jobs=-1,
)
grid_search.fit(X, y)

Out[7]:

GridSearchCV(cv=5, estimator=StochTreeBARTRegressor(), n_jobs=-1,
             param_grid={'mean_forest_params': [{'alpha': 0.95, 'beta': 2.0,
                                                 'num_trees': 50},
                                                {'alpha': 0.9, 'beta': 1.5,
                                                 'num_trees': 100},
                                                {'alpha': 0.85, 'beta': 1.0,
                                                 'num_trees': 200}],
                         'num_gfr': [10, 40], 'num_mcmc': [0, 1000]},
             scoring='r2')

In a Jupyter environment, please rerun this cell to show the HTML representation or trust the notebook.
On GitHub, the HTML representation is unable to render, please try loading this page with nbviewer.org.

In [8]:

Copied!





cv_best_ind = np.argwhere(grid_search.cv_results_['rank_test_score'] == 1).item(0)
best_num_gfr = grid_search.cv_results_['param_num_gfr'][cv_best_ind].item(0)
best_num_mcmc = grid_search.cv_results_['param_num_mcmc'][cv_best_ind].item(0)
best_mean_forest_params = grid_search.cv_results_['param_mean_forest_params'][cv_best_ind]
best_num_trees = best_mean_forest_params['num_trees']
best_alpha = best_mean_forest_params['alpha']
best_beta = best_mean_forest_params['beta']
print_message = f"""
Hyperparameters chosen by grid search: 
  num_gfr: {best_num_gfr} 
  num_mcmc: {best_num_mcmc} 
  num_trees: {best_num_trees} 
  alpha: {best_alpha} 
  beta: {best_beta}
"""
print(print_message)
cv_best_ind = np.argwhere(grid_search.cv_results_['rank_test_score'] == 1).item(0)
best_num_gfr = grid_search.cv_results_['param_num_gfr'][cv_best_ind].item(0)
best_num_mcmc = grid_search.cv_results_['param_num_mcmc'][cv_best_ind].item(0)
best_mean_forest_params = grid_search.cv_results_['param_mean_forest_params'][cv_best_ind]
best_num_trees = best_mean_forest_params['num_trees']
best_alpha = best_mean_forest_params['alpha']
best_beta = best_mean_forest_params['beta']
print_message = f"""
Hyperparameters chosen by grid search: 
  num_gfr: {best_num_gfr} 
  num_mcmc: {best_num_mcmc} 
  num_trees: {best_num_trees} 
  alpha: {best_alpha} 
  beta: {best_beta}
"""
print(print_message)

Hyperparameters chosen by grid search: 
  num_gfr: 10 
  num_mcmc: 1000 
  num_trees: 50 
  alpha: 0.95 
  beta: 2.0

BART Classification via `sklearn` Estimator¶

Now, we demonstrate the same functionality with binary and categorical outcomes, which require working with the StochTreeBARTBinaryClassifier class (and a wrapper for multi-class outcomes).

First, we load a dataset from sklearn with a binary outcome.

In [9]:

Copied!

dataset = load_breast_cancer()
X = dataset.data
y = dataset.target
dataset = load_breast_cancer()
X = dataset.data
y = dataset.target

And we fit a binary classification model as follows

In [10]:

Copied!

clf = StochTreeBARTBinaryClassifier(general_params={"random_seed": random_seed})
clf.fit(X=X, y=y)
clf = StochTreeBARTBinaryClassifier(general_params={"random_seed": random_seed})
clf.fit(X=X, y=y)

Out[10]:

StochTreeBARTBinaryClassifier(general_params={'random_seed': 1234})

In a Jupyter environment, please rerun this cell to show the HTML representation or trust the notebook.
On GitHub, the HTML representation is unable to render, please try loading this page with nbviewer.org.

In addition to class predictions, we can compute and visualize the predicted probability of each class via predict_proba().

In [11]:

Copied!

probs = clf.predict_proba(X)
plt.hist(probs[:, 1], bins=30)
plt.show()
probs = clf.predict_proba(X)
plt.hist(probs[:, 1], bins=30)
plt.show()

Now, we load a multi-class classification dataset from sklearn.

In [12]:

Copied!

dataset = load_wine()
X = dataset.data
y = dataset.target
dataset = load_wine()
X = dataset.data
y = dataset.target

And fit a multi-class classification model by wrapping a OneVsRestClassifier around StochTreeBARTBinaryClassifier.

In [13]:

Copied!





clf = OneVsRestClassifier(
    StochTreeBARTBinaryClassifier(general_params={"random_seed": random_seed})
)
clf.fit(X=X, y=y)
clf = OneVsRestClassifier(
    StochTreeBARTBinaryClassifier(general_params={"random_seed": random_seed})
)
clf.fit(X=X, y=y)

Out[13]:

OneVsRestClassifier(estimator=StochTreeBARTBinaryClassifier(general_params={'random_seed': 1234}))

In a Jupyter environment, please rerun this cell to show the HTML representation or trust the notebook.
On GitHub, the HTML representation is unable to render, please try loading this page with nbviewer.org.

And we visualize the histogram of predicted probabilities for each outcome category.

In [14]:

Copied!





fig, (ax1, ax2, ax3) = plt.subplots(3, 1)
fig.tight_layout(pad=3.0)
probs = clf.predict_proba(X)
ax1.hist(probs[y == 0, 0], bins=30)
ax1.set_title("Predicted Probabilities for Class 0")
ax1.set_xlim(0, 1)
ax2.hist(probs[y == 1, 1], bins=30)
ax2.set_title("Predicted Probabilities for Class 1")
ax2.set_xlim(0, 1)
ax3.hist(probs[y == 2, 2], bins=30)
ax3.set_title("Predicted Probabilities for Class 2")
ax3.set_xlim(0, 1)
plt.show()
fig, (ax1, ax2, ax3) = plt.subplots(3, 1)
fig.tight_layout(pad=3.0)
probs = clf.predict_proba(X)
ax1.hist(probs[y == 0, 0], bins=30)
ax1.set_title("Predicted Probabilities for Class 0")
ax1.set_xlim(0, 1)
ax2.hist(probs[y == 1, 1], bins=30)
ax2.set_title("Predicted Probabilities for Class 1")
ax2.set_xlim(0, 1)
ax3.hist(probs[y == 2, 2], bins=30)
ax3.set_title("Predicted Probabilities for Class 2")
ax3.set_xlim(0, 1)
plt.show()

	estimator estimator: estimator object This is assumed to implement the scikit-learn estimator interface. Either estimator needs to provide a ``score`` function, or ``scoring`` must be passed.	StochTreeBARTRegressor()
	param_grid param_grid: dict or list of dictionaries Dictionary with parameters names (`str`) as keys and lists of parameter settings to try as values, or a list of such dictionaries, in which case the grids spanned by each dictionary in the list are explored. This enables searching over any sequence of parameter settings.	{'mean_forest_params': [{'alpha': 0.95, 'beta': 2.0, 'num_trees': 50}, {'alpha': 0.9, 'beta': 1.5, 'num_trees': 100}, ...], 'num_gfr': [10, 40], 'num_mcmc': [0, 1000]}
	scoring scoring: str, callable, list, tuple or dict, default=None Strategy to evaluate the performance of the cross-validated model on the test set. If `scoring` represents a single score, one can use: - a single string (see :ref:`scoring_string_names`); - a callable (see :ref:`scoring_callable`) that returns a single value; - `None`, the `estimator`'s :ref:`default evaluation criterion ` is used. If `scoring` represents multiple scores, one can use: - a list or tuple of unique strings; - a callable returning a dictionary where the keys are the metric names and the values are the metric scores; - a dictionary with metric names as keys and callables as values. See :ref:`multimetric_grid_search` for an example.	'r2'
	n_jobs n_jobs: int, default=None Number of jobs to run in parallel. ``None`` means 1 unless in a :obj:`joblib.parallel_backend` context. ``-1`` means using all processors. See :term:`Glossary ` for more details. .. versionchanged:: v0.20 `n_jobs` default changed from 1 to None	-1
	refit refit: bool, str, or callable, default=True Refit an estimator using the best found parameters on the whole dataset. For multiple metric evaluation, this needs to be a `str` denoting the scorer that would be used to find the best parameters for refitting the estimator at the end. Where there are considerations other than maximum score in choosing a best estimator, ``refit`` can be set to a function which returns the selected ``best_index_`` given ``cv_results_``. In that case, the ``best_estimator_`` and ``best_params_`` will be set according to the returned ``best_index_`` while the ``best_score_`` attribute will not be available. The refitted estimator is made available at the ``best_estimator_`` attribute and permits using ``predict`` directly on this ``GridSearchCV`` instance. Also for multiple metric evaluation, the attributes ``best_index_``, ``best_score_`` and ``best_params_`` will only be available if ``refit`` is set and all of them will be determined w.r.t this specific scorer. See ``scoring`` parameter to know more about multiple metric evaluation. See :ref:`sphx_glr_auto_examples_model_selection_plot_grid_search_digits.py` to see how to design a custom selection strategy using a callable via `refit`. See :ref:`this example ` for an example of how to use ``refit=callable`` to balance model complexity and cross-validated score. .. versionchanged:: 0.20 Support for callable added.	True
	cv cv: int, cross-validation generator or an iterable, default=None Determines the cross-validation splitting strategy. Possible inputs for cv are: - None, to use the default 5-fold cross validation, - integer, to specify the number of folds in a `(Stratified)KFold`, - :term:`CV splitter`, - An iterable yielding (train, test) splits as arrays of indices. For integer/None inputs, if the estimator is a classifier and ``y`` is either binary or multiclass, :class:`StratifiedKFold` is used. In all other cases, :class:`KFold` is used. These splitters are instantiated with `shuffle=False` so the splits will be the same across calls. Refer :ref:`User Guide ` for the various cross-validation strategies that can be used here. .. versionchanged:: 0.22 ``cv`` default value if None changed from 3-fold to 5-fold.	5
	verbose verbose: int Controls the verbosity: the higher, the more messages. - >1 : the computation time for each fold and parameter candidate is displayed; - >2 : the score is also displayed; - >3 : the fold and candidate parameter indexes are also displayed together with the starting time of the computation.	0
	pre_dispatch pre_dispatch: int, or str, default='2n_jobs' Controls the number of jobs that get dispatched during parallel execution. Reducing this number can be useful to avoid an explosion of memory consumption when more jobs get dispatched than CPUs can process. This parameter can be: - None, in which case all the jobs are immediately created and spawned. Use this for lightweight and fast-running jobs, to avoid delays due to on-demand spawning of the jobs - An int, giving the exact number of total jobs that are spawned - A str, giving an expression as a function of n_jobs, as in '2n_jobs'	'2*n_jobs'
	error_score error_score: 'raise' or numeric, default=np.nan Value to assign to the score if an error occurs in estimator fitting. If set to 'raise', the error is raised. If a numeric value is given, FitFailedWarning is raised. This parameter does not affect the refit step, which will always raise the error.	nan
	return_train_score return_train_score: bool, default=False If ``False``, the ``cv_results_`` attribute will not include training scores. Computing training scores is used to get insights on how different parameter settings impact the overfitting/underfitting trade-off. However computing the scores on the training set can be computationally expensive and is not strictly required to select the parameters that yield the best generalization performance. .. versionadded:: 0.19 .. versionchanged:: 0.21 Default value was changed from ``True`` to ``False``	False

	estimator estimator: estimator object A regressor or a classifier that implements :term:`fit`. When a classifier is passed, :term:`decision_function` will be used in priority and it will fallback to :term:`predict_proba` if it is not available. When a regressor is passed, :term:`predict` is used.	StochTreeBART..._seed': 1234})
	n_jobs n_jobs: int, default=None The number of jobs to use for the computation: the `n_classes` one-vs-rest problems are computed in parallel. ``None`` means 1 unless in a :obj:`joblib.parallel_backend` context. ``-1`` means using all processors. See :term:`Glossary ` for more details. .. versionchanged:: 0.20 `n_jobs` default changed from 1 to None	None
	verbose verbose: int, default=0 The verbosity level, if non zero, progress messages are printed. Below 50, the output is sent to stderr. Otherwise, the output is sent to stdout. The frequency of the messages increases with the verbosity level, reporting all iterations at 10. See :class:`joblib.Parallel` for more details. .. versionadded:: 1.1	0

	num_gfr	10
	num_burnin	0
	num_mcmc	100
	general_params	{'random_seed': 1234}
	mean_forest_params	None
	variance_forest_params	None
	rfx_params	None

Scikit-Learn Estimator Wrappers for BART¶

BART Regression via sklearn Estimator¶

Cross-Validating a BART Model¶

BART Classification via sklearn Estimator¶

BART Regression via `sklearn` Estimator¶

BART Classification via `sklearn` Estimator¶