bartz 0.4.1__py3-none-any.whl → 0.5.0__py3-none-any.whl

bartz/BART.py

@@ -1,6 +1,6 @@
 # bartz/src/bartz/BART.py
 #
-# Copyright (c) 2024, Giacomo Petrillo
+# Copyright (c) 2024-2025, Giacomo Petrillo
 #
 # This file is part of bartz.
 #
@@ -27,11 +27,8 @@ import functools
 import jax
 import jax.numpy as jnp
 
-from . import jaxext
-from . import grove
-from . import mcmcstep
-from . import mcmcloop
-from . import prepcovars
+from . import grove, jaxext, mcmcloop, mcmcstep, prepcovars
+
 
 class gbart:
     """
@@ -53,10 +50,11 @@ class gbart:
         Whether to use predictors quantiles instead of a uniform grid to bin
         predictors.
     sigest : float, optional
-        An estimate of the residual standard deviation on `y_train`, used to
-        set `lamda`. If not specified, it is estimated by linear regression.
-        If `y_train` has less than two elements, it is set to 1. If n <= p, it
-        is set to the variance of `y_train`. Ignored if `lamda` is specified.
+        An estimate of the residual standard deviation on `y_train`, used to set
+        `lamda`. If not specified, it is estimated by linear regression (with
+        intercept, and without taking into account `w`). If `y_train` has less
+        than two elements, it is set to 1. If n <= p, it is set to the standard
+        deviation of `y_train`. Ignored if `lamda` is specified.
     sigdf : int, default 3
         The degrees of freedom of the scaled inverse-chisquared prior on the
         noise variance.
@@ -82,6 +80,12 @@ class gbart:
     offset : float, optional
         The prior mean of the latent mean function. If not specified, it is set
         to the mean of `y_train`. If `y_train` is empty, it is set to 0.
+    w : array (n,), optional
+        Coefficients that rescale the error standard deviation on each
+        datapoint. Not specifying `w` is equivalent to setting it to 1 for all
+        datapoints. Note: `w` is ignored in the automatic determination of
+        `sigest`, so either the weights should be O(1), or `sigest` should be
+        specified by the user.
     ntree : int, default 200
         The number of trees used to represent the latent mean function.
     numcut : int, default 255
@@ -108,6 +112,8 @@ class gbart:
         The number of iterations (including skipped ones) between each log.
     seed : int or jax random key, default 0
         The seed for the random number generator.
+    initkw : dict
+        Additional arguments passed to `mcmcstep.init`.
 
     Attributes
     ----------
@@ -135,8 +141,6 @@ class gbart:
         The number of trees.
     maxdepth : int
         The maximum depth of the trees.
-    initkw : dict
-        Additional arguments passed to `mcmcstep.init`.
 
     Methods
     -------
@@ -158,10 +162,13 @@ class gbart:
     - A lot of functionality is missing (variable selection, discrete response).
     - There are some additional attributes, and some missing.
 
-    The linear regression used to set `sigest` adds an intercept.
     """
 
-    def __init__(self, x_train, y_train, *,
+    def __init__(
+        self,
+        x_train,
+        y_train,
+        *,
         x_test=None,
         usequants=False,
         sigest=None,
@@ -173,6 +180,7 @@ class gbart:
         maxdepth=6,
         lamda=None,
         offset=None,
+        w=None,
         ntree=200,
         numcut=255,
         ndpost=1000,
@@ -180,26 +188,41 @@ class gbart:
         keepevery=1,
         printevery=100,
         seed=0,
-        initkw={},
-        ):
-
+        initkw=None,
+    ):
         x_train, x_train_fmt = self._process_predictor_input(x_train)
-
-        y_train, y_train_fmt = self._process_response_input(y_train)
+        y_train, _ = self._process_response_input(y_train)
         self._check_same_length(x_train, y_train)
+        if w is not None:
+            w, _ = self._process_response_input(w)
+            self._check_same_length(x_train, w)
 
         offset = self._process_offset_settings(y_train, offset)
         scale = self._process_scale_settings(y_train, k)
-        lamda, sigest = self._process_noise_variance_settings(x_train, y_train, sigest, sigdf, sigquant, lamda, offset)
+        lamda, sigest = self._process_noise_variance_settings(
+            x_train, y_train, sigest, sigdf, sigquant, lamda, offset
+        )
 
         splits, max_split = self._determine_splits(x_train, usequants, numcut)
         x_train = self._bin_predictors(x_train, splits)
-
-        y_train = self._transform_input(y_train, offset, scale)
-        lamda_scaled = lamda / (scale * scale)
-
-        mcmc_state = self._setup_mcmc(x_train, y_train, max_split, lamda_scaled, sigdf, power, base, maxdepth, ntree, initkw)
-        final_state, burnin_trace, main_trace = self._run_mcmc(mcmc_state, ndpost, nskip, keepevery, printevery, seed)
+        y_train, lamda_scaled = self._transform_input(y_train, lamda, offset, scale)
+
+        mcmc_state = self._setup_mcmc(
+            x_train,
+            y_train,
+            w,
+            max_split,
+            lamda_scaled,
+            sigdf,
+            power,
+            base,
+            maxdepth,
+            ntree,
+            initkw,
+        )
+        final_state, burnin_trace, main_trace = self._run_mcmc(
+            mcmc_state, ndpost, nskip, keepevery, printevery, seed
+        )
 
         sigma = self._extract_sigma(main_trace, scale)
         first_sigma = self._extract_sigma(burnin_trace, scale)
@@ -239,7 +262,7 @@ class gbart:
 
         Parameters
         ----------
-        x_test : array (m, p) or DataFrame
+        x_test : array (p, m) or DataFrame
             The test predictors.
 
         Returns
@@ -285,7 +308,9 @@ class gbart:
         assert get_length(x1) == get_length(x2)
 
     @staticmethod
-    def _process_noise_variance_settings(x_train, y_train, sigest, sigdf, sigquant, lamda, offset):
+    def _process_noise_variance_settings(
+        x_train, y_train, sigest, sigdf, sigquant, lamda, offset
+    ):
         if lamda is not None:
             return lamda, None
         else:
@@ -298,7 +323,7 @@ class gbart:
             else:
                 x_centered = x_train.T - x_train.mean(axis=1)
                 y_centered = y_train - y_train.mean()
-                    # centering is equivalent to adding an intercept column
+                # centering is equivalent to adding an intercept column
                 _, chisq, rank, _ = jnp.linalg.lstsq(x_centered, y_centered)
                 chisq = chisq.squeeze(0)
                 dof = len(y_train) - rank
@@ -336,11 +361,25 @@ class gbart:
         return prepcovars.bin_predictors(x, splits)
 
     @staticmethod
-    def _transform_input(y, offset, scale):
-        return (y - offset) / scale
+    def _transform_input(y, lamda, offset, scale):
+        y = (y - offset) / scale
+        lamda = lamda / (scale * scale)
+        return y, lamda
 
     @staticmethod
-    def _setup_mcmc(x_train, y_train, max_split, lamda, sigdf, power, base, maxdepth, ntree, initkw):
+    def _setup_mcmc(
+        x_train,
+        y_train,
+        w,
+        max_split,
+        lamda,
+        sigdf,
+        power,
+        base,
+        maxdepth,
+        ntree,
+        initkw,
+    ):
         depth = jnp.arange(maxdepth - 1)
         p_nonterminal = base / (1 + depth).astype(float) ** power
         sigma2_alpha = sigdf / 2
@@ -348,6 +387,7 @@ class gbart:
         kw = dict(
             X=x_train,
             y=y_train,
+            error_scale=w,
             max_split=max_split,
             num_trees=ntree,
             p_nonterminal=p_nonterminal,
@@ -355,17 +395,20 @@ class gbart:
             sigma2_beta=sigma2_beta,
             min_points_per_leaf=5,
         )
-        kw.update(initkw)
+        if initkw is not None:
+            kw.update(initkw)
         return mcmcstep.init(**kw)
 
     @staticmethod
     def _run_mcmc(mcmc_state, ndpost, nskip, keepevery, printevery, seed):
-        if isinstance(seed, jax.Array) and jnp.issubdtype(seed.dtype, jax.dtypes.prng_key):
+        if isinstance(seed, jax.Array) and jnp.issubdtype(
+            seed.dtype, jax.dtypes.prng_key
+        ):
             key = seed
         else:
             key = jax.random.key(seed)
         callback = mcmcloop.make_simple_print_callback(printevery)
-        return mcmcloop.run_mcmc(mcmc_state, nskip, ndpost, keepevery, callback, key)
+        return mcmcloop.run_mcmc(key, mcmc_state, nskip, ndpost, keepevery, callback)
 
     @staticmethod
     def _predict(trace, x):
@@ -379,9 +422,9 @@ class gbart:
     def _extract_sigma(trace, scale):
         return scale * jnp.sqrt(trace['sigma2'])
 
-
     def _show_tree(self, i_sample, i_tree, print_all=False):
         from . import debug
+
         trace = self._main_trace
         leaf_tree = trace['leaf_trees'][i_sample, i_tree]
         var_tree = trace['var_trees'][i_sample, i_tree]
@@ -396,7 +439,9 @@ class gbart:
         else:
             resid = bart['resid']
             alpha = bart['sigma2_alpha'] + resid.size / 2
-            norm2 = jnp.dot(resid, resid, preferred_element_type=bart['sigma2_beta'].dtype)
+            norm2 = jnp.dot(
+                resid, resid, preferred_element_type=bart['sigma2_beta'].dtype
+            )
             beta = bart['sigma2_beta'] + norm2 / 2
         sigma2 = beta / alpha
         return jnp.sqrt(sigma2) * self.scale
@@ -404,22 +449,32 @@ class gbart:
     def _compare_resid(self):
         bart = self._mcmc_state
         resid1 = bart['resid']
-        yhat = grove.evaluate_forest(bart['X'], bart['leaf_trees'], bart['var_trees'], bart['split_trees'], jnp.float32)
+        yhat = grove.evaluate_forest(
+            bart['X'],
+            bart['leaf_trees'],
+            bart['var_trees'],
+            bart['split_trees'],
+            jnp.float32,
+        )
         resid2 = bart['y'] - yhat
         return resid1, resid2
 
     def _avg_acc(self):
         trace = self._main_trace
+
         def acc(prefix):
             acc = trace[f'{prefix}_acc_count']
             prop = trace[f'{prefix}_prop_count']
             return acc.sum() / prop.sum()
+
         return acc('grow'), acc('prune')
 
     def _avg_prop(self):
         trace = self._main_trace
+
         def prop(prefix):
             return trace[f'{prefix}_prop_count'].sum()
+
         pgrow = prop('grow')
         pprune = prop('prune')
         total = pgrow + pprune
@@ -432,16 +487,21 @@ class gbart:
 
     def _depth_distr(self):
         from . import debug
+
         trace = self._main_trace
         split_trees = trace['split_trees']
         return debug.trace_depth_distr(split_trees)
 
     def _points_per_leaf_distr(self):
         from . import debug
-        return debug.trace_points_per_leaf_distr(self._main_trace, self._mcmc_state['X'])
+
+        return debug.trace_points_per_leaf_distr(
+            self._main_trace, self._mcmc_state['X']
+        )
 
     def _check_trees(self):
         from . import debug
+
         return debug.check_trace(self._main_trace, self._mcmc_state)
 
     def _tree_goes_bad(self):

bartz/__init__.py

@@ -1,6 +1,6 @@
 # bartz/src/bartz/__init__.py
 #
-# Copyright (c) 2024, Giacomo Petrillo
+# Copyright (c) 2024-2025, Giacomo Petrillo
 #
 # This file is part of bartz.
 #
@@ -28,13 +28,5 @@ Super-fast BART (Bayesian Additive Regression Trees) in Python
 See the manual at https://gattocrucco.github.io/bartz/docs
 """
 
-from ._version import __version__
-
-from . import BART
-
-from . import debug
-from . import grove
-from . import mcmcstep
-from . import mcmcloop
-from . import prepcovars
-from . import jaxext
+from . import BART, debug, grove, jaxext, mcmcloop, mcmcstep, prepcovars  # noqa: F401
+from ._version import __version__  # noqa: F401

bartz/_version.py

@@ -1 +1 @@
-__version__ = '0.4.1'
+__version__ = '0.5.0'

bartz/debug.py

@@ -1,21 +1,19 @@
 import functools
 
 import jax
-from jax import numpy as jnp
 from jax import lax
+from jax import numpy as jnp
 
-from . import grove
-from . import mcmcstep
-from . import jaxext
+from . import grove, jaxext
 
-def print_tree(leaf_tree, var_tree, split_tree, print_all=False):
 
+def print_tree(leaf_tree, var_tree, split_tree, print_all=False):
     tee = '├──'
     corner = '└──'
     join = '│  '
     space = '   '
     down = '┐'
-    bottom = '╢' # '┨' # 
+    bottom = '╢'  # '┨' #
 
     def traverse_tree(index, depth, indent, first_indent, next_indent, unused):
         if index >= len(leaf_tree):
@@ -58,7 +56,7 @@ def print_tree(leaf_tree, var_tree, split_tree, print_all=False):
 
         indent += next_indent
         unused = unused or is_leaf
-        
+
         if unused and not print_all:
             return
 
@@ -67,58 +65,80 @@ def print_tree(leaf_tree, var_tree, split_tree, print_all=False):
 
     traverse_tree(1, 0, '', '', '', False)
 
+
 def tree_actual_depth(split_tree):
     is_leaf = grove.is_actual_leaf(split_tree, add_bottom_level=True)
     depth = grove.tree_depths(is_leaf.size)
     depth = jnp.where(is_leaf, depth, 0)
     return jnp.max(depth)
 
+
 def forest_depth_distr(split_trees):
     depth = grove.tree_depth(split_trees) + 1
     depths = jax.vmap(tree_actual_depth)(split_trees)
     return jnp.bincount(depths, length=depth)
 
+
 def trace_depth_distr(split_trees_trace):
     return jax.vmap(forest_depth_distr)(split_trees_trace)
 
+
 def points_per_leaf_distr(var_tree, split_tree, X):
     traverse_tree = jax.vmap(grove.traverse_tree, in_axes=(1, None, None))
     indices = traverse_tree(X, var_tree, split_tree)
-    count_tree = jnp.zeros(2 * split_tree.size, dtype=jaxext.minimal_unsigned_dtype(indices.size))
+    count_tree = jnp.zeros(
+        2 * split_tree.size, dtype=jaxext.minimal_unsigned_dtype(indices.size)
+    )
     count_tree = count_tree.at[indices].add(1)
     is_leaf = grove.is_actual_leaf(split_tree, add_bottom_level=True).view(jnp.uint8)
     return jnp.bincount(count_tree, is_leaf, length=X.shape[1] + 1)
 
+
 def forest_points_per_leaf_distr(bart, X):
     distr = jnp.zeros(X.shape[1] + 1, int)
     trees = bart['var_trees'], bart['split_trees']
+
     def loop(distr, tree):
         return distr + points_per_leaf_distr(*tree, X), None
+
     distr, _ = lax.scan(loop, distr, trees)
     return distr
 
+
 def trace_points_per_leaf_distr(bart, X):
     def loop(_, bart):
         return None, forest_points_per_leaf_distr(bart, X)
+
     _, distr = lax.scan(loop, None, bart)
     return distr
 
+
 def check_types(leaf_tree, var_tree, split_tree, max_split):
     expected_var_dtype = jaxext.minimal_unsigned_dtype(max_split.size - 1)
     expected_split_dtype = max_split.dtype
-    return var_tree.dtype == expected_var_dtype and split_tree.dtype == expected_split_dtype
+    return (
+        var_tree.dtype == expected_var_dtype
+        and split_tree.dtype == expected_split_dtype
+    )
+
 
 def check_sizes(leaf_tree, var_tree, split_tree, max_split):
     return leaf_tree.size == 2 * var_tree.size == 2 * split_tree.size
 
+
 def check_unused_node(leaf_tree, var_tree, split_tree, max_split):
     return (var_tree[0] == 0) & (split_tree[0] == 0)
 
+
 def check_leaf_values(leaf_tree, var_tree, split_tree, max_split):
     return jnp.all(jnp.isfinite(leaf_tree))
 
+
 def check_stray_nodes(leaf_tree, var_tree, split_tree, max_split):
-    index = jnp.arange(2 * split_tree.size, dtype=jaxext.minimal_unsigned_dtype(2 * split_tree.size - 1))
+    index = jnp.arange(
+        2 * split_tree.size,
+        dtype=jaxext.minimal_unsigned_dtype(2 * split_tree.size - 1),
+    )
     parent_index = index >> 1
     is_not_leaf = split_tree.at[index].get(mode='fill', fill_value=0) != 0
     parent_is_leaf = split_tree[parent_index] == 0
@@ -126,6 +146,7 @@ def check_stray_nodes(leaf_tree, var_tree, split_tree, max_split):
     stray = stray.at[1].set(False)
     return ~jnp.any(stray)
 
+
 check_functions = [
     check_types,
     check_sizes,
@@ -134,6 +155,7 @@ check_functions = [
     check_stray_nodes,
 ]
 
+
 def check_tree(leaf_tree, var_tree, split_tree, max_split):
     error_type = jaxext.minimal_unsigned_dtype(2 ** len(check_functions) - 1)
     error = error_type(0)
@@ -144,15 +166,19 @@ def check_tree(leaf_tree, var_tree, split_tree, max_split):
         error |= bit
     return error
 
+
 def describe_error(error):
-    return [
-        func.__name__
-        for i, func in enumerate(check_functions)
-        if error & (1 << i)
-    ]
+    return [func.__name__ for i, func in enumerate(check_functions) if error & (1 << i)]
+
 
 check_forest = jax.vmap(check_tree, in_axes=(0, 0, 0, None))
 
+
 @functools.partial(jax.vmap, in_axes=(0, None))
 def check_trace(trace, state):
-    return check_forest(trace['leaf_trees'], trace['var_trees'], trace['split_trees'], state['max_split'])
+    return check_forest(
+        trace['leaf_trees'],
+        trace['var_trees'],
+        trace['split_trees'],
+        state['max_split'],
+    )

bartz/grove.py

@@ -1,6 +1,6 @@
 # bartz/src/bartz/grove.py
 #
-# Copyright (c) 2024, Giacomo Petrillo
+# Copyright (c) 2024-2025, Giacomo Petrillo
 #
 # This file is part of bartz.
 #
@@ -43,12 +43,12 @@ Since the nodes at the bottom can only be leaves and not decision nodes, the 'va
 import functools
 import math
 
-import jax
-from jax import numpy as jnp
 from jax import lax
+from jax import numpy as jnp
 
 from . import jaxext
 
+
 def make_tree(depth, dtype):
     """
     Make an array to represent a binary tree.
@@ -66,7 +66,8 @@ def make_tree(depth, dtype):
     tree : array
         An array of zeroes with shape (2 ** depth,).
     """
-    return jnp.zeros(2 ** depth, dtype)
+    return jnp.zeros(2**depth, dtype)
+
 
 def tree_depth(tree):
     """
@@ -85,6 +86,7 @@ def tree_depth(tree):
     """
     return int(round(math.log2(tree.shape[-1])))
 
+
 def traverse_tree(x, var_tree, split_tree):
     """
     Find the leaf where a point falls into.
@@ -125,6 +127,7 @@ def traverse_tree(x, var_tree, split_tree):
     (_, index), _ = lax.scan(loop, carry, None, depth, unroll=16)
     return index
 
+
 @functools.partial(jaxext.vmap_nodoc, in_axes=(None, 0, 0))
 @functools.partial(jaxext.vmap_nodoc, in_axes=(1, None, None))
 def traverse_forest(X, var_trees, split_trees):
@@ -147,6 +150,7 @@ def traverse_forest(X, var_trees, split_trees):
     """
     return traverse_tree(X, var_trees, split_trees)
 
+
 def evaluate_forest(X, leaf_trees, var_trees, split_trees, dtype=None, sum_trees=True):
     """
     Evaluate a ensemble of trees at an array of points.
@@ -178,11 +182,12 @@ def evaluate_forest(X, leaf_trees, var_trees, split_trees, dtype=None, sum_trees
     leaves = leaf_trees[tree_index[:, None], indices]
     if sum_trees:
         return jnp.sum(leaves, axis=0, dtype=dtype)
-            # this sum suggests to swap the vmaps, but I think it's better for X
-            # copying to keep it that way
+    # this sum suggests to swap the vmaps, but I think it's better for X
+    # copying to keep it that way
     else:
         return leaves
 
+
 def is_actual_leaf(split_tree, *, add_bottom_level=False):
     """
     Return a mask indicating the leaf nodes in a tree.
@@ -211,6 +216,7 @@ def is_actual_leaf(split_tree, *, add_bottom_level=False):
     parent_nonleaf = parent_nonleaf.at[1].set(True)
     return is_leaf & parent_nonleaf
 
+
 def is_leaves_parent(split_tree):
     """
     Return a mask indicating the nodes with leaf (and only leaf) children.
@@ -225,14 +231,17 @@ def is_leaves_parent(split_tree):
     is_leaves_parent : bool array (2 ** (d - 1),)
         The mask indicating which nodes have leaf children.
     """
-    index = jnp.arange(split_tree.size, dtype=jaxext.minimal_unsigned_dtype(2 * split_tree.size - 1))
-    left_index = index << 1 # left child
-    right_index = left_index + 1 # right child
+    index = jnp.arange(
+        split_tree.size, dtype=jaxext.minimal_unsigned_dtype(2 * split_tree.size - 1)
+    )
+    left_index = index << 1  # left child
+    right_index = left_index + 1  # right child
     left_leaf = split_tree.at[left_index].get(mode='fill', fill_value=0) == 0
     right_leaf = split_tree.at[right_index].get(mode='fill', fill_value=0) == 0
     is_not_leaf = split_tree.astype(bool)
     return is_not_leaf & left_leaf & right_leaf
-        # the 0-th item has split == 0, so it's not counted
+    # the 0-th item has split == 0, so it's not counted
+
 
 def tree_depths(tree_length):
     """
@@ -253,7 +262,7 @@ def tree_depths(tree_length):
     depths = []
     depth = 0
     for i in range(tree_length):
-        if i == 2 ** depth:
+        if i == 2**depth:
             depth += 1
         depths.append(depth - 1)
     depths[0] = 0