Expand source code
from .sklearnmodel import BART, ShrunkBART, ShrunkBARTCV

# re-exported for callers; listed so the intent is explicit
__all__ = [
    "BART", "ShrunkBART", "ShrunkBARTCV",
]

Sub-modules

imodels.experimental.bartpy.data
imodels.experimental.bartpy.diagnostics
imodels.experimental.bartpy.errors
imodels.experimental.bartpy.extensions
imodels.experimental.bartpy.features
imodels.experimental.bartpy.initializers
imodels.experimental.bartpy.model
imodels.experimental.bartpy.mutation
imodels.experimental.bartpy.node
imodels.experimental.bartpy.plotting
imodels.experimental.bartpy.runner

A model for running many instances of BartPy models in parallel

imodels.experimental.bartpy.samplers
imodels.experimental.bartpy.sigma
imodels.experimental.bartpy.sklearnmodel
imodels.experimental.bartpy.split
imodels.experimental.bartpy.splitcondition
imodels.experimental.bartpy.trace
imodels.experimental.bartpy.tree

Classes

class BART (n_trees: int = 200, n_chains: int = 4, sigma_a: float = 0.001, sigma_b: float = 0.001, n_samples: int = 200, n_burn: int = 200, thin: float = 0.1, alpha: float = 0.95, beta: float = 2.0, store_in_sample_predictions: bool = False, store_acceptance_trace: bool = False, tree_sampler: TreeMutationSampler = <imodels.experimental.bartpy.samplers.unconstrainedtree.treemutation.UnconstrainedTreeMutationSampler object>, initializer: Initializer | None = None, n_jobs=-1, classification: bool = False, max_rules=None)

The main access point to building BART models in BartPy

Parameters

n_trees : int
the number of trees to use, more trees will make a smoother fit, but slow training and fitting
n_chains : int
the number of independent chains to run more chains will improve the quality of the samples, but will require more computation
sigma_a : float
shape parameter of the prior on sigma
sigma_b : float
scale parameter of the prior on sigma
n_samples : int
how many recorded samples to take
n_burn : int
how many samples to run without recording to reach convergence
thin : float
percentage of samples to store. use this to save memory when running large models
p_grow : float
probability of choosing a grow mutation in tree mutation sampling
p_prune : float
probability of choosing a prune mutation in tree mutation sampling
alpha : float
prior parameter on tree structure
beta : float
prior parameter on tree structure
store_in_sample_predictions : bool
whether to store full prediction samples set to False if you don't need in sample results - saves a lot of memory
store_acceptance_trace : bool
whether to store acceptance rates of the gibbs samples unless you're very memory constrained, you wouldn't want to set this to false useful for diagnostics
tree_sampler : TreeMutationSampler
Method of sampling used on trees defaults to bartpy.samplers.unconstrainedtree
initializer : Initializer
Class that handles the initialization of tree structure and leaf values
n_jobs : int
how many cores to use when computing MCMC samples set to -1 to use all cores
Expand source code
class BART(SklearnModel):

    @staticmethod
    def _get_n_nodes(trees):
        nodes = 0
        for tree in trees:
            nodes += len(tree.decision_nodes)
        return nodes

    @property
    def sample_complexity(self):
        # samples = self._model_samples
        # trees = [s.trees for s in samples]
        complexities = [self._get_n_nodes(t) for t in self.trees]
        return np.sum(complexities)

    @staticmethod
    def sub_forest(trees, n_nodes):
        nodes = 0
        for i, tree in enumerate(trees):
            nodes += len(tree.decision_nodes)
            if nodes >= n_nodes:
                return trees[0:i + 1]

    @property
    def trees(self):
        trs = [s.trees for s in self._model_samples]
        return trs

    def update_complexity(self, i):
        samples_complexity = [self._get_n_nodes(t) for t in self.trees]

        # complexity_sum = 0
        arg_sort_complexity = np.argsort(samples_complexity)
        self._model_samples = self._model_samples[arg_sort_complexity[:i + 1]]

        return self

Ancestors

  • SklearnModel
  • sklearn.base.BaseEstimator
  • sklearn.utils._estimator_html_repr._HTMLDocumentationLinkMixin
  • sklearn.utils._metadata_requests._MetadataRequester
  • sklearn.base.RegressorMixin

Static methods

def sub_forest(trees, n_nodes)
Expand source code
@staticmethod
def sub_forest(trees, n_nodes):
    nodes = 0
    for i, tree in enumerate(trees):
        nodes += len(tree.decision_nodes)
        if nodes >= n_nodes:
            return trees[0:i + 1]

Instance variables

var sample_complexity
Expand source code
@property
def sample_complexity(self):
    # samples = self._model_samples
    # trees = [s.trees for s in samples]
    complexities = [self._get_n_nodes(t) for t in self.trees]
    return np.sum(complexities)
var trees
Expand source code
@property
def trees(self):
    trs = [s.trees for s in self._model_samples]
    return trs

Methods

def update_complexity(self, i)
Expand source code
def update_complexity(self, i):
    samples_complexity = [self._get_n_nodes(t) for t in self.trees]

    # complexity_sum = 0
    arg_sort_complexity = np.argsort(samples_complexity)
    self._model_samples = self._model_samples[arg_sort_complexity[:i + 1]]

    return self

Inherited members

class ShrunkBART (estimator_, reg_param, scheme)

Base class for all estimators in scikit-learn.

Inheriting from this class provides default implementations of:

  • setting and getting parameters used by GridSearchCV and friends;
  • textual and HTML representation displayed in terminals and IDEs;
  • estimator serialization;
  • parameters validation;
  • data validation;
  • feature names validation.

Read more in the :ref:User Guide <rolling_your_own_estimator>.

Notes

All estimators should specify all the parameters that can be set at the class level in their __init__ as explicit keyword arguments (no *args or **kwargs).

Examples

>>> import numpy as np
>>> from sklearn.base import BaseEstimator
>>> class MyEstimator(BaseEstimator):
...     def __init__(self, *, param=1):
...         self.param = param
...     def fit(self, X, y=None):
...         self.is_fitted_ = True
...         return self
...     def predict(self, X):
...         return np.full(shape=X.shape[0], fill_value=self.param)
>>> estimator = MyEstimator(param=2)
>>> estimator.get_params()
{'param': 2}
>>> X = np.array([[1, 2], [2, 3], [3, 4]])
>>> y = np.array([1, 0, 1])
>>> estimator.fit(X, y).predict(X)
array([2, 2, 2])
>>> estimator.set_params(param=3).fit(X, y).predict(X)
array([3, 3, 3])
Expand source code
class ShrunkBART(ImputedBART):

    def __init__(self, estimator_, reg_param, scheme):
        super(ShrunkBART, self).__init__(estimator_)
        self.reg_param = reg_param
        self.scheme = scheme

    def shrink_tree(self, tree):
        root = get_root_node(tree)
        tree_d_node = shrink_node(root, self.reg_param, parent_val=None, parent_num=None, cum_sum=0, scheme=self.scheme,
                                  constant=np.mean(self.estimator_.data.y.values))
        d, l = get_nodes(tree_d_node)
        tree._nodes = d + l
        return tree

    def fit(self, *args, **kwargs):
        if not self.estimator_.fitted:
            self.estimator_.fit(*args, **kwargs)
        samples = []
        for s in self.estimator_.model_samples:
            for i, tree in enumerate(s._trees):
                s_tree = self.shrink_tree(expand_tree(copy.deepcopy(tree)))
                s._trees[i] = s_tree
            samples.append(s)
        self.estimator_._model_samples = samples
        self.fitted_ = True

Ancestors

  • ImputedBART
  • sklearn.base.BaseEstimator
  • sklearn.utils._estimator_html_repr._HTMLDocumentationLinkMixin
  • sklearn.utils._metadata_requests._MetadataRequester

Subclasses

Methods

def fit(self, *args, **kwargs)
Expand source code
def fit(self, *args, **kwargs):
    if not self.estimator_.fitted:
        self.estimator_.fit(*args, **kwargs)
    samples = []
    for s in self.estimator_.model_samples:
        for i, tree in enumerate(s._trees):
            s_tree = self.shrink_tree(expand_tree(copy.deepcopy(tree)))
            s._trees[i] = s_tree
        samples.append(s)
    self.estimator_._model_samples = samples
    self.fitted_ = True
def shrink_tree(self, tree)
Expand source code
def shrink_tree(self, tree):
    root = get_root_node(tree)
    tree_d_node = shrink_node(root, self.reg_param, parent_val=None, parent_num=None, cum_sum=0, scheme=self.scheme,
                              constant=np.mean(self.estimator_.data.y.values))
    d, l = get_nodes(tree_d_node)
    tree._nodes = d + l
    return tree
class ShrunkBARTCV (estimator_: sklearn.base.BaseEstimator, scheme: str, reg_param_list: List[float] = [0.1, 1, 10, 50, 100, 500], cv: int = 3, scoring=None)

Base class for all estimators in scikit-learn.

Inheriting from this class provides default implementations of:

  • setting and getting parameters used by GridSearchCV and friends;
  • textual and HTML representation displayed in terminals and IDEs;
  • estimator serialization;
  • parameters validation;
  • data validation;
  • feature names validation.

Read more in the :ref:User Guide <rolling_your_own_estimator>.

Notes

All estimators should specify all the parameters that can be set at the class level in their __init__ as explicit keyword arguments (no *args or **kwargs).

Examples

>>> import numpy as np
>>> from sklearn.base import BaseEstimator
>>> class MyEstimator(BaseEstimator):
...     def __init__(self, *, param=1):
...         self.param = param
...     def fit(self, X, y=None):
...         self.is_fitted_ = True
...         return self
...     def predict(self, X):
...         return np.full(shape=X.shape[0], fill_value=self.param)
>>> estimator = MyEstimator(param=2)
>>> estimator.get_params()
{'param': 2}
>>> X = np.array([[1, 2], [2, 3], [3, 4]])
>>> y = np.array([1, 0, 1])
>>> estimator.fit(X, y).predict(X)
array([2, 2, 2])
>>> estimator.set_params(param=3).fit(X, y).predict(X)
array([3, 3, 3])
Expand source code
class ShrunkBARTCV(ShrunkBART):
    def __init__(self, estimator_: BaseEstimator, scheme: str,
                 reg_param_list: List[float] = [0.1, 1, 10, 50, 100, 500],
                 cv: int = 3, scoring=None):
        super(ShrunkBARTCV, self).__init__(estimator_, None, scheme)
        self.reg_param_list = np.array(reg_param_list)
        self.cv = cv
        self.scoring = scoring

    def fit(self, X, y, *args, **kwargs):
        self.scores_ = []
        for reg_param in self.reg_param_list:
            est = ShrunkBART(deepcopy(self.estimator_), reg_param, self.scheme)
            cv_scores = cross_val_score(est, X, y, cv=self.cv, scoring=self.scoring)
            self.scores_.append(np.mean(cv_scores))
        self.reg_param = self.reg_param_list[np.argmax(self.scores_)]
        super().fit(X=X, y=y)

Ancestors

  • ShrunkBART
  • ImputedBART
  • sklearn.base.BaseEstimator
  • sklearn.utils._estimator_html_repr._HTMLDocumentationLinkMixin
  • sklearn.utils._metadata_requests._MetadataRequester

Methods

def fit(self, X, y, *args, **kwargs)
Expand source code
def fit(self, X, y, *args, **kwargs):
    self.scores_ = []
    for reg_param in self.reg_param_list:
        est = ShrunkBART(deepcopy(self.estimator_), reg_param, self.scheme)
        cv_scores = cross_val_score(est, X, y, cv=self.cv, scoring=self.scoring)
        self.scores_.append(np.mean(cv_scores))
    self.reg_param = self.reg_param_list[np.argmax(self.scores_)]
    super().fit(X=X, y=y)