From d37df69d11e568f52b03303c94daf8e5950cb317 Mon Sep 17 00:00:00 2001 From: DanSava Date: Tue, 23 Nov 2021 14:28:45 +0200 Subject: [PATCH] Account for ert and ert3 discrepancy when uploading parameter data to storage. --- webviz_ert/data_loader/__init__.py | 50 ++++++++++++++++++++++------ webviz_ert/models/ensemble_model.py | 25 +++++++++++--- webviz_ert/models/parameter_model.py | 5 ++- 3 files changed, 64 insertions(+), 16 deletions(-) diff --git a/webviz_ert/data_loader/__init__.py b/webviz_ert/data_loader/__init__.py index a9d20efe..e08d28f2 100644 --- a/webviz_ert/data_loader/__init__.py +++ b/webviz_ert/data_loader/__init__.py @@ -191,32 +191,62 @@ def get_ensemble_userdata(self, ensemble_id: str) -> dict: def get_ensemble_parameters(self, ensemble_id: str) -> list: return self._get(url=f"ensembles/{ensemble_id}/parameters").json() + def get_record_labels(self, ensemble_id: str, name: str) -> list: + return self._get(url=f"ensembles/{ensemble_id}/records/{name}/labels").json() + def get_experiment_priors(self, experiment_id: str) -> dict: return json.loads( self._query(GET_PRIORS, id=experiment_id)["experiment"]["priors"] ) def get_ensemble_parameter_data( - self, ensemble_id: str, parameter_name: str + self, + ensemble_id: str, + parameter_name: str, + realizations: list, ) -> pd.DataFrame: - resp = self._get( - url=f"ensembles/{ensemble_id}/records/{parameter_name}", - headers={"accept": "application/x-parquet"}, - ) - stream = io.BytesIO(resp.content) - df = pd.read_parquet(stream) - return df + try: + if "::" in parameter_name: + name, label = parameter_name.split("::", 1) + resp = self._get( + url=f"ensembles/{ensemble_id}/records/{name}", + headers={"accept": "application/x-parquet"}, + params={"label": label}, + ) + else: + resp = self._get( + url=f"ensembles/{ensemble_id}/records/{parameter_name}", + headers={"accept": "application/x-parquet"}, + ) + stream = io.BytesIO(resp.content) + df = pd.read_parquet(stream) + return df + except DataLoaderException as e: + logger.error(e) + return self.get_ensemble_record_data( + ensemble_id=ensemble_id, + record_name=name, + active_realizations=realizations, + label=label, + ) def get_ensemble_record_data( - self, ensemble_id: str, record_name: str, active_realizations: List[int] + self, + ensemble_id: str, + record_name: str, + active_realizations: List[int], + label: Optional[str] = None, ) -> pd.DataFrame: dfs = [] for rel_idx in active_realizations: try: + params = {"realization_index": rel_idx} + if label: + params["label"] = label resp = self._get( url=f"ensembles/{ensemble_id}/records/{record_name}", headers={"accept": "application/x-parquet"}, - params={"realization_index": rel_idx}, + params=params, ) stream = io.BytesIO(resp.content) df = pd.read_parquet(stream).transpose().sort_index() diff --git a/webviz_ert/models/ensemble_model.py b/webviz_ert/models/ensemble_model.py index 3e850bcd..d0e6857b 100644 --- a/webviz_ert/models/ensemble_model.py +++ b/webviz_ert/models/ensemble_model.py @@ -1,15 +1,17 @@ import json import pandas as pd from typing import Mapping, List, Dict, Union, Any, Optional -from webviz_ert.data_loader import ( - get_data_loader, -) +from webviz_ert.data_loader import get_data_loader, DataLoaderException from webviz_ert.models import Response, PriorModel, ParametersModel def _create_parameter_models( - parameters_names: list, priors: dict, ensemble_id: str, project_id: str + parameters_names: list, + priors: dict, + ensemble_id: str, + project_id: str, + realizations: list, ) -> Optional[Mapping[str, ParametersModel]]: parameters = {} for param in parameters_names: @@ -30,6 +32,7 @@ def _create_parameter_models( param_id="", # TODO? project_id=project_id, ensemble_id=ensemble_id, + realizations=realizations, ) return parameters @@ -98,7 +101,18 @@ def parameters( self, ) -> Optional[Mapping[str, ParametersModel]]: if not self._parameters: - parameter_names = self._data_loader.get_ensemble_parameters(self._id) + parameter_names = [] + # TODO Remove except code path after PR https://github.com/equinor/ert-storage/pull/180 is merged + try: + for param_name in self._data_loader.get_ensemble_parameters(self._id): + labels = self._data_loader.get_record_labels(self._id, param_name) + if len(labels) > 0: + for label in labels: + parameter_names.append(f"{param_name}::{label}") + else: + parameter_names.append(param_name) + except DataLoaderException as e: + parameter_names = self._data_loader.get_ensemble_parameters(self._id) parameter_priors = ( self._data_loader.get_experiment_priors(self._experiment_id) if not self._parent @@ -109,6 +123,7 @@ def parameters( parameter_priors, ensemble_id=self._id, project_id=self._project_id, + realizations=self._active_realizations, ) return self._parameters diff --git a/webviz_ert/models/parameter_model.py b/webviz_ert/models/parameter_model.py index a7d7c8b1..dc221dd3 100644 --- a/webviz_ert/models/parameter_model.py +++ b/webviz_ert/models/parameter_model.py @@ -28,9 +28,12 @@ def __init__(self, **kwargs: Any): self._data_loader = get_data_loader(self._project_id) def data_df(self) -> pd.DataFrame: + realization = [] if self._realizations is None else self._realizations if self._data_df.empty: _data_df = self._data_loader.get_ensemble_parameter_data( - ensemble_id=self._ensemble_id, parameter_name=self.key + ensemble_id=self._ensemble_id, + parameter_name=self.key, + realizations=realization, ) if _data_df is not None: _data_df = _data_df.transpose()