Split out optimizer observer logic

equinor · Nov 18, 2024 · d9acf08 · d9acf08
1 parent 74c1250
commit d9acf08
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 17 deletions.
diff --git a/src/ert/run_models/everest_run_model.py b/src/ert/run_models/everest_run_model.py
@@ -402,13 +402,12 @@ def run_experiment(
         optimizer = self._create_optimizer(simulator)
 
         self.ever_storage = EverestStorage(
-            simulator,
-            optimizer,
             Path(self.everest_config.optimization_output_dir)
             / "dakota"
             / "OPT_DEFAULT.out",
             output_dir=Path(self.everest_config.optimization_output_dir),
         )
+        self.ever_storage.observe_optimizer(optimizer)
 
         # The SqliteStorage object is used to store optimization results from
         # Seba in an sqlite database. It reacts directly to events emitted by

diff --git a/src/everest/everest_storage.py b/src/everest/everest_storage.py
@@ -23,8 +23,6 @@
 from ropt.results import FunctionResults, GradientResults, convert_to_maximize
 from seba_sqlite import sqlite_storage
 
-from everest.simulator import Simulator
-
 logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
@@ -163,9 +161,7 @@ def write_to_experiment(self, experiment: _OptimizerOnlyExperiment):
             for df_key, df in batch_data.existing_dataframes.items():
                 df.write_parquet(experiment.optimizer_mount_point / f"{df_key}.parquet")
 
-    def read_from_experiment(
-        self, experiment: _OptimizerOnlyExperiment
-    ) -> EverestStorageDataFrames:
+    def read_from_experiment(self, experiment: _OptimizerOnlyExperiment) -> None:
         self.initial_values = polars.read_parquet(
             experiment.optimizer_mount_point / "initial_values.parquet"
         )
@@ -245,20 +241,19 @@ def get_ensemble_by_name(self, name: str) -> _OptimizerOnlyEnsemble:
 class EverestStorage:
     def __init__(
         self,
-        simulator: Simulator,
-        optimizer: BasicOptimizer,
         merit_file: Path,
         output_dir: Path,
     ) -> None:
-        self._simulator = simulator
-        self._optimizer = optimizer
         self._initialized = False
         self._control_ensemble_id = 0
         self._gradient_ensemble_id = 0
-        self._simulator_results = None
 
         self._merit_file = merit_file
 
+        self._dataframes = EverestStorageDataFrames()
+        self._experiment = _OptimizerOnlyExperiment(output_dir)
+
+    def observe_optimizer(self, optimizer: BasicOptimizer) -> None:
         optimizer.add_observer(EventType.START_OPTIMIZER_STEP, self._initialize)
         optimizer.add_observer(
             EventType.FINISHED_EVALUATION, self._handle_finished_batch_event
@@ -268,9 +263,6 @@ def __init__(
             self._handle_finished_event,
         )
 
-        self._dataframes = EverestStorageDataFrames()
-        self._experiment = _OptimizerOnlyExperiment(output_dir)
-
     @property
     def experiment(self) -> _OptimizerOnlyExperiment:
         # Should be replaced with ERT experiment
@@ -515,7 +507,7 @@ def _rename_columns(df: polars.DataFrame):
 
     def _store_gradient_results(
         self, config: EnOptConfig, results: FunctionResults
-    ) -> Tuple[polars.DataFrame, polars.DataFrame]:
+    ) -> Tuple[polars.DataFrame, polars.DataFrame, polars.DataFrame, polars.DataFrame]:
         obj_gradient_per_real = polars.from_pandas(
             results.to_dataframe(config, "evaluations").reset_index()
         )
@@ -672,7 +664,7 @@ def _handle_finished_batch_event(self, event: Event):
         converted_results = tuple(
             convert_to_maximize(result) for result in event.results
         )
-        results = []
+        results: List[FunctionResults | GradientResults] = []
         best_value = -np.inf
         best_results = None
         for item in converted_results: