Group observations by response (WIP)

src/ert/analysis/_es_update.py

@@ -295,21 +295,35 @@ def _get_obs_and_measure_data(
 
         df = filtered_response.to_dataframe().reset_index()
 
-        observation_keys.append(df["name"].to_list())
-        observation_values.append(df["observations"].to_list())
-        observation_errors.append(df["std"].to_list())
+        observation_keys.extend(
+            observation.to_dataframe()
+            .reset_index()["name"]
+            .to_numpy()
+            .reshape((len(filtered_response.name), -1))
+            .tolist()
+        )
+        observation_values.extend(
+            observation["observations"]
+            .to_dataframe()
+            .to_numpy()
+            .reshape((1, -1))
+            .tolist()
+        )
+        observation_errors.extend(
+            observation["std"].to_dataframe().to_numpy().reshape((1, -1)).tolist()
+        )
 
-        measured_data.append(
+        measured_data.extend(
             filtered_response["values"]
             .transpose(..., "realization")
-            .values.reshape((-1, len(filtered_response.realization)))
+            .values.reshape(
+                (len(filtered_response.name), -1, len(filtered_response.realization))
+            )
         )
     source_fs.load_responses.cache_clear()
 
-    # Measured_data, an array of 3 dimensions
-    # Outer dimension: One array per observation
-    # Mid dimension is ??? Sometimes length 1, sometimes nreals?
-    # Inner dimension: value is "values", index is realization
+    # Measured_data, an array of nd arrays with shape (1, nreals)
+    # Each inner array has 1 dimension containing obs key, and nreals "values"
     return (
         np.concatenate(measured_data, axis=0),
         np.concatenate(observation_values),

src/ert/dark_storage/common.py

@@ -1,5 +1,5 @@
 import contextlib
-from typing import Any, Dict, Iterator, List, Sequence, Union
+from typing import Any, Dict, Iterator, List, Sequence, Union, Callable
 from uuid import UUID
 
 import numpy as np
@@ -148,19 +148,62 @@ def data_for_key(
     return pd.DataFrame()
 
 
+def _transform_time_pandas(x):
+    return pd.Timestamp(x).isoformat()
 def get_all_observations(experiment: ExperimentReader) -> List[Dict[str, Any]]:
     observations = []
-    for key, dataset in experiment.observations.items():
-        observation = {
-            "name": key,
-            "values": list(dataset["observations"].values.flatten()),
-            "errors": list(dataset["std"].values.flatten()),
-        }
-        if "time" in dataset.coords:
-            observation["x_axis"] = _prepare_x_axis(dataset["time"].values.flatten())  # type: ignore
-        else:
-            observation["x_axis"] = _prepare_x_axis(dataset["index"].values.flatten())  # type: ignore
-        observations.append(observation)
+    for response_key, dataset in experiment.observations.items():
+        # observation = {
+        #     "name": response_key,
+        #     "values": list(dataset["observations"].values.flatten()),
+        #     "errors": list(dataset["std"].values.flatten()),
+        # }
+
+
+        fn: Callable[[str], any]
+        # x_axis_fn = None
+        x_coord_key = "time" if "time" in dataset.coords else "index"
+
+        # sample = dataset[x_coord_key].values.item()
+        # if isinstance(sample, pd.Timestamp):
+        #    x_axis_fn = _transform_time_pandas
+        # else:
+        #    x_axis_fn = str
+
+        # assert x_axis_fn is not None
+
+        for obs_name in dataset["name"].values.flatten():
+            ds_for_name = dataset.sel(name=obs_name)
+            df_for_name = ds_for_name.reset_index()
+            observations.append({
+                "name": obs_name,
+                "values": df_for_name["values"].to_list(),
+                "errors": df_for_name["std"].to_list(),
+                "x_axis": _prepare_x_axis(df_for_name[x_coord_key].to_list())
+            })
+
+        #dataset = dataset.assign(x_axis=lambda x: x_axis_fn(getattr(x,x_coord_key)))
+
+
+
+        # if "time" in dataset.coords:
+        #     #observation["x_axis"] = _prepare_x_axis(dataset["time"].values.flatten())  # type: ignore
+        #     x_coord_key = "time"
+        #     sample = dataset["time"].values.item()
+        #     if isinstance(sample, pd.Timestamp):
+        #         fn = _transform_time_pandas
+        #     else:
+        #         fn = str
+        # else:
+        #     x_coord_key = "time"
+        #     sample = dataset["index"].values.item()
+        #     if isinstance(sample, pd.Timestamp):
+        #         fn = _transform_time_pandas
+        #     else:
+        #         fn = str
+            #observation["x_axis"] = _prepare_x_axis(dataset["index"].values.flatten())  # type: ignore
+        #observations.append(observation)
+        observations.extend(dataset)
 
     observations.sort(key=lambda x: x["x_axis"])  # type: ignore
     return observations