{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# 二维&三维Poisson问题\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/physics_driven/mindspore_poisson_geometry.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/physics_driven/mindspore_poisson_geometry.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/physics_driven/poisson_geometry.ipynb)\n", "\n", "本案例要求**MindSpore版本 >= 2.0.0**调用如下接口: *mindspore.jit，mindspore.jit_class，mindspore.jacrev*。" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 问题描述\n", "\n", "本案例演示如何利用PINNs在不同几何体下求解二维和三维泊松方程。二维泊松方程定义为\n", "\n", "$$\n", "\\Delta u = -\\sin(4\\pi x)\\sin(4\\pi y),\n", "$$\n", "\n", "而三维方程定义为\n", "\n", "$$\n", "\\Delta u = -\\sin(4\\pi x)\\sin(4\\pi y)\\sin(4\\pi z),\n", "$$\n", "\n", "很容易验证，以下函数分别满足二维和三维泊松方程\n", "\n", "$$\n", "u = \\frac{1}{32\\pi^2} \\sin(4\\pi x)\\sin(4\\pi y), \\\\\n", "u = \\frac{1}{48\\pi^2} \\sin(4\\pi x)\\sin(4\\pi y)\\sin(4\\pi z).\n", "$$\n", "\n", "如果在几何体边界按以上函数取狄利克雷边界条件，那么这些函数就是我们想要得到的解。因而，我们可以利用以上函数来验证结果。\n", "对于二维问题，本例演示在矩形，圆形，三角形和五边形区域求解方程，而对于三维问题，我们将在四面体，圆柱和圆锥区域内求解方程。\n", "\n", "## 技术路径\n", "\n", "MindSpore Flow求解该问题的具体流程如下：\n", "\n", "1. 创建数据集。\n", "2. 构建模型。\n", "3. 优化器。\n", "4. Poisson。\n", "5. 模型训练。\n", "6. 模型评估。" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "下述`poisson_cfg.yaml`配置文件可以在[applications/physics_driven/poisson/point_source/poisson_cfg.yaml](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/physics_driven/poisson/point_source/poisson_cfg.yaml)下载, `src`包可以在[applications/physics_driven/poisson/point_source/src](https://gitee.com/mindspore/mindscience/tree/master/MindFlow/applications/physics_driven/poisson/point_source/src)下载。\n" ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "import time\n", "\n", "import mindspore as ms\n", "from mindspore import nn, ops, jit\n", "from mindflow import load_yaml_config\n", "\n", "from src.model import create_model\n", "from src.lr_scheduler import OneCycleLR\n", "from src.dataset import create_dataset\n", "\n", "\n", "ms.set_context(mode=ms.GRAPH_MODE, save_graphs=False, device_target=\"GPU\")\n", "\n", "# Load config\n", "file_cfg = \"poisson_cfg.yaml\"\n", "config = load_yaml_config(file_cfg)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 创建数据集\n", "\n", "本案例在求解域及边值条件进行随机采样，生成训练数据集与测试数据集。具体方法见``src/dataset.py``。设值``geom_name``来选择几何体，可选择rectangle、disk、triangle、pentagon、tetrahedon、cylinder和cone。" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "geom_name = \"triangle\"\n", "ds_train, n_dim = create_dataset(geom_name, config)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 构建模型\n", "\n", "本案例采用带3个隐藏层的多层感知器，并带有以下特点:\n", "\n", "- 采用激活函数： $f(x) = x \\exp(-x^2/(2e))$\n", "\n", "- 最后一层线性层使用weight normalization。\n", "\n", "- 所有权重都采用``mindspore``的``HeUniform``初始化。\n", "\n", "具体定义见``src/model.py``。" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [], "source": [ "model = create_model(**config['model'][f'{n_dim}d'])" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Poisson\n", "\n", "在利用``mindflow``求解PDE时，我们需要写一个``mindflow.PDEWithLloss``的子类来定义控制方程，边界条件和损失函数。在求解区域内和边界上均采用L2损失，并利用``mindflow``的``MTLWeightedLossCell``多目标损失函数将两个损失结合起来。" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "poisson: sin(4*pi*x)*sin(4*pi*y) + Derivative(u(x, y), (x, 2)) + Derivative(u(x, y), (y, 2))\n", " Item numbers of current derivative formula nodes: 3\n", "bc: u(x, y) - sin(4*pi*x)*sin(4*pi*y)/(32*pi**2)\n", " Item numbers of current derivative formula nodes: 2\n" ] } ], "source": [ "import sympy\n", "from mindspore import numpy as ms_np\n", "from mindflow import PDEWithLoss, MTLWeightedLossCell, sympy_to_mindspore\n", "\n", "\n", "class Poisson(PDEWithLoss):\n", " \"\"\"Define the loss of the Poisson equation.\"\"\"\n", " def __init__(self, model, n_dim):\n", " if n_dim == 2:\n", " var_str = 'x y'\n", " elif n_dim == 3:\n", " var_str = 'x y z'\n", " else:\n", " raise ValueError(\"`n_dim` can only be 2 or 3.\")\n", " self.in_vars = sympy.symbols(var_str)\n", " self.out_vars = (sympy.Function('u')(*self.in_vars),)\n", " super(Poisson, self).__init__(model, self.in_vars, self.out_vars)\n", " self.bc_nodes = sympy_to_mindspore(self.bc(n_dim), self.in_vars, self.out_vars)\n", " self.loss_fn = MTLWeightedLossCell(num_losses=2)\n", "\n", " def pde(self):\n", " \"\"\"Define the gonvering equation.\"\"\"\n", " poisson = 0\n", " src_term = 1\n", " sym_u = self.out_vars[0]\n", " for var in self.in_vars:\n", " poisson += sympy.diff(sym_u, (var, 2))\n", " src_term *= sympy.sin(4*sympy.pi*var)\n", " poisson += src_term\n", " equations = {\"poisson\": poisson}\n", " return equations\n", "\n", " def bc(self, n_dim):\n", " \"\"\"Define the boundary condition.\"\"\"\n", " bc_term = 1\n", " for var in self.in_vars:\n", " bc_term *= sympy.sin(4*sympy.pi*var)\n", " bc_term *= 1/(16*n_dim*sympy.pi*sympy.pi)\n", " bc_eq = self.out_vars[0] - bc_term\n", " equations = {\"bc\": bc_eq}\n", " return equations\n", "\n", " def get_loss(self, pde_data, bc_data):\n", " \"\"\"Define the loss function.\"\"\"\n", " res_pde = self.parse_node(self.pde_nodes, inputs=pde_data)\n", " res_bc = self.parse_node(self.bc_nodes, inputs=bc_data)\n", " loss_pde = ms_np.mean(ms_np.square(res_pde[0]))\n", " loss_bc = ms_np.mean(ms_np.square(res_bc[0]))\n", " return self.loss_fn((loss_pde, loss_bc))\n", "\n", "# Create the problem\n", "problem = Poisson(model, n_dim)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 优化器\n", "\n", "本案例采用Adam优化器，并配合[Super-Convergence: Very Fast Training of Neural Networks Using Large Learning Rates](https://arxiv.org/abs/1708.07120)提出的动态学习率进行训练。动态学习率定义参见``src/lr_scheduler.py``。" ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [], "source": [ "n_epochs = 50\n", "\n", "params = model.trainable_params() + problem.loss_fn.trainable_params()\n", "steps_per_epoch = config['data']['domain']['size']//config['batch_size']\n", "learning_rate = OneCycleLR(total_steps=steps_per_epoch*n_epochs, **config['optimizer'])\n", "opt = nn.Adam(params, learning_rate=learning_rate)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 模型训练\n", "\n", "使用MindSpore>= 2.0.0的版本，可以使用函数式编程范式训练神经网络。" ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [], "source": [ "def train():\n", " # Create\n", " grad_fn = ms.value_and_grad(problem.get_loss, None, opt.parameters, has_aux=False)\n", "\n", " @jit\n", " def train_step(pde_data, bc_data):\n", " loss, grads = grad_fn(pde_data, bc_data)\n", " loss = ops.depend(loss, opt(grads))\n", " return loss\n", "\n", " def train_epoch(model, dataset, i_epoch):\n", " n_step = dataset.get_dataset_size()\n", " model.set_train()\n", " for i_step, (pde_data, bc_data) in enumerate(dataset):\n", " local_time_beg = time.time()\n", " loss = train_step(pde_data, bc_data)\n", "\n", " if i_step%50 == 0 or i_step + 1 == n_step:\n", " print(\"\\repoch: {}, loss: {:>f}, time elapsed: {:.1f}ms [{}/{}]\".format(\n", " i_epoch, float(loss), (time.time() - local_time_beg)*1000, i_step + 1, n_step))\n", "\n", " for i_epoch in range(n_epochs):\n", " train_epoch(model, ds_train, i_epoch)" ] }, { "cell_type": "code", "execution_count": 7, "metadata": { "scrolled": true }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "epoch: 0, loss: 1.527029, time elapsed: 12050.1ms [1/200]\n", "epoch: 0, loss: 1.468655, time elapsed: 52.4ms [51/200]\n", "epoch: 0, loss: 1.442717, time elapsed: 52.3ms [101/200]\n", "epoch: 0, loss: 1.430150, time elapsed: 52.4ms [151/200]\n", "epoch: 0, loss: 1.420228, time elapsed: 53.4ms [200/200]\n", "epoch: 1, loss: 1.419910, time elapsed: 53.0ms [1/200]\n", "epoch: 1, loss: 1.407040, time elapsed: 52.5ms [51/200]\n", "epoch: 1, loss: 1.386505, time elapsed: 52.4ms [101/200]\n", "epoch: 1, loss: 1.362307, time elapsed: 52.4ms [151/200]\n", "epoch: 1, loss: 1.349054, time elapsed: 52.5ms [200/200]\n", "epoch: 2, loss: 1.349143, time elapsed: 53.7ms [1/200]\n", "epoch: 2, loss: 1.336657, time elapsed: 52.7ms [51/200]\n", "epoch: 2, loss: 1.323158, time elapsed: 52.6ms [101/200]\n", "epoch: 2, loss: 1.307419, time elapsed: 52.9ms [151/200]\n", "epoch: 2, loss: 1.289993, time elapsed: 52.7ms [200/200]\n", "epoch: 3, loss: 1.289594, time elapsed: 53.5ms [1/200]\n", "epoch: 3, loss: 1.270476, time elapsed: 52.4ms [51/200]\n", "epoch: 3, loss: 1.246817, time elapsed: 52.6ms [101/200]\n", "epoch: 3, loss: 1.222093, time elapsed: 52.6ms [151/200]\n", "epoch: 3, loss: 1.194862, time elapsed: 52.3ms [200/200]\n", "epoch: 4, loss: 1.194533, time elapsed: 52.5ms [1/200]\n", "epoch: 4, loss: 1.164445, time elapsed: 52.6ms [51/200]\n", "epoch: 4, loss: 1.134136, time elapsed: 52.5ms [101/200]\n", "epoch: 4, loss: 1.100014, time elapsed: 52.6ms [151/200]\n", "epoch: 4, loss: 1.064941, time elapsed: 52.4ms [200/200]\n", "...\n", "epoch: 45, loss: 0.001281, time elapsed: 53.0ms [1/200]\n", "epoch: 45, loss: 0.001264, time elapsed: 52.6ms [51/200]\n", "epoch: 45, loss: 0.001263, time elapsed: 52.5ms [101/200]\n", "epoch: 45, loss: 0.001236, time elapsed: 52.6ms [151/200]\n", "epoch: 45, loss: 0.001237, time elapsed: 52.5ms [200/200]\n", "epoch: 46, loss: 0.001218, time elapsed: 52.7ms [1/200]\n", "epoch: 46, loss: 0.001209, time elapsed: 52.6ms [51/200]\n", "epoch: 46, loss: 0.001191, time elapsed: 52.6ms [101/200]\n", "epoch: 46, loss: 0.001202, time elapsed: 52.7ms [151/200]\n", "epoch: 46, loss: 0.001182, time elapsed: 52.9ms [200/200]\n", "epoch: 47, loss: 0.001174, time elapsed: 53.0ms [1/200]\n", "epoch: 47, loss: 0.001186, time elapsed: 52.7ms [51/200]\n", "epoch: 47, loss: 0.001182, time elapsed: 52.6ms [101/200]\n", "epoch: 47, loss: 0.001169, time elapsed: 52.8ms [151/200]\n", "epoch: 47, loss: 0.001172, time elapsed: 52.7ms [200/200]\n", "epoch: 48, loss: 0.001165, time elapsed: 52.7ms [1/200]\n", "epoch: 48, loss: 0.001168, time elapsed: 52.6ms [51/200]\n", "epoch: 48, loss: 0.001148, time elapsed: 52.5ms [101/200]\n", "epoch: 48, loss: 0.001159, time elapsed: 52.7ms [151/200]\n", "epoch: 48, loss: 0.001171, time elapsed: 52.5ms [200/200]\n", "epoch: 49, loss: 0.001156, time elapsed: 52.7ms [1/200]\n", "epoch: 49, loss: 0.001155, time elapsed: 52.6ms [51/200]\n", "epoch: 49, loss: 0.001148, time elapsed: 52.6ms [101/200]\n", "epoch: 49, loss: 0.001159, time elapsed: 52.9ms [151/200]\n", "epoch: 49, loss: 0.001153, time elapsed: 52.6ms [200/200]\n", "End-to-End total time: 584.182409286499 s\n" ] } ], "source": [ "time_beg = time.time()\n", "train()\n", "print(\"End-to-End total time: {} s\".format(time.time() - time_beg))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 模型评估\n", "\n", "可通过以下函数来计算模型的L2相对误差。" ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Relative L2 error (domain): 0.0310\n", "Relative L2 error (bc): 0.0833\n", "\n" ] } ], "source": [ "from eval import calculate_l2_error\n", "\n", "n_samps = 5000 # Number of test samples\n", "ds_test, _ = create_dataset(geom_name, config, n_samps)\n", "calculate_l2_error(model, ds_test, n_dim)" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "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.7.0" } }, "nbformat": 4, "nbformat_minor": 2 }