From adc6f9a33ff0aa1a7ab3152b40e526b7262a70f7 Mon Sep 17 00:00:00 2001
From: Feda Curic <feda.curic@gmail.com>
Date: Thu, 5 Sep 2024 08:57:56 +0200
Subject: [PATCH] Dont use get_ensemble_by_name from LocalStorage

Co-authored by: larsevj <levje@equinor.com>

erts api has been updated such that get_ensemble_by_name
is only available via LocalExperiment.
---
 pyproject.toml                                |  2 +-
 .../workflows/ahm_analysis/ahmanalysis.py     | 19 +++++++++++++------
 .../ahm_analysis/test_ahm_analysis.py         |  3 ++-
 .../ahm_analysis/test_integration.py          |  6 ++++--
 4 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 9df7e3ff..6887779e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -29,7 +29,7 @@ classifiers = [
 dynamic = ["version"]
 dependencies = [
     "resdata",
-    "ert>=10.2.0b13",
+    "ert>=11.0.0-b0",
     "importlib_metadata",
     "importlib_resources",
     "numpy",
diff --git a/src/semeio/workflows/ahm_analysis/ahmanalysis.py b/src/semeio/workflows/ahm_analysis/ahmanalysis.py
index 3caba191..2adaaf79 100644
--- a/src/semeio/workflows/ahm_analysis/ahmanalysis.py
+++ b/src/semeio/workflows/ahm_analysis/ahmanalysis.py
@@ -125,16 +125,24 @@ def run(
             prior_name,
             target_name,
         )
+
+        prior_ensemble = None
         # Get the prior scalar parameter distributions
-        ensemble = self.storage.get_ensemble_by_name(prior_name)
-        prior_data = ensemble.load_all_gen_kw_data()
+        for experiment in self.storage.experiments:
+            try:
+                prior_ensemble = experiment.get_ensemble_by_name(prior_name)
+                break
+            except KeyError:
+                pass
+
+        if prior_ensemble is None:
+            raise KeyError(f"No ensembles named {prior_name} found in storage.")
+        prior_data = prior_ensemble.load_all_gen_kw_data()
         try:
             raise_if_empty(
                 dataframes=[
                     prior_data,
-                    self.facade.load_all_misfit_data(
-                        self.storage.get_ensemble_by_name(prior_name)
-                    ),
+                    self.facade.load_all_misfit_data(prior_ensemble),
                 ],
                 messages=[
                     "Empty prior ensemble",
@@ -177,7 +185,6 @@ def run(
                 "tmp_storage", "w"
             ) as storage:
                 try:
-                    prior_ensemble = self.storage.get_ensemble_by_name(prior_name)
                     prev_experiment = prior_ensemble.experiment
                     experiment = storage.create_experiment(
                         parameters=prev_experiment.parameter_configuration.values(),
diff --git a/tests/workflows/ahm_analysis/test_ahm_analysis.py b/tests/workflows/ahm_analysis/test_ahm_analysis.py
index bdcd572f..b7683dd2 100644
--- a/tests/workflows/ahm_analysis/test_ahm_analysis.py
+++ b/tests/workflows/ahm_analysis/test_ahm_analysis.py
@@ -17,7 +17,8 @@ def test_make_update_log_df(snake_oil_facade, snapshot):
     snapshots are correct, they are just documenting the current behavior.
     """
     with open_storage(snake_oil_facade.enspath, "w") as storage:
-        prior_ens = storage.get_ensemble_by_name("default")
+        experiment = storage.get_experiment_by_name("ensemble-experiment")
+        prior_ens = experiment.get_ensemble_by_name("default")
         posterior_ens = storage.create_ensemble(
             prior_ens.experiment_id,
             ensemble_size=prior_ens.ensemble_size,
diff --git a/tests/workflows/ahm_analysis/test_integration.py b/tests/workflows/ahm_analysis/test_integration.py
index c4c9b81a..ea7f4519 100644
--- a/tests/workflows/ahm_analysis/test_integration.py
+++ b/tests/workflows/ahm_analysis/test_integration.py
@@ -14,10 +14,11 @@
 def test_ahmanalysis_run(snake_oil_facade):
     """test data_set with only scalar parameters"""
     with open_storage(snake_oil_facade.enspath, "w") as storage:
+        experiment = storage.get_experiment_by_name("ensemble-experiment")
         snake_oil_facade.run_ertscript(
             ahmanalysis.AhmAnalysisJob,
             storage,
-            storage.get_ensemble_by_name("default"),
+            experiment.get_ensemble_by_name("default"),
         )
 
     # assert that this returns/generates a KS csv file
@@ -70,10 +71,11 @@ def test_ahmanalysis_run_deactivated_obs(snake_oil_facade, snapshot, caplog):
     with open_storage(snake_oil_facade.enspath, "w") as storage, caplog.at_level(
         logging.WARNING
     ):
+        experiment = storage.get_experiment_by_name("ensemble-experiment")
         snake_oil_facade.run_ertscript(
             ahmanalysis.AhmAnalysisJob,
             storage,
-            storage.get_ensemble_by_name("default"),
+            experiment.get_ensemble_by_name("default"),
         )
     assert "Analysis failed for" in caplog.text