From 80f65ece77c2cb591c2d0d9214b9e3813cc7c1bb Mon Sep 17 00:00:00 2001
From: Gert Mertes <gert.mertes@ecmwf.int>
Date: Wed, 20 Nov 2024 12:00:09 +0000
Subject: [PATCH] Support opendata to gaussian grid interpolation

---
 pyproject.toml                      |  2 +-
 src/ai_models/inputs/interpolate.py | 13 ++++++++++++-
 src/ai_models/inputs/opendata.py    |  4 ++--
 src/ai_models/outputs/__init__.py   |  7 +++++--
 4 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 7bf13af..f81aade 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -44,7 +44,7 @@ dependencies = [
     "earthkit-data>=0.10.3",
     "earthkit-meteo",
     "earthkit-regrid",
-    "eccodes>=2.37",
+    "eccodes>=2.38.3",
     "ecmwf-api-client",
     "ecmwf-opendata",
     "entrypoints",
diff --git a/src/ai_models/inputs/interpolate.py b/src/ai_models/inputs/interpolate.py
index c5e3145..3fe04dd 100644
--- a/src/ai_models/inputs/interpolate.py
+++ b/src/ai_models/inputs/interpolate.py
@@ -21,6 +21,8 @@ def __init__(self, grid, source, metadata):
         self.source = list(source) if isinstance(source, tuple) else source
         self.metadata = metadata
 
+        self._reduced_gg = metadata.get("gridType") == "reduced_gg"
+
     def __call__(self, ds):
         tmp = temp_file()
 
@@ -29,7 +31,16 @@ def __call__(self, ds):
         result = []
         for f in tqdm.tqdm(ds, delay=0.5, desc="Interpolating", leave=False):
             data = ekr.interpolate(f.to_numpy(), dict(grid=self.source), dict(grid=self.grid))
-            out.write(data, template=f, **self.metadata)
+            template = f
+
+            if self._reduced_gg:
+                # template is missing the pl matrix required to output a reduced_gg grib
+                # this is a hack to let earthkit generate it for us
+                template = None
+                keys = ("shortName", "levelist", "date", "time", "step", "number")
+                self.metadata.update({key: f._metadata[key] for key in keys if key in f._metadata})
+
+            out.write(data, template=template, **self.metadata)
 
         out.close()
 
diff --git a/src/ai_models/inputs/opendata.py b/src/ai_models/inputs/opendata.py
index 14813c5..718c444 100644
--- a/src/ai_models/inputs/opendata.py
+++ b/src/ai_models/inputs/opendata.py
@@ -45,8 +45,8 @@
             Nj=1801,
         ),
     ),
-    # "N320": ("0p25", (0.25, 0.25), True, False, dict(gridType='reduced_gg')),
-    # "O96": ("0p25", (0.25, 0.25), True, False, dict(gridType='reduced_gg', )),
+    "N320": ("0p25", (0.25, 0.25), True, False, dict(gridType="reduced_gg", N=320)),
+    "O96": ("0p25", (0.25, 0.25), True, False, dict(gridType="reduced_gg", N=96)),
 }
 
 
diff --git a/src/ai_models/outputs/__init__.py b/src/ai_models/outputs/__init__.py
index 2e9a616..7260f8b 100644
--- a/src/ai_models/outputs/__init__.py
+++ b/src/ai_models/outputs/__init__.py
@@ -58,6 +58,9 @@ def output(self):
 
     def write(self, data, *args, check=False, **kwargs):
 
+        if kwargs.get("param") in ("cp", "tp"):
+            kwargs["edition"] = 1
+
         try:
             handle, path = self.output.write(data, *args, **kwargs)
 
@@ -85,9 +88,9 @@ def write(self, data, *args, check=False, **kwargs):
             # Check that the GRIB keys are as expected
 
             if kwargs.get("expver") is None:
-                ignore = ("template", "check_nans", "expver", "class", "type", "stream")
+                ignore = ("edition", "template", "check_nans", "expver", "class", "type", "stream")
             else:
-                ignore = ("template", "check_nans")
+                ignore = ("edition", "template", "check_nans")
 
             for key, value in itertools.chain(self.grib_keys.items(), kwargs.items()):
                 if key in ignore: