From 9aa8efa4821fa91509b8b4b30a643332c784fb23 Mon Sep 17 00:00:00 2001 From: Baudouin Raoult Date: Thu, 19 Sep 2024 06:32:58 +0000 Subject: [PATCH 1/3] fix diagnotics bug --- .../inference/checkpoint/metadata/__init__.py | 56 +++++++++++++++++-- src/anemoi/inference/runner.py | 38 ++++++++----- 2 files changed, 75 insertions(+), 19 deletions(-) diff --git a/src/anemoi/inference/checkpoint/metadata/__init__.py b/src/anemoi/inference/checkpoint/metadata/__init__.py index 8bab431..f961a98 100644 --- a/src/anemoi/inference/checkpoint/metadata/__init__.py +++ b/src/anemoi/inference/checkpoint/metadata/__init__.py @@ -6,6 +6,7 @@ # nor does it submit to any jurisdiction. +import json import logging from functools import cached_property @@ -155,7 +156,12 @@ def select(self): # order = self._dataset["order_by"] return dict( # valid_datetime="ascending", - param_level=self.variables, + param_level=sorted( + set(self.variables) + - set(self.computed_constants) + - set(self.computed_forcings) + - set(self.diagnostic_params) + ), # ensemble=self.checkpoint.ordering('ensemble'), remapping={"param_level": "{param}_{levelist}"}, ) @@ -204,7 +210,7 @@ def _computed_constants(self): LOG.debug("computed_constants data_mask: %s", data_mask) LOG.debug("computed_constants model_mask: %s", model_mask) - LOG.info("Computed constants: %s", names) + LOG.debug("Computed constants: %s", names) return data_mask, model_mask, names @@ -230,8 +236,6 @@ def _computed_forcings(self): ] ) - print("FORCINGS", self._forcing_params()) - constants = set(self._forcing_params()) - set(self.constants_from_input) - set(self.computed_constants) if constants - known: @@ -241,7 +245,7 @@ def _computed_forcings(self): LOG.debug("computed_forcing data_mask: %s", data_mask) LOG.debug("computed_forcing model_mask: %s", model_mask) - LOG.info("Computed forcings: %s", names) + # LOG.info("Computed forcings: %s", names) return data_mask, model_mask, names @@ -265,7 +269,7 @@ def _constants_from_input(self): LOG.debug("constants_from_input: %s", data_mask) LOG.debug("constants_from_input: %s", model_mask) - LOG.info("Constants from input: %s", names) + LOG.debug("Constants from input: %s", names) return data_mask, model_mask, names @@ -307,6 +311,11 @@ def diagnostic_params(self): def prognostic_params(self): return [self.index_to_variable[i] for i in self._indices["data"]["input"]["prognostic"]] + @cached_property + def accumulations_params(self): + # We assume that accumulations are the ones that are forecasts + return sorted(p[0] for p in self.param_step_sfc_pairs) + ########################################################################### @cached_property def precision(self): @@ -349,3 +358,38 @@ def predict_step_shape(self): self.number_of_grid_points, # Grid points self.num_input_features, # Fields ) + + ########################################################################### + def summary(self): + + print(f"Prognostics: ({len(self.prognostic_params)})") + print(sorted(self.prognostic_params)) + print() + + print(f"Diagnostics: ({len(self.diagnostic_params)})") + print(sorted(self.diagnostic_params)) + print() + + print(f"Retrieved constants: ({len(self.constants_from_input)})") + print(sorted(self.constants_from_input)) + print() + + print(f"Computed constants: ({len(self.computed_constants)})") + print(sorted(self.computed_constants)) + print() + + print(f"Computed forcings: ({len(self.computed_forcings)})") + print(sorted(self.computed_forcings)) + print() + + print(f"Accumulations: ({len(self.accumulations_params)})") + print(sorted(self.accumulations_params)) + print() + + print("Select:") + print(json.dumps(self.select, indent=2)) + print() + + print("Order by:") + print(json.dumps(self.order_by, indent=2)) + print() diff --git a/src/anemoi/inference/runner.py b/src/anemoi/inference/runner.py index ab3faf4..1b60eae 100644 --- a/src/anemoi/inference/runner.py +++ b/src/anemoi/inference/runner.py @@ -98,6 +98,8 @@ def run( _description_ """ + self.checkpoint.summary() + if autocast is None: autocast = self.checkpoint.precision @@ -310,7 +312,7 @@ def get_most_recent_datetime(input_fields): most_recent_datetime = get_most_recent_datetime(input_fields) reference_fields = [f for f in input_fields if f.datetime()["valid_time"] == most_recent_datetime] - precip_template = reference_fields[self.checkpoint.variable_to_index["lsm"]] + prognostic_template = reference_fields[self.checkpoint.variable_to_index["lsm"]] accumulated_output = np.zeros( shape=(len(diagnostic_output_mask), number_of_grid_points), @@ -321,13 +323,14 @@ def get_most_recent_datetime(input_fields): output_callback( input_fields, self.checkpoint.diagnostic_params, - precip_template, + prognostic_template, accumulated_output[0].shape, ) else: output_callback(input_fields) prognostic_params = self.checkpoint.prognostic_params + accumulations_params = self.checkpoint.accumulations_params # with self.stepper(self.hour_steps) as stepper: @@ -362,19 +365,28 @@ def get_most_recent_datetime(input_fields): if len(diagnostic_output_mask): for n, param in enumerate(self.checkpoint.diagnostic_params): accumulated_output[n] += np.maximum(0, diagnostic_fields_numpy[:, n]) - assert precip_template.datetime()["valid_time"] == most_recent_datetime, ( - precip_template.datetime()["valid_time"], + assert prognostic_template.datetime()["valid_time"] == most_recent_datetime, ( + prognostic_template.datetime()["valid_time"], most_recent_datetime, ) - output_callback( - accumulated_output[n], - stepType="accum", - template=precip_template, - startStep=0, - endStep=step, - param=param, - check_nans=True, # param in can_be_missing, - ) + + if param in accumulations_params: + output_callback( + accumulated_output[n], + stepType="accum", + template=prognostic_template, + startStep=0, + endStep=step, + param=param, + check_nans=True, # param in can_be_missing, + ) + else: + output_callback( + diagnostic_fields_numpy[:, n], + template=prognostic_template, + step=step, + check_nans=True, # param in can_be_missing, + ) # Next step From 89db0c979f9f8aae10ec199ca6b52e842b3188ab Mon Sep 17 00:00:00 2001 From: Baudouin Raoult Date: Fri, 20 Sep 2024 10:51:40 +0000 Subject: [PATCH 2/3] remove diagnostic param from mars request --- .../inference/checkpoint/metadata/__init__.py | 12 +++++----- src/anemoi/inference/plugin.py | 6 ++--- src/anemoi/inference/runner.py | 24 +++++++++++++++++++ 3 files changed, 33 insertions(+), 9 deletions(-) diff --git a/src/anemoi/inference/checkpoint/metadata/__init__.py b/src/anemoi/inference/checkpoint/metadata/__init__.py index f961a98..433ab35 100644 --- a/src/anemoi/inference/checkpoint/metadata/__init__.py +++ b/src/anemoi/inference/checkpoint/metadata/__init__.py @@ -386,10 +386,10 @@ def summary(self): print(sorted(self.accumulations_params)) print() - print("Select:") - print(json.dumps(self.select, indent=2)) - print() + # print("Select:") + # print(json.dumps(self.select, indent=2)) + # print() - print("Order by:") - print(json.dumps(self.order_by, indent=2)) - print() + # print("Order by:") + # print(json.dumps(self.order_by, indent=2)) + # print() diff --git a/src/anemoi/inference/plugin.py b/src/anemoi/inference/plugin.py index 086e458..52a2bdc 100644 --- a/src/anemoi/inference/plugin.py +++ b/src/anemoi/inference/plugin.py @@ -80,15 +80,15 @@ def _output(self, *args, **kwargs): @property def param_sfc(self): - return self.runner.checkpoint.param_sfc + return self.runner.param_sfc @property def param_level_pl(self): - return self.runner.checkpoint.param_level_pl + return self.runner.param_level_pl @property def param_level_ml(self): - return self.runner.checkpoint.param_level_ml + return self.runner.param_level_ml @property def constant_fields(self): diff --git a/src/anemoi/inference/runner.py b/src/anemoi/inference/runner.py index 1b60eae..1890592 100644 --- a/src/anemoi/inference/runner.py +++ b/src/anemoi/inference/runner.py @@ -419,6 +419,30 @@ def lagged(self): result = [-s * self.hour_steps for s in result] return sorted(result) + @property + def param_sfc(self): + param_sfc = self.checkpoint.param_sfc + + # Remove diagnostic params + + param_sfc = [p for p in param_sfc if p not in self.checkpoint.diagnostic_params] + + return param_sfc + + @property + def param_level_pl(self): + + # To do remove diagnostic params + + return self.checkpoint.param_level_pl + + @property + def param_level_ml(self): + + # To do remove diagnostic params + + return self.checkpoint.param_level_ml + class DefaultRunner(Runner): """_summary_ From b37188a8bbfb5bdf07403e1b926cfb6ec46334a1 Mon Sep 17 00:00:00 2001 From: Gert Mertes Date: Tue, 24 Sep 2024 10:16:01 +0000 Subject: [PATCH 3/3] chore: changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 99e4049..2bdcd58 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,7 @@ Keep it human-readable, your future self will thank you! - ci-hpc-config ### Changed +- Fix: diagnostics bug when fields are non-accumulated, remove diagnostics from mars request [#18](https://github.com/ecmwf/anemoi-inference/pull/18) - ci: updated workflows on PR and releases to use reusable actions ### Removed