bartz 0.5.0__tar.gz → 0.6.0__tar.gz

{bartz-0.5.0 → bartz-0.6.0}/PKG-INFO

@@ -1,14 +1,16 @@
 Metadata-Version: 2.4
 Name: bartz
-Version: 0.5.0
+Version: 0.6.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-Dist: equinox>=0.12.2
+Requires-Dist: jax>=0.4.35
+Requires-Dist: jaxlib>=0.4.35
+Requires-Dist: jaxtyping>=0.3.2
+Requires-Dist: numpy>=1.25.2
+Requires-Dist: scipy>=1.11.4
 Requires-Python: >=3.10
 Project-URL: Documentation, https://gattocrucco.github.io/bartz/docs-dev
 Project-URL: Homepage, https://github.com/Gattocrucco/bartz

{bartz-0.5.0 → bartz-0.6.0}/pyproject.toml

@@ -28,7 +28,7 @@ build-backend = "uv_build"
 
 [project]
 name = "bartz"
-version = "0.5.0"
+version = "0.6.0"
 description = "Super-fast BART (Bayesian Additive Regression Trees) in Python"
 authors = [
     {name = "Giacomo Petrillo", email = "info@giacomopetrillo.com"},
@@ -36,14 +36,13 @@ authors = [
 license = "MIT"
 readme = "README.md"
 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",
+    "equinox>=0.12.2",
+    "jax>=0.4.35",
+    "jaxlib>=0.4.35",
+    "jaxtyping>=0.3.2",
+    "numpy>=1.25.2",
+    "scipy>=1.11.4",
 ]
 
 [project.urls]
@@ -57,8 +56,8 @@ only-local = [
     "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",
+    "pydoclint>=0.6.6",
     "ruff>=0.11.9",
     "scikit-learn>=1.6.1",
     "tomli>=2.2.1",
@@ -71,12 +70,15 @@ ci = [
     "myst-parser>=4.0.1",
     "numpydoc>=1.8.0",
     "packaging>=25.0",
+    "polars[pandas,pyarrow]>=1.29.0",
     "pytest>=8.3.5",
     "pytest-timeout>=2.4.0",
     "sphinx>=8.1.3",
+    "sphinx-autodoc-typehints>=3.0.1",
 ]
 
 [tool.pytest.ini_options]
+cache_dir = "config/pytest_cache"
 testpaths = ["tests"]
 filterwarnings = [
     'error:scatter inputs have incompatible types.*',
@@ -85,8 +87,9 @@ addopts = [
     "-r xXfE",
     "--pdbcls=IPython.terminal.debugger:TerminalPdb",
     "--durations=3",
+    "--verbose",
 ]
-timeout = 32
+timeout = 64
 timeout_method = "thread"  # when jax hangs, signals do not work
 
 # I wanted to use `--import-mode=importlib`, but it breaks importing submodules,
@@ -101,6 +104,7 @@ show_missing = true
 
 [tool.coverage.html]
 show_contexts = true
+directory = "_site/coverage"
 
 [tool.coverage.paths]
 # the first path in each list must be the source directory in the machine that's
@@ -129,6 +133,7 @@ local = [
 
 [tool.ruff]
 exclude = [".asv", "*.ipynb"]
+cache-dir = "config/ruff_cache"
 
 [tool.ruff.format]
 quote-style = "single"
@@ -138,12 +143,43 @@ select = [
     "B", # bugbear: grab bag of additional stuff
     "UP", # pyupgrade: fix some outdated idioms
     "I", # isort: sort and reformat import statements
-    "F", # flake8
+    "F", # pyflakes
+    "D", # pydocstyle
+    "PT", # flake8-pytest-style
 ]
 ignore = [
-    "B028", # warn with stacklevel = 2
+    "B028",  # warn with stacklevel = 2
+    "D105",  # Missing docstring in magic method
+    "F722",  # Syntax error in forward annotation. I ignore this because jaxtyping uses strings for shapes instead of for deferred annotations.
+    "F821",  # Undefined name. I ignore this because strings in jaxtyping.
+    "UP037",  # Remove quotes from type annotation. Ignore because jaxtyping.
 ]
 
+[tool.ruff.lint.per-file-ignores]
+"{config/*,benchmarks/*,docs/*,src/bartz/debug.py,tests/rbartpackages/*,tests/__init__.py}" = [
+    "D100",  # Missing docstring in public module
+    "D101",  # Missing docstring in public class
+    "D102",  # Missing docstring in public method
+    "D103",  # Missing docstring in public function
+    "D104",  # Missing docstring in public package
+]
+
+[tool.ruff.lint.pydocstyle]
+convention = "numpy"
+ignore-decorators = ["functools.cached_property"]
+
+[tool.pydoclint]
+baseline = "config/pydoclint-baseline.txt"
+auto-regenerate-baseline = true
+arg-type-hints-in-signature = true
+arg-type-hints-in-docstring = false
+check-return-types = false
+check-yield-types = false
+treat-property-methods-as-class-attributes = true
+check-style-mismatch = true
+show-filenames-in-every-violation-message = true
+check-class-attributes = false
+
 [tool.uv]
 python-downloads = "never"
 python-preference = "only-system"

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

@@ -22,13 +22,21 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+"""Implement a user interface that mimics the R BART package."""
+
 import functools
+import math
+from typing import Any, Literal
 
 import jax
 import jax.numpy as jnp
+from jax.scipy.special import ndtri
+from jaxtyping import Array, Bool, Float, Float32
 
 from . import grove, jaxext, mcmcloop, mcmcstep, prepcovars
 
+FloatLike = float | Float[Any, '']
+
 
 class gbart:
     """
@@ -46,6 +54,9 @@ class gbart:
         The training responses.
     x_test : array (p, m) or DataFrame, optional
         The test predictors.
+    type
+        The type of regression. 'wbart' for continuous regression, 'pbart' for
+        binary regression with probit link.
     usequants : bool, default False
         Whether to use predictors quantiles instead of a uniform grid to bin
         predictors.
@@ -70,16 +81,20 @@ class gbart:
         Parameters of the prior on tree node generation. The probability that a
         node at depth `d` (0-based) is non-terminal is ``base / (1 + d) **
         power``.
-    maxdepth : int, default 6
-        The maximum depth of the trees. This is 1-based, so with the default
-        ``maxdepth=6``, the depths of the levels range from 0 to 5.
-    lamda : float, optional
-        The scale of the prior on the noise variance. If ``lamda==1``, the
-        prior is an inverse chi-squared scaled to have harmonic mean 1. If
-        not specified, it is set based on `sigest` and `sigquant`.
-    offset : float, optional
+    lamda
+        The prior harmonic mean of the error variance. (The harmonic mean of x
+        is 1/mean(1/x).) If not specified, it is set based on `sigest` and
+        `sigquant`.
+    tau_num
+        The numerator in the expression that determines the prior standard
+        deviation of leaves. If not specified, default to ``(max(y_train) -
+        min(y_train)) / 2`` (or 1 if `y_train` has less than two elements) for
+        continuous regression, and 3 for binary regression.
+    offset
         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.
+        to the mean of `y_train` for continuous regression, and to
+        ``Phi^-1(mean(y_train))`` for binary regression. If `y_train` is empty,
+        `offset` 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
@@ -108,12 +123,24 @@ class gbart:
         The number of initial MCMC samples to discard as burn-in.
     keepevery : int, default 1
         The thinning factor for the MCMC samples, after burn-in.
-    printevery : int, default 100
-        The number of iterations (including skipped ones) between each log.
+    printevery : int or None, default 100
+        The number of iterations (including thinned-away ones) between each log
+        line. Set to `None` to disable logging.
+
+        `printevery` has a few unexpected side effects. On cpu, interrupting
+        with ^C halts the MCMC only on the next log. And the total number of
+        iterations is a multiple of `printevery`, so if ``nskip + keepevery *
+        ndpost`` is not a multiple of `printevery`, some of the last iterations
+        will not be saved.
     seed : int or jax random key, default 0
         The seed for the random number generator.
-    initkw : dict
+    maxdepth : int, default 6
+        The maximum depth of the trees. This is 1-based, so with the default
+        ``maxdepth=6``, the depths of the levels range from 0 to 5.
+    init_kw : dict
         Additional arguments passed to `mcmcstep.init`.
+    run_mcmc_kw : dict
+        Additional arguments passed to `mcmcloop.run_mcmc`.
 
     Attributes
     ----------
@@ -131,20 +158,8 @@ class gbart:
         The standard deviation of the error in the burn-in phase.
     offset : float
         The prior mean of the latent mean function.
-    scale : float
-        The prior standard deviation of the latent mean function.
-    lamda : float
-        The prior harmonic mean of the error variance.
     sigest : float or None
         The estimated standard deviation of the error used to set `lamda`.
-    ntree : int
-        The number of trees.
-    maxdepth : int
-        The maximum depth of the trees.
-
-    Methods
-    -------
-    predict
 
     Notes
     -----
@@ -153,14 +168,17 @@ class gbart:
 
     - If `x_train` and `x_test` are matrices, they have one predictor per row
       instead of per column.
+    - If `type` is not specified, it is determined solely based on the data type
+      of `y_train`, and not on whether it contains only two unique values.
     - If ``usequants=False``, R BART switches to quantiles anyway if there are
       less predictor values than the required number of bins, while bartz
       always follows the specification.
     - The error variance parameter is called `lamda` instead of `lambda`.
     - `rm_const` is always `False`.
     - The default `numcut` is 255 instead of 100.
-    - A lot of functionality is missing (variable selection, discrete response).
+    - A lot of functionality is missing (e.g., variable selection).
     - There are some additional attributes, and some missing.
+    - The trees have a maximum depth.
 
     """
 
@@ -170,6 +188,7 @@ class gbart:
         y_train,
         *,
         x_test=None,
+        type: Literal['wbart', 'pbart'] = 'wbart',
         usequants=False,
         sigest=None,
         sigdf=3,
@@ -177,9 +196,9 @@ class gbart:
         k=2,
         power=2,
         base=0.95,
-        maxdepth=6,
-        lamda=None,
-        offset=None,
+        lamda: FloatLike | None = None,
+        tau_num: FloatLike | None = None,
+        offset: FloatLike | None = None,
         w=None,
         ntree=200,
         numcut=255,
@@ -188,7 +207,9 @@ class gbart:
         keepevery=1,
         printevery=100,
         seed=0,
-        initkw=None,
+        maxdepth=6,
+        init_kw=None,
+        run_mcmc_kw=None,
     ):
         x_train, x_train_fmt = self._process_predictor_input(x_train)
         y_train, _ = self._process_response_input(y_train)
@@ -197,42 +218,41 @@ class gbart:
             w, _ = self._process_response_input(w)
             self._check_same_length(x_train, w)
 
+        y_train = self._process_type_settings(y_train, type, w)
+        # from here onwards, the type is determined by y_train.dtype == bool
         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
+        sigma_mu = self._process_leaf_sdev_settings(y_train, k, ntree, tau_num)
+        lamda, sigest = self._process_error_variance_settings(
+            x_train, y_train, sigest, sigdf, sigquant, lamda
         )
 
         splits, max_split = self._determine_splits(x_train, usequants, numcut)
         x_train = self._bin_predictors(x_train, splits)
-        y_train, lamda_scaled = self._transform_input(y_train, lamda, offset, scale)
 
         mcmc_state = self._setup_mcmc(
             x_train,
             y_train,
+            offset,
             w,
             max_split,
-            lamda_scaled,
+            lamda,
+            sigma_mu,
             sigdf,
             power,
             base,
             maxdepth,
             ntree,
-            initkw,
+            init_kw,
         )
         final_state, burnin_trace, main_trace = self._run_mcmc(
-            mcmc_state, ndpost, nskip, keepevery, printevery, seed
+            mcmc_state, ndpost, nskip, keepevery, printevery, seed, run_mcmc_kw
         )
 
-        sigma = self._extract_sigma(main_trace, scale)
-        first_sigma = self._extract_sigma(burnin_trace, scale)
+        sigma = self._extract_sigma(main_trace)
+        first_sigma = self._extract_sigma(burnin_trace)
 
-        self.offset = offset
-        self.scale = scale
-        self.lamda = lamda
+        self.offset = final_state.offset  # from the state because of buffer donation
         self.sigest = sigest
-        self.ntree = ntree
-        self.maxdepth = maxdepth
         self.sigma = sigma
         self.first_sigma = first_sigma
 
@@ -248,9 +268,8 @@ class gbart:
 
     @functools.cached_property
     def yhat_train(self):
-        x_train = self._mcmc_state['X']
-        yhat_train = self._predict(self._main_trace, x_train)
-        return self._transform_output(yhat_train, self.offset, self.scale)
+        x_train = self._mcmc_state.X
+        return self._predict(self._main_trace, x_train)
 
     @functools.cached_property
     def yhat_train_mean(self):
@@ -269,12 +288,19 @@ class gbart:
         -------
         yhat_test : array (ndpost, m)
             The conditional posterior mean at `x_test` for each MCMC iteration.
+
+        Raises
+        ------
+        ValueError
+            If `x_test` has a different format than `x_train`.
         """
         x_test, x_test_fmt = self._process_predictor_input(x_test)
-        self._check_compatible_formats(x_test_fmt, self._x_train_fmt)
+        if x_test_fmt != self._x_train_fmt:
+            raise ValueError(
+                f'Input format mismatch: {x_test_fmt=} != x_train_fmt={self._x_train_fmt!r}'
+            )
         x_test = self._bin_predictors(x_test, self._splits)
-        yhat_test = self._predict(self._main_trace, x_test)
-        return self._transform_output(yhat_test, self.offset, self.scale)
+        return self._predict(self._main_trace, x_test)
 
     @staticmethod
     def _process_predictor_input(x):
@@ -287,10 +313,6 @@ class gbart:
         assert x.ndim == 2
         return x, fmt
 
-    @staticmethod
-    def _check_compatible_formats(fmt1, fmt2):
-        assert fmt1 == fmt2
-
     @staticmethod
     def _process_response_input(y):
         if hasattr(y, 'to_numpy'):
@@ -308,18 +330,26 @@ class gbart:
         assert get_length(x1) == get_length(x2)
 
     @staticmethod
-    def _process_noise_variance_settings(
-        x_train, y_train, sigest, sigdf, sigquant, lamda, offset
-    ):
-        if lamda is not None:
+    def _process_error_variance_settings(
+        x_train, y_train, sigest, sigdf, sigquant, lamda
+    ) -> tuple[Float32[Array, ''] | None, ...]:
+        if y_train.dtype == bool:
+            if sigest is not None:
+                raise ValueError('Let `sigest=None` for binary regression')
+            if lamda is not None:
+                raise ValueError('Let `lamda=None` for binary regression')
+            return None, None
+        elif lamda is not None:
+            if sigest is not None:
+                raise ValueError('Let `sigest=None` if `lamda` is specified')
             return lamda, None
         else:
             if sigest is not None:
-                sigest2 = sigest * sigest
+                sigest2 = jnp.square(sigest)
             elif y_train.size < 2:
                 sigest2 = 1
             elif y_train.size <= x_train.shape[0]:
-                sigest2 = jnp.var(y_train - offset)
+                sigest2 = jnp.var(y_train)
             else:
                 x_centered = x_train.T - x_train.mean(axis=1)
                 y_centered = y_train - y_train.mean()
@@ -334,20 +364,62 @@ class gbart:
             return sigest2 / invchi2rid, jnp.sqrt(sigest2)
 
     @staticmethod
-    def _process_offset_settings(y_train, offset):
+    def _process_type_settings(y_train, type, w):
+        match type:
+            case 'wbart':
+                if y_train.dtype != jnp.float32:
+                    raise TypeError(
+                        'Continuous regression requires y_train.dtype=float32,'
+                        f' got {y_train.dtype=} instead.'
+                    )
+            case 'pbart':
+                if w is not None:
+                    raise ValueError(
+                        'Binary regression does not support weights, set `w=None`'
+                    )
+                if y_train.dtype != bool:
+                    raise TypeError(
+                        'Binary regression requires y_train.dtype=bool,'
+                        f' got {y_train.dtype=} instead.'
+                    )
+            case _:
+                raise ValueError(f'Invalid {type=}')
+
+        return y_train
+
+    @staticmethod
+    def _process_offset_settings(
+        y_train: Float32[Array, 'n'] | Bool[Array, 'n'],
+        offset: float | Float32[Any, ''] | None,
+    ) -> Float32[Array, '']:
         if offset is not None:
-            return offset
+            return jnp.asarray(offset)
         elif y_train.size < 1:
-            return 0
+            return jnp.array(0.0)
         else:
-            return y_train.mean()
+            mean = y_train.mean()
 
-    @staticmethod
-    def _process_scale_settings(y_train, k):
-        if y_train.size < 2:
-            return 1
+        if y_train.dtype == bool:
+            return ndtri(mean)
         else:
-            return (y_train.max() - y_train.min()) / (2 * k)
+            return mean
+
+    @staticmethod
+    def _process_leaf_sdev_settings(
+        y_train: Float32[Array, 'n'] | Bool[Array, 'n'],
+        k: float,
+        ntree: int,
+        tau_num: FloatLike | None,
+    ):
+        if tau_num is None:
+            if y_train.dtype == bool:
+                tau_num = 3.0
+            elif y_train.size < 2:
+                tau_num = 1.0
+            else:
+                tau_num = (y_train.max() - y_train.min()) / 2
+
+        return tau_num / (k * math.sqrt(ntree))
 
     @staticmethod
     def _determine_splits(x_train, usequants, numcut):
@@ -360,67 +432,83 @@ class gbart:
     def _bin_predictors(x, splits):
         return prepcovars.bin_predictors(x, splits)
 
-    @staticmethod
-    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,
+        offset,
         w,
         max_split,
         lamda,
+        sigma_mu,
         sigdf,
         power,
         base,
         maxdepth,
         ntree,
-        initkw,
+        init_kw,
     ):
         depth = jnp.arange(maxdepth - 1)
         p_nonterminal = base / (1 + depth).astype(float) ** power
-        sigma2_alpha = sigdf / 2
-        sigma2_beta = lamda * sigma2_alpha
+
+        if y_train.dtype == bool:
+            sigma2_alpha = None
+            sigma2_beta = None
+        else:
+            sigma2_alpha = sigdf / 2
+            sigma2_beta = lamda * sigma2_alpha
+
         kw = dict(
             X=x_train,
-            y=y_train,
+            # copy y_train because it's going to be donated in the mcmc loop
+            y=jnp.array(y_train),
+            offset=offset,
             error_scale=w,
             max_split=max_split,
             num_trees=ntree,
             p_nonterminal=p_nonterminal,
+            sigma_mu2=jnp.square(sigma_mu),
             sigma2_alpha=sigma2_alpha,
             sigma2_beta=sigma2_beta,
             min_points_per_leaf=5,
         )
-        if initkw is not None:
-            kw.update(initkw)
+        if init_kw is not None:
+            kw.update(init_kw)
         return mcmcstep.init(**kw)
 
     @staticmethod
-    def _run_mcmc(mcmc_state, ndpost, nskip, keepevery, printevery, seed):
+    def _run_mcmc(mcmc_state, ndpost, nskip, keepevery, printevery, seed, run_mcmc_kw):
         if isinstance(seed, jax.Array) and jnp.issubdtype(
             seed.dtype, jax.dtypes.prng_key
         ):
-            key = seed
+            key = seed.copy()
+            # copy because the inner loop in run_mcmc will donate the buffer
         else:
             key = jax.random.key(seed)
-        callback = mcmcloop.make_simple_print_callback(printevery)
-        return mcmcloop.run_mcmc(key, mcmc_state, nskip, ndpost, keepevery, callback)
 
-    @staticmethod
-    def _predict(trace, x):
-        return mcmcloop.evaluate_trace(trace, x)
+        kw = dict(
+            n_burn=nskip,
+            n_skip=keepevery,
+            inner_loop_length=printevery,
+            allow_overflow=True,
+        )
+        if printevery is not None:
+            kw.update(mcmcloop.make_print_callbacks())
+        if run_mcmc_kw is not None:
+            kw.update(run_mcmc_kw)
+
+        return mcmcloop.run_mcmc(key, mcmc_state, ndpost, **kw)
 
     @staticmethod
-    def _transform_output(y, offset, scale):
-        return offset + scale * y
+    def _extract_sigma(trace) -> Float32[Array, 'trace_length'] | None:
+        if trace['sigma2'] is None:
+            return None
+        else:
+            return jnp.sqrt(trace['sigma2'])
 
     @staticmethod
-    def _extract_sigma(trace, scale):
-        return scale * jnp.sqrt(trace['sigma2'])
+    def _predict(trace, x):
+        return mcmcloop.evaluate_trace(trace, x)
 
     def _show_tree(self, i_sample, i_tree, print_all=False):
         from . import debug
@@ -444,19 +532,26 @@ class gbart:
             )
             beta = bart['sigma2_beta'] + norm2 / 2
         sigma2 = beta / alpha
-        return jnp.sqrt(sigma2) * self.scale
+        return jnp.sqrt(sigma2)
 
     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,
+        resid1 = bart.resid
+
+        trees = grove.evaluate_forest(
+            bart.X,
+            bart.forest.leaf_trees,
+            bart.forest.var_trees,
+            bart.forest.split_trees,
+            jnp.float32,  # TODO remove these configurable dtypes around
         )
-        resid2 = bart['y'] - yhat
+
+        if bart.z is not None:
+            ref = bart.z
+        else:
+            ref = bart.y
+        resid2 = ref - (trees + bart.offset)
+
         return resid1, resid2
 
     def _avg_acc(self):
@@ -495,9 +590,7 @@ class gbart:
     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

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

@@ -23,7 +23,7 @@
 # SOFTWARE.
 
 """
-Super-fast BART (Bayesian Additive Regression Trees) in Python
+Super-fast BART (Bayesian Additive Regression Trees) in Python.
 
 See the manual at https://gattocrucco.github.io/bartz/docs
 """

bartz-0.6.0/src/bartz/_version.py

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

{bartz-0.5.0 → bartz-0.6.0}/src/bartz/debug.py

@@ -180,5 +180,5 @@ def check_trace(trace, state):
         trace['leaf_trees'],
         trace['var_trees'],
         trace['split_trees'],
-        state['max_split'],
+        state.max_split,
     )