Implement consistent handling of ensembles/realizations

webviz_subsurface/plugins/_co2_migration/_co2volume.py

@@ -0,0 +1,43 @@
+import pandas
+import numpy as np
+import pathlib
+from typing import Dict
+
+
+def _read_co2_volumes(realization_paths: Dict[str, str]):
+    records = []
+    for rz_name, rz_path in realization_paths.items():
+        try:
+            df = pandas.read_csv(
+                pathlib.Path(rz_path) / "share" / "results" / "tables" / "co2_volumes.csv",
+            )
+        except FileNotFoundError:
+            continue
+        last = df.iloc[np.argmax(df["date"])]
+        label = str(rz_name)
+        records += [
+            (label, last["co2_inside"], "inside", 0.0),
+            (label, last["co2_outside"], "outside", last["co2_outside"]),
+        ]
+    df = pandas.DataFrame.from_records(records, columns=["ensemble", "volume", "containment", "volume_outside"])
+    df.sort_values("volume_outside", inplace=True, ascending=False)
+    return df
+
+
+def generate_co2_volume_figure(realization_paths: Dict[str, str], height):
+    import plotly.express as px
+    df = _read_co2_volumes(realization_paths)
+    fig = px.bar(df, y="ensemble", x="volume", color="containment", title="End-state CO2 containment [m³]", orientation="h")
+    fig.layout.height = height
+    fig.layout.legend.title.text = ""
+    fig.layout.legend.orientation = "h"
+    fig.layout.yaxis.title = ""
+    fig.layout.yaxis.tickangle = -90
+    fig.layout.xaxis.exponentformat = "power"
+    fig.layout.xaxis.title = ""
+    fig.layout.paper_bgcolor = "rgba(0,0,0,0)"
+    fig.layout.margin.b = 10
+    fig.layout.margin.t = 60
+    fig.layout.margin.l = 10
+    fig.layout.margin.r = 10
+    return fig

webviz_subsurface/plugins/_co2_migration/_utils.py

@@ -1,51 +1,19 @@
-import pandas
-import pathlib
-import numpy as np
-from typing import Optional
+from typing import Dict
 from enum import Enum
+from fmu.ensemble import ScratchEnsemble
 
 
 FAULT_POLYGON_ATTRIBUTE = "dl_extracted_faultlines"
 
 
 class MapAttribute(Enum):
-    # TODO: change to upper case
-    MigrationTime = "migration-time"
-    MaxSaturation = "max-saturation"
+    MIGRATION_TIME = "migration-time"
+    MAX_SATURATION = "max-saturation"
 
 
-def _read_co2_volumes(ensemble_paths):
-    records = []
-    for ens in ensemble_paths:
-        try:
-            df = pandas.read_csv(
-                pathlib.Path(ens) / "share" / "results" / "tables" / "co2_volumes.csv",
-            )
-        except FileNotFoundError:
-            continue
-        last = df.iloc[np.argmax(df["date"])]
-        label = pathlib.Path(ens).name
-        records += [
-            (label, last["co2_inside"], "inside"),
-            (label, last["co2_outside"], "outside"),
-        ]
-    return pandas.DataFrame.from_records(records, columns=["ensemble", "volume", "containment"])
-
-
-def generate_co2_volume_figure(ensemble_paths, height):
-    import plotly.express as px
-    df = _read_co2_volumes(ensemble_paths)
-    fig = px.bar(df, y="ensemble", x="volume", color="containment", title="End-state CO2 containment [m³]", orientation="h")
-    fig.layout.height = height
-    fig.layout.legend.title.text = ""
-    fig.layout.legend.orientation = "h"
-    fig.layout.yaxis.title = ""
-    fig.layout.yaxis.tickangle = -90
-    fig.layout.xaxis.exponentformat = "power"
-    fig.layout.xaxis.title = ""
-    fig.layout.paper_bgcolor = "rgba(0,0,0,0)"
-    fig.layout.margin.b = 10
-    fig.layout.margin.t = 60
-    fig.layout.margin.l = 10
-    fig.layout.margin.r = 10
-    return fig
+def realization_paths(ensemble_path) -> Dict[str, str]:
+    scratch = ScratchEnsemble("tmp", paths=ensemble_path)
+    return {
+        r: s.runpath()
+        for r, s in scratch.realizations.items()
+    }

webviz_subsurface/plugins/_co2_migration/callbacks.py

@@ -1,10 +1,8 @@
 import json
 import dash
+from typing import Callable, Dict, List, Optional
 from dash import callback, Output, Input, State
 from dash.exceptions import PreventUpdate
-from typing import Callable, Dict, List, Optional
-from ._utils import MapAttribute, FAULT_POLYGON_ATTRIBUTE
-from .layout import LayoutElements
 # TODO: tmp?
 from webviz_subsurface.plugins._map_viewer_fmu._tmp_well_pick_provider import (
     WellPickProvider,
@@ -28,17 +26,41 @@
     SurfaceAddress,
     SurfaceServer,
 )
+from ._utils import MapAttribute, FAULT_POLYGON_ATTRIBUTE, realization_paths
+from ._co2volume import generate_co2_volume_figure
+from .layout import LayoutElements, LayoutStyle
 
 
 def plugin_callbacks(
     get_uuid: Callable,
+    ensemble_paths: Dict[str, str],  # TODO: To be replaced by table provider or similar
     ensemble_surface_providers: Dict[str, EnsembleSurfaceProvider],
     surface_server: SurfaceServer,
     ensemble_fault_polygons_providers: Dict[str, EnsembleFaultPolygonsProvider],
     fault_polygons_server: FaultPolygonsServer,
     license_boundary_file: Optional[str],
     well_pick_provider: Optional[WellPickProvider],
 ):
+    @callback(
+        Output(get_uuid(LayoutElements.REALIZATIONINPUT), "options"),
+        Output(get_uuid(LayoutElements.REALIZATIONINPUT), "value"),
+        Output(get_uuid(LayoutElements.ENSEMBLEBARPLOT), "figure"),
+        Input(get_uuid(LayoutElements.ENSEMBLEINPUT), "value"),
+    )
+    def set_ensemble(ensemble):
+        rz_paths = realization_paths(ensemble_paths[ensemble])
+        realizations = [
+            dict(label=r, value=r)
+            for r in sorted(rz_paths.keys())
+        ]
+        # TODO: get realization names elsewhere?
+        # TODO: volumes should probably be read through a table provider or similar instead
+        fig = generate_co2_volume_figure(
+            rz_paths,
+            LayoutStyle.ENSEMBLEBARPLOTHEIGHT,
+        )
+        return realizations, realizations[0]["value"], fig
+
     # TODO: Verify optional parameters behave correctly when not provided
     # TODO: sync zone/horizon names across data types?
     @callback(
@@ -48,14 +70,17 @@ def plugin_callbacks(
         Output(get_uuid(LayoutElements.WELLPICKZONEINPUT), 'options'),
         Output(get_uuid(LayoutElements.MAPZONEINPUT), 'options'),
         Input(get_uuid(LayoutElements.ENSEMBLEINPUT), 'value'),
+        Input(get_uuid(LayoutElements.REALIZATIONINPUT), 'value'),
         Input(get_uuid(LayoutElements.PROPERTY), 'value'),
     )
-    def set_ensemble(ensemble, prop):
+    def set_realization(ensemble, realization, prop):
+        if realization is None:
+            raise PreventUpdate
         if ensemble is None:
-            return [], [], [], []
+            return [], [], [], [], []
         # Dates
         surface_provider = ensemble_surface_providers[ensemble]
-        date_list = surface_provider.surface_dates_for_attribute(MapAttribute.MaxSaturation.value)
+        date_list = surface_provider.surface_dates_for_attribute(MapAttribute.MAX_SATURATION.value)
         if date_list is None:
             dates = {}
             initial_date = dash.no_update
@@ -104,12 +129,13 @@ def set_ensemble(ensemble, prop):
         Input(get_uuid(LayoutElements.FAULTPOLYGONINPUT), "value"),
         Input(get_uuid(LayoutElements.WELLPICKZONEINPUT), "value"),
         Input(get_uuid(LayoutElements.MAPZONEINPUT), "value"),
+        Input(get_uuid(LayoutElements.REALIZATIONINPUT), "value"),
         State(get_uuid(LayoutElements.ENSEMBLEINPUT), "value"),
     )
-    def update_map_attribute(attribute, date, polygon_name, well_pick_horizon, surface_name, ensemble):
+    def update_map_attribute(attribute, date, polygon_name, well_pick_horizon, surface_name, realization, ensemble):
         if ensemble is None:
             raise PreventUpdate
-        if MapAttribute(attribute) == MapAttribute.MaxSaturation and date is None:
+        if MapAttribute(attribute) == MapAttribute.MAX_SATURATION and date is None:
             raise PreventUpdate
         date = str(date)
         if surface_name is None:
@@ -118,10 +144,10 @@ def update_map_attribute(attribute, date, polygon_name, well_pick_horizon, surfa
         layer_model, viewport_bounds = create_layer_model(
             surface_server=surface_server,
             surface_provider=ensemble_surface_providers[ensemble],
-            colormap_address=_derive_colormap_address(surface_name, attribute, date),
+            colormap_address=_derive_colormap_address(surface_name, attribute, date, realization),
             fault_polygons_server=fault_polygons_server,
             polygon_provider=ensemble_fault_polygons_providers[ensemble],
-            polygon_address=_derive_fault_polygon_address(polygon_name),
+            polygon_address=_derive_fault_polygon_address(polygon_name, realization),
             license_boundary_file=license_boundary_file,
             well_pick_provider=well_pick_provider,
             well_pick_horizon=well_pick_horizon,
@@ -216,31 +242,31 @@ def _parse_polygon_file(filename: str):
     return as_geojson
 
 
-def _derive_colormap_address(surface_name: str, attribute: str, date):
+def _derive_colormap_address(surface_name: str, attribute: str, date, realization: int):
     attribute = MapAttribute(attribute)
-    if attribute == MapAttribute.MigrationTime:
+    if attribute == MapAttribute.MIGRATION_TIME:
         return SimulatedSurfaceAddress(
-            attribute=MapAttribute.MigrationTime.value,
+            attribute=MapAttribute.MIGRATION_TIME.value,
             name=surface_name,
             datestr=None,
-            realization=0,  # TODO
+            realization=realization,
         )
-    elif attribute == MapAttribute.MaxSaturation:
+    elif attribute == MapAttribute.MAX_SATURATION:
         return SimulatedSurfaceAddress(
-            attribute=MapAttribute.MaxSaturation.value,
+            attribute=MapAttribute.MAX_SATURATION.value,
             name=surface_name,
             datestr=date,
-            realization=0,  # TODO
+            realization=realization,
         )
     else:
         raise NotImplementedError
 
 
-def _derive_fault_polygon_address(polygon_name):
+def _derive_fault_polygon_address(polygon_name, realization):
     return SimulatedFaultPolygonsAddress(
         attribute=FAULT_POLYGON_ATTRIBUTE,
         name=polygon_name,
-        realization=0,
+        realization=realization,
     )
 
 

webviz_subsurface/plugins/_co2_migration/co2_migration.py

@@ -1,7 +1,6 @@
-import pathlib
+from typing import List, Optional
 import pandas
 from dash import Dash
-from typing import Dict, List, Optional
 from webviz_config import WebvizPluginABC, WebvizSettings
 from webviz_subsurface._providers import (
     EnsembleFaultPolygonsProviderFactory,
@@ -48,12 +47,12 @@ def layout(self):
         return main_layout(
             get_uuid=self.uuid,
             ensembles=list(self._ensemble_surface_providers.keys()),
-            ensemble_paths=self._ensemble_paths,
         )
 
     def set_callbacks(self):
         plugin_callbacks(
             get_uuid=self.uuid,
+            ensemble_paths=self._ensemble_paths,
             ensemble_surface_providers=self._ensemble_surface_providers,
             surface_server=self._surface_server,
             ensemble_fault_polygons_providers=self._ensemble_fault_polygons_providers,

webviz_subsurface/plugins/_co2_migration/layout.py

@@ -1,10 +1,10 @@
-import json
 from typing import Callable, List, Any, Dict
+from enum import unique, Enum
+import plotly.graph_objects as go
 from dash import html, dcc
 import webviz_core_components as wcc
 from webviz_subsurface_components import DeckGLMap
-from enum import unique, Enum
-from ._utils import MapAttribute, generate_co2_volume_figure
+from ._utils import MapAttribute
 
 
 @unique
@@ -18,6 +18,7 @@ class LayoutElements(str, Enum):
 
     PROPERTY = "property"
     ENSEMBLEINPUT = "ensembleinput"
+    REALIZATIONINPUT = "realizationinput"
     DATEINPUT = "dateinput"
     FAULTPOLYGONINPUT = "faultpolygoninput"
     WELLPICKZONEINPUT = "wellpickzoneinput"
@@ -42,14 +43,14 @@ class LayoutStyle:
     }
 
 
-def main_layout(get_uuid: Callable, ensembles: List[str], ensemble_paths: Dict[str, str]) -> html.Div:
+def main_layout(get_uuid: Callable, ensembles: List[str]) -> html.Div:
     return html.Div(
         [
             wcc.Frame(
                 style=LayoutStyle.SIDEBAR,
                 children=[
                     PropertySelectorLayout(get_uuid),
-                    EnsembleSelectorLayout(get_uuid, ensembles, ensemble_paths),
+                    EnsembleSelectorLayout(get_uuid, ensembles),
                     DateSelectorLayout(get_uuid),
                     ZoneSelectorLayout(get_uuid),
                 ]
@@ -90,30 +91,34 @@ def __init__(self, children: List[Any]) -> None:
 
 class PropertySelectorLayout(html.Div):
     def __init__(self, get_uuid: Callable):
-        super().__init__(children=[
-            html.Div(
-                [
-                    wcc.Dropdown(
-                        label="Property",
-                        id=get_uuid(LayoutElements.PROPERTY),
-                        options=[
-                            dict(label=m.value, value=m.value)
-                            for m in MapAttribute
-                        ],
-                        value=MapAttribute.MigrationTime,
-                        clearable=False,
-                    )
-                ] 
-            )
-        ])
+        super().__init__(children=wcc.Selectors(
+            label="Property",
+            open_details=True,
+            children=[
+                html.Div(
+                    [
+                        wcc.Dropdown(
+                            id=get_uuid(LayoutElements.PROPERTY),
+                            options=[
+                                dict(label=m.value, value=m.value)
+                                for m in MapAttribute
+                            ],
+                            value=MapAttribute.MIGRATION_TIME,
+                            clearable=False,
+                        )
+                    ] 
+                )
+            ]
+        ))
 
 
 class EnsembleSelectorLayout(html.Div):
-    def __init__(self, get_uuid: Callable, ensembles: List[str], ensemble_paths: Dict[str, str]):
+    def __init__(self, get_uuid: Callable, ensembles: List[str]):
         super().__init__(children=wcc.Selectors(
             label="Ensemble",
             open_details=True,
             children=[
+                "Ensemble",
                 wcc.Dropdown(
                     id=get_uuid(LayoutElements.ENSEMBLEINPUT),
                     options=[
@@ -122,13 +127,14 @@ def __init__(self, get_uuid: Callable, ensembles: List[str], ensemble_paths: Dic
                     ],
                     value=ensembles[0]
                 ),
+                "Realization",
+                wcc.Dropdown(
+                    id=get_uuid(LayoutElements.REALIZATIONINPUT),
+                ),
                 # TODO: check that this does not yield an error with missing csv files
                 dcc.Graph(
                     id=get_uuid(LayoutElements.ENSEMBLEBARPLOT),
-                    figure=generate_co2_volume_figure(
-                        [ensemble_paths[ens] for ens in ensembles],
-                        LayoutStyle.ENSEMBLEBARPLOTHEIGHT,
-                    ),
+                    figure=go.Figure(),
                     config={
                         "displayModeBar": False,
                     }
@@ -169,7 +175,7 @@ def __init__(self, get_uuid: Callable):
                     wcc.Dropdown(
                         id=get_uuid(LayoutElements.WELLPICKZONEINPUT)
                     ),
-                    "Map",
+                    "Color Map",
                     wcc.Dropdown(
                         id=get_uuid(LayoutElements.MAPZONEINPUT),
                     ),