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

bartz/prepcovars.py

@@ -24,40 +24,82 @@
 
 """Functions to preprocess data."""
 
-import functools
+from functools import partial
 
-import jax
+from jax import jit, vmap
 from jax import numpy as jnp
+from jaxtyping import Array, Float, Integer, Real, UInt
 
-from . import jaxext
+from bartz.jaxext import autobatch, minimal_unsigned_dtype, unique
 
 
-@functools.partial(jax.jit, static_argnums=(1,))
-def quantilized_splits_from_matrix(X, max_bins):
+def parse_xinfo(
+    xinfo: Float[Array, 'p m'],
+) -> tuple[Float[Array, 'p m'], UInt[Array, ' p']]:
+    """Parse pre-defined splits in the format of the R package BART.
+
+    Parameters
+    ----------
+    xinfo
+        A matrix with the cutpoins to use to bin each predictor. Each row shall
+        contain a sorted list of cutpoints for a predictor. If there are less
+        cutpoints than the number of columns in the matrix, fill the remaining
+        cells with NaN.
+
+        `xinfo` shall be a matrix even if `x_train` is a dataframe.
+
+    Returns
+    -------
+    splits : Float[Array, 'p m']
+        `xinfo` modified by replacing nan with a large value.
+    max_split : UInt[Array, 'p']
+        The number of non-nan elements in each row of `xinfo`.
+    """
+    is_not_nan = ~jnp.isnan(xinfo)
+    max_split = jnp.sum(is_not_nan, axis=1)
+    max_split = max_split.astype(minimal_unsigned_dtype(xinfo.shape[1]))
+    huge = _huge_value(xinfo)
+    splits = jnp.where(is_not_nan, xinfo, huge)
+    return splits, max_split
+
+
+@partial(jit, static_argnums=(1,))
+def quantilized_splits_from_matrix(
+    X: Real[Array, 'p n'], max_bins: int
+) -> tuple[Real[Array, 'p m'], UInt[Array, ' p']]:
     """
     Determine bins that make the distribution of each predictor uniform.
 
     Parameters
     ----------
-    X : array (p, n)
+    X
         A matrix with `p` predictors and `n` observations.
-    max_bins : int
+    max_bins
         The maximum number of bins to produce.
 
     Returns
     -------
-    splits : array (p, m)
+    splits : Real[Array, 'p m']
         A matrix containing, for each predictor, the boundaries between bins.
         `m` is ``min(max_bins, n) - 1``, which is an upper bound on the number
         of splits. Each predictor may have a different number of splits; unused
         values at the end of each row are filled with the maximum value
         representable in the type of `X`.
-    max_split : array (p,)
+    max_split : UInt[Array, ' p']
         The number of actually used values in each row of `splits`.
+
+    Raises
+    ------
+    ValueError
+        If `X` has no columns or if `max_bins` is less than 1.
     """
     out_length = min(max_bins, X.shape[1]) - 1
 
-    @functools.partial(jaxext.autobatch, max_io_nbytes=2**29)
+    if out_length < 0:
+        msg = f'{X.shape[1]=} and {max_bins=}, they should be both at least 1.'
+        raise ValueError(msg)
+
+    @partial(autobatch, max_io_nbytes=2**29)
     def quantilize(X):
         # wrap this function because autobatch needs traceable args
         return _quantilized_splits_from_matrix(X, out_length)
@@ -65,23 +107,28 @@ def quantilized_splits_from_matrix(X, max_bins):
     return quantilize(X)
 
 
-@functools.partial(jax.vmap, in_axes=(0, None))
-def _quantilized_splits_from_matrix(x, out_length):
-    huge = jaxext.huge_value(x)
-    u, actual_length = jaxext.unique(x, size=x.size, fill_value=huge)
-    actual_length -= 1
+@partial(vmap, in_axes=(0, None))
+def _quantilized_splits_from_matrix(
+    x: Real[Array, 'p n'], out_length: int
+) -> tuple[Real[Array, 'p m'], UInt[Array, ' p']]:
+    # find the sorted unique values in x
+    huge = _huge_value(x)
+    u, actual_length = unique(x, size=x.size, fill_value=huge)
+
+    # compute the midpoints between each unique value
     if jnp.issubdtype(x.dtype, jnp.integer):
-        midpoints = u[:-1] + jaxext.ensure_unsigned(u[1:] - u[:-1]) // 2
-        indices = jnp.arange(
-            midpoints.size, dtype=jaxext.minimal_unsigned_dtype(midpoints.size - 1)
-        )
-        midpoints = jnp.where(indices < actual_length, midpoints, huge)
+        midpoints = u[:-1] + _ensure_unsigned(u[1:] - u[:-1]) // 2
     else:
-        midpoints = (u[1:] + u[:-1]) / 2
+        midpoints = u[:-1] + (u[1:] - u[:-1]) / 2
+        # using x_i + (x_i+1 - x_i) / 2 instead of (x_i + x_i+1) / 2 is to
+        # avoid overflow
+    actual_length -= 1
+    if midpoints.size:
+        midpoints = midpoints.at[actual_length].set(huge)
+
+    # take a subset of the midpoints if there are more than the requested maximum
     indices = jnp.linspace(-1, actual_length, out_length + 2)[1:-1]
-    indices = jnp.around(indices).astype(
-        jaxext.minimal_unsigned_dtype(midpoints.size - 1)
-    )
+    indices = jnp.around(indices).astype(minimal_unsigned_dtype(midpoints.size - 1))
     # indices calculation with float rather than int to avoid potential
     # overflow with int32, and to round to nearest instead of rounding down
     decimated_midpoints = midpoints[indices]
@@ -90,41 +137,92 @@ def _quantilized_splits_from_matrix(x, out_length):
         actual_length > out_length, decimated_midpoints, truncated_midpoints
     )
     max_split = jnp.minimum(actual_length, out_length)
-    max_split = max_split.astype(jaxext.minimal_unsigned_dtype(out_length))
+    max_split = max_split.astype(minimal_unsigned_dtype(out_length))
     return splits, max_split
 
 
-@functools.partial(jax.jit, static_argnums=(1,))
-def uniform_splits_from_matrix(X, num_bins):
+def _huge_value(x: Array) -> int | float:
+    """
+    Return the maximum value that can be stored in `x`.
+
+    Parameters
+    ----------
+    x
+        A numerical numpy or jax array.
+
+    Returns
+    -------
+    The maximum value allowed by `x`'s type (finite for floats).
+    """
+    if jnp.issubdtype(x.dtype, jnp.integer):
+        return jnp.iinfo(x.dtype).max
+    else:
+        return float(jnp.finfo(x.dtype).max)
+
+
+def _ensure_unsigned(x: Integer[Array, '*shape']) -> UInt[Array, '*shape']:
+    """If x has signed integer type, cast it to the unsigned dtype of the same size."""
+    return x.astype(_signed_to_unsigned(x.dtype))
+
+
+def _signed_to_unsigned(int_dtype: jnp.dtype) -> jnp.dtype:
+    """
+    Map a signed integer type to its unsigned counterpart.
+
+    Unsigned types are passed through.
+    """
+    assert jnp.issubdtype(int_dtype, jnp.integer)
+    if jnp.issubdtype(int_dtype, jnp.unsignedinteger):
+        return int_dtype
+    match int_dtype:
+        case jnp.int8:
+            return jnp.uint8
+        case jnp.int16:
+            return jnp.uint16
+        case jnp.int32:
+            return jnp.uint32
+        case jnp.int64:
+            return jnp.uint64
+        case _:
+            msg = f'unexpected integer type {int_dtype}'
+            raise TypeError(msg)
+
+
+@partial(jit, static_argnums=(1,))
+def uniform_splits_from_matrix(
+    X: Real[Array, 'p n'], num_bins: int
+) -> tuple[Real[Array, 'p m'], UInt[Array, ' p']]:
     """
     Make an evenly spaced binning grid.
 
     Parameters
     ----------
-    X : array (p, n)
+    X
         A matrix with `p` predictors and `n` observations.
-    num_bins : int
+    num_bins
         The number of bins to produce.
 
     Returns
     -------
-    splits : array (p, num_bins - 1)
+    splits : Real[Array, 'p m']
         A matrix containing, for each predictor, the boundaries between bins.
         The excluded endpoints are the minimum and maximum value in each row of
         `X`.
-    max_split : array (p,)
+    max_split : UInt[Array, ' p']
         The number of cutpoints in each row of `splits`, i.e., ``num_bins - 1``.
     """
     low = jnp.min(X, axis=1)
     high = jnp.max(X, axis=1)
     splits = jnp.linspace(low, high, num_bins + 1, axis=1)[:, 1:-1]
     assert splits.shape == (X.shape[0], num_bins - 1)
-    max_split = jnp.full(*splits.shape, jaxext.minimal_unsigned_dtype(num_bins - 1))
+    max_split = jnp.full(*splits.shape, minimal_unsigned_dtype(num_bins - 1))
     return splits, max_split
 
 
-@functools.partial(jax.jit, static_argnames=('method',))
-def bin_predictors(X, splits, **kw):
+@partial(jit, static_argnames=('method',))
+def bin_predictors(
+    X: Real[Array, 'p n'], splits: Real[Array, 'p m'], **kw
+) -> UInt[Array, 'p n']:
     """
     Bin the predictors according to the given splits.
 
@@ -132,27 +230,25 @@ def bin_predictors(X, splits, **kw):
 
     Parameters
     ----------
-    X : array (p, n)
+    X
         A matrix with `p` predictors and `n` observations.
-    splits : array (p, m)
+    splits
         A matrix containing, for each predictor, the boundaries between bins.
         `m` is the maximum number of splits; each row may have shorter
         actual length, marked by padding unused locations at the end of the
         row with the maximum value allowed by the type.
-    **kw : dict
+    **kw
         Additional arguments are passed to `jax.numpy.searchsorted`.
 
     Returns
     -------
-    X_binned : int array (p, n)
-        A matrix with `p` predictors and `n` observations, where each predictor
-        has been replaced by the index of the bin it falls into.
+    `X` but with each value replaced by the index of the bin it falls into.
     """
 
-    @functools.partial(jaxext.autobatch, max_io_nbytes=2**29)
-    @jax.vmap
+    @partial(autobatch, max_io_nbytes=2**29)
+    @vmap
     def bin_predictors(x, splits):
-        dtype = jaxext.minimal_unsigned_dtype(splits.size)
+        dtype = minimal_unsigned_dtype(splits.size)
         return jnp.searchsorted(splits, x, **kw).astype(dtype)
 
     return bin_predictors(X, splits)

{bartz-0.6.0.dist-info → bartz-0.7.0.dist-info}/METADATA

@@ -1,13 +1,12 @@
 Metadata-Version: 2.4
 Name: bartz
-Version: 0.6.0
+Version: 0.7.0
 Summary: Super-fast BART (Bayesian Additive Regression Trees) in Python
 Author: Giacomo Petrillo
 Author-email: Giacomo Petrillo <info@giacomopetrillo.com>
 License-Expression: MIT
 Requires-Dist: equinox>=0.12.2
-Requires-Dist: jax>=0.4.35
-Requires-Dist: jaxlib>=0.4.35
+Requires-Dist: jax>=0.5.3
 Requires-Dist: jaxtyping>=0.3.2
 Requires-Dist: numpy>=1.25.2
 Requires-Dist: scipy>=1.11.4

bartz-0.7.0.dist-info/RECORD

@@ -0,0 +1,17 @@
+bartz/.DS_Store,sha256=7191af46d7b8c0d4c03c502f94eb01353bc2e615d75c45b3af0e31ab238034b5,6148
+bartz/BART.py,sha256=6b129e20a258d724e0cba5ffaf377b5c4d62e545c1f3a737602c2a0be5d84b96,29601
+bartz/__init__.py,sha256=98c579136a8755390210ada33e713290749e4d7fca58550c791f77f192f4b4a1,1436
+bartz/_version.py,sha256=d3d868979a2f2fa02b20b248259f8c8ac7273a329ae918da1139ad7602695b67,22
+bartz/debug.py,sha256=5082c2dd07f6d3f8353c57491c857cca4aea913647e24502cc375f9aadc84975,43736
+bartz/grove.py,sha256=f64505623feec7edcec96930909f9a9326f290976d7e9bdb7b9abcb32fe425fe,10559
+bartz/jaxext/__init__.py,sha256=6cd2e7c23ccc4f0399fb3f7989312a500fd8bd3f7f07eb85ce537f2c8873f35a,6705
+bartz/jaxext/_autobatch.py,sha256=b5dbaec52e39b4b32c824fde47eaaf33f496a1574cc336d3a79fa71f4c5e348a,7116
+bartz/jaxext/scipy/__init__.py,sha256=a1f5990a75c1c73908565be4cd5fa1c07278ad3a02b78b21e2f1225b388ab6b5,1227
+bartz/jaxext/scipy/special.py,sha256=f0e777c29a77d46d55ff2f0169b3d8581de24d8580b8d6f1c7c1473910fb62d9,8111
+bartz/jaxext/scipy/stats.py,sha256=703beb9fcfe606a195fe9a3143eceed367c87a31bd75763ecdb6e509c2e87f53,1483
+bartz/mcmcloop.py,sha256=e841c076bf5392b1ad8eab955c70a63d18ac321c249d255c7424b62e52c127c1,22134
+bartz/mcmcstep.py,sha256=3b2a4bdad3bbf836efc7cc6f52adac3ba23dce00a1cca4a99e5a8d79d10c3e68,84092
+bartz/prepcovars.py,sha256=50334621ca6ec7a6e35d21ce5ff8b96d4cbf0e696c9bcd10d86268e301a820a7,8728
+bartz-0.7.0.dist-info/WHEEL,sha256=607c46fee47e440c91332c738096ff0f5e54ca3b0818ee85462dd5172a38e793,79
+bartz-0.7.0.dist-info/METADATA,sha256=758e12296acf815c9bae50601ff502e6dee97c9ac0d817232e45e3dfb39665bc,2815
+bartz-0.7.0.dist-info/RECORD,,

{bartz-0.6.0.dist-info → bartz-0.7.0.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: uv 0.7.8
+Generator: uv 0.7.19
 Root-Is-Purelib: true
 Tag: py3-none-any