fix(ecmwf): Full ensemble conversion fixes

openclimatefix · Jan 7, 2025 · 2c4ff74 · 2c4ff74
1 parent f2f9a57
commit 2c4ff74
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 11 deletions.
diff --git a/src/nwp_consumer/internal/entities/modelmetadata.py b/src/nwp_consumer/internal/entities/modelmetadata.py
@@ -170,7 +170,7 @@ class Models:
         resolution="0.1 degrees",
         expected_coordinates=NWPDimensionCoordinateMap(
             init_time=[],
-            step=list(range(0, 48, 1)),
+            step=list(range(0, 3, 1)),
             variable=[
                 Parameter.DOWNWARD_LONGWAVE_RADIATION_FLUX_GL,
                 Parameter.DOWNWARD_SHORTWAVE_RADIATION_FLUX_GL,
@@ -184,7 +184,7 @@ class Models:
                 Parameter.TEMPERATURE_SL,
                 Parameter.TOTAL_PRECIPITATION_RATE_GL,
             ],
-            ensemble_member=list(range(1, 51)),
+            ensemble_member=list(range(1, 3)),
             latitude=[v/10 for v in range(900, -900, -1)],
             longitude=[v/10 for v in range(-1800, 1800, 1)],
         ),

diff --git a/src/nwp_consumer/internal/repositories/raw_repositories/ecmwf_mars.py b/src/nwp_consumer/internal/repositories/raw_repositories/ecmwf_mars.py
@@ -265,18 +265,18 @@ def fetch_init_data(
             params=self.model().expected_coordinates.variable,
             init_time=it,
             steps=self.model().expected_coordinates.step,
-            nwse=",".join([str(ord) for ord in self.model().expected_coordinates.nwse()]),
+            nwse="/".join([str(ord) for ord in self.model().expected_coordinates.nwse()]),
             number=self.model().expected_coordinates.ensemble_member,
         )
 
-        match self.model().name:
-            case "ECMWF_HRES_IFS_0P1DEGREE":
-                yield delayed(self._download_and_convert)(req.as_operational_request())
-            case "ECMWF_ENS_STAT_0P1DEGREE":
-                for stat_req in [req.as_ensemble_mean_request(), req.as_ensemble_std_request()]:
-                    yield delayed(self._download_and_convert)(stat_req)
-            case "ECMWF_ENS_OP1DEGREE":
-                yield delayed(self._download_and_convert)(req.as_full_ensemble_request())
+        # Yield the download and convert function with the appropriate request type
+        if self.model().expected_coordinates.ensemble_stat is not None:
+            for stat_req in [req.as_ensemble_mean_request(), req.as_ensemble_std_request()]:
+                yield delayed(self._download_and_convert)(stat_req)
+        elif self.model().expected_coordinates.ensemble_member is not None:
+            yield delayed(self._download_and_convert)(req.as_full_ensemble_request())
+        else:
+            yield delayed(self._download_and_convert)(req.as_operational_request())
 
     def _download_and_convert(self, mr: _MARSRequest) -> ResultE[list[xr.DataArray]]:
         """Download and convert data from the ECMWF MARS server.
@@ -353,6 +353,8 @@ def _convert(path: pathlib.Path) -> ResultE[list[xr.DataArray]]:
                 .expand_dims("init_time")
                 .to_dataarray(name=ECMWFMARSRawRepository.model().name)
             )
+            if "number" in da.coords:
+                da = da.rename({"number": "ensemble_member"})
             da = (
                 da.drop_vars(
                     names=[