Repository: ecmwf/earthkit-data
Pull request #281 (2024-01-17 12:28:03)

Support subhourly steps in grib


Patch:
diff --git a/earthkit/data/sources/list_of_dicts.py b/earthkit/data/sources/list_of_dicts.py
index 3417d6c..782b7da 100644
--- a/earthkit/data/sources/list_of_dicts.py
+++ b/earthkit/data/sources/list_of_dicts.py
@@ -16,6 +16,7 @@ from earthkit.data.core.metadata import RawMetadata
 from earthkit.data.readers.grib.index import GribFieldList
 from earthkit.data.readers.grib.metadata import GribMetadata
 from earthkit.data.utils.bbox import BoundingBox
+from earthkit.data.utils.dates import step_to_delta
 from earthkit.data.utils.projections import Projection
 
 LOG = logging.getLogger(__name__)
@@ -43,16 +44,15 @@ class VirtualGribMetadata(RawMetadata):
         ("dataTime", "time"),
         ("level", "levelist"),
         ("step", "endStep", "stepRange"),
+        ("param", "shortName"),
     }
 
     def __init__(self, m):
         super().__init__(m)
 
-    def _get_internal_key(self, key, astype=None, default=None, raise_on_missing=True):
+    def _get(self, key, astype=None, default=None, raise_on_missing=True):
         def _key_name(key):
-            if key == "param":
-                key = "shortName"
-            elif key == "_param_id":
+            if key == "_param_id":
                 key = "paramId"
             return key
 
@@ -71,13 +71,16 @@ class VirtualGribMetadata(RawMetadata):
                             key = k
                             break
 
+            if key not in self and key in ("step", "endStep", "stepRange"):
+                return self._default_step(astype=astype)
+
         if key_type is None:
             key_type = astype
 
         if key == "stepRange" and key_type is None:
             key_type = str
 
-        return super()._get_internal_key(
+        return super()._get(
             key, astype=key_type, default=default, raise_on_missing=raise_on_missing
         )
 
@@ -132,8 +135,8 @@ class VirtualGribMetadata(RawMetadata):
         }
 
     def _base_datetime(self):
-        date = self.get("date", None)
-        time = self.get("time", None)
+        date = int(self.get("date", None))
+        time = int(self.get("time", None))
         return datetime.datetime(
             date // 10000,
             date % 10000 // 100,
@@ -144,7 +147,14 @@ class VirtualGribMetadata(RawMetadata):
 
     def _valid_datetime(self):
         step = self.get("endStep", None)
-        return self._base_datetime() + datetime.timedelta(hours=step)
+        return self._base_datetime() + step_to_delta(step)
+
+    def _default_step(self, astype=None):
+        step = "0"
+        if astype is None:
+            return step
+        else:
+            return astype(step)
 
     def projection(self):
         return Projection.from_proj_string(self.get("projTargetString", None))
diff --git a/earthkit/data/utils/dates.py b/earthkit/data/utils/dates.py
index 5e809b6..7bd7898 100644
--- a/earthkit/data/utils/dates.py
+++ b/earthkit/data/utils/dates.py
@@ -8,11 +8,16 @@
 #
 
 import datetime
+import re
 
 import numpy as np
 
 from earthkit.data.wrappers import get_wrapper
 
+ECC_SECONDS_FACTORS = {"s": 1, "m": 60, "h": 3600}
+NUM_STEP_PATTERN = re.compile(r"\d+")
+SUFFIX_STEP_PATTERN = re.compile(r"\d+[a-zA-Z]{1}")
+
 
 def to_datetime(dt):
     if isinstance(dt, datetime.datetime):
@@ -99,3 +104,21 @@ def to_datetime_list(datetimes):  # noqa C901
 
 def to_date_list(obj):
     return sorted(set(to_datetime_list(obj)))
+
+
+def step_to_delta(step):
+    # TODO: make it work for all the ecCodes step formats
+    if isinstance(step, int):
+        return datetime.timedelta(hours=step)
+    elif isinstance(step, str):
+        if re.fullmatch(NUM_STEP_PATTERN, step):
+            sec = int(step) * 3600
+            return datetime.timedelta(seconds=sec)
+        elif re.fullmatch(SUFFIX_STEP_PATTERN, step):
+            factor = ECC_SECONDS_FACTORS.get(step[-1], None)
+            if factor is None:
+                raise ValueError(f"Unsupported ecCodes step units in step: {step}")
+            sec = int(step[:-1]) * factor
+            return datetime.timedelta(seconds=sec)
+
+    raise ValueError(f"Unsupported ecCodes step: {step}")
diff --git a/earthkit/data/utils/message.py b/earthkit/data/utils/message.py
index 979b713..bf3a37a 100644
--- a/earthkit/data/utils/message.py
+++ b/earthkit/data/utils/message.py
@@ -20,6 +20,8 @@ from earthkit.data.core.caching import CACHE, auxiliary_cache_file
 
 LOG = logging.getLogger(__name__)
 
+os.environ["ECCODES_GRIB_SHOW_HOUR_STEPUNIT"] = "1"
+
 # For some reason, cffi can get stuck in the GC if that function
 # needs to be called defined for the first time in a GC thread.
 try:


Test patch:
diff --git a/tests/grib/test_grib_output.py b/tests/grib/test_grib_output.py
index 8034cd9..a7ffa27 100644
--- a/tests/grib/test_grib_output.py
+++ b/tests/grib/test_grib_output.py
@@ -128,7 +128,7 @@ def test_grib_output_mars_labeling():
 
         assert ds[0].metadata("date") == 20010101
         assert ds[0].metadata("edition") == 2
-        assert ds[0].metadata("step") == 24
+        assert ds[0].metadata("step", astype=int) == 24
         assert ds[0].metadata("expver") == "test"
         assert ds[0].metadata("levtype") == "sfc"
         assert ds[0].metadata("param") == "msl"
@@ -189,7 +189,7 @@ def test_grib_output_tp():
         assert ds[0].metadata("param") == "tp"
         assert ds[0].metadata("levtype") == "sfc"
         assert ds[0].metadata("edition") == 1
-        assert ds[0].metadata("step") == 48
+        assert ds[0].metadata("step", astype=int) == 48
         assert ds[0].metadata("generatingProcessIdentifier") == 255
 
         assert np.allclose(ds[0].to_numpy(), data, rtol=EPSILON, atol=EPSILON)
diff --git a/tests/sources/test_list_of_dicts.py b/tests/sources/test_list_of_dicts.py
new file mode 100644
index 0000000..b2c749d
--- /dev/null
+++ b/tests/sources/test_list_of_dicts.py
@@ -0,0 +1,60 @@
+#!/usr/bin/env python3
+
+# (C) Copyright 2020 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+
+import datetime
+
+import pytest
+
+from earthkit.data import from_source
+
+
+@pytest.fixture
+def lod():
+    prototype = {
+        "gridType": "regular_ll",
+        "Nx": 2,
+        "Ny": 3,
+        "distinctLatitudes": [-10.0, 0.0, 10.0],
+        "distinctLongitudes": [0.0, 10.0],
+        "_param_id": 167,
+        "values": [[1, 2], [3, 4], [5, 6]],
+        "date": "20180801",
+        "time": "1200",
+    }
+    return [
+        {"param": "t", "levelist": 500, **prototype},
+        {"param": "t", "levelist": 850, **prototype},
+        {"param": "u", "levelist": 500, **prototype},
+        {"param": "u", "levelist": 850, **prototype},
+        {"param": "d", "levelist": 850, **prototype},
+        {"param": "d", "levelist": 600, **prototype},
+    ]
+
+
+def test_list_of_dicts(lod):
+    ds = from_source("list-of-dicts", lod)
+
+    assert len(ds) == 6
+    ref = [("t", 500), ("t", 850), ("u", 500), ("u", 850), ("d", 850), ("d", 600)]
+    assert ds.metadata("param", "levelist") == ref
+
+    assert ds[0].metadata("step") == "0"
+
+    assert ds[0].datetime() == {
+        "base_time": datetime.datetime(2018, 8, 1, 12, 0),
+        "valid_time": datetime.datetime(2018, 8, 1, 12, 0),
+    }
+
+
+if __name__ == "__main__":
+    from earthkit.data.testing import main
+
+    main()
diff --git a/tests/utils/test_dates.py b/tests/utils/test_dates.py
new file mode 100644
index 0000000..031c70e
--- /dev/null
+++ b/tests/utils/test_dates.py
@@ -0,0 +1,47 @@
+#!/usr/bin/env python3
+
+# (C) Copyright 2020 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+
+from datetime import timedelta
+
+import pytest
+
+from earthkit.data.utils.dates import step_to_delta
+
+
+@pytest.mark.parametrize(
+    "step,expected_delta,error",
+    [
+        (12, timedelta(hours=12), None),
+        ("12h", timedelta(hours=12), None),
+        ("12s", timedelta(seconds=12), None),
+        ("12m", timedelta(minutes=12), None),
+        ("1m", timedelta(minutes=1), None),
+        ("", None, ValueError),
+        ("m", None, ValueError),
+        ("1Z", None, ValueError),
+        ("m1", None, ValueError),
+        ("-1", None, ValueError),
+        ("-1s", None, ValueError),
+        ("1.1s", None, ValueError),
+    ],
+)
+def test_step_to_delta(step, expected_delta, error):
+    if error is None:
+        assert step_to_delta(step) == expected_delta
+    else:
+        with pytest.raises(error):
+            step_to_delta(step)
+
+
+if __name__ == "__main__":
+    from earthkit.data.testing import main
+
+    main(__file__)