{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Weak imposition of Dirichlet conditions for the Poisson problem\n",
    "Author: Jørgen S. Dokken\n",
    "\n",
    "In this section, we will go through how to solve the Poisson problem from the\n",
    "[Fundamentals](./fundamentals_code.ipynb) tutorial using Nitsche's method {cite}`nitsche-Nitsche1971`.\n",
    "The idea of weak imposition is that we add additional terms to the variational formulation to\n",
    "impose the boundary condition, instead of modifying the matrix system using strong imposition (lifting).\n",
    "\n",
    "We start by importing the required modules and creating the mesh and function space for our solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dolfinx import fem, mesh, plot, default_scalar_type\n",
    "from dolfinx.fem.petsc import LinearProblem\n",
    "import numpy\n",
    "from mpi4py import MPI\n",
    "from ufl import (\n",
    "    Circumradius,\n",
    "    FacetNormal,\n",
    "    SpatialCoordinate,\n",
    "    TrialFunction,\n",
    "    TestFunction,\n",
    "    div,\n",
    "    dx,\n",
    "    ds,\n",
    "    grad,\n",
    "    inner,\n",
    ")\n",
    "\n",
    "N = 8\n",
    "domain = mesh.create_unit_square(MPI.COMM_WORLD, N, N)\n",
    "V = fem.functionspace(domain, (\"Lagrange\", 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "Next, we create a function containing the exact solution (which will also be used in the Dirichlet boundary condition)\n",
    "and the corresponding source function for the right hand side. Note that we use {py:class}`ufl.SpatialCoordinate`\n",
    "to define the exact solution, which in turn is interpolated into {py:class}`uD<dolfinx.fem.Function>`\n",
    "by wrapping the {py:class}`ufl-expression<ufl.core.expr.Expr>` as a {py:class}`dolfinx.fem.Expression`.\n",
    "For the source function, we use the symbolic differentiation capabilities of UFL to\n",
    "compute the negative Laplacian of the exact solution, which we can use directly in the variational formulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "uD = fem.Function(V)\n",
    "x = SpatialCoordinate(domain)\n",
    "u_ex = 1 + x[0] ** 2 + 2 * x[1] ** 2\n",
    "uD.interpolate(fem.Expression(u_ex, V.element.interpolation_points))\n",
    "f = -div(grad(u_ex))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "As opposed to the first tutorial, we now have to have another look at the variational form.\n",
    "We start by integrating the problem by parts, to obtain\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "  \\int_{\\Omega} \\nabla u \\cdot \\nabla v~\\mathrm{d}x\n",
    "  - \\int_{\\partial\\Omega}\\nabla u \\cdot n v~\\mathrm{d}s = \\int_{\\Omega} f v~\\mathrm{d}x.\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "As we are not using strong enforcement, we do not set the trace of the test function to $0$ on the outer boundary.\n",
    "Instead, we add the following two terms to the variational formulation\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "  -\\int_{\\partial\\Omega} \\nabla  v \\cdot n (u-u_D)~\\mathrm{d}s\n",
    "  + \\frac{\\alpha}{h} \\int_{\\partial\\Omega} (u-u_D)v~\\mathrm{d}s.\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "where the first term enforces symmetry to the bilinear form, while the latter term enforces coercivity.\n",
    "$u_D$ is the known Dirichlet condition, and $h$ is the diameter of the circumscribed sphere of the mesh element.\n",
    "We create bilinear and linear form, $a$ and $L$\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "    a(u, v) &= \\int_{\\Omega} \\nabla u \\cdot \\nabla v~\\mathrm{d}x + \\int_{\\partial\\Omega}-(n \\cdot\\nabla u) v - (n \\cdot \\nabla v) u + \\frac{\\alpha}{h} uv~\\mathrm{d}s,\\\\\n",
    "    L(v) &= \\int_{\\Omega} fv~\\mathrm{d}x + \\int_{\\partial\\Omega} -(n \\cdot \\nabla v) u_D + \\frac{\\alpha}{h} u_Dv~\\mathrm{d}s\n",
    "\\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "u = TrialFunction(V)\n",
    "v = TestFunction(V)\n",
    "n = FacetNormal(domain)\n",
    "h = 2 * Circumradius(domain)\n",
    "alpha = fem.Constant(domain, default_scalar_type(10))\n",
    "a = inner(grad(u), grad(v)) * dx - inner(n, grad(u)) * v * ds\n",
    "a += -inner(n, grad(v)) * u * ds + alpha / h * inner(u, v) * ds\n",
    "L = inner(f, v) * dx\n",
    "L += -inner(n, grad(v)) * uD * ds + alpha / h * inner(uD, v) * ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "As we now have the variational form, we can solve the arising\n",
    "{py:class}`linear problem<dolfinx.fem.petsc.LinearProblem>`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "problem = LinearProblem(a, L, petsc_options_prefix=\"nitsche_poisson\")\n",
    "uh = problem.solve()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "We compute the error of the computation by comparing it to the analytical solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "error_form = fem.form(inner(uh - uD, uh - uD) * dx)\n",
    "error_local = fem.assemble_scalar(error_form)\n",
    "errorL2 = numpy.sqrt(domain.comm.allreduce(error_local, op=MPI.SUM))\n",
    "if domain.comm.rank == 0:\n",
    "    print(rf\"$L^2$-error: {errorL2:.2e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "We observe that the $L^2$-error is of the same magnitude as in the first tutorial.\n",
    "As in the previous tutorial, we also compute the maximal error for all the degrees of freedom."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "error_max = domain.comm.allreduce(\n",
    "    numpy.max(numpy.abs(uD.x.array - uh.x.array)), op=MPI.MAX\n",
    ")\n",
    "if domain.comm.rank == 0:\n",
    "    print(f\"Error_max : {error_max:.2e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "We observe that as we weakly impose the boundary condition,\n",
    "we no longer fullfill the equation to machine precision at the mesh vertices.\n",
    "We also plot the solution using {py:mod}`pyvista`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyvista\n",
    "\n",
    "grid = pyvista.UnstructuredGrid(*plot.vtk_mesh(V))\n",
    "grid.point_data[\"u\"] = uh.x.array.real\n",
    "grid.set_active_scalars(\"u\")\n",
    "plotter = pyvista.Plotter()\n",
    "plotter.add_mesh(grid, show_edges=True, show_scalar_bar=True)\n",
    "plotter.view_xy()\n",
    "if not pyvista.OFF_SCREEN:\n",
    "    plotter.show()\n",
    "else:\n",
    "    figure = plotter.screenshot(\"nitsche.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "```{bibliography}\n",
    "   :filter: cited\n",
    "   :labelprefix:\n",
    "   :keyprefix: nitsche-\n",
    "```"
   ]
  }
 ],
 "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": 5
}