Source code for meteodatalab.mars

"""Mars request helper class."""

# Standard library
import dataclasses as dc
import json
import typing
from collections.abc import Iterable
from enum import Enum
from functools import cache
from importlib.resources import files

# Third-party
import pydantic
import yaml
from pydantic import dataclasses as pdc

ValidationError = pydantic.ValidationError


[docs] class Class(str, Enum): OPERATIONAL_DATA = "od"
[docs] class LevType(str, Enum): MODEL_LEVEL = "ml" PRESSURE_LEVEL = "pl" SURFACE = "sfc" SURFACE_OTHER = "sol" POT_VORTICITY = "pv" POT_TEMPERATURE = "pt" DEPTH = "dp"
[docs] class Model(str, Enum): COSMO_1E = "COSMO-1E" COSMO_2E = "COSMO-2E" KENDA_1 = "KENDA-1" SNOWPOLINO = "SNOWPOLINO" ICON_CH1_EPS = "ICON-CH1-EPS" ICON_CH2_EPS = "ICON-CH2-EPS" KENDA_CH1 = "KENDA-CH1"
[docs] class Stream(str, Enum): ENS_DATA_ASSIMIL = "enda" ENS_FORECAST = "enfo"
[docs] class Type(str, Enum): DETERMINISTIC = "det" ENS_MEMBER = "ememb" ENS_MEAN = "emean" ENS_STD_DEV = "estdv"
[docs] class FeatureType(str, Enum): TIMESERIES = "timeseries"
[docs] class Point(typing.NamedTuple): lat: float lon: float
[docs] @dc.dataclass(frozen=True) class TimeseriesFeature: type: FeatureType = FeatureType.TIMESERIES points: list[Point] = dc.field(default_factory=list) start: int = 0 end: int = 0 @pydantic.validator("type") @classmethod def validate_type(cls, v: str) -> str: if v != FeatureType.TIMESERIES: raise ValueError("Wrong type") return v
@cache def _load_mapping(): mapping_path = files("meteodatalab.data").joinpath("field_mappings.yml") return yaml.safe_load(mapping_path.open()) N_LVL = { Model.COSMO_1E: 80, Model.COSMO_2E: 60, } @pdc.dataclass( frozen=True, config=pydantic.ConfigDict(use_enum_values=True), ) class Request: param: str | tuple[str, ...] date: str | None = None # YYYYMMDD time: str | None = None # hhmm expver: str = "0001" levelist: int | tuple[int, ...] | None = None number: int | tuple[int, ...] | None = None step: int | tuple[int, ...] | None = None class_: Class = dc.field( default=Class.OPERATIONAL_DATA, metadata=dict(alias="class"), ) levtype: LevType = LevType.MODEL_LEVEL model: Model = Model.COSMO_1E stream: Stream = Stream.ENS_FORECAST type: Type = Type.ENS_MEMBER feature: TimeseriesFeature | None = None def dump(self): if pydantic.__version__.startswith("1"): json_str = json.dumps(self, default=pydantic.json.pydantic_encoder) obj = json.loads(json_str.replace("class_", "class")) return {key: value for key, value in obj.items() if value is not None} root = pydantic.RootModel(self) return root.model_dump( mode="json", by_alias=True, exclude_none=True, ) def _param_id(self): mapping = _load_mapping() if isinstance(self.param, Iterable) and not isinstance(self.param, str): return [mapping[param]["cosmo"]["paramId"] for param in self.param] return mapping[self.param]["cosmo"]["paramId"] def _staggered(self): mapping = _load_mapping() if isinstance(self.param, Iterable) and not isinstance(self.param, str): return any( mapping[param]["cosmo"].get("vertStag", False) for param in self.param ) return mapping[self.param]["cosmo"].get("vertStag", False) def to_fdb(self) -> dict[str, typing.Any]: if self.date is None or self.time is None: raise RuntimeError("date and time are required fields for FDB.") if self.levelist is None and self.levtype == LevType.MODEL_LEVEL: n_lvl = N_LVL[self.model] if self._staggered(): n_lvl += 1 levelist: int | tuple[int, ...] | None = tuple(range(1, n_lvl + 1)) else: levelist = self.levelist obj = dc.replace(self, levelist=levelist) out = typing.cast(dict[str, typing.Any], obj.dump()) return out | {"param": self._param_id()} def to_polytope(self) -> dict[str, typing.Any]: result = self.to_fdb() if isinstance(result["param"], list): param: str | list[str] = [str(p) for p in result["param"]] else: param = str(result["param"]) return result | {"param": param}