Update README.md
Browse files
README.md
CHANGED
@@ -54,28 +54,20 @@ and direct observational data.
|
|
54 |
|
55 |
## How to Get Started with the Model
|
56 |
|
57 |
-
To
|
58 |
-
|
59 |
-
`anemoi` as `model-name` and then pass the path to the checkpoint (`aifs_single_v0.2.1.ckpt`) and the initial conditions.
|
60 |
-
You can find an example of a set of initial conditions in the GRIB file `example_20241107_12_n320.grib`.
|
61 |
|
62 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
63 |
|
64 |
-
```
|
65 |
-
# 1st - create the conda environment
|
66 |
-
export CONDA_ENV=aifs-env
|
67 |
-
conda create -n ${CONDA_ENV} python=3.10
|
68 |
-
conda activate ${CONDA_ENV}
|
69 |
-
|
70 |
-
pip install torch=2.4
|
71 |
-
pip install anemoi-inference[plugin] anemoi-models==0.2
|
72 |
-
pip install ninja
|
73 |
-
pip install flash-attn --no-build-isolation
|
74 |
-
|
75 |
-
# 2nd - Run ai-models to generate weather forecast
|
76 |
-
|
77 |
-
ai-models anemoi --checkpoint aifs_single_v0.2.1.ckpt --file example_20241107_12_n320.grib
|
78 |
-
```
|
79 |
|
80 |
🚨 **Note** we train AIFS using `flash_attention` (https://github.com/Dao-AILab/flash-attention).
|
81 |
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).
|
@@ -85,51 +77,6 @@ Additonally the use of 'Flash Attention' package also imposes certain requiremen
|
|
85 |
🚨 **Note** the `aifs_single_v0.2.1.ckpt` checkpoint just contains the model’s weights.
|
86 |
That file does not contain any information about the optimizer states, lr-scheduler states, etc.
|
87 |
|
88 |
-
**Note** By default, when running `ai-models` the model will be run for a 10-day lead time (240 hours).
|
89 |
-
It is possible to extend or modify the lead time to for example 15 days by doing `ai-models --lead-time 360`
|
90 |
-
Please refer to `ai-models` documentation for more information regarding defaults and available command line options.
|
91 |
-
|
92 |
-
|
93 |
-
After running the `ai-models` command the output of the forecast should be written into `anemoi.grib`
|
94 |
-
Below you can find an example to read that file and load it as numpy array or xarray.
|
95 |
-
|
96 |
-
```
|
97 |
-
|
98 |
-
import earthkit.data as ekd
|
99 |
-
import cartopy.crs as ccrs
|
100 |
-
import numpy as np
|
101 |
-
from plotting_utils import map_scatter
|
102 |
-
|
103 |
-
source_filename='anemoi.grib'
|
104 |
-
aifs_forecast = ekd.from_source('file',source_filename)
|
105 |
-
|
106 |
-
# to load a variable as a numpy array
|
107 |
-
aifs_forecast.sel(param="2t").to_numpy()
|
108 |
-
|
109 |
-
# to load a variables as an xarray array
|
110 |
-
aifs_forecast.sel(param="2t").to_xarray()
|
111 |
-
|
112 |
-
aifs_forecast_2t_xr = aifs_forecast.sel(param="2t").to_xarray()
|
113 |
-
|
114 |
-
fig, axes = plt.subplots(subplot_kw={'projection': ccrs.Robinson()}, figsize=(16,6))
|
115 |
-
|
116 |
-
map_scatter(data_v021_2t_xr.latitude, data_v021_2t_xr.longitude, data_v021_2t_xr.t2m[0,1,0,:],
|
117 |
-
axes=axes,
|
118 |
-
proj=ccrs.Robinson(),
|
119 |
-
vmin=np.percentile(data_v021_2t_xr.t2m[0,1,0,:], 2),
|
120 |
-
vmax=np.percentile(data_v021_2t_xr.t2m[0,1,0,:], 98),
|
121 |
-
cmap="RdYlBu_r",
|
122 |
-
mask_ocean=False,
|
123 |
-
plot_title='2-Meter Temperature Forecast',
|
124 |
-
cmap_label="Temperature (K)"
|
125 |
-
)
|
126 |
-
|
127 |
-
```
|
128 |
-
<div style="display: flex; justify-content: center;">
|
129 |
-
<img src="2t_aifs_v021.png" alt="AIFS 2m Temperature" style="width: 50%;"/>
|
130 |
-
</div>
|
131 |
-
|
132 |
-
|
133 |
|
134 |
## Training Details
|
135 |
|
|
|
54 |
|
55 |
## How to Get Started with the Model
|
56 |
|
57 |
+
To generate a new forecast using AIFS, you can use [anemoi-inference](https://github.com/ecmwf/anemoi-inference). In the [following notebook](run_AIFS_v0_2_1.ipynb), a
|
58 |
+
step-by-step workflow is specified to run the AIFS using the HuggingFace model:
|
|
|
|
|
59 |
|
60 |
+
- Install the required packages
|
61 |
+
- Select a date
|
62 |
+
- Get the data from the [ECMWF Open Data API](https://www.ecmwf.int/en/forecasts/datasets/open-data)
|
63 |
+
- Get input fields
|
64 |
+
- Add the single levels fields and pressure levels fields
|
65 |
+
- Convert geopotential height into greopotential
|
66 |
+
- Create the initial state
|
67 |
+
- Create a runner
|
68 |
+
- Run the forecast
|
69 |
+
- Plot a field
|
70 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
71 |
|
72 |
🚨 **Note** we train AIFS using `flash_attention` (https://github.com/Dao-AILab/flash-attention).
|
73 |
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).
|
|
|
77 |
🚨 **Note** the `aifs_single_v0.2.1.ckpt` checkpoint just contains the model’s weights.
|
78 |
That file does not contain any information about the optimizer states, lr-scheduler states, etc.
|
79 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
80 |
|
81 |
## Training Details
|
82 |
|