Including stratigraphy in the well completion plugin

tests/unit_tests/plugin_tests/test_well_completions.py

@@ -0,0 +1,91 @@
+import pytest
+
+from webviz_subsurface.plugins._well_completions import extract_stratigraphy
+from webviz_subsurface._datainput.well_completions import remove_invalid_colors
+
+
+def test_remove_invalid_colors():
+    """Tests that colors that are not 6 digit hexadecimal are removed from the lyr
+    parse zone list
+    """
+    zonelist = [
+        {"name": "ZoneA", "color": "#FFFFFF"},
+        {"name": "ZoneB", "color": "#FFF"},
+    ]
+    assert remove_invalid_colors(zonelist) == [
+        {"name": "ZoneA", "color": "#FFFFFF"},
+        {
+            "name": "ZoneB",
+        },
+    ]
+
+
+def test_extract_stratigraphy():
+    """Checks that the merging of the layer_zone_mapping and the stratigraphy is
+    correct and that the colors are added following the correct prioritization
+    rules
+    """
+    layer_zone_mapping = {1: "ZoneA.1", 2: "ZoneA.2", 3: "ZoneB.1"}
+    stratigraphy = [
+        {
+            "name": "ZoneA",
+            "color": "#000000",
+            "subzones": [{"name": "ZoneA.1"}, {"name": "ZoneA.2"}, {"name": "ZoneA.3"}],
+        },
+        {
+            "name": "ZoneB",
+            "subzones": [{"name": "ZoneB.1", "subzones": [{"name": "ZoneB.1.1"}]}],
+        },
+        {
+            "name": "ZoneC",
+        },
+    ]
+    zone_color_mapping = {"ZoneA": "#111111", "ZoneA.1": "#222222"}
+    theme_colors = ["#FFFFFF"]
+
+    result = extract_stratigraphy(
+        layer_zone_mapping, stratigraphy, zone_color_mapping, theme_colors
+    )
+    assert result == [
+        {
+            "name": "ZoneA",
+            "color": "#000000",  # colors given in the stratigraphy has priority
+            "subzones": [
+                {
+                    "name": "ZoneA.1",
+                    "color": "#222222",  # color from zone_color_mapping
+                },
+                {"name": "ZoneA.2", "color": "#FFFFFF"}  # color from theme_colors
+                # ZoneA.3 is not here because it is not in the layer_zone_mapping
+            ],
+        },
+        {
+            "name": "ZoneB",
+            "color": "#808080",  # Since it's not a leaf, color is set to grey
+            "subzones": [
+                {
+                    "name": "ZoneB.1",
+                    "color": "#FFFFFF"  # color from theme_colors
+                    # No subzones here because ZoneB.1.1 is not in the layer_zone_mapping
+                }
+            ],
+        },
+        # ZoneC is removed since it's not in the layer_zone_mapping
+    ]
+
+
+def test_extract_stratigraphy_errors():
+    """Checks that a ValueError is raised when a Zone is in the layer_zone_mapping, but
+    not in the stratigraphy.
+    """
+    layer_zone_mapping = {1: "ZoneA", 2: "ZoneB", 3: "ZoneD"}
+    stratigraphy = [
+        {
+            "name": "ZoneA",
+        },
+        {
+            "name": "ZoneB",
+        },
+    ]
+    with pytest.raises(ValueError):
+        extract_stratigraphy(layer_zone_mapping, stratigraphy, {}, [])

webviz_subsurface/_datainput/well_completions.py

@@ -1,21 +1,47 @@
-from typing import Optional, Dict
+from typing import Optional, Dict, List, Tuple, Any
 import json
+import re
 from pathlib import Path
 import glob
+import logging
 
-import ecl2df
+from ecl2df import common
+
+
+def remove_invalid_colors(zonelist: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Removes colors in the zonelist from the lyr file that is not 6 digit
+    hexadecimal.
+    """
+    for zonedict in zonelist:
+        if "color" in zonedict and not re.match(
+            "^#([A-Fa-f0-9]{6})", zonedict["color"]
+        ):
+            logging.warning(
+                f"""The zone color {zonedict["color"]} will be ignored. """
+                "Only 6 digit hexadecimal colors are accepted in the well completions plugin."
+            )
+            zonedict.pop("color")
+    return zonelist
 
 
 def read_zone_layer_mapping(
     ensemble_path: str, zone_layer_mapping_file: str
-) -> Optional[Dict[int, str]]:
+) -> Tuple[Optional[Dict[int, str]], Optional[Dict[str, str]]]:
     """Searches for a zone layer mapping file (lyr format) on the scratch disk. \
     If one file is found it is parsed using functionality from the ecl2df \
     library.
     """
     for filename in glob.glob(f"{ensemble_path}/{zone_layer_mapping_file}"):
-        return ecl2df.EclFiles("").get_zonemap(filename=filename)
-    return None
+        zonelist = common.parse_lyrfile(filename=filename)
+        layer_zone_mapping = common.convert_lyrlist_to_zonemap(zonelist)
+        zonelist = remove_invalid_colors(zonelist)
+        zone_color_mapping = {
+            zonedict["name"]: zonedict["color"]
+            for zonedict in zonelist
+            if "color" in zonedict
+        }
+        return layer_zone_mapping, zone_color_mapping
+    return None, None
 
 
 def read_well_attributes(
@@ -61,3 +87,34 @@ def read_well_attributes(
             for well_data in file_content["wells"]
         }
     return None
+
+
+def read_stratigraphy(
+    ensemble_path: str, stratigraphy_file: str
+) -> Optional[List[Dict]]:
+    """Searches for a stratigraphy json file on the scratch disk. \
+    If a file is found the content is returned as a list of dicts.
+    """
+    for filename in glob.glob(f"{ensemble_path}/{stratigraphy_file}"):
+        return json.loads(Path(filename).read_text())
+    return None
+
+
+def get_ecl_unit_system(ensemble_path: str) -> Optional[str]:
+    """Returns the unit system of an eclipse deck. The options are \
+    METRIC, FIELD, LAB and PVT-M.
+
+    If none of these are found, the function returns None, even though \
+    the default unit system is METRIC. This is because the unit system \
+    keyword could be in an include file.
+    """
+    unit_systems = ["METRIC", "FIELD", "LAB", "PVT-M"]
+    for filename in glob.glob(f"{ensemble_path}/eclipse/model/*.DATA"):
+        with open(filename, "r") as handle:
+            ecl_data_lines = handle.readlines()
+
+        for unit_system in unit_systems:
+            if any(line.startswith(unit_system) for line in ecl_data_lines):
+                return unit_system
+        return None
+    return None