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

@@ -6,47 +6,42 @@
 ###############################################################################
 """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):
         """Receives in input a dictionary of retrieved nodes. Does all the logic here."""
 
         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

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()]

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,16 +160,18 @@ 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")
 @click.argument("codelabel")
 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)