bartz.stochtree.MeanForestParams

class bartz.stochtree.MeanForestParams(*, num_trees=200, alpha=0.95, beta=2.0, min_samples_leaf=5, max_depth=10, sample_sigma2_leaf=True, sigma2_leaf_init=None, sigma2_leaf_shape=3.0, sigma2_leaf_scale=None)

Mirror of stochtree’s mean_forest_params dict, restricted to the keys bartz handles.

num_trees: int = 200

Number of trees in the conditional mean ensemble.

alpha: float | Float[Array, ''] | Float[ndarray, ''] = 0.95

Tree split prior base.

beta: float | Float[Array, ''] | Float[ndarray, ''] = 2.0

Tree split prior decay.

min_samples_leaf: int = 5

Minimum number of training samples at a leaf.

max_depth: int = 10

Maximum tree depth. Must be a non-negative integer at most 16.

sample_sigma2_leaf: bool = True

Whether to sample the leaf variance with an inverse-gamma prior, or hold it fixed at sigma2_leaf_init.

sigma2_leaf_init: float | Float[Array, ''] | Float[ndarray, ''] | None = None

Initial leaf variance. If None, matches stochtree’s defaults of var(resid_train) / num_trees for continuous and 2 / num_trees for probit.

sigma2_leaf_shape: float | Float[Array, ''] | Float[ndarray, ''] = 3.0

Twice the shape parameter of the inverse-gamma prior on the leaf variance, i.e., the prior is IG(sigma2_leaf_shape / 2, sigma2_leaf_scale / 2).

sigma2_leaf_scale: float | Float[Array, ''] | Float[ndarray, ''] | None = None

Twice the scale parameter of the inverse-gamma prior on the leaf variance. If None, defaults to var(resid_train) / num_trees for continuous and 1 / num_trees for probit.