From 0d2d3b61abd22268a356ca8ea2390db4fe382ca1 Mon Sep 17 00:00:00 2001 From: Aliaksandr Yakutovich Date: Wed, 31 Jan 2024 11:18:31 +0100 Subject: [PATCH] Fix: retrieve the last structure even if a calculation has failed (#205) I also modified one of the tests to ensure this behaviour is checked. --- aiida_cp2k/parsers/__init__.py | 60 +++++++++---------- aiida_cp2k/utils/parser.py | 7 +-- .../single_calculations/example_restart.py | 33 +++++----- 3 files changed, 46 insertions(+), 54 deletions(-) diff --git a/aiida_cp2k/parsers/__init__.py b/aiida_cp2k/parsers/__init__.py index 26e133d..0da7111 100644 --- a/aiida_cp2k/parsers/__init__.py +++ b/aiida_cp2k/parsers/__init__.py @@ -6,23 +6,16 @@ ############################################################################### """AiiDA-CP2K output parser.""" -from aiida.common import exceptions -from aiida.engine import ExitCode -from aiida.orm import Dict +import ase +from aiida import common, engine, orm, parsers, plugins -# + -from aiida.parsers import Parser -from aiida.plugins import DataFactory +from .. import utils -from aiida_cp2k import utils +StructureData = plugins.DataFactory("core.structure") +BandsData = plugins.DataFactory("core.array.bands") -# - -StructureData = DataFactory("core.structure") -BandsData = DataFactory("core.array.bands") - - -class Cp2kBaseParser(Parser): +class Cp2kBaseParser(parsers.Parser): """Basic AiiDA parser for the output of CP2K.""" def parse(self, **kwargs): @@ -30,23 +23,25 @@ def parse(self, **kwargs): try: _ = self.retrieved - except exceptions.NotExistent: + except common.NotExistent: return self.exit_codes.ERROR_NO_RETRIEVED_FOLDER exit_code = self._parse_stdout() - if exit_code is not None: - return exit_code + # Even though the simpulation might have failed, we still want to parse the output structure. try: - returned = self._parse_trajectory() - if isinstance(returned, StructureData): - self.out("output_structure", returned) - else: # in case this is an error code - return returned - except exceptions.NotExistent: - pass + 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.") - return ExitCode(0) + if exit_code is not None: + return exit_code + if isinstance(last_structure, engine.ExitCode): + return last_structure + return engine.ExitCode(0) def _parse_stdout(self): """Basic CP2K output file parser.""" @@ -63,19 +58,16 @@ def _parse_stdout(self): # Parse the standard output. result_dict = utils.parse_cp2k_output(output_string) - self.out("output_parameters", Dict(dict=result_dict)) + self.out("output_parameters", orm.Dict(dict=result_dict)) return None - def _parse_trajectory(self): + def _parse_final_structure(self): """CP2K trajectory parser.""" - - from ase import Atoms - fname = self.node.process_class._DEFAULT_RESTART_FILE_NAME # Check if the restart file is present. if fname not in self.retrieved.base.repository.list_object_names(): - raise exceptions.NotExistent( + raise common.NotExistent( "No restart file available, so the output trajectory can't be extracted" ) @@ -85,7 +77,9 @@ def _parse_trajectory(self): except OSError: return self.exit_codes.ERROR_OUTPUT_STDOUT_READ - return StructureData(ase=Atoms(**utils.parse_cp2k_trajectory(output_string))) + return StructureData( + ase=ase.Atoms(**utils.parse_cp2k_trajectory(output_string)) + ) def _check_stdout_for_errors(self, output_string): """This function checks the CP2K output file for some basic errors.""" @@ -169,7 +163,7 @@ def _parse_stdout(self): ) self.out("output_bands", bnds) - self.out("output_parameters", Dict(dict=result_dict)) + self.out("output_parameters", orm.Dict(dict=result_dict)) return None @@ -208,5 +202,5 @@ def _parse_stdout(self): except KeyError: pass - self.out("output_parameters", Dict(dict=result_dict)) + self.out("output_parameters", orm.Dict(dict=result_dict)) return None diff --git a/aiida_cp2k/utils/parser.py b/aiida_cp2k/utils/parser.py index 774f4f1..aca942b 100644 --- a/aiida_cp2k/utils/parser.py +++ b/aiida_cp2k/utils/parser.py @@ -9,6 +9,8 @@ import math import re +import numpy as np + def parse_cp2k_output(fstring): """Parse CP2K output into a dictionary.""" @@ -323,8 +325,6 @@ def _parse_bands_cp2k_greater_81(lines, line_n): def _parse_bands(lines, n_start, cp2k_version): """Parse band structure from the CP2K output.""" - import numpy as np - kpoints = [] labels = [] bands_s1 = [] @@ -377,9 +377,6 @@ def _parse_bands(lines, n_start, cp2k_version): def parse_cp2k_trajectory(content): """CP2K trajectory parser.""" - - import numpy as np - # Parse coordinate section match = re.search(r"\n\s*&COORD\n(.*?)\n\s*&END COORD\n", content, re.DOTALL) coord_lines = [line.strip().split() for line in match.group(1).splitlines()] diff --git a/examples/single_calculations/example_restart.py b/examples/single_calculations/example_restart.py index 2c0d049..6b9d3cd 100644 --- a/examples/single_calculations/example_restart.py +++ b/examples/single_calculations/example_restart.py @@ -13,12 +13,9 @@ import ase.io import click -from aiida.common import NotExistent -from aiida.engine import run, run_get_node -from aiida.orm import Dict, SinglefileData, load_code -from aiida.plugins import DataFactory +from aiida import common, engine, orm, plugins -StructureData = DataFactory("core.structure") +StructureData = plugins.DataFactory("core.structure") def example_restart(cp2k_code): @@ -34,17 +31,17 @@ def example_restart(cp2k_code): ) # Basis set. - basis_file = SinglefileData( + basis_file = orm.SinglefileData( file=os.path.join(thisdir, "..", "files", "BASIS_MOLOPT") ) # Pseudopotentials. - pseudo_file = SinglefileData( + pseudo_file = orm.SinglefileData( file=os.path.join(thisdir, "..", "files", "GTH_POTENTIALS") ) # CP2K input. - params1 = Dict( + params1 = orm.Dict( { "GLOBAL": { "RUN_TYPE": "GEO_OPT", @@ -100,7 +97,6 @@ def example_restart(cp2k_code): } ) - # ------------------------------------------------------------------------------ # Construct process builder. builder = cp2k_code.get_builder() @@ -119,7 +115,7 @@ def example_restart(cp2k_code): builder.metadata.options.max_wallclock_seconds = 1 * 2 * 60 print("Submitted calculation 1.") - calc1_outputs, calc1 = run_get_node(builder) + calc1_outputs, calc1 = engine.run_get_node(builder) # Check walltime exceeded. if calc1.exit_status == 400: @@ -128,7 +124,10 @@ def example_restart(cp2k_code): print("FAIL, walltime wasn't exceeded as it should.") sys.exit(1) - # ------------------------------------------------------------------------------ + print(calc1_outputs) + assert "output_structure" in calc1_outputs, "The output_structure is missing." + print("OK, output_structure is present, even though the calculation has failed.") + # Set up and start the second calculation. # Parameters. @@ -139,7 +138,7 @@ def example_restart(cp2k_code): params2["FORCE_EVAL"]["DFT"]["RESTART_FILE_NAME"] = restart_wfn_fn params2["FORCE_EVAL"]["DFT"]["SCF"]["SCF_GUESS"] = "RESTART" params2["EXT_RESTART"] = {"RESTART_FILE_NAME": "./parent_calc/aiida-1.restart"} - params2 = Dict(params2) + params2 = orm.Dict(params2) # Structure. atoms2 = ase.io.read(os.path.join(thisdir, "..", "files", "h2o.xyz")) @@ -152,7 +151,7 @@ def example_restart(cp2k_code): builder.parent_calc_folder = calc1_outputs["remote_folder"] print("Submitted calculation 2.") - calc2 = run(builder) + calc2 = engine.run(builder) # Check energy. expected_energy = -17.1566455959 @@ -161,7 +160,9 @@ def example_restart(cp2k_code): # Ensure that this warning originates from overwritting coordinates. output = calc2["retrieved"].base.repository.get_object_content("aiida.out") - assert re.search("WARNING .* :: Overwriting coordinates", output) + assert re.search( + "WARNING .* :: Overwriting coordinates", output + ), "No warning about overwritting coordinates." @click.command("cli") @@ -169,8 +170,8 @@ def example_restart(cp2k_code): def cli(codelabel): """Click interface.""" try: - code = load_code(codelabel) - except NotExistent: + code = orm.load_code(codelabel) + except common.NotExistent: print(f"The code '{codelabel}' does not exist.") sys.exit(1) example_restart(code)