Support for step zero

ai_models_graphcast/model.py

@@ -192,9 +192,12 @@ def start_date(self) -> "datetime":
         return self.all_fields.order_by(valid_datetime="descending")[0].datetime
 
     def run(self):
+        # We ignore 'tp' so that we make sure that step 0 is a field of zero values
+        self.write_input_fields(self.fields_sfc, ignore=["tp"], accumulations=["tp"])
+        self.write_input_fields(self.fields_pl)
+
         with self.timer("Building model"):
             self.load_model()
-        # all_fields = self.all_fields.to_xarray()
 
         with self.timer("Creating input data (total)"):
             with self.timer("Creating training data"):

setup.py

@@ -43,9 +43,9 @@ def read(fname):
     include_package_data=True,
     # JAX requirements are in requirements.txt
     install_requires=[
-        "ai-models>=0.2.3",
-        "earthkit-meteo",
+        "ai-models>=0.3.6",
         "dm-tree",
+        "dm-haiku==0.0.10",
     ],
     zip_safe=True,
     keywords="tool",