From ac09579e437fdedd5a6c0655b49aaacc8e1bfbbc Mon Sep 17 00:00:00 2001
From: "Yngve S. Kristiansen" <ynk@equinor.com>
Date: Tue, 14 Jan 2025 15:37:03 +0100
Subject: [PATCH] (wip) Store ert realization mapping in experiment

---
 src/ert/run_models/everest_run_model.py |  8 +++++
 src/everest/everest_storage.py          | 29 +++++++++++++++++
 tests/everest/test_everest_storage.py   | 43 +++++++++++++++++++++++++
 3 files changed, 80 insertions(+)
 create mode 100644 tests/everest/test_everest_storage.py

diff --git a/src/ert/run_models/everest_run_model.py b/src/ert/run_models/everest_run_model.py
index c9d6871efe3..93c5a0a3c67 100644
--- a/src/ert/run_models/everest_run_model.py
+++ b/src/ert/run_models/everest_run_model.py
@@ -165,10 +165,17 @@ 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:
@@ -176,6 +183,7 @@ def run_experiment(
 
         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,
diff --git a/src/everest/everest_storage.py b/src/everest/everest_storage.py
index e6859e7aef6..e00209f1bae 100644
--- a/src/everest/everest_storage.py
+++ b/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()
diff --git a/tests/everest/test_everest_storage.py b/tests/everest/test_everest_storage.py
new file mode 100644
index 00000000000..c6962ae578d
--- /dev/null
+++ b/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