license: mit
FourCastNet: a global data-driven high-resolution weather model
This is a global data-driven high-resolution weather model implemented, trained and open sourced by High-Flyer AI. It is the first AI weather model, which can compare with the ECMWF Integrated Forecasting System (IFS).
See also: Github repository and High-flyer AI's blog
Typhoon track prediction:
Water vapour prediction:
For more cases about FourCastNet prediction, please have a look at HF-Earth, a daily updated demo released by High-Flyer AI.
Inference
You can load the weights backbone.pt
and precipitation.pt
to generate weather predictions, as shown in the following pseudocode. The complete code is released at ./infer2img.py
.
import xarray as xr
import cartopy.crs as ccrs
from afnonet import AFNONet # download the code from https://github.com/HFAiLab/FourCastNet/blob/master/model/afnonet.py
backbone_model = AFNONet(img_size=[720, 1440], in_chans=20, out_chans=20, norm_layer=partial(nn.LayerNorm, eps=1e-6))
backbone_model.load('./backbone.pt')
precip_model = AFNONet(img_size=[720, 1440], in_chans=20, out_chans=1, norm_layer=partial(nn.LayerNorm, eps=1e-6))
precip_model.load('./precipitation.pt')
input_x = get_data('2023-01-01 00:00:00')
pred_x = backbone_model(input_x) # input Xt, output Xt+1
pred_p = precip_model(pred_x) # input Xt+1, output Pt+1
plot_data = xr.Dataset([pred_x, pred_p])
ax = plt.axes(projection=ccrs.PlateCarree())
plot_data.plot(ax=ax, transform=ccrs.PlateCarree(), add_colorbar=False, add_labels=False, rasterized=True)
ax.coastlines(resolution='110m')
plt.savefig('img.png')
FourCastNet can predict 7 surface variables, plus 5 atmospheric variables at each of 3 or 4 pressure levels, for 21 variables total. The details of these variables follow the paper.
Description of Files
backbone.pt
- the weights of backbone model, 191MB, which is trained on 20 atmospheric variables from
1979-01
to2022-12
.
precipitation.pt
- the weights of precipitation model, 187MB, which is trained on the variable
total_precipitation
from1979-01
to2022-12
.
infer2img.py
- Case code: load the above two weights to generate images of global weather prediction.