Introduce new restart_incomplete_calculation handler, improve parsi…

…ng (#156)

- Add `_read_stdout` and `_check_stdout_for_errors` methods of the base parser that
are generic and can be reused in all the derived parser. This makes the code shorter.
- Add `add_ext_restart_section` and `add_wfn_restart_section` capable of adding only
the corresponding sections.
- Add a bunch of exit codes to the `Cp2kBaseWorkChain`.
- Introduce a new `restart_incomplete_calculation` handler as a replacement for the
`resubmit_unconverged_geometry` one. Added a deprecation warning in the docstring
of the latter one. Note: every time `restart_incomplete_calculation` detects the ability
to restart a work chain, it launches `add_wfn_restart_section` and (sometimes)
`add_ext_restart_section`. This happens always even if the corresponding sections are
already present in the input dictionary.
- Adapt examples to the changes in the code.

aiida_cp2k/parsers/__init__.py

@@ -11,12 +11,16 @@
 import os
 from aiida.common import exceptions
 
+# +
 from aiida.parsers import Parser
 from aiida.common import OutputParsingError, NotExistent
 from aiida.engine import ExitCode
 from aiida.orm import Dict
 from aiida.plugins import DataFactory
 
+from aiida_cp2k import utils
+# -
+
 StructureData = DataFactory('structure')  # pylint: disable=invalid-name
 BandsData = DataFactory('array.bands')  # pylint: disable=invalid-name
 
@@ -50,33 +54,25 @@ def parse(self, **kwargs):
     def _parse_stdout(self):
         """Basic CP2K output file parser."""
 
-        from aiida_cp2k.utils import parse_cp2k_output
-
-        fname = self.node.get_attribute('output_filename')
-
-        if fname not in self.retrieved.list_object_names():
-            return self.exit_codes.ERROR_OUTPUT_STDOUT_MISSING
-
-        try:
-            output_string = self.retrieved.get_object_content(fname)
-        except IOError:
-            return self.exit_codes.ERROR_OUTPUT_STDOUT_READ
-
-        result_dict = parse_cp2k_output(output_string)
+        # Read the standard output of CP2K.
+        exit_code, output_string = self._read_stdout()
+        if exit_code:
+            return exit_code
 
-        if "aborted" in result_dict:
-            return self.exit_codes.ERROR_OUTPUT_CONTAINS_ABORT
+        # Check the standard output for errors.
+        exit_code = self._check_stdout_for_errors(output_string)
+        if exit_code:
+            return exit_code
 
+        # Parse the standard output.
+        result_dict = utils.parse_cp2k_output(output_string)
         self.out("output_parameters", Dict(dict=result_dict))
-
         return None
 
     def _parse_trajectory(self):
         """CP2K trajectory parser."""
 
         from ase import Atoms
-        from aiida_cp2k.utils import parse_cp2k_trajectory
-
         fname = self.node.process_class._DEFAULT_RESTART_FILE_NAME  # pylint: disable=protected-access
 
         # Check if the restart file is present.
@@ -89,35 +85,55 @@ def _parse_trajectory(self):
         except IOError:
             return self.exit_codes.ERROR_OUTPUT_STDOUT_READ
 
-        return StructureData(ase=Atoms(**parse_cp2k_trajectory(output_string)))
+        return StructureData(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."""
 
-class Cp2kAdvancedParser(Cp2kBaseParser):
-    """Advanced AiiDA parser class for the output of CP2K."""
+        if "ABORT" in output_string:
+            return self.exit_codes.ERROR_OUTPUT_CONTAINS_ABORT
 
-    def _parse_stdout(self):
-        """Advanced CP2K output file parser."""
+        if "exceeded requested execution time" in output_string:
+            return self.exit_codes.ERROR_OUT_OF_WALLTIME
 
-        from aiida_cp2k.utils import parse_cp2k_output_advanced
+        if "PROGRAM STOPPED IN" not in output_string:
+            return self.exit_codes.ERROR_OUTPUT_INCOMPLETE
 
-        fname = self.node.process_class._DEFAULT_OUTPUT_FILE  # pylint: disable=protected-access
-        if fname not in self.retrieved.list_object_names():
-            raise OutputParsingError("CP2K output file not retrieved.")
+        return None
+
+    def _read_stdout(self):
+        """Read the standard output file. If impossible, return a non-zero exit code."""
+
+        fname = self.node.get_attribute('output_filename')
 
+        if fname not in self.retrieved.list_object_names():
+            return self.exit_codes.ERROR_OUTPUT_STDOUT_MISSING, None
         try:
             output_string = self.retrieved.get_object_content(fname)
         except IOError:
-            return self.exit_codes.ERROR_OUTPUT_STDOUT_READ
+            return self.exit_codes.ERROR_OUTPUT_STDOUT_READ, None
 
-        result_dict = parse_cp2k_output_advanced(output_string)
+        return None, output_string
 
-        # nwarnings is the last thing to be printed in th eCP2K output file:
-        # if it is not there, CP2K didn't finish properly
-        if 'nwarnings' not in result_dict:
-            raise OutputParsingError("CP2K did not finish properly.")
 
-        if "aborted" in result_dict:
-            return self.exit_codes.ERROR_OUTPUT_CONTAINS_ABORT
+class Cp2kAdvancedParser(Cp2kBaseParser):
+    """Advanced AiiDA parser class for the output of CP2K."""
+
+    def _parse_stdout(self):
+        """Advanced CP2K output file parser."""
+
+        # Read the standard output of CP2K.
+        exit_code, output_string = self._read_stdout()
+        if exit_code:
+            return exit_code
+
+        # Check the standard output for errors.
+        exit_code = self._check_stdout_for_errors(output_string)
+        if exit_code:
+            return exit_code
+
+        # Parse the standard output.
+        result_dict = utils.parse_cp2k_output_advanced(output_string)
 
         # Compute the bandgap for Spin1 and Spin2 if eigen was parsed (works also with smearing!)
         if 'eigen_spin1_au' in result_dict:
@@ -161,14 +177,17 @@ def _parse_stdout(self):
 
         from cp2k_output_tools import parse_iter
 
-        fname = self.node.process_class._DEFAULT_OUTPUT_FILE  # pylint: disable=protected-access
-        if fname not in self.retrieved.list_object_names():
-            raise OutputParsingError("CP2K output file not retrieved.")
+        # Read the standard output of CP2K.
+        exit_code, output_string = self._read_stdout()
+        if exit_code:
+            return exit_code
 
-        try:
-            output_string = self.retrieved.get_object_content(fname)
-        except IOError:
-            return self.exit_codes.ERROR_OUTPUT_STDOUT_READ
+        # Check the standard output for errors.
+        exit_code = self._check_stdout_for_errors(output_string)
+        if exit_code:
+            return exit_code
+
+        # Parse the standard output.
 
         result_dict = {}
 
@@ -177,14 +196,6 @@ def _parse_stdout(self):
         for match in parse_iter(output_string, key_mangling=True):
             result_dict.update(match)
 
-        # nwarnings is the last thing to be printed in the CP2K output file:
-        # if it is not there, CP2K didn't finish properly most likely
-        if 'nwarnings' not in result_dict:
-            raise OutputParsingError("CP2K did not finish properly")
-
-        if "aborted" in result_dict:
-            return self.exit_codes.ERROR_OUTPUT_CONTAINS_ABORT
-
         try:
             # the cp2k-output-tools parser is more hierarchical, be compatible with
             # the basic parser here and provide the total force eval energy as energy

aiida_cp2k/utils/__init__.py

@@ -7,7 +7,10 @@
 ###############################################################################
 """AiiDA-CP2K utils"""
 
-from .input_generator import Cp2kInput, add_restart_sections
+from .input_generator import Cp2kInput
+from .input_generator import add_ext_restart_section
+from .input_generator import add_restart_sections
+from .input_generator import add_wfn_restart_section
 from .parser import parse_cp2k_output
 from .parser import parse_cp2k_output_advanced
 from .parser import parse_cp2k_trajectory

aiida_cp2k/utils/input_generator.py

@@ -8,7 +8,7 @@
 """AiiDA-CP2K input generator."""
 
 from copy import deepcopy
-from collections.abc import Mapping, Sequence, MutableSequence
+from collections.abc import Mapping, Sequence, MutableSequence, MutableMapping
 
 from aiida.orm import Dict
 from aiida.engine import calcfunction
@@ -155,7 +155,7 @@ def _render_section(output, params, indent=0):
             if key.startswith(("@", "$")):
                 raise ValueError("CP2K preprocessor directives not supported.")
 
-            if isinstance(val, Mapping):
+            if isinstance(val, MutableMapping):
                 line = f"{' ' * indent}&{key}"
                 if "_" in val:  # if there is a section parameter, add it
                     line += f" {val.pop('_')}"
@@ -191,5 +191,35 @@ def add_restart_sections(input_dict):
         },
     }
     merge_dict(params, restart_wfn_dict)
+
+    # overwrite the complete EXT_RESTART section if present
+    params['EXT_RESTART'] = {'RESTART_FILE_NAME': './parent_calc/aiida-1.restart'}
+    return Dict(dict=params)
+
+
+@calcfunction
+def add_wfn_restart_section(input_dict, is_kpoints):
+    """Add wavefunction restart section to the input dictionary."""
+    params = input_dict.get_dict()
+    fname = './parent_calc/aiida-RESTART.kp' if is_kpoints else './parent_calc/aiida-RESTART.wfn'
+    restart_wfn_dict = {
+        'FORCE_EVAL': {
+            'DFT': {
+                'RESTART_FILE_NAME': fname,
+                'SCF': {
+                    'SCF_GUESS': 'RESTART',
+                },
+            },
+        },
+    }
+    merge_dict(params, restart_wfn_dict)
+    return Dict(dict=params)
+
+
+@calcfunction
+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'}
     return Dict(dict=params)

aiida_cp2k/utils/parser.py

@@ -23,10 +23,6 @@ def parse_cp2k_output(fstring):
             result_dict["energy_units"] = "a.u."
         if "The number of warnings for this run is" in line:
             result_dict["nwarnings"] = int(line.split()[-1])
-        if "exceeded requested execution time" in line:
-            result_dict["exceeded_walltime"] = True
-        if "ABORT" in line:
-            result_dict["aborted"] = True
 
     return result_dict
 
@@ -56,10 +52,6 @@ def parse_cp2k_output_advanced(fstring):  # pylint: disable=too-many-locals, too
             result_dict['energy_scf'] = energy_scf
         if 'The number of warnings for this run is' in line:
             result_dict['nwarnings'] = int(line.split()[-1])
-        if 'exceeded requested execution time' in line:
-            result_dict['exceeded_walltime'] = True
-        if "ABORT" in line:
-            result_dict["aborted"] = True
         if "KPOINTS| Band Structure Calculation" in line:
             kpoints, labels, bands = _parse_bands(lines, i_line, cp2k_version)
             result_dict["kpoint_data"] = {