Coverage for src/bartz/mcmcstep/_step.py: 99%
546 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 11:03 +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 var_histogram
52from bartz.mcmcstep._moves import Moves, propose_moves, split_range
53from bartz.mcmcstep._reduction import ReductionConfig
54from bartz.mcmcstep._state import (
55 Forest,
56 State,
57 StepConfig,
58 chol_with_gersh,
59 round_to_pow2,
60 scaled_error_cov_inv,
61 shard_map_state,
62 split_key_for_chains,
63 vmap_chains,
64)
67@jit(donate_argnums=(1,))
68@split_key_for_chains
69@shard_map_state
70@vmap_chains
71@float32_matmuls
72def step(key: Key[Array, ''], state: State) -> State:
73 """
74 Do one MCMC step.
76 Parameters
77 ----------
78 key
79 A jax random key.
80 state
81 A BART mcmc state, as created by `init`.
83 Returns
84 -------
85 The new BART mcmc state.
87 Notes
88 -----
89 The memory of the input state is re-used for the output state, so the input
90 state can not be used any more after calling `step`. All this applies
91 outside of `jax.jit`.
92 """
93 keys = split(key, 4)
95 state = step_resid_inexact_integral(state)
96 state = step_trees(keys.pop(), state)
98 if state.z is not None:
99 state = step_z(keys.pop(), state)
101 if state.error_cov_inv.nu is not None:
102 state = step_error_cov_inv(keys.pop(), state)
104 state = step_sparse(keys.pop(), state)
105 return step_config(state)
108@named_call
109def step_resid_inexact_integral(state: State) -> State:
110 """Accumulate the mean square of the residuals subject to rounding.
112 Adds to `State.resid_inexact_integral` the current mean square of the
113 residuals at or above ``2^(leaf_quantization+1)`` (in
114 `State.resid_eff_scale` units), below which the running residual updates
115 are exact because the `State.resid` dtype spacing still matches the leaf
116 quantum. Without quantization every update rounds and all residuals count.
117 """
118 # filter away residuals below quantization threshold
119 q = state.config.leaf_quantization
120 resid = state.resid
121 if q is None:
122 inexact = resid
123 else:
124 # the quantization grid follows `resid_eff_scale` while `resid` stays
125 # in `resid_unit` units, so the threshold carries their (power of
126 # two, hence exact) ratio; it is cast to the resid dtype to keep the
127 # comparison from widening the n-sized array
128 exact_below = 2.0 ** (q + 1) * state.resid_eff_scale / state.resid_unit
129 inexact = jnp.where(
130 jnp.abs(resid) >= exact_below[..., None].astype(resid.dtype), resid, 0
131 )
133 # mask missing datapoints if any
134 if state.inv_sdev_scale is not None:
135 inexact = jnp.where(state.inv_sdev_scale != 0, inexact, 0)
137 # compute mean square residuals above threshold
138 ms = jnp.einsum(
139 '...n,...n->...', inexact, inexact, preferred_element_type=jnp.float32
140 )
141 if state.config.data_sharded:
142 ms = lax.psum(ms, 'data')
143 ms /= jnp.maximum(state.n_non_missing, 1)
145 # accumulate to the running sum
146 return replace(state, resid_inexact_integral=state.resid_inexact_integral + ms)
149@named_call
150def step_trees(key: Key[Array, ''], state: State) -> State:
151 """
152 Forest sampling step of BART MCMC.
154 Parameters
155 ----------
156 key
157 A jax random key.
158 state
159 A BART mcmc state, as created by `init`.
161 Returns
162 -------
163 The new BART mcmc state.
165 Notes
166 -----
167 This function zeroes the proposal counters.
168 """
169 keys = split(key)
170 moves = propose_moves(keys.pop(), state.forest)
171 return accept_moves_and_sample_leaves(keys.pop(), state, moves)
174@named_call
175def accept_moves_and_sample_leaves(
176 key: Key[Array, ''], state: State, moves: Moves
177) -> State:
178 """
179 Accept or reject the proposed moves and sample the new leaf values.
181 Parameters
182 ----------
183 key
184 A jax random key.
185 state
186 A valid BART mcmc state.
187 moves
188 The proposed moves, see `propose_moves`.
190 Returns
191 -------
192 A new (valid) BART mcmc state.
193 """
194 pso = accept_moves_parallel_stage(key, state, moves)
195 state, moves = accept_moves_sequential_stage(pso)
196 return accept_moves_final_stage(state, moves)
199class Counts(Module):
200 """Number of datapoints in the nodes involved in proposed moves for each tree."""
202 lrt: UInt[Array, '*num_trees 3']
203 """Number of datapoints in the left child, right child, and parent
204 (``= left + right``), stacked along the trailing axis."""
207class PreLkV(Module):
208 """Non-sequential terms of the likelihood ratio for each tree.
210 These terms are derived from the leaf precompute terms (`PreLf`) gathered
211 at the nodes involved in each move. The terms for the left child, right
212 child, and their join (the parent node) are stacked along the axis right
213 after the tree axis. Each term is, in the univariate case, the scalar
215 ``error_cov_inv^2 / (leaf_prior_cov_inv + n * error_cov_inv)``.
217 In the multivariate homoskedastic or scalar-error-scale case, this is the matrix term
219 ``error_cov_inv @ inv(leaf_prior_cov_inv + n * error_cov_inv) @ error_cov_inv``.
221 In the multivariate case with per-component error scales, this is instead
223 ``chol(leaf_prior_cov_inv + n * error_cov_inv)``
225 ``n`` is the number of datapoints in the node, or the likelihood precision
226 scale in the heteroskedastic case.
227 """
229 # `log_sqrt_term` is declared before `lrt` so its single (union-free)
230 # annotation binds the variadic `*num_trees` axis first; otherwise the
231 # runtime typechecker can greedily mis-bind `*num_trees` against the `k`
232 # axis of the `... | ... k k` union (the multivariate and univariate
233 # layouts are rank-ambiguous).
234 log_sqrt_term: Float32[Array, '*num_trees']
235 """The logarithm of the square root term of the likelihood ratio."""
237 lrt: Float32[Array, '*num_trees 3'] | Float32[Array, '*num_trees 3 k k']
238 """Scaled full conditional variance, scaled covariance, or precision
239 cholesky, for the left child, right child, and their join."""
242class PreLf(Module):
243 """Pre-computed terms used to sample leaves from their posterior.
245 These terms can be computed in parallel across trees.
247 For each tree and leaf, the terms are scalars in the univariate case
248 (`PreLfUV`), and matrices/vectors in the multivariate case (`PreLfMV`,
249 `PreLfMVHet`).
251 Abstract base: the layouts differ in rank, so they live in concrete
252 subclasses with union-free annotations; a single class carrying a shape
253 union would make the greedy variadic mis-bind against the ``k`` axes under
254 the runtime typechecker. The concrete class also tags the meaning of
255 `mean_factor`, which drives the dispatch in `precompute_likelihood_terms`
256 and in the sequential stage. The ``num_trees`` axis is variadic so the same
257 annotations also match a per-element layout if vmapped over trees.
258 """
260 mean_factor: AbstractVar[
261 Float32[Array, '*num_trees tree_size']
262 | Float32[Array, '*num_trees k k tree_size']
263 ]
264 """The factor to be right-multiplied by the sum of the scaled residuals to
265 obtain the posterior mean."""
267 centered_leaves: AbstractVar[
268 Float32[Array, '*num_trees tree_size']
269 | Float32[Array, '*num_trees k tree_size']
270 ]
271 """The mean-zero normal values to be added to the posterior mean to
272 obtain the posterior leaf samples."""
275class PreLfUV(PreLf):
276 """`PreLf` for the univariate case."""
278 mean_factor: Float32[Array, '*num_trees tree_size']
279 """``error_cov_inv / prec``, where ``prec`` is the posterior precision of
280 the leaf."""
282 centered_leaves: Float32[Array, '*num_trees tree_size']
283 """Zero-mean normal draws with the posterior variance of each leaf."""
286class PreLfMV(PreLf):
287 """`PreLf` for the multivariate homoskedastic or scalar-error-scale case."""
289 mean_factor: Float32[Array, '*num_trees k k tree_size']
290 """``error_cov_inv @ inv(prec)``, where ``prec`` is the posterior precision
291 of the leaf."""
293 centered_leaves: Float32[Array, '*num_trees k tree_size']
294 """Zero-mean normal draws with the posterior covariance of each leaf."""
296 logdet_prec: Float32[Array, '*num_trees tree_size']
297 """The log-determinant of the posterior precision of each leaf."""
300class PreLfMVHet(PreLf):
301 """`PreLf` for the multivariate case with per-component error scales."""
303 mean_factor: Float32[Array, '*num_trees k k tree_size']
304 """The lower Cholesky factor of the posterior precision of each leaf; the
305 mean solve happens downstream in the sequential stage."""
307 centered_leaves: Float32[Array, '*num_trees k tree_size']
308 """Zero-mean normal draws with the posterior covariance of each leaf."""
311class ParallelStageOut(Module):
312 """The output of `accept_moves_parallel_stage`."""
314 state: State
315 """A partially updated BART mcmc state."""
317 moves: Moves
318 """The proposed moves, with `partial_ratio` set to `None` and
319 `log_trans_prior_ratio` set to its final value."""
321 # `num_trees` stays a fixed (non-variadic) axis: `ParallelStageOut` is always
322 # built with the tree axis present (never per tree under vmap), so the union
323 # is disambiguated by rank/dtype and needs no anchor (cf. `PreLf`).
324 prec_trees: (
325 Float32[Array, 'num_trees tree_size']
326 | UInt32[Array, 'num_trees tree_size']
327 | Float32[Array, 'num_trees k k tree_size']
328 )
329 """The likelihood precision scale in each potential or actual leaf node."""
331 prelkv: PreLkV
332 """Object with pre-computed terms of the likelihood ratios."""
334 prelf: PreLf
335 """Object with pre-computed terms of the leaf samples."""
338@named_call
339def accept_moves_parallel_stage(
340 key: Key[Array, ''], state: State, moves: Moves
341) -> ParallelStageOut:
342 """
343 Pre-compute quantities used to accept moves, in parallel across trees.
345 Parameters
346 ----------
347 key
348 A jax random key.
349 state
350 A BART mcmc state.
351 moves
352 The proposed moves, see `propose_moves`.
354 Returns
355 -------
356 An object with all that could be done in parallel.
357 """
358 # apply the prunes pending from the previous step to the leaf indices,
359 # then, where the new move is grow, modify the state like the move was
360 # accepted.
361 leaf_indices = apply_moves_to_leaf_indices(
362 state.forest.leaf_indices, state.forest.to_prune, state.forest.move_node
363 )
364 state = replace(
365 state,
366 forest=replace(
367 state.forest,
368 var_tree=moves.var_tree,
369 leaf_indices=apply_grow_to_indices(moves, leaf_indices, state.X),
370 leaf_tree=adapt_leaf_trees_to_grow_indices(state.forest.leaf_tree, moves),
371 ),
372 )
374 # update the cached number of datapoints per leaf at the nodes involved
375 # in the moves
376 if (
377 state.forest.min_points_per_decision_node is not None
378 or state.forest.min_points_per_leaf is not None
379 or state.prec_scale is None
380 ):
381 assert state.forest.count_tree is not None
382 count_trees, move_counts = compute_count_trees(
383 state.forest.count_tree, state.forest.leaf_indices, moves, state.config
384 )
385 state = replace(state, forest=replace(state.forest, count_tree=count_trees))
387 # affluence of the nodes touched by each move: whether they would be
388 # growable as leaves (admissible rule + enough datapoints). The children
389 # must also lie within the heap, i.e. not be at the bottom level; the
390 # parent always does. These feed the transition ratio and the final
391 # `affluence_tree` update.
392 _, half = state.forest.var_tree.shape
393 lrt_affluent = (moves.lrt_nodes < half) & moves.lrt_growable
394 if state.forest.min_points_per_decision_node is not None:
395 lrt_affluent &= move_counts.lrt >= state.forest.min_points_per_decision_node
396 moves = replace(moves, lrt_affluent=lrt_affluent)
398 # veto grove move if new leaves don't have enough datapoints
399 if state.forest.min_points_per_leaf is not None:
400 moves = replace(
401 moves,
402 allowed=moves.allowed
403 & jnp.all(
404 move_counts.lrt[..., :2] >= state.forest.min_points_per_leaf, axis=-1
405 ),
406 )
408 # update the cached number of datapoints per leaf, weighted by error
409 # precision scale, at the nodes involved in the moves
410 if state.prec_scale is None:
411 prec_trees = count_trees
412 else:
413 assert state.forest.prec_tree is not None
414 prec_trees = compute_prec_trees(
415 state.forest.prec_tree,
416 state.prec_scale,
417 state.forest.leaf_indices,
418 moves,
419 state.config,
420 )
421 state = replace(state, forest=replace(state.forest, prec_tree=prec_trees))
423 # compute some missing information about moves
424 moves = complete_ratio(moves, state.forest.p_nonterminal)
425 save_ratios = state.forest.log_likelihood is not None
426 state = replace(
427 state,
428 forest=replace(
429 state.forest,
430 grow_prop_count=jnp.sum(moves.grow),
431 prune_prop_count=jnp.sum(moves.allowed & ~moves.grow),
432 log_trans_prior=moves.log_trans_prior_ratio if save_ratios else None,
433 ),
434 )
436 # `prec_trees` and the per-leaf residual sums are in stored `prec_scale`
437 # units, so their common `inv_sdev_unit ** 2` factor is folded into the
438 # error precision once here instead
439 error_cov_inv = scaled_error_cov_inv(state)
440 assert state.forest.leaf_prior_cov_inv is not None
441 prelf = precompute_leaf_terms(
442 key, prec_trees, error_cov_inv, state.forest.leaf_prior_cov_inv
443 )
444 prelkv = precompute_likelihood_terms(
445 error_cov_inv, state.forest.leaf_prior_cov_inv, prelf, moves
446 )
448 return ParallelStageOut(
449 state=state, moves=moves, prec_trees=prec_trees, prelkv=prelkv, prelf=prelf
450 )
453@named_call
454def apply_grow_to_indices(
455 moves: Moves, leaf_indices: UInt[Array, 'num_trees n'], X: UInt[Array, 'p n']
456) -> UInt[Array, 'num_trees n']:
457 """
458 Update the leaf indices to apply a grow move.
460 Parameters
461 ----------
462 moves
463 The proposed moves, see `propose_moves`.
464 leaf_indices
465 The index of the leaf each datapoint falls into, with the prunes
466 pending from the previous step already applied.
467 X
468 The predictors matrix.
470 Returns
471 -------
472 The updated leaf indices.
473 """
474 return _apply_grow_to_indices(moves, leaf_indices, X)
477@partial(vmap_nodoc, in_axes=(0, 0, None))
478def _apply_grow_to_indices(
479 moves: Moves, leaf_indices: UInt[Array, ' n'], X: UInt[Array, 'p n']
480) -> UInt[Array, ' n']:
481 """Implement `apply_grow_to_indices`."""
482 left_child = moves.lrt_nodes[0].astype(leaf_indices.dtype)
483 x: UInt[Array, ' n'] = X[moves.grow_var, :]
484 go_right = x >= moves.grow_split
485 tree_size = jnp.array(2 * moves.var_tree.size)
486 node_to_update = jnp.where(moves.grow, moves.lrt_nodes[2], tree_size)
487 return jnp.where(
488 leaf_indices == node_to_update, left_child + go_right, leaf_indices
489 )
492def _fill_lrt_total(lrt: Shaped[Array, '*k_k 3']) -> Shaped[Array, '*k_k 3']:
493 """Set the total slot of stacked (left, right, total) values to left + right.
495 The left and right slots pass through unchanged, the stale value in the
496 total slot is ignored. Implemented with fusable elementwise operations.
497 """
498 total = lrt[..., 0] + lrt[..., 1]
499 return jnp.where(jnp.arange(3) == 2, total[..., None], lrt)
502@overload
503def _compute_count_or_prec_trees(
504 prec_scale: None,
505 trees: UInt32[Array, 'num_trees tree_size'],
506 leaf_indices: UInt[Array, 'num_trees n'],
507 moves: Moves,
508 config: StepConfig,
509) -> tuple[UInt32[Array, 'num_trees tree_size'], Counts]: ...
512@overload
513def _compute_count_or_prec_trees(
514 prec_scale: Float[Array, ' n'] | Float[Array, 'k k n'],
515 trees: Float32[Array, 'num_trees tree_size']
516 | Float32[Array, 'num_trees k k tree_size'],
517 leaf_indices: UInt[Array, 'num_trees n'],
518 moves: Moves,
519 config: StepConfig,
520) -> (
521 tuple[Float32[Array, 'num_trees tree_size'], None]
522 | tuple[Float32[Array, 'num_trees k k tree_size'], None]
523): ...
526def _compute_count_or_prec_trees(
527 prec_scale: Float[Array, ' n'] | Float[Array, 'k k n'] | None,
528 trees: UInt32[Array, 'num_trees tree_size']
529 | Float32[Array, 'num_trees tree_size']
530 | Float32[Array, 'num_trees k k tree_size'],
531 leaf_indices: UInt[Array, 'num_trees n'],
532 moves: Moves,
533 config: StepConfig,
534) -> (
535 tuple[UInt32[Array, 'num_trees tree_size'], Counts]
536 | tuple[Float32[Array, 'num_trees tree_size'], None]
537 | tuple[Float32[Array, 'num_trees k k tree_size'], None]
538):
539 """Implement `compute_count_trees` and `compute_prec_trees`."""
541 def all_trees() -> (
542 tuple[UInt32[Array, 'num_trees tree_size'], Counts]
543 | tuple[Float32[Array, 'num_trees tree_size'], None]
544 | tuple[Float32[Array, 'num_trees k k tree_size'], None]
545 ):
546 compute = vmap(_compute_count_or_prec_tree, in_axes=(None, 0, 0, 0, None))
547 return compute(prec_scale, trees, leaf_indices, moves, config)
549 if config.prec_count_num_trees is None:
550 return all_trees()
552 batch_size = config.prec_count_num_trees
554 def compute(
555 args: tuple[
556 UInt32[Array, ' tree_size']
557 | Float32[Array, ' tree_size']
558 | Float32[Array, 'k k tree_size'],
559 UInt[Array, ' n'],
560 Moves,
561 ],
562 ) -> (
563 tuple[UInt32[Array, ' tree_size'], Counts]
564 | tuple[Float32[Array, ' tree_size'], None]
565 | tuple[Float32[Array, 'k k tree_size'], None]
566 ):
567 tree, leaf_indices, moves = args
568 return _compute_count_or_prec_tree(
569 prec_scale, tree, leaf_indices, moves, config
570 )
572 def tree_batches() -> (
573 tuple[UInt32[Array, 'num_trees tree_size'], Counts]
574 | tuple[Float32[Array, 'num_trees tree_size'], None]
575 | tuple[Float32[Array, 'num_trees k k tree_size'], None]
576 ):
577 # sliced_map instead of lax.map because under the chain vmap lax.map's
578 # reshape of the tree axis becomes a transpose of `leaf_indices` that
579 # xla materializes in full (cf. `SeqStageInAllTrees.leaf_indices`)
580 return sliced_map(compute, (trees, leaf_indices, moves), batch_size=batch_size)
582 # the tree batching bounds the reduction temporaries on cpu; on gpu those
583 # fuse anyway
584 return lax.platform_dependent(cpu=tree_batches, cuda=all_trees)
587def _compute_count_or_prec_tree(
588 prec_scale: Float[Array, ' n'] | Float[Array, 'k k n'] | None,
589 tree: UInt32[Array, ' tree_size']
590 | Float32[Array, ' tree_size']
591 | Float32[Array, 'k k tree_size'],
592 leaf_indices: UInt[Array, ' n'],
593 moves: Moves,
594 config: StepConfig,
595) -> (
596 tuple[UInt32[Array, ' tree_size'], Counts]
597 | tuple[Float32[Array, ' tree_size'], None]
598 | tuple[Float32[Array, 'k k tree_size'], None]
599):
600 """Update the cached count or precision tree for a single tree."""
601 (tree_size,) = moves.var_tree.shape
602 tree_size *= 2
604 if prec_scale is None:
605 value = 1
606 dtype = jnp.uint32
607 reduction_config = config.count_reduction_config
608 else:
609 value = prec_scale
610 dtype = jnp.float32
611 reduction_config = config.prec_reduction_config
613 # the cached tree is valid at the leaves, and the move only changes the
614 # values at the nodes it involves, so reduce into the move's children alone:
615 # the contiguous pair (left, right) = (2 * node, 2 * node + 1) = lrt_nodes[:2]
616 lr = reduction_config._reduce( # noqa: SLF001
617 value,
618 leaf_indices,
619 size=tree_size,
620 subset_start=moves.lrt_nodes[0],
621 subset_length=2,
622 dtype=dtype,
623 data_sharded=config.data_sharded,
624 )
626 # write the children sums into the cache along with their total at the
627 # parent node (a non-leaf in the post-grow indexing the reduce runs on);
628 # the precision-scaled version of the counts is not needed because the likelihood
629 # terms are derived from the leaf terms
630 total = lr[..., 0] + lr[..., 1]
631 lrt = jnp.concatenate([lr, total[..., None]], axis=-1)
632 tree = tree.at[..., moves.lrt_nodes].set(lrt)
634 if prec_scale is None:
635 return tree, Counts(lrt=lrt)
636 else:
637 return tree, None
640@named_call
641def compute_count_trees(
642 count_trees: UInt32[Array, 'num_trees tree_size'],
643 leaf_indices: UInt[Array, 'num_trees n'],
644 moves: Moves,
645 config: StepConfig,
646) -> tuple[UInt32[Array, 'num_trees tree_size'], Counts]:
647 """
648 Update the cached number of datapoints per leaf at the moves' nodes.
650 Parameters
651 ----------
652 count_trees
653 The cached number of points in each leaf; valid at the leaves of the
654 pre-move trees.
655 leaf_indices
656 The index of the leaf each datapoint falls into, with the deeper version
657 of the tree (post-GROW, pre-PRUNE).
658 moves
659 The proposed moves, see `propose_moves`.
660 config
661 The MCMC configuration.
663 Returns
664 -------
665 count_trees : UInt32[Array, 'num_trees tree_size']
666 The updated cache, valid in each potential or actual leaf node.
667 counts : Counts
668 The counts of the number of points in the leaves grown or pruned by the
669 moves.
670 """
671 return _compute_count_or_prec_trees(None, count_trees, leaf_indices, moves, config)
674@named_call
675def compute_prec_trees(
676 prec_trees: Float32[Array, 'num_trees tree_size']
677 | Float32[Array, 'num_trees k k tree_size'],
678 prec_scale: Float[Array, ' n'] | Float[Array, 'k k n'],
679 leaf_indices: UInt[Array, 'num_trees n'],
680 moves: Moves,
681 config: StepConfig,
682) -> Float32[Array, 'num_trees tree_size'] | Float32[Array, 'num_trees k k tree_size']:
683 """
684 Update the cached per-leaf likelihood precision scale at the moves' nodes.
686 Parameters
687 ----------
688 prec_trees
689 The cached likelihood precision scale in each leaf; valid at the leaves
690 of the pre-move trees.
691 prec_scale
692 The scale of the precision of the error on each datapoint.
693 leaf_indices
694 The index of the leaf each datapoint falls into, with the deeper version
695 of the tree (post-GROW, pre-PRUNE).
696 moves
697 The proposed moves, see `propose_moves`.
698 config
699 The MCMC configuration.
701 Returns
702 -------
703 The updated cache, valid in each potential or actual leaf node.
704 """
705 trees, _ = _compute_count_or_prec_trees(
706 prec_scale, prec_trees, leaf_indices, moves, config
707 )
708 return trees
711@partial(vmap_nodoc, in_axes=(0, None))
712def complete_ratio(moves: Moves, p_nonterminal: Float32[Array, ' tree_size']) -> Moves:
713 """
714 Complete non-likelihood MH ratio calculation.
716 This function adds the probability of choosing a prune move over the grow
717 move in the inverse transition, and the prior odds that the modified node
718 is nonterminal with terminal children.
720 Parameters
721 ----------
722 moves
723 The proposed moves. Must have already been updated to keep into account
724 the thresholds on the number of datapoints per node, this happens in
725 `accept_moves_parallel_stage`.
726 p_nonterminal
727 The a priori probability of each node being nonterminal conditional on
728 its ancestors, including at the maximum depth where it should be zero.
730 Returns
731 -------
732 The updated moves, with `partial_ratio=None` and `log_trans_prior_ratio` set.
733 """
734 assert moves.lrt_affluent is not None
736 # can the children be grown by the proposal? `lrt_affluent` already folds
737 # in the `min_points_per_decision_node` threshold, because the grow
738 # proposal draws from the pool of leaves that pass it. This enters only the
739 # transition probability.
741 # p_prune if grow
742 other_growable_leaves = moves.num_growable >= 2
743 grow_again_allowed = other_growable_leaves | jnp.any(moves.lrt_affluent[:2])
744 grow_p_prune = jnp.where(grow_again_allowed, 0.5, 1.0)
746 # p_prune if prune
747 prune_p_prune = jnp.where(moves.num_growable, 0.5, 1)
749 # select p_prune
750 p_prune = jnp.where(moves.grow, grow_p_prune, prune_p_prune)
752 # prior odds of the node being nonterminal, times the prior probability of
753 # both children being terminal. The children terminality uses the
754 # admissibility ignoring counts, because the standard BART prior conditions
755 # the non-terminal probability only on the existence of available decision
756 # rules, not on the count thresholds (which are a bartz proposal-efficiency
757 # device, not part of the target distribution). The fill value avoids a 0
758 # and then an inf in the log if the move is not allowed and the indices are
759 # out of bounds.
760 pnt = p_nonterminal.at[moves.lrt_nodes].get(mode='fill', fill_value=0.5)
761 prior_ratio = pnt[2] / (1 - pnt[2]) * jnp.prod(1 - pnt[:2] * moves.lrt_growable[:2])
763 assert moves.partial_ratio is not None
764 return replace(
765 moves,
766 log_trans_prior_ratio=jnp.log(moves.partial_ratio * prior_ratio * p_prune),
767 partial_ratio=None,
768 )
771@named_call
772def adapt_leaf_trees_to_grow_indices(
773 leaf_trees: Float[Array, 'num_trees tree_size']
774 | Float[Array, 'num_trees k tree_size'],
775 moves: Moves,
776) -> Float[Array, 'num_trees tree_size'] | Float[Array, 'num_trees k tree_size']:
777 """
778 Modify leaves such that post-grow indices work on the original tree.
780 The value of the leaf to grow is copied to what would be its children if the
781 grow move was accepted.
783 Parameters
784 ----------
785 leaf_trees
786 The leaf values.
787 moves
788 The proposed moves, see `propose_moves`.
790 Returns
791 -------
792 The modified leaf values.
793 """
794 return _adapt_leaf_trees_to_grow_indices(leaf_trees, moves)
797@vmap_nodoc
798def _adapt_leaf_trees_to_grow_indices(
799 leaf_trees: Float[Array, ' tree_size'] | Float[Array, ' k tree_size'], moves: Moves
800) -> Float[Array, ' tree_size'] | Float[Array, ' k tree_size']:
801 """Implement `adapt_leaf_trees_to_grow_indices`."""
802 # the parent slot is written back unchanged to share a single scatter
803 values_at_node = leaf_trees[..., moves.lrt_nodes[2]]
804 return leaf_trees.at[
805 ..., jnp.where(moves.grow, moves.lrt_nodes, leaf_trees.size)
806 ].set(values_at_node[..., None])
809def _logdet_from_chol(L: Float32[Array, '... k k']) -> Float32[Array, '...']:
810 """Compute logdet of A = LL' via Cholesky (sum of log of diag^2)."""
811 diags: Float32[Array, '... k'] = jnp.diagonal(L, axis1=-2, axis2=-1)
812 return 2.0 * jnp.sum(jnp.log(diags), axis=-1)
815def compute_B(
816 error_cov_inv: Float32[Array, 'k k'], resid: Float32[Array, 'k k *tree_size']
817) -> Float32[Array, ' k *tree_size']:
818 """Compute the leaf score from the leaf precision-scaled sum of residuals."""
819 return jnp.einsum('ab,ab...->a...', error_cov_inv, resid)
822def _precompute_leaf_terms_uv(
823 key: Key[Array, ''],
824 prec_trees: Float32[Array, 'num_trees tree_size']
825 | UInt32[Array, 'num_trees tree_size'],
826 error_cov_inv: Float32[Array, ''],
827 leaf_prior_cov_inv: Float32[Array, ''],
828 z: Float32[Array, 'num_trees tree_size'] | None = None,
829) -> PreLfUV:
830 prec_lk = prec_trees * error_cov_inv
831 var_post = jnp.reciprocal(prec_lk + leaf_prior_cov_inv)
832 if z is None:
833 z = random.normal(key, prec_trees.shape, error_cov_inv.dtype)
834 return PreLfUV(
835 mean_factor=var_post * error_cov_inv,
836 # | mean = mean_lk * prec_lk * var_post
837 # | resid_tree = mean_lk * prec_tree -->
838 # | --> mean_lk = resid_tree / prec_tree (kind of)
839 # | mean_factor =
840 # | = mean / resid_tree =
841 # | = resid_tree / prec_tree * prec_lk * var_post / resid_tree =
842 # | = 1 / prec_tree * prec_tree / sigma2 * var_post =
843 # | = var_post / sigma2
844 centered_leaves=z * jnp.sqrt(var_post),
845 )
848def _precompute_leaf_terms_mv(
849 key: Key[Array, ''],
850 prec_trees: Float32[Array, 'num_trees tree_size']
851 | UInt32[Array, 'num_trees tree_size'],
852 error_cov_inv: Float32[Array, 'k k'],
853 leaf_prior_cov_inv: Float32[Array, 'k k'],
854 z: Float32[Array, 'num_trees k tree_size'] | None = None,
855) -> PreLfMV:
856 num_trees, tree_size = prec_trees.shape
857 k, _ = error_cov_inv.shape
858 if z is None: 858 ↛ 861line 858 didn't jump to line 861 because the condition on line 858 was always true
859 z = random.normal(key, (num_trees, k, tree_size))
861 def per_leaf(
862 prec: Float32[Array, ''] | UInt32[Array, ''], z: Float32[Array, ' k']
863 ) -> tuple[Float32[Array, 'k k'], Float32[Array, ' k'], Float32[Array, '']]:
864 L_prec = chol_with_gersh(leaf_prior_cov_inv + prec * error_cov_inv)
865 Y = solve_triangular(L_prec, error_cov_inv, lower=True)
866 mean_factor = solve_triangular(L_prec, Y, trans='T', lower=True).mT
867 centered = solve_triangular(L_prec, z[:, None], trans='T', lower=True).squeeze(
868 -1
869 )
870 # only a few leaves per tree end up using their logdet, but reducing
871 # right away is lighter on memory than storing diagonals for later
872 return mean_factor, centered, _logdet_from_chol(L_prec)
874 # vmap over trees then over leaves; the leaf axis is trailing in both
875 # `prec_trees`/`z` (in_axes) and the stored output (out_axes=-1)
876 return PreLfMV(*vmap(vmap(per_leaf, in_axes=(0, -1), out_axes=-1))(prec_trees, z))
879def _precompute_leaf_terms_mv_het(
880 key: Key[Array, ''],
881 prec_trees: Float32[Array, 'num_trees k k tree_size'],
882 error_cov_inv: Float32[Array, 'k k'],
883 leaf_prior_cov_inv: Float32[Array, 'k k'],
884 z: Float32[Array, 'num_trees k tree_size'] | None = None,
885) -> PreLfMVHet:
886 num_trees, k, _, tree_size = prec_trees.shape
887 if z is None: 887 ↛ 890line 887 didn't jump to line 890 because the condition on line 887 was always true
888 z = random.normal(key, (num_trees, k, tree_size))
890 def per_leaf(
891 prec: Float32[Array, 'k k'], z: Float32[Array, ' k']
892 ) -> tuple[Float32[Array, 'k k'], Float32[Array, ' k']]:
893 # mean_factor stores the precision cholesky itself; the mean solve happens
894 # downstream in `accept_move_and_sample_leaves`
895 L_prec = chol_with_gersh(leaf_prior_cov_inv + error_cov_inv * prec)
896 centered = solve_triangular(L_prec, z[:, None], trans='T', lower=True).squeeze(
897 -1
898 )
899 return L_prec, centered
901 # vmap over trees then over leaves; the leaf axis is trailing in both
902 # `prec_trees`/`z` (in_axes=-1) and the stored output (out_axes=-1)
903 return PreLfMVHet(
904 *vmap(vmap(per_leaf, in_axes=(-1, -1), out_axes=-1))(prec_trees, z)
905 )
908@named_call
909def precompute_leaf_terms(
910 key: Key[Array, ''],
911 prec_trees: Float32[Array, 'num_trees tree_size']
912 | UInt32[Array, 'num_trees tree_size']
913 | Float32[Array, 'num_trees k k tree_size'],
914 error_cov_inv: Float32[Array, ''] | Float32[Array, 'k k'],
915 leaf_prior_cov_inv: Float32[Array, ''] | Float32[Array, 'k k'],
916 z: Float32[Array, 'num_trees tree_size']
917 | Float32[Array, 'num_trees k tree_size']
918 | None = None,
919) -> PreLf:
920 """
921 Pre-compute terms used to sample leaves from their posterior.
923 Handles both univariate and multivariate cases based on the shape of the
924 input arrays.
926 Parameters
927 ----------
928 key
929 A jax random key.
930 prec_trees
931 The likelihood precision scale in each potential or actual leaf node.
932 error_cov_inv
933 The inverse error variance (univariate) or the inverse of error
934 covariance matrix (multivariate). If `prec_scale` is set, this is the
935 global error precision factor, with the squared `inv_sdev_unit` folded in
936 (see `bartz.mcmcstep.State.inv_sdev_unit`).
937 leaf_prior_cov_inv
938 The inverse prior variance of each leaf (univariate) or the inverse of
939 prior covariance matrix of each leaf (multivariate).
940 z
941 Optional standard normal noise to use for sampling the centered leaves.
942 This is intended for testing purposes only.
944 Returns
945 -------
946 Pre-computed terms for leaf sampling.
947 """
948 if error_cov_inv.ndim == 0:
949 return _precompute_leaf_terms_uv(
950 key, prec_trees, error_cov_inv, leaf_prior_cov_inv, z
951 )
952 elif prec_trees.ndim == 4:
953 return _precompute_leaf_terms_mv_het(
954 key, prec_trees, error_cov_inv, leaf_prior_cov_inv, z
955 )
956 else:
957 return _precompute_leaf_terms_mv(
958 key, prec_trees, error_cov_inv, leaf_prior_cov_inv, z
959 )
962@vmap_nodoc
963def _gather_lrt(
964 leaf_values: Float32[Array, '*k_k tree_size'], lrt_nodes: Int32[Array, ' 3']
965) -> Float32[Array, ' 3 *k_k']:
966 """Gather per-tree leaf values at the left child, right child, and parent."""
967 return jnp.moveaxis(leaf_values[..., lrt_nodes], -1, 0)
970def _precompute_likelihood_terms_uv(
971 error_cov_inv: Float32[Array, ''],
972 leaf_prior_cov_inv: Float32[Array, ''],
973 prelf: PreLfUV,
974 lrt_nodes: Int32[Array, 'num_trees 3'],
975) -> PreLkV:
976 # mean_factor is error_cov_inv / prec, complete the sandwich
977 lrt = error_cov_inv * _gather_lrt(prelf.mean_factor, lrt_nodes)
978 # the same value with the prior-only precision, computed with the same
979 # operations as in `_precompute_leaf_terms_uv` such that it matches `lrt`
980 # bitwise on empty nodes and the ratio is exactly 1 without data
981 prior_lrt = error_cov_inv * (jnp.reciprocal(leaf_prior_cov_inv) * error_cov_inv)
982 log_sqrt_term = jnp.log(lrt[..., 0] * lrt[..., 1] / (prior_lrt * lrt[..., 2])) / 2
983 return PreLkV(lrt=lrt, log_sqrt_term=log_sqrt_term)
986def _precompute_likelihood_terms_mv(
987 error_cov_inv: Float32[Array, 'k k'],
988 leaf_prior_cov_inv: Float32[Array, 'k k'],
989 prelf: PreLfMV,
990 lrt_nodes: Int32[Array, 'num_trees 3'],
991) -> PreLkV:
992 logdet_prior = _logdet_from_chol(chol_with_gersh(leaf_prior_cov_inv))
993 logdet_prec = _gather_lrt(prelf.logdet_prec, lrt_nodes)
994 log_sqrt_term = (logdet_prior + logdet_prec @ jnp.array([-1.0, -1.0, 1.0])) / 2
996 # mean_factor is error_cov_inv @ inv(prec), complete the sandwich
997 mean_factor = _gather_lrt(prelf.mean_factor, lrt_nodes) # (num_trees, 3, k, k)
998 return PreLkV(lrt=mean_factor @ error_cov_inv, log_sqrt_term=log_sqrt_term)
1001def _precompute_likelihood_terms_mv_het(
1002 leaf_prior_cov_inv: Float32[Array, 'k k'],
1003 prelf: PreLfMVHet,
1004 lrt_nodes: Int32[Array, 'num_trees 3'],
1005) -> PreLkV:
1006 logdet_prior = _logdet_from_chol(chol_with_gersh(leaf_prior_cov_inv))
1008 # mean_factor is the precision cholesky itself
1009 L = _gather_lrt(prelf.mean_factor, lrt_nodes) # (num_trees, 3, k, k)
1010 log_sqrt_term = (
1011 logdet_prior + _logdet_from_chol(L) @ jnp.array([-1.0, -1.0, 1.0])
1012 ) / 2
1013 return PreLkV(lrt=L, log_sqrt_term=log_sqrt_term)
1016@named_call
1017def precompute_likelihood_terms(
1018 error_cov_inv: Float32[Array, ''] | Float32[Array, 'k k'],
1019 leaf_prior_cov_inv: Float32[Array, ''] | Float32[Array, 'k k'],
1020 prelf: PreLf,
1021 moves: Moves,
1022) -> PreLkV:
1023 """
1024 Pre-compute terms used in the likelihood ratio of the acceptance step.
1026 The likelihood ratio terms are mostly a subset of the leaf sampling terms,
1027 so they are derived from `prelf`, gathered at the nodes involved in the
1028 moves.
1030 Parameters
1031 ----------
1032 error_cov_inv
1033 The inverse error variance (univariate) or the inverse of the error
1034 covariance matrix (multivariate). If `prec_scale` is set, this is the
1035 global error precision factor, with the squared `inv_sdev_unit` folded in
1036 (see `bartz.mcmcstep.State.inv_sdev_unit`).
1037 leaf_prior_cov_inv
1038 The inverse prior variance of each leaf (univariate) or the inverse of
1039 prior covariance matrix of each leaf (multivariate).
1040 prelf
1041 The pre-computed terms of the leaf sampling, see `precompute_leaf_terms`.
1042 moves
1043 The proposed moves, see `propose_moves`.
1045 Returns
1046 -------
1047 Pre-computed terms of the likelihood ratio, one per tree.
1048 """
1049 if isinstance(prelf, PreLfUV):
1050 return _precompute_likelihood_terms_uv(
1051 error_cov_inv, leaf_prior_cov_inv, prelf, moves.lrt_nodes
1052 )
1053 elif isinstance(prelf, PreLfMVHet):
1054 return _precompute_likelihood_terms_mv_het(
1055 leaf_prior_cov_inv, prelf, moves.lrt_nodes
1056 )
1057 else:
1058 assert isinstance(prelf, PreLfMV)
1059 return _precompute_likelihood_terms_mv(
1060 error_cov_inv, leaf_prior_cov_inv, prelf, moves.lrt_nodes
1061 )
1064@named_call
1065def accept_moves_sequential_stage(pso: ParallelStageOut) -> tuple[State, Moves]:
1066 """
1067 Accept/reject the moves one tree at a time.
1069 This is the most performance-sensitive function because it contains all and
1070 only the parts of the algorithm that can not be parallelized across trees.
1072 Parameters
1073 ----------
1074 pso
1075 The output of `accept_moves_parallel_stage`.
1077 Returns
1078 -------
1079 state : State
1080 A partially updated BART mcmc state.
1081 moves : Moves
1082 The accepted/rejected moves, with `acc` and `to_prune` set.
1083 """
1085 def loop(
1086 resid: Float[Array, ' n'] | Float[Array, ' k n'], pt: SeqStageInPerTree
1087 ) -> tuple[
1088 Float[Array, ' n'] | Float[Array, ' k n'],
1089 tuple[
1090 Float[Array, ' tree_size'] | Float[Array, ' k tree_size'],
1091 Bool[Array, ''],
1092 Bool[Array, ''],
1093 Float32[Array, ''] | None,
1094 ],
1095 ]:
1096 resid, leaf_tree, acc, to_prune, lkratio = accept_move_and_sample_leaves(
1097 resid,
1098 SeqStageInAllTrees(
1099 pso.state.X,
1100 pso.state.forest.leaf_indices,
1101 pso.state.config.resid_reduction_config,
1102 pso.state.config.data_sharded,
1103 pso.state.prec_scale,
1104 pso.state.forest.log_likelihood is not None,
1105 scaled_error_cov_inv(pso.state)
1106 if isinstance(pso.prelf, PreLfMVHet)
1107 else None,
1108 pso.state.forest.leaf_unit,
1109 pso.state.resid_unit,
1110 pso.state.resid_eff_scale,
1111 pso.state.config.leaf_quantization,
1112 ),
1113 pt,
1114 )
1115 return resid, (leaf_tree, acc, to_prune, lkratio)
1117 num_trees, _ = pso.state.forest.leaf_indices.shape
1118 pts = SeqStageInPerTree(
1119 pso.state.forest.leaf_tree,
1120 pso.prec_trees,
1121 pso.moves,
1122 jnp.arange(num_trees, dtype=minimal_unsigned_dtype(num_trees - 1)),
1123 pso.prelkv,
1124 pso.prelf,
1125 )
1126 resid, (leaf_trees, acc, to_prune, lkratio) = lax.scan(
1127 loop, pso.state.resid, pts, unroll=pso.state.config.sequential_unroll
1128 )
1130 state = replace(
1131 pso.state,
1132 resid=resid,
1133 forest=replace(pso.state.forest, leaf_tree=leaf_trees, log_likelihood=lkratio),
1134 )
1135 moves = replace(pso.moves, acc=acc, to_prune=to_prune)
1137 return state, moves
1140class SeqStageInAllTrees(Module):
1141 """The inputs to `accept_move_and_sample_leaves` that are shared by all trees."""
1143 X: UInt[Array, 'p n']
1144 """The predictors."""
1146 leaf_indices: UInt[Array, 'num_trees n']
1147 """The leaf indices for the largest version of each tree compatible with
1148 the move. Kept whole here and sliced per tree with
1149 `SeqStageInPerTree.tree_index`, instead of consumed as scan xs, because
1150 the scan batching rules move the chain axis of closed-over values to the
1151 front, matching the storage layout, while the xs rule would transpose it
1152 to after the tree axis."""
1154 resid_reduction_config: ReductionConfig
1155 """How to sum the residuals in each leaf."""
1157 data_sharded: bool = field(static=True)
1158 """Whether the data axis is sharded across devices."""
1160 prec_scale: Float[Array, ' n'] | Float[Array, 'k k n'] | None
1161 """The scale of the precision of the error on each datapoint. If None, it
1162 is assumed to be 1."""
1164 save_ratios: bool = field(static=True)
1165 """Whether to save the acceptance ratios."""
1167 error_cov_inv: Float32[Array, 'k k'] | None
1168 """The global error precision scale, with the squared `inv_sdev_unit` folded
1169 in (see `bartz.mcmcstep.State.inv_sdev_unit`). Set only in the
1170 multivariate case with per-component error scales, where the sequential
1171 stage needs it to compute the leaf scores."""
1173 leaf_unit: Float32[Array, ''] | Float32[Array, ' k']
1174 """The storage unit of the leaf values, see `bartz.mcmcstep.Forest.leaf_unit`."""
1176 resid_unit: Float32[Array, ''] | Float32[Array, ' k']
1177 """The storage unit of the residuals, see `bartz.mcmcstep.State.resid_unit`."""
1179 resid_eff_scale: Float32[Array, ''] | Float32[Array, ' k']
1180 """The measured scale of the residuals, see `bartz.mcmcstep.State.resid_eff_scale`."""
1182 leaf_quantization: Int32[Array, ''] | None
1183 """Leaf quantization setting, see `bartz.mcmcstep.StepConfig`."""
1186class SeqStageInPerTree(Module):
1187 """The inputs to `accept_move_and_sample_leaves` that are separate for each tree."""
1189 # Although consumed one tree at a time by `lax.scan`, this object is only
1190 # ever constructed in the stacked (batched) form fed to the scan, so
1191 # `num_trees` stays a fixed (non-variadic) leading axis disambiguated by
1192 # rank/dtype (cf. `ParallelStageOut`); the per-tree slices reach `loop` via
1193 # scan, which does not re-run `__init__`.
1194 leaf_tree: (
1195 Float[Array, 'num_trees tree_size'] | Float[Array, 'num_trees k tree_size']
1196 )
1197 """The leaf values of the trees."""
1199 prec_tree: (
1200 Float32[Array, 'num_trees tree_size']
1201 | UInt32[Array, 'num_trees tree_size']
1202 | Float32[Array, 'num_trees k k tree_size']
1203 )
1204 """The likelihood precision scale in each potential or actual leaf node."""
1206 move: Moves
1207 """The proposed move, see `propose_moves`."""
1209 tree_index: UInt[Array, ' num_trees']
1210 """The index of the tree, used to slice `SeqStageInAllTrees.leaf_indices`."""
1212 prelkv: PreLkV
1213 """The pre-computed terms of the likelihood ratio which are specific to the tree."""
1215 prelf: PreLf
1216 """The pre-computed terms of the leaf sampling which are specific to the tree."""
1219@named_call
1220def accept_move_and_sample_leaves(
1221 resid: Float[Array, ' n'] | Float[Array, ' k n'],
1222 at: SeqStageInAllTrees,
1223 pt: SeqStageInPerTree,
1224) -> tuple[
1225 Float[Array, ' n'] | Float[Array, ' k n'],
1226 Float[Array, ' tree_size'] | Float[Array, ' k tree_size'],
1227 Bool[Array, ''],
1228 Bool[Array, ''],
1229 Float32[Array, ''] | None,
1230]:
1231 """
1232 Accept or reject a proposed move and sample the new leaf values.
1234 Parameters
1235 ----------
1236 resid
1237 The residuals (data minus forest value), in units of `at.resid_unit`
1238 and a possibly narrow dtype (see `State.resid`). The reduction over the
1239 residuals runs in these units; the per-leaf sums are scaled back to data
1240 units afterwards.
1241 at
1242 The inputs that are the same for all trees.
1243 pt
1244 The inputs that are separate for each tree.
1246 Returns
1247 -------
1248 resid : Float[Array, 'n'] | Float[Array, ' k n']
1249 The updated residuals, in the same stored units and dtype as the input.
1250 leaf_tree : Float[Array, 'tree_size'] | Float[Array, ' k tree_size']
1251 The new leaf values of the tree.
1252 acc : Bool[Array, '']
1253 Whether the move was accepted.
1254 to_prune : Bool[Array, '']
1255 Whether, to reflect the acceptance status of the move, the state should
1256 be updated by pruning the leaves involved in the move.
1257 log_lk_ratio : Float32[Array, ''] | None
1258 The logarithm of the likelihood ratio for the move. `None` if not to be
1259 saved.
1260 """
1261 leaf_indices = at.leaf_indices[pt.tree_index, :]
1263 # sum residuals in each leaf, in tree proposed by grow move
1264 if at.prec_scale is None:
1265 scaled_resid = resid
1266 else:
1267 scaled_resid = resid * at.prec_scale
1269 tree_size = pt.leaf_tree.shape[-1] # 2**d
1271 resid_tree = sum_resid(
1272 scaled_resid,
1273 leaf_indices,
1274 tree_size,
1275 at.resid_reduction_config,
1276 at.data_sharded,
1277 )
1279 # the residuals are stored and reduced in units of `resid_unit` (and a
1280 # possibly narrow dtype) for bandwidth; the float32 per-leaf sums are scaled
1281 # back to data units here, after the reduction and before everything else.
1282 # with error scales, the sums keep the stored `prec_scale` units, whose
1283 # `inv_sdev_unit ** 2` factor is folded into the error precision instead
1284 resid_tree *= at.resid_unit[..., None]
1286 # convert the starting tree to data units; multiplying by the float32
1287 # scale also takes care of upcasting narrow leaf dtypes
1288 prev_leaf_tree = at.leaf_unit[..., None] * pt.leaf_tree
1290 # subtract starting tree from function
1291 resid_tree += pt.prec_tree * prev_leaf_tree
1293 # sum residuals in parent node modified by move and compute likelihood;
1294 # the children slots are written back unchanged to share a single scatter
1295 assert pt.move.lrt_nodes.dtype == jnp.int32
1296 resid_lrt = _fill_lrt_total(resid_tree[..., pt.move.lrt_nodes])
1297 resid_tree = resid_tree.at[..., pt.move.lrt_nodes].set(resid_lrt)
1299 log_lk_ratio = compute_likelihood_ratio(resid_lrt, pt.prelkv, at.error_cov_inv)
1301 # calculate accept/reject ratio; the ratio is filled in by `complete_ratio`
1302 assert pt.move.log_trans_prior_ratio is not None
1303 log_ratio = pt.move.log_trans_prior_ratio + log_lk_ratio
1304 log_ratio = jnp.where(pt.move.grow, log_ratio, -log_ratio)
1305 if not at.save_ratios:
1306 log_lk_ratio = None
1308 # determine whether to accept the move
1309 acc = pt.move.allowed & (pt.move.logu <= log_ratio)
1311 # compute leaves posterior and sample leaves
1312 if at.error_cov_inv is not None:
1313 # multivariate w/ per-component error scales
1314 b_tree = compute_B(at.error_cov_inv, resid_tree) # (k, 2**d)
1315 l_lead = jnp.moveaxis(pt.prelf.mean_factor, -1, 0) # (2**d, k, k)
1316 b_lead = b_tree.T[:, :, None] # (2**d, k, 1)
1317 y = solve_triangular(l_lead, b_lead, lower=True)
1318 mu = solve_triangular(l_lead, y, lower=True, trans='T').squeeze(-1)
1319 mean_post = mu.T # (k, 2**d)
1320 elif resid.ndim > 1:
1321 # multivariate homoskedastic or scalar error scales
1322 mean_post = jnp.einsum('kil,kl->il', pt.prelf.mean_factor, resid_tree)
1323 else:
1324 # univariate
1325 mean_post = resid_tree * pt.prelf.mean_factor
1326 leaf_tree = mean_post + pt.prelf.centered_leaves
1328 # copy leaves around such that the leaf indices point to the correct leaf;
1329 # the parent slot is written back unchanged to share a single scatter.
1330 # this mirroring persists into the output state, where it keeps the
1331 # not-yet-pruned `leaf_indices` valid to evaluate the trees
1332 to_prune = acc ^ pt.move.grow
1333 leaf_tree = leaf_tree.at[
1334 ..., jnp.where(to_prune, pt.move.lrt_nodes, tree_size)
1335 ].set(leaf_tree[..., pt.move.lrt_nodes[2], None])
1337 # round the new leaves to the storage units and dtype; the residuals are
1338 # then updated with the rounded values to stay consistent with the trees
1339 leaf_tree = leaf_tree / at.leaf_unit[..., None]
1340 if at.leaf_quantization is not None:
1341 # quantize the leaves such that the residual updates below are mostly
1342 # exact. the target grid is the spacing of `resid.dtype` values of
1343 # magnitude 2^leaf_quantization in units of the measured residual
1344 # scale, so 2^(leaf_quantization - nmant) * resid_eff_scale in data
1345 # units; dividing by leaf_unit converts it to the leaf units
1346 # `leaf_tree` is in here. the scales are powers of two, so the quantum
1347 # is one too and the leaf storage dtype's own (coarser) rounding below
1348 # preserves multiples of it.
1349 nmant = jnp.finfo(resid.dtype).nmant
1350 quantum = (at.resid_eff_scale / at.leaf_unit) * 2.0 ** (
1351 at.leaf_quantization - nmant
1352 )
1353 leaf_tree = jnp.round(leaf_tree / quantum[..., None]) * quantum[..., None]
1354 # round to the storage dtype explicitly before converting: xla's
1355 # excess-precision optimization (on by default) may otherwise elide the
1356 # narrowing conversion in the `leaf_delta` computation below, updating the
1357 # residuals with unrounded leaves while the trees store rounded ones
1358 finfo = jnp.finfo(pt.leaf_tree.dtype)
1359 leaf_tree = lax.reduce_precision(leaf_tree, finfo.nexp, finfo.nmant)
1360 leaf_tree = leaf_tree.astype(pt.leaf_tree.dtype)
1362 # replace old tree with new tree in function values; the per-leaf data-unit
1363 # delta is converted back to the stored residual units and dtype *before* the
1364 # scatter, so the n-sized update stays in the narrow `resid` storage
1365 leaf_delta = prev_leaf_tree - at.leaf_unit[..., None] * leaf_tree
1366 delta = (leaf_delta / at.resid_unit[..., None]).astype(resid.dtype)
1367 resid += delta[..., leaf_indices]
1369 return resid, leaf_tree, acc, to_prune, log_lk_ratio
1372@named_call
1373def sum_resid(
1374 scaled_resid: Float[Array, ' n'] | Float[Array, 'k n'] | Float[Array, 'k k n'],
1375 leaf_indices: UInt[Array, ' n'],
1376 tree_size: int,
1377 reduction_config: ReductionConfig,
1378 data_sharded: bool,
1379) -> (
1380 Float32[Array, ' {tree_size}']
1381 | Float32[Array, 'k {tree_size}']
1382 | Float32[Array, 'k k {tree_size}']
1383):
1384 """
1385 Sum the residuals in each leaf.
1387 Parameters
1388 ----------
1389 scaled_resid
1390 The residuals (data minus forest value) multiplied by the error
1391 precision scale.
1392 leaf_indices
1393 The leaf indices of the tree (in which leaf each data point falls into).
1394 tree_size
1395 The size of the tree array (2 ** d).
1396 reduction_config
1397 How to sum the residuals in each leaf.
1398 data_sharded
1399 Whether the data axis is sharded; if true, the result is psum-reduced
1400 across the ``'data'`` axis of the enclosing `shard_map`.
1402 Returns
1403 -------
1404 The per-leaf sum, with the same leading dimensions as ``scaled_resid`` and a trailing axis over the leaves.
1405 """
1406 return reduction_config._reduce( # noqa: SLF001
1407 scaled_resid,
1408 leaf_indices,
1409 size=tree_size,
1410 dtype=jnp.float32,
1411 data_sharded=data_sharded,
1412 )
1415def _compute_likelihood_ratio_uv(
1416 resid_lrt: Float32[Array, ' 3'], prelkv: PreLkV
1417) -> Float32[Array, '']:
1418 # quadratic form r * v * r for each of the (left, right, total) terms
1419 qf = resid_lrt * resid_lrt * prelkv.lrt
1420 exp_term = 0.5 * (qf @ jnp.array([1.0, 1.0, -1.0]))
1421 return prelkv.log_sqrt_term + exp_term
1424def _compute_likelihood_ratio_mv(
1425 resid_lrt: Float32[Array, 'k 3'], prelkv: PreLkV
1426) -> Float32[Array, '']:
1427 # quadratic form r' M r for each of the (left, right, total) terms
1428 qf = jnp.einsum('it,tij,jt->t', resid_lrt, prelkv.lrt, resid_lrt)
1429 exp_term = 0.5 * (qf @ jnp.array([1.0, 1.0, -1.0]))
1430 return prelkv.log_sqrt_term + exp_term
1433def _compute_likelihood_ratio_mv_het(
1434 resid_lrt: Float32[Array, 'k k 3'],
1435 error_cov_inv: Float32[Array, 'k k'],
1436 prelkv: PreLkV,
1437) -> Float32[Array, '']:
1438 b = compute_B(error_cov_inv, resid_lrt) # (k, 3)
1439 y = solve_triangular(prelkv.lrt, b.T[..., None], lower=True).squeeze(-1) # (3, k)
1440 qf = jnp.einsum('ti,ti->t', y, y)
1441 exp_term = 0.5 * (qf @ jnp.array([1.0, 1.0, -1.0]))
1442 return prelkv.log_sqrt_term + exp_term
1445@named_call
1446def compute_likelihood_ratio(
1447 resid_lrt: (Float32[Array, ' 3'] | Float32[Array, 'k 3'] | Float32[Array, 'k k 3']),
1448 prelkv: PreLkV,
1449 error_cov_inv: Float32[Array, 'k k'] | None,
1450) -> Float32[Array, '']:
1451 """
1452 Compute the likelihood ratio of a grow move.
1454 Parameters
1455 ----------
1456 resid_lrt
1457 The sum of the residuals (scaled by error precision scale) of the
1458 datapoints falling in the left child, right child, and parent node
1459 involved in the move, stacked along the trailing axis.
1460 prelkv
1461 The pre-computed terms of the likelihood ratio, see
1462 `precompute_likelihood_terms`.
1463 error_cov_inv
1464 The global error precision scale, with the squared `inv_sdev_unit` folded
1465 in (see `bartz.mcmcstep.State.inv_sdev_unit`). Set only in the
1466 multivariate case with per-component error scales.
1468 Returns
1469 -------
1470 The log-likelihood ratio log P(data | new tree) - log P(data | old tree).
1471 """
1472 if error_cov_inv is not None:
1473 return _compute_likelihood_ratio_mv_het(resid_lrt, error_cov_inv, prelkv)
1474 elif resid_lrt.ndim > 1:
1475 return _compute_likelihood_ratio_mv(resid_lrt, prelkv)
1476 else:
1477 return _compute_likelihood_ratio_uv(resid_lrt, prelkv)
1480@named_call
1481def accept_moves_final_stage(state: State, moves: Moves) -> State:
1482 """
1483 Post-process the mcmc state after accepting/rejecting the moves.
1485 This function is separate from `accept_moves_sequential_stage` to signal it
1486 can work in parallel across trees.
1488 The prunes are not applied to `bartz.mcmcstep.Forest.leaf_indices` here;
1489 they are recorded in `to_prune`/`move_node` and applied at the beginning of
1490 the next step.
1492 Parameters
1493 ----------
1494 state
1495 A partially updated BART mcmc state.
1496 moves
1497 The proposed moves (see `propose_moves`) as updated by
1498 `accept_moves_sequential_stage`.
1500 Returns
1501 -------
1502 The fully updated BART mcmc state.
1503 """
1504 assert moves.acc is not None
1505 assert moves.to_prune is not None
1506 return replace(
1507 state,
1508 forest=replace(
1509 state.forest,
1510 grow_acc_count=jnp.sum(moves.acc & moves.grow),
1511 prune_acc_count=jnp.sum(moves.acc & ~moves.grow),
1512 to_prune=moves.to_prune,
1513 move_node=moves.lrt_nodes[..., 2],
1514 split_tree=apply_moves_to_split_trees(state.forest.split_tree, moves),
1515 affluence_tree=apply_moves_to_affluence_trees(
1516 state.forest.affluence_tree, moves
1517 ),
1518 ),
1519 )
1522@named_call
1523def apply_moves_to_leaf_indices(
1524 leaf_indices: UInt[Array, 'num_trees n'],
1525 to_prune: Bool[Array, ' num_trees'],
1526 move_node: Int32[Array, ' num_trees'],
1527) -> UInt[Array, 'num_trees n']:
1528 """
1529 Apply the prunes pending from the previous step to the leaf indices.
1531 Parameters
1532 ----------
1533 leaf_indices
1534 The index of the leaf each datapoint falls into, in the largest version
1535 of each tree compatible with the last moves (see
1536 `Forest.leaf_indices`).
1537 to_prune
1538 Whether the last move on each tree ended in a prune yet to be applied.
1539 move_node
1540 The node the last move on each tree operated on.
1542 Returns
1543 -------
1544 The updated leaf indices.
1545 """
1546 return _apply_moves_to_leaf_indices(leaf_indices, to_prune, move_node)
1549@vmap_nodoc
1550def _apply_moves_to_leaf_indices(
1551 leaf_indices: UInt[Array, ' n'],
1552 to_prune: Bool[Array, ''],
1553 move_node: Int32[Array, ''],
1554) -> UInt[Array, ' n']:
1555 """Implement `apply_moves_to_leaf_indices`."""
1556 mask = ~jnp.array(1, leaf_indices.dtype) # ...1111111110
1557 is_child = (leaf_indices & mask) == (move_node << 1)
1558 return jnp.where(
1559 is_child & to_prune, move_node.astype(leaf_indices.dtype), leaf_indices
1560 )
1563@named_call
1564def apply_moves_to_split_trees(
1565 split_tree: UInt[Array, 'num_trees half_tree_size'], moves: Moves
1566) -> UInt[Array, 'num_trees half_tree_size']:
1567 """
1568 Update the split trees to match the accepted move.
1570 Parameters
1571 ----------
1572 split_tree
1573 The cutpoints of the decision nodes in the initial trees.
1574 moves
1575 The proposed moves (see `propose_moves`), as updated by
1576 `accept_moves_sequential_stage`.
1578 Returns
1579 -------
1580 The updated split trees.
1581 """
1582 return _apply_moves_to_split_trees(split_tree, moves)
1585@vmap_nodoc
1586def _apply_moves_to_split_trees(
1587 split_tree: UInt[Array, ' half_tree_size'], moves: Moves
1588) -> UInt[Array, ' half_tree_size']:
1589 """Implement `apply_moves_to_split_trees`."""
1590 assert moves.to_prune is not None
1591 # a single scatter serves both cases: an accepted grow writes the new
1592 # cutpoint, while pruning (accepted prune or rejected grow) zeroes the node
1593 return split_tree.at[
1594 jnp.where(moves.grow | moves.to_prune, moves.lrt_nodes[2], split_tree.size)
1595 ].set(jnp.where(moves.to_prune, 0, moves.grow_split).astype(split_tree.dtype))
1598@named_call
1599def apply_moves_to_affluence_trees(
1600 affluence_tree: Bool[Array, 'num_trees half_tree_size'], moves: Moves
1601) -> Bool[Array, 'num_trees half_tree_size']:
1602 """
1603 Update the affluence trees to match the accepted move.
1605 The affluence tree marks the growable leaves; this restores that invariant
1606 after the move by re-marking only the nodes it touched, starting from the
1607 clean pre-move mask.
1609 Parameters
1610 ----------
1611 affluence_tree
1612 The mask of the growable leaves in the initial trees.
1613 moves
1614 The proposed moves (see `propose_moves`), as updated by
1615 `accept_moves_sequential_stage`.
1617 Returns
1618 -------
1619 The updated affluence trees.
1620 """
1621 return _apply_moves_to_affluence_trees(affluence_tree, moves)
1624@vmap_nodoc
1625def _apply_moves_to_affluence_trees(
1626 affluence_tree: Bool[Array, ' half_tree_size'], moves: Moves
1627) -> Bool[Array, ' half_tree_size']:
1628 """Implement `apply_moves_to_affluence_trees`."""
1629 assert moves.to_prune is not None
1630 assert moves.lrt_affluent is not None
1631 # GROW: node becomes internal, children become leaves with their affluence.
1632 # PRUNE (accepted prune or rejected grow): node becomes a leaf with its
1633 # affluence, children are deleted. Either way all three nodes are written:
1634 # the mask keeps the affluence of the nodes that become leaves and zeroes
1635 # the rest. If no move is applied (a rejected prune), the indices resolve
1636 # to `size` and the writes drop.
1637 becomes_leaf = moves.to_prune ^ jnp.array([True, True, False])
1638 return affluence_tree.at[
1639 jnp.where(moves.grow | moves.to_prune, moves.lrt_nodes, affluence_tree.size)
1640 ].set(moves.lrt_affluent & becomes_leaf)
1643@jit
1644def _sample_wishart_bartlett(
1645 key: Key[Array, ''],
1646 df: Float32[Array, ''] | float,
1647 scale_inv: Float32[Array, 'k k'],
1648) -> Float32[Array, 'k k']:
1649 """
1650 Sample a precision matrix W ~ Wishart(df, scale_inv^-1) using Bartlett decomposition.
1652 Parameters
1653 ----------
1654 key
1655 A JAX random key
1656 df
1657 Degrees of freedom
1658 scale_inv
1659 Scale matrix of the corresponding Inverse Wishart distribution
1661 Returns
1662 -------
1663 A sample from Wishart(df, scale)
1664 """
1665 keys = split(key)
1667 # Diagonal elements: A_ii ~ sqrt(chi^2(df - i)), with chi^2(k) = Gamma(k/2, scale=2).
1668 # sqrt(2 * Gamma) = sqrt(2) * exp(loggamma / 2), folding the sqrt into the exp.
1669 k, _ = scale_inv.shape
1670 df_vector = df - jnp.arange(k)
1671 diag_A = jnp.sqrt(2.0) * jnp.exp(loggamma(keys.pop(), df_vector / 2.0) / 2.0)
1673 off_diag_A = random.normal(keys.pop(), (k, k))
1674 A = jnp.tril(off_diag_A, -1) + jnp.diag(diag_A)
1675 L = chol_with_gersh(scale_inv, absolute_eps=True)
1676 T = solve_triangular(L, A, lower=True, trans='T')
1678 return T @ T.T
1681def step_resid_eff_scale(
1682 state: State,
1683 norm2: Float32[Array, ''] | Float32[Array, ' k'],
1684 prec_sum: Float32[Array, ''] | Float32[Array, ' k'],
1685) -> State:
1686 """Update `State.resid_eff_scale` with the measured scale of the residuals.
1688 Parameters
1689 ----------
1690 state
1691 A BART MCMC state.
1692 norm2
1693 The squared (precision-scaled, masked) norm of the residuals in data units.
1694 prec_sum
1695 The sum of the precision scales (unmasked-datapoint count without
1696 error scales), matching the scaling of `norm2`.
1698 Returns
1699 -------
1700 The state with `resid_eff_scale` set to the residual rms rounded to a power of two.
1701 """
1702 scale = round_to_pow2(jnp.sqrt(norm2 / prec_sum))
1703 # keep the previous scale if the residuals vanish
1704 scale = jnp.where(scale == 0, state.resid_eff_scale, scale)
1705 return replace(state, resid_eff_scale=scale)
1708def _step_error_cov_inv_mv(key: Key[Array, ''], state: State) -> State:
1709 assert state.error_cov_inv.nu is not None
1710 assert state.error_cov_inv.rate is not None
1712 # keep the residuals in their stored (narrow) dtype and resid_unit units;
1713 # the reduction accumulates in float32 and its (k, k) result is rescaled to
1714 # data units, so no n-sized float32 array is ever materialized
1715 resid = state.resid
1716 if state.inv_sdev_scale is not None:
1717 # 2-D inv_sdev_scale dispatches to the diagonal path, so here it is 1-D
1718 resid *= state.inv_sdev_scale
1719 df_post = state.error_cov_inv.nu + state.n_non_missing
1720 # unit of the stored precision-scaled residuals: `resid` is in `resid_unit`
1721 # units and `inv_sdev_scale` in `inv_sdev_unit` units (a scalar here,
1722 # matching the 1-D `inv_sdev_scale`; 1 without error scales)
1723 scale = state.resid_unit * state.inv_sdev_unit
1724 rrt = jnp.einsum(
1725 'an,bn->ab', resid, resid, preferred_element_type=jnp.float32
1726 ) * jnp.outer(scale, scale)
1727 if state.config.data_sharded:
1728 rrt = lax.psum(rrt, 'data')
1729 scale_post = state.error_cov_inv.rate + rrt
1731 prec = _sample_wishart_bartlett(key, df_post, scale_post)
1732 state = step_resid_eff_scale(state, jnp.diagonal(rrt), state.sum_diag_prec_scale)
1733 return replace(state, error_cov_inv=replace(state.error_cov_inv, value=prec))
1736def _step_error_cov_inv_diag(key: Key[Array, ''], state: State) -> State:
1737 """Per-component inverse-gamma update for univariate, mixed, and partial-missing paths."""
1738 assert state.error_cov_inv.rate is not None
1739 assert state.error_cov_inv.nu is not None
1741 # keep the residuals in their stored (narrow) dtype and resid_unit units;
1742 # the reduction accumulates in float32 and its small result is rescaled to
1743 # data units, so no n-sized float32 array is ever materialized
1744 resid = state.resid
1745 if state.inv_sdev_scale is not None:
1746 resid *= state.inv_sdev_scale
1748 # alpha
1749 alpha = state.error_cov_inv.nu / 2 + state.n_non_missing / 2
1751 # beta; `resid` is stored in `resid_unit` units and `inv_sdev_scale` in
1752 # `inv_sdev_unit` units (1 without error scales)
1753 norm2 = jnp.einsum(
1754 '...n,...n->...', resid, resid, preferred_element_type=jnp.float32
1755 ) * jnp.square(state.resid_unit * state.inv_sdev_unit)
1756 if state.config.data_sharded:
1757 norm2 = lax.psum(norm2, 'data')
1758 scale = state.error_cov_inv.rate
1759 kshape = resid.shape[:-1]
1760 if kshape:
1761 scale = jnp.diag(scale)
1762 beta = scale / 2 + norm2 / 2
1764 # draw the gamma from the first of a split, mirroring the Bartlett sampler
1765 # in the multivariate path so the two branches coincide at k=1
1766 keys = split(key)
1767 samples = jnp.exp(loggamma(keys.pop(), alpha, kshape))
1768 prec = samples / beta
1769 if state.binary_indices is not None:
1770 prec = prec.at[state.binary_indices].set(1.0)
1771 if kshape:
1772 prec = jnp.diag(prec)
1773 state = step_resid_eff_scale(state, norm2, state.sum_diag_prec_scale)
1774 return replace(state, error_cov_inv=replace(state.error_cov_inv, value=prec))
1777@named_call
1778def step_error_cov_inv(key: Key[Array, ''], state: State) -> State:
1779 """MCMC-update the inverse error covariance."""
1780 if (
1781 state.error_cov_inv.value.ndim == 2
1782 and state.binary_indices is None
1783 and (state.inv_sdev_scale is None or state.inv_sdev_scale.ndim == 1)
1784 ):
1785 return _step_error_cov_inv_mv(key, state)
1786 else:
1787 return _step_error_cov_inv_diag(key, state)
1790@named_call
1791def step_z(key: Key[Array, ''], state: State) -> State:
1792 """
1793 MCMC-update the latent variable for binary regression.
1795 Parameters
1796 ----------
1797 key
1798 A jax random key.
1799 state
1800 A BART MCMC state.
1802 Returns
1803 -------
1804 The updated BART MCMC state.
1805 """
1806 assert state.z is not None
1808 inv_sdev = state.inv_sdev_scale
1809 inv_sdev_unit = state.inv_sdev_unit
1810 err_scale = state.error_scale
1811 if state.binary_indices is not None:
1812 resid_unit = state.resid_unit[state.binary_indices]
1813 resid = state.resid[state.binary_indices, :]
1814 binary_y = state.y[state.binary_indices, :] != 0
1815 # per-component scales (2-D) are restricted to the binary components;
1816 # a scalar-per-datapoint scale (1-D) is shared and broadcasts as-is
1817 if inv_sdev is not None and inv_sdev.ndim > 1:
1818 inv_sdev = inv_sdev[state.binary_indices, :]
1819 inv_sdev_unit = inv_sdev_unit[state.binary_indices]
1820 if err_scale is not None and err_scale.ndim > 1:
1821 err_scale = err_scale[state.binary_indices, :]
1822 else:
1823 resid_unit = state.resid_unit
1824 resid = state.resid
1825 binary_y = state.y != 0
1827 trees_plus_offset = state.z - resid * resid_unit[..., None]
1828 if state.config.data_sharded:
1829 # decorrelate the seed across data shards; the seed is replicated
1830 # because the trees and most of the algorithm are replicated
1831 key = random.fold_in(key, lax.axis_index('data'))
1833 if err_scale is None:
1834 # homoskedastic probit: the latent error has unit scale. A missingness
1835 # mask (without error scales) needs no handling here: masked points are
1836 # dropped from the likelihood, so their latent is sampled like any other.
1837 resid = truncated_normal_onesided(key, (), ~binary_y, -trees_plus_offset)
1838 else:
1839 # heteroskedastic probit: the latent error has per-datapoint scale
1840 # `err_scale`, so threshold in standardized units, then rescale.
1841 assert inv_sdev is not None
1842 # masked datapoints draw garbage, but every consumer of `resid` scales
1843 # it by the (zero) `inv_sdev`, so the value is irrelevant as long as finite
1844 resid = err_scale * truncated_normal_onesided(
1845 key, (), ~binary_y, -trees_plus_offset * inv_sdev * inv_sdev_unit[..., None]
1846 )
1847 z = trees_plus_offset + resid
1849 resid = (resid / resid_unit[..., None]).astype(state.resid.dtype)
1850 if state.binary_indices is not None:
1851 resid = state.resid.at[state.binary_indices, :].set(resid)
1853 return replace(state, z=z, resid=resid)
1856def _blocked_mass_tree(
1857 key: Key[Array, ''],
1858 var_tree: UInt[Array, ' half_tree_size'],
1859 split_tree: UInt[Array, ' half_tree_size'],
1860 max_split: UInt[Array, ' p'],
1861 s: Float32[Array, ' p'],
1862) -> Float32[Array, ' p']:
1863 """Per-variable data-augmentation mass blocked by a single tree.
1865 At each internal node, draws the latent augmentation weight ``lambda / e``
1866 (``lambda`` exponential, ``e`` the eligible split probability mass at the
1867 node) and adds it to every variable ineligible at that node.
1869 Parameters
1870 ----------
1871 key
1872 Random key for sampling.
1873 var_tree
1874 The splitting axes of the tree.
1875 split_tree
1876 The splitting points of the tree.
1877 max_split
1878 The maximum split index for each variable.
1879 s
1880 Split probabilities normalized over selectable variables.
1882 Returns
1883 -------
1884 The blocked mass for each variable.
1885 """
1886 (half_tree_size,) = split_tree.shape
1887 d_minus_1 = half_tree_size.bit_length() - 1 # number of decision-node levels
1888 p = max_split.size
1889 nodes = jnp.arange(half_tree_size)
1890 split = split_tree.astype(jnp.int32)
1891 is_internal = split_tree.astype(bool)
1893 # Range [lo, hi) of cutpoints still available for each node's own splitting
1894 # variable, given the constraints inherited from the ancestors.
1895 lo, hi = vmap(split_range, in_axes=(None, None, None, 0, 0))(
1896 var_tree, split_tree, max_split, nodes, var_tree
1897 )
1899 # An internal node exhausts its own variable for a child when its cutpoint
1900 # sits at the matching end of the available range, so the variable becomes
1901 # ineligible throughout that child's subtree. Row 0 is the left child (low
1902 # end lo), row 1 the right child (high end hi - 1).
1903 blocks = is_internal & (split == jnp.stack([lo, hi - 1]))
1905 # A node can block at most its own splitting variable, so the per-variable
1906 # totals are recovered from these per-node blocks via top-down/bottom-up
1907 # accumulation over depth levels, rather than scanning each node's ancestors.
1909 # Ineligible mass per node: the s-mass of the variables blocked along the
1910 # path from the root. Each variable is blocked at exactly one node per path,
1911 # so summing the per-node increments top-down reproduces the per-node sum
1912 # over distinct ineligible variables.
1913 parent = nodes >> 1
1914 side = nodes & 1 # 0 if the node is a left child, 1 if a right child
1915 parent_blocks = blocks[side, parent]
1916 # var_tree[parent] is a valid index wherever parent_blocks holds (the parent
1917 # is then internal); elsewhere the clamped gather is masked away
1918 ineligible_mass = jnp.where(parent_blocks, s[var_tree[parent]], 0.0)
1919 for level in range(1, d_minus_1):
1920 lhs, rhs = 1 << level, 1 << (level + 1)
1921 parent_mass = jnp.repeat(ineligible_mass[lhs >> 1 : rhs >> 1], 2)
1922 ineligible_mass = ineligible_mass.at[lhs:rhs].add(parent_mass)
1924 # Per-node augmentation weight lambda_b / e_b, zero at non-internal nodes. The
1925 # eligible mass is positive at internal nodes (the split variable is eligible);
1926 # the floor only guards against round-off, and is unused where weight is zero.
1927 eligible_mass = jnp.maximum(1.0 - ineligible_mass, jnp.finfo(jnp.float32).eps)
1928 weight = jnp.where(is_internal, random.exponential(key, (half_tree_size,)), 0.0)
1929 weight /= eligible_mass
1931 # Subtree weight: total weight of each node's internal descendants and itself,
1932 # accumulated bottom-up. The children of the deepest decision level are leaves
1933 # and contribute nothing.
1934 subtree_weight = weight
1935 for level in range(d_minus_1 - 2, -1, -1):
1936 lhs, rhs = 1 << level, 1 << (level + 1)
1937 children = subtree_weight[2 * lhs : 2 * rhs].reshape(-1, 2).sum(axis=1)
1938 subtree_weight = subtree_weight.at[lhs:rhs].add(children)
1940 # A variable blocked by node b at one of its children is ineligible in that
1941 # whole child subtree, so it accumulates the subtree weight; scatter it onto
1942 # the splitting variable. Only the upper half of nodes have internal children
1943 # (the deepest decision nodes block into leaves, contributing nothing); their
1944 # children are exactly subtree_weight reshaped into [left, right] pairs.
1945 half = half_tree_size // 2
1946 contrib = (blocks[:, :half] * subtree_weight.reshape(-1, 2).T).sum(axis=0)
1947 scatter_var = jnp.where(is_internal, var_tree, p)[:half]
1948 return jnp.zeros(p).at[scatter_var].add(contrib)
1951def sample_s_augmentation(key: Key[Array, ''], forest: Forest) -> Int32[Array, ' p']:
1952 """Sample the data-augmentation counts for the exact full conditional of `s`.
1954 At each internal node, the variables with no available cutpoint given the
1955 ancestors (plus the globally blocked ones) cannot be split on, so the plain
1956 Dirichlet update for `s` is only approximate. This samples, for each
1957 variable, the number of ineligible draws discarded before each realized
1958 split, to be added to the variable usage counts.
1960 Parameters
1961 ----------
1962 key
1963 Random key for sampling.
1964 forest
1965 The forest, providing the trees and the current `log_s`.
1967 Returns
1968 -------
1969 The discarded-draws count for each variable.
1970 """
1971 assert forest.log_s is not None
1972 keys = split(key)
1973 (num_trees, _) = forest.var_tree.shape
1975 # split probabilities normalized over the selectable (non-blocked) variables
1976 selectable = forest.max_split > 0
1977 s = softmax(forest.log_s, where=selectable)
1979 # blocked_mass[j] = sum over internal nodes where j is ineligible of
1980 # lambda_b / e_b, with lambda_b ~ Exponential(1)
1981 blocked_mass = vmap(_blocked_mass_tree, in_axes=(0, 0, 0, None, None))(
1982 keys.pop(num_trees), forest.var_tree, forest.split_tree, forest.max_split, s
1983 ).sum(axis=0) # shape (p,)
1985 # the per-node discarded-draw counts are negative-multinomial, with no
1986 # closed form when summed over nodes, but their Gamma-Poisson mixture does:
1987 # A_j | {lambda_b} ~ Poisson(s_j * blocked_mass[j]), independent across j
1988 return poisson(keys.pop(), s * blocked_mass, dtype=jnp.int32)
1991@named_call
1992def step_s(key: Key[Array, ''], state: State) -> State:
1993 """
1994 Update `log_s` using Dirichlet sampling.
1996 The prior is s ~ Dirichlet(theta/p, ..., theta/p), and the posterior
1997 is s ~ Dirichlet(theta/p + varcount, ..., theta/p + varcount), where
1998 varcount is the count of how many times each variable is used in the
1999 current forest.
2001 Parameters
2002 ----------
2003 key
2004 Random key for sampling.
2005 state
2006 The current BART state.
2008 Returns
2009 -------
2010 Updated BART state with re-sampled `log_s`.
2012 Notes
2013 -----
2014 By default this full conditional is approximate, because it ignores the
2015 decision rules forbidden by the ancestors of each node. If
2016 ``state.config.augment`` is set, the forbidden rules are accounted for
2017 exactly with the data augmentation of `sample_s_augmentation`.
2018 """
2019 assert state.forest.theta is not None
2021 # reserve the Dirichlet draw key first and unconditionally, so it does not
2022 # depend on whether augmentation is on; then the two modes draw identically
2023 # when there are no forbidden rules, since the augmentation is exactly zero
2024 keys = split(key)
2025 log_s_key = keys.pop()
2027 # histogram current variable usage
2028 p = state.forest.max_split.size
2029 varcount = var_histogram(
2030 p, state.forest.var_tree, state.forest.split_tree, sum_batch_axis=-1
2031 )
2033 # the Dirichlet posterior concentration, optionally completed with the exact
2034 # accounting of forbidden rules via data augmentation
2035 alpha = state.forest.theta / p + varcount
2036 if state.config.augment:
2037 alpha = alpha + sample_s_augmentation(keys.pop(), state.forest)
2039 # sample from the Dirichlet posterior and update the forest with the new s
2040 log_s = loggamma(log_s_key, alpha)
2041 return replace(state, forest=replace(state.forest, log_s=log_s))
2044@named_call
2045def step_theta(key: Key[Array, ''], state: State, *, num_grid: int = 1000) -> State:
2046 """
2047 Update `theta`.
2049 The prior is theta / (theta + rho) ~ Beta(a, b).
2051 Parameters
2052 ----------
2053 key
2054 Random key for sampling.
2055 state
2056 The current BART state.
2057 num_grid
2058 The number of points in the evenly-spaced grid used to sample
2059 theta / (theta + rho).
2061 Returns
2062 -------
2063 Updated BART state with re-sampled `theta`.
2064 """
2065 assert state.forest.log_s is not None
2066 assert state.forest.rho is not None
2067 assert state.forest.a is not None
2068 assert state.forest.b is not None
2070 # the grid points are the midpoints of num_grid bins in (0, 1)
2071 padding = 1 / (2 * num_grid)
2072 lambda_grid = jnp.linspace(padding, 1 - padding, num_grid)
2074 # normalize s
2075 log_s = state.forest.log_s - logsumexp(state.forest.log_s)
2077 # sample lambda
2078 logp, theta_grid = _log_p_lambda(
2079 lambda_grid, log_s, state.forest.rho, state.forest.a, state.forest.b
2080 )
2081 i = random.categorical(key, logp)
2082 theta = theta_grid[i]
2084 return replace(state, forest=replace(state.forest, theta=theta))
2087def _log_p_lambda(
2088 lambda_: Float32[Array, ' num_grid'],
2089 log_s: Float32[Array, ' p'],
2090 rho: Float32[Array, ''],
2091 a: Float32[Array, ''],
2092 b: Float32[Array, ''],
2093) -> tuple[Float32[Array, ' num_grid'], Float32[Array, ' num_grid']]:
2094 # in the following I use lambda_[::-1] == 1 - lambda_
2095 theta = rho * lambda_ / lambda_[::-1]
2096 p = log_s.size
2097 return (
2098 (a - 1) * jnp.log1p(-lambda_[::-1]) # log(lambda)
2099 + (b - 1) * jnp.log1p(-lambda_) # log(1 - lambda)
2100 + gammaln(theta)
2101 - p * gammaln(theta / p)
2102 + theta / p * jnp.sum(log_s)
2103 ), theta
2106@named_call
2107def step_sparse(key: Key[Array, ''], state: State) -> State:
2108 """
2109 Update the sparsity parameters.
2111 This invokes `step_s`, and then `step_theta` only if the parameters of
2112 the theta prior are defined.
2114 Parameters
2115 ----------
2116 key
2117 Random key for sampling.
2118 state
2119 The current BART state.
2121 Returns
2122 -------
2123 Updated BART state with re-sampled `log_s` and `theta`.
2124 """
2125 if state.config.sparse_on_at is not None:
2126 state = lax.cond(
2127 state.config.steps_done < state.config.sparse_on_at,
2128 lambda _key, state: state,
2129 _step_sparse,
2130 key,
2131 state,
2132 )
2133 return state
2136def _step_sparse(key: Key[Array, ''], state: State) -> State:
2137 keys = split(key)
2138 state = step_s(keys.pop(), state)
2139 if state.forest.rho is not None:
2140 state = step_theta(keys.pop(), state)
2141 return state
2144@named_call
2145def step_config(state: State) -> State:
2146 config = state.config
2147 config = replace(config, steps_done=config.steps_done + 1)
2148 return replace(state, config=config)