bartz 0.4.0__tar.gz → 0.5.0__tar.gz

bartz-0.5.0/PKG-INFO

@@ -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/README.md

@@ -0,0 +1,31 @@
+[![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.4.0 → bartz-0.5.0}/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,51 +23,58 @@
 # SOFTWARE.
 
 [build-system]
-requires = ["poetry-core"]
-build-backend = "poetry.core.masonry.api"
+requires = ["uv_build>=0.7.3,<0.8.0"]
+build-backend = "uv_build"
 
-[tool.poetry]
+[project]
 name = "bartz"
-version = "0.4.0"
-description = "A JAX implementation of BART"
-authors = ["Giacomo Petrillo <info@giacomopetrillo.com>"]
+version = "0.5.0"
+description = "Super-fast BART (Bayesian Additive Regression Trees) in Python"
+authors = [
+    {name = "Giacomo Petrillo", email = "info@giacomopetrillo.com"},
+]
 license = "MIT"
 readme = "README.md"
-repository = "https://github.com/Gattocrucco/bartz"
+requires-python = ">=3.10"
 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",
+]
+
+[project.urls]
+Homepage = "https://github.com/Gattocrucco/bartz"
+Documentation = "https://gattocrucco.github.io/bartz/docs-dev"
+Issues = "https://github.com/Gattocrucco/bartz/issues"
 
-[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"
-matplotlib = "^3.8.3"
-appnope = "^0.1.4"
-tomli = "^2.0.1"
-packaging = "^24.0"
-xgboost = "^2.0.3"
-pre-commit = "^3.7.0"
-polars = "^0.20.19"
-scikit-learn = "^1.4.2"
-matplotlib-label-lines = "^0.7.0"
-
-[tool.poetry.group.test.dependencies]
-coverage = "^7.4.3"
-pytest = "^8.1.1"
-
-[tool.poetry.group.docs.dependencies]
-Sphinx = "^7.2.6"
-numpydoc = "^1.6.0"
-myst-parser = "^2.0.0"
+[dependency-groups]
+only-local = [
+    "appnope>=0.1.4",
+    "ipython>=8.36.0",
+    "matplotlib>=3.10.3",
+    "matplotlib-label-lines>=0.8.1",
+    "polars[pandas,pyarrow]>=1.29.0",
+    "pre-commit>=4.2.0",
+    "ruff>=0.11.9",
+    "scikit-learn>=1.6.1",
+    "tomli>=2.2.1",
+    "virtualenv>=20.31.2",
+    "xgboost>=3.0.0",
+]
+ci = [
+    "asv>=0.6.4",
+    "coverage>=7.8.0",
+    "myst-parser>=4.0.1",
+    "numpydoc>=1.8.0",
+    "packaging>=25.0",
+    "pytest>=8.3.5",
+    "pytest-timeout>=2.4.0",
+    "sphinx>=8.1.3",
+]
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]
@@ -79,6 +86,8 @@ addopts = [
     "--pdbcls=IPython.terminal.debugger:TerminalPdb",
     "--durations=3",
 ]
+timeout = 32
+timeout_method = "thread"  # when jax hangs, signals do not work
 
 # I wanted to use `--import-mode=importlib`, but it breaks importing submodules,
 # in particular `from . import util`.
@@ -117,3 +126,24 @@ local = [
     '/opt/hostedtoolcache/Python/*/*/lib/python*/site-packages/bartz/',
     'C:\hostedtoolcache\windows\Python\*\*\Lib\site-packages\bartz\',
 ]
+
+[tool.ruff]
+exclude = [".asv", "*.ipynb"]
+
+[tool.ruff.format]
+quote-style = "single"
+
+[tool.ruff.lint]
+select = [
+    "B", # bugbear: grab bag of additional stuff
+    "UP", # pyupgrade: fix some outdated idioms
+    "I", # isort: sort and reformat import statements
+    "F", # flake8
+]
+ignore = [
+    "B028", # warn with stacklevel = 2
+]
+
+[tool.uv]
+python-downloads = "never"
+python-preference = "only-system"

{bartz-0.4.0 → bartz-0.5.0}/src/bartz/BART.py

@@ -1,6 +1,6 @@
 # bartz/src/bartz/BART.py
 #
-# Copyright (c) 2024, Giacomo Petrillo
+# Copyright (c) 2024-2025, Giacomo Petrillo
 #
 # This file is part of bartz.
 #
@@ -27,11 +27,8 @@ import functools
 import jax
 import jax.numpy as jnp
 
-from . import jaxext
-from . import grove
-from . import mcmcstep
-from . import mcmcloop
-from . import prepcovars
+from . import grove, jaxext, mcmcloop, mcmcstep, prepcovars
+
 
 class gbart:
     """
@@ -53,10 +50,11 @@ class gbart:
         Whether to use predictors quantiles instead of a uniform grid to bin
         predictors.
     sigest : float, optional
-        An estimate of the residual standard deviation on `y_train`, used to
-        set `lamda`. If not specified, it is estimated by linear regression.
-        If `y_train` has less than two elements, it is set to 1. If n <= p, it
-        is set to the variance of `y_train`. Ignored if `lamda` is specified.
+        An estimate of the residual standard deviation on `y_train`, used to set
+        `lamda`. If not specified, it is estimated by linear regression (with
+        intercept, and without taking into account `w`). If `y_train` has less
+        than two elements, it is set to 1. If n <= p, it is set to the standard
+        deviation of `y_train`. Ignored if `lamda` is specified.
     sigdf : int, default 3
         The degrees of freedom of the scaled inverse-chisquared prior on the
         noise variance.
@@ -82,6 +80,12 @@ class gbart:
     offset : float, optional
         The prior mean of the latent mean function. If not specified, it is set
         to the mean of `y_train`. If `y_train` is empty, it is set to 0.
+    w : array (n,), optional
+        Coefficients that rescale the error standard deviation on each
+        datapoint. Not specifying `w` is equivalent to setting it to 1 for all
+        datapoints. Note: `w` is ignored in the automatic determination of
+        `sigest`, so either the weights should be O(1), or `sigest` should be
+        specified by the user.
     ntree : int, default 200
         The number of trees used to represent the latent mean function.
     numcut : int, default 255
@@ -108,6 +112,8 @@ class gbart:
         The number of iterations (including skipped ones) between each log.
     seed : int or jax random key, default 0
         The seed for the random number generator.
+    initkw : dict
+        Additional arguments passed to `mcmcstep.init`.
 
     Attributes
     ----------
@@ -135,8 +141,6 @@ class gbart:
         The number of trees.
     maxdepth : int
         The maximum depth of the trees.
-    initkw : dict
-        Additional arguments passed to `mcmcstep.init`.
 
     Methods
     -------
@@ -158,10 +162,13 @@ class gbart:
     - A lot of functionality is missing (variable selection, discrete response).
     - There are some additional attributes, and some missing.
 
-    The linear regression used to set `sigest` adds an intercept.
     """
 
-    def __init__(self, x_train, y_train, *,
+    def __init__(
+        self,
+        x_train,
+        y_train,
+        *,
         x_test=None,
         usequants=False,
         sigest=None,
@@ -173,6 +180,7 @@ class gbart:
         maxdepth=6,
         lamda=None,
         offset=None,
+        w=None,
         ntree=200,
         numcut=255,
         ndpost=1000,
@@ -180,26 +188,41 @@ class gbart:
         keepevery=1,
         printevery=100,
         seed=0,
-        initkw={},
-        ):
-
+        initkw=None,
+    ):
         x_train, x_train_fmt = self._process_predictor_input(x_train)
-
-        y_train, y_train_fmt = self._process_response_input(y_train)
+        y_train, _ = self._process_response_input(y_train)
         self._check_same_length(x_train, y_train)
+        if w is not None:
+            w, _ = self._process_response_input(w)
+            self._check_same_length(x_train, w)
 
         offset = self._process_offset_settings(y_train, offset)
         scale = self._process_scale_settings(y_train, k)
-        lamda, sigest = self._process_noise_variance_settings(x_train, y_train, sigest, sigdf, sigquant, lamda, offset)
+        lamda, sigest = self._process_noise_variance_settings(
+            x_train, y_train, sigest, sigdf, sigquant, lamda, offset
+        )
 
         splits, max_split = self._determine_splits(x_train, usequants, numcut)
         x_train = self._bin_predictors(x_train, splits)
-
-        y_train = self._transform_input(y_train, offset, scale)
-        lamda_scaled = lamda / (scale * scale)
-
-        mcmc_state = self._setup_mcmc(x_train, y_train, max_split, lamda_scaled, sigdf, power, base, maxdepth, ntree, initkw)
-        final_state, burnin_trace, main_trace = self._run_mcmc(mcmc_state, ndpost, nskip, keepevery, printevery, seed)
+        y_train, lamda_scaled = self._transform_input(y_train, lamda, offset, scale)
+
+        mcmc_state = self._setup_mcmc(
+            x_train,
+            y_train,
+            w,
+            max_split,
+            lamda_scaled,
+            sigdf,
+            power,
+            base,
+            maxdepth,
+            ntree,
+            initkw,
+        )
+        final_state, burnin_trace, main_trace = self._run_mcmc(
+            mcmc_state, ndpost, nskip, keepevery, printevery, seed
+        )
 
         sigma = self._extract_sigma(main_trace, scale)
         first_sigma = self._extract_sigma(burnin_trace, scale)
@@ -239,7 +262,7 @@ class gbart:
 
         Parameters
         ----------
-        x_test : array (m, p) or DataFrame
+        x_test : array (p, m) or DataFrame
             The test predictors.
 
         Returns
@@ -285,7 +308,9 @@ class gbart:
         assert get_length(x1) == get_length(x2)
 
     @staticmethod
-    def _process_noise_variance_settings(x_train, y_train, sigest, sigdf, sigquant, lamda, offset):
+    def _process_noise_variance_settings(
+        x_train, y_train, sigest, sigdf, sigquant, lamda, offset
+    ):
         if lamda is not None:
             return lamda, None
         else:
@@ -298,7 +323,7 @@ class gbart:
             else:
                 x_centered = x_train.T - x_train.mean(axis=1)
                 y_centered = y_train - y_train.mean()
-                    # centering is equivalent to adding an intercept column
+                # centering is equivalent to adding an intercept column
                 _, chisq, rank, _ = jnp.linalg.lstsq(x_centered, y_centered)
                 chisq = chisq.squeeze(0)
                 dof = len(y_train) - rank
@@ -336,11 +361,25 @@ class gbart:
         return prepcovars.bin_predictors(x, splits)
 
     @staticmethod
-    def _transform_input(y, offset, scale):
-        return (y - offset) / scale
+    def _transform_input(y, lamda, offset, scale):
+        y = (y - offset) / scale
+        lamda = lamda / (scale * scale)
+        return y, lamda
 
     @staticmethod
-    def _setup_mcmc(x_train, y_train, max_split, lamda, sigdf, power, base, maxdepth, ntree, initkw):
+    def _setup_mcmc(
+        x_train,
+        y_train,
+        w,
+        max_split,
+        lamda,
+        sigdf,
+        power,
+        base,
+        maxdepth,
+        ntree,
+        initkw,
+    ):
         depth = jnp.arange(maxdepth - 1)
         p_nonterminal = base / (1 + depth).astype(float) ** power
         sigma2_alpha = sigdf / 2
@@ -348,6 +387,7 @@ class gbart:
         kw = dict(
             X=x_train,
             y=y_train,
+            error_scale=w,
             max_split=max_split,
             num_trees=ntree,
             p_nonterminal=p_nonterminal,
@@ -355,17 +395,20 @@ class gbart:
             sigma2_beta=sigma2_beta,
             min_points_per_leaf=5,
         )
-        kw.update(initkw)
+        if initkw is not None:
+            kw.update(initkw)
         return mcmcstep.init(**kw)
 
     @staticmethod
     def _run_mcmc(mcmc_state, ndpost, nskip, keepevery, printevery, seed):
-        if isinstance(seed, jax.Array) and jnp.issubdtype(seed.dtype, jax.dtypes.prng_key):
+        if isinstance(seed, jax.Array) and jnp.issubdtype(
+            seed.dtype, jax.dtypes.prng_key
+        ):
             key = seed
         else:
             key = jax.random.key(seed)
         callback = mcmcloop.make_simple_print_callback(printevery)
-        return mcmcloop.run_mcmc(mcmc_state, nskip, ndpost, keepevery, callback, key)
+        return mcmcloop.run_mcmc(key, mcmc_state, nskip, ndpost, keepevery, callback)
 
     @staticmethod
     def _predict(trace, x):
@@ -379,9 +422,9 @@ class gbart:
     def _extract_sigma(trace, scale):
         return scale * jnp.sqrt(trace['sigma2'])
 
-
     def _show_tree(self, i_sample, i_tree, print_all=False):
         from . import debug
+
         trace = self._main_trace
         leaf_tree = trace['leaf_trees'][i_sample, i_tree]
         var_tree = trace['var_trees'][i_sample, i_tree]
@@ -396,7 +439,9 @@ class gbart:
         else:
             resid = bart['resid']
             alpha = bart['sigma2_alpha'] + resid.size / 2
-            norm2 = jnp.dot(resid, resid, preferred_element_type=bart['sigma2_beta'].dtype)
+            norm2 = jnp.dot(
+                resid, resid, preferred_element_type=bart['sigma2_beta'].dtype
+            )
             beta = bart['sigma2_beta'] + norm2 / 2
         sigma2 = beta / alpha
         return jnp.sqrt(sigma2) * self.scale
@@ -404,22 +449,32 @@ class gbart:
     def _compare_resid(self):
         bart = self._mcmc_state
         resid1 = bart['resid']
-        yhat = grove.evaluate_forest(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
 
     def _avg_acc(self):
         trace = self._main_trace
+
         def acc(prefix):
             acc = trace[f'{prefix}_acc_count']
             prop = trace[f'{prefix}_prop_count']
             return acc.sum() / prop.sum()
+
         return acc('grow'), acc('prune')
 
     def _avg_prop(self):
         trace = self._main_trace
+
         def prop(prefix):
             return trace[f'{prefix}_prop_count'].sum()
+
         pgrow = prop('grow')
         pprune = prop('prune')
         total = pgrow + pprune
@@ -432,16 +487,21 @@ class gbart:
 
     def _depth_distr(self):
         from . import debug
+
         trace = self._main_trace
         split_trees = trace['split_trees']
         return debug.trace_depth_distr(split_trees)
 
     def _points_per_leaf_distr(self):
         from . import debug
-        return debug.trace_points_per_leaf_distr(self._main_trace, self._mcmc_state['X'])
+
+        return debug.trace_points_per_leaf_distr(
+            self._main_trace, self._mcmc_state['X']
+        )
 
     def _check_trees(self):
         from . import debug
+
         return debug.check_trace(self._main_trace, self._mcmc_state)
 
     def _tree_goes_bad(self):

{bartz-0.4.0 → bartz-0.5.0}/src/bartz/__init__.py

@@ -1,6 +1,6 @@
 # bartz/src/bartz/__init__.py
 #
-# Copyright (c) 2024, Giacomo Petrillo
+# Copyright (c) 2024-2025, Giacomo Petrillo
 #
 # This file is part of bartz.
 #
@@ -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,18 +23,10 @@
 # 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
 """
 
-from ._version import __version__
-
-from . import BART
-
-from . import debug
-from . import grove
-from . import mcmcstep
-from . import mcmcloop
-from . import prepcovars
-from . import jaxext
+from . import BART, debug, grove, jaxext, mcmcloop, mcmcstep, prepcovars  # noqa: F401
+from ._version import __version__  # noqa: F401

bartz-0.5.0/src/bartz/_version.py

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