diff --git a/ecml_tools/create/input.py b/ecml_tools/create/input.py index 8633da9..0e4902c 100644 --- a/ecml_tools/create/input.py +++ b/ecml_tools/create/input.py @@ -688,6 +688,41 @@ def __repr__(self): return super().__repr__(content) +class IncludeResult(Result): + def __init__(self, context, action_path, dates, result, results): + super().__init__(context, action_path, dates) + # result is the content of the include + self.result = result + # results is the list of the included results + self.results = results + + @cached_property + def datasource(self): + for i in self.results: + # for each include trigger the datasource to be computed + # and saved in context but drop it + i.datasource + # then return the content of the result + # which can use the datasources of the included results + return self.result.datasource + + +class IncludeAction(ActionWithList): + def __init__(self, context, action_path, includes, content): + super().__init__(context, ["include"], *includes) + self.content = action_factory(content, context, ["input"]) + + def select(self, dates): + results = [a.select(dates) for a in self.actions] + return IncludeResult( + self.context, + self.action_path, + dates, + self.content.select(dates), + results, + ) + + class ConcatAction(Action): def __init__(self, context, action_path, *configs): super().__init__(context, action_path, *configs) @@ -745,7 +780,7 @@ def action_factory(config, context, action_path): cls = dict( date_shift=DateShiftAction, # date_filter=DateFilterAction, - # include=IncludeAction, + include=IncludeAction, concat=ConcatAction, join=JoinAction, pipe=PipeAction, @@ -811,8 +846,17 @@ def __init__(self, /, order_by, flatten_grid, remapping): class InputBuilder: - def __init__(self, config, **kwargs): + def __init__(self, config, include, **kwargs): self.kwargs = kwargs + + config = deepcopy(config) + if include: + config = dict( + include=dict( + includes=include, + content=config, + ) + ) self.config = config self.action_path = ["input"] diff --git a/ecml_tools/create/loaders.py b/ecml_tools/create/loaders.py index d487d9c..a4101ca 100644 --- a/ecml_tools/create/loaders.py +++ b/ecml_tools/create/loaders.py @@ -81,6 +81,7 @@ def build_input(self): builder = build_input( self.main_config.input, + include=self.main_config.get("include", {}), order_by=self.output.order_by, flatten_grid=self.output.flatten_grid, remapping=build_remapping(self.output.remapping), diff --git a/tests/create-include.yaml b/tests/create-include.yaml new file mode 100644 index 0000000..ad3e53f --- /dev/null +++ b/tests/create-include.yaml @@ -0,0 +1,41 @@ +description: "develop version of the dataset for a few days and a few variables, once data on mars is cached it should take a few seconds to generate the dataset" +dataset_status: testing +purpose: aifs +name: test-small +config_format_version: 2 + +common: + mars_request: &mars_request + expver: "0001" + class: ea + grid: 20./20. + +dates: + start: 2020-12-30 00:00:00 + end: 2021-01-03 12:00:00 + frequency: 12h + group_by: monthly + +include: + - mars: + <<: *mars_request + param: [2t] + levtype: sfc + stream: oper + type: an + +input: + constants: + template: ${include.0.mars} + param: + - cos_latitude + +output: + chunking: { dates: 1, ensembles: 1 } + dtype: float32 + flatten_grid: True + order_by: [valid_datetime, param_level, number] + remapping: + param_level: "{param}_{levelist}" + statistics: param_level + statistics_end: 2021 diff --git a/tests/create-missing.yaml b/tests/create-missing.yaml index 61b9d89..8e5b89f 100644 --- a/tests/create-missing.yaml +++ b/tests/create-missing.yaml @@ -17,34 +17,19 @@ dates: group_by: monthly missing: ['2021-01-03 00:00:00'] +include: + - mars: + <<: *mars_request + param: [2t] + levtype: sfc + stream: oper + type: an input: - join: - - mars: - <<: *mars_request - param: [2t] - levtype: sfc - stream: oper - type: an - - - mars: - <<: *mars_request - param: [q, t] - levtype: pl - level: [50, 100] - stream: oper - type: an - - - accumulations: - <<: *mars_request - levtype: sfc - param: [cp, tp] - # accumulation_period: 6h - - - constants: - template: ${input.join.0.mars} - param: - - cos_latitude + constants: + template: ${include.0.mars} + param: + - cos_latitude output: chunking: { dates: 1, ensembles: 1 }