Update README.md

README.md

@@ -2,55 +2,58 @@
 license: mit
 ---
 
-# Global Data-driven High-resolution Weather Model
+# FourCastNet: a global data-driven high-resolution weather model
 
-This is a global data-driven high-resolution weather model implemented and open sourced by [High-Flyer AI](https://www.high-flyer.cn/). It is the first AI weather model, which can compare with the ECMWF Integrated Forecasting System (IFS).
+This is a global data-driven high-resolution weather model implemented, trained and open sourced by [High-Flyer AI](https://www.high-flyer.cn/en/). It is the first AI weather model, which can compare with the ECMWF Integrated Forecasting System (IFS).
 
-Typhoon track comparison:
+See also: [Github repository](https://github.com/HFAiLab/FourCastNet) and [High-flyer AI's blog](https://www.high-flyer.cn/blog/fourcastnet/)
+
+Typhoon track prediction:
 
 ![](./img/wind_small.gif)
 
-Water vapour comparison:
+Water vapour prediction:
 
 ![](./img/precipitation_small.gif)
 
-## Requirements
+For more cases about FourCastNet prediction, please have a look at [HF-Earth](https://www.high-flyer.cn/hf-earth/), a daily updated demo released by [High-Flyer AI](https://www.high-flyer.cn/en/).
 
-- [hfai](https://doc.hfai.high-flyer.cn/index.html)
-- torch >=1.8
+## 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`.
 
-## Training
-The raw data is from the public dataset, [ERA5](https://www.ecmwf.int/en/forecasts/datasets/reanalysis-datasets/era5) , which is integrated into the dataset warehouse, `hfai.datasets`.
+```python
+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
 
-submit these tasks to Yinghuo HPC:
+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')
 
-1. pretrain `backbone.pt`
+input_x = get_data('2023-01-01 00:00:00')
 
-   ```shell
-    hfai python train/pretrain.py -- -n 8 -p 30
-   ```
+pred_x = backbone_model(input_x)    # input Xt, output Xt+1
+pred_p = precip_model(pred_x)       # input Xt+1, output Pt+1
 
-2. finetune `backbone.pt`
+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')
+```
 
-   ```shell
-    hfai python train/fine_tune.py -- -n 8 -p 30
-   ```
+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](https://arxiv.org/abs/2202.11214).
 
-3. train `precipitation.pt`
 
-   ```shell
-    hfai python train/precipitation.py -- -n 8 -p 30
-   ```
+## Description of Files
 
+`backbone.pt` 
++ the weights of backbone model, 191MB, which is trained on 20 atmospheric variables from `1979-01` to `2022-12`.
 
-## Citation
+`precipitation.pt`
++ the weights of precipitation model, 187MB, which is trained on the variable `total_precipitation` from `1979-01` to `2022-12`.
 
-```bibtex
-@article{pathak2022fourcastnet,
-  title={Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators},
-  author={Pathak, Jaideep and Subramanian, Shashank and Harrington, Peter and Raja, Sanjeev and Chattopadhyay, Ashesh and Mardani, Morteza and Kurth, Thorsten and Hall, David and Li, Zongyi and Azizzadenesheli, Kamyar and others},
-  journal={arXiv preprint arXiv:2202.11214},
-  year={2022}
-}
-```
+`infer2img.py`
++ Case code: load the above two weights to generate images of global weather prediction.