better grib2 support

ecmwf-lab · Sep 23, 2024 · 33cae5b · 33cae5b
1 parent cf4b994
commit 33cae5b
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 13 deletions.
diff --git a/src/ai_models/inputs/opendata.py b/src/ai_models/inputs/opendata.py
@@ -61,6 +61,9 @@ def __init__(self, owner, **kwargs):
         self.owner = owner
 
     def _adjust(self, kwargs):
+
+        kwargs.setdefault("step", 0)
+
         if "level" in kwargs:
             # OpenData uses levelist instead of level
             kwargs["levelist"] = kwargs.pop("level")
@@ -105,7 +108,7 @@ def pl_load_source(self, **kwargs):
             kwargs["param"] = param
             gh_to_z = make_z_from_gh
 
-        logging.debug("load source ecmwf-open-data %s", kwargs)
+        logging.info("ecmwf-open-data %s", kwargs)
 
         opendata = recenter(ekd.from_source("ecmwf-open-data", **kwargs))
         opendata = gh_to_z(opendata)

diff --git a/src/ai_models/model.py b/src/ai_models/model.py
@@ -5,6 +5,7 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+import base64
 import datetime
 import json
 import logging
@@ -541,6 +542,8 @@ def write_input_fields(
         if ignore is None:
             ignore = []
 
+        fields.save("input.grib")
+
         with self.timer("Writing step 0"):
             for field in fields:
                 if field.metadata("shortName") in ignore:
@@ -560,18 +563,49 @@ def write_input_fields(
                 if accumulations_shape is None:
                     accumulations_shape = accumulations_template.shape
 
-                for param in accumulations:
-                    self.write(
-                        np.zeros(accumulations_shape, dtype=np.float32),
-                        stepType="accum",
-                        template=accumulations_template,
-                        param=param,
-                        startStep=0,
-                        endStep=0,
-                        date=int(self.start_datetime.strftime("%Y%m%d")),
-                        time=int(self.start_datetime.strftime("%H%M")),
-                        check=True,
-                    )
+                if accumulations_template.metadata("edition") == 1:
+                    for param in accumulations:
+
+                        self.write(
+                            np.zeros(accumulations_shape, dtype=np.float32),
+                            stepType="accum",
+                            template=accumulations_template,
+                            param=param,
+                            startStep=0,
+                            endStep=0,
+                            date=int(self.start_datetime.strftime("%Y%m%d")),
+                            time=int(self.start_datetime.strftime("%H%M")),
+                            check=True,
+                        )
+                else:
+                    # # TODO: Remove this when accumulations are supported for GRIB edition 2
+
+                    template = """
+                    R1JJQv//AAIAAAAAAAAA3AAAABUBAGIAABsBAQfoCRYGAAAAAQAAABECAAEAAQAJBAIwMDAxAAAA
+                    SAMAAA/XoAAAAAAG////////////////////AAAFoAAAAtEAAAAA/////wVdSoAAAAAAMIVdSoAV
+                    cVlwAAPQkAAD0JAAAAAAOgQAAAAIAcEC//8AAAABAAAAAAH//////////////wfoCRYGAAABAAAA
+                    AAECAQAAAAD/AAAAAAAAABUFAA/XoAAAAAAAAIAKAAAAAAAAAAYG/wAAAAUHNzc3N0dSSUL//wAC
+                    AAAAAAAAANwAAAAVAQBiAAAbAQEH6AkWDAAAAAEAAAARAgABAAEACQQBMDAwMQAAAEgDAAAP16AA
+                    AAAABv///////////////////wAABaAAAALRAAAAAP////8FXUqAAAAAADCFXUqAFXFZcAAD0JAA
+                    A9CQAAAAADoEAAAACAHBAv//AAAAAQAAAAAB//////////////8H6AkWDAAAAQAAAAABAgEAAAAA
+                    /wAAAAAAAAAVBQAP16AAAAAAAACACgAAAAAAAAAGBv8AAAAFBzc3Nzc=
+                    """
+
+                    template = base64.b64decode(template)
+                    accumulations_template = ekd.from_source("memory", template)[0]
+
+                    for param in accumulations:
+                        self.write(
+                            np.zeros(accumulations_shape, dtype=np.float32),
+                            stepType="accum",
+                            template=accumulations_template,
+                            param=param,
+                            startStep=0,
+                            endStep=0,
+                            date=int(self.start_datetime.strftime("%Y%m%d")),
+                            time=int(self.start_datetime.strftime("%H%M")),
+                            check=True,
+                        )
 
 
 def load_model(name, **kwargs):