bartz 0.6.0__tar.gz → 0.8.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (34) hide show
  1. {bartz-0.6.0 → bartz-0.8.0}/PKG-INFO +18 -13
  2. {bartz-0.6.0 → bartz-0.8.0}/README.md +9 -7
  3. bartz-0.8.0/pyproject.toml +250 -0
  4. bartz-0.8.0/src/bartz/BART/__init__.py +27 -0
  5. bartz-0.8.0/src/bartz/BART/_gbart.py +522 -0
  6. {bartz-0.6.0 → bartz-0.8.0}/src/bartz/__init__.py +6 -4
  7. bartz-0.8.0/src/bartz/_interface.py +937 -0
  8. bartz-0.8.0/src/bartz/_profiler.py +318 -0
  9. bartz-0.8.0/src/bartz/_version.py +1 -0
  10. bartz-0.8.0/src/bartz/debug.py +1319 -0
  11. bartz-0.8.0/src/bartz/grove.py +412 -0
  12. bartz-0.8.0/src/bartz/jaxext/__init__.py +287 -0
  13. bartz-0.8.0/src/bartz/jaxext/_autobatch.py +444 -0
  14. bartz-0.8.0/src/bartz/jaxext/scipy/__init__.py +25 -0
  15. bartz-0.8.0/src/bartz/jaxext/scipy/special.py +239 -0
  16. bartz-0.8.0/src/bartz/jaxext/scipy/stats.py +36 -0
  17. bartz-0.8.0/src/bartz/mcmcloop.py +859 -0
  18. bartz-0.8.0/src/bartz/mcmcstep/__init__.py +35 -0
  19. bartz-0.8.0/src/bartz/mcmcstep/_moves.py +904 -0
  20. bartz-0.8.0/src/bartz/mcmcstep/_state.py +1114 -0
  21. bartz-0.8.0/src/bartz/mcmcstep/_step.py +1603 -0
  22. bartz-0.8.0/src/bartz/prepcovars.py +254 -0
  23. bartz-0.8.0/src/bartz/testing/__init__.py +29 -0
  24. bartz-0.8.0/src/bartz/testing/_dgp.py +442 -0
  25. bartz-0.6.0/pyproject.toml +0 -185
  26. bartz-0.6.0/src/bartz/BART.py +0 -603
  27. bartz-0.6.0/src/bartz/_version.py +0 -1
  28. bartz-0.6.0/src/bartz/debug.py +0 -184
  29. bartz-0.6.0/src/bartz/grove.py +0 -310
  30. bartz-0.6.0/src/bartz/jaxext.py +0 -423
  31. bartz-0.6.0/src/bartz/mcmcloop.py +0 -511
  32. bartz-0.6.0/src/bartz/mcmcstep.py +0 -2335
  33. bartz-0.6.0/src/bartz/prepcovars.py +0 -158
  34. {bartz-0.6.0 → bartz-0.8.0}/src/bartz/.DS_Store +0 -0
@@ -1,20 +1,23 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: bartz
3
- Version: 0.6.0
3
+ Version: 0.8.0
4
4
  Summary: Super-fast BART (Bayesian Additive Regression Trees) in Python
5
5
  Author: Giacomo Petrillo
6
6
  Author-email: Giacomo Petrillo <info@giacomopetrillo.com>
7
7
  License-Expression: MIT
8
8
  Requires-Dist: equinox>=0.12.2
9
- Requires-Dist: jax>=0.4.35
10
- Requires-Dist: jaxlib>=0.4.35
9
+ Requires-Dist: jax>=0.5.3
11
10
  Requires-Dist: jaxtyping>=0.3.2
12
11
  Requires-Dist: numpy>=1.25.2
13
12
  Requires-Dist: scipy>=1.11.4
13
+ Requires-Dist: jax[cuda12] ; extra == 'cuda12'
14
+ Requires-Dist: jax[cuda13] ; extra == 'cuda13'
14
15
  Requires-Python: >=3.10
15
- Project-URL: Documentation, https://gattocrucco.github.io/bartz/docs-dev
16
- Project-URL: Homepage, https://github.com/Gattocrucco/bartz
17
- Project-URL: Issues, https://github.com/Gattocrucco/bartz/issues
16
+ Project-URL: Documentation, https://bartz-org.github.io/bartz/docs-dev
17
+ Project-URL: Homepage, https://github.com/bartz-org/bartz
18
+ Project-URL: Issues, https://github.com/bartz-org/bartz/issues
19
+ Provides-Extra: cuda12
20
+ Provides-Extra: cuda13
18
21
  Description-Content-Type: text/markdown
19
22
 
20
23
  [![PyPI](https://img.shields.io/pypi/v/bartz)](https://pypi.org/project/bartz/)
@@ -32,16 +35,18 @@ This Python module provides an implementation of BART that runs on GPU, to proce
32
35
 
33
36
  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.
34
37
 
35
- [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.
38
+ [This Colab notebook](https://colab.research.google.com/github/bartz-org/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 10 minutes.
39
+
40
+ BART is a very flexible method with many variants. This implementation provides only a small subset of the possible features. If you need a feature from [another BART implementation](https://bartz-org.github.io/bartz/docs-dev/pkglist.html) or from the BART literature, please [open an issue on github](https://github.com/bartz-org/bartz/issues).
36
41
 
37
42
  ## Links
38
43
 
39
- - [Documentation (latest release)](https://gattocrucco.github.io/bartz/docs)
40
- - [Documentation (development version)](https://gattocrucco.github.io/bartz/docs-dev)
41
- - [Repository](https://github.com/Gattocrucco/bartz)
42
- - [Code coverage](https://gattocrucco.github.io/bartz/coverage)
43
- - [Benchmarks](https://gattocrucco.github.io/bartz/benchmarks)
44
- - [List of BART packages](https://gattocrucco.github.io/bartz/docs-dev/pkglist.html)
44
+ - [Documentation (latest release)](https://bartz-org.github.io/bartz/docs)
45
+ - [Documentation (development version)](https://bartz-org.github.io/bartz/docs-dev)
46
+ - [Repository](https://github.com/bartz-org/bartz)
47
+ - [Code coverage](https://bartz-org.github.io/bartz/coverage)
48
+ - [Benchmarks](https://bartz-org.github.io/bartz/benchmarks)
49
+ - [List of BART packages](https://bartz-org.github.io/bartz/docs-dev/pkglist.html)
45
50
 
46
51
  ## Citing bartz
47
52
 
@@ -13,16 +13,18 @@ This Python module provides an implementation of BART that runs on GPU, to proce
13
13
 
14
14
  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.
15
15
 
16
- [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.
16
+ [This Colab notebook](https://colab.research.google.com/github/bartz-org/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 10 minutes.
17
+
18
+ BART is a very flexible method with many variants. This implementation provides only a small subset of the possible features. If you need a feature from [another BART implementation](https://bartz-org.github.io/bartz/docs-dev/pkglist.html) or from the BART literature, please [open an issue on github](https://github.com/bartz-org/bartz/issues).
17
19
 
18
20
  ## Links
19
21
 
20
- - [Documentation (latest release)](https://gattocrucco.github.io/bartz/docs)
21
- - [Documentation (development version)](https://gattocrucco.github.io/bartz/docs-dev)
22
- - [Repository](https://github.com/Gattocrucco/bartz)
23
- - [Code coverage](https://gattocrucco.github.io/bartz/coverage)
24
- - [Benchmarks](https://gattocrucco.github.io/bartz/benchmarks)
25
- - [List of BART packages](https://gattocrucco.github.io/bartz/docs-dev/pkglist.html)
22
+ - [Documentation (latest release)](https://bartz-org.github.io/bartz/docs)
23
+ - [Documentation (development version)](https://bartz-org.github.io/bartz/docs-dev)
24
+ - [Repository](https://github.com/bartz-org/bartz)
25
+ - [Code coverage](https://bartz-org.github.io/bartz/coverage)
26
+ - [Benchmarks](https://bartz-org.github.io/bartz/benchmarks)
27
+ - [List of BART packages](https://bartz-org.github.io/bartz/docs-dev/pkglist.html)
26
28
 
27
29
  ## Citing bartz
28
30
 
@@ -0,0 +1,250 @@
1
+ # bartz/pyproject.toml
2
+ #
3
+ # Copyright (c) 2024-2026, The Bartz Contributors
4
+ #
5
+ # This file is part of bartz.
6
+ #
7
+ # Permission is hereby granted, free of charge, to any person obtaining a copy
8
+ # of this software and associated documentation files (the "Software"), to deal
9
+ # in the Software without restriction, including without limitation the rights
10
+ # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
11
+ # copies of the Software, and to permit persons to whom the Software is
12
+ # furnished to do so, subject to the following conditions:
13
+ #
14
+ # The above copyright notice and this permission notice shall be included in all
15
+ # copies or substantial portions of the Software.
16
+ #
17
+ # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
18
+ # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
19
+ # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
20
+ # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
21
+ # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
22
+ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
23
+ # SOFTWARE.
24
+
25
+ [build-system]
26
+ requires = ["uv_build>=0.9.5,<0.10.0"]
27
+ build-backend = "uv_build"
28
+
29
+ [project]
30
+ name = "bartz"
31
+ version = "0.8.0"
32
+ description = "Super-fast BART (Bayesian Additive Regression Trees) in Python"
33
+ authors = [{ name = "Giacomo Petrillo", email = "info@giacomopetrillo.com" }]
34
+ license = "MIT"
35
+ readme = "README.md"
36
+ requires-python = ">=3.10"
37
+ dependencies = [
38
+ "equinox>=0.12.2",
39
+ "jax>=0.5.3",
40
+ "jaxtyping>=0.3.2",
41
+ "numpy>=1.25.2",
42
+ "scipy>=1.11.4",
43
+ ]
44
+
45
+ [project.optional-dependencies]
46
+ cuda12 = ["jax[cuda12]"]
47
+ cuda13 = ["jax[cuda13]"]
48
+
49
+ [project.urls]
50
+ Homepage = "https://github.com/bartz-org/bartz"
51
+ Documentation = "https://bartz-org.github.io/bartz/docs-dev"
52
+ Issues = "https://github.com/bartz-org/bartz/issues"
53
+
54
+ [dependency-groups]
55
+ dev = [
56
+ "appnope>=0.1.4",
57
+ "ipython>=8.36.0",
58
+ "matplotlib>=3.10.3",
59
+ "matplotlib-label-lines>=0.8.1",
60
+ "pre-commit>=4.2.0",
61
+ "rich>=13.9.4",
62
+ "snakeviz>=2.2.2",
63
+ "tqdm>=4.67.1",
64
+ "virtualenv>=20.31.2",
65
+ "asv>=0.6.4",
66
+ "beartype>=0.20.2",
67
+ "flaky>=3.8.1",
68
+ "gitpython>=3.1.43",
69
+ "packaging>=25.0",
70
+ "polars[pandas,pyarrow]>=1.29.0",
71
+ "pytest>=8.3.5",
72
+ "pytest-cov>=6.1.1",
73
+ "pytest-subtests>=0.14.1",
74
+ "pytest-timeout>=2.4.0",
75
+ "pytest-xdist>=3.6.1",
76
+ "rpy2>=3.5.17",
77
+ "sphinx>=8.1.3",
78
+ "sphinx-autodoc-typehints>=3.0.1",
79
+ "tomli>=2.2.1",
80
+ "myst-nb>=1.2.0",
81
+ "jupyterlab>=4.4.2",
82
+ ]
83
+
84
+ [tool.pytest.ini_options]
85
+ cache_dir = "config/pytest_cache"
86
+ testpaths = ["tests"]
87
+ addopts = [
88
+ "-r xXfE",
89
+ "--pdbcls=IPython.terminal.debugger:TerminalPdb",
90
+ "--durations=3",
91
+ "--verbose",
92
+ "--import-mode=importlib",
93
+ ]
94
+ filterwarnings = [
95
+ "ignore:unclosed database:ResourceWarning",
96
+ ]
97
+ timeout = 512
98
+ timeout_method = "thread" # when jax hangs, signals do not work
99
+
100
+ [tool.coverage.run]
101
+ branch = true
102
+ source_pkgs = ["bartz", "tests"]
103
+
104
+ [tool.coverage.report]
105
+ show_missing = true
106
+
107
+ [tool.coverage.html]
108
+ show_contexts = true
109
+ directory = "_site/coverage"
110
+
111
+ [tool.coverage.paths]
112
+ # the first path in each list must be the source directory in the machine that's
113
+ # generating the coverage report
114
+
115
+ github = [
116
+ '/home/runner/work/bartz/bartz/src/bartz/',
117
+ '/Users/runner/work/bartz/bartz/src/bartz/',
118
+ 'D:\a\bartz\bartz\src\bartz\',
119
+ '/Library/Frameworks/Python.framework/Versions/*/lib/python*/site-packages/bartz/',
120
+ '/Users/runner/hostedtoolcache/Python/*/*/lib/python*/site-packages/bartz/',
121
+ '/opt/hostedtoolcache/Python/*/*/lib/python*/site-packages/bartz/',
122
+ 'C:\hostedtoolcache\windows\Python\*\*\Lib\site-packages\bartz\',
123
+ ]
124
+
125
+ local = [
126
+ 'src/bartz/',
127
+ '/home/runner/work/bartz/bartz/src/bartz/',
128
+ '/Users/runner/work/bartz/bartz/src/bartz/',
129
+ 'D:\a\bartz\bartz\src\bartz\',
130
+ '/Library/Frameworks/Python.framework/Versions/*/lib/python*/site-packages/bartz/',
131
+ '/Users/runner/hostedtoolcache/Python/*/*/lib/python*/site-packages/bartz/',
132
+ '/opt/hostedtoolcache/Python/*/*/lib/python*/site-packages/bartz/',
133
+ 'C:\hostedtoolcache\windows\Python\*\*\Lib\site-packages\bartz\',
134
+ ]
135
+
136
+ [tool.ruff]
137
+ exclude = [".asv", "*.ipynb", "benchmarks/latest_bartz"]
138
+ cache-dir = "config/ruff_cache"
139
+
140
+ [tool.ruff.format]
141
+ quote-style = "single"
142
+ skip-magic-trailing-comma = true
143
+
144
+ [tool.ruff.lint.isort]
145
+ split-on-trailing-comma = false
146
+
147
+ [tool.ruff.lint]
148
+ select = [
149
+ "ERA", # eradicate
150
+ "S", # flake8-bandit
151
+ "BLE", # flake8-blind-except
152
+ "B", # bugbear
153
+ "A", # flake8-builtins
154
+ "C4", # flake8-comprehensions
155
+ "CPY", # flake8-copyright
156
+ "DTZ", # flake8-datetimez
157
+ "T10", # flake8-debugger
158
+ "EM", # flake8-errmsg
159
+ "EXE", # flake8-executable
160
+ "FIX", # flake8-fixme
161
+ "ISC", # flake8-implicit-str-concat
162
+ "INP", # flake8-no-pep420
163
+ "PIE", # flake8-pie
164
+ "T20", # flake8-print
165
+ "PT", # flake8-pytest-style
166
+ "RSE", # flake8-raise
167
+ "RET", # flake8-return
168
+ "SLF", # flake8-self
169
+ "SIM", # flake8-simplify
170
+ "TID", # flake8-tidy-imports
171
+ "ARG", # flake8-unused-arguments
172
+ "PTH", # flake8-use-pathlib
173
+ "FLY", # flynt
174
+ "I", # isort
175
+ "C90", # mccabe
176
+ "NPY", # NumPy-specific rules
177
+ "PERF", # Perflint
178
+ "W", # pycodestyle Warning
179
+ "F", # pyflakes
180
+ "D", # pydocstyle
181
+ "PGH", # pygrep-hooks
182
+ "PLC", # Pylint Convention
183
+ "PLE", # Pylint Error
184
+ "PLR", # Pylint Refactor
185
+ "PLW", # Pyling Warning
186
+ "UP", # pyupgrade
187
+ "FURB", # refurb
188
+ "RUF", # Ruff-specific rules
189
+ "TRY", # tryceratops
190
+ ]
191
+ ignore = [
192
+ "B028", # warn with stacklevel = 2
193
+ "C408", # Unnecessary `dict()` call (rewrite as a literal), it's too convenient for kwargs
194
+ "D105", # Missing docstring in magic method
195
+ "F722", # Syntax error in forward annotation. I ignore this because jaxtyping uses strings for shapes instead of for deferred annotations.
196
+ "PIE790", # Unnecessary ... or pass. Ignored because sometimes I use ... as sentinel to tell the rest of ruff and pyright that an implementation is a stub.
197
+ "PLR0912", # Too many branches; ignore bc C901 already handles this
198
+ "PLR0913", # Too many arguments in function definition. Maybe I should do something about this?
199
+ "PLR2004", # Magic value used in comparison, consider replacing `*` with a constant variable
200
+ "RET505", # Unnecessary `{branch}` after `return` statement. I ignore this because I like to keep branches for readability.
201
+ "RET506", # Unnecessary `else` after `raise` statement. I ignore this because I like to keep branches for readability.
202
+ "S101", # Use of `assert` detected. Too annoying.
203
+ "SIM108", # SIM108 Use ternary operator `*` instead of `if`-`else`-block, I find blocks more readable
204
+ "UP037", # Remove quotes from type annotation. Ignore because jaxtyping.
205
+ ]
206
+
207
+ [tool.ruff.lint.per-file-ignores]
208
+ "{config/*,docs/*}" = [
209
+ "D100", # Missing docstring in public module
210
+ "D101", # Missing docstring in public class
211
+ "D102", # Missing docstring in public method
212
+ "D103", # Missing docstring in public function
213
+ "D104", # Missing docstring in public package
214
+ "INP001", # File * is part of an implicit namespace package. Add an `__init__.py`.
215
+ ]
216
+ "src/bartz/_version.py" = [
217
+ "CPY001", # Missing copyright notice at top of file
218
+ ]
219
+ "{config/*,docs/*,tests/*,benchmarks/_vendor_latest_bartz.py,scripts/*}" = [
220
+ "T201", # `print` found
221
+ ]
222
+ "{tests/*,benchmarks/*}" = [
223
+ "SLF001", # Private member accessed: `*`
224
+ "TID253", # `{module}` is banned at the module level
225
+ ]
226
+ "docs/conf.py" = [
227
+ "S607", # Starting a process with a partial executable path. Ignored because for a build script it makes more sense to use PATH.
228
+ ]
229
+
230
+ [tool.ruff.lint.pydocstyle]
231
+ convention = "numpy"
232
+
233
+ [tool.ruff.lint.flake8-copyright]
234
+ min-file-size = 1
235
+
236
+ [tool.ruff.lint.flake8-tidy-imports]
237
+ banned-module-level-imports = ["bartz.debug"]
238
+ ban-relative-imports = "all"
239
+
240
+ [tool.pydoclint]
241
+ arg-type-hints-in-signature = true
242
+ arg-type-hints-in-docstring = false
243
+ check-return-types = false
244
+ check-yield-types = false
245
+ treat-property-methods-as-class-attributes = true
246
+ check-style-mismatch = true
247
+ show-filenames-in-every-violation-message = true
248
+ check-class-attributes = false
249
+ # do not check class attributes because in dataclasses I document them as
250
+ # init parameters because they are duplicated in the html docs otherwise.
@@ -0,0 +1,27 @@
1
+ # bartz/src/bartz/BART/__init__.py
2
+ #
3
+ # Copyright (c) 2026, The Bartz Contributors
4
+ #
5
+ # This file is part of bartz.
6
+ #
7
+ # Permission is hereby granted, free of charge, to any person obtaining a copy
8
+ # of this software and associated documentation files (the "Software"), to deal
9
+ # in the Software without restriction, including without limitation the rights
10
+ # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
11
+ # copies of the Software, and to permit persons to whom the Software is
12
+ # furnished to do so, subject to the following conditions:
13
+ #
14
+ # The above copyright notice and this permission notice shall be included in all
15
+ # copies or substantial portions of the Software.
16
+ #
17
+ # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
18
+ # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
19
+ # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
20
+ # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
21
+ # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
22
+ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
23
+ # SOFTWARE.
24
+
25
+ """Implement classes `mc_gbart` and `gbart` that mimic the R BART3 package."""
26
+
27
+ from bartz.BART._gbart import gbart, mc_gbart # noqa: F401