Implement trajectory parser, provide example of MD simulations.

aiida_cp2k/calculations/__init__.py

@@ -40,6 +40,9 @@ class Cp2kCalculation(CalcJob):
     _DEFAULT_PROJECT_NAME = "aiida"
     _DEFAULT_RESTART_FILE_NAME = _DEFAULT_PROJECT_NAME + "-1.restart"
     _DEFAULT_TRAJECT_FILE_NAME = _DEFAULT_PROJECT_NAME + "-pos-1.dcd"
+    _DEFAULT_TRAJECT_XYZ_FILE_NAME = _DEFAULT_PROJECT_NAME + "-pos-1.xyz"
+    _DEFAULT_TRAJECT_FORCES_FILE_NAME = _DEFAULT_PROJECT_NAME + "-frc-1.xyz"
+    _DEFAULT_TRAJECT_CELL_FILE_NAME = _DEFAULT_PROJECT_NAME + "-1.cell"
     _DEFAULT_PARENT_CALC_FLDR_NAME = "parent_calc/"
     _DEFAULT_COORDS_FILE_NAME = "aiida.coords.xyz"
     _DEFAULT_PARSER = "cp2k_base_parser"
@@ -162,6 +165,24 @@ def define(cls, spec):
             "ERROR_STRUCTURE_PARSE",
             message="The output structure could not be parsed.",
         )
+        spec.exit_code(
+            321,
+            "ERROR_COORDINATES_TRAJECTORY_READ",
+            message="The coordinates trajectory file could not be read.",
+        )
+
+        spec.exit_code(
+            323,
+            "ERROR_FORCES_TRAJECTORY_READ",
+            message="The forces trajectory file could not be read.",
+        )
+
+        spec.exit_code(
+            325,
+            "ERROR_CELLS_TRAJECTORY_READ",
+            message="The cells trajectory file could not be read.",
+        )
+
         spec.exit_code(
             350,
             "ERROR_UNEXPECTED_PARSER_EXCEPTION",
@@ -329,6 +350,9 @@ def prepare_for_submission(self, folder):
             self._DEFAULT_OUTPUT_FILE,
             self._DEFAULT_RESTART_FILE_NAME,
             self._DEFAULT_TRAJECT_FILE_NAME,
+            self._DEFAULT_TRAJECT_XYZ_FILE_NAME,
+            self._DEFAULT_TRAJECT_FORCES_FILE_NAME,
+            self._DEFAULT_TRAJECT_CELL_FILE_NAME,
         ]
         calcinfo.retrieve_list += settings.pop("additional_retrieve_list", [])
 

aiida_cp2k/parsers/__init__.py

@@ -7,6 +7,7 @@
 """AiiDA-CP2K output parser."""
 
 import ase
+import numpy as np
 from aiida import common, engine, orm, parsers, plugins
 
 from .. import utils
@@ -29,18 +30,30 @@ def parse(self, **kwargs):
         exit_code = self._parse_stdout()
 
         # Even though the simpulation might have failed, we still want to parse the output structure.
+        last_structure = None
         try:
             last_structure = self._parse_final_structure()
             if isinstance(last_structure, StructureData):
                 self.out("output_structure", last_structure)
         except common.NotExistent:
-            last_structure = None
-            self.logger.warning("No Restart file found in the retrieved folder.")
+            self.logger.warning("No restart file found in the retrieved folder.")
+
+        trajectory = None
+        try:
+            if last_structure is not None:
+                trajectory = self._parse_trajectory(last_structure)
+                if isinstance(trajectory, orm.TrajectoryData):
+                    self.out("output_trajectory", trajectory)
+        except common.NotExistent:
+            self.logger.warning("No trajectory file found in the retrieved folder.")
 
         if exit_code is not None:
             return exit_code
         if isinstance(last_structure, engine.ExitCode):
             return last_structure
+        if isinstance(trajectory, engine.ExitCode):
+            return trajectory
+
         return engine.ExitCode(0)
 
     def _parse_stdout(self):
@@ -108,10 +121,80 @@ def _read_stdout(self):
         try:
             output_string = self.retrieved.base.repository.get_object_content(fname)
         except OSError:
-            return self.exit_codes.ERROR_OUTPUT_STDOUT_READ, None
+            return self.exit_codes.ERROR_OUTPUT_READ, None
 
         return None, output_string
 
+    def _parse_trajectory(self, structure):
+        """CP2K trajectory parser."""
+
+        symbols = [str(site.kind_name) for site in structure.sites]
+
+        # Handle the positions trajectory
+        xyz_traj_fname = self.node.process_class._DEFAULT_TRAJECT_XYZ_FILE_NAME
+
+        # Read the trajectory file.
+        try:
+            output_xyz_pos = self.retrieved.base.repository.get_object_content(
+                xyz_traj_fname
+            )
+        except OSError:
+            return self.exit_codes.ERROR_COORDINATES_TRAJECTORY_READ
+
+        from cp2k_output_tools.trajectories.xyz import parse
+
+        positions_traj = []
+        stepids_traj = []
+        for frame in parse(output_xyz_pos):
+            _, positions = zip(*frame["atoms"])
+            positions_traj.append(positions)
+            stepids_traj.append(int(frame["comment"].split()[2][:-1]))
+        positions_traj = np.array(positions_traj)
+        stepids_traj = np.array(stepids_traj)
+
+        cell_traj = None
+        cell_traj_fname = self.node.process_class._DEFAULT_TRAJECT_CELL_FILE_NAME
+        try:
+            if cell_traj_fname in self.retrieved.base.repository.list_object_names():
+                output_cell_pos = self.retrieved.base.repository.get_object_content(
+                    cell_traj_fname
+                )
+                cell_traj = np.array(
+                    [
+                        np.fromstring(line, sep=" ")[2:-1].reshape(3, 3)
+                        for line in output_cell_pos.splitlines()[1:]
+                    ]
+                )
+        except OSError:
+            return self.exit_codes.ERROR_CELLS_TRAJECTORY_READ
+
+        forces_traj = None
+        forces_traj_fname = self.node.process_class._DEFAULT_TRAJECT_FORCES_FILE_NAME
+        try:
+            if forces_traj_fname in self.retrieved.base.repository.list_object_names():
+                output_forces = self.retrieved.base.repository.get_object_content(
+                    forces_traj_fname
+                )
+                forces_traj = []
+                for frame in parse(output_forces):
+                    _, forces = zip(*frame["atoms"])
+                    forces_traj.append(forces)
+                forces_traj = np.array(forces_traj)
+        except OSError:
+            return self.exit_codes.ERROR_FORCES_TRAJECTORY_READ
+
+        trajectory = orm.TrajectoryData()
+        trajectory.set_trajectory(
+            stepids=stepids_traj,
+            cells=cell_traj,
+            symbols=symbols,
+            positions=positions_traj,
+        )
+        if forces_traj is not None:
+            trajectory.set_array("forces", forces_traj)
+
+        return trajectory
+
 
 class Cp2kAdvancedParser(Cp2kBaseParser):
     """Advanced AiiDA parser class for the output of CP2K."""

examples/single_calculations/example_mm_md.py

@@ -0,0 +1,153 @@
+###############################################################################
+# Copyright (c), The AiiDA-CP2K authors.                                      #
+# SPDX-License-Identifier: MIT                                                #
+# AiiDA-CP2K is hosted on GitHub at https://github.com/aiidateam/aiida-cp2k   #
+# For further information on the license, see the LICENSE.txt file.           #
+###############################################################################
+"""Run molecular dynamics calculation."""
+
+import os
+import sys
+
+import ase.io
+import click
+from aiida import common, engine, orm
+
+
+def example_mm(cp2k_code):
+    """Run molecular mechanics calculation."""
+
+    print("Testing CP2K ENERGY on H2O (MM) ...")
+
+    # Force field.
+    with open(os.path.join("/tmp", "water.pot"), "w") as f:
+        f.write(
+            """BONDS
+    H    H       0.000     1.5139
+    O    H     450.000     0.9572
+
+    ANGLES
+    H    O    H      55.000   104.5200
+
+    DIHEDRALS
+
+    IMPROPER
+
+    NONBONDED
+    H      0.000000  -0.046000     0.224500
+    O      0.000000  -0.152100     1.768200
+
+    HBOND CUTHB 0.5
+
+    END"""
+        )
+
+    water_pot = orm.SinglefileData(file=os.path.join("/tmp", "water.pot"))
+
+    thisdir = os.path.dirname(os.path.realpath(__file__))
+
+    # structure using pdb format, because it also carries topology information
+    atoms = ase.io.read(os.path.join(thisdir, "..", "files", "h2o.xyz"))
+    atoms.center(vacuum=10.0)
+    atoms.write(os.path.join("/tmp", "coords.pdb"), format="proteindatabank")
+    coords_pdb = orm.SinglefileData(file=os.path.join("/tmp", "coords.pdb"))
+
+    # Parameters.
+    # Based on cp2k/tests/Fist/regtest-1-1/water_1.inp
+    parameters = orm.Dict(
+        {
+            "FORCE_EVAL": {
+                "METHOD": "fist",
+                "STRESS_TENSOR": "analytical",
+                "MM": {
+                    "FORCEFIELD": {
+                        "PARM_FILE_NAME": "water.pot",
+                        "PARMTYPE": "CHM",
+                        "CHARGE": [
+                            {"ATOM": "O", "CHARGE": -0.8476},
+                            {"ATOM": "H", "CHARGE": 0.4238},
+                        ],
+                    },
+                    "POISSON": {
+                        "EWALD": {
+                            "EWALD_TYPE": "spme",
+                            "ALPHA": 0.44,
+                            "GMAX": 24,
+                            "O_SPLINE": 6,
+                        }
+                    },
+                },
+                "SUBSYS": {
+                    "CELL": {
+                        "ABC": "%f  %f  %f" % tuple(atoms.cell.diagonal()),
+                    },
+                    "TOPOLOGY": {
+                        "COORD_FILE_NAME": "coords.pdb",
+                        "COORD_FILE_FORMAT": "PDB",
+                    },
+                },
+            },
+            "MOTION": {
+                "CONSTRAINT": {},
+                "MD": {
+                    "THERMOSTAT": {"CSVR": {}, "TYPE": "csvr"},
+                    "BAROSTAT": {},
+                    "STEPS": 1000,
+                    "ENSEMBLE": "npt_f",
+                    "TEMPERATURE": 300.0,
+                },
+                "PRINT": {
+                    "TRAJECTORY": {"EACH": {"MD": 5}},
+                    "RESTART": {"EACH": {"MD": 5}},
+                    "RESTART_HISTORY": {"_": "OFF"},
+                    "CELL": {"EACH": {"MD": 5}},
+                    "FORCES": {"EACH": {"MD": 5}, "FORMAT": "XYZ"},
+                },
+            },
+            "GLOBAL": {
+                "CALLGRAPH": "master",
+                "CALLGRAPH_FILE_NAME": "runtime",
+                "PRINT_LEVEL": "medium",
+                "RUN_TYPE": "MD",
+            },
+        }
+    )
+
+    # Construct process builder.
+    builder = cp2k_code.get_builder()
+    builder.parameters = parameters
+    builder.code = cp2k_code
+    builder.file = {
+        "water_pot": water_pot,
+        "coords_pdb": coords_pdb,
+    }
+    builder.metadata.options.resources = {
+        "num_machines": 1,
+        "num_mpiprocs_per_machine": 1,
+    }
+    builder.metadata.options.max_wallclock_seconds = 1 * 3 * 60
+
+    print("Submitted calculation...")
+    results = engine.run(builder)
+    assert "output_trajectory" in results, "Output trajectory not found among results."
+    traj = results["output_trajectory"]
+
+    assert traj.get_cells().shape == (201, 3, 3), "Unexpected shape of cells."
+    assert traj.get_positions().shape == (201, 3, 3), "Unexpected shape of positions."
+    assert traj.get_array("forces").shape == (201, 3, 3), "Unexpected shape of forces."
+
+
+@click.command("cli")
+@click.argument("codelabel")
+def cli(codelabel):
+    """Click interface."""
+    try:
+        code = orm.load_code(codelabel)
+    except common.NotExistent:
+        print(f"The code '{codelabel}' does not exist.")
+        sys.exit(1)
+    example_mm(code)
+
+
+if __name__ == "__main__":
+    cli()