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")
 
 
@@ -45,6 +47,8 @@ class Cp2kCalculation(CalcJob):
     _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_INPUT_TRAJECT_XYZ_FILE_NAME = "trajectory.xyz"
+    _DEFAULT_INPUT_CELL_FILE_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="Possible input trajectory fro REFTRAJMD.",
+        )
         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,16 @@ def prepare_for_submission(self, folder):
                 conflicting_keys=["COORDINATE"],
             )
 
+        # Craeate input trajectory files
+        if "trajectory" in self.inputs:
+            self._write_trajectories(
+                self.inputs.structure,
+                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 +414,16 @@ 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(structure, 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, structure.get_ase())
+        with open(folder.get_abs_path(name_pos), mode="w", encoding="utf-8") as fobj:
+            fobj.write(xyz)
+        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.
@@ -415,3 +451,29 @@ def _atoms_to_xyz(atoms):
     xyz = f"{len(elem_coords)}\n\n"
     xyz += "\n".join(map(add, elem_symbols, elem_coords))
     return xyz
+
+
+def _trajectory_to_xyz_and_cell(trajectory, atoms):
+    """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 = "#   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"
+    xyz = ""
+    elem_symbols = kind_names(atoms)
+
+    for i, step in enumerate(trajectory.get_array("positions")):
+        elem_coords = [f"{p[0]:25.16f} {p[1]:25.16f} {p[2]:25.16f}" for p in step]
+        xyz += f"{len(elem_coords)}\n"
+        xyz += f"i = {i+1} , time = {(i+1)*0.5} \n"
+        xyz += "\n".join(map(add, elem_symbols, elem_coords))
+        xyz += "\n"
+    if "cells" in trajectory.get_arraynames():
+        cell_vecs = [
+            f"{i+1} {(i+1)*0.5:6.3f} {p[0][0]:25.16f} {p[0][1]:25.16f} {p[0][2]:25.16f} {p[1][0]:25.16f} {p[1][1]:25.16f} {p[1][2]:25.16f} {p[2][0]:25.16f} {p[2][1]:25.16f} {p[2][2]:25.16f} {np.dot(p[0],np.cross(p[1],p[2]))}"
+            for (i, p) in enumerate(trajectory.get_array("cells"))
+        ]
+        cell += "\n".join(cell_vecs)
+    return xyz, cell

aiida_cp2k/utils/input_generator.py

@@ -207,9 +207,19 @@ 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"}
+    if params["GLOBAL"]["RUN_TYPE"] == "MD":
+        params["EXT_RESTART"]["RESTART_DEFAULT"] = ".TRUE."
+        params["EXT_RESTART"]["RESTART_COUNTERS"] = ".TRUE."
+        params["EXT_RESTART"]["RESTART_POS"] = ".TRUE."
+        params["EXT_RESTART"]["RESTART_VEL"] = ".TRUE."
+        params["EXT_RESTART"]["RESTART_CELL"] = ".TRUE."
+        params["EXT_RESTART"]["RESTART_THERMOSTAT"] = ".TRUE."
+        params["EXT_RESTART"]["RESTART_CONSTRAINT"] = ".FALSE."
     return Dict(params)

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/single_calculations/example_dft_md_reftraj.py

@@ -0,0 +1,179 @@
+###############################################################################
+# 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 simple DFT calculation."""
+
+import os
+import random
+import sys
+
+import ase.io
+import click
+import numpy as np
+from aiida.common import NotExistent
+from aiida.engine import run
+from aiida.orm import Dict, SinglefileData, load_code
+from aiida.plugins import DataFactory
+from ase import Atoms
+
+StructureData = DataFactory("core.structure")
+TrajectoryData = DataFactory("core.array.trajectory")
+
+
+def example_dft_md_reftraj(cp2k_code):
+    """Run simple DFT calculation."""
+
+    print("Testing CP2K MD REFTRAJ on H2 (DFT)...")
+
+    thisdir = os.path.dirname(os.path.realpath(__file__))
+
+    # Structure.
+    structure = StructureData(
+        ase=ase.io.read(os.path.join(thisdir, "..", "files", "h2.xyz"))
+    )
+
+    # Trajectory.
+    steps = 5
+    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)
+
+    # Basis set.
+    basis_file = SinglefileData(
+        file=os.path.join(thisdir, "..", "files", "BASIS_MOLOPT")
+    )
+
+    # Pseudopotentials.
+    pseudo_file = SinglefileData(
+        file=os.path.join(thisdir, "..", "files", "GTH_POTENTIALS")
+    )
+
+    # Parameters.
+    parameters = Dict(
+        {
+            "GLOBAL": {
+                "RUN_TYPE": "MD",
+                "PRINT_LEVEL": "LOW",
+                "WALLTIME": 100,
+                "PROJECT": "aiida",
+            },
+            "MOTION": {
+                "MD": {
+                    "ENSEMBLE": "REFTRAJ",
+                    "STEPS": steps,
+                    "REFTRAJ": {
+                        "FIRST_SNAPSHOT": 1,
+                        "LAST_SNAPSHOT": steps,
+                        "EVAL_FORCES": ".TRUE.",
+                        "TRAJ_FILE_NAME": "trajectory.xyz",
+                        "CELL_FILE_NAME": "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 = cp2k_code.get_builder()
+    builder.structure = structure
+    builder.trajectory = trajectory
+    builder.parameters = parameters
+    builder.code = cp2k_code
+    builder.file = {
+        "basis": basis_file,
+        "pseudo": pseudo_file,
+    }
+    builder.metadata.options.resources = {
+        "num_machines": 1,
+        "num_mpiprocs_per_machine": 1,
+    }
+    builder.metadata.options.max_wallclock_seconds = 1 * 3 * 60
+
+    print("Submitted calculation...")
+    run(builder)
+
+
+@click.command("cli")
+@click.argument("codelabel")
+def cli(codelabel):
+    """Click interface."""
+    try:
+        code = load_code(codelabel)
+    except NotExistent:
+        print(f"The code '{codelabel}' does not exist.")
+        sys.exit(1)
+    example_dft_md_reftraj(code)
+
+
+if __name__ == "__main__":
+    cli()