modify yinglong (#937)

* modify yinglong

* refine

Author: zhiminzhang0830

examples/yinglong/README.md

@@ -0,0 +1,73 @@
+# A Study of Data-driven Limited Area Model for Weather Forecasting
+
+Recently, artificial intelligence-based models for forecasting global weather have been rapidly developed. Most of the global models are trained on reanalysis datasets with a spatial resolution of 0.25◦ × 0.25◦. However, study on artificial intelligence-based limited area weather forecasting models is still limited. In this study, an artificial intelligence-based limited area weather forecasting model (YingLong) is developed. YingLong utilizes a parallel structure of global and local blocks to capture multiscale meteorological features. Its predictability on surface temperature, humidity and wind speed is comparable to the predictability of the dynamical limited area model WRF-ARW, but with a much faster running speed. YingLong is also applied to investigate the issues related to the lateral boundary condition of artificial intelligence-based limited area models. The difference between artificial intelligence-based limited area models and dynamical limited area models is also discussed.
+
+This code is the implementation of YingLong. We select the southeastern region of the United States, which is around the range of 80-110W, 30-42N, with 440 × 408 grid points in Lambert projection.
+
+<div align=center>
+    <img src="doc/fig_arch1.jpeg" width="70%" height="auto" >
+</div>
+
+## Installation
+
+### 1. Install PaddlePaddle
+
+Please install the <font color="red"><b>2.6.0</b></font>  or <font color="red"><b>develop</b></font> version of PaddlePaddle according to your environment on the official website of [PaddlePaddle](https://www.paddlepaddle.org.cn/en/install/quick?docurl=/documentation/docs/en/develop/install/pip/linux-pip_en.html).
+
+For example, if your environment is linux and CUDA 11.2, you can install PaddlePaddle by the following command.
+
+``` shell
+python -m pip install paddlepaddle-gpu==2.6.0.post112 -f https://www.paddlepaddle.org.cn/whl/linux/mkl/avx/stable.html
+```
+
+After installation, run the following command to verify if PaddlePaddle has been successfully installed.
+
+``` shell
+python -c "import paddle; paddle.utils.run_check()"
+```
+
+If `"PaddlePaddle is installed successfully! Let's start deep learning with PaddlePaddle now."` appears, to verify that the installation was successful.
+
+### 2. Install PaddleScience
+
+Clone the code of PaddleScience from [here](https://github.com/PaddlePaddle/PaddleScience.git) and install requirements.
+
+``` shell
+git clone -b develop https://github.com/PaddlePaddle/PaddleScience.git
+cd PaddleScience
+pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
+export PYTHONPATH=$PWD
+```
+
+## Example Usage
+
+### 1. Download the data and model weights
+
+``` shell
+cd examples/yinglong
+wget https://paddle-org.bj.bcebos.com/paddlescience/datasets/yinglong/western_valid_data.tar
+tar -xvf western_valid_data.tar
+wget https://paddle-org.bj.bcebos.com/paddlescience/datasets/yinglong/eastern_valid_data.tar
+tar -xvf eastern_valid_data.tar
+wget https://paddle-org.bj.bcebos.com/paddlescience/models/yinglong/inference.tar
+tar -xvf inference.tar
+```
+
+### 2. Run the code
+
+The following code runs the YingLong model, and the model output will be saved in `outputs_yinglong_eastern(western)/result.npy`.
+
+``` shell
+# yinglong_eastern
+python predict.py -cn=yinglong_eastern.yaml
+# yinglong_western
+python predict.py -cn=yinglong_western.yaml
+```
+
+We also visualized the predicted wind speed at 10 meters above ground level, with an initial field of 0:00 on January 1, 2022. Click [eastern](https://paddle-org.bj.bcebos.com/paddlescience/docs/Yinglong/result_eastern.gif)/[western](https://paddle-org.bj.bcebos.com/paddlescience/docs/Yinglong/result_western.gif) to view the prediction results.
+
+## License
+
+YingLong was released by Shanghai Zhangjiang Institute of Mathematics, Baidu inc.
+
+The commercial use of these models is forbidden.

examples/yinglong1/conf/yinglong_12.yaml renamed to examples/yinglong/conf/yinglong_eastern.yaml

@@ -2,7 +2,7 @@ hydra:
   run:
     # dynamic output directory according to running time and override name
     # dir: outputs_yinglong/${now:%Y-%m-%d}/${now:%H-%M-%S}/${hydra.job.override_dirname}
-    dir: ./outputs_yinglong_12
+    dir: ./outputs_yinglong_eastern
   job:
     name: ${mode} # name of logfile
     chdir: false # keep current working directory unchanged
@@ -33,7 +33,7 @@ log_freq: 20
 # inference settings
 INFER:
   pretrained_model_path: null
-  export_path: ./inference/yinglong_12
+  export_path: inference/yinglong_eastern
   pdmodel_path: ${INFER.export_path}.pdmodel
   pdiparams_path: ${INFER.export_path}.pdiparams
   onnx_path: ${INFER.export_path}.onnx
@@ -43,13 +43,14 @@ INFER:
   ir_optim: false
   min_subgraph_size: 30
   gpu_mem: 100
-  gpu_id: 0
+  gpu_id: 3
   max_batch_size: 1
   num_cpu_threads: 10
   batch_size: 1
-  mean_path: ./hrrr_example_24vars/stat/mean_crop.npy
-  std_path: ./hrrr_example_24vars/stat/std_crop.npy
-  input_file: ./hrrr_example_24vars/valid/2022/01/01.h5
+  mean_path: ./eastern_valid_data/stat/mean_crop.npy
+  std_path: ./eastern_valid_data/stat/std_crop.npy
+  input_file: [./eastern_valid_data/valid/2022/01/01.h5, ./eastern_valid_data/valid/2022/01/02.h5, ./eastern_valid_data/valid/2022/01/03.h5]
   init_time: 2022/01/01/00
-  nwp_file: ./hrrr_example_24vars/nwp_convert/2022/01/01/00.h5
-  num_timestamps: 22
+  nwp_file: [./eastern_valid_data/hrrr_nwp_h5/2022/01/01.h5, ./eastern_valid_data/hrrr_nwp_h5/2022/01/02.h5, ./eastern_valid_data/hrrr_nwp_h5/2022/01/03.h5]
+  geo_file: ./eastern_valid_data/geo.h5
+  num_timestamps: 48

examples/yinglong1/conf/yinglong_24.yaml renamed to examples/yinglong/conf/yinglong_western.yaml

@@ -2,7 +2,7 @@ hydra:
   run:
     # dynamic output directory according to running time and override name
     # dir: outputs_yinglong/${now:%Y-%m-%d}/${now:%H-%M-%S}/${hydra.job.override_dirname}
-    dir: ./outputs_yinglong_24
+    dir: ./outputs_yinglong_western
   job:
     name: ${mode} # name of logfile
     chdir: false # keep current working directory unchanged
@@ -33,7 +33,7 @@ log_freq: 20
 # inference settings
 INFER:
   pretrained_model_path: null
-  export_path: ./inference/yinglong_24
+  export_path: inference/yinglong_western
   pdmodel_path: ${INFER.export_path}.pdmodel
   pdiparams_path: ${INFER.export_path}.pdiparams
   onnx_path: ${INFER.export_path}.onnx
@@ -43,13 +43,14 @@ INFER:
   ir_optim: false
   min_subgraph_size: 30
   gpu_mem: 100
-  gpu_id: 0
+  gpu_id: 3
   max_batch_size: 1
   num_cpu_threads: 10
   batch_size: 1
-  mean_path: ./hrrr_example_69vars/stat/mean_crop.npy
-  std_path: ./hrrr_example_69vars/stat/std_crop.npy
-  input_file: ./hrrr_example_69vars/valid/2022/01/01.h5
+  mean_path: ./western_valid_data/stat/mean_crop.npy
+  std_path: ./western_valid_data/stat/std_crop.npy
+  input_file: [./western_valid_data/valid/2022/01/01.h5, ./western_valid_data/valid/2022/01/02.h5, ./western_valid_data/valid/2022/01/03.h5]
   init_time: 2022/01/01/00
-  nwp_file: ./hrrr_example_69vars/nwp_convert/2022/01/01/00.h5
-  num_timestamps: 22
+  nwp_file: [./western_valid_data/hrrr_nwp_h5/2022/01/01.h5, ./western_valid_data/hrrr_nwp_h5/2022/01/02.h5, ./western_valid_data/hrrr_nwp_h5/2022/01/03.h5]
+  geo_file: ./western_valid_data/geo.h5
+  num_timestamps: 48

examples/yinglong/doc/fig_arch1.jpeg

@@ -23,7 +23,7 @@
 import matplotlib
 import numpy as np
 import paddle
-from matplotlib import cm
+from matplotlib import colormaps as cm
 from matplotlib import pyplot as plt
 from matplotlib.legend_handler import HandlerBase
 from matplotlib.patches import Rectangle
@@ -102,7 +102,7 @@ def plot_weather(
         vmin,
         vmax,
         log_norm,
-        cmap=cm.get_cmap("turbo", 1000),
+        cmap=cm.get_cmap("turbo"),
     ):
         ax.title.set_text(title_text)
         ax.set_yticks(yticks)

examples/yinglong1/plot.py renamed to examples/yinglong/plot.py

@@ -20,13 +20,28 @@
 import paddle
 import pandas as pd
 from omegaconf import DictConfig
+from omegaconf import ListConfig
 from packaging import version
 
-from examples.yinglong1.plot import save_plot_weather_from_dict
-from examples.yinglong1.predictor import YingLongPredictor
+from examples.yinglong.plot import save_plot_weather_from_dict
+from examples.yinglong.predictor import YingLongPredictor
 from ppsci.utils import logger
 
 
+def read_h5py(filename):
+    """Read hdf5 file."""
+    if isinstance(filename, str):
+        data = np.array(h5py.File(filename, "r")["fields"])
+    elif isinstance(filename, ListConfig):
+        data = [h5py.File(x, "r")["fields"] for x in filename]
+        data = np.concatenate(data, axis=0)
+    else:
+        raise TypeError(
+            f"The type of INFER.input_file must be str or list, but got {type(filename)}"
+        )
+    return data
+
+
 def inference(cfg: DictConfig):
     # log paddlepaddle's version
     if version.Version(paddle.__version__) != version.Version("0.0.0"):
@@ -50,15 +65,16 @@ def inference(cfg: DictConfig):
     # The variable name is 'z50', 'z500', 'z850', 'z1000', 't50', 't500', 't850', 'z1000',
     # 's50', 's500', 's850', 's1000', 'u50', 'u500', 'u850', 'u1000', 'v50', 'v500',
     # 'v850', 'v1000', 'mslp', 'u10', 'v10', 't2m'.
-    input_file = h5py.File(cfg.INFER.input_file, "r")["fields"]
-    nwp_file = h5py.File(cfg.INFER.nwp_file, "r")["fields"]
+    input_data = read_h5py(cfg.INFER.input_file)
+    nwp_data = read_h5py(cfg.INFER.nwp_file)
+    geo_data = read_h5py(cfg.INFER.geo_file)
 
     # input_data.shape: (1, 24, 440, 408)
-    input_data = input_file[0:1]
+    input_data_0 = input_data[0:1]
     # nwp_data.shape: # (num_timestamps, 24, 440, 408)
-    nwp_data = nwp_file[0:num_timestamps]
+    nwp_data = nwp_data[0:num_timestamps]
     # ground_truth.shape: (num_timestamps, 24, 440, 408)
-    ground_truth = input_file[1 : num_timestamps + 1]
+    ground_truth = input_data[1 : num_timestamps + 1]
 
     # create time stamps
     cur_time = pd.to_datetime(cfg.INFER.init_time, format="%Y/%m/%d/%H")
@@ -68,7 +84,7 @@ def inference(cfg: DictConfig):
         time_stamps.append([cur_time])
 
     # run predictor
-    pred_data = predictor.predict(input_data, time_stamps, nwp_data)
+    pred_data = predictor.predict(input_data_0, time_stamps, nwp_data, geo_data)
     pred_data = pred_data.squeeze(axis=1)  # (num_timestamps, 24, 440, 408)
 
     # save predict data
@@ -103,12 +119,14 @@ def inference(cfg: DictConfig):
         vmin=0,
         vmax=15,
         colorbar_label="m/s",
-        num_timestamps=12,  # only plot 12 timestamps
+        num_timestamps=48,
     )
     logger.info(f"Save plot to {cfg.output_dir}")
 
 
-@hydra.main(version_base=None, config_path="./conf", config_name="yinglong_12.yaml")
+@hydra.main(
+    version_base=None, config_path="./conf", config_name="yinglong_eastern.yaml"
+)
 def main(cfg: DictConfig):
     if cfg.mode == "infer":
         inference(cfg)

examples/yinglong1/predict_12layers.py renamed to examples/yinglong/predict.py

@@ -21,7 +21,7 @@
 from omegaconf import DictConfig
 
 from deploy.python_infer import base
-from examples.yinglong1.timefeatures import time_features
+from examples.yinglong.timefeatures import time_features
 from ppsci.utils import logger
 
 
@@ -57,6 +57,7 @@ def __init__(
         self.input_data_handle = self.predictor.get_input_handle(self.input_names[0])
         self.time_stamps_handle = self.predictor.get_input_handle(self.input_names[1])
         self.nwp_data_handle = self.predictor.get_input_handle(self.input_names[2])
+        self.geo_data_handle = self.predictor.get_input_handle(self.input_names[3])
 
         # get output names and data handles
         self.output_names = self.predictor.get_output_names()
@@ -95,6 +96,7 @@ def predict(
         input_data: np.ndarray,
         time_stamps: List[List[pd.Timestamp]],
         nwp_data: np.ndarray,
+        geo_data: np.ndarray,
         batch_size: int = 1,
     ) -> np.ndarray:
         """Predicts the output of the yinglong model for the given input.
@@ -103,6 +105,7 @@ def predict(
             input_data (np.ndarray): Input data of shape (N, T, H, W).
             time_stamps (List[List[pd.Timestamp]]): Timestamps data.
             nwp_data (np.ndarray): NWP data.
+            geo_data (np.ndarray): Geographic data.
             batch_size (int, optional): Batch size, now only support 1. Defaults to 1.
 
         Returns:
@@ -118,6 +121,7 @@ def predict(
             self.input_names[0]: self.input_data_handle,
             self.input_names[1]: self.time_stamps_handle,
             self.input_names[2]: self.nwp_data_handle,
+            self.input_names[3]: self.geo_data_handle,
         }
         # prepare output handle(s)
         output_handles = {self.output_names[0]: self.output_handle}
@@ -144,6 +148,7 @@ def predict(
                 self.input_names[0]: input_data,
                 self.input_names[1]: time_stamps,
                 self.input_names[2]: nwp_data,
+                self.input_names[3]: geo_data,
             }
 
             # send batch input data to input handle(s)