{ "cells": [ { "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "# 基于FNO求解二维Navier-Stokes\n", "\n", "[![下载Notebook](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.0.0-alpha/resource/_static/logo_notebook.png)](https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/r2.0.0-alpha/mindflow/zh_cn/data_driven/mindspore_FNO2D.ipynb) [![下载样例代码](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.0.0-alpha/resource/_static/logo_download_code.png)](https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/r2.0.0-alpha/mindflow/zh_cn/data_driven/mindspore_FNO2D.py) [![查看源文件](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.0.0-alpha/resource/_static/logo_source.png)](https://gitee.com/mindspore/docs/blob/r2.0.0-alpha/docs/mindflow/docs/source_zh_cn/data_driven/FNO2D.ipynb)" ] }, { "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "## 概述\n", "\n", "计算流体力学是21世纪流体力学领域的重要技术之一，其通过使用数值方法在计算机中对流体力学的控制方程进行求解，从而实现流动的分析、预测和控制。传统的有限元法（finite element method，FEM）和有限差分法（finite difference method，FDM）常用于复杂的仿真流程（物理建模、网格划分、数值离散、迭代求解等）和较高的计算成本，往往效率低下。因此，借助AI提升流体仿真效率是十分必要的。\n", "\n", "近年来，随着神经网络的迅猛发展，为科学计算提供了新的范式。经典的神经网络是在有限维度的空间进行映射，只能学习与特定离散化相关的解。与经典神经网络不同，傅里叶神经算子（Fourier Neural Operator，FNO）是一种能够学习无限维函数空间映射的新型深度学习架构。该架构可直接学习从任意函数参数到解的映射，用于解决一类偏微分方程的求解问题，具有更强的泛化能力。更多信息可参考[Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/abs/2010.08895)。\n", "\n", "本案例教程介绍利用傅里叶神经算子的纳维-斯托克斯方程（Navier-Stokes equation）求解方法。" ] }, { "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "## 纳维-斯托克斯方程（Navier-Stokes equation）\n", "\n", "纳维-斯托克斯方程（Navier-Stokes equation）是计算流体力学领域的经典方程，是一组描述流体动量守恒的偏微分方程，简称N-S方程。它在二维不可压缩流动中的涡度形式如下：\n", "\n", "$$\n", "\\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]\n", "$$\n", "\n", "$$\n", "\\nabla \\cdot u(x, t)=0, \\quad x \\in(0,1)^2, t \\in[0, T]\n", "$$\n", "\n", "$$\n", "w(x, 0)=w_0(x), \\quad x \\in(0,1)^2\n", "$$\n", "\n", "其中$u$表示速度场，$w=\\nabla \\times u$表示涡度，$w_0(x)$表示初始条件，$\\nu$表示粘度系数，$f(x)$为外力合力项。\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 问题描述\n", "\n", "本案例利用Fourier Neural Operator学习某一个时刻对应涡度到下一时刻涡度的映射，实现二维不可压缩N-S方程的求解：\n", "\n", "$$\n", "w_t \\mapsto w(\\cdot, t+1)\n", "$$" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 技术路径\n", "\n", "MindFlow求解该问题的具体流程如下：\n", "\n", "1. 创建数据集。\n", "2. 构建模型。\n", "3. 优化器与损失函数。\n", "4. 定义求解器。\n", "5. 定义回调函数。\n", "6. 模型训练。" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Fourier Neural Operator\n", "\n", "Fourier Neural Operator模型构架如下图所示。图中$w_0(x)$表示初始涡度，通过Lifting Layer实现输入向量的高维映射，然后将映射结果作为Fourier Layer的输入，进行频域信息的非线性变换，最后由Decoding Layer将变换结果映射至最终的预测结果$w_1(x)$。\n", "\n", "Lifting Layer、Fourier Layer以及Decoding Layer共同组成了Fourier Neural Operator。\n", "\n", "![Fourier Neural Operator模型构架](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.0.0-alpha/docs/mindflow/docs/source_zh_cn/data_driven/images/FNO.png)\n", "\n", "Fourier Layer网络结构如下图所示。图中V表示输入向量，上框表示向量经过傅里叶变换后，经过线性变换R，过滤高频信息，然后进行傅里叶逆变换；另一分支经过线性变换W，最后通过激活函数，得到Fourier Layer输出向量。\n", "\n", "![Fourier Layer网络结构](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.0.0-alpha/docs/mindflow/docs/source_zh_cn/data_driven/images/FNO-2.png)" ] }, { "cell_type": "code", "execution_count": 11, "metadata": { "pycharm": { "name": "#%%\n" } }, "outputs": [], "source": [ "import os\n", "import numpy as np\n", "\n", "from mindspore import nn, context, Tensor, set_seed\n", "from mindspore import DynamicLossScaleManager, LossMonitor, TimeMonitor, CheckpointConfig, ModelCheckpoint\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "下述`src`包可以在[applications/data_driven/navier_stokes/src](https://gitee.com/mindspore/mindscience/tree/r0.2.0-alpha/MindFlow/applications/data_driven/navier_stokes/src)下载。\n" ] }, { "cell_type": "code", "execution_count": 12, "metadata": { "pycharm": { "name": "#%%\n" } }, "outputs": [], "source": [ "from mindflow import FNO2D, RelativeRMSELoss, Solver, load_yaml_config, get_warmup_cosine_annealing_lr\n", "\n", "from src import PredictCallback, create_training_dataset\n", "\n", "set_seed(0)\n", "np.random.seed(0)" ] }, { "cell_type": "code", "execution_count": 13, "metadata": { "pycharm": { "name": "#%%\n" } }, "outputs": [], "source": [ "context.set_context(mode=context.GRAPH_MODE, save_graphs=False, device_target='GPU')" ] }, { "cell_type": "code", "execution_count": 14, "metadata": { "pycharm": { "name": "#%%\n" } }, "outputs": [], "source": [ "config = load_yaml_config('navier_stokes_2d.yaml')\n", "data_params = config[\"data\"]\n", "model_params = config[\"model\"]\n", "optimizer_params = config[\"optimizer\"]\n", "callback_params = config[\"callback\"]" ] }, { "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "## 创建数据集\n", "\n", "训练与测试数据下载: [data_driven/navier_stokes/dataset](https://download.mindspore.cn/mindscience/mindflow/dataset/applications/data_driven/navier_stokes/dataset/) .\n", "\n", "本案例根据Zongyi Li在 [Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/pdf/2010.08895.pdf) 一文中对数据集的设置生成训练数据集与测试数据集。具体设置如下：\n", "\n", "基于周期性边界，生成满足如下分布的初始条件$w_0(x)$：\n", "\n", "$$\n", "w_0 \\sim \\mu, \\mu=\\mathcal{N}\\left(0,7^{3 / 2}(-\\Delta+49 I)^{-2.5}\\right)\n", "$$\n", "\n", "外力项设置为：\n", "\n", "$$\n", "f(x)=0.1\\left(\\sin \\left(2 \\pi\\left(x_1+x_2\\right)\\right)+\\right.\\cos(2 \\pi(x_1+x_2)))\n", "$$\n", "\n", "采用`Crank-Nicolson`方法生成数据，时间步长设置为1e-4，最终数据以每 t = 1 个时间单位记录解。所有数据均在256×256的网格上生成，并被下采样至64×64网格。本案例选取粘度系数$\\nu=1e−5$，训练集样本量为19000个，测试集样本量为3800个。" ] }, { "cell_type": "code", "execution_count": 15, "metadata": { "pycharm": { "name": "#%%\n" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Data preparation finished\n" ] } ], "source": [ "train_dataset = create_training_dataset(data_params, input_resolution=model_params[\"input_resolution\"], shuffle=True)\n", "test_input = np.load(os.path.join(data_params[\"path\"], \"test/inputs.npy\"))\n", "test_label = np.load(os.path.join(data_params[\"path\"], \"test/label.npy\"))" ] }, { "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "## 构建模型\n", "\n", "网络由1层Lifting layer、多层Fourier Layer以及1层Decoding layer叠加组成：\n", "\n", "- Lifting layer对应样例代码中`FNO2D.fc0`，将输出数据$x$映射至高维；\n", "\n", "- 多层Fourier Layer的叠加对应样例代码中`FNO2D.fno_seq`，本案例采用离散傅里叶变换实现时域与频域的转换；\n", "\n", "- Decoding layer对应代码中`FNO2D.fc1`与`FNO2D.fc2`，获得最终的预测值。" ] }, { "cell_type": "code", "execution_count": 16, "metadata": { "pycharm": { "name": "#%%\n" } }, "outputs": [], "source": [ "model = FNO2D(in_channels=model_params[\"in_channels\"],\n", " out_channels=model_params[\"out_channels\"],\n", " resolution=model_params[\"input_resolution\"],\n", " modes=model_params[\"modes\"],\n", " channels=model_params[\"width\"],\n", " depth=model_params[\"depth\"]\n", " )" ] }, { "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "## 优化器与损失函数\n", "\n", "使用相对均方根误差作为网络训练损失函数：" ] }, { "cell_type": "code", "execution_count": 17, "metadata": { "pycharm": { "name": "#%%\n" } }, "outputs": [], "source": [ "steps_per_epoch = train_dataset.get_dataset_size()\n", "lr = get_warmup_cosine_annealing_lr(lr_init=optimizer_params[\"initial_lr\"],\n", " last_epoch=optimizer_params[\"train_epochs\"],\n", " steps_per_epoch=steps_per_epoch,\n", " warmup_epochs=optimizer_params[\"warmup_epochs\"])\n", "\n", "optimizer = nn.Adam(model.trainable_params(), learning_rate=Tensor(lr))\n", "loss_scale = DynamicLossScaleManager()\n", "\n", "# prepare loss function\n", "loss_fn = RelativeRMSELoss()" ] }, { "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "## 定义求解器\n", "\n", "Solver类是模型训练和推理的接口。输入优化器、网络模型、损失函数、损失缩放策略等，即可定义求解器对象solver。代码中optimizer_params、model_params对应各项参数均在[配置文件](https://gitee.com/mindspore/mindscience/blob/r0.2.0-alpha/MindFlow/applications/data_driven/navier_stokes/navier_stokes_2d.yaml)中修改。" ] }, { "cell_type": "code", "execution_count": 18, "metadata": { "pycharm": { "name": "#%%\n" } }, "outputs": [], "source": [ "solver = Solver(model,\n", " optimizer=optimizer,\n", " loss_scale_manager=loss_scale,\n", " loss_fn=loss_fn,\n", " )" ] }, { "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "## 定义回调函数" ] }, { "cell_type": "code", "execution_count": 19, "metadata": { "pycharm": { "name": "#%%\n" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "./FNO2D\n", "check test dataset shape: (200, 19, 64, 64, 1), (200, 19, 64, 64, 1)\n" ] } ], "source": [ "summary_dir = os.path.join(callback_params[\"summary_dir\"], 'FNO2D')\n", "print(summary_dir)\n", "pred_cb = PredictCallback(model=model,\n", " inputs=test_input,\n", " label=test_label,\n", " config=callback_params,\n", " summary_dir=summary_dir)\n", "\n", "ckpt_config = CheckpointConfig(save_checkpoint_steps=callback_params[\"save_checkpoint_steps\"] * steps_per_epoch,\n", " keep_checkpoint_max=callback_params[\"keep_checkpoint_max\"])\n", "ckpt_dir = os.path.join(summary_dir, \"ckpt\")\n", "ckpt_cb = ModelCheckpoint(prefix=model_params[\"name\"],\n", " directory=ckpt_dir,\n", " config=ckpt_config)" ] }, { "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "## 模型训练\n", "\n", "调用求解器接口进行模型训练，调用回调接口进行评估。" ] }, { "cell_type": "code", "execution_count": 20, "metadata": { "pycharm": { "name": "#%%\n" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "epoch: 1 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 2.07)\n", "Train epoch time: 36526.785 ms, per step time: 36.527 ms\n", "epoch: 2 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 2.00379)\n", "Train epoch time: 29215.492 ms, per step time: 29.215 ms\n", "epoch: 3 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 1.40253)\n", "Train epoch time: 29217.016 ms, per step time: 29.217 ms\n", "epoch: 4 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 1.79683)\n", "Train epoch time: 29243.756 ms, per step time: 29.244 ms\n", "epoch: 5 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 1.42917)\n", "Train epoch time: 29197.400 ms, per step time: 29.197 ms\n", "epoch: 6 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 1.24265)\n", "Train epoch time: 29199.672 ms, per step time: 29.200 ms\n", "epoch: 7 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 1.48525)\n", "Train epoch time: 29193.341 ms, per step time: 29.193 ms\n", "epoch: 8 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 1.2069)\n", "Train epoch time: 29198.366 ms, per step time: 29.198 ms\n", "epoch: 9 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 1.17752)\n", "Train epoch time: 29210.540 ms, per step time: 29.211 ms\n", "epoch: 10 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 1.25935)\n", "Train epoch time: 29180.896 ms, per step time: 29.181 ms\n", "================================Start Evaluation================================\n", "mean rel_rmse_error: 0.14311510016862303\n", "=================================End Evaluation=================================\n", "predict total time: 16.270995616912842 s\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "...\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "epoch: 141 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 0.667819)\n", "Train epoch time: 29181.800 ms, per step time: 29.182 ms\n", "epoch: 142 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 0.610858)\n", "Train epoch time: 29203.687 ms, per step time: 29.204 ms\n", "epoch: 143 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 0.616083)\n", "Train epoch time: 29199.107 ms, per step time: 29.199 ms\n", "epoch: 144 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 0.609115)\n", "Train epoch time: 29302.156 ms, per step time: 29.302 ms\n", "epoch: 145 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 0.518936)\n", "Train epoch time: 29234.649 ms, per step time: 29.235 ms\n", "epoch: 146 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 0.822775)\n", "Train epoch time: 29228.318 ms, per step time: 29.228 ms\n", "epoch: 147 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 0.802282)\n", "Train epoch time: 29231.589 ms, per step time: 29.232 ms\n", "epoch: 148 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 0.669333)\n", "Train epoch time: 29285.277 ms, per step time: 29.285 ms\n", "epoch: 149 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 0.615759)\n", "Train epoch time: 29311.589 ms, per step time: 29.312 ms\n", "epoch: 150 step: 1000, loss is Tensor(shape=[], dtype=Float32, value= 0.75713)\n", "Train epoch time: 29280.815 ms, per step time: 29.281 ms\n", "================================Start Evaluation================================\n", "mean rel_rmse_error: 0.06585168887209147\n", "=================================End Evaluation=================================\n", "predict total time: 12.599207639694214 s\n" ] } ], "source": [ "solver.train(epoch=optimizer_params[\"train_epochs\"],\n", " train_dataset=train_dataset,\n", " callbacks=[LossMonitor(), TimeMonitor(), pred_cb, ckpt_cb],\n", " dataset_sink_mode=True)" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3.9.0 ('py39')", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.9.0" }, "vscode": { "interpreter": { "hash": "57ace93c29d9374277a79956c3f1b916d7d9a05468d906842f9921d0d494a29f" } } }, "nbformat": 4, "nbformat_minor": 1 }