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

bartz/prepcovars.py

@@ -1,6 +1,6 @@
 # bartz/src/bartz/prepcovars.py
 #
-# Copyright (c) 2024, Giacomo Petrillo
+# Copyright (c) 2024-2025, Giacomo Petrillo
 #
 # This file is part of bartz.
 #
@@ -28,7 +28,7 @@ import jax
 from jax import numpy as jnp
 
 from . import jaxext
-from . import grove
+
 
 @functools.partial(jax.jit, static_argnums=(1,))
 def quantilized_splits_from_matrix(X, max_bins):
@@ -54,12 +54,15 @@ def quantilized_splits_from_matrix(X, max_bins):
         The number of actually used values in each row of `splits`.
     """
     out_length = min(max_bins, X.shape[1]) - 1
+
     # return _quantilized_splits_from_matrix(X, out_length)
-    @functools.partial(jaxext.autobatch, max_io_nbytes=2 ** 29)
+    @functools.partial(jaxext.autobatch, max_io_nbytes=2**29)
     def quantilize(X):
         return _quantilized_splits_from_matrix(X, out_length)
+
     return quantilize(X)
 
+
 @functools.partial(jax.vmap, in_axes=(0, None))
 def _quantilized_splits_from_matrix(x, out_length):
     huge = jaxext.huge_value(x)
@@ -67,21 +70,28 @@ def _quantilized_splits_from_matrix(x, out_length):
     actual_length -= 1
     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))
+        indices = jnp.arange(
+            midpoints.size, dtype=jaxext.minimal_unsigned_dtype(midpoints.size - 1)
+        )
         midpoints = jnp.where(indices < actual_length, midpoints, huge)
     else:
         midpoints = (u[1:] + u[:-1]) / 2
     indices = jnp.linspace(-1, actual_length, out_length + 2)[1:-1]
-    indices = jnp.around(indices).astype(jaxext.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
+    indices = jnp.around(indices).astype(
+        jaxext.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]
     truncated_midpoints = midpoints[:out_length]
-    splits = jnp.where(actual_length > out_length, decimated_midpoints, truncated_midpoints)
+    splits = jnp.where(
+        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))
     return splits, max_split
 
+
 @functools.partial(jax.jit, static_argnums=(1,))
 def uniform_splits_from_matrix(X, num_bins):
     """
@@ -110,6 +120,7 @@ def uniform_splits_from_matrix(X, num_bins):
     max_split = jnp.full(*splits.shape, jaxext.minimal_unsigned_dtype(num_bins - 1))
     return splits, max_split
 
+
 @functools.partial(jax.jit, static_argnames=('method',))
 def bin_predictors(X, splits, **kw):
     """
@@ -135,9 +146,11 @@ def bin_predictors(X, splits, **kw):
         A matrix with `p` predictors and `n` observations, where each predictor
         has been replaced by the index of the bin it falls into.
     """
-    @functools.partial(jaxext.autobatch, max_io_nbytes=2 ** 29)
+
+    @functools.partial(jaxext.autobatch, max_io_nbytes=2**29)
     @jax.vmap
     def bin_predictors(x, splits):
         dtype = jaxext.minimal_unsigned_dtype(splits.size)
         return jnp.searchsorted(splits, x, **kw).astype(dtype)
+
     return bin_predictors(X, splits)

bartz-0.5.0.dist-info/METADATA

@@ -0,0 +1,48 @@
+Metadata-Version: 2.4
+Name: bartz
+Version: 0.5.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: jax>=0.4.35,<1
+Requires-Dist: jaxlib>=0.4.35,<1
+Requires-Dist: numpy>=1.25.2,<3
+Requires-Dist: scipy>=1.11.4,<2
+Requires-Python: >=3.10
+Project-URL: Documentation, https://gattocrucco.github.io/bartz/docs-dev
+Project-URL: Homepage, https://github.com/Gattocrucco/bartz
+Project-URL: Issues, https://github.com/Gattocrucco/bartz/issues
+Description-Content-Type: text/markdown
+
+[![PyPI](https://img.shields.io/pypi/v/bartz)](https://pypi.org/project/bartz/)
+[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.13931477.svg)](https://doi.org/10.5281/zenodo.13931477)
+
+# BART vectoriZed
+
+An implementation of Bayesian Additive Regression Trees (BART) in JAX.
+
+If you don't know what BART is, but know XGBoost, consider BART as a sort of Bayesian XGBoost. bartz makes BART run as fast as XGBoost.
+
+BART is a nonparametric Bayesian regression technique. Given training predictors $X$ and responses $y$, BART finds a function to predict $y$ given $X$. The result of the inference is a sample of possible functions, representing the uncertainty over the determination of the function.
+
+This Python module provides an implementation of BART that runs on GPU, to process large datasets faster. It is also good on CPU. Most other implementations of BART are for R, and run on CPU only.
+
+On CPU, bartz runs at the speed of dbarts (the fastest implementation I know of) if n > 20,000, but using 1/20 of the memory. On GPU, the speed premium depends on sample size; it is convenient over CPU only for n > 10,000. The maximum speedup is currently 200x, on an Nvidia A100 and with at least 2,000,000 observations.
+
+[This Colab notebook](https://colab.research.google.com/github/Gattocrucco/bartz/blob/main/docs/examples/basic_simdata.ipynb) runs bartz with n = 100,000 observations, p = 1000 predictors, 10,000 trees, for 1000 MCMC iterations, in 5 minutes.
+
+## Links
+
+- [Documentation (latest release)](https://gattocrucco.github.io/bartz/docs)
+- [Documentation (development version)](https://gattocrucco.github.io/bartz/docs-dev)
+- [Repository](https://github.com/Gattocrucco/bartz)
+- [Code coverage](https://gattocrucco.github.io/bartz/coverage)
+- [Benchmarks](https://gattocrucco.github.io/bartz/benchmarks)
+- [List of BART packages](https://gattocrucco.github.io/bartz/docs-dev/pkglist.html)
+
+## Citing bartz
+
+Article: Petrillo (2024), "Very fast Bayesian Additive Regression Trees on GPU", [arXiv:2410.23244](https://arxiv.org/abs/2410.23244).
+
+To cite the software directly, including the specific version, use [zenodo](https://doi.org/10.5281/zenodo.13931477).

bartz-0.5.0.dist-info/RECORD

@@ -0,0 +1,13 @@
+bartz/.DS_Store,sha256=7191af46d7b8c0d4c03c502f94eb01353bc2e615d75c45b3af0e31ab238034b5,6148
+bartz/BART.py,sha256=50faa27d82bc1c5ab6eb676a006bd7a95e2cbe10a541c5de7a100f068a76a820,17898
+bartz/__init__.py,sha256=c3b1d04b62e52220f060d28f38403a536fcbe2cea06d0a4087dc0143b05cb1fd,1433
+bartz/_version.py,sha256=d0f688d9e48e0a9e6490d7292a95126cb1dfebaacbf714331696322c6a44b723,22
+bartz/debug.py,sha256=0ab2b06e53245b763d4c684653031c912ed794e328d2f427dbfb017c975840af,5359
+bartz/grove.py,sha256=13d05ca6877b073cfc2d0e887fc03f74be35eadb6b587b83a0a0d41de1accb10,8494
+bartz/jaxext.py,sha256=a439363fd44be9a26770a65dd30c3d140a877ec0ddcdd65ba8fed4471c629cd0,11552
+bartz/mcmcloop.py,sha256=15f7de40d1b3a57183f90db2972b66ab7b455f4829cca22c5686736b154c1fac,9216
+bartz/mcmcstep.py,sha256=48aaf9d4a32aa4b84a73b4ec365fe712204dd872e5df3de5cc3715a37e362430,60504
+bartz/prepcovars.py,sha256=e76c05c5218a47096c368f28c2df0178b44b428d26ba9c15fa6dd6c09fd2f2a1,5849
+bartz-0.5.0.dist-info/WHEEL,sha256=11319513b72fe7d4379f6f9f7984b3cca85f688201c09325839a04908e107b75,78
+bartz-0.5.0.dist-info/METADATA,sha256=d875077205f7d03c4c4f8c84acfe3079c4808be32ea3ab358640fe581b6f3b46,2795
+bartz-0.5.0.dist-info/RECORD,,

bartz-0.5.0.dist-info/WHEEL

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

bartz-0.4.0.dist-info/LICENSE

@@ -1,21 +0,0 @@
-MIT License
-
-Copyright (c) 2024 Giacomo Petrillo
-
-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.

bartz-0.4.0.dist-info/METADATA

@@ -1,77 +0,0 @@
-Metadata-Version: 2.1
-Name: bartz
-Version: 0.4.0
-Summary: A JAX implementation of BART
-Home-page: https://github.com/Gattocrucco/bartz
-License: MIT
-Author: Giacomo Petrillo
-Author-email: info@giacomopetrillo.com
-Requires-Python: >=3.10,<4.0
-Classifier: License :: OSI Approved :: MIT License
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
-Requires-Dist: jax (>=0.4.23,<0.5.0)
-Requires-Dist: jaxlib (>=0.4.23,<0.5.0)
-Requires-Dist: numpy (>=1.25.2,<2.0.0)
-Requires-Dist: scipy (>=1.11.4,<2.0.0)
-Project-URL: Bug Tracker, https://github.com/Gattocrucco/bartz/issues
-Project-URL: Repository, https://github.com/Gattocrucco/bartz
-Description-Content-Type: text/markdown
-
-[![PyPI](https://img.shields.io/pypi/v/bartz)](https://pypi.org/project/bartz/)
-
-# BART vectoriZed
-
-A branchless vectorized implementation of Bayesian Additive Regression Trees (BART) in JAX.
-
-BART is a nonparametric Bayesian regression technique. Given predictors $X$ and responses $y$, BART finds a function to predict $y$ given $X$. The result of the inference is a sample of possible functions, representing the uncertainty over the determination of the function.
-
-This Python module provides an implementation of BART that runs on GPU, to process large datasets faster. It is also good on CPU. Most other implementations of BART are for R, and run on CPU only.
-
-On CPU, bartz runs at the speed of dbarts (the fastest implementation I know of), but using half the memory. On GPU, the speed premium depends on sample size; with 50000 datapoints and 5000 trees, on an Nvidia Tesla V100 GPU it's 12 times faster than an Apple M1 CPU, and this factor is linearly proportional to the number of datapoints.
-
-## Links
-
-- [Documentation (latest release)](https://gattocrucco.github.io/bartz/docs)
-- [Documentation (development version)](https://gattocrucco.github.io/bartz/docs-dev)
-- [Repository](https://github.com/Gattocrucco/bartz)
-- [Code coverage](https://gattocrucco.github.io/bartz/coverage)
-
-## Other BART packages
-
-- [stochtree](https://github.com/StochasticTree) C++ library with R and Python bindings taylored to researchers who want to make their own BART variants
-- [bnptools](https://github.com/rsparapa/bnptools) Feature-rich R packages for BART and some variants
-- [dbarts](https://github.com/vdorie/dbarts) Fast R package
-- [bartMachine](https://github.com/kapelner/bartMachine) Fast R package, supports missing predictors imputation
-- [SoftBART](https://github.com/theodds/SoftBART) R package with a smooth version of BART
-- [bcf](https://github.com/jaredsmurray/bcf) R package for a version of BART for causal inference
-- [flexBART](https://github.com/skdeshpande91/flexBART) Fast R package, supports categorical predictors
-- [flexBCF](https://github.com/skdeshpande91/flexBCF) R package, version of bcf optimized for large datasets
-- [XBART](https://github.com/JingyuHe/XBART) R/Python package, XBART is a faster variant of BART
-- [BART](https://github.com/JingyuHe/BART) R package, BART warm-started with XBART
-- [XBCF](https://github.com/socket778/XBCF)
-- [BayesTree](https://cran.r-project.org/package=BayesTree) R package, original BART implementation
-- [bartCause](https://github.com/vdorie/bartCause) R package, pre-made BART-based workflows for causal inference
-- [stan4bart](https://github.com/vdorie/stan4bart)
-- [VCBART](https://github.com/skdeshpande91/VCBART)
-- [monbart](https://github.com/jaredsmurray/monbart)
-- [mBART](https://github.com/remcc/mBART_shlib)
-- [SequentialBART](https://github.com/mjdaniels/SequentialBART)
-- [sparseBART](https://github.com/cspanbauer/sparseBART)
-- [pymc-bart](https://github.com/pymc-devs/pymc-bart)
-- [semibart](https://github.com/zeldow/semibart)
-- [CSP-BART](https://github.com/ebprado/CSP-BART)
-- [AMBARTI](https://github.com/ebprado/AMBARTI)
-- [MOTR-BART](https://github.com/ebprado/MOTR-BART)
-- [bcfbma](https://github.com/EoghanONeill/bcfbma)
-- [bartBMAnew](https://github.com/EoghanONeill/bartBMAnew)
-- [BART-BMA](https://github.com/BelindaHernandez/BART-BMA) (superseded by bartBMAnew)
-- [gpbart](https://github.com/MateusMaiaDS/gpbart)
-- [GPBART](https://github.com/nchenderson/GPBART)
-- [bartpy](https://github.com/JakeColtman/bartpy)
-- [BayesTreePrior](https://github.com/AlexiaJM/BayesTreePrior)
-- [BayesTree.jl](https://github.com/mathcg/BayesTree.jl)
-- [longbet](https://github.com/google/longbet)
-

bartz-0.4.0.dist-info/RECORD

@@ -1,13 +0,0 @@
-bartz/BART.py,sha256=CbGzFWtYw5u38Z9-Hy3CbDXpKOOvPFAAkSqu2HZl8no,16862
-bartz/__init__.py,sha256=E96vsP0bZ8brejpZmEmRoXuMsUdinO_B_SKUUl1rLsg,1448
-bartz/_version.py,sha256=2eiWQI55fd-roDdkt4Hvl9WzrTJ4xQo33VzFud6D03U,22
-bartz/debug.py,sha256=9ZH-JfwZVu5OPhHBEyXQHAU5H9KIu1vxLK7yNv4m4Ew,5314
-bartz/grove.py,sha256=x_6NK_l-hrXfy1PhssYNJkX41-w_WqjDziww0E7YRS8,8500
-bartz/jaxext.py,sha256=RcVWTCGS8lXF7GBsNbKrpuA4MTcokItq0CpWm3s7CGk,12033
-bartz/mcmcloop.py,sha256=lKDszvniNXka99X3e9RCrTgvEAZHA7ZbVXEgxUYvKMY,7634
-bartz/mcmcstep.py,sha256=diI9vHXHMvu_Lk_bSJ-a038OnEbXDpNEikVPhRcxEys,54987
-bartz/prepcovars.py,sha256=mMgfL-LGJ_8QpOL6iy7yfkL8A7FrT7Zfn5M3voyNwSQ,5818
-bartz-0.4.0.dist-info/LICENSE,sha256=heuIJZQK9IexJYC-fYHoLUrgj8HG8yS3G072EvKh-94,1073
-bartz-0.4.0.dist-info/METADATA,sha256=K86CVXT6ayPnc2hjhreYGMEeYWfYJIZdDkKuBB0-FYA,4500
-bartz-0.4.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-bartz-0.4.0.dist-info/RECORD,,

bartz-0.4.0.dist-info/WHEEL

@@ -1,4 +0,0 @@
-Wheel-Version: 1.0
-Generator: poetry-core 1.9.0
-Root-Is-Purelib: true
-Tag: py3-none-any