{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Combining Dirichlet and Neumann conditions\n",
    "Author: Jørgen S. Dokken\n",
    "\n",
    "Let's return to the Poisson problem from the [Fundamentals chapter](./../chapter1/fundamentals.md) and see how to extend the mathematics and the implementation to handle Dirichlet condition in combination with a Neumann condition.\n",
    "The domain is still the unit square, but now we set the Dirichlet condition $u=u_D$ at the left and right sides, while the Neumann condition \n",
    "\n",
    "$$\n",
    "-\\frac{\\partial u}{\\partial n}=g\n",
    "$$\n",
    "\n",
    "is applied to the  remaining sides $y=0$ and $y=1$.\n",
    "\n",
    "## The PDE problem\n",
    "Let $\\Lambda_D$ and $\\Lambda_N$ denote parts of the boundary $\\partial \\Omega$ where the Dirichlet and Neumann conditions apply, respectively.\n",
    "The complete boundary-value problem can be written as\n",
    "\n",
    "$$\n",
    "-\\nabla^2 u =f \\qquad \\text{in } \\Omega,\n",
    "$$\n",
    "$$\n",
    "u=u_D \\qquad\\text{on } \\Lambda_D,\n",
    "$$\n",
    "$$\n",
    "-\\frac{\\partial u}{\\partial n}=g \\qquad \\text{on }\\Lambda_N\n",
    "$$\n",
    "\n",
    "Again, we choose $u=1+x^2+2y^2$ as the exact solution and adjust $f, g,$ and $u_D$ accordingly\n",
    "\n",
    "$$\n",
    "f(x,y)=-6,\n",
    "$$\n",
    "$$\n",
    "g(x,y)=\\begin{cases}\n",
    "0, & y=0,\\\\\n",
    "-4, & y=1,\n",
    "\\end{cases}\n",
    "$$\n",
    "$$\n",
    "u_D(x,y)=1+x^2+2y^2.\n",
    "$$\n",
    "\n",
    "For the ease of programming, we define $g$ as a function over the whole domain $\\Omega$ such that $g$ takes on the correct values at $y=0$ and $y=1$. One possible extension is\n",
    "$$\n",
    " g(x,y)=-4y.\n",
    "$$\n",
    "## The variational formulation\n",
    "The first task is to derive the variational formulatin. This time we cannot omit the boundary term arising from integration by parts, because $v$ is only zero on $\\Lambda_D$. We have\n",
    "\n",
    "$$\n",
    "-\\int_\\Omega (\\nabla^2u)v~\\mathrm{d} x = \\int_\\Omega \\nabla u \\cdot \\nabla v ~\\mathrm{d} x - \\int_{\\partial\\Omega}\\frac{\\partial u}{\\partial n}v~\\mathrm{d}s,\n",
    "$$\n",
    "\n",
    "and since $v=0$ on $\\Lambda_D$,\n",
    "\n",
    "$$\n",
    "- \\int_{\\partial\\Omega}\\frac{\\partial u}{\\partial n}v~\\mathrm{d}s= - \\int_{\\Lambda_N}\\frac{\\partial u}{\\partial n}v~\\mathrm{d}s =\\int_{\\Lambda_N} gv~\\mathrm{d}s,\n",
    "$$\n",
    "\n",
    "by applying the boundary condition on $\\Lambda_N$.\n",
    "The resulting weak from reads\n",
    "\n",
    "$$\n",
    "    \\int_\\Omega \\nabla u \\cdot \\nabla v~\\mathrm{d} x = \\int_\\Omega fv~\\mathrm{d} x - \\int_{\\Lambda_N}gv~\\mathrm{d}s.\n",
    "$$\n",
    "Expressing this equation in the standard notation $a(u,v)=L(v)$ is straight-forward with \n",
    "\n",
    "$$\n",
    "    a(u,v) = \\int_{\\Omega} \\nabla u \\cdot \\nabla v ~\\mathrm{d} x,\\\\\n",
    "$$\n",
    "$$\n",
    "L(v) = \\int_{\\Omega} fv ~\\mathrm{d} x - \\int_{\\Lambda_N} gv~\\mathrm{d} s.\n",
    "$$\n",
    "\n",
    "## Implementation\n",
    "As in the previous example, we define our mesh,function space and bilinear form $a(u,v)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pyvista\n",
    "\n",
    "from dolfinx.fem import (Constant, Function, FunctionSpace, \n",
    "                         assemble_scalar, dirichletbc, form, locate_dofs_geometrical)\n",
    "from dolfinx.fem.petsc import LinearProblem\n",
    "from dolfinx.mesh import create_unit_square\n",
    "from mpi4py import MPI\n",
    "from petsc4py.PETSc import ScalarType\n",
    "from ufl import SpatialCoordinate, TestFunction, TrialFunction, dot, ds, dx, grad\n",
    "\n",
    "mesh = create_unit_square(MPI.COMM_WORLD, 10, 10)\n",
    "V = FunctionSpace(mesh, (\"CG\", 1))\n",
    "u = TrialFunction(V)\n",
    "v = TestFunction(V)\n",
    "a = dot(grad(u), grad(v)) * dx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we get to the Neumann and Dirichlet boundary condition. As previously, we use a Python-function to define the boundary where we should have a Dirichlet condition. Then, with this function, we locate degrees of freedom that fullfils this condition. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def u_exact(x):\n",
    "    return 1 + x[0]**2 + 2*x[1]**2\n",
    "\n",
    "def boundary_D(x):\n",
    "    return np.logical_or(np.isclose(x[0], 0), np.isclose(x[0],1))\n",
    "\n",
    "dofs_D = locate_dofs_geometrical(V, boundary_D)\n",
    "u_bc = Function(V)\n",
    "u_bc.interpolate(u_exact)\n",
    "bc = dirichletbc(u_bc, dofs_D)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next step is to define the Neumann condition. We first define $g$ uses `UFL`s `SpatialCoordinate`-function, and then in turn create a boundary integration measure `ds`. As the test function $v$ is zero on the boundary integrals over the Dirichlet boundary dissapears, and wee can integrate `g*v*ds` over the entire boundary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = SpatialCoordinate(mesh)\n",
    "g = -4 * x[1]\n",
    "f = Constant(mesh, ScalarType(-6))\n",
    "L = f * v * dx - g * v * ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now assemble and solve the linear system of equations "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error_L2 : 5.27e-03\n",
      "Error_max : 2.44e-15\n"
     ]
    }
   ],
   "source": [
    "problem = LinearProblem(a, L, bcs=[bc], petsc_options={\"ksp_type\": \"preonly\", \"pc_type\": \"lu\"})\n",
    "uh = problem.solve()\n",
    "\n",
    "V2 = FunctionSpace(mesh, (\"CG\", 2))\n",
    "uex = Function(V2)\n",
    "uex.interpolate(u_exact)\n",
    "error_L2 = assemble_scalar(form((uh - uex)**2 * dx))\n",
    "error_L2 = np.sqrt(MPI.COMM_WORLD.allreduce(error_L2, op=MPI.SUM))\n",
    "\n",
    "u_vertex_values = uh.x.array\n",
    "uex_1 = Function(V)\n",
    "uex_1.interpolate(uex)\n",
    "u_ex_vertex_values = uex_1.x.array\n",
    "error_max = np.max(np.abs(u_vertex_values - u_ex_vertex_values))\n",
    "error_max = MPI.COMM_WORLD.allreduce(error_max, op=MPI.MAX)\n",
    "print(f\"Error_L2 : {error_L2:.2e}\")\n",
    "print(f\"Error_max : {error_max:.2e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualization\n",
    "To look at the actual solution, run the script as a python script with `off_screen=True` or as a Jupyter notebook with `off_screen=False`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[0m\u001b[2m2022-08-02 20:57:18.646 (   2.918s) [        3A132000]    vtkExtractEdges.cxx:435   INFO| \u001b[0mExecuting edge extractor: points are renumbered\u001b[0m\n",
      "\u001b[0m\u001b[2m2022-08-02 20:57:18.646 (   2.919s) [        3A132000]    vtkExtractEdges.cxx:551   INFO| \u001b[0mCreated 320 edges\u001b[0m\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4be044a073774ccaa5011c3bd955c546",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Renderer(camera=PerspectiveCamera(aspect=1.3333333333333333, children=(DirectionalLight(intensity=0.25, positi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pyvista.set_jupyter_backend(\"pythreejs\")\n",
    "\n",
    "from dolfinx.plot import create_vtk_mesh\n",
    "pyvista_cells, cell_types, geometry = create_vtk_mesh(V)\n",
    "grid = pyvista.UnstructuredGrid(pyvista_cells, cell_types, geometry)\n",
    "grid.point_data[\"u\"] = uh.x.array\n",
    "grid.set_active_scalars(\"u\")\n",
    "\n",
    "plotter = pyvista.Plotter()\n",
    "plotter.add_text(\"uh\", position=\"upper_edge\", font_size=14, color=\"black\")\n",
    "plotter.add_mesh(grid, show_edges=True)\n",
    "plotter.view_xy()\n",
    "\n",
    "if not pyvista.OFF_SCREEN:\n",
    "    plotter.show()\n",
    "else:\n",
    "    pyvista.start_xvfb()\n",
    "    figure = plotter.screenshot(\"neumann_dirichlet.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}