Fixed prediction bug for Python random effects models with "non-default" random effects IDs

Author: andrewherren

src/py_stochtree.cpp

@@ -1573,6 +1573,18 @@ class RandomEffectsLabelMapperCpp {
   StochTree::LabelMapper* GetLabelMapper() {
     return rfx_label_mapper_.get();
   }
+  int MapGroupIdToArrayIndex(int original_label) {
+    return rfx_label_mapper_->CategoryNumber(original_label);
+  }
+  py::array_t<int> MapMultipleGroupIdsToArrayIndices(py::array_t<int> original_labels) {
+    int output_size = original_labels.size();
+    auto result = py::array_t<int>(py::detail::any_container<py::ssize_t>({output_size}));
+    auto accessor = result.mutable_unchecked<1>();
+    for (int i = 0; i < output_size; i++) {
+      accessor(i) = rfx_label_mapper_->CategoryNumber(original_labels.at(i));
+    }
+    return result;
+  }
 
  private:
   std::unique_ptr<StochTree::LabelMapper> rfx_label_mapper_;
@@ -2410,7 +2422,9 @@ PYBIND11_MODULE(stochtree_cpp, m) {
     .def("DumpJsonString", &RandomEffectsLabelMapperCpp::DumpJsonString)
     .def("LoadFromJsonString", &RandomEffectsLabelMapperCpp::LoadFromJsonString)
     .def("LoadFromJson", &RandomEffectsLabelMapperCpp::LoadFromJson)
-    .def("GetLabelMapper", &RandomEffectsLabelMapperCpp::GetLabelMapper);
+    .def("GetLabelMapper", &RandomEffectsLabelMapperCpp::GetLabelMapper)
+    .def("MapGroupIdToArrayIndex", &RandomEffectsLabelMapperCpp::MapGroupIdToArrayIndex)
+    .def("MapMultipleGroupIdsToArrayIndices", &RandomEffectsLabelMapperCpp::MapMultipleGroupIdsToArrayIndices);
 
   py::class_<RandomEffectsModelCpp>(m, "RandomEffectsModelCpp")
     .def(py::init<int, int>())

stochtree/bart.py

@@ -1988,6 +1988,12 @@ def predict(
                     "Random effects basis has a different dimension than the basis used to train this model"
                 )
 
+        # Convert rfx_group_ids to their corresponding array position indices in the random effects parameter sample arrays
+        if rfx_group_ids is not None:
+            rfx_group_id_indices = self.rfx_container.map_group_ids_to_array_indices(
+                rfx_group_ids
+            )
+        
         # Random effects predictions
         if predict_rfx or predict_rfx_intermediate:
             if rfx_basis is not None:
@@ -2017,7 +2023,7 @@ def predict(
                     )
                     for i in range(n_train):
                         rfx_predictions_raw[i, 0, :] = rfx_beta_draws[
-                            rfx_group_ids[i], :
+                            rfx_group_id_indices[i], :
                         ]
                 rfx_predictions = np.squeeze(rfx_predictions_raw[:, 0, :])
 

stochtree/bcf.py

@@ -3053,6 +3053,12 @@ def predict(
                 raise ValueError(
                     "rfx_basis must have the same number of columns as the random effects basis used to sample this model"
                 )
+            
+        # Convert rfx_group_ids to their corresponding array position indices in the random effects parameter sample arrays
+        if rfx_group_ids is not None:
+            rfx_group_id_indices = self.rfx_container.map_group_ids_to_array_indices(
+                rfx_group_ids
+            )
 
         # Random effects predictions
         if predict_rfx or predict_rfx_intermediate:
@@ -3073,14 +3079,14 @@ def predict(
                 )
                 for i in range(X.shape[0]):
                     rfx_predictions_raw[i, :, :] = rfx_beta_draws[
-                        :, rfx_group_ids[i], :
+                        :, rfx_group_id_indices[i], :
                     ]
             elif rfx_beta_draws.ndim == 2:
                 rfx_predictions_raw = np.empty(
                     shape=(X.shape[0], 1, rfx_beta_draws.shape[1])
                 )
                 for i in range(X.shape[0]):
-                    rfx_predictions_raw[i, 0, :] = rfx_beta_draws[rfx_group_ids[i], :]
+                    rfx_predictions_raw[i, 0, :] = rfx_beta_draws[rfx_group_id_indices[i], :]
             else:
                 raise ValueError(
                     "Unexpected number of dimensions in extracted random effects samples"

stochtree/random_effects.py

@@ -418,6 +418,38 @@ def extract_parameter_samples(self) -> dict[str, np.ndarray]:
             "sigma_samples": sigma_samples,
         }
         return output
+    
+    def map_group_id_to_array_index(self, group_id: int) -> int:
+        """
+        Map an integer-valued random effects group ID to its group's corresponding position in the arrays that store random effects parameter samples.
+
+        Parameters
+        ----------
+        group_id : int
+            Group identifier to be converted to an array position.
+
+        Returns
+        -------
+        int
+            The position of `group_id` in the parameter sample arrays underlying the random effects container.
+        """
+        return self.rfx_label_mapper_cpp.MapGroupIdToArrayIndex(group_id)
+    
+    def map_group_ids_to_array_indices(self, group_ids: np.ndarray) -> np.ndarray:
+        """
+        Map an array of integer-valued random effects group IDs to their groups' corresponding positions in the arrays that store random effects parameter samples.
+
+        Parameters
+        ----------
+        group_ids : np.ndarray
+            Array of group identifiers (integer-valued) to be converted to an array position.
+
+        Returns
+        -------
+        np.ndarray
+            Numpy array of the position of `group_id` in the parameter sample arrays underlying the random effects container.
+        """
+        return self.rfx_label_mapper_cpp.MapMultipleGroupIdsToArrayIndices(group_ids)
 
 
 class RandomEffectsModel: