bartz.mcmcloop.MainTraceWithTrainPred

class bartz.mcmcloop.MainTraceWithTrainPred(has_chains, mesh, grow_prop_count, grow_acc_count, prune_prop_count, prune_acc_count, error_cov_inv, theta, log_likelihood, log_trans_prior, leaf_tree, var_tree, split_tree, offset, leaf_unit, varprob, train_pred)

Main trace that also stores the latent predictions at the training points.

Computed cheaply from the running State.resid (as ref - resid * resid_unit, with ref the response y or the binary latent z), so bartz.Bart.predict can return them without re-evaluating the trees.

leaf_tree: Float[Array, '*chains_and_samples num_trees tree_size'] | Float[Array, '*chains_and_samples num_trees k tree_size']

The leaf values, in units of leaf_unit.

var_tree: UInt[Array, '*chains_and_samples num_trees tree_size//2']

The decision axes.

split_tree: UInt[Array, '*chains_and_samples num_trees tree_size//2']

The decision boundaries.

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

Constant shift added to the scaled sum of trees.

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

The storage unit of the leaf values. Predictions are offset + leaf_unit * (sum of leaf values over trees).

varprob: Float32[Array, '*chains_and_samples p'] | None

The probability of choosing each variable for a decision rule, normalized over variables, or None when variable selection is off.

has_chains: bool

Whether the trace carries an explicit chain axis.

mesh: Mesh | None

The device mesh the trace arrays are sharded on, or None.

grow_prop_count: Int32[Array, '*chains_and_samples']

The number of grow proposals made during one full MCMC cycle.

grow_acc_count: Int32[Array, '*chains_and_samples']

The number of grow moves accepted during one full MCMC cycle.

prune_prop_count: Int32[Array, '*chains_and_samples']

The number of prune proposals made during one full MCMC cycle.

prune_acc_count: Int32[Array, '*chains_and_samples']

The number of prune moves accepted during one full MCMC cycle.

error_cov_inv: Float32[Array, '*chains_and_samples'] | Float32[Array, '*chains_and_samples k k']

The inverse error covariance (scalar for univariate, matrix for multivariate). Identity in binary regression.

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

The concentration parameter of the Dirichlet prior on the variable split probabilities, or None if it was not sampled.

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

The log likelihood ratio of the proposed move on each tree, or None.

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

The log transition and prior Metropolis-Hastings ratio of the proposed move on each tree, or None.

train_pred: Float32[Array, '*chains_and_samples n'] | Float32[Array, '*chains_and_samples k n']

The latent predictions at the training points, offset plus the scaled sum of trees.

classmethod from_state(state)

Create a single-item main trace with train predictions from a MCMC state.

Return type:

MainTraceWithTrainPred