Add GraphCast 1 degree model (#256)

* Added GraphCast small model
* Added install group
* Update WB2 data source with needed variables
* Added tests and docs

---------

Co-authored-by: Rodrigo Almeida <[email protected]>
Co-authored-by: Oliver <[email protected]>
Co-authored-by: Nicholas Geneva <[email protected]>

Author: 4 people

.gitlab-ci.yml

@@ -86,6 +86,10 @@ pytest:
   extends:
     - .earth2studio:test:base
     - .test:cache
+  variables:
+    XLA_PYTHON_CLIENT_ALLOCATOR: platform
+    XLA_PYTHON_CLIENT_PREALLOCATE: false
+    # https://docs.jax.dev/en/latest/gpu_memory_allocation.html
   rules:
     - if: $PIPELINE_TEST_ALL == "true"
       when: on_success

CHANGELOG.md

@@ -10,6 +10,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Added Graphcast 1 degree model
+
 ### Changed
 
 ### Deprecated

Makefile

@@ -5,7 +5,7 @@ install:
 .PHONY: install-docker
 install-docker:
 	uv pip install --system --break-system-packages .[all,aifs] --group dev
-	uv pip install --system --break-system-packages zarr==3.0.7
+	uv pip install --system --break-system-packages zarr~=3.0
 
 .PHONY: setup-ci
 setup-ci:

docs/modules/models.rst

@@ -33,6 +33,7 @@ Thus are typically used to generate forecast predictions.
    models.px.FCN
    models.px.FengWu
    models.px.FuXi
+   models.px.GraphCastSmall
    models.px.InterpModAFNO
    models.px.Pangu24
    models.px.Pangu6

docs/userguide/about/install.md

@@ -241,6 +241,26 @@ pip install earth2studio[fuxi]
 uv add earth2studio --extra fuxi
 ```
 
+:::
+::::
+:::::
+:::::{tab-item} GraphCast
+Notes: Requires JAX with CUDA support and the GraphCast package.
+
+::::{tab-set}
+:::{tab-item} pip
+
+```bash
+pip install earth2studio[graphcast]
+```
+
+:::
+:::{tab-item} uv
+
+```bash
+uv add earth2studio --extra graphcast
+```
+
 :::
 ::::
 :::::

earth2studio/lexicon/arco.py

@@ -94,6 +94,7 @@ class ARCOLexicon(metaclass=LexiconType):
         "msl": "mean_sea_level_pressure::",
         "tcwv": "total_column_water_vapour::",
         "tp": "total_precipitation::",
+        "lsm": "land_sea_mask::",
         "z": "geopotential_at_surface::",
     }
     VOCAB.update({f"u{level}": f"u_component_of_wind::{level}" for level in LEVELS})

earth2studio/lexicon/base.py

@@ -45,6 +45,7 @@ def __getitem__(cls, val: str) -> tuple[str, Callable]:
     "fg10m": "maximum 10 m wind gust since previous post-processing",
     "lsm": "land-sea mask",
     "zsl": "geopotential at mean sea-level",
+    "z": "geopotential at surface",
     "uvcossza": "cosine of solar zenith angle",
     "refc": "Maximum/Composite radar reflectivity (dB)",
     "csnow": "categorical snow",

earth2studio/lexicon/wb2.py

@@ -40,6 +40,8 @@ class WB2Lexicon(metaclass=LexiconType):
         "v10m": "10m_v_component_of_wind::",
         "t2m": "2m_temperature::",
         "sp": "surface_pressure::",
+        "lsm": "land_sea_mask::",
+        "z": "geopotential_at_surface::",
         "msl": "mean_sea_level_pressure::",
         "sst": "sea_surface_temperature::",
         "tcwv": "total_column_water_vapour::",

earth2studio/models/px/__init__.py

@@ -24,6 +24,7 @@
 from earth2studio.models.px.fcn import FCN
 from earth2studio.models.px.fengwu import FengWu
 from earth2studio.models.px.fuxi import FuXi
+from earth2studio.models.px.graphcast_small import GraphCastSmall
 from earth2studio.models.px.interpmodafno import InterpModAFNO
 from earth2studio.models.px.pangu import Pangu3, Pangu6, Pangu24
 from earth2studio.models.px.persistence import Persistence