{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# 数据驱动(FNO2D和UNET2D两种backbone)下跨声速翼型复杂流场的多时间步预测\n",
    "\n",
    "[![下载Notebook](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/resource/_static/logo_notebook.svg)](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/master/mindflow/zh_cn/data_driven/mindspore_2D_unsteady.ipynb) [![下载样例代码](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/resource/_static/logo_download_code.svg)](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/master/mindflow/zh_cn/data_driven/mindspore_2D_unsteady.py) [![查看源文件](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/resource/_static/logo_source.svg)](https://gitee.com/mindspore/docs/blob/master/docs/mindflow/docs/source_zh_cn/data_driven/2D_unsteady.ipynb)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **概述**\n",
    "\n",
    "高精度非定常流动模拟是计算流体力学中的关键课题，有着广泛的应用场景和广阔的市场价值。然而，传统方法存在着算不快、算不准、算不稳等问题，通过AI方法探索流场演化规律为此提供了新的视角。\n",
    "\n",
    "本案例在二维跨声速翼型场景下提供了端到端的非定常复杂流场预测解决方案。案例中搭建了傅里叶神经算子(Fourier Neural Operator，FNO)和Unet两种网络结构。可以在保证一定精度的前提下，根据输入的*k*个时间步的流场，稳定预测出接续的*m*个时间步的流场。\n",
    "\n",
    "本案例中，流场来流马赫数达到了*Ma*=0.73。通过本案例可以验证深度学习方法在存在激波等复杂流动结构场景中，对多物理参数变化下非定常流场预测的有效性。\n",
    "\n",
    "![img_1.png](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/docs/mindflow/docs/source_zh_cn/data_driven/images/img_1.png)\n",
    "\n",
    "## 问题描述\n",
    "\n",
    "本案例利用*k*个时刻的流场学习接续的*m*个时刻的流场，实现二维可压缩非定常流场的预测：\n",
    "\n",
    "$$\n",
    "u_{[t_0\\sim t_{k-1}]} \\mapsto u_{[t_k\\sim t_{k+m}]}.\n",
    "$$\n",
    "\n",
    "## 技术路径\n",
    "\n",
    "求解该问题的具体流程如下：\n",
    "\n",
    "1. 构建数据集。\n",
    "2. 构建模型。\n",
    "3. 优化器与损失函数。\n",
    "4. 模型训练。\n",
    "\n",
    "## 准备环节\n",
    "\n",
    "实践前，确保已经正确安装合适版本的MindSpore。如果没有，可以通过：\n",
    "\n",
    "* [MindSpore安装页面](https://www.mindspore.cn/install) 安装MindSpore。\n",
    "\n",
    "## 二维翼型非定常流场预测的实现\n",
    "\n",
    "二维翼型非定常流场预测的实现分为以下7个步骤：\n",
    "\n",
    "1. 配置网络与训练参数\n",
    "2. 数据集的准备\n",
    "3. 模型构建\n",
    "4. 损失函数与优化器\n",
    "5. 训练函数\n",
    "6. 模型训练\n",
    "7. 结果可视化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "\n",
    "import numpy as np\n",
    "from mindspore import nn, Tensor, context, ops, jit, set_seed, data_sink, save_checkpoint\n",
    "from mindspore import dtype as mstype\n",
    "from mindflow.common import get_warmup_cosine_annealing_lr\n",
    "from mindflow.loss import RelativeRMSELoss\n",
    "from mindflow.utils import load_yaml_config, print_log\n",
    "\n",
    "from src import Trainer, init_dataset, init_model, plt_log, check_file_path, count_params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "set_seed(0)\n",
    "np.random.seed(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "context.set_context(mode=context.GRAPH_MODE,\n",
    "                    save_graphs=False,\n",
    "                    device_target=\"Ascend\",\n",
    "                    device_id=0)\n",
    "use_ascend = context.get_context(\"device_target\") == \"Ascend\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 配置网络与训练参数\n",
    "\n",
    "从配置文件中读取四类参数，分别为模型相关参数（model）、数据相关参数（data）、优化器相关参数（optimizer)和回调相关参数(callback)。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "config = load_yaml_config(\"./config/2D_unsteady.yaml\")\n",
    "data_params = config[\"data\"]\n",
    "model_params = config[\"model\"]\n",
    "optimizer_params = config[\"optimizer\"]\n",
    "summary_params = config[\"summary\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 数据集的准备\n",
    "\n",
    "数据集下载地址：[data_driven/airfoil/2D_unsteady](https://download.mindspore.cn/mindscience/mindflow/dataset/applications/data_driven/airfoil/2D_unsteady/)\n",
    "\n",
    "数据为npz类型文件，其维度(*t*, *H*, *W*, *C*)为(9997, 128, 128, 3)。其中，*t*为时间步数，*H*和*W*为流场分辨率，*C*为通道数，3个通道分别为速度*U*、*V*和压力*P*。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input size (3560, 8, 128, 128, 3)\n",
      "label size (3560, 32, 128, 128, 3)\n",
      "train_batch_size : 8\n",
      "train dataset size: 2967\n",
      "test dataset size: 593\n",
      "train batch dataset size: 370\n",
      "test batch dataset size: 74\n"
     ]
    }
   ],
   "source": [
    "train_dataset, test_dataset, means, stds = init_dataset(data_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 模型构建\n",
    "\n",
    "通过initial_model()函数调用，调用之前，需要先针对硬件定制loss_scaler和compute_dtype。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "compute_dtype_of_FNO Float16\n"
     ]
    }
   ],
   "source": [
    "if use_ascend:\n",
    "    from mindspore.amp import DynamicLossScaler, all_finite, auto_mixed_precision\n",
    "    loss_scaler = DynamicLossScaler(1024, 2, 100)\n",
    "    compute_dtype = mstype.float16\n",
    "    model = init_model(\"fno2d\", data_params, model_params, compute_dtype=compute_dtype)\n",
    "    auto_mixed_precision(model, optimizer_params[\"amp_level\"][\"fno2d\"])\n",
    "else:\n",
    "    context.set_context(enable_graph_kernel=True)\n",
    "    loss_scaler = None\n",
    "    compute_dtype = mstype.float32\n",
    "    model = init_model(\"fno2d\", data_params, model_params, compute_dtype=compute_dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 损失函数与优化器\n",
    "\n",
    "当前案例中的损失函数采用了RelativeRMSELoss，优化器则选择了AdamWeightDecay，其中，学习率衰减采用了warmup_cosine_annealing_lr的策略。用户也可以根据需要定制适合的损失函数与优化器。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model parameter count: 9464259\n",
      "learing rate: 0.001, T_in: 8, T_out: 32\n"
     ]
    }
   ],
   "source": [
    "loss_fn = RelativeRMSELoss()\n",
    "summary_dir = os.path.join(summary_params[\"summary_dir\"], \"Exp01\", \"fno2d\")\n",
    "ckpt_dir = os.path.join(summary_dir, \"ckpt_dir\")\n",
    "check_file_path(ckpt_dir)\n",
    "check_file_path(os.path.join(ckpt_dir, 'img'))\n",
    "print_log('model parameter count:', count_params(model.trainable_params()))\n",
    "print_log(\n",
    "    f'learing rate: {optimizer_params[\"lr\"][\"fno2d\"]}, T_in: {data_params[\"T_in\"]}, T_out: {data_params[\"T_out\"]}')\n",
    "steps_per_epoch = train_dataset.get_dataset_size()\n",
    "\n",
    "lr = get_warmup_cosine_annealing_lr(optimizer_params[\"lr\"][\"fno2d\"], steps_per_epoch,\n",
    "                                    optimizer_params[\"epochs\"], optimizer_params[\"warm_up_epochs\"])\n",
    "optimizer = nn.AdamWeightDecay(model.trainable_params(),\n",
    "                               learning_rate=Tensor(lr),\n",
    "                               weight_decay=optimizer_params[\"weight_decay\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 训练函数\n",
    "\n",
    "使用**MindSpore>= 2.0.0**的版本，可以使用函数式编程范式训练神经网络，单步训练函数使用jit装饰。数据下沉函数data_sink，传入单步训练函数和训练数据集。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(model, data_params, loss_fn, means, stds)\n",
    "\n",
    "def forward_fn(inputs, labels):\n",
    "    loss, _, _, _ = trainer.get_loss(inputs, labels)\n",
    "    if use_ascend:\n",
    "        loss = loss_scaler.scale(loss)\n",
    "    return loss\n",
    "\n",
    "grad_fn = ops.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=False)\n",
    "\n",
    "@jit\n",
    "def train_step(inputs, labels):\n",
    "    loss, grads = grad_fn(inputs, labels)\n",
    "    if use_ascend:\n",
    "        loss = loss_scaler.unscale(loss)\n",
    "        if all_finite(grads):\n",
    "            grads = loss_scaler.unscale(grads)\n",
    "    loss_new = ops.depend(loss, optimizer(grads))\n",
    "    return loss_new\n",
    "\n",
    "def test_step(inputs, labels):\n",
    "    return trainer.get_loss(inputs, labels)\n",
    "\n",
    "train_size = train_dataset.get_dataset_size()\n",
    "test_size = test_dataset.get_dataset_size()\n",
    "train_sink = data_sink(train_step, train_dataset, sink_size=1)\n",
    "test_sink = data_sink(test_step, test_dataset, sink_size=1)\n",
    "test_interval = summary_params[\"test_interval\"]\n",
    "save_ckpt_interval = summary_params[\"save_ckpt_interval\"]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 模型训练\n",
    "\n",
    "模型训练过程中边训练边推理。用户可以直接加载测试数据集，每训练test_interval个epoch后输出一次测试集上的推理精度并保存可视化结果。同时，还可以每隔save_checkpoint_interval保存一次checkpoint文件。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1, step time: 2.652332, loss: 0.733017\n",
      "epoch: 2, step time: 0.688175, loss: 0.203251\n",
      "epoch: 3, step time: 0.686817, loss: 0.128816\n",
      "epoch: 4, step time: 0.685909, loss: 0.109786\n",
      "epoch: 5, step time: 0.688545, loss: 0.093725\n",
      "epoch: 6, step time: 0.685986, loss: 0.076027\n",
      "epoch: 7, step time: 0.686459, loss: 0.069847\n",
      "epoch: 8, step time: 0.688228, loss: 0.058694\n",
      "epoch: 9, step time: 0.688053, loss: 0.060886\n",
      "epoch: 10, step time: 0.692221, loss: 0.065305\n",
      "---------------------------start test-------------------------\n",
      " test epoch: 10, loss: 0.03798117920381923\n",
      "---------------------------end test---------------------------\n",
      "...\n",
      "epoch: 191, step time: 0.693253, loss: 0.007012\n",
      "epoch: 192, step time: 0.691330, loss: 0.007043\n",
      "epoch: 193, step time: 0.692804, loss: 0.006986\n",
      "epoch: 194, step time: 0.690053, loss: 0.006973\n",
      "epoch: 195, step time: 0.692159, loss: 0.006967\n",
      "epoch: 196, step time: 0.690170, loss: 0.006944\n",
      "epoch: 197, step time: 0.690344, loss: 0.006930\n",
      "epoch: 198, step time: 0.690674, loss: 0.006911\n",
      "epoch: 199, step time: 0.690877, loss: 0.006904\n",
      "epoch: 200, step time: 0.689170, loss: 0.006892\n",
      "---------------------------start test-------------------------\n",
      " test epoch: 200, loss: 0.005457837492891436\n",
      "---------------------------end test---------------------------\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(1, optimizer_params[\"epochs\"] + 1):\n",
    "    time_beg = time.time()\n",
    "    train_l2_step = 0.0\n",
    "    model.set_train()\n",
    "    for step in range(1, train_size + 1):\n",
    "        loss = train_sink()\n",
    "        train_l2_step += loss.asnumpy()\n",
    "    train_l2_step = train_l2_step / train_size / data_params[\"T_out\"]\n",
    "    print_log(\n",
    "        f\"epoch: {epoch}, step time: {(time.time() - time_beg) / steps_per_epoch:>7f}, loss: {train_l2_step:>7f}\")\n",
    "\n",
    "    if epoch % test_interval == 0:\n",
    "        model.set_train(False)\n",
    "        test_l2_by_step = [0.0 for _ in range(data_params[\"T_out\"])]\n",
    "        print_log(\"---------------------------start test-------------------------\")\n",
    "        for step in range(test_size):\n",
    "            _, pred, truth, step_losses = test_sink()\n",
    "            for i in range(data_params[\"T_out\"]):\n",
    "                test_l2_by_step[i] += step_losses[i].asnumpy()\n",
    "        test_l2_by_step = [error / test_size for error in test_l2_by_step]\n",
    "        test_l2_step = np.mean(test_l2_by_step)\n",
    "        print_log(f' test epoch: {epoch}, loss: {test_l2_step}')\n",
    "        print_log(\"---------------------------end test---------------------------\")\n",
    "\n",
    "        plt_log(predicts=pred.asnumpy(),\n",
    "                labels=truth.asnumpy(),\n",
    "                img_dir=os.path.join(ckpt_dir, 'img'),\n",
    "                epoch=epoch\n",
    "                )\n",
    "\n",
    "    if epoch % save_ckpt_interval == 0:\n",
    "        save_checkpoint(model, ckpt_file_name=os.path.join(ckpt_dir, 'airfoil2D_unsteady.ckpt'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## **结果可视化**\n",
    "\n",
    "UNET2D backbone下，不同时刻压力P的实际值、预测值和误差在流场中的分布如下图：\n",
    "\n",
    "![Unet_P.png](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/docs/mindflow/docs/source_zh_cn/data_driven/images/Unet_P.png)\n",
    "\n",
    "UNET2D backbone下，不同时刻速度U的实际值、预测值和误差在流场中的分布如下图：\n",
    "\n",
    "![Unet_U.png](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/docs/mindflow/docs/source_zh_cn/data_driven/images/Unet_U.png)\n",
    "\n",
    "UNET2D backbone下，不同时刻速度V的实际值、预测值和误差在流场中的分布如下图：\n",
    "\n",
    "![Unet_V.png](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/docs/mindflow/docs/source_zh_cn/data_driven/images/Unet_V.png)\n",
    "\n",
    "FNO2D backbone下，不同时刻压力P的实际值、预测值和误差在流场中的分布如下图：\n",
    "\n",
    "![FNO_P.png](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/docs/mindflow/docs/source_zh_cn/data_driven/images/FNO_P.png)\n",
    "\n",
    "FNO2D backbone下，不同时刻速度U的实际值、预测值和误差在流场中的分布如下图：\n",
    "\n",
    "![FNO_U.png](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/docs/mindflow/docs/source_zh_cn/data_driven/images/FNO_U.png)\n",
    "\n",
    "FNO2D backbone下，不同时刻速度V的实际值、预测值和误差在流场中的分布如下图：\n",
    "\n",
    "![FNO_V.png](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/docs/mindflow/docs/source_zh_cn/data_driven/images/FNO_V.png)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:zys_test]",
   "language": "python",
   "name": "conda-env-zys_test-py"
  },
  "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.7.16"
  },
  "vscode": {
   "interpreter": {
    "hash": "b9063439a3781aed32d6b0dd4804a0c8b51ecec7893a0f31b99846bc91ef39eb"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}