diff --git a/src/nwp_consumer/internal/entities/coordinates.py b/src/nwp_consumer/internal/entities/coordinates.py index f77d730..0937ae9 100644 --- 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 index be0214f..776d87a 100644 --- 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 index d426ced..9767c7d 100644 --- 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 index ea297f4..1aeed18 100644 --- 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 index 7064d66..ab02e66 100644 --- 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(