Skip to content

Commit

Permalink
feat(nwp-consumer): Add west-europe option for IFS
Browse files Browse the repository at this point in the history
  • Loading branch information
devsjc committed Jan 17, 2025
1 parent 2bb79a9 commit 21783fd
Show file tree
Hide file tree
Showing 4 changed files with 41 additions and 14 deletions.
8 changes: 8 additions & 0 deletions src/nwp_consumer/internal/entities/coordinates.py
Original file line number Diff line number Diff line change
Expand Up @@ -510,6 +510,14 @@ def crop(
" and both must sit between 180 and -180 degrees.",
))

if (south < self.latitude[0] or north > self.latitude[-1]
or west < self.longitude[0] or east > self.longitude[-1]):
return Failure(ValueError(
"Cannot crop coordinates to a region outside the bounds of the map. "
f"Crop region '{north, west, south, east}' not in "
f"map bounds '{self.nwse()}'.",
))

# Determine the indices of the region in the latitude and longitude lists
lat_indices = [
i for i, lat in enumerate(self.latitude)
Expand Down
18 changes: 12 additions & 6 deletions src/nwp_consumer/internal/entities/modelmetadata.py
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,12 @@ def with_region(self, region: str) -> "ModelMetadata":
).map(lambda coords: dataclasses.replace(
self, name=f"{self.name}_india", expected_coordinates=coords,
)).unwrap()
case "west-europe":
return self.expected_coordinates.crop(
north=63, west=-12, south=35, east=26,
).map(lambda coords: dataclasses.replace(
self, name=f"{self.name}_west-europe", expected_coordinates=coords,
)).unwrap()
case _:
log.warning(f"Unknown region '{region}', not cropping expected coordinates.")
return self
Expand All @@ -115,7 +121,7 @@ class Models:
"""Namespace containing known models."""

ECMWF_HRES_IFS_0P1DEGREE: ModelMetadata = ModelMetadata(
name="HRES-IFS",
name="hres-ifs",
resolution="0.1 degrees",
expected_coordinates=NWPDimensionCoordinateMap(
init_time=[],
Expand Down Expand Up @@ -147,7 +153,7 @@ class Models:
"""ECMWF's High Resolution Integrated Forecast System."""

ECMWF_ENS_STAT_0P1DEGREE: ModelMetadata = ModelMetadata(
name="ENS-Stat",
name="ens-stat",
resolution="0.1 degrees",
expected_coordinates=NWPDimensionCoordinateMap(
init_time=[],
Expand All @@ -166,7 +172,7 @@ class Models:
"""Summary statistics from ECMWF's Ensemble Forecast System."""

ECMWF_ENS_0P1DEGREE: ModelMetadata = ModelMetadata(
name="ENS",
name="ens",
resolution="0.1 degrees",
expected_coordinates=NWPDimensionCoordinateMap(
init_time=[],
Expand All @@ -193,7 +199,7 @@ class Models:
"""Full ensemble data from ECMWF's Ensemble Forecast System."""

NCEP_GFS_1DEGREE: ModelMetadata = ModelMetadata(
name="NCEP-GFS",
name="ncep-gfs",
resolution="1 degree",
expected_coordinates=NWPDimensionCoordinateMap(
init_time=[],
Expand Down Expand Up @@ -224,7 +230,7 @@ class Models:
"""NCEP's Global Forecast System."""

MO_UM_GLOBAL_17KM: ModelMetadata = ModelMetadata(
name="UM-Global",
name="um-global",
resolution="17km",
expected_coordinates = NWPDimensionCoordinateMap(
init_time=[],
Expand Down Expand Up @@ -259,7 +265,7 @@ class Models:
"""MetOffice's Unified Model, in the Global configuration, at a resolution of 17km."""

MO_UM_GLOBAL_10KM: ModelMetadata = ModelMetadata(
name="UM-Global",
name="um-global",
resolution="10km",
expected_coordinates=NWPDimensionCoordinateMap(
init_time=[],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -222,12 +222,14 @@ def repository() -> entities.RawRepositoryMetadata:
"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-west-europe": entities.Models.ECMWF_HRES_IFS_0P1DEGREE\
.with_region("west-europe"),
"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")\
.with_chunk_count_overrides(
{"latitude": 2, "longitude": 2, "variable": 1, "ensemble_member": 5},
),
.with_chunk_count_overrides(
{"latitude": 2, "longitude": 2, "variable": 1, "ensemble_member": 5},
),
},
)

Expand Down Expand Up @@ -331,6 +333,7 @@ def _convert(path: pathlib.Path) -> ResultE[list[xr.DataArray]]:
f"Error context: {e}",
))

processed_das: list[xr.DataArray] = []
try:
# Merge the datasets back into one
ds: xr.Dataset = xr.merge(
Expand All @@ -340,7 +343,7 @@ def _convert(path: pathlib.Path) -> ResultE[list[xr.DataArray]]:
)
del dss

if "ens" in ECMWFMARSRawRepository.model().name.lower():
if "ens-stat" in path.as_posix():
# Add in missing coordinates for mean/std data
if "enfo-es" in path.name:
ds = ds.expand_dims(dim={"ensemble_stat": ["std"]})
Expand All @@ -356,7 +359,7 @@ def _convert(path: pathlib.Path) -> ResultE[list[xr.DataArray]]:
.expand_dims("init_time")
.to_dataarray(name=ECMWFMARSRawRepository.model().name)
)
if "number" in da.coords:
if "ens" in path.as_posix():
da = da.rename({"number": "ensemble_member"})
da = (
da.drop_vars(
Expand All @@ -370,11 +373,21 @@ def _convert(path: pathlib.Path) -> ResultE[list[xr.DataArray]]:
.sortby(variables=["step", "variable", "longitude"])
.sortby(variables="latitude", ascending=False)
)
# Put each variable into its own DataArray:
# * Each raw file does not contain a full set of parameters
# * and so may not produce a contiguous subset of the expected coordinates.
processed_das.extend(
[
da.where(cond=da.coords["variable"] == v, drop=True)
for v in da.coords["variable"].values
],
)


except Exception as e:
return Failure(ValueError(
f"Error processing DataArray for path '{path!s}'. Error context: {e}",
))

return Success([da])
return Success(processed_das)

Original file line number Diff line number Diff line change
Expand Up @@ -267,8 +267,8 @@ def _convert(path: pathlib.Path) -> ResultE[list[xr.DataArray]]:
# * and so may not produce a contiguous subset of the expected coordinates.
processed_das.extend(
[
da.where(cond=da["variable"] == v, drop=True)
for v in da["variable"].values
da.where(cond=da.coords["variable"] == v, drop=True)
for v in da.coords["variable"].values
],
)

Expand Down

0 comments on commit 21783fd

Please sign in to comment.