{
 "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)\n",
    "and see how to extend the mathematics and the implementation to handle Dirichlet condition\n",
    "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,\n",
    "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$\n",
    "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\n",
    "$g$ takes on the correct values at $y=0$ and $y=1$. One possible extension is\n",
    "\n",
    "$$\n",
    " g(x,y)=-4y.\n",
    "$$\n",
    "\n",
    "## The variational formulation\n",
    "The first task is to derive the variational formulation.\n",
    "This time we cannot omit the boundary term arising from integration by parts,\n",
    "because $v$ is only zero on $\\Lambda_D$. We have\n",
    "\n",
    "$$\n",
    "-\\int_\\Omega (\\nabla^2u)v~\\mathrm{d} x =\n",
    "\\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=\n",
    "- \\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 =\n",
    "\\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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dolfinx import default_scalar_type\n",
    "from dolfinx.fem import (\n",
    "    Constant,\n",
    "    Function,\n",
    "    functionspace,\n",
    "    assemble_scalar,\n",
    "    dirichletbc,\n",
    "    form,\n",
    "    locate_dofs_geometrical,\n",
    ")\n",
    "from dolfinx.fem.petsc import LinearProblem\n",
    "from dolfinx.mesh import create_unit_square\n",
    "from dolfinx.plot import vtk_mesh\n",
    "\n",
    "from mpi4py import MPI\n",
    "from ufl import SpatialCoordinate, TestFunction, TrialFunction, dot, ds, dx, grad\n",
    "\n",
    "import numpy as np\n",
    "import pyvista\n",
    "\n",
    "mesh = create_unit_square(MPI.COMM_WORLD, 10, 10)\n",
    "V = functionspace(mesh, (\"Lagrange\", 1))\n",
    "u = TrialFunction(V)\n",
    "v = TestFunction(V)\n",
    "a = dot(grad(u), grad(v)) * dx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "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 fulfill this condition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def u_exact(x):\n",
    "    return 1 + x[0] ** 2 + 2 * x[1] ** 2\n",
    "\n",
    "\n",
    "def boundary_D(x):\n",
    "    return np.logical_or(np.isclose(x[0], 0), np.isclose(x[0], 1))\n",
    "\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 disappears, and we can integrate `g*v*ds` over the entire boundary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = SpatialCoordinate(mesh)\n",
    "g = -4 * x[1]\n",
    "f = Constant(mesh, default_scalar_type(-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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "problem = LinearProblem(\n",
    "    a,\n",
    "    L,\n",
    "    bcs=[bc],\n",
    "    petsc_options={\"ksp_type\": \"preonly\", \"pc_type\": \"lu\"},\n",
    "    petsc_options_prefix=\"neumann_dirichlet_\",\n",
    ")\n",
    "uh = problem.solve()\n",
    "\n",
    "V2 = functionspace(mesh, (\"Lagrange\", 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": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "pyvista_cells, cell_types, geometry = 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",
    "    figure = plotter.screenshot(\"neumann_dirichlet.png\")"
   ]
  }
 ],
 "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
}