bartz.mcmcloop.MainTrace

class bartz.mcmcloop.MainTrace(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, varprob)

MCMC trace with trees and diagnostic values.

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.

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

The leaf values.

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 sum of 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.

classmethod from_state(state)

Create a single-item main trace from a MCMC state.

Return type:

MainTrace