102 add md calcjob

README.md

@@ -17,7 +17,7 @@ machine learning interatomic potentials aiida plugin
   - CHGNET
 - [x] Single point calculations
 - [x] Geometry optimisation
-- [ ] Molecular Dynamics:
+- [x] Molecular Dynamics:
   - NVE
   - NVT (Langevin(Eijnden/Ciccotti flavour) and Nosé-Hoover (Melchionna flavour))
   - NPT (Nosé-Hoover (Melchiona flavour))
@@ -86,9 +86,12 @@ See the [developer guide](https://stfc.github.io/aiida-mlip/developer_guide/inde
   * [`calculations/`](aiida_mlip/calculations/): Plugin `Calcjob` classes
     * [`singlepoint.py`](aiida_mlip/calculations/singlepoint.py): `Calcjob` class to run single point calculations using mlips
     * [`geomopt.py`](aiida_mlip/calculations/geomopt.py): `Calcjob` class to perform geometry optimization using mlips
+    * [`base.py`](aiida_mlip/calculations/base.py): Base `Calcjob` class for other calculations
+    * [`md.py`](aiida_mlip/calculations/md.py): `Calcjob` class to perform molecular dynamics using mlips
   * [`parsers/`](aiida_mlip/parsers/): `Parsers` for the calculations
     * [`sp_parser.py`](aiida_mlip/parsers/sp_parser.py): `Parser` for `Singlepoint` calculation.
     * [`opt_parser.py`](aiida_mlip/parsers/opt_parser.py): `Parser` for `Geomopt` calculation.
+    * [`md_parser.py`](aiida_mlip/parsers/md_parser.py): `Parser` for `MD` calculation.
   * [`helpers/`](aiida_mlip/helpers/): `Helpers` to run calculations.
 * [`docs/`](docs/source/): Code documentation
   * [`apidoc/`](docs/source/apidoc/): API documentation

aiida_mlip/calculations/base.py

@@ -0,0 +1,222 @@
+"""Base class for features common to most calculations."""
+
+from ase.io import write
+
+from aiida.common import datastructures
+import aiida.common.folders
+from aiida.engine import CalcJob, CalcJobProcessSpec
+import aiida.engine.processes
+from aiida.orm import SinglefileData, Str, StructureData
+
+from aiida_mlip.data.model import ModelData
+
+
+class BaseJanus(CalcJob):  # numpydoc ignore=PR01
+    """
+    Calcjob implementation to run single point calculations using mlips.
+
+    Attributes
+    ----------
+    _DEFAULT_OUTPUT_FILE : str
+        Default stdout file name.
+    _DEFAULT_INPUT_FILE : str
+        Default input file name.
+    _LOG_FILE : str
+        Default log file name.
+
+    Methods
+    -------
+    define(spec: CalcJobProcessSpec) -> None:
+        Define the process specification, its inputs, outputs and exit codes.
+    validate_inputs(value: dict, port_namespace: PortNamespace) -> Optional[str]:
+        Check if the inputs are valid.
+    prepare_for_submission(folder: Folder) -> CalcInfo:
+        Create the input files for the `CalcJob`.
+    """
+
+    _DEFAULT_OUTPUT_FILE = "aiida-stdout.txt"
+    _DEFAULT_INPUT_FILE = "aiida.xyz"
+    _LOG_FILE = "aiida.log"
+
+    @classmethod
+    def define(cls, spec: CalcJobProcessSpec) -> None:
+        """
+        Define the process specification, its inputs, outputs and exit codes.
+
+        Parameters
+        ----------
+        spec : `aiida.engine.CalcJobProcessSpec`
+            The calculation job process spec to define.
+        """
+        super().define(spec)
+
+        # Define inputs
+        spec.input(
+            "architecture",
+            valid_type=Str,
+            default=lambda: Str("mace"),
+            help="Mlip architecture to use for calculation, defaults to mace",
+        )
+        spec.input(
+            "model",
+            valid_type=ModelData,
+            required=False,
+            help="Mlip model used for calculation",
+        )
+        spec.input("structure", valid_type=StructureData, help="The input structure.")
+        spec.input("precision", valid_type=Str, help="Precision level for calculation")
+        spec.input(
+            "device",
+            valid_type=Str,
+            required=False,
+            default=lambda: Str("cpu"),
+            help="Device on which to run calculation (cpu, cuda or mps)",
+        )
+
+        spec.input(
+            "log_filename",
+            valid_type=Str,
+            required=False,
+            default=lambda: Str(cls._LOG_FILE),
+            help="Name of the log output file",
+        )
+        spec.input(
+            "metadata.options.output_filename",
+            valid_type=str,
+            default=cls._DEFAULT_OUTPUT_FILE,
+        )
+        spec.input(
+            "metadata.options.input_filename",
+            valid_type=str,
+            default=cls._DEFAULT_INPUT_FILE,
+        )
+        spec.input(
+            "metadata.options.scheduler_stdout",
+            valid_type=str,
+            default="_scheduler-stdout.txt",
+            help="Filename to which the content of stdout of the scheduler is written.",
+        )
+        spec.inputs.validator = cls.validate_inputs
+
+        spec.output("std_output", valid_type=SinglefileData)
+        spec.output("log_output", valid_type=SinglefileData)
+        # Exit codes
+        spec.exit_code(
+            305,
+            "ERROR_MISSING_OUTPUT_FILES",
+            message="Some output files missing or cannot be read",
+        )
+
+    @classmethod
+    def validate_inputs(
+        cls, inputs: dict, port_namespace: aiida.engine.processes.ports.PortNamespace
+    ):
+        """
+        Check if the inputs are valid.
+
+        Parameters
+        ----------
+        inputs : dict
+            The inputs dictionary.
+
+        port_namespace : `aiida.engine.processes.ports.PortNamespace`
+            An instance of aiida's `PortNameSpace`.
+
+        Raises
+        ------
+        ValueError
+            Error message if validation fails, None otherwise.
+        """
+        # Wrapping processes may choose to exclude certain input ports
+        # If the ports have been excluded, skip the validation.
+        if "structure" not in port_namespace:
+            raise ValueError("'Structure' namespaces is required.")
+
+        if "input_filename" in inputs:
+            if not inputs["input_filename"].value.endswith(".xyz"):
+                raise ValueError(
+                    "The parameter 'input_filename' must end with '.xyz'"
+                )  #
+
+    # pylint: disable=too-many-locals
+    def prepare_for_submission(
+        self, folder: aiida.common.folders.Folder
+    ) -> datastructures.CalcInfo:
+        """
+        Create the input files for the `Calcjob`.
+
+        Parameters
+        ----------
+        folder : aiida.common.folders.Folder
+            An `aiida.common.folders.Folder` to temporarily write files on disk.
+
+        Returns
+        -------
+        aiida.common.datastructures.CalcInfo
+            An instance of `aiida.common.datastructures.CalcInfo`.
+        """
+
+        # Create needed inputs
+        # Define architecture from model if model is given,
+        # otherwise get architecture from inputs and download default model
+        architecture = (
+            str((self.inputs.model).architecture)
+            if self.inputs.model
+            else str(self.inputs.architecture.value)
+        )
+        if self.inputs.model:
+            model_path = self.inputs.model.filepath
+        else:
+            model_path = ModelData.download(
+                "https://github.com/stfc/janus-core/raw/main/tests/models/mace_mp_small.model",  # pylint:disable=line-too-long
+                architecture,
+            ).filepath
+
+        # The inputs are saved in the node, but we want their value as a string
+        precision = (self.inputs.precision).value
+        device = (self.inputs.device).value
+        input_filename = self.inputs.metadata.options.input_filename
+        log_filename = (self.inputs.log_filename).value
+
+        # Transform the structure data in xyz file called input_filename
+        structure = self.inputs.structure
+
+        atoms = structure.get_ase()
+        with folder.open(input_filename, "w", encoding="utf-8") as inputfile:
+            write(inputfile, images=atoms)
+
+        cmd_line = {
+            "arch": architecture,
+            "struct": input_filename,
+            "device": device,
+            "log": log_filename,
+            "calc-kwargs": {"model": model_path, "default_dtype": precision},
+        }
+
+        codeinfo = datastructures.CodeInfo()
+
+        # Initialize cmdline_params as an empty list
+        codeinfo.cmdline_params = []
+
+        # Adding command line params for when we run janus
+        codeinfo.cmdline_params.append("calculation")
+        for flag, value in cmd_line.items():
+            codeinfo.cmdline_params += [f"--{flag}", str(value)]
+
+        # Node where the code is saved
+        codeinfo.code_uuid = self.inputs.code.uuid
+        # Save name of output as you need it for running the code
+        codeinfo.stdout_name = self.metadata.options.output_filename
+
+        calcinfo = datastructures.CalcInfo()
+        calcinfo.codes_info = [codeinfo]
+        # Save the info about the node where the calc is stored
+        calcinfo.uuid = str(self.uuid)
+        # Retrieve output files
+        calcinfo.retrieve_list = [
+            self.metadata.options.output_filename,
+            self.uuid,
+            log_filename,
+        ]
+
+        return calcinfo

aiida_mlip/calculations/md.py

@@ -0,0 +1,111 @@
+"""Class to run md calculations."""
+
+from aiida.common import datastructures
+import aiida.common.folders
+from aiida.engine import CalcJobProcessSpec
+import aiida.engine.processes
+from aiida.orm import Dict, SinglefileData, Str, StructureData, TrajectoryData
+
+from aiida_mlip.calculations.base import BaseJanus
+
+
+class MD(BaseJanus):  # numpydoc ignore=PR01
+    """
+    Calcjob implementation to run geometry MD calculations using mlips.
+
+    Methods
+    -------
+    define(spec: CalcJobProcessSpec) -> None:
+        Define the process specification, its inputs, outputs and exit codes.
+    prepare_for_submission(folder: Folder) -> CalcInfo:
+        Create the input files for the `CalcJob`.
+    """
+
+    _DEFAULT_TRAJ_FILE = "aiida-traj.xyz"
+    _DEFAULT_STATS_FILE = "aiida-stats.dat"
+
+    @classmethod
+    def define(cls, spec: CalcJobProcessSpec) -> None:
+        """
+        Define the process specification, its inputs, outputs and exit codes.
+
+        Parameters
+        ----------
+        spec : `aiida.engine.CalcJobProcessSpec`
+            The calculation job process spec to define.
+        """
+        super().define(spec)
+
+        # Additional inputs for molecula dynamics
+        spec.input(
+            "ensemble",
+            valid_type=Str,
+            required=True,
+            help="Name for thermodynamic ensemble",
+        )
+
+        spec.input(
+            "md_dict",
+            valid_type=Dict,
+            required=False,
+            default=lambda: Dict(
+                {
+                    "traj-file": cls._DEFAULT_TRAJ_FILE,
+                    "stats-file": cls._DEFAULT_STATS_FILE,
+                }
+            ),
+            help="Keywords for molecular dynamics",
+        )
+
+        spec.inputs["metadata"]["options"]["parser_name"].default = "janus.md_parser"
+
+        spec.output("stats_file", valid_type=SinglefileData)
+        spec.output("traj_file", valid_type=SinglefileData)
+        spec.output("traj_output", valid_type=TrajectoryData)
+        spec.output("final_structure", valid_type=StructureData)
+
+    def prepare_for_submission(
+        self, folder: aiida.common.folders.Folder
+    ) -> datastructures.CalcInfo:
+        """
+        Create the input files for the `Calcjob`.
+
+        Parameters
+        ----------
+        folder : aiida.common.folders.Folder
+            An `aiida.common.folders.Folder` to temporarily write files on disk.
+
+        Returns
+        -------
+        aiida.common.datastructures.CalcInfo
+            An instance of `aiida.common.datastructures.CalcInfo`.
+        """
+
+        # Call the parent class method to prepare common inputs
+        calcinfo = super().prepare_for_submission(folder)
+        codeinfo = calcinfo.codes_info[0]
+
+        md_dictionary = self.inputs.md_dict.get_dict()
+
+        if not "traj-file" in md_dictionary:
+            md_dictionary["traj-file"] = str(self._DEFAULT_TRAJ_FILE)
+        if not "stats-file" in md_dictionary:
+            md_dictionary["stats-file"] = str(self._DEFAULT_STATS_FILE)
+
+        ensemble = self.inputs.ensemble.value.lower()
+        codeinfo.cmdline_params[0] = "md"
+
+        codeinfo.cmdline_params += ["--ensemble", ensemble]
+
+        for flag, value in md_dictionary.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]
+
+        calcinfo.retrieve_list.append(md_dictionary["traj-file"])
+        calcinfo.retrieve_list.append(md_dictionary["stats-file"])
+
+        return calcinfo