Add retrieve

ecmwf-lab · Aug 16, 2023 · ba36b21 · ba36b21
1 parent 5ec4fbe
commit ba36b21
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/ai_models/inputs/__init__.py b/ai_models/inputs/__init__.py
@@ -29,6 +29,7 @@ def fields_sfc(self):
                     param=self.owner.param_sfc,
                     grid=self.owner.grid,
                     area=self.owner.area,
+                    **self.owner.retrieve,
                 )
                 for date, time in self.owner.datetimes()
             ],

diff --git a/ai_models/model.py b/ai_models/model.py
@@ -52,6 +52,7 @@ def add(self, field):
 class Model:
     lagged = False
     assets_extra_dir = None
+    retrieve = {}  # Extra parameters for retrieve
 
     def __init__(self, input, output, download_assets, **kwargs):
         self.input = get_input(input, self, **kwargs)
@@ -252,6 +253,8 @@ def print_requests(self):
             grid=self.grid,
             area=self.area,
         )
+        first.update(self.retrieve)
+
         for date, time in self.datetimes():  # noqa F402
             param, level = self.param_level_pl