{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# FNO for 2D Navier-Stokes\n",
    "\n",
    "[![DownloadNotebook](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.1/resource/_static/logo_notebook_en.svg)](https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/r2.1/mindflow/en/data_driven/mindspore_navier_stokes_FNO2D.ipynb) [![DownloadCode](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.1/resource/_static/logo_download_code_en.svg)](https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/r2.1/mindflow/en/data_driven/mindspore_navier_stokes_FNO2D.py) [![ViewSource](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.1/resource/_static/logo_source_en.svg)](https://gitee.com/mindspore/docs/blob/r2.1/docs/mindflow/docs/source_en/data_driven/navier_stokes_FNO2D.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Overview\n",
    "\n",
    "Computational fluid dynamics is one of the most important techniques in the field of fluid mechanics in the 21st century. The flow analysis, prediction and control can be realized by solving the governing equations of fluid mechanics by numerical method. Traditional finite element method (FEM) and finite difference method (FDM) are inefficient because of the complex simulation process (physical modeling, meshing, numerical discretization, iterative solution, etc.) and high computing costs. Therefore, it is necessary to improve the efficiency of fluid simulation with AI.\n",
    "\n",
    "Machine learning methods provide a new paradigm for scientific computing by providing a fast solver similar to traditional methods. Classical neural networks learn mappings between finite dimensional spaces and can only learn solutions related to specific discretizations. Different from traditional neural networks, Fourier Neural Operator (FNO) is a new deep learning architecture that can learn mappings between infinite-dimensional function spaces. It directly learns mappings from arbitrary function parameters to solutions to solve a class of partial differential equations.  Therefore, it has a stronger generalization capability. More information can be found in the paper, [Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/abs/2010.08895).\n",
    "\n",
    "This tutorial describes how to solve the Navier-Stokes equation using Fourier neural operator."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Navier-Stokes equation\n",
    "\n",
    "Navier-Stokes equation is a classical equation in computational fluid dynamics. It is a set of partial differential equations describing the conservation of fluid momentum, called N-S equation for short. Its vorticity form in two-dimensional incompressible flows is as follows:\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",
    "where $u$ is the velocity field, $w=\\nabla \\times u$ is the vorticity, $w_0(x)$ is the initial vorticity, $\\nu$ is the viscosity coefficient, $f(x)$ is the forcing function.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem Description\n",
    "\n",
    "We aim to solve two-dimensional incompressible N-S equation by learning the operator mapping from each time step to the next time step:\n",
    "\n",
    "$$\n",
    "w_t \\mapsto w(\\cdot, t+1)\n",
    "$$"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Technology Path\n",
    "\n",
    "MindSpore Flow solves the problem as follows:\n",
    "\n",
    "1. Training Dataset Construction.\n",
    "2. Model Construction.\n",
    "3. Optimizer and Loss Function.\n",
    "4. Model Training."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fourier Neural Operator\n",
    "\n",
    "The following figure shows the architecture of the Fourier Neural Operator model. In the figure, $w_0(x)$ represents the initial vorticity.  The input vector is lifted to higher dimension channel space by the lifting layer. Then the mapping result is used as the input of the Fourier layer to perform nonlinear transformation of the frequency domain information. Finally, the decoding layer maps the transformation result to the final prediction result $w_1(x)$.\n",
    "\n",
    "The Fourier Neural Operator consists of the lifting Layer, Fourier Layers, and the decoding Layer.\n",
    "\n",
    "![Fourier Neural Operator model structure](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.1/docs/mindflow/docs/source_en/data_driven/images/FNO.png)\n",
    "\n",
    "Fourier layers: Start from input V. On top: apply the Fourier transform $\\mathcal{F}$; a linear transform R on the lower Fourier modes and filters out the higher modes; then apply the inverse Fourier transform $\\mathcal{F}^{-1}$. On the bottom: apply a local linear transform W.  Finally, the Fourier Layer output vector is obtained through the activation function.\n",
    "\n",
    "![Fourier Layer structure](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.1/docs/mindflow/docs/source_en/data_driven/images/FNO-2.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import numpy as np\n",
    "\n",
    "import mindspore\n",
    "from mindspore import nn, ops, Tensor, jit, set_seed"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following `src` pacakage can be downloaded in [applications/data_driven/navier_stokes_fno/src](https://gitee.com/mindspore/mindscience/tree/r0.3/MindFlow/applications/data_driven/navier_stokes_fno/src). Parameters can be modified in [config](https://gitee.com/mindspore/mindscience/blob/r0.3/MindFlow/applications/data_driven/navier_stokes_fno/navier_stokes_2d.yaml)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from mindflow.cell import FNO2D\n",
    "from mindflow.common import get_warmup_cosine_annealing_lr\n",
    "from mindflow.loss import RelativeRMSELoss\n",
    "from mindflow.utils import load_yaml_config\n",
    "from mindflow.pde import UnsteadyFlowWithLoss\n",
    "from src import calculate_l2_error, create_training_dataset\n",
    "\n",
    "set_seed(0)\n",
    "np.random.seed(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# set context for training: using graph mode for high performance training with GPU acceleration\n",
    "mindspore.set_context(mode=mindspore.GRAPH_MODE, device_target='GPU', device_id=2)\n",
    "use_ascend = mindspore.get_context(attr_key='device_target') == \"Ascend\"\n",
    "config = load_yaml_config('navier_stokes_2d.yaml')\n",
    "data_params = config[\"data\"]\n",
    "model_params = config[\"model\"]\n",
    "optimizer_params = config[\"optimizer\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Training Dataset Construction\n",
    "\n",
    "Download the training and test dataset: [data_driven/navier_stokes/dataset](https://download.mindspore.cn/mindscience/mindflow/dataset/applications/data_driven/navier_stokes/dataset/) .\n",
    "\n",
    "In this case, training data sets and test data sets are generated according to Zongyi Li's data set in [Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/pdf/2010.08895.pdf) . The settings are as follows:\n",
    "\n",
    "The initial condition $w_0(x)$ is generated according to periodic boundary conditions:\n",
    "\n",
    "$$\n",
    "w_0 \\sim \\mu, \\mu=\\mathcal{N}\\left(0,7^{3 / 2}(-\\Delta+49 I)^{-2.5}\\right)\n",
    "$$\n",
    "\n",
    "The forcing function is defined as:\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",
    "We use a time-step of 1e-4 for the Crank–Nicolson scheme in the data-generated process where we record the solution every t = 1 time units.  All data are generated on a 256 × 256 grid and are downsampled to 64 × 64.  In this case, the viscosity coefficient $\\nu=1e-5$, the number of samples in the training set is 19000, and the number of samples in the test set is 3800."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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": [
    "## Model Construction\n",
    "\n",
    "The network is composed of 1 lifting layer, multiple Fourier layers and 1 decoding layer:\n",
    "\n",
    "- The Lifting layer corresponds to the `FNO2D.fc0` in the case, and maps the output data $x$ to the high dimension;\n",
    "\n",
    "- Multi-layer Fourier Layer corresponds to the `FNO2D.fno_seq` in the case. Discrete Fourier transform is used to realize the conversion between time domain and frequency domain;\n",
    "\n",
    "- The Decoding layer corresponds to `FNO2D.fc1` and `FNO2D.fc2` in the case to obtain the final predictive value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "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",
    "              depths=model_params[\"depth\"]\n",
    "              )\n",
    "\n",
    "model_params_list = []\n",
    "for k, v in model_params.items():\n",
    "    model_params_list.append(f\"{k}-{v}\")\n",
    "model_name = \"_\".join(model_params_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Optimizer and Loss Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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",
    "\n",
    "problem = UnsteadyFlowWithLoss(model, loss_fn=RelativeRMSELoss(), data_format=\"NHWC\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Model Training\n",
    "\n",
    "With **MindSpore version >= 2.0.0**, we can use the functional programming for training neural networks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def train():\n",
    "    if use_ascend:\n",
    "        from mindspore.amp import DynamicLossScaler, auto_mixed_precision, all_finite\n",
    "        loss_scaler = DynamicLossScaler(1024, 2, 100)\n",
    "        auto_mixed_precision(model, 'O3')\n",
    "\n",
    "    def forward_fn(train_inputs, train_label):\n",
    "        loss = problem.get_loss(train_inputs, train_label)\n",
    "        if use_ascend:\n",
    "            loss = loss_scaler.scale(loss)\n",
    "        return loss\n",
    "\n",
    "    grad_fn = mindspore.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=False)\n",
    "\n",
    "    @jit\n",
    "    def train_step(train_inputs, train_label):\n",
    "        loss, grads = grad_fn(train_inputs, train_label)\n",
    "        if use_ascend:\n",
    "            loss = loss_scaler.unscale(loss)\n",
    "            if all_finite(grads):\n",
    "                grads = loss_scaler.unscale(grads)\n",
    "                loss = ops.depend(loss, optimizer(grads))\n",
    "        else:\n",
    "            loss = ops.depend(loss, optimizer(grads))\n",
    "        return loss\n",
    "\n",
    "    sink_process = mindspore.data_sink(train_step, train_dataset, sink_size=1)\n",
    "    summary_dir = os.path.join(config[\"summary_dir\"], model_name)\n",
    "\n",
    "    for cur_epoch in range(optimizer_params[\"train_epochs\"]):\n",
    "        local_time_beg = time.time()\n",
    "        model.set_train()\n",
    "        cur_loss = 0.0\n",
    "        for _ in range(steps_per_epoch):\n",
    "            cur_loss = sink_process()\n",
    "\n",
    "        print(\"epoch: %s, loss is %s\" % (cur_epoch + 1, cur_loss), flush=True)\n",
    "        local_time_end = time.time()\n",
    "        epoch_seconds = (local_time_end - local_time_beg) * 1000\n",
    "        step_seconds = epoch_seconds / steps_per_epoch\n",
    "        print(\"Train epoch time: {:5.3f} ms, per step time: {:5.3f} ms\".format\n",
    "              (epoch_seconds, step_seconds), flush=True)\n",
    "\n",
    "        if (cur_epoch + 1) % config[\"save_checkpoint_epoches\"] == 0:\n",
    "            ckpt_dir = os.path.join(summary_dir, \"ckpt\")\n",
    "            if not os.path.exists(ckpt_dir):\n",
    "                os.makedirs(ckpt_dir)\n",
    "            mindspore.save_checkpoint(model, os.path.join(ckpt_dir, model_params[\"name\"]))\n",
    "\n",
    "        if (cur_epoch + 1) % config['eval_interval'] == 0:\n",
    "            calculate_l2_error(model, test_input, test_label, config[\"test_batch_size\"])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1, loss is 1.7631323\n",
      "Train epoch time: 50405.954 ms, per step time: 50.406 ms\n",
      "epoch: 2, loss is 1.9283392\n",
      "Train epoch time: 36591.429 ms, per step time: 36.591 ms\n",
      "epoch: 3, loss is 1.4265916\n",
      "Train epoch time: 35085.079 ms, per step time: 35.085 ms\n",
      "epoch: 4, loss is 1.8609437\n",
      "Train epoch time: 34407.280 ms, per step time: 34.407 ms\n",
      "epoch: 5, loss is 1.5222052\n",
      "Train epoch time: 34596.965 ms, per step time: 34.597 ms\n",
      "epoch: 6, loss is 1.3424721\n",
      "Train epoch time: 33847.209 ms, per step time: 33.847 ms\n",
      "epoch: 7, loss is 1.607729\n",
      "Train epoch time: 33106.981 ms, per step time: 33.107 ms\n",
      "epoch: 8, loss is 1.3308442\n",
      "Train epoch time: 33051.339 ms, per step time: 33.051 ms\n",
      "epoch: 9, loss is 1.3169765\n",
      "Train epoch time: 33901.816 ms, per step time: 33.902 ms\n",
      "epoch: 10, loss is 1.4149593\n",
      "Train epoch time: 33908.748 ms, per step time: 33.909 ms\n",
      "================================Start Evaluation================================\n",
      "mean rel_rmse_error: 0.15500953359901906\n",
      "=================================End Evaluation=================================\n",
      "...\n",
      "epoch: 141, loss is 0.777328\n",
      "Train epoch time: 32549.911 ms, per step time: 32.550 ms\n",
      "epoch: 142, loss is 0.7008966\n",
      "Train epoch time: 32522.572 ms, per step time: 32.523 ms\n",
      "epoch: 143, loss is 0.72377646\n",
      "Train epoch time: 32566.685 ms, per step time: 32.567 ms\n",
      "epoch: 144, loss is 0.72175145\n",
      "Train epoch time: 32435.932 ms, per step time: 32.436 ms\n",
      "epoch: 145, loss is 0.6235678\n",
      "Train epoch time: 32463.707 ms, per step time: 32.464 ms\n",
      "epoch: 146, loss is 0.9351083\n",
      "Train epoch time: 32448.413 ms, per step time: 32.448 ms\n",
      "epoch: 147, loss is 0.9283789\n",
      "Train epoch time: 32472.401 ms, per step time: 32.472 ms\n",
      "epoch: 148, loss is 0.7655642\n",
      "Train epoch time: 32604.642 ms, per step time: 32.605 ms\n",
      "epoch: 149, loss is 0.7233772\n",
      "Train epoch time: 32649.832 ms, per step time: 32.650 ms\n",
      "epoch: 150, loss is 0.86825275\n",
      "Train epoch time: 32589.243 ms, per step time: 32.589 ms\n",
      "================================Start Evaluation================================\n",
      "mean rel_rmse_error: 0.07437102290522307\n",
      "=================================End Evaluation=================================\n",
      "predict total time: 15.212349653244019 s\n"
     ]
    }
   ],
   "source": [
    "train()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "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 (default, Nov 15 2020, 14:28:56) \n[GCC 7.3.0]"
  },
  "vscode": {
   "interpreter": {
    "hash": "e0054f2ca9303a347df1df0a145044d583b96ff24b3a906b6fe09c25d291b073"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}