跳转至

2D-LDC(2D Lid Driven Cavity Flow)

AI Studio快速体验

# linux
wget -nc -P ./data/ \
    https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re100.mat \
    https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re400.mat \
    https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re1000.mat \
    https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re3200.mat
# windows
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re100.mat --create-dirs -o ./data/ldc_Re100.mat
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re400.mat --create-dirs -o ./data/ldc_Re400.mat
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re1000.mat --create-dirs -o ./data/ldc_Re1000.mat
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re3200.mat --create-dirs -o ./data/ldc_Re3200.mat
python ldc_2d_Re3200_sota.py

# linux
wget -nc -P ./data/ https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re1000.mat
# windows
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re1000.mat --create-dirs -o ./data/ldc_Re1000.mat
python ldc_2d_Re3200_sota.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/ldc/ldc_re1000_sota_pretrained.pdparams

python ldc_2d_Re3200_sota.py mode=export

# linux
wget -nc -P ./data/ https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re1000.mat
# windows
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re1000.mat --create-dirs -o ./data/ldc_Re1000.mat
python ldc_2d_Re3200_sota.py mode=infer
预训练模型 \(Re\) 指标
- 100 U_validator/loss: 0.00017
U_validator/L2Rel.U: 0.04875
- 400 U_validator/loss: 0.00047
U_validator/L2Rel.U: 0.07554
ldc_re1000_sota_pretrained.pdparams 1000 U_validator/loss: 0.00053
U_validator/L2Rel.U: 0.07777
- 3200 U_validator/loss: 0.00227
U_validator/L2Rel.U: 0.15440		 
# linux
wget -nc -P ./data/ \
    https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re100.mat \
    https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re400.mat \
    https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re1000.mat \
    https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re1600.mat \
    https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re3200.mat
# windows
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re100.mat --create-dirs -o ./data/ldc_Re100.mat
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re400.mat --create-dirs -o ./data/ldc_Re400.mat
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re1000.mat --create-dirs -o ./data/ldc_Re1000.mat
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re1600.mat --create-dirs -o ./data/ldc_Re1600.mat
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re3200.mat --create-dirs -o ./data/ldc_Re3200.mat
python ldc_2d_Re3200_piratenet.py

# linux
wget -nc -P ./data/ https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re3200.mat
# windows
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re3200.mat --create-dirs -o ./data/ldc_Re3200.mat
python ldc_2d_Re3200_piratenet.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/ldc/ldc_re3200_piratenet_pretrained.pdparams

python ldc_2d_Re3200_piratenet.py mode=export

# linux
wget -nc -P ./data/ https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re3200.mat
# windows
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/ldc/ldc_Re3200.mat --create-dirs -o ./data/ldc_Re3200.mat
python ldc_2d_Re3200_piratenet.py mode=infer
预训练模型 \(Re\) 指标
- 100 U_validator/loss: 0.00016
U_validator/L2Rel.U: 0.04741
- 400 U_validator/loss: 0.00071
U_validator/L2Rel.U: 0.09288
- 1000 U_validator/loss: 0.00191
U_validator/L2Rel.U: 0.14797
- 1600 U_validator/loss: 0.00276
U_validator/L2Rel.U: 0.17360
ldc_re3200_piratenet_pretrained.pdparams 3200 U_validator/loss: 0.00016
U_validator/L2Rel.U: 0.04166

说明

本案例仅提供 \(Re=1000/3200\) 两种情况下的预训练模型，若需要其他雷诺数下的预训练模型，请执行训练命令手动训练即可得到各雷诺数下的模型权重。

1. 背景简介

顶盖方腔驱动流LDC问题在许多领域中都有应用。例如，这个问题可以用于计算流体力学（CFD）领域中验证计算方法的有效性。虽然这个问题的边界条件相对简单，但是其流动特性却非常复杂。在顶盖驱动流LDC中，顶壁朝x方向以U=1的速度移动，而其他三个壁则被定义为无滑移边界条件，即速度为零。

此外，顶盖方腔驱动流LDC问题也被用于研究和预测空气动力学中的流动现象。例如，在汽车工业中，通过模拟和分析车体内部的空气流动，可以帮助优化车辆的设计和性能。

总的来说，顶盖方腔驱动流LDC问题在计算流体力学、空气动力学以及相关领域中都有广泛的应用，对于研究和预测流动现象、优化产品设计等方面都起到了重要的作用。

2. 问题定义

本案例假设 \(Re=3200\)，计算域为一个长宽均为 1 的方腔，应用以下公式进行顶盖驱动方腔流研究稳态流场问题：

质量守恒：

\[ \dfrac{\partial u}{\partial x} + \dfrac{\partial v}{\partial y} = 0 \]

\(x\) 动量守恒：

\[ u\dfrac{\partial u}{\partial x} + v\dfrac{\partial u}{\partial y} = -\dfrac{1}{\rho}\dfrac{\partial p}{\partial x} + \nu(\dfrac{\partial ^2 u}{\partial x ^2} + \dfrac{\partial ^2 u}{\partial y ^2}) \]

\(y\) 动量守恒：

\[ u\dfrac{\partial v}{\partial x} + v\dfrac{\partial v}{\partial y} = -\dfrac{1}{\rho}\dfrac{\partial p}{\partial y} + \nu(\dfrac{\partial ^2 v}{\partial x ^2} + \dfrac{\partial ^2 v}{\partial y ^2}) \]

令：

\(t^* = \dfrac{L}{U_0}\)

\(x^*=y^* = L\)

\(u^*=v^* = U_0\)

\(p^* = \rho {U_0}^2\)

定义：

无量纲坐标 \(x：X = \dfrac{x}{x^*}\)；无量纲坐标 \(y：Y = \dfrac{y}{y^*}\)

无量纲速度 \(x：U = \dfrac{u}{u^*}\)；无量纲速度 \(y：V = \dfrac{v}{u^*}\)

无量纲压力 \(P = \dfrac{p}{p^*}\)

雷诺数 \(Re = \dfrac{L U_0}{\nu}\)

则可获得如下无量纲Navier-Stokes方程，施加于方腔内部：

质量守恒：

\[ \dfrac{\partial U}{\partial X} + \dfrac{\partial U}{\partial Y} = 0 \]

\(x\) 动量守恒：

\[ U\dfrac{\partial U}{\partial X} + V\dfrac{\partial U}{\partial Y} = -\dfrac{\partial P}{\partial X} + \dfrac{1}{Re}(\dfrac{\partial ^2 U}{\partial X^2} + \dfrac{\partial ^2 U}{\partial Y^2}) \]

\(y\) 动量守恒：

\[ U\dfrac{\partial V}{\partial X} + V\dfrac{\partial V}{\partial Y} = -\dfrac{\partial P}{\partial Y} + \dfrac{1}{Re}(\dfrac{\partial ^2 V}{\partial X^2} + \dfrac{\partial ^2 V}{\partial Y^2}) \]

对于方腔边界，则需施加 Dirichlet 边界条件：

上边界：

\[ u(x, y) = 1 − \dfrac{\cosh (C_0(x − 0.5))} {\cosh (0.5C_0)} , \]

左边界、下边界、右边界：

\[ u=0, v=0 \]

3. 问题求解

接下来开始讲解如何将问题一步一步地转化为 PaddleScience 代码，用深度学习的方法求解该问题。 为了快速理解 PaddleScience，接下来仅对模型构建、方程构建、计算域构建等关键步骤进行阐述，而其余细节请参考 API文档。

3.1 模型构建

在 2D-LDC 问题中，每一个已知的坐标点 \((x, y)\) 都有自身的横向速度 \(u\)、纵向速度 \(v\)、压力 \(p\) 三个待求解的未知量，我们在这里使用适合于 PINN 任务的 PirateNet 来表示 \((x, y)\) 到 \((u, v, p)\) 的映射函数 \(f: \mathbb{R}^2 \to \mathbb{R}^3\) ，即：

\[ u, v, p = f(x, y) \]

上式中 \(f\) 即为 PirateNet 模型本身，用 PaddleScience 代码表示如下

41
42
# set model
model = ppsci.arch.PirateNet(**cfg.MODEL)

其中 cfg.MODEL 配置如下所示：

38
39
40
41
# model settings
MODEL:
  input_keys: ["x", "y"]
  output_keys: ["u", "v", "p"]

为了在计算时，准确快速地访问具体变量的值，我们在这里指定网络模型的输入变量名是 ["x", "y"]，输出变量名是 ["u", "v", "p"]，这些命名与后续代码保持一致。

如上所示，通过指定 PirateNet 的层数、神经元个数以及激活函数，我们就实例化出了一个拥有 12 层隐藏神经元，每层神经元数为 256，使用 "tanh" 作为激活函数的神经网络模型 model。

3.2 Curriculum Learning

为了加快收敛速度，我们使用 Curriculum learning 的方法来训练模型，即先训练模型在低雷诺数下，然后逐步增加雷诺数，最终达到高雷诺数下的收敛。

210
211
for idx in range(len(cfg.Re)):
    train_curriculum(cfg, idx)

3.3 方程构建

由于 2D-LDC 使用的是 Navier-Stokes 方程的2维稳态形式，因此可以直接使用 PaddleScience 内置的 NavierStokes。

88
89
90
91
# set equation
equation = {
    "NavierStokes": ppsci.equation.NavierStokes(1 / Re, 1, dim=2, time=False)
}

在课程学习的函数中，我们在实例化 NavierStokes 类时需指定必要的参数：动力粘度 \(\nu=\frac{1}{Re}\), 流体密度 \(\rho=1.0\)，其中 \(Re\) 是一个随着训练过程中会逐步增大的变量。

3.4 计算域构建

本文中 2D-LDC 问题训练、评估所需的数据，通过读取对应雷诺数的文件得到。

 93
 94
 95
 96
 97
 98
 99
100
101
102
103
# load data
data = sio.loadmat(f"./data/ldc_Re{Re}.mat")
u_ref = data["u"].astype(dtype)
v_ref = data["v"].astype(dtype)
U_ref = np.sqrt(u_ref**2 + v_ref**2).reshape(-1, 1)
x_star = data["x"].flatten().astype(dtype)
y_star = data["y"].flatten().astype(dtype)
x0 = x_star[0]
x1 = x_star[-1]
y0 = y_star[0]
y1 = y_star[-1]

3.5 约束构建

根据 2. 问题定义 得到的无量纲公式和和边界条件，对应了在计算域中指导模型训练的两个约束条件，即：

  1. 施加在矩形内部点上的无量纲 Navier-Stokes 方程约束（经过简单移项）

    \[ \dfrac{\partial U}{\partial X} + \dfrac{\partial U}{\partial Y} = 0 \]
    \[ U\dfrac{\partial U}{\partial X} + V\dfrac{\partial U}{\partial Y} + \dfrac{\partial P}{\partial X} - \dfrac{1}{Re}(\dfrac{\partial ^2 U}{\partial X^2} + \dfrac{\partial ^2 U}{\partial Y^2}) = 0 \]
    \[ U\dfrac{\partial V}{\partial X} + V\dfrac{\partial V}{\partial Y} + \dfrac{\partial P}{\partial Y} - \dfrac{1}{Re}(\dfrac{\partial ^2 V}{\partial X^2} + \dfrac{\partial ^2 V}{\partial Y^2}) = 0 \]

    为了方便获取中间变量，NavierStokes 类内部将上式左侧的结果分别命名为 continuity, momentum_x, momentum_y。

  2. 施加在矩形上、下、左、右边界上的 Dirichlet 边界条件约束

    上边界：

    \[ u(x, y) = 1 − \dfrac{\cosh (C_0(x − 0.5))} {\cosh (0.5C_0)} , \]

    左边界、下边界、右边界：

    \[ u=0, v=0 \]

接下来使用 PaddleScience 内置的 SupervisedConstraint 构建上述两种约束条件。

3.5.1 内部点约束

以作用在矩形内部点上的 SupervisedConstraint 为例，代码如下：

105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
# set N-S pde constraint
def gen_input_batch():
    tx = np.random.uniform(
        [x0, y0],
        [x1, y1],
        (cfg_t.TRAIN.batch_size.pde, 2),
    ).astype(dtype)
    return {"x": tx[:, 0:1], "y": tx[:, 1:2]}

def gen_label_batch(input_batch):
    return {
        "continuity": np.zeros([cfg_t.TRAIN.batch_size.pde, 1], dtype),
        "momentum_x": np.zeros([cfg_t.TRAIN.batch_size.pde, 1], dtype),
        "momentum_y": np.zeros([cfg_t.TRAIN.batch_size.pde, 1], dtype),
    }

pde_constraint = ppsci.constraint.SupervisedConstraint(
    {
        "dataset": {
            "name": "ContinuousNamedArrayDataset",
            "input": gen_input_batch,
            "label": gen_label_batch,
        },
    },
    output_expr=equation["NavierStokes"].equations,
    loss=ppsci.loss.MSELoss("mean"),
    name="PDE",
)

SupervisedConstraint 的第一个参数是数据集配置，用于描述如何构建输入数据，此处填入输入数据和标签数据的构造函数函数 gen_input_batch 和 gen_label_batch，以及数据集的名称 ContinuousNamedArrayDataset；

第二个参数是约束变量的目标值，此处填入在 3.3 方程构建 章节中实例化好的 equation["NavierStokes"].equations；

第三个参数是损失函数，此处我们选用常用的 MSE 函数，且 reduction 设置为 "mean"，即我们会将参与计算的所有数据点产生的损失项平均；

第四个参数是约束条件的名字，我们需要给每一个约束条件命名，方便后续对其索引。此处我们命名为 "PDE" 即可。

3.5.2 边界约束

将上边界的标签数据按照上述对应公式进行处理，其余点的标签数据设置为 0。然后继续构建方腔边界的 Dirichlet 约束，我们仍然使用 SupervisedConstraint 类。

134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
# set boundary conditions
x_bc = sample_points_on_square_boundary(
    cfg_t.TRAIN.batch_size.bc, eps=0.0
).astype(
    dtype
)  # avoid singularity a right corner for u velocity
v_bc = np.zeros((cfg_t.TRAIN.batch_size.bc * 4, 1), dtype)
u_bc = copy.deepcopy(v_bc)
lid_bc_fn = lambda x: 1 - np.cosh(50 * (x - 0.5)) / np.cosh(50 * 0.5)
u_bc[: cfg_t.TRAIN.batch_size.bc] = lid_bc_fn(
    x_bc[: cfg_t.TRAIN.batch_size.bc, 0:1]
)
bc = ppsci.constraint.SupervisedConstraint(
    {
        "dataset": {
            "name": "IterableNamedArrayDataset",
            "input": {
                "x": x_bc[:, 0:1],
                "y": x_bc[:, 1:2],
            },
            "label": {"u": u_bc, "v": v_bc},
        },
    },
    output_expr={"u": lambda out: out["u"], "v": lambda out: out["v"]},
    loss=ppsci.loss.MSELoss("mean"),
    name="BC",
)

在微分方程约束、边界约束、初值约束构建完毕之后，以我们刚才的命名为关键字，封装到一个字典中，方便后续访问。

161
162
163
164
165
# wrap constraints together
constraint = {
    pde_constraint.name: pde_constraint,
    bc.name: bc,
}

3.6 超参数设定

接下来需要在配置文件中指定训练轮数，分别在 Re=100, 400, 1000, 1600, 3200 上训练 10, 20, 50, 50, 500 轮，每轮迭代次数为 1000，

33
34
35
# working conditions
Re: [100, 400, 1000, 1600, 3200]
epochs: [10, 20, 50, 50, 500]

其次，设置合适的学习率衰减策略，

52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
# training settings
TRAIN:
  epochs: 10
  iters_per_epoch: 1000
  save_freq: 100
  eval_during_train: true
  eval_freq: 1
  lr_scheduler:
    epochs: ${sum:${epochs}}
    iters_per_epoch: ${TRAIN.iters_per_epoch}
    learning_rate: 1.0e-3
    gamma: 0.9
    decay_steps: 10000
    warmup_epoch: 5
    by_epoch: false

最后，设置训练过程中损失自动平衡策略为 GradNorm，

72
73
74
75
grad_norm:
  update_freq: 1000
  momentum: 0.9
  init_weights: [10, 1, 1, 100, 100]

3.7 优化器构建

训练过程会调用优化器来更新模型参数，此处选择较为常用的 Adam 优化器。

44
45
46
47
48
# set optimizer
lr_scheduler = ppsci.optimizer.lr_scheduler.ExponentialDecay(
    **cfg.TRAIN.lr_scheduler
)()
optimizer = ppsci.optimizer.Adam(lr_scheduler)(model)

3.8 评估器构建

在训练过程中通常会按一定轮数间隔，用验证集（测试集）评估当前模型的训练情况，因此使用 ppsci.validate.SupervisedValidator 构建评估器。

167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
# set validator
xy_star = misc.cartesian_product(x_star, y_star).astype(dtype)
eval_data = {"x": xy_star[:, 0:1], "y": xy_star[:, 1:2]}
eval_label = {"U": U_ref.reshape([-1, 1])}
U_validator = ppsci.validate.SupervisedValidator(
    {
        "dataset": {
            "name": "NamedArrayDataset",
            "input": eval_data,
            "label": eval_label,
        },
        "batch_size": cfg_t.EVAL.batch_size,
    },
    ppsci.loss.MSELoss("mean"),
    {"U": lambda out: (out["u"] ** 2 + out["v"] ** 2).sqrt()},
    metric={"L2Rel": ppsci.metric.L2Rel()},
    name="U_validator",
)
validator = {U_validator.name: U_validator}

此处计算 \(U=\sqrt{u^2+v^2}\) 的预测误差；

评价指标 metric 选择 ppsci.metric.L2Rel 即可；

其余配置与 约束构建 的设置类似。

3.9 模型训练、评估与可视化

完成上述设置之后，只需要将上述实例化的对象按顺序传递给 ppsci.solver.Solver，然后启动训练、评估

187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
# initialize solver
solver = ppsci.solver.Solver(
    model,
    constraint,
    optimizer=optimizer,
    equation=equation,
    validator=validator,
    loss_aggregator=grad_norm,
    cfg=cfg_t,
)
# train model
solver.train()
# evaluate after finished training
solver.eval()
# visualize prediction after finished training
pred_dict = solver.predict(
    eval_data, batch_size=cfg_t.EVAL.batch_size, return_numpy=True
)
U_pred = np.sqrt(pred_dict["u"] ** 2 + pred_dict["v"] ** 2).reshape(
    [len(x_star), len(y_star)]
)
plot(U_pred, cfg_t.output_dir)

4. 完整代码

ldc_2d_Re3200_piratenet.py
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
"""
Reference: https://github.com/PredictiveIntelligenceLab/jaxpi/tree/main/examples/ldc
"""

from __future__ import annotations

import copy
import os
from os import path as osp

import hydra
import numpy as np
import paddle
import scipy.io as sio
from matplotlib import pyplot as plt
from omegaconf import DictConfig

import ppsci
from ppsci.loss import mtl
from ppsci.utils import misc

dtype = paddle.get_default_dtype()


def plot(U_pred: np.ndarray, output_dir: str):
    os.makedirs(output_dir, exist_ok=True)
    fig_path = osp.join(output_dir, "ac.png")

    fig = plt.figure()
    plt.pcolor(U_pred.T, cmap="jet")
    plt.xlabel("x")
    plt.ylabel("y")
    plt.colorbar()
    plt.title(r"Prediction of $U=\sqrt{{u^2+v^2}}$")
    fig.savefig(fig_path, bbox_inches="tight")
    ppsci.utils.logger.info(f"Saving figure to {fig_path}")
    plt.close()


def train(cfg: DictConfig):
    # set model
    model = ppsci.arch.PirateNet(**cfg.MODEL)

    # set optimizer
    lr_scheduler = ppsci.optimizer.lr_scheduler.ExponentialDecay(
        **cfg.TRAIN.lr_scheduler
    )()
    optimizer = ppsci.optimizer.Adam(lr_scheduler)(model)
    grad_norm = mtl.GradNorm(
        model,
        5,
        update_freq=cfg.TRAIN.grad_norm.update_freq,
        momentum=cfg.TRAIN.grad_norm.momentum,
        init_weights=list(cfg.TRAIN.grad_norm.init_weights),
    )

    def sample_points_on_square_boundary(num_pts_per_side, eps):
        # Sample points along the top side (x=1 to x=0, y=1)
        top_coords = np.linspace(0, 1, num_pts_per_side)
        top = np.column_stack((top_coords, np.ones_like(top_coords)))

        # Sample points along the bottom side (x=0 to x=1, y=0)
        bottom_coords = np.linspace(0, 1, num_pts_per_side)
        bottom = np.column_stack((bottom_coords, np.zeros_like(bottom_coords)))

        # Sample points along the left side (x=0, y=1 to y=0)
        left_coords = np.linspace(0, 1 - eps, num_pts_per_side)
        left = np.column_stack((np.zeros_like(left_coords), left_coords))

        # Sample points along the right side (x=1, y=0 to y=1)
        right_coords = np.linspace(0, 1 - eps, num_pts_per_side)
        right = np.column_stack((np.ones_like(right_coords), right_coords))

        # Combine the points from all sides
        points = np.vstack((top, bottom, left, right))

        return points

    def train_curriculum(cfg, idx):
        cfg_t = copy.deepcopy(cfg)
        Re = cfg_t.Re[idx]
        cfg_t.output_dir = osp.join(cfg_t.output_dir, f"Re_{int(Re)}")
        cfg_t.TRAIN.epochs = cfg_t.epochs[idx]
        ppsci.utils.logger.message(
            f"Training curriculum {idx + 1}/{len(cfg_t.epochs)} Re={Re:.5g} epochs={cfg_t.epochs[idx]}"
        )

        # set equation
        equation = {
            "NavierStokes": ppsci.equation.NavierStokes(1 / Re, 1, dim=2, time=False)
        }

        # load data
        data = sio.loadmat(f"./data/ldc_Re{Re}.mat")
        u_ref = data["u"].astype(dtype)
        v_ref = data["v"].astype(dtype)
        U_ref = np.sqrt(u_ref**2 + v_ref**2).reshape(-1, 1)
        x_star = data["x"].flatten().astype(dtype)
        y_star = data["y"].flatten().astype(dtype)
        x0 = x_star[0]
        x1 = x_star[-1]
        y0 = y_star[0]
        y1 = y_star[-1]

        # set N-S pde constraint
        def gen_input_batch():
            tx = np.random.uniform(
                [x0, y0],
                [x1, y1],
                (cfg_t.TRAIN.batch_size.pde, 2),
            ).astype(dtype)
            return {"x": tx[:, 0:1], "y": tx[:, 1:2]}

        def gen_label_batch(input_batch):
            return {
                "continuity": np.zeros([cfg_t.TRAIN.batch_size.pde, 1], dtype),
                "momentum_x": np.zeros([cfg_t.TRAIN.batch_size.pde, 1], dtype),
                "momentum_y": np.zeros([cfg_t.TRAIN.batch_size.pde, 1], dtype),
            }

        pde_constraint = ppsci.constraint.SupervisedConstraint(
            {
                "dataset": {
                    "name": "ContinuousNamedArrayDataset",
                    "input": gen_input_batch,
                    "label": gen_label_batch,
                },
            },
            output_expr=equation["NavierStokes"].equations,
            loss=ppsci.loss.MSELoss("mean"),
            name="PDE",
        )

        # set boundary conditions
        x_bc = sample_points_on_square_boundary(
            cfg_t.TRAIN.batch_size.bc, eps=0.0
        ).astype(
            dtype
        )  # avoid singularity a right corner for u velocity
        v_bc = np.zeros((cfg_t.TRAIN.batch_size.bc * 4, 1), dtype)
        u_bc = copy.deepcopy(v_bc)
        lid_bc_fn = lambda x: 1 - np.cosh(50 * (x - 0.5)) / np.cosh(50 * 0.5)
        u_bc[: cfg_t.TRAIN.batch_size.bc] = lid_bc_fn(
            x_bc[: cfg_t.TRAIN.batch_size.bc, 0:1]
        )
        bc = ppsci.constraint.SupervisedConstraint(
            {
                "dataset": {
                    "name": "IterableNamedArrayDataset",
                    "input": {
                        "x": x_bc[:, 0:1],
                        "y": x_bc[:, 1:2],
                    },
                    "label": {"u": u_bc, "v": v_bc},
                },
            },
            output_expr={"u": lambda out: out["u"], "v": lambda out: out["v"]},
            loss=ppsci.loss.MSELoss("mean"),
            name="BC",
        )
        # wrap constraints together
        constraint = {
            pde_constraint.name: pde_constraint,
            bc.name: bc,
        }

        # set validator
        xy_star = misc.cartesian_product(x_star, y_star).astype(dtype)
        eval_data = {"x": xy_star[:, 0:1], "y": xy_star[:, 1:2]}
        eval_label = {"U": U_ref.reshape([-1, 1])}
        U_validator = ppsci.validate.SupervisedValidator(
            {
                "dataset": {
                    "name": "NamedArrayDataset",
                    "input": eval_data,
                    "label": eval_label,
                },
                "batch_size": cfg_t.EVAL.batch_size,
            },
            ppsci.loss.MSELoss("mean"),
            {"U": lambda out: (out["u"] ** 2 + out["v"] ** 2).sqrt()},
            metric={"L2Rel": ppsci.metric.L2Rel()},
            name="U_validator",
        )
        validator = {U_validator.name: U_validator}

        # initialize solver
        solver = ppsci.solver.Solver(
            model,
            constraint,
            optimizer=optimizer,
            equation=equation,
            validator=validator,
            loss_aggregator=grad_norm,
            cfg=cfg_t,
        )
        # train model
        solver.train()
        # evaluate after finished training
        solver.eval()
        # visualize prediction after finished training
        pred_dict = solver.predict(
            eval_data, batch_size=cfg_t.EVAL.batch_size, return_numpy=True
        )
        U_pred = np.sqrt(pred_dict["u"] ** 2 + pred_dict["v"] ** 2).reshape(
            [len(x_star), len(y_star)]
        )
        plot(U_pred, cfg_t.output_dir)

    for idx in range(len(cfg.Re)):
        train_curriculum(cfg, idx)


def evaluate(cfg: DictConfig):
    # set model
    model = ppsci.arch.PirateNet(**cfg.MODEL)

    data = sio.loadmat(cfg.EVAL_DATA_PATH)
    data = dict(data)
    u_ref = data["u"].astype(dtype)
    v_ref = data["v"].astype(dtype)
    U_ref = np.sqrt(u_ref**2 + v_ref**2).reshape(-1, 1)
    x_star = data["x"].flatten().astype(dtype)  # [nx, ]
    y_star = data["y"].flatten().astype(dtype)  # [ny, ]

    # set validator
    xy_star = misc.cartesian_product(x_star, y_star).astype(dtype)
    eval_data = {"x": xy_star[:, 0:1], "y": xy_star[:, 1:2]}
    eval_label = {"U": U_ref.reshape([-1, 1])}
    U_validator = ppsci.validate.SupervisedValidator(
        {
            "dataset": {
                "name": "NamedArrayDataset",
                "input": eval_data,
                "label": eval_label,
            },
            "batch_size": cfg.EVAL.batch_size,
        },
        ppsci.loss.MSELoss("mean"),
        {"U": lambda out: (out["u"] ** 2 + out["v"] ** 2).sqrt()},
        metric={"L2Rel": ppsci.metric.L2Rel()},
        name="U_validator",
    )
    validator = {U_validator.name: U_validator}

    # initialize solver
    solver = ppsci.solver.Solver(
        model,
        validator=validator,
        cfg=cfg,
    )

    # evaluate after finished training
    solver.eval()
    # visualize prediction after finished training
    pred_dict = solver.predict(
        eval_data, batch_size=cfg.EVAL.batch_size, return_numpy=True
    )
    U_pred = np.sqrt(pred_dict["u"] ** 2 + pred_dict["v"] ** 2).reshape(
        [len(x_star), len(y_star)]
    )
    # plot
    plot(U_pred, cfg.output_dir)


def export(cfg: DictConfig):
    # set model
    model = ppsci.arch.PirateNet(**cfg.MODEL)

    # initialize solver
    solver = ppsci.solver.Solver(model, cfg=cfg)
    # export model
    from paddle.static import InputSpec

    input_spec = [
        {key: InputSpec([None, 1], "float32", name=key) for key in model.input_keys},
    ]
    solver.export(input_spec, cfg.INFER.export_path, with_onnx=False)


def inference(cfg: DictConfig):
    from deploy.python_infer import pinn_predictor

    predictor = pinn_predictor.PINNPredictor(cfg)
    data = sio.loadmat(cfg.EVAL_DATA_PATH)
    data = dict(data)
    x_star = data["x"].flatten().astype(dtype)  # [nx, ]
    y_star = data["y"].flatten().astype(dtype)  # [ny, ]
    xy_star = misc.cartesian_product(x_star, y_star).astype(dtype)
    input_dict = {"x": xy_star[:, 0:1], "y": xy_star[:, 1:2]}

    output_dict = predictor.predict(input_dict, cfg.INFER.batch_size)
    # mapping data to cfg.INFER.output_keys
    output_dict = {
        store_key: output_dict[infer_key]
        for store_key, infer_key in zip(cfg.MODEL.output_keys, output_dict.keys())
    }
    U_pred = np.sqrt(output_dict["u"] ** 2 + output_dict["v"] ** 2).reshape(
        [len(x_star), len(y_star)]
    )
    plot(U_pred, cfg.output_dir)


@hydra.main(
    version_base=None, config_path="./conf", config_name="ldc_2d_Re3200_piratenet.yaml"
)
def main(cfg: DictConfig):
    if cfg.mode == "train":
        train(cfg)
    elif cfg.mode == "eval":
        evaluate(cfg)
    elif cfg.mode == "export":
        export(cfg)
    elif cfg.mode == "infer":
        inference(cfg)
    else:
        raise ValueError(
            f"cfg.mode should in ['train', 'eval', 'export', 'infer'], but got '{cfg.mode}'"
        )


if __name__ == "__main__":
    main()

5. 结果展示

下方展示了模型对于边长为 1 的正方形计算域的内部点进行预测的结果 \(U=\sqrt{u^2+v^2}\)。

ldc_re1000_sota_ac

可以看到在 \(Re=1000\) 下，预测结果与求解器的结果基本相同（L2 相对误差为 7.7%）。

ldc_re3200_piratenet_ac

可以看到在 \(Re=3200\) 下，预测结果与求解器的结果基本相同（L2 相对误差为 4.1%）。

6. 参考资料