python/demo_stokes_stable_pairs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f06ef76d",
   "metadata": {},
   "source": [
    "# Stokes equations with various stable pairs of finite elements\n",
    "\n",
    "## Equation and problem definition\n",
    "\n",
    "### Strong formulation\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "-\\nabla \\cdot (\\nabla u + p I) &= f \\; {\\rm in} \\; \\Omega, \\\\\n",
    "\\nabla \\cdot u &= 0 \\; {\\rm in} \\; \\Omega\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "\n",
    "A typical set of boundary conditions on the boundary\n",
    "$\\partial\\Omega = \\Gamma_{D} \\cup \\Gamma_{N}$ can be:\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "u &= u_0 \\; {\\rm on} \\; \\Gamma_{D}, \\\\\n",
    "\\nabla u \\cdot n + p n &= g \\; {\\rm on} \\; \\Gamma_{N}.\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "### Weak formulation\n",
    "\n",
    "We formulate the Stokes equations' mixed variational form; that is, a\n",
    "form where the two variables, the velocity and the pressure, are\n",
    "approximated. We have the problem: find $(u, p) \\in W$ such that\n",
    "\n",
    "$$\n",
    "a((u, p), (v, q)) = L((v, q))\n",
    "$$\n",
    "\n",
    "for all $(v, q) \\in W$, where\n",
    "\n",
    "$$\n",
    "a((u, p), (v, q)) := \\int_{\\Omega} \\nabla u \\cdot \\nabla v\n",
    "- \\nabla \\cdot v \\; p + \\nabla \\cdot u \\; q \\, {\\rm d} x,\n",
    "$$\n",
    "\n",
    "$$\n",
    "L((v, q)) := \\int_{\\Omega} f \\cdot v \\, {\\rm d} x\n",
    "+ \\int_{\\partial \\Omega_N} g \\cdot v \\, {\\rm d} s.\n",
    "$$\n",
    "\n",
    "The space $W$ is a mixed (product) function space $W = V \\times Q$,\n",
    "such that $u \\in V$ and $q \\in Q$.\n",
    "\n",
    "### Domain and boundary conditions\n",
    "\n",
    "We define the lid-driven cavity problem with the following\n",
    "domain and boundary conditions:\n",
    "\n",
    "- $\\Omega = [0,1]\\times[0,1]$ (a unit square)\n",
    "- $\\Gamma_D = \\partial \\Omega$\n",
    "- $u_0 = (1, 0)^\\top$ at $x_1 = 1$ and $u_0 = (0,\n",
    "  0)^\\top$ otherwise\n",
    "- $f = (0, 0)^\\top$\n",
    "\n",
    "### Discretization\n",
    "\n",
    "There are many ways to choose the discretized function spaces $V_h \\subset V$\n",
    "and $Q_h \\subset Q$ for the mixed variational form.\n",
    "Care must be taken that the combination of $V_h$ and $Q_h$ is stable\n",
    "in the sense of the inf-sup condition.\n",
    "Common are the following stable pairs:\n",
    "\n",
    "1. $(\\mathcal{P}_2, \\mathcal{P}_1$):\n",
    "The [Taylor-Hood element](https://defelement.com/elements/taylor-hood.html)\n",
    "for $k=2$\n",
    "2. $(\\mathcal{P}_1 + \\mathcal{B}_3, \\mathcal{P}_1)$:\n",
    "The [MINI element](https://defelement.com/elements/mini.html)\n",
    "3. $(\\mathcal{P}_1^{\\rm CR}, \\mathcal{P}_0)$:\n",
    "The [non-conforming Crouzeix-Raviart element]\n",
    "(https://defelement.com/elements/crouzeix-raviart.html)\n",
    "4. $(\\mathcal{P}_4, \\mathcal{P}_3^{-1})$:\n",
    "The Scott-Vogelius element for $k=4$\n",
    "\n",
    "In the following, the Stokes equations in the lid-driven cavity setting\n",
    "are solved using each of these stable pairs.\n",
    "\n",
    "## Implementation\n",
    "\n",
    "We first import the modules and functions that the program uses:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29cc5da9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import ufl\n",
    "from dolfinx import fem\n",
    "from dolfinx.fem import (Constant, Function, FunctionSpace, dirichletbc, form,\n",
    "                         locate_dofs_topological)\n",
    "from dolfinx.io import XDMFFile\n",
    "from dolfinx.mesh import (CellType, GhostMode, create_rectangle,\n",
    "                          locate_entities_boundary)\n",
    "from ufl import div, dx, grad, inner\n",
    "\n",
    "from mpi4py import MPI\n",
    "from petsc4py import PETSc\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24f67ef7",
   "metadata": {},
   "source": [
    "We create a {py:class}`Mesh <dolfinx.mesh.Mesh>`, define functions to\n",
    "geometrically locate subsets of its boundary and define a function\n",
    "describing the velocity to be imposed as a boundary condition in a lid\n",
    "driven cavity problem:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f424364",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create mesh\n",
    "msh = create_rectangle(MPI.COMM_WORLD,\n",
    "                       [np.array([0, 0]), np.array([1, 1])],\n",
    "                       [32, 32],\n",
    "                       CellType.triangle, GhostMode.none)\n",
    "\n",
    "\n",
    "# Function to mark x = 0, x = 1 and y = 0\n",
    "def noslip_boundary(x):\n",
    "    return np.logical_or(np.logical_or(np.isclose(x[0], 0.0),\n",
    "                                       np.isclose(x[0], 1.0)),\n",
    "                         np.isclose(x[1], 0.0))\n",
    "\n",
    "\n",
    "# Function to mark the lid (y = 1)\n",
    "def lid(x):\n",
    "    return np.isclose(x[1], 1.0)\n",
    "\n",
    "\n",
    "# Lid velocity\n",
    "def lid_velocity_expression(x):\n",
    "    return np.stack((np.ones(x.shape[1]), np.zeros(x.shape[1])))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edc604ec",
   "metadata": {},
   "source": [
    "### 1. $(\\mathcal{P}_2, \\mathcal{P}_1$): The Taylor-Hood element for $k=2$\n",
    "\n",
    "For the Taylor-Hood, the discrete function spaces are chosen as\n",
    "$V_h = \\mathcal{P}_k$ and $Q_h = \\mathcal{P}_{k-1}$ with $k \\geq 2$.\n",
    "\n",
    "For $k = 2$ we have:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20f11378",
   "metadata": {},
   "outputs": [],
   "source": [
    "P2 = ufl.VectorElement(\"Lagrange\", msh.ufl_cell(), 2)\n",
    "P1 = ufl.FiniteElement(\"Lagrange\", msh.ufl_cell(), 1)\n",
    "V, Q = FunctionSpace(msh, P2), FunctionSpace(msh, P1)\n",
    "\n",
    "(u, p) = ufl.TrialFunction(V), ufl.TrialFunction(Q)\n",
    "(v, q) = ufl.TestFunction(V), ufl.TestFunction(Q)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb99e91c",
   "metadata": {},
   "source": [
    "First, we have to define the boundary conditions for the velocity field\n",
    "in the problem setting. In the case of the lid driven cavity scenario,\n",
    "there is a driving velocity condition on the top boundary\n",
    "and a no-slip boundary condition on the remaining boundary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0371bb3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def define_bcs(V):\n",
    "    # No-slip boundary condition for velocity field (`V`) on boundaries\n",
    "    # where x = 0, x = 1, and y = 0\n",
    "    noslip = np.zeros(msh.geometry.dim, dtype=PETSc.ScalarType)\n",
    "    facets = locate_entities_boundary(msh, 1, noslip_boundary)\n",
    "    bc0 = dirichletbc(noslip, locate_dofs_topological(V, 1, facets), V)\n",
    "\n",
    "    # Driving velocity condition u = (1, 0) on top boundary (y = 1)\n",
    "    lid_velocity = Function(V)\n",
    "    lid_velocity.interpolate(lid_velocity_expression)\n",
    "    facets = locate_entities_boundary(msh, 1, lid)\n",
    "    bc1 = dirichletbc(lid_velocity, locate_dofs_topological(V, 1, facets))\n",
    "\n",
    "    # Collect Dirichlet boundary conditions\n",
    "    bcs = [bc0, bc1]\n",
    "    return bcs\n",
    "\n",
    "\n",
    "bcs = define_bcs(V)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efddfd87",
   "metadata": {},
   "source": [
    "The variational problem for the Stokes equations can be defined as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d195c801",
   "metadata": {},
   "outputs": [],
   "source": [
    "def define_weak_form(u, p, v, q):\n",
    "    # Define variational problem\n",
    "    f = Constant(msh, (PETSc.ScalarType(0), PETSc.ScalarType(0)))\n",
    "\n",
    "    a = form([[inner(grad(u), grad(v)) * dx, inner(p, div(v)) * dx],\n",
    "              [inner(div(u), q) * dx, None]])\n",
    "    L = form([inner(f, v) * dx,\n",
    "              inner(Constant(msh, PETSc.ScalarType(0)), q) * dx])\n",
    "    return a, L\n",
    "\n",
    "\n",
    "a, L = define_weak_form(u, p, v, q)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0938a98",
   "metadata": {},
   "source": [
    "Before the solution can be obtained, the linear system has to be assembled.\n",
    "We set a null space to account for the homogeneous Neumann\n",
    "boundary conditions of the pressure field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcb0de33",
   "metadata": {},
   "outputs": [],
   "source": [
    "def assemble_system(a, L, bcs, V):\n",
    "    A = fem.petsc.assemble_matrix_block(a, bcs=bcs)\n",
    "    A.assemble()\n",
    "    b = fem.petsc.assemble_vector_block(L, a, bcs=bcs)\n",
    "\n",
    "    # Set near nullspace for pressure\n",
    "    null_vec = A.createVecLeft()\n",
    "    offset = V.dofmap.index_map.size_local * V.dofmap.index_map_bs\n",
    "    null_vec.array[offset:] = 1.0\n",
    "    null_vec.normalize()\n",
    "    nsp = PETSc.NullSpace().create(vectors=[null_vec])\n",
    "    assert nsp.test(A)\n",
    "    A.setNullSpace(nsp)\n",
    "    return A, b\n",
    "\n",
    "\n",
    "A, b = assemble_system(a, L, bcs, V)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d47f663c",
   "metadata": {},
   "source": [
    "Finally, the linear system can be solved to obtain the solution functions\n",
    "$u$ and $p$. Here, a direct solver is used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c35aa40",
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_system(A, b, msh, V):\n",
    "    # Create LU solver\n",
    "    ksp = PETSc.KSP().create(msh.comm)\n",
    "    ksp.setOperators(A)\n",
    "    ksp.setType(\"preonly\")\n",
    "    ksp.getPC().setType(\"lu\")\n",
    "    ksp.getPC().setFactorSolverType(\"superlu_dist\")\n",
    "\n",
    "    # Compute solution\n",
    "    x = A.createVecLeft()\n",
    "    ksp.solve(b, x)\n",
    "\n",
    "    # Create Functions and scatter x solution\n",
    "    u, p = Function(V), Function(Q)\n",
    "    V_map = V.dofmap.index_map\n",
    "    offset = V_map.size_local * V.dofmap.index_map_bs\n",
    "    u.x.array[:offset] = x.array_r[:offset]\n",
    "    p.x.array[:(len(x.array_r) - offset)] = x.array_r[offset:]\n",
    "    return u, p\n",
    "\n",
    "\n",
    "u, p = solve_system(A, b, msh, V)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d34f07c9",
   "metadata": {},
   "source": [
    "Print the norm of the solved velocity and pressure coefficient vectors:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0905c2a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def l2_norm(sol):\n",
    "    comm = sol.function_space.mesh.comm\n",
    "    error = form(sol**2 * ufl.dx)\n",
    "    return np.sqrt(comm.allreduce(fem.assemble_scalar(error), MPI.SUM))\n",
    "\n",
    "\n",
    "coef_norm_u_1 = u.x.norm()\n",
    "coef_norm_p_1 = p.x.norm()\n",
    "l2_norm_u_1 = l2_norm(u)\n",
    "l2_norm_p_1 = l2_norm(p)\n",
    "if MPI.COMM_WORLD.rank == 0:\n",
    "    print(\"(1) Norm of velocity coeff. vector \"\n",
    "          \"with the Taylor-Hood element:      {}\".format(coef_norm_u_1))\n",
    "    print(\"(1) Norm of pressure coeff. vector \"\n",
    "          \"with the Taylor-Hood element:      {}\".format(coef_norm_p_1))\n",
    "    print(\"(1) L2 Norm of the velocity field \"\n",
    "          \"with the Taylor-Hood element:       {}\".format(l2_norm_u_1))\n",
    "    print(\"(1) L2 Norm of pressure field \"\n",
    "          \"with the Taylor-Hood element:           {}\".format(l2_norm_p_1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25b1bbb2",
   "metadata": {},
   "source": [
    "The solved velocity and pressure fields are saved as XDMF files\n",
    "and can be visualized e.g. in Paraview."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05212959",
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_solution(sol, file_name):\n",
    "    with XDMFFile(MPI.COMM_WORLD, file_name, \"w\") as file_xdmf:\n",
    "        sol.x.scatter_forward()\n",
    "        file_xdmf.write_mesh(msh)\n",
    "        file_xdmf.write_function(sol)\n",
    "\n",
    "\n",
    "save_solution(u, \"out_stokes_stable_pairs/1_velocity.xdmf\")\n",
    "save_solution(p, \"out_stokes_stable_pairs/1_pressure.xdmf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cd2b421",
   "metadata": {},
   "source": [
    "### 2. $(\\mathcal{P}_1 + \\mathcal{B}_3, \\mathcal{P}_1)$: The MINI element"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81a970cc",
   "metadata": {},
   "source": [
    "For the so-called [MINI element](https://defelement.com/elements/mini.html)\n",
    "the finite element for the velocity field is chosen as\n",
    "the Lagrange element of degree 1 enriched with a\n",
    "[bubble element](https://defelement.com/elements/bubble.html) of degree 3.\n",
    "The finite element for the pressure field is chosen as\n",
    "Lagrange element of degree 1.\n",
    "\n",
    "In `dolfinx`, enriched finite elements can be obtained simply\n",
    "by using the `+` operator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7480a883",
   "metadata": {},
   "outputs": [],
   "source": [
    "P1 = ufl.FiniteElement(\"Lagrange\", msh.ufl_cell(), 1)\n",
    "B = ufl.FiniteElement(\"Bubble\", msh.ufl_cell(), 3)\n",
    "V_enriched = ufl.VectorElement(P1 + B)\n",
    "V, Q = FunctionSpace(msh, V_enriched), FunctionSpace(msh, P1)\n",
    "\n",
    "(u, p) = ufl.TrialFunction(V), ufl.TrialFunction(Q)\n",
    "(v, q) = ufl.TestFunction(V), ufl.TestFunction(Q)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60ea203b",
   "metadata": {},
   "source": [
    "We solve the Stokes equations as before but this time using the MINI element.\n",
    "Subsequently, the solved velocity and pressure fields are saved."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72d7cca1",
   "metadata": {},
   "outputs": [],
   "source": [
    "bcs = define_bcs(V)\n",
    "a, L = define_weak_form(u, p, v, q)\n",
    "A, b = assemble_system(a, L, bcs, V)\n",
    "u, p = solve_system(A, b, msh, V)\n",
    "coef_norm_u_2 = u.x.norm()\n",
    "coef_norm_p_2 = p.x.norm()\n",
    "l2_norm_u_2 = l2_norm(u)\n",
    "l2_norm_p_2 = l2_norm(p)\n",
    "if MPI.COMM_WORLD.rank == 0:\n",
    "    print(\"(2) Norm of velocity coeff. vector \"\n",
    "          \"with the MINI element:             {}\".format(coef_norm_u_2))\n",
    "    print(\"(2) Norm of pressure coeff. vector \"\n",
    "          \"with the MINI element:             {}\".format(coef_norm_p_2))\n",
    "    print(\"(2) L2 Norm of the velocity field \"\n",
    "          \"with the MINI element:              {}\".format(l2_norm_u_2))\n",
    "    print(\"(2) L2 Norm of pressure field \"\n",
    "          \"with the MINI element:                  {}\".format(l2_norm_p_2))\n",
    "save_solution(u, \"out_stokes_stable_pairs/2_velocity.xdmf\")\n",
    "save_solution(p, \"out_stokes_stable_pairs/2_pressure.xdmf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52a3efa1",
   "metadata": {},
   "source": [
    "### 3. $(\\mathcal{P}_1^{\\rm CR}, \\mathcal{P}_0)$:\n",
    "The non-conforming Crouzeix-Raviart element\n",
    "\n",
    "Another possibility for chosing the discretized function spaces\n",
    "is by using the [Crouzeix-Raviart element]\n",
    "(https://defelement.com/elements/crouzeix-raviart.html).\n",
    "Here, a non-conforming variant of the Lagrange element of degree 1\n",
    "is used for the velocity field,\n",
    "whereas the DG element is used for the pressure field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7421e6d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "P1CR = ufl.VectorElement(\"Crouzeix-Raviart\", msh.ufl_cell(), 1)\n",
    "P0 = ufl.FiniteElement(\"DG\", msh.ufl_cell(), 0)\n",
    "V, Q = FunctionSpace(msh, P1CR), FunctionSpace(msh, P1)\n",
    "\n",
    "(u, p) = ufl.TrialFunction(V), ufl.TrialFunction(Q)\n",
    "(v, q) = ufl.TestFunction(V), ufl.TestFunction(Q)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a271b65",
   "metadata": {},
   "source": [
    "We solve the Stokes equations as before but this time using\n",
    "the Crouzeix-Raviart element.\n",
    "Subsequently, the solved velocity and pressure fields are saved."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07d2a4b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "bcs = define_bcs(V)\n",
    "a, L = define_weak_form(u, p, v, q)\n",
    "A, b = assemble_system(a, L, bcs, V)\n",
    "u, p = solve_system(A, b, msh, V)\n",
    "coef_norm_u_3 = u.x.norm()\n",
    "coef_norm_p_3 = p.x.norm()\n",
    "l2_norm_u_3 = l2_norm(u)\n",
    "l2_norm_p_3 = l2_norm(p)\n",
    "if MPI.COMM_WORLD.rank == 0:\n",
    "    print(\"(3) Norm of velocity coeff. vector \"\n",
    "          \"with the Crouzeix-Raviart element: {}\".format(coef_norm_u_3))\n",
    "    print(\"(3) Norm of pressure coeff. vector \"\n",
    "          \"with the Crouzeix-Raviart element: {}\".format(coef_norm_p_3))\n",
    "    print(\"(3) L2 Norm of the velocity field \"\n",
    "          \"with the Crouzeix-Raviart element:  {}\".format(l2_norm_u_3))\n",
    "    print(\"(3) L2 Norm of pressure field \"\n",
    "          \"with the Crouzeix-Raviart element:      {}\".format(l2_norm_p_3))\n",
    "save_solution(u, \"out_stokes_stable_pairs/3_velocity.xdmf\")\n",
    "save_solution(p, \"out_stokes_stable_pairs/3_pressure.xdmf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f2bd61b",
   "metadata": {},
   "source": [
    "### 4. $(\\mathcal{P}_4, \\mathcal{P}_{3}^{-1})$:\n",
    "The Scott-Vogelius element for $k=4$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15a56429",
   "metadata": {},
   "source": [
    "A possible to solve the Stokes equatoins using the Scott-Vogelius\n",
    "element is the iterated penalty method\n",
    "by Scott and Vogelius (1985). This iterative procedure is also\n",
    "explained in the FEniCS book, chapter 20.3.\n",
    "\n",
    "In each iteration, we have to find $u^n \\in V$ such that\n",
    "\n",
    "$$\n",
    "a(u^n, v) + r (\\nabla \\cdot u^n, \\nabla \\cdot v)\n",
    "= (f, v) - (\\nabla \\cdot v, \\nabla \\cdot w^n)\n",
    "$$\n",
    "\n",
    "for all $v \\in V$, where\n",
    "\n",
    "$$\n",
    "w^{n+1} = w^n + \\rho u^n.\n",
    "$$\n",
    "\n",
    "The parameters $r \\in \\mathbb{R}$ and $\\rho \\in \\mathbb{R}^+$\n",
    "can be chosen to achieve good convergence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1a746a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define function space\n",
    "P4 = ufl.VectorElement(\"Lagrange\", msh.ufl_cell(), 4)\n",
    "P3 = ufl.FiniteElement(\"DG\", msh.ufl_cell(), 3)\n",
    "V, Q = FunctionSpace(msh, P4), FunctionSpace(msh, P3)\n",
    "\n",
    "# Define trial and test functions\n",
    "u = ufl.TrialFunction(V)\n",
    "v = ufl.TestFunction(V)\n",
    "\n",
    "# Define auxiliary function and parameters\n",
    "w = Function(V)\n",
    "rho = 1.0e3\n",
    "r = -rho\n",
    "\n",
    "# Define the variational problem\n",
    "f = Constant(msh, (PETSc.ScalarType(0), PETSc.ScalarType(0)))\n",
    "a = inner(grad(u), grad(v)) * dx + r * inner(div(u), div(v)) * dx\n",
    "L = inner(f, v) * dx + inner(div(w), div(v)) * dx\n",
    "\n",
    "bcs = define_bcs(V)\n",
    "problem = fem.petsc.LinearProblem(a, L, bcs=bcs,\n",
    "                                  petsc_options={\"ksp_type\": \"preonly\",\n",
    "                                                 \"pc_type\": \"lu\"})\n",
    "# Iterate to fix point\n",
    "iters = 0\n",
    "max_iters = 100\n",
    "u_m_u = 1\n",
    "u_old_vec = None\n",
    "while iters < max_iters and u_m_u > 1e-8:\n",
    "    u = problem.solve()\n",
    "    w.vector.axpy(rho, u.vector)\n",
    "    if iters != 0:\n",
    "        u_m_u = (u.vector - u_old_vec).norm(2)\n",
    "    u_old_vec = u.vector.copy()\n",
    "    iters += 1\n",
    "    print(f'iteration {iters}: u_m_u = {u_m_u}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa8761a3",
   "metadata": {},
   "source": [
    "After that, we have to reconstruct the pressure field by computing $div(w)$\n",
    "and projecting the function onto the function space\n",
    "for the pressure field $Q$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd57bddc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reconstruct pressure field\n",
    "def project(v, target_func, bcs=[]):\n",
    "    # Ensure we have a mesh and attach to measure\n",
    "    V = target_func.function_space\n",
    "    dx = ufl.dx(V.mesh)\n",
    "\n",
    "    # Define variational problem for projection\n",
    "    w = ufl.TestFunction(V)\n",
    "    Pv = ufl.TrialFunction(V)\n",
    "    a = fem.form(ufl.inner(Pv, w) * dx)\n",
    "    L = fem.form(ufl.inner(v, w) * dx)\n",
    "\n",
    "    # Assemble linear system\n",
    "    A = fem.petsc.assemble_matrix(a, bcs)\n",
    "    A.assemble()\n",
    "    b = fem.petsc.assemble_vector(L)\n",
    "    fem.petsc.apply_lifting(b, [a], [bcs])\n",
    "    b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)\n",
    "    fem.petsc.set_bc(b, bcs)\n",
    "\n",
    "    # Solve linear system\n",
    "    solver = PETSc.KSP().create(A.getComm())\n",
    "    solver.setOperators(A)\n",
    "    solver.solve(b, target_func.vector)\n",
    "\n",
    "\n",
    "p = Function(Q)\n",
    "project(div(w), p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77a9640e",
   "metadata": {},
   "outputs": [],
   "source": [
    "coef_norm_u_4 = u.x.norm()\n",
    "coef_norm_p_4 = p.x.norm()\n",
    "l2_norm_u_4 = l2_norm(u)\n",
    "l2_norm_p_4 = l2_norm(p)\n",
    "if MPI.COMM_WORLD.rank == 0:\n",
    "    print(\"(4) Norm of velocity coeff. vector \"\n",
    "          \"with the Scott-Vogelius element:   {}\".format(coef_norm_u_4))\n",
    "    print(\"(4) Norm of pressure coeff. vector \"\n",
    "          \"with the Scott-Vogelius element:   {}\".format(coef_norm_p_4))\n",
    "    print(\"(4) L2 Norm of the velocity field \"\n",
    "          \"with the Scott-Vogelius element:    {}\".format(l2_norm_u_4))\n",
    "    print(\"(4) L2 Norm of pressure field \"\n",
    "          \"with the Scott-Vogelius element:        {}\".format(l2_norm_p_4))\n",
    "save_solution(u, \"out_stokes_stable_pairs/4_velocity.xdmf\")\n",
    "save_solution(p, \"out_stokes_stable_pairs/4_pressure.xdmf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3a3782a",
   "metadata": {},
   "source": [
    "## Interpretation\n",
    "\n",
    "We solved the Stokes equations for the lid driven cavity setting using\n",
    "different stable pairs of finite elements.\n",
    "Due to the rather coarse discretization, slight differences in the\n",
    "velocity field can be observed (e.g. in ParaView),\n",
    "but these gradually disappear with finer discretization.\n",
    "The pressure field is only determined except to a constant\n",
    "due to the homogeneous Neumann boundary conditions and therefore\n",
    "the L2 norms of the pressure field can differ\n",
    "from each other a lot. Due to the different basis functions,\n",
    "the solved coefficients are completely different\n",
    "despite the almost coinciding velocity field."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}