From 49e94eff9203a5782ec3cac5c49a3736fa1a1d9e Mon Sep 17 00:00:00 2001 From: LucR31 Date: Wed, 20 Dec 2023 14:49:20 +0000 Subject: [PATCH 1/3] Postprocessing plugin --- aiida_flexpart/calculations/flexpart_post.py | 55 +++++++++++++++++++ aiida_flexpart/parsers/flexpart_post.py | 56 ++++++++++++++++++++ 2 files changed, 111 insertions(+) create mode 100644 aiida_flexpart/calculations/flexpart_post.py create mode 100644 aiida_flexpart/parsers/flexpart_post.py diff --git a/aiida_flexpart/calculations/flexpart_post.py b/aiida_flexpart/calculations/flexpart_post.py new file mode 100644 index 0000000..aeb5090 --- /dev/null +++ b/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 diff --git a/aiida_flexpart/parsers/flexpart_post.py b/aiida_flexpart/parsers/flexpart_post.py new file mode 100644 index 0000000..db13747 --- /dev/null +++ b/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) From 7828b31e5287009f9a5b1455fa01152a8762f7ef Mon Sep 17 00:00:00 2001 From: LucR31 <94859181+LucR31@users.noreply.github.com> Date: Wed, 20 Dec 2023 15:53:13 +0100 Subject: [PATCH 2/3] Apply suggestions from code review Co-authored-by: Aliaksandr Yakutovich --- aiida_flexpart/calculations/flexpart_post.py | 4 +--- aiida_flexpart/parsers/flexpart_post.py | 6 +----- 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/aiida_flexpart/calculations/flexpart_post.py b/aiida_flexpart/calculations/flexpart_post.py index aeb5090..ea19bb3 100644 --- a/aiida_flexpart/calculations/flexpart_post.py +++ b/aiida_flexpart/calculations/flexpart_post.py @@ -3,9 +3,7 @@ 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 +from aiida import orm, common, engine class PostProcessingCalculation(CalcJob): diff --git a/aiida_flexpart/parsers/flexpart_post.py b/aiida_flexpart/parsers/flexpart_post.py index db13747..967bc1f 100644 --- a/aiida_flexpart/parsers/flexpart_post.py +++ b/aiida_flexpart/parsers/flexpart_post.py @@ -4,11 +4,7 @@ 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 +from aiida import engine, parsers, plugins, common, orm FlexpartCalculation = CalculationFactory('flexpart.post') From 1fff3fd6ed69edde37897172adce4d068446445f Mon Sep 17 00:00:00 2001 From: LucR31 Date: Wed, 20 Dec 2023 14:57:42 +0000 Subject: [PATCH 3/3] change requests --- aiida_flexpart/calculations/flexpart_post.py | 6 +++--- aiida_flexpart/parsers/flexpart_post.py | 17 +++++++++-------- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/aiida_flexpart/calculations/flexpart_post.py b/aiida_flexpart/calculations/flexpart_post.py index ea19bb3..8d8fe55 100644 --- a/aiida_flexpart/calculations/flexpart_post.py +++ b/aiida_flexpart/calculations/flexpart_post.py @@ -6,7 +6,7 @@ from aiida import orm, common, engine -class PostProcessingCalculation(CalcJob): +class PostProcessingCalculation(engine.CalcJob): """AiiDA calculation plugin for post processing.""" @classmethod def define(cls, spec): @@ -39,14 +39,14 @@ def prepare_for_submission(self, folder): if 'input_offline_dir' in self.inputs: params += ['-n',self.inputs.input_offline_dir.get_remote_path()] - codeinfo = datastructures.CodeInfo() + codeinfo = common.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 = common.CalcInfo() calcinfo.codes_info = [codeinfo] calcinfo.retrieve_list = ['grid_time_*.nc', 'boundary_sensitivity_*.nc', 'aiida.out'] diff --git a/aiida_flexpart/parsers/flexpart_post.py b/aiida_flexpart/parsers/flexpart_post.py index 967bc1f..5ca2168 100644 --- a/aiida_flexpart/parsers/flexpart_post.py +++ b/aiida_flexpart/parsers/flexpart_post.py @@ -6,10 +6,10 @@ """ from aiida import engine, parsers, plugins, common, orm -FlexpartCalculation = CalculationFactory('flexpart.post') +FlexpartCalculation = plugins.CalculationFactory('flexpart.post') -class FlexpartPostParser(Parser): +class FlexpartPostParser(parsers.Parser): """ Parser class for parsing output of calculation. """ @@ -24,7 +24,7 @@ def __init__(self, node): """ super().__init__(node) if not issubclass(node.process_class, FlexpartCalculation): - raise exceptions.ParsingError('Can only parse FlexpartCalculation') + raise common.ParsingError('Can only parse FlexpartCalculation') def parse(self, **kwargs): """ @@ -39,14 +39,15 @@ def parse(self, **kwargs): 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)) + self.logger.error( + f"Found files '{files_retrieved}', expected to find '{files_expected}'" + ) return self.exit_codes.ERROR_MISSING_OUTPUT_FILES # add output file - self.logger.info("Parsing '{}'".format(output_filename)) + self.logger.info(f"Parsing '{output_filename}'") with self.retrieved.open(output_filename, 'rb') as handle: - output_node = SinglefileData(file=handle) + output_node = orm.SinglefileData(file=handle) self.out('output_file', output_node) - return ExitCode(0) + return engine.ExitCode(0)