Merge pull request #155 from StochasticTree/sampler-cleanup-hotfix

Fix post-sampler cleanup issues in BART and BCF for R and Python

Author: andrewherren

R/bart.R

@@ -826,13 +826,13 @@ bart <- function(X_train, y_train, leaf_basis_train = NULL, rfx_group_ids_train
     if ((!keep_gfr) && (num_gfr > 0)) {
         for (i in 1:num_gfr) {
             if (include_mean_forest) {
-                forest_samples_mean$delete_sample(i-1)
+                forest_samples_mean$delete_sample(0)
             }
             if (include_variance_forest) {
-                forest_samples_variance$delete_sample(i-1)
+                forest_samples_variance$delete_sample(0)
             }
             if (has_rfx) {
-                rfx_samples$delete_sample(i-1)
+                rfx_samples$delete_sample(0)
             }
         }
         if (sample_sigma_global) {

R/bcf.R

@@ -1214,13 +1214,13 @@ bcf <- function(X_train, Z_train, y_train, propensity_train = NULL, rfx_group_id
     # Remove GFR samples if they are not to be retained
     if ((!keep_gfr) && (num_gfr > 0)) {
         for (i in 1:num_gfr) {
-            forest_samples_mu$delete_sample(i-1)
-            forest_samples_tau$delete_sample(i-1)
+            forest_samples_mu$delete_sample(0)
+            forest_samples_tau$delete_sample(0)
             if (include_variance_forest) {
-                forest_samples_variance$delete_sample(i-1)
+                forest_samples_variance$delete_sample(0)
             }
             if (has_rfx) {
-                rfx_samples$delete_sample(i-1)
+                rfx_samples$delete_sample(0)
             }
         }
         if (sample_sigma_global) {

stochtree/bart.py

@@ -226,6 +226,10 @@ def sample(self, X_train: Union[np.array, pd.DataFrame], y_train: np.array, basi
         keep_vars_variance = variance_forest_params_updated['keep_vars']
         drop_vars_variance = variance_forest_params_updated['drop_vars']
 
+        # Override keep_gfr if there are no MCMC samples
+        if num_mcmc == 0:
+            keep_gfr = True
+        
         # Check that num_chains >= 1
         if not isinstance(num_chains, Integral) or num_chains < 1:
             raise ValueError("num_chains must be an integer greater than 0")
@@ -680,9 +684,9 @@ def sample(self, X_train: Union[np.array, pd.DataFrame], y_train: np.array, basi
         if not keep_gfr and num_gfr > 0:
             for i in range(num_gfr):
                 if self.include_mean_forest:
-                    self.forest_container_mean.delete_sample(i)
+                    self.forest_container_mean.delete_sample(0)
                 if self.include_variance_forest:
-                    self.forest_container_variance.delete_sample(i)
+                    self.forest_container_variance.delete_sample(0)
             if self.sample_sigma_global:
                 self.global_var_samples = self.global_var_samples[num_gfr:]
             if self.sample_sigma_leaf:

stochtree/bcf.py

@@ -299,6 +299,10 @@ def sample(self, X_train: Union[pd.DataFrame, np.array], Z_train: np.array, y_tr
         keep_vars_variance = variance_forest_params_updated['keep_vars']
         drop_vars_variance = variance_forest_params_updated['drop_vars']
 
+        # Override keep_gfr if there are no MCMC samples
+        if num_mcmc == 0:
+            keep_gfr = True
+        
         # Variable weight preprocessing (and initialization if necessary)
         if variable_weights is None:
             if X_train.ndim > 1:
@@ -1051,10 +1055,10 @@ def sample(self, X_train: Union[pd.DataFrame, np.array], Z_train: np.array, y_tr
         # Remove GFR samples if they are not to be retained
         if not keep_gfr and num_gfr > 0:
             for i in range(num_gfr):
-                self.forest_container_mu.delete_sample(i)
-                self.forest_container_tau.delete_sample(i)
+                self.forest_container_mu.delete_sample(0)
+                self.forest_container_tau.delete_sample(0)
                 if self.include_variance_forest:
-                    self.forest_container_variance.delete_sample(i)
+                    self.forest_container_variance.delete_sample(0)
             if self.adaptive_coding:
                 self.b1_samples = self.b1_samples[num_gfr:]
                 self.b0_samples = self.b0_samples[num_gfr:]