(wip) Store ert realization mapping in experiment

src/ert/run_models/everest_run_model.py

@@ -165,17 +165,25 @@ def run_experiment(
     ) -> None:
         self.log_at_startup()
         self._eval_server_cfg = evaluator_server_config
+
+        realization_mapping = self._create_ert_realization_mapping(
+            geo_realizations=self._everest_config.model.realizations,
+            num_perturbations=self._everest_config.optimization.perturbation_num,
+        )
+
         self._experiment = self._storage.create_experiment(
             name=f"EnOpt@{datetime.datetime.now().strftime('%Y-%m-%d@%H:%M:%S')}",
             parameters=self.ert_config.ensemble_config.parameter_configuration,
             responses=self.ert_config.ensemble_config.response_configuration,
+            simulation_arguments={"realization_mapping": realization_mapping},
         )
 
         # Initialize the ropt optimizer:
         optimizer = self._create_optimizer()
 
         self.ever_storage = EverestStorage(
             output_dir=Path(self._everest_config.optimization_output_dir),
+            ens_path=self.ert_config.ens_path,
         )
         self.ever_storage.observe_optimizer(
             optimizer,

src/everest/everest_storage.py

@@ -64,6 +64,11 @@ def existing_dataframes(self) -> dict[str, polars.DataFrame]:
         }
 
 
+class _EverestInfoForErtRealization(TypedDict, total=False):
+    geo_realization: int
+    perturbation: int | None  # None means it is an unperturbed evaluation
+
+
 @dataclass
 class EverestStorageDataFrames:
     batches: list[BatchDataFrames] = field(default_factory=list)
@@ -72,6 +77,28 @@ class EverestStorageDataFrames:
     nonlinear_constraints: polars.DataFrame | None = None
     realization_weights: polars.DataFrame | None = None
 
+    @staticmethod
+    def _create_ert_realization_mapping(
+        geo_realizations: list[int], num_perturbations: int
+    ) -> dict[int, _EverestInfoForErtRealization]:
+        mapping = {}
+        ert_realization_counter = 0
+        for geo_real in geo_realizations:
+            mapping[ert_realization_counter] = _EverestInfoForErtRealization(
+                geo_realization=geo_real
+            )
+            ert_realization_counter += 1
+
+        for geo_real in geo_realizations:
+            for p in range(num_perturbations):
+                mapping[ert_realization_counter] = _EverestInfoForErtRealization(
+                    geo_realization=geo_real,
+                    perturbation=p + 1,  # 1-indexed
+                )
+                ert_realization_counter += 1
+
+        return mapping
+
     @property
     def simulation_to_geo_realization_map(self) -> dict[int, int]:
         """
@@ -265,11 +292,13 @@ class _MeritValue:
 class EverestStorage:
     def __init__(
         self,
+        ens_path: str,
         output_dir: Path,
     ) -> None:
         self._control_ensemble_id = 0
         self._gradient_ensemble_id = 0
 
+        self._ens_path = ens_path
         self._output_dir = output_dir
         self._merit_file: Path | None = None
         self.data = EverestStorageDataFrames()

tests/everest/test_everest_storage.py

@@ -0,0 +1,43 @@
+import pytest
+
+from everest.everest_storage import EverestStorage
+
+
+@pytest.mark.parametrize(
+    "geo_realizations, num_perturbations, expected",
+    [
+        (
+            [0, 2],
+            7,
+            {
+                0: {"geo_realization": 0},
+                1: {"geo_realization": 2},
+                2: {"geo_realization": 0, "perturbation": 1},
+                3: {"geo_realization": 0, "perturbation": 2},
+                4: {"geo_realization": 0, "perturbation": 3},
+                5: {"geo_realization": 0, "perturbation": 4},
+                6: {"geo_realization": 0, "perturbation": 5},
+                7: {"geo_realization": 0, "perturbation": 6},
+                8: {"geo_realization": 0, "perturbation": 7},
+                9: {"geo_realization": 2, "perturbation": 1},
+                10: {"geo_realization": 2, "perturbation": 2},
+                11: {"geo_realization": 2, "perturbation": 3},
+                12: {"geo_realization": 2, "perturbation": 4},
+                13: {"geo_realization": 2, "perturbation": 5},
+                14: {"geo_realization": 2, "perturbation": 6},
+                15: {"geo_realization": 2, "perturbation": 7},
+            },
+        ),
+        ([0], 0, {0: {"geo_realization": 0}}),
+        (
+            [0],
+            1,
+            {0: {"geo_realization": 0}, 1: {"geo_realization": 0, "perturbation": 1}},
+        ),
+    ],
+)
+def test_realization_mapping(geo_realizations, num_perturbations, expected):
+    mapping = EverestStorage._create_ert_realization_mapping(
+        geo_realizations, num_perturbations
+    )
+    assert mapping == expected