{
"cells": [
{
"attachments": {},
"cell_type": "markdown",
"metadata": {
"pycharm": {
"name": "#%% md\n"
}
},
"source": [
"# FNO for 1D Burgers\n",
"\n",
"[](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/r2.2/mindflow/en/data_driven/mindspore_burgers_FNO1D.ipynb) [](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/r2.2/mindflow/en/data_driven/mindspore_burgers_FNO1D.py) [](https://gitee.com/mindspore/docs/blob/r2.2/docs/mindflow/docs/source_en/data_driven/burgers_FNO1D.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 a specific discretization. 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 1-d Burgers' equation using Fourier neural operator."
]
},
{
"cell_type": "markdown",
"metadata": {
"pycharm": {
"name": "#%% md\n"
}
},
"source": [
"## Burgers' equation\n",
"\n",
"The 1-d Burgers’ equation is a non-linear PDE with various applications including modeling the one\n",
"dimensional flow of a viscous fluid. It takes the form\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",
"where $u$ is the velocity field, $u_0$ is the initial condition and $\\nu$ is the viscosity coefficient.\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Problem Description\n",
"\n",
"We aim to learn the operator mapping the initial condition to the solution at time one:\n",
"\n",
"$$\n",
"u_0 \\mapsto u(\\cdot, 1)\n",
"$$"
]
},
{
"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 Fourier Neural Operator model architecture is shown below. In the figure, $w_0(x)$ denotes the initial vorticity, the high-dimensional mapping of the input vectors is realized by Lifting Layer, and then the mapping result is used as the input of Fourier Layer to carry out the non-linear transformation of the frequency-domain information, and finally, the mapping result of the transformation is mapped to the final prediction result $w_1(x)$ by Decoding Layer.\n",
"\n",
"The Fourier Neural Operator consists of the Lifting Layer, Fourier Layers, and the Decoding Layer.\n",
"\n",
"\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",
""
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {
"pycharm": {
"name": "#%%\n"
}
},
"outputs": [],
"source": [
"import os\n",
"import time\n",
"import numpy as np\n",
"import mindspore as ms\n",
"\n",
"from mindspore.amp import DynamicLossScaler, auto_mixed_precision, all_finite\n",
"from mindspore import nn, Tensor, set_seed, ops, data_sink, jit, save_checkpoint\n",
"from mindspore import dtype as mstype\n",
"from mindflow import FNO1D, RelativeRMSELoss, load_yaml_config, get_warmup_cosine_annealing_lr\n",
"from mindflow.pde import UnsteadyFlowWithLoss"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"The following `src` pacakage can be downloaded in [applications/data_driven/burgers/fno1d/src](https://gitee.com/mindspore/mindscience/tree/r0.6/MindFlow/applications/data_driven/burgers/fno1d/src)."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {
"pycharm": {
"name": "#%%\n"
}
},
"outputs": [],
"source": [
"from src.dataset import create_training_dataset\n",
"\n",
"set_seed(0)\n",
"np.random.seed(0)\n",
"\n",
"ms.set_context(mode=ms.GRAPH_MODE, device_target=\"GPU\", device_id=5)\n",
"use_ascend = ms.get_context(attr_key='device_target') == \"Ascend\""
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"You can get parameters of model, data and optimizer from [config](https://gitee.com/mindspore/mindscience/blob/r0.6/MindFlow/applications/data_driven/burgers/fno1d/configs/fno1d.yaml)."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {
"pycharm": {
"name": "#%%\n"
}
},
"outputs": [],
"source": [
"config = load_yaml_config('fno1d.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/burgers/dataset](https://download.mindspore.cn/mindscience/mindflow/dataset/applications/data_driven/burgers/dataset/).\n",
"\n",
"In this case, training datasets and test datasets are generated according to Zongyi Li's dataset 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 $u_0(x)$ is generated according to periodic boundary conditions:\n",
"\n",
"$$\n",
"u_0 \\sim \\mu, \\mu=\\mathcal{N}\\left(0,625(-\\Delta+25 I)^{-2}\\right)\n",
"$$\n",
"\n",
"We set the viscosity to $\\nu=0.1$ and solve the equation using a split step method where the heat equation part is solved exactly in Fourier space then the non-linear part is advanced, again in Fourier space, using a very fine forward Euler method. The number of samples in the training set is 1000, and the number of samples in the test set is 200.\n"
]
},
{
"cell_type": "code",
"execution_count": 4,
"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",
"test_input = Tensor(np.expand_dims(test_input, -2), mstype.float32)\n",
"test_label = Tensor(np.expand_dims(test_label, -2), mstype.float32)"
]
},
{
"attachments": {},
"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 `FNO1D.fc0` in the case, and maps the output data $x$ to the high dimension;\n",
"\n",
"- Multi-layer Fourier Layer corresponds to the `FNO1D.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 `FNO1D.fc1` and `FNO1D.fc2` in the case to obtain the final predictive value.\n",
"\n",
"The initialization of the model based on the network above, parameters can be modified in [configuration file](https://gitee.com/mindspore/mindscience/blob/r0.6/MindFlow/applications/data_driven/burgers/fno1d/configs/fno1d.yaml)."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {
"pycharm": {
"name": "#%%\n"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"name:FNO1D_in_channels:1_out_channels:1_resolution:1024_modes:16_width:64_depth:4\n"
]
}
],
"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",
" depths=model_params[\"depth\"])\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)\n",
"print(model_name)"
]
},
{
"cell_type": "markdown",
"metadata": {
"pycharm": {
"name": "#%% md\n"
}
},
"source": [
"## Optimizer and Loss Function\n",
"\n",
"Use the relative root mean square error as the network training loss function:\n"
]
},
{
"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=1)\n",
"optimizer = nn.Adam(model.trainable_params(), learning_rate=Tensor(lr))\n",
"\n",
"if use_ascend:\n",
" loss_scaler = DynamicLossScaler(1024, 2, 100)\n",
" auto_mixed_precision(model, 'O1')\n",
"else:\n",
" loss_scaler = None"
]
},
{
"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. `MindSpore Flow` provides a training interface for unsteady problems `UnsteadyFlowWithLoss` for model training and evaluation."
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {
"pycharm": {
"name": "#%%\n"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"./summary/name:FNO1D_in_channels:1_out_channels:1_resolution:1024_modes:16_width:64_depth:4\n",
"epoch: 1, time elapsed: 21747.305870056152ms, loss: 2.167046070098877\n",
"epoch: 2, time elapsed: 5525.397539138794ms, loss: 0.5935954451560974\n",
"epoch: 3, time elapsed: 5459.984540939331ms, loss: 0.7349425554275513\n",
"epoch: 4, time elapsed: 4948.82869720459ms, loss: 0.6338694095611572\n",
"epoch: 5, time elapsed: 5571.3865756988525ms, loss: 0.3174982964992523\n",
"epoch: 6, time elapsed: 5712.041616439819ms, loss: 0.3099440038204193\n",
"epoch: 7, time elapsed: 5218.639135360718ms, loss: 0.3117891848087311\n",
"epoch: 8, time elapsed: 4819.460153579712ms, loss: 0.1810857653617859\n",
"epoch: 9, time elapsed: 4968.810081481934ms, loss: 0.1386510729789734\n",
"epoch: 10, time elapsed: 4849.36785697937ms, loss: 0.2102256715297699\n",
"================================Start Evaluation================================\n",
"mean rms_error: 0.027940063\n",
"=================================End Evaluation=================================\n",
"...\n",
"epoch: 91, time elapsed: 4398.104429244995ms, loss: 0.019643772393465042\n",
"epoch: 92, time elapsed: 5479.56109046936ms, loss: 0.0641067773103714\n",
"epoch: 93, time elapsed: 5549.5476722717285ms, loss: 0.02199840545654297\n",
"epoch: 94, time elapsed: 6238.730907440186ms, loss: 0.024467874318361282\n",
"epoch: 95, time elapsed: 5434.457778930664ms, loss: 0.025712188333272934\n",
"epoch: 96, time elapsed: 6481.106281280518ms, loss: 0.02247200347483158\n",
"epoch: 97, time elapsed: 6303.435325622559ms, loss: 0.026637140661478043\n",
"epoch: 98, time elapsed: 5162.56856918335ms, loss: 0.030040305107831955\n",
"epoch: 99, time elapsed: 5364.72225189209ms, loss: 0.02589748054742813\n",
"epoch: 100, time elapsed: 5902.378797531128ms, loss: 0.028599221259355545\n",
"================================Start Evaluation================================\n",
"mean rms_error: 0.0037017763\n",
"=================================End Evaluation=================================\n"
]
}
],
"source": [
"problem = UnsteadyFlowWithLoss(model, loss_fn=RelativeRMSELoss(), data_format=\"NHWTC\")\n",
"\n",
"summary_dir = os.path.join(config[\"summary_dir\"], model_name)\n",
"print(summary_dir)\n",
"\n",
"def forward_fn(data, label):\n",
" loss = problem.get_loss(data, label)\n",
" return loss\n",
"\n",
"grad_fn = ms.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=False)\n",
"\n",
"@jit\n",
"def train_step(data, label):\n",
" loss, grads = grad_fn(data, 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",
" return loss\n",
"\n",
"sink_process = data_sink(train_step, train_dataset, 1)\n",
"summary_dir = os.path.join(config[\"summary_dir\"], model_name)\n",
"\n",
"for epoch in range(1, config[\"epochs\"] + 1):\n",
" model.set_train()\n",
" local_time_beg = time.time()\n",
" for _ in range(steps_per_epoch):\n",
" cur_loss = sink_process()\n",
" print(\"epoch: {}, time elapsed: {}ms, loss: {}\".format(epoch, (time.time() - local_time_beg) * 1000, cur_loss.asnumpy()))\n",
"\n",
" if epoch % config['eval_interval'] == 0:\n",
" model.set_train(False)\n",
" print(\"================================Start Evaluation================================\")\n",
" rms_error = problem.get_loss(test_input, test_label)/test_input.shape[0]\n",
" print(\"mean rms_error:\", rms_error)\n",
" print(\"=================================End Evaluation=================================\")\n",
" ckpt_dir = os.path.join(summary_dir, \"ckpt\")\n",
" if not os.path.exists(ckpt_dir):\n",
" os.makedirs(ckpt_dir)\n",
" save_checkpoint(model, os.path.join(ckpt_dir, model_params[\"name\"] + '_epoch' + str(epoch)))"
]
}
],
"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
}