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 new file mode 100644 index 0000000..db02942 --- /dev/null +++ b/src/mxalign/loaders/ifs_forecast.py @@ -0,0 +1,67 @@ +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 Exception: + raise ImportError("Please install the cfgrib package to load IFS-Forecasts") + + name = "ifs-forecast" + + space = Space.GRID + time = Time.FORECAST + 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, + "step": -1, + "values": -1, + }, + **kwargs, + ) + + 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 and "number" in ds.coords: + rename_dims["number"] = "member" + 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}) + + 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.transpose("reference_time", "lead_time", ...)