Combine _SummaryType and SummaryKeyType

src/ert/config/_read_summary.py

@@ -23,73 +23,7 @@
 import resfo
 from pydantic import PositiveInt
 
-SPECIAL_KEYWORDS = [
-    "NEWTON",
-    "NAIMFRAC",
-    "NLINEARS",
-    "NLINSMIN",
-    "NLINSMAX",
-    "ELAPSED",
-    "MAXDPR",
-    "MAXDSO",
-    "MAXDSG",
-    "MAXDSW",
-    "STEPTYPE",
-    "WNEWTON",
-]
-
-
-class _SummaryType(Enum):
-    AQUIFER = auto()
-    BLOCK = auto()
-    COMPLETION = auto()
-    FIELD = auto()
-    GROUP = auto()
-    LOCAL_BLOCK = auto()
-    LOCAL_COMPLETION = auto()
-    LOCAL_WELL = auto()
-    NETWORK = auto()
-    SEGMENT = auto()
-    WELL = auto()
-    REGION = auto()
-    INTER_REGION = auto()
-    OTHER = auto()
-
-    @classmethod
-    def from_keyword(cls, summary_keyword: str) -> _SummaryType:
-        KEYWORD_TYPE_MAPPING = {
-            "A": cls.AQUIFER,
-            "B": cls.BLOCK,
-            "C": cls.COMPLETION,
-            "F": cls.FIELD,
-            "G": cls.GROUP,
-            "LB": cls.LOCAL_BLOCK,
-            "LC": cls.LOCAL_COMPLETION,
-            "LW": cls.LOCAL_WELL,
-            "N": cls.NETWORK,
-            "S": cls.SEGMENT,
-            "W": cls.WELL,
-        }
-        if not summary_keyword:
-            raise ValueError("Got empty summary keyword")
-        if any(special in summary_keyword for special in SPECIAL_KEYWORDS):
-            return cls.OTHER
-        if summary_keyword[0] in KEYWORD_TYPE_MAPPING:
-            return KEYWORD_TYPE_MAPPING[summary_keyword[0]]
-        if summary_keyword[0:2] in KEYWORD_TYPE_MAPPING:
-            return KEYWORD_TYPE_MAPPING[summary_keyword[0:2]]
-        if summary_keyword == "RORFR":
-            return cls.REGION
-
-        if any(
-            re.fullmatch(pattern, summary_keyword)
-            for pattern in [r"R.FT.*", r"R..FT.*", r"R.FR.*", r"R..FR.*", r"R.F"]
-        ):
-            return cls.INTER_REGION
-        if summary_keyword[0] == "R":
-            return cls.REGION
-
-        return cls.OTHER
+from ert.summary_key_type import SummaryKeyType
 
 
 def _cell_index(
@@ -128,53 +62,53 @@ def make_summary_key(
     lj: Optional[int] = None,
     lk: Optional[int] = None,
 ) -> Optional[str]:
-    sum_type = _SummaryType.from_keyword(keyword)
+    sum_type = SummaryKeyType.from_keyword(keyword)
     if sum_type in [
-        _SummaryType.FIELD,
-        _SummaryType.OTHER,
+        SummaryKeyType.FIELD,
+        SummaryKeyType.OTHER,
     ]:
         return keyword
     if sum_type in [
-        _SummaryType.REGION,
-        _SummaryType.AQUIFER,
+        SummaryKeyType.REGION,
+        SummaryKeyType.AQUIFER,
     ]:
         return f"{keyword}:{number}"
-    if sum_type == _SummaryType.BLOCK:
+    if sum_type == SummaryKeyType.BLOCK:
         nx, ny = _check_if_missing("block", "dimens", nx, ny)
         (number,) = _check_if_missing("block", "nums", number)
         i, j, k = _cell_index(number - 1, nx, ny)
         return f"{keyword}:{i},{j},{k}"
     if sum_type in [
-        _SummaryType.GROUP,
-        _SummaryType.WELL,
+        SummaryKeyType.GROUP,
+        SummaryKeyType.WELL,
     ]:
         return f"{keyword}:{name}"
-    if sum_type == _SummaryType.SEGMENT:
+    if sum_type == SummaryKeyType.SEGMENT:
         return f"{keyword}:{name}:{number}"
-    if sum_type == _SummaryType.COMPLETION:
+    if sum_type == SummaryKeyType.COMPLETION:
         nx, ny = _check_if_missing("completion", "dimens", nx, ny)
         (number,) = _check_if_missing("completion", "nums", number)
         i, j, k = _cell_index(number - 1, nx, ny)
         return f"{keyword}:{name}:{i},{j},{k}"
-    if sum_type == _SummaryType.INTER_REGION:
+    if sum_type == SummaryKeyType.INTER_REGION:
         (number,) = _check_if_missing("inter region", "nums", number)
         r1 = number % 32768
         r2 = ((number - r1) // 32768) - 10
         return f"{keyword}:{r1}-{r2}"
-    if sum_type == _SummaryType.LOCAL_WELL:
+    if sum_type == SummaryKeyType.LOCAL_WELL:
         (name,) = _check_if_missing("local well", "WGNAMES", name)
         (lgr_name,) = _check_if_missing("local well", "LGRS", lgr_name)
         return f"{keyword}:{lgr_name}:{name}"
-    if sum_type == _SummaryType.LOCAL_BLOCK:
+    if sum_type == SummaryKeyType.LOCAL_BLOCK:
         li, lj, lk = _check_if_missing("local block", "NUMLX", li, lj, lk)
         (lgr_name,) = _check_if_missing("local block", "LGRS", lgr_name)
         return f"{keyword}:{lgr_name}:{li},{lj},{lk}"
-    if sum_type == _SummaryType.LOCAL_COMPLETION:
+    if sum_type == SummaryKeyType.LOCAL_COMPLETION:
         li, lj, lk = _check_if_missing("local completion", "NUMLX", li, lj, lk)
         (name,) = _check_if_missing("local completion", "WGNAMES", name)
         (lgr_name,) = _check_if_missing("local completion", "LGRS", lgr_name)
         return f"{keyword}:{lgr_name}:{name}:{li},{lj},{lk}"
-    if sum_type == _SummaryType.NETWORK:
+    if sum_type == SummaryKeyType.NETWORK:
         (name,) = _check_if_missing("network", "WGNAMES", name)
         return f"{keyword}:{name}"
     raise ValueError(f"Unexpected keyword type: {sum_type}")

src/ert/gui/plottery/plots/ensemble.py

@@ -7,7 +7,7 @@
 
 from ert.gui.plottery.plots.history import plotHistory
 from ert.gui.tools.plot.plot_api import EnsembleObject
-from ert.shared.storage.summary_key_utils import is_rate
+from ert.summary_key_type import is_rate
 
 from .observations import plotObservations
 from .plot_tools import PlotTools

src/ert/shared/storage/summary_key_utils.py → src/ert/summary_key_type.py

@@ -1,7 +1,12 @@
+from __future__ import annotations
+
+import re
 from enum import Enum, auto
-from typing import List
+from typing import (
+    List,
+)
 
-special_keys = [
+SPECIAL_KEYWORDS = [
     "NAIMFRAC",
     "NBAKFL",
     "NBYTOT",
@@ -29,6 +34,61 @@
     "STEPTYPE",
     "WNEWTON",
 ]
+
+
+class SummaryKeyType(Enum):
+    AQUIFER = auto()
+    BLOCK = auto()
+    COMPLETION = auto()
+    FIELD = auto()
+    GROUP = auto()
+    LOCAL_BLOCK = auto()
+    LOCAL_COMPLETION = auto()
+    LOCAL_WELL = auto()
+    NETWORK = auto()
+    SEGMENT = auto()
+    WELL = auto()
+    REGION = auto()
+    INTER_REGION = auto()
+    OTHER = auto()
+
+    @classmethod
+    def from_keyword(cls, summary_keyword: str) -> SummaryKeyType:
+        KEYWORD_TYPE_MAPPING = {
+            "A": cls.AQUIFER,
+            "B": cls.BLOCK,
+            "C": cls.COMPLETION,
+            "F": cls.FIELD,
+            "G": cls.GROUP,
+            "LB": cls.LOCAL_BLOCK,
+            "LC": cls.LOCAL_COMPLETION,
+            "LW": cls.LOCAL_WELL,
+            "N": cls.NETWORK,
+            "S": cls.SEGMENT,
+            "W": cls.WELL,
+        }
+        if not summary_keyword:
+            raise ValueError("Got empty summary keyword")
+        if any(special in summary_keyword for special in SPECIAL_KEYWORDS):
+            return cls.OTHER
+        if summary_keyword[0] in KEYWORD_TYPE_MAPPING:
+            return KEYWORD_TYPE_MAPPING[summary_keyword[0]]
+        if summary_keyword[0:2] in KEYWORD_TYPE_MAPPING:
+            return KEYWORD_TYPE_MAPPING[summary_keyword[0:2]]
+        if summary_keyword == "RORFR":
+            return cls.REGION
+
+        if any(
+            re.fullmatch(pattern, summary_keyword)
+            for pattern in [r"R.FT.*", r"R..FT.*", r"R.FR.*", r"R..FR.*", r"R.F"]
+        ):
+            return cls.INTER_REGION
+        if summary_keyword[0] == "R":
+            return cls.REGION
+
+        return cls.OTHER
+
+
 rate_keys = [
     "OPR",
     "OIR",
@@ -100,62 +160,6 @@
 ]
 
 
-class SummaryKeyType(Enum):
-    INVALID = auto()
-    FIELD = auto()
-    REGION = auto()
-    GROUP = auto()
-    WELL = auto()
-    SEGMENT = auto()
-    BLOCK = auto()
-    AQUIFER = auto()
-    COMPLETION = auto()
-    NETWORK = auto()
-    REGION_2_REGION = auto()
-    LOCAL_BLOCK = auto()
-    LOCAL_COMPLETION = auto()
-    LOCAL_WELL = auto()
-    MISC = auto()
-
-    @staticmethod
-    def determine_key_type(key: str) -> "SummaryKeyType":
-        if key in special_keys:
-            return SummaryKeyType.MISC
-
-        if key.startswith("L"):
-            secondary = key[1] if len(key) > 1 else ""
-            return {
-                "B": SummaryKeyType.LOCAL_BLOCK,
-                "C": SummaryKeyType.LOCAL_COMPLETION,
-                "W": SummaryKeyType.LOCAL_WELL,
-            }.get(secondary, SummaryKeyType.MISC)
-
-        if key.startswith("R"):
-            if len(key) == 3 and key[2] == "F":
-                return SummaryKeyType.REGION_2_REGION
-            if key == "RNLF":
-                return SummaryKeyType.REGION_2_REGION
-            if key == "RORFR":
-                return SummaryKeyType.REGION
-            if len(key) >= 4 and key[2] == "F" and key[3] in {"T", "R"}:
-                return SummaryKeyType.REGION_2_REGION
-            if len(key) >= 5 and key[3] == "F" and key[4] in {"T", "R"}:
-                return SummaryKeyType.REGION_2_REGION
-            return SummaryKeyType.REGION
-
-        # default cases or miscellaneous if not matched
-        return {
-            "A": SummaryKeyType.AQUIFER,
-            "B": SummaryKeyType.BLOCK,
-            "C": SummaryKeyType.COMPLETION,
-            "F": SummaryKeyType.FIELD,
-            "G": SummaryKeyType.GROUP,
-            "N": SummaryKeyType.NETWORK,
-            "S": SummaryKeyType.SEGMENT,
-            "W": SummaryKeyType.WELL,
-        }.get(key[0], SummaryKeyType.MISC)
-
-
 def _match_keyword_vector(start: int, rate_keys: List[str], keyword: str) -> bool:
     if len(keyword) < start:
         return False
@@ -169,7 +173,7 @@ def _match_keyword_string(start: int, rate_string: str, keyword: str) -> bool:
 
 
 def is_rate(key: str) -> bool:
-    key_type = SummaryKeyType.determine_key_type(key)
+    key_type = SummaryKeyType.from_keyword(key)
     if key_type in {
         SummaryKeyType.WELL,
         SummaryKeyType.GROUP,
@@ -191,7 +195,7 @@ def is_rate(key: str) -> bool:
     if key_type == SummaryKeyType.SEGMENT:
         return _match_keyword_vector(1, seg_rate_keys, key)
 
-    if key_type == SummaryKeyType.REGION_2_REGION:
+    if key_type == SummaryKeyType.INTER_REGION:
         # Region to region rates are identified by R*FR or R**FR
         if _match_keyword_string(2, "FR", key):
             return True

tests/unit_tests/config/summary_generator.py

@@ -17,7 +17,7 @@
 from pydantic import PositiveInt, conint
 from typing_extensions import Self
 
-from ert.config._read_summary import SPECIAL_KEYWORDS
+from ert.summary_key_type import SPECIAL_KEYWORDS
 
 from .egrid_generator import GrdeclKeyword