Implement support for the REFTRAJ simpulation

aiida_cp2k/calculations/__init__.py

@@ -8,6 +8,7 @@
 
 from operator import add
 
+import numpy as np
 from aiida.common import CalcInfo, CodeInfo, InputValidationError
 from aiida.engine import CalcJob
 from aiida.orm import Dict, RemoteData, SinglefileData
@@ -25,6 +26,7 @@
 
 BandsData = DataFactory("core.array.bands")
 StructureData = DataFactory("core.structure")
+TrajectoryData = DataFactory("core.array.trajectory")
 KpointsData = DataFactory("core.array.kpoints")
 
 
@@ -44,7 +46,9 @@ class Cp2kCalculation(CalcJob):
     _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_COORDS_FILE_NAME = _DEFAULT_PROJECT_NAME + ".coords.xyz"
+    _DEFAULT_INPUT_TRAJECT_XYZ_FILE_NAME = _DEFAULT_PROJECT_NAME + "-reftraj.xyz"
+    _DEFAULT_INPUT_CELL_FILE_NAME = _DEFAULT_PROJECT_NAME + "-reftraj.cell"
     _DEFAULT_PARSER = "cp2k_base_parser"
 
     @classmethod
@@ -59,6 +63,12 @@ def define(cls, spec):
             required=False,
             help="The main input structure.",
         )
+        spec.input(
+            "trajectory",
+            valid_type=TrajectoryData,
+            required=False,
+            help="Input trajectory for a REFTRAJ simulation.",
+        )
         spec.input(
             "settings",
             valid_type=Dict,
@@ -219,6 +229,12 @@ def define(cls, spec):
             required=False,
             help="The relaxed output structure.",
         )
+        spec.output(
+            "output_trajectory",
+            valid_type=TrajectoryData,
+            required=False,
+            help="The output trajectory.",
+        )
         spec.output(
             "output_bands",
             valid_type=BandsData,
@@ -270,6 +286,15 @@ def prepare_for_submission(self, folder):
                 conflicting_keys=["COORDINATE"],
             )
 
+        # Create input trajectory files
+        if "trajectory" in self.inputs:
+            self._write_trajectories(
+                self.inputs.trajectory,
+                folder,
+                self._DEFAULT_INPUT_TRAJECT_XYZ_FILE_NAME,
+                self._DEFAULT_INPUT_CELL_FILE_NAME,
+            )
+
         if "basissets" in self.inputs:
             validate_basissets(
                 inp,
@@ -388,6 +413,19 @@ def _write_structure(structure, folder, name):
         with open(folder.get_abs_path(name), mode="w", encoding="utf-8") as fobj:
             fobj.write(xyz)
 
+    @staticmethod
+    def _write_trajectories(trajectory, folder, name_pos, name_cell):
+        """Function that writes a structure and takes care of element tags."""
+
+        (xyz, cell) = _trajectory_to_xyz_and_cell(trajectory)
+        with open(folder.get_abs_path(name_pos), mode="w", encoding="utf-8") as fobj:
+            fobj.write(xyz)
+        if cell is not None:
+            with open(
+                folder.get_abs_path(name_cell), mode="w", encoding="utf-8"
+            ) as fobj:
+                fobj.write(cell)
+
 
 def kind_names(atoms):
     """Get atom kind names from ASE atoms based on tags.
@@ -402,7 +440,7 @@ def kind_names(atoms):
     return list(map(add, atoms.get_chemical_symbols(), elem_tags))
 
 
-def _atoms_to_xyz(atoms):
+def _atoms_to_xyz(atoms, infoline="No info"):
     """Converts ASE atoms to string, taking care of element tags.
 
     :param atoms: ASE Atoms instance
@@ -412,6 +450,33 @@ def _atoms_to_xyz(atoms):
     elem_coords = [
         f"{p[0]:25.16f} {p[1]:25.16f} {p[2]:25.16f}" for p in atoms.get_positions()
     ]
-    xyz = f"{len(elem_coords)}\n\n"
+    xyz = f"{len(elem_coords)}\n"
+    xyz += f"{infoline}\n"
     xyz += "\n".join(map(add, elem_symbols, elem_coords))
     return xyz
+
+
+def _trajectory_to_xyz_and_cell(trajectory):
+    """Converts postions and cell from a TrajectoryData  to string, taking care of element tags from ASE atoms.
+
+    :param atoms: ASE Atoms instance
+    :param trajectory: TrajectoryData instance
+    :returns: positions str (in xyz format) and cell str
+    """
+    cell = None
+    xyz = ""
+    stepids = trajectory.get_stepids()
+    for i, step in enumerate(stepids):
+        xyz += _atoms_to_xyz(
+            trajectory.get_step_structure(i).get_ase(),
+            infoline=f"i = {step+1} , time = {(step+1)*0.5}",  # reftraj trajectories cannot start from STEP 0
+        )
+        xyz += "\n"
+    if "cells" in trajectory.get_arraynames():
+        cell = "#   Step   Time [fs]       Ax [Angstrom]       Ay [Angstrom]       Az [Angstrom]       Bx [Angstrom]       By [Angstrom]       Bz [Angstrom]       Cx [Angstrom]       Cy [Angstrom]       Cz [Angstrom]      Volume [Angstrom^3]\n"
+        cell_vecs = [
+            f"{stepid+1} {(stepid+1)*0.5:6.3f} {cellvec[0][0]:25.16f} {cellvec[0][1]:25.16f} {cellvec[0][2]:25.16f} {cellvec[1][0]:25.16f} {cellvec[1][1]:25.16f} {cellvec[1][2]:25.16f} {cellvec[2][0]:25.16f} {cellvec[2][1]:25.16f} {cellvec[2][2]:25.16f} {np.dot(cellvec[0],np.cross(cellvec[1],cellvec[2]))}"
+            for (stepid, cellvec) in zip(stepids, trajectory.get_array("cells"))
+        ]
+        cell += "\n".join(cell_vecs)
+    return xyz, cell

aiida_cp2k/utils/input_generator.py

@@ -207,9 +207,20 @@ def add_wfn_restart_section(input_dict, is_kpoints):
 
 
 @calcfunction
-def add_ext_restart_section(input_dict):
+def add_ext_restart_section(input_dict, first_snapshot=None):
     """Add external restart section to the input dictionary."""
     params = input_dict.get_dict()
+    if first_snapshot is not None:
+        params["MOTION"]["MD"]["REFTRAJ"]["FIRST_SNAPSHOT"] = first_snapshot
     # overwrite the complete EXT_RESTART section if present
-    params["EXT_RESTART"] = {"RESTART_FILE_NAME": "./parent_calc/aiida-1.restart"}
+    params["EXT_RESTART"] = {
+        "RESTART_FILE_NAME": "./parent_calc/aiida-1.restart",
+        "RESTART_DEFAULT": ".TRUE.",
+        "RESTART_COUNTERS": ".TRUE.",
+        "RESTART_POS": ".TRUE.",
+        "RESTART_VEL": ".TRUE.",
+        "RESTART_CELL": ".TRUE.",
+        "RESTART_THERMOSTAT": ".TRUE.",
+        "RESTART_CONSTRAINT": ".FALSE.",
+    }
     return Dict(params)

aiida_cp2k/utils/parser.py

@@ -124,7 +124,7 @@ def parse_cp2k_output_advanced(
 
         # If a tag has been detected, now read the following line knowing what they are
         if line_is in ["eigen_spin1_au", "eigen_spin2_au"]:
-            if "------" in line:
+            if "------" in line or "*** WARNING" in line:
                 continue
             splitted_line = line.split()
             try:

aiida_cp2k/workchains/base.py

@@ -72,7 +72,7 @@ def restart_incomplete_calculation(self, calc):
         content_string = calc.outputs.retrieved.base.repository.get_object_content(calc.base.attributes.get('output_filename'))
 
         # CP2K was updating geometry - continue with that.
-        restart_geometry_transformation = "Max. gradient              =" in content_string
+        restart_geometry_transformation = "Max. gradient              =" in content_string or "MD| Step number" in content_string
         end_inner_scf_loop = "Total energy: " in content_string
         # The message is written in the log file when the CP2K input parameter `LOG_PRINT_KEY` is set to True.
         if not (restart_geometry_transformation or end_inner_scf_loop or "Writing RESTART" in content_string):
@@ -89,7 +89,13 @@ def restart_incomplete_calculation(self, calc):
         params = add_wfn_restart_section(params, Bool('kpoints' in self.ctx.inputs))
 
         if restart_geometry_transformation:
-            params = add_ext_restart_section(params)
+            # Check if we need to fix restart snapshot in REFTRAJ MD
+            first_snapshot = None
+            try:
+                first_snapshot = int(params['MOTION']['MD']['REFTRAJ']['FIRST_SNAPSHOT']) + calc.outputs.output_trajectory.get_shape('positions')[0]
+            except KeyError:
+                pass
+            params = add_ext_restart_section(params, first_snapshot=first_snapshot)
 
         self.ctx.inputs.parameters = params  # params (new or old ones) that include the necessary restart information.
         self.report(

examples/workchains/example_base_md_reftraj_restart.py

@@ -0,0 +1,206 @@
+###############################################################################
+# 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.           #
+###############################################################################
+"""An example testing the restart calculation handler for geo_opt run in CP2K."""
+
+import os
+import random
+import sys
+
+import ase.io
+import click
+import numpy as np
+from aiida import common, engine, orm, plugins
+
+Cp2kBaseWorkChain = plugins.WorkflowFactory("cp2k.base")
+StructureData = plugins.DataFactory("core.structure")
+TrajectoryData = plugins.DataFactory("core.array.trajectory")
+
+
+def example_base(cp2k_code):
+    """Run simple DFT calculation through a workchain."""
+
+    thisdir = os.path.dirname(os.path.realpath(__file__))
+
+    print("Testing CP2K MD REFTRAJ on H2 (DFT) through a workchain...")
+
+    # Basis set.
+    basis_file = orm.SinglefileData(
+        file=os.path.join(thisdir, "..", "files", "BASIS_MOLOPT")
+    )
+
+    # Pseudopotentials.
+    pseudo_file = orm.SinglefileData(
+        file=os.path.join(thisdir, "..", "files", "GTH_POTENTIALS")
+    )
+
+    # Structure.
+    structure = StructureData(
+        ase=ase.io.read(os.path.join(thisdir, "..", "files", "h2.xyz"))
+    )
+
+    # Trajectory.
+    steps = 20
+    positions = np.array(
+        [[[2, 2, 2.73 + 0.05 * random.random()], [2, 2, 2]] for i in range(steps)]
+    )
+    cells = np.array(
+        [
+            [[4, 0, 0], [0, 4, 0], [0, 0, 4.75 + 0.05 * random.random()]]
+            for i in range(steps)
+        ]
+    )
+    symbols = ["H", "H"]
+    trajectory = TrajectoryData()
+    trajectory.set_trajectory(symbols, positions, cells=cells)
+
+    # Parameters.
+    parameters = orm.Dict(
+        {
+            "GLOBAL": {
+                "RUN_TYPE": "MD",
+                "PRINT_LEVEL": "LOW",
+                "WALLTIME": 4,
+                "PROJECT": "aiida",
+            },
+            "MOTION": {
+                "MD": {
+                    "ENSEMBLE": "REFTRAJ",
+                    "STEPS": steps,
+                    "REFTRAJ": {
+                        "FIRST_SNAPSHOT": 1,
+                        "LAST_SNAPSHOT": steps,
+                        "EVAL_FORCES": ".TRUE.",
+                        "TRAJ_FILE_NAME": "aiida-reftraj.xyz",
+                        "CELL_FILE_NAME": "aiida-reftraj.cell",
+                        "VARIABLE_VOLUME": ".TRUE.",
+                    },
+                },
+                "PRINT": {
+                    "RESTART": {
+                        "EACH": {
+                            "MD": 1,
+                        },
+                    },
+                    "FORCES": {
+                        "EACH": {
+                            "MD": 1,
+                        },
+                    },
+                    "CELL": {
+                        "EACH": {
+                            "MD": 1,
+                        },
+                    },
+                },
+            },
+            "FORCE_EVAL": {
+                "METHOD": "Quickstep",
+                "DFT": {
+                    "BASIS_SET_FILE_NAME": "BASIS_MOLOPT",
+                    "POTENTIAL_FILE_NAME": "GTH_POTENTIALS",
+                    "QS": {
+                        "EPS_DEFAULT": 1.0e-12,
+                        "WF_INTERPOLATION": "ps",
+                        "EXTRAPOLATION_ORDER": 3,
+                    },
+                    "MGRID": {
+                        "NGRIDS": 4,
+                        "CUTOFF": 280,
+                        "REL_CUTOFF": 30,
+                    },
+                    "XC": {
+                        "XC_FUNCTIONAL": {
+                            "_": "LDA",
+                        },
+                    },
+                    "POISSON": {
+                        "PERIODIC": "none",
+                        "PSOLVER": "MT",
+                    },
+                },
+                "SUBSYS": {
+                    "KIND": [
+                        {
+                            "_": "O",
+                            "BASIS_SET": "DZVP-MOLOPT-SR-GTH",
+                            "POTENTIAL": "GTH-LDA-q6",
+                        },
+                        {
+                            "_": "H",
+                            "BASIS_SET": "DZVP-MOLOPT-SR-GTH",
+                            "POTENTIAL": "GTH-LDA-q1",
+                        },
+                    ],
+                },
+            },
+        }
+    )
+
+    # Construct process builder.
+    builder = Cp2kBaseWorkChain.get_builder()
+
+    # Switch on resubmit_unconverged_geometry disabled by default.
+    builder.handler_overrides = orm.Dict(
+        {"restart_incomplete_calculation": {"enabled": True}}
+    )
+
+    # Input structure.
+    builder.cp2k.structure = structure
+    builder.cp2k.trajectory = trajectory
+    builder.cp2k.parameters = parameters
+    builder.cp2k.code = cp2k_code
+    builder.cp2k.file = {
+        "basis": basis_file,
+        "pseudo": pseudo_file,
+    }
+    builder.cp2k.metadata.options.resources = {
+        "num_machines": 1,
+        "num_mpiprocs_per_machine": 1,
+    }
+
+    print("Submitted calculation...")
+    calc = engine.run(builder)
+
+    if "EXT_RESTART" in calc["final_input_parameters"].dict:
+        print("OK, EXT_RESTART section is present in the final_input_parameters.")
+    else:
+        print(
+            "ERROR, EXT_RESTART section is NOT present in the final_input_parameters."
+        )
+        sys.exit(1)
+
+    stepids = np.concatenate(
+        [
+            called.outputs.output_trajectory.get_stepids()
+            for called in calc.called
+            if isinstance(called, orm.CalcJobNode)
+        ]
+    )
+
+    if np.all(stepids == np.arange(1, steps + 1)):
+        print("OK, stepids are correct.")
+    else:
+        print(
+            f"ERROR, stepids are NOT correct. Expected: {np.arange(1, steps + 1)} but got:  {stepids}"
+        )
+        sys.exit(1)
+
+
+@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_base(code)
+
+
+if __name__ == "__main__":
+    cli()