Coverage for src/bartz/_interface.py: 96%
657 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 11:03 +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, 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.
211 y_train
212 The training responses. For univariate regression, a 1D array of shape
213 `(n,)`. For multivariate regression, a 2D array of shape `(k, n)` where
214 `k` is the number of response components, as introduced in [3]_. For
215 binary regression, the convention is that non-zero values mean 1, zero
216 mean 0, like booleans.
217 outcome_type
218 The type of regression. ``'continuous'`` for continuous regression,
219 ``'binary'`` for binary regression with probit link. For multivariate
220 regression, a scalar value applies to all components; alternatively, a
221 sequence of per-component types (e.g., ``['binary', 'continuous']``)
222 specifies mixed outcome types. Binary components in multivariate
223 outcomes follow the multivariate probit BART formulation of [4]_.
224 sparse
225 A `SparseConfig` for the sparsity-inducing variable selection prior of
226 [1]_. Disabled by default; pass a `SparseConfig` to enable it.
227 varprob
228 The probability distribution over the `p` predictors for choosing a
229 predictor to split on in a decision node a priori. Must be > 0. It does
230 not need to be normalized to sum to 1. If not specified, use a uniform
231 distribution. If `sparse` is enabled, this is used as initial value for
232 the MCMC.
233 binner
234 A callable that, given the training predictors and a random key,
235 returns a `~bartz.prepcovars.Binner` instance. The default is
236 `~bartz.prepcovars.UniqueQuantileBinner`, which places cutpoints at
237 the quantiles of each predictor. Other built-in options are
238 `~bartz.prepcovars.RangeEvenBinner` (evenly-spaced cutpoints over the
239 observed range) and `~bartz.prepcovars.GivenSplitsBinner` (R BART
240 ``xinfo`` format). To pass options, use `functools.partial`, e.g.
241 ``binner=partial(UniqueQuantileBinner, max_bins=128)``.
242 rm_const
243 How to treat predictors with no associated decision rules (i.e., there
244 are no available cutpoints for that predictor). If `True` (default),
245 they are ignored. If `False`, an error is raised if there are any.
246 sigma_df
247 The degrees of freedom of the prior on the error precision. For
248 multivariate regression with `k` components, the Wishart degrees of
249 freedom are set to ``sigma_df + k - 1``.
250 sigma_scale
251 Sets the scale of the prior on the error precision. If 'auto'
252 (default), the prior is scaled so that the error precision equals
253 ``diag(1 / var(y_train))`` in expectation, where with `error_scale` the
254 variance is a precision-weighted one that estimates the error variance
255 at unit error scale. Otherwise, ``square(sigma_scale)`` is the prior
256 harmonic mean of the error variance; for multivariate regression a
257 scalar is broadcast to all components. For mixed outcome types, binary
258 components are ignored.
259 sigma_init
260 The initial value of the error standard deviation in the MCMC. If
261 'auto' (default), the initial error precision is set to ``diag(1 /
262 var(y_train))``, with the same precision-weighted variance as
263 `sigma_scale` when `error_scale` is given. Otherwise, the initial
264 precision is ``diag(1 / square(sigma_init))``; for multivariate
265 regression a scalar is broadcast to all components. For mixed outcome
266 types, binary components are ignored.
267 k
268 The inverse scale of the prior standard deviation on the latent mean
269 function, relative to half the observed range of `y_train`. If `y_train`
270 has less than two elements, `k` is ignored and the scale is set to 1.
271 power
272 base
273 Parameters of the prior on tree node generation. The probability that a
274 node at depth `d` (0-based) is non-terminal is ``base / (1 + d) **
275 power``.
276 tau_num
277 The numerator in the expression that determines the prior standard
278 deviation of leaves. If not specified, default to ``(max(y_train) -
279 min(y_train)) / 2`` (or 1 if `y_train` has less than two elements) for
280 continuous regression, and 3 for binary regression. For multivariate
281 regression, the range is computed per component. For mixed outcome
282 types, each component uses the default for its type.
283 offset
284 The prior mean of the latent mean function. If not specified, it is set
285 to the mean of `y_train` for continuous regression, and to
286 ``Phi^-1(mean(y_train != 0))`` for binary regression. If `y_train` is
287 empty, `offset` is set to 0. With binary regression, if `y_train` is
288 all zero or all non-zero, it is set to ``Phi^-1(1/(n+1))`` or
289 ``Phi^-1(n/(n+1))``, respectively. For multivariate regression, can be
290 a scalar (broadcast to all components) or a `(k,)` vector. If not
291 specified, it is set to the per-component mean of `y_train`. For mixed
292 outcome types, each component uses the default for its type.
293 error_scale
294 Coefficients that rescale the error standard deviation on each
295 datapoint ("w" in BART3). Not specifying `error_scale` is equivalent to
296 setting it to 1 for all datapoints. Shape ``(n,)`` applies the same
297 scalar scale to every outcome component; for multivariate regression,
298 ``(k, n)`` instead supplies a per-component scale per datapoint.
299 Supported with binary (probit) outcomes, where it scales the latent
300 error so the success probability is ``Phi(latent / error_scale)``,
301 including the binary components of a mixed regression.
302 missing
303 Boolean mask with the same shape as `y_train`; `True` marks entries
304 to be ignored by the MCMC. The values of `y_train` at masked
305 positions are ignored and may be anything, even non-finite. If 2-D,
306 the error covariance must be diagonal.
307 num_trees
308 The number of trees used to represent the latent mean function.
309 n_save
310 The number of MCMC samples to save, after burn-in, per chain. The
311 total trace length across all chains is ``num_chains * n_save``.
312 n_burn
313 The number of initial MCMC samples to discard as burn-in. This number
314 of samples is discarded from each chain.
315 n_skip
316 The thinning factor for the MCMC samples, after burn-in.
317 printevery
318 The number of iterations (including thinned-away ones) between each log
319 line. Set to `None` to disable progress reporting entirely (this ignores
320 `pbar`). ^C interrupts the MCMC only every `printevery` iterations, so
321 with reporting disabled it's impossible to kill the MCMC conveniently.
322 pbar
323 If `True`, show a `tqdm` progress bar instead of printing log lines. The
324 bar advances every iteration and refreshes the acceptance statistics
325 every `printevery` iterations. Ignored if `printevery` is `None`.
326 num_chains
327 The number of independent Markov chains to run.
329 The difference between ``num_chains=None`` and ``num_chains=1`` is that
330 in the latter case in the object attributes and some methods there will
331 be an explicit chain axis of size 1.
332 num_chain_devices
333 The number of devices to spread the chains across. Must be a divisor of
334 `num_chains`. Each device will run a fraction of the chains. If 'auto'
335 (default) and running on cpu, the number of devices is picked
336 automatically based on the number of cores and the number of available
337 devices (all the virtual jax cpu devices, or the `devices` list if set).
338 num_data_devices
339 The number of devices to split datapoints across. Must be a divisor of
340 `n`. This is useful only with very high `n`, about > 1000_000. `predict`
341 parallelizes across the same devices, splitting the test points; the
342 number of test points must be a multiple of `num_data_devices` as well.
344 If both num_chain_devices and num_data_devices are specified, the total
345 number of devices used is the product of the two.
346 devices
347 One or more devices used to run the MCMC on. If not specified, the
348 computation will follow the placement of the input arrays. If a list of
349 devices, this argument can be longer than the number of devices needed.
350 seed
351 The seed for the random number generator.
352 maxdepth
353 The maximum depth of the trees. This is 1-based, so with the default
354 ``maxdepth=6``, the depths of the levels range from 0 to 5.
355 precompute_predict_train
356 If `True`, compute the predictions at the training points during the
357 MCMC. Off by default; makes ``predict('train', ...)`` faster at the cost
358 of more memory.
359 init_kw
360 Additional arguments passed to `bartz.mcmcstep.init`.
361 run_mcmc_kw
362 Additional arguments passed to `bartz.mcmcloop.run_mcmc`.
364 Notes
365 -----
366 On gpu, the leaves of the trees are stored in float16 (see the
367 ``leaf_dtype`` argument of `bartz.mcmcstep.init`), which limits the
368 signal-to-noise ratio of the fit to less than about 1000, a limit that
369 should always hold in realistic usage.
371 References
372 ----------
373 .. [1] Linero, Antonio R. (2018). “Bayesian Regression Trees for
374 High-Dimensional Prediction and Variable Selection”. In: Journal of the
375 American Statistical Association 113.522, pp. 626-636.
376 .. [2] Hugh A. Chipman, Edward I. George, Robert E. McCulloch "BART:
377 Bayesian additive regression trees," The Annals of Applied Statistics,
378 Ann. Appl. Stat. 4(1), 266-298, (March 2010).
379 .. [3] Um, Seungha, Antonio R. Linero, Debajyoti Sinha, and Dipankar
380 Bandyopadhyay (2023). "Bayesian additive regression trees for
381 multivariate skewed responses". In: Statistics in Medicine 42.3,
382 pp. 246-263.
383 .. [4] Goh, Yong Chen, Wuu Kuang Soh, Andrew C. Parnell, and Keefe
384 Murphy (2024). "Joint Models for Handling Non-Ignorable Missing
385 Data using Bayesian Additive Regression Trees: Application to
386 Leaf Photosynthetic Traits Data". arXiv:2412.14946 [stat.ME].
388 """
390 _main_trace: MainTrace
391 _burnin_trace: BurninTrace
392 _mcmc_state: State
393 _binner: Binner
394 _binary_mask: Bool[Array, ''] | Bool[Array, ' k']
395 # WORKAROUND(jax<0.9.1): use `jax.tree.static` instead of `field(static=True)`
396 _x_train_fmt: Any = field(static=True)
397 _device: Device | None = field(static=True)
399 def __init__(
400 self,
401 x_train: Real[ArrayLike, 'p n'] | DataFrame,
402 y_train: Float32[ArrayLike, ' n']
403 | Float32[ArrayLike, 'k n']
404 | Series
405 | DataFrame,
406 *,
407 outcome_type: OutcomeType | str | Sequence[OutcomeType | str] = 'continuous',
408 sparse: SparseConfig = SparseConfig(enabled=False),
409 varprob: Float[ArrayLike, ' p'] | None = None,
410 binner: BinnerFactory = UniqueQuantileBinner,
411 rm_const: bool = True,
412 sigma_df: FloatLike = 3.0,
413 sigma_scale: FloatLike | Float[ArrayLike, ' k'] | Literal['auto'] = 'auto',
414 sigma_init: FloatLike | Float[ArrayLike, ' k'] | Literal['auto'] = 'auto',
415 k: FloatLike = 2.0,
416 power: FloatLike = 2.0,
417 base: FloatLike = 0.95,
418 tau_num: FloatLike | None = None,
419 offset: FloatLike | Float[ArrayLike, ' k'] | None = None,
420 error_scale: Float[ArrayLike, ' n']
421 | Float[ArrayLike, 'k n']
422 | Series
423 | DataFrame
424 | None = None,
425 missing: Bool[ArrayLike, ' n']
426 | Bool[ArrayLike, 'k n']
427 | Series
428 | DataFrame
429 | None = None,
430 num_trees: int = 200,
431 n_save: int = 1000,
432 n_burn: int = 1000,
433 n_skip: int = 1,
434 printevery: int | None = 100,
435 pbar: bool = True,
436 num_chains: int | None = 4,
437 num_chain_devices: int | None | Literal['auto'] = 'auto',
438 num_data_devices: int | None = None,
439 devices: Literal['cpu', 'gpu'] | Device | Sequence[Device] | None = None,
440 seed: int | Key[Array, ''] = 0,
441 maxdepth: int = 6,
442 precompute_predict_train: bool = False,
443 init_kw: Mapping = MappingProxyType({}),
444 run_mcmc_kw: Mapping = MappingProxyType({}),
445 ) -> None:
446 # check data and put it in the right format
447 x_train, x_train_fmt = _process_predictor_input(x_train)
448 y_train = _process_response_input(y_train)
449 _check_same_length(x_train, y_train)
451 if error_scale is not None:
452 # `error_scale` is donated downstream as `init`'s `error_scale`, which
453 # keeps it (sharded) as `State.error_scale` for prediction
454 error_scale = _process_response_input(error_scale)
455 _check_same_length(x_train, error_scale)
457 if missing is not None:
458 missing = _process_response_input(missing, dtype=jnp.bool_)
459 _check_same_length(x_train, missing)
461 # check data types are correct for continuous/binary/multivariate regression
462 outcome_type, binary_mask = _check_type_settings(
463 y_train, outcome_type, error_scale
464 )
466 # process "standardization" settings
467 offset = _process_offset_settings(
468 y_train,
469 binary_mask,
470 missing,
471 None if offset is None else jnp.asarray(offset, jnp.float32),
472 )
473 leaf_prior_cov_inv = _process_leaf_variance_settings(
474 y_train,
475 binary_mask,
476 missing,
477 jnp.asarray(k, jnp.float32),
478 num_trees,
479 None if tau_num is None else jnp.asarray(tau_num, jnp.float32),
480 )
481 error_cov_inv = _process_error_variance_settings(
482 y_train,
483 outcome_type,
484 binary_mask,
485 missing,
486 sigma_df,
487 sigma_scale,
488 sigma_init,
489 error_scale,
490 )
492 # split the user-provided seed into an mcmc key and a binner key
493 if not is_key(seed):
494 seed = random.key(seed)
495 keys = split(seed)
497 # construct the binner and bin x_train
498 binner_obj = binner(x_train, key=keys.pop())
499 x_train = binner_obj.bin(x_train)
500 # copy max_split because `mcmcstep.init` donates it
501 max_split = jnp.array(binner_obj.max_split)
503 # setup and run mcmc
504 initial_state, mcmc_key, device, check_platform = _setup_mcmc(
505 x_train,
506 y_train,
507 outcome_type,
508 offset,
509 error_scale,
510 missing,
511 max_split,
512 leaf_prior_cov_inv,
513 error_cov_inv,
514 power,
515 base,
516 maxdepth,
517 num_trees,
518 init_kw,
519 rm_const,
520 sparse,
521 varprob,
522 num_chains,
523 num_chain_devices,
524 num_data_devices,
525 devices,
526 n_burn,
527 keys.pop(),
528 )
529 result = _run_mcmc(
530 initial_state,
531 n_save,
532 n_burn,
533 n_skip,
534 printevery,
535 pbar,
536 mcmc_key,
537 precompute_predict_train,
538 run_mcmc_kw,
539 check_platform,
540 )
542 # set private attributes
543 self._main_trace = result.main_trace
544 self._burnin_trace = result.burnin_trace
545 self._mcmc_state = result.final_state
546 self._binner = binner_obj
547 self._x_train_fmt = x_train_fmt
548 self._binary_mask = binary_mask
549 self._device = device
551 def predict(
552 self,
553 x_test: Real[ArrayLike, 'p m'] | DataFrame | str,
554 *,
555 kind: PredictKind | str = 'mean',
556 key: Key[Array, ''] | None = None,
557 error_scale: Float[ArrayLike, ' m']
558 | Float[ArrayLike, 'k m']
559 | Series
560 | DataFrame
561 | None = None,
562 ) -> (
563 Float32[Array, ' m']
564 | Float32[Array, 'k m']
565 | Float32[Array, 'ndpost m']
566 | Float32[Array, 'ndpost k m']
567 ):
568 """
569 Compute predictions at `x_test`.
571 Parameters
572 ----------
573 x_test
574 The test predictors, or the string ``'train'`` to compute
575 predictions on the training data.
576 kind
577 The kind of output. See `PredictKind` for details.
578 key
579 Jax random key, required when ``kind='outcome_samples'``.
580 error_scale
581 Per-observation error scale. Used with ``kind='outcome_samples'``,
582 and also with ``kind='mean'`` or ``'mean_samples'`` for binary
583 outcomes (since the success probability is ``Phi(latent /
584 error_scale)``). Required when the model was fit with `error_scale`
585 and ``x_test`` is new data. Shape matches the shape used at
586 fitting: ``(m,)`` for scalar scales, ``(k, m)`` for multivariate
587 per-component scales.
589 Returns
590 -------
591 Predictions at `x_test` in the requested format.
593 Raises
594 ------
595 ValueError
596 If `x_test` has a different format than `x_train`, or if `error_scale`
597 is specified when it should be `None`, or if `error_scale` is not
598 specified when it is required, or if the model splits datapoints
599 across devices (`num_data_devices`) and the number of test points
600 is not a multiple of the number of data devices.
602 Notes
603 -----
604 If the model splits datapoints across devices (`num_data_devices`),
605 the test points and the returned predictions are split the same way.
606 """
607 # parse arguments
608 kind = PredictKind(kind)
609 if kind is PredictKind.outcome_samples and key is None: 609 ↛ 610line 609 didn't jump to line 610 because the condition on line 609 was never true
610 msg = '`key` not specified'
611 raise ValueError(msg)
612 error_scale = self._process_error_scale_test(x_test, kind, error_scale)
613 x_test_is_train = isinstance(x_test, str) and x_test == 'train'
615 # use the predictions precomputed during the MCMC, if available
616 if x_test_is_train and isinstance(self._main_trace, MainTraceWithTrainPred):
617 return predict_train(
618 key,
619 self._main_trace,
620 error_scale,
621 self._mcmc_state.binary_indices,
622 self._mcmc_state.z is not None,
623 kind,
624 )
626 x_test = self._process_x_test(x_test, error_scale)
628 # place new test data on the devices of the model; the training data
629 # is already in place
630 if not x_test_is_train:
631 x_test, error_scale = self._device_put_test(x_test, error_scale)
633 # invoke jitted implementation
634 return predict(
635 key,
636 self._main_trace,
637 x_test,
638 error_scale,
639 self._mcmc_state.binary_indices,
640 self._mcmc_state.z is not None,
641 kind,
642 # the test points are sharded over the mesh 'data' axis (when
643 # there is one): the training data at `init`, new test data by
644 # `_device_put_test`. `evaluate_trace` can't detect this on its
645 # own at trace time, so declare it.
646 'shard_and_autobatch',
647 )
649 def _drop_device_info(self) -> 'Bart':
650 """Return a copy of the model without device placement metadata.
652 Clear the meshes in the MCMC state config and in the traces, and the
653 explicitly requested device. Only this static metadata is dropped: the
654 arrays keep their actual placement.
655 """
656 config = replace(self._mcmc_state.config, mesh=None)
657 main_trace = replace(self._main_trace, mesh=None)
658 burnin_trace = replace(self._burnin_trace, mesh=None)
659 obj = tree_at(
660 lambda b: (b._mcmc_state.config, b._main_trace, b._burnin_trace), # noqa: SLF001
661 self,
662 (config, main_trace, burnin_trace),
663 )
664 # `_device` is a static field, out of `tree_at`'s reach, so modify the
665 # fresh copy in place
666 object.__setattr__(obj, '_device', None)
667 return obj
669 def dump(self, path: str | PathLike) -> None:
670 """Serialize the fitted model to a file with `pickle`.
672 Parameters
673 ----------
674 path
675 The file to write to.
677 Notes
678 -----
679 Intended for short-term storage (e.g. caching across processes), not
680 long-term archival: the format depends on the versions of bartz, jax and
681 equinox. The arrays are copied to host memory and all device/sharding
682 placement is dropped; `load` reconstructs a single-device model.
683 """
684 # drop all device info (`Device` objects are not picklable), then
685 # gather any sharded arrays to host (dropping their sharding); the
686 # reload is single-device
687 obj = self._drop_device_info()
688 obj = jax.device_get(obj)
689 with Path(path).open('wb') as file:
690 pickle.dump(obj, file, protocol=pickle.HIGHEST_PROTOCOL)
692 @classmethod
693 def load(cls, path: str | PathLike) -> 'Bart':
694 """Load a model saved with `dump`.
696 Parameters
697 ----------
698 path
699 The file to read from.
701 Returns
702 -------
703 The deserialized model, on host memory with no device placement.
705 Raises
706 ------
707 TypeError
708 If the file does not contain a `Bart` instance.
709 """
710 with Path(path).open('rb') as file:
711 obj = pickle.load(file) # noqa: S301, the user owns the file
712 if not isinstance(obj, cls):
713 msg = f'unpickled a {type(obj).__name__}, not a {cls.__name__}'
714 raise TypeError(msg)
715 return obj
717 @property
718 def offset(self) -> Float32[Array, ''] | Float32[Array, ' k']:
719 """The prior mean of the latent mean function."""
720 return self._mcmc_state.forest.offset
722 @property
723 def n_save(self) -> int:
724 """The number of posterior samples after burn-in saved per chain."""
725 sample_axis = trace_sample_axes(self._main_trace).grow_prop_count
726 return self._main_trace.grow_prop_count.shape[sample_axis]
728 @property
729 def num_chains(self) -> int | None:
730 """The number of chains, `None` if scalar."""
731 return self._mcmc_state.num_chains()
733 @property
734 def ndpost(self) -> int:
735 """The total number of posterior samples after burn-in across all chains."""
736 return self._main_trace.grow_prop_count.size
738 @property
739 def num_trees(self) -> int:
740 """Return the number of trees used in the model."""
741 forest = self._mcmc_state.forest
742 chain_axis = chain_vmap_axes(forest).split_tree
743 # chainless split_tree is (num_trees, half_tree_size); num_trees is core axis 0
744 axis = chainful_axis(0, chain_axis)
745 return forest.split_tree.shape[axis]
747 def get_latent_prec(
748 self, only_continuous: bool = False
749 ) -> (
750 Float32[Array, ' n_burn_plus_n_save']
751 | Float32[Array, 'n_burn_plus_n_save k k']
752 | Float32[Array, 'num_chains n_burn_plus_n_save']
753 | Float32[Array, 'num_chains n_burn_plus_n_save k k']
754 ):
755 """Return the posterior samples of the latent error precision matrix.
757 Parameters
758 ----------
759 only_continuous
760 If `True` and the model has mixed binary-continuous outcomes,
761 return only the submatrix for the continuous components.
763 Returns
764 -------
765 MCMC samples of the error precision matrix.
767 Raises
768 ------
769 ValueError
770 If `only_continuous` is `True` but the model has only binary
771 outcomes, so there is no continuous submatrix to return.
773 Notes
774 -----
775 This method is meant to check for convergence, so it returns the full
776 MCMC trace and does not concatenate chains together. For probit
777 regression, this returns the precision of the latent error term, not
778 the Bernoulli precision for the binary outcome. For heteroskedastic
779 regression, the returned precision is the global precision parameter,
780 that would have to be divided by a squared error scale to get the
781 precision on a given datapoint.
782 """
783 binary_indices = self._mcmc_state.binary_indices
784 if (
785 only_continuous
786 and binary_indices is None
787 and self._mcmc_state.z is not None
788 ):
789 msg = 'Model has only binary outcomes, so there is no continuous submatrix to return.'
790 raise ValueError(msg)
792 return get_latent_prec(
793 self._burnin_trace,
794 self._main_trace,
795 binary_indices,
796 only_continuous=only_continuous,
797 )
799 def get_error_sdev(
800 self, mean: bool = False
801 ) -> (
802 Float32[Array, ' ndpost']
803 | Float32[Array, 'ndpost k']
804 | Float32[Array, '']
805 | Float32[Array, ' k']
806 ):
807 """Return the error standard deviation, post-burnin, chains concatenated.
809 Parameters
810 ----------
811 mean
812 If `True`, average the error covariance matrix across samples before
813 taking the square root, returning a single scalar or vector instead
814 of posterior samples.
816 Returns
817 -------
818 Posterior samples (or single estimate) of the error standard deviation; NaN for binary outcomes.
820 Notes
821 -----
822 Binary outcomes do have a standard deviation of course, but it's not
823 returned by this method because that would require to evaluate
824 predictions on a given X, since the Bernoulli variance is p(1-p).
825 """
826 # binary outcomes are filled with NaN, so disable the NaN check
827 with debug_nans(False):
828 return get_error_sdev(self._main_trace, self._binary_mask, mean=mean)
830 @cached_property
831 def accept(
832 self,
833 ) -> (
834 Float32[Array, ' n_burn_plus_n_save']
835 | Float32[Array, 'num_chains n_burn_plus_n_save']
836 ):
837 """Fraction of trees with an accepted grow or prune move at each iteration.
839 Includes the burn-in iterations and does not concatenate the chains, to
840 allow checking convergence. The iterations thinned away by `n_skip` are
841 not recorded.
842 """
843 return get_accept(self._burnin_trace, self._main_trace, self.num_trees)
845 @cached_property
846 def varcount(self) -> Int32[Array, 'ndpost p']:
847 """Histogram of predictor usage for decision rules in the trees."""
848 p = self._mcmc_state.forest.max_split.size
849 return varcount(p, self._main_trace)
851 @cached_property
852 def varcount_mean(self) -> Float32[Array, ' p']:
853 """Average of `varcount` across MCMC iterations."""
854 return self.varcount.mean(axis=0)
856 @cached_property
857 def varprob(self) -> Float32[Array, 'ndpost p']:
858 """Posterior samples of the probability of choosing each predictor for a decision rule."""
859 return varprob(self._mcmc_state.forest.max_split, self._main_trace)
861 @cached_property
862 def varprob_mean(self) -> Float32[Array, ' p']:
863 """The marginal posterior probability of each predictor being chosen for a decision rule."""
864 return self.varprob.mean(axis=0)
866 def _process_error_scale_test(
867 self,
868 x_test: Real[ArrayLike, 'p m'] | DataFrame | str,
869 kind: PredictKind,
870 error_scale: Float[ArrayLike, ' m']
871 | Float[ArrayLike, 'k m']
872 | Series
873 | DataFrame
874 | None,
875 ) -> Float32[Array, ' m'] | Float32[Array, 'k m'] | None:
876 """Validate and resolve the error scales for prediction.
878 Parameters
879 ----------
880 x_test
881 The raw (not yet processed) test predictors, or ``'train'``.
882 kind
883 The prediction kind.
884 error_scale
885 User-provided per-observation error scale, or `None`.
887 Returns
888 -------
889 The resolved error scale as a float32 array, or `None` if not applicable.
891 Raises
892 ------
893 ValueError
894 If `error_scale` is specified when it should be `None`, or missing
895 when required.
896 """
897 x_test_is_train = isinstance(x_test, str) and x_test == 'train'
898 train_error_scale = self._mcmc_state.error_scale
899 has_train_error_scale = train_error_scale is not None
900 is_binary = self._mcmc_state.z is not None
901 # the error scales enter the outcome samples of any outcome type, and
902 # also the mean of binary outcomes, since P(y=1) = Phi(latent / scale)
903 needs_error_scale = has_train_error_scale and (
904 kind is PredictKind.outcome_samples
905 or (is_binary and kind in (PredictKind.mean, PredictKind.mean_samples))
906 )
908 if not needs_error_scale:
909 if error_scale is not None: 909 ↛ 910line 909 didn't jump to line 910 because the condition on line 909 was never true
910 msg = (
911 '`error_scale` must be `None` in this configuration (error'
912 " scales are used with kind='outcome_samples', and with"
913 " kind='mean' or 'mean_samples' for binary outcomes, and"
914 ' only when the model was fit with `error_scale`)'
915 )
916 raise ValueError(msg)
917 return None
919 if x_test_is_train:
920 if error_scale is not None: 920 ↛ 921line 920 didn't jump to line 921 because the condition on line 920 was never true
921 msg = (
922 "`error_scale` must be `None` when x_test='train'"
923 ' (the training error scales are used automatically)'
924 )
925 raise ValueError(msg)
926 return train_error_scale
928 # new test data, model was fit with error scales
929 if error_scale is None:
930 msg = (
931 '`error_scale` is required because the model was fit with'
932 ' `error_scale` and x_test is new data'
933 )
934 raise ValueError(msg)
935 error_scale_test = _process_response_input(error_scale)
936 assert train_error_scale is not None # implied by needs_error_scale
937 # the per-observation axis is checked separately, against x_test
938 if error_scale_test.shape[:-1] != train_error_scale.shape[:-1]:
939 msg = (
940 f'`error_scale` shape mismatch with the training error scales: '
941 f'got {error_scale_test.shape=}, but the leading dimensions must '
942 f'match the training {train_error_scale.shape=} (only the '
943 f'per-observation axis may differ).'
944 )
945 raise ValueError(msg)
946 return error_scale_test
948 def _process_x_test(
949 self,
950 x_test: Real[ArrayLike, 'p m'] | DataFrame | str,
951 error_scale: Float32[Array, ' m'] | Float32[Array, 'k m'] | None,
952 ) -> UInt[Array, 'p m']:
953 """Convert x_test to binned format suitable for prediction."""
954 if isinstance(x_test, str):
955 if x_test != 'train': 955 ↛ 956line 955 didn't jump to line 956 because the condition on line 955 was never true
956 msg = (
957 f"x_test must be an array, a DataFrame, or 'train', got {x_test!r}"
958 )
959 raise ValueError(msg)
960 return self._mcmc_state.X
961 x_test, x_test_fmt = _process_predictor_input(x_test)
962 if x_test_fmt != self._x_train_fmt:
963 msg = f'Input format mismatch: {x_test_fmt=} != x_train_fmt={self._x_train_fmt!r}'
964 raise ValueError(msg)
965 if error_scale is not None:
966 _check_same_length(error_scale, x_test)
967 return self._binner.bin(x_test)
969 def _device_put_test(
970 self,
971 x_test: UInt[Array, 'p m'],
972 error_scale: Float32[Array, ' m'] | Float32[Array, 'k m'] | None,
973 ) -> tuple[UInt[Array, 'p m'], Float32[Array, ' m'] | Float32[Array, 'k m'] | None]:
974 """Place new test data on the devices of the model.
976 Mirror the placement of the training data done at fit time: shard over
977 the mesh if there is one (the observation axis over 'data'), else move
978 to the device requested explicitly at construction, if any. The inputs
979 are donated, so they must not be used elsewhere.
980 """
981 mesh = self._mcmc_state.config.mesh
982 if mesh is not None:
983 put = lambda a: device_put(
984 a,
985 NamedSharding(mesh, leaf_partition_spec(a.ndim, None, -1, mesh)),
986 donate=True,
987 )
988 elif self._device is not None:
989 put = lambda a: device_put(a, self._device, donate=True)
990 else:
991 return x_test, error_scale
992 if error_scale is None:
993 return put(x_test), None
994 else:
995 return put(x_test), put(error_scale)
997 def _check_trees(
998 self, error: bool = False
999 ) -> UInt[Array, 'num_chains n_save num_trees']:
1000 """Apply `bartz.grove.check_trace` to all the tree draws.
1002 Parameters
1003 ----------
1004 error
1005 If `True`, throw an error if any invalid trees are found.
1007 Returns
1008 -------
1009 An array where non-zero entries indicate invalid trees.
1011 Raises
1012 ------
1013 RuntimeError
1014 If `error` is `True` and any invalid trees are found.
1015 """
1016 out = check_trees(self._main_trace, self._mcmc_state.forest.max_split)
1017 if error:
1018 bad_count = jnp.count_nonzero(out).item()
1019 if bad_count > 0:
1020 msg = f'Found {bad_count} invalid trees in the MCMC trace.'
1021 raise RuntimeError(msg)
1022 return out
1024 def _tree_goes_bad(self) -> Bool[Array, 'num_chains n_save num_trees']:
1025 """Find iterations where a tree becomes invalid.
1027 Returns
1028 -------
1029 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``.
1030 """
1031 return tree_goes_bad(self._main_trace, self._mcmc_state.forest.max_split)
1033 def _check_replicated_trees(self) -> None:
1034 """Check that the trees are equal across data-sharded devices.
1036 If the data is sharded across devices, verify that the trees (which
1037 should be replicated) are identical on all shards.
1039 Raises
1040 ------
1041 RuntimeError
1042 If the trees differ across devices.
1043 """
1044 state = self._mcmc_state
1045 mesh = state.config.mesh
1046 if mesh is not None and 'data' in mesh.axis_names:
1047 # drop the data-sharded `leaf_indices` (not replicated) before the
1048 # cross-shard equality check; `None` is a deliberately off-type
1049 # placeholder, so use `tree_at`, which (unlike `dataclasses.replace`)
1050 # bypasses the `__init__` type checks
1051 replicated_forest = tree_at(lambda f: f.leaf_indices, state.forest, None)
1052 equal = equal_shards(
1053 replicated_forest, 'data', in_specs=PartitionSpec(), mesh=mesh
1054 )
1055 equal_array = jnp.stack(tree.leaves(equal))
1056 all_equal = jnp.all(equal_array)
1057 if not all_equal.item(): 1057 ↛ 1058line 1057 didn't jump to line 1058 because the condition on line 1057 was never true
1058 msg = 'The trees differ across data-sharded devices.'
1059 raise RuntimeError(msg)
1061 def _compare_resid(
1062 self,
1063 ) -> tuple[
1064 Float32[Array, '*num_chains n'] | Float32[Array, '*num_chains k n'],
1065 Float32[Array, '*num_chains n'] | Float32[Array, '*num_chains k n'],
1066 ]:
1067 """Re-compute residuals to compare them with the updated ones.
1069 Returns
1070 -------
1071 resid1
1072 The final state of the residuals updated during the MCMC.
1073 resid2
1074 The residuals computed from the final state of the trees.
1075 """
1076 return compare_resid(self._mcmc_state)
1078 def _depth_distr(self) -> Int32[Array, '*num_chains n_save d']:
1079 """Histogram of tree depths for each state of the trees.
1081 Returns
1082 -------
1083 A matrix where each row contains a histogram of tree depths.
1084 """
1085 return depth_distr(self._main_trace)
1087 def _points_per_node_distr(
1088 self, node_type: Literal['leaf', 'leaf-parent']
1089 ) -> Int32[Array, '*num_chains n_save n_plus_1']:
1090 return points_per_node_distr_trace(
1091 self._mcmc_state.X, self._main_trace, node_type
1092 )
1094 def _points_per_decision_node_distr(
1095 self,
1096 ) -> Int32[Array, '*num_chains n_save n_plus_1']:
1097 """Histogram of number of points belonging to parent-of-leaf nodes.
1099 Returns
1100 -------
1101 For each chain, a matrix where each row contains a histogram of number of points.
1102 """
1103 return self._points_per_node_distr('leaf-parent')
1105 def _points_per_leaf_distr(self) -> Int32[Array, '*num_chains n_save n_plus_1']:
1106 """Histogram of number of points belonging to leaves.
1108 Returns
1109 -------
1110 A matrix where each row contains a histogram of number of points.
1111 """
1112 return self._points_per_node_distr('leaf')
1114 def _print_tree(
1115 self, i_chain: int, i_sample: int, i_tree: int, print_all: bool = False
1116 ) -> None:
1117 """Print a single tree in human-readable format.
1119 Parameters
1120 ----------
1121 i_chain
1122 The index of the MCMC chain.
1123 i_sample
1124 The index of the (post-burnin) sample in the chain.
1125 i_tree
1126 The index of the tree in the sample.
1127 print_all
1128 If `True`, also print the content of unused node slots.
1129 """
1130 trace = self._main_trace
1131 trees = _trees_chain_first(trace)
1132 index = (i_chain, i_sample, i_tree) if trace.has_chains else (i_sample, i_tree)
1133 # index the heap arrays, leaving the (unbatched) leaf scale alone; the
1134 # trailing ellipsis covers the extra `k` axis of multivariate leaves
1135 trees = tree_at(
1136 lambda t: (t.var_tree, t.split_tree, t.leaf_tree),
1137 trees,
1138 replace_fn=lambda x: x[(*index, ...)],
1139 )
1140 s = format_tree(trees, print_all=print_all)
1141 print(s) # noqa: T201, this method is intended for debug
1144def _process_predictor_input(
1145 x: Real[ArrayLike, 'p n'] | DataFrame,
1146) -> tuple[Shaped[Array, 'p n'], Any]:
1147 if isinstance(x, DataFrame):
1148 fmt = dict(kind='dataframe', columns=x.columns)
1149 x = x.to_numpy().T
1150 else:
1151 fmt = dict(kind='array', num_covar=x.shape[0])
1152 x = jnp.asarray(x)
1153 assert x.ndim == 2
1154 return x, fmt
1157def _process_response_input(
1158 arr: Shaped[ArrayLike, ' n'] | Shaped[ArrayLike, 'k n'] | Series | DataFrame,
1159 /,
1160 *,
1161 dtype: DTypeLike = jnp.float32,
1162) -> Shaped[Array, ' n'] | Shaped[Array, 'k n']:
1163 if isinstance(arr, DataFrame):
1164 arr = arr.to_numpy().T
1165 elif isinstance(arr, Series):
1166 arr = arr.to_numpy()
1167 # one unconditional copy, safe to donate downstream
1168 arr = jnp.array(arr, dtype, copy=True)
1169 if arr.ndim < 1 or arr.ndim > 2: 1169 ↛ 1170line 1169 didn't jump to line 1170 because the condition on line 1169 was never true
1170 msg = f'response-like input must be 1D (n,) or 2D (k, n). Got {arr.ndim=}.'
1171 raise ValueError(msg)
1172 return arr
1175def _check_same_length(x1: Shaped[Array, '... n'], x2: Shaped[Array, '... n']) -> None:
1176 get_length = lambda x: x.shape[-1]
1177 assert get_length(x1) == get_length(x2)
1180def _check_type_settings(
1181 y_train: Float32[Array, ' n'] | Float32[Array, 'k n'],
1182 outcome_type: OutcomeType | str | Sequence[OutcomeType | str],
1183 error_scale: Float[Array, ' n'] | Float[Array, 'k n'] | None,
1184) -> tuple[OutcomeType | tuple[OutcomeType, ...], Bool[Array, ''] | Bool[Array, ' k']]:
1185 # standardize outcome_type to OutcomeType or tuple[OutcomeType, ...]
1186 if isinstance(outcome_type, Sequence) and not isinstance(outcome_type, str):
1187 outcome_type = tuple(OutcomeType(t) for t in outcome_type)
1188 num_types = len(outcome_type)
1189 if len(set(outcome_type)) == 1:
1190 outcome_type = outcome_type[0]
1191 else:
1192 num_types = None
1193 outcome_type = OutcomeType(outcome_type)
1195 # validation
1196 if num_types is not None and (y_train.ndim != 2 or num_types != y_train.shape[0]):
1197 msg = (
1198 f'Sequence outcome_type of length {num_types}'
1199 f' requires y_train.shape=({num_types}, n),'
1200 f' found {y_train.shape=}.'
1201 )
1202 raise ValueError(msg)
1203 if ( 1203 ↛ 1208line 1203 didn't jump to line 1208 because the condition on line 1203 was never true
1204 error_scale is not None
1205 and error_scale.ndim == 2
1206 and (y_train.ndim != 2 or error_scale.shape[0] != y_train.shape[0])
1207 ):
1208 msg = (
1209 f'2D error_scale (per-component scales) requires y_train of '
1210 f'shape (k, n) with matching k; got {error_scale.shape=}, '
1211 f'{y_train.shape=}.'
1212 )
1213 raise ValueError(msg)
1215 if isinstance(outcome_type, tuple):
1216 binary_mask = jnp.array([t is OutcomeType.binary for t in outcome_type])
1217 else:
1218 binary_mask = jnp.bool_(outcome_type is OutcomeType.binary)
1219 binary_mask = jnp.broadcast_to(binary_mask, y_train.shape[:-1])
1221 return outcome_type, binary_mask
1224def _process_sparsity_settings(
1225 x_train: Real[Array, 'p n'], sparse: SparseConfig
1226) -> (
1227 tuple[None, None, None, None]
1228 | tuple[FloatLike, None, None, None]
1229 | tuple[None, FloatLike, FloatLike, FloatLike]
1230):
1231 """Return (theta, a, b, rho)."""
1232 if not sparse.enabled:
1233 return None, None, None, None
1234 elif sparse.theta is not None:
1235 return sparse.theta, None, None, None
1236 else:
1237 rho = sparse.rho
1238 if rho is None:
1239 p, _ = x_train.shape
1240 rho = float(p)
1241 return None, sparse.a, sparse.b, rho
1244@jit
1245def _process_offset_settings(
1246 y_train: Float32[Array, ' n'] | Float32[Array, 'k n'],
1247 binary_mask: Bool[Array, ''] | Bool[Array, ' k'],
1248 missing: Bool[Array, ' n'] | Bool[Array, 'k n'] | None,
1249 offset: Float32[Array, ''] | Float32[Array, ' k'] | None,
1250) -> Float32[Array, ''] | Float32[Array, ' k']:
1251 """Determine and return offset."""
1252 if offset is not None:
1253 return jnp.broadcast_to(offset, y_train.shape[:-1])
1254 if y_train.shape[-1] < 1:
1255 return jnp.zeros(y_train.shape[:-1])
1257 if missing is None:
1258 *_, n_valid = y_train.shape
1259 prop = (y_train != 0).mean(-1)
1260 mean = y_train.mean(-1)
1261 else:
1262 *_, n = y_train.shape
1263 n_valid = n - jnp.count_nonzero(missing, axis=-1)
1264 safe_n = jnp.maximum(n_valid, 1)
1265 prop = jnp.where(missing, 0, y_train != 0).sum(-1) / safe_n
1266 mean = jnp.where(missing, 0.0, y_train).sum(-1) / safe_n
1268 bound = jnp.reciprocal(1.0 + n_valid)
1269 binary_offset = ndtri(jnp.clip(prop, bound, 1 - bound))
1270 offset = jnp.where(binary_mask, binary_offset, mean)
1271 return jnp.where(n_valid > 0, offset, 0.0)
1274@jit(static_argnums=(4,))
1275def _process_leaf_variance_settings(
1276 y_train: Float32[Array, ' n'] | Float32[Array, 'k n'],
1277 binary_mask: Bool[Array, ''] | Bool[Array, ' k'],
1278 missing: Bool[Array, ' n'] | Bool[Array, 'k n'] | None,
1279 k: Float[Array, ''],
1280 num_trees: int,
1281 tau_num: Float[Array, ''] | None,
1282) -> Float32[Array, ''] | Float32[Array, 'k k']:
1283 """Return `leaf_prior_cov_inv`."""
1284 # determine `tau_num` if not specified
1285 *kshape, n = y_train.shape
1286 if tau_num is None:
1287 if n < 2:
1288 continuous_tau = jnp.ones(kshape)
1289 elif missing is None:
1290 continuous_tau = (y_train.max(-1) - y_train.min(-1)) / 2
1291 else:
1292 n_valid = n - jnp.count_nonzero(missing, axis=-1)
1293 ymax = jnp.where(missing, -jnp.inf, y_train).max(-1)
1294 ymin = jnp.where(missing, jnp.inf, y_train).min(-1)
1295 continuous_tau = jnp.where(n_valid >= 2, (ymax - ymin) / 2, 1.0)
1296 tau_num = jnp.where(binary_mask, 3.0, continuous_tau)
1298 # leaf prior standard deviation
1299 sigma_mu = tau_num / (k * math.sqrt(num_trees))
1301 # leaf prior precision matrix
1302 leaf_prior_cov_inv = jnp.reciprocal(jnp.square(sigma_mu))
1303 if y_train.ndim == 2:
1304 leaf_prior_cov_inv = jnp.diag(jnp.broadcast_to(leaf_prior_cov_inv, kshape))
1305 return leaf_prior_cov_inv
1308def _process_error_variance_settings(
1309 y_train: Float32[Array, ' n'] | Float32[Array, 'k n'],
1310 outcome_type: OutcomeType | tuple[OutcomeType, ...],
1311 binary_mask: Bool[Array, ''] | Bool[Array, ' k'],
1312 missing: Bool[Array, ' n'] | Bool[Array, 'k n'] | None,
1313 sigma_df: FloatLike,
1314 sigma_scale: FloatLike | Float[ArrayLike, ' k'] | Literal['auto'],
1315 sigma_init: FloatLike | Float[ArrayLike, ' k'] | Literal['auto'],
1316 error_scale: Float32[Array, ' n'] | Float32[Array, 'k n'] | None,
1317) -> Wishart | None:
1318 """Build the error precision prior from the user settings."""
1319 if outcome_type is OutcomeType.binary:
1320 if not isinstance(sigma_scale, str) or not isinstance(sigma_init, str):
1321 msg = (
1322 'Do not set `sigma_scale` or `sigma_init` for binary regression, '
1323 'they are ignored'
1324 )
1325 raise ValueError(msg)
1326 return None
1328 *kdims, _ = y_train.shape # () or (k,)
1329 k = kdims[0] if kdims else 1
1330 nu = jnp.asarray(sigma_df, jnp.float32) + (k - 1)
1332 # guarded per-component variance of y_train, computed only when an 'auto'
1333 # spec needs it (this function is not jitted, so it would not be elided)
1334 if isinstance(sigma_scale, str) or isinstance(sigma_init, str):
1335 vary = _guarded_response_variance(y_train, error_scale, missing)
1336 else:
1337 vary = None
1339 # prior rate: E[precision] = nu / rate, so rate = nu * var per component
1340 rate_diag = jnp.where(
1341 binary_mask, 0.0, nu * _resolve_error_variance(sigma_scale, vary, kdims)
1342 )
1344 # initial precision = 1 / var per component (1 for binary components)
1345 init_var = _resolve_error_variance(sigma_init, vary, kdims)
1346 init_diag = jnp.where(binary_mask, 1.0, jnp.reciprocal(init_var))
1348 if y_train.ndim == 2:
1349 rate, init = jnp.diag(rate_diag), jnp.diag(init_diag)
1350 else:
1351 rate, init = rate_diag, init_diag
1352 return make_error_cov_prior(nu, rate, init, outcome_type, missing)
1355@jit
1356def _guarded_response_variance(
1357 y_train: Float32[Array, ' n'] | Float32[Array, 'k n'],
1358 error_scale: Float32[Array, ' n'] | Float32[Array, 'k n'] | None,
1359 missing: Bool[Array, ' n'] | Bool[Array, 'k n'] | None,
1360) -> Float32[Array, '*k']:
1361 """Per-component variance of `y_train`, used by the 'auto' error scale.
1363 A precision-weighted variance (precision ``1 / error_scale ** 2``) estimates
1364 ``sigma ** 2`` at unit error scale; `missing` entries are dropped. The
1365 variance is guarded to 1 when undefined (fewer than 2 valid points) or
1366 non-positive.
1367 """
1368 if error_scale is None and missing is None:
1369 vary = jnp.var(y_train, axis=-1)
1370 return jnp.where(vary > 0, vary, 1.0)
1371 else:
1372 prec = (
1373 jnp.ones(())
1374 if error_scale is None
1375 else jnp.reciprocal(jnp.square(error_scale))
1376 )
1377 if missing is not None:
1378 prec = jnp.where(missing, 0.0, prec)
1379 y_train = jnp.where(missing, 0.0, y_train)
1380 n_valid = jnp.count_nonzero(prec, axis=-1)
1381 wmean = jnp.sum(prec * y_train, axis=-1) / jnp.sum(prec, axis=-1)
1382 sqdev = prec * jnp.square(y_train - wmean[..., None])
1383 vary = jnp.sum(sqdev, axis=-1) / n_valid
1384 # guard on n_valid too: with a single valid point the variance is 0 in
1385 # exact arithmetic, but float rounding in wmean can leave a tiny
1386 # positive vary that would slip past the `vary > 0` guard
1387 return jnp.where((n_valid > 1) & (vary > 0), vary, 1.0)
1390def _resolve_error_variance(
1391 spec: FloatLike | Float[ArrayLike, ' k'] | Literal['auto'],
1392 vary: Float32[Array, '*k'] | None,
1393 shape: Sequence[int],
1394) -> Float32[Array, '*k']:
1395 """Per-component error variance from a scale spec ('auto' uses var(y))."""
1396 if isinstance(spec, str):
1397 if spec != 'auto':
1398 msg = f"unrecognized value {spec!r}, expected 'auto' or a number"
1399 raise ValueError(msg)
1400 assert vary is not None # computed iff some spec is 'auto'
1401 return vary
1402 else:
1403 return jnp.broadcast_to(jnp.square(jnp.asarray(spec, jnp.float32)), shape)
1406def make_error_cov_prior(
1407 nu: Float32[Array, ''],
1408 rate: Float32[Array, ''] | Float32[Array, 'k k'],
1409 value: Float32[Array, ''] | Float32[Array, 'k k'],
1410 outcome_type: OutcomeType | tuple[OutcomeType, ...],
1411 missing: Bool[Array, ' n'] | Bool[Array, 'k n'] | None,
1412) -> Wishart:
1413 """Build the error precision prior, diagonal-constrained where required.
1415 Mixed binary-continuous and partial-missing (2-D mask) regression restrict
1416 the error covariance to diagonal, so they take a `DiagWishart`; the dense
1417 cases take a `Wishart`. `init` re-checks this choice. `value` is the initial
1418 value of the precision.
1419 """
1420 if isinstance(outcome_type, tuple):
1421 binary = [t is OutcomeType.binary for t in outcome_type]
1422 is_mixed = any(binary) and not all(binary)
1423 else:
1424 is_mixed = False
1425 # a 2-D missingness mask only occurs with multivariate y (checked in `init`)
1426 partial_missing = missing is not None and missing.ndim == 2
1427 if is_mixed or partial_missing:
1428 return DiagWishart(nu=nu, rate=rate, value=value)
1429 else:
1430 return Wishart(nu=nu, rate=rate, value=value)
1433def _setup_mcmc(
1434 x_train: Real[Array, 'p n'],
1435 y_train: Float32[Array, ' n'] | Float32[Array, 'k n'],
1436 outcome_type: OutcomeType | tuple[OutcomeType, ...],
1437 offset: Float32[Array, ''] | Float32[Array, ' k'],
1438 error_scale: Float[Array, ' n'] | Float[Array, 'k n'] | None,
1439 missing: Bool[Array, ' n'] | Bool[Array, 'k n'] | None,
1440 max_split: UInt[Array, ' p'],
1441 leaf_prior_cov_inv: Float32[Array, ''] | Float32[Array, 'k k'],
1442 error_cov_inv: Wishart | None,
1443 power: FloatLike,
1444 base: FloatLike,
1445 maxdepth: int,
1446 num_trees: int,
1447 init_kw: Mapping[str, Any],
1448 rm_const: bool,
1449 sparse: SparseConfig,
1450 varprob: Float[ArrayLike, ' p'] | None,
1451 num_chains: int | None,
1452 num_chain_devices: int | None | Literal['auto'],
1453 num_data_devices: int | None,
1454 devices: Literal['cpu', 'gpu'] | Device | Sequence[Device] | None,
1455 n_burn: int,
1456 mcmc_key: Key[Array, ''],
1457) -> tuple[State, Key[Array, ''], Device | None, Literal['cpu', 'gpu'] | None]:
1458 theta, a, b, rho = _process_sparsity_settings(x_train, sparse)
1460 device_kw, device, check_platform = process_device_settings(
1461 y_train, num_chains, num_chain_devices, num_data_devices, devices
1462 )
1464 kw: dict = dict(
1465 X=x_train,
1466 y=y_train,
1467 outcome_type=outcome_type,
1468 offset=offset,
1469 error_scale=error_scale,
1470 missing=missing,
1471 max_split=max_split,
1472 num_trees=num_trees,
1473 p_nonterminal=make_p_nonterminal(maxdepth, base, power),
1474 leaf_prior_cov_inv=leaf_prior_cov_inv,
1475 error_cov_inv=error_cov_inv,
1476 min_points_per_decision_node=10,
1477 filter_splitless_vars=jnp.sum(max_split == 0).item() if rm_const else 0,
1478 log_s=process_varprob(varprob, max_split),
1479 theta=theta,
1480 a=a,
1481 b=b,
1482 rho=rho,
1483 sparse_on_at=n_burn // 2 if sparse.enabled else None,
1484 augment=sparse.augment,
1485 **device_kw,
1486 )
1488 kw = dict(kw, **init_kw)
1490 state = init(**kw)
1492 # put state and mcmc key on device if requested explicitly by the user
1493 if device is not None:
1494 mcmc_key, state = device_put((mcmc_key, state), device, donate=True)
1496 return state, mcmc_key, device, check_platform
1499def _run_mcmc(
1500 mcmc_state: State,
1501 n_save: int,
1502 n_burn: int,
1503 n_skip: int,
1504 printevery: int | None,
1505 pbar: bool,
1506 key: Key[Array, ''],
1507 precompute_predict_train: bool,
1508 run_mcmc_kw: Mapping,
1509 check_platform: Literal['cpu', 'gpu'] | None,
1510) -> RunMCMCResult:
1511 # fill list of callbacks
1512 callbacks = ()
1513 if printevery is not None:
1514 if pbar:
1515 callbacks += (make_tqdm_callback(mcmc_state, report_every=printevery),)
1516 else:
1517 callbacks += (
1518 make_print_callback(
1519 mcmc_state,
1520 dot_every=None if printevery == 1 else 1,
1521 report_every=printevery,
1522 ),
1523 )
1524 if check_platform is not None:
1525 callbacks += (CheckPlatformCallback(check_platform),)
1527 # prepare arguments
1528 kw: dict = dict(
1529 n_burn=n_burn,
1530 n_skip=n_skip,
1531 inner_loop_length=printevery,
1532 callback=CallbackTuple(callbacks),
1533 )
1534 if precompute_predict_train:
1535 kw = dict(**kw, main_trace_type=MainTraceWithTrainPred)
1536 kw = dict(kw, **run_mcmc_kw)
1538 return run_mcmc(key, mcmc_state, n_save, **kw)
1541@jit(static_argnames='p')
1542# this is jitted such that lax.collapse below does not create a copy
1543def varcount(p: int, trace: MainTrace) -> Int32[Array, 'ndpost p']:
1544 """Histogram of predictor usage for decision rules in the trees, squashing chains."""
1545 varcount: Int32[Array, '*chains samples p']
1546 varcount = compute_varcount(p, trace, out_chain_axis=0)
1547 return lax.collapse(varcount, 0, -1)
1550@jit(static_argnames='mean')
1551def get_error_sdev(
1552 trace: MainTrace,
1553 binary_mask: Bool[Array, ''] | Bool[Array, ' k'],
1554 *,
1555 mean: bool = False,
1556) -> (
1557 Float32[Array, ' ndpost']
1558 | Float32[Array, 'ndpost k']
1559 | Float32[Array, '']
1560 | Float32[Array, ' k']
1561):
1562 """Error standard deviation, post-burnin, chains concatenated."""
1563 prec = trace.error_cov_inv
1564 if trace.has_chains:
1565 # shape (chains, samples) or (chains, samples, k, k), concatenate chains
1566 prec = chain_to_axis(prec, chain_vmap_axes(trace).error_cov_inv)
1567 prec = lax.collapse(prec, 0, 2)
1568 is_uv = prec.ndim == 1
1569 if is_uv:
1570 # univariate case, reshape to 1x1 matrix
1571 prec = prec[..., None, None]
1573 # invert precision to covariance, then take diagonal variance
1574 cov = inv_via_chol_with_gersh(prec)
1575 var = jnp.diagonal(cov, axis1=-2, axis2=-1)
1576 if mean:
1577 var = var.mean(0)
1578 sdev = jnp.sqrt(var)
1579 if is_uv:
1580 sdev = sdev.squeeze(-1)
1581 return jnp.where(binary_mask, jnp.nan, sdev)
1584@jit(static_argnames='only_continuous')
1585def get_latent_prec(
1586 burnin_trace: BurninTrace,
1587 main_trace: MainTrace,
1588 binary_indices: Int32[Array, ' kb'] | None,
1589 *,
1590 only_continuous: bool = False,
1591) -> (
1592 Float32[Array, ' n_burn_plus_n_save']
1593 | Float32[Array, 'n_burn_plus_n_save k k']
1594 | Float32[Array, 'num_chains n_burn_plus_n_save']
1595 | Float32[Array, 'num_chains n_burn_plus_n_save k k']
1596):
1597 """Latent error precision trace, burn-in + main concatenated."""
1598 burnin = burnin_trace.error_cov_inv
1599 main = main_trace.error_cov_inv
1600 sample_axis = trace_sample_axes(main_trace).error_cov_inv
1601 prec = jnp.concatenate([burnin, main], axis=sample_axis)
1602 prec = chain_to_axis(prec, chain_vmap_axes(main_trace).error_cov_inv)
1603 if only_continuous and binary_indices is not None:
1604 *_, k, _ = prec.shape
1605 kc = k - binary_indices.size
1606 mask = jnp.ones(k, dtype=bool).at[binary_indices].set(False)
1607 (cont_indices,) = jnp.nonzero(mask, size=kc)
1608 prec = prec[..., cont_indices[:, None], cont_indices[None, :]]
1609 return prec
1612@jit(static_argnames='num_trees')
1613def get_accept(
1614 burnin_trace: BurninTrace, main_trace: MainTrace, num_trees: int
1615) -> (
1616 Float32[Array, ' n_burn_plus_n_save']
1617 | Float32[Array, 'num_chains n_burn_plus_n_save']
1618):
1619 """Fraction of trees with an accepted move, burn-in + main concatenated."""
1620 sample_axis = trace_sample_axes(main_trace).grow_acc_count
1621 acc = jnp.concatenate(
1622 [
1623 burnin_trace.grow_acc_count + burnin_trace.prune_acc_count,
1624 main_trace.grow_acc_count + main_trace.prune_acc_count,
1625 ],
1626 axis=sample_axis,
1627 )
1628 acc = chain_to_axis(acc, chain_vmap_axes(main_trace).grow_acc_count)
1629 return acc / num_trees
1632@jit
1633def varprob(
1634 max_split: UInt[Array, ' p'], trace: MainTrace
1635) -> Float32[Array, 'ndpost p']:
1636 """Posterior samples of predictor selection probability, chains concatenated."""
1637 p = max_split.size
1638 varprob = trace.varprob
1639 if varprob is None:
1640 ndpost = trace.grow_prop_count.size
1641 peff = jnp.count_nonzero(max_split)
1642 out = jnp.where(max_split, 1 / peff, 0)
1643 return jnp.broadcast_to(out, (ndpost, p))
1644 varprob = chain_to_axis(varprob, chain_vmap_axes(trace).varprob)
1645 return varprob.reshape(-1, p)
1648def _trees_chain_first(obj: TreeHeaps) -> TreesTrace:
1649 """Extract `obj`'s heap arrays, moving any chain axis to the front.
1651 Returns a `TreesTrace` whose leading axis is the chain axis when `obj`
1652 carries one, and the bare per-object heap arrays otherwise.
1653 """
1654 trees = project(TreesTrace, obj)
1655 if get_has_chains(obj):
1656 axes = trees.axes_from_dataclass(chain_vmap_axes(obj))
1657 # WORKAROUND(python<3.14): use operator.is_none
1658 trees = tree.map(chain_to_axis, trees, axes, is_leaf=lambda x: x is None)
1659 return trees
1662@jit
1663def check_trees(
1664 trace: MainTrace, max_split: UInt[Array, ' p']
1665) -> UInt[Array, 'num_chains n_save num_trees']:
1666 """Apply `bartz.grove.check_trace` to all the tree draws."""
1667 trees = _trees_chain_first(trace)
1668 out: UInt[Array, '*chains samples num_trees']
1669 out = check_trace(trees, max_split)
1670 if out.ndim < 3:
1671 out = out[None, :, :]
1672 return out
1675@jit
1676def tree_goes_bad(
1677 trace: MainTrace, max_split: UInt[Array, ' p']
1678) -> Bool[Array, 'num_chains n_save num_trees']:
1679 """Find iterations where a tree becomes invalid."""
1680 bad = check_trees(trace, max_split).astype(bool)
1681 bad_before = jnp.pad(bad[:, :-1, :], [(0, 0), (1, 0), (0, 0)])
1682 return bad & ~bad_before
1685@jit
1686def compare_resid(
1687 state: State,
1688) -> tuple[
1689 Float32[Array, '*num_chains n'] | Float32[Array, '*num_chains k n'],
1690 Float32[Array, '*num_chains n'] | Float32[Array, '*num_chains k n'],
1691]:
1692 """Re-compute residuals to compare them with the updated ones."""
1693 chain_axes = chain_vmap_axes(state)
1694 resid1 = chain_to_axis(state.resid * state.resid_unit[..., None], chain_axes.resid)
1695 z = chain_to_axis(state.z, chain_axes.z) if state.z is not None else None
1697 forests = _trees_chain_first(state.forest)
1698 trees = evaluate_forest(state.X, forests, sum_batch_axis=-1)
1700 if state.binary_indices is not None:
1701 # mixed binary-continuous: z has only binary rows, y has all rows
1702 assert z is not None
1703 ref = jnp.broadcast_to(state.y, resid1.shape)
1704 ref = ref.at[..., state.binary_indices, :].set(z)
1705 elif z is not None:
1706 ref = z
1707 else:
1708 ref = state.y
1709 resid2 = ref - (trees + state.forest.offset[..., None])
1711 return resid1, resid2
1714@jit
1715def depth_distr(trace: MainTrace) -> Int32[Array, '*num_chains n_save d']:
1716 """Histogram of tree depths for each state of the trees."""
1717 split_tree = chain_to_axis(trace.split_tree, chain_vmap_axes(trace).split_tree)
1718 out: Int32[Array, '*chains samples d']
1719 out = forest_depth_distr(split_tree)
1720 if out.ndim < 3: 1720 ↛ 1722line 1720 didn't jump to line 1722 because the condition on line 1720 was always true
1721 out = out[None, :, :]
1722 return out
1725@jit(static_argnames='node_type')
1726def points_per_node_distr_trace(
1727 X: UInt[Array, 'p n'], trace: MainTrace, node_type: Literal['leaf', 'leaf-parent']
1728) -> Int32[Array, '*num_chains n_save n+1']:
1729 """Histogram of number of points per node, for every tree draw in the trace."""
1730 chain_axes = chain_vmap_axes(trace)
1731 var_tree = chain_to_axis(trace.var_tree, chain_axes.var_tree)
1732 split_tree = chain_to_axis(trace.split_tree, chain_axes.split_tree)
1733 out: Int32[Array, '*chains samples n+1']
1734 out = points_per_node_distr(X, var_tree, split_tree, node_type, sum_batch_axis=-1)
1735 if out.ndim < 3:
1736 out = out[None, :, :]
1737 return out
1740class DeviceKwArgs(TypedDict):
1741 num_chains: int | None
1742 mesh: Mesh | None
1743 leaf_dtype: DTypeLike
1744 prec_scale_dtype: DTypeLike
1747def process_device_settings(
1748 y_train: Float32[Array, ' n'] | Float32[Array, 'k n'],
1749 num_chains: int | None,
1750 num_chain_devices: int | None | Literal['auto'],
1751 num_data_devices: int | None,
1752 devices: Literal['cpu', 'gpu'] | Device | Sequence[Device] | None,
1753) -> tuple[DeviceKwArgs, Device | None, Literal['cpu', 'gpu'] | None]:
1754 """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."""
1755 # whether the user pinned a concrete pool of devices (vs. inheriting all of
1756 # the platform's devices); the auto chain sharding may not exceed that pool
1757 explicit_devices = devices is not None and not isinstance(devices, str)
1759 # the platform is deduced (rather than fixed by the user) only when neither a
1760 # device pool nor a platform string is passed
1761 platform_deduced = devices is None
1763 platform, device, devices = _determine_devices(y_train, devices)
1764 check_platform = platform if platform_deduced else None
1765 num_chain_devices = _determine_num_chain_devices(
1766 platform,
1767 num_chains,
1768 num_chain_devices,
1769 num_data_devices,
1770 len(devices),
1771 explicit_devices,
1772 )
1773 mesh, device = _determine_mesh(num_chain_devices, num_data_devices, device, devices)
1775 # prepare arguments to `init`
1776 dtype = jnp.float16 if platform == 'gpu' else jnp.float32
1777 settings = DeviceKwArgs(
1778 num_chains=num_chains, mesh=mesh, leaf_dtype=dtype, prec_scale_dtype=dtype
1779 )
1781 return settings, device, check_platform
1784def _determine_devices(
1785 y_train: Float32[Array, ' n'] | Float32[Array, 'k n'],
1786 devices: Literal['cpu', 'gpu'] | Device | Sequence[Device] | None,
1787) -> tuple[Literal['cpu', 'gpu'], Device | None, Sequence[Device]]:
1788 """Determine the target platform and set of devices for the MCMC, and possibly a single target device."""
1789 if isinstance(devices, str):
1790 platform: Literal['cpu', 'gpu'] = devices # ty:ignore[invalid-assignment]
1791 devices = jax.devices(platform)
1792 return platform, devices[0], devices
1793 elif devices is not None:
1794 if not hasattr(devices, '__len__'):
1795 devices = (devices,)
1796 device = devices[0]
1797 return device.platform, device, devices
1798 elif hasattr(y_train, 'platform'): 1798 ↛ 1805line 1798 didn't jump to line 1805 because the condition on line 1798 was always true
1799 # set device=None because if the devices were not specified explicitly
1800 # we may be in the case where computation will follow data placement,
1801 # do not disturb jax as the user may be playing with vmap, jit, reshard...
1802 platform = cast(Literal['cpu', 'gpu'], y_train.platform())
1803 return platform, None, jax.devices(platform)
1804 else:
1805 msg = 'not possible to infer device from `y_train`, please set `devices`'
1806 raise ValueError(msg)
1809def _largest_divisor_at_most(n: int, cap: int) -> int:
1810 """Return the largest divisor of `n` in [1, cap]."""
1811 for d in range(cap, 0, -1): 1811 ↛ 1814line 1811 didn't jump to line 1814 because the loop on line 1811 didn't complete
1812 if n % d == 0:
1813 return d
1814 return 1 # unreachable: 1 always divides n
1817def _determine_num_chain_devices(
1818 platform: str,
1819 num_chains: int | None,
1820 num_chain_devices: int | None | Literal['auto'],
1821 num_data_devices: int | None,
1822 num_devices: int,
1823 explicit_devices: bool,
1824) -> int | None:
1825 """Resolve and validate `num_chain_devices`, returning the chain mesh axis size or `None`."""
1826 if num_chain_devices == 'auto':
1827 num_chain_devices = _auto_num_chain_devices(
1828 platform, num_chains, num_data_devices, num_devices, explicit_devices
1829 )
1831 # an explicit value must be a positive divisor of the number of chains
1832 if num_chain_devices is not None:
1833 effective_chains = 1 if num_chains is None else num_chains
1834 if num_chain_devices < 1 or effective_chains % num_chain_devices:
1835 chains_desc = (
1836 'a single chain (num_chains=None)'
1837 if num_chains is None
1838 else f'num_chains={num_chains}'
1839 )
1840 msg = (
1841 f'num_chain_devices={num_chain_devices} must be a positive '
1842 f'divisor of the number of chains ({chains_desc})'
1843 )
1844 raise ValueError(msg)
1846 # there is no chain axis to shard when the chains are scalar
1847 if num_chains is None:
1848 return None
1849 return num_chain_devices
1852def _auto_num_chain_devices(
1853 platform: str,
1854 num_chains: int | None,
1855 num_data_devices: int | None,
1856 num_devices: int,
1857 explicit_devices: bool,
1858) -> int | None:
1859 """Pick `num_chain_devices` automatically for multi-chain cpu runs.
1861 `num_data_devices` reserves devices for the data axis, so the chain axis can
1862 only use a fraction of them; this keeps the ``chains x data`` mesh within the
1863 `num_devices` available devices.
1864 """
1865 if num_chains is None or num_chains == 1 or platform != 'cpu':
1866 return None
1867 data_devices = num_data_devices or 1
1868 num_cores = cpu_count()
1869 assert num_cores is not None, 'could not determine number of cpu cores'
1871 # devices available for the chain axis after reserving for the data axis
1872 core_budget = max(1, num_cores // data_devices)
1873 num_shards = _largest_divisor_at_most(num_chains, core_budget)
1875 if num_shards > 1:
1876 # the mesh draws from `num_devices` devices, whether those are all the
1877 # platform's devices or an explicit subset passed by the user
1878 device_budget = max(1, num_devices // data_devices)
1879 if device_budget < num_shards:
1880 new_num_shards = _largest_divisor_at_most(num_chains, device_budget)
1881 warn(
1882 _auto_chain_devices_warning(
1883 num_chains,
1884 num_shards,
1885 new_num_shards,
1886 device_budget,
1887 num_devices,
1888 num_data_devices,
1889 explicit_devices,
1890 )
1891 )
1892 num_shards = new_num_shards
1894 return num_shards if num_shards > 1 else None
1897def _auto_chain_devices_warning(
1898 num_chains: int,
1899 desired: int,
1900 actual: int,
1901 device_budget: int,
1902 num_devices: int,
1903 num_data_devices: int | None,
1904 explicit_devices: bool,
1905) -> str:
1906 """Compose the warning shown when auto chain sharding is capped by the device count."""
1907 if explicit_devices:
1908 pool = f'the {num_devices} devices passed in `devices`'
1909 few = f'only {num_devices} devices were passed in `devices`'
1910 advice = ''
1911 else:
1912 pool = f'the {num_devices} jax cpu devices'
1913 few = f'jax is set up with only {num_devices} cpu devices'
1914 advice = (
1915 ' To enable more parallelization, increase the limit with '
1916 '`jax.config.update("jax_num_cpu_devices", <num_devices>)`.'
1917 )
1918 if num_data_devices:
1919 limit = (
1920 f'only {device_budget} of {pool} are free for chains '
1921 f'(num_data_devices={num_data_devices} reserves the rest)'
1922 )
1923 else:
1924 limit = few
1925 return (
1926 f'`Bart` would like to shard {num_chains} chains across {desired} '
1927 f'devices, but {limit}, so it will use {actual} devices for chains '
1928 f'instead.{advice}'
1929 )
1932def _determine_mesh(
1933 num_chain_devices: int | None,
1934 num_data_devices: int | None,
1935 device: Device | None,
1936 devices: Sequence[Device],
1937) -> tuple[Mesh, None] | tuple[None, Device | None]:
1938 """Create a jax device mesh for `mcmcstep.init()`."""
1939 if num_chain_devices is None and num_data_devices is None:
1940 return None, device
1941 else:
1942 mesh = {}
1943 if num_chain_devices is not None:
1944 mesh.update(chains=num_chain_devices)
1945 if num_data_devices is not None:
1946 mesh.update(data=num_data_devices)
1947 mesh = make_mesh(
1948 tuple(mesh.values()),
1949 tuple(mesh),
1950 axis_types=(AxisType.Auto,) * len(mesh),
1951 devices=devices,
1952 )
1953 return mesh, None
1954 # set device=None because `mcmcstep.init` will `device_put` with the
1955 # mesh already, we don't want to undo its work
1958def process_varprob(
1959 varprob: Float[ArrayLike, ' p'] | None, max_split: UInt[Array, ' p']
1960) -> Float32[Array, ' p'] | None:
1961 """Convert varprob to log_s."""
1962 if varprob is None:
1963 return None
1964 varprob = jnp.asarray(varprob)
1965 assert varprob.shape == max_split.shape, 'varprob must have shape (p,)'
1966 varprob = error_if(varprob, varprob <= 0, 'varprob must be > 0')
1967 return jnp.log(varprob)
1970def predict_latent(
1971 x: UInt[Array, 'p m'],
1972 trace: MainTrace,
1973 test_points: Literal['none', 'autobatch', 'shard_and_autobatch'] = 'none',
1974) -> Float32[Array, 'ndpost m'] | Float32[Array, 'ndpost k m']:
1975 """Evaluate trees on already quantized `x`, and squash chains."""
1976 return evaluate_trace(x, trace, flatten_chains=True, test_points=test_points)
1979@jit(static_argnums=(5, 6, 7))
1980def predict(
1981 key: Key[Array, ''] | None,
1982 trace: MainTrace,
1983 x_test: UInt[Array, 'p m'],
1984 error_scale: Float[Array, ' m'] | Float[Array, 'k m'] | None,
1985 binary_indices: Int32[Array, ' kb'] | None,
1986 has_binary: bool,
1987 kind: PredictKind | str,
1988 test_points: Literal['none', 'autobatch', 'shard_and_autobatch'],
1989 /,
1990) -> (
1991 Float32[Array, ' m']
1992 | Float32[Array, 'k m']
1993 | Float32[Array, 'ndpost m']
1994 | Float32[Array, 'ndpost k m']
1995):
1996 """Implement `Bart.predict` by evaluating the trees on `x_test`."""
1997 # get latent i.e. bare sum-of-trees predictions
1998 latent = predict_latent(x_test, trace, test_points)
1999 return _predict_from_latent(
2000 key, trace, latent, error_scale, binary_indices, has_binary, kind
2001 )
2004@jit(static_argnums=(4, 5))
2005def predict_train(
2006 key: Key[Array, ''] | None,
2007 trace: MainTraceWithTrainPred,
2008 error_scale: Float[Array, ' n'] | Float[Array, 'k n'] | None,
2009 binary_indices: Int32[Array, ' kb'] | None,
2010 has_binary: bool,
2011 kind: PredictKind | str,
2012 /,
2013) -> (
2014 Float32[Array, ' n']
2015 | Float32[Array, 'k n']
2016 | Float32[Array, 'ndpost n']
2017 | Float32[Array, 'ndpost k n']
2018):
2019 """Implement `Bart.predict('train')` from the precomputed predictions."""
2020 latent = _flatten_chain_sample(
2021 trace.train_pred,
2022 chain_vmap_axes(trace).train_pred,
2023 trace_sample_axes(trace).train_pred,
2024 )
2025 return _predict_from_latent(
2026 key, trace, latent, error_scale, binary_indices, has_binary, kind
2027 )
2030def _predict_from_latent(
2031 key: Key[Array, ''] | None,
2032 trace: MainTrace,
2033 latent: Float32[Array, 'ndpost m'] | Float32[Array, 'ndpost k m'],
2034 error_scale: Float[Array, ' m'] | Float[Array, 'k m'] | None,
2035 binary_indices: Int32[Array, ' kb'] | None,
2036 has_binary: bool,
2037 kind: PredictKind | str,
2038) -> (
2039 Float32[Array, ' m']
2040 | Float32[Array, 'k m']
2041 | Float32[Array, 'ndpost m']
2042 | Float32[Array, 'ndpost k m']
2043):
2044 """Turn the sum-of-trees `latent` samples into the requested prediction kind."""
2045 if kind is PredictKind.latent_samples:
2046 return latent
2048 # sample posterior (uses latent directly, no probit squash needed)
2049 if kind is PredictKind.outcome_samples:
2050 assert key is not None
2051 return sample_outcome(
2052 key, trace, latent, error_scale, binary_indices, has_binary
2053 )
2055 # squash predictions to (0, 1) if probit; with heteroskedastic error scales
2056 # P(y=1) = Phi(latent / error_scale)
2057 if has_binary: # self._mcmc_state.z is not None
2058 # error_scale is (m,) or (k, m), so it broadcasts against latent
2059 arg = latent if error_scale is None else latent / error_scale
2060 if binary_indices is not None:
2061 # mixed: squash only the binary rows, leaving continuous rows as-is
2062 indexing = jnp.s_[..., binary_indices, :]
2063 mean_samples = latent.at[indexing].set(ndtr(arg[indexing]))
2064 else:
2065 mean_samples = ndtr(arg)
2066 else:
2067 mean_samples = latent
2069 # take mean or return samples
2070 if kind is PredictKind.mean:
2071 return mean_samples.mean(axis=0)
2072 return mean_samples
2075def _flatten_chain_sample(
2076 arr: Float[Array, '*shape'], chain_axis: int | None, sample_axis: int
2077) -> Float[Array, '*flat_shape']:
2078 """Fold the chain axis into the sample axis, matching `predict_latent`'s layout."""
2079 if chain_axis is None:
2080 return arr
2081 arr = jnp.moveaxis(arr, (chain_axis, sample_axis), (0, 1))
2082 return lax.collapse(arr, 0, 2)
2085@jit(static_argnums=(5,))
2086def sample_outcome(
2087 key: Key[Array, ''],
2088 trace: MainTrace,
2089 latent: Float32[Array, 'ndpost m'] | Float32[Array, 'ndpost k m'],
2090 error_scale: Float32[Array, ' m'] | Float32[Array, 'k m'] | None,
2091 binary_indices: Int32[Array, ' kb'] | None,
2092 has_binary: bool,
2093 /,
2094) -> Float32[Array, 'ndpost m'] | Float32[Array, 'ndpost k m']:
2095 """Sample from the posterior predictive distribution."""
2096 # move error_cov_inv chain axis to 0
2097 prec = chain_to_axis(trace.error_cov_inv, chain_vmap_axes(trace).error_cov_inv)
2099 if latent.ndim > 2: # multivariate case
2100 error_cov_inv = lax.collapse(prec, 0, -2)
2102 # Cholesky of precision: error_cov_inv = L @ L^T
2103 L = chol_with_gersh(error_cov_inv) # (ndpost, k, k)
2105 # Sample z ~ N(0, I) and solve L^T @ error = z
2106 # so error = L^{-T} z ~ N(0, L^{-T} L^{-1}) = N(0, Sigma)
2107 z = random.normal(key, latent.shape) # (ndpost, k, m)
2108 error = solve_triangular(L, z, trans='T', lower=True) # (ndpost, k, m)
2109 if error_scale is not None: 2109 ↛ 2122line 2109 didn't jump to line 2122 because the condition on line 2109 was always true
2110 # error_scale is (m,) or (k, m) so it always broadcasts right
2111 error *= error_scale
2112 else: # univariate
2113 # pure binary probit has unit-scale latent error; continuous scales it
2114 # by `sigma`. Either way, optionally rescaled per datapoint by error_scale.
2115 error = random.normal(key, latent.shape)
2116 if not has_binary:
2117 sigma = jnp.sqrt(jnp.reciprocal(prec)).reshape(-1)
2118 error *= sigma[..., None]
2119 if error_scale is not None: 2119 ↛ 2122line 2119 didn't jump to line 2122 because the condition on line 2119 was always true
2120 error *= error_scale[None, :]
2122 outcome = latent + error
2124 # convert binary outcomes via latent probit thresholding
2125 if binary_indices is not None:
2126 idx = jnp.s_[..., binary_indices, :]
2127 outcome = outcome.at[idx].set(jnp.where(outcome[idx] > 0, 1.0, 0.0))
2128 elif has_binary:
2129 outcome = jnp.where(outcome > 0, 1.0, 0.0)
2131 return outcome