From 3a7028bec6b86fb02dbde9f31edc58e708846e37 Mon Sep 17 00:00:00 2001
From: Harrison Cook <harrison.cook@ecmwf.int>
Date: Mon, 9 Dec 2024 16:02:00 +0000
Subject: [PATCH] Allow for model to set GRIB edition

---
 src/ai_models/model.py            | 5 +++--
 src/ai_models/outputs/__init__.py | 2 +-
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/ai_models/model.py b/src/ai_models/model.py
index 01eda32..c15ac10 100644
--- 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
index 2e9a616..fcb40aa 100644
--- 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,