Allow for model to set GRIB edition

ecmwf-lab · Dec 9, 2024 · 3a7028b · 3a7028b
1 parent c5847cd
commit 3a7028b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/src/ai_models/model.py b/src/ai_models/model.py
@@ -63,6 +63,7 @@ class Model:
     assets_extra_dir = None
     retrieve = {}  # Extra parameters for retrieve
     version = 1  # To be overriden in subclasses
+    grib_edition = 2  # Default GRIB edition
     grib_extra_metadata = {}  # Extra metadata for grib files
 
     param_level_ml = ([], [])  # param, level
@@ -260,7 +261,7 @@ def _datetimes(self, dates):
         assert isinstance(time, int)
         if time < 100:
             time *= 100
-        assert time in (0, 600, 1200, 1800), time
+        # assert time in (0, 600, 1200, 1800), time
 
         lagged = self.lagged
         if not lagged:
@@ -300,7 +301,7 @@ def datetimes(self, step=0):
         if time < 100:
             time *= 100
 
-        assert time in (0, 600, 1200, 1800), time
+        # assert time in (0, 600, 1200, 1800), time
 
         full = datetime.datetime(
             date // 10000,

diff --git a/src/ai_models/outputs/__init__.py b/src/ai_models/outputs/__init__.py
@@ -38,7 +38,7 @@ def __init__(self, owner, path, metadata, **kwargs):
 
     @cached_property
     def grib_keys(self):
-        edition = self.metadata.pop("edition", 2)
+        edition = self.metadata.pop("edition", self.owner.grib_edition)
 
         _grib_keys = dict(
             edition=edition,