Aifs crps

README.md

@@ -139,6 +139,7 @@ Available models include but are not limited to:
 | Aurora | 0.25° | Transformer | 6h | Global |
 | FuXi | 0.25° | Transformer | 6h | Global |
 | AIFS | 0.25° | Transformer | 6h | Global |
+| AIFS Ensemble | 0.25° | Transformer Ensemble | 6h | Global |
 | StormCast | 3km | Diffusion + Regression | 1h | Regional (US) |
 | SFNO | 0.25° | Neural Operator | 6h | Global |
 | DLESyM | 0.25° | Convolutional | 6h | Global |

docs/modules/models.rst

@@ -27,6 +27,7 @@ Thus are typically used to generate forecast predictions.
    :template: prognostic.rst
 
    AIFS
+   AIFSENS
    Aurora
    CBottleVideo
    DiagnosticWrapper

docs/userguide/about/install.md

@@ -126,6 +126,30 @@ pip install earth2studio[aifs] --no-build-isolation
 uv add earth2studio --extra aifs
 ```
 
+:::
+::::
+:::::
+:::::{tab-item} AIFS Ensemble
+Notes: The AIFS Ensemble model relies on updated ECMWF checkpoints with ensemble
+sampling support. Similar to the deterministic AIFS variant this extra depends on
+[flash-attention](https://github.com/Dao-AILab/flash-attention), which can take a long
+time to compile. See the [troubleshooting docs](https://nvidia.github.io/earth2studio/userguide/support/troubleshooting.html)
+for compilation tips.
+
+::::{tab-set}
+:::{tab-item} pip
+
+```bash
+pip install earth2studio[aifsens] --no-build-isolation
+```
+
+:::
+:::{tab-item} uv
+
+```bash
+uv add earth2studio --extra aifsens
+```
+
 :::
 ::::
 :::::

docs/userguide/support/troubleshooting.md

@@ -80,9 +80,10 @@ For more information see the [NGC CLI docs](https://docs.ngc.nvidia.com/cli/inde
 
 Still having some problems? Open an issue.
 
-## Flash Attention has long build time for AIFS
+## Flash Attention has long build time for AIFS models
 
-This is a known issue with the library with several [issues](https://github.com/Dao-AILab/flash-attention/issues/1038)
+Both the deterministic AIFS and AIFS Ensemble extras depend on Flash Attention. This is
+a known issue with the library with several [issues](https://github.com/Dao-AILab/flash-attention/issues/1038)
 on the subject.
 There are a few options to try outside of just waiting for the build to complete.
 

earth2studio/lexicon/base.py

@@ -183,6 +183,7 @@ def __getitem__(cls, val: str) -> tuple[str, Callable]:
     "swvl2": "volumetric soil water layer 2 (m³/m³)",
     "stl1": "soil temperature level 1 (K)",
     "stl2": "soil temperature level 2 (K)",
+<<<<<<< HEAD
     "abi01c": "GOES ABI Channel 1 - Blue (0.47 µm)",
     "abi02c": "GOES ABI Channel 2 - Red (0.64 µm)",
     "abi03c": "GOES ABI Channel 3 - Vegetation (0.86 µm)",
@@ -234,6 +235,24 @@ def __getitem__(cls, val: str) -> tuple[str, Callable]:
     "ccl": "cloud cover fraction by atmospheric layers - low/mid/high clouds (0 - 1)",
     "aod": "aerosol optical depth at 550nm - atmospheric particulate loading",
     "vash": "volcanic ash detection and concentration from thermal IR signature",
+=======
+    "vis047": "Visible band at 0.47 µm (Blue)",
+    "vis064": "Visible band at 0.64 µm (Red)",
+    "nir086": "Near-IR band at 0.86 µm (Vegetation)",
+    "nir137": "Near-IR band at 1.37 µm (Cirrus)",
+    "nir161": "Near-IR band at 1.61 µm (Snow/Ice)",
+    "nir224": "Near-IR band at 2.24 µm (Cloud Particle Size)",
+    "ir390": "IR band at 3.90 µm (Shortwave Window)",
+    "ir619": "IR band at 6.19 µm (Upper-level Water Vapor)",
+    "ir695": "IR band at 6.95 µm (Mid-level Water Vapor)",
+    "ir734": "IR band at 7.34 µm (Lower-level Water Vapor)",
+    "ir850": "IR band at 8.50 µm (Cloud-top Phase)",
+    "ir961": "IR band at 9.61 µm (Ozone)",
+    "ir1035": "IR band at 10.35 µm (Clean IR Longwave Window)",
+    "ir1120": "IR band at 11.20 µm (IR Longwave Window)",
+    "ir1230": "IR band at 12.30 µm (Dirty IR Longwave Window)",
+    "ir1330": "IR band at 13.30 µm (CO2 Longwave IR)",
+>>>>>>> 8bc8969b (add variables to lexicon)
     "hcc": "high cloud cover (0 - 1)",
     "mcc": "medium cloud cover (0 - 1)",
     "lcc": "low cloud cover (0 - 1)",

earth2studio/models/px/__init__.py

@@ -17,6 +17,7 @@
 import warnings
 
 from earth2studio.models.px.aifs import AIFS
+from earth2studio.models.px.aifsens import AIFSENS
 from earth2studio.models.px.aurora import Aurora
 from earth2studio.models.px.base import PrognosticModel
 from earth2studio.models.px.cbottle_video import CBottleVideo