Skip to content

Commit

Permalink
feat(model-repositories): Use available-models pattern
Browse files Browse the repository at this point in the history
  • Loading branch information
@devsjc
devsjc committed Dec 17, 2024
1 parent 0e5bdf0 commit abe151e
Show file tree
Hide file tree
Showing 12 changed files with 299 additions and 372 deletions.
3 changes: 2 additions & 1 deletion src/nwp_consumer/internal/entities/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
"""

from .repometadata import ModelRepositoryMetadata
from .modelmetadata import ModelMetadata
from .modelmetadata import ModelMetadata, Models
from .tensorstore import ParameterScanResult, TensorStore
from .postprocess import PostProcessOptions, CodecOptions
from .notification import PerformanceMetadata, StoreCreatedNotification, StoreAppendedNotification
Expand All @@ -27,6 +27,7 @@
__all__ = [
"ModelRepositoryMetadata",
"ModelMetadata",
"Models",
"ParameterScanResult",
"TensorStore",
"PostProcessOptions",
Expand Down
10 changes: 10 additions & 0 deletions src/nwp_consumer/internal/entities/coordinates.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -508,3 +508,13 @@ def crop(
longitude=[self.longitude[i] for i in lon_indices],
))

def nwse(self) -> tuple[float, float, float, float]:
"""Return the north, west, south, and east bounds of the map.
Returns:
A tuple of the north, west, south, and east bounds of the map.
"""
if self.latitude is not None and self.longitude is not None:
return self.latitude[0], self.longitude[0], self.latitude[-1], self.longitude[-1]
return 90.0, -180.0, -90.0, 180.0 # TODO: Cross this bridge when we come to it

280 changes: 140 additions & 140 deletions src/nwp_consumer/internal/entities/modelmetadata.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@
import logging

import numpy as np
from returns.result import ResultE, Success

from .coordinates import NWPDimensionCoordinateMap
from .parameters import Parameter
Expand Down Expand Up @@ -58,7 +57,7 @@ def __str__(self) -> str:
))
return pretty

def with_region(self, region: str) -> ResultE["ModelMetadata"]:
def with_region(self, region: str) -> "ModelMetadata":
"""Returns metadata for the given model cropped to the given region.
If an unknown region is given, the metadata is returned unchanged,
Expand All @@ -70,167 +69,168 @@ def with_region(self, region: str) -> ResultE["ModelMetadata"]:
north=62, west=-12, south=48, east=3,
).map(lambda coords: dataclasses.replace(
self, name=f"{self.name}_uk", expected_coordinates=coords,
))
)).unwrap()
case "india":
return self.expected_coordinates.crop(
north=35, west=67, south=6, east=97,
).map(lambda coords: dataclasses.replace(
self, name=f"{self.name}_india", expected_coordinates=coords,
))
)).unwrap()
case _:
log.warning(f"Unknown region '{region}', not cropping expected coordinates.")
return Success(self)

# --- Models --- #

ECMWF_HRES_IFS_0P1DEGREE: ModelMetadata = ModelMetadata(
name="HRES-IFS",
resolution="0.1 degrees",
expected_coordinates=NWPDimensionCoordinateMap(
init_time=[],
step=list(range(0, 85, 1)),
variable=[
Parameter.WIND_U_COMPONENT_10m,
Parameter.WIND_V_COMPONENT_10m,
Parameter.WIND_U_COMPONENT_100m,
Parameter.WIND_V_COMPONENT_100m,
Parameter.WIND_U_COMPONENT_200m,
Parameter.WIND_V_COMPONENT_200m,
Parameter.TEMPERATURE_SL,
Parameter.TOTAL_PRECIPITATION_RATE_GL,
Parameter.DOWNWARD_SHORTWAVE_RADIATION_FLUX_GL,
Parameter.DOWNWARD_LONGWAVE_RADIATION_FLUX_GL,
Parameter.CLOUD_COVER_HIGH,
Parameter.CLOUD_COVER_MEDIUM,
Parameter.CLOUD_COVER_LOW,
Parameter.CLOUD_COVER_TOTAL,
Parameter.SNOW_DEPTH_GL,
Parameter.VISIBILITY_SL,
Parameter.DIRECT_SHORTWAVE_RADIATION_FLUX_GL,
Parameter.DOWNWARD_ULTRAVIOLET_RADIATION_FLUX_GL,
],
latitude=[float(f"{lat / 10:.2f}") for lat in range(900, -900 - 1, -1)],
longitude=[float(f"{lon / 10:.2f}") for lon in range(-1800, 1800 + 1, 1)],
),
)
"""ECMWF's High Resolution Integrated Forecast System."""


ECMWF_ENS_0P1DEGREE: ModelMetadata = ModelMetadata(
name="ENS",
resolution="0.1 degrees",
expected_coordinates=NWPDimensionCoordinateMap(
init_time=[],
step=list(range(0, 85, 3)),
variable=[
Parameter.PRESSURE_MSL,
Parameter.WIND_SPEED_10m,
Parameter.WIND_SPEED_100m,
Parameter.TEMPERATURE_SL,
],
ensemble_stat=["mean", "std", "P10", "P25", "P75", "P90"],
latitude=[v/10 for v in range(900, -900, -1)],
longitude=[v/10 for v in range(-1800, 1800, 1)],
),
)
"""ECMWF's Ensemble Forecast System."""

NCEP_GFS_1DEGREE: ModelMetadata = ModelMetadata(
name="NCEP-GFS",
resolution="1 degree",
expected_coordinates=NWPDimensionCoordinateMap(
init_time=[],
step=list(range(0, 49, 3)),
variable=sorted(
[
return self


class Models:
"""Namespace containing known models."""

ECMWF_HRES_IFS_0P1DEGREE: ModelMetadata = ModelMetadata(
name="HRES-IFS",
resolution="0.1 degrees",
expected_coordinates=NWPDimensionCoordinateMap(
init_time=[],
step=list(range(0, 85, 1)),
variable=[
Parameter.WIND_U_COMPONENT_10m,
Parameter.WIND_V_COMPONENT_10m,
Parameter.WIND_U_COMPONENT_100m,
Parameter.WIND_V_COMPONENT_100m,
Parameter.WIND_U_COMPONENT_200m,
Parameter.WIND_V_COMPONENT_200m,
Parameter.TEMPERATURE_SL,
Parameter.CLOUD_COVER_TOTAL,
Parameter.TOTAL_PRECIPITATION_RATE_GL,
Parameter.DOWNWARD_SHORTWAVE_RADIATION_FLUX_GL,
Parameter.DOWNWARD_LONGWAVE_RADIATION_FLUX_GL,
Parameter.CLOUD_COVER_HIGH,
Parameter.CLOUD_COVER_MEDIUM,
Parameter.CLOUD_COVER_LOW,
Parameter.DOWNWARD_SHORTWAVE_RADIATION_FLUX_GL,
Parameter.DOWNWARD_LONGWAVE_RADIATION_FLUX_GL,
Parameter.TOTAL_PRECIPITATION_RATE_GL,
Parameter.CLOUD_COVER_TOTAL,
Parameter.SNOW_DEPTH_GL,
Parameter.RELATIVE_HUMIDITY_SL,
Parameter.VISIBILITY_SL,
Parameter.WIND_U_COMPONENT_10m,
Parameter.WIND_V_COMPONENT_10m,
Parameter.WIND_U_COMPONENT_100m,
Parameter.WIND_V_COMPONENT_100m,
Parameter.DIRECT_SHORTWAVE_RADIATION_FLUX_GL,
Parameter.DOWNWARD_ULTRAVIOLET_RADIATION_FLUX_GL,
],
latitude=[float(f"{lat / 10:.2f}") for lat in range(900, -900 - 1, -1)],
longitude=[float(f"{lon / 10:.2f}") for lon in range(-1800, 1800 + 1, 1)],
),
latitude=[float(lat) for lat in range(90, -90 - 1, -1)],
longitude=[float(lon) for lon in range(-180, 180 + 1, 1)],
),
)
"""NCEP's Global Forecast System."""

MO_UM_GLOBAL_17KM: ModelMetadata = ModelMetadata(
name="UM-Global",
resolution="17km",
expected_coordinates = NWPDimensionCoordinateMap(
init_time=[],
step=list(range(0, 48, 1)),
variable=[
Parameter.DOWNWARD_SHORTWAVE_RADIATION_FLUX_GL,
Parameter.CLOUD_COVER_TOTAL,
Parameter.CLOUD_COVER_HIGH,
Parameter.CLOUD_COVER_LOW,
Parameter.CLOUD_COVER_MEDIUM,
Parameter.RELATIVE_HUMIDITY_SL,
Parameter.SNOW_DEPTH_GL,
Parameter.TEMPERATURE_SL,
Parameter.WIND_U_COMPONENT_10m,
Parameter.WIND_V_COMPONENT_10m,
Parameter.VISIBILITY_SL,
],
latitude=[
float(f"{lat:.4f}") for lat in np.arange(89.856, -89.856 - 0.156, -0.156)
],
longitude=[
float(f"{lon:.4f}") for lon in np.concatenate([
np.arange(-45, 45, 0.234),
np.arange(45, 135, 0.234),
np.arange(135, 225, 0.234),
np.arange(225, 315, 0.234),
])
],
# TODO: Change to -180 -> 180
),
)
"""MetOffice's Unified Model, in the Global configuration, at a resolution of 17km."""

MO_UM_GLOBAL_10KM: ModelMetadata = ModelMetadata(
name="UM-Global",
resolution="10km",
expected_coordinates=NWPDimensionCoordinateMap(
init_time=[],
step=list(range(0, 55)),
variable=sorted(
[
)
"""ECMWF's High Resolution Integrated Forecast System."""

ECMWF_ENS_STAT_0P1DEGREE: ModelMetadata = ModelMetadata(
name="ENS-Stat",
resolution="0.1 degrees",
expected_coordinates=NWPDimensionCoordinateMap(
init_time=[],
step=list(range(0, 85, 3)),
variable=[
Parameter.PRESSURE_MSL,
Parameter.WIND_SPEED_10m,
Parameter.WIND_SPEED_100m,
Parameter.TEMPERATURE_SL,
],
ensemble_stat=["mean", "std", "P10", "P25", "P75", "P90"],
latitude=[v/10 for v in range(900, -900, -1)],
longitude=[v/10 for v in range(-1800, 1800, 1)],
),
)
"""Summary statistics from ECMWF's Ensemble Forecast System."""

NCEP_GFS_1DEGREE: ModelMetadata = ModelMetadata(
name="NCEP-GFS",
resolution="1 degree",
expected_coordinates=NWPDimensionCoordinateMap(
init_time=[],
step=list(range(0, 49, 3)),
variable=sorted(
[
Parameter.TEMPERATURE_SL,
Parameter.CLOUD_COVER_TOTAL,
Parameter.CLOUD_COVER_HIGH,
Parameter.CLOUD_COVER_MEDIUM,
Parameter.CLOUD_COVER_LOW,
Parameter.DOWNWARD_SHORTWAVE_RADIATION_FLUX_GL,
Parameter.DOWNWARD_LONGWAVE_RADIATION_FLUX_GL,
Parameter.TOTAL_PRECIPITATION_RATE_GL,
Parameter.SNOW_DEPTH_GL,
Parameter.RELATIVE_HUMIDITY_SL,
Parameter.VISIBILITY_SL,
Parameter.WIND_U_COMPONENT_10m,
Parameter.WIND_V_COMPONENT_10m,
Parameter.WIND_U_COMPONENT_100m,
Parameter.WIND_V_COMPONENT_100m,
],
),
latitude=[float(lat) for lat in range(90, -90 - 1, -1)],
longitude=[float(lon) for lon in range(-180, 180 + 1, 1)],
),
)
"""NCEP's Global Forecast System."""

MO_UM_GLOBAL_17KM: ModelMetadata = ModelMetadata(
name="UM-Global",
resolution="17km",
expected_coordinates = NWPDimensionCoordinateMap(
init_time=[],
step=list(range(0, 48, 1)),
variable=[
Parameter.DOWNWARD_SHORTWAVE_RADIATION_FLUX_GL,
Parameter.CLOUD_COVER_TOTAL,
Parameter.CLOUD_COVER_HIGH,
Parameter.CLOUD_COVER_MEDIUM,
Parameter.CLOUD_COVER_LOW,
Parameter.VISIBILITY_SL,
Parameter.CLOUD_COVER_MEDIUM,
Parameter.RELATIVE_HUMIDITY_SL,
Parameter.SNOW_DEPTH_GL,
Parameter.DOWNWARD_SHORTWAVE_RADIATION_FLUX_GL,
Parameter.TEMPERATURE_SL,
Parameter.WIND_U_COMPONENT_10m,
Parameter.WIND_V_COMPONENT_10m,
Parameter.VISIBILITY_SL,
],
latitude=[
float(f"{lat:.4f}") for lat in np.arange(89.856, -89.856 - 0.156, -0.156)
],
longitude=[
float(f"{lon:.4f}") for lon in np.concatenate([
np.arange(-45, 45, 0.234),
np.arange(45, 135, 0.234),
np.arange(135, 225, 0.234),
np.arange(225, 315, 0.234),
])
],
# TODO: Change to -180 -> 180
),
)
"""MetOffice's Unified Model, in the Global configuration, at a resolution of 17km."""

MO_UM_GLOBAL_10KM: ModelMetadata = ModelMetadata(
name="UM-Global",
resolution="10km",
expected_coordinates=NWPDimensionCoordinateMap(
init_time=[],
step=list(range(0, 55)),
variable=sorted(
[
Parameter.CLOUD_COVER_TOTAL,
Parameter.CLOUD_COVER_HIGH,
Parameter.CLOUD_COVER_MEDIUM,
Parameter.CLOUD_COVER_LOW,
Parameter.VISIBILITY_SL,
Parameter.RELATIVE_HUMIDITY_SL,
Parameter.SNOW_DEPTH_GL,
Parameter.DOWNWARD_SHORTWAVE_RADIATION_FLUX_GL,
Parameter.TEMPERATURE_SL,
Parameter.WIND_U_COMPONENT_10m,
Parameter.WIND_V_COMPONENT_10m,
],
),
latitude=[
float(f"{lat:.4f}")
for lat in np.arange(89.953125, -89.953125 - 0.09375, -0.09375)
],
longitude=[
float(f"{lon:.4f}")
for lon in np.arange(-179.929687, 179.929688 + 0.140625, 0.140625)
],
),
latitude=[
float(f"{lat:.4f}")
for lat in np.arange(89.953125, -89.953125 - 0.09375, -0.09375)
],
longitude=[
float(f"{lon:.4f}")
for lon in np.arange(-179.929687, 179.929688 + 0.140625, 0.140625)
],
),
)
"""MetOffice's Unified Model, in the Global configuration, at a resolution of 10km."""
)
"""MetOffice's Unified Model, in the Global configuration, at a resolution of 10km."""

5 changes: 5 additions & 0 deletions src/nwp_consumer/internal/entities/repometadata.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

import pandas as pd

from .modelmetadata import ModelMetadata
from .postprocess import PostProcessOptions


Expand Down Expand Up @@ -70,6 +71,9 @@ class ModelRepositoryMetadata:
postprocess_options: PostProcessOptions
"""Options for post-processing the data."""

available_models: dict[str, ModelMetadata]
"""A dictionary of available models and their metadata."""

def determine_latest_it_from(self, t: dt.datetime) -> dt.datetime:
"""Determine the latest available initialization time from a given time.
Expand Down Expand Up @@ -117,3 +121,4 @@ def __str__(self) -> str:
"\n".join(f"\t\t{var}={val}" for var, val in self.optional_env.items()),
))
return pretty

Loading

0 comments on commit abe151e

Please sign in to comment.