{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# 2D & 3D Poisson\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/physics_driven/mindspore_poisson_geometry.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/physics_driven/mindspore_poisson_geometry.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/physics_driven/poisson_geometry.ipynb)\n", "\n", "This notebook requires MindSpore version >= 2.0.0 to support new APIs including: mindspore.jit, mindspore.jit_class, mindspore.jacrev." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Problem description\n", "\n", "This example demonstrates how to use PINNs to solve 2D and 3D Poisson equations in different geometries. The 2D equation is defined by\n", "\n", "$$\n", "\\Delta u = -\\sin(4\\pi x)\\sin(4\\pi y),\n", "$$\n", "\n", "and the 3D equation is given by\n", "\n", "$$\n", "\\Delta u = -\\sin(4\\pi x)\\sin(4\\pi y)\\sin(4\\pi z),\n", "$$\n", "\n", "The following functions satisfy the 2D and 3D equations respectively,\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", "If we set the Dirichlet boundary condition using these functions, they become the solution to the problems. In this example, we solve the 2D equation in rectangle, disk, triangle and pentagon. In the 3D case, we solve the problem in a tetrahedron, cylinder and cone.\n", "\n", "## Method\n", "\n", "MindSpore Flow solves the problem as follows:\n", "\n", "1. Dataset Construction.\n", "2. Model Construction.\n", "3. Poisson.\n", "4. Optimizer.\n", "5. Model Training.\n", "6. Model Evaluation." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The `poisson_cfg.yaml` file can be downloaded at [applications/physics_driven/poisson/poisson_cfg.yaml](https://gitee.com/mindspore/mindscience/blob/r0.3/MindFlow/applications/physics_driven/poisson/poisson_cfg.yaml), and the `src` package can be downloaded at [applications/physics_driven/poisson/src](https://gitee.com/mindspore/mindscience/tree/r0.3/MindFlow/applications/physics_driven/poisson/src)." ] }, { "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": [ "## Dataset Construction\n", "\n", "This example creates the dataset by ramdom sampling in the domain and on the boundaries (See ``src/dataset.py`` for the implementation). Set ``geom_name`` to select geometry, which can be 'triangle', 'pentagon', 'tetrahedron', 'cylinder', and '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": [ "## Model Construction\n", "\n", "This example adopts a MLP with 3 hidden layers with the following features:\n", "\n", "- Using $f(x) = x \\exp(-x^2/(2e))$ as the activation function.\n", "- Using weight normalization for the last layer.\n", "- Employing the ``HeUniform`` from ``mindspore`` for the initialization of all weights.\n", "\n", "See ``src/model.py`` for the details." ] }, { "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", "When applying ``mindflow`` to solving PDEs, we need to write a subclass of ``mindflow.PDEWithLloss`` to define the governing equation, boundary condition and loss function. In this example, we adopt the L2 loss in the domain and on the boundaries. Both losses are combined using the multitarget loss defined by ``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": [ "## Optimizer\n", "\n", "This example applies the Adam optimizer, and adopts the dynamic learning rate proposed by [Super-Convergence: Very Fast Training of Neural Networks Using Large Learning Rates](https://arxiv.org/abs/1708.07120). See ``src/lr_scheduler.py`` for the implementation of the dynamic learning rate." ] }, { "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": [ "## 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": 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": [ "## Model Evaluation\n", "\n", "We can use the following function to calculate the relative L2 error." ] }, { "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.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": 2 }