Updated serialization methods to reflect python BART updates

Author: andrewherren

stochtree/bart.py

@@ -2386,13 +2386,16 @@ def to_json(self) -> str:
         bart_json.add_boolean("include_mean_forest", self.include_mean_forest)
         bart_json.add_boolean("include_variance_forest", self.include_variance_forest)
         bart_json.add_boolean("has_rfx", self.has_rfx)
+        bart_json.add_boolean("has_rfx_basis", self.has_rfx_basis)
+        bart_json.add_scalar("num_rfx_basis", self.num_rfx_basis)
         bart_json.add_integer("num_gfr", self.num_gfr)
         bart_json.add_integer("num_burnin", self.num_burnin)
         bart_json.add_integer("num_mcmc", self.num_mcmc)
         bart_json.add_integer("num_samples", self.num_samples)
         bart_json.add_integer("num_basis", self.num_basis)
         bart_json.add_boolean("requires_basis", self.has_basis)
         bart_json.add_boolean("probit_outcome_model", self.probit_outcome_model)
+        bart_json.add_string("rfx_model_spec", self.rfx_model_spec)
 
         # Add parameter samples
         if self.sample_sigma2_global:
@@ -2427,6 +2430,8 @@ def from_json(self, json_string: str) -> None:
         self.include_mean_forest = bart_json.get_boolean("include_mean_forest")
         self.include_variance_forest = bart_json.get_boolean("include_variance_forest")
         self.has_rfx = bart_json.get_boolean("has_rfx")
+        self.has_rfx_basis = bart_json.get_boolean("has_rfx_basis")
+        self.num_rfx_basis = bart_json.get_scalar("num_rfx_basis")
         if self.include_mean_forest:
             # TODO: don't just make this a placeholder that we overwrite
             self.forest_container_mean = ForestContainer(0, 0, False, False)
@@ -2465,6 +2470,7 @@ def from_json(self, json_string: str) -> None:
         self.num_basis = bart_json.get_integer("num_basis")
         self.has_basis = bart_json.get_boolean("requires_basis")
         self.probit_outcome_model = bart_json.get_boolean("probit_outcome_model")
+        self.rfx_model_spec = bart_json.get_string("rfx_model_spec")
 
         # Unpack parameter samples
         if self.sample_sigma2_global:
@@ -2550,6 +2556,8 @@ def from_json_string_list(self, json_string_list: list[str]) -> None:
 
         # Unpack random effects
         self.has_rfx = json_object_default.get_boolean("has_rfx")
+        self.has_rfx_basis = json_object_default.get_boolean("has_rfx_basis")
+        self.num_rfx_basis = json_object_default.get_scalar("num_rfx_basis")
         if self.has_rfx:
             self.rfx_container = RandomEffectsContainer()
             for i in range(len(json_object_list)):
@@ -2575,6 +2583,9 @@ def from_json_string_list(self, json_string_list: list[str]) -> None:
         self.probit_outcome_model = json_object_default.get_boolean(
             "probit_outcome_model"
         )
+        self.rfx_model_spec = json_object_default.get_string(
+            "rfx_model_spec"
+        )
 
         # Unpack number of samples
         for i in range(len(json_object_list)):

test/python/test_bart.py

@@ -1123,6 +1123,7 @@ def conditional_stddev(X):
 
         # Specify scalar rfx parameters
         rfx_params = {
+            "model_spec": "custom",
             "working_parameter_prior_mean": 1.,
             "group_parameter_prior_mean": 1.,
             "working_parameter_prior_cov": 1.,
@@ -1144,11 +1145,12 @@ def conditional_stddev(X):
             num_gfr=num_gfr,
             num_burnin=num_burnin,
             num_mcmc=num_mcmc,
-            rfx_params=rfx_params,
+            random_effects_params=rfx_params,
         )
 
         # Specify all relevant rfx parameters as vectors
         rfx_params = {
+            "model_spec": "custom",
             "working_parameter_prior_mean": np.repeat(1., num_rfx_basis),
             "group_parameter_prior_mean": np.repeat(1., num_rfx_basis),
             "working_parameter_prior_cov": np.identity(num_rfx_basis),
@@ -1170,5 +1172,32 @@ def conditional_stddev(X):
             num_gfr=num_gfr,
             num_burnin=num_burnin,
             num_mcmc=num_mcmc,
-            rfx_params=rfx_params,
+            random_effects_params=rfx_params,
         )
+
+        # Fit a simpler intercept-only RFX model
+        rfx_params = {
+            "model_spec": "intercept_only"
+        }
+        bart_model_4 = BARTModel()
+        bart_model_4.sample(
+            X_train=X_train,
+            y_train=y_train,
+            leaf_basis_train=basis_train,
+            rfx_group_ids_train=group_labels_train,
+            X_test=X_test,
+            leaf_basis_test=basis_test,
+            rfx_group_ids_test=group_labels_test,
+            num_gfr=num_gfr,
+            num_burnin=num_burnin,
+            num_mcmc=num_mcmc,
+            random_effects_params=rfx_params,
+        )
+        preds = bart_model_4.predict(
+            covariates=X_test,
+            basis=basis_test,
+            rfx_group_ids=group_labels_test,
+            type="posterior",
+            terms="rfx"
+        )
+        assert preds.shape == (n_test, num_mcmc)