chapter2/diffusion_code.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Diffusion of a Gaussian function\n",
    "\n",
    "Author: Jørgen S. Dokken\n",
    "\n",
    "Let us now solve a more interesting problem, namely the diffusion of a Gaussian hill. We take the initial value to be\n",
    "\\begin{align}\n",
    "    u_0(x,y)&= e^{-ax^2-ay^2}\n",
    "\\end{align}\n",
    "for $a=5$ on the domain $[-2,2]\\times[-2,2]$. For this problem we will use homogeneous Dirichlet boundary conditions ($u_D=0$).\n",
    "\n",
    "The first difference from the previous problem is that we are not using a unit square. We create the rectangular domain with `dolfinx.mesh.create_rectangle`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib as mpl\n",
    "import pyvista\n",
    "import ufl\n",
    "import numpy as np\n",
    "\n",
    "from petsc4py import PETSc\n",
    "from mpi4py import MPI\n",
    "\n",
    "from dolfinx import fem, mesh, io, plot\n",
    "from dolfinx.fem.petsc import assemble_vector, assemble_matrix, create_vector, apply_lifting, set_bc\n",
    "\n",
    "# Define temporal parameters\n",
    "t = 0  # Start time\n",
    "T = 1.0  # Final time\n",
    "num_steps = 50\n",
    "dt = T / num_steps  # time step size\n",
    "\n",
    "# Define mesh\n",
    "nx, ny = 50, 50\n",
    "domain = mesh.create_rectangle(MPI.COMM_WORLD, [np.array([-2, -2]), np.array([2, 2])],\n",
    "                               [nx, ny], mesh.CellType.triangle)\n",
    "V = fem.functionspace(domain, (\"Lagrange\", 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that we have used a much higher resolution than before to better resolve features of the solution.\n",
    "We also easily update the intial and boundary conditions. Instead of using a class to define the initial condition, we simply use a function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create initial condition\n",
    "def initial_condition(x, a=5):\n",
    "    return np.exp(-a * (x[0]**2 + x[1]**2))\n",
    "\n",
    "\n",
    "u_n = fem.Function(V)\n",
    "u_n.name = \"u_n\"\n",
    "u_n.interpolate(initial_condition)\n",
    "\n",
    "# Create boundary condition\n",
    "fdim = domain.topology.dim - 1\n",
    "boundary_facets = mesh.locate_entities_boundary(\n",
    "    domain, fdim, lambda x: np.full(x.shape[1], True, dtype=bool))\n",
    "bc = fem.dirichletbc(PETSc.ScalarType(0), fem.locate_dofs_topological(V, fdim, boundary_facets), V)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Time-dependent output\n",
    "To visualize the solution in an external program such as Paraview, we create a an `XDMFFile` which we can store multiple solutions in. The main advantage with an XDMFFile is that we only need to store the mesh once and that we can append multiple solutions to the same grid, reducing the storage space.\n",
    "The first argument to the XDMFFile is which communicator should be used to store the data. As we would like one output, independent of the number of processors, we use the `COMM_WORLD`. The second argument is the file name of the output file, while the third argument is the state of the file,\n",
    "this could be read (`\"r\"`), write (`\"w\"`) or append (`\"a\"`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xdmf = io.XDMFFile(domain.comm, \"diffusion.xdmf\", \"w\")\n",
    "xdmf.write_mesh(domain)\n",
    "\n",
    "# Define solution variable, and interpolate initial solution for visualization in Paraview\n",
    "uh = fem.Function(V)\n",
    "uh.name = \"uh\"\n",
    "uh.interpolate(initial_condition)\n",
    "xdmf.write_function(uh, t)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Variational problem and solver\n",
    "As in the previous example, we prepare objects for time dependent problems, such that we do not have to recreate data-structures."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u, v = ufl.TrialFunction(V), ufl.TestFunction(V)\n",
    "f = fem.Constant(domain, PETSc.ScalarType(0))\n",
    "a = u * v * ufl.dx + dt * ufl.dot(ufl.grad(u), ufl.grad(v)) * ufl.dx\n",
    "L = (u_n + dt * f) * v * ufl.dx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparing linear algebra structures for time dependent problems\n",
    "We note that even if `u_n` is time dependent, we will reuse the same function for `f` and `u_n` at every time step. We therefore call `dolfinx.fem.form` to generate assembly kernels for the matrix and vector."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bilinear_form = fem.form(a)\n",
    "linear_form = fem.form(L)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We observe that the left hand side of the system, the matrix $A$ does not change from one time step to another, thus we only need to assemble it once. However, the right hand side, which is dependent on the previous time step `u_n`, we have to assemble it every time step. Therefore, we only create a vector `b` based on `L`, which we will reuse at every time step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A = assemble_matrix(bilinear_form, bcs=[bc])\n",
    "A.assemble()\n",
    "b = create_vector(linear_form)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using petsc4py to create a linear solver\n",
    "As we have already assembled `a` into the matrix `A`, we can no longer use the `dolfinx.fem.petsc.LinearProblem` class to solve the problem. Therefore, we create a linear algebra solver using PETSc, assign the matrix `A` to the solver, and choose the solution strategy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "solver = PETSc.KSP().create(domain.comm)\n",
    "solver.setOperators(A)\n",
    "solver.setType(PETSc.KSP.Type.PREONLY)\n",
    "solver.getPC().setType(PETSc.PC.Type.LU)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualization of time dependent problem using pyvista\n",
    "We use the DOLFINx plotting functionality, which is based on pyvista to plot the solution at every $15$th time step. We would also like to visualize a colorbar reflecting the minimal and maximum value of $u$ at each time step. We use the following convenience function `plot_function` for this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista.start_xvfb()\n",
    "\n",
    "grid = pyvista.UnstructuredGrid(*plot.vtk_mesh(V))\n",
    "\n",
    "plotter = pyvista.Plotter()\n",
    "plotter.open_gif(\"u_time.gif\", fps=10)\n",
    "\n",
    "grid.point_data[\"uh\"] = uh.x.array\n",
    "warped = grid.warp_by_scalar(\"uh\", factor=1)\n",
    "\n",
    "viridis = mpl.colormaps.get_cmap(\"viridis\").resampled(25)\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",
    "renderer = plotter.add_mesh(warped, show_edges=True, lighting=False,\n",
    "                            cmap=viridis, scalar_bar_args=sargs,\n",
    "                            clim=[0, max(uh.x.array)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Updating the solution and right hand side per time step\n",
    "To be able to solve the variation problem at each time step, we have to assemble the right hand side and apply the boundary condition before calling\n",
    "`solver.solve(b, uh.x.petsc_vec)`. We start by resetting the values in `b` as we are reusing the vector at every time step.\n",
    "The next step is to assemble the vector calling `dolfinx.fem.petsc.assemble_vector(b, L)`, which means that we are assembling the linear form `L(v)` into the vector `b`. Note that we do not supply the boundary conditions for assembly, as opposed to the left hand side.\n",
    "This is because we want to use lifting to apply the boundary condition, which preserves symmetry of the matrix $A$ in the bilinear form $a(u,v)=a(v,u)$ without Dirichlet boundary conditions.\n",
    "When we have applied the boundary condition, we can solve the linear system and update values that are potentially shared between processors.\n",
    "Finally, before moving to the next time step, we update the solution at the previous time step to the solution at this time step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(num_steps):\n",
    "    t += dt\n",
    "\n",
    "    # Update the right hand side reusing the initial vector\n",
    "    with b.localForm() as loc_b:\n",
    "        loc_b.set(0)\n",
    "    assemble_vector(b, linear_form)\n",
    "\n",
    "    # Apply Dirichlet boundary condition to the vector\n",
    "    apply_lifting(b, [bilinear_form], [[bc]])\n",
    "    b.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)\n",
    "    set_bc(b, [bc])\n",
    "\n",
    "    # Solve linear problem\n",
    "    solver.solve(b, uh.x.petsc_vec)\n",
    "    uh.x.scatter_forward()\n",
    "\n",
    "    # Update solution at previous time step (u_n)\n",
    "    u_n.x.array[:] = uh.x.array\n",
    "\n",
    "    # Write solution to file\n",
    "    xdmf.write_function(uh, t)\n",
    "    # Update plot\n",
    "    new_warped = grid.warp_by_scalar(\"uh\", factor=1)\n",
    "    warped.points[:, :] = new_warped.points\n",
    "    warped.point_data[\"uh\"][:] = uh.x.array\n",
    "    plotter.write_frame()\n",
    "plotter.close()\n",
    "xdmf.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"./u_time.gif\" alt=\"gif\" class=\"bg-primary mb-1\" width=\"800px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Animation with Paraview\n",
    "We can also use Paraview to create an animation. We open the file in paraview with `File->Open`, and then press `Apply` in the properties panel.\n",
    "\n",
    "Then, we add a time-annotation to the figure, pressing: `Sources->Alphabetical->Annotate Time` and `Apply` in the properties panel. It Is also a good idea to select an output resolution, by pressing `View->Preview->1280 x 720 (HD)`.\n",
    "\n",
    "Then finally, click `File->Save Animation`, and save the animation to the desired format, such as `avi`, `ogv` or a sequence of `png`s. Make sure to set the frame rate to something sensible, in the range of $5-10$ frames per second."
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,py:light"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}