Skip to content

Commit

Permalink
Some fixes to the MOM6 input parser:
Browse files Browse the repository at this point in the history
 - correctly handle the case of keys that were deleted.
 - change all added keys to uppercase when doing round-trip parsing, as use of uppercase is a soft convention of the MOM6 inputs.
  • Loading branch information
micaeljtoliveira committed Apr 9, 2024
1 parent faeb614 commit 7c08d32
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 3 deletions.
19 changes: 18 additions & 1 deletion om3utils/mom6_input.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,7 @@ def _patch_mom6_input_str(mom6_input_str: str) -> tuple[str, dict]:
block_pattern = re.compile(r"KPP%|%KPP|CVMix_CONVECTION%|%CVMix_CONVECTION|CVMIX_DDIFF%|%CVMIX_DDIFF")
override_directive_pattern = re.compile(r"^(#override\s*?)")
incorrect_directive_pattern = re.compile(r"^(#\s+)")
comment_directive_pattern = re.compile(r"^#(?:(?!override)\w+\b\s*=\s*\w+$)")
comment_directive_pattern = re.compile(r"^#((?!override)\w+\b\s*=\s*\w+$)")

# Modify the input while recording the changes
patch = {}
Expand Down Expand Up @@ -229,6 +229,9 @@ class Mom6Input(dict):
# A record of all the changes done to the dictionary that can be passed to f90nml to do round-trip parsing
_nml_patch = None

# A record of keys that have been deleted from the dictionary
_deleted_keys = []

def __init__(self, file_name: str = None):
"""Read NOM6 parameters from file.
Expand Down Expand Up @@ -262,6 +265,10 @@ def __setitem__(self, key, value):
namelist patch used for round-trip parsing.
"""
super().__setitem__(key.upper(), value)

if key.upper() in self._deleted_keys:
self._deleted_keys.remove(key.upper())

if self._nml_patch:
self._nml_patch["mom6"][key.upper()] = value

Expand All @@ -271,6 +278,7 @@ def __getitem__(self, key):

def __delitem__(self, key):
"""Override method to delete item from dict, so that all keys are stored in uppercase."""
self._deleted_keys.append(key.upper())
super().__delitem__(key.upper())

def write(self, file: Path):
Expand All @@ -286,6 +294,15 @@ def write(self, file: Path):
parser = f90nml.Parser()
parser.read(nml_file, self._nml_patch, tmp_file)
mom6_input_str = _unpatch_mom6_input_str(tmp_file.getvalue(), self._file_patch)

# Change keys to uppercase using a regex substitution, as there seems to be no way of doing this with f90nml
# when applying a nml patch.
mom6_input_str = re.sub(r"((?<=^)|(?<=\n))(\w+)", lambda pat: pat.group(2).upper(), mom6_input_str)

# Explicitly removed keys from string
for key in self._deleted_keys:
mom6_input_str = re.sub(r"\s*" + f"{key}" + r"\s*=\s*\S*\s*\n", r"\n", mom6_input_str)

file.write_text(mom6_input_str)

def _keys_to_upper(self):
Expand Down
7 changes: 5 additions & 2 deletions tests/test_mom6_input.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ def complex_mom6_input_file(tmp_path):
DT = 1800.0 ! This is a comment
! This is another comment
!COMMENTED_VAR = 3
TO_BE_REMOVED = 10.0
BOOL = True
"""
return MockFile(file, mom6_input_str)
Expand All @@ -72,7 +73,7 @@ def modified_mom6_input_file(tmp_path):
BOOL = True
added_var = 32
ADDED_VAR = 32
"""
return MockFile(file, mom6_input_str)

Expand All @@ -94,8 +95,10 @@ def test_round_trip_mom6_input(tmp_path, complex_mom6_input_file, modified_mom6_
mom6_input_from_file = Mom6Input(file_name=complex_mom6_input_file.file)
mom6_input_from_file["dt"] = 900.0
mom6_input_from_file["ADDED_VAR"] = 1
del mom6_input_from_file["ADDED_VAR"]
mom6_input_from_file["ADDED_VAR"] = 32
del mom6_input_from_file["N_SMOOTH"]
mom6_input_from_file["N_SMOOTH"] = 4
del mom6_input_from_file["TO_BE_REMOVED"]

write_mom6_input(mom6_input_from_file, tmp_path / "MOM_input_new")

Expand Down

0 comments on commit 7c08d32

Please sign in to comment.