webviz-subsurface 0.2.36__py3-none-any.whl → 0.2.37__py3-none-any.whl

webviz_subsurface/plugins/_co2_leakage/_utilities/co2volume.py

@@ -9,6 +9,9 @@ import plotly.graph_objects as go
 
 from webviz_subsurface._providers import EnsembleTableProvider
 from webviz_subsurface._utils.enum_shim import StrEnum
+from webviz_subsurface.plugins._co2_leakage._utilities.containment_data_provider import (
+    ContainmentDataProvider,
+)
 from webviz_subsurface.plugins._co2_leakage._utilities.generic import (
     Co2MassScale,
     Co2VolumeScale,
@@ -41,6 +44,26 @@ _COLOR_ZONES = [
     "#34b36f",
 ]
 
+_LIGHTER_COLORS = {
+    "black": "#909090",
+    "#222222": "#909090",
+    "#00aa00": "#55ff55",
+    "#006ddd": "#6eb6ff",
+    "#dd4300": "#ff9a6e",
+    "#e91451": "#f589a8",
+    "#daa218": "#f2d386",
+    "#208eb7": "#81cde9",
+    "#84bc04": "#cdfc63",
+    "#b74532": "#e19e92",
+    "#9a89b4": "#ccc4d9",
+    "#8d30ba": "#c891e3",
+    "#256b33": "#77d089",
+    "#95704d": "#cfb7a1",
+    "#1357ca": "#7ba7f3",
+    "#f75ef0": "#fbaef7",
+    "#34b36f": "#93e0b7",
+}
+
 
 def _read_dataframe(
     table_provider: EnsembleTableProvider,
@@ -54,40 +77,6 @@ def _read_dataframe(
     return df
 
 
-def read_menu_options(
-    table_provider: EnsembleTableProvider,
-    realization: int,
-    relpath: str,
-) -> Dict[str, List[str]]:
-    col_names = table_provider.column_names()
-    df = table_provider.get_column_data(col_names, [realization])
-    required_columns = ["date", "amount", "phase", "containment", "zone", "region"]
-    missing_columns = [col for col in required_columns if col not in col_names]
-    if len(missing_columns) > 0:
-        raise KeyError(
-            f"Missing expected columns {', '.join(missing_columns)} in {relpath}"
-            f" in realization {realization} (and possibly other csv-files). "
-            f"Provided files are likely from an old version of ccs-scripts."
-        )
-    zones = ["all"]
-    for zone in list(df["zone"]):
-        if zone not in zones:
-            zones.append(zone)
-    regions = ["all"]
-    for region in list(df["region"]):
-        if region not in regions:
-            regions.append(region)
-    if "free_gas" in list(df["phase"]):
-        phases = ["total", "free_gas", "trapped_gas", "aqueous"]
-    else:
-        phases = ["total", "gas", "aqueous"]
-    return {
-        "zones": zones if len(zones) > 1 else [],
-        "regions": regions if len(regions) > 1 else [],
-        "phases": phases,
-    }
-
-
 def _get_colors(num_cols: int = 3, split: str = "zone") -> List[str]:
     if split == "containment":
         return [_COLOR_HAZARDOUS, _COLOR_OUTSIDE, _COLOR_CONTAINED]
@@ -106,7 +95,7 @@ def _get_marks(num_marks: int, mark_choice: str) -> List[str]:
         return [""] * num_marks
     if mark_choice == "containment":
         return ["x", "/", ""]
-    if mark_choice in ["zone", "region"]:
+    if mark_choice in ["zone", "region", "plume_group"]:
         base_pattern = ["", "/", "x", "-", "\\", "+", "|", "."]
         if num_marks > len(base_pattern):
             base_pattern *= int(np.ceil(num_marks / len(base_pattern)))
@@ -123,7 +112,7 @@ def _get_line_types(mark_options: List[str], mark_choice: str) -> List[str]:
         return ["solid"]
     if mark_choice == "containment":
         return ["dash", "dot", "solid"]
-    if mark_choice in ["zone", "region"]:
+    if mark_choice in ["zone", "region", "plume_group"]:
         if len(mark_options) > 8:
             warnings.warn(
                 f"Large number of {mark_choice}s might make it hard "
@@ -167,6 +156,93 @@ def _prepare_pattern_and_color_options(
     return cat_ord, colors, marks
 
 
+def _prepare_pattern_and_color_options_statistics_plot(
+    df: pd.DataFrame,
+    containment_info: Dict,
+    color_choice: str,
+    mark_choice: str,
+) -> Tuple[Dict, List, List]:
+    mark_options = [] if mark_choice == "none" else containment_info[f"{mark_choice}s"]
+    color_options = containment_info[f"{color_choice}s"]
+    num_colors = len(color_options)
+    num_marks = num_colors if mark_choice == "none" else len(mark_options)
+    line_types = _get_line_types(mark_options, mark_choice)
+    colors = _get_colors(num_colors, color_choice)
+
+    if mark_choice == "phase":
+        mark_options = ["total"] + mark_options
+        line_types = ["solid"] + line_types
+        num_marks += 1
+    if color_choice == "containment":
+        color_options = ["total"] + color_options
+        colors = ["black"] + colors
+        num_colors += 1
+
+    filter_mark = mark_choice != "phase"
+    filter_color = color_choice not in ["phase", "containment"]
+    _filter_rows(df, color_choice, mark_choice, filter_mark, filter_color)
+
+    if mark_choice == "none":
+        cat_ord = {"type": color_options}
+        df["type"] = df[color_choice]
+        return cat_ord, colors, line_types
+    df["type"] = [", ".join((c, m)) for c, m in zip(df[color_choice], df[mark_choice])]
+
+    if containment_info["sorting"] == "color":
+        cat_ord = {
+            "type": [", ".join((c, m)) for c in color_options for m in mark_options],
+        }
+        colors = [c for c in colors for _ in range(num_marks)]
+        line_types = line_types * num_colors
+    else:
+        cat_ord = {
+            "type": [", ".join((c, m)) for m in mark_options for c in color_options],
+        }
+        colors = colors * num_marks
+        line_types = [m for m in line_types for _ in range(num_colors)]
+
+    for m in mark_options + ["total", "all"]:
+        df["type"] = df["type"].replace(f"total, {m}", m)
+        df["type"] = df["type"].replace(f"all, {m}", m)
+    for m in color_options:
+        df["type"] = df["type"].replace(f"{m}, total", m)
+        df["type"] = df["type"].replace(f"{m}, all", m)
+    cat_ord["type"] = [
+        label.replace("total, ", "") if "total, " in label else label
+        for label in cat_ord["type"]
+    ]
+    cat_ord["type"] = [
+        label.replace("all, ", "") if "all, " in label else label
+        for label in cat_ord["type"]
+    ]
+    cat_ord["type"] = [
+        label.replace(", total", "") if ", total" in label else label
+        for label in cat_ord["type"]
+    ]
+    cat_ord["type"] = [
+        label.replace(", all", "") if ", all" in label else label
+        for label in cat_ord["type"]
+    ]
+
+    return cat_ord, colors, line_types
+
+
+def _find_default_option_statistics_figure(
+    df: pd.DataFrame, categories: List[str]
+) -> str:
+    if "hazardous" in categories:
+        default_option = "hazardous"
+    else:
+        max_value = -999.9
+        default_option = categories[0]
+        for category in categories:
+            df_filtered = df[df["type"] == category]
+            if df_filtered["amount"].max() > max_value:
+                max_value = df_filtered["amount"].max()
+                default_option = category
+    return default_option
+
+
 def _prepare_line_type_and_color_options(
     df: pd.DataFrame,
     containment_info: Dict,
@@ -180,6 +256,7 @@ def _prepare_line_type_and_color_options(
     num_colors = len(color_options)
     line_types = _get_line_types(mark_options, mark_choice)
     colors = _get_colors(num_colors, color_choice)
+
     filter_mark = True
     if mark_choice == "phase":
         mark_options = ["total"] + mark_options
@@ -225,22 +302,8 @@ def _prepare_line_type_and_color_options(
     return options
 
 
-def _find_scale_factor(
-    table_provider: EnsembleTableProvider,
-    scale: Union[Co2MassScale, Co2VolumeScale],
-) -> float:
-    if scale in (Co2MassScale.KG, Co2VolumeScale.CUBIC_METERS):
-        return 1.0
-    if scale in (Co2MassScale.MTONS, Co2VolumeScale.BILLION_CUBIC_METERS):
-        return 1e9
-    if scale in (Co2MassScale.NORMALIZE, Co2VolumeScale.NORMALIZE):
-        df = table_provider.get_column_data(table_provider.column_names())
-        return df["total"].max()
-    return 1.0
-
-
 def _read_terminal_co2_volumes(
-    table_provider: EnsembleTableProvider,
+    table_provider: ContainmentDataProvider,
     realizations: List[int],
     scale: Union[Co2MassScale, Co2VolumeScale],
     containment_info: Dict[str, Union[str, None, List[str]]],
@@ -258,11 +321,10 @@ def _read_terminal_co2_volumes(
     records[color_choice] = []
     if mark_choice != "none":
         records[mark_choice] = []
-    scale_factor = _find_scale_factor(table_provider, scale)
     data_frame = None
     for real in realizations:
-        df = _read_dataframe(table_provider, real, scale_factor)
-        df = df[df["date"] == np.max(df["date"])]
+        df = table_provider.extract_dataframe(real, scale)
+        df = df[df["date"] == containment_info["date_option"]]
         _add_sort_key_and_real(df, str(real), containment_info)
         _filter_columns(df, color_choice, mark_choice, containment_info)
         _filter_rows(df, color_choice, mark_choice)
@@ -285,7 +347,7 @@ def _filter_columns(
 ) -> None:
     filter_columns = [
         col
-        for col in ["phase", "containment", "zone", "region"]
+        for col in ["phase", "containment", "zone", "region", "plume_group"]
         if col not in [mark_choice, color_choice]
     ]
     for col in filter_columns:
@@ -298,8 +360,10 @@ def _filter_rows(
     color_choice: str,
     mark_choice: str,
     filter_mark: bool = True,
+    filter_color: bool = True,
 ) -> None:
-    df.query(f'{color_choice} not in ["total", "all"]', inplace=True)
+    if filter_color:
+        df.query(f'{color_choice} not in ["total", "all"]', inplace=True)
     if mark_choice != "none" and filter_mark:
         df.query(f'{mark_choice} not in ["total", "all"]', inplace=True)
 
@@ -315,6 +379,7 @@ def _add_sort_key_and_real(
             & (df["containment"] == "hazardous")
             & (df["zone"] == containment_info["zone"])
             & (df["region"] == containment_info["region"])
+            & (df["plume_group"] == containment_info["plume_group"])
         ]["amount"]
     )
     sort_value_secondary = np.sum(
@@ -323,6 +388,7 @@ def _add_sort_key_and_real(
             & (df["containment"] == "outside")
             & (df["zone"] == containment_info["zone"])
             & (df["region"] == containment_info["region"])
+            & (df["plume_group"] == containment_info["plume_group"])
         ]["amount"]
     )
     df["real"] = [label] * df.shape[0]
@@ -331,14 +397,13 @@ def _add_sort_key_and_real(
 
 
 def _read_co2_volumes(
-    table_provider: EnsembleTableProvider,
+    table_provider: ContainmentDataProvider,
     realizations: List[int],
     scale: Union[Co2MassScale, Co2VolumeScale],
 ) -> pd.DataFrame:
-    scale_factor = _find_scale_factor(table_provider, scale)
     return pd.concat(
         [
-            _read_dataframe(table_provider, r, scale_factor).assign(realization=r)
+            table_provider.extract_dataframe(r, scale).assign(realization=r)
             for r in realizations
         ]
     )
@@ -357,15 +422,21 @@ def _change_names(
         df["name"] = df["name"].replace(f"{m}, all", m)
 
 
-def _adjust_figure(fig: go.Figure) -> None:
+def _adjust_figure(fig: go.Figure, plot_title: Optional[str] = None) -> None:
     fig.layout.legend.orientation = "v"
     fig.layout.legend.title.text = ""
     fig.layout.legend.itemwidth = 40
     fig.layout.xaxis.exponentformat = "power"
-    fig.layout.title.x = 0.5
+    if plot_title is not None:
+        fig.layout.title.text = plot_title
+        fig.layout.title.font = {"size": 14}
+        fig.layout.margin.t = 40
+        fig.layout.title.y = 0.95
+    else:
+        fig.layout.margin.t = 15
+    fig.layout.title.x = 0.4
     fig.layout.paper_bgcolor = "rgba(0,0,0,0)"
     fig.layout.margin.b = 6
-    fig.layout.margin.t = 15
     fig.layout.margin.l = 10
     fig.layout.margin.r = 10
     fig.update_layout(
@@ -402,7 +473,7 @@ def _add_prop_to_df(
 
 
 def generate_co2_volume_figure(
-    table_provider: EnsembleTableProvider,
+    table_provider: ContainmentDataProvider,
     realizations: List[int],
     scale: Union[Co2MassScale, Co2VolumeScale],
     containment_info: Dict[str, Any],
@@ -429,17 +500,21 @@ def generate_co2_volume_figure(
         pattern_shape_sequence=marks,
         orientation="h",
         category_orders=cat_ord,
-        hover_data={"prop": True, "real": False},
+        custom_data=["type", "prop"],
+    )
+    fig.update_traces(
+        hovertemplate="Type: %{customdata[0]}<br>Amount: %{x:.3f}<br>"
+        "Realization: %{y}<br>Proportion: %{customdata[1]}<extra></extra>",
     )
     fig.layout.yaxis.title = "Realization"
     fig.layout.xaxis.title = scale.value
-    _adjust_figure(fig)
+    _adjust_figure(fig, plot_title=containment_info["date_option"])
     return fig
 
 
 # pylint: disable=too-many-locals
 def generate_co2_time_containment_one_realization_figure(
-    table_provider: EnsembleTableProvider,
+    table_provider: ContainmentDataProvider,
     scale: Union[Co2MassScale, Co2VolumeScale],
     time_series_realization: int,
     y_limits: List[Optional[float]],
@@ -477,10 +552,11 @@ def generate_co2_time_containment_one_realization_figure(
         pattern_shape_sequence=marks,
         category_orders=cat_ord,
         range_y=y_limits,
-        hover_data={
-            "prop": True,
-            "amount": ":.3f",
-        },
+        custom_data=["type", "prop"],
+    )
+    fig.update_traces(
+        hovertemplate="Type: %{customdata[0]}<br>Date: %{x}<br>"
+        "Amount: %{y:.3f}<br>Proportion: %{customdata[1]}<extra></extra>",
     )
     _add_hover_info_in_field(fig, df, cat_ord, colors)
     fig.layout.yaxis.range = y_limits
@@ -535,15 +611,15 @@ def _add_hover_info_in_field(
             p15 = prev_val + 0.15 * amount
             p85 = prev_val + 0.85 * amount
             y_vals = np.linspace(p15, p85, 8).tolist() * len(date_dict[date])
-            y_vals.sort()
+            y_vals.sort()  # type: ignore[attr-defined]
             fig.add_trace(
                 go.Scatter(
                     x=date_dict[date] * 8,
                     y=y_vals,
                     mode="lines",
                     line=go.scatter.Line(color=color),
-                    text=f"type={name}<br>date={date_strings[date]}<br>"
-                    f"amount={amount:.3f}<br>prop={prop}",
+                    text=f"Type: {name}<br>Date: {date_strings[date]}<br>"
+                    f"Amount: {amount:.3f}<br>Proportion: {prop}",
                     opacity=0,
                     hoverinfo="text",
                     hoveron="points",
@@ -553,9 +629,82 @@ def _add_hover_info_in_field(
             prev_vals[date] = prev_val + amount
 
 
+def _connect_plume_groups(
+    df: pd.DataFrame,
+    color_choice: str,
+    mark_choice: str,
+) -> None:
+    col_list = ["realization"]
+    if color_choice == "plume_group" and mark_choice != "none":
+        col_list.append(mark_choice)
+    elif mark_choice == "plume_group":
+        col_list.append(color_choice)
+
+    cols: Union[List[str], str] = col_list
+    if len(col_list) == 1:
+        cols = col_list[0]
+    # Find points where plumes start or end, to connect the lines
+    end_points = []
+    start_points = []
+    for plume_name, df_sub in df.groupby("plume_group"):
+        if plume_name == "undetermined":
+            continue
+        for _, df_sub2 in df_sub.groupby(cols):
+            # Assumes the data frame is sorted on date
+            mask_end = (
+                (df_sub2["amount"] == 0.0)
+                & (df_sub2["amount"].shift(1) > 0.0)
+                & (df_sub2.index > 0)
+            )
+            mask_start = (
+                (df_sub2["amount"] > 0.0)
+                & (df_sub2["amount"].shift(1) == 0.0)
+                & (df_sub2.index > 0)
+            )
+            first_index_end = mask_end.idxmax() if mask_end.any() else None
+            first_index_start = mask_start.idxmax() if mask_start.any() else None
+            transition_row_end = (
+                df_sub2.loc[first_index_end] if first_index_end is not None else None
+            )
+            transition_row_start = (
+                df_sub2.loc[first_index_start]
+                if first_index_start is not None
+                else None
+            )
+            if transition_row_end is not None:
+                end_points.append(transition_row_end)
+                # Replace 0 with np.nan for all dates after this
+                date = str(transition_row_end["date"])
+                df.loc[
+                    (df["plume_group"] == plume_name)
+                    & (df["amount"] == 0.0)
+                    & (df["date"] > date),
+                    "amount",
+                ] = np.nan
+            if transition_row_start is not None:
+                start_points.append(transition_row_start)
+    for end_point in end_points:
+        plume1 = end_point["plume_group"]
+        row1 = end_point.drop(["amount", "plume_group", "name"])
+        for start_point in start_points:
+            plume2 = start_point["plume_group"]
+            if plume1 in plume2 and len(plume1) < len(plume2):
+                row2 = start_point.drop(["amount", "plume_group", "name"])
+                if row1.equals(row2):
+                    row_to_change = df.eq(end_point).all(axis=1)
+                    if sum(row_to_change) == 1:
+                        df.loc[row_to_change, "amount"] = start_point["amount"]
+    df["is_merged"] = ["+" in x for x in df["plume_group"].values]
+    df.loc[
+        (df["plume_group"] != "all") & (df["is_merged"]) & (df["amount"] == 0.0),
+        "amount",
+    ] = np.nan
+    df.drop(columns="is_merged", inplace=True)
+
+
 # pylint: disable=too-many-locals
 def generate_co2_time_containment_figure(
-    table_provider: EnsembleTableProvider,
+    table_provider: ContainmentDataProvider,
     realizations: List[int],
     scale: Union[Co2MassScale, Co2VolumeScale],
     containment_info: Dict[str, Any],
@@ -570,6 +719,12 @@ def generate_co2_time_containment_figure(
     active_cols_at_startup = list(
         options[options["line_type"].isin(["solid", "0px"])]["name"]
     )
+    if "plume_group" in df:
+        try:
+            _connect_plume_groups(df, color_choice, mark_choice)
+        except ValueError:
+            pass
+
     fig = go.Figure()
     # Generate dummy scatters for legend entries
     dummy_args = {"x": df["date"], "mode": "lines", "hoverinfo": "none"}
@@ -585,28 +740,65 @@ def generate_co2_time_containment_figure(
         if name not in active_cols_at_startup:
             args["visible"] = "legendonly"
         fig.add_scatter(y=[0.0], **dummy_args, **args)
-    for rlz in realizations:
-        sub_df = df[df["realization"] == rlz].copy().reset_index()
-        _add_prop_to_df(
-            sub_df, np.unique(df["date"]), "date", [color_choice, mark_choice]
+
+    hover_template = (
+        "Type: %{meta[1]}<br>Date: %{x}<br>Amount: %{y:.3f}<br>"
+        "Realization: %{meta[0]}<br>Proportion: %{customdata}"
+    )
+
+    if containment_info["use_stats"]:
+        df_no_real = df.drop(columns=["REAL", "realization"]).reset_index(drop=True)
+        if mark_choice == "none":
+            df_grouped = df_no_real.groupby(
+                ["date", "name", color_choice], as_index=False
+            )
+        else:
+            df_grouped = df_no_real.groupby(
+                ["date", "name", color_choice, mark_choice], as_index=False
+            )
+        df_mean = df_grouped.agg(np.mean)
+        df_mean["realization"] = ["mean"] * df_mean.shape[0]
+        df_p10 = df_grouped.agg(lambda x: np.quantile(x, 0.1))
+        df_p10["realization"] = ["p10"] * df_p10.shape[0]
+        df_p90 = df_grouped.agg(lambda x: np.quantile(x, 0.9))
+        df_p90["realization"] = ["p90"] * df_p90.shape[0]
+        df = (
+            pd.concat([df_mean, df_p10, df_p90])
+            .sort_values(["name", "date"])
+            .reset_index(drop=True)
+        )
+        realizations = ["p10", "mean", "p90"]  # type: ignore
+        hover_template = (
+            "Type: %{meta[1]}<br>Date: %{x}<br>Amount: %{y:.3f}<br>"
+            "Statistic: %{meta[0]}"
         )
+    for rlz in realizations:
+        lwd = 1.5 if rlz in ["p10", "p90"] else 2.5
+        sub_df = df[df["realization"] == rlz].copy().reset_index(drop=True)
+        if not containment_info["use_stats"]:
+            _add_prop_to_df(
+                sub_df, np.unique(df["date"]), "date", [color_choice, mark_choice]
+            )
         common_args = {
             "x": sub_df["date"],
-            "hovertemplate": "%{x}: %{y}<br>Realization: %{meta[0]}<br>Prop: %{customdata}%",
-            "meta": [rlz],
             "showlegend": False,
         }
         for name, color, line_type in zip(
             options["name"], options["color"], options["line_type"]
         ):
-            # NBNB-AS: Check this, mypy complains:
             args = {
                 "line_dash": line_type,
-                "marker_color": color,
+                "line_width": lwd,
+                "marker_color": (
+                    _LIGHTER_COLORS[color] if rlz in ["p10", "p90"] else color
+                ),
                 "legendgroup": name,
-                "name": name,
-                "customdata": sub_df[sub_df["name"] == name]["prop"],  # type: ignore
+                "name": "",
+                "meta": [rlz, name],
+                "hovertemplate": hover_template,
             }
+            if not containment_info["use_stats"]:
+                args["customdata"] = sub_df[sub_df["name"] == name]["prop"]
             if name not in active_cols_at_startup:
                 args["visible"] = "legendonly"
             fig.add_scatter(
@@ -618,3 +810,56 @@ def generate_co2_time_containment_figure(
     fig.layout.yaxis.autorange = True
     _adjust_figure(fig)
     return fig
+
+
+def generate_co2_statistics_figure(
+    table_provider: ContainmentDataProvider,
+    realizations: List[int],
+    scale: Union[Co2MassScale, Co2VolumeScale],
+    containment_info: Dict[str, Any],
+) -> go.Figure:
+    date_option = containment_info["date_option"]
+    df = _read_co2_volumes(table_provider, realizations, scale)
+    df = df[df["date"] == date_option]
+    df = df.drop(columns=["date"]).reset_index(drop=True)
+    color_choice = containment_info["color_choice"]
+    mark_choice = containment_info["mark_choice"]
+    _filter_columns(df, color_choice, mark_choice, containment_info)
+    cat_ord, colors, line_types = _prepare_pattern_and_color_options_statistics_plot(
+        df,
+        containment_info,
+        color_choice,
+        mark_choice,
+    )
+
+    # Remove if we want realization as label?
+    df = df.drop(columns=["REAL", "realization"]).reset_index(drop=True)
+    fig = px.ecdf(
+        df,
+        x="amount",
+        ecdfmode="reversed",
+        ecdfnorm="probability",
+        markers=True,
+        color="type",
+        color_discrete_sequence=colors,
+        line_dash="type" if mark_choice != "none" else None,
+        line_dash_sequence=line_types,
+        category_orders=cat_ord,
+    )
+
+    default_option = _find_default_option_statistics_figure(df, cat_ord["type"])
+    for trace in fig.data:
+        if trace.name != default_option:
+            trace.visible = "legendonly"
+
+    fig.update_traces(
+        hovertemplate="Type: %{data.name}<br>Amount: %{x:.3f}<br>"
+        "Probability: %{y:.3f}<extra></extra>",
+    )
+    fig.layout.yaxis.range = [-0.02, 1.02]
+    fig.layout.legend.tracegroupgap = 0
+    fig.layout.xaxis.title = scale.value
+    fig.layout.yaxis.title = "Probability"
+    _adjust_figure(fig, plot_title=containment_info["date_option"])
+
+    return fig