Update README.md

README.md

@@ -27,6 +27,10 @@ are available to the public under ECMWF’s open data policy. (https://www.ecmwf
 AIFS is based on a graph neural network (GNN) encoder and decoder, and a sliding window transformer processor, 
 and is trained on ECMWF’s ERA5 re-analysis and ECMWF’s operational numerical weather prediction (NWP) analyses. 
 
+<div style="display: flex; justify-content: center;">
+  <img src="aifs_diagram.png" alt="High-level AIFS diagram" style="width: 50%;"/>
+</div>
+
 <div style="display: flex;">
   <img src="encoder_graph.jpeg" alt="Encoder graph" style="width: 50%;"/>
   <img src="decoder_graph.jpeg" alt="Decoder graph" style="width: 50%;"/>
@@ -70,14 +74,32 @@ pip install anemoi-inference[plugin] anemoi-models==0.2
 pip install ninja
 pip install flash-attn --no-build-isolation
 
-# Run ai-models 
+# 2nd Run ai-models to generate weather forecast
 
 ai-models anemoi --checkpoint aifs_single_v0.2.1.ckpt --file example_20241107_12_n320.grib
 ```
-! ISSUE WITH PYTORCH 2.5 and flash_attention, and cuda 12.4. For now keep it to torch 2.4
 
-The above command will write the forecast results to anemoi.grib
-#missing - example how to plot/open the anemoi.grib
+**Note** we train AIFS using `flash_attention` (https://github.com/Dao-AILab/flash-attention)
+There are currently some issues when trying to install flash attention with the latest PyTorch version 2.5 and CUDA 12.4 (https://github.com/Dao-AILab/flash-attention/issues/1330)
+For that reason we recommen you install PyTorch 2.4
+
+After running the `ai-models` command the output of the forecast should be written into `anemoi.grib` 
+Below you can find an example to read that file and load it as numpy array or xarray. 
+
+```
+
+import earthkit.data as ekd
+
+source_filename='anemoi.grib'
+aifs_forecast = ekd.from_source('file',source_filename)
+
+# to load it as a numpy array
+aifs_forecast.to_numpy()
+
+# to load it as a xarray array
+aifs_forecast.to_xarray()
+
+```
 
 ## Training Details