diff --git a/stochtree/bcf.py b/stochtree/bcf.py index 10bf3f3c..2e1b2ad2 100644 --- a/stochtree/bcf.py +++ b/stochtree/bcf.py @@ -932,16 +932,15 @@ def sample( if sample_sigma2_leaf_tau is not None: if not isinstance(sample_sigma2_leaf_tau, bool): raise ValueError("sample_sigma2_leaf_tau must be a bool") - if propensity_covariate is not None: - if propensity_covariate not in [ - "prognostic", - "treatment_effect", - "both", - "none", - ]: - raise ValueError( - "propensity_covariate must be one of 'prognostic', 'treatment_effect', 'both', or 'none'" - ) + if propensity_covariate not in [ + "prognostic", + "treatment_effect", + "both", + "none", + ]: + raise ValueError( + "propensity_covariate must be one of 'prognostic', 'treatment_effect', 'both', or 'none'" + ) if b_0 is not None: b_0 = check_scalar( x=b_0, @@ -1663,15 +1662,6 @@ def sample( ] = 0 # Update covariates to include propensities if requested - if propensity_covariate not in [ - "none", - "prognostic", - "treatment_effect", - "both", - ]: - raise ValueError( - "propensity_covariate must equal one of 'none', 'prognostic', 'treatment_effect', or 'both'" - ) if propensity_covariate != "none": feature_types = np.append( feature_types, np.repeat(0, propensity_train.shape[1]) @@ -1700,9 +1690,10 @@ def sample( variable_weights_tau = np.append( variable_weights_tau, np.repeat(1 / num_cov_orig, propensity_train.shape[1]) ) - variable_weights_variance = np.append( - variable_weights_variance, np.repeat(0.0, propensity_train.shape[1]) - ) + # For now, propensities are not included in the variance forest + variable_weights_variance = np.append( + variable_weights_variance, np.repeat(0.0, propensity_train.shape[1]) + ) # Renormalize variable weights variable_weights_mu = variable_weights_mu / np.sum(variable_weights_mu) diff --git a/test/python/test_bcf.py b/test/python/test_bcf.py index eca2a5ff..c5a1446f 100644 --- a/test/python/test_bcf.py +++ b/test/python/test_bcf.py @@ -194,6 +194,19 @@ def test_binary_bcf(self): # Check treatment effect prediction method tau_hat = bcf_model.predict(X=X_test, Z=Z_test, terms="cate") + # Check that we can run BCF without propensities + bcf_model = BCFModel() + general_params = {"propensity_covariate": "none"} + bcf_model.sample( + X_train=X_train, + Z_train=Z_train, + y_train=y_train, + num_gfr=num_gfr, + num_burnin=num_burnin, + num_mcmc=num_mcmc, + general_params=general_params, + ) + def test_continuous_univariate_bcf(self): # RNG random_seed = 101