bartz 0.4.1__py3-none-any.whl → 0.6.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.
 #
@@ -22,13 +22,15 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+"""Functions to preprocess data."""
+
 import functools
 
 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 +56,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):
+        # wrap this function because autobatch needs traceable args
         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 +72,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 +122,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 +148,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.4.1.dist-info → bartz-0.6.0.dist-info}/METADATA

@@ -1,22 +1,20 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: bartz
-Version: 0.4.1
+Version: 0.6.0
 Summary: Super-fast BART (Bayesian Additive Regression Trees) in Python
-License: MIT
 Author: Giacomo Petrillo
-Author-email: info@giacomopetrillo.com
+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: jaxtyping>=0.3.2
+Requires-Dist: numpy>=1.25.2
+Requires-Dist: scipy>=1.11.4
 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
+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/)
@@ -42,6 +40,7 @@ On CPU, bartz runs at the speed of dbarts (the fastest implementation I know of)
 - [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
@@ -49,4 +48,3 @@ On CPU, bartz runs at the speed of dbarts (the fastest implementation I know of)
 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.6.0.dist-info/RECORD

@@ -0,0 +1,13 @@
+bartz/.DS_Store,sha256=7191af46d7b8c0d4c03c502f94eb01353bc2e615d75c45b3af0e31ab238034b5,6148
+bartz/BART.py,sha256=4de980d640f694e4b65fb17d7c5fe605e15c10f730ee842d39e45501ae8b6829,21606
+bartz/__init__.py,sha256=b57b8ecb9f97ced8a543fa65e191ebda4addd79ef12a8fb357e4f6b23ac2ca72,1434
+bartz/_version.py,sha256=0816378ec0bef470a6ede67a08a0feb182c27a3a8e275a5858f40ce0b1885dc2,22
+bartz/debug.py,sha256=6e06cccfe1c5d6858e62200cd679b73a528af27780d33ba6c164a3a6d96660a5,5356
+bartz/grove.py,sha256=955f03269f373e7bc4f917c6d2ba598e363e8f910adea7f2195693c8443ad381,9636
+bartz/jaxext.py,sha256=ac9aff86e4fb9ff43b5d98995df2b9416a715c82f87270eb1c547fae8973c6dd,13122
+bartz/mcmcloop.py,sha256=e2ad8ae5fb0a49ad3c874e25b6bf2ea966f4d678c0b99bf575393cbe388514bf,17695
+bartz/mcmcstep.py,sha256=5ef4303dfe06be65c3bdfaf44217c74c3e6244444b03dcc5d6eacdee8c9b1d34,72891
+bartz/prepcovars.py,sha256=0db7601a7a73cba508ed4d8c1ca9ac07baacd4f756978bbc80cf173c9257a5ed,5894
+bartz-0.6.0.dist-info/WHEEL,sha256=ec89e25a684ec3ad8fb2a9be806e47666a27f7dcdcf9b55bff3f1d3ed5e60e0a,78
+bartz-0.6.0.dist-info/METADATA,sha256=c35ffe8f0dab8de368e277041105f64650b26f9f4582015cb1b334f88a0e72d6,2846
+bartz-0.6.0.dist-info/RECORD,,

bartz-0.6.0.dist-info/WHEEL

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

bartz-0.4.1.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.1.dist-info/RECORD

@@ -1,13 +0,0 @@
-bartz/BART.py,sha256=CbGzFWtYw5u38Z9-Hy3CbDXpKOOvPFAAkSqu2HZl8no,16862
-bartz/__init__.py,sha256=CcWbDdelQoObG4MW7_jit0Y9fQumU2_Yv3LYtBi1Ag8,1480
-bartz/_version.py,sha256=qBs4HqYsPn6yUHWEuCN4rGUC05gABbl800d8LBd8h9w,22
-bartz/debug.py,sha256=9ZH-JfwZVu5OPhHBEyXQHAU5H9KIu1vxLK7yNv4m4Ew,5314
-bartz/grove.py,sha256=x_6NK_l-hrXfy1PhssYNJkX41-w_WqjDziww0E7YRS8,8500
-bartz/jaxext.py,sha256=UiAHw6pQMz8rjtn_JChFAIdPgUF393S5m-Dduyndnbo,11409
-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.1.dist-info/LICENSE,sha256=heuIJZQK9IexJYC-fYHoLUrgj8HG8yS3G072EvKh-94,1073
-bartz-0.4.1.dist-info/METADATA,sha256=nOOpuP_b-eAJSR1sgyzJaGK4I1iuEWVFMJ9D3JMkY5M,2890
-bartz-0.4.1.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-bartz-0.4.1.dist-info/RECORD,,

bartz-0.4.1.dist-info/WHEEL

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