Skip to content

Commit

Permalink
Add key util function is_rate
Browse files Browse the repository at this point in the history
  • Loading branch information
xjules committed Sep 10, 2024
1 parent cbbd64d commit bcb3c1f
Show file tree
Hide file tree
Showing 2 changed files with 215 additions and 1 deletion.
3 changes: 2 additions & 1 deletion src/ert/gui/plottery/plots/ensemble.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

from ert.gui.plottery.plots.history import plotHistory
from ert.gui.tools.plot.plot_api import EnsembleObject
from ert.shared.storage.key_utils import is_rate

from .observations import plotObservations
from .plot_tools import PlotTools
Expand Down Expand Up @@ -45,7 +46,7 @@ def plot(
plot_context.deactivateDateSupport()
plot_context.x_axis = plot_context.INDEX_AXIS

draw_style = "steps" if plot_context.history_data else None
draw_style = "steps" if is_rate(plot_context.key) else None

Check failure on line 49 in src/ert/gui/plottery/plots/ensemble.py

View workflow job for this annotation

GitHub Actions / type-checking (3.12)

Argument 1 to "is_rate" has incompatible type "Callable[[], str]"; expected "str"
self._plotLines(
axes,
config,
Expand Down
213 changes: 213 additions & 0 deletions src/ert/shared/storage/key_utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,213 @@
from enum import Enum, auto

special_vars = [
"NAIMFRAC",
"NBAKFL",
"NBYTOT",
"NCPRLINS",
"NEWTFL",
"NEWTON",
"NLINEARP",
"NLINEARS",
"NLINSMAX",
"NLINSMIN",
"NLRESMAX",
"NLRESSUM",
"NMESSAGE",
"NNUMFL",
"NNUMST",
"NTS",
"NTSECL",
"NTSMCL",
"NTSPCL",
"ELAPSED",
"MAXDPR",
"MAXDSO",
"MAXDSG",
"MAXDSW",
"STEPTYPE",
"WNEWTON",
]
rate_vars = [
"OPR",
"OIR",
"OVPR",
"OVIR",
"OFR",
"OPP",
"OPI",
"OMR",
"GPR",
"GIR",
"GVPR",
"GVIR",
"GFR",
"GPP",
"GPI",
"GMR",
"WGPR",
"WGIR",
"WPR",
"WIR",
"WVPR",
"WVIR",
"WFR",
"WPP",
"WPI",
"WMR",
"LPR",
"LFR",
"VPR",
"VIR",
"VFR",
"GLIR",
"RGR",
"EGR",
"EXGR",
"SGR",
"GSR",
"FGR",
"GIMR",
"GCR",
"NPR",
"NIR",
"CPR",
"CIR",
"SIR",
"SPR",
"TIR",
"TPR",
"GOR",
"WCT",
"OGR",
"WGR",
"GLR",
]

seg_rate_vars = [
"OFR",
"GFR",
"WFR",
"CFR",
"SFR",
"TFR",
"CVPR",
"WCT",
"GOR",
"OGR",
"WGR",
]


class SummaryVarType(Enum):
RD_SMSPEC_INVALID_VAR = auto()
RD_SMSPEC_FIELD_VAR = auto()
RD_SMSPEC_REGION_VAR = auto()
RD_SMSPEC_GROUP_VAR = auto()
RD_SMSPEC_WELL_VAR = auto()
RD_SMSPEC_SEGMENT_VAR = auto()
RD_SMSPEC_BLOCK_VAR = auto()
RD_SMSPEC_AQUIFER_VAR = auto()
RD_SMSPEC_COMPLETION_VAR = auto()
RD_SMSPEC_NETWORK_VAR = auto()
RD_SMSPEC_REGION_2_REGION_VAR = auto()
RD_SMSPEC_LOCAL_BLOCK_VAR = auto()
RD_SMSPEC_LOCAL_COMPLETION_VAR = auto()
RD_SMSPEC_LOCAL_WELL_VAR = auto()
RD_SMSPEC_MISC_VAR = auto()

@staticmethod
def determine_var_type(var):

Check failure on line 120 in src/ert/shared/storage/key_utils.py

View workflow job for this annotation

GitHub Actions / type-checking (3.12)

Function is missing a type annotation
# Default case for single characters or unknown vars
default_case = {
"A": SummaryVarType.RD_SMSPEC_AQUIFER_VAR,
"B": SummaryVarType.RD_SMSPEC_BLOCK_VAR,
"C": SummaryVarType.RD_SMSPEC_COMPLETION_VAR,
"F": SummaryVarType.RD_SMSPEC_FIELD_VAR,
"G": SummaryVarType.RD_SMSPEC_GROUP_VAR,
"N": SummaryVarType.RD_SMSPEC_NETWORK_VAR,
"S": SummaryVarType.RD_SMSPEC_SEGMENT_VAR,
"W": SummaryVarType.RD_SMSPEC_WELL_VAR,
}

# Special case for 'L' with nested logic
if var.startswith("L"):
secondary = var[1] if len(var) > 1 else None
return {
"B": SummaryVarType.RD_SMSPEC_LOCAL_BLOCK_VAR,
"C": SummaryVarType.RD_SMSPEC_LOCAL_COMPLETION_VAR,
"W": SummaryVarType.RD_SMSPEC_LOCAL_WELL_VAR,
}.get(secondary, SummaryVarType.RD_SMSPEC_MISC_VAR)

# Special case for 'R' with more complex nested logic
if var.startswith("R"):
if len(var) == 3 and var[2] == "F":
return SummaryVarType.RD_SMSPEC_REGION_2_REGION_VAR
if var == "RNLF":
return SummaryVarType.RD_SMSPEC_REGION_2_REGION_VAR
if var == "RORFR":
return SummaryVarType.RD_SMSPEC_REGION_VAR
if len(var) >= 4 and var[2] == "F" and var[3] in {"T", "R"}:
return SummaryVarType.RD_SMSPEC_REGION_2_REGION_VAR
if len(var) >= 5 and var[3] == "F" and var[4] in {"T", "R"}:
return SummaryVarType.RD_SMSPEC_REGION_2_REGION_VAR
return SummaryVarType.RD_SMSPEC_REGION_VAR

# Fallback to default cases or miscellaneous if not matched
return default_case.get(var[0], SummaryVarType.RD_SMSPEC_MISC_VAR)


def identify_var_type(key: str) -> SummaryVarType:
if key not in special_vars:
return SummaryVarType.determine_var_type(key)

Check failure on line 162 in src/ert/shared/storage/key_utils.py

View workflow job for this annotation

GitHub Actions / type-checking (3.12)

Call to untyped function "determine_var_type" of "SummaryVarType" in typed context
return SummaryVarType.RD_SMSPEC_MISC_VAR


def match_keyword_vector(start, rate_vars, keyword):

Check failure on line 166 in src/ert/shared/storage/key_utils.py

View workflow job for this annotation

GitHub Actions / type-checking (3.12)

Function is missing a type annotation
"""
Check if 'keyword' matches any item in 'rate_vars' starting from
index 'start' onwards.
"""
# Get the suffix of keyword starting from 'start' index (if not out of range)
suffix = keyword[start:] if len(keyword) > start else ""
return any(suffix.startswith(var) for var in rate_vars)


def match_keyword_string(start, rate_string, keyword):

Check failure on line 176 in src/ert/shared/storage/key_utils.py

View workflow job for this annotation

GitHub Actions / type-checking (3.12)

Function is missing a type annotation
"""
Check if 'keyword' matches 'rate_string' starting from index 'start'.
"""
return keyword[start:].startswith(rate_string)


def is_rate(key: str) -> bool:
var_type = identify_var_type(key)
if var_type in {
SummaryVarType.RD_SMSPEC_WELL_VAR,
SummaryVarType.RD_SMSPEC_GROUP_VAR,
SummaryVarType.RD_SMSPEC_FIELD_VAR,
SummaryVarType.RD_SMSPEC_REGION_VAR,
SummaryVarType.RD_SMSPEC_COMPLETION_VAR,
SummaryVarType.RD_SMSPEC_LOCAL_WELL_VAR,
SummaryVarType.RD_SMSPEC_LOCAL_COMPLETION_VAR,
SummaryVarType.RD_SMSPEC_NETWORK_VAR,
}:
# Process local and network vars differently
if var_type in {
SummaryVarType.RD_SMSPEC_LOCAL_WELL_VAR,
SummaryVarType.RD_SMSPEC_LOCAL_COMPLETION_VAR,
SummaryVarType.RD_SMSPEC_NETWORK_VAR,
}:
return match_keyword_vector(2, rate_vars, key)

Check failure on line 201 in src/ert/shared/storage/key_utils.py

View workflow job for this annotation

GitHub Actions / type-checking (3.12)

Call to untyped function "match_keyword_vector" in typed context
return match_keyword_vector(1, rate_vars, key)

Check failure on line 202 in src/ert/shared/storage/key_utils.py

View workflow job for this annotation

GitHub Actions / type-checking (3.12)

Call to untyped function "match_keyword_vector" in typed context

if var_type == SummaryVarType.RD_SMSPEC_SEGMENT_VAR:
return match_keyword_vector(1, seg_rate_vars, key)

Check failure on line 205 in src/ert/shared/storage/key_utils.py

View workflow job for this annotation

GitHub Actions / type-checking (3.12)

Call to untyped function "match_keyword_vector" in typed context

if var_type == SummaryVarType.RD_SMSPEC_REGION_2_REGION_VAR:
# Region to region rates are identified by R*FR or R**FR
if match_keyword_string(2, "FR", key):

Check failure on line 209 in src/ert/shared/storage/key_utils.py

View workflow job for this annotation

GitHub Actions / type-checking (3.12)

Call to untyped function "match_keyword_string" in typed context
return True
return match_keyword_string(3, "FR", key)

Check failure on line 211 in src/ert/shared/storage/key_utils.py

View workflow job for this annotation

GitHub Actions / type-checking (3.12)

Call to untyped function "match_keyword_string" in typed context

return False

0 comments on commit bcb3c1f

Please sign in to comment.