From cec1365ca2bbd25b452d6d7e98f75aa0a667f8c3 Mon Sep 17 00:00:00 2001
From: Stagiair user Niels Carlier <nicarlier@kili.oma.be>
Date: Tue, 12 May 2026 09:16:01 +0200
Subject: [PATCH 1/2] updated ifs forecast loader to also handle ensemble data

---
 src/mxalign/loaders/ifs_forecast.py | 83 +++++++++++++++++++----------
 1 file changed, 56 insertions(+), 27 deletions(-)

diff --git a/src/mxalign/loaders/ifs_forecast.py b/src/mxalign/loaders/ifs_forecast.py
index 70964b0..b291d36 100644
--- a/src/mxalign/loaders/ifs_forecast.py
+++ b/src/mxalign/loaders/ifs_forecast.py
@@ -1,52 +1,81 @@
-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")
+
+        if "lead_time" in ds.coords and np.issubdtype(ds["lead_time"].dtype, np.timedelta64):
+            ds = ds.assign_coords(
+                lead_time=(ds["lead_time"].values / np.timedelta64(1, "h")).astype(int)
+            )
+            ds["lead_time"].attrs["units"] = "h"
+
+        return ds
+
+    def _get_properties(self, ds):
+        if "member" in ds.dims:
+            uncertainty = Uncertainty.ENSEMBLE
+        elif "quantile" in ds.dims:
+            uncertainty = Uncertainty.QUANTILE
+        else:
+            uncertainty = Uncertainty.DETERMINISTIC
+
+        return Properties(
+            space=Space.GRID,
+            time=Time.FORECAST,
+            uncertainty=uncertainty,
+        )
\ No newline at end of file

From 7ee849956f64ebba6f0dea839db741d87142c743 Mon Sep 17 00:00:00 2001
From: Stagiair user Niels Carlier <nicarlier@kili.oma.be>
Date: Tue, 12 May 2026 13:00:43 +0200
Subject: [PATCH 2/2] updates after review

---
 src/mxalign/loaders/ifs_forecast.py | 25 ++++++-------------------
 1 file changed, 6 insertions(+), 19 deletions(-)

diff --git a/src/mxalign/loaders/ifs_forecast.py b/src/mxalign/loaders/ifs_forecast.py
index b291d36..22ae4f9 100644
--- a/src/mxalign/loaders/ifs_forecast.py
+++ b/src/mxalign/loaders/ifs_forecast.py
@@ -58,24 +58,11 @@ def _load(self):
         if "surface" in ds.variables:
             ds = ds.drop_vars("surface")
 
-        if "lead_time" in ds.coords and np.issubdtype(ds["lead_time"].dtype, np.timedelta64):
-            ds = ds.assign_coords(
-                lead_time=(ds["lead_time"].values / np.timedelta64(1, "h")).astype(int)
-            )
-            ds["lead_time"].attrs["units"] = "h"
-
         return ds
 
-    def _get_properties(self, ds):
-        if "member" in ds.dims:
-            uncertainty = Uncertainty.ENSEMBLE
-        elif "quantile" in ds.dims:
-            uncertainty = Uncertainty.QUANTILE
-        else:
-            uncertainty = Uncertainty.DETERMINISTIC
-
-        return Properties(
-            space=Space.GRID,
-            time=Time.FORECAST,
-            uncertainty=uncertainty,
-        )
\ No newline at end of file
+    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