bartz.mcmcstep.init

bartz.mcmcstep.init(*, X, y, outcome_type='continuous', offset, max_split, num_trees, p_nonterminal, leaf_prior_cov_inv, leaf_dtype=<class 'jax.numpy.float32'>, prec_scale_dtype=<class 'jax.numpy.float32'>, resid_dtype=<class 'jax.numpy.float32'>, leaf_quantization=None, error_cov_inv=None, error_scale=None, missing=None, min_points_per_decision_node=None, min_points_per_leaf=None, resid_reduction_config=AutoBatchedReduction(), count_reduction_config=AutoOneHotReduction(), prec_reduction_config=AutoOneHotReduction(), prec_count_num_trees='auto', sequential_unroll=2, save_ratios=False, filter_splitless_vars=0, log_s=None, theta=None, a=None, b=None, rho=None, sparse_on_at=None, augment=True, num_chains=None, mesh=None)[source]

Make a BART posterior sampling MCMC initial state.

Parameters:

  • X (UInt[Array, 'p n'] | UInt[ndarray, 'p n']) – The predictors. Note this is trasposed compared to the usual convention.

  • y (Float32[Array, 'n'] | Float32[ndarray, 'n'] | Float32[Array, 'k n'] | Float32[ndarray, 'k n']) – The response. If two-dimensional, the outcome is multivariate with the first axis indicating the component. For binary data, non-zero means 1, zero means 0.

  • outcome_type (OutcomeType | str | Sequence[OutcomeType | str], default: 'continuous') – Whether the regression is continuous or binary (probit). Can also be a sequence of OutcomeType values, one per outcome component, for mixed binary-continuous multivariate regression.

  • offset (float | Float[Array, ''] | Float[ndarray, ''] | Float[Array, 'k'] | Float[ndarray, 'k']) – Constant shift added to the sum of trees. 0 if not specified.

  • max_split (UInt[Array, 'p'] | UInt[ndarray, 'p']) – The maximum split index for each variable. All split ranges start at 1.

  • num_trees (int) – The number of trees in the forest.

  • p_nonterminal (Float32[Array, 'd_minus_1'] | Float32[ndarray, 'd_minus_1']) – The probability of a nonterminal node at each depth. The maximum depth of trees is fixed by the length of this array. Use make_p_nonterminal to set it with the conventional formula.

  • leaf_prior_cov_inv (float | Float[Array, ''] | Float[ndarray, ''] | Float[Array, 'k k'] | Float[ndarray, 'k k']) – The prior precision matrix of a leaf, see Forest.leaf_prior_cov_inv.

  • leaf_dtype (str | type[Any] | dtype | SupportsDType, default: <class 'jax.numpy.float32'>)

  • prec_scale_dtype (str | type[Any] | dtype | SupportsDType, default: <class 'jax.numpy.float32'>)

  • resid_dtype (str | type[Any] | dtype | SupportsDType, default: <class 'jax.numpy.float32'>) – Storage dtypes for, respectively: leaves (Forest.leaf_tree), derived error scales (State.prec_scale and State.inv_sdev_scale; the raw State.error_scale stays float32), and running residuals (State.resid). These quantities are stored in scaled units so that narrow dtypes do not under/overflow (see Forest.leaf_unit), though float16 error scales may still overflow if their dynamic range is high enough. Leaf full conditionals are always computed and sampled in float32. Narrow residual dtypes may easily break the MCMC, resid_dtype is an experimental setting.

  • leaf_quantization (int | Integer[Array, ''] | Integer[ndarray, ''] | None, default: None) – Quantize the leaves to (almost) stop the numerical drift of the running residuals, see StepConfig.leaf_quantization.

  • error_cov_inv (Wishart | None, default: None) – The prior and initial value of the error term precision matrix (see Wishart). Leave it unspecified for binary regression. Mixed binary-continuous and per-outcome-component missingness require a DiagWishart.

  • error_scale (Float32[Array, 'n'] | Float32[ndarray, 'n'] | Float32[Array, 'k n'] | Float32[ndarray, 'k n'] | None, default: None) – Per-datapoint error scales (called w in the R package BART3); see State.error_scale. An unspecified error_scale is equivalent to error_scale = 1 for all datapoints.

  • missing (Bool[Array, 'n'] | Bool[ndarray, 'n'] | Bool[Array, 'k n'] | Bool[ndarray, 'k n'] | None, default: None) – Boolean mask indicating which datapoints are missing. True marks entries to be ignored by the MCMC. The masked values of y may be anything, even non-finite. If missing is 2-D, error_cov_inv must be a DiagWishart.

  • min_points_per_decision_node (int | Integer[Array, ''] | Integer[ndarray, ''] | None, default: None)

  • min_points_per_leaf (int | Integer[Array, ''] | Integer[ndarray, ''] | None, default: None) – The minimum number of datapoints in a decision node and in a leaf, respectively; 0 if not specified. The leaf constraint is not taken into account in the proposal distribution because that would be expensive. The two are independent and not checked for coherence; it makes sense to set min_points_per_decision_node >= 2 * min_points_per_leaf.

  • resid_reduction_config (ReductionConfig, default: AutoBatchedReduction())

  • count_reduction_config (ReductionConfig, default: AutoOneHotReduction())

  • prec_reduction_config (ReductionConfig, default: AutoOneHotReduction()) – How to sum the residuals, count the datapoints, and sum the likelihood precisions in each leaf, respectively. See ReductionConfig and its subclasses.

  • prec_count_num_trees (int | None | Literal['auto'], default: 'auto') – The number of trees to process at a time when counting datapoints or computing the likelihood precision. If None, do all trees at once, which may use too much memory on cpu. If ‘auto’ (default), it’s chosen automatically.

  • sequential_unroll (int | bool, default: 2) – See StepConfig.sequential_unroll. Unrolling may speed up the MCMC at the cost of longer compilation; 1 means no unrolling.

  • save_ratios (bool, default: False) – Whether to save the Metropolis-Hastings ratios.

  • filter_splitless_vars (int, default: 0) – The maximum number of variables without splits that can be ignored. If there are more, init raises an exception.

  • log_s (Float32[Array, 'p'] | Float32[ndarray, 'p'] | None, default: None)

  • theta (float | Float[Array, ''] | Float[ndarray, ''] | None, default: None)

  • a (float | Float[Array, ''] | Float[ndarray, ''] | None, default: None)

  • b (float | Float[Array, ''] | Float[ndarray, ''] | None, default: None)

  • rho (float | Float[Array, ''] | Float[ndarray, ''] | None, default: None) – Sparsity (variable selection) parameters, see Forest.log_s and Forest.theta. If rho, a, b are set, an unspecified theta is initialized to rho; if theta is set, an unspecified log_s is initialized to uniform.

  • sparse_on_at (int | Integer[Array, ''] | Integer[ndarray, ''] | None, default: None) – After how many MCMC steps to turn on variable selection.

  • augment (bool, default: True) – See StepConfig.augment. If disabled, forbidden decision rules are ignored when counting variable usage, which may be faster but is an approximation.

  • num_chains (int | None, default: None) – The number of independent MCMC chains. Single chain with scalar values if not specified.

  • mesh (Mesh | dict[str, int] | None, default: None) –

    A jax mesh used to shard data and computation across multiple devices. If it has a ‘chains’ axis, that axis is used to shard the chains. If it has a ‘data’ axis, that axis is used to shard the datapoints.

    As a shorthand, if a dictionary mapping axis names to axis size is passed, the corresponding mesh is created, e.g., dict(chains=4, data=2) will let jax pick 8 devices to split chains (which must be a multiple of 4) across 4 pairs of devices, where in each pair the data is split in two.

    Note: if a mesh is passed, the arrays are always sharded according to it. In particular even if the mesh has no ‘chains’ or ‘data’ axis, the arrays will be replicated on all devices in the mesh.

Returns:

State – An initialized BART MCMC state.

Raises:

ValueError – If arguments unused in binary regression are set.

Notes

In decision nodes, the values in X[i, :] are compared to a cutpoint out of the range [1, 2, ..., max_split[i]]. A point belongs to the left child iff X[i, j] < cutpoint. Thus it makes sense for X[i, :] to be integers in the range [0, 1, ..., max_split[i]].

In general the arrays passed to this function as arguments may be donated, invalidating them. Create copies before passing them to init if this happens and you need them again.