From bd7a8898dbc54c8d998d61373b4de52addd79869 Mon Sep 17 00:00:00 2001 From: Aliaksandr Yakutovich Date: Wed, 7 Feb 2024 14:19:22 +0000 Subject: [PATCH] Implement trajectory parser, provide example of MD simulations. --- aiida_cp2k/calculations/__init__.py | 24 +++ aiida_cp2k/parsers/__init__.py | 66 ++++++- examples/single_calculations/example_mm_md.py | 170 ++++++++++++++++++ 3 files changed, 259 insertions(+), 1 deletion(-) create mode 100644 examples/single_calculations/example_mm_md.py diff --git a/aiida_cp2k/calculations/__init__.py b/aiida_cp2k/calculations/__init__.py index e495158..5dea1bf 100644 --- a/aiida_cp2k/calculations/__init__.py +++ b/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", []) diff --git a/aiida_cp2k/parsers/__init__.py b/aiida_cp2k/parsers/__init__.py index 0da7111..3ae6d2c 100644 --- a/aiida_cp2k/parsers/__init__.py +++ b/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 @@ -36,11 +37,21 @@ def parse(self, **kwargs): except common.NotExistent: last_structure = None self.logger.warning("No Restart file found in the retrieved folder.") + + try: + trajectory = self._parse_trajectory(last_structure) + if isinstance(trajectory, orm.TrajectoryData): + self.out("output_trajectory", trajectory) + except common.NotExistent: + trajectory = None + 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,9 +119,62 @@ 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])) + + + cell_traj_fname = self.node.process_class._DEFAULT_TRAJECT_CELL_FILE_NAME + try: + output_cell_pos = self.retrieved.base.repository.get_object_content(cell_traj_fname) + except OSError: + return self.exit_codes.ERROR_CELLS_TRAJECTORY_READ + cell_traj = np.array([np.fromstring(l, sep=" ")[2:-1].reshape(3,3) for l in output_cell_pos.splitlines()[1:]]) + + + forces_traj_fname = self.node.process_class._DEFAULT_TRAJECT_FORCES_FILE_NAME + try: + output_forces = self.retrieved.base.repository.get_object_content(forces_traj_fname) + except OSError: + return self.exit_codes.ERROR_FORCES_TRAJECTORY_READ + forces_traj = [] + for frame in parse(output_forces): + _, forces = zip(*frame["atoms"]) + forces_traj.append(forces) + + trajectory = orm.TrajectoryData() + trajectory.set_trajectory( + stepids=np.array(stepids_traj), + cells=cell_traj, + symbols=symbols, + positions=np.array(positions_traj), + ) + trajectory.set_array("forces", np.array(forces_traj)) + + return trajectory class Cp2kAdvancedParser(Cp2kBaseParser): diff --git a/examples/single_calculations/example_mm_md.py b/examples/single_calculations/example_mm_md.py new file mode 100644 index 0000000..c197869 --- /dev/null +++ b/examples/single_calculations/example_mm_md.py @@ -0,0 +1,170 @@ +############################################################################### +# 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...") + calc = engine.run(builder) + + +@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()