From 7f24efe30fe1037ce638c38f7ab8c52542ae25fd Mon Sep 17 00:00:00 2001
From: OmerRonen <omerronen10@gmail.com>
Date: Sun, 12 Dec 2021 15:48:37 +0200
Subject: [PATCH] final

---
 .idea/deployment.xml                          |  21 --
 .idea/inspectionProfiles/Project_Default.xml  |  24 +-
 .idea/markdown-navigator.xml                  |  72 ----
 .../markdown-navigator/profiles_settings.xml  |   3 -
 bartpy/__init__.py                            |   1 +
 bartpy/model.py                               |  28 +-
 bartpy/node.py                                |  31 +-
 bartpy/samplers/leafnode.py                   |   7 +-
 bartpy/samplers/modelsampler.py               |  19 +-
 bartpy/sigma.py                               |   6 +-
 bartpy/sklearnmodel.py                        | 354 +++++++++++++++++-
 11 files changed, 431 insertions(+), 135 deletions(-)
 delete mode 100644 .idea/deployment.xml
 delete mode 100644 .idea/markdown-navigator.xml
 delete mode 100644 .idea/markdown-navigator/profiles_settings.xml

diff --git a/.idea/deployment.xml b/.idea/deployment.xml
deleted file mode 100644
index 9aeac29..0000000
--- a/.idea/deployment.xml
+++ /dev/null
@@ -1,21 +0,0 @@
-<?xml version="1.0" encoding="UTF-8"?>
-<project version="4">
-  <component name="PublishConfigData">
-    <serverData>
-      <paths name="Copy of project-level server 'test'">
-        <serverdata>
-          <mappings>
-            <mapping local="$PROJECT_DIR$" web="/" />
-          </mappings>
-        </serverdata>
-      </paths>
-      <paths name="Copy of project-level server 'test' (1)">
-        <serverdata>
-          <mappings>
-            <mapping local="$PROJECT_DIR$" web="/" />
-          </mappings>
-        </serverdata>
-      </paths>
-    </serverData>
-  </component>
-</project>
\ No newline at end of file
diff --git a/.idea/inspectionProfiles/Project_Default.xml b/.idea/inspectionProfiles/Project_Default.xml
index 4f208a9..09d1f7e 100644
--- a/.idea/inspectionProfiles/Project_Default.xml
+++ b/.idea/inspectionProfiles/Project_Default.xml
@@ -1,18 +1,24 @@
 <component name="InspectionProjectProfileManager">
   <profile version="1.0">
     <option name="myName" value="Project Default" />
+    <inspection_tool class="PyPackageRequirementsInspection" enabled="true" level="WARNING" enabled_by_default="true">
+      <option name="ignoredPackages">
+        <value>
+          <list size="6">
+            <item index="0" class="java.lang.String" itemvalue="tensorflow" />
+            <item index="1" class="java.lang.String" itemvalue="googleapiclient" />
+            <item index="2" class="java.lang.String" itemvalue="pandas" />
+            <item index="3" class="java.lang.String" itemvalue="biopython" />
+            <item index="4" class="java.lang.String" itemvalue="pyvcf" />
+            <item index="5" class="java.lang.String" itemvalue="spliceai" />
+          </list>
+        </value>
+      </option>
+    </inspection_tool>
     <inspection_tool class="PyPep8NamingInspection" enabled="true" level="WEAK WARNING" enabled_by_default="true">
       <option name="ignoredErrors">
         <list>
-          <option value="N803" />
-          <option value="N806" />
-        </list>
-      </option>
-      <option name="ignoredBaseClasses">
-        <list>
-          <option value="unittest.TestCase" />
-          <option value="unittest.case.TestCase" />
-          <option value="bartpy.data.Data" />
+          <option value="N802" />
         </list>
       </option>
     </inspection_tool>
diff --git a/.idea/markdown-navigator.xml b/.idea/markdown-navigator.xml
deleted file mode 100644
index 20f3a39..0000000
--- a/.idea/markdown-navigator.xml
+++ /dev/null
@@ -1,72 +0,0 @@
-<?xml version="1.0" encoding="UTF-8"?>
-<project version="4">
-  <component name="MarkdownProjectSettings">
-    <PreviewSettings splitEditorLayout="SPLIT" splitEditorPreview="PREVIEW" useGrayscaleRendering="false" zoomFactor="2.0" maxImageWidth="0" showGitHubPageIfSynced="false" allowBrowsingInPreview="false" synchronizePreviewPosition="true" highlightPreviewType="NONE" highlightFadeOut="5" highlightOnTyping="true" synchronizeSourcePosition="true" verticallyAlignSourceAndPreviewSyncPosition="true" showSearchHighlightsInPreview="false" showSelectionInPreview="true" openRemoteLinks="true">
-      <PanelProvider>
-        <provider providerId="com.vladsch.idea.multimarkdown.editor.swing.html.panel" providerName="Default - Swing" />
-      </PanelProvider>
-    </PreviewSettings>
-    <ParserSettings gitHubSyntaxChange="false">
-      <PegdownExtensions>
-        <option name="ABBREVIATIONS" value="false" />
-        <option name="ANCHORLINKS" value="true" />
-        <option name="ASIDE" value="false" />
-        <option name="ATXHEADERSPACE" value="true" />
-        <option name="AUTOLINKS" value="true" />
-        <option name="DEFINITIONS" value="false" />
-        <option name="DEFINITION_BREAK_DOUBLE_BLANK_LINE" value="false" />
-        <option name="FENCED_CODE_BLOCKS" value="true" />
-        <option name="FOOTNOTES" value="false" />
-        <option name="HARDWRAPS" value="false" />
-        <option name="HTML_DEEP_PARSER" value="false" />
-        <option name="INSERTED" value="false" />
-        <option name="QUOTES" value="false" />
-        <option name="RELAXEDHRULES" value="true" />
-        <option name="SMARTS" value="false" />
-        <option name="STRIKETHROUGH" value="true" />
-        <option name="SUBSCRIPT" value="false" />
-        <option name="SUPERSCRIPT" value="false" />
-        <option name="SUPPRESS_HTML_BLOCKS" value="false" />
-        <option name="SUPPRESS_INLINE_HTML" value="false" />
-        <option name="TABLES" value="true" />
-        <option name="TASKLISTITEMS" value="true" />
-        <option name="TOC" value="false" />
-        <option name="WIKILINKS" value="true" />
-      </PegdownExtensions>
-      <ParserOptions>
-        <option name="COMMONMARK_LISTS" value="true" />
-        <option name="DUMMY" value="false" />
-        <option name="EMOJI_SHORTCUTS" value="true" />
-        <option name="FLEXMARK_FRONT_MATTER" value="false" />
-        <option name="GFM_LOOSE_BLANK_LINE_AFTER_ITEM_PARA" value="false" />
-        <option name="GFM_TABLE_RENDERING" value="true" />
-        <option name="GITBOOK_URL_ENCODING" value="false" />
-        <option name="GITHUB_EMOJI_URL" value="false" />
-        <option name="GITHUB_LISTS" value="false" />
-        <option name="GITHUB_WIKI_LINKS" value="true" />
-        <option name="JEKYLL_FRONT_MATTER" value="false" />
-        <option name="SIM_TOC_BLANK_LINE_SPACER" value="true" />
-      </ParserOptions>
-    </ParserSettings>
-    <HtmlSettings headerTopEnabled="false" headerBottomEnabled="false" bodyTopEnabled="false" bodyBottomEnabled="false" embedUrlContent="false" addPageHeader="true" embedImages="false" embedHttpImages="false">
-      <GeneratorProvider>
-        <provider providerId="com.vladsch.idea.multimarkdown.editor.swing.html.generator" providerName="Default Swing HTML Generator" />
-      </GeneratorProvider>
-      <headerTop />
-      <headerBottom />
-      <bodyTop />
-      <bodyBottom />
-    </HtmlSettings>
-    <CssSettings previewScheme="UI_SCHEME" cssUri="" isCssUriEnabled="false" isCssTextEnabled="false" isDynamicPageWidth="true">
-      <StylesheetProvider>
-        <provider providerId="com.vladsch.idea.multimarkdown.editor.swing.html.css" providerName="Default Swing Stylesheet" />
-      </StylesheetProvider>
-      <ScriptProviders />
-      <cssText />
-    </CssSettings>
-    <HtmlExportSettings updateOnSave="false" parentDir="$ProjectFileDir$" targetDir="$ProjectFileDir$" cssDir="" scriptDir="" plainHtml="false" imageDir="" copyLinkedImages="false" imageUniquifyType="0" targetExt="" useTargetExt="false" noCssNoScripts="false" linkToExportedHtml="true" exportOnSettingsChange="true" regenerateOnProjectOpen="false" linkFormatType="HTTP_ABSOLUTE" />
-    <LinkMapSettings>
-      <textMaps />
-    </LinkMapSettings>
-  </component>
-</project>
\ No newline at end of file
diff --git a/.idea/markdown-navigator/profiles_settings.xml b/.idea/markdown-navigator/profiles_settings.xml
deleted file mode 100644
index 57927c5..0000000
--- a/.idea/markdown-navigator/profiles_settings.xml
+++ /dev/null
@@ -1,3 +0,0 @@
-<component name="MarkdownNavigator.ProfileManager">
-  <settings default="" pdf-export="" />
-</component>
\ No newline at end of file
diff --git a/bartpy/__init__.py b/bartpy/__init__.py
index e69de29..90c624a 100644
--- a/bartpy/__init__.py
+++ b/bartpy/__init__.py
@@ -0,0 +1 @@
+from .sklearnmodel import BART, ShrunkBART, ShrunkBARTCV
\ No newline at end of file
diff --git a/bartpy/model.py b/bartpy/model.py
index b1f1742..7b57e88 100644
--- a/bartpy/model.py
+++ b/bartpy/model.py
@@ -17,12 +17,13 @@ class Model:
     def __init__(self,
                  data: Optional[Data],
                  sigma: Sigma,
-                 trees: Optional[List[Tree]]=None,
-                 n_trees: int=50,
-                 alpha: float=0.95,
-                 beta: float=2.,
-                 k: int=2.,
-                 initializer: Initializer=SklearnTreeInitializer()):
+                 trees: Optional[List[Tree]] = None,
+                 n_trees: int = 50,
+                 alpha: float = 0.95,
+                 beta: float = 2.,
+                 k: int = 2.,
+                 initializer: Initializer = SklearnTreeInitializer(),
+                 classification: bool = False):
 
         self.data = deepcopy(data)
         self.alpha = float(alpha)
@@ -31,6 +32,7 @@ def __init__(self,
         self._sigma = sigma
         self._prediction = None
         self._initializer = initializer
+        self.classification = classification
 
         if trees is None:
             self.n_trees = n_trees
@@ -53,7 +55,7 @@ def residuals(self) -> np.ndarray:
     def unnormalized_residuals(self) -> np.ndarray:
         return self.data.y.unnormalized_y - self.data.y.unnormalize_y(self.predict())
 
-    def predict(self, X: np.ndarray=None) -> np.ndarray:
+    def predict(self, X: np.ndarray = None) -> np.ndarray:
         if X is not None:
             return self._out_of_sample_predict(X)
         return np.sum([tree.predict() for tree in self.trees], axis=0)
@@ -77,8 +79,20 @@ def refreshed_trees(self) -> Generator[Tree, None, None]:
             yield tree
             self._prediction += tree.predict()
 
+    def update_z_values(self, y):
+        if not self.classification:
+            return
+        z = np.random.normal(loc=self.predict(self.data.X.values))
+        one_label = np.maximum(z[y == 1], 0)
+        zero_label = np.minimum(z[y == 0], 0)
+        z[y == 1] = one_label
+        z[y == 0] = zero_label
+        self.data.update_y(z)
+
     @property
     def sigma_m(self) -> float:
+        if self.classification:
+            return 3 / (self.k * np.power(self.n_trees, 0.5))
         return 0.5 / (self.k * np.power(self.n_trees, 0.5))
 
     @property
diff --git a/bartpy/node.py b/bartpy/node.py
index 80bd841..253a70a 100644
--- a/bartpy/node.py
+++ b/bartpy/node.py
@@ -11,7 +11,8 @@ class TreeNode(object):
         - Data relevant for the node
         - Links to children nodes
     """
-    def __init__(self, split: Split, depth: int, left_child: 'TreeNode'=None, right_child: 'TreeNode'=None):
+
+    def __init__(self, split: Split, depth: int, left_child: 'TreeNode' = None, right_child: 'TreeNode' = None):
         self.depth = depth
         self._split = split
         self._left_child = left_child
@@ -55,6 +56,9 @@ def __init__(self, split: Split, depth=0, value=0.0):
     def set_value(self, value: float) -> None:
         self._value = value
 
+    def set_mean_response(self, value: float) -> None:
+        self._mean_response = value
+
     @property
     def current_value(self):
         return self._value
@@ -65,6 +69,14 @@ def predict(self) -> float:
     def is_splittable(self) -> bool:
         return self.data.X.is_at_least_one_splittable_variable()
 
+    @property
+    def n_obs(self):
+        return self.data.X.n_obsv
+
+    @property
+    def mean_response(self):
+        return self._mean_response
+
 
 class DecisionNode(TreeNode):
     """
@@ -81,6 +93,23 @@ def is_prunable(self) -> bool:
     def most_recent_split_condition(self) -> SplitCondition:
         return self.left_child.split.most_recent_split_condition()
 
+    @property
+    def n_obs(self):
+        n_l = self.left_child.n_obs
+        n_r = self.right_child.n_obs
+        return n_l + n_r
+
+
+    @property
+    def current_value(self):
+        n_l = self.left_child.n_obs
+        n_r = self.right_child.n_obs
+        l_val = self.left_child.current_value if type(self.left_child) == DecisionNode else self.left_child.mean_response
+        r_val = self.right_child.current_value if type(self.right_child) == DecisionNode else self.right_child.mean_response
+        l_sum = l_val * n_l
+        r_sum = r_val * n_r
+        return (l_sum + r_sum) / self.n_obs
+
 
 def split_node(node: LeafNode, split_conditions: Tuple[SplitCondition, SplitCondition]) -> DecisionNode:
     """
diff --git a/bartpy/samplers/leafnode.py b/bartpy/samplers/leafnode.py
index d70b702..c853453 100644
--- a/bartpy/samplers/leafnode.py
+++ b/bartpy/samplers/leafnode.py
@@ -28,10 +28,11 @@ def sample(self, model: Model, node: LeafNode) -> float:
         n = node.data.X.n_obsv
         likihood_var = (model.sigma.current_value() ** 2) / n
         likihood_mean = node.data.y.summed_y() / n
+        node.set_mean_response(likihood_mean)
         posterior_variance = 1. / (1. / prior_var + 1. / likihood_var)
         posterior_mean = likihood_mean * (prior_var / (likihood_var + prior_var))
-        return posterior_mean + (self._scalar_sampler.sample() * np.power(posterior_variance / model.n_trees, 0.5))
-
+        val = posterior_mean + (self._scalar_sampler.sample() * np.power(posterior_variance / model.n_trees, 0.5))
+        return val
 
 # class VectorizedLeafNodeSampler(Sampler):
 
@@ -45,5 +46,3 @@ def sample(self, model: Model, node: LeafNode) -> float:
 #         prior_var = model.sigma_m ** 2
 #         n_s = []
 #         sum_s = []
-        
-
diff --git a/bartpy/samplers/modelsampler.py b/bartpy/samplers/modelsampler.py
index 4d8b997..37d09a3 100644
--- a/bartpy/samplers/modelsampler.py
+++ b/bartpy/samplers/modelsampler.py
@@ -1,3 +1,4 @@
+import copy
 from collections import defaultdict
 from typing import List, Mapping, Union, Any, Type
 
@@ -16,9 +17,11 @@ class ModelSampler(Sampler):
 
     def __init__(self,
                  schedule: SampleSchedule,
-                 trace_logger_class: Type[TraceLogger]=TraceLogger):
+                 trace_logger_class: Type[TraceLogger]=TraceLogger,
+                 n_rules: int=None):
         self.schedule = schedule
         self.trace_logger_class = trace_logger_class
+        self.n_rules = n_rules
 
     def step(self, model: Model, trace_logger: TraceLogger):
         step_result = defaultdict(list)
@@ -35,21 +38,27 @@ def samples(self, model: Model,
                 thin: float=0.1,
                 store_in_sample_predictions: bool=True,
                 store_acceptance: bool=True) -> Chain:
-        print("Starting burn")
+        # print("Starting burn")
 
         trace_logger = self.trace_logger_class()
+        y = copy.deepcopy(model.data.y.unnormalized_y)
 
-        for _ in tqdm(range(n_burn)):
+        for _ in range(n_burn):
+            model.update_z_values(y)
             self.step(model, trace_logger)
+
         trace = []
         model_trace = []
         acceptance_trace = []
-        print("Starting sampling")
+        # print("Starting sampling")
 
         thin_inverse = 1. / thin
 
-        for ss in tqdm(range(n_samples)):
+        for ss in range(n_samples):
+            model.update_z_values(y)
             step_trace_dict = self.step(model, trace_logger)
+            # print(step_trace_dict)
+
             if ss % thin_inverse == 0:
                 if store_in_sample_predictions:
                     in_sample_log = trace_logger["In Sample Prediction"](model.predict())
diff --git a/bartpy/sigma.py b/bartpy/sigma.py
index b289527..4a254df 100644
--- a/bartpy/sigma.py
+++ b/bartpy/sigma.py
@@ -1,5 +1,4 @@
 
-
 class Sigma:
     """
     A representation of the sigma term in the model.
@@ -18,16 +17,19 @@ class Sigma:
 
     """
 
-    def __init__(self, alpha: float, beta: float, scaling_factor: float):
+    def __init__(self, alpha: float, beta: float, scaling_factor: float, classification :bool = False):
         self.alpha = alpha
         self.beta = beta
         self._current_value = 1.0
         self.scaling_factor = scaling_factor
+        self._classification = classification
 
     def set_value(self, value: float) -> None:
         self._current_value = value
 
     def current_value(self) -> float:
+        if self._classification:
+            return 1
         return self._current_value
 
     def current_unnormalized_value(self) -> float:
diff --git a/bartpy/sklearnmodel.py b/bartpy/sklearnmodel.py
index 7e92b1a..3876401 100644
--- a/bartpy/sklearnmodel.py
+++ b/bartpy/sklearnmodel.py
@@ -1,15 +1,22 @@
+import copy
 from copy import deepcopy
 from typing import List, Callable, Mapping, Union, Optional
 
+import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
+import scipy.stats
 from joblib import Parallel, delayed
 from sklearn.base import RegressorMixin, BaseEstimator
+from sklearn.model_selection import cross_val_score
+from sklearn.tree import DecisionTreeClassifier
+from sklearn import datasets, model_selection
 
 from bartpy.data import Data
 from bartpy.initializers.initializer import Initializer
 from bartpy.initializers.sklearntreeinitializer import SklearnTreeInitializer
 from bartpy.model import Model
+from bartpy.node import LeafNode, DecisionNode
 from bartpy.samplers.leafnode import LeafNodeSampler
 from bartpy.samplers.modelsampler import ModelSampler, Chain
 from bartpy.samplers.schedule import SampleSchedule
@@ -24,7 +31,17 @@ def run_chain(model: 'SklearnModel', X: np.ndarray, y: np.ndarray):
     Run a single chain for a model
     Primarily used as a building block for constructing a parallel run of multiple chains
     """
+
+    # TODO: support for classification F^{-1} (y) ~ N(G(x), 1)
+
+    # if model.classification:
+    #     z = np.random.normal(loc=model.predict(X))
+    #     z[y == 1] = np.maximum(z[y == 1], 0)
+    #     z[y == 0] = np.minimum(z[y == 0], 0)
+    #     y = z
+
     model.model = model._construct_model(X, y)
+
     return model.sampler.samples(model.model,
                                  model.n_samples,
                                  model.n_burn,
@@ -37,6 +54,92 @@ def delayed_run_chain():
     return run_chain
 
 
+def get_nodes(root_node):
+    decision_nodes = []
+    leaf_nodes = []
+
+    def _add_nodes(n):
+        if type(n) == LeafNode:
+            leaf_nodes.append(n)
+        elif type(n) == DecisionNode:
+            decision_nodes.append(n)
+        if n.left_child:
+            _add_nodes(n.left_child)
+        if n.right_child:
+            _add_nodes(n.right_child)
+
+    _add_nodes(root_node)
+    return decision_nodes, leaf_nodes
+
+
+def get_root_node(tree):
+    for n in tree.decision_nodes:
+        if n.depth == 0:
+            return n
+
+
+def shrink_tree(tree, reg_param):
+    root = get_root_node(tree)
+    tree_d_node = shrink_node(root, reg_param)
+    d, l = get_nodes(tree_d_node)
+    tree._nodes = d + l
+    return tree
+
+
+def expand_node(node):
+    mask_int = 1 - node.data.mask.astype(int)
+    # y = node.data.y.values
+    val = np.sum(node.data.y.values
+                 * mask_int) / np.sum(mask_int)
+    node.set_value(val)
+    return node
+
+
+def expand_tree(tree):
+    decision, leaves = get_nodes(get_root_node(tree))
+    leaves_new = [expand_node(n) for n in leaves]
+    tree._nodes = decision + leaves_new
+    return tree
+
+
+def shrink_node(node, reg_param, parent_val, parent_num, cum_sum, scheme, constant):
+    """Shrink the tree
+    """
+
+    # node.set_value(node.mean_response)
+
+    left = node.left_child
+    right = node.right_child
+    is_leaf = type(node) == LeafNode
+    # if self.prediction_task == 'regression':
+    val = node.current_value
+    is_root = parent_val is None and parent_num is None
+    n_samples = node.n_obs if (scheme != "leaf_based" or is_root) else parent_num
+
+    if is_root:
+        val_new = val
+
+    else:
+        reg_term = reg_param if scheme == "constant" else reg_param / parent_num
+
+        val_new = (val - parent_val) / (1 + reg_term)
+
+    cum_sum += val_new
+
+    if is_leaf:
+        if scheme == "leaf_based":
+            v = constant + (val - constant) / (1 + reg_param / node.n_obs)
+            node.set_value(v)
+        else:
+            node.set_value(cum_sum)
+
+    else:
+        shrink_node(left, reg_param, val, parent_num=n_samples, cum_sum=cum_sum, scheme=scheme, constant=constant)
+        shrink_node(right, reg_param, val, parent_num=n_samples, cum_sum=cum_sum, scheme=scheme, constant=constant)
+
+    return node
+
+
 class SklearnModel(BaseEstimator, RegressorMixin):
     """
     The main access point to building BART models in BartPy
@@ -94,11 +197,12 @@ def __init__(self,
                  thin: float = 0.1,
                  alpha: float = 0.95,
                  beta: float = 2.,
-                 store_in_sample_predictions: bool=False,
-                 store_acceptance_trace: bool=False,
-                 tree_sampler: TreeMutationSampler=get_tree_sampler(0.5, 0.5),
-                 initializer: Optional[Initializer]=None,
-                 n_jobs=-1):
+                 store_in_sample_predictions: bool = False,
+                 store_acceptance_trace: bool = False,
+                 tree_sampler: TreeMutationSampler = get_tree_sampler(0.5, 0.5),
+                 initializer: Optional[Initializer] = None,
+                 n_jobs=-1,
+                 classification: bool = False):
         self.n_trees = n_trees
         self.n_chains = n_chains
         self.sigma_a = sigma_a
@@ -118,6 +222,7 @@ def __init__(self,
         self.initializer = initializer
         self.schedule = SampleSchedule(self.tree_sampler, LeafNodeSampler(), SigmaSampler())
         self.sampler = ModelSampler(self.schedule)
+        self.classification = classification
 
         self.sigma, self.data, self.model, self._prediction_samples, self._model_samples, self.extract = [None] * 6
 
@@ -140,10 +245,18 @@ def fit(self, X: Union[np.ndarray, pd.DataFrame], y: np.ndarray) -> 'SklearnMode
         self.model = self._construct_model(X, y)
         self.extract = Parallel(n_jobs=self.n_jobs)(self.f_delayed_chains(X, y))
         self.combined_chains = self._combine_chains(self.extract)
-        self._model_samples, self._prediction_samples = self.combined_chains["model"], self.combined_chains["in_sample_predictions"]
+        self._model_samples, self._prediction_samples = self.combined_chains["model"], self.combined_chains[
+            "in_sample_predictions"]
         self._acceptance_trace = self.combined_chains["acceptance"]
+        self.fitted_ = True
         return self
 
+    @property
+    def fitted(self):
+        if hasattr(self, "fitted_"):
+            return self.fitted_
+        return False
+
     @staticmethod
     def _combine_chains(extract: List[Chain]) -> Chain:
         keys = list(extract[0].keys())
@@ -164,13 +277,14 @@ def _construct_model(self, X: np.ndarray, y: np.ndarray) -> Model:
         if len(X) == 0 or X.shape[1] == 0:
             raise ValueError("Empty covariate matrix passed")
         self.data = self._convert_covariates_to_data(X, y)
-        self.sigma = Sigma(self.sigma_a, self.sigma_b, self.data.y.normalizing_scale)
+        self.sigma = Sigma(self.sigma_a, self.sigma_b, self.data.y.normalizing_scale, self.classification)
         self.model = Model(self.data,
                            self.sigma,
                            n_trees=self.n_trees,
                            alpha=self.alpha,
                            beta=self.beta,
-                           initializer=self.initializer)
+                           initializer=self.initializer,
+                           classification=self.classification)
         return self.model
 
     def f_delayed_chains(self, X: np.ndarray, y: np.ndarray):
@@ -205,7 +319,7 @@ def f_chains(self) -> List[Callable[[], Chain]]:
         """
         return [delayed_run_chain() for _ in range(self.n_chains)]
 
-    def predict(self, X: np.ndarray=None) -> np.ndarray:
+    def predict(self, X: np.ndarray = None) -> np.ndarray:
         """
         Predict the target corresponding to the provided covariate matrix
         If X is None, will predict based on training covariates
@@ -228,7 +342,14 @@ def predict(self, X: np.ndarray=None) -> np.ndarray:
             raise ValueError(
                 "In sample predictions only possible if model.store_in_sample_predictions is `True`.  Either set the parameter to True or pass a non-None X parameter")
         else:
-            return self._out_of_sample_predict(X)
+            predictions = self._out_of_sample_predict(X)
+            if self.classification:
+                return np.round(predictions, 0)
+            return predictions
+
+    def predict_proba(self, X: np.ndarray = None) -> np.ndarray:
+        return self._out_of_sample_predict(X)
+
 
     def residuals(self, X=None, y=None) -> np.ndarray:
         """
@@ -288,7 +409,12 @@ def rmse(self, X, y) -> float:
         return np.sqrt(np.sum(self.l2_error(X, y)))
 
     def _out_of_sample_predict(self, X):
-        return self.data.y.unnormalize_y(np.mean([x.predict(X) for x in self._model_samples], axis=0))
+        samples = self._model_samples
+        predictions_transformed = [x.predict(X) for x in samples]
+        predictions = self.data.y.unnormalize_y(np.mean(predictions_transformed, axis=0))
+        if self.classification:
+            predictions = scipy.stats.norm.cdf(predictions)
+        return predictions
 
     def fit_predict(self, X, y):
         self.fit(X, y)
@@ -364,3 +490,209 @@ def from_extract(self, extract: List[Chain], X: np.ndarray, y: np.ndarray) -> 'S
         self._acceptance_trace = combined_chain["acceptance"]
         new_model.data = self._convert_covariates_to_data(X, y)
         return new_model
+
+
+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
+
+
+class ImputedBART(BaseEstimator):
+    def __init__(self, estimator_):
+        # super(ShrunkBARTRegressor, self).__init__()
+        self.estimator_ = estimator_
+
+    def predict(self, *args, **kwargs):
+        return self.estimator_.predict(*args, **kwargs)
+
+    def predict_proba(self, *args, **kwargs):
+        if hasattr(self.estimator_, 'predict_proba'):
+            return self.estimator_.predict_proba(*args, **kwargs)
+        else:
+            return NotImplemented
+
+    def score(self, *args, **kwargs):
+        if hasattr(self.estimator_, 'score'):
+            return self.estimator_.score(*args, **kwargs)
+        else:
+            return NotImplemented
+
+
+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
+
+
+class ExpandedBART(ImputedBART):
+
+    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 = expand_tree(copy.deepcopy(tree))
+                s._trees[i] = s_tree
+            samples.append(s)
+        self.estimator_._model_samples = samples
+        self.fitted_ = True
+
+
+# class ExpandedBARTRegressor(ImputedBARTRegressor):
+#
+#     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 = expend_tree(copy.deepcopy(tree), args[1])
+#                 s._trees[i] = s_tree
+#             samples.append(s)
+#         self.estimator_._model_samples = samples
+#         self.fitted_ = True
+
+
+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)
+
+
+def main():
+    # iris = datasets.load_iris()
+    # idx = np.logical_or(iris.target == 0, iris.target == 1)
+    # X, y = iris.data[idx, ...], iris.target[idx]
+    X, y = datasets.load_diabetes(return_X_y=True)
+
+    X_train, X_test, y_train, y_test = model_selection.train_test_split(
+        X, y, test_size=0.3, random_state=1)
+    bart = BART(classification=False)
+    bart.fit(X_train, y_train)
+    preds_org = bart.predict(X_test)
+    mse = np.linalg.norm(preds_org - y_test)
+    print(mse)
+    # tree = DecisionTreeClassifier()
+    # tree.fit(X, y)
+    # preds_tree = tree.predict_proba(X)
+    # bart_s = ShrunkBARTCV(copy.deepcopy(bart), scheme="node_based")
+    # bart_s.fit(X, y)
+    #
+    # # bart_s_c = ShrunkBART(copy.deepcopy(bart), reg_param=2, scheme="constant")
+    # # bart_s_c.fit(X, y)
+    # #
+    # # bart_s_l = ShrunkBART(copy.deepcopy(bart), reg_param=2, scheme="leaf_based")
+    # # bart_s_l.fit(X, y)
+    # # bart_s_cv = ShrunkBARTRegressorCV(estimator_=copy.deepcopy(bart))
+    # # bart_s_cv.fit(X, y)
+    # e_bart = ExpandedBART(estimator_=copy.deepcopy(bart))
+    # e_bart.fit(X, y)
+    #
+    # preds = bart_s.predict(X)
+    #
+    # # preds_c = bart_s_c.predict(X)
+    # # preds_l = bart_s_l.predict(X)
+    # # # preds_cv = bart_s_cv.predict(X)
+    # preds_bart_e = e_bart.predict(X)
+    # fig, ax = plt.subplots(1)
+    #
+    # ax.scatter(np.arange(len(y)), preds_org, c="orange", label="bart")
+    # ax.scatter(np.arange(len(y)), preds, c="purple", alpha=0.3, label="shrunk node")
+    # # ax.scatter(np.arange(len(y)), preds_c, c="blue", alpha=0.3, label="shrunk constant")
+    # # ax.scatter(np.arange(len(y)), preds_l, c="red", alpha=0.3, label="shrunk leaf")
+    # ax.scatter(np.arange(len(y)), preds_bart_e, c="green", alpha=0.3, label="average")
+    # # preds_all = [preds_org, preds_c, preds_l, preds_bart_e, preds]
+    # # shift = 0.5
+    # # rng = (np.min([np.min(p) for p in preds_all]) - shift, np.max([np.max(p) for p in preds_all]) + shift)
+    # # n_bins = 200
+    # # alpha = 0.8
+    # # ax.hist(preds_org, color="orange", alpha=alpha, label="bart", bins=n_bins, range=rng)
+    # # ax.hist(preds, color="purple", alpha=alpha, label="shrunk node", bins=n_bins, range=rng)
+    # # ax.hist(preds_c, color="blue", alpha=alpha, label="shrunk constant", bins=n_bins, range=rng)
+    # # ax.hist(preds_l, color="red", alpha=alpha, label="shrunk leaf", bins=n_bins, range=rng)
+    # # ax.hist(preds_bart_e, color="green", alpha=alpha, label="average", bins=n_bins, range=rng)
+    # #
+    # # ax.set_xlabel("Predicted Value")
+    # # ax.set_ylabel("Count")
+    #
+    # plt.title(np.mean(y))
+    #
+    # plt.legend(loc="upper left")
+    # plt.savefig("bart_shrink.png")
+    # # plt.show()
+    # #
+    # # plt.close()
+    #
+
+if __name__ == '__main__':
+    main()