EverestRunModel: sanitize realization var names

equinor · Dec 20, 2024 · 7318071 · 7318071
1 parent 347da4d
commit 7318071
Showing 1 changed file with 18 additions and 14 deletions.
diff --git a/src/ert/run_models/everest_run_model.py b/src/ert/run_models/everest_run_model.py
@@ -407,13 +407,13 @@ def _get_cached_results(
         cached_results: dict[int, Any] = {}
         if self._simulator_cache is not None:
             assert evaluator_context.config.realizations.names is not None
-            for sim_idx, realization in enumerate(evaluator_context.realizations):
+            for control_idx, real_idx in enumerate(evaluator_context.realizations):
                 cached_data = self._simulator_cache.get(
-                    evaluator_context.config.realizations.names[realization],
-                    control_values[sim_idx, :],
+                    int(evaluator_context.config.realizations.names[real_idx]),
+                    control_values[control_idx, :],
                 )
                 if cached_data is not None:
-                    cached_results[sim_idx] = cached_data
+                    cached_results[control_idx] = cached_data
         return cached_results
 
     @staticmethod
@@ -522,9 +522,12 @@ def _get_run_args(
         self.active_realizations = [True] * len(batch_data)
         assert evaluator_context.config.realizations.names is not None
         for sim_id, control_idx in enumerate(batch_data.keys()):
-            realization = evaluator_context.realizations[control_idx]
             substitutions[f"<GEO_ID_{sim_id}_0>"] = str(
-                evaluator_context.config.realizations.names[realization]
+                int(
+                    evaluator_context.config.realizations.names[
+                        evaluator_context.realizations[control_idx]
+                    ]
+                )
             )
         run_paths = Runpaths(
             jobname_format=self.ert_config.model_config.jobname_format_string,
@@ -654,9 +657,12 @@ def _add_results_to_cache(
         if self._simulator_cache is not None:
             assert evaluator_context.config.realizations.names is not None
             for control_idx in batch_data:
-                realization = evaluator_context.realizations[control_idx]
                 self._simulator_cache.add(
-                    evaluator_context.config.realizations.names[realization],
+                    int(
+                        evaluator_context.config.realizations.names[
+                            evaluator_context.realizations[control_idx]
+                        ]
+                    ),
                     control_values[control_idx, ...],
                     objectives[control_idx, ...],
                     None if constraints is None else constraints[control_idx, ...],
@@ -761,12 +767,12 @@ def __init__(self) -> None:
 
     def add(
         self,
-        realization_id: int,
+        realization: int,
         control_values: NDArray[np.float64],
         objectives: NDArray[np.float64],
         constraints: NDArray[np.float64] | None,
     ) -> None:
-        self._data[realization_id].append(
+        self._data[realization].append(
             (
                 control_values.copy(),
                 objectives.copy(),
@@ -775,11 +781,9 @@ def add(
         )
 
     def get(
-        self, realization_id: int, controls: NDArray[np.float64]
+        self, realization: int, controls: NDArray[np.float64]
     ) -> tuple[NDArray[np.float64], NDArray[np.float64] | None] | None:
-        for control_values, objectives, constraints in self._data.get(
-            realization_id, []
-        ):
+        for control_values, objectives, constraints in self._data.get(realization, []):
             if np.allclose(controls, control_values, rtol=0.0, atol=self.EPS):
                 return objectives, constraints
         return None