fix(ecmwf): Enable working with full ensembles

openclimatefix · Jan 10, 2025 · 5ad56d0 · 5ad56d0
1 parent 2c4ff74
commit 5ad56d0
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 13 deletions.
diff --git a/src/nwp_consumer/internal/entities/coordinates.py b/src/nwp_consumer/internal/entities/coordinates.py
@@ -39,6 +39,7 @@
 import datetime as dt
 import json
 import logging
+import math
 from importlib.metadata import PackageNotFoundError, version
 
 import dask.array
@@ -224,6 +225,8 @@ def from_pandas(
                 ],
                 ensemble_stat=pd_indexes["ensemble_stat"].to_list() \
                     if "ensemble_stat" in pd_indexes else None,
+                ensemble_member=pd_indexes["ensemble_member"].to_list() \
+                    if "ensemble_member" in pd_indexes else None,
                 latitude=pd_indexes["latitude"].to_list() \
                     if "latitude" in pd_indexes else None,
                 longitude=pd_indexes["longitude"].to_list() \
@@ -335,7 +338,7 @@ def determine_region(
             return Failure(
                 KeyError(
                     "Cannot find slices in non-matching coordinate mappings: "
-                    "both objects must have identical dimensions (rank and labels)."
+                    "both objects must have identical dimensions (rank and labels). "
                     f"Got: {inner.dims} (inner) and {self.dims} (outer).",
                 ),
             )
@@ -416,17 +419,17 @@ def chunking(self, chunk_count_overrides: dict[str, int]) -> dict[str, int]:
             chunk_count_overrides: A dictionary mapping dimension labels to the
                 number of chunks to split the dimension into.
         """
-        out_dict: dict[str, int] = {
-            "init_time": 1,
-            "step": 1,
-            "variable": 1,
-        } | {
-            dim: len(getattr(self, dim)) // chunk_count_overrides.get(dim, 2)
-            if len(getattr(self, dim)) > 8 else 1
+        default_dict: dict[str, int] = {
+            dim: 1
+                if len(getattr(self, dim)) <= 8 or dim in ["init_time", "step", "variable"]
+                else math.ceil(len(getattr(self, dim)))
             for dim in self.dims
-            if dim not in ["init_time", "step", "variable"]
         }
 
+        out_dict = {}
+        for key in default_dict:
+            out_dict[key] = chunk_count_overrides.get(key, default_dict[key])
+
         return out_dict
 
 

diff --git a/src/nwp_consumer/internal/entities/modelmetadata.py b/src/nwp_consumer/internal/entities/modelmetadata.py
@@ -170,10 +170,11 @@ class Models:
         resolution="0.1 degrees",
         expected_coordinates=NWPDimensionCoordinateMap(
             init_time=[],
-            step=list(range(0, 3, 1)),
+            step=list(range(0, 48, 1)),
             variable=[
                 Parameter.DOWNWARD_LONGWAVE_RADIATION_FLUX_GL,
                 Parameter.DOWNWARD_SHORTWAVE_RADIATION_FLUX_GL,
+                Parameter.DOWNWARD_ULTRAVIOLET_RADIATION_FLUX_GL,
                 Parameter.WIND_U_COMPONENT_10m,
                 Parameter.WIND_V_COMPONENT_10m,
                 Parameter.SNOW_DEPTH_GL,
@@ -184,7 +185,7 @@ class Models:
                 Parameter.TEMPERATURE_SL,
                 Parameter.TOTAL_PRECIPITATION_RATE_GL,
             ],
-            ensemble_member=list(range(1, 3)),
+            ensemble_member=list(range(1, 51)),
             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
@@ -224,7 +224,8 @@ def repository() -> entities.RawRepositoryMetadata:
                 "hres-ifs-india": entities.Models.ECMWF_HRES_IFS_0P1DEGREE.with_region("india"),
                 "ens-stat-india": entities.Models.ECMWF_ENS_STAT_0P1DEGREE.with_region("india"),
                 "ens-stat-uk": entities.Models.ECMWF_ENS_STAT_0P1DEGREE.with_region("uk"),
-                "ens-uk": entities.Models.ECMWF_ENS_0P1DEGREE.with_region("uk"),
+                "ens-uk": entities.Models.ECMWF_ENS_0P1DEGREE.with_region("uk")\
+                    .with_chunk_count_overrides({"latitude": 1, "longitude": 1, "variable": 1}),
             },
         )
 

diff --git a/src/nwp_consumer/internal/repositories/raw_repositories/ecmwf_realtime.py b/src/nwp_consumer/internal/repositories/raw_repositories/ecmwf_realtime.py
@@ -85,7 +85,8 @@ def repository() -> entities.RawRepositoryMetadata:
             available_models={
                 "default": entities.Models.ECMWF_HRES_IFS_0P1DEGREE.with_region("uk"),
                 "hres-ifs-uk": entities.Models.ECMWF_HRES_IFS_0P1DEGREE.with_region("uk"),
-                "hres-ifs-india": entities.Models.ECMWF_HRES_IFS_0P1DEGREE.with_region("india"),
+                "hres-ifs-india": entities.Models.ECMWF_HRES_IFS_0P1DEGREE.with_region("india")\
+                    .with_chunk_count_overrides({"variable": 1}),
             },
         )
 

diff --git a/src/nwp_consumer/internal/services/consumer_service.py b/src/nwp_consumer/internal/services/consumer_service.py
@@ -146,6 +146,7 @@ def _create_suitable_store(
                     year=multiple_its.year,
                     month=multiple_its.month,
                 )
+        its = [it.replace(tzinfo=dt.UTC) for it in its]
 
         # Create a store for the data with the relevant init time coordinates
         return entities.TensorStore.initialize_empty_store(