Implement support for the REFTRAJ simpulation (#207)

- Add optional `trajectory` input of type`TrajectoryData` to the inputs to the cp2k calculation,
which will further be transformed into `aiida-reftraj.xyz` and `aiida-reftraj.cell`.
- Update the restart handler that specifies the `EXT_RESTART` sections explicitly.
- Update the restart handler to make it understand that the MD simulation produced some steps.
- Add an example of a reftraj calculation that also does a restart.

---------
Co-authored-by: Aliaksandr Yakutovich <[email protected]>

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/__init__.py

@@ -6,7 +6,12 @@
 ###############################################################################
 """AiiDA-CP2K utils"""
 
-from .input_generator import Cp2kInput, add_ext_restart_section, add_wfn_restart_section
+from .input_generator import (
+    Cp2kInput,
+    add_ext_restart_section,
+    add_first_snapshot_in_reftraj_section,
+    add_wfn_restart_section,
+)
 from .parser import parse_cp2k_output, parse_cp2k_output_advanced, parse_cp2k_trajectory
 from .workchains import (
     HARTREE2EV,
@@ -23,6 +28,7 @@
 __all__ = [
     "Cp2kInput",
     "add_ext_restart_section",
+    "add_first_snapshot_in_reftraj_section",
     "add_wfn_restart_section",
     "parse_cp2k_output",
     "parse_cp2k_output_advanced",

aiida_cp2k/utils/input_generator.py

@@ -211,5 +211,22 @@ def add_ext_restart_section(input_dict):
     """Add external restart section to the input dictionary."""
     params = input_dict.get_dict()
     # 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)
+
+
+@calcfunction
+def add_first_snapshot_in_reftraj_section(input_dict, first_snapshot):
+    """Add first_snapshot in REFTRAJ section to the input dictionary."""
+    params = input_dict.get_dict()
+    params["MOTION"]["MD"]["REFTRAJ"]["FIRST_SNAPSHOT"] = first_snapshot
     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

@@ -1,21 +1,13 @@
 """Base work chain to run a CP2K calculation."""
 
-from aiida.common import AttributeDict
-from aiida.engine import (
-    BaseRestartWorkChain,
-    ProcessHandlerReport,
-    process_handler,
-    while_,
-)
-from aiida.orm import Bool, Dict
-from aiida.plugins import CalculationFactory
+from aiida import common, engine, orm, plugins
 
-from ..utils import add_ext_restart_section, add_wfn_restart_section
+from .. import utils
 
-Cp2kCalculation = CalculationFactory('cp2k')
+Cp2kCalculation = plugins.CalculationFactory('cp2k')
 
 
-class Cp2kBaseWorkChain(BaseRestartWorkChain):
+class Cp2kBaseWorkChain(engine.BaseRestartWorkChain):
     """Workchain to run a CP2K calculation with automated error handling and restarts."""
 
     _process_class = Cp2kCalculation
@@ -28,7 +20,7 @@ def define(cls, spec):
 
         spec.outline(
             cls.setup,
-            while_(cls.should_run_process)(
+            engine.while_(cls.should_run_process)(
                 cls.run_process,
                 cls.inspect_process,
                 cls.overwrite_input_structure,
@@ -37,7 +29,7 @@ def define(cls, spec):
         )
 
         spec.expose_outputs(Cp2kCalculation)
-        spec.output('final_input_parameters', valid_type=Dict, required=False,
+        spec.output('final_input_parameters', valid_type=orm.Dict, required=False,
                     help='The input parameters used for the final calculation.')
         spec.exit_code(400, 'NO_RESTART_DATA', message="The calculation didn't produce any data to restart from.")
         spec.exit_code(300, 'ERROR_UNRECOVERABLE_FAILURE',
@@ -52,7 +44,7 @@ def setup(self):
         internal loop.
         """
         super().setup()
-        self.ctx.inputs = AttributeDict(self.exposed_inputs(Cp2kCalculation, 'cp2k'))
+        self.ctx.inputs = common.AttributeDict(self.exposed_inputs(Cp2kCalculation, 'cp2k'))
 
     def results(self):
         super().results()
@@ -63,7 +55,7 @@ def overwrite_input_structure(self):
         if "output_structure" in self.ctx.children[self.ctx.iteration-1].outputs:
             self.ctx.inputs.structure = self.ctx.children[self.ctx.iteration-1].outputs.output_structure
 
-    @process_handler(priority=401, exit_codes=[
+    @engine.process_handler(priority=401, exit_codes=[
         Cp2kCalculation.exit_codes.ERROR_OUT_OF_WALLTIME,
         Cp2kCalculation.exit_codes.ERROR_OUTPUT_INCOMPLETE,
     ], enabled=False)
@@ -72,7 +64,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):
@@ -81,18 +73,26 @@ def restart_incomplete_calculation(self, calc):
                         "Sending a signal to stop the Base work chain.")
 
             # Signaling to the base work chain that the problem could not be recovered.
-            return ProcessHandlerReport(True, self.exit_codes.NO_RESTART_DATA)
+            return engine.ProcessHandlerReport(True, self.exit_codes.NO_RESTART_DATA)
 
         self.ctx.inputs.parent_calc_folder = calc.outputs.remote_folder
         params = self.ctx.inputs.parameters
 
-        params = add_wfn_restart_section(params, Bool('kpoints' in self.ctx.inputs))
+        params = utils.add_wfn_restart_section(params, orm.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]
+                if first_snapshot:
+                    params = utils.add_first_snapshot_in_reftraj_section(params, first_snapshot)
+            except KeyError:
+                pass
+            params = utils.add_ext_restart_section(params)
 
         self.ctx.inputs.parameters = params  # params (new or old ones) that include the necessary restart information.
         self.report(
             "The CP2K calculation wasn't completed. The restart of the calculation might be able to "
             "fix the problem.")
-        return ProcessHandlerReport(False)
+        return engine.ProcessHandlerReport(False)