Skip to content

Commit

Permalink
Flexibly add stream:enfo to metadata
Browse files Browse the repository at this point in the history
  • Loading branch information
HCookie committed Feb 11, 2025
1 parent 56f99ff commit dc5e13b
Showing 1 changed file with 4 additions and 1 deletion.
5 changes: 4 additions & 1 deletion src/ai_models_gencast/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ class GenCast(Model):
download_url = "https://storage.googleapis.com/dm_graphcast/gencast/{file}"

grib_edition = 1
grib_extra_metadata = {"type": "pf", "stream": "enfo"}
grib_extra_metadata = {"type": "pf"}

# Download
download_files = SHARED_DOWNLOAD_FILES
Expand Down Expand Up @@ -109,6 +109,9 @@ def __init__(self, **kwargs):
f"Number of ensemble members must match `member_number`,\nNot {self.num_ensemble_members=} and {self.member_number=}"
)

if "stream" not in getattr(self, "metadata", {}):
self.grib_extra_metadata["stream"] = "enfo"

# Jax doesn't seem to like passing configs as args through the jit. Passing it
# in via partial (instead of capture by closure) forces jax to invalidate the
# jit cache if you change configs.
Expand Down

0 comments on commit dc5e13b

Please sign in to comment.