{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Multi-timestep Complicated Flow Field Prediction with Transonic Airfoil under Data Driven Mode(with Two Kinds of Backbones: FNO2D and UNET2D)\n",
    "\n",
    "[![DownloadNotebook](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.5.0/resource/_static/logo_notebook_en.svg)](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/r2.5.0/mindflow/en/data_driven/mindspore_2D_unsteady.ipynb) [![DownloadCode](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.5.0/resource/_static/logo_download_code_en.svg)](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/r2.5.0/mindflow/en/data_driven/mindspore_2D_unsteady.py) [![View Source On Gitee](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.5.0/resource/_static/logo_source_en.svg)](https://gitee.com/mindspore/docs/blob/r2.5.0/docs/mindflow/docs/source_en/data_driven/2D_unsteady.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Introduction**\n",
    "\n",
    "High precision unsteady flow simulation is among the key topics in computational fluid dynamics(CFD), with a wide range of application scenarios and broad market value. However, traditional methods encountered problems such as long time-consumption, poor precision and hard convergence. AI methods provides a new perspective to explore the evolution mechanism of flow fields.\n",
    "\n",
    "The current application provides an end-to-end solution for predicting unsteady and complex flow fields in a two-dimensional transonic airfoil scenario. Two network backbones, Fourier Neural Operator (FNO) and Unet, are constructed in order to stably predict the flow field for subsequent *m* time steps based on the input flow field of *k* time steps, while ensuring a certain level of accuracy.\n",
    "\n",
    "The Mach number of the incoming flow reached *Ma*=0.73 in the current application. Thus it can verify the effectiveness of deep learning methods in predicting unsteady flow fields under multiple physical parameter changes in complex flow structures such as shock waves.\n",
    "\n",
    "![img_1.png](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.5.0/docs/mindflow/docs/source_en/data_driven/images/img_1_unsteady.png)\n",
    "\n",
    "## Problem Description\n",
    "\n",
    "We aim to predict next *k* steps of two-dimensional compressible unsteady flow by learning from past *m* steps flow field in the current application:\n",
    "\n",
    "$$\n",
    "u_{[t_0\\sim t_{k-1}]} \\mapsto u_{[t_k\\sim t_{k+m}]}\n",
    "$$\n",
    "\n",
    "## Technology Path\n",
    "\n",
    "The solution steps to the problem is presented as follows:\n",
    "\n",
    "1. Training Dataset Construction.\n",
    "2. Model Construction.\n",
    "3. Optimizer and Loss Function.\n",
    "4. Model Training.\n",
    "\n",
    "## **Preparation**\n",
    "\n",
    "Before practice, ensure that MindSpore of suitable version has been correctly installed. If not, you can run the following command:\n",
    "\n",
    "* [MindSpore installation page](https://www.mindspore.cn/install) Install MindSpore.\n",
    "\n",
    "## 2D Airfoil Unsteady Flow Field Prediction Implementation\n",
    "\n",
    "The implementation of 2D airfoil unsteady flow field prediction consists of the following seven steps:\n",
    "\n",
    "1. Configure network and training parameters\n",
    "2. Datasets preparation\n",
    "3. Model building\n",
    "4. Loss function and optimizer\n",
    "5. Train function\n",
    "6. Model training\n",
    "7. Result visualization"
   ]
  },
  {
   "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": [
    "## Configure Network and Training Parameters\n",
    "\n",
    "Load four types of parameters from the configuration file, which are model-related parameters (model), data-related parameters (data), optimizer-related parameters (optimizer) and callback-related parameters(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": [
    "## Datasets Preparation\n",
    "\n",
    "Dataset download link: [data_driven/airfoil/2D_unsteady](https://download.mindspore.cn/mindscience/mindflow/dataset/applications/data_driven/airfoil/2D_unsteady/)\n",
    "\n",
    "The data is numpy npz file with a dimension(*t*, *H*, *W*, *C*) of (9997, 128, 128, 3), where *t* is the total time steps, *H* and *W* are flow field resolution, *C* is the channel number(3 channels are velocity *U*, *V* and pressure *P*, respectively).\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": [
    "## Model Building\n",
    "\n",
    "The model is built by calling initial_model() function. Before calling function, loss_scaler and compute_dtype should be customized according to the corresponding hardware.\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": [
    "## Loss Function and optimizer\n",
    "\n",
    "In the current application, RelativeRMSELoss and AdamWeightDecay are selected to be the loss function and optimizer, respectively. Meanwhile the warmup_cosine_annealing_lr strategy is selected as the learning rate scheduler. Users could also customize their own loss function and learning rate scheduler."
   ]
  },
  {
   "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": [
    "## Train Function\n",
    "\n",
    "With **MindSpore>= 2.0.0**, users could train neural networks using functional programming paradigms, and single-step training functions are decorated with jit. The data_sink function is used to transfer the step-by-step training function and training dataset."
   ]
  },
  {
   "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": [
    "## Model Training\n",
    "\n",
    "Inference is performed during model training. Users could directly load the test data set, output the inference precision and save visualization results on the test set every test_interval epochs. Simultaneously, users could save checkpoint file every save_checkpoint_interval steps."
   ]
  },
  {
   "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": [
    "## **Result visualization**\n",
    "\n",
    "The label value, prediction value and error of pressure P distribution in the flow field of different time steps under UNET2D backbone:\n",
    "\n",
    "![Unet_P.png](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.5.0/docs/mindflow/docs/source_en/data_driven/images/Unet_P.png)\n",
    "\n",
    "The label value, prediction value and error of velocity U distribution in the flow field of different time steps under UNET2D backbone:\n",
    "\n",
    "![Unet_U.png](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.5.0/docs/mindflow/docs/source_en/data_driven/images/Unet_U.png)\n",
    "\n",
    "The label value, prediction value and error of velocity V distribution in the flow field of different time steps under UNET2D backbone:\n",
    "\n",
    "![Unet_V.png](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.5.0/docs/mindflow/docs/source_en/data_driven/images/Unet_V.png)\n",
    "\n",
    "The label value, prediction value and error of pressure P distribution in the flow field of different time steps under FNO2D backbone:\n",
    "\n",
    "![FNO_P.png](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.5.0/docs/mindflow/docs/source_en/data_driven/images/FNO_P.png)\n",
    "\n",
    "The label value, prediction value and error of velocity U distribution in the flow field of different time steps under FNO2D backbone:\n",
    "\n",
    "![FNO_U.png](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.5.0/docs/mindflow/docs/source_en/data_driven/images/FNO_U.png)\n",
    "\n",
    "The label value, prediction value and error of velocity V distribution in the flow field of different time steps under FNO2D backbone:\n",
    "\n",
    "![FNO_V.png](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.5.0/docs/mindflow/docs/source_en/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
}