Graph Machine Learning
AnemoI
English
hcookie129 commited on
Commit
3eb31a0
·
0 Parent(s):

AIFS-ENS 1.0 Release

Browse files
.gitattributes ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ *.7z filter=lfs diff=lfs merge=lfs -text
2
+ *.arrow filter=lfs diff=lfs merge=lfs -text
3
+ *.bin filter=lfs diff=lfs merge=lfs -text
4
+ *.bz2 filter=lfs diff=lfs merge=lfs -text
5
+ *.ckpt filter=lfs diff=lfs merge=lfs -text
6
+ *.ftz filter=lfs diff=lfs merge=lfs -text
7
+ *.gz filter=lfs diff=lfs merge=lfs -text
8
+ *.h5 filter=lfs diff=lfs merge=lfs -text
9
+ *.joblib filter=lfs diff=lfs merge=lfs -text
10
+ *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
+ *.mlmodel filter=lfs diff=lfs merge=lfs -text
12
+ *.model filter=lfs diff=lfs merge=lfs -text
13
+ *.msgpack filter=lfs diff=lfs merge=lfs -text
14
+ *.npy filter=lfs diff=lfs merge=lfs -text
15
+ *.npz filter=lfs diff=lfs merge=lfs -text
16
+ *.onnx filter=lfs diff=lfs merge=lfs -text
17
+ *.ot filter=lfs diff=lfs merge=lfs -text
18
+ *.parquet filter=lfs diff=lfs merge=lfs -text
19
+ *.pb filter=lfs diff=lfs merge=lfs -text
20
+ *.pickle filter=lfs diff=lfs merge=lfs -text
21
+ *.pkl filter=lfs diff=lfs merge=lfs -text
22
+ *.pt filter=lfs diff=lfs merge=lfs -text
23
+ *.pth filter=lfs diff=lfs merge=lfs -text
24
+ *.rar filter=lfs diff=lfs merge=lfs -text
25
+ *.safetensors filter=lfs diff=lfs merge=lfs -text
26
+ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
27
+ *.tar.* filter=lfs diff=lfs merge=lfs -text
28
+ *.tar filter=lfs diff=lfs merge=lfs -text
29
+ *.tflite filter=lfs diff=lfs merge=lfs -text
30
+ *.tgz filter=lfs diff=lfs merge=lfs -text
31
+ *.wasm filter=lfs diff=lfs merge=lfs -text
32
+ *.xz filter=lfs diff=lfs merge=lfs -text
33
+ *.zip filter=lfs diff=lfs merge=lfs -text
34
+ *.zst filter=lfs diff=lfs merge=lfs -text
35
+ *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ aifs_diagram.png filter=lfs diff=lfs merge=lfs -text
37
+ decoder_graph.jpeg filter=lfs diff=lfs merge=lfs -text
38
+ encoder_graph.jpeg filter=lfs diff=lfs merge=lfs -text
39
+ assets/aifs-crps-schematic.png filter=lfs diff=lfs merge=lfs -text
README.md ADDED
@@ -0,0 +1,149 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: cc-by-4.0
3
+ pipeline_tag: graph-ml
4
+ language:
5
+ - en
6
+ library_name: anemoi
7
+ ---
8
+
9
+ # AIFS ENS - v1.0
10
+
11
+ <!-- Provide a quick summary of what the model is/does. -->
12
+
13
+ Here, we introduce the Artificial Intelligence Forecasting System (AIFS) ENS,
14
+ an ensemble-based probablistic data driven forecast model developed by the
15
+ [European Centre for Medium-Range Weather Forecasts (ECMWF)](https://www.ecmwf.int/).
16
+
17
+ AIFS produces highly skilled forecasts for upper-air variables, surface weather parameters and
18
+ tropical cyclone tracks. AIFS ENS is run four times daily alongside ECMWF’s physics-based NWP model and forecasts
19
+ are available to the public under ECMWF’s [open data policy](https://www.ecmwf.int/en/forecasts/datasets/open-data).
20
+ Note that due to the non-determinism of GPUs, users will be unable to exactly reproduce an official AIFS ENS forecast
21
+ when running AIFS ENS themselves.
22
+
23
+ For more details please refer to our [documentation](https://confluence.ecmwf.int/display/FCST/Implementation+of+AIFS+ENS+v1).
24
+
25
+ ## Data Details
26
+
27
+ ### Data parameters
28
+
29
+
30
+ **Note**
31
+ Regarding precipitation units, it's worth noting that AIFS model was trained on \\(m^{3}/m^{2}\\) and will therefore produce precip in that units.
32
+ If one wants to retrieve precipitation from Open data, the units will be \\(mm\\).
33
+
34
+
35
+ ## Model Details
36
+
37
+ ### Description
38
+
39
+ <!-- Provide a longer summary of what this model is. -->
40
+
41
+ For full details of the AIFS ENS, please see [Lang et al. 2024](https://arxiv.org/abs/2412.15832).
42
+
43
+ AIFS is based on a graph neural network (GNN) encoder and decoder, and a sliding window transformer processor,
44
+ and is trained on ECMWF’s ERA5 re-analysis and ECMWF’s operational numerical weather prediction (NWP) analyses.
45
+
46
+ <div style="display: flex; justify-content: center;">
47
+ <img src="assets/encoder_graph.jpeg" alt="Encoder graph" style="width: 50%;"/>
48
+ <img src="assets/decoder_graph.jpeg" alt="Decoder graph" style="width: 50%;"/>
49
+ </div>
50
+
51
+ It has a flexible and modular design and supports several levels of parallelism to enable training on
52
+ high resolution input data. AIFS forecast skill is assessed by comparing its forecasts to NWP analyses
53
+ and direct observational data.
54
+
55
+ AIFS ENS was trained to minimise the Continuous Ranked Probability Score (CRPS) calculated over
56
+ a small ensemble.
57
+ <div style="display: flex; justify-content: center;">
58
+ <img src="assets/aifs-crps-schematic.png" alt="Training schematic" style="width: 75%;"/>
59
+ </div>
60
+
61
+ - **Developed by:** ECMWF
62
+ - **Model type:** Encoder-processor-decoder model
63
+ - **License:** These model weights are published under a Creative Commons Attribution 4.0 International (CC BY 4.0).
64
+ To view a copy of this licence, visit https://creativecommons.org/licenses/by/4.0/.
65
+ The notebooks and other script files are published under an Apache 2.0 licence, to view a copy of this license, visit https://www.apache.org/licenses/LICENSE-2.0.txt.
66
+
67
+ ### Resolution
68
+
69
+ | | Component | Horizontal Resolution [kms] | Vertical Resolution [levels] |
70
+ |---|:---:|:---:|:---:|
71
+ | Atmosphere | AIFS ENS v1.0 | ~ 31 | 13 |
72
+
73
+
74
+ ### Sources
75
+
76
+ <!-- Provide the basic links for the model. -->
77
+
78
+ - **Repository:** [Anemoi](https://anemoi.readthedocs.io/en/latest/index.html) is an open-source framework for
79
+ creating machine learning (ML) weather forecasting systems, co-developed by ECMWF and a range of national meteorological
80
+ services across Europe.
81
+ - **Paper:** https://arxiv.org/abs/2412.15832
82
+
83
+ ## Get started
84
+
85
+ 🚨 **Note** `aifs-ens-1.0` requires `38GB` of GPU memory by default, it is possible to reduce this by setting `ANEMOI_INFERENCE_NUM_CHUNKS` to a multiple of 2.
86
+ For example setting it to `16` reduces the memory requirements to `24GB` at the cost of throughput. Therefore issues may arise with lower end GPUs.
87
+
88
+
89
+ To generate a new forecast using AIFS ENS, you can use [anemoi-inference](https://github.com/ecmwf/anemoi-inference). In the [example notebook](run_AIFS_ENS_v1.ipynb), a
90
+ step-by-step workflow is specified to run AIFS ENS using the HuggingFace model:
91
+
92
+ 1. **Install Required Packages and Imports**
93
+ - `torch==2.5.0 anemoi-inference[huggingface]==0.6.0 anemoi-models==0.6.0 anemoi-graphs==0.6.0 anemoi-datasets==0.5.23`
94
+ - `earthkit-regrid==0.4.0 'ecmwf-opendata>=0.3.19'`
95
+ - `flash_attn`
96
+ 2. **Retrieve Initial Conditions from ECMWF Open Data**
97
+ - Select a date
98
+ - Get the data from the [ECMWF Open Data API](https://www.ecmwf.int/en/forecasts/datasets/open-data)
99
+ - Get input fields
100
+ - Add the single levels fields and pressure levels fields
101
+ - Convert geopotential height into geopotential
102
+ - Create the initial state
103
+ 3. **Load the Model and Run the Forecast**
104
+ - Download the Model's Checkpoint from Hugging Face
105
+ - Create a runner
106
+ - Run the forecast using anemoi-inference
107
+ 4. **Inspect the generated forecast**
108
+ - Plot a field
109
+
110
+ 🚨 **Note** we train AIFS using [`flash_attention`](https://github.com/Dao-AILab/flash-attention).
111
+ The use of the 'Flash Attention' package also imposes certain requirements in terms of software and hardware, which can be found under
112
+ [Installation and Features](https://github.com/Dao-AILab/flash-attention?tab=readme-ov-file#installation-and-features) on the GitHub repo.
113
+
114
+ 🚨 **Note** the `aifs_ens_v1.0.ckpt` checkpoint just contains the model’s weights. It does not contain any information about the optimizer states,
115
+ lr-scheduler states, etc.
116
+
117
+ ## Citation
118
+
119
+ <!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
120
+
121
+ If you use this model in your work, please cite it as follows:
122
+
123
+ **BibTeX:**
124
+
125
+ ```
126
+ @misc{lang2024aifscrps,
127
+ title={AIFS-CRPS: Ensemble forecasting using a model trained with a loss function based on the Continuous Ranked Probability Score},
128
+ author={Simon Lang and Mihai Alexe and Mariana C. A. Clare and Christopher Roberts and Rilwan Adewoyin and Zied Ben Bouallègue and Matthew Chantry and Jesper Dramsch and Peter D. Dueben and Sara Hahner and Pedro Maciel and Ana Prieto-Nemesio and Cathal O'Brien and Florian Pinault and Jan Polster and Baudouin Raoult and Steffen Tietsche and Martin Leutbecher},
129
+ year={2024},
130
+ eprint={2412.15832},
131
+ archivePrefix={arXiv},
132
+ primaryClass={physics.ao-ph},
133
+ url={https://arxiv.org/abs/2412.15832},
134
+ }
135
+
136
+ ```
137
+
138
+ **APA:**
139
+
140
+ ```
141
+ Lang, S., Alexe, M., Clare, M. C., Roberts, C., Adewoyin, R., Bouallègue, Z. B., ... & Leutbecher, M. (2024). AIFS-CRPS: Ensemble forecasting using a model trained with a loss function based on the Continuous Ranked Probability Score. arXiv preprint arXiv:2412.15832.
142
+
143
+ ```
144
+
145
+
146
+ ## More Information
147
+
148
+ [Find the paper here](https://arxiv.org/pdf/2412.15832)
149
+
aifs-ens-crps-1.0.ckpt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:5f176e2e018f2c8f54cbb8927f7ee64a948bfdb6960a4867645304e1f6e3134a
3
+ size 921591109
assets/aifs-crps-schematic.png ADDED

Git LFS Details

  • SHA256: e0845756a1aca4f637ef2be762d6148d8ffede60457c8c423b6729f6b6d3e0d4
  • Pointer size: 132 Bytes
  • Size of remote file: 2.84 MB
assets/decoder_graph.jpeg ADDED

Git LFS Details

  • SHA256: a306a3f914ed55b70bb2d8fe89b7070857e963115f1bfa8d58f6d8910b820b3c
  • Pointer size: 131 Bytes
  • Size of remote file: 215 kB
assets/encoder_graph.jpeg ADDED

Git LFS Details

  • SHA256: fe96aa418e231a623dba209aa048248b6506d8c3106df39ac82e3bef2098ba1f
  • Pointer size: 131 Bytes
  • Size of remote file: 168 kB
run_AIFS_ENS_v1.ipynb ADDED
The diff for this file is too large to render. See raw diff