CCS-213: Handle plume_groups (equinor#36)

* CCS-213: Start using plume group from containment results.

* CCS-213: Cont using plume group from containment results.

* CCS-213: Fixes for plume groups.

* CCS-213: Add plume groups to filter.

* CCS-213: Some minor fixes.

* CCS-213: Minor fix.

* Some tests.

* CCS-213: Join lines before/after merge when coloring by plume groups.

* CCS-213: Some refactoring.

* CCS-213: Fix combine plume group with none as mark.

* CCS-213: Minor fix.

* CCS-213: Remove some prints.

* CCS-213: Use starting 0-values for unmerged wells.

* CCS-213: Fix for multiple realization.

* CCS-213: Fix case without plume groups.

* CCS-213: Minor fix.

* CCS-213: Fix black.

* CCS-213: Some refactoring.

* CCS-213: Minor change, remove temp code.

Author: AudunSektnanNR

webviz_subsurface/plugins/_co2_leakage/_plugin.py

@@ -291,6 +291,7 @@ def _set_callbacks(self) -> None:
                 Input(self._settings_component(ViewSettings.Ids.REGION), "value"),
                 Input(self._settings_component(ViewSettings.Ids.PHASE), "value"),
                 Input(self._settings_component(ViewSettings.Ids.CONTAINMENT), "value"),
+                Input(self._settings_component(ViewSettings.Ids.PLUME_GROUP), "value"),
                 Input(self._settings_component(ViewSettings.Ids.COLOR_BY), "value"),
                 Input(self._settings_component(ViewSettings.Ids.MARK_BY), "value"),
                 Input(self._settings_component(ViewSettings.Ids.SORT_PLOT), "value"),
@@ -309,6 +310,7 @@ def update_graphs(
                 region: Optional[str],
                 phase: str,
                 containment: str,
+                plume_group: str,
                 color_choice: str,
                 mark_choice: Optional[str],
                 sorting: str,
@@ -320,6 +322,7 @@ def update_graphs(
                     region,
                     phase,
                     containment,
+                    plume_group,
                     color_choice,
                     mark_choice,
                     sorting,

webviz_subsurface/plugins/_co2_leakage/_utilities/callbacks.py

@@ -441,6 +441,7 @@ def process_containment_info(
     region: Optional[str],
     phase: str,
     containment: str,
+    plume_group: str,
     color_choice: str,
     mark_choice: Optional[str],
     sorting: str,
@@ -450,10 +451,22 @@ def process_containment_info(
         mark_choice = "phase"
     zones = menu_options["zones"]
     regions = menu_options["regions"]
+    plume_groups = menu_options["plume_groups"]
     if len(zones) > 0:
         zones = [zone_name for zone_name in zones if zone_name != "all"]
     if len(regions) > 0:
         regions = [reg_name for reg_name in regions if reg_name != "all"]
+    if len(plume_groups) > 0:
+        plume_groups = [pg_name for pg_name in plume_groups if pg_name != "all"]
+
+        def plume_sort_key(name: str):
+            if name == "?":
+                return 999
+            else:
+                return name.count("+")
+
+        plume_groups = sorted(plume_groups, key=plume_sort_key)
+
     containments = ["hazardous", "outside", "contained"]
     phases = [phase for phase in menu_options["phases"] if phase != "total"]
     if "zone" in [mark_choice, color_choice]:
@@ -467,11 +480,13 @@ def process_containment_info(
         "regions": regions,
         "phase": phase,
         "containment": containment,
+        "plume_group": plume_group,
         "color_choice": color_choice,
         "mark_choice": mark_choice,
         "sorting": sorting,
         "phases": phases,
         "containments": containments,
+        "plume_groups": plume_groups,
     }
 
 
@@ -491,12 +506,18 @@ def set_plot_ids(
             if containment_info["region"] is not None
             else "None"
         )
+        plume_group_str = (
+            containment_info["plume_group"]
+            if containment_info["plume_group"] is not None
+            else "None"
+        )
         plot_id = "-".join(
             (
                 source,
                 scale,
                 zone_str,
                 region_str,
+                plume_group_str,
                 str(containment_info["phase"]),
                 str(containment_info["containment"]),
                 containment_info["color_choice"],

webviz_subsurface/plugins/_co2_leakage/_utilities/co2volume.py

@@ -75,7 +75,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)))
@@ -92,7 +92,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 "
@@ -239,7 +239,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:
@@ -269,6 +269,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(
@@ -277,6 +278,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]
@@ -511,6 +513,73 @@ 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):
+    cols = ["realization"]
+    if color_choice == "plume_group" and mark_choice != "none":
+        cols.append(mark_choice)
+    elif mark_choice == "plume_group":
+        cols.append(color_choice)
+    # 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 == "?":
+            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 == True, "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"] == True)
+        & (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: ContainmentDataProvider,
@@ -528,6 +597,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 Exception:
+            pass
+
     fig = go.Figure()
     # Generate dummy scatters for legend entries
     dummy_args = {"x": df["date"], "mode": "lines", "hoverinfo": "none"}

webviz_subsurface/plugins/_co2_leakage/_utilities/containment_data_provider.py

@@ -45,7 +45,11 @@ def extract_condensed_dataframe(
         co2_scale: Union[Co2MassScale, Co2VolumeScale],
     ) -> pd.DataFrame:
         df = self._provider.get_column_data(self._provider.column_names())
-        df = df.loc[(df["zone"] == "all") & (df["region"] == "all")]
+        df = df.loc[
+            (df["zone"] == "all")
+            & (df["region"] == "all")
+            & (df["plume_group"] == "all")
+        ]
         if co2_scale == Co2MassScale.MTONS:
             df.loc[:, "amount"] /= 1e9
         elif co2_scale == Co2MassScale.NORMALIZE:
@@ -82,6 +86,19 @@ def _get_menu_options(provider: EnsembleTableProvider) -> MenuOptions:
         for region in list(df["region"]):
             if region not in regions:
                 regions.append(region)
+        plume_groups = ["all"]
+        for plume_group in list(df["plume_group"]):
+            if plume_group not in plume_groups:
+                plume_groups.append(plume_group)
+
+        def plume_sort_key(name: str):
+            if name == "?":
+                return 999
+            else:
+                return name.count("+")
+
+        plume_groups = sorted(plume_groups, key=plume_sort_key)
+
         if "free_gas" in list(df["phase"]):
             phases = ["total", "free_gas", "trapped_gas", "aqueous"]
         else:
@@ -90,12 +107,21 @@ def _get_menu_options(provider: EnsembleTableProvider) -> MenuOptions:
             "zones": zones if len(zones) > 1 else [],
             "regions": regions if len(regions) > 1 else [],
             "phases": phases,
+            "plume_groups": plume_groups if len(plume_groups) > 1 else [],
         }
 
     @staticmethod
     def _validate(provider: EnsembleTableProvider) -> None:
         col_names = provider.column_names()
-        required_columns = ["date", "amount", "phase", "containment", "zone", "region"]
+        required_columns = [
+            "date",
+            "amount",
+            "phase",
+            "containment",
+            "zone",
+            "region",
+            "plume_group",
+        ]
         missing_columns = [col for col in required_columns if col not in col_names]
         realization = provider.realizations()[0]
         if len(missing_columns) == 0:

webviz_subsurface/plugins/_co2_leakage/_utilities/generic.py

@@ -177,6 +177,7 @@ class MenuOptions(TypedDict):
     zones: List[str]
     regions: List[str]
     phases: List[str]
+    plume_groups: List[str]
 
 
 class MapThresholds:

webviz_subsurface/plugins/_co2_leakage/_utilities/initialization.py

@@ -274,6 +274,7 @@ def init_dictionary_of_content(
     content["maps"] = has_maps
     content["zones"] = False
     content["regions"] = False
+    content["plume_groups"] = False
     if content["mass"] or content["volume"]:
         content["zones"] = max(
             len(inner_dict["zones"]) > 0
@@ -285,6 +286,11 @@ def init_dictionary_of_content(
             for outer_dict in menu_options.values()
             for inner_dict in outer_dict.values()
         )
+        content["plume_groups"] = max(
+            len(inner_dict["plume_groups"]) > 0
+            for outer_dict in menu_options.values()
+            for inner_dict in outer_dict.values()
+        )
     return content