Override keep_gfr if no MCMC samples

Author: andrewherren

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")

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: