ENGLISH | 简体中文

Contents

• PINNs for Shallow Water Equations
• Dataset
• Environment Requirements
• Quick Start
• Script Description
  • Script and Sample Code
  • Script Parameters
  • Training Process
  • Evaluation Process

PINNs for Shallow Water Equations

Shallow-Water Equation (SWE) are the canonical test bed for developing noval algorithms for the dynamical core of weather and climate prediction models. Bihlo proposes a Physics-Informed Neural Networks (PINNs) to solve the SWE on the rotating sphere. One drawback of PINNs is that the number of training collocation points increases with the larger problem domain. To overcome the large time domain in the test cases, the author proposes to split the large time domain into several non-overlapping subintervals and solve the SWE in each subinterval consecutively, by training a new neural networks for each interval.

In this repository, the scenario of cosine bell advection around the sphere is illustrated.

paper: Bihlo, Alex, and Roman O. Popovych. Physics-Informed Neural Networks for the Shallow-Water Equations on the Sphere. arXiv.org, February 12, 2022.

Dataset

The dataset used for training is randomly generated by function collocation_points in ./src/process.py. User should give the range of data, defined by the boundaries of days, lambda and theta. By default, the ranges are:

• days: [0, 12]
• lambda: [- pi, pi]
• theta: [- pi / 2, pi / 2]

The size of dataset depends on the number of samples. which are controlled by n_pde and n_iv in config.yaml, the default values are 100000 and 10000, respectively.

The pretrained checkpoint files will be downloaded automatically at the first launch. If you need to download the checkpoint files manually, please visit this link.

Environment Requirements

• Hardware(Ascend/GPU)
  • Prepare hardware environment with Ascend or GPU processor.
• Framework
  • MindSpore
• For more information, please check the resources below:
  • MindSpore Tutorials
  • MindSpore Python API

Quick Start

After installing MindSpore via the official website, you can start training and evaluation as follows:

• running on Ascend or on GPU

Default:

python train.py

Full command:

python train.py \
    --layers 4 20 20 20 20 1 \
    --load_ckpt false \
    --save_ckpt_path ./checkpoints \
    --load_ckpt_path ./checkpoints \
    --save_fig true \
    --figures_path ./figures \
    --log_path ./logs \
    --lr 1e-3 \
    --epochs 18000 \
    --n_pde 100000 \
    --n_iv 10000 \
    --u 1 \
    --h 1000 \
    --days 12 \
    --download_data pinns_swe \
    --force_download false \
    --amp_level O3 \
    --device_id 0 \
    --mode 0

Script Description

Script and Sample Code

├── pinns_swe
│   ├── checkpoints                       # checkpoints files
│   ├── data                              # data files
│   ├── figures                           # plot figures
│   ├── logs                              # log files
│   ├── src                               # source codes
│   │   ├── network.py                    # network architecture
│   │   ├── plot.py                       # plotting results
│   │   ├── process.py                    # data process
│   │   ├── problem.py                    # training process definition
│   │   └── linear_advection_sphere.py    # loss definition according to advection equation
│   ├── config.yaml                       # hyper-parameters configuration
│   ├── README.md                         # English model descriptions
│   ├── README_CN.md                      # Chinese model description
│   ├── train.py                          # python training script
│   └── eval.py                           # python evaluation script

Script Parameters

Important parameters in train.py are as follows:

Training Process

• running on GPU/Ascend

  python train.py

  The loss values during training will be printed in the console, which can also be inspected after training in log file.

  # grep "loss" log
  PDE loss, IC loss in 0th epoch: 0.04082385, 0.13227281, interval 28.69731688, total: 28.69731688
  PDE loss, IC loss in 1th epoch: 0.02216472, 0.05938588, interval 3.24713469, total: 31.94445157
  PDE loss, IC loss in 2th epoch: 0.01156821, 0.02318317, interval 3.31807733, total: 35.26252890
  PDE loss, IC loss in 3th epoch: 0.00694417, 0.00913251, interval 3.22263527, total: 38.48516417
  PDE loss, IC loss in 4th epoch: 0.00577628, 0.00795174, interval 3.32371068, total: 41.80887485
  PDE loss, IC loss in 5th epoch: 0.00556142, 0.01145195, interval 3.30852318, total: 45.11739802
  PDE loss, IC loss in 6th epoch: 0.00492313, 0.01358479, interval 3.31264329, total: 48.43004131
  PDE loss, IC loss in 7th epoch: 0.00375959, 0.01274938, interval 3.32251096, total: 51.75255227
  ...

• After training, you can still review the training process through the log file saved in log_path, ./logs directory by default.

• The model checkpoint will be saved in save_ckpt_path, ./checkpoint directory by default.

Evaluation Process

Before running the command below, please check the checkpoint loading path load_ckpt_path specified in config.yaml for evaluation.

• running on GPU/Ascend

  python eval.py

  You can view the process and results through the log_path, ./logs by default. The result pictures are saved in figures_path, ./figures by default.