Make command line args conversion a function (#225)

Co-authored-by: Elliott Kasoar <[email protected]>

Author: oerc0122

aiida_mlip/calculations/base.py

@@ -13,6 +13,7 @@
 
 from aiida_mlip.data.config import JanusConfigfile
 from aiida_mlip.data.model import ModelData
+from aiida_mlip.helpers.converters import kwarg_to_param
 
 
 def validate_inputs(
@@ -261,10 +262,7 @@ def prepare_for_submission(
         codeinfo = datastructures.CodeInfo()
 
         # Initialize cmdline_params with a placeholder "calculation" command
-        codeinfo.cmdline_params = ["calculation"]
-
-        for flag, value in cmd_line.items():
-            codeinfo.cmdline_params += [f"--{flag}", str(value)]
+        codeinfo.cmdline_params = ["calculation", *kwarg_to_param(cmd_line)]
 
         # Node where the code is saved
         codeinfo.code_uuid = self.inputs.code.uuid

aiida_mlip/calculations/descriptors.py

@@ -9,6 +9,7 @@
 from aiida.orm import Bool
 
 from aiida_mlip.calculations.singlepoint import Singlepoint
+from aiida_mlip.helpers.converters import kwarg_to_param
 
 
 class Descriptors(Singlepoint):  # numpydoc ignore=PR01
@@ -89,20 +90,17 @@ def prepare_for_submission(
 
         # Adding command line params for when we run janus
         # descriptors is overwriting the placeholder "calculation" from the base.py file
-        codeinfo.cmdline_params[0] = "descriptors"
 
         cmdline_options = {
             key.replace("_", "-"): getattr(self.inputs, key).value
             for key in ("invariants_only", "calc_per_element", "calc_per_atom")
             if key in self.inputs
         }
 
-        for flag, value in cmdline_options.items():
-            if isinstance(value, bool):
-                # Add boolean flags without value if True
-                if value:
-                    codeinfo.cmdline_params.append(f"--{flag}")
-            else:
-                codeinfo.cmdline_params += [f"--{flag}", value]
+        codeinfo.cmdline_params = [
+            "descriptors",
+            *codeinfo.cmdline_params[1:],
+            *kwarg_to_param(cmdline_options),
+        ]
 
         return calcinfo

aiida_mlip/calculations/geomopt.py

@@ -18,6 +18,7 @@
 from plumpy.utils import AttributesFrozendict
 
 from aiida_mlip.calculations.singlepoint import Singlepoint
+from aiida_mlip.helpers.converters import kwarg_to_param
 
 
 class GeomOpt(Singlepoint):  # numpydoc ignore=PR01
@@ -124,15 +125,11 @@ def prepare_for_submission(
 
         # Adding command line params for when we run janus
         # 'geomopt' is overwriting the placeholder "calculation" from the base.py file
-        codeinfo.cmdline_params[0] = "geomopt"
-
-        for flag, value in geom_opt_cmdline.items():
-            if isinstance(value, bool):
-                # Add boolean flags without value if True
-                if value:
-                    codeinfo.cmdline_params.append(f"--{flag}")
-            else:
-                codeinfo.cmdline_params += [f"--{flag}", value]
+        codeinfo.cmdline_params = [
+            "geomopt",
+            *codeinfo.cmdline_params[1:],
+            *kwarg_to_param(geom_opt_cmdline),
+        ]
 
         calcinfo.retrieve_list.append(minimize_kwargs["traj_kwargs"]["filename"])
 

aiida_mlip/calculations/md.py

@@ -9,6 +9,7 @@
 from aiida.orm import Dict, SinglefileData, Str, StructureData, TrajectoryData
 
 from aiida_mlip.calculations.base import BaseJanus
+from aiida_mlip.helpers.converters import kwarg_to_param
 
 
 class MD(BaseJanus):  # numpydoc ignore=PR01
@@ -108,16 +109,13 @@ def prepare_for_submission(
             raise ValueError("'ensemble' not provided.")
 
         # md is overwriting the placeholder "calculation" from the base.py file
-        codeinfo.cmdline_params[0] = "md"
-
-        codeinfo.cmdline_params += ["--ensemble", ensemble]
-
-        for flag, value in md_dictionary.items():
-            # Add boolean flags without value if True
-            if isinstance(value, bool) and value:
-                codeinfo.cmdline_params.append(f"--{flag}")
-            else:
-                codeinfo.cmdline_params += [f"--{flag}", value]
+        codeinfo.cmdline_params = [
+            "md",
+            *codeinfo.cmdline_params[1:],
+            "--ensemble",
+            ensemble,
+            *kwarg_to_param(md_dictionary),
+        ]
 
         calcinfo.retrieve_list.append(md_dictionary["traj-file"])
         calcinfo.retrieve_list.append(md_dictionary["stats-file"])

aiida_mlip/calculations/singlepoint.py

@@ -97,11 +97,15 @@ def prepare_for_submission(
 
         # Adding command line params for when we run janus
         # singlepoint is overwriting the placeholder "calculation" from the base.py file
-        codeinfo.cmdline_params[0] = "singlepoint"
 
         # The inputs are saved in the node, but we want their value as a string
         xyz_filename = (self.inputs.out).value
-        codeinfo.cmdline_params += ["--out", xyz_filename]
+        codeinfo.cmdline_params = [
+            "singlepoint",
+            *codeinfo.cmdline_params[1:],
+            "--out",
+            xyz_filename,
+        ]
 
         if "properties" in self.inputs:
             properties = self.inputs.properties.value

aiida_mlip/calculations/train.py

@@ -13,6 +13,7 @@
 
 from aiida_mlip.data.config import JanusConfigfile
 from aiida_mlip.data.model import ModelData
+from aiida_mlip.helpers.converters import kwarg_to_param
 
 
 def validate_inputs(
@@ -190,18 +191,16 @@ def prepare_for_submission(
         codeinfo = datastructures.CodeInfo()
 
         # Initialize cmdline_params with train command
-        codeinfo.cmdline_params = ["train"]
         # Create the rest of the command line
-        cmd_line = {"mlip-config": config_copy}
-        if self.inputs.fine_tune:
-            cmd_line["fine-tune"] = None
-
-        # Add cmd line params to codeinfo
-        for flag, value in cmd_line.items():
-            if value is None:
-                codeinfo.cmdline_params += [f"--{flag}"]
-            else:
-                codeinfo.cmdline_params += [f"--{flag}", str(value)]
+        cmd_line = {
+            "mlip-config": config_copy,
+            "fine-tune": bool(self.inputs.fine_tune),
+        }
+
+        codeinfo.cmdline_params = [
+            "train",
+            *kwarg_to_param(cmd_line),
+        ]
 
         # Node where the code is saved
         codeinfo.code_uuid = self.inputs.code.uuid

aiida_mlip/helpers/converters.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from pathlib import Path
+from typing import Any
 
 from aiida.orm import Bool, Dict, Str, StructureData, TrajectoryData, load_code
 from ase.io import read
@@ -99,3 +100,41 @@ def convert_to_nodes(dictionary: dict, convert_all: bool = False) -> dict:
             continue
         new_dict[key] = value
     return new_dict
+
+
+def kwarg_to_param(params: dict[str, Any]) -> list[str]:
+    """
+    Convert a dictionary of kwargs to a set of commandline flags.
+
+    Bools are converted as though ``store_true`` flag keys.
+
+    Parameters
+    ----------
+    params : dict[str, Any]
+        Dictionary of arguments to convert.
+
+    Returns
+    -------
+    list[str]
+        Commandline arguments as flags.
+
+    Examples
+    --------
+    >>> kwarg_to_param({"name": "Geoff", "key": True})
+    ['--name', 'Geoff', '--key']
+    >>> kwarg_to_param({"value": 6, "falsey": False})
+    ['--value', '6', '--no-falsey']
+    """
+    cmdline_params = []
+
+    for key, val in params.items():
+        key = key.replace("_", "-")
+        match val:
+            case bool() if val:
+                cmdline_params.append(f"--{key}")
+            case bool():
+                cmdline_params.append(f"--no-{key}")
+            case _:
+                cmdline_params.extend((f"--{key}", str(val)))
+
+    return cmdline_params