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