bartz 0.3.0__tar.gz → 0.4.1__tar.gz

bartz-0.4.1/PKG-INFO

@@ -0,0 +1,52 @@
+Metadata-Version: 2.3
+Name: bartz
+Version: 0.4.1
+Summary: Super-fast BART (Bayesian Additive Regression Trees) in Python
+License: MIT
+Author: Giacomo Petrillo
+Author-email: info@giacomopetrillo.com
+Requires-Python: >=3.10
+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
+Classifier: Programming Language :: Python :: 3.13
+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)
+Project-URL: Bug Tracker, 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)
+- [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.4.1/README.md

@@ -0,0 +1,30 @@
+[![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)
+- [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.3.0 → bartz-0.4.1}/pyproject.toml

@@ -1,6 +1,6 @@
 # bartz/pyproject.toml
 #
-# Copyright (c) 2024, Giacomo Petrillo
+# Copyright (c) 2024-2025, Giacomo Petrillo
 #
 # This file is part of bartz.
 #
@@ -23,48 +23,53 @@
 # SOFTWARE.
 
 [build-system]
-requires = ["poetry-core"]
+requires = ["poetry-core>=2.0.0,<3.0.0"]
 build-backend = "poetry.core.masonry.api"
 
-[tool.poetry]
+[project]
 name = "bartz"
-version = "0.3.0"
-description = "A JAX implementation of BART"
-authors = ["Giacomo Petrillo <info@giacomopetrillo.com>"]
+version = "0.4.1"
+description = "Super-fast BART (Bayesian Additive Regression Trees) in Python"
+authors = [
+    {name = "Giacomo Petrillo", email = "info@giacomopetrillo.com"},
+]
 license = "MIT"
 readme = "README.md"
+requires-python = ">=3.10"
 repository = "https://github.com/Gattocrucco/bartz"
 packages = [
     { include = "bartz", from = "src" },
 ]
+dependencies = [
+    "jax >=0.4.35,<1",
+    "jaxlib >=0.4.35,<1",
+    "numpy >=1.25.2,<3",
+    "scipy >=1.11.4,<2",
+]
 
 [tool.poetry.urls]
 "Bug Tracker" = "https://github.com/Gattocrucco/bartz/issues"
 
-[tool.poetry.dependencies]
-python = "^3.10"
-jax = "^0.4.23"
-jaxlib = "^0.4.23"
-numpy = "^1.25.2"
-scipy = "^1.11.4"
-
 [tool.poetry.group.dev.dependencies]
-ipython = "^8.22.2"
+ipython = ">=8.22.2"
 matplotlib = "^3.8.3"
 appnope = "^0.1.4"
 tomli = "^2.0.1"
 packaging = "^24.0"
 xgboost = "^2.0.3"
 pre-commit = "^3.7.0"
+scikit-learn = "^1.4.2"
+matplotlib-label-lines = "^0.7.0"
+polars = {version = "^1.0.0", extras = ["pandas", "pyarrow"]}
 
 [tool.poetry.group.test.dependencies]
 coverage = "^7.4.3"
 pytest = "^8.1.1"
 
 [tool.poetry.group.docs.dependencies]
-Sphinx = "^7.2.6"
+Sphinx = ">=7.2.6,<9"
 numpydoc = "^1.6.0"
-myst-parser = "^2.0.0"
+myst-parser = "^4.0.1"
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]

{bartz-0.3.0 → bartz-0.4.1}/src/bartz/__init__.py

@@ -10,10 +10,10 @@
 # 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
@@ -23,7 +23,7 @@
 # SOFTWARE.
 
 """
-A jax implementation of BART
+Super-fast BART (Bayesian Additive Regression Trees) in Python
 
 See the manual at https://gattocrucco.github.io/bartz/docs
 """

bartz-0.4.1/src/bartz/_version.py

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

{bartz-0.3.0 → bartz-0.4.1}/src/bartz/jaxext.py

@@ -1,6 +1,6 @@
 # bartz/src/bartz/jaxext.py
 #
-# Copyright (c) 2024, Giacomo Petrillo
+# Copyright (c) 2024-2025, Giacomo Petrillo
 #
 # This file is part of bartz.
 #
@@ -45,27 +45,6 @@ def castto(func, type):
         return func(*args, **kw).astype(type)
     return newfunc
 
-def pure_callback_ufunc(callback, dtype, *args, excluded=None, **kwargs):
-    """ version of `jax.pure_callback` that deals correctly with ufuncs,
-    see `<https://github.com/google/jax/issues/17187>`_ """
-    if excluded is None:
-        excluded = ()
-    shape = jnp.broadcast_shapes(*(
-        a.shape
-        for i, a in enumerate(args)
-        if i not in excluded
-    ))
-    ndim = len(shape)
-    padded_args = [
-        a if i in excluded
-        else jnp.expand_dims(a, tuple(range(ndim - a.ndim)))
-        for i, a in enumerate(args)
-    ]
-    result = jax.ShapeDtypeStruct(shape, dtype)
-    return jax.pure_callback(callback, result, *padded_args, vectorized=True, **kwargs)
-
-    # TODO when jax solves this, check version and piggyback on original if new
-
 class scipy:
 
     class special:
@@ -74,9 +53,11 @@ class scipy:
         def gammainccinv(a, y):
             a = jnp.asarray(a)
             y = jnp.asarray(y)
+            shape = jnp.broadcast_shapes(a.shape, y.shape)
             dtype = float_type(a.dtype, y.dtype)
+            dummy = jax.ShapeDtypeStruct(shape, dtype)
             ufunc = castto(special.gammainccinv, dtype)
-            return pure_callback_ufunc(ufunc, dtype, a, y)
+            return jax.pure_callback(ufunc, dummy, a, y, vmap_method='expand_dims')
 
     class stats: