Coverage for src/bartz/_interface.py: 96%
676 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-10-15 10:04 +0000
1# bartz/src/bartz/_interface.py
2#
3# Copyright (c) 2025-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.
25"""Main high-level interface of the package."""
27import math
28import pickle
29from collections.abc import Collection, Hashable, Mapping, Sequence
30from dataclasses import replace
31from enum import Enum
32from functools import cached_property
33from os import PathLike, cpu_count
34from pathlib import Path
36# WORKAROUND(python<3.15): use frozendict instead of MappingProxyType
37from types import MappingProxyType
38from typing import Any, Literal, Protocol, TypedDict, cast, runtime_checkable
39from warnings import warn
41import jax
42import jax.numpy as jnp
43from equinox import Module, error_if, field, tree_at
44from jax import Device, debug_nans, device_put, lax, make_mesh, random, tree
45from jax.scipy.linalg import solve_triangular
46from jax.scipy.special import ndtr, ndtri
47from jax.sharding import AxisType, Mesh, NamedSharding, PartitionSpec
48from jax.typing import DTypeLike
49from jaxtyping import Array, Bool, Float, Float32, Int32, Key, Real, Shaped, UInt
50from numpy import ndarray
52from bartz._jaxext import equal_shards, is_key, jit, jit_active, project, split
53from bartz.grove import (
54 TreeHeaps,
55 TreesTrace,
56 check_trace,
57 evaluate_forest,
58 forest_depth_distr,
59 format_tree,
60 points_per_node_distr,
61)
62from bartz.mcmcloop import (
63 BurninTrace,
64 CallbackTuple,
65 CheckPlatformCallback,
66 MainTrace,
67 MainTraceWithTrainPred,
68 RunMCMCResult,
69 compute_varcount,
70 evaluate_trace,
71 make_print_callback,
72 make_tqdm_callback,
73 run_mcmc,
74)
75from bartz.mcmcstep import DiagWishart, OutcomeType, Wishart, make_p_nonterminal
76from bartz.mcmcstep._axes import (
77 chain_to_axis,
78 chain_vmap_axes,
79 chainful_axis,
80 get_has_chains,
81 trace_sample_axes,
82)
83from bartz.mcmcstep._state import (
84 ArrayLike,
85 FloatLike,
86 State,
87 chol_with_gersh,
88 init,
89 inv_via_chol_with_gersh,
90 leaf_partition_spec,
91)
92from bartz.prepcovars import Binner, BinnerFactory, UniqueQuantileBinner
95class PredictKind(Enum):
96 """Kind of output of `Bart.predict`."""
98 mean = 'mean'
99 """The posterior mean of the conditional mean, shape ``(m,)`` (or
100 ``(k, m)`` for multivariate regression)."""
102 mean_samples = 'mean_samples'
103 """Per-sample conditional mean, shape ``(num_chains * n_save, m)``
104 (or ``(num_chains * n_save, k, m)``). For binary regression, this is
105 the probit-transformed sum-of-trees, divided by the error scale `w`
106 first if the model is heteroskedastic."""
108 outcome_samples = 'outcome_samples'
109 """Samples of the outcome variable, shape ``(num_chains * n_save,
110 m)`` (or ``(num_chains * n_save, k, m)``). For binary regression,
111 these are Bernoulli draws. For continuous regression, these are
112 Gaussian draws with the posterior noise variance."""
114 latent_samples = 'latent_samples'
115 """Raw sum-of-trees values, shape ``(num_chains * n_save, m)`` (or
116 ``(num_chains * n_save, k, m)``)."""
119@runtime_checkable
120class DataFrame(Protocol):
121 """DataFrame duck-type for `Bart`."""
123 @property
124 def columns(self) -> Collection[str]:
125 """The names of the columns."""
126 ...
128 def to_numpy(self) -> Shaped[ndarray, '*shape']:
129 """Convert the dataframe to a 2d numpy array with columns on the second axis."""
130 ...
133@runtime_checkable
134class Series(Protocol):
135 """Series duck-type for `Bart`."""
137 @property
138 def name(self) -> Hashable:
139 """The name of the series."""
140 ...
142 def to_numpy(self) -> Shaped[ndarray, '*shape']:
143 """Convert the series to a 1d numpy array."""
144 ...
147class SparseConfig(Module):
148 R"""
149 Configuration of a sparsity-inducing variable selection prior.
151 This is the prior of [1]_. Pass an instance to the `sparse` argument of
152 `Bart` to activate variable selection on the predictors. The prior on the
153 choice of predictor for each decision rule is
155 .. math::
156 (s_1, \ldots, s_p) \sim
157 \operatorname{Dirichlet}(\mathtt{theta}/p, \ldots, \mathtt{theta}/p).
159 If `theta` is not specified, it's a priori distributed according to
161 .. math::
162 \frac{\mathtt{theta}}{\mathtt{theta} + \mathtt{rho}} \sim
163 \operatorname{Beta}(\mathtt{a}, \mathtt{b}).
165 References
166 ----------
167 .. [1] Linero, Antonio R. (2018). “Bayesian Regression Trees for
168 High-Dimensional Prediction and Variable Selection”. In: Journal of the
169 American Statistical Association 113.522, pp. 626-636.
170 """
172 theta: FloatLike | None = None
173 """Concentration of the Dirichlet prior. If not specified, it is sampled
174 from a Beta prior parametrized by `a`, `b` and `rho`. If set directly, it
175 should be in the ballpark of the predictor count p or lower."""
177 a: FloatLike = 0.5
178 """Shape parameter of the Beta prior on ``theta / (theta + rho)``."""
180 b: FloatLike = 1.0
181 """Shape parameter of the Beta prior on ``theta / (theta + rho)``."""
183 rho: FloatLike | None = None
184 """Scale of the Beta prior on `theta`. If not specified, set to the number
185 of predictors p. Lower values prefer more sparsity."""
187 augment: bool = field(static=True, default=True)
188 """Whether to account exactly for the decision rules forbidden by the
189 ancestors of each node when updating the variable selection probabilities,
190 using data augmentation. On by default. Setting it to `False` ignores the
191 forbidden rules, which is faster but only approximate. This matters most
192 with few predictors with few cutpoints each, where the same predictor
193 cannot be re-used down a branch."""
195 enabled: bool = field(static=True, default=True)
196 """Whether variable selection is active."""
199class Bart(Module):
200 R"""
201 Nonparametric regression with Bayesian Additive Regression Trees (BART).
203 Regress `y_train` on `x_train` with a latent mean function represented as
204 a sum of decision trees [2]_. The inference is carried out by sampling the
205 posterior distribution of the tree ensemble with an MCMC.
207 Parameters
208 ----------
209 x_train
210 The training predictors. Missing values are not supported: NaN raises
211 an error.
212 y_train
213 The training responses. For univariate regression, a 1D array of shape
214 `(n,)`. For multivariate regression, a 2D array of shape `(k, n)` where
215 `k` is the number of response components, as introduced in [3]_. For
216 binary regression, the convention is that non-zero values mean 1, zero
217 mean 0, like booleans.
218 outcome_type
219 The type of regression. ``'continuous'`` for continuous regression,
220 ``'binary'`` for binary regression with probit link. For multivariate
221 regression, a scalar value applies to all components; alternatively, a
222 sequence of per-component types (e.g., ``['binary', 'continuous']``)
223 specifies mixed outcome types. Binary components in multivariate
224 outcomes follow the multivariate probit BART formulation of [4]_.
225 sparse
226 A `SparseConfig` for the sparsity-inducing variable selection prior of
227 [1]_. Disabled by default; pass a `SparseConfig` to enable it.
228 varprob
229 The probability distribution over the `p` predictors for choosing a
230 predictor to split on in a decision node a priori. Must be > 0. It does
231 not need to be normalized to sum to 1. If not specified, use a uniform
232 distribution. If `sparse` is enabled, this is used as initial value for
233 the MCMC.
234 binner
235 A callable that, given the training predictors and a random key,
236 returns a `~bartz.prepcovars.Binner` instance. The default is
237 `~bartz.prepcovars.UniqueQuantileBinner`, which places cutpoints at
238 the quantiles of each predictor. Other built-in options are
239 `~bartz.prepcovars.RangeEvenBinner` (evenly-spaced cutpoints over the
240 observed range) and `~bartz.prepcovars.GivenSplitsBinner` (R BART
241 ``xinfo`` format). To pass options, use `functools.partial`, e.g.
242 ``binner=partial(UniqueQuantileBinner, max_bins=128)``.
243 rm_const
244 How to treat predictors with no associated decision rules (i.e., there
245 are no available cutpoints for that predictor). If `True` (default),
246 they are ignored. If `False`, an error is raised if there are any.
247 sigma_df
248 The degrees of freedom of the prior on the error precision. For
249 multivariate regression with `k` components, the Wishart degrees of
250 freedom are set to ``sigma_df + k - 1``, such that each marginal error
251 variance keeps ``sigma_df`` degrees of freedom; the same holds when the
252 error covariance is restricted to diagonal.
253 sigma_scale
254 Sets the scale of the prior on the error precision. If 'auto'
255 (default), the prior is scaled so that the prior harmonic mean of the
256 error variance is ``var(y_train)``, where with `error_scale` the
257 variance is a precision-weighted one that estimates the error variance
258 at unit error scale. Otherwise, ``square(sigma_scale)`` is the prior
259 harmonic mean of the error variance. For multivariate regression this
260 applies to the marginal variance of each component, whatever `k`, and
261 a scalar is broadcast to all components. For mixed outcome types,
262 binary components are ignored.
263 sigma_init
264 The initial value of the error standard deviation in the MCMC. If
265 'auto' (default), the initial error precision is set to ``diag(1 /
266 var(y_train))``, with the same precision-weighted variance as
267 `sigma_scale` when `error_scale` is given. Otherwise, the initial
268 precision is ``diag(1 / square(sigma_init))``; for multivariate
269 regression a scalar is broadcast to all components. For mixed outcome
270 types, binary components are ignored.
271 k
272 The inverse scale of the prior standard deviation on the latent mean
273 function, relative to half the observed range of `y_train`. If `y_train`
274 has less than two elements, `k` is ignored and the scale is set to 1.
275 power
276 base
277 Parameters of the prior on tree node generation. The probability that a
278 node at depth `d` (0-based) is non-terminal is ``base / (1 + d) **
279 power``.
280 tau_num
281 The numerator in the expression that determines the prior standard
282 deviation of leaves. If not specified, default to ``(max(y_train) -
283 min(y_train)) / 2`` (or 1 if `y_train` has less than two elements) for
284 continuous regression, and 3 for binary regression. For multivariate
285 regression, the range is computed per component. For mixed outcome
286 types, each component uses the default for its type.
287 sigma_mu_df
288 If specified, the prior precision of the leaves is sampled in the MCMC
289 instead of being held fixed. The prior is Wishart (Gamma in the
290 univariate case) with ``sigma_mu_df + k - 1`` degrees of freedom for
291 `k` outcome components, such that each marginal leaf variance keeps
292 ``sigma_mu_df`` degrees of freedom, like `sigma_df`. The prior harmonic
293 mean of each marginal leaf variance, which is also its initial value,
294 is the fixed value otherwise set by `tau_num` and the `k` argument (the
295 leaf scale, not the number of components).
296 offset
297 The prior mean of the latent mean function. If not specified, it is set
298 to the mean of `y_train` for continuous regression, and to
299 ``Phi^-1(mean(y_train != 0))`` for binary regression. If `y_train` is
300 empty, `offset` is set to 0. With binary regression, if `y_train` is
301 all zero or all non-zero, it is set to ``Phi^-1(1/(n+1))`` or
302 ``Phi^-1(n/(n+1))``, respectively. For multivariate regression, can be
303 a scalar (broadcast to all components) or a `(k,)` vector. If not
304 specified, it is set to the per-component mean of `y_train`. For mixed
305 outcome types, each component uses the default for its type.
306 error_scale
307 Coefficients that rescale the error standard deviation on each
308 datapoint ("w" in BART3). Not specifying `error_scale` is equivalent to
309 setting it to 1 for all datapoints. Shape ``(n,)`` applies the same
310 scalar scale to every outcome component; for multivariate regression,
311 ``(k, n)`` instead supplies a per-component scale per datapoint.
312 Supported with binary (probit) outcomes, where it scales the latent
313 error so the success probability is ``Phi(latent / error_scale)``,
314 including the binary components of a mixed regression.
315 missing
316 Boolean mask with the same shape as `y_train`; `True` marks entries
317 to be ignored by the MCMC. The values of `y_train` at masked
318 positions are ignored and may be anything, even non-finite. If 2-D,
319 the error covariance must be diagonal.
320 num_trees
321 The number of trees used to represent the latent mean function.
322 n_save
323 The number of MCMC samples to save, after burn-in, per chain. The
324 total trace length across all chains is ``num_chains * n_save``.
325 n_burn
326 The number of initial MCMC samples to discard as burn-in. This number
327 of samples is discarded from each chain.
328 n_skip
329 The thinning factor for the MCMC samples, after burn-in.
330 printevery
331 The number of iterations (including thinned-away ones) between each log
332 line. Set to `None` to disable progress reporting entirely (this ignores
333 `pbar`). ^C interrupts the MCMC only every `printevery` iterations, so
334 with reporting disabled it's impossible to kill the MCMC conveniently.
335 pbar
336 If `True`, show a `tqdm` progress bar instead of printing log lines. The
337 bar advances every iteration and refreshes the acceptance statistics
338 every `printevery` iterations. Ignored if `printevery` is `None`.
339 num_chains
340 The number of independent Markov chains to run.
342 The difference between ``num_chains=None`` and ``num_chains=1`` is that
343 in the latter case in the object attributes and some methods there will
344 be an explicit chain axis of size 1.
345 num_chain_devices
346 The number of devices to spread the chains across. Must be a divisor of
347 `num_chains`. Each device will run a fraction of the chains. If 'auto'
348 (default) and running on cpu, the number of devices is picked
349 automatically based on the number of cores and the number of available
350 devices (all the virtual jax cpu devices, or the `devices` list if set).
351 num_data_devices
352 The number of devices to split datapoints across. Must be a divisor of
353 `n`. This is useful only with very high `n`, about > 1000_000. `predict`
354 parallelizes across the same devices, splitting the test points; the
355 number of test points must be a multiple of `num_data_devices` as well.
357 If both num_chain_devices and num_data_devices are specified, the total
358 number of devices used is the product of the two.
359 devices
360 One or more devices used to run the MCMC on. If not specified, the
361 computation will follow the placement of the input arrays. If a list of
362 devices, this argument can be longer than the number of devices needed.
363 seed
364 The seed for the random number generator.
365 maxdepth
366 The maximum depth of the trees. This is 1-based, so with the default
367 ``maxdepth=6``, the depths of the levels range from 0 to 5.
368 precompute_predict_train
369 If `True`, compute the predictions at the training points during the
370 MCMC. Off by default; makes ``predict('train', ...)`` faster at the cost
371 of more memory.
372 init_kw
373 Additional arguments passed to `bartz.mcmcstep.init`.
374 run_mcmc_kw
375 Additional arguments passed to `bartz.mcmcloop.run_mcmc`.
377 Notes
378 -----
379 On gpu, the leaves of the trees are stored in float16 (see the
380 ``leaf_dtype`` argument of `bartz.mcmcstep.init`), which limits the
381 signal-to-noise ratio of the fit to less than about 1000, a limit that
382 should always hold in realistic usage.
384 References
385 ----------
386 .. [1] Linero, Antonio R. (2018). “Bayesian Regression Trees for
387 High-Dimensional Prediction and Variable Selection”. In: Journal of the
388 American Statistical Association 113.522, pp. 626-636.
389 .. [2] Hugh A. Chipman, Edward I. George, Robert E. McCulloch "BART:
390 Bayesian additive regression trees," The Annals of Applied Statistics,
391 Ann. Appl. Stat. 4(1), 266-298, (March 2010).
392 .. [3] Um, Seungha, Antonio R. Linero, Debajyoti Sinha, and Dipankar
393 Bandyopadhyay (2023). "Bayesian additive regression trees for
394 multivariate skewed responses". In: Statistics in Medicine 42.3,
395 pp. 246-263.
396 .. [4] Goh, Yong Chen, Wuu Kuang Soh, Andrew C. Parnell, and Keefe
397 Murphy (2024). "Joint Models for Handling Non-Ignorable Missing
398 Data using Bayesian Additive Regression Trees: Application to
399 Leaf Photosynthetic Traits Data". arXiv:2412.14946 [stat.ME].
401 """
403 _main_trace: MainTrace
404 _burnin_trace: BurninTrace
405 _mcmc_state: State
406 _binner: Binner
407 _binary_mask: Bool[Array, ''] | Bool[Array, ' k']
408 # WORKAROUND(jax<0.9.1): use `jax.tree.static` instead of `field(static=True)`
409 _x_train_fmt: Any = field(static=True)
410 _device: Device | None = field(static=True)
412 def __init__(
413 self,
414 x_train: Real[ArrayLike, 'p n'] | DataFrame,
415 y_train: Float32[ArrayLike, ' n']
416 | Float32[ArrayLike, 'k n']
417 | Series
418 | DataFrame,
419 *,
420 outcome_type: OutcomeType | str | Sequence[OutcomeType | str] = 'continuous',
421 sparse: SparseConfig = SparseConfig(enabled=False),
422 varprob: Float[ArrayLike, ' p'] | None = None,
423 binner: BinnerFactory = UniqueQuantileBinner,
424 rm_const: bool = True,
425 sigma_df: FloatLike = 3.0,
426 sigma_scale: FloatLike | Float[ArrayLike, ' k'] | Literal['auto'] = 'auto',
427 sigma_init: FloatLike | Float[ArrayLike, ' k'] | Literal['auto'] = 'auto',
428 k: FloatLike = 2.0,
429 power: FloatLike = 2.0,
430 base: FloatLike = 0.95,
431 tau_num: FloatLike | None = None,
432 sigma_mu_df: FloatLike | None = None,
433 offset: FloatLike | Float[ArrayLike, ' k'] | None = None,
434 error_scale: Float[ArrayLike, ' n']
435 | Float[ArrayLike, 'k n']
436 | Series
437 | DataFrame
438 | None = None,
439 missing: Bool[ArrayLike, ' n']
440 | Bool[ArrayLike, 'k n']
441 | Series
442 | DataFrame
443 | None = None,
444 num_trees: int = 200,
445 n_save: int = 1000,
446 n_burn: int = 1000,
447 n_skip: int = 1,
448 printevery: int | None = 100,
449 pbar: bool = True,
450 num_chains: int | None = 4,
451 num_chain_devices: int | None | Literal['auto'] = 'auto',
452 num_data_devices: int | None = None,
453 devices: Literal['cpu', 'gpu'] | Device | Sequence[Device] | None = None,
454 seed: int | Key[Array, ''] = 0,
455 maxdepth: int = 6,
456 precompute_predict_train: bool = False,
457 init_kw: Mapping = MappingProxyType({}),
458 run_mcmc_kw: Mapping = MappingProxyType({}),
459 ) -> None:
460 # check data and put it in the right format
461 x_train, x_train_fmt = _process_predictor_input(x_train)
462 y_train = _process_response_input(y_train)
463 _check_same_length(x_train, y_train)
465 if error_scale is not None:
466 # `error_scale` is donated downstream as `init`'s `error_scale`, which
467 # keeps it (sharded) as `State.error_scale` for prediction
468 error_scale = _process_response_input(error_scale)
469 _check_same_length(x_train, error_scale)
471 if missing is not None:
472 missing = _process_response_input(missing, dtype=jnp.bool_)
473 _check_same_length(x_train, missing)
475 # check data types are correct for continuous/binary/multivariate regression
476 outcome_type, binary_mask = _check_type_settings(
477 y_train, outcome_type, error_scale
478 )
480 # process "standardization" settings
481 offset = _process_offset_settings(
482 y_train,
483 binary_mask,
484 missing,
485 None if offset is None else jnp.asarray(offset, jnp.float32),
486 )
487 leaf_prior_cov_inv = _process_leaf_variance_settings(
488 y_train,
489 binary_mask,
490 missing,
491 k=jnp.asarray(k, jnp.float32),
492 num_trees=num_trees,
493 tau_num=None if tau_num is None else jnp.asarray(tau_num, jnp.float32),
494 sigma_mu_df=None
495 if sigma_mu_df is None
496 else jnp.asarray(sigma_mu_df, jnp.float32),
497 )
498 error_cov_inv = _process_error_variance_settings(
499 y_train,
500 binary_mask,
501 missing,
502 outcome_type=outcome_type,
503 sigma_df=sigma_df,
504 sigma_scale=sigma_scale,
505 sigma_init=sigma_init,
506 error_scale=error_scale,
507 )
509 # split the user-provided seed into an mcmc key and a binner key
510 if not is_key(seed):
511 seed = random.key(seed)
512 keys = split(seed)
514 # construct the binner and bin x_train
515 binner_obj = binner(x_train, key=keys.pop())
516 x_train = binner_obj.bin(x_train)
517 # copy max_split because `mcmcstep.init` donates it
518 max_split = jnp.array(binner_obj.max_split)
520 warn_byte_boundaries(x_train, max_split, maxdepth)
522 # setup and run mcmc
523 initial_state, mcmc_key, device, check_platform = _setup_mcmc(
524 keys.pop(),
525 x_train,
526 y_train,
527 outcome_type=outcome_type,
528 offset=offset,
529 error_scale=error_scale,
530 missing=missing,
531 max_split=max_split,
532 leaf_prior_cov_inv=leaf_prior_cov_inv,
533 error_cov_inv=error_cov_inv,
534 power=power,
535 base=base,
536 maxdepth=maxdepth,
537 num_trees=num_trees,
538 init_kw=init_kw,
539 rm_const=rm_const,
540 sparse=sparse,
541 varprob=varprob,
542 num_chains=num_chains,
543 num_chain_devices=num_chain_devices,
544 num_data_devices=num_data_devices,
545 devices=devices,
546 n_burn=n_burn,
547 )
548 result = _run_mcmc(
549 mcmc_key,
550 initial_state,
551 n_save,
552 n_burn=n_burn,
553 n_skip=n_skip,
554 printevery=printevery,
555 pbar=pbar,
556 precompute_predict_train=precompute_predict_train,
557 run_mcmc_kw=run_mcmc_kw,
558 check_platform=check_platform,
559 )
561 # set private attributes
562 self._main_trace = result.main_trace
563 self._burnin_trace = result.burnin_trace
564 self._mcmc_state = result.final_state
565 self._binner = binner_obj
566 self._x_train_fmt = x_train_fmt
567 self._binary_mask = binary_mask
568 self._device = device
570 def predict(
571 self,
572 x_test: Real[ArrayLike, 'p m'] | DataFrame | str,
573 *,
574 kind: PredictKind | str = 'mean',
575 key: Key[Array, ''] | None = None,
576 error_scale: Float[ArrayLike, ' m']
577 | Float[ArrayLike, 'k m']
578 | Series
579 | DataFrame
580 | None = None,
581 ) -> (
582 Float32[Array, ' m']
583 | Float32[Array, 'k m']
584 | Float32[Array, 'ndpost m']
585 | Float32[Array, 'ndpost k m']
586 ):
587 """
588 Compute predictions at `x_test`.
590 Parameters
591 ----------
592 x_test
593 The test predictors, or the string ``'train'`` to compute
594 predictions on the training data. Missing values are not supported:
595 NaN raises an error.
596 kind
597 The kind of output. See `PredictKind` for details.
598 key
599 Jax random key, required when ``kind='outcome_samples'``.
600 error_scale
601 Per-observation error scale. Used with ``kind='outcome_samples'``,
602 and also with ``kind='mean'`` or ``'mean_samples'`` for binary
603 outcomes (since the success probability is ``Phi(latent /
604 error_scale)``). Required when the model was fit with `error_scale`
605 and ``x_test`` is new data. Shape matches the shape used at
606 fitting: ``(m,)`` for scalar scales, ``(k, m)`` for multivariate
607 per-component scales.
609 Returns
610 -------
611 Predictions at `x_test` in the requested format.
613 Raises
614 ------
615 ValueError
616 If `x_test` contains NaN or has a different format than `x_train`,
617 or if `error_scale` is specified when it should be `None`, or if
618 `error_scale` is not specified when it is required, or if the model
619 splits datapoints across devices (`num_data_devices`) and the number
620 of test points is not a multiple of the number of data devices.
622 Notes
623 -----
624 If the model splits datapoints across devices (`num_data_devices`),
625 the test points and the returned predictions are split the same way.
626 """
627 # parse arguments
628 kind = PredictKind(kind)
629 if kind is PredictKind.outcome_samples and key is None: 629 ↛ 630line 629 didn't jump to line 630 because the condition on line 629 was never true
630 msg = '`key` not specified'
631 raise ValueError(msg)
632 error_scale = self._process_error_scale_test(x_test, kind, error_scale)
633 x_test_is_train = isinstance(x_test, str) and x_test == 'train'
635 # use the predictions precomputed during the MCMC, if available
636 if x_test_is_train and isinstance(self._main_trace, MainTraceWithTrainPred):
637 return predict_train(
638 key,
639 self._main_trace,
640 error_scale=error_scale,
641 binary_indices=self._mcmc_state.binary_indices,
642 has_binary=self._mcmc_state.z is not None,
643 kind=kind,
644 )
646 x_test = self._process_x_test(x_test, error_scale)
648 # place new test data on the devices of the model; the training data
649 # is already in place
650 if not x_test_is_train:
651 x_test, error_scale = self._device_put_test(x_test, error_scale)
653 # invoke jitted implementation
654 return predict(
655 key,
656 self._main_trace,
657 x_test,
658 error_scale=error_scale,
659 binary_indices=self._mcmc_state.binary_indices,
660 has_binary=self._mcmc_state.z is not None,
661 kind=kind,
662 # the test points are sharded over the mesh 'data' axis (when
663 # there is one): the training data at `init`, new test data by
664 # `_device_put_test`. `evaluate_trace` can't detect this on its
665 # own at trace time, so declare it.
666 test_points='shard_and_autobatch',
667 )
669 def _drop_device_info(self) -> 'Bart':
670 """Return a copy of the model without device placement metadata.
672 Clear the meshes in the MCMC state config and in the traces, and the
673 explicitly requested device. Only this static metadata is dropped: the
674 arrays keep their actual placement.
675 """
676 config = replace(self._mcmc_state.config, mesh=None)
677 main_trace = replace(self._main_trace, mesh=None)
678 burnin_trace = replace(self._burnin_trace, mesh=None)
679 obj = tree_at(
680 lambda b: (b._mcmc_state.config, b._main_trace, b._burnin_trace), # noqa: SLF001
681 self,
682 (config, main_trace, burnin_trace),
683 )
684 # `_device` is a static field, out of `tree_at`'s reach, so modify the
685 # fresh copy in place
686 object.__setattr__(obj, '_device', None)
687 return obj
689 def dump(self, path: str | PathLike) -> None:
690 """Serialize the fitted model to a file with `pickle`.
692 Parameters
693 ----------
694 path
695 The file to write to.
697 Notes
698 -----
699 Intended for short-term storage (e.g. caching across processes), not
700 long-term archival: the format depends on the versions of bartz, jax and
701 equinox. The arrays are copied to host memory and all device/sharding
702 placement is dropped; `load` reconstructs a single-device model.
703 """
704 # drop all device info (`Device` objects are not picklable), then
705 # gather any sharded arrays to host (dropping their sharding); the
706 # reload is single-device
707 obj = self._drop_device_info()
708 obj = jax.device_get(obj)
709 with Path(path).open('wb') as file:
710 pickle.dump(obj, file, protocol=pickle.HIGHEST_PROTOCOL)
712 @classmethod
713 def load(cls, path: str | PathLike) -> 'Bart':
714 """Load a model saved with `dump`.
716 Parameters
717 ----------
718 path
719 The file to read from.
721 Returns
722 -------
723 The deserialized model, on host memory with no device placement.
725 Raises
726 ------
727 TypeError
728 If the file does not contain a `Bart` instance.
729 """
730 with Path(path).open('rb') as file:
731 obj = pickle.load(file) # noqa: S301, the user owns the file
732 if not isinstance(obj, cls):
733 msg = f'unpickled a {type(obj).__name__}, not a {cls.__name__}'
734 raise TypeError(msg)
735 return obj
737 @property
738 def offset(self) -> Float32[Array, ''] | Float32[Array, ' k']:
739 """The prior mean of the latent mean function."""
740 return self._mcmc_state.forest.offset
742 @property
743 def n_save(self) -> int:
744 """The number of posterior samples after burn-in saved per chain."""
745 sample_axis = trace_sample_axes(self._main_trace).grow_prop_count
746 return self._main_trace.grow_prop_count.shape[sample_axis]
748 @property
749 def num_chains(self) -> int | None:
750 """The number of chains, `None` if scalar."""
751 return self._mcmc_state.num_chains()
753 @property
754 def ndpost(self) -> int:
755 """The total number of posterior samples after burn-in across all chains."""
756 return self._main_trace.grow_prop_count.size
758 @property
759 def num_trees(self) -> int:
760 """Return the number of trees used in the model."""
761 forest = self._mcmc_state.forest
762 chain_axis = chain_vmap_axes(forest).split_tree
763 # chainless split_tree is (num_trees, half_tree_size); num_trees is core axis 0
764 axis = chainful_axis(0, chain_axis)
765 return forest.split_tree.shape[axis]
767 def get_latent_prec(
768 self, only_continuous: bool = False
769 ) -> (
770 Float32[Array, ' n_burn_plus_n_save']
771 | Float32[Array, 'n_burn_plus_n_save k k']
772 | Float32[Array, 'num_chains n_burn_plus_n_save']
773 | Float32[Array, 'num_chains n_burn_plus_n_save k k']
774 ):
775 """Return the posterior samples of the latent error precision matrix.
777 Parameters
778 ----------
779 only_continuous
780 If `True` and the model has mixed binary-continuous outcomes,
781 return only the submatrix for the continuous components.
783 Returns
784 -------
785 MCMC samples of the error precision matrix.
787 Raises
788 ------
789 ValueError
790 If `only_continuous` is `True` but the model has only binary
791 outcomes, so there is no continuous submatrix to return.
793 Notes
794 -----
795 This method is meant to check for convergence, so it returns the full
796 MCMC trace and does not concatenate chains together. For probit
797 regression, this returns the precision of the latent error term, not
798 the Bernoulli precision for the binary outcome. For heteroskedastic
799 regression, the returned precision is the global precision parameter,
800 that would have to be divided by a squared error scale to get the
801 precision on a given datapoint.
802 """
803 binary_indices = self._mcmc_state.binary_indices
804 if (
805 only_continuous
806 and binary_indices is None
807 and self._mcmc_state.z is not None
808 ):
809 msg = 'Model has only binary outcomes, so there is no continuous submatrix to return.'
810 raise ValueError(msg)
812 return get_latent_prec(
813 self._burnin_trace,
814 self._main_trace,
815 binary_indices,
816 only_continuous=only_continuous,
817 )
819 def get_error_sdev(
820 self, mean: bool = False
821 ) -> (
822 Float32[Array, ' ndpost']
823 | Float32[Array, 'ndpost k']
824 | Float32[Array, '']
825 | Float32[Array, ' k']
826 ):
827 """Return the error standard deviation, post-burnin, chains concatenated.
829 Parameters
830 ----------
831 mean
832 If `True`, average the error covariance matrix across samples before
833 taking the square root, returning a single scalar or vector instead
834 of posterior samples.
836 Returns
837 -------
838 Posterior samples (or single estimate) of the error standard deviation; NaN for binary outcomes.
840 Notes
841 -----
842 Binary outcomes do have a standard deviation of course, but it's not
843 returned by this method because that would require to evaluate
844 predictions on a given X, since the Bernoulli variance is p(1-p).
845 """
846 # binary outcomes are filled with NaN, so disable the NaN check
847 with debug_nans(False):
848 return get_error_sdev(self._main_trace, self._binary_mask, mean=mean)
850 @cached_property
851 def accept(
852 self,
853 ) -> (
854 Float32[Array, ' n_burn_plus_n_save']
855 | Float32[Array, 'num_chains n_burn_plus_n_save']
856 ):
857 """Fraction of trees with an accepted grow or prune move at each iteration.
859 Includes the burn-in iterations and does not concatenate the chains, to
860 allow checking convergence. The iterations thinned away by `n_skip` are
861 not recorded.
862 """
863 return get_accept(self._burnin_trace, self._main_trace, self.num_trees)
865 @cached_property
866 def varcount(self) -> Int32[Array, 'ndpost p']:
867 """Histogram of predictor usage for decision rules in the trees."""
868 p = self._mcmc_state.forest.max_split.size
869 return varcount(p, self._main_trace)
871 @cached_property
872 def varcount_mean(self) -> Float32[Array, ' p']:
873 """Average of `varcount` across MCMC iterations."""
874 return self.varcount.mean(axis=0)
876 @cached_property
877 def varprob(self) -> Float32[Array, 'ndpost p']:
878 """Posterior samples of the probability of choosing each predictor for a decision rule."""
879 return varprob(self._mcmc_state.forest.max_split, self._main_trace)
881 @cached_property
882 def varprob_mean(self) -> Float32[Array, ' p']:
883 """The marginal posterior probability of each predictor being chosen for a decision rule."""
884 return self.varprob.mean(axis=0)
886 def _process_error_scale_test(
887 self,
888 x_test: Real[ArrayLike, 'p m'] | DataFrame | str,
889 kind: PredictKind,
890 error_scale: Float[ArrayLike, ' m']
891 | Float[ArrayLike, 'k m']
892 | Series
893 | DataFrame
894 | None,
895 ) -> Float32[Array, ' m'] | Float32[Array, 'k m'] | None:
896 """Validate and resolve the error scales for prediction.
898 Parameters
899 ----------
900 x_test
901 The raw (not yet processed) test predictors, or ``'train'``.
902 kind
903 The prediction kind.
904 error_scale
905 User-provided per-observation error scale, or `None`.
907 Returns
908 -------
909 The resolved error scale as a float32 array, or `None` if not applicable.
911 Raises
912 ------
913 ValueError
914 If `error_scale` is specified when it should be `None`, or missing
915 when required.
916 """
917 x_test_is_train = isinstance(x_test, str) and x_test == 'train'
918 train_error_scale = self._mcmc_state.error_scale
919 has_train_error_scale = train_error_scale is not None
920 is_binary = self._mcmc_state.z is not None
921 # the error scales enter the outcome samples of any outcome type, and
922 # also the mean of binary outcomes, since P(y=1) = Phi(latent / scale)
923 needs_error_scale = has_train_error_scale and (
924 kind is PredictKind.outcome_samples
925 or (is_binary and kind in (PredictKind.mean, PredictKind.mean_samples))
926 )
928 if not needs_error_scale:
929 if error_scale is not None: 929 ↛ 930line 929 didn't jump to line 930 because the condition on line 929 was never true
930 msg = (
931 '`error_scale` must be `None` in this configuration (error'
932 " scales are used with kind='outcome_samples', and with"
933 " kind='mean' or 'mean_samples' for binary outcomes, and"
934 ' only when the model was fit with `error_scale`)'
935 )
936 raise ValueError(msg)
937 return None
939 if x_test_is_train:
940 if error_scale is not None: 940 ↛ 941line 940 didn't jump to line 941 because the condition on line 940 was never true
941 msg = (
942 "`error_scale` must be `None` when x_test='train'"
943 ' (the training error scales are used automatically)'
944 )
945 raise ValueError(msg)
946 return train_error_scale
948 # new test data, model was fit with error scales
949 if error_scale is None:
950 msg = (
951 '`error_scale` is required because the model was fit with'
952 ' `error_scale` and x_test is new data'
953 )
954 raise ValueError(msg)
955 error_scale_test = _process_response_input(error_scale)
956 assert train_error_scale is not None # implied by needs_error_scale
957 # the per-observation axis is checked separately, against x_test
958 if error_scale_test.shape[:-1] != train_error_scale.shape[:-1]:
959 msg = (
960 f'`error_scale` shape mismatch with the training error scales: '
961 f'got {error_scale_test.shape=}, but the leading dimensions must '
962 f'match the training {train_error_scale.shape=} (only the '
963 f'per-observation axis may differ).'
964 )
965 raise ValueError(msg)
966 return error_scale_test
968 def _process_x_test(
969 self,
970 x_test: Real[ArrayLike, 'p m'] | DataFrame | str,
971 error_scale: Float32[Array, ' m'] | Float32[Array, 'k m'] | None,
972 ) -> UInt[Array, 'p m']:
973 """Convert x_test to binned format suitable for prediction."""
974 if isinstance(x_test, str):
975 if x_test != 'train': 975 ↛ 976line 975 didn't jump to line 976 because the condition on line 975 was never true
976 msg = (
977 f"x_test must be an array, a DataFrame, or 'train', got {x_test!r}"
978 )
979 raise ValueError(msg)
980 return self._mcmc_state.X
981 x_test, x_test_fmt = _process_predictor_input(x_test)
982 if x_test_fmt != self._x_train_fmt:
983 msg = f'Input format mismatch: {x_test_fmt=} != x_train_fmt={self._x_train_fmt!r}'
984 raise ValueError(msg)
985 if error_scale is not None:
986 _check_same_length(error_scale, x_test)
987 return self._binner.bin(x_test)
989 def _device_put_test(
990 self,
991 x_test: UInt[Array, 'p m'],
992 error_scale: Float32[Array, ' m'] | Float32[Array, 'k m'] | None,
993 ) -> tuple[UInt[Array, 'p m'], Float32[Array, ' m'] | Float32[Array, 'k m'] | None]:
994 """Place new test data on the devices of the model.
996 Mirror the placement of the training data done at fit time: shard over
997 the mesh if there is one (the observation axis over 'data'), else move
998 to the device requested explicitly at construction, if any. The inputs
999 are donated, so they must not be used elsewhere.
1000 """
1001 mesh = self._mcmc_state.config.mesh
1002 if mesh is not None:
1003 put = lambda a: device_put(
1004 a,
1005 NamedSharding(mesh, leaf_partition_spec(a.ndim, None, -1, mesh)),
1006 donate=True,
1007 )
1008 elif self._device is not None:
1009 put = lambda a: device_put(a, self._device, donate=True)
1010 else:
1011 return x_test, error_scale
1012 if error_scale is None:
1013 return put(x_test), None
1014 else:
1015 return put(x_test), put(error_scale)
1017 def _check_trees(
1018 self, error: bool = False
1019 ) -> UInt[Array, 'num_chains n_save num_trees']:
1020 """Apply `bartz.grove.check_trace` to all the tree draws.
1022 Parameters
1023 ----------
1024 error
1025 If `True`, throw an error if any invalid trees are found.
1027 Returns
1028 -------
1029 An array where non-zero entries indicate invalid trees.
1031 Raises
1032 ------
1033 RuntimeError
1034 If `error` is `True` and any invalid trees are found.
1035 """
1036 out = check_trees(self._main_trace, self._mcmc_state.forest.max_split)
1037 if error:
1038 bad_count = jnp.count_nonzero(out).item()
1039 if bad_count > 0:
1040 msg = f'Found {bad_count} invalid trees in the MCMC trace.'
1041 raise RuntimeError(msg)
1042 return out
1044 def _tree_goes_bad(self) -> Bool[Array, 'num_chains n_save num_trees']:
1045 """Find iterations where a tree becomes invalid.
1047 Returns
1048 -------
1049 An array where ``(i, j, k)`` is `True` if tree `k` is invalid at iteration `j` in chain `i` but not at iteration ``j - 1``.
1050 """
1051 return tree_goes_bad(self._main_trace, self._mcmc_state.forest.max_split)
1053 def _check_replicated_trees(self) -> None:
1054 """Check that the trees are equal across data-sharded devices.
1056 If the data is sharded across devices, verify that the trees (which
1057 should be replicated) are identical on all shards.
1059 Raises
1060 ------
1061 RuntimeError
1062 If the trees differ across devices.
1063 """
1064 state = self._mcmc_state
1065 mesh = state.config.mesh
1066 if mesh is not None and 'data' in mesh.axis_names:
1067 # drop the data-sharded `leaf_indices` (not replicated) before the
1068 # cross-shard equality check; `None` is a deliberately off-type
1069 # placeholder, so use `tree_at`, which (unlike `dataclasses.replace`)
1070 # bypasses the `__init__` type checks
1071 replicated_forest = tree_at(lambda f: f.leaf_indices, state.forest, None)
1072 equal = equal_shards(
1073 replicated_forest, 'data', in_specs=PartitionSpec(), mesh=mesh
1074 )
1075 equal_array = jnp.stack(tree.leaves(equal))
1076 all_equal = jnp.all(equal_array)
1077 if not all_equal.item(): 1077 ↛ 1078line 1077 didn't jump to line 1078 because the condition on line 1077 was never true
1078 msg = 'The trees differ across data-sharded devices.'
1079 raise RuntimeError(msg)
1081 def _compare_resid(
1082 self,
1083 ) -> tuple[
1084 Float32[Array, '*num_chains n'] | Float32[Array, '*num_chains k n'],
1085 Float32[Array, '*num_chains n'] | Float32[Array, '*num_chains k n'],
1086 ]:
1087 """Re-compute residuals to compare them with the updated ones.
1089 Returns
1090 -------
1091 resid1
1092 The final state of the residuals updated during the MCMC.
1093 resid2
1094 The residuals computed from the final state of the trees.
1095 """
1096 return compare_resid(self._mcmc_state)
1098 def _depth_distr(self) -> Int32[Array, '*num_chains n_save d']:
1099 """Histogram of tree depths for each state of the trees.
1101 Returns
1102 -------
1103 A matrix where each row contains a histogram of tree depths.
1104 """
1105 return depth_distr(self._main_trace)
1107 def _points_per_node_distr(
1108 self, node_type: Literal['leaf', 'leaf-parent']
1109 ) -> Int32[Array, '*num_chains n_save n_plus_1']:
1110 return points_per_node_distr_trace(
1111 self._mcmc_state.X, self._main_trace, node_type
1112 )
1114 def _points_per_decision_node_distr(
1115 self,
1116 ) -> Int32[Array, '*num_chains n_save n_plus_1']:
1117 """Histogram of number of points belonging to parent-of-leaf nodes.
1119 Returns
1120 -------
1121 For each chain, a matrix where each row contains a histogram of number of points.
1122 """
1123 return self._points_per_node_distr('leaf-parent')
1125 def _points_per_leaf_distr(self) -> Int32[Array, '*num_chains n_save n_plus_1']:
1126 """Histogram of number of points belonging to leaves.
1128 Returns
1129 -------
1130 A matrix where each row contains a histogram of number of points.
1131 """
1132 return self._points_per_node_distr('leaf')
1134 def _print_tree(
1135 self, i_chain: int, i_sample: int, i_tree: int, print_all: bool = False
1136 ) -> None:
1137 """Print a single tree in human-readable format.
1139 Parameters
1140 ----------
1141 i_chain
1142 The index of the MCMC chain.
1143 i_sample
1144 The index of the (post-burnin) sample in the chain.
1145 i_tree
1146 The index of the tree in the sample.
1147 print_all
1148 If `True`, also print the content of unused node slots.
1149 """
1150 trace = self._main_trace
1151 trees = _trees_chain_first(trace)
1152 index = (i_chain, i_sample, i_tree) if trace.has_chains else (i_sample, i_tree)
1153 # index the heap arrays, leaving the (unbatched) leaf scale alone; the
1154 # trailing ellipsis covers the extra `k` axis of multivariate leaves
1155 trees = tree_at(
1156 lambda t: (t.var_tree, t.split_tree, t.leaf_tree),
1157 trees,
1158 replace_fn=lambda x: x[(*index, ...)],
1159 )
1160 s = format_tree(trees, print_all=print_all)
1161 print(s) # noqa: T201, this method is intended for debug
1164def _process_predictor_input(
1165 x: Real[ArrayLike, 'p n'] | DataFrame,
1166) -> tuple[Shaped[Array, 'p n'], Any]:
1167 if isinstance(x, DataFrame):
1168 fmt = dict(kind='dataframe', columns=x.columns)
1169 x = x.to_numpy().T
1170 else:
1171 fmt = dict(kind='array', num_covar=x.shape[0])
1172 x = jnp.asarray(x)
1173 assert x.ndim == 2
1174 if jnp.issubdtype(x.dtype, jnp.floating) and not jit_active() and any_nan(x):
1175 msg = 'predictors contain NaN; missing predictor values are not supported'
1176 raise ValueError(msg)
1177 return x, fmt
1180@jit
1181def any_nan(x: Float[Array, '*shape']) -> Bool[Array, '']:
1182 # A sum of non-negative terms is nan iff a term is nan, provided the
1183 # accumulator has inf to overflow to, hence float32 (some float8 types lack
1184 # inf). Unlike isnan or min/max, this fuses with the reduction and reliably
1185 # propagates nan on cpu.
1186 return jnp.isnan(jnp.sum(jnp.abs(x), dtype=jnp.float32))
1189def _process_response_input(
1190 arr: Shaped[ArrayLike, ' n'] | Shaped[ArrayLike, 'k n'] | Series | DataFrame,
1191 /,
1192 *,
1193 dtype: DTypeLike = jnp.float32,
1194) -> Shaped[Array, ' n'] | Shaped[Array, 'k n']:
1195 if isinstance(arr, DataFrame):
1196 arr = arr.to_numpy().T
1197 elif isinstance(arr, Series):
1198 arr = arr.to_numpy()
1199 # one unconditional copy, safe to donate downstream
1200 arr = jnp.array(arr, dtype, copy=True)
1201 if arr.ndim < 1 or arr.ndim > 2: 1201 ↛ 1202line 1201 didn't jump to line 1202 because the condition on line 1201 was never true
1202 msg = f'response-like input must be 1D (n,) or 2D (k, n). Got {arr.ndim=}.'
1203 raise ValueError(msg)
1204 return arr
1207def _check_same_length(x1: Shaped[Array, '... n'], x2: Shaped[Array, '... n']) -> None:
1208 get_length = lambda x: x.shape[-1]
1209 assert get_length(x1) == get_length(x2)
1212def warn_byte_boundaries(
1213 x_train: UInt[Array, 'p n'], max_split: UInt[Array, ' p'], maxdepth: int
1214) -> None:
1215 """Warn if settings make arrays wider than one byte per element."""
1216 itemsize = max(x_train.dtype.itemsize, max_split.dtype.itemsize)
1217 if itemsize > 1:
1218 msg = (
1219 'More than 255 cutpoints per predictor make the binned predictors '
1220 f'and split points {8 * itemsize}-bit instead of 8-bit, increasing '
1221 'memory usage and running time. Consider using at most 255 cutpoints.'
1222 )
1223 warn(msg)
1224 if maxdepth > 8:
1225 msg = (
1226 f'maxdepth={maxdepth} > 8 makes the leaf indices 16-bit instead of '
1227 '8-bit, increasing memory usage and running time. Consider using '
1228 'maxdepth <= 8. Through bartz.stochtree.BARTModel, the equivalent '
1229 'is mean_forest_params=dict(max_depth=7) in sample().'
1230 )
1231 warn(msg)
1234def _check_type_settings(
1235 y_train: Float32[Array, ' n'] | Float32[Array, 'k n'],
1236 outcome_type: OutcomeType | str | Sequence[OutcomeType | str],
1237 error_scale: Float[Array, ' n'] | Float[Array, 'k n'] | None,
1238) -> tuple[OutcomeType | tuple[OutcomeType, ...], Bool[Array, ''] | Bool[Array, ' k']]:
1239 # standardize outcome_type to OutcomeType or tuple[OutcomeType, ...]
1240 if isinstance(outcome_type, Sequence) and not isinstance(outcome_type, str):
1241 outcome_type = tuple(OutcomeType(t) for t in outcome_type)
1242 num_types = len(outcome_type)
1243 if len(set(outcome_type)) == 1:
1244 outcome_type = outcome_type[0]
1245 else:
1246 num_types = None
1247 outcome_type = OutcomeType(outcome_type)
1249 # validation
1250 if num_types is not None and (y_train.ndim != 2 or num_types != y_train.shape[0]):
1251 msg = (
1252 f'Sequence outcome_type of length {num_types}'
1253 f' requires y_train.shape=({num_types}, n),'
1254 f' found {y_train.shape=}.'
1255 )
1256 raise ValueError(msg)
1257 if ( 1257 ↛ 1262line 1257 didn't jump to line 1262 because the condition on line 1257 was never true
1258 error_scale is not None
1259 and error_scale.ndim == 2
1260 and (y_train.ndim != 2 or error_scale.shape[0] != y_train.shape[0])
1261 ):
1262 msg = (
1263 f'2D error_scale (per-component scales) requires y_train of '
1264 f'shape (k, n) with matching k; got {error_scale.shape=}, '
1265 f'{y_train.shape=}.'
1266 )
1267 raise ValueError(msg)
1269 if isinstance(outcome_type, tuple):
1270 binary_mask = jnp.array([t is OutcomeType.binary for t in outcome_type])
1271 else:
1272 binary_mask = jnp.bool_(outcome_type is OutcomeType.binary)
1273 binary_mask = jnp.broadcast_to(binary_mask, y_train.shape[:-1])
1275 return outcome_type, binary_mask
1278def _process_sparsity_settings(
1279 x_train: Real[Array, 'p n'], sparse: SparseConfig
1280) -> (
1281 tuple[None, None, None, None]
1282 | tuple[FloatLike, None, None, None]
1283 | tuple[None, FloatLike, FloatLike, FloatLike]
1284):
1285 """Return (theta, a, b, rho)."""
1286 if not sparse.enabled:
1287 return None, None, None, None
1288 elif sparse.theta is not None:
1289 return sparse.theta, None, None, None
1290 else:
1291 rho = sparse.rho
1292 if rho is None:
1293 p, _ = x_train.shape
1294 rho = float(p)
1295 return None, sparse.a, sparse.b, rho
1298@jit
1299def _process_offset_settings(
1300 y_train: Float32[Array, ' n'] | Float32[Array, 'k n'],
1301 binary_mask: Bool[Array, ''] | Bool[Array, ' k'],
1302 missing: Bool[Array, ' n'] | Bool[Array, 'k n'] | None,
1303 offset: Float32[Array, ''] | Float32[Array, ' k'] | None,
1304) -> Float32[Array, ''] | Float32[Array, ' k']:
1305 """Determine and return offset."""
1306 if offset is not None:
1307 return jnp.broadcast_to(offset, y_train.shape[:-1])
1308 if y_train.shape[-1] < 1:
1309 return jnp.zeros(y_train.shape[:-1])
1311 if missing is None:
1312 *_, n_valid = y_train.shape
1313 prop = (y_train != 0).mean(-1)
1314 mean = y_train.mean(-1)
1315 else:
1316 *_, n = y_train.shape
1317 n_valid = n - jnp.count_nonzero(missing, axis=-1)
1318 safe_n = jnp.maximum(n_valid, 1)
1319 prop = jnp.where(missing, 0, y_train != 0).sum(-1) / safe_n
1320 mean = jnp.where(missing, 0.0, y_train).sum(-1) / safe_n
1322 bound = jnp.reciprocal(1.0 + n_valid)
1323 binary_offset = ndtri(jnp.clip(prop, bound, 1 - bound))
1324 offset = jnp.where(binary_mask, binary_offset, mean)
1325 return jnp.where(n_valid > 0, offset, 0.0)
1328@jit(static_argnames='num_trees')
1329def _process_leaf_variance_settings(
1330 y_train: Float32[Array, ' n'] | Float32[Array, 'k n'],
1331 binary_mask: Bool[Array, ''] | Bool[Array, ' k'],
1332 missing: Bool[Array, ' n'] | Bool[Array, 'k n'] | None,
1333 *,
1334 k: Float[Array, ''],
1335 num_trees: int,
1336 tau_num: Float[Array, ''] | None,
1337 sigma_mu_df: Float[Array, ''] | None,
1338) -> Wishart:
1339 """Return `leaf_prior_cov_inv`, with a prior iff `sigma_mu_df` is set."""
1340 # determine `tau_num` if not specified
1341 *kshape, n = y_train.shape
1342 if tau_num is None:
1343 if n < 2:
1344 continuous_tau = jnp.ones(kshape)
1345 elif missing is None:
1346 continuous_tau = (y_train.max(-1) - y_train.min(-1)) / 2
1347 else:
1348 n_valid = n - jnp.count_nonzero(missing, axis=-1)
1349 ymax = jnp.where(missing, -jnp.inf, y_train).max(-1)
1350 ymin = jnp.where(missing, jnp.inf, y_train).min(-1)
1351 continuous_tau = jnp.where(n_valid >= 2, (ymax - ymin) / 2, 1.0)
1352 tau_num = jnp.where(binary_mask, 3.0, continuous_tau)
1354 # leaf prior standard deviation
1355 sigma_mu = tau_num / (k * math.sqrt(num_trees))
1357 # leaf prior precision matrix
1358 sigma2_mu = jnp.broadcast_to(jnp.square(sigma_mu), kshape)
1359 leaf_prior_cov_inv = jnp.reciprocal(sigma2_mu)
1360 if y_train.ndim == 2:
1361 leaf_prior_cov_inv = jnp.diag(leaf_prior_cov_inv)
1363 if sigma_mu_df is None:
1364 return Wishart(nu=None, rate=None, value=leaf_prior_cov_inv)
1365 else:
1366 # each marginal variance is inverse-gamma with alpha = sigma_mu_df / 2
1367 # and beta = rate_ii / 2, so rate_ii = sigma_mu_df * sigma2_mu makes
1368 # sigma2_mu its harmonic mean for any k
1369 rate = sigma_mu_df * sigma2_mu
1370 if y_train.ndim == 2:
1371 rate = jnp.diag(rate)
1372 return Wishart.from_inv_wishart_marginal_nu(
1373 sigma_mu_df, rate, leaf_prior_cov_inv
1374 )
1377def _process_error_variance_settings(
1378 y_train: Float32[Array, ' n'] | Float32[Array, 'k n'],
1379 binary_mask: Bool[Array, ''] | Bool[Array, ' k'],
1380 missing: Bool[Array, ' n'] | Bool[Array, 'k n'] | None,
1381 *,
1382 outcome_type: OutcomeType | tuple[OutcomeType, ...],
1383 sigma_df: FloatLike,
1384 sigma_scale: FloatLike | Float[ArrayLike, ' k'] | Literal['auto'],
1385 sigma_init: FloatLike | Float[ArrayLike, ' k'] | Literal['auto'],
1386 error_scale: Float32[Array, ' n'] | Float32[Array, 'k n'] | None,
1387) -> Wishart | None:
1388 """Build the error precision prior from the user settings."""
1389 if outcome_type is OutcomeType.binary:
1390 if not isinstance(sigma_scale, str) or not isinstance(sigma_init, str):
1391 msg = (
1392 'Do not set `sigma_scale` or `sigma_init` for binary regression, '
1393 'they are ignored'
1394 )
1395 raise ValueError(msg)
1396 return None
1398 *kdims, _ = y_train.shape # () or (k,)
1399 sigma_df = jnp.asarray(sigma_df, jnp.float32)
1401 # guarded per-component variance of y_train, computed only when an 'auto'
1402 # spec needs it (this function is not jitted, so it would not be elided)
1403 if isinstance(sigma_scale, str) or isinstance(sigma_init, str):
1404 vary = _guarded_response_variance(y_train, error_scale, missing)
1405 else:
1406 vary = None
1408 # prior rate: each marginal variance is inverse-gamma with alpha =
1409 # sigma_df / 2 and beta = rate_ii / 2, so rate_ii = sigma_df * var makes var
1410 # its harmonic mean for any k
1411 rate_diag = jnp.where(
1412 binary_mask, 0.0, sigma_df * _resolve_error_variance(sigma_scale, vary, kdims)
1413 )
1415 # initial precision = 1 / var per component (1 for binary components)
1416 init_var = _resolve_error_variance(sigma_init, vary, kdims)
1417 init_diag = jnp.where(binary_mask, 1.0, jnp.reciprocal(init_var))
1419 if y_train.ndim == 2:
1420 rate, init = jnp.diag(rate_diag), jnp.diag(init_diag)
1421 else:
1422 rate, init = rate_diag, init_diag
1423 return make_error_cov_prior(sigma_df, rate, init, outcome_type, missing)
1426@jit
1427def _guarded_response_variance(
1428 y_train: Float32[Array, ' n'] | Float32[Array, 'k n'],
1429 error_scale: Float32[Array, ' n'] | Float32[Array, 'k n'] | None,
1430 missing: Bool[Array, ' n'] | Bool[Array, 'k n'] | None,
1431) -> Float32[Array, '*k']:
1432 """Per-component variance of `y_train`, used by the 'auto' error scale.
1434 A precision-weighted variance (precision ``1 / error_scale ** 2``) estimates
1435 ``sigma ** 2`` at unit error scale; `missing` entries are dropped. The
1436 variance is guarded to 1 when undefined (fewer than 2 valid points) or
1437 non-positive.
1438 """
1439 if error_scale is None and missing is None:
1440 vary = jnp.var(y_train, axis=-1)
1441 return jnp.where(vary > 0, vary, 1.0)
1442 else:
1443 prec = (
1444 jnp.ones(())
1445 if error_scale is None
1446 else jnp.reciprocal(jnp.square(error_scale))
1447 )
1448 if missing is not None:
1449 prec = jnp.where(missing, 0.0, prec)
1450 y_train = jnp.where(missing, 0.0, y_train)
1451 n_valid = jnp.count_nonzero(prec, axis=-1)
1452 wmean = jnp.sum(prec * y_train, axis=-1) / jnp.sum(prec, axis=-1)
1453 sqdev = prec * jnp.square(y_train - wmean[..., None])
1454 vary = jnp.sum(sqdev, axis=-1) / n_valid
1455 # guard on n_valid too: with a single valid point the variance is 0 in
1456 # exact arithmetic, but float rounding in wmean can leave a tiny
1457 # positive vary that would slip past the `vary > 0` guard
1458 return jnp.where((n_valid > 1) & (vary > 0), vary, 1.0)
1461def _resolve_error_variance(
1462 spec: FloatLike | Float[ArrayLike, ' k'] | Literal['auto'],
1463 vary: Float32[Array, '*k'] | None,
1464 shape: Sequence[int],
1465) -> Float32[Array, '*k']:
1466 """Per-component error variance from a scale spec ('auto' uses var(y))."""
1467 if isinstance(spec, str):
1468 if spec != 'auto':
1469 msg = f"unrecognized value {spec!r}, expected 'auto' or a number"
1470 raise ValueError(msg)
1471 assert vary is not None # computed iff some spec is 'auto'
1472 return vary
1473 else:
1474 return jnp.broadcast_to(jnp.square(jnp.asarray(spec, jnp.float32)), shape)
1477def make_error_cov_prior(
1478 marginal_nu: Float32[Array, ''],
1479 rate: Float32[Array, ''] | Float32[Array, 'k k'],
1480 value: Float32[Array, ''] | Float32[Array, 'k k'],
1481 outcome_type: OutcomeType | tuple[OutcomeType, ...],
1482 missing: Bool[Array, ' n'] | Bool[Array, 'k n'] | None,
1483) -> Wishart:
1484 """Build the error precision prior, diagonal-constrained where required.
1486 Mixed binary-continuous and partial-missing (2-D mask) regression restrict
1487 the error covariance to diagonal, so they take a `DiagWishart`; the dense
1488 cases take a `Wishart`. `init` re-checks this choice. `marginal_nu` is the
1489 degrees of freedom of each marginal variance, `value` the initial value of
1490 the precision.
1491 """
1492 if isinstance(outcome_type, tuple):
1493 binary = [t is OutcomeType.binary for t in outcome_type]
1494 is_mixed = any(binary) and not all(binary)
1495 else:
1496 is_mixed = False
1497 # a 2-D missingness mask only occurs with multivariate y (checked in `init`)
1498 partial_missing = missing is not None and missing.ndim == 2
1499 cls = DiagWishart if is_mixed or partial_missing else Wishart
1500 return cls.from_inv_wishart_marginal_nu(marginal_nu, rate, value)
1503def _setup_mcmc(
1504 mcmc_key: Key[Array, ''],
1505 x_train: Real[Array, 'p n'],
1506 y_train: Float32[Array, ' n'] | Float32[Array, 'k n'],
1507 *,
1508 outcome_type: OutcomeType | tuple[OutcomeType, ...],
1509 offset: Float32[Array, ''] | Float32[Array, ' k'],
1510 error_scale: Float[Array, ' n'] | Float[Array, 'k n'] | None,
1511 missing: Bool[Array, ' n'] | Bool[Array, 'k n'] | None,
1512 max_split: UInt[Array, ' p'],
1513 leaf_prior_cov_inv: Wishart,
1514 error_cov_inv: Wishart | None,
1515 power: FloatLike,
1516 base: FloatLike,
1517 maxdepth: int,
1518 num_trees: int,
1519 init_kw: Mapping[str, Any],
1520 rm_const: bool,
1521 sparse: SparseConfig,
1522 varprob: Float[ArrayLike, ' p'] | None,
1523 num_chains: int | None,
1524 num_chain_devices: int | None | Literal['auto'],
1525 num_data_devices: int | None,
1526 devices: Literal['cpu', 'gpu'] | Device | Sequence[Device] | None,
1527 n_burn: int,
1528) -> tuple[State, Key[Array, ''], Device | None, Literal['cpu', 'gpu'] | None]:
1529 theta, a, b, rho = _process_sparsity_settings(x_train, sparse)
1531 device_kw, device, check_platform = process_device_settings(
1532 y_train, num_chains, num_chain_devices, num_data_devices, devices
1533 )
1535 kw: dict = dict(
1536 X=x_train,
1537 y=y_train,
1538 outcome_type=outcome_type,
1539 offset=offset,
1540 error_scale=error_scale,
1541 missing=missing,
1542 max_split=max_split,
1543 num_trees=num_trees,
1544 p_nonterminal=make_p_nonterminal(maxdepth, base, power),
1545 leaf_prior_cov_inv=leaf_prior_cov_inv,
1546 error_cov_inv=error_cov_inv,
1547 min_points_per_decision_node=10,
1548 filter_splitless_vars=jnp.sum(max_split == 0).item() if rm_const else 0,
1549 log_s=process_varprob(varprob, max_split),
1550 theta=theta,
1551 a=a,
1552 b=b,
1553 rho=rho,
1554 sparse_on_at=n_burn // 2 if sparse.enabled else None,
1555 augment=sparse.augment,
1556 **device_kw,
1557 )
1559 kw = dict(kw, **init_kw)
1561 state = init(**kw)
1563 # put state and mcmc key on device if requested explicitly by the user
1564 if device is not None:
1565 mcmc_key, state = device_put((mcmc_key, state), device, donate=True)
1567 return state, mcmc_key, device, check_platform
1570def _run_mcmc(
1571 key: Key[Array, ''],
1572 mcmc_state: State,
1573 n_save: int,
1574 *,
1575 n_burn: int,
1576 n_skip: int,
1577 printevery: int | None,
1578 pbar: bool,
1579 precompute_predict_train: bool,
1580 run_mcmc_kw: Mapping,
1581 check_platform: Literal['cpu', 'gpu'] | None,
1582) -> RunMCMCResult:
1583 # fill list of callbacks
1584 callbacks = ()
1585 if printevery is not None:
1586 if pbar:
1587 callbacks += (make_tqdm_callback(mcmc_state, report_every=printevery),)
1588 else:
1589 callbacks += (
1590 make_print_callback(
1591 mcmc_state,
1592 dot_every=None if printevery == 1 else 1,
1593 report_every=printevery,
1594 ),
1595 )
1596 if check_platform is not None:
1597 callbacks += (CheckPlatformCallback(check_platform),)
1599 # prepare arguments
1600 kw: dict = dict(
1601 n_burn=n_burn,
1602 n_skip=n_skip,
1603 inner_loop_length=printevery,
1604 callback=CallbackTuple(callbacks),
1605 )
1606 if precompute_predict_train:
1607 kw = dict(**kw, main_trace_type=MainTraceWithTrainPred)
1608 kw = dict(kw, **run_mcmc_kw)
1610 return run_mcmc(key, mcmc_state, n_save, **kw)
1613@jit(static_argnames='p')
1614# this is jitted such that lax.collapse below does not create a copy
1615def varcount(p: int, trace: MainTrace) -> Int32[Array, 'ndpost p']:
1616 """Histogram of predictor usage for decision rules in the trees, squashing chains."""
1617 varcount: Int32[Array, '*chains samples p']
1618 varcount = compute_varcount(p, trace, out_chain_axis=0)
1619 return lax.collapse(varcount, 0, -1)
1622@jit(static_argnames='mean')
1623def get_error_sdev(
1624 trace: MainTrace,
1625 binary_mask: Bool[Array, ''] | Bool[Array, ' k'],
1626 *,
1627 mean: bool = False,
1628) -> (
1629 Float32[Array, ' ndpost']
1630 | Float32[Array, 'ndpost k']
1631 | Float32[Array, '']
1632 | Float32[Array, ' k']
1633):
1634 """Error standard deviation, post-burnin, chains concatenated."""
1635 prec = trace.error_cov_inv
1636 if trace.has_chains:
1637 # shape (chains, samples) or (chains, samples, k, k), concatenate chains
1638 prec = chain_to_axis(prec, chain_vmap_axes(trace).error_cov_inv)
1639 prec = lax.collapse(prec, 0, 2)
1640 is_uv = prec.ndim == 1
1641 if is_uv:
1642 # univariate case, reshape to 1x1 matrix
1643 prec = prec[..., None, None]
1645 # invert precision to covariance, then take diagonal variance
1646 cov = inv_via_chol_with_gersh(prec)
1647 var = jnp.diagonal(cov, axis1=-2, axis2=-1)
1648 if mean:
1649 var = var.mean(0)
1650 sdev = jnp.sqrt(var)
1651 if is_uv:
1652 sdev = sdev.squeeze(-1)
1653 return jnp.where(binary_mask, jnp.nan, sdev)
1656@jit(static_argnames='only_continuous')
1657def get_latent_prec(
1658 burnin_trace: BurninTrace,
1659 main_trace: MainTrace,
1660 binary_indices: Int32[Array, ' kb'] | None,
1661 *,
1662 only_continuous: bool = False,
1663) -> (
1664 Float32[Array, ' n_burn_plus_n_save']
1665 | Float32[Array, 'n_burn_plus_n_save k k']
1666 | Float32[Array, 'num_chains n_burn_plus_n_save']
1667 | Float32[Array, 'num_chains n_burn_plus_n_save k k']
1668):
1669 """Latent error precision trace, burn-in + main concatenated."""
1670 burnin = burnin_trace.error_cov_inv
1671 main = main_trace.error_cov_inv
1672 sample_axis = trace_sample_axes(main_trace).error_cov_inv
1673 prec = jnp.concatenate([burnin, main], axis=sample_axis)
1674 prec = chain_to_axis(prec, chain_vmap_axes(main_trace).error_cov_inv)
1675 if only_continuous and binary_indices is not None:
1676 *_, k, _ = prec.shape
1677 kc = k - binary_indices.size
1678 mask = jnp.ones(k, dtype=bool).at[binary_indices].set(False)
1679 (cont_indices,) = jnp.nonzero(mask, size=kc)
1680 prec = prec[..., cont_indices[:, None], cont_indices[None, :]]
1681 return prec
1684@jit(static_argnames='num_trees')
1685def get_accept(
1686 burnin_trace: BurninTrace, main_trace: MainTrace, num_trees: int
1687) -> (
1688 Float32[Array, ' n_burn_plus_n_save']
1689 | Float32[Array, 'num_chains n_burn_plus_n_save']
1690):
1691 """Fraction of trees with an accepted move, burn-in + main concatenated."""
1692 sample_axis = trace_sample_axes(main_trace).grow_acc_count
1693 acc = jnp.concatenate(
1694 [
1695 burnin_trace.grow_acc_count + burnin_trace.prune_acc_count,
1696 main_trace.grow_acc_count + main_trace.prune_acc_count,
1697 ],
1698 axis=sample_axis,
1699 )
1700 acc = chain_to_axis(acc, chain_vmap_axes(main_trace).grow_acc_count)
1701 return acc / num_trees
1704@jit
1705def varprob(
1706 max_split: UInt[Array, ' p'], trace: MainTrace
1707) -> Float32[Array, 'ndpost p']:
1708 """Posterior samples of predictor selection probability, chains concatenated."""
1709 p = max_split.size
1710 varprob = trace.varprob
1711 if varprob is None:
1712 ndpost = trace.grow_prop_count.size
1713 peff = jnp.count_nonzero(max_split)
1714 out = jnp.where(max_split, 1 / peff, 0)
1715 return jnp.broadcast_to(out, (ndpost, p))
1716 varprob = chain_to_axis(varprob, chain_vmap_axes(trace).varprob)
1717 return varprob.reshape(-1, p)
1720def _trees_chain_first(obj: TreeHeaps) -> TreesTrace:
1721 """Extract `obj`'s heap arrays, moving any chain axis to the front.
1723 Returns a `TreesTrace` whose leading axis is the chain axis when `obj`
1724 carries one, and the bare per-object heap arrays otherwise.
1725 """
1726 trees = project(TreesTrace, obj)
1727 if get_has_chains(obj):
1728 axes = trees.axes_from_dataclass(chain_vmap_axes(obj))
1729 # WORKAROUND(python<3.14): use operator.is_none
1730 trees = tree.map(chain_to_axis, trees, axes, is_leaf=lambda x: x is None)
1731 return trees
1734@jit
1735def check_trees(
1736 trace: MainTrace, max_split: UInt[Array, ' p']
1737) -> UInt[Array, 'num_chains n_save num_trees']:
1738 """Apply `bartz.grove.check_trace` to all the tree draws."""
1739 trees = _trees_chain_first(trace)
1740 out: UInt[Array, '*chains samples num_trees']
1741 out = check_trace(trees, max_split)
1742 if out.ndim < 3:
1743 out = out[None, :, :]
1744 return out
1747@jit
1748def tree_goes_bad(
1749 trace: MainTrace, max_split: UInt[Array, ' p']
1750) -> Bool[Array, 'num_chains n_save num_trees']:
1751 """Find iterations where a tree becomes invalid."""
1752 bad = check_trees(trace, max_split).astype(bool)
1753 bad_before = jnp.pad(bad[:, :-1, :], [(0, 0), (1, 0), (0, 0)])
1754 return bad & ~bad_before
1757@jit
1758def compare_resid(
1759 state: State,
1760) -> tuple[
1761 Float32[Array, '*num_chains n'] | Float32[Array, '*num_chains k n'],
1762 Float32[Array, '*num_chains n'] | Float32[Array, '*num_chains k n'],
1763]:
1764 """Re-compute residuals to compare them with the updated ones."""
1765 chain_axes = chain_vmap_axes(state)
1766 resid1 = chain_to_axis(state.resid * state.resid_unit[..., None], chain_axes.resid)
1767 z = chain_to_axis(state.z, chain_axes.z) if state.z is not None else None
1769 forests = _trees_chain_first(state.forest)
1770 trees = evaluate_forest(state.X, forests, sum_batch_axis=-1)
1772 if state.binary_indices is not None:
1773 # mixed binary-continuous: z has only binary rows, y has all rows
1774 assert z is not None
1775 ref = jnp.broadcast_to(state.y, resid1.shape)
1776 ref = ref.at[..., state.binary_indices, :].set(z)
1777 elif z is not None:
1778 ref = z
1779 else:
1780 ref = state.y
1781 resid2 = ref - (trees + state.forest.offset[..., None])
1783 return resid1, resid2
1786@jit
1787def depth_distr(trace: MainTrace) -> Int32[Array, '*num_chains n_save d']:
1788 """Histogram of tree depths for each state of the trees."""
1789 split_tree = chain_to_axis(trace.split_tree, chain_vmap_axes(trace).split_tree)
1790 out: Int32[Array, '*chains samples d']
1791 out = forest_depth_distr(split_tree)
1792 if out.ndim < 3: 1792 ↛ 1794line 1792 didn't jump to line 1794 because the condition on line 1792 was always true
1793 out = out[None, :, :]
1794 return out
1797@jit(static_argnames='node_type')
1798def points_per_node_distr_trace(
1799 X: UInt[Array, 'p n'], trace: MainTrace, node_type: Literal['leaf', 'leaf-parent']
1800) -> Int32[Array, '*num_chains n_save n+1']:
1801 """Histogram of number of points per node, for every tree draw in the trace."""
1802 chain_axes = chain_vmap_axes(trace)
1803 var_tree = chain_to_axis(trace.var_tree, chain_axes.var_tree)
1804 split_tree = chain_to_axis(trace.split_tree, chain_axes.split_tree)
1805 out: Int32[Array, '*chains samples n+1']
1806 out = points_per_node_distr(X, var_tree, split_tree, node_type, sum_batch_axis=-1)
1807 if out.ndim < 3:
1808 out = out[None, :, :]
1809 return out
1812class DeviceKwArgs(TypedDict):
1813 num_chains: int | None
1814 mesh: Mesh | None
1815 leaf_dtype: DTypeLike
1816 prec_scale_dtype: DTypeLike
1819def process_device_settings(
1820 y_train: Float32[Array, ' n'] | Float32[Array, 'k n'],
1821 num_chains: int | None,
1822 num_chain_devices: int | None | Literal['auto'],
1823 num_data_devices: int | None,
1824 devices: Literal['cpu', 'gpu'] | Device | Sequence[Device] | None,
1825) -> tuple[DeviceKwArgs, Device | None, Literal['cpu', 'gpu'] | None]:
1826 """Return the arguments for `mcmcstep.init` related to devices, an optional device where to put the state, and the platform to check at runtime iff it was deduced."""
1827 # whether the user pinned a concrete pool of devices (vs. inheriting all of
1828 # the platform's devices); the auto chain sharding may not exceed that pool
1829 explicit_devices = devices is not None and not isinstance(devices, str)
1831 # the platform is deduced (rather than fixed by the user) only when neither a
1832 # device pool nor a platform string is passed
1833 platform_deduced = devices is None
1835 platform, device, devices = _determine_devices(y_train, devices)
1836 check_platform = platform if platform_deduced else None
1837 num_chain_devices = _determine_num_chain_devices(
1838 platform,
1839 num_chains,
1840 num_chain_devices,
1841 num_data_devices=num_data_devices,
1842 num_devices=len(devices),
1843 explicit_devices=explicit_devices,
1844 )
1845 mesh, device = _determine_mesh(num_chain_devices, num_data_devices, device, devices)
1847 # prepare arguments to `init`
1848 dtype = jnp.float16 if platform == 'gpu' else jnp.float32
1849 settings = DeviceKwArgs(
1850 num_chains=num_chains, mesh=mesh, leaf_dtype=dtype, prec_scale_dtype=dtype
1851 )
1853 return settings, device, check_platform
1856def _determine_devices(
1857 y_train: Float32[Array, ' n'] | Float32[Array, 'k n'],
1858 devices: Literal['cpu', 'gpu'] | Device | Sequence[Device] | None,
1859) -> tuple[Literal['cpu', 'gpu'], Device | None, Sequence[Device]]:
1860 """Determine the target platform and set of devices for the MCMC, and possibly a single target device."""
1861 if isinstance(devices, str):
1862 platform: Literal['cpu', 'gpu'] = devices # ty:ignore[invalid-assignment]
1863 devices = jax.devices(platform)
1864 return platform, devices[0], devices
1865 elif devices is not None:
1866 if not hasattr(devices, '__len__'):
1867 devices = (devices,)
1868 device = devices[0]
1869 return device.platform, device, devices
1870 elif hasattr(y_train, 'platform'): 1870 ↛ 1877line 1870 didn't jump to line 1877 because the condition on line 1870 was always true
1871 # set device=None because if the devices were not specified explicitly
1872 # we may be in the case where computation will follow data placement,
1873 # do not disturb jax as the user may be playing with vmap, jit, reshard...
1874 platform = cast(Literal['cpu', 'gpu'], y_train.platform())
1875 return platform, None, jax.devices(platform)
1876 else:
1877 msg = 'not possible to infer device from `y_train`, please set `devices`'
1878 raise ValueError(msg)
1881def _largest_divisor_at_most(n: int, cap: int) -> int:
1882 """Return the largest divisor of `n` in [1, cap]."""
1883 for d in range(cap, 0, -1): 1883 ↛ 1886line 1883 didn't jump to line 1886 because the loop on line 1883 didn't complete
1884 if n % d == 0:
1885 return d
1886 return 1 # unreachable: 1 always divides n
1889def _determine_num_chain_devices(
1890 platform: str,
1891 num_chains: int | None,
1892 num_chain_devices: int | None | Literal['auto'],
1893 *,
1894 num_data_devices: int | None,
1895 num_devices: int,
1896 explicit_devices: bool,
1897) -> int | None:
1898 """Resolve and validate `num_chain_devices`, returning the chain mesh axis size or `None`."""
1899 if num_chain_devices == 'auto':
1900 num_chain_devices = _auto_num_chain_devices(
1901 platform, num_chains, num_data_devices, num_devices, explicit_devices
1902 )
1904 # an explicit value must be a positive divisor of the number of chains
1905 if num_chain_devices is not None:
1906 effective_chains = 1 if num_chains is None else num_chains
1907 if num_chain_devices < 1 or effective_chains % num_chain_devices:
1908 chains_desc = (
1909 'a single chain (num_chains=None)'
1910 if num_chains is None
1911 else f'num_chains={num_chains}'
1912 )
1913 msg = (
1914 f'num_chain_devices={num_chain_devices} must be a positive '
1915 f'divisor of the number of chains ({chains_desc})'
1916 )
1917 raise ValueError(msg)
1919 # there is no chain axis to shard when the chains are scalar
1920 if num_chains is None:
1921 return None
1922 return num_chain_devices
1925def _auto_num_chain_devices(
1926 platform: str,
1927 num_chains: int | None,
1928 num_data_devices: int | None,
1929 num_devices: int,
1930 explicit_devices: bool,
1931) -> int | None:
1932 """Pick `num_chain_devices` automatically for multi-chain cpu runs.
1934 `num_data_devices` reserves devices for the data axis, so the chain axis can
1935 only use a fraction of them; this keeps the ``chains x data`` mesh within the
1936 `num_devices` available devices.
1937 """
1938 if num_chains is None or num_chains == 1 or platform != 'cpu':
1939 return None
1940 data_devices = num_data_devices or 1
1941 num_cores = cpu_count()
1942 assert num_cores is not None, 'could not determine number of cpu cores'
1944 # devices available for the chain axis after reserving for the data axis
1945 core_budget = max(1, num_cores // data_devices)
1946 num_shards = _largest_divisor_at_most(num_chains, core_budget)
1948 if num_shards > 1:
1949 # the mesh draws from `num_devices` devices, whether those are all the
1950 # platform's devices or an explicit subset passed by the user
1951 device_budget = max(1, num_devices // data_devices)
1952 if device_budget < num_shards:
1953 new_num_shards = _largest_divisor_at_most(num_chains, device_budget)
1954 warn(
1955 _auto_chain_devices_warning(
1956 num_chains=num_chains,
1957 desired=num_shards,
1958 actual=new_num_shards,
1959 device_budget=device_budget,
1960 num_devices=num_devices,
1961 num_data_devices=num_data_devices,
1962 explicit_devices=explicit_devices,
1963 )
1964 )
1965 num_shards = new_num_shards
1967 return num_shards if num_shards > 1 else None
1970def _auto_chain_devices_warning(
1971 *,
1972 num_chains: int,
1973 desired: int,
1974 actual: int,
1975 device_budget: int,
1976 num_devices: int,
1977 num_data_devices: int | None,
1978 explicit_devices: bool,
1979) -> str:
1980 """Compose the warning shown when auto chain sharding is capped by the device count."""
1981 if explicit_devices:
1982 pool = f'the {num_devices} devices passed in `devices`'
1983 few = f'only {num_devices} devices were passed in `devices`'
1984 advice = ''
1985 else:
1986 pool = f'the {num_devices} jax cpu devices'
1987 few = f'jax is set up with only {num_devices} cpu devices'
1988 advice = (
1989 ' To enable more parallelization, increase the limit with '
1990 '`jax.config.update("jax_num_cpu_devices", <num_devices>)`.'
1991 )
1992 if num_data_devices:
1993 limit = (
1994 f'only {device_budget} of {pool} are free for chains '
1995 f'(num_data_devices={num_data_devices} reserves the rest)'
1996 )
1997 else:
1998 limit = few
1999 return (
2000 f'`Bart` would like to shard {num_chains} chains across {desired} '
2001 f'devices, but {limit}, so it will use {actual} devices for chains '
2002 f'instead.{advice}'
2003 )
2006def _determine_mesh(
2007 num_chain_devices: int | None,
2008 num_data_devices: int | None,
2009 device: Device | None,
2010 devices: Sequence[Device],
2011) -> tuple[Mesh, None] | tuple[None, Device | None]:
2012 """Create a jax device mesh for `mcmcstep.init()`."""
2013 if num_chain_devices is None and num_data_devices is None:
2014 return None, device
2015 else:
2016 mesh = {}
2017 if num_chain_devices is not None:
2018 mesh.update(chains=num_chain_devices)
2019 if num_data_devices is not None:
2020 mesh.update(data=num_data_devices)
2021 mesh = make_mesh(
2022 tuple(mesh.values()),
2023 tuple(mesh),
2024 axis_types=(AxisType.Auto,) * len(mesh),
2025 devices=devices,
2026 )
2027 return mesh, None
2028 # set device=None because `mcmcstep.init` will `device_put` with the
2029 # mesh already, we don't want to undo its work
2032def process_varprob(
2033 varprob: Float[ArrayLike, ' p'] | None, max_split: UInt[Array, ' p']
2034) -> Float32[Array, ' p'] | None:
2035 """Convert varprob to log_s."""
2036 if varprob is None:
2037 return None
2038 varprob = jnp.asarray(varprob)
2039 assert varprob.shape == max_split.shape, 'varprob must have shape (p,)'
2040 varprob = error_if(varprob, varprob <= 0, 'varprob must be > 0')
2041 return jnp.log(varprob)
2044def predict_latent(
2045 x: UInt[Array, 'p m'],
2046 trace: MainTrace,
2047 test_points: Literal['none', 'autobatch', 'shard_and_autobatch'] = 'none',
2048) -> Float32[Array, 'ndpost m'] | Float32[Array, 'ndpost k m']:
2049 """Evaluate trees on already quantized `x`, and squash chains."""
2050 return evaluate_trace(x, trace, flatten_chains=True, test_points=test_points)
2053@jit(static_argnames=('has_binary', 'kind', 'test_points'))
2054def predict(
2055 key: Key[Array, ''] | None,
2056 trace: MainTrace,
2057 x_test: UInt[Array, 'p m'],
2058 *,
2059 error_scale: Float[Array, ' m'] | Float[Array, 'k m'] | None,
2060 binary_indices: Int32[Array, ' kb'] | None,
2061 has_binary: bool,
2062 kind: PredictKind | str,
2063 test_points: Literal['none', 'autobatch', 'shard_and_autobatch'],
2064) -> (
2065 Float32[Array, ' m']
2066 | Float32[Array, 'k m']
2067 | Float32[Array, 'ndpost m']
2068 | Float32[Array, 'ndpost k m']
2069):
2070 """Implement `Bart.predict` by evaluating the trees on `x_test`."""
2071 # get latent i.e. bare sum-of-trees predictions
2072 latent = predict_latent(x_test, trace, test_points)
2073 return _predict_from_latent(
2074 key,
2075 trace,
2076 latent,
2077 error_scale=error_scale,
2078 binary_indices=binary_indices,
2079 has_binary=has_binary,
2080 kind=kind,
2081 )
2084@jit(static_argnames=('has_binary', 'kind'))
2085def predict_train(
2086 key: Key[Array, ''] | None,
2087 trace: MainTraceWithTrainPred,
2088 *,
2089 error_scale: Float[Array, ' n'] | Float[Array, 'k n'] | None,
2090 binary_indices: Int32[Array, ' kb'] | None,
2091 has_binary: bool,
2092 kind: PredictKind | str,
2093) -> (
2094 Float32[Array, ' n']
2095 | Float32[Array, 'k n']
2096 | Float32[Array, 'ndpost n']
2097 | Float32[Array, 'ndpost k n']
2098):
2099 """Implement `Bart.predict('train')` from the precomputed predictions."""
2100 latent = _flatten_chain_sample(
2101 trace.train_pred,
2102 chain_vmap_axes(trace).train_pred,
2103 trace_sample_axes(trace).train_pred,
2104 )
2105 return _predict_from_latent(
2106 key,
2107 trace,
2108 latent,
2109 error_scale=error_scale,
2110 binary_indices=binary_indices,
2111 has_binary=has_binary,
2112 kind=kind,
2113 )
2116def _predict_from_latent(
2117 key: Key[Array, ''] | None,
2118 trace: MainTrace,
2119 latent: Float32[Array, 'ndpost m'] | Float32[Array, 'ndpost k m'],
2120 *,
2121 error_scale: Float[Array, ' m'] | Float[Array, 'k m'] | None,
2122 binary_indices: Int32[Array, ' kb'] | None,
2123 has_binary: bool,
2124 kind: PredictKind | str,
2125) -> (
2126 Float32[Array, ' m']
2127 | Float32[Array, 'k m']
2128 | Float32[Array, 'ndpost m']
2129 | Float32[Array, 'ndpost k m']
2130):
2131 """Turn the sum-of-trees `latent` samples into the requested prediction kind."""
2132 if kind is PredictKind.latent_samples:
2133 return latent
2135 # sample posterior (uses latent directly, no probit squash needed)
2136 if kind is PredictKind.outcome_samples:
2137 assert key is not None
2138 return sample_outcome(
2139 key,
2140 trace,
2141 latent,
2142 error_scale=error_scale,
2143 binary_indices=binary_indices,
2144 has_binary=has_binary,
2145 )
2147 # squash predictions to (0, 1) if probit; with heteroskedastic error scales
2148 # P(y=1) = Phi(latent / error_scale)
2149 if has_binary: # self._mcmc_state.z is not None
2150 # error_scale is (m,) or (k, m), so it broadcasts against latent
2151 arg = latent if error_scale is None else latent / error_scale
2152 if binary_indices is not None:
2153 # mixed: squash only the binary rows, leaving continuous rows as-is
2154 indexing = jnp.s_[..., binary_indices, :]
2155 mean_samples = latent.at[indexing].set(ndtr(arg[indexing]))
2156 else:
2157 mean_samples = ndtr(arg)
2158 else:
2159 mean_samples = latent
2161 # take mean or return samples
2162 if kind is PredictKind.mean:
2163 return mean_samples.mean(axis=0)
2164 return mean_samples
2167def _flatten_chain_sample(
2168 arr: Float[Array, '*shape'], chain_axis: int | None, sample_axis: int
2169) -> Float[Array, '*flat_shape']:
2170 """Fold the chain axis into the sample axis, matching `predict_latent`'s layout."""
2171 if chain_axis is None:
2172 return arr
2173 arr = jnp.moveaxis(arr, (chain_axis, sample_axis), (0, 1))
2174 return lax.collapse(arr, 0, 2)
2177@jit(static_argnames='has_binary')
2178def sample_outcome(
2179 key: Key[Array, ''],
2180 trace: MainTrace,
2181 latent: Float32[Array, 'ndpost m'] | Float32[Array, 'ndpost k m'],
2182 *,
2183 error_scale: Float32[Array, ' m'] | Float32[Array, 'k m'] | None,
2184 binary_indices: Int32[Array, ' kb'] | None,
2185 has_binary: bool,
2186) -> Float32[Array, 'ndpost m'] | Float32[Array, 'ndpost k m']:
2187 """Sample from the posterior predictive distribution."""
2188 # move error_cov_inv chain axis to 0
2189 prec = chain_to_axis(trace.error_cov_inv, chain_vmap_axes(trace).error_cov_inv)
2191 if latent.ndim > 2: # multivariate case
2192 error_cov_inv = lax.collapse(prec, 0, -2)
2194 # Cholesky of precision: error_cov_inv = L @ L^T
2195 L = chol_with_gersh(error_cov_inv) # (ndpost, k, k)
2197 # Sample z ~ N(0, I) and solve L^T @ error = z
2198 # so error = L^{-T} z ~ N(0, L^{-T} L^{-1}) = N(0, Sigma)
2199 z = random.normal(key, latent.shape) # (ndpost, k, m)
2200 error = solve_triangular(L, z, trans='T', lower=True) # (ndpost, k, m)
2201 if error_scale is not None: 2201 ↛ 2214line 2201 didn't jump to line 2214 because the condition on line 2201 was always true
2202 # error_scale is (m,) or (k, m) so it always broadcasts right
2203 error *= error_scale
2204 else: # univariate
2205 # pure binary probit has unit-scale latent error; continuous scales it
2206 # by `sigma`. Either way, optionally rescaled per datapoint by error_scale.
2207 error = random.normal(key, latent.shape)
2208 if not has_binary:
2209 sigma = jnp.sqrt(jnp.reciprocal(prec)).reshape(-1)
2210 error *= sigma[..., None]
2211 if error_scale is not None: 2211 ↛ 2214line 2211 didn't jump to line 2214 because the condition on line 2211 was always true
2212 error *= error_scale[None, :]
2214 outcome = latent + error
2216 # convert binary outcomes via latent probit thresholding
2217 if binary_indices is not None:
2218 idx = jnp.s_[..., binary_indices, :]
2219 outcome = outcome.at[idx].set(jnp.where(outcome[idx] > 0, 1.0, 0.0))
2220 elif has_binary:
2221 outcome = jnp.where(outcome > 0, 1.0, 0.0)
2223 return outcome