Add tool to upgrade collisions/closure in Hermes-3

Author: cmacmackin

src/boutupgrader/__init__.py

@@ -11,6 +11,7 @@
 from .bout_v5_xzinterpolation_upgrader import add_parser as add_xzinterp_parser
 from .bout_v6_coordinates_upgrader import add_parser as add_v6_coordinates_parser
 from .bout_v6_input_file_upgrader import add_parser as add_v6_input_parser
+from .hermes_collisions_input_file_upgrader import add_parser as add_hermes_input_parser
 
 try:
     # This gives the version if the boututils package was installed
@@ -63,6 +64,10 @@ def main():
         "v6", help="BOUT++ v6 upgrades"
     ).add_subparsers(title="v6 subcommands", required=True)
 
+    hermes_subcommand = subcommand.add_parser(
+        "hermes", help="Hermes-3 subcommands"
+    ).add_subparsers(title="hermes subcommands", required=True)
+
     add_3to4_parser(v4_subcommand, common_args, files_args)
     add_factory_parser(v5_subcommand, common_args, files_args)
     add_format_parser(v5_subcommand, common_args, files_args)
@@ -73,6 +78,7 @@ def main():
     add_xzinterp_parser(v5_subcommand, common_args, files_args)
     add_v6_coordinates_parser(v6_subcommand, common_args, files_args)
     add_v6_input_parser(v6_subcommand, common_args, files_args)
+    add_hermes_input_parser(hermes_subcommand, common_args, files_args)
 
     args = parser.parse_args()
     args.func(args)

src/boutupgrader/bout_v5_input_file_upgrader.py

@@ -203,7 +203,7 @@ def remove_deleted(deleted, options_file):
             del options_file[key]
 
 
-def apply_fixes(replacements, deleted, options_file):
+def apply_fixes(replacements, deleted, options_file, additional_modifications):
     """Apply all fixes in this module"""
 
     modified = copy.deepcopy(options_file)
@@ -212,6 +212,8 @@ def apply_fixes(replacements, deleted, options_file):
 
     remove_deleted(deleted, modified)
 
+    additional_modifications(modified)
+
     return modified
 
 
@@ -236,14 +238,14 @@ def possibly_apply_patch(patch, options_file, quiet=False, force=False):
     return make_change
 
 
-def add_parser_general(subcommand, default_args, files_args, run):
+def add_parser_general(subcommand, default_args, files_args, run, name):
     parser = subcommand.add_parser(
         "input",
         formatter_class=argparse.RawDescriptionHelpFormatter,
         help="Fix input files",
         description=textwrap.dedent(
-            """\
-            Fix input files for BOUT++ v5+
+            f"""\
+            Fix input files for {name}
 
             Please note that this will only fix input options in sections with
             standard or default names. You may also need to fix options in custom
@@ -290,7 +292,7 @@ def add_parser_general(subcommand, default_args, files_args, run):
     parser.set_defaults(func=run)
 
 
-def run_general(REPLACEMENTS, DELETED, args):
+def run_general(REPLACEMENTS, DELETED, args, *, additional_modifications=None):
     from boutdata.data import BoutOptions, BoutOptionsFile
 
     # Monkey-patch BoutOptions to make sure it's case sensitive
@@ -324,7 +326,7 @@ def run_general(REPLACEMENTS, DELETED, args):
             continue
 
         try:
-            modified = apply_fixes(REPLACEMENTS, DELETED, original)
+            modified = apply_fixes(REPLACEMENTS, DELETED, original, additional_modifications)
         except RuntimeError as e:
             print(e)
             continue
@@ -347,4 +349,4 @@ def run(args):
 
 
 def add_parser(subcommand, default_args, files_args):
-    return add_parser_general(subcommand, default_args, files_args, run)
+    return add_parser_general(subcommand, default_args, files_args, run, "BOUT++ v5+")

src/boutupgrader/bout_v6_input_file_upgrader.py

@@ -19,4 +19,4 @@ def run(args):
 
 
 def add_parser(subcommand, default_args, files_args):
-    return add_parser_general(subcommand, default_args, files_args, run)
+    return add_parser_general(subcommand, default_args, files_args, run, "BOUT++ v6+")

src/boutupgrader/hermes_collisions_input_file_upgrader.py

@@ -0,0 +1,59 @@
+from itertools import chain
+from typing import TypedDict
+import re
+
+from .bout_v5_input_file_upgrader import add_parser_general, run_general
+from boutdata.data import BoutOptionsFile
+
+class Replacement(TypedDict):
+    old: str
+    new: str
+
+REPLACEMENTS = []
+DELETED = []
+
+NEW_NAMES = {
+    "collisions": ["braginskii_collisions", "braginskii_friction", "braginskii_heat_exchange"],
+    "electron_viscosity": ["braginskii_electron_viscosity"],
+    "ion_viscosity": ["braginskii_ion_viscosity"],
+    "thermal_force": ["braginskii_thermal_force"]
+}
+
+COMPONENT_RE = re.compile(r"[+\-\w]+")
+
+
+def update_component_names(options_file: BoutOptionsFile) -> None:
+    """Change the names of closure-related components to reflect the refactor"""
+    has_collisions = False
+    recycling_component = ""
+    components = re.findall(COMPONENT_RE, options_file["hermes:components"])
+    for section_name in components:
+        section = options_file.getSection(section_name)
+        explicit_types = "type" in section
+        types = [tname.strip() for tname in section["type"].split(",")] if explicit_types else [section_name]
+        open_paren = types[0][0] == "("
+        close_paren = types[-1][-1] == ")"
+        if open_paren:
+            types[0] = types[0][1:]
+        if close_paren:
+            types[-1] = types[-1][:-1]
+        new_types = list(chain.from_iterable(NEW_NAMES.get(t, [t]) for t in types))
+        has_collisions = has_collisions or "collisions" in types
+        if "recycling" in new_types:
+            recycling_component = section_name
+        if new_types != types:
+            section["type"] =  ("(" if open_paren else "") + ", ".join(new_types) + (")" if close_paren else "")
+    # Add braginskii_conduction to the end of the list of components
+    if has_collisions:
+        components.append("braginskii_conduction")
+        # Make sure recycling is evaluated after conduction
+        if recycling_component != "":
+            components.remove(recycling_component)
+            components.append(recycling_component)
+    options_file["hermes:components"] = "(" + ", ".join(components) + ")"
+
+def run(args) -> None:
+    run_general(REPLACEMENTS, DELETED, args, additional_modifications=update_component_names)
+    
+def add_parser(subcommand, default_args, files_args):
+    return add_parser_general(subcommand, default_args, files_args, run, "Hermes-3")