Replace opt_kwargs with minimize_kwargs (#180)

Co-authored-by: ElliottKasoar <[email protected]>

Author: MM0hsin

aiida_mlip/calculations/geomopt.py

@@ -5,17 +5,17 @@
 from aiida.common import datastructures
 import aiida.common.folders
 from aiida.engine import CalcJobProcessSpec
-import aiida.engine.processes
 from aiida.orm import (
     Bool,
     Dict,
     Float,
     Int,
     SinglefileData,
-    Str,
     StructureData,
     TrajectoryData,
 )
+from aiida.orm.utils.managers import NodeLinksManager
+from plumpy.utils import AttributesFrozendict
 
 from aiida_mlip.calculations.singlepoint import Singlepoint
 
@@ -48,13 +48,6 @@ def define(cls, spec: CalcJobProcessSpec) -> None:
         super().define(spec)
 
         # Additional inputs for geometry optimisation
-        spec.input(
-            "traj",
-            valid_type=Str,
-            required=False,
-            default=lambda: Str(cls.DEFAULT_TRAJ_FILE),
-            help="Path to save optimisation frames to",
-        )
         spec.input(
             "opt_cell_fully",
             valid_type=Bool,
@@ -82,10 +75,10 @@ def define(cls, spec: CalcJobProcessSpec) -> None:
         )
 
         spec.input(
-            "opt_kwargs",
+            "minimize_kwargs",
             valid_type=Dict,
             required=False,
-            help="Other optimisation keywords",
+            help="All other keyword arguments to pass to geometry optimizer",
         )
 
         spec.inputs["metadata"]["options"]["parser_name"].default = "mlip.opt_parser"
@@ -114,17 +107,12 @@ def prepare_for_submission(
         calcinfo = super().prepare_for_submission(folder)
         codeinfo = calcinfo.codes_info[0]
 
-        minimize_kwargs = (
-            f"{{'traj_kwargs': {{'filename': '{self.inputs.traj.value}'}}}}"
-        )
+        minimize_kwargs = self.set_minimize_kwargs(self.inputs)
 
         geom_opt_cmdline = {
             "minimize-kwargs": minimize_kwargs,
             "write-traj": True,
         }
-        if "opt_kwargs" in self.inputs:
-            opt_kwargs = self.inputs.opt_kwargs.get_dict()
-            geom_opt_cmdline["opt-kwargs"] = opt_kwargs
         if "opt_cell_fully" in self.inputs:
             geom_opt_cmdline["opt-cell-fully"] = self.inputs.opt_cell_fully.value
         if "opt_cell_lengths" in self.inputs:
@@ -146,6 +134,36 @@ def prepare_for_submission(
             else:
                 codeinfo.cmdline_params += [f"--{flag}", value]
 
-        calcinfo.retrieve_list.append(self.inputs.traj.value)
+        calcinfo.retrieve_list.append(minimize_kwargs["traj_kwargs"]["filename"])
 
         return calcinfo
+
+    @classmethod
+    def set_minimize_kwargs(
+        cls, inputs: AttributesFrozendict | NodeLinksManager
+    ) -> dict[str, dict[str, str]]:
+        """
+        Set minimize kwargs from CalcJob inputs.
+
+        Parameters
+        ----------
+        inputs : x
+            CalcJob inputs.
+
+        Returns
+        -------
+        dict[str, dict[str, str]]
+            Set minimize_kwargs dict with trajectory filename extracted from `traj`,
+            the config file, or set as the default.
+        """
+        if "minimize_kwargs" in inputs:
+            minimize_kwargs = inputs.minimize_kwargs.get_dict()
+        elif "config" in inputs:
+            minimize_kwargs = inputs.config.as_dictionary.get("minimize_kwargs", {})
+        else:
+            minimize_kwargs = {}
+
+        minimize_kwargs.setdefault("traj_kwargs", {})
+        minimize_kwargs["traj_kwargs"].setdefault("filename", cls.DEFAULT_TRAJ_FILE)
+
+        return minimize_kwargs

aiida_mlip/parsers/opt_parser.py

@@ -10,6 +10,7 @@
 from aiida.orm.nodes.process.process import ProcessNode
 from aiida.plugins import CalculationFactory
 
+from aiida_mlip.calculations.geomopt import GeomOpt
 from aiida_mlip.helpers.converters import xyz_to_aiida_traj
 from aiida_mlip.parsers.sp_parser import SPParser
 
@@ -75,7 +76,9 @@ def parse(self, **kwargs) -> ExitCode:
         exit_code = super().parse(**kwargs)
 
         if exit_code == ExitCode(0):
-            traj_file = (self.node.inputs.traj).value
+            traj_file = GeomOpt.set_minimize_kwargs(self.node.inputs)["traj_kwargs"][
+                "filename"
+            ]
 
             # Parse the trajectory file and save it as `SingleFileData`
             with self.retrieved.open(traj_file, "rb") as handle:

docs/source/user_guide/tutorial.rst

@@ -65,6 +65,7 @@ The other inputs can be set up as AiiDA ``Str``. There is a default for every in
         "fmax": Float(0.1),
         "opt_cell_lengths": Bool(False),
         "opt_cell_fully": Bool(True),
+        "minimize_kwargs": Dict({"filter_kwargs": {"constant_volume": True}}),
         "metadata": {"options": {"resources": {"num_machines": 1}}},
         }
 
@@ -168,7 +169,7 @@ The calculation can also be interacted with through verdi cli. Use ``verdi proce
     max_force        1124  Float
     model            1119  ModelData
     structure        1120  StructureData
-    traj             1129  Str
+    minimize_kwargs  1129  Dict
     opt_cell_lengths 1125  Bool
     xyz_output_name  1127  Str
 

examples/calculations/submit_geomopt.py

@@ -45,14 +45,17 @@ def geomopt(params: dict) -> None:
         "fmax": Float(params["fmax"]),
         "opt_cell_lengths": Bool(params["opt_cell_lengths"]),
         "opt_cell_fully": Bool(params["opt_cell_fully"]),
-        # "opt_kwargs": Dict({"restart": "rest.pkl"}),
         "steps": Int(params["steps"]),
     }
 
     # Only calc_kwargs add if set
     if params["calc_kwargs"]:
         inputs["calc_kwargs"] = Dict(params["calc_kwargs"])
 
+    # Only minimize_kwargs add if set
+    if params["minimize_kwargs"]:
+        inputs["minimize_kwargs"] = Dict(params["minimize_kwargs"])
+
     # Run calculation
     result, node = run_get_node(GeomoptCalc, **inputs)
     print(f"Printing results from calculation: {result}")
@@ -105,6 +108,15 @@ def geomopt(params: dict) -> None:
 @click.option(
     "--steps", default=1000, type=int, help="Maximum number of optimisation steps."
 )
+@click.option(
+    "--minimize-kwargs",
+    default="{}",
+    type=str,
+    help=(
+        "Keyword arguments to pass to geometry optimizer, including 'opt_kwargs', "
+        "'filter_kwargs', and 'traj_kwargs'."
+    ),
+)
 def cli(
     codelabel,
     struct,
@@ -116,9 +128,11 @@ def cli(
     opt_cell_lengths,
     opt_cell_fully,
     steps,
+    minimize_kwargs,
 ) -> None:
     """Click interface."""
     calc_kwargs = ast.literal_eval(calc_kwargs)
+    minimize_kwargs = ast.literal_eval(minimize_kwargs)
 
     try:
         code = load_code(codelabel)
@@ -137,6 +151,7 @@ def cli(
         "opt_cell_lengths": opt_cell_lengths,
         "opt_cell_fully": opt_cell_fully,
         "steps": steps,
+        "minimize_kwargs": minimize_kwargs,
     }
 
     # Submit single point

examples/tutorials/calculations/geometry-optimisation.ipynb

@@ -153,6 +153,7 @@
     "        \"fmax\": Float(0.1), \n",
     "        \"opt_cell_lengths\": Bool(False), \n",
     "        \"opt_cell_fully\": Bool(True),\n",
+    "        \"minimize_kwargs\": Dict({\"filter_kwargs\": {\"constant_volume\": True}}),\n",
     "        \"metadata\": {\"options\": {\"resources\": {\"num_machines\": 1}}},\n",
     "    }"
    ]

tests/calculations/configs/config_janus_geomopt.yml

@@ -0,0 +1,9 @@
+arch: mace_mp
+struct: "NaCl.cif"
+calc_kwargs:
+  dispersion: True
+minimize_kwargs:
+  opt_kwargs:
+    alpha: 100
+fmax: 0.1
+opt_cell_lengths: True

tests/calculations/configs/config_janus_geomopt_traj.yml

@@ -0,0 +1,11 @@
+arch: mace_mp
+struct: "NaCl.cif"
+calc_kwargs:
+  dispersion: True
+minimize_kwargs:
+  filter_kwargs:
+    constant_volume: True
+  traj_kwargs:
+    filename: "test-traj.xyz"
+fmax: 0.1
+opt_cell_lengths: True