OpenCastKit: an open-source solutions of global data-driven high-resolution weather forecasting

This is an open-source solutions of global data-driven high-resolution weather forecasting, implemented and improved by High-Flyer AI. It can compare with the ECMWF Integrated Forecasting System (IFS).

Several cases:

For more cases about FourCastNet/GraphCast prediction, please have a look at HF-Earth, a daily updated demo released by High-Flyer AI.

Inference

FourCastNet

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/OpenCastKit/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('./weights/fourcastnet/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('./weights/fourcastnet/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.