Merge pull request #313 from karlwx/graphcast

Add GFS GraphCast AI Model
blaylockbk · May 15, 2024 · 92a7b36 · 92a7b36
2 parents 4498b3f + c19d56d
commit 92a7b36
Show file tree

Hide file tree

Showing 3 changed files with 504 additions and 151 deletions.
diff --git a/docs/user_guide/tutorial/model_notebooks/gfs.ipynb b/docs/user_guide/tutorial/model_notebooks/gfs.ipynb
diff --git a/herbie/models/gfs.py b/herbie/models/gfs.py
@@ -56,6 +56,21 @@ def template(self):
         self.IDX_SUFFIX = [".idx"]
         self.LOCALFILE = f"{self.get_remoteFileName}"
 
+class graphcast:
+    def template(self):
+        self.DESCRIPTION = "GraphCast Global Forecast System (EXPERIMENTAL)"
+        self.DETAILS = {
+            "aws document": "https://registry.opendata.aws/noaa-nws-graphcastgfs-pds/",
+        }
+        self.PRODUCTS = {
+            "pgrb2.0p25": "common fields, 0.25 degree resolution",
+        }
+        self.SOURCES = {
+            "aws": f"https://noaa-nws-graphcastgfs-pds.s3.amazonaws.com/graphcastgfs.{self.date:%Y%m%d/%H}/forecasts_13_levels/graphcastgfs.t{self.date:%H}z.{self.product}.f{self.fxx:03d}",
+        }
+        self.IDX_SUFFIX = [".idx"]
+        self.LOCALFILE = f"{self.get_remoteFileName}"
+
 
 class gfs_wave:
     def template(self):

diff --git a/tests/test_gfs.py b/tests/test_gfs.py
@@ -33,3 +33,24 @@ def test_gfs():
     f.unlink()
 
     H.xarray(filter)
+
+def test_graphcast():
+    H = Herbie(
+        today-timedelta(hours=12),
+        priority="aws",
+        product="pgrb2.0p25",
+        model="graphcast",
+        fxx=6,
+        save_dir=save_dir,
+        overwrite=True,
+    )
+
+    assert H.grib, "GFS grib2 file not found"
+    assert H.idx, "GFS index file not found"
+
+    filter = "HGT:500 mb"
+    f = H.download(filter)
+    assert H.get_localFilePath(filter).exists(), "File doesn't exist!"
+    f.unlink()
+
+    H.xarray(filter)