Merge pull request #327 from StochasticTree/python_docs_rework

Moving parameter list documentation to notes section in python docstrings for BART and BCF and change BCF adaptive coding default to False

Author: andrewherren

R/bcf.R

@@ -102,7 +102,7 @@ NULL
 #'   - `sigma2_global_scale` Scale parameter in the `IG(sigma2_global_shape, sigma2_global_scale)` global error variance model. Default: `0`.
 #'   - `variable_weights` Numeric weights reflecting the relative probability of splitting on each variable. Does not need to sum to 1 but cannot be negative. Defaults to `rep(1/ncol(X_train), ncol(X_train))` if not set here. Note that if the propensity score is included as a covariate in either forest, its weight will default to `1/ncol(X_train)`. A workaround if you wish to provide a custom weight for the propensity score is to include it as a column in `X_train` and then set `propensity_covariate` to `'none'` adjust `keep_vars` accordingly for the `prognostic` or `treatment_effect` forests.
 #'   - `propensity_covariate` Whether to include the propensity score as a covariate in either or both of the forests. Enter `"none"` for neither, `"prognostic"` for the prognostic forest, `"treatment_effect"` for the treatment forest, and `"both"` for both forests. If this is not `"none"` and a propensity score is not provided, it will be estimated from (`X_train`, `Z_train`) using `stochtree::bart()`. Default: `"mu"`.
-#'   - `adaptive_coding` Whether or not to use an "adaptive coding" scheme in which a binary treatment variable is not coded manually as (0,1) or (-1,1) but learned via parameters `b_0` and `b_1` that attach to the outcome model `[b_0 (1-Z) + b_1 Z] tau(X)`. This is ignored when Z is not binary. Default: `TRUE`.
+#'   - `adaptive_coding` Whether or not to use an "adaptive coding" scheme in which a binary treatment variable is not coded manually as (0,1) or (-1,1) but learned via parameters `b_0` and `b_1` that attach to the outcome model `[b_0 (1-Z) + b_1 Z] tau(X)`. This is ignored when Z is not binary. Default: `FALSE`.
 #'   - `control_coding_init` Initial value of the "control" group coding parameter. This is ignored when Z is not binary. Default: `-0.5`.
 #'   - `treated_coding_init` Initial value of the "treatment" group coding parameter. This is ignored when Z is not binary. Default: `0.5`.
 #'   - `rfx_prior_var` Prior on the (diagonals of the) covariance of the additive group-level random regression coefficients. Must be a vector of length `ncol(rfx_basis_train)`. Default: `rep(1, ncol(rfx_basis_train))`
@@ -258,7 +258,7 @@ bcf <- function(
     sigma2_global_scale = 0,
     variable_weights = NULL,
     propensity_covariate = "prognostic",
-    adaptive_coding = TRUE,
+    adaptive_coding = FALSE,
     control_coding_init = -0.5,
     treated_coding_init = 0.5,
     rfx_prior_var = NULL,

stochtree/bart.py

@@ -224,7 +224,7 @@ test_that("BCF multi-chain: sample counts with GFR warm-start", {
     propensity_train = d$pi_train,
     X_test = d$X_test, Z_test = d$Z_test, propensity_test = d$pi_test,
     num_gfr = n_gfr, num_burnin = 5, num_mcmc = n_mcmc,
-    general_params = list(num_chains = n_chains, num_threads = 1)
+    general_params = list(num_chains = n_chains, num_threads = 1, adaptive_coding = TRUE)
   )
   expected <- n_chains * n_mcmc
   expect_length(m$sigma2_global_samples, expected)
@@ -279,7 +279,7 @@ test_that("BCF multi-chain: all samples finite with GFR + multiple chains", {
     propensity_train = d$pi_train,
     X_test = d$X_test, Z_test = d$Z_test, propensity_test = d$pi_test,
     num_gfr = 6, num_burnin = 20, num_mcmc = 10,
-    general_params = list(num_chains = 3, num_threads = 1)
+    general_params = list(num_chains = 3, num_threads = 1, adaptive_coding = TRUE)
   )
   expect_true(all(is.finite(m$sigma2_global_samples)),
               label = "sigma2 samples must be finite (no chain-transition blowup)")

stochtree/bcf.py

@@ -234,7 +234,7 @@ test_that("BCF print method", {
   pi_train <- pi_X[train_inds]; pi_test <- pi_X[test_inds]
   y_train <- y[train_inds]; y_test <- y[test_inds]
 
-  # --- User-provided propensity, binary treatment, adaptive coding (defaults) ---
+  # --- User-provided propensity, binary treatment, default coding (defaults) ---
   bcf_model <- bcf(
     X_train = X_train, y_train = y_train, Z_train = Z_train,
     propensity_train = pi_train,
@@ -249,7 +249,7 @@ test_that("BCF print method", {
   expect_true(any(grepl("prognostic forest", out, fixed = TRUE)))
   expect_true(any(grepl("treatment effect forest", out, fixed = TRUE)))
   expect_true(any(grepl("User-provided propensity scores", out, fixed = TRUE)))
-  expect_true(any(grepl("adaptive coding", out, fixed = TRUE)))
+  expect_true(any(grepl("default coding", out, fixed = TRUE)))
   expect_true(any(grepl("1 chain of", out, fixed = TRUE)))
   expect_true(any(grepl("retaining every iteration", out, fixed = TRUE)))
 
@@ -327,7 +327,7 @@ test_that("BCF summary method", {
     propensity_train = pi_train,
     X_test = X_test, Z_test = Z_test, propensity_test = pi_test,
     num_gfr = 0, num_burnin = 10, num_mcmc = 10,
-    general_params = list(sample_sigma2_global = TRUE),
+    general_params = list(sample_sigma2_global = TRUE, adaptive_coding = TRUE),
     prognostic_forest_params = list(sample_sigma2_leaf = TRUE),
     treatment_effect_forest_params = list(sample_sigma2_leaf = TRUE)
   )

test/R/testthat/test-multi-chain.R

@@ -193,7 +193,8 @@ test_that("BCF JSON uses canonical field names (sigma2_init, b1_samples, b0_samp
   bcf_model <- bcf(
     X_train = X, Z_train = Z, y_train = y,
     propensity_train = pi_x,
-    num_gfr = 0, num_burnin = 0, num_mcmc = 10
+    num_gfr = 0, num_burnin = 0, num_mcmc = 10,
+    general_params = list(adaptive_coding = TRUE)
   )
   json_string <- saveBCFModelToJsonString(bcf_model)
 
@@ -224,7 +225,8 @@ test_that("BCF JSON deserialization handles legacy field names with warnings", {
   bcf_model <- bcf(
     X_train = X, Z_train = Z, y_train = y,
     propensity_train = pi_x,
-    num_gfr = 0, num_burnin = 0, num_mcmc = 10
+    num_gfr = 0, num_burnin = 0, num_mcmc = 10,
+    general_params = list(adaptive_coding = TRUE)
   )
   preds_orig <- predict(bcf_model, X_test, Z_test, pi_test)
 

test/R/testthat/test-print-summary.R

@@ -247,6 +247,7 @@ def test_continuous_univariate_bcf(self):
         num_mcmc = 10
 
         # Run BCF with test set and propensity score
+        # adaptive_coding=True triggers a UserWarning for non-binary treatment
         with pytest.warns(UserWarning):
             bcf_model = BCFModel()
             variance_forest_params = {"num_trees": 0}
@@ -261,6 +262,7 @@ def test_continuous_univariate_bcf(self):
                 num_gfr=num_gfr,
                 num_burnin=num_burnin,
                 num_mcmc=num_mcmc,
+                general_params={"adaptive_coding": True},
                 variance_forest_params=variance_forest_params,
             )
 
@@ -304,6 +306,7 @@ def test_continuous_univariate_bcf(self):
                 num_gfr=num_gfr,
                 num_burnin=num_burnin,
                 num_mcmc=num_mcmc,
+                general_params={"adaptive_coding": True},
                 variance_forest_params=variance_forest_params,
             )
 
@@ -375,6 +378,7 @@ def test_continuous_univariate_bcf(self):
                 num_gfr=num_gfr,
                 num_burnin=num_burnin,
                 num_mcmc=num_mcmc,
+                general_params={"adaptive_coding": True},
                 variance_forest_params=variance_forest_params,
             )
 
@@ -413,6 +417,7 @@ def test_continuous_univariate_bcf(self):
                 num_gfr=num_gfr,
                 num_burnin=num_burnin,
                 num_mcmc=num_mcmc,
+                general_params={"adaptive_coding": True},
                 variance_forest_params=variance_forest_params,
             )
 
@@ -452,6 +457,7 @@ def test_continuous_univariate_bcf(self):
                 num_gfr=num_gfr,
                 num_burnin=num_burnin,
                 num_mcmc=num_mcmc,
+                general_params={"adaptive_coding": True},
                 variance_forest_params=variance_forest_params,
             )
 
@@ -486,6 +492,7 @@ def test_continuous_univariate_bcf(self):
                 num_gfr=num_gfr,
                 num_burnin=num_burnin,
                 num_mcmc=num_mcmc,
+                general_params={"adaptive_coding": True},
                 variance_forest_params=variance_forest_params,
             )
 
@@ -576,6 +583,7 @@ def test_multivariate_bcf(self):
         num_mcmc = 10
 
         # Run BCF with test set and propensity score
+        # adaptive_coding=True triggers a UserWarning for non-binary treatment
         with pytest.warns(UserWarning):
             bcf_model = BCFModel()
             variance_forest_params = {"num_trees": 0}
@@ -590,6 +598,7 @@ def test_multivariate_bcf(self):
                 num_gfr=num_gfr,
                 num_burnin=num_burnin,
                 num_mcmc=num_mcmc,
+                general_params={"adaptive_coding": True},
                 variance_forest_params=variance_forest_params,
             )
 
@@ -630,6 +639,7 @@ def test_multivariate_bcf(self):
                 num_gfr=num_gfr,
                 num_burnin=num_burnin,
                 num_mcmc=num_mcmc,
+                general_params={"adaptive_coding": True},
                 variance_forest_params=variance_forest_params,
             )
 
@@ -668,6 +678,7 @@ def test_multivariate_bcf(self):
                 num_gfr=num_gfr,
                 num_burnin=num_burnin,
                 num_mcmc=num_mcmc,
+                general_params={"adaptive_coding": True},
                 variance_forest_params=variance_forest_params,
             )
 
@@ -682,6 +693,7 @@ def test_multivariate_bcf(self):
                 num_gfr=num_gfr,
                 num_burnin=num_burnin,
                 num_mcmc=num_mcmc,
+                general_params={"adaptive_coding": True},
                 variance_forest_params=variance_forest_params,
             )
 

test/R/testthat/test-serialization.R

@@ -215,7 +215,14 @@ def test_sample_counts_with_gfr(self, bcf_data):
         n_chains = self.NUM_CHAINS
         n_mcmc = self.NUM_MCMC
         n_gfr = self.NUM_GFR
-        m = _bcf(bcf_data, num_gfr=n_gfr, num_burnin=5, num_mcmc=n_mcmc, num_chains=n_chains)
+        m = _bcf(
+            bcf_data,
+            num_gfr=n_gfr,
+            num_burnin=5,
+            num_mcmc=n_mcmc,
+            num_chains=n_chains,
+            adaptive_coding=True,
+        )
         expected = n_chains * n_mcmc
         assert m.global_var_samples.shape == (expected,)
         # BCF-specific samples
@@ -254,6 +261,7 @@ def test_samples_finite_gfr_multi_chain(self, bcf_data):
             num_burnin=20,
             num_mcmc=self.NUM_MCMC,
             num_chains=self.NUM_CHAINS,
+            adaptive_coding=True,
         )
         assert np.all(np.isfinite(m.global_var_samples)), (
             "sigma2 samples contain non-finite values; possible chain-transition blowup."

test/python/test_bcf.py

@@ -332,7 +332,7 @@ def test_unsampled_model(self):
         assert "Empty BCFModel()" in str(BCFModel())
 
     def test_default_model(self, bcf_data):
-        """Binary treatment, user propensity, adaptive coding (defaults): 2 base terms."""
+        """Binary treatment, user propensity, default coding (defaults): 2 base terms."""
         model = BCFModel()
         model.sample(
             X_train=bcf_data["X_train"],
@@ -350,7 +350,7 @@ def test_default_model(self, bcf_data):
         assert "BCFModel run with prognostic forest" in s
         assert "treatment effect forest" in s
         assert "User-provided propensity scores" in s
-        assert "adaptive coding" in s
+        assert "default coding" in s
         assert "1 chain" in s
         assert "retaining every iteration" in s