Coverage for src/bartz/mcmcstep/_step.py: 99%
570 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-10-15 10:04 +0000
1# bartz/src/bartz/mcmcstep/_step.py
2#
3# Copyright (c) 2024-2026, The Bartz Contributors
4#
5# This file is part of bartz.
6#
7# Permission is hereby granted, free of charge, to any person obtaining a copy
8# of this software and associated documentation files (the "Software"), to deal
9# in the Software without restriction, including without limitation the rights
10# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
11# copies of the Software, and to permit persons to whom the Software is
12# furnished to do so, subject to the following conditions:
13#
14# The above copyright notice and this permission notice shall be included in all
15# copies or substantial portions of the Software.
16#
17# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
18# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
19# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
20# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
21# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
22# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
23# SOFTWARE.
25"""Implement `step`, `step_trees`, and the accept-reject logic."""
27from dataclasses import replace
28from functools import partial
29from typing import overload
31from equinox import AbstractVar
32from jax import lax, named_call, random, vmap
33from jax import numpy as jnp
34from jax.nn import softmax
35from jax.scipy.linalg import solve_triangular
36from jax.scipy.special import gammaln, logsumexp
37from jaxtyping import Array, Bool, Float, Float32, Int32, Key, Shaped, UInt, UInt32
39from bartz._jaxext import (
40 Module,
41 field,
42 float32_matmuls,
43 jit,
44 minimal_unsigned_dtype,
45 sliced_map,
46 split,
47 truncated_normal_onesided,
48 vmap_nodoc,
49)
50from bartz._jaxext.random import loggamma, poisson
51from bartz.grove import is_actual_leaf, var_histogram
52from bartz.mcmcstep._moves import Moves, propose_moves, split_range
53from bartz.mcmcstep._reduction import ReductionConfig
54from bartz.mcmcstep._state import (
55 DiagWishart,
56 Forest,
57 State,
58 StepConfig,
59 Wishart,
60 chol_with_gersh,
61 round_to_pow2,
62 scaled_error_cov_inv,
63 shard_map_state,
64 split_key_for_chains,
65 vmap_chains,
66)
69@jit(donate_argnums=(1,))
70@split_key_for_chains
71@shard_map_state
72@vmap_chains
73@float32_matmuls
74def step(key: Key[Array, ''], state: State) -> State:
75 """
76 Do one MCMC step.
78 Parameters
79 ----------
80 key
81 A jax random key.
82 state
83 A BART mcmc state, as created by `init`.
85 Returns
86 -------
87 The new BART mcmc state.
89 Notes
90 -----
91 The memory of the input state is re-used for the output state, so the input
92 state can not be used any more after calling `step`. All this applies
93 outside of `jax.jit`.
94 """
95 keys = split(key, 5)
97 state = step_resid_inexact_integral(state)
98 state = step_z(keys.pop(), state)
99 state = step_trees(keys.pop(), state)
101 state = step_leaf_prior_cov_inv(keys.pop(), state)
102 state = step_error_cov_inv(keys.pop(), state)
103 state = step_sparse(keys.pop(), state)
104 return step_config(state)
107@named_call
108def step_resid_inexact_integral(state: State) -> State:
109 """Accumulate the mean square of the residuals subject to rounding.
111 Adds to `State.resid_inexact_integral` the current mean square of the
112 residuals at or above ``2^(leaf_quantization+1)`` (in
113 `State.resid_eff_scale` units), below which the running residual updates
114 are exact because the `State.resid` dtype spacing still matches the leaf
115 quantum. Without quantization every update rounds and all residuals count.
116 """
117 # filter away residuals below quantization threshold
118 q = state.config.leaf_quantization
119 resid = state.resid
120 if q is None:
121 inexact = resid
122 else:
123 # the quantization grid follows `resid_eff_scale` while `resid` stays
124 # in `resid_unit` units, so the threshold carries their (power of
125 # two, hence exact) ratio; it is cast to the resid dtype to keep the
126 # comparison from widening the n-sized array
127 exact_below = 2.0 ** (q + 1) * state.resid_eff_scale / state.resid_unit
128 inexact = jnp.where(
129 jnp.abs(resid) >= exact_below[..., None].astype(resid.dtype), resid, 0
130 )
132 # mask missing datapoints if any
133 if state.inv_sdev_scale is not None:
134 inexact = jnp.where(state.inv_sdev_scale != 0, inexact, 0)
136 # compute mean square residuals above threshold
137 ms = jnp.einsum(
138 '...n,...n->...', inexact, inexact, preferred_element_type=jnp.float32
139 )
140 if state.config.data_sharded:
141 ms = lax.psum(ms, 'data')
142 ms /= jnp.maximum(state.n_non_missing, 1)
144 # accumulate to the running sum
145 return replace(state, resid_inexact_integral=state.resid_inexact_integral + ms)
148@named_call
149def step_trees(key: Key[Array, ''], state: State) -> State:
150 """
151 Forest sampling step of BART MCMC.
153 Parameters
154 ----------
155 key
156 A jax random key.
157 state
158 A BART mcmc state, as created by `init`.
160 Returns
161 -------
162 The new BART mcmc state.
164 Notes
165 -----
166 This function zeroes the proposal counters.
167 """
168 keys = split(key)
169 moves = propose_moves(keys.pop(), state.forest)
170 return accept_moves_and_sample_leaves(keys.pop(), state, moves)
173@named_call
174def accept_moves_and_sample_leaves(
175 key: Key[Array, ''], state: State, moves: Moves
176) -> State:
177 """
178 Accept or reject the proposed moves and sample the new leaf values.
180 Parameters
181 ----------
182 key
183 A jax random key.
184 state
185 A valid BART mcmc state.
186 moves
187 The proposed moves, see `propose_moves`.
189 Returns
190 -------
191 A new (valid) BART mcmc state.
192 """
193 pso = accept_moves_parallel_stage(key, state, moves)
194 state, moves = accept_moves_sequential_stage(pso)
195 return accept_moves_final_stage(state, moves)
198class Counts(Module):
199 """Number of datapoints in the nodes involved in proposed moves for each tree."""
201 lrt: UInt[Array, '*num_trees 3']
202 """Number of datapoints in the left child, right child, and parent
203 (``= left + right``), stacked along the trailing axis."""
206class PreLkV(Module):
207 """Non-sequential terms of the likelihood ratio for each tree.
209 These terms are derived from the leaf precompute terms (`PreLf`) gathered
210 at the nodes involved in each move. The terms for the left child, right
211 child, and their join (the parent node) are stacked along the axis right
212 after the tree axis. Each term is, in the univariate case, the scalar
214 ``error_cov_inv^2 / (leaf_prior_cov_inv + n * error_cov_inv)``.
216 In the multivariate homoskedastic or scalar-error-scale case, this is the matrix term
218 ``error_cov_inv @ inv(leaf_prior_cov_inv + n * error_cov_inv) @ error_cov_inv``.
220 In the multivariate case with per-component error scales, this is instead
222 ``chol(leaf_prior_cov_inv + n * error_cov_inv)``
224 ``n`` is the number of datapoints in the node, or the likelihood precision
225 scale in the heteroskedastic case.
226 """
228 # `log_sqrt_term` is declared before `lrt` so its single (union-free)
229 # annotation binds the variadic `*num_trees` axis first; otherwise the
230 # runtime typechecker can greedily mis-bind `*num_trees` against the `k`
231 # axis of the `... | ... k k` union (the multivariate and univariate
232 # layouts are rank-ambiguous).
233 log_sqrt_term: Float32[Array, '*num_trees']
234 """The logarithm of the square root term of the likelihood ratio."""
236 lrt: Float32[Array, '*num_trees 3'] | Float32[Array, '*num_trees 3 k k']
237 """Scaled full conditional variance, scaled covariance, or precision
238 cholesky, for the left child, right child, and their join."""
241class PreLf(Module):
242 """Pre-computed terms used to sample leaves from their posterior.
244 These terms can be computed in parallel across trees.
246 For each tree and leaf, the terms are scalars in the univariate case
247 (`PreLfUV`), and matrices/vectors in the multivariate case (`PreLfMV`,
248 `PreLfMVHet`).
250 Abstract base: the layouts differ in rank, so they live in concrete
251 subclasses with union-free annotations; a single class carrying a shape
252 union would make the greedy variadic mis-bind against the ``k`` axes under
253 the runtime typechecker. The concrete class also tags the meaning of
254 `mean_factor`, which drives the dispatch in `precompute_likelihood_terms`
255 and in the sequential stage. The ``num_trees`` axis is variadic so the same
256 annotations also match a per-element layout if vmapped over trees.
257 """
259 mean_factor: AbstractVar[
260 Float32[Array, '*num_trees tree_size']
261 | Float32[Array, '*num_trees k k tree_size']
262 ]
263 """The factor to be right-multiplied by the sum of the scaled residuals to
264 obtain the posterior mean."""
266 centered_leaves: AbstractVar[
267 Float32[Array, '*num_trees tree_size']
268 | Float32[Array, '*num_trees k tree_size']
269 ]
270 """The mean-zero normal values to be added to the posterior mean to
271 obtain the posterior leaf samples."""
274class PreLfUV(PreLf):
275 """`PreLf` for the univariate case."""
277 mean_factor: Float32[Array, '*num_trees tree_size']
278 """``error_cov_inv / prec``, where ``prec`` is the posterior precision of
279 the leaf."""
281 centered_leaves: Float32[Array, '*num_trees tree_size']
282 """Zero-mean normal draws with the posterior variance of each leaf."""
285class PreLfMV(PreLf):
286 """`PreLf` for the multivariate homoskedastic or scalar-error-scale case."""
288 mean_factor: Float32[Array, '*num_trees k k tree_size']
289 """``error_cov_inv @ inv(prec)``, where ``prec`` is the posterior precision
290 of the leaf."""
292 centered_leaves: Float32[Array, '*num_trees k tree_size']
293 """Zero-mean normal draws with the posterior covariance of each leaf."""
295 logdet_prec: Float32[Array, '*num_trees tree_size']
296 """The log-determinant of the posterior precision of each leaf."""
299class PreLfMVHet(PreLf):
300 """`PreLf` for the multivariate case with per-component error scales."""
302 mean_factor: Float32[Array, '*num_trees k k tree_size']
303 """The lower Cholesky factor of the posterior precision of each leaf; the
304 mean solve happens downstream in the sequential stage."""
306 centered_leaves: Float32[Array, '*num_trees k tree_size']
307 """Zero-mean normal draws with the posterior covariance of each leaf."""
310class ParallelStageOut(Module):
311 """The output of `accept_moves_parallel_stage`."""
313 state: State
314 """A partially updated BART mcmc state."""
316 moves: Moves
317 """The proposed moves, with `partial_ratio` set to `None` and
318 `log_trans_prior_ratio` set to its final value."""
320 # `num_trees` stays a fixed (non-variadic) axis: `ParallelStageOut` is always
321 # built with the tree axis present (never per tree under vmap), so the union
322 # is disambiguated by rank/dtype and needs no anchor (cf. `PreLf`).
323 prec_trees: (
324 Float32[Array, 'num_trees tree_size']
325 | UInt32[Array, 'num_trees tree_size']
326 | Float32[Array, 'num_trees k k tree_size']
327 )
328 """The likelihood precision scale in each potential or actual leaf node."""
330 prelkv: PreLkV
331 """Object with pre-computed terms of the likelihood ratios."""
333 prelf: PreLf
334 """Object with pre-computed terms of the leaf samples."""
337@named_call
338def accept_moves_parallel_stage(
339 key: Key[Array, ''], state: State, moves: Moves
340) -> ParallelStageOut:
341 """
342 Pre-compute quantities used to accept moves, in parallel across trees.
344 Parameters
345 ----------
346 key
347 A jax random key.
348 state
349 A BART mcmc state.
350 moves
351 The proposed moves, see `propose_moves`.
353 Returns
354 -------
355 An object with all that could be done in parallel.
356 """
357 # apply the prunes pending from the previous step to the leaf indices,
358 # then, where the new move is grow, modify the state like the move was
359 # accepted.
360 leaf_indices = apply_moves_to_leaf_indices(
361 state.forest.leaf_indices, state.forest.to_prune, state.forest.move_node
362 )
363 state = replace(
364 state,
365 forest=replace(
366 state.forest,
367 var_tree=moves.var_tree,
368 leaf_indices=apply_grow_to_indices(moves, leaf_indices, state.X),
369 leaf_tree=adapt_leaf_trees_to_grow_indices(state.forest.leaf_tree, moves),
370 ),
371 )
373 # update the cached number of datapoints per leaf at the nodes involved
374 # in the moves
375 if (
376 state.forest.min_points_per_decision_node is not None
377 or state.forest.min_points_per_leaf is not None
378 or state.prec_scale is None
379 ):
380 assert state.forest.count_tree is not None
381 count_trees, move_counts = compute_count_trees(
382 state.forest.count_tree, state.forest.leaf_indices, moves, state.config
383 )
384 state = replace(state, forest=replace(state.forest, count_tree=count_trees))
386 # affluence of the nodes touched by each move: whether they would be
387 # growable as leaves (admissible rule + enough datapoints). The children
388 # must also lie within the heap, i.e. not be at the bottom level; the
389 # parent always does. These feed the transition ratio and the final
390 # `affluence_tree` update.
391 _, half = state.forest.var_tree.shape
392 lrt_affluent = (moves.lrt_nodes < half) & moves.lrt_growable
393 if state.forest.min_points_per_decision_node is not None:
394 lrt_affluent &= move_counts.lrt >= state.forest.min_points_per_decision_node
395 moves = replace(moves, lrt_affluent=lrt_affluent)
397 # veto grove move if new leaves don't have enough datapoints
398 if state.forest.min_points_per_leaf is not None:
399 moves = replace(
400 moves,
401 allowed=moves.allowed
402 & jnp.all(
403 move_counts.lrt[..., :2] >= state.forest.min_points_per_leaf, axis=-1
404 ),
405 )
407 # update the cached number of datapoints per leaf, weighted by error
408 # precision scale, at the nodes involved in the moves
409 if state.prec_scale is None:
410 prec_trees = count_trees
411 else:
412 assert state.forest.prec_tree is not None
413 prec_trees = compute_prec_trees(
414 state.forest.prec_tree,
415 state.prec_scale,
416 state.forest.leaf_indices,
417 moves,
418 state.config,
419 )
420 state = replace(state, forest=replace(state.forest, prec_tree=prec_trees))
422 # compute some missing information about moves
423 moves = complete_ratio(moves, state.forest.p_nonterminal)
424 save_ratios = state.forest.log_likelihood is not None
425 state = replace(
426 state,
427 forest=replace(
428 state.forest,
429 grow_prop_count=jnp.sum(moves.grow),
430 prune_prop_count=jnp.sum(moves.allowed & ~moves.grow),
431 log_trans_prior=moves.log_trans_prior_ratio if save_ratios else None,
432 ),
433 )
435 # `prec_trees` and the per-leaf residual sums are in stored `prec_scale`
436 # units, so their common `inv_sdev_unit ** 2` factor is folded into the
437 # error precision once here instead
438 error_cov_inv = scaled_error_cov_inv(state)
439 leaf_prior_cov_inv = state.forest.leaf_prior_cov_inv.value
440 prelf = precompute_leaf_terms(key, prec_trees, error_cov_inv, leaf_prior_cov_inv)
441 prelkv = precompute_likelihood_terms(
442 error_cov_inv, leaf_prior_cov_inv, prelf, moves
443 )
445 return ParallelStageOut(
446 state=state, moves=moves, prec_trees=prec_trees, prelkv=prelkv, prelf=prelf
447 )
450@named_call
451def apply_grow_to_indices(
452 moves: Moves, leaf_indices: UInt[Array, 'num_trees n'], X: UInt[Array, 'p n']
453) -> UInt[Array, 'num_trees n']:
454 """
455 Update the leaf indices to apply a grow move.
457 Parameters
458 ----------
459 moves
460 The proposed moves, see `propose_moves`.
461 leaf_indices
462 The index of the leaf each datapoint falls into, with the prunes
463 pending from the previous step already applied.
464 X
465 The predictors matrix.
467 Returns
468 -------
469 The updated leaf indices.
470 """
471 return _apply_grow_to_indices(moves, leaf_indices, X)
474@partial(vmap_nodoc, in_axes=(0, 0, None))
475def _apply_grow_to_indices(
476 moves: Moves, leaf_indices: UInt[Array, ' n'], X: UInt[Array, 'p n']
477) -> UInt[Array, ' n']:
478 """Implement `apply_grow_to_indices`."""
479 left_child = moves.lrt_nodes[0].astype(leaf_indices.dtype)
480 x: UInt[Array, ' n'] = X[moves.grow_var, :]
481 go_right = x >= moves.grow_split
482 tree_size = jnp.array(2 * moves.var_tree.size)
483 node_to_update = jnp.where(moves.grow, moves.lrt_nodes[2], tree_size)
484 return jnp.where(
485 leaf_indices == node_to_update, left_child + go_right, leaf_indices
486 )
489def _fill_lrt_total(lrt: Shaped[Array, '*k_k 3']) -> Shaped[Array, '*k_k 3']:
490 """Set the total slot of stacked (left, right, total) values to left + right.
492 The left and right slots pass through unchanged, the stale value in the
493 total slot is ignored. Implemented with fusable elementwise operations.
494 """
495 total = lrt[..., 0] + lrt[..., 1]
496 return jnp.where(jnp.arange(3) == 2, total[..., None], lrt)
499@overload
500def _compute_count_or_prec_trees(
501 prec_scale: None,
502 trees: UInt32[Array, 'num_trees tree_size'],
503 leaf_indices: UInt[Array, 'num_trees n'],
504 moves: Moves,
505 config: StepConfig,
506) -> tuple[UInt32[Array, 'num_trees tree_size'], Counts]: ...
509@overload
510def _compute_count_or_prec_trees(
511 prec_scale: Float[Array, ' n'] | Float[Array, 'k k n'],
512 trees: Float32[Array, 'num_trees tree_size']
513 | Float32[Array, 'num_trees k k tree_size'],
514 leaf_indices: UInt[Array, 'num_trees n'],
515 moves: Moves,
516 config: StepConfig,
517) -> (
518 tuple[Float32[Array, 'num_trees tree_size'], None]
519 | tuple[Float32[Array, 'num_trees k k tree_size'], None]
520): ...
523def _compute_count_or_prec_trees(
524 prec_scale: Float[Array, ' n'] | Float[Array, 'k k n'] | None,
525 trees: UInt32[Array, 'num_trees tree_size']
526 | Float32[Array, 'num_trees tree_size']
527 | Float32[Array, 'num_trees k k tree_size'],
528 leaf_indices: UInt[Array, 'num_trees n'],
529 moves: Moves,
530 config: StepConfig,
531) -> (
532 tuple[UInt32[Array, 'num_trees tree_size'], Counts]
533 | tuple[Float32[Array, 'num_trees tree_size'], None]
534 | tuple[Float32[Array, 'num_trees k k tree_size'], None]
535):
536 """Implement `compute_count_trees` and `compute_prec_trees`."""
538 def all_trees() -> (
539 tuple[UInt32[Array, 'num_trees tree_size'], Counts]
540 | tuple[Float32[Array, 'num_trees tree_size'], None]
541 | tuple[Float32[Array, 'num_trees k k tree_size'], None]
542 ):
543 compute = vmap(_compute_count_or_prec_tree, in_axes=(None, 0, 0, 0, None))
544 return compute(prec_scale, trees, leaf_indices, moves, config)
546 if config.prec_count_num_trees is None:
547 return all_trees()
549 batch_size = config.prec_count_num_trees
551 def compute(
552 args: tuple[
553 UInt32[Array, ' tree_size']
554 | Float32[Array, ' tree_size']
555 | Float32[Array, 'k k tree_size'],
556 UInt[Array, ' n'],
557 Moves,
558 ],
559 ) -> (
560 tuple[UInt32[Array, ' tree_size'], Counts]
561 | tuple[Float32[Array, ' tree_size'], None]
562 | tuple[Float32[Array, 'k k tree_size'], None]
563 ):
564 tree, leaf_indices, moves = args
565 return _compute_count_or_prec_tree(
566 prec_scale, tree, leaf_indices, moves, config
567 )
569 def tree_batches() -> (
570 tuple[UInt32[Array, 'num_trees tree_size'], Counts]
571 | tuple[Float32[Array, 'num_trees tree_size'], None]
572 | tuple[Float32[Array, 'num_trees k k tree_size'], None]
573 ):
574 # sliced_map instead of lax.map because under the chain vmap lax.map's
575 # reshape of the tree axis becomes a transpose of `leaf_indices` that
576 # xla materializes in full (cf. `SeqStageInAllTrees.leaf_indices`)
577 return sliced_map(compute, (trees, leaf_indices, moves), batch_size=batch_size)
579 # the tree batching bounds the reduction temporaries on cpu; on gpu those
580 # fuse anyway
581 return lax.platform_dependent(cpu=tree_batches, cuda=all_trees)
584def _compute_count_or_prec_tree(
585 prec_scale: Float[Array, ' n'] | Float[Array, 'k k n'] | None,
586 tree: UInt32[Array, ' tree_size']
587 | Float32[Array, ' tree_size']
588 | Float32[Array, 'k k tree_size'],
589 leaf_indices: UInt[Array, ' n'],
590 moves: Moves,
591 config: StepConfig,
592) -> (
593 tuple[UInt32[Array, ' tree_size'], Counts]
594 | tuple[Float32[Array, ' tree_size'], None]
595 | tuple[Float32[Array, 'k k tree_size'], None]
596):
597 """Update the cached count or precision tree for a single tree."""
598 (tree_size,) = moves.var_tree.shape
599 tree_size *= 2
601 if prec_scale is None:
602 value = 1
603 dtype = jnp.uint32
604 reduction_config = config.count_reduction_config
605 else:
606 value = prec_scale
607 dtype = jnp.float32
608 reduction_config = config.prec_reduction_config
610 # the cached tree is valid at the leaves, and the move only changes the
611 # values at the nodes it involves, so reduce into the move's children alone:
612 # the contiguous pair (left, right) = (2 * node, 2 * node + 1) = lrt_nodes[:2]
613 lr = reduction_config._reduce( # noqa: SLF001
614 value,
615 leaf_indices,
616 size=tree_size,
617 subset_start=moves.lrt_nodes[0],
618 subset_length=2,
619 dtype=dtype,
620 data_sharded=config.data_sharded,
621 )
623 # write the children sums into the cache along with their total at the
624 # parent node (a non-leaf in the post-grow indexing the reduce runs on);
625 # the precision-scaled version of the counts is not needed because the likelihood
626 # terms are derived from the leaf terms
627 total = lr[..., 0] + lr[..., 1]
628 lrt = jnp.concatenate([lr, total[..., None]], axis=-1)
629 tree = tree.at[..., moves.lrt_nodes].set(lrt)
631 if prec_scale is None:
632 return tree, Counts(lrt=lrt)
633 else:
634 return tree, None
637@named_call
638def compute_count_trees(
639 count_trees: UInt32[Array, 'num_trees tree_size'],
640 leaf_indices: UInt[Array, 'num_trees n'],
641 moves: Moves,
642 config: StepConfig,
643) -> tuple[UInt32[Array, 'num_trees tree_size'], Counts]:
644 """
645 Update the cached number of datapoints per leaf at the moves' nodes.
647 Parameters
648 ----------
649 count_trees
650 The cached number of points in each leaf; valid at the leaves of the
651 pre-move trees.
652 leaf_indices
653 The index of the leaf each datapoint falls into, with the deeper version
654 of the tree (post-GROW, pre-PRUNE).
655 moves
656 The proposed moves, see `propose_moves`.
657 config
658 The MCMC configuration.
660 Returns
661 -------
662 count_trees : UInt32[Array, 'num_trees tree_size']
663 The updated cache, valid in each potential or actual leaf node.
664 counts : Counts
665 The counts of the number of points in the leaves grown or pruned by the
666 moves.
667 """
668 return _compute_count_or_prec_trees(None, count_trees, leaf_indices, moves, config)
671@named_call
672def compute_prec_trees(
673 prec_trees: Float32[Array, 'num_trees tree_size']
674 | Float32[Array, 'num_trees k k tree_size'],
675 prec_scale: Float[Array, ' n'] | Float[Array, 'k k n'],
676 leaf_indices: UInt[Array, 'num_trees n'],
677 moves: Moves,
678 config: StepConfig,
679) -> Float32[Array, 'num_trees tree_size'] | Float32[Array, 'num_trees k k tree_size']:
680 """
681 Update the cached per-leaf likelihood precision scale at the moves' nodes.
683 Parameters
684 ----------
685 prec_trees
686 The cached likelihood precision scale in each leaf; valid at the leaves
687 of the pre-move trees.
688 prec_scale
689 The scale of the precision of the error on each datapoint.
690 leaf_indices
691 The index of the leaf each datapoint falls into, with the deeper version
692 of the tree (post-GROW, pre-PRUNE).
693 moves
694 The proposed moves, see `propose_moves`.
695 config
696 The MCMC configuration.
698 Returns
699 -------
700 The updated cache, valid in each potential or actual leaf node.
701 """
702 trees, _ = _compute_count_or_prec_trees(
703 prec_scale, prec_trees, leaf_indices, moves, config
704 )
705 return trees
708@partial(vmap_nodoc, in_axes=(0, None))
709def complete_ratio(moves: Moves, p_nonterminal: Float32[Array, ' tree_size']) -> Moves:
710 """
711 Complete non-likelihood MH ratio calculation.
713 This function adds the probability of choosing a prune move over the grow
714 move in the inverse transition, and the prior odds that the modified node
715 is nonterminal with terminal children.
717 Parameters
718 ----------
719 moves
720 The proposed moves. Must have already been updated to keep into account
721 the thresholds on the number of datapoints per node, this happens in
722 `accept_moves_parallel_stage`.
723 p_nonterminal
724 The a priori probability of each node being nonterminal conditional on
725 its ancestors, including at the maximum depth where it should be zero.
727 Returns
728 -------
729 The updated moves, with `partial_ratio=None` and `log_trans_prior_ratio` set.
730 """
731 assert moves.lrt_affluent is not None
733 # can the children be grown by the proposal? `lrt_affluent` already folds
734 # in the `min_points_per_decision_node` threshold, because the grow
735 # proposal draws from the pool of leaves that pass it. This enters only the
736 # transition probability.
738 # p_prune if grow
739 other_growable_leaves = moves.num_growable >= 2
740 grow_again_allowed = other_growable_leaves | jnp.any(moves.lrt_affluent[:2])
741 grow_p_prune = jnp.where(grow_again_allowed, 0.5, 1.0)
743 # p_prune if prune
744 prune_p_prune = jnp.where(moves.num_growable, 0.5, 1)
746 # select p_prune
747 p_prune = jnp.where(moves.grow, grow_p_prune, prune_p_prune)
749 # prior odds of the node being nonterminal, times the prior probability of
750 # both children being terminal. The children terminality uses the
751 # admissibility ignoring counts, because the count thresholds restrict the
752 # support of the prior rather than the conditional non-terminal
753 # probabilities, so they cancel in the prior ratio between two trees in
754 # the support. The fill value avoids a 0 and then an inf in the log if the
755 # move is not allowed and the indices are out of bounds.
756 pnt = p_nonterminal.at[moves.lrt_nodes].get(mode='fill', fill_value=0.5)
757 prior_ratio = pnt[2] / (1 - pnt[2]) * jnp.prod(1 - pnt[:2] * moves.lrt_growable[:2])
759 assert moves.partial_ratio is not None
760 return replace(
761 moves,
762 log_trans_prior_ratio=jnp.log(moves.partial_ratio * prior_ratio * p_prune),
763 partial_ratio=None,
764 )
767@named_call
768def adapt_leaf_trees_to_grow_indices(
769 leaf_trees: Float[Array, 'num_trees tree_size']
770 | Float[Array, 'num_trees k tree_size'],
771 moves: Moves,
772) -> Float[Array, 'num_trees tree_size'] | Float[Array, 'num_trees k tree_size']:
773 """
774 Modify leaves such that post-grow indices work on the original tree.
776 The value of the leaf to grow is copied to what would be its children if the
777 grow move was accepted.
779 Parameters
780 ----------
781 leaf_trees
782 The leaf values.
783 moves
784 The proposed moves, see `propose_moves`.
786 Returns
787 -------
788 The modified leaf values.
789 """
790 return _adapt_leaf_trees_to_grow_indices(leaf_trees, moves)
793@vmap_nodoc
794def _adapt_leaf_trees_to_grow_indices(
795 leaf_trees: Float[Array, ' tree_size'] | Float[Array, ' k tree_size'], moves: Moves
796) -> Float[Array, ' tree_size'] | Float[Array, ' k tree_size']:
797 """Implement `adapt_leaf_trees_to_grow_indices`."""
798 # the parent slot is written back unchanged to share a single scatter
799 values_at_node = leaf_trees[..., moves.lrt_nodes[2]]
800 return leaf_trees.at[
801 ..., jnp.where(moves.grow, moves.lrt_nodes, leaf_trees.size)
802 ].set(values_at_node[..., None])
805def _logdet_from_chol(L: Float32[Array, '... k k']) -> Float32[Array, '...']:
806 """Compute logdet of A = LL' via Cholesky (sum of log of diag^2)."""
807 diags: Float32[Array, '... k'] = jnp.diagonal(L, axis1=-2, axis2=-1)
808 return 2.0 * jnp.sum(jnp.log(diags), axis=-1)
811def compute_B(
812 error_cov_inv: Float32[Array, 'k k'], resid: Float32[Array, 'k k *tree_size']
813) -> Float32[Array, ' k *tree_size']:
814 """Compute the leaf score from the leaf precision-scaled sum of residuals."""
815 return jnp.einsum('ab,ab...->a...', error_cov_inv, resid)
818def _precompute_leaf_terms_uv(
819 key: Key[Array, ''],
820 prec_trees: Float32[Array, 'num_trees tree_size']
821 | UInt32[Array, 'num_trees tree_size'],
822 error_cov_inv: Float32[Array, ''],
823 leaf_prior_cov_inv: Float32[Array, ''],
824 z: Float32[Array, 'num_trees tree_size'] | None = None,
825) -> PreLfUV:
826 prec_lk = prec_trees * error_cov_inv
827 var_post = jnp.reciprocal(prec_lk + leaf_prior_cov_inv)
828 if z is None:
829 z = random.normal(key, prec_trees.shape, error_cov_inv.dtype)
830 return PreLfUV(
831 mean_factor=var_post * error_cov_inv,
832 # | mean = mean_lk * prec_lk * var_post
833 # | resid_tree = mean_lk * prec_tree -->
834 # | --> mean_lk = resid_tree / prec_tree (kind of)
835 # | mean_factor =
836 # | = mean / resid_tree =
837 # | = resid_tree / prec_tree * prec_lk * var_post / resid_tree =
838 # | = 1 / prec_tree * prec_tree / sigma2 * var_post =
839 # | = var_post / sigma2
840 centered_leaves=z * jnp.sqrt(var_post),
841 )
844def _precompute_leaf_terms_mv(
845 key: Key[Array, ''],
846 prec_trees: Float32[Array, 'num_trees tree_size']
847 | UInt32[Array, 'num_trees tree_size'],
848 error_cov_inv: Float32[Array, 'k k'],
849 leaf_prior_cov_inv: Float32[Array, 'k k'],
850 z: Float32[Array, 'num_trees k tree_size'] | None = None,
851) -> PreLfMV:
852 num_trees, tree_size = prec_trees.shape
853 k, _ = error_cov_inv.shape
854 if z is None: 854 ↛ 857line 854 didn't jump to line 857 because the condition on line 854 was always true
855 z = random.normal(key, (num_trees, k, tree_size))
857 def per_leaf(
858 prec: Float32[Array, ''] | UInt32[Array, ''], z: Float32[Array, ' k']
859 ) -> tuple[Float32[Array, 'k k'], Float32[Array, ' k'], Float32[Array, '']]:
860 L_prec = chol_with_gersh(leaf_prior_cov_inv + prec * error_cov_inv)
861 Y = solve_triangular(L_prec, error_cov_inv, lower=True)
862 mean_factor = solve_triangular(L_prec, Y, trans='T', lower=True).mT
863 centered = solve_triangular(L_prec, z[:, None], trans='T', lower=True).squeeze(
864 -1
865 )
866 # only a few leaves per tree end up using their logdet, but reducing
867 # right away is lighter on memory than storing diagonals for later
868 return mean_factor, centered, _logdet_from_chol(L_prec)
870 # vmap over trees then over leaves; the leaf axis is trailing in both
871 # `prec_trees`/`z` (in_axes) and the stored output (out_axes=-1)
872 return PreLfMV(*vmap(vmap(per_leaf, in_axes=(0, -1), out_axes=-1))(prec_trees, z))
875def _precompute_leaf_terms_mv_het(
876 key: Key[Array, ''],
877 prec_trees: Float32[Array, 'num_trees k k tree_size'],
878 error_cov_inv: Float32[Array, 'k k'],
879 leaf_prior_cov_inv: Float32[Array, 'k k'],
880 z: Float32[Array, 'num_trees k tree_size'] | None = None,
881) -> PreLfMVHet:
882 num_trees, k, _, tree_size = prec_trees.shape
883 if z is None: 883 ↛ 886line 883 didn't jump to line 886 because the condition on line 883 was always true
884 z = random.normal(key, (num_trees, k, tree_size))
886 def per_leaf(
887 prec: Float32[Array, 'k k'], z: Float32[Array, ' k']
888 ) -> tuple[Float32[Array, 'k k'], Float32[Array, ' k']]:
889 # mean_factor stores the precision cholesky itself; the mean solve happens
890 # downstream in `accept_move_and_sample_leaves`
891 L_prec = chol_with_gersh(leaf_prior_cov_inv + error_cov_inv * prec)
892 centered = solve_triangular(L_prec, z[:, None], trans='T', lower=True).squeeze(
893 -1
894 )
895 return L_prec, centered
897 # vmap over trees then over leaves; the leaf axis is trailing in both
898 # `prec_trees`/`z` (in_axes=-1) and the stored output (out_axes=-1)
899 return PreLfMVHet(
900 *vmap(vmap(per_leaf, in_axes=(-1, -1), out_axes=-1))(prec_trees, z)
901 )
904@named_call
905def precompute_leaf_terms(
906 key: Key[Array, ''],
907 prec_trees: Float32[Array, 'num_trees tree_size']
908 | UInt32[Array, 'num_trees tree_size']
909 | Float32[Array, 'num_trees k k tree_size'],
910 error_cov_inv: Float32[Array, ''] | Float32[Array, 'k k'],
911 leaf_prior_cov_inv: Float32[Array, ''] | Float32[Array, 'k k'],
912 z: Float32[Array, 'num_trees tree_size']
913 | Float32[Array, 'num_trees k tree_size']
914 | None = None,
915) -> PreLf:
916 """
917 Pre-compute terms used to sample leaves from their posterior.
919 Handles both univariate and multivariate cases based on the shape of the
920 input arrays.
922 Parameters
923 ----------
924 key
925 A jax random key.
926 prec_trees
927 The likelihood precision scale in each potential or actual leaf node.
928 error_cov_inv
929 The inverse error variance (univariate) or the inverse of error
930 covariance matrix (multivariate). If `prec_scale` is set, this is the
931 global error precision factor, with the squared `inv_sdev_unit` folded in
932 (see `bartz.mcmcstep.State.inv_sdev_unit`).
933 leaf_prior_cov_inv
934 The inverse prior variance of each leaf (univariate) or the inverse of
935 prior covariance matrix of each leaf (multivariate).
936 z
937 Optional standard normal noise to use for sampling the centered leaves.
938 This is intended for testing purposes only.
940 Returns
941 -------
942 Pre-computed terms for leaf sampling.
943 """
944 if error_cov_inv.ndim == 0:
945 return _precompute_leaf_terms_uv(
946 key, prec_trees, error_cov_inv, leaf_prior_cov_inv, z
947 )
948 elif prec_trees.ndim == 4:
949 return _precompute_leaf_terms_mv_het(
950 key, prec_trees, error_cov_inv, leaf_prior_cov_inv, z
951 )
952 else:
953 return _precompute_leaf_terms_mv(
954 key, prec_trees, error_cov_inv, leaf_prior_cov_inv, z
955 )
958@vmap_nodoc
959def _gather_lrt(
960 leaf_values: Float32[Array, '*k_k tree_size'], lrt_nodes: Int32[Array, ' 3']
961) -> Float32[Array, ' 3 *k_k']:
962 """Gather per-tree leaf values at the left child, right child, and parent."""
963 return jnp.moveaxis(leaf_values[..., lrt_nodes], -1, 0)
966def _precompute_likelihood_terms_uv(
967 error_cov_inv: Float32[Array, ''],
968 leaf_prior_cov_inv: Float32[Array, ''],
969 prelf: PreLfUV,
970 lrt_nodes: Int32[Array, 'num_trees 3'],
971) -> PreLkV:
972 # mean_factor is error_cov_inv / prec, complete the sandwich
973 lrt = error_cov_inv * _gather_lrt(prelf.mean_factor, lrt_nodes)
974 # the same value with the prior-only precision, computed with the same
975 # operations as in `_precompute_leaf_terms_uv` such that it matches `lrt`
976 # bitwise on empty nodes and the ratio is exactly 1 without data
977 prior_lrt = error_cov_inv * (jnp.reciprocal(leaf_prior_cov_inv) * error_cov_inv)
978 log_sqrt_term = jnp.log(lrt[..., 0] * lrt[..., 1] / (prior_lrt * lrt[..., 2])) / 2
979 return PreLkV(lrt=lrt, log_sqrt_term=log_sqrt_term)
982def _precompute_likelihood_terms_mv(
983 error_cov_inv: Float32[Array, 'k k'],
984 leaf_prior_cov_inv: Float32[Array, 'k k'],
985 prelf: PreLfMV,
986 lrt_nodes: Int32[Array, 'num_trees 3'],
987) -> PreLkV:
988 logdet_prior = _logdet_from_chol(chol_with_gersh(leaf_prior_cov_inv))
989 logdet_prec = _gather_lrt(prelf.logdet_prec, lrt_nodes)
990 log_sqrt_term = (logdet_prior + logdet_prec @ jnp.array([-1.0, -1.0, 1.0])) / 2
992 # mean_factor is error_cov_inv @ inv(prec), complete the sandwich
993 mean_factor = _gather_lrt(prelf.mean_factor, lrt_nodes) # (num_trees, 3, k, k)
994 return PreLkV(lrt=mean_factor @ error_cov_inv, log_sqrt_term=log_sqrt_term)
997def _precompute_likelihood_terms_mv_het(
998 leaf_prior_cov_inv: Float32[Array, 'k k'],
999 prelf: PreLfMVHet,
1000 lrt_nodes: Int32[Array, 'num_trees 3'],
1001) -> PreLkV:
1002 logdet_prior = _logdet_from_chol(chol_with_gersh(leaf_prior_cov_inv))
1004 # mean_factor is the precision cholesky itself
1005 L = _gather_lrt(prelf.mean_factor, lrt_nodes) # (num_trees, 3, k, k)
1006 log_sqrt_term = (
1007 logdet_prior + _logdet_from_chol(L) @ jnp.array([-1.0, -1.0, 1.0])
1008 ) / 2
1009 return PreLkV(lrt=L, log_sqrt_term=log_sqrt_term)
1012@named_call
1013def precompute_likelihood_terms(
1014 error_cov_inv: Float32[Array, ''] | Float32[Array, 'k k'],
1015 leaf_prior_cov_inv: Float32[Array, ''] | Float32[Array, 'k k'],
1016 prelf: PreLf,
1017 moves: Moves,
1018) -> PreLkV:
1019 """
1020 Pre-compute terms used in the likelihood ratio of the acceptance step.
1022 The likelihood ratio terms are mostly a subset of the leaf sampling terms,
1023 so they are derived from `prelf`, gathered at the nodes involved in the
1024 moves.
1026 Parameters
1027 ----------
1028 error_cov_inv
1029 The inverse error variance (univariate) or the inverse of the error
1030 covariance matrix (multivariate). If `prec_scale` is set, this is the
1031 global error precision factor, with the squared `inv_sdev_unit` folded in
1032 (see `bartz.mcmcstep.State.inv_sdev_unit`).
1033 leaf_prior_cov_inv
1034 The inverse prior variance of each leaf (univariate) or the inverse of
1035 prior covariance matrix of each leaf (multivariate).
1036 prelf
1037 The pre-computed terms of the leaf sampling, see `precompute_leaf_terms`.
1038 moves
1039 The proposed moves, see `propose_moves`.
1041 Returns
1042 -------
1043 Pre-computed terms of the likelihood ratio, one per tree.
1044 """
1045 if isinstance(prelf, PreLfUV):
1046 return _precompute_likelihood_terms_uv(
1047 error_cov_inv, leaf_prior_cov_inv, prelf, moves.lrt_nodes
1048 )
1049 elif isinstance(prelf, PreLfMVHet):
1050 return _precompute_likelihood_terms_mv_het(
1051 leaf_prior_cov_inv, prelf, moves.lrt_nodes
1052 )
1053 else:
1054 assert isinstance(prelf, PreLfMV)
1055 return _precompute_likelihood_terms_mv(
1056 error_cov_inv, leaf_prior_cov_inv, prelf, moves.lrt_nodes
1057 )
1060@named_call
1061def accept_moves_sequential_stage(pso: ParallelStageOut) -> tuple[State, Moves]:
1062 """
1063 Accept/reject the moves one tree at a time.
1065 This is the most performance-sensitive function because it contains all and
1066 only the parts of the algorithm that can not be parallelized across trees.
1068 Parameters
1069 ----------
1070 pso
1071 The output of `accept_moves_parallel_stage`.
1073 Returns
1074 -------
1075 state : State
1076 A partially updated BART mcmc state.
1077 moves : Moves
1078 The accepted/rejected moves, with `acc` and `to_prune` set.
1079 """
1081 def loop(
1082 resid: Float[Array, ' n'] | Float[Array, ' k n'], pt: SeqStageInPerTree
1083 ) -> tuple[
1084 Float[Array, ' n'] | Float[Array, ' k n'],
1085 tuple[
1086 Float[Array, ' tree_size'] | Float[Array, ' k tree_size'],
1087 Bool[Array, ''],
1088 Bool[Array, ''],
1089 Float32[Array, ''] | None,
1090 ],
1091 ]:
1092 resid, leaf_tree, acc, to_prune, lkratio = accept_move_and_sample_leaves(
1093 resid,
1094 SeqStageInAllTrees(
1095 pso.state.X,
1096 pso.state.forest.leaf_indices,
1097 pso.state.config.resid_reduction_config,
1098 pso.state.config.data_sharded,
1099 pso.state.prec_scale,
1100 pso.state.forest.log_likelihood is not None,
1101 scaled_error_cov_inv(pso.state)
1102 if isinstance(pso.prelf, PreLfMVHet)
1103 else None,
1104 pso.state.forest.leaf_unit,
1105 pso.state.resid_unit,
1106 pso.state.resid_eff_scale,
1107 pso.state.config.leaf_quantization,
1108 ),
1109 pt,
1110 )
1111 return resid, (leaf_tree, acc, to_prune, lkratio)
1113 num_trees, _ = pso.state.forest.leaf_indices.shape
1114 pts = SeqStageInPerTree(
1115 pso.state.forest.leaf_tree,
1116 pso.prec_trees,
1117 pso.moves,
1118 jnp.arange(num_trees, dtype=minimal_unsigned_dtype(num_trees - 1)),
1119 pso.prelkv,
1120 pso.prelf,
1121 )
1122 resid, (leaf_trees, acc, to_prune, lkratio) = lax.scan(
1123 loop, pso.state.resid, pts, unroll=pso.state.config.sequential_unroll
1124 )
1126 state = replace(
1127 pso.state,
1128 resid=resid,
1129 forest=replace(pso.state.forest, leaf_tree=leaf_trees, log_likelihood=lkratio),
1130 )
1131 moves = replace(pso.moves, acc=acc, to_prune=to_prune)
1133 return state, moves
1136class SeqStageInAllTrees(Module):
1137 """The inputs to `accept_move_and_sample_leaves` that are shared by all trees."""
1139 X: UInt[Array, 'p n']
1140 """The predictors."""
1142 leaf_indices: UInt[Array, 'num_trees n']
1143 """The leaf indices for the largest version of each tree compatible with
1144 the move. Kept whole here and sliced per tree with
1145 `SeqStageInPerTree.tree_index`, instead of consumed as scan xs, because
1146 the scan batching rules move the chain axis of closed-over values to the
1147 front, matching the storage layout, while the xs rule would transpose it
1148 to after the tree axis."""
1150 resid_reduction_config: ReductionConfig
1151 """How to sum the residuals in each leaf."""
1153 data_sharded: bool = field(static=True)
1154 """Whether the data axis is sharded across devices."""
1156 prec_scale: Float[Array, ' n'] | Float[Array, 'k k n'] | None
1157 """The scale of the precision of the error on each datapoint. If None, it
1158 is assumed to be 1."""
1160 save_ratios: bool = field(static=True)
1161 """Whether to save the acceptance ratios."""
1163 error_cov_inv: Float32[Array, 'k k'] | None
1164 """The global error precision scale, with the squared `inv_sdev_unit` folded
1165 in (see `bartz.mcmcstep.State.inv_sdev_unit`). Set only in the
1166 multivariate case with per-component error scales, where the sequential
1167 stage needs it to compute the leaf scores."""
1169 leaf_unit: Float32[Array, ''] | Float32[Array, ' k']
1170 """The storage unit of the leaf values, see `bartz.mcmcstep.Forest.leaf_unit`."""
1172 resid_unit: Float32[Array, ''] | Float32[Array, ' k']
1173 """The storage unit of the residuals, see `bartz.mcmcstep.State.resid_unit`."""
1175 resid_eff_scale: Float32[Array, ''] | Float32[Array, ' k']
1176 """The measured scale of the residuals, see `bartz.mcmcstep.State.resid_eff_scale`."""
1178 leaf_quantization: Int32[Array, ''] | None
1179 """Leaf quantization setting, see `bartz.mcmcstep.StepConfig`."""
1182class SeqStageInPerTree(Module):
1183 """The inputs to `accept_move_and_sample_leaves` that are separate for each tree."""
1185 # Although consumed one tree at a time by `lax.scan`, this object is only
1186 # ever constructed in the stacked (batched) form fed to the scan, so
1187 # `num_trees` stays a fixed (non-variadic) leading axis disambiguated by
1188 # rank/dtype (cf. `ParallelStageOut`); the per-tree slices reach `loop` via
1189 # scan, which does not re-run `__init__`.
1190 leaf_tree: (
1191 Float[Array, 'num_trees tree_size'] | Float[Array, 'num_trees k tree_size']
1192 )
1193 """The leaf values of the trees."""
1195 prec_tree: (
1196 Float32[Array, 'num_trees tree_size']
1197 | UInt32[Array, 'num_trees tree_size']
1198 | Float32[Array, 'num_trees k k tree_size']
1199 )
1200 """The likelihood precision scale in each potential or actual leaf node."""
1202 move: Moves
1203 """The proposed move, see `propose_moves`."""
1205 tree_index: UInt[Array, ' num_trees']
1206 """The index of the tree, used to slice `SeqStageInAllTrees.leaf_indices`."""
1208 prelkv: PreLkV
1209 """The pre-computed terms of the likelihood ratio which are specific to the tree."""
1211 prelf: PreLf
1212 """The pre-computed terms of the leaf sampling which are specific to the tree."""
1215@named_call
1216def accept_move_and_sample_leaves(
1217 resid: Float[Array, ' n'] | Float[Array, ' k n'],
1218 at: SeqStageInAllTrees,
1219 pt: SeqStageInPerTree,
1220) -> tuple[
1221 Float[Array, ' n'] | Float[Array, ' k n'],
1222 Float[Array, ' tree_size'] | Float[Array, ' k tree_size'],
1223 Bool[Array, ''],
1224 Bool[Array, ''],
1225 Float32[Array, ''] | None,
1226]:
1227 """
1228 Accept or reject a proposed move and sample the new leaf values.
1230 Parameters
1231 ----------
1232 resid
1233 The residuals (data minus forest value), in units of `at.resid_unit`
1234 and a possibly narrow dtype (see `State.resid`). The reduction over the
1235 residuals runs in these units; the per-leaf sums are scaled back to data
1236 units afterwards.
1237 at
1238 The inputs that are the same for all trees.
1239 pt
1240 The inputs that are separate for each tree.
1242 Returns
1243 -------
1244 resid : Float[Array, 'n'] | Float[Array, ' k n']
1245 The updated residuals, in the same stored units and dtype as the input.
1246 leaf_tree : Float[Array, 'tree_size'] | Float[Array, ' k tree_size']
1247 The new leaf values of the tree.
1248 acc : Bool[Array, '']
1249 Whether the move was accepted.
1250 to_prune : Bool[Array, '']
1251 Whether, to reflect the acceptance status of the move, the state should
1252 be updated by pruning the leaves involved in the move.
1253 log_lk_ratio : Float32[Array, ''] | None
1254 The logarithm of the likelihood ratio for the move. `None` if not to be
1255 saved.
1256 """
1257 leaf_indices = at.leaf_indices[pt.tree_index, :]
1259 # sum residuals in each leaf, in tree proposed by grow move
1260 if at.prec_scale is None:
1261 scaled_resid = resid
1262 else:
1263 scaled_resid = resid * at.prec_scale
1265 tree_size = pt.leaf_tree.shape[-1] # 2**d
1267 resid_tree = sum_resid(
1268 scaled_resid,
1269 leaf_indices,
1270 tree_size,
1271 at.resid_reduction_config,
1272 at.data_sharded,
1273 )
1275 # the residuals are stored and reduced in units of `resid_unit` (and a
1276 # possibly narrow dtype) for bandwidth; the float32 per-leaf sums are scaled
1277 # back to data units here, after the reduction and before everything else.
1278 # with error scales, the sums keep the stored `prec_scale` units, whose
1279 # `inv_sdev_unit ** 2` factor is folded into the error precision instead
1280 resid_tree *= at.resid_unit[..., None]
1282 # convert the starting tree to data units; multiplying by the float32
1283 # scale also takes care of upcasting narrow leaf dtypes
1284 prev_leaf_tree = at.leaf_unit[..., None] * pt.leaf_tree
1286 # subtract starting tree from function
1287 resid_tree += pt.prec_tree * prev_leaf_tree
1289 # sum residuals in parent node modified by move and compute likelihood;
1290 # the children slots are written back unchanged to share a single scatter
1291 assert pt.move.lrt_nodes.dtype == jnp.int32
1292 resid_lrt = _fill_lrt_total(resid_tree[..., pt.move.lrt_nodes])
1293 resid_tree = resid_tree.at[..., pt.move.lrt_nodes].set(resid_lrt)
1295 log_lk_ratio = compute_likelihood_ratio(resid_lrt, pt.prelkv, at.error_cov_inv)
1297 # calculate accept/reject ratio; the ratio is filled in by `complete_ratio`
1298 assert pt.move.log_trans_prior_ratio is not None
1299 log_ratio = pt.move.log_trans_prior_ratio + log_lk_ratio
1300 log_ratio = jnp.where(pt.move.grow, log_ratio, -log_ratio)
1301 if not at.save_ratios:
1302 log_lk_ratio = None
1304 # determine whether to accept the move
1305 acc = pt.move.allowed & (pt.move.logu <= log_ratio)
1307 # compute leaves posterior and sample leaves
1308 if at.error_cov_inv is not None:
1309 # multivariate w/ per-component error scales
1310 b_tree = compute_B(at.error_cov_inv, resid_tree) # (k, 2**d)
1311 l_lead = jnp.moveaxis(pt.prelf.mean_factor, -1, 0) # (2**d, k, k)
1312 b_lead = b_tree.T[:, :, None] # (2**d, k, 1)
1313 y = solve_triangular(l_lead, b_lead, lower=True)
1314 mu = solve_triangular(l_lead, y, lower=True, trans='T').squeeze(-1)
1315 mean_post = mu.T # (k, 2**d)
1316 elif resid.ndim > 1:
1317 # multivariate homoskedastic or scalar error scales
1318 mean_post = jnp.einsum('kil,kl->il', pt.prelf.mean_factor, resid_tree)
1319 else:
1320 # univariate
1321 mean_post = resid_tree * pt.prelf.mean_factor
1322 leaf_tree = mean_post + pt.prelf.centered_leaves
1324 # copy leaves around such that the leaf indices point to the correct leaf;
1325 # the parent slot is written back unchanged to share a single scatter.
1326 # this mirroring persists into the output state, where it keeps the
1327 # not-yet-pruned `leaf_indices` valid to evaluate the trees
1328 to_prune = acc ^ pt.move.grow
1329 leaf_tree = leaf_tree.at[
1330 ..., jnp.where(to_prune, pt.move.lrt_nodes, tree_size)
1331 ].set(leaf_tree[..., pt.move.lrt_nodes[2], None])
1333 # round the new leaves to the storage units and dtype; the residuals are
1334 # then updated with the rounded values to stay consistent with the trees
1335 leaf_tree = leaf_tree / at.leaf_unit[..., None]
1336 if at.leaf_quantization is not None:
1337 # quantize the leaves such that the residual updates below are mostly
1338 # exact. the target grid is the spacing of `resid.dtype` values of
1339 # magnitude 2^leaf_quantization in units of the measured residual
1340 # scale, so 2^(leaf_quantization - nmant) * resid_eff_scale in data
1341 # units; dividing by leaf_unit converts it to the leaf units
1342 # `leaf_tree` is in here. the scales are powers of two, so the quantum
1343 # is one too and the leaf storage dtype's own (coarser) rounding below
1344 # preserves multiples of it.
1345 nmant = jnp.finfo(resid.dtype).nmant
1346 quantum = (at.resid_eff_scale / at.leaf_unit) * 2.0 ** (
1347 at.leaf_quantization - nmant
1348 )
1349 leaf_tree = jnp.round(leaf_tree / quantum[..., None]) * quantum[..., None]
1350 # round to the storage dtype explicitly before converting: xla's
1351 # excess-precision optimization (on by default) may otherwise elide the
1352 # narrowing conversion in the `leaf_delta` computation below, updating the
1353 # residuals with unrounded leaves while the trees store rounded ones
1354 finfo = jnp.finfo(pt.leaf_tree.dtype)
1355 leaf_tree = lax.reduce_precision(leaf_tree, finfo.nexp, finfo.nmant)
1356 leaf_tree = leaf_tree.astype(pt.leaf_tree.dtype)
1358 # replace old tree with new tree in function values; the per-leaf data-unit
1359 # delta is converted back to the stored residual units and dtype *before* the
1360 # scatter, so the n-sized update stays in the narrow `resid` storage
1361 leaf_delta = prev_leaf_tree - at.leaf_unit[..., None] * leaf_tree
1362 delta = (leaf_delta / at.resid_unit[..., None]).astype(resid.dtype)
1363 # on cpu, materialize the per-leaf delta, else xla fuses its computation
1364 # into the gather and repeats it for each datapoint; on gpu that is cheaper
1365 # than the extra kernel launch
1366 delta = lax.platform_dependent(
1367 delta, cpu=lax.optimization_barrier, default=lambda x: x
1368 )
1369 resid += delta[..., leaf_indices]
1371 return resid, leaf_tree, acc, to_prune, log_lk_ratio
1374@named_call
1375def sum_resid(
1376 scaled_resid: Float[Array, ' n'] | Float[Array, 'k n'] | Float[Array, 'k k n'],
1377 leaf_indices: UInt[Array, ' n'],
1378 tree_size: int,
1379 reduction_config: ReductionConfig,
1380 data_sharded: bool,
1381) -> (
1382 Float32[Array, ' {tree_size}']
1383 | Float32[Array, 'k {tree_size}']
1384 | Float32[Array, 'k k {tree_size}']
1385):
1386 """
1387 Sum the residuals in each leaf.
1389 Parameters
1390 ----------
1391 scaled_resid
1392 The residuals (data minus forest value) multiplied by the error
1393 precision scale.
1394 leaf_indices
1395 The leaf indices of the tree (in which leaf each data point falls into).
1396 tree_size
1397 The size of the tree array (2 ** d).
1398 reduction_config
1399 How to sum the residuals in each leaf.
1400 data_sharded
1401 Whether the data axis is sharded; if true, the result is psum-reduced
1402 across the ``'data'`` axis of the enclosing `shard_map`.
1404 Returns
1405 -------
1406 The per-leaf sum, with the same leading dimensions as ``scaled_resid`` and a trailing axis over the leaves.
1407 """
1408 return reduction_config._reduce( # noqa: SLF001
1409 scaled_resid,
1410 leaf_indices,
1411 size=tree_size,
1412 dtype=jnp.float32,
1413 data_sharded=data_sharded,
1414 )
1417def _compute_likelihood_ratio_uv(
1418 resid_lrt: Float32[Array, ' 3'], prelkv: PreLkV
1419) -> Float32[Array, '']:
1420 # quadratic form r * v * r for each of the (left, right, total) terms
1421 qf = resid_lrt * resid_lrt * prelkv.lrt
1422 exp_term = 0.5 * (qf @ jnp.array([1.0, 1.0, -1.0]))
1423 return prelkv.log_sqrt_term + exp_term
1426def _compute_likelihood_ratio_mv(
1427 resid_lrt: Float32[Array, 'k 3'], prelkv: PreLkV
1428) -> Float32[Array, '']:
1429 # quadratic form r' M r for each of the (left, right, total) terms
1430 qf = jnp.einsum('it,tij,jt->t', resid_lrt, prelkv.lrt, resid_lrt)
1431 exp_term = 0.5 * (qf @ jnp.array([1.0, 1.0, -1.0]))
1432 return prelkv.log_sqrt_term + exp_term
1435def _compute_likelihood_ratio_mv_het(
1436 resid_lrt: Float32[Array, 'k k 3'],
1437 error_cov_inv: Float32[Array, 'k k'],
1438 prelkv: PreLkV,
1439) -> Float32[Array, '']:
1440 b = compute_B(error_cov_inv, resid_lrt) # (k, 3)
1441 y = solve_triangular(prelkv.lrt, b.T[..., None], lower=True).squeeze(-1) # (3, k)
1442 qf = jnp.einsum('ti,ti->t', y, y)
1443 exp_term = 0.5 * (qf @ jnp.array([1.0, 1.0, -1.0]))
1444 return prelkv.log_sqrt_term + exp_term
1447@named_call
1448def compute_likelihood_ratio(
1449 resid_lrt: (Float32[Array, ' 3'] | Float32[Array, 'k 3'] | Float32[Array, 'k k 3']),
1450 prelkv: PreLkV,
1451 error_cov_inv: Float32[Array, 'k k'] | None,
1452) -> Float32[Array, '']:
1453 """
1454 Compute the likelihood ratio of a grow move.
1456 Parameters
1457 ----------
1458 resid_lrt
1459 The sum of the residuals (scaled by error precision scale) of the
1460 datapoints falling in the left child, right child, and parent node
1461 involved in the move, stacked along the trailing axis.
1462 prelkv
1463 The pre-computed terms of the likelihood ratio, see
1464 `precompute_likelihood_terms`.
1465 error_cov_inv
1466 The global error precision scale, with the squared `inv_sdev_unit` folded
1467 in (see `bartz.mcmcstep.State.inv_sdev_unit`). Set only in the
1468 multivariate case with per-component error scales.
1470 Returns
1471 -------
1472 The log-likelihood ratio log P(data | new tree) - log P(data | old tree).
1473 """
1474 if error_cov_inv is not None:
1475 return _compute_likelihood_ratio_mv_het(resid_lrt, error_cov_inv, prelkv)
1476 elif resid_lrt.ndim > 1:
1477 return _compute_likelihood_ratio_mv(resid_lrt, prelkv)
1478 else:
1479 return _compute_likelihood_ratio_uv(resid_lrt, prelkv)
1482@named_call
1483def accept_moves_final_stage(state: State, moves: Moves) -> State:
1484 """
1485 Post-process the mcmc state after accepting/rejecting the moves.
1487 This function is separate from `accept_moves_sequential_stage` to signal it
1488 can work in parallel across trees.
1490 The prunes are not applied to `bartz.mcmcstep.Forest.leaf_indices` here;
1491 they are recorded in `to_prune`/`move_node` and applied at the beginning of
1492 the next step.
1494 Parameters
1495 ----------
1496 state
1497 A partially updated BART mcmc state.
1498 moves
1499 The proposed moves (see `propose_moves`) as updated by
1500 `accept_moves_sequential_stage`.
1502 Returns
1503 -------
1504 The fully updated BART mcmc state.
1505 """
1506 assert moves.acc is not None
1507 assert moves.to_prune is not None
1508 return replace(
1509 state,
1510 forest=replace(
1511 state.forest,
1512 grow_acc_count=jnp.sum(moves.acc & moves.grow),
1513 prune_acc_count=jnp.sum(moves.acc & ~moves.grow),
1514 to_prune=moves.to_prune,
1515 move_node=moves.lrt_nodes[..., 2],
1516 split_tree=apply_moves_to_split_trees(state.forest.split_tree, moves),
1517 affluence_tree=apply_moves_to_affluence_trees(
1518 state.forest.affluence_tree, moves
1519 ),
1520 ),
1521 )
1524@named_call
1525def apply_moves_to_leaf_indices(
1526 leaf_indices: UInt[Array, 'num_trees n'],
1527 to_prune: Bool[Array, ' num_trees'],
1528 move_node: Int32[Array, ' num_trees'],
1529) -> UInt[Array, 'num_trees n']:
1530 """
1531 Apply the prunes pending from the previous step to the leaf indices.
1533 Parameters
1534 ----------
1535 leaf_indices
1536 The index of the leaf each datapoint falls into, in the largest version
1537 of each tree compatible with the last moves (see
1538 `Forest.leaf_indices`).
1539 to_prune
1540 Whether the last move on each tree ended in a prune yet to be applied.
1541 move_node
1542 The node the last move on each tree operated on.
1544 Returns
1545 -------
1546 The updated leaf indices.
1547 """
1548 return _apply_moves_to_leaf_indices(leaf_indices, to_prune, move_node)
1551@vmap_nodoc
1552def _apply_moves_to_leaf_indices(
1553 leaf_indices: UInt[Array, ' n'],
1554 to_prune: Bool[Array, ''],
1555 move_node: Int32[Array, ''],
1556) -> UInt[Array, ' n']:
1557 """Implement `apply_moves_to_leaf_indices`."""
1558 mask = ~jnp.array(1, leaf_indices.dtype) # ...1111111110
1559 is_child = (leaf_indices & mask) == (move_node << 1)
1560 return jnp.where(
1561 is_child & to_prune, move_node.astype(leaf_indices.dtype), leaf_indices
1562 )
1565@named_call
1566def apply_moves_to_split_trees(
1567 split_tree: UInt[Array, 'num_trees half_tree_size'], moves: Moves
1568) -> UInt[Array, 'num_trees half_tree_size']:
1569 """
1570 Update the split trees to match the accepted move.
1572 Parameters
1573 ----------
1574 split_tree
1575 The cutpoints of the decision nodes in the initial trees.
1576 moves
1577 The proposed moves (see `propose_moves`), as updated by
1578 `accept_moves_sequential_stage`.
1580 Returns
1581 -------
1582 The updated split trees.
1583 """
1584 return _apply_moves_to_split_trees(split_tree, moves)
1587@vmap_nodoc
1588def _apply_moves_to_split_trees(
1589 split_tree: UInt[Array, ' half_tree_size'], moves: Moves
1590) -> UInt[Array, ' half_tree_size']:
1591 """Implement `apply_moves_to_split_trees`."""
1592 assert moves.to_prune is not None
1593 # a single scatter serves both cases: an accepted grow writes the new
1594 # cutpoint, while pruning (accepted prune or rejected grow) zeroes the node
1595 return split_tree.at[
1596 jnp.where(moves.grow | moves.to_prune, moves.lrt_nodes[2], split_tree.size)
1597 ].set(jnp.where(moves.to_prune, 0, moves.grow_split).astype(split_tree.dtype))
1600@named_call
1601def apply_moves_to_affluence_trees(
1602 affluence_tree: Bool[Array, 'num_trees half_tree_size'], moves: Moves
1603) -> Bool[Array, 'num_trees half_tree_size']:
1604 """
1605 Update the affluence trees to match the accepted move.
1607 The affluence tree marks the growable leaves; this restores that invariant
1608 after the move by re-marking only the nodes it touched, starting from the
1609 clean pre-move mask.
1611 Parameters
1612 ----------
1613 affluence_tree
1614 The mask of the growable leaves in the initial trees.
1615 moves
1616 The proposed moves (see `propose_moves`), as updated by
1617 `accept_moves_sequential_stage`.
1619 Returns
1620 -------
1621 The updated affluence trees.
1622 """
1623 return _apply_moves_to_affluence_trees(affluence_tree, moves)
1626@vmap_nodoc
1627def _apply_moves_to_affluence_trees(
1628 affluence_tree: Bool[Array, ' half_tree_size'], moves: Moves
1629) -> Bool[Array, ' half_tree_size']:
1630 """Implement `apply_moves_to_affluence_trees`."""
1631 assert moves.to_prune is not None
1632 assert moves.lrt_affluent is not None
1633 # GROW: node becomes internal, children become leaves with their affluence.
1634 # PRUNE (accepted prune or rejected grow): node becomes a leaf with its
1635 # affluence, children are deleted. Either way all three nodes are written:
1636 # the mask keeps the affluence of the nodes that become leaves and zeroes
1637 # the rest. If no move is applied (a rejected prune), the indices resolve
1638 # to `size` and the writes drop.
1639 becomes_leaf = moves.to_prune ^ jnp.array([True, True, False])
1640 return affluence_tree.at[
1641 jnp.where(moves.grow | moves.to_prune, moves.lrt_nodes, affluence_tree.size)
1642 ].set(moves.lrt_affluent & becomes_leaf)
1645@jit
1646def _sample_wishart_bartlett(
1647 key: Key[Array, ''],
1648 df: Float32[Array, ''] | float,
1649 scale_inv: Float32[Array, 'k k'],
1650) -> Float32[Array, 'k k']:
1651 """
1652 Sample a precision matrix W ~ Wishart(df, scale_inv^-1) using Bartlett decomposition.
1654 Parameters
1655 ----------
1656 key
1657 A JAX random key
1658 df
1659 Degrees of freedom
1660 scale_inv
1661 Scale matrix of the corresponding Inverse Wishart distribution
1663 Returns
1664 -------
1665 A sample from Wishart(df, scale)
1666 """
1667 keys = split(key)
1669 # Diagonal elements: A_ii ~ sqrt(chi^2(df - i)), with chi^2(k) = Gamma(k/2, scale=2).
1670 # sqrt(2 * Gamma) = sqrt(2) * exp(loggamma / 2), folding the sqrt into the exp.
1671 k, _ = scale_inv.shape
1672 df_vector = df - jnp.arange(k)
1673 diag_A = jnp.sqrt(2.0) * jnp.exp(loggamma(keys.pop(), df_vector / 2.0) / 2.0)
1675 off_diag_A = random.normal(keys.pop(), (k, k))
1676 A = jnp.tril(off_diag_A, -1) + jnp.diag(diag_A)
1677 L = chol_with_gersh(scale_inv, absolute_eps=True)
1678 T = solve_triangular(L, A, lower=True, trans='T')
1680 return T @ T.T
1683def sample_wishart_posterior(
1684 key: Key[Array, ''],
1685 prior: Wishart,
1686 count: Int32[Array, ''] | Int32[Array, ' k'],
1687 scatter: Float32[Array, ''] | Float32[Array, ' k'] | Float32[Array, 'k k'],
1688) -> Float32[Array, ''] | Float32[Array, 'k k']:
1689 """Draw a precision matrix from its conjugate posterior.
1691 Parameters
1692 ----------
1693 key
1694 A jax random key.
1695 prior
1696 The prior on the precision; `Wishart.nu` and `Wishart.rate` must be set.
1697 count
1698 The number of observations. It can vary per component only with a
1699 `DiagWishart` prior. Fractional counts are not supported.
1700 scatter
1701 The sum of the outer products of the observations: the full matrix with
1702 a dense `Wishart` prior, only its diagonal with a `DiagWishart`, a
1703 scalar sum of squares in the univariate case.
1705 Returns
1706 -------
1707 A draw from the posterior, with the same shape as `Wishart.rate`.
1708 """
1709 assert prior.nu is not None
1710 assert prior.rate is not None
1712 if prior.rate.ndim == 2 and not isinstance(prior, DiagWishart):
1713 assert scatter.shape == prior.rate.shape
1714 return _sample_wishart_bartlett(key, prior.nu + count, prior.rate + scatter)
1715 else:
1716 # independent gamma per component; the variance df matches the inverse
1717 # wishart marginals
1718 marginal_nu = prior.inv_wishart_marginal_nu
1719 assert marginal_nu is not None
1720 alpha = marginal_nu / 2 + count / 2
1722 kshape = prior.rate.shape[:1]
1723 assert scatter.shape == kshape
1724 rate = jnp.diagonal(prior.rate) if kshape else prior.rate
1725 beta = rate / 2 + scatter / 2
1727 # draw the gamma from the first of a split, mirroring the Bartlett
1728 # sampler in the dense path so the two branches coincide at k=1
1729 keys = split(key)
1730 prec = jnp.exp(loggamma(keys.pop(), alpha, kshape)) / beta
1731 if kshape:
1732 prec = jnp.diag(prec)
1733 return prec
1736def step_resid_eff_scale(
1737 state: State,
1738 norm2: Float32[Array, ''] | Float32[Array, ' k'],
1739 prec_sum: Float32[Array, ''] | Float32[Array, ' k'],
1740) -> State:
1741 """Update `State.resid_eff_scale` with the measured scale of the residuals.
1743 Parameters
1744 ----------
1745 state
1746 A BART MCMC state.
1747 norm2
1748 The squared (precision-scaled, masked) norm of the residuals in data units.
1749 prec_sum
1750 The sum of the precision scales (unmasked-datapoint count without
1751 error scales), matching the scaling of `norm2`.
1753 Returns
1754 -------
1755 The state with `resid_eff_scale` set to the residual rms rounded to a power of two.
1756 """
1757 scale = round_to_pow2(jnp.sqrt(norm2 / prec_sum))
1758 # keep the previous scale if the residuals vanish
1759 scale = jnp.where(scale == 0, state.resid_eff_scale, scale)
1760 return replace(state, resid_eff_scale=scale)
1763def _step_error_cov_inv_mv(key: Key[Array, ''], state: State) -> State:
1764 # keep the residuals in their stored (narrow) dtype and resid_unit units;
1765 # the reduction accumulates in float32 and its (k, k) result is rescaled to
1766 # data units, so no n-sized float32 array is ever materialized
1767 resid = state.resid
1768 if state.inv_sdev_scale is not None:
1769 resid *= state.inv_sdev_scale
1771 # unit of the stored precision-scaled residuals: `resid` is in `resid_unit`
1772 # units and `inv_sdev_scale` in `inv_sdev_unit` units (1 without error
1773 # scales, per-component with 2-D `inv_sdev_scale`)
1774 scale = state.resid_unit * state.inv_sdev_unit
1776 rrt = jnp.einsum(
1777 'an,bn->ab', resid, resid, preferred_element_type=jnp.float32
1778 ) * jnp.outer(scale, scale)
1779 if state.config.data_sharded:
1780 rrt = lax.psum(rrt, 'data')
1782 # we take the max as a way to take any of the in this case equal values in
1783 # n_non_missing
1784 count = jnp.max(state.n_non_missing)
1786 prec = sample_wishart_posterior(key, state.error_cov_inv, count, rrt)
1787 state = step_resid_eff_scale(state, jnp.diagonal(rrt), state.sum_diag_prec_scale)
1788 return replace(state, error_cov_inv=replace(state.error_cov_inv, value=prec))
1791def _step_error_cov_inv_diag(key: Key[Array, ''], state: State) -> State:
1792 """Per-component inverse-gamma update for univariate, mixed, and partial-missing paths."""
1793 # keep the residuals in their stored (narrow) dtype and resid_unit units;
1794 # the reduction accumulates in float32 and its small result is rescaled to
1795 # data units, so no n-sized float32 array is ever materialized
1796 resid = state.resid
1797 if state.inv_sdev_scale is not None:
1798 resid *= state.inv_sdev_scale
1800 # `resid` is stored in `resid_unit` units and `inv_sdev_scale` in
1801 # `inv_sdev_unit` units (1 without error scales)
1802 norm2 = jnp.einsum(
1803 '...n,...n->...', resid, resid, preferred_element_type=jnp.float32
1804 ) * jnp.square(state.resid_unit * state.inv_sdev_unit)
1805 if state.config.data_sharded:
1806 norm2 = lax.psum(norm2, 'data')
1808 prec = sample_wishart_posterior(
1809 key, state.error_cov_inv, state.n_non_missing, norm2
1810 )
1811 if state.binary_indices is not None:
1812 # the binary components have unit error variance, they have no prior
1813 prec = prec.at[state.binary_indices, state.binary_indices].set(1.0)
1814 state = step_resid_eff_scale(state, norm2, state.sum_diag_prec_scale)
1815 return replace(state, error_cov_inv=replace(state.error_cov_inv, value=prec))
1818@named_call
1819def step_error_cov_inv(key: Key[Array, ''], state: State) -> State:
1820 """MCMC-update the inverse error covariance, if it has a prior."""
1821 if state.error_cov_inv.nu is None:
1822 return state
1823 elif state.error_cov_inv.value.ndim == 2 and not isinstance(
1824 state.error_cov_inv, DiagWishart
1825 ):
1826 return _step_error_cov_inv_mv(key, state)
1827 else:
1828 return _step_error_cov_inv_diag(key, state)
1831def leaf_scatter(
1832 forest: Forest,
1833) -> tuple[Int32[Array, ''], Float32[Array, ''] | Float32[Array, 'k k']]:
1834 """Count the leaves of the forest and sum their outer products.
1836 Parameters
1837 ----------
1838 forest
1839 The forest, with no chain axis.
1841 Returns
1842 -------
1843 count : Int32[Array, '']
1844 The total number of leaves over all trees.
1845 scatter : Float32[Array, ''] | Float32[Array, 'k k']
1846 The sum over the leaves of the outer product of the values, in data units.
1847 """
1848 # the values on the dangling children of pruned nodes mirror the parent's,
1849 # so mask on the actual leaves only
1850 is_leaf = vmap(partial(is_actual_leaf, add_bottom_level=True))(forest.split_tree)
1851 count = jnp.sum(is_leaf)
1852 leaf_tree = forest.leaf_tree
1853 if leaf_tree.ndim == 3:
1854 scatter = jnp.einsum(
1855 'tan,tbn,tn->ab',
1856 leaf_tree,
1857 leaf_tree,
1858 is_leaf,
1859 preferred_element_type=jnp.float32,
1860 )
1861 scatter *= jnp.outer(forest.leaf_unit, forest.leaf_unit)
1862 else:
1863 scatter = jnp.einsum(
1864 'tn,tn,tn->',
1865 leaf_tree,
1866 leaf_tree,
1867 is_leaf,
1868 preferred_element_type=jnp.float32,
1869 )
1870 scatter *= jnp.square(forest.leaf_unit)
1871 return count, scatter
1874@named_call
1875def step_leaf_prior_cov_inv(key: Key[Array, ''], state: State) -> State:
1876 """MCMC-update the leaf prior precision from its conjugate posterior, if it has a prior."""
1877 prior = state.forest.leaf_prior_cov_inv
1878 if prior.nu is None:
1879 return state
1880 else:
1881 count, scatter = leaf_scatter(state.forest)
1882 prec = sample_wishart_posterior(key, prior, count, scatter)
1883 return replace(
1884 state,
1885 forest=replace(state.forest, leaf_prior_cov_inv=replace(prior, value=prec)),
1886 )
1889@named_call
1890def step_z(key: Key[Array, ''], state: State) -> State:
1891 """
1892 MCMC-update the latent variable for binary regression, if any.
1894 Parameters
1895 ----------
1896 key
1897 A jax random key.
1898 state
1899 A BART MCMC state.
1901 Returns
1902 -------
1903 The updated BART MCMC state.
1904 """
1905 if state.z is None:
1906 return state
1908 inv_sdev = state.inv_sdev_scale
1909 inv_sdev_unit = state.inv_sdev_unit
1910 err_scale = state.error_scale
1911 if state.binary_indices is not None:
1912 resid_unit = state.resid_unit[state.binary_indices]
1913 resid = state.resid[state.binary_indices, :]
1914 binary_y = state.y[state.binary_indices, :] != 0
1915 # per-component scales (2-D) are restricted to the binary components;
1916 # a scalar-per-datapoint scale (1-D) is shared and broadcasts as-is
1917 if inv_sdev is not None and inv_sdev.ndim > 1:
1918 inv_sdev = inv_sdev[state.binary_indices, :]
1919 inv_sdev_unit = inv_sdev_unit[state.binary_indices]
1920 if err_scale is not None and err_scale.ndim > 1:
1921 err_scale = err_scale[state.binary_indices, :]
1922 else:
1923 resid_unit = state.resid_unit
1924 resid = state.resid
1925 binary_y = state.y != 0
1927 trees_plus_offset = state.z - resid * resid_unit[..., None]
1928 if state.config.data_sharded:
1929 # decorrelate the seed across data shards; the seed is replicated
1930 # because the trees and most of the algorithm are replicated
1931 key = random.fold_in(key, lax.axis_index('data'))
1933 if err_scale is None:
1934 # homoskedastic probit: the latent error has unit scale. A missingness
1935 # mask (without error scales) needs no handling here: masked points are
1936 # dropped from the likelihood, so their latent is sampled like any other.
1937 resid = truncated_normal_onesided(key, (), ~binary_y, -trees_plus_offset)
1938 else:
1939 # heteroskedastic probit: the latent error has per-datapoint scale
1940 # `err_scale`, so threshold in standardized units, then rescale.
1941 assert inv_sdev is not None
1942 # masked datapoints draw garbage, but every consumer of `resid` scales
1943 # it by the (zero) `inv_sdev`, so the value is irrelevant as long as finite
1944 resid = err_scale * truncated_normal_onesided(
1945 key, (), ~binary_y, -trees_plus_offset * inv_sdev * inv_sdev_unit[..., None]
1946 )
1947 z = trees_plus_offset + resid
1949 resid = (resid / resid_unit[..., None]).astype(state.resid.dtype)
1950 if state.binary_indices is not None:
1951 resid = state.resid.at[state.binary_indices, :].set(resid)
1953 return replace(state, z=z, resid=resid)
1956def _blocked_mass_tree(
1957 key: Key[Array, ''],
1958 var_tree: UInt[Array, ' half_tree_size'],
1959 split_tree: UInt[Array, ' half_tree_size'],
1960 max_split: UInt[Array, ' p'],
1961 s: Float32[Array, ' p'],
1962) -> Float32[Array, ' p']:
1963 """Per-variable data-augmentation mass blocked by a single tree.
1965 At each internal node, draws the latent augmentation weight ``lambda / e``
1966 (``lambda`` exponential, ``e`` the eligible split probability mass at the
1967 node) and adds it to every variable ineligible at that node.
1969 Parameters
1970 ----------
1971 key
1972 Random key for sampling.
1973 var_tree
1974 The splitting axes of the tree.
1975 split_tree
1976 The splitting points of the tree.
1977 max_split
1978 The maximum split index for each variable.
1979 s
1980 Split probabilities normalized over selectable variables.
1982 Returns
1983 -------
1984 The blocked mass for each variable.
1985 """
1986 (half_tree_size,) = split_tree.shape
1987 d_minus_1 = half_tree_size.bit_length() - 1 # number of decision-node levels
1988 p = max_split.size
1989 nodes = jnp.arange(half_tree_size)
1990 split = split_tree.astype(jnp.int32)
1991 is_internal = split_tree.astype(bool)
1993 # Range [lo, hi) of cutpoints still available for each node's own splitting
1994 # variable, given the constraints inherited from the ancestors.
1995 lo, hi = vmap(split_range, in_axes=(None, None, None, 0, 0))(
1996 var_tree, split_tree, max_split, nodes, var_tree
1997 )
1999 # An internal node exhausts its own variable for a child when its cutpoint
2000 # sits at the matching end of the available range, so the variable becomes
2001 # ineligible throughout that child's subtree. Row 0 is the left child (low
2002 # end lo), row 1 the right child (high end hi - 1).
2003 blocks = is_internal & (split == jnp.stack([lo, hi - 1]))
2005 # A node can block at most its own splitting variable, so the per-variable
2006 # totals are recovered from these per-node blocks via top-down/bottom-up
2007 # accumulation over depth levels, rather than scanning each node's ancestors.
2009 # Ineligible mass per node: the s-mass of the variables blocked along the
2010 # path from the root. Each variable is blocked at exactly one node per path,
2011 # so summing the per-node increments top-down reproduces the per-node sum
2012 # over distinct ineligible variables.
2013 parent = nodes >> 1
2014 side = nodes & 1 # 0 if the node is a left child, 1 if a right child
2015 parent_blocks = blocks[side, parent]
2016 # var_tree[parent] is a valid index wherever parent_blocks holds (the parent
2017 # is then internal); elsewhere the clamped gather is masked away
2018 ineligible_mass = jnp.where(parent_blocks, s[var_tree[parent]], 0.0)
2019 for level in range(1, d_minus_1):
2020 lhs, rhs = 1 << level, 1 << (level + 1)
2021 parent_mass = jnp.repeat(ineligible_mass[lhs >> 1 : rhs >> 1], 2)
2022 ineligible_mass = ineligible_mass.at[lhs:rhs].add(parent_mass)
2024 # Per-node augmentation weight lambda_b / e_b, zero at non-internal nodes. The
2025 # eligible mass is positive at internal nodes (the split variable is eligible);
2026 # the floor only guards against round-off, and is unused where weight is zero.
2027 eligible_mass = jnp.maximum(1.0 - ineligible_mass, jnp.finfo(jnp.float32).eps)
2028 weight = jnp.where(is_internal, random.exponential(key, (half_tree_size,)), 0.0)
2029 weight /= eligible_mass
2031 # Subtree weight: total weight of each node's internal descendants and itself,
2032 # accumulated bottom-up. The children of the deepest decision level are leaves
2033 # and contribute nothing.
2034 subtree_weight = weight
2035 for level in range(d_minus_1 - 2, -1, -1):
2036 lhs, rhs = 1 << level, 1 << (level + 1)
2037 children = subtree_weight[2 * lhs : 2 * rhs].reshape(-1, 2).sum(axis=1)
2038 subtree_weight = subtree_weight.at[lhs:rhs].add(children)
2040 # A variable blocked by node b at one of its children is ineligible in that
2041 # whole child subtree, so it accumulates the subtree weight; scatter it onto
2042 # the splitting variable. Only the upper half of nodes have internal children
2043 # (the deepest decision nodes block into leaves, contributing nothing); their
2044 # children are exactly subtree_weight reshaped into [left, right] pairs.
2045 half = half_tree_size // 2
2046 contrib = (blocks[:, :half] * subtree_weight.reshape(-1, 2).T).sum(axis=0)
2047 scatter_var = jnp.where(is_internal, var_tree, p)[:half]
2048 return jnp.zeros(p).at[scatter_var].add(contrib)
2051def sample_s_augmentation(key: Key[Array, ''], forest: Forest) -> Int32[Array, ' p']:
2052 """Sample the data-augmentation counts for the exact full conditional of `s`.
2054 At each internal node, the variables with no available cutpoint given the
2055 ancestors (plus the globally blocked ones) cannot be split on, so the plain
2056 Dirichlet update for `s` is only approximate. This samples, for each
2057 variable, the number of ineligible draws discarded before each realized
2058 split, to be added to the variable usage counts.
2060 Parameters
2061 ----------
2062 key
2063 Random key for sampling.
2064 forest
2065 The forest, providing the trees and the current `log_s`.
2067 Returns
2068 -------
2069 The discarded-draws count for each variable.
2070 """
2071 assert forest.log_s is not None
2072 keys = split(key)
2073 (num_trees, _) = forest.var_tree.shape
2075 # split probabilities normalized over the selectable (non-blocked) variables
2076 selectable = forest.max_split > 0
2077 s = softmax(forest.log_s, where=selectable)
2079 # blocked_mass[j] = sum over internal nodes where j is ineligible of
2080 # lambda_b / e_b, with lambda_b ~ Exponential(1)
2081 blocked_mass = vmap(_blocked_mass_tree, in_axes=(0, 0, 0, None, None))(
2082 keys.pop(num_trees), forest.var_tree, forest.split_tree, forest.max_split, s
2083 ).sum(axis=0) # shape (p,)
2085 # the per-node discarded-draw counts are negative-multinomial, with no
2086 # closed form when summed over nodes, but their Gamma-Poisson mixture does:
2087 # A_j | {lambda_b} ~ Poisson(s_j * blocked_mass[j]), independent across j
2088 return poisson(keys.pop(), s * blocked_mass, dtype=jnp.int32)
2091@named_call
2092def step_s(key: Key[Array, ''], state: State) -> State:
2093 """
2094 Update `log_s` using Dirichlet sampling.
2096 The prior is s ~ Dirichlet(theta/p, ..., theta/p), and the posterior
2097 is s ~ Dirichlet(theta/p + varcount, ..., theta/p + varcount), where
2098 varcount is the count of how many times each variable is used in the
2099 current forest.
2101 Parameters
2102 ----------
2103 key
2104 Random key for sampling.
2105 state
2106 The current BART state.
2108 Returns
2109 -------
2110 Updated BART state with re-sampled `log_s`.
2112 Notes
2113 -----
2114 By default this full conditional is approximate, because it ignores the
2115 decision rules forbidden by the ancestors of each node. If
2116 ``state.config.augment`` is set, the forbidden rules are accounted for
2117 exactly with the data augmentation of `sample_s_augmentation`.
2118 """
2119 assert state.forest.theta is not None
2121 # reserve the Dirichlet draw key first and unconditionally, so it does not
2122 # depend on whether augmentation is on; then the two modes draw identically
2123 # when there are no forbidden rules, since the augmentation is exactly zero
2124 keys = split(key)
2125 log_s_key = keys.pop()
2127 # histogram current variable usage
2128 p = state.forest.max_split.size
2129 varcount = var_histogram(
2130 p, state.forest.var_tree, state.forest.split_tree, sum_batch_axis=-1
2131 )
2133 # the Dirichlet posterior concentration, optionally completed with the exact
2134 # accounting of forbidden rules via data augmentation
2135 alpha = state.forest.theta / p + varcount
2136 if state.config.augment:
2137 alpha = alpha + sample_s_augmentation(keys.pop(), state.forest)
2139 # sample from the Dirichlet posterior and update the forest with the new s
2140 log_s = loggamma(log_s_key, alpha)
2141 return replace(state, forest=replace(state.forest, log_s=log_s))
2144@named_call
2145def step_theta(key: Key[Array, ''], state: State, *, num_grid: int = 1000) -> State:
2146 """
2147 Update `theta`.
2149 The prior is theta / (theta + rho) ~ Beta(a, b).
2151 Parameters
2152 ----------
2153 key
2154 Random key for sampling.
2155 state
2156 The current BART state.
2157 num_grid
2158 The number of points in the evenly-spaced grid used to sample
2159 theta / (theta + rho).
2161 Returns
2162 -------
2163 Updated BART state with re-sampled `theta`.
2164 """
2165 assert state.forest.log_s is not None
2166 assert state.forest.rho is not None
2167 assert state.forest.a is not None
2168 assert state.forest.b is not None
2170 # the grid points are the midpoints of num_grid bins in (0, 1)
2171 padding = 1 / (2 * num_grid)
2172 lambda_grid = jnp.linspace(padding, 1 - padding, num_grid)
2174 # normalize s
2175 log_s = state.forest.log_s - logsumexp(state.forest.log_s)
2177 # sample lambda
2178 logp, theta_grid = _log_p_lambda(
2179 lambda_grid, log_s, state.forest.rho, state.forest.a, state.forest.b
2180 )
2181 i = random.categorical(key, logp)
2182 theta = theta_grid[i]
2184 return replace(state, forest=replace(state.forest, theta=theta))
2187def _log_p_lambda(
2188 lambda_: Float32[Array, ' num_grid'],
2189 log_s: Float32[Array, ' p'],
2190 rho: Float32[Array, ''],
2191 a: Float32[Array, ''],
2192 b: Float32[Array, ''],
2193) -> tuple[Float32[Array, ' num_grid'], Float32[Array, ' num_grid']]:
2194 # in the following I use lambda_[::-1] == 1 - lambda_
2195 theta = rho * lambda_ / lambda_[::-1]
2196 p = log_s.size
2197 return (
2198 (a - 1) * jnp.log1p(-lambda_[::-1]) # log(lambda)
2199 + (b - 1) * jnp.log1p(-lambda_) # log(1 - lambda)
2200 + gammaln(theta)
2201 - p * gammaln(theta / p)
2202 + theta / p * jnp.sum(log_s)
2203 ), theta
2206@named_call
2207def step_sparse(key: Key[Array, ''], state: State) -> State:
2208 """
2209 Update the sparsity parameters.
2211 This invokes `step_s`, and then `step_theta` only if the parameters of
2212 the theta prior are defined.
2214 Parameters
2215 ----------
2216 key
2217 Random key for sampling.
2218 state
2219 The current BART state.
2221 Returns
2222 -------
2223 Updated BART state with re-sampled `log_s` and `theta`.
2224 """
2225 if state.config.sparse_on_at is not None:
2226 state = lax.cond(
2227 state.config.steps_done < state.config.sparse_on_at,
2228 lambda _key, state: state,
2229 _step_sparse,
2230 key,
2231 state,
2232 )
2233 return state
2236def _step_sparse(key: Key[Array, ''], state: State) -> State:
2237 keys = split(key)
2238 state = step_s(keys.pop(), state)
2239 if state.forest.rho is not None:
2240 state = step_theta(keys.pop(), state)
2241 return state
2244@named_call
2245def step_config(state: State) -> State:
2246 config = state.config
2247 config = replace(config, steps_done=config.steps_done + 1)
2248 return replace(state, config=config)