bartz 0.4.1__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.4.1.dist-info → bartz-0.5.0.dist-info}/METADATA

@@ -1,22 +1,18 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: bartz
-Version: 0.4.1
+Version: 0.5.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: 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
-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 +38,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 +46,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.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.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