基于FNO求解一维Burgers

概述

计算流体力学是21世纪流体力学领域的重要技术之一，其通过使用数值方法在计算机中对流体力学的控制方程进行求解，从而实现流动的分析、预测和控制。传统的有限元法（finite element method，FEM）和有限差分法（finite difference method，FDM）常用于复杂的仿真流程（物理建模、网格划分、数值离散、迭代求解等）和较高的计算成本，往往效率低下。因此，借助AI提升流体仿真效率是十分必要的。

近年来，随着神经网络的迅猛发展，为科学计算提供了新的范式。经典的神经网络是在有限维度的空间进行映射，只能学习与特定离散化相关的解。与经典神经网络不同，傅里叶神经算子（Fourier Neural Operator，FNO）是一种能够学习无限维函数空间映射的新型深度学习架构。该架构可直接学习从任意函数参数到解的映射，用于解决一类偏微分方程的求解问题，具有更强的泛化能力。更多信息可参考Fourier Neural Operator for Parametric Partial Differential Equations。

本案例教程介绍利用傅里叶神经算子的1-d Burgers方程求解方法。

伯格斯方程（Burgers’ equation）

一维伯格斯方程（1-d Burgers’ equation）是一个非线性偏微分方程，具有广泛应用，包括一维粘性流体流动建模。它的形式如下：

\[\partial_t u(x, t)+\partial_x (u^2(x, t)/2)=\nu \partial_{xx} u(x, t), \quad x \in(0,1), t \in(0, 1]\]

\[u(x, 0)=u_0(x), \quad x \in(0,1)\]

其中\(u\)表示速度场，\(u_0\)表示初始条件，\(\nu\)表示粘度系数。

问题描述

本案例利用Fourier Neural Operator学习初始状态到下一时刻状态的映射，实现一维Burgers’方程的求解：

\[u_0 \mapsto u(\cdot, 1)\]

技术路径

MindFlow求解该问题的具体流程如下：

创建数据集。
构建模型。
优化器与损失函数。
定义求解器。
定义回调函数。
模型训练。

Fourier Neural Operator

Fourier Neural Operator模型构架如下图所示。图中\(w_0(x)\)表示初始涡度，通过Lifting Layer实现输入向量的高维映射，然后将映射结果作为Fourier Layer的输入，进行频域信息的非线性变换，最后由Decoding Layer将变换结果映射至最终的预测结果\(w_1(x)\)。

Lifting Layer、Fourier Layer以及Decoding Layer共同组成了Fourier Neural Operator。

Fourier Neural Operator模型构架

Fourier Layer网络结构如下图所示。图中V表示输入向量，上框表示向量经过傅里叶变换后，经过线性变换R，过滤高频信息，然后进行傅里叶逆变换；另一分支经过线性变换W，最后通过激活函数，得到Fourier Layer输出向量。

Fourier Layer网络结构

[1]:

import os

import numpy as np

from mindspore import context, nn, Tensor, set_seed
from mindspore import DynamicLossScaleManager, LossMonitor, TimeMonitor

下述src包可以在applications/data_driven/burgers/src下载。

[ ]:

from mindflow import FNO1D, RelativeRMSELoss, Solver, load_yaml_config, get_warmup_cosine_annealing_lr

from src import PredictCallback, create_training_dataset

set_seed(0)
np.random.seed(0)

[3]:

context.set_context(mode=context.GRAPH_MODE, device_target='GPU', device_id=4)

[4]:

config = load_yaml_config("burgers1d.yaml")
data_params = config["data"]
model_params = config["model"]
optimizer_params = config["optimizer"]
callback_params = config["callback"]

创建数据集

下载训练与测试数据集: data_driven/burgers/dataset。

本案例根据Zongyi Li在 Fourier Neural Operator for Parametric Partial Differential Equations 一文中对数据集的设置生成训练数据集与测试数据集。具体设置如下：基于周期性边界，生成满足如下分布的初始条件\(u_0(x)\)：

\[u_0 \sim \mu, \mu=\mathcal{N}\left(0,625(-\Delta+25 I)^{-2}\right)\]

本案例选取粘度系数\(\nu=0.1\)，并使用分步法求解方程，其中热方程部分在傅里叶空间中精确求解，然后使用前向欧拉方法求解非线性部分。训练集样本量为1000个，测试集样本量为200个。

[5]:

# create training dataset
train_dataset = create_training_dataset(data_params, shuffle=True)

# create test dataset
test_input, test_label = np.load(os.path.join(data_params["path"], "test/inputs.npy")), \
                         np.load(os.path.join(data_params["path"], "test/label.npy"))

Data preparation finished
input_path:  (1000, 1024, 1)
label_path:  (1000, 1024)

构建模型

网络由1层Lifting layer、1层Decoding layer以及多层Fourier Layer叠加组成：

Lifting layer对应样例代码中FNO1D.fc0，将输出数据\(x\)映射至高维；
多层Fourier Layer的叠加对应样例代码中FNO1D.fno_seq，本案例采用离散傅里叶变换实现时域与频域的转换；
Decoding layer对应代码中FNO1D.fc1与FNO1D.fc2，获得最终的预测值。

基于上述网络结构，进行模型初始化，其中模型参数可在配置文件中修改。

[6]:

model = FNO1D(in_channels=model_params["in_channels"],
              out_channels=model_params["out_channels"],
              resolution=model_params["resolution"],
              modes=model_params["modes"],
              channels=model_params["width"],
              depth=model_params["depth"])

优化器与损失函数

使用相对均方根误差作为网络训练损失函数：

[7]:

steps_per_epoch = train_dataset.get_dataset_size()
lr = get_warmup_cosine_annealing_lr(lr_init=optimizer_params["initial_lr"],
                                    last_epoch=optimizer_params["train_epochs"],
                                    steps_per_epoch=steps_per_epoch,
                                    warmup_epochs=1)
optimizer = nn.Adam(model.trainable_params(), learning_rate=Tensor(lr))
loss_scale = DynamicLossScaleManager()

loss_fn = RelativeRMSELoss()

定义求解器

Solver类是模型训练和推理的接口。输入优化器、网络模型、损失函数、损失缩放策略等，即可定义求解器对象solver。代码中optimizer_params、model_params对应各项参数均在配置文件中修改。

[8]:

solver = Solver(model,
                optimizer=optimizer,
                loss_scale_manager=loss_scale,
                loss_fn=loss_fn,
                )

定义回调函数

[9]:

summary_dir = os.path.join(callback_params["summary_dir"], "FNO1D")
print(summary_dir)
pred_cb = PredictCallback(model=model,
                          inputs=test_input,
                          label=test_label,
                          config=config,
                          summary_dir=summary_dir)

./FNO1D
check test dataset shape: (200, 1024, 1), (200, 1024)

模型训练

调用求解器接口进行模型训练，调用回调接口进行评估。

[10]:

solver.train(epoch=optimizer_params["train_epochs"],
             train_dataset=train_dataset,
             callbacks=[LossMonitor(), TimeMonitor(), pred_cb],
             dataset_sink_mode=True)

epoch: 1 step: 125, loss is 2.377823
Train epoch time: 5782.938 ms, per step time: 46.264 ms
epoch: 2 step: 125, loss is 0.88470775
Train epoch time: 1150.446 ms, per step time: 9.204 ms
epoch: 3 step: 125, loss is 0.98071647
Train epoch time: 1135.464 ms, per step time: 9.084 ms
epoch: 4 step: 125, loss is 0.5404751
Train epoch time: 1114.245 ms, per step time: 8.914 ms
epoch: 5 step: 125, loss is 0.39976493
Train epoch time: 1125.107 ms, per step time: 9.001 ms
epoch: 6 step: 125, loss is 0.508416
Train epoch time: 1127.477 ms, per step time: 9.020 ms
epoch: 7 step: 125, loss is 0.42839915
Train epoch time: 1125.775 ms, per step time: 9.006 ms
epoch: 8 step: 125, loss is 0.28270185
Train epoch time: 1118.428 ms, per step time: 8.947 ms
epoch: 9 step: 125, loss is 0.24137405
Train epoch time: 1121.705 ms, per step time: 8.974 ms
epoch: 10 step: 125, loss is 0.22623646
Train epoch time: 1118.699 ms, per step time: 8.950 ms
================================Start Evaluation================================
mean rms_error: 0.03270653011277318
=================================End Evaluation=================================
...
predict total time: 0.5012176036834717 s
epoch: 91 step: 125, loss is 0.026378194
Train epoch time: 1119.095 ms, per step time: 8.953 ms
epoch: 92 step: 125, loss is 0.057838168
Train epoch time: 1116.712 ms, per step time: 8.934 ms
epoch: 93 step: 125, loss is 0.034773324
Train epoch time: 1107.931 ms, per step time: 8.863 ms
epoch: 94 step: 125, loss is 0.029720988
Train epoch time: 1109.336 ms, per step time: 8.875 ms
epoch: 95 step: 125, loss is 0.02933883
Train epoch time: 1111.804 ms, per step time: 8.894 ms
epoch: 96 step: 125, loss is 0.03140598
Train epoch time: 1116.788 ms, per step time: 8.934 ms
epoch: 97 step: 125, loss is 0.03695058
Train epoch time: 1115.020 ms, per step time: 8.920 ms
epoch: 98 step: 125, loss is 0.039841708
Train epoch time: 1120.316 ms, per step time: 8.963 ms
epoch: 99 step: 125, loss is 0.039001673
Train epoch time: 1134.618 ms, per step time: 9.077 ms
epoch: 100 step: 125, loss is 0.038434036
Train epoch time: 1116.549 ms, per step time: 8.932 ms
================================Start Evaluation================================
mean rms_error: 0.005707952339434996
=================================End Evaluation=================================
predict total time: 0.5055065155029297 s