Iterate ensembles to look for history to plot

Author: andreas-el

src/ert/gui/tools/plot/plot_api.py

@@ -209,25 +209,28 @@ def observations_for_key(self, ensemble_name, key) -> pd.DataFrame:
             }
             return pd.DataFrame(data_struct).T
 
-    def history_data(self, key, ensemble=None) -> pd.DataFrame:
+    def history_data(self, key, ensembles: Optional[List[str]]) -> pd.DataFrame:
         """Returns a pandas DataFrame with the data points for the history for a
         given data key, if any.  The row index is the index/date and the column
         index is the key."""
-
-        if ":" in key:
-            head, tail = key.split(":", 2)
-            history_key = f"{head}H:{tail}"
-        else:
-            history_key = f"{key}H"
-
-        df = self.data_for_key(ensemble, history_key)
-
-        if not df.empty:
-            df = df.T
-            # Drop columns with equal data
-            duplicate_cols = [
-                cc[0] for cc in combi(df.columns, r=2) if (df[cc[0]] == df[cc[1]]).all()
-            ]
-            return df.drop(columns=duplicate_cols)
+        if ensembles:
+            for ensemble in ensembles:
+                if ":" in key:
+                    head, tail = key.split(":", 2)
+                    history_key = f"{head}H:{tail}"
+                else:
+                    history_key = f"{key}H"
+
+                df = self.data_for_key(ensemble, history_key)
+
+                if not df.empty:
+                    df = df.T
+                    # Drop columns with equal data
+                    duplicate_cols = [
+                        cc[0]
+                        for cc in combi(df.columns, r=2)
+                        if (df[cc[0]] == df[cc[1]]).all()
+                    ]
+                    return df.drop(columns=duplicate_cols)
 
         return pd.DataFrame()

src/ert/gui/tools/plot/plot_window.py

@@ -208,15 +208,17 @@ def currentPlotChanged(self):
             plot_config = PlotConfig.createCopy(self._plot_customizer.getPlotConfig())
             plot_context = PlotContext(plot_config, ensembles, key)
 
-            ensemble = plot_context.ensembles()[0] if plot_context.ensembles() else None
-
             # Check if key is a history key.
             # If it is it already has the data it needs
             if str(key).endswith("H") or "H:" in str(key):
                 plot_context.history_data = DataFrame()
             else:
                 try:
-                    plot_context.history_data = self._api.history_data(key, ensemble)
+                    plot_context.history_data = self._api.history_data(
+                        key,
+                        plot_context.ensembles(),
+                    )
+
                 except (RequestError, TimeoutError) as e:
                     logger.exception(e)
                     msg = f"{e}"

tests/unit_tests/gui/tools/plot/test_plot_api.py

@@ -101,12 +101,24 @@ def test_all_data_type_keys(api):
 
 
 def test_load_history_data(api):
-    df = api.history_data(ensemble="default_0", key="FOPR")
+    df = api.history_data(ensembles=["default_0"], key="FOPR")
     assert_frame_equal(
         df, pd.DataFrame({1: [0.2, 0.2, 1.2], 3: [1.0, 1.1, 1.2], 4: [1.0, 1.1, 1.3]})
     )
 
 
+def test_load_history_data_searches_until_history_found(api):
+    df = api.history_data(ensembles=["no-history", "default_0"], key="FOPR")
+    assert_frame_equal(
+        df, pd.DataFrame({1: [0.2, 0.2, 1.2], 3: [1.0, 1.1, 1.2], 4: [1.0, 1.1, 1.3]})
+    )
+
+
+def test_load_history_data_returns_empty_frame_if_no_history(api):
+    df = api.history_data(ensembles=["no-history", "still-no-history"], key="FOPR")
+    assert_frame_equal(df, pd.DataFrame())
+
+
 def test_plot_api_request_errors_all_data_type_keys(api, mocker):
     # Mock the experiment name to be something unexpected
     mocker.patch(