{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# 基于FNO求解一维Burgers\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_FNO1D.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_FNO1D.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/FNO1D.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",
    "本案例教程介绍利用傅里叶神经算子的1-d Burgers方程求解方法。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 伯格斯方程（Burgers' equation）\n",
    "\n",
    "一维伯格斯方程（1-d Burgers' equation）是一个非线性偏微分方程，具有广泛应用，包括一维粘性流体流动建模。它的形式如下：\n",
    "\n",
    "$$\n",
    "\\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]\n",
    "$$\n",
    "\n",
    "$$\n",
    "u(x, 0)=u_0(x), \\quad x \\in(0,1)\n",
    "$$\n",
    "\n",
    "其中$u$表示速度场，$u_0$表示初始条件，$\\nu$表示粘度系数。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 问题描述\n",
    "\n",
    "本案例利用Fourier Neural Operator学习初始状态到下一时刻状态的映射，实现一维Burgers'方程的求解：\n",
    "\n",
    "$$\n",
    "u_0 \\mapsto u(\\cdot, 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": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from mindspore import context, nn, Tensor, set_seed\n",
    "from mindspore import DynamicLossScaleManager, LossMonitor, TimeMonitor\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "下述`src`包可以在[applications/data_driven/burgers/src](https://gitee.com/mindspore/mindscience/tree/r0.2.0-alpha/MindFlow/applications/data_driven/burgers/src)下载。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from mindflow import FNO1D, RelativeRMSELoss, Solver, load_yaml_config, get_warmup_cosine_annealing_lr\n",
    "\n",
    "from src import PredictCallback, create_training_dataset\n",
    "\n",
    "\n",
    "set_seed(0)\n",
    "np.random.seed(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "context.set_context(mode=context.GRAPH_MODE, device_target='GPU', device_id=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "config = load_yaml_config(\"burgers1d.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/burgers/dataset](https://download.mindspore.cn/mindscience/mindflow/dataset/applications/data_driven/burgers/dataset/)。\n",
    "\n",
    "本案例根据Zongyi Li在 [Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/pdf/2010.08895.pdf) 一文中对数据集的设置生成训练数据集与测试数据集。具体设置如下：\n",
    "基于周期性边界，生成满足如下分布的初始条件$u_0(x)$：\n",
    "\n",
    "$$\n",
    "u_0 \\sim \\mu, \\mu=\\mathcal{N}\\left(0,625(-\\Delta+25 I)^{-2}\\right)\n",
    "$$\n",
    "\n",
    "本案例选取粘度系数$\\nu=0.1$，并使用分步法求解方程，其中热方程部分在傅里叶空间中精确求解，然后使用前向欧拉方法求解非线性部分。训练集样本量为1000个，测试集样本量为200个。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data preparation finished\n",
      "input_path:  (1000, 1024, 1)\n",
      "label_path:  (1000, 1024)\n"
     ]
    }
   ],
   "source": [
    "# create training dataset\n",
    "train_dataset = create_training_dataset(data_params, shuffle=True)\n",
    "\n",
    "# create test dataset\n",
    "test_input, test_label = np.load(os.path.join(data_params[\"path\"], \"test/inputs.npy\")), \\\n",
    "                         np.load(os.path.join(data_params[\"path\"], \"test/label.npy\"))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 构建模型\n",
    "\n",
    "网络由1层Lifting layer、1层Decoding layer以及多层Fourier Layer叠加组成：\n",
    "\n",
    "- Lifting layer对应样例代码中`FNO1D.fc0`，将输出数据$x$映射至高维；\n",
    "\n",
    "- 多层Fourier Layer的叠加对应样例代码中`FNO1D.fno_seq`，本案例采用离散傅里叶变换实现时域与频域的转换；\n",
    "\n",
    "- Decoding layer对应代码中`FNO1D.fc1`与`FNO1D.fc2`，获得最终的预测值。\n",
    "\n",
    "基于上述网络结构，进行模型初始化，其中模型参数可在[配置文件](https://gitee.com/mindspore/mindscience/blob/r0.2.0-alpha/MindFlow/applications/data_driven/burgers/burgers1d.yaml)中修改。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "model = FNO1D(in_channels=model_params[\"in_channels\"],\n",
    "              out_channels=model_params[\"out_channels\"],\n",
    "              resolution=model_params[\"resolution\"],\n",
    "              modes=model_params[\"modes\"],\n",
    "              channels=model_params[\"width\"],\n",
    "              depth=model_params[\"depth\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 优化器与损失函数\n",
    "\n",
    "使用相对均方根误差作为网络训练损失函数："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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=1)\n",
    "optimizer = nn.Adam(model.trainable_params(), learning_rate=Tensor(lr))\n",
    "loss_scale = DynamicLossScaleManager()\n",
    "\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/burgers/burgers1d.yaml)中修改。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "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": 9,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./FNO1D\n",
      "check test dataset shape: (200, 1024, 1), (200, 1024)\n"
     ]
    }
   ],
   "source": [
    "summary_dir = os.path.join(callback_params[\"summary_dir\"], \"FNO1D\")\n",
    "print(summary_dir)\n",
    "pred_cb = PredictCallback(model=model,\n",
    "                          inputs=test_input,\n",
    "                          label=test_label,\n",
    "                          config=config,\n",
    "                          summary_dir=summary_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 模型训练\n",
    "\n",
    "调用求解器接口进行模型训练，调用回调接口进行评估。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1 step: 125, loss is 2.377823\n",
      "Train epoch time: 5782.938 ms, per step time: 46.264 ms\n",
      "epoch: 2 step: 125, loss is 0.88470775\n",
      "Train epoch time: 1150.446 ms, per step time: 9.204 ms\n",
      "epoch: 3 step: 125, loss is 0.98071647\n",
      "Train epoch time: 1135.464 ms, per step time: 9.084 ms\n",
      "epoch: 4 step: 125, loss is 0.5404751\n",
      "Train epoch time: 1114.245 ms, per step time: 8.914 ms\n",
      "epoch: 5 step: 125, loss is 0.39976493\n",
      "Train epoch time: 1125.107 ms, per step time: 9.001 ms\n",
      "epoch: 6 step: 125, loss is 0.508416\n",
      "Train epoch time: 1127.477 ms, per step time: 9.020 ms\n",
      "epoch: 7 step: 125, loss is 0.42839915\n",
      "Train epoch time: 1125.775 ms, per step time: 9.006 ms\n",
      "epoch: 8 step: 125, loss is 0.28270185\n",
      "Train epoch time: 1118.428 ms, per step time: 8.947 ms\n",
      "epoch: 9 step: 125, loss is 0.24137405\n",
      "Train epoch time: 1121.705 ms, per step time: 8.974 ms\n",
      "epoch: 10 step: 125, loss is 0.22623646\n",
      "Train epoch time: 1118.699 ms, per step time: 8.950 ms\n",
      "================================Start Evaluation================================\n",
      "mean rms_error: 0.03270653011277318\n",
      "=================================End Evaluation=================================\n",
      "...\n",
      "predict total time: 0.5012176036834717 s\n",
      "epoch: 91 step: 125, loss is 0.026378194\n",
      "Train epoch time: 1119.095 ms, per step time: 8.953 ms\n",
      "epoch: 92 step: 125, loss is 0.057838168\n",
      "Train epoch time: 1116.712 ms, per step time: 8.934 ms\n",
      "epoch: 93 step: 125, loss is 0.034773324\n",
      "Train epoch time: 1107.931 ms, per step time: 8.863 ms\n",
      "epoch: 94 step: 125, loss is 0.029720988\n",
      "Train epoch time: 1109.336 ms, per step time: 8.875 ms\n",
      "epoch: 95 step: 125, loss is 0.02933883\n",
      "Train epoch time: 1111.804 ms, per step time: 8.894 ms\n",
      "epoch: 96 step: 125, loss is 0.03140598\n",
      "Train epoch time: 1116.788 ms, per step time: 8.934 ms\n",
      "epoch: 97 step: 125, loss is 0.03695058\n",
      "Train epoch time: 1115.020 ms, per step time: 8.920 ms\n",
      "epoch: 98 step: 125, loss is 0.039841708\n",
      "Train epoch time: 1120.316 ms, per step time: 8.963 ms\n",
      "epoch: 99 step: 125, loss is 0.039001673\n",
      "Train epoch time: 1134.618 ms, per step time: 9.077 ms\n",
      "epoch: 100 step: 125, loss is 0.038434036\n",
      "Train epoch time: 1116.549 ms, per step time: 8.932 ms\n",
      "================================Start Evaluation================================\n",
      "mean rms_error: 0.005707952339434996\n",
      "=================================End Evaluation=================================\n",
      "predict total time: 0.5055065155029297 s\n"
     ]
    }
   ],
   "source": [
    "solver.train(epoch=optimizer_params[\"train_epochs\"],\n",
    "             train_dataset=train_dataset,\n",
    "             callbacks=[LossMonitor(), TimeMonitor(), pred_cb],\n",
    "             dataset_sink_mode=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.0 64-bit",
   "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.8.0"
  },
  "vscode": {
   "interpreter": {
    "hash": "df0893f56f349688326838aaeea0de204df53a132722cbd565e54b24a8fec5f6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}