Updated python BART interface for user-provided RFX parameters and unit tests

Author: andrewherren

stochtree/bart.py

@@ -23,7 +23,7 @@
 )
 from .sampler import RNG, ForestSampler, GlobalVarianceModel, LeafVarianceModel
 from .serialization import JSONSerializer
-from .utils import NotSampledError, _expand_dims_1d, _expand_dims_2d_diag
+from .utils import NotSampledError, _expand_dims_1d, _expand_dims_2d, _expand_dims_2d_diag
 
 
 class BARTModel:
@@ -260,6 +260,12 @@ def sample(
         keep_every = general_params_updated["keep_every"]
         num_chains = general_params_updated["num_chains"]
         self.probit_outcome_model = general_params_updated["probit_outcome_model"]
+        rfx_working_parameter_prior_mean = general_params_updated["rfx_working_parameter_prior_mean"]
+        rfx_group_parameter_prior_mean = general_params_updated["rfx_group_parameter_prior_mean"]
+        rfx_working_parameter_prior_cov = general_params_updated["rfx_working_parameter_prior_cov"]
+        rfx_group_parameter_prior_cov = general_params_updated["rfx_group_parameter_prior_cov"]
+        rfx_variance_prior_shape = general_params_updated["rfx_variance_prior_shape"]
+        rfx_variance_prior_scale = general_params_updated["rfx_variance_prior_scale"]
 
         # 2. Mean forest parameters
         num_trees_mean = mean_forest_params_updated["num_trees"]
@@ -291,14 +297,6 @@ def sample(
         drop_vars_variance = variance_forest_params_updated["drop_vars"]
         num_features_subsample_variance = variance_forest_params_updated["num_features_subsample"]
 
-        # 4. Random effects parameters
-        rfx_working_parameter_prior_mean = general_params_updated["rfx_working_parameter_prior_mean"]
-        rfx_group_parameter_prior_mean = general_params_updated["rfx_group_parameter_prior_mean"]
-        rfx_working_parameter_prior_cov = general_params_updated["rfx_working_parameter_prior_cov"]
-        rfx_group_parameter_prior_cov = general_params_updated["rfx_group_parameter_prior_cov"]
-        rfx_variance_prior_shape = general_params_updated["rfx_variance_prior_shape"]
-        rfx_variance_prior_scale = general_params_updated["rfx_variance_prior_scale"]
-
         # Override keep_gfr if there are no MCMC samples
         if num_mcmc == 0:
             keep_gfr = True
@@ -993,10 +991,7 @@ def sample(
             if rfx_group_parameter_prior_mean is None:
                 xi_init = np.tile(np.expand_dims(alpha_init, 1), (1, num_rfx_groups))
             else:
-                xi_init = _expand_dims_1d(rfx_group_parameter_prior_mean, num_rfx_components)
-                # If it's a vector, expand to matrix
-                if xi_init.ndim == 1:
-                    xi_init = np.tile(np.expand_dims(xi_init, 1), (1, num_rfx_groups))
+                xi_init = _expand_dims_2d(rfx_group_parameter_prior_mean, num_rfx_components, num_rfx_groups)
 
             if rfx_working_parameter_prior_cov is None:
                 sigma_alpha_init = np.identity(num_rfx_components)

test/python/test_bart.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pytest
 from sklearn.model_selection import train_test_split
 
 from stochtree import BARTModel
@@ -1009,3 +1010,157 @@ def conditional_stddev(X):
         np.testing.assert_allclose(
             rfx_preds_train_3[:, num_mcmc : (2 * num_mcmc)], rfx_preds_train_2
         )
+
+    def test_bart_rfx_parameters(self):
+        # RNG
+        random_seed = 101
+        rng = np.random.default_rng(random_seed)
+
+        # Generate covariates and basis
+        n = 100
+        p_X = 10
+        p_W = 1
+        X = rng.uniform(0, 1, (n, p_X))
+        W = rng.uniform(0, 1, (n, p_W))
+
+        # Generate RFX group labels and basis term
+        num_rfx_basis = 2
+        num_rfx_groups = 4
+        group_labels = rng.choice(num_rfx_groups, size=n)
+        rfx_basis = np.empty((n, num_rfx_basis))
+        rfx_basis[:, 0] = 1.0
+        if num_rfx_basis > 1:
+            rfx_basis[:, 1:] = rng.uniform(-1, 1, (n, num_rfx_basis - 1))
+
+        # Define the outcome mean function
+        def outcome_mean(X, W):
+            return np.where(
+                (X[:, 0] >= 0.0) & (X[:, 0] < 0.25),
+                -7.5 * W[:, 0],
+                np.where(
+                    (X[:, 0] >= 0.25) & (X[:, 0] < 0.5),
+                    -2.5 * W[:, 0],
+                    np.where(
+                        (X[:, 0] >= 0.5) & (X[:, 0] < 0.75),
+                        2.5 * W[:, 0],
+                        7.5 * W[:, 0],
+                    ),
+                ),
+            )
+
+        # Define the group rfx function
+        def rfx_term(group_labels, basis):
+            return np.where(
+                group_labels == 0, -5 + 1. * basis[:,1], 5 - 1. * basis[:,1]
+            )
+
+        # Define the conditional standard deviation function
+        def conditional_stddev(X):
+            return np.where(
+                (X[:, 0] >= 0.0) & (X[:, 0] < 0.25),
+                0.25,
+                np.where(
+                    (X[:, 0] >= 0.25) & (X[:, 0] < 0.5),
+                    0.5,
+                    np.where((X[:, 0] >= 0.5) & (X[:, 0] < 0.75), 1, 2),
+                ),
+            )
+
+        # Generate outcome
+        epsilon = rng.normal(0, 1, n)
+        y = (
+            outcome_mean(X, W)
+            + rfx_term(group_labels, rfx_basis)
+            + epsilon * conditional_stddev(X)
+        )
+
+        # Test-train split
+        sample_inds = np.arange(n)
+        train_inds, test_inds = train_test_split(sample_inds, test_size=0.5)
+        X_train = X[train_inds, :]
+        X_test = X[test_inds, :]
+        basis_train = W[train_inds, :]
+        basis_test = W[test_inds, :]
+        group_labels_train = group_labels[train_inds]
+        group_labels_test = group_labels[test_inds]
+        rfx_basis_train = rfx_basis[train_inds, :]
+        rfx_basis_test = rfx_basis[test_inds, :]
+        y_train = y[train_inds]
+        n_train = X_train.shape[0]
+        n_test = X_test.shape[0]
+
+        # BART settings
+        num_gfr = 10
+        num_burnin = 0
+        num_mcmc = 10
+
+        # Specify no rfx parameters
+        general_params = {}
+        bart_model = BARTModel()
+        bart_model.sample(
+            X_train=X_train,
+            y_train=y_train,
+            leaf_basis_train=basis_train,
+            rfx_group_ids_train=group_labels_train,
+            rfx_basis_train=rfx_basis_train,
+            X_test=X_test,
+            leaf_basis_test=basis_test,
+            rfx_group_ids_test=group_labels_test,
+            rfx_basis_test=rfx_basis_test,
+            num_gfr=num_gfr,
+            num_burnin=num_burnin,
+            num_mcmc=num_mcmc,
+            general_params=general_params,
+        )
+
+        # Specify scalar rfx parameters
+        general_params = {
+            "rfx_working_parameter_prior_mean": 1.,
+            "rfx_group_parameter_prior_mean": 1.,
+            "rfx_working_parameter_prior_cov": 1.,
+            "rfx_group_parameter_prior_cov": 1.,
+            "rfx_variance_prior_shape": 1,
+            "rfx_variance_prior_scale": 1
+        }
+        bart_model_2 = BARTModel()
+        bart_model_2.sample(
+            X_train=X_train,
+            y_train=y_train,
+            leaf_basis_train=basis_train,
+            rfx_group_ids_train=group_labels_train,
+            rfx_basis_train=rfx_basis_train,
+            X_test=X_test,
+            leaf_basis_test=basis_test,
+            rfx_group_ids_test=group_labels_test,
+            rfx_basis_test=rfx_basis_test,
+            num_gfr=num_gfr,
+            num_burnin=num_burnin,
+            num_mcmc=num_mcmc,
+            general_params=general_params,
+        )
+
+        # Specify all relevant rfx parameters as vectors
+        general_params = {
+            "rfx_working_parameter_prior_mean": np.repeat(1., num_rfx_basis),
+            "rfx_group_parameter_prior_mean": np.repeat(1., num_rfx_basis),
+            "rfx_working_parameter_prior_cov": np.identity(num_rfx_basis),
+            "rfx_group_parameter_prior_cov": np.identity(num_rfx_basis),
+            "rfx_variance_prior_shape": 1,
+            "rfx_variance_prior_scale": 1
+        }
+        bart_model_2 = BARTModel()
+        bart_model_2.sample(
+            X_train=X_train,
+            y_train=y_train,
+            leaf_basis_train=basis_train,
+            rfx_group_ids_train=group_labels_train,
+            rfx_basis_train=rfx_basis_train,
+            X_test=X_test,
+            leaf_basis_test=basis_test,
+            rfx_group_ids_test=group_labels_test,
+            rfx_basis_test=rfx_basis_test,
+            num_gfr=num_gfr,
+            num_burnin=num_burnin,
+            num_mcmc=num_mcmc,
+            general_params=general_params,
+        )