bartz 0.0.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

bartz/{interface.py → BART.py}

@@ -1,4 +1,4 @@
-# bartz/src/bartz/interface.py
+# bartz/src/bartz/BART.py
 #
 # Copyright (c) 2024, Giacomo Petrillo
 #
@@ -33,12 +33,12 @@ from . import mcmcstep
 from . import mcmcloop
 from . import prepcovars
 
-class BART:
+class gbart:
     """
     Nonparametric regression with Bayesian Additive Regression Trees (BART).
 
     Regress `y_train` on `x_train` with a latent mean function represented as
-    a sum of decision trees. The inference is carried out by estimating the
+    a sum of decision trees. The inference is carried out by sampling the
     posterior distribution of the tree ensemble with an MCMC.
 
     Parameters
@@ -86,7 +86,7 @@ class BART:
         predictor is binned such that its distribution in `x_train` is
         approximately uniform across bins. The number of bins is at most the
         number of unique values appearing in `x_train`, or ``numcut + 1``.
-        Before running the algorithm, the predictors are compressed to th
+        Before running the algorithm, the predictors are compressed to the
         smallest integer type that fits the bin indices, so `numcut` is best set
         to the maximum value of an unsigned integer type.
     ndpost : int, default 1000
@@ -133,7 +133,7 @@ class BART:
 
     Notes
     -----
-    This interface imitates the function `wbart` from the R package `BART
+    This interface imitates the function `gbart` from the R package `BART
     <https://cran.r-project.org/package=BART>`_, but with these differences:
 
     - If `x_train` and `x_test` are matrices, they have one predictor per row
@@ -142,6 +142,7 @@ class BART:
     - `usequants` is always `True`.
     - `rm_const` is always `False`.
     - The default `numcut` is 255 instead of 100.
+    - A lot of functionality is missing (variable selection, discrete response).
     - There are some additional attributes, and some missing.
     """
 
@@ -321,7 +322,7 @@ class BART:
         p_nonterminal = base / (1 + depth).astype(float) ** power
         sigma2_alpha = sigdf / 2
         sigma2_beta = lamda * sigma2_alpha
-        return mcmcstep.make_bart(
+        return mcmcstep.init(
             X=x_train,
             y=y_train,
             max_split=max_split,
@@ -354,13 +355,6 @@ class BART:
         return scale * jnp.sqrt(trace['sigma2'])
 
     
-    def _predict_debug(self, x_test):
-        from . import debug
-        x_test, x_test_fmt = self._process_predictor_input(x_test)
-        self._check_compatible_formats(x_test_fmt, self._x_train_fmt)
-        x_test = self._bin_predictors(x_test, self._splits)
-        return debug.trace_evaluate_trees(self._main_trace, x_test)
-
     def _show_tree(self, i_sample, i_tree, print_all=False):
         from . import debug
         trace = self._main_trace
@@ -385,7 +379,7 @@ class BART:
     def _compare_resid(self):
         bart = self._mcmc_state
         resid1 = bart['resid']
-        yhat = grove.evaluate_tree_vmap_x(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
 
@@ -427,7 +421,5 @@ class BART:
 
     def _tree_goes_bad(self):
         bad = self._check_trees().astype(bool)
-        bad_before = bad[:-1]
-        bad_after = bad[1:]
-        goes_bad = bad_after & ~bad_before
-        return jnp.pad(goes_bad, [(1, 0), (0, 0)])
+        bad_before = jnp.pad(bad[:-1], [(1, 0), (0, 0)])
+        return bad & ~bad_before

bartz/__init__.py

@@ -28,8 +28,13 @@ A jax implementation of BART
 See the manual at https://gattocrucco.github.io/bartz/docs
 """
 
-__version__ = '0.0.1'
+from ._version import __version__
 
-from .interface import BART
+from . import BART
 
 from . import debug
+from . import grove
+from . import mcmcstep
+from . import mcmcloop
+from . import prepcovars
+from . import jaxext

bartz/_version.py

@@ -0,0 +1 @@
+__version__ = '0.2.0'

bartz/debug.py

@@ -6,22 +6,7 @@ from jax import lax
 
 from . import grove
 from . import mcmcstep
-
-def trace_evaluate_trees(bart, X):
-    """
-    Evaluate all trees, for all samples, at all x. Out axes:
-        0: mcmc sample
-        1: tree
-        2: X
-    """
-    def loop(_, bart):
-        return None, evaluate_all_trees(X, bart['leaf_trees'], bart['var_trees'], bart['split_trees'])
-    _, y = lax.scan(loop, None, bart)
-    return y
-
-@functools.partial(jax.vmap, in_axes=(None, 0, 0, 0)) # vectorize over forest
-def evaluate_all_trees(X, leaf_trees, var_trees, split_trees):
-    return grove.evaluate_tree_vmap_x(X, leaf_trees, var_trees, split_trees, jnp.float32)
+from . import jaxext
 
 def print_tree(leaf_tree, var_tree, split_tree, print_all=False):
 
@@ -97,8 +82,10 @@ 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):
-    dummy = jnp.ones(X.shape[1], jnp.uint8)
-    _, count_tree = mcmcstep.agg_values(X, var_tree, split_tree, dummy, dummy.dtype)
+    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 = 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)
 
@@ -117,7 +104,7 @@ def trace_points_per_leaf_distr(bart, X):
     return distr
 
 def check_types(leaf_tree, var_tree, split_tree, max_split):
-    expected_var_dtype = grove.minimal_unsigned_dtype(max_split.size - 1)
+    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
 
@@ -125,13 +112,13 @@ 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 (leaf_tree[0] == 0) & (var_tree[0] == 0) & (split_tree[0] == 0)
+    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=grove.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
@@ -148,7 +135,7 @@ check_functions = [
 ]
 
 def check_tree(leaf_tree, var_tree, split_tree, max_split):
-    error_type = grove.minimal_unsigned_dtype(2 ** len(check_functions) - 1)
+    error_type = jaxext.minimal_unsigned_dtype(2 ** len(check_functions) - 1)
     error = error_type(0)
     for i, func in enumerate(check_functions):
         ok = func(leaf_tree, var_tree, split_tree, max_split)

bartz/grove.py

@@ -28,13 +28,15 @@ Functions to create and manipulate binary trees.
 
 A tree is represented with arrays as a heap. The root node is at index 1. The children nodes of a node at index :math:`i` are at indices :math:`2i` (left child) and :math:`2i + 1` (right child). The array element at index 0 is unused.
 
-A decision tree is represented by tree arrays: 'leaf', 'var', and 'split'. The 'leaf' array contains the values in the leaves. The 'var' array contains the axes along which the decision nodes operate. The 'split' array contains the decision boundaries.
+A decision tree is represented by tree arrays: 'leaf', 'var', and 'split'.
 
-Whether a node is a leaf is indicated by the corresponding 'split' element being 0.
+The 'leaf' array contains the values in the leaves.
 
-Since the nodes at the bottom can only be leaves and not decision nodes, the 'var' and 'split' arrays have half the length of the 'leaf' array.
+The 'var' array contains the axes along which the decision nodes operate.
+
+The 'split' array contains the decision boundaries. The boundaries are open on the right, i.e., a point belongs to the left child iff x < split. Whether a node is a leaf is indicated by the corresponding 'split' element being 0.
 
-The unused array element at index 0 is always fixed to 0 by convention.
+Since the nodes at the bottom can only be leaves and not decision nodes, the 'var' and 'split' arrays have half the length of the 'leaf' array.
 
 """
 
@@ -63,24 +65,18 @@ def make_tree(depth, dtype):
     -------
     tree : array
         An array of zeroes with shape (2 ** depth,).
-
-    Notes
-    -----
-    The tree is represented as a heap, with the root node at index 1, and the
-    children of the node at index i at indices 2 * i and 2 * i + 1. The element
-    at index 0 is unused.
     """
     return jnp.zeros(2 ** depth, dtype)
 
 def tree_depth(tree):
     """
-    Return the maximum depth of a binary tree created by `make_tree`.
+    Return the maximum depth of a tree.
 
     Parameters
     ----------
     tree : array
-        A binary tree created by `make_tree`. If the array is ND, the tree
-        structure is assumed to be along the last axis.
+        A tree created by `make_tree`. If the array is ND, the tree structure is
+        assumed to be along the last axis.
 
     Returns
     -------
@@ -89,120 +85,98 @@ def tree_depth(tree):
     """
     return int(round(math.log2(tree.shape[-1])))
 
-def evaluate_tree(X, leaf_trees, var_trees, split_trees, out_dtype):
+def traverse_tree(x, var_tree, split_tree):
     """
-    Evaluate a decision tree or forest.
+    Find the leaf where a point falls into.
 
     Parameters
     ----------
-    X : array (p,)
+    x : array (p,)
         The coordinates to evaluate the tree at.
-    leaf_trees : array (n,) or (m, n)
-        The leaf values of the tree or forest. If the input is a forest, the
-        first axis is the tree index, and the values are summed.
-    var_trees : array (n,) or (m, n)
-        The variable indices of the tree or forest. Each index is in [0, p) and
-        indicates which value of `X` to consider.
-    split_trees : array (n,) or (m, n)
-        The split values of the tree or forest. Leaf nodes are indicated by the
-        condition `split == 0`. If non-zero, the node has children, and its left
-        children is assigned points which satisfy `x < split`.
-    out_dtype : dtype
-        The dtype of the output.
+    var_tree : array (2 ** (d - 1),)
+        The decision axes of the tree.
+    split_tree : array (2 ** (d - 1),)
+        The decision boundaries of the tree.
 
     Returns
     -------
-    out : scalar
-        The value of the tree or forest at the given point.
+    index : int
+        The index of the leaf.
     """
 
-    is_forest = leaf_trees.ndim == 2
-    if is_forest:
-        m, _ = leaf_trees.shape
-        forest_shape = m,
-        tree_index = jnp.arange(m, dtype=minimal_unsigned_dtype(m - 1)),
-    else:
-        forest_shape = ()
-        tree_index = ()
-
     carry = (
-        jnp.zeros(forest_shape, bool),
-        jnp.zeros((), out_dtype),
-        jnp.ones(forest_shape, minimal_unsigned_dtype(leaf_trees.shape[-1] - 1))
+        jnp.zeros((), bool),
+        jnp.ones((), jaxext.minimal_unsigned_dtype(2 * var_tree.size - 1)),
     )
 
     def loop(carry, _):
-        leaf_found, out, node_index = carry
-
-        is_leaf = split_trees.at[tree_index + (node_index,)].get(mode='fill', fill_value=0) == 0
-        leaf_value = leaf_trees[tree_index + (node_index,)]
-        if is_forest:
-            leaf_sum = jnp.sum(leaf_value, where=is_leaf) # TODO set dtype to large float
-                # alternative: dot(is_leaf, leaf_value):
-                # - maybe faster
-                # - maybe less accurate
-                # - fucked by nans
-        else:
-            leaf_sum = jnp.where(is_leaf, leaf_value, 0)
-        out += leaf_sum
-        leaf_found |= is_leaf
-        
-        split = split_trees.at[tree_index + (node_index,)].get(mode='fill', fill_value=0)
-        var = var_trees.at[tree_index + (node_index,)].get(mode='fill', fill_value=0)
-        x = X[var]
+        leaf_found, index = carry
+
+        split = split_tree[index]
+        var = var_tree[index]
         
-        node_index <<= 1
-        node_index += x >= split
-        node_index = jnp.where(leaf_found, 0, node_index)
+        leaf_found |= split == 0
+        child_index = (index << 1) + (x[var] >= split)
+        index = jnp.where(leaf_found, index, child_index)
 
-        carry = leaf_found, out, node_index
-        return carry, _
+        return (leaf_found, index), None
 
-    depth = tree_depth(leaf_trees)
-    (_, out, _), _ = lax.scan(loop, carry, None, depth)
-    return out
+    depth = tree_depth(var_tree)
+    (_, index), _ = lax.scan(loop, carry, None, depth, unroll=16)
+    return index
 
-def minimal_unsigned_dtype(max_value):
+@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):
     """
-    Return the smallest unsigned integer dtype that can represent a given
-    maximum value.
+    Find the leaves where points fall into.
+
+    Parameters
+    ----------
+    X : array (p, n)
+        The coordinates to evaluate the trees at.
+    var_trees : array (m, 2 ** (d - 1))
+        The decision axes of the trees.
+    split_trees : array (m, 2 ** (d - 1))
+        The decision boundaries of the trees.
+
+    Returns
+    -------
+    indices : array (m, n)
+        The indices of the leaves.
     """
-    if max_value < 2 ** 8:
-        return jnp.uint8
-    if max_value < 2 ** 16:
-        return jnp.uint16
-    if max_value < 2 ** 32:
-        return jnp.uint32
-    return jnp.uint64
-
-@functools.partial(jaxext.vmap_nodoc, in_axes=(1, None, None, None, None), out_axes=0)
-def evaluate_tree_vmap_x(X, leaf_trees, var_trees, split_trees, out_dtype):
+    return traverse_tree(X, var_trees, split_trees)
+
+def evaluate_forest(X, leaf_trees, var_trees, split_trees, dtype):
     """
-    Evaluate a decision tree or forest over multiple points.
+    Evaluate a ensemble of trees at an array of points.
 
     Parameters
     ----------
     X : array (p, n)
-        The points to evaluate the tree at.
-    leaf_trees : array (n,) or (m, n)
+        The coordinates to evaluate the trees at.
+    leaf_trees : array (m, 2 ** d)
         The leaf values of the tree or forest. If the input is a forest, the
         first axis is the tree index, and the values are summed.
-    var_trees : array (n,) or (m, n)
-        The variable indices of the tree or forest. Each index is in [0, p) and
-        indicates which value of `X` to consider.
-    split_trees : array (n,) or (m, n)
-        The split values of the tree or forest. Leaf nodes are indicated by the
-        condition `split == 0`. If non-zero, the node has children, and its left
-        children is assigned points which satisfy `x < split`.
-    out_dtype : dtype
+    var_trees : array (m, 2 ** (d - 1))
+        The decision axes of the trees.
+    split_trees : array (m, 2 ** (d - 1))
+        The decision boundaries of the trees.
+    dtype : dtype
         The dtype of the output.
 
     Returns
     -------
-    out : (n,)
-        The value of the tree or forest at each point.
+    out : array (n,)
+        The sum of the values of the trees at the points in `X`.
     """
-    return evaluate_tree(X, leaf_trees, var_trees, split_trees, out_dtype)
+    indices = traverse_forest(X, var_trees, split_trees)
+    ntree, _ = leaf_trees.shape
+    tree_index = jnp.arange(ntree, dtype=jaxext.minimal_unsigned_dtype(ntree - 1))[:, None]
+    leaves = leaf_trees[tree_index, indices]
+    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
 
 def is_actual_leaf(split_tree, *, add_bottom_level=False):
     """
@@ -226,7 +200,7 @@ def is_actual_leaf(split_tree, *, add_bottom_level=False):
     if add_bottom_level:
         size *= 2
         is_leaf = jnp.concatenate([is_leaf, jnp.ones_like(is_leaf)])
-    index = jnp.arange(size, dtype=minimal_unsigned_dtype(size - 1))
+    index = jnp.arange(size, dtype=jaxext.minimal_unsigned_dtype(size - 1))
     parent_index = index >> 1
     parent_nonleaf = split_tree[parent_index].astype(bool)
     parent_nonleaf = parent_nonleaf.at[1].set(True)
@@ -239,14 +213,14 @@ def is_leaves_parent(split_tree):
     Parameters
     ----------
     split_tree : int array (2 ** (d - 1),)
-        The splitting points of the tree.
+        The decision boundaries of the tree.
 
     Returns
     -------
     is_leaves_parent : bool array (2 ** (d - 1),)
         The mask indicating which nodes have leaf children.
     """
-    index = jnp.arange(split_tree.size, dtype=minimal_unsigned_dtype(2 * split_tree.size - 1))
+    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
@@ -278,25 +252,4 @@ def tree_depths(tree_length):
             depth += 1
         depths.append(depth - 1)
     depths[0] = 0
-    return jnp.array(depths, minimal_unsigned_dtype(max(depths)))
-
-def index_depth(index, tree_length):
-    """
-    Return the depth of a node in a binary tree.
-
-    Parameters
-    ----------
-    index : int
-        The index of the node.
-    tree_length : int
-        The length of the tree array, i.e., 2 ** d.
-
-    Returns
-    -------
-    depth : int
-        The depth of the node. The root node (index 1) has depth 0. The depth is
-        the position of the most significant non-zero bit in the index. If
-        ``index == 0``, return -1.
-    """
-    depths = tree_depths(tree_length)
-    return depths[index]
+    return jnp.array(depths, jaxext.minimal_unsigned_dtype(max(depths)))