From e2de37d7cadbe5fdc00dcf66b4c86a35b6b57667 Mon Sep 17 00:00:00 2001 From: dietervdb-meteo Date: Fri, 3 Jan 2025 15:25:49 +0200 Subject: [PATCH 1/8] fix typo --- src/anemoi/inference/checkpoint.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/anemoi/inference/checkpoint.py b/src/anemoi/inference/checkpoint.py index 2d62f8d..de09388 100644 --- a/src/anemoi/inference/checkpoint.py +++ b/src/anemoi/inference/checkpoint.py @@ -312,7 +312,7 @@ def mars_requests(self, *, variables, dates, use_grib_paramid=False, **kwargs): @cached_property def _supporting_arrays(self): - return self._metadata.supporting_arrays + return self._metadata._supporting_arrays @property def name(self): From 66ed6f16628a5d3448577445492d238afb3a3ba6 Mon Sep 17 00:00:00 2001 From: dietervdb-meteo Date: Fri, 3 Jan 2025 16:21:57 +0200 Subject: [PATCH 2/8] implement grid_indices mask --- src/anemoi/inference/inputs/dataset.py | 23 ++++++++++++++++++++--- src/anemoi/inference/metadata.py | 2 ++ 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/src/anemoi/inference/inputs/dataset.py b/src/anemoi/inference/inputs/dataset.py index c7122e3..2696fab 100644 --- a/src/anemoi/inference/inputs/dataset.py +++ b/src/anemoi/inference/inputs/dataset.py @@ -28,12 +28,26 @@ class DatasetInput(Input): def __init__(self, context, args, kwargs): super().__init__(context) + + grid_indices=kwargs.pop("grid_indices", None) + self.args, self.kwargs = args, kwargs if context.verbosity > 0: LOG.info( "Opening dataset with\nargs=%s\nkwargs=%s", json.dumps(args, indent=4), json.dumps(kwargs, indent=4) ) + if grid_indices is None and "grid_indices" in context.checkpoint._supporting_arrays: + grid_indices = context.checkpoint.load_supporting_array("grid_indices") + if context.verbosity > 0: + LOG.info( + "Loading supporting array `grid_indices` from checkpoint, \ + the input grid will be reduced accordingly." + ) + + self.grid_indices = slice(None) if grid_indices is None else grid_indices + + @cached_property def ds(self): from anemoi.datasets import open_dataset @@ -48,11 +62,13 @@ def create_input_state(self, *, date=None): raise ValueError("`date` must be provided") date = to_datetime(date) + latitudes=self.ds.latitudes + longitudes=self.ds.longitudes input_state = dict( date=date, - latitudes=self.ds.latitudes, - longitudes=self.ds.longitudes, + latitudes=latitudes[self.grid_indices], + longitudes=longitudes[self.grid_indices], fields=dict(), ) @@ -69,7 +85,8 @@ def create_input_state(self, *, date=None): if variable not in requested_variables: continue # Squeeze the data to remove the ensemble dimension - fields[variable] = np.squeeze(data[:, i], axis=1) + values = np.squeeze(data[:, i], axis=1) + fields[variable] = values[:,self.grid_indices] return input_state diff --git a/src/anemoi/inference/metadata.py b/src/anemoi/inference/metadata.py index 29a5791..9f7d8bb 100644 --- a/src/anemoi/inference/metadata.py +++ b/src/anemoi/inference/metadata.py @@ -139,6 +139,8 @@ def output_tensor_index_to_variable(self): @cached_property def number_of_grid_points(self): """Return the number of grid points per fields""" + if "grid_indices" in self._supporting_arrays: + return len(self.load_supporting_array("grid_indices")) try: return self._metadata.dataset.shape[-1] except AttributeError: From 22d0d336041f7f4b611a09de0d589d3740bf9446 Mon Sep 17 00:00:00 2001 From: dietervdb-meteo Date: Fri, 3 Jan 2025 16:47:13 +0200 Subject: [PATCH 3/8] adapt forcings --- src/anemoi/inference/forcings.py | 6 ++++-- src/anemoi/inference/inputs/dataset.py | 2 ++ 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/anemoi/inference/forcings.py b/src/anemoi/inference/forcings.py index 984d90a..4b7d042 100644 --- a/src/anemoi/inference/forcings.py +++ b/src/anemoi/inference/forcings.py @@ -136,8 +136,10 @@ def __init__(self, context, input, variables, variables_mask): self.variables_mask = variables_mask assert isinstance(input, DatasetInput), "Currently only boundary forcings from dataset supported." self.input = input - num_lam, num_other = input.ds.grids - self.spatial_mask = np.array([False] * num_lam + [True] * num_other, dtype=bool) + if "output_mask" in context.checkpoint._supporting_arrays: + self.spatial_mask= ~context.checkpoint.load_supporting_array("output_mask") + else: + self.spatial_mask=np.array([False] * len(input["latitudes"]) , dtype=bool) self.kinds = dict(retrieved=True) # Used for debugging def __repr__(self): diff --git a/src/anemoi/inference/inputs/dataset.py b/src/anemoi/inference/inputs/dataset.py index 2696fab..e0fcfee 100644 --- a/src/anemoi/inference/inputs/dataset.py +++ b/src/anemoi/inference/inputs/dataset.py @@ -99,6 +99,8 @@ def load_forcings(self, *, variables, dates): data = np.squeeze(data, axis=2) # Reorder the dimensions to (variable, date, values) data = np.swapaxes(data, 0, 1) + # apply reduction to `grid_indices` + data=data[...,self.grid_indices] return data def _load_dates(self, dates): From 2c488bb7985bb57b8d0869d8f627e68c3782e9be Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 3 Jan 2025 16:39:32 +0000 Subject: [PATCH 4/8] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/anemoi/inference/forcings.py | 4 ++-- src/anemoi/inference/inputs/dataset.py | 11 +++++------ 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/src/anemoi/inference/forcings.py b/src/anemoi/inference/forcings.py index 4b7d042..6d3f1c8 100644 --- a/src/anemoi/inference/forcings.py +++ b/src/anemoi/inference/forcings.py @@ -137,9 +137,9 @@ def __init__(self, context, input, variables, variables_mask): assert isinstance(input, DatasetInput), "Currently only boundary forcings from dataset supported." self.input = input if "output_mask" in context.checkpoint._supporting_arrays: - self.spatial_mask= ~context.checkpoint.load_supporting_array("output_mask") + self.spatial_mask = ~context.checkpoint.load_supporting_array("output_mask") else: - self.spatial_mask=np.array([False] * len(input["latitudes"]) , dtype=bool) + self.spatial_mask = np.array([False] * len(input["latitudes"]), dtype=bool) self.kinds = dict(retrieved=True) # Used for debugging def __repr__(self): diff --git a/src/anemoi/inference/inputs/dataset.py b/src/anemoi/inference/inputs/dataset.py index e0fcfee..68bdea3 100644 --- a/src/anemoi/inference/inputs/dataset.py +++ b/src/anemoi/inference/inputs/dataset.py @@ -29,7 +29,7 @@ class DatasetInput(Input): def __init__(self, context, args, kwargs): super().__init__(context) - grid_indices=kwargs.pop("grid_indices", None) + grid_indices = kwargs.pop("grid_indices", None) self.args, self.kwargs = args, kwargs if context.verbosity > 0: @@ -47,7 +47,6 @@ def __init__(self, context, args, kwargs): self.grid_indices = slice(None) if grid_indices is None else grid_indices - @cached_property def ds(self): from anemoi.datasets import open_dataset @@ -62,8 +61,8 @@ def create_input_state(self, *, date=None): raise ValueError("`date` must be provided") date = to_datetime(date) - latitudes=self.ds.latitudes - longitudes=self.ds.longitudes + latitudes = self.ds.latitudes + longitudes = self.ds.longitudes input_state = dict( date=date, @@ -86,7 +85,7 @@ def create_input_state(self, *, date=None): continue # Squeeze the data to remove the ensemble dimension values = np.squeeze(data[:, i], axis=1) - fields[variable] = values[:,self.grid_indices] + fields[variable] = values[:, self.grid_indices] return input_state @@ -100,7 +99,7 @@ def load_forcings(self, *, variables, dates): # Reorder the dimensions to (variable, date, values) data = np.swapaxes(data, 0, 1) # apply reduction to `grid_indices` - data=data[...,self.grid_indices] + data = data[..., self.grid_indices] return data def _load_dates(self, dates): From 3c45423d3c329a0f41f105103cf1624bac8a9eb7 Mon Sep 17 00:00:00 2001 From: dietervdb-meteo Date: Tue, 7 Jan 2025 10:51:01 +0200 Subject: [PATCH 5/8] fix _get_datasets_full_paths --- src/anemoi/inference/metadata.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/anemoi/inference/metadata.py b/src/anemoi/inference/metadata.py index 9f7d8bb..e5dc0ab 100644 --- a/src/anemoi/inference/metadata.py +++ b/src/anemoi/inference/metadata.py @@ -512,14 +512,12 @@ def _find(x): _find(y) if isinstance(x, dict): - if "dataset" in x: + if "dataset" in x and isinstance(x["dataset"], str): result.append(x["dataset"]) for k, v in x.items(): _find(v) - _find(self._config.dataloader.training.dataset) - return result def open_dataset_args_kwargs(self, *, use_original_paths, from_dataloader=None): From f8ed73e1df568408881fcbc48feee45edb02f72d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 7 Jan 2025 08:52:11 +0000 Subject: [PATCH 6/8] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/anemoi/inference/metadata.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/anemoi/inference/metadata.py b/src/anemoi/inference/metadata.py index e5dc0ab..5eb674a 100644 --- a/src/anemoi/inference/metadata.py +++ b/src/anemoi/inference/metadata.py @@ -512,11 +512,12 @@ def _find(x): _find(y) if isinstance(x, dict): - if "dataset" in x and isinstance(x["dataset"], str): + if "dataset" in x and isinstance(x["dataset"], str): result.append(x["dataset"]) for k, v in x.items(): _find(v) + _find(self._config.dataloader.training.dataset) return result From 0ca9a9156f3c4972d02f12014c4de5a5342c17d3 Mon Sep 17 00:00:00 2001 From: dietervdb-meteo Date: Tue, 7 Jan 2025 17:40:40 +0200 Subject: [PATCH 7/8] change boundary forcings trigger --- src/anemoi/inference/metadata.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/anemoi/inference/metadata.py b/src/anemoi/inference/metadata.py index 5eb674a..c12b833 100644 --- a/src/anemoi/inference/metadata.py +++ b/src/anemoi/inference/metadata.py @@ -718,9 +718,7 @@ def boundary_forcings_inputs(self, context, input_state): result = [] - output_mask = self._config_model.get("output_mask", None) - if output_mask is not None: - assert output_mask == "cutout", "Currently only cutout as output mask supported." + if "output_mask" in self._supporting_arrays: result.append( context.create_boundary_forcings( self.prognostic_variables, From 2abd3ac8151e771ef85e9865f7e1c2e8c3dc4e4d Mon Sep 17 00:00:00 2001 From: Dieter Van den Bleeken Date: Tue, 7 Jan 2025 17:41:36 +0100 Subject: [PATCH 8/8] Update CHANGELOG.md --- CHANGELOG.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0fbdcae..6b2a84c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,9 @@ Keep it human-readable, your future self will thank you! ## [Unreleased] ### Added +- Add support for models with unconnected nodes dropped from input [#95](https://github.com/ecmwf/anemoi-inference/pull/95). +- Change trigger for boundary forcings [#95](https://github.com/ecmwf/anemoi-inference/pull/95). +- Add support for automatic loading of anemoi-datasets of more general type [#95](https://github.com/ecmwf/anemoi-inference/pull/95). - Add initial state output in netcdf format - Fix: Enable inference when no constant forcings are used - Add anemoi-transform link to documentation