From 9f8fd2627e342ec7b44471aa6d41d9e31fba0cfa Mon Sep 17 00:00:00 2001 From: mpvginde Date: Tue, 7 Apr 2026 11:21:21 +0200 Subject: [PATCH 1/7] port ifs-forecast loader --- src/mxalign/loaders/ifs_forecast.py | 52 +++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) create mode 100644 src/mxalign/loaders/ifs_forecast.py diff --git a/src/mxalign/loaders/ifs_forecast.py b/src/mxalign/loaders/ifs_forecast.py new file mode 100644 index 0000000..70964b0 --- /dev/null +++ b/src/mxalign/loaders/ifs_forecast.py @@ -0,0 +1,52 @@ +from pathlib import Path +import xarray as xr + +from .registry import register_loader +from ..properties.properties import Space, Time, Uncertainty +from .base import BaseLoader + +@register_loader +class IFSForecastLoader(BaseLoader): + try: + import cfgrib + except: + ImportError("Please install the cfgrib package to load IFS-Forecasts") + + name = "ifs-forecast" + + space = Space.GRID + time = Time.FORECAST + uncertainty = Uncertainty.DETERMINISTIC + + def _load(self): + + kwargs = self.kwargs.copy() + + files = [self.files] if isinstance(self.files, str) else self.files + + ds = xr.open_mfdataset( + files, + combine="nested", + concat_dim="time", + chunks={ + "time" : 1, + "step": -1, + "values": -1 + }, + **kwargs + ) + + ds.coords["longitude"] = (ds.coords["longitude"] + 180.) % 360. -180. + + ds_out = ds.rename_dims( + time="reference_time", + step="lead_time", + values="grid_index" + ).rename_vars( + time="reference_time", + step="lead_time" + ).drop_vars( + ["number","surface"] + ) + + return ds_out \ No newline at end of file From ce10399ced8537f5f029c6b33a4649d9efb5efb6 Mon Sep 17 00:00:00 2001 From: nielscarlier <136123206+nielscarlier@users.noreply.github.com> Date: Tue, 12 May 2026 15:35:13 +0200 Subject: [PATCH 2/7] updated ifs forecast loader to also handle ensemble data (#19) * updated ifs forecast loader to also handle ensemble data Co-authored-by: Stagiair user Niels Carlier --- src/mxalign/loaders/ifs_forecast.py | 70 ++++++++++++++++++----------- 1 file changed, 43 insertions(+), 27 deletions(-) diff --git a/src/mxalign/loaders/ifs_forecast.py b/src/mxalign/loaders/ifs_forecast.py index 70964b0..22ae4f9 100644 --- a/src/mxalign/loaders/ifs_forecast.py +++ b/src/mxalign/loaders/ifs_forecast.py @@ -1,52 +1,68 @@ -from pathlib import Path +import numpy as np import xarray as xr from .registry import register_loader -from ..properties.properties import Space, Time, Uncertainty +from ..properties.properties import Properties, Space, Time, Uncertainty from .base import BaseLoader + @register_loader class IFSForecastLoader(BaseLoader): try: - import cfgrib - except: - ImportError("Please install the cfgrib package to load IFS-Forecasts") - + import cfgrib + except Exception: + raise ImportError("Please install the cfgrib package to load IFS-Forecasts") + name = "ifs-forecast" space = Space.GRID time = Time.FORECAST - uncertainty = Uncertainty.DETERMINISTIC + uncertainty = None def _load(self): - kwargs = self.kwargs.copy() - files = [self.files] if isinstance(self.files, str) else self.files - + ds = xr.open_mfdataset( files, combine="nested", concat_dim="time", chunks={ - "time" : 1, + "time": 1, "step": -1, - "values": -1 + "values": -1, }, - **kwargs - ) - - ds.coords["longitude"] = (ds.coords["longitude"] + 180.) % 360. -180. - - ds_out = ds.rename_dims( - time="reference_time", - step="lead_time", - values="grid_index" - ).rename_vars( - time="reference_time", - step="lead_time" - ).drop_vars( - ["number","surface"] + **kwargs, ) - return ds_out \ No newline at end of file + ds.coords["longitude"] = (ds.coords["longitude"] + 180.0) % 360.0 - 180.0 + + rename_dims = { + "time": "reference_time", + "step": "lead_time", + "values": "grid_index", + } + rename_vars = { + "time": "reference_time", + "step": "lead_time", + } + + if "number" in ds.dims: + rename_dims["number"] = "ensemble_member" + if "number" in ds.coords: + rename_vars["number"] = "ensemble_member" + + ds = ds.rename_dims({k: v for k, v in rename_dims.items() if k in ds.dims}) + ds = ds.rename_vars({k: v for k, v in rename_vars.items() if k in ds.variables}) + + if "surface" in ds.variables: + ds = ds.drop_vars("surface") + + return ds + + if "member" in ds.dims: + uncertainty = Uncertainty.ENSEMBLE + elif "quantile" in ds.dims: + uncertainty = Uncertainty.QUANTILE + else: + uncertainty = Uncertainty.DETERMINISTIC \ No newline at end of file From 87cd64ab84c7ec1d4e49837e99483d69525ba6ff Mon Sep 17 00:00:00 2001 From: Michiel Van Ginderachter Date: Tue, 12 May 2026 13:47:34 +0000 Subject: [PATCH 3/7] fix linting etc --- ...ADR-001_seperate-validation-and-loading.md | 6 +- doc/adr/ADR-002_mxalign-loader-interface.md | 2 +- doc/adr/template.md | 2 +- src/mxalign/loaders/anemoi_inference.py | 67 +++++++++---------- 4 files changed, 36 insertions(+), 41 deletions(-) diff --git a/doc/adr/ADR-001_seperate-validation-and-loading.md b/doc/adr/ADR-001_seperate-validation-and-loading.md index 47bbdbc..237349e 100644 --- a/doc/adr/ADR-001_seperate-validation-and-loading.md +++ b/doc/adr/ADR-001_seperate-validation-and-loading.md @@ -10,7 +10,7 @@ informed: Francois B., Buurman S. ## Context -Currently `mxalign` implements dataset loading functionality, validation functionality (checking if the loaded dataset has all the correct metadata and) and alignment functionality. However, dataset-loading and -validation fall outside the main scope of the `mxalign` package. +Currently `mxalign` implements dataset loading functionality, validation functionality (checking if the loaded dataset has all the correct metadata and) and alignment functionality. However, dataset-loading and -validation fall outside the main scope of the `mxalign` package. ## Decision Drivers @@ -27,7 +27,7 @@ Currently `mxalign` implements dataset loading functionality, validation functio ## Decision Outcome -Chosen option 1., because it allows for most flexibility and provides a clear entry point for users who want to bring their own loader. +Chosen option 1., because it allows for most flexibility and provides a clear entry point for users who want to bring their own loader. ### Consequences @@ -36,4 +36,4 @@ Chosen option 1., because it allows for most flexibility and provides a clear en * `mxalign` is now only responsible for the alignment tasks ## More Information -Currently the interface between dataset loaded with an `mlwp-data-loaders` loader and `mxalign` is not defined. Ideally `mxalign` should know the traits of the dataset to correctly align dataset. How do we inform `mxalign` on the traits? See [ADR-002](./ADR-002_mxalign-loader-interface.md) for possible options. +Currently the interface between dataset loaded with an `mlwp-data-loaders` loader and `mxalign` is not defined. Ideally `mxalign` should know the traits of the dataset to correctly align dataset. How do we inform `mxalign` on the traits? See [ADR-002](./ADR-002_mxalign-loader-interface.md) for possible options. diff --git a/doc/adr/ADR-002_mxalign-loader-interface.md b/doc/adr/ADR-002_mxalign-loader-interface.md index 0e359df..08dac30 100644 --- a/doc/adr/ADR-002_mxalign-loader-interface.md +++ b/doc/adr/ADR-002_mxalign-loader-interface.md @@ -71,4 +71,4 @@ Chosen option: "{title of option 1}", because {justification. e.g., only option, ## More Information -{You might want to provide additional evidence/confidence for the decision outcome here and/or document the team agreement on the decision and/or define when/how this decision the decision should be realized and if/when it should be re-visited. Links to other decisions and resources might appear here as well.} \ No newline at end of file +{You might want to provide additional evidence/confidence for the decision outcome here and/or document the team agreement on the decision and/or define when/how this decision the decision should be realized and if/when it should be re-visited. Links to other decisions and resources might appear here as well.} diff --git a/doc/adr/template.md b/doc/adr/template.md index 0e359df..08dac30 100644 --- a/doc/adr/template.md +++ b/doc/adr/template.md @@ -71,4 +71,4 @@ Chosen option: "{title of option 1}", because {justification. e.g., only option, ## More Information -{You might want to provide additional evidence/confidence for the decision outcome here and/or document the team agreement on the decision and/or define when/how this decision the decision should be realized and if/when it should be re-visited. Links to other decisions and resources might appear here as well.} \ No newline at end of file +{You might want to provide additional evidence/confidence for the decision outcome here and/or document the team agreement on the decision and/or define when/how this decision the decision should be realized and if/when it should be re-visited. Links to other decisions and resources might appear here as well.} diff --git a/src/mxalign/loaders/anemoi_inference.py b/src/mxalign/loaders/anemoi_inference.py index e72e391..68f8799 100644 --- a/src/mxalign/loaders/anemoi_inference.py +++ b/src/mxalign/loaders/anemoi_inference.py @@ -5,17 +5,14 @@ from ..properties.properties import Space, Time, Uncertainty from .base import BaseLoader -DEFAULTS_NETCDF = { - "chunks": "auto", - "engine": "h5netcdf", - "parallel": True -} +DEFAULTS_NETCDF = {"chunks": "auto", "engine": "h5netcdf", "parallel": True} DEFAULTS_ZARR = { "chunks": "auto", "storage_options": {"anon": True}, } + @register_loader class AnemoiInferenceLoader(BaseLoader): name = "anemoi-inference" @@ -26,75 +23,73 @@ class AnemoiInferenceLoader(BaseLoader): def _load(self): - kwargs = self.kwargs.copy() - - if isinstance(self.files,str): + + if isinstance(self.files, str): if Path(self.files).suffix.lower() == ".zarr": files = self.files - for k, v in DEFAULTS_ZARR.items(): - kwargs[k] = self.kwargs.get(k,v) + kwargs[k] = self.kwargs.get(k, v) loader = _open_zarr else: files = [self.files] for k, v in DEFAULTS_NETCDF.items(): - kwargs[k] = self.kwargs.get(k,v) + kwargs[k] = self.kwargs.get(k, v) loader = _open_mf_dataset else: files = self.files - if Path(files[0]).suffix.lower() == ".zarr": + if Path(files[0]).suffix.lower() == ".zarr": for k, v in DEFAULTS_ZARR.items(): - kwargs[k] = self.kwargs.get(k,v) + kwargs[k] = self.kwargs.get(k, v) kwargs["engine"] = "zarr" - else: + else: for k, v in DEFAULTS_NETCDF.items(): - kwargs[k] = self.kwargs.get(k,v) + kwargs[k] = self.kwargs.get(k, v) loader = _open_mf_dataset - ds = loader(files, **kwargs) return ds + def _open_mf_dataset(files, **kwargs): - times = xr.open_dataset(files[0], engine=kwargs["engine"], chunks=kwargs["chunks"])["time"].values - lead_times = times - times[0] + times = xr.open_dataset(files[0], engine=kwargs["engine"], chunks=kwargs["chunks"])[ + "time" + ].values + lead_times = times - times[0] - ds = xr.open_mfdataset( - files, - preprocess=_preprocess, - **kwargs - ) + ds = xr.open_mfdataset(files, preprocess=_preprocess, **kwargs) - ds_out = ds.\ - assign_coords({"lead_time": ("time", lead_times)}).\ - rename_dims({"values": "grid_index"}).\ - swap_dims({"time": "lead_time"}) + ds_out = ( + ds.assign_coords({"lead_time": ("time", lead_times)}) + .rename_dims({"values": "grid_index"}) + .swap_dims({"time": "lead_time"}) + ) return ds_out + def _open_zarr(files, **kwargs): ds = xr.open_zarr(files, **kwargs) times = ds["time"].values - lead_times = times - times[0] - + lead_times = times - times[0] + ds_out = _preprocess(ds) - - ds_out = ds_out.\ - assign_coords({"lead_time": ("time", lead_times)}).\ - rename_dims({"values": "grid_index"}).\ - swap_dims({"time": "lead_time"}) - - return ds_out + ds_out = ( + ds_out.assign_coords({"lead_time": ("time", lead_times)}) + .rename_dims({"values": "grid_index"}) + .swap_dims({"time": "lead_time"}) + ) + + return ds_out def _preprocess(ds): From 1132266a7f45d271ff806c33d0cc3cc31b3f9064 Mon Sep 17 00:00:00 2001 From: Michiel Van Ginderachter Date: Tue, 12 May 2026 13:48:02 +0000 Subject: [PATCH 4/7] fix property setting --- src/mxalign/loaders/ifs_forecast.py | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/src/mxalign/loaders/ifs_forecast.py b/src/mxalign/loaders/ifs_forecast.py index 22ae4f9..6175c37 100644 --- a/src/mxalign/loaders/ifs_forecast.py +++ b/src/mxalign/loaders/ifs_forecast.py @@ -1,15 +1,14 @@ -import numpy as np import xarray as xr from .registry import register_loader -from ..properties.properties import Properties, Space, Time, Uncertainty +from ..properties.properties import Space, Time, Uncertainty from .base import BaseLoader @register_loader class IFSForecastLoader(BaseLoader): try: - import cfgrib + import cfgrib except Exception: raise ImportError("Please install the cfgrib package to load IFS-Forecasts") @@ -48,9 +47,10 @@ def _load(self): } if "number" in ds.dims: - rename_dims["number"] = "ensemble_member" + rename_dims["number"] = "member" + self if "number" in ds.coords: - rename_vars["number"] = "ensemble_member" + rename_vars["number"] = "member" ds = ds.rename_dims({k: v for k, v in rename_dims.items() if k in ds.dims}) ds = ds.rename_vars({k: v for k, v in rename_vars.items() if k in ds.variables}) @@ -58,11 +58,12 @@ def _load(self): if "surface" in ds.variables: ds = ds.drop_vars("surface") + if "member" in ds.dims: + self.uncertainty = Uncertainty.ENSEMBLE + elif "quantile" in ds.dims: + self.uncertainty = Uncertainty.QUANTILE + else: + self.uncertainty = Uncertainty.DETERMINISTIC return ds - if "member" in ds.dims: - uncertainty = Uncertainty.ENSEMBLE - elif "quantile" in ds.dims: - uncertainty = Uncertainty.QUANTILE - else: - uncertainty = Uncertainty.DETERMINISTIC \ No newline at end of file + From 03348ade0f3950000181a1c1545affee4a5e9048 Mon Sep 17 00:00:00 2001 From: Michiel Van Ginderachter Date: Tue, 12 May 2026 13:57:44 +0000 Subject: [PATCH 5/7] add import to register loader --- src/mxalign/loaders/__init__.py | 2 ++ src/mxalign/loaders/ifs_forecast.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/mxalign/loaders/__init__.py b/src/mxalign/loaders/__init__.py index f04b80e..329bfa3 100644 --- a/src/mxalign/loaders/__init__.py +++ b/src/mxalign/loaders/__init__.py @@ -1,11 +1,13 @@ from . import anemoi_datasets from . import anemoi_inference from . import harp_obstable +from . import ifs_forecast from . import base __all__ = [ "anemoi_datasets", "anemoi_inference", + "ifs_forecast", "harp_obstable", "base", ] diff --git a/src/mxalign/loaders/ifs_forecast.py b/src/mxalign/loaders/ifs_forecast.py index 6175c37..7fd9fae 100644 --- a/src/mxalign/loaders/ifs_forecast.py +++ b/src/mxalign/loaders/ifs_forecast.py @@ -64,6 +64,6 @@ def _load(self): self.uncertainty = Uncertainty.QUANTILE else: self.uncertainty = Uncertainty.DETERMINISTIC - return ds + return ds.transpose("reference_time", "lead_time", ...) From b4665303094594c353bdc57ecc11c9daa4493b9f Mon Sep 17 00:00:00 2001 From: Michiel Van Ginderachter Date: Tue, 12 May 2026 13:58:26 +0000 Subject: [PATCH 6/7] pre-commit fixes --- src/mxalign/loaders/ifs_forecast.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/mxalign/loaders/ifs_forecast.py b/src/mxalign/loaders/ifs_forecast.py index 7fd9fae..8001d03 100644 --- a/src/mxalign/loaders/ifs_forecast.py +++ b/src/mxalign/loaders/ifs_forecast.py @@ -65,5 +65,3 @@ def _load(self): else: self.uncertainty = Uncertainty.DETERMINISTIC return ds.transpose("reference_time", "lead_time", ...) - - From 43713cd8d49cad04c2f2e03e1faafa54afa6e69c Mon Sep 17 00:00:00 2001 From: Michiel Van Ginderachter Date: Tue, 12 May 2026 14:07:30 +0000 Subject: [PATCH 7/7] remove members coord when det. --- src/mxalign/loaders/ifs_forecast.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mxalign/loaders/ifs_forecast.py b/src/mxalign/loaders/ifs_forecast.py index 8001d03..db02942 100644 --- a/src/mxalign/loaders/ifs_forecast.py +++ b/src/mxalign/loaders/ifs_forecast.py @@ -46,11 +46,11 @@ def _load(self): "step": "lead_time", } - if "number" in ds.dims: + if "number" in ds.dims and "number" in ds.coords: rename_dims["number"] = "member" - self - if "number" in ds.coords: rename_vars["number"] = "member" + else: + ds = ds.drop_vars("number") ds = ds.rename_dims({k: v for k, v in rename_dims.items() if k in ds.dims}) ds = ds.rename_vars({k: v for k, v in rename_vars.items() if k in ds.variables})