bartz 0.6.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

bartz/grove.py

@@ -22,93 +22,113 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-"""
+"""Functions to create and manipulate binary decision trees."""
 
-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.
+import math
+from functools import partial
+from typing import Protocol
 
-The 'var' array contains the axes along which the decision nodes operate.
+import jax
+from jax import jit, lax
+from jax import numpy as jnp
+from jaxtyping import Array, Bool, DTypeLike, Float32, Int32, Real, Shaped, UInt
 
-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. Unused nodes also have split set to 0.
+from bartz.jaxext import minimal_unsigned_dtype, vmap_nodoc
 
-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.
 
-"""
+class TreeHeaps(Protocol):
+    """A protocol for dataclasses that represent trees.
 
-import functools
-import math
+    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.
 
-import jax
-from jax import lax
-from jax import numpy as jnp
+    Parameters
+    ----------
+    leaf_tree
+        The values in the leaves of the trees. This array can be dirty, i.e.,
+        unused nodes can have whatever value.
+    var_tree
+        The axes along which the decision nodes operate. This array can be
+        dirty but for the always unused node at index 0 which must be set to 0.
+    split_tree
+        The decision boundaries of the trees. 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. Unused nodes also have split set to 0. This array can't be dirty.
+
+    Notes
+    -----
+    Since the nodes at the bottom can only be leaves and not decision nodes,
+    `var_tree` and `split_tree` are half as long as `leaf_tree`.
+    """
 
-from . import jaxext
+    leaf_tree: Float32[Array, '* 2**d']
+    var_tree: UInt[Array, '* 2**(d-1)']
+    split_tree: UInt[Array, '* 2**(d-1)']
 
 
-def make_tree(depth, dtype):
+def make_tree(depth: int, dtype: DTypeLike) -> Shaped[Array, ' 2**{depth}']:
     """
     Make an array to represent a binary tree.
 
     Parameters
     ----------
-    depth : int
+    depth
         The maximum depth of the tree. Depth 1 means that there is only a root
         node.
-    dtype : dtype
+    dtype
         The dtype of the array.
 
     Returns
     -------
-    tree : array
-        An array of zeroes with shape (2 ** depth,).
+    An array of zeroes with the appropriate shape.
     """
     return jnp.zeros(2**depth, dtype)
 
 
-def tree_depth(tree):
+def tree_depth(tree: Shaped[Array, '* 2**d']) -> int:
     """
     Return the maximum depth of a tree.
 
     Parameters
     ----------
-    tree : array
+    tree
         A tree created by `make_tree`. If the array is ND, the tree structure is
         assumed to be along the last axis.
 
     Returns
     -------
-    depth : int
-        The maximum depth of the tree.
+    The maximum depth of the tree.
     """
-    return int(round(math.log2(tree.shape[-1])))
+    return round(math.log2(tree.shape[-1]))
 
 
-def traverse_tree(x, var_tree, split_tree):
+def traverse_tree(
+    x: Real[Array, ' p'],
+    var_tree: UInt[Array, ' 2**(d-1)'],
+    split_tree: UInt[Array, ' 2**(d-1)'],
+) -> Int32[Array, '']:
     """
     Find the leaf where a point falls into.
 
     Parameters
     ----------
-    x : array (p,)
+    x
         The coordinates to evaluate the tree at.
-    var_tree : array (2 ** (d - 1),)
+    var_tree
         The decision axes of the tree.
-    split_tree : array (2 ** (d - 1),)
+    split_tree
         The decision boundaries of the tree.
 
     Returns
     -------
-    index : int
-        The index of the leaf.
+    The index of the leaf.
     """
     carry = (
         jnp.zeros((), bool),
-        jnp.ones((), jaxext.minimal_unsigned_dtype(2 * var_tree.size - 1)),
+        jnp.ones((), minimal_unsigned_dtype(2 * var_tree.size - 1)),
     )
 
     def loop(carry, _):
@@ -128,111 +148,107 @@ def traverse_tree(x, var_tree, split_tree):
     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):
+@partial(vmap_nodoc, in_axes=(None, 0, 0))
+@partial(vmap_nodoc, in_axes=(1, None, None))
+def traverse_forest(
+    X: Real[Array, 'p n'],
+    var_trees: UInt[Array, 'm 2**(d-1)'],
+    split_trees: UInt[Array, 'm 2**(d-1)'],
+) -> Int32[Array, 'm n']:
     """
     Find the leaves where points fall into.
 
     Parameters
     ----------
-    X : array (p, n)
+    X
         The coordinates to evaluate the trees at.
-    var_trees : array (m, 2 ** (d - 1))
+    var_trees
         The decision axes of the trees.
-    split_trees : array (m, 2 ** (d - 1))
+    split_trees
         The decision boundaries of the trees.
 
     Returns
     -------
-    indices : array (m, n)
-        The indices of the leaves.
+    The indices of the leaves.
     """
     return traverse_tree(X, var_trees, split_trees)
 
 
-def evaluate_forest(X, leaf_trees, var_trees, split_trees, dtype=None, sum_trees=True):
+def evaluate_forest(
+    X: UInt[Array, 'p n'], trees: TreeHeaps, *, sum_trees: bool = True
+) -> Float32[Array, ' n'] | Float32[Array, 'm n']:
     """
     Evaluate a ensemble of trees at an array of points.
 
     Parameters
     ----------
-    X : array (p, n)
+    X
         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 (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, optional
-        The dtype of the output. Ignored if `sum_trees` is `False`.
-    sum_trees : bool, default True
+    trees
+        The tree heaps, with batch shape (m,).
+    sum_trees
         Whether to sum the values across trees.
 
     Returns
     -------
-    out : array (n,) or (m, n)
-        The (sum of) the values of the trees at the points in `X`.
+    The (sum of) the values of the trees at the points in `X`.
     """
-    indices = traverse_forest(X, var_trees, split_trees)
-    ntree, _ = leaf_trees.shape
-    tree_index = jnp.arange(ntree, dtype=jaxext.minimal_unsigned_dtype(ntree - 1))
-    leaves = leaf_trees[tree_index[:, None], indices]
+    indices = traverse_forest(X, trees.var_tree, trees.split_tree)
+    ntree, _ = trees.leaf_tree.shape
+    tree_index = jnp.arange(ntree, dtype=minimal_unsigned_dtype(ntree - 1))
+    leaves = trees.leaf_tree[tree_index[:, None], indices]
     if sum_trees:
-        return jnp.sum(leaves, axis=0, dtype=dtype)
+        return jnp.sum(leaves, axis=0, dtype=jnp.float32)
     # 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):
+def is_actual_leaf(
+    split_tree: UInt[Array, ' 2**(d-1)'], *, add_bottom_level: bool = False
+) -> Bool[Array, ' 2**(d-1)'] | Bool[Array, ' 2**d']:
     """
     Return a mask indicating the leaf nodes in a tree.
 
     Parameters
     ----------
-    split_tree : int array (2 ** (d - 1),)
+    split_tree
         The splitting points of the tree.
-    add_bottom_level : bool, default False
+    add_bottom_level
         If True, the bottom level of the tree is also considered.
 
     Returns
     -------
-    is_actual_leaf : bool array (2 ** (d - 1) or 2 ** d,)
-        The mask indicating the leaf nodes. The length is doubled if
-        `add_bottom_level` is True.
+    The mask marking the leaf nodes. Length doubled if `add_bottom_level` is True.
     """
     size = split_tree.size
     is_leaf = split_tree == 0
     if add_bottom_level:
         size *= 2
         is_leaf = jnp.concatenate([is_leaf, jnp.ones_like(is_leaf)])
-    index = jnp.arange(size, dtype=jaxext.minimal_unsigned_dtype(size - 1))
+    index = jnp.arange(size, dtype=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)
     return is_leaf & parent_nonleaf
 
 
-def is_leaves_parent(split_tree):
+def is_leaves_parent(split_tree: UInt[Array, ' 2**(d-1)']) -> Bool[Array, ' 2**(d-1)']:
     """
     Return a mask indicating the nodes with leaf (and only leaf) children.
 
     Parameters
     ----------
-    split_tree : int array (2 ** (d - 1),)
+    split_tree
         The decision boundaries of the tree.
 
     Returns
     -------
-    is_leaves_parent : bool array (2 ** (d - 1),)
-        The mask indicating which nodes have leaf children.
+    The mask indicating which nodes have leaf children.
     """
     index = jnp.arange(
-        split_tree.size, dtype=jaxext.minimal_unsigned_dtype(2 * split_tree.size - 1)
+        split_tree.size, dtype=minimal_unsigned_dtype(2 * split_tree.size - 1)
     )
     left_index = index << 1  # left child
     right_index = left_index + 1  # right child
@@ -243,21 +259,24 @@ def is_leaves_parent(split_tree):
     # the 0-th item has split == 0, so it's not counted
 
 
-def tree_depths(tree_length):
+def tree_depths(tree_length: int) -> Int32[Array, ' {tree_length}']:
     """
     Return the depth of each node in a binary tree.
 
     Parameters
     ----------
-    tree_length : int
+    tree_length
         The length of the tree array, i.e., 2 ** d.
 
     Returns
     -------
-    depth : array (tree_length,)
-        The depth of each node. The root node (index 1) has depth 0. The depth
-        is the position of the most significant non-zero bit in the index. The
-        first element (the unused node) is marked as depth 0.
+    The depth of each node.
+
+    Notes
+    -----
+    The root node (index 1) has depth 0. The depth is the position of the most
+    significant non-zero bit in the index. The first element (the unused node)
+    is marked as depth 0.
     """
     depths = []
     depth = 0
@@ -266,22 +285,21 @@ def tree_depths(tree_length):
             depth += 1
         depths.append(depth - 1)
     depths[0] = 0
-    return jnp.array(depths, jaxext.minimal_unsigned_dtype(max(depths)))
+    return jnp.array(depths, minimal_unsigned_dtype(max(depths)))
 
 
-def is_used(split_tree):
+def is_used(split_tree: UInt[Array, ' 2**(d-1)']) -> Bool[Array, ' 2**d']:
     """
     Return a mask indicating the used nodes in a tree.
 
     Parameters
     ----------
-    split_tree : int array (2 ** (d - 1),)
+    split_tree
         The decision boundaries of the tree.
 
     Returns
     -------
-    is_used : bool array (2 ** d,)
-        A mask indicating which nodes are actually used.
+    A mask indicating which nodes are actually used.
     """
     internal_node = split_tree.astype(bool)
     internal_node = jnp.concatenate([internal_node, jnp.zeros_like(internal_node)])
@@ -289,22 +307,50 @@ def is_used(split_tree):
     return internal_node | actual_leaf
 
 
-def forest_fill(split_trees):
+@jit
+def forest_fill(split_tree: UInt[Array, 'num_trees 2**(d-1)']) -> Float32[Array, '']:
     """
     Return the fraction of used nodes in a set of trees.
 
     Parameters
     ----------
-    split_trees : array (m, 2 ** (d - 1),)
+    split_tree
         The decision boundaries of the trees.
 
     Returns
     -------
-    fill : float
-        The number of tree nodes in the forest over the maximum number that
-        could be stored in the arrays.
+    Number of tree nodes over the maximum number that could be stored.
     """
-    m, _ = split_trees.shape
-    used = jax.vmap(is_used)(split_trees)
+    num_trees, _ = split_tree.shape
+    used = jax.vmap(is_used)(split_tree)
     count = jnp.count_nonzero(used)
-    return count / (used.size - m)
+    return count / (used.size - num_trees)
+
+
+def var_histogram(
+    p: int, var_tree: UInt[Array, '* 2**(d-1)'], split_tree: UInt[Array, '* 2**(d-1)']
+) -> Int32[Array, ' {p}']:
+    """
+    Count how many times each variable appears in a tree.
+
+    Parameters
+    ----------
+    p
+        The number of variables (the maximum value that can occur in
+        `var_tree` is ``p - 1``).
+    var_tree
+        The decision axes of the tree.
+    split_tree
+        The decision boundaries of the tree.
+
+    Returns
+    -------
+    The histogram of the variables used in the tree.
+
+    Notes
+    -----
+    If there are leading axes in the tree arrays (i.e., multiple trees), the
+    returned counts are cumulative over trees.
+    """
+    is_internal = split_tree.astype(bool)
+    return jnp.zeros(p, int).at[var_tree].add(is_internal)

bartz/jaxext/__init__.py

@@ -0,0 +1,213 @@
+# bartz/src/bartz/jaxext/__init__.py
+#
+# Copyright (c) 2024-2025, Giacomo Petrillo
+#
+# This file is part of bartz.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Additions to jax."""
+
+import functools
+import math
+from collections.abc import Sequence
+
+import jax
+from jax import numpy as jnp
+from jax import random
+from jax.lax import scan
+from jax.scipy.special import ndtr
+from jaxtyping import Array, Bool, Float32, Key, Scalar, Shaped
+
+from bartz.jaxext._autobatch import autobatch  # noqa: F401
+from bartz.jaxext.scipy.special import ndtri
+
+
+def vmap_nodoc(fun, *args, **kw):
+    """
+    Acts like `jax.vmap` but preserves the docstring of the function unchanged.
+
+    This is useful if the docstring already takes into account that the
+    arguments have additional axes due to vmap.
+    """
+    doc = fun.__doc__
+    fun = jax.vmap(fun, *args, **kw)
+    fun.__doc__ = doc
+    return fun
+
+
+def minimal_unsigned_dtype(value):
+    """Return the smallest unsigned integer dtype that can represent `value`."""
+    if value < 2**8:
+        return jnp.uint8
+    if value < 2**16:
+        return jnp.uint16
+    if value < 2**32:
+        return jnp.uint32
+    return jnp.uint64
+
+
+@functools.partial(jax.jit, static_argnums=(1,))
+def unique(
+    x: Shaped[Array, ' _'], size: int, fill_value: Scalar
+) -> tuple[Shaped[Array, ' {size}'], int]:
+    """
+    Restricted version of `jax.numpy.unique` that uses less memory.
+
+    Parameters
+    ----------
+    x
+        The input array.
+    size
+        The length of the output.
+    fill_value
+        The value to fill the output with if `size` is greater than the number
+        of unique values in `x`.
+
+    Returns
+    -------
+    out : Shaped[Array, '{size}']
+        The unique values in `x`, sorted, and right-padded with `fill_value`.
+    actual_length : int
+        The number of used values in `out`.
+    """
+    if x.size == 0:
+        return jnp.full(size, fill_value, x.dtype), 0
+    if size == 0:
+        return jnp.empty(0, x.dtype), 0
+    x = jnp.sort(x)
+
+    def loop(carry, x):
+        i_out, last, out = carry
+        i_out = jnp.where(x == last, i_out, i_out + 1)
+        out = out.at[i_out].set(x)
+        return (i_out, x, out), None
+
+    carry = 0, x[0], jnp.full(size, fill_value, x.dtype)
+    (actual_length, _, out), _ = scan(loop, carry, x[:size])
+    return out, actual_length + 1
+
+
+class split:
+    """
+    Split a key into `num` keys.
+
+    Parameters
+    ----------
+    key
+        The key to split.
+    num
+        The number of keys to split into.
+    """
+
+    def __init__(self, key: Key[Array, ''], num: int = 2):
+        self._keys = random.split(key, num)
+
+    def __len__(self):
+        return self._keys.size
+
+    def pop(self, shape: int | tuple[int, ...] | None = None) -> Key[Array, '*']:
+        """
+        Pop one or more keys from the list.
+
+        Parameters
+        ----------
+        shape
+            The shape of the keys to pop. If `None`, a single key is popped.
+            If an integer, that many keys are popped. If a tuple, the keys are
+            reshaped to that shape.
+
+        Returns
+        -------
+        The popped keys as a jax array with the requested shape.
+
+        Raises
+        ------
+        IndexError
+            If `shape` is larger than the number of keys left in the list.
+
+        Notes
+        -----
+        The keys are popped from the beginning of the list, so for example
+        ``list(keys.pop(2))`` is equivalent to ``[keys.pop(), keys.pop()]``.
+        """
+        if shape is None:
+            shape = ()
+        elif not isinstance(shape, tuple):
+            shape = (shape,)
+        size_to_pop = math.prod(shape)
+        if size_to_pop > self._keys.size:
+            msg = f'Cannot pop {size_to_pop} keys from {self._keys.size} keys'
+            raise IndexError(msg)
+        popped_keys = self._keys[:size_to_pop]
+        self._keys = self._keys[size_to_pop:]
+        return popped_keys.reshape(shape)
+
+
+def truncated_normal_onesided(
+    key: Key[Array, ''],
+    shape: Sequence[int],
+    upper: Bool[Array, '*'],
+    bound: Float32[Array, '*'],
+) -> Float32[Array, '*']:
+    """
+    Sample from a one-sided truncated standard normal distribution.
+
+    Parameters
+    ----------
+    key
+        JAX random key.
+    shape
+        Shape of output array, broadcasted with other inputs.
+    upper
+        True for (-∞, bound], False for [bound, ∞).
+    bound
+        The truncation boundary.
+
+    Returns
+    -------
+    Array of samples from the truncated normal distribution.
+    """
+    # Pseudocode:
+    # | if upper:
+    # |     if bound < 0:
+    # |         ndtri(uniform(0, ndtr(bound))) =
+    # |         ndtri(ndtr(bound) * u)
+    # |     if bound > 0:
+    # |         -ndtri(uniform(ndtr(-bound), 1)) =
+    # |         -ndtri(ndtr(-bound) + ndtr(bound) * (1 - u))
+    # | if not upper:
+    # |     if bound < 0:
+    # |         ndtri(uniform(ndtr(bound), 1)) =
+    # |         ndtri(ndtr(bound) + ndtr(-bound) * (1 - u))
+    # |     if bound > 0:
+    # |         -ndtri(uniform(0, ndtr(-bound))) =
+    # |         -ndtri(ndtr(-bound) * u)
+    shape = jnp.broadcast_shapes(shape, upper.shape, bound.shape)
+    bound_pos = bound > 0
+    ndtr_bound = ndtr(bound)
+    ndtr_neg_bound = ndtr(-bound)
+    scale = jnp.where(upper, ndtr_bound, ndtr_neg_bound)
+    shift = jnp.where(upper, ndtr_neg_bound, ndtr_bound)
+    u = random.uniform(key, shape)
+    left_u = scale * (1 - u)  # ~ uniform in (0, ndtr(±bound)]
+    right_u = shift + scale * u  # ~ uniform in [ndtr(∓bound), 1)
+    truncated_u = jnp.where(upper ^ bound_pos, left_u, right_u)
+    truncated_norm = ndtri(truncated_u)
+    return jnp.where(bound_pos, -truncated_norm, truncated_norm)