bartz.mcmcstep.State

class bartz.mcmcstep.State(_chain_anchor, X, y, z, binary_indices, resid, resid_unit, resid_eff_scale, resid_inexact_integral, error_cov_inv, error_scale, prec_scale, inv_sdev_scale, inv_sdev_unit, n_non_missing, sum_diag_prec_scale, forest, config)

Represents the MCMC state of BART.

X: UInt[Array, 'p n']

The predictors.

y: Float32[Array, 'n'] | Float32[Array, 'k n']

The response, in data units. Binary components are stored as 0/1. Missing values are replaced by Forest.offset.

z: None | Float32[Array, '*chains n'] | Float32[Array, '*chains kb n']

The latent outcomes for binary regression. None in continuous regression. In the mixed binary-continuous case, only the binary outcome components are stored.

binary_indices: None | Int32[Array, 'kb']

The indices of binary outcome components in the full list of outcome components. None when there are no binary components.

resid: Float[Array, '*chains n'] | Float[Array, '*chains k n']

The residuals, resid_unit * resid = (y or z) - sum of trees.

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

The storage unit of resid (see init’s resid_dtype), same scheme as Forest.leaf_unit. Equal to leaf_unit * sqrt(num_trees) (the marginal prior standard deviation of the sum of trees) rounded to a power of two.

resid_eff_scale: Float32[Array, '*chains'] | Float32[Array, '*chains k']

The measured scale of the residuals (their precision-weighted root mean square in data units), rounded to a power of two. Sets the leaf quantization grid (see StepConfig.leaf_quantization). Initialized to resid_unit, then tracks the MCMC (while the storage unit of resid stays fixed at resid_unit). Not updated in purely binary regression, where the scale of the latent residuals should be stable.

resid_inexact_integral: Float32[Array, '*chains'] | Float32[Array, '*chains k']

Sum over the MCMC steps done of the mean square of the residuals (in resid_unit units) large enough that their running updates round. Used by sum_trees_eps to estimate the accumulated rounding drift.

error_cov_inv: Wishart

The inverse error covariance (error_cov_inv.value, scalar for univariate) with its Wishart prior. Fixed at the identity with no prior in binary regression.

error_scale: Float32[Array, 'n'] | Float32[Array, 'k n'] | None

The per-datapoint error scales (the error_scale argument of init). The error precision on a datapoint is error_cov_inv.value / outer(error_scale, error_scale). For binary components the (fixed, unit) probit latent error is scaled instead, so the success probability is Phi(sum of trees / error_scale). inv_sdev_scale and prec_scale are derived from this and the missingness mask.

prec_scale: Float[Array, 'n'] | Float[Array, 'k k n'] | None

The scale on the error precision, prec_scale = outer(inv_sdev_scale, inv_sdev_scale) per datapoint (inv_sdev_scale ** 2 for scalar scales), so it’s in units of inv_sdev_unit ** 2. Stored, like inv_sdev_scale, in init’s prec_scale_dtype.

inv_sdev_scale: Float[Array, 'n'] | Float[Array, 'k n'] | None

inv_sdev_scale * inv_sdev_unit = 1 / error_scale, zeroed at missing datapoints. Not None when missing is set even if fit without error scales.

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

The storage unit of inv_sdev_scale, to avoid under/overflow with short dtypes; same scheme as Forest.leaf_unit. Set to the root mean square of 1 / error_scale over non-missing datapoints, rounded to a power of two; 1 if fit without error scales. Constant along the MCMC.

n_non_missing: Int32[Array, ''] | Int32[Array, 'k']

The number of non-missing datapoints.

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

sum(1 / error_scale ** 2) over non-missing datapoints; equal to n_non_missing if fit without error scales.

forest: Forest

The sum of trees model.

config: StepConfig

Metadata and configurations for the MCMC step.

property has_chains: bool

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

num_chains()

Return the number of chains, or None if not multichain.

Return type:

int | None

sum_trees_eps()

Estimate the absolute accuracy limit of the sum of trees (in data units).

The analogue of finfo(dtype).eps for the sum of trees. This combines three terms: floating point resolution, random walk accumulation of numerical error on running residuals, and if leaf quantization is active, breakdown of the mcmc due to the quantization being too coarse. The latter term is currently broken actually, sorry about that.

Return type:

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