aiwq-checkpoints

README.md

@@ -13,6 +13,11 @@ individual trained checkpoints for the atmosphere and ocean components, which
 can be used to improve model variability in ensembles. The model architecture
 is a U-Net with padding operations modified to support using the HEALPix grid.
 
+The default package provided for this model contains the checkpoints used
+in the ECMWF AI Weather Quest S2S competition. These checkpoints are trained
+with CRPS loss and use sampled random noise to produce ensemble variability
+each forward pass, seeded with the `set_rng` method.
+
 For training recipies see [PhysicsNeMo](https://github.com/NVIDIA/physicsnemo/tree/main/examples/weather/dlwp_healpix), for inference [Earth2Studio](https://nvidia.github.io/earth2studio/examples/14_dlesym_example.html#sphx-glr-examples-14-dlesym-example-py).
 
 This model is ready for commercial/non-commercial use.
@@ -178,4 +183,4 @@ corresponds to an approximate 1 degree lat/lon grid. <br>
 
 NVIDIA believes Trustworthy AI is a shared responsibility and we have established policies and practices to enable development for a wide array of AI applications.  When downloaded or used in accordance with our terms of service, developers should work with their internal model team to ensure this model meets requirements for the relevant industry and use case and addresses unforeseen product misuse.  
 
-Please report security vulnerabilities or NVIDIA AI Concerns [here](https://www.nvidia.com/en-us/support/submit-security-vulnerability/).
+Please report security vulnerabilities or NVIDIA AI Concerns [here](https://www.nvidia.com/en-us/support/submit-security-vulnerability/).

checkpoints/v1.0/config.yaml

@@ -0,0 +1,59 @@
+io:
+  atmos_variables: ["z500", "tau300-700", "z1000", "t2m", "tcwv", "t850", "z250", "ws10m"]
+  ocean_variables: ["sst"]
+  atmos_coupling_variables: ["sst"]
+  ocean_coupling_variables: ["z1000", "ws10m"]
+  atmos_diagnostic_variables: []
+  ocean_diagnostic_variables: []
+  atmos_input_times: [-18, -12, -6, 0]
+  ocean_input_times: [-48, 0]
+  atmos_output_times: [6, 12, 18, 24, 30, 36, 42, 48, 54, 60, 66, 72, 78, 84, 90, 96]
+  ocean_output_times: [48, 96]
+  atmos_constants:
+    - land_sea_mask
+    - topography
+  ocean_constants:
+    - land_sea_mask
+models:
+  atmos_model_checkpoints:
+  - atmos_model_0.mdlus
+  - atmos_model_1.mdlus
+  - atmos_model_2.mdlus
+  - atmos_model_3.mdlus
+  ocean_model_checkpoints:
+  - ocean_model_0.mdlus
+  - ocean_model_1.mdlus
+  - ocean_model_2.mdlus
+  - ocean_model_3.mdlus
+data:
+  nside: 64
+  scaling:
+    tcwv:
+      mean: 24.147533416748047
+      std: 16.7200984954834
+    tau300-700:
+      mean: 61808.91796875
+      std: 2609.878662109375
+    sst:
+      mean: 291.37762451171875
+      std: 10.202408790588379
+    t2m:
+      mean: 287.40704345703125
+      std: 15.39192008972168
+    t850:
+      mean: 281.1323547363281
+      std: 12.35011100769043
+    ws10m:
+      mean: 6.144942283630371
+      std: 3.6602368354797363
+    z1000:
+      mean: 937.166748046875
+      std: 901.9874877929688
+    z500:
+      mean: 55511.10546875
+      std: 2737.783447265625
+    z250:
+      mean: 103580.8984375
+      std: 4610.26611328125
+cln:
+  enabled: false

checkpoints/v1.1/atmos/atmos_model_0.mdlus

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d3752ccd163c0157263f1dcdb722824cb4a2cca7368b07050172a8ba632a0950
+size 20650834

checkpoints/v1.1/config.yaml

@@ -0,0 +1,63 @@
+io:
+  atmos_variables: ["z500", "tau300-700", "z1000", "t2m", "tcwv", "t850", "z250", "ws10m", "ttr-3h"]
+  ocean_variables: ["sst"]
+  atmos_coupling_variables: ["sst"]
+  ocean_coupling_variables: ["z1000", "ws10m"]
+  atmos_input_times: [-18, -12, -6, 0]
+  ocean_input_times: [-48, 0]
+  atmos_output_times: [6, 12, 18, 24, 30, 36, 42, 48, 54, 60, 66, 72, 78, 84, 90, 96]
+  ocean_output_times: [48, 96]
+  atmos_constants:
+    - land_sea_mask
+    - topography
+  ocean_constants:
+    - land_sea_mask
+  atmos_diagnostic_variables: ["tp6", "msl"]
+  ocean_diagnostic_variables: []
+models:
+  atmos_model_checkpoints:
+  - atmos_model_0.mdlus
+  ocean_model_checkpoints:
+  - ocean_model_0.mdlus
+data:
+  nside: 64
+  scaling:
+    z500:
+      mean: 55482.71875
+      std: 2740.59130859375
+    tau300-700:
+      mean: 61743.29296875
+      std: 2610.57080078125
+    z1000:
+      mean: 936.1904907226562
+      std: 903.8828125
+    t2m:
+      mean: 287.40704345703125
+      std: 15.39192008972168
+    tcwv:
+      mean: 24.147533416748047
+      std: 16.7200984954834
+    t850:
+      mean: 280.5975036621094
+      std: 12.360532760620117
+    z250:
+      mean: 103530.8125
+      std: 4612.06005859375
+    ws10m:
+      mean: 6.1527934074401855
+      std: 3.650933265686035
+    ttr-3h:
+      mean: -2617476.9631782947
+      std: 494939.84
+    sst:
+      mean: 290.95892333984375
+      std: 10.623672485351562
+    tp6:
+      mean: 7.1346015
+      std: 4.5873095
+    msl:
+      mean: 101142.414
+      std: 1112.5691
+cln:
+  enabled: true
+  condition_shape: 32

checkpoints/v1.1/ocean/ocean_model_0.mdlus

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0f88299adef56d7fbfba500c2656ebba3bc13da70b150a6dbb00fed82f8715e4
+size 19555332

config.yaml

@@ -1,55 +1,18 @@
-io:
-  atmos_variables: ["z500", "tau300-700", "z1000", "t2m", "tcwv", "t850", "z250", "ws10m"]
-  ocean_variables: ["sst"]
-  atmos_coupling_variables: ["sst"]
-  ocean_coupling_variables: ["z1000", "ws10m"]
-  atmos_input_times: [-18, -12, -6, 0]
-  ocean_input_times: [-48, 0]
-  atmos_output_times: [6, 12, 18, 24, 30, 36, 42, 48, 54, 60, 66, 72, 78, 84, 90, 96]
-  ocean_output_times: [48, 96]
-  atmos_constants:
-    - land_sea_mask
-    - topography
-  ocean_constants:
-    - land_sea_mask
-models:
-  atmos_model_checkpoints:
-  - atmos_model_0.mdlus
-  - atmos_model_1.mdlus
-  - atmos_model_2.mdlus
-  - atmos_model_3.mdlus
-  ocean_model_checkpoints:
-  - ocean_model_0.mdlus
-  - ocean_model_1.mdlus
-  - ocean_model_2.mdlus
-  - ocean_model_3.mdlus
-data:
-  nside: 64
-  scaling:
-    tcwv:
-      mean: 24.147533416748047
-      std: 16.7200984954834
-    tau300-700:
-      mean: 61808.91796875
-      std: 2609.878662109375
-    sst:
-      mean: 291.37762451171875
-      std: 10.202408790588379
-    t2m:
-      mean: 287.40704345703125
-      std: 15.39192008972168
-    t850:
-      mean: 281.1323547363281
-      std: 12.35011100769043
-    ws10m:
-      mean: 6.144942283630371
-      std: 3.6602368354797363
-    z1000:
-      mean: 937.166748046875
-      std: 901.9874877929688
-    z500:
-      mean: 55511.10546875
-      std: 2737.783447265625
-    z250:
-      mean: 103580.8984375
-      std: 4610.26611328125
+# default_version is the checkpoint submitted to the ECMWF AI Weather Quest
+# competition (https://aiweatherquest.ecmwf.int/): conditional layer norm
+# (single set of weights, sampled per forecast) + diagnostic outputs
+# (tp6, msl). v1.0 is the prior ERA5 checkpoint (multiple
+# checkpoint pairs for ensembling, prognostic-only output), kept for
+# reproducibility and marked deprecated -- pass version="v1.0"
+# explicitly to load it.
+default_version: v1.1
+versions:
+  v1.0:
+    path: checkpoints/v1.0
+    description: "Prior ERA5 checkpoint; multiple checkpoint pairs for ensembling; prognostic-only output."
+    deprecated: true
+    deprecation_message: "DLESyM checkpoint version 'v1.0' will be removed in a future release; use 'v1.1' instead."
+  v1.1:
+    path: checkpoints/v1.1
+    description: "AI Weather Quest submission (https://aiweatherquest.ecmwf.int/); conditional layer norm ensembling + diagnostic outputs (tp6, msl)."
+    deprecated: false