Handle sub-hourly steps (ecmwf#370)

Metamess · Apr 20, 2024 · 392cad7 · 392cad7
1 parent 33804b9
commit 392cad7
Show file tree

Hide file tree

Showing 5 changed files with 29 additions and 3 deletions.
diff --git a/cfgrib/cfmessage.py b/cfgrib/cfmessage.py
@@ -107,6 +107,17 @@ def to_grib_step(message, step_ns, step_unit=1, step_key="endStep:int", step_uni
     message[step_unit_key] = step_unit
 
 
+def from_grib_step_units(message):
+    # type: (abc.Field) -> float
+    # we always index steps in hours
+    return 1
+
+
+def to_grib_step_units(message, step_unit=1, step_unit_key="stepUnits:int"):
+    # type: (abc.MutableField, int, int, str, str) -> None
+    message[step_unit_key] = step_unit
+
+
 def from_grib_month(message, verifying_month_key="verifyingMonth", epoch=DEFAULT_EPOCH):
     # type: (abc.Field, str, datetime.datetime) -> int
     date = message[verifying_month_key]
@@ -149,6 +160,8 @@ def build_valid_time(time, step):
 COMPUTED_KEYS = {
     "time": (from_grib_date_time, to_grib_date_time),
     "step": (from_grib_step, to_grib_step),
+    "endStep": (from_grib_step, to_grib_step),
+    "stepUnits": (from_grib_step_units, to_grib_step_units),
     "valid_time": (
         functools.partial(from_grib_date_time, date_key="validityDate", time_key="validityTime"),
         functools.partial(to_grib_date_time, date_key="validityDate", time_key="validityTime"),

diff --git a/tests/sample-data/step_60m.grib b/tests/sample-data/step_60m.grib
diff --git a/tests/test_25_cfmessage.py b/tests/test_25_cfmessage.py
@@ -28,7 +28,7 @@ def test_to_grib_date_time() -> None:
 
 
 def test_from_grib_step() -> None:
-    message = {"endStep": 1, "stepUnits": 1}
+    message = {"endStep:int": 1, "stepUnits:int": 1}
     step_seconds = cfmessage.from_grib_step(message)
 
     assert step_seconds == 1
@@ -40,8 +40,8 @@ def test_to_grib_step() -> None:
 
     cfmessage.to_grib_step(message, step_ns, step_unit=1)
 
-    assert message["endStep"] == 1
-    assert message["stepUnits"] == 1
+    assert message["endStep:int"] == 1
+    assert message["stepUnits:int"] == 1
 
     with pytest.raises(ValueError):
         cfmessage.to_grib_step(message, 0, step_unit=3)

diff --git a/tests/test_40_xarray_store.py b/tests/test_40_xarray_store.py
@@ -2,6 +2,7 @@
 
 import gribapi  # type: ignore
 import numpy as np
+import pandas as pd
 import pytest
 
 xr = pytest.importorskip("xarray")  # noqa
@@ -17,6 +18,7 @@
 TEST_DATA_MULTIPLE_FIELDS = os.path.join(SAMPLE_DATA_FOLDER, "regular_gg_ml_g2.grib")
 TEST_DATA_DIFFERENT_STEP_TYPES = os.path.join(SAMPLE_DATA_FOLDER, "cfrzr_and_cprat.grib")
 TEST_DATA_DIFFERENT_STEP_TYPES_ZEROS = os.path.join(SAMPLE_DATA_FOLDER, "cfrzr_and_cprat_0s.grib")
+TEST_DATA_STEPS_IN_MINUTES = os.path.join(SAMPLE_DATA_FOLDER, "step_60m.grib")
 TEST_DATA_ALTERNATE_ROWS_MERCATOR = os.path.join(SAMPLE_DATA_FOLDER, "ds.waveh.5.grib")
 
 
@@ -157,6 +159,16 @@ def test_open_datasets_differet_step_types_zeros() -> None:
     assert res[1].cfrzr.attrs["GRIB_stepType"] == "avg"
 
 
+def test_open_dataset_steps_in_minutes() -> None:
+    res = xarray_store.open_dataset(TEST_DATA_STEPS_IN_MINUTES)
+
+    var = res["t2m"]
+    steps = var.step
+    assert steps[0] == pd.Timedelta("0 hours")
+    assert steps[1] == pd.Timedelta("1 hours")
+    assert steps[5] == pd.Timedelta("5 hours")
+
+
 def test_alternating_scanning_mercator() -> None:
     ds = xarray_store.open_dataset(TEST_DATA_ALTERNATE_ROWS_MERCATOR)
     values = ds.variables["shww"].data

diff --git a/tests/test_50_sample_data.py b/tests/test_50_sample_data.py
@@ -29,6 +29,7 @@
         "single_gridpoint",
         "spherical_harmonics",
         "t_analysis_and_fc_0",
+        "step_60m",
     ],
 )
 def test_open_dataset(grib_name: str) -> None: