Graph Machine Learning
AnemoI
English
jpxkqx commited on
Commit
acf6ad6
1 Parent(s): 1e9affa

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +12 -65
README.md CHANGED
@@ -54,28 +54,20 @@ and direct observational data.
54
 
55
  ## How to Get Started with the Model
56
 
57
- To be able to run AIFS to generate a new forecast, you can use ai-models https://github.com/ecmwf-lab/ai-models.
58
- `ai-models` command can be used to run different models, since in this case we are looking at using AIFS we need to specify
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
- Use the code below to get started with the model.
 
 
 
 
 
 
 
 
 
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