Merge pull request #212 from ecmwf/feature/grib-output-from-climetlab

Update grib output code from climetlab
ecmwf · Oct 3, 2023 · c03643a · c03643a
2 parents f578185 + ce0479d
commit c03643a
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 12 deletions.
diff --git a/earthkit/data/readers/grib/output.py b/earthkit/data/readers/grib/output.py
@@ -20,6 +20,22 @@
 
 ACCUMULATIONS = {("tp", 2): {"productDefinitionTemplateNumber": 8}}
 
+_ORDER = (
+    "edition",
+    "setLocalDefinition",
+    "typeOfGeneratingProcess",
+    "productDefinitionTemplateNumber",
+)
+
+ORDER = {}
+for i, k in enumerate(_ORDER):
+    ORDER[k] = i
+
+
+def order(key):
+    ORDER.setdefault(key, len(ORDER))
+    return ORDER[key]
+
 
 class Combined:
     def __init__(self, handle, metadata):
@@ -89,13 +105,9 @@ def write(
         else:
             handle = template.handle.clone()
 
+        # print("->", metadata)
         self.update_metadata(handle, metadata, compulsory)
-
-        other = {}
-
-        for k, v in list(metadata.items()):
-            if not isinstance(v, (int, float, str)):
-                other[k] = metadata.pop(k)
+        # print("<-", metadata)
 
         if check_nans:
             import numpy as np
@@ -107,11 +119,13 @@ def write(
                 metadata["missingValue"] = missing_value
                 metadata["bitmapPresent"] = 1
 
-        LOG.debug("GribOutput.metadata %s, other %s", metadata, other)
+        metadata = {
+            k: v for k, v in sorted(metadata.items(), key=lambda x: order(x[0]))
+        }
 
-        handle.set_multiple(metadata)
+        LOG.debug("GribOutput.metadata %s", metadata)
 
-        for k, v in other.items():
+        for k, v in metadata.items():
             handle.set(k, v)
 
         handle.set_values(values)
@@ -165,6 +179,13 @@ def update_metadata(self, handle, metadata, compulsory):
 
         if "number" in metadata:
             compulsory += ("numberOfForecastsInEnsemble",)
+            productDefinitionTemplateNumber = {"tp": 11}
+            metadata[
+                "productDefinitionTemplateNumber"
+            ] = productDefinitionTemplateNumber.get(handle.get("shortName"), 1)
+
+        if metadata.get("type") in ("pf", "cf"):
+            metadata.setdefault("typeOfGeneratingProcess", 4)
 
         if "levelist" in metadata:
             metadata.setdefault("levtype", "pl")

diff --git a/tests/grib/test_grib_output.py b/tests/grib/test_grib_output.py
@@ -163,15 +163,15 @@ def test_grib_output_pl():
     sys.version_info < (3, 10),
     reason="ignore_cleanup_errors requires Python 3.10 or later",
 )
-@pytest.mark.skipif(True, reason="grib_set_values fails")
 def test_grib_output_tp():
     data = np.random.random((181, 360))
 
     with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmp:
         path = os.path.join(tmp, "a.grib")
 
         f = earthkit.data.new_grib_output(path, date=20010101)
-        f.write(data, param="tp", step=48)
+        # TODO: make it work for edition=2
+        f.write(data, param="tp", step=48, edition=1)
         f.close()
 
         ds = earthkit.data.from_source("file", path)
@@ -180,7 +180,7 @@ def test_grib_output_tp():
         assert ds[0].metadata("date") == 20010101
         assert ds[0].metadata("param") == "tp"
         assert ds[0].metadata("levtype") == "sfc"
-        assert ds[0].metadata("edition") == 2
+        assert ds[0].metadata("edition") == 1
         assert ds[0].metadata("step") == 48
 
         assert np.allclose(ds[0].to_numpy(), data, rtol=EPSILON, atol=EPSILON)