Postprocessing plugin

aiida_flexpart/calculations/flexpart_post.py

@@ -0,0 +1,55 @@
+# -*- coding: utf-8 -*-
+"""
+Calculations provided by aiida_flexpart.
+Register calculations via the "aiida.calculations" entry point in setup.json.
+"""
+from aiida import orm
+from aiida.common import datastructures
+from aiida.engine import CalcJob
+
+
+class PostProcessingCalculation(CalcJob):
+    """AiiDA calculation plugin for post processing."""
+    @classmethod
+    def define(cls, spec):
+        """Define inputs and outputs of the calculation."""
+        # yapf: disable
+        super().define(spec)
+
+        # set default values for AiiDA options
+        spec.inputs['metadata']['options']['resources'].default = {
+            'num_machines': 1,
+            'num_mpiprocs_per_machine': 1,
+        }
+
+        #INPUTS
+        spec.input('metadata.options.parser_name', valid_type=str, default='flexpart.post')
+        spec.input('input_dir', valid_type = orm.RemoteData, required=True,
+                   help = 'main FLEXPART output dir')
+        spec.input('input_offline_dir', valid_type = orm.RemoteData, required=False,
+                   help = 'offline-nested FLEXPART output dir')
+        spec.input('metadata.options.output_filename', valid_type=str, default='aiida.out', required=True)
+        #exit codes
+        spec.outputs.dynamic = True
+        spec.exit_code(300, 'ERROR_MISSING_OUTPUT_FILES', message='Calculation did not produce all expected output files.')
+
+    def prepare_for_submission(self, folder):
+
+        params  = ['-m',self.inputs.input_dir.get_remote_path(),
+                   '-r','./'
+                  ]
+        if 'input_offline_dir' in self.inputs:
+            params += ['-n',self.inputs.input_offline_dir.get_remote_path()]
+
+        codeinfo = datastructures.CodeInfo()
+        codeinfo.cmdline_params = params
+        codeinfo.code_uuid = self.inputs.code.uuid
+        codeinfo.stdout_name = self.metadata.options.output_filename
+        codeinfo.withmpi = self.inputs.metadata.options.withmpi
+
+        # Prepare a `CalcInfo` to be returned to the engine
+        calcinfo = datastructures.CalcInfo()
+        calcinfo.codes_info = [codeinfo]
+        calcinfo.retrieve_list = ['grid_time_*.nc', 'boundary_sensitivity_*.nc', 'aiida.out']
+
+        return calcinfo

aiida_flexpart/parsers/flexpart_post.py

@@ -0,0 +1,56 @@
+# -*- coding: utf-8 -*-
+"""
+Parsers provided by aiida_flexpart.
+
+Register parsers via the "aiida.parsers" entry point in setup.json.
+"""
+from aiida.engine import ExitCode
+from aiida.parsers.parser import Parser
+from aiida.plugins import CalculationFactory
+from aiida.common import exceptions
+from aiida.orm import SinglefileData
+
+FlexpartCalculation = CalculationFactory('flexpart.post')
+
+
+class FlexpartPostParser(Parser):
+    """
+    Parser class for parsing output of calculation.
+    """
+    def __init__(self, node):
+        """
+        Initialize Parser instance
+
+        Checks that the ProcessNode being passed was produced by a FlexpartCalculation.
+
+        :param node: ProcessNode of calculation
+        :param type node: :class:`aiida.orm.ProcessNode`
+        """
+        super().__init__(node)
+        if not issubclass(node.process_class, FlexpartCalculation):
+            raise exceptions.ParsingError('Can only parse FlexpartCalculation')
+
+    def parse(self, **kwargs):
+        """
+        Parse outputs, store results in database.
+
+        :returns: an exit code, if parsing fails (or nothing if parsing succeeds)
+        """
+        output_filename = self.node.get_option('output_filename')
+
+        # Check that folder content is as expected
+        files_retrieved = self.retrieved.list_object_names()
+        files_expected = [output_filename]
+        # Note: set(A) <= set(B) checks whether A is a subset of B
+        if not set(files_expected) <= set(files_retrieved):
+            self.logger.error("Found files '{}', expected to find '{}'".format(
+                files_retrieved, files_expected))
+            return self.exit_codes.ERROR_MISSING_OUTPUT_FILES
+
+        # add output file
+        self.logger.info("Parsing '{}'".format(output_filename))
+        with self.retrieved.open(output_filename, 'rb') as handle:
+            output_node = SinglefileData(file=handle)
+        self.out('output_file', output_node)
+
+        return ExitCode(0)