Skip to content

Commit

Permalink
Iterate ensembles to look for history to plot
Browse files Browse the repository at this point in the history
  • Loading branch information
andreas-el committed May 2, 2024
1 parent 4d16ba1 commit e0c0f66
Show file tree
Hide file tree
Showing 3 changed files with 38 additions and 21 deletions.
37 changes: 20 additions & 17 deletions src/ert/gui/tools/plot/plot_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
8 changes: 5 additions & 3 deletions src/ert/gui/tools/plot/plot_window.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}"
Expand Down
14 changes: 13 additions & 1 deletion tests/unit_tests/gui/tools/plot/test_plot_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down

0 comments on commit e0c0f66

Please sign in to comment.