Add mars module (#91)

## Purpose

Add module that provides a helper class to build valid MARS requests in preparation for the FDB integration. Indeed, the FDB requests are a subset of MARS that requires the paramId. On top of that, the values for the stream, class, etc are relatively short and cryptic.

## Code changes:

- Add the mars module.

src/idpi/data/field_mappings.yml

@@ -25,6 +25,9 @@ ETADOT:
   ifs:
     name: etadot
     paramId: 77
+P:
+  cosmo:
+    paramId: 500001
 T:
   ifs:
     name: t
@@ -37,6 +40,12 @@ QV:
     paramId: 133
   cosmo:
     paramId: 500035
+QC:
+  cosmo:
+    paramId: 500100
+QI:
+  cosmo:
+    paramId: 500101
 U:
   ifs:
     name: u
@@ -49,6 +58,10 @@ V:
     paramId: 132
   cosmo:
     paramId: 500030
+W:
+  cosmo:
+    paramId: 500032
+    vertStag: true
 PS:
   ifs:
     name: sp
@@ -134,3 +147,10 @@ NSSS:
 OMEGA:
   cosmo:
     paramId: 500031
+HHL:
+  cosmo:
+    paramId: 500008
+    vertStag: true
+HSURF:
+  cosmo:
+    paramId: 500007

src/idpi/mars.py

@@ -0,0 +1,109 @@
+"""Mars request helper class."""
+
+# Standard library
+import dataclasses as dc
+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
+
+
+class Class(str, Enum):
+    OPERATIONAL_DATA = "od"
+
+
+class LevType(str, Enum):
+    MODEL_LEVEL = "ml"
+    PRESSURE_LEVEL = "pl"
+    SURFACE = "sfc"
+    SURFACE_OTHER = "sol"
+    POT_VORTICITY = "pv"
+    POT_TEMPERATURE = "pt"
+    DEPTH = "dp"
+
+
+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"
+
+
+class Stream(str, Enum):
+    ENS_DATA_ASSIMIL = "enda"
+    ENS_FORECAST = "enfo"
+
+
+class Type(str, Enum):
+    DETERMINISTIC = "det"
+    ENS_MEMBER = "ememb"
+    ENS_MEAN = "emean"
+    ENS_STD_DEV = "estdv"
+
+
+@cache
+def _load_mapping():
+    mapping_path = files("idpi.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 | int
+    date: str | None = None  # YYYYMMDD
+    time: str | None = None  # hhmm
+
+    expver: str = "0001"
+    levelist: int | tuple[int, ...] | None = None
+    number: int | tuple[int, ...] = 1
+    step: int | tuple[int, ...] = 0
+
+    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
+
+    def dump(self, exclude_defaults: bool = False):
+        root = pydantic.RootModel(self)
+        return root.model_dump(
+            mode="json",
+            by_alias=True,
+            exclude_none=True,
+            exclude_defaults=exclude_defaults,
+        )
+
+    def to_fdb(self):
+        mapping = _load_mapping()
+        param_id = mapping[self.param]["cosmo"]["paramId"]
+        staggered = mapping[self.param]["cosmo"].get("vertStag", False)
+
+        if self.levelist is None and self.levtype == LevType.MODEL_LEVEL:
+            n_lvl = N_LVL[self.model]
+            if staggered:
+                n_lvl += 1
+            levelist = tuple(range(1, n_lvl + 1))
+        else:
+            levelist = self.levelist
+
+        obj = dc.replace(self, param=param_id, levelist=levelist)
+        return obj.dump(exclude_defaults=False)

tests/test_idpi/test_mars.py

@@ -0,0 +1,68 @@
+# Third-party
+import pytest
+
+# First-party
+from idpi import mars
+
+
+@pytest.fixture
+def sample():
+    # for param=U, date=20200101, time=0000
+    return {
+        "class": "od",
+        "date": "20200101",
+        "expver": "0001",
+        "levtype": "ml",
+        "levelist": list(range(1, 81)),
+        "model": "COSMO-1E",
+        "number": 1,
+        "stream": "enfo",
+        "param": 500028,  # U
+        "time": "0000",
+        "type": "ememb",
+        "step": 0,
+    }
+
+
+def test_fdb_defaults(sample):
+    observed = mars.Request("U", date="20200101", time="0000").to_fdb()
+    expected = sample
+
+    assert observed == expected
+
+
+def test_fdb_c2e(sample):
+    observed = mars.Request(
+        "U",
+        date="20200101",
+        time="0000",
+        model=mars.Model.COSMO_2E,
+    ).to_fdb()
+    expected = sample | {"model": "COSMO-2E", "levelist": list(range(1, 61))}
+
+    assert observed == expected
+
+
+def test_fdb_sfc(sample):
+    observed = mars.Request(
+        "HSURF",
+        date="20200101",
+        time="0000",
+        levtype=mars.LevType.SURFACE,
+    ).to_fdb()
+    sample.pop("levelist")
+    expected = sample | {"param": 500007, "levtype": "sfc"}
+
+    assert observed == expected
+
+
+def test_request_raises():
+    with pytest.raises(ValueError):
+        mars.Request("U", date="20200101", time="0000", model="undef")
+
+
+def test_no_defaults():
+    observed = mars.Request("U").dump(exclude_defaults=True)
+    expected = {"param": "U"}
+
+    assert observed == expected