bartz.mcmcstep.Forest

class bartz.mcmcstep.Forest(var_tree, split_tree, affluence_tree, leaf_tree, leaf_unit, offset, grow_prop_count, prune_prop_count, grow_acc_count, prune_acc_count, max_split, blocked_vars, p_nonterminal, p_propose_grow, leaf_indices, to_prune, move_node, count_tree, prec_tree, min_points_per_decision_node, min_points_per_leaf, log_trans_prior, log_likelihood, leaf_prior_cov_inv, log_s, theta, a, b, rho)

Represents the MCMC state of a sum of trees.

var_tree: UInt[Array, '*chains num_trees half_tree_size']

Variables/predictors/axes of decision rules.

split_tree: UInt[Array, '*chains num_trees half_tree_size']

Cutpoints/boundaries of decision rules.

affluence_tree: Bool[Array, '*chains num_trees half_tree_size']

Marks leaves that can be grown.

leaf_tree: Float[Array, '*chains num_trees 2*half_tree_size'] | Float[Array, '*chains num_trees k 2*half_tree_size']

The leaf values, in units of leaf_unit. The function computed by the forest is offset + leaf_unit * (sum of leaves).

leaf_unit: Float32[Array, ''] | Float32[Array, 'k']

The storage unit of the leaves. Keeps the stored values O(1) whatever the data units, so they do not under/overflow narrow leaf_tree dtypes. Set to the marginal prior standard deviation of a leaf, rounded to a power of two so converting to and from data units is exact.

offset: Float32[Array, ''] | Float32[Array, 'k']

Constant shift added to the scaled sum of trees, see leaf_tree.

grow_prop_count: Int32[Array, '*chains']

The number of grow proposals made during the last MCMC step.

prune_prop_count: Int32[Array, '*chains']

The number of prune proposals made during the last MCMC step.

grow_acc_count: Int32[Array, '*chains']

The number of grow moves accepted during the last MCMC step.

prune_acc_count: Int32[Array, '*chains']

The number of prune moves accepted during the last MCMC step.

max_split: UInt[Array, 'p']

The maximum split index for each predictor.

blocked_vars: UInt[Array, 'q'] | None

Indices of not to be used variables/predictors. This shall include at least all the i that yield max_split[i] == 0, otherwise behavior is undefined.

p_nonterminal: Float32[Array, '2*half_tree_size']

The prior probability of each node being nonterminal (conditional on its ancestors leaving at least one available decision rule). Includes the nodes at maximum depth which shall be set to 0.

p_propose_grow: Float32[Array, 'half_tree_size']

The unnormalized probability of picking a leaf for a grow proposal.

leaf_indices: UInt[Array, '*chains num_trees n']

The index of the leaf each datapoint falls into, for each tree, in the largest version of the tree compatible with the last moves.

A pending prune (accepted prune or rejected grow, marked per-tree by to_prune) is not yet applied to the indices; step folds it in at the beginning of the next iteration. Evaluating the trees at these indices is correct anyway because leaf_tree mirrors the value of a pruned node onto its dangling children.

to_prune: Bool[Array, '*chains num_trees']

Whether the last move on each tree ended in a prune (accepted prune or rejected grow) whose application to leaf_indices is still pending.

move_node: Int32[Array, '*chains num_trees']

The node the last move on each tree operated on (the leaf to grow or the node to prune). Meaningful only where to_prune is set.

count_tree: UInt32[Array, '*chains num_trees 2*half_tree_size'] | None

The number of datapoints per leaf. Valid at the leaves and at the nodes involved in the latest moves, dirty elsewhere. None if there are per-datapoint error scales and no minimum-points-per-node constraints, which makes the counts unused.

prec_tree: Float32[Array, '*chains num_trees 2*half_tree_size'] | Float32[Array, '*chains num_trees k k 2*half_tree_size'] | None

The sum of State.prec_scale over the datapoints in each leaf, in the same units; valid/dirty like count_tree. None if there are no per-datapoint error scales, in which case count_tree takes its place.

min_points_per_decision_node: Int32[Array, ''] | None

The minimum number of data points in a decision node.

min_points_per_leaf: Int32[Array, ''] | None

The minimum number of data points in a leaf node.

log_trans_prior: Float32[Array, '*chains num_trees'] | None

The log transition and prior Metropolis-Hastings ratio for the proposed move on each tree.

log_likelihood: Float32[Array, '*chains num_trees'] | None

The log likelihood ratio.

leaf_prior_cov_inv: Float32[Array, ''] | Float32[Array, 'k k'] | None

The prior precision matrix of a leaf, conditional on the tree structure (a scalar inverse variance for univariate). The prior mean of a leaf is zero; the prior covariance of the sum of trees is num_trees * leaf_prior_cov_inv^-1.

log_s: Float32[Array, '*chains p'] | None

The logarithm of the prior probability for choosing a variable to split along in a decision rule. Not normalized. Variables that do not have any available decision rule at a given node are masked away. If None, use a uniform distribution.

theta: Float32[Array, '*chains'] | None

The concentration parameter for the Dirichlet prior on the variable distribution s. If not set, log_s is left constant.

a: Float32[Array, ''] | None

Parameter of the prior on theta. If not set, theta is left constant.

b: Float32[Array, ''] | None

Parameter of the prior on theta. If not set, theta is left constant.

rho: Float32[Array, ''] | None

Parameter of the prior on theta. If not set, theta is left constant.

property has_chains: bool

Whether non-constant attributes in the forest carry an explicit chain axis.