helmholtz.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e5450a52-a8e9-4616-8cd0-41cfac7c097f",
   "metadata": {
    "incorrectly_encoded_metadata": "slideshow={\"slide_type\": \"slide\"} tags=[] jp-MarkdownHeadingCollapsed=true",
    "lines_to_next_cell": 0,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "# The Helmholtz equation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7f39062-2993-455a-9c35-207f3fd06234",
   "metadata": {},
   "source": [
    "In this tutorial, we will learn:\n",
    "\n",
    " - How to solve PDEs with complex-valued fields,\n",
    " - How to import and use high-order meshes from Gmsh,\n",
    " - How to use high order discretizations,\n",
    " - How to use UFL expressions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca874c12-82db-46d4-b826-b6aea843bf06",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "## Problem statement\n",
    "We will solve the Helmholtz equation subject to a first order absorbing boundary condition:\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "\\Delta u + k^2 u &= 0 && \\text{in } \\Omega,\\\\\n",
    "\\nabla u \\cdot \\mathbf{n} - \\mathrm{j}ku &= g && \\text{on } \\partial\\Omega,\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "where $k$ is a piecewise constant wavenumber, $\\mathrm{j}=\\sqrt{-1}$, and $g$ is the boundary source term computed as\n",
    "\n",
    "$$g = \\nabla u_\\text{inc} \\cdot \\mathbf{n} - \\mathrm{j}ku_\\text{inc}$$\n",
    "\n",
    "for an incoming plane wave $u_\\text{inc}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "433d9338-15c9-4ed5-9fa0-c83281da39ee",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from mpi4py import MPI\n",
    "\n",
    "import dolfinx.fem.petsc\n",
    "import ufl"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de959fc8-6fce-4ecb-83ab-b2ca6edb0475",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": []
   },
   "source": [
    "This example is designed to be executed with complex-valued degrees of freedom. To be able to solve this problem, we use the complex build of PETSc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "257a7024-7fb4-4d10-9dee-eacd607bea97",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using <class 'numpy.complex128'>.\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "from petsc4py import PETSc\n",
    "\n",
    "if not np.issubdtype(PETSc.ScalarType, np.complexfloating):\n",
    "    print(\"This tutorial requires complex number support\")\n",
    "    sys.exit(0)\n",
    "else:\n",
    "    print(f\"Using {PETSc.ScalarType}.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da41ebf9-34dd-4238-bbcc-393648210608",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": []
   },
   "source": [
    "## Defining model parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "dd8366af-6cc1-4fee-80b0-a499b7c4cb8f",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# wavenumber in free space (air)\n",
    "k0 = 10 * np.pi\n",
    "\n",
    "# Corresponding wavelength\n",
    "lmbda = 2 * np.pi / k0\n",
    "\n",
    "# Polynomial degree\n",
    "degree = 6\n",
    "\n",
    "# Mesh order\n",
    "mesh_order = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bfd09b8-9a20-4294-8ab1-781ab2c08cd8",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "## Interfacing with GMSH"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cd01dce-8a93-4753-a7ea-0eac1fed9ada",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": []
   },
   "source": [
    "We will use [Gmsh](http://gmsh.info/) to generate the computational domain (mesh) for this example. As long as Gmsh has been installed (including its Python API), DOLFINx supports direct input of Gmsh models (generated on one process). DOLFINx will then in turn distribute the mesh over all processes in the communicator passed to `dolfinx.io.gmshio.model_to_mesh`.\n",
    "\n",
    "The function `generate_mesh` creates a Gmsh model on rank 0 of `MPI.COMM_WORLD`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "eec83fe0-eced-437c-b4de-116792338d07",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Info    : Reading 'domain.msh'...\n",
      "Info    : 15 entities\n",
      "Info    : 2985 nodes\n",
      "Info    : 1444 elements\n",
      "Info    : Done reading 'domain.msh'\n"
     ]
    }
   ],
   "source": [
    "from dolfinx.io import gmshio\n",
    "from mesh_generation import generate_mesh\n",
    "\n",
    "# MPI communicator\n",
    "comm = MPI.COMM_WORLD\n",
    "\n",
    "file_name = \"domain.msh\"\n",
    "generate_mesh(file_name, lmbda, order=mesh_order)\n",
    "mesh, cell_tags, _ = gmshio.read_from_msh(file_name, comm,\n",
    "                                          rank=0, gdim=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3566dea-f954-484c-b41d-d6787c44f1af",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "## Material parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb727dcf-b350-4707-8ba6-63eb5ce7f739",
   "metadata": {
    "slideshow": {
     "slide_type": "notes"
    },
    "tags": []
   },
   "source": [
    "In this problem, the wave number in the different parts of the domain depends on cell markers, inputted through `cell_tags`.\n",
    "We use the fact that a discontinuous Lagrange space of order 0 (cell-wise constants) has a one-to-one mapping with the cells local to the process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d6b41f72-c9ae-4111-b268-0a7e53c7d459",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "W = dolfinx.fem.functionspace(mesh, (\"DG\", 0))\n",
    "k = dolfinx.fem.Function(W)\n",
    "k.x.array[:] = k0\n",
    "k.x.array[cell_tags.find(1)] = 3 * k0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "58dd0acb-43f3-4aa3-9fcf-441793cfab3d",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "import pyvista\n",
    "import matplotlib.pyplot as plt\n",
    "from dolfinx.plot import vtk_mesh\n",
    "pyvista.start_xvfb()\n",
    "pyvista.set_plot_theme(\"paraview\")\n",
    "sargs = dict(title_font_size=25, label_font_size=20, fmt=\"%.2e\", color=\"black\",\n",
    "             position_x=0.1, position_y=0.8, width=0.8, height=0.1)\n",
    "\n",
    "def export_function(grid, name, show_mesh=False, tessellate=False):\n",
    "    grid.set_active_scalars(name)\n",
    "    plotter = pyvista.Plotter(window_size=(700, 700))\n",
    "    t_grid = grid.tessellate() if tessellate else grid\n",
    "    renderer = plotter.add_mesh(t_grid, show_edges=False, scalar_bar_args=sargs)\n",
    "    if show_mesh:\n",
    "        V = dolfinx.fem.functionspace(mesh, (\"Lagrange\", 1))\n",
    "        grid_mesh = pyvista.UnstructuredGrid(*vtk_mesh(V))\n",
    "        renderer = plotter.add_mesh(grid_mesh, style=\"wireframe\", line_width=0.1, color=\"k\")\n",
    "        plotter.view_xy()\n",
    "    plotter.view_xy()\n",
    "    plotter.camera.zoom(1.3)\n",
    "    plotter.export_html(f\"./{name}.html\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6e33792d-50b9-4e0c-92d6-c43745243439",
   "metadata": {
    "lines_to_next_cell": 2,
    "tags": [
     "ignore-output",
     "hide-input",
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "grid = pyvista.UnstructuredGrid(*vtk_mesh(mesh))\n",
    "grid.cell_data[\"wavenumber\"] = k.x.array.real"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1841452a-799c-429c-96ce-b0a638d56e8f",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "export_function(grid, \"wavenumber\", show_mesh=True, tessellate=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b623a5ff-72b3-48a7-a3de-42d614a101e7",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<iframe src='./wavenumber.html' height=\"700px\" width=\"700px\"></iframe>  <!--  # noqa, -->\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%html\n",
    "<iframe src='./wavenumber.html' height=\"700px\" width=\"700px\"></iframe>  <!--  # noqa, -->"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27cc978a-bc0d-4a77-856d-f17e27f06094",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "## Boundary source term\n",
    "\n",
    "$$g = \\nabla u_{inc} \\cdot \\mathbf{n} - \\mathrm{j}ku_{inc}$$\n",
    "\n",
    "where $u_{inc} = e^{-jkx}$, the incoming wave, is a plane wave propagating in the $x$ direction."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a72cc1c-1043-46ab-bb7b-2d7f89823449",
   "metadata": {
    "slideshow": {
     "slide_type": "notes"
    },
    "tags": []
   },
   "source": [
    "\n",
    "Next, we define the boundary source term by using `ufl.SpatialCoordinate`. When using this function, all quantities using this expression will be evaluated at quadrature points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "cf99d1fe-55df-4630-910e-a3caf118153b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "n = ufl.FacetNormal(mesh)\n",
    "x = ufl.SpatialCoordinate(mesh)\n",
    "uinc = ufl.exp(1j * k * x[0])\n",
    "g = ufl.dot(ufl.grad(uinc), n) - 1j * k * uinc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eae8efd8-d18c-4075-b3fd-db1b3d142335",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "## Variational form"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cef49f19-5c6e-4437-b6b0-2af9e3392763",
   "metadata": {
    "slideshow": {
     "slide_type": "notes"
    },
    "tags": []
   },
   "source": [
    "Next, we define the variational problem using a 4th order Lagrange space. Note that as we are using complex valued functions, we have to use the appropriate inner product; see [DOLFINx tutorial: Complex numbers](https://jorgensd.github.io/dolfinx-tutorial/chapter1/complex_mode.html) for more information."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e36a8dc7-ff02-4b72-ad79-6bff7b40a2cf",
   "metadata": {},
   "source": [
    "$$ -\\int_\\Omega \\nabla u \\cdot \\nabla \\bar{v} ~ dx + \\int_\\Omega k^2 u \\,\\bar{v}~ dx - j\\int_{\\partial \\Omega} ku  \\bar{v} ~ ds = \\int_{\\partial \\Omega} g \\, \\bar{v}~ ds \\qquad \\forall v \\in \\widehat{V}. $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f7b6e171-6832-4174-8804-b0315f5a3c11",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import basix.ufl\n",
    "element = basix.ufl.element(\"Lagrange\", mesh.topology.cell_name(), degree)\n",
    "V = dolfinx.fem.functionspace(mesh, element)\n",
    "\n",
    "u = ufl.TrialFunction(V)\n",
    "v = ufl.TestFunction(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b4327cad-a8f7-4009-9f17-e50341fc4082",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "a = - ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx \\\n",
    "    + k**2 * ufl.inner(u, v) * ufl.dx \\\n",
    "    - 1j * k * ufl.inner(u, v) * ufl.ds\n",
    "L = ufl.inner(g, v) * ufl.ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "078f094e-4ae9-4403-98ea-bc2dcf053483",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "## Linear solver"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9547f3e1-38f3-43f2-9a0b-b6ea297dfa59",
   "metadata": {
    "slideshow": {
     "slide_type": "notes"
    },
    "tags": []
   },
   "source": [
    "Next, we will solve the problem using a direct solver (LU)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6c8ea5ef-77df-4a44-a06e-85ed2f186940",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "opt = {\"ksp_type\": \"preonly\", \"pc_type\": \"lu\"}\n",
    "problem = dolfinx.fem.petsc.LinearProblem(a, L, petsc_options=opt)\n",
    "uh = problem.solve()\n",
    "uh.name = \"u\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31456b13-bec6-4e82-9e33-975d92b0be55",
   "metadata": {
    "lines_to_next_cell": 0,
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": []
   },
   "source": [
    "## Postprocessing\n",
    "\n",
    "### Visualising using PyVista"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "006df3cc-5436-4ba2-ae2e-4205e5a1f117",
   "metadata": {
    "lines_to_next_cell": 2,
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "topology, cells, geometry = vtk_mesh(V)\n",
    "grid = pyvista.UnstructuredGrid(topology, cells, geometry)\n",
    "grid.point_data[\"Abs(u)\"] = np.abs(uh.x.array)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "34817139-1016-44cc-8d40-c16f239d9ab4",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "export_function(grid, \"Abs(u)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "358a3a09-18dd-457d-a8ab-b08ade039f43",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<iframe src='./Abs(u).html' width=\"700px\" height=\"700px\">></iframe>  <!--  # noqa, -->\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%html\n",
    "<iframe src='./Abs(u).html' width=\"700px\" height=\"700px\">></iframe>  <!--  # noqa, -->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "6a966a99-0f93-45be-bae0-4d0676fe2fc3",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "# Open a gif\n",
    "plotter = pyvista.Plotter()\n",
    "plotter.open_gif(\"wave.gif\")\n",
    "boring_cmap = plt.cm.get_cmap(\"coolwarm\", 25)\n",
    "\n",
    "pts = grid.points.copy()\n",
    "warped = grid.warp_by_scalar()\n",
    "renderer = plotter.add_mesh(\n",
    "    warped, show_edges=False, clim=[-2, 2.5], lighting=False, cmap=boring_cmap)\n",
    "\n",
    "nframe = 27\n",
    "for phase in np.linspace(0, 2 * np.pi, nframe + 1)[:nframe]:\n",
    "    data = uh.x.array * np.exp(1j * phase)\n",
    "    plotter.update_scalars(data.real, render=False)\n",
    "    warped = grid.warp_by_scalar(factor=0.1)\n",
    "    plotter.update_coordinates(warped.points.copy(), render=False)\n",
    "\n",
    "    # Write a frame. This triggers a render.\n",
    "    plotter.write_frame()\n",
    "plotter.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "067b6793-6bba-407a-b468-3c251082b610",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "<img src=\"./wave.gif\" alt=\"gif\" class=\"bg-primary mb-1\" width=\"800px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cbf8d9b-04e1-4013-95ff-cce689897a64",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "### Post-processing with Paraview"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "e93e9893-5111-4389-8ba2-e615b393cdb4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from dolfinx.io import XDMFFile, VTXWriter\n",
    "u_abs = dolfinx.fem.Function(V, dtype=np.float64)\n",
    "u_abs.x.array[:] = np.abs(uh.x.array)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "018960ee-5fac-45a0-8d24-6c4ce881cb72",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "ignore-output"
    ]
   },
   "source": [
    "### VTXWriter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "e23cbd75-d089-4fa6-8d68-478f36a1ff32",
   "metadata": {},
   "outputs": [],
   "source": [
    "# VTX can write higher order functions\n",
    "with VTXWriter(comm, \"out_high_order.bp\", [u_abs], engine=\"BP4\") as f:\n",
    "    f.write(0.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48734442-c1bf-434b-921c-d3f2a677d63d",
   "metadata": {},
   "source": [
    "<img src=\"./img/vtx.png\" alt=\"vtx\" class=\"bg-primary mb-1\" width=\"500px\">"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb"
  },
  "kernelspec": {
   "display_name": "Python 3 (DOLFINx complex)",
   "language": "python",
   "name": "python3-complex"
  },
  "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.10.12"
  },
  "vscode": {
   "interpreter": {
    "hash": "aee8b7b246df8f9039afb4144a1f6fd8d2ca17a180786b69acc140d282b71a49"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}