基于FNO求解二维Navier-Stokes

下载Notebook下载样例代码查看源文件

概述

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

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

本案例教程介绍利用傅里叶神经算子的纳维-斯托克斯方程(Navier-Stokes equation)求解方法。

纳维-斯托克斯方程(Navier-Stokes equation)

纳维-斯托克斯方程(Navier-Stokes equation)是计算流体力学领域的经典方程,是一组描述流体动量守恒的偏微分方程,简称N-S方程。它在二维不可压缩流动中的涡度形式如下:

\[\partial_t w(x, t)+u(x, t) \cdot \nabla w(x, t)=\nu \Delta w(x, t)+f(x), \quad x \in(0,1)^2, t \in(0, T]\]
\[\nabla \cdot u(x, t)=0, \quad x \in(0,1)^2, t \in[0, T]\]
\[w(x, 0)=w_0(x), \quad x \in(0,1)^2\]

其中\(u\)表示速度场,\(w=\nabla \times u\)表示涡度,\(w_0(x)\)表示初始条件,\(\nu\)表示粘度系数,\(f(x)\)为外力合力项。

问题描述

本案例利用Fourier Neural Operator学习某一个时刻对应涡度到下一时刻涡度的映射,实现二维不可压缩N-S方程的求解:

\[w_t \mapsto w(\cdot, t+1)\]

技术路径

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

  1. 创建数据集。

  2. 构建模型。

  3. 优化器与损失函数。

  4. 定义求解器。

  5. 定义回调函数。

  6. 模型训练。

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网络结构

[11]:
import os
import numpy as np

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

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

[12]:
from mindflow import FNO2D, RelativeRMSELoss, Solver, load_yaml_config, get_warmup_cosine_annealing_lr

from src import PredictCallback, create_training_dataset

set_seed(0)
np.random.seed(0)
[13]:
context.set_context(mode=context.GRAPH_MODE, save_graphs=False, device_target='GPU')
[14]:
config = load_yaml_config('navier_stokes_2d.yaml')
data_params = config["data"]
model_params = config["model"]
optimizer_params = config["optimizer"]
callback_params = config["callback"]

创建数据集

训练与测试数据下载: data_driven/navier_stokes/dataset .

本案例根据Zongyi Li在 Fourier Neural Operator for Parametric Partial Differential Equations 一文中对数据集的设置生成训练数据集与测试数据集。具体设置如下:

基于周期性边界,生成满足如下分布的初始条件\(w_0(x)\)

\[w_0 \sim \mu, \mu=\mathcal{N}\left(0,7^{3 / 2}(-\Delta+49 I)^{-2.5}\right)\]

外力项设置为:

\[f(x)=0.1\left(\sin \left(2 \pi\left(x_1+x_2\right)\right)+\right.\cos(2 \pi(x_1+x_2)))\]

采用Crank-Nicolson方法生成数据,时间步长设置为1e-4,最终数据以每 t = 1 个时间单位记录解。所有数据均在256×256的网格上生成,并被下采样至64×64网格。本案例选取粘度系数\(\nu=1e−5\),训练集样本量为19000个,测试集样本量为3800个。

[15]:
train_dataset = create_training_dataset(data_params, input_resolution=model_params["input_resolution"], shuffle=True)
test_input = np.load(os.path.join(data_params["path"], "test/inputs.npy"))
test_label = np.load(os.path.join(data_params["path"], "test/label.npy"))
Data preparation finished

构建模型

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

  • Lifting layer对应样例代码中FNO2D.fc0,将输出数据\(x\)映射至高维;

  • 多层Fourier Layer的叠加对应样例代码中FNO2D.fno_seq,本案例采用离散傅里叶变换实现时域与频域的转换;

  • Decoding layer对应代码中FNO2D.fc1FNO2D.fc2,获得最终的预测值。

[16]:
model = FNO2D(in_channels=model_params["in_channels"],
              out_channels=model_params["out_channels"],
              resolution=model_params["input_resolution"],
              modes=model_params["modes"],
              channels=model_params["width"],
              depth=model_params["depth"]
              )

优化器与损失函数

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

[17]:
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=optimizer_params["warmup_epochs"])

optimizer = nn.Adam(model.trainable_params(), learning_rate=Tensor(lr))
loss_scale = DynamicLossScaleManager()

# prepare loss function
loss_fn = RelativeRMSELoss()

定义求解器

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

[18]:
solver = Solver(model,
                optimizer=optimizer,
                loss_scale_manager=loss_scale,
                loss_fn=loss_fn,
                )

定义回调函数

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

ckpt_config = CheckpointConfig(save_checkpoint_steps=callback_params["save_checkpoint_steps"] * steps_per_epoch,
                               keep_checkpoint_max=callback_params["keep_checkpoint_max"])
ckpt_dir = os.path.join(summary_dir, "ckpt")
ckpt_cb = ModelCheckpoint(prefix=model_params["name"],
                          directory=ckpt_dir,
                          config=ckpt_config)
./FNO2D
check test dataset shape: (200, 19, 64, 64, 1), (200, 19, 64, 64, 1)

模型训练

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

[20]:
solver.train(epoch=optimizer_params["train_epochs"],
             train_dataset=train_dataset,
             callbacks=[LossMonitor(), TimeMonitor(), pred_cb, ckpt_cb],
             dataset_sink_mode=True)
epoch: 1 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 2.07)
Train epoch time: 36526.785 ms, per step time: 36.527 ms
epoch: 2 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 2.00379)
Train epoch time: 29215.492 ms, per step time: 29.215 ms
epoch: 3 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 1.40253)
Train epoch time: 29217.016 ms, per step time: 29.217 ms
epoch: 4 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 1.79683)
Train epoch time: 29243.756 ms, per step time: 29.244 ms
epoch: 5 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 1.42917)
Train epoch time: 29197.400 ms, per step time: 29.197 ms
epoch: 6 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 1.24265)
Train epoch time: 29199.672 ms, per step time: 29.200 ms
epoch: 7 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 1.48525)
Train epoch time: 29193.341 ms, per step time: 29.193 ms
epoch: 8 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 1.2069)
Train epoch time: 29198.366 ms, per step time: 29.198 ms
epoch: 9 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 1.17752)
Train epoch time: 29210.540 ms, per step time: 29.211 ms
epoch: 10 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 1.25935)
Train epoch time: 29180.896 ms, per step time: 29.181 ms
================================Start Evaluation================================
mean rel_rmse_error: 0.14311510016862303
=================================End Evaluation=================================
predict total time: 16.270995616912842 s
...
epoch: 141 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 0.667819)
Train epoch time: 29181.800 ms, per step time: 29.182 ms
epoch: 142 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 0.610858)
Train epoch time: 29203.687 ms, per step time: 29.204 ms
epoch: 143 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 0.616083)
Train epoch time: 29199.107 ms, per step time: 29.199 ms
epoch: 144 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 0.609115)
Train epoch time: 29302.156 ms, per step time: 29.302 ms
epoch: 145 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 0.518936)
Train epoch time: 29234.649 ms, per step time: 29.235 ms
epoch: 146 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 0.822775)
Train epoch time: 29228.318 ms, per step time: 29.228 ms
epoch: 147 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 0.802282)
Train epoch time: 29231.589 ms, per step time: 29.232 ms
epoch: 148 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 0.669333)
Train epoch time: 29285.277 ms, per step time: 29.285 ms
epoch: 149 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 0.615759)
Train epoch time: 29311.589 ms, per step time: 29.312 ms
epoch: 150 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 0.75713)
Train epoch time: 29280.815 ms, per step time: 29.281 ms
================================Start Evaluation================================
mean rel_rmse_error: 0.06585168887209147
=================================End Evaluation=================================
predict total time: 12.599207639694214 s