Add calc kwargs input (#220)

Author: ElliottKasoar

aiida_mlip/calculations/base.py

@@ -8,7 +8,7 @@
 import aiida.common.folders
 from aiida.engine import CalcJob, CalcJobProcessSpec
 import aiida.engine.processes
-from aiida.orm import SinglefileData, Str, StructureData
+from aiida.orm import Dict, SinglefileData, Str, StructureData
 from ase.io import read, write
 
 from aiida_mlip.data.config import JanusConfigfile
@@ -134,12 +134,6 @@ def define(cls, spec: CalcJobProcessSpec) -> None:
             required=False,
             help="The input structure.",
         )
-        spec.input(
-            "precision",
-            valid_type=Str,
-            required=False,
-            help="Precision level for calculation",
-        )
         spec.input(
             "device",
             valid_type=Str,
@@ -171,6 +165,13 @@ def define(cls, spec: CalcJobProcessSpec) -> None:
             help="Filename to which the content of stdout of the scheduler is written.",
         )
 
+        spec.input(
+            "calc_kwargs",
+            valid_type=Dict,
+            required=False,
+            help="Keyword arguments to pass to selected calculator.",
+        )
+
         spec.input(
             "config",
             valid_type=JanusConfigfile,
@@ -223,13 +224,14 @@ def prepare_for_submission(
         }
 
         # The inputs are saved in the node, but we want their value as a string
-        if "precision" in self.inputs:
-            precision = (self.inputs.precision).value
-            cmd_line["calc-kwargs"] = {"default_dtype": precision}
         if "device" in self.inputs:
-            device = (self.inputs.device).value
+            device = self.inputs.device.value
             cmd_line["device"] = device
 
+        # Set calc_kwargs from dict and specific stored inputs
+        if "calc_kwargs" in self.inputs:
+            cmd_line["calc-kwargs"] = self.inputs.calc_kwargs.get_dict()
+
         # Define architecture from model if model is given,
         # otherwise get architecture from inputs and download default model
         self._add_arch_to_cmdline(cmd_line)
@@ -310,7 +312,6 @@ def _add_model_to_cmdline(
         dict
             Dictionary containing the cmd line keys updated with the model.
         """
-        model_path = None
         if "model" in self.inputs:
             # Raise error if model is None (different than model not given as input)
             if self.inputs.model is None:
@@ -322,5 +323,4 @@ def _add_model_to_cmdline(
             ):
                 shutil.copyfileobj(source, target)
 
-            model_path = "mlff.model"
-            cmd_line.setdefault("calc-kwargs", {})["model"] = model_path
+            cmd_line["model"] = "mlff.model"

docs/source/user_guide/calculations.rst

@@ -40,8 +40,8 @@ The inputs can be grouped into a dictionary:
             "architecture": Str,
             "structure": StructureData,
             "model": ModelData,
-            "precision": Str,
             "device": Str,
+            "calc_kwargs": Dict,
         }
     SinglePointCalculation = CalculationFactory("mlip.sp")
     submit(SinglePointCalculation, **inputs)
@@ -58,6 +58,9 @@ The config file contains the parameters in yaml format:
     device: "cpu"
     struct: "path/to/structure.cif"
     model: "path/to/model.model"
+    calc_kwargs:
+      dispersion: True
+
 
 And it is used as shown below. Note that some parameters, which are specific to AiiDA, need to be given individually.
 
@@ -101,8 +104,6 @@ In this case  the structure used is going to be "path/to/structure2.xyz" rather
 Refer to the API documentation for additional parameters that can be passed.
 Some parameters are not required and don't have a default value set in aiida-mlip. In that case the default values will be the same as `janus <https://stfc.github.io/janus-core/>`_
 The only default parameters defined in aiida-mlip are the names of the input and output files, as they do not affect the results of the calculation itself, and are needed in AiiDA to parse the results.
-For example in the code above the parameter "precision" is never defined, neither in the config nor in the run_get_node function.
-The parameter will default to the janus default, which is "float64"
 
 
 Submission
@@ -120,7 +121,7 @@ They will be converted to AiiDA data types by the script itself.
 
 .. code-block:: python
 
-    verdi run submit_singlepoint.py "janus@localhost" --structure "path/to/structure" --model "path/to/model" --precision "float64" --device "cpu"
+    verdi run submit_singlepoint.py "janus@localhost" --structure "path/to/structure" --model "path/to/model" --device "cpu"
 
 The submit_using_config.py script can be used to facilitate submission using a config file.
 
@@ -157,7 +158,7 @@ They will be converted to AiiDA data types by the script itself.
 
 .. code-block:: python
 
-    verdi run submit_geomopt.py "janus@localhost" --structure "path/to/structure" --model "path/to/model" --precision "float64" --device "cpu"
+    verdi run submit_geomopt.py "janus@localhost" --structure "path/to/structure" --model "path/to/model" --device "cpu"
 
 
 

docs/source/user_guide/tutorial.rst

@@ -53,14 +53,15 @@ The other inputs can be set up as AiiDA Str. There is a default for every input
 
 .. code-block:: python
 
-    from aiida.orm import Bool, Float, Str
+    from aiida.orm import Bool, Dict, Float, Str
+
     inputs = {
         "code": code,
         "model": model,
         "structure": structure,
         "architecture": Str(model.architecture),
-        "precision": Str("float64"),
         "device": Str("cpu"),
+        "calc_kwargs": Dict({}),
         "max_force": Float(0.1), # Specific to geometry optimisation: convergence criteria
         "opt_cell_lengths": Bool(False), # Specific to geometry optimisation
         "opt_cell_fully": Bool(True), # Specific to geometry optimisation: to optimise the cell
@@ -159,13 +160,13 @@ The calculation can also be interacted with through verdi cli. Use `verdi proces
     Inputs             PK  Type
     ---------------  ----  -------------
     architecture     1121  Str
+    calc_kwargs      1122  Dict
     code                2  InstalledCode
     device           1123  Str
     opt_cell_fully   1126  Bool
     log_filename     1128  Str
     max_force        1124  Float
     model            1119  ModelData
-    precision        1122  Str
     structure        1120  StructureData
     traj             1129  Str
     opt_cell_lengths 1125  Bool

examples/calculations/submit_descriptors.py

@@ -2,9 +2,11 @@
 
 from __future__ import annotations
 
+import ast
+
 from aiida.common import NotExistent
 from aiida.engine import run_get_node
-from aiida.orm import Bool, Str, load_code
+from aiida.orm import Bool, Dict, Str, load_code
 from aiida.plugins import CalculationFactory
 import click
 
@@ -39,13 +41,16 @@ def descriptors(params: dict) -> None:
         "arch": Str(params["arch"]),
         "struct": structure,
         "model": model,
-        "precision": Str(params["precision"]),
         "device": Str(params["device"]),
         "invariants_only": Bool(params["invariants_only"]),
         "calc_per_element": Bool(params["calc_per_element"]),
         "calc_per_atom": Bool(params["calc_per_atom"]),
     }
 
+    # Only calc_kwargs add if set
+    if params["calc_kwargs"]:
+        inputs["calc_kwargs"] = Dict(params["calc_kwargs"])
+
     # Run calculation
     result, node = run_get_node(DescriptorsCalc, **inputs)
     print(f"Printing results from calculation: {result}")
@@ -77,7 +82,10 @@ def descriptors(params: dict) -> None:
     "--device", default="cpu", type=str, help="Device to run calculations on."
 )
 @click.option(
-    "--precision", default="float64", type=str, help="Chosen level of precision."
+    "--calc-kwargs",
+    default="{}",
+    type=str,
+    help="Keyword arguments to pass to calculator.",
 )
 @click.option(
     "--invariants-only",
@@ -103,12 +111,14 @@ def cli(
     model,
     arch,
     device,
-    precision,
+    calc_kwargs,
     invariants_only,
     calc_per_element,
     calc_per_atom,
 ) -> None:
     """Click interface."""
+    calc_kwargs = ast.literal_eval(calc_kwargs)
+
     try:
         code = load_code(codelabel)
     except NotExistent as exc:
@@ -121,7 +131,7 @@ def cli(
         "model": model,
         "arch": arch,
         "device": device,
-        "precision": precision,
+        "calc_kwargs": calc_kwargs,
         "invariants_only": invariants_only,
         "calc_per_element": calc_per_element,
         "calc_per_atom": calc_per_atom,

examples/calculations/submit_geomopt.py

@@ -2,9 +2,11 @@
 
 from __future__ import annotations
 
+import ast
+
 from aiida.common import NotExistent
 from aiida.engine import run_get_node
-from aiida.orm import Bool, Float, Int, Str, load_code
+from aiida.orm import Bool, Dict, Float, Int, Str, load_code
 from aiida.plugins import CalculationFactory
 import click
 
@@ -39,7 +41,6 @@ def geomopt(params: dict) -> None:
         "arch": Str(params["arch"]),
         "struct": structure,
         "model": model,
-        "precision": Str(params["precision"]),
         "device": Str(params["device"]),
         "fmax": Float(params["fmax"]),
         "opt_cell_lengths": Bool(params["opt_cell_lengths"]),
@@ -48,6 +49,10 @@ def geomopt(params: dict) -> None:
         "steps": Int(params["steps"]),
     }
 
+    # Only calc_kwargs add if set
+    if params["calc_kwargs"]:
+        inputs["calc_kwargs"] = Dict(params["calc_kwargs"])
+
     # Run calculation
     result, node = run_get_node(GeomoptCalc, **inputs)
     print(f"Printing results from calculation: {result}")
@@ -79,7 +84,10 @@ def geomopt(params: dict) -> None:
     "--device", default="cpu", type=str, help="Device to run calculations on."
 )
 @click.option(
-    "--precision", default="float64", type=str, help="Chosen level of precision."
+    "--calc-kwargs",
+    default="{}",
+    type=str,
+    help="Keyword arguments to pass to calculator.",
 )
 @click.option("--fmax", default=0.1, type=float, help="Maximum force for convergence.")
 @click.option(
@@ -103,13 +111,15 @@ def cli(
     model,
     arch,
     device,
-    precision,
+    calc_kwargs,
     fmax,
     opt_cell_lengths,
     opt_cell_fully,
     steps,
 ) -> None:
     """Click interface."""
+    calc_kwargs = ast.literal_eval(calc_kwargs)
+
     try:
         code = load_code(codelabel)
     except NotExistent as exc:
@@ -122,7 +132,7 @@ def cli(
         "model": model,
         "arch": arch,
         "device": device,
-        "precision": precision,
+        "calc_kwargs": calc_kwargs,
         "fmax": fmax,
         "opt_cell_lengths": opt_cell_lengths,
         "opt_cell_fully": opt_cell_fully,

examples/calculations/submit_md.py

@@ -41,12 +41,15 @@ def md(params: dict) -> None:
         "arch": Str(params["arch"]),
         "struct": structure,
         "model": model,
-        "precision": Str(params["precision"]),
         "device": Str(params["device"]),
         "ensemble": Str(params["ensemble"]),
         "md_kwargs": Dict(params["md_dict"]),
     }
 
+    # Only calc_kwargs add if set
+    if params["calc_kwargs"]:
+        inputs["calc_kwargs"] = Dict(params["calc_kwargs"])
+
     # Run calculation
     result, node = run_get_node(MDCalc, **inputs)
     print(f"Printing results from calculation: {result}")
@@ -78,7 +81,10 @@ def md(params: dict) -> None:
     "--device", default="cpu", type=str, help="Device to run calculations on."
 )
 @click.option(
-    "--precision", default="float64", type=str, help="Chosen level of precision."
+    "--calc-kwargs",
+    default="{}",
+    type=str,
+    help="Keyword arguments to pass to calculator.",
 )
 @click.option(
     "--ensemble", default="nve", type=str, help="Name of thermodynamic ensemble."
@@ -90,10 +96,12 @@ def md(params: dict) -> None:
     help="String containing a dictionary with other md parameters",
 )
 def cli(
-    codelabel, struct, model, arch, device, precision, ensemble, md_dict_str
+    codelabel, struct, model, arch, device, calc_kwargs, ensemble, md_dict_str
 ) -> None:
     """Click interface."""
     md_dict = ast.literal_eval(md_dict_str)
+    calc_kwargs = ast.literal_eval(calc_kwargs)
+
     try:
         code = load_code(codelabel)
     except NotExistent as exc:
@@ -106,7 +114,7 @@ def cli(
         "model": model,
         "arch": arch,
         "device": device,
-        "precision": precision,
+        "calc_kwargs": calc_kwargs,
         "ensemble": ensemble,
         "md_dict": md_dict,
     }

examples/calculations/submit_singlepoint.py

@@ -2,9 +2,11 @@
 
 from __future__ import annotations
 
+import ast
+
 from aiida.common import NotExistent
 from aiida.engine import run_get_node
-from aiida.orm import Str, load_code
+from aiida.orm import Dict, Str, load_code
 from aiida.plugins import CalculationFactory
 import click
 
@@ -39,10 +41,13 @@ def singlepoint(params: dict) -> None:
         "arch": Str(params["arch"]),
         "struct": structure,
         "model": model,
-        "precision": Str(params["precision"]),
         "device": Str(params["device"]),
     }
 
+    # Only calc_kwargs add if set
+    if params["calc_kwargs"]:
+        inputs["calc_kwargs"] = Dict(params["calc_kwargs"])
+
     # Run calculation
     result, node = run_get_node(SinglepointCalc, **inputs)
     print(f"Printing results from calculation: {result}")
@@ -74,10 +79,15 @@ def singlepoint(params: dict) -> None:
     "--device", default="cpu", type=str, help="Device to run calculations on."
 )
 @click.option(
-    "--precision", default="float64", type=str, help="Chosen level of precision."
+    "--calc-kwargs",
+    default="{}",
+    type=str,
+    help="Keyword arguments to pass to calculator.",
 )
-def cli(codelabel, struct, model, arch, device, precision) -> None:
+def cli(codelabel, struct, model, arch, device, calc_kwargs) -> None:
     """Click interface."""
+    calc_kwargs = ast.literal_eval(calc_kwargs)
+
     try:
         code = load_code(codelabel)
     except NotExistent as exc:
@@ -90,7 +100,7 @@ def cli(codelabel, struct, model, arch, device, precision) -> None:
         "model": model,
         "arch": arch,
         "device": device,
-        "precision": precision,
+        "calc_kwargs": calc_kwargs,
     }
 
     # Submit single point