bartz.bcf.bcf

class bartz.bcf.bcf(x_train, y_train, z_train, *, pihat_train=None, x_test=None, z_test=None, pihat_test=None, include_pihat_in_mu=True, include_pihat_in_tau=False, num_trees_mu=250, num_trees_tau=100, ndpost=1000, nskip=100, num_chains=None, k_mu=2.0, k_tau=10.0, sigma_df=3.0, sigma_scale='auto', sigma_init='auto', leaf_prior_cov_inv_mu=None, leaf_prior_cov_inv_tau=None, min_points_per_leaf_mu=5, min_points_per_leaf_tau=5, tau_0_prior_var=None, sample_intercept=True, adaptive_coding=False, sample_sigma2_leaf_mu=True, sigma2_leaf_shape_mu=3.0, sigma2_leaf_scale_mu=None, sample_sigma2_leaf_tau=False, sigma2_leaf_shape_tau=3.0, sigma2_leaf_scale_tau=None, standardize=True, outcome_type='continuous', delta_max=0.9, seed=0)[source]

Bayesian Causal Forests (BCF).

Regress y_train on x_train and z_train (treatment) with two latent mean functions represented as sums of decision trees: Y = mu(X, pihat) + tau(X, pihat) * Z + error

For continuous outcomes, the hyperparameters with units are on the scale of the standardized response, see standardize, while the outputs are on the scale of y_train.

Parameters:

  • x_train (Real[Array, 'n p'] | Real[ndarray, 'n p'] | DataFrame) – The training predictors (confounders/modifiers).

  • y_train (Float32[Array, 'n'] | Float32[ndarray, 'n'] | Series) – The training responses.

  • z_train (Float32[Array, 'n'] | Float32[ndarray, 'n'] | Series) – The binary treatment assignment (0 or 1).

  • pihat_train (Float32[Array, 'n'] | Float32[ndarray, 'n'] | Series | None, default: None) – The estimated propensity scores. If provided, appended to x_train.

  • x_test (Real[Array, 'm p'] | Real[ndarray, 'm p'] | DataFrame | None, default: None) – The test predictors. If provided, predictions at these points are computed and stored in mu_test and tau_test.

  • z_test (Float32[Array, 'm'] | Float32[ndarray, 'm'] | Series | None, default: None) – The test treatment assignment (0 or 1). If provided together with x_test, the outcome prediction is stored in yhat_test.

  • pihat_test (Float32[Array, 'm'] | Float32[ndarray, 'm'] | Series | None, default: None) – The test propensity scores. Must be passed together with pihat_train.

  • include_pihat_in_mu (bool, default: True) – Whether to include propensity scores in the prognostic forest.

  • include_pihat_in_tau (bool, default: False) – Whether to include propensity scores in the treatment effect forest.

  • num_trees_mu (int, default: 250) – The number of trees used for the prognostic forest mu.

  • num_trees_tau (int, default: 100) – The number of trees used for the treatment effect forest tau.

  • ndpost (int, default: 1000) – The number of MCMC samples to save, after burn-in, per chain. The posterior samples of all chains are concatenated, for a total of num_chains * ndpost.

  • nskip (int, default: 100) – The number of initial MCMC samples to discard as burn-in, per chain.

  • num_chains (int | None, default: None) – The number of independent MCMC chains. None for a single chain without an explicit chain axis, which is equivalent to 1.

  • k_mu (float, default: 2.0) – Prior parameter k for prognostic forest.

  • k_tau (float, default: 10.0) – Prior parameter k for treatment forest.

  • sigma_df (float, default: 3.0) – Prior degrees of freedom for error variance.

  • sigma_scale (float | Literal['auto'], default: 'auto') – Prior scale of the error standard deviation.

  • sigma_init (float | Literal['auto'], default: 'auto') – Initial value of the error standard deviation.

  • leaf_prior_cov_inv_mu (float | Float[Array, ''] | Float[ndarray, ''] | None, default: None) – Custom leaf prior precision for the prognostic forest.

  • leaf_prior_cov_inv_tau (float | Float[Array, ''] | Float[ndarray, ''] | None, default: None) – Custom leaf prior precision for the treatment effect forest.

  • min_points_per_leaf_mu (int, default: 5) – Minimum data points per leaf for prognostic forest.

  • min_points_per_leaf_tau (int, default: 5) – Minimum data points per leaf for treatment forest.

  • tau_0_prior_var (float | None, default: None) – Prior variance for global intercept tau_0.

  • sample_intercept (bool, default: True) – Whether to sample a global treatment intercept tau_0.

  • adaptive_coding (bool, default: False) – Whether to use adaptive coding for the treatment effect.

  • sample_sigma2_leaf_mu (bool, default: True) – Whether to sample the leaf parameter variance for the prognostic forest.

  • sigma2_leaf_shape_mu (float, default: 3.0) – The shape parameter for the Inverse-Gamma prior on the prognostic forest leaf variance.

  • sigma2_leaf_scale_mu (float | Float[Array, ''] | Float[ndarray, ''] | None, default: None) – The scale parameter for the Inverse-Gamma prior on the prognostic forest leaf variance.

  • sample_sigma2_leaf_tau (bool, default: False) – Whether to sample the leaf parameter variance for the treatment effect forest.

  • sigma2_leaf_shape_tau (float, default: 3.0) – The shape parameter for the Inverse-Gamma prior on the treatment effect forest leaf variance.

  • sigma2_leaf_scale_tau (float | Float[Array, ''] | Float[ndarray, ''] | None, default: None) – The scale parameter for the Inverse-Gamma prior on the treatment effect forest leaf variance.

  • standardize (bool, default: True) – Whether to standardize y_train internally, ignored for binary outcomes. If False, the hyperparameters with units (sigma_scale, sigma_init, leaf_prior_cov_inv_*, tau_0_prior_var, sigma2_leaf_scale_*) are on the scale of y_train instead of the standardized one.

  • outcome_type (Literal['continuous', 'binary'], default: 'continuous') – Either ‘continuous’ or ‘binary’ (probit link).

  • delta_max (float, default: 0.9) – Maximum plausible treatment effect on the probability scale for binary probit models.

  • seed (int | Key[Array, ''], default: 0) – The seed for the random number generator.

Raises:

ValueError – If binary outcome is specified but y_train contains values other than 0 or 1, or if the format of x_test does not match x_train format, or if z_test or pihat_test is passed without x_test, or if only one of pihat_train and pihat_test is passed, or if z_test or pihat_test does not match the length of x_test.

save_npz(path)[source]

Save the loaded BCF traces to an NPZ archive.

Parameters:

path (str | Path) – The file path to save the NPZ archive.

Return type:

None

classmethod load_npz(path)[source]

Load BCF traces from an NPZ archive, bypassing __init__ MCMC.

Parameters:

path (str | Path) – The file path to the saved NPZ archive.

Returns:

bcf – A reconstituted bcf object ready for prediction.

Raises:

ValueError – If the schema version in the archive is unsupported.

predict(x_test, *, pihat_test=None, include_pihat_in_mu=True, include_pihat_in_tau=False)[source]

Compute predictions for both mu and tau forests at x_test.

Parameters:

  • x_test (Real[Array, 'm p'] | Real[ndarray, 'm p'] | DataFrame) – The test predictors.

  • pihat_test (Float32[Array, 'm'] | Float32[ndarray, 'm'] | Series | None, default: None) – The test propensity scores.

  • include_pihat_in_mu (bool, default: True) – Whether to include propensity scores in prognostic forest prediction.

  • include_pihat_in_tau (bool, default: False) – Whether to include propensity scores in treatment effect forest prediction.

Returns:

dict – A dictionary with “mu” and “tau” containing the posterior samples of the respective forests evaluated at x_test. Shapes are (num_chains * ndpost, m).

Raises:

ValueError – If the format of x_test does not match x_train format.

property sigma_trace: Float32[Array, 'num_samples'][source]

The posterior trace of residual standard deviation on the outcome scale, chains concatenated.

property mu_test: Float32[Array, 'num_samples m'] | None[source]

The control mean at x_test for each MCMC iteration.

On the latent probit scale for binary outcomes.

property tau_test: Float32[Array, 'num_samples m'] | None[source]

The treatment effect at x_test for each MCMC iteration.

On the latent probit scale for binary outcomes.

property yhat_test: Float32[Array, 'num_samples m'] | None[source]

The outcome at x_test under z_test for each MCMC iteration.

On the latent probit scale for binary outcomes; see prob_test.

property prob_test: Float32[Array, 'num_samples m'] | None[source]

The probability of y being True at x_test under z_test.

None unless the outcome is binary and x_test and z_test were passed to the constructor.

predict_potential_outcomes(x_test, *, pihat_test=None, rho=0.5, key=None, include_pihat_in_mu=True, include_pihat_in_tau=False)[source]

Sample joint posterior predictive potential outcomes Y(0), Y(1), and lift.

Parameters:

  • x_test (Real[Array, 'm p'] | Real[ndarray, 'm p'] | DataFrame) – The test predictors.

  • pihat_test (Float32[Array, 'm'] | Float32[ndarray, 'm'] | Series | None, default: None) – Optional test propensity scores.

  • rho (float, default: 0.5) – Cross-world counterfactual noise correlation in [0, 1].

  • key (Key[Array, ''] | int | None, default: None) – JAX PRNG key or integer seed for stochastic noise sampling.

  • include_pihat_in_mu (bool, default: True) – Whether to include propensity scores in prognostic forest prediction.

  • include_pihat_in_tau (bool, default: False) – Whether to include propensity scores in treatment effect forest prediction.

Returns:

dict – A dictionary mapping outcome names to posterior predictive arrays.

Raises:

ValueError – If rho is not within [0, 1].