ch11-code-listing (1).ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chapter 11: Partial differential equations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Robert Johansson\n",
    "\n",
    "Source code listings for [Numerical Python - A Practical Techniques Approach for Industry](http://www.apress.com/9781484205549) (ISBN 978-1-484205-54-9).\n",
    "\n",
    "The source code listings can be downloaded from http://www.apress.com/9781484205549"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format='retina'\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mpl_toolkits.mplot3d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.sparse as sp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.sparse.linalg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.linalg as la"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finite-difference method"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1d example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Heat equation: \n",
    "\n",
    "$$-5 = u_{xx}, u(x=0) = 1, u(x=1) = 2$$\n",
    "\n",
    "$$\n",
    "u_{xx}[n] = (u[n-1] - 2u[n] + u[n+1])/dx^2\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u0 = 1\n",
    "u1 = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dx = 1.0 / (N + 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A = (np.eye(N, k=-1) - 2 * np.eye(N) + np.eye(N, k=1)) / dx**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d = -5 * np.ones(N)\n",
    "d[0] -= u0 / dx**2\n",
    "d[N-1] -= u1 / dx**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u = np.linalg.solve(A, d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(0, 1, N+2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "U = np.hstack([[u0], u, [u1]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(8, 4))\n",
    "ax.plot(x, U)\n",
    "ax.plot(x[1:-1], u, 'ks')\n",
    "ax.set_xlim(0, 1)\n",
    "ax.set_xlabel(r\"$x$\", fontsize=18)\n",
    "ax.set_ylabel(r\"$u(x)$\", fontsize=18)\n",
    "fig.savefig(\"ch11-fdm-1d.pdf\")\n",
    "fig.tight_layout();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2d example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "laplace equation: $u_{xx} + u_{yy} = 0$\n",
    "\n",
    "on boundary:\n",
    "\n",
    "$$\n",
    "u(x=0) = u(x=1) = u(y = 0) = u(y = 1) = 10\n",
    "$$\n",
    "\n",
    "$$\n",
    "u_{xx}[m, n] = (u[m-1, n] - 2u[m,n] + u[m+1,n])/dx^2\n",
    "$$\n",
    "\n",
    "$$\n",
    "u_{yy}[m, n] = (u[m, n-1] - 2u[m,n] + u[m,n+1])/dy^2\n",
    "$$\n",
    "\n",
    "final equation\n",
    "\n",
    "$$\n",
    "0 \n",
    "= \n",
    "(u[m-1 + N n] - 2u[m + N n] + u[m+1 + N n])/dx^2\n",
    "+\n",
    "(u[m + N *(n-1)] - 2u[m + N* n] + u[m + N(n+1]))/dy^2\n",
    "= \n",
    "(u[m + N n -1] - 2u[m + N n] + u[m + N n + 1])/dx^2\n",
    "+\n",
    "(u[m + N n -N)] - 2u[m + N n] + u[m + N n + N]))/dy^2\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u0_t, u0_b = 5, -5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u0_l, u0_r = 3, -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dx = 1. / (N+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A_1d = (sp.eye(N, k=-1) + sp.eye(N, k=1) - 4 * sp.eye(N))/dx**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A = sp.kron(sp.eye(N), A_1d) + (sp.eye(N**2, k=-N) + sp.eye(N**2, k=N))/dx**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A.nnz * 1.0/ np.prod(A.shape) * 2000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d = np.zeros((N, N))\n",
    "\n",
    "d[0, :] += -u0_b \n",
    "d[-1, :] += -u0_t\n",
    "d[:, 0] += -u0_l\n",
    "d[:, -1] += -u0_r\n",
    "\n",
    "d = d.reshape(N**2) / dx**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u = sp.linalg.spsolve(A, d).reshape(N, N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "U = np.vstack([np.ones((1, N+2)) * u0_b,\n",
    "               np.hstack([np.ones((N, 1)) * u0_l, u, np.ones((N, 1)) * u0_r]),\n",
    "               np.ones((1, N+2)) * u0_t])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(8, 6))\n",
    "\n",
    "x = np.linspace(0, 1, N+2)\n",
    "X, Y = np.meshgrid(x, x)\n",
    "\n",
    "c = ax.pcolor(X, Y, U, vmin=-5, vmax=5, cmap=mpl.cm.get_cmap('RdBu_r'))\n",
    "cb = plt.colorbar(c, ax=ax)\n",
    "\n",
    "ax.set_xlabel(r\"$x_1$\", fontsize=18)\n",
    "ax.set_ylabel(r\"$x_2$\", fontsize=18)\n",
    "cb.set_label(r\"$u(x_1, x_2)$\", fontsize=18)\n",
    "fig.savefig(\"ch11-fdm-2d.pdf\")\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(0, 1, N+2)\n",
    "X, Y = np.meshgrid(x, x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(12, 5.5))\n",
    "cmap = mpl.cm.get_cmap('RdBu_r')\n",
    "\n",
    "ax = fig.add_subplot(1, 2, 1)\n",
    "p = ax.pcolor(X, Y, U, vmin=-5, vmax=5, cmap=cmap)\n",
    "ax.set_xlabel(r\"$x_1$\", fontsize=18)\n",
    "ax.set_ylabel(r\"$x_2$\", fontsize=18)\n",
    "\n",
    "ax = fig.add_subplot(1, 2, 2, projection='3d')\n",
    "p = ax.plot_surface(X, Y, U, vmin=-5, vmax=5, rstride=3, cstride=3, linewidth=0, cmap=cmap)\n",
    "ax.set_xlabel(r\"$x_1$\", fontsize=16)\n",
    "ax.set_ylabel(r\"$x_2$\", fontsize=16)\n",
    "cb = plt.colorbar(p, ax=ax, shrink=0.75)\n",
    "cb.set_label(r\"$u(x_1, x_2)$\", fontsize=18)\n",
    "\n",
    "fig.savefig(\"ch11-fdm-2d.pdf\")\n",
    "fig.savefig(\"ch11-fdm-2d.png\")\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compare performance when using dense/sparse matrices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A_dense = A.todense()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit np.linalg.solve(A_dense, d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit la.solve(A_dense, d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit sp.linalg.spsolve(A, d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "10.8 / 31.9e-3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2d example with source term"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d = - np.ones((N, N))\n",
    "d = d.reshape(N**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u = sp.linalg.spsolve(A, d).reshape(N, N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "U = np.vstack([np.zeros((1, N+2)),\n",
    "               np.hstack([np.zeros((N, 1)), u, np.zeros((N, 1))]),\n",
    "               np.zeros((1, N+2))])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(0, 1, N+2)\n",
    "X, Y = np.meshgrid(x, x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(8, 6), subplot_kw={'projection': '3d'})\n",
    "\n",
    "p = ax.plot_surface(X, Y, U, rstride=4, cstride=4, linewidth=0, cmap=mpl.cm.get_cmap(\"Reds\"))\n",
    "cb = fig.colorbar(p, shrink=0.5)\n",
    "\n",
    "ax.set_xlabel(r\"$x_1$\", fontsize=18)\n",
    "ax.set_ylabel(r\"$x_2$\", fontsize=18)\n",
    "cb.set_label(r\"$u(x_1, x_2)$\", fontsize=18)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## FEniCS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dolfin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mshr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dolfin.parameters[\"reorder_dofs_serial\"] = False\n",
    "dolfin.parameters[\"allow_extrapolation\"] = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N1 = N2 = 75"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = dolfin.RectangleMesh(dolfin.Point(0, 0), dolfin.Point(1, 1), N1, N2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dolfin.RectangleMesh(dolfin.Point(0, 0), dolfin.Point(1, 1), 10, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function space from mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = dolfin.FunctionSpace(mesh, 'Lagrange', 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Variational problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u = dolfin.TrialFunction(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v = dolfin.TestFunction(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = dolfin.inner(dolfin.nabla_grad(u), dolfin.nabla_grad(v)) * dolfin.dx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f1 = dolfin.Constant(1.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "L1 = f1 * v * dolfin.dx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f2 = dolfin.Constant(1.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f2 = dolfin.Expression(\"x[0]*x[0] + x[1]*x[1]\", degree=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "L2 = f2 * v * dolfin.dx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Boundary conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u0 = dolfin.Constant(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def u0_boundary(x, on_boundary):\n",
    "    # try to pin down the function at some interior region:\n",
    "    #if np.sqrt((x[0]-0.5)**2 + (x[1]-0.5)**2) < 0.1:\n",
    "    #    return True\n",
    "    return on_boundary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bc = dolfin.DirichletBC(V, u0, u0_boundary)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solve the problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A = dolfin.assemble(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "b = dolfin.assemble(L1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bc.apply(A, b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_sol1 = dolfin.Function(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dolfin.solve(A, u_sol1.vector(), b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_sol2 = dolfin.Function(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dolfin.solve(a == L2, u_sol2, bc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dolfin plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dolfin.plot(u_sol1)\n",
    "# dolfin.interactive()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save VTK files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dolfin.File('u_sol1.pvd') << u_sol1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dolfin.File('u_sol2.pvd') << u_sol2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = dolfin.File('combined.pvd')\n",
    "f << mesh\n",
    "f << u_sol1\n",
    "f << u_sol2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_sol1([0.21, 0.67])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Obtain NumPy arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_mat1 = np.array(u_sol1.vector()).reshape(N1+1, N2+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_mat2 = np.array(u_sol2.vector()).reshape(N1+1, N2+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X, Y = np.meshgrid(np.linspace(0, 1, N1+2), np.linspace(0, 1, N2+2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))\n",
    "cmap = mpl.cm.get_cmap('Reds')\n",
    "\n",
    "c = ax1.pcolor(X, Y, u_mat1, cmap=cmap)\n",
    "cb = plt.colorbar(c, ax=ax1)\n",
    "ax1.set_xlabel(r\"$x$\", fontsize=18)\n",
    "ax1.set_ylabel(r\"$y$\", fontsize=18)\n",
    "cb.set_label(r\"$u(x, y)$\", fontsize=18)\n",
    "cb.set_ticks([0.0, 0.02, 0.04, 0.06])\n",
    "\n",
    "c = ax2.pcolor(X, Y, u_mat2, cmap=cmap)\n",
    "cb = plt.colorbar(c, ax=ax2)\n",
    "ax1.set_xlabel(r\"$x$\", fontsize=18)\n",
    "ax1.set_ylabel(r\"$y$\", fontsize=18)\n",
    "cb.set_label(r\"$u(x, y)$\", fontsize=18)\n",
    "cb.set_ticks([0.0, 0.02, 0.04])\n",
    "\n",
    "fig.savefig(\"ch11-fdm-2d-ex1.pdf\")\n",
    "fig.savefig(\"ch11-fdm-2d-ex1.png\")\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X, Y = np.meshgrid(np.linspace(0, 1, N1+1), np.linspace(0, 1, N2+1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6), subplot_kw={'projection': '3d'})\n",
    "\n",
    "p = ax1.plot_surface(X, Y, u_mat1, rstride=4, cstride=4, linewidth=0, cmap=mpl.cm.get_cmap(\"Reds\"))\n",
    "cb = fig.colorbar(p, ax=ax1, shrink=0.5)\n",
    "ax1.set_xlabel(r\"$x_1$\", fontsize=18)\n",
    "ax1.set_ylabel(r\"$x_2$\", fontsize=18)\n",
    "cb.set_label(r\"$u(x_1, x_2)$\", fontsize=18)\n",
    "\n",
    "p = ax2.plot_surface(X, Y, u_mat2, rstride=4, cstride=4, linewidth=0, cmap=mpl.cm.get_cmap(\"Reds\"))\n",
    "cb = fig.colorbar(p, ax=ax2, shrink=0.5)\n",
    "ax2.set_xlabel(r\"$x_1$\", fontsize=18)\n",
    "ax2.set_ylabel(r\"$x_2$\", fontsize=18)\n",
    "cb.set_label(r\"$u(x_1, x_2)$\", fontsize=18)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Triangulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mesh_triangulation(mesh):\n",
    "    coordinates = mesh.coordinates()\n",
    "    triangles = mesh.cells()\n",
    "    triangulation = mpl.tri.Triangulation(coordinates[:, 0], coordinates[:, 1], triangles)\n",
    "    return triangulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "triangulation = mesh_triangulation(mesh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))\n",
    "\n",
    "ax1.triplot(triangulation)\n",
    "ax1.set_xlabel(r\"$x$\", fontsize=18)\n",
    "ax1.set_ylabel(r\"$y$\", fontsize=18)\n",
    "\n",
    "c = ax2.tripcolor(triangulation, np.array(u_sol2.vector()), cmap=cmap)\n",
    "cb = plt.colorbar(c, ax=ax2)\n",
    "ax2.set_xlabel(r\"$x$\", fontsize=18)\n",
    "ax2.set_ylabel(r\"$y$\", fontsize=18)\n",
    "cb.set_label(r\"$u(x, y)$\", fontsize=18)\n",
    "cb.set_ticks([0.0, 0.02, 0.04])\n",
    "\n",
    "fig.savefig(\"ch11-fdm-2d-ex2.pdf\")\n",
    "fig.savefig(\"ch11-fdm-2d-ex2.png\")\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dirichlet boundary conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N1 = N2 = 75"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = dolfin.RectangleMesh(dolfin.Point(0, 0), dolfin.Point(1, 1), N1, N2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = dolfin.FunctionSpace(mesh, 'Lagrange', 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u = dolfin.TrialFunction(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v = dolfin.TestFunction(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = dolfin.inner(dolfin.nabla_grad(u), dolfin.nabla_grad(v)) * dolfin.dx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = dolfin.Constant(0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "L = f * v * dolfin.dx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def u0_top_boundary(x, on_boundary):\n",
    "    return on_boundary and abs(x[1]-1) < 1e-8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def u0_bottom_boundary(x, on_boundary):\n",
    "    return on_boundary and abs(x[1]) < 1e-8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def u0_left_boundary(x, on_boundary):\n",
    "    return on_boundary and abs(x[0]) < 1e-8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def u0_right_boundary(x, on_boundary):\n",
    "    return on_boundary and abs(x[0]-1) < 1e-8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bc_t = dolfin.DirichletBC(V, dolfin.Constant(5), u0_top_boundary)\n",
    "bc_b = dolfin.DirichletBC(V, dolfin.Constant(-5), u0_bottom_boundary)\n",
    "bc_l = dolfin.DirichletBC(V, dolfin.Constant(3), u0_left_boundary)\n",
    "bc_r = dolfin.DirichletBC(V, dolfin.Constant(-1), u0_right_boundary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bcs = [bc_t, bc_b, bc_r, bc_l]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_sol = dolfin.Function(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dolfin.solve(a == L, u_sol, bcs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_mat = np.array(u_sol.vector()).reshape(N1+1, N2+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(0, 1, N1+2)\n",
    "y = np.linspace(0, 1, N1+2)\n",
    "X, Y = np.meshgrid(x, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(8, 6))\n",
    "\n",
    "c = ax.pcolor(X, Y, u_mat, vmin=-5, vmax=5, cmap=mpl.cm.get_cmap('RdBu_r'))\n",
    "cb = plt.colorbar(c, ax=ax)\n",
    "ax.set_xlabel(r\"$x_1$\", fontsize=18)\n",
    "ax.set_ylabel(r\"$x_2$\", fontsize=18)\n",
    "cb.set_label(r\"$u(x_1, x_2)$\", fontsize=18)\n",
    "fig.savefig(\"ch11-fdm-2d-ex3.pdf\")\n",
    "fig.savefig(\"ch11-fdm-2d-ex3.png\")\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Circular geometry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r_outer = 1\n",
    "r_inner = 0.25\n",
    "r_middle = 0.1\n",
    "x0, y0 = 0.4, 0.4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "domain = mshr.Circle(dolfin.Point(.0, .0), r_outer) \\\n",
    "    - mshr.Circle(dolfin.Point(.0, .0), r_inner) \\\n",
    "    - mshr.Circle(dolfin.Point( x0,  y0), r_middle) \\\n",
    "    - mshr.Circle(dolfin.Point( x0, -y0), r_middle) \\\n",
    "    - mshr.Circle(dolfin.Point(-x0,  y0), r_middle) \\\n",
    "    - mshr.Circle(dolfin.Point(-x0, -y0), r_middle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = mshr.generate_mesh(domain, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = dolfin.FunctionSpace(mesh, 'Lagrange', 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u = dolfin.TrialFunction(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v = dolfin.TestFunction(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = dolfin.inner(dolfin.nabla_grad(u), dolfin.nabla_grad(v)) * dolfin.dx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = dolfin.Constant(1.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "L = f * v * dolfin.dx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def u0_outer_boundary(x, on_boundary):\n",
    "    x, y = x[0], x[1]\n",
    "    return on_boundary and abs(np.sqrt(x**2 + y**2) - r_outer) < 5e-2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def u0_inner_boundary(x, on_boundary):\n",
    "    x, y = x[0], x[1]\n",
    "    return on_boundary and abs(np.sqrt(x**2 + y**2) - r_inner) < 5e-2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def u0_middle_boundary(x, on_boundary):\n",
    "    x, y = x[0], x[1]\n",
    "    if on_boundary:\n",
    "        for _x0 in [-x0, x0]:\n",
    "            for _y0 in [-y0, y0]:\n",
    "                if abs(np.sqrt((x+_x0)**2 + (y+_y0)**2) - r_middle) < 5e-2:\n",
    "                    return True\n",
    "    \n",
    "    return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bc_inner = dolfin.DirichletBC(V, dolfin.Constant(15), u0_inner_boundary)\n",
    "bc_middle = dolfin.DirichletBC(V, dolfin.Constant(0), u0_middle_boundary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bcs = [bc_inner, bc_middle]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_sol = dolfin.Function(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dolfin.solve(a == L, u_sol, bcs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "triangulation = mesh_triangulation(mesh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))\n",
    "\n",
    "ax1.triplot(triangulation)\n",
    "ax1.set_xlabel(r\"$x$\", fontsize=18)\n",
    "ax1.set_ylabel(r\"$y$\", fontsize=18)\n",
    "\n",
    "c = ax2.tripcolor(triangulation, np.array(u_sol.vector()), cmap=mpl.cm.get_cmap(\"Reds\"))\n",
    "cb = plt.colorbar(c, ax=ax2)\n",
    "ax2.set_xlabel(r\"$x$\", fontsize=18)\n",
    "ax2.set_ylabel(r\"$y$\", fontsize=18)\n",
    "cb.set_label(r\"$u(x, y)$\", fontsize=18)\n",
    "cb.set_ticks([0.0, 5, 10, 15])\n",
    "\n",
    "fig.savefig(\"ch11-fdm-2d-ex4.pdf\")\n",
    "fig.savefig(\"ch11-fdm-2d-ex4.png\")\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Post processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "outer_boundary = dolfin.AutoSubDomain(lambda x, on_bnd: on_bnd and abs(np.sqrt(x[0]**2 + x[1]**2) - r_outer) < 5e-2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bc_outer = dolfin.DirichletBC(V, 1, outer_boundary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mask_outer = dolfin.Function(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bc_outer.apply(mask_outer.vector())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_outer = u_sol.vector()[mask_outer.vector() == 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_outer = mesh.coordinates()[mask_outer.vector() == 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "phi = np.angle(x_outer[:, 0] + 1j * x_outer[:, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "order = np.argsort(phi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(8, 4))\n",
    "ax.plot(phi[order], u_outer[order], 's-', lw=2)\n",
    "ax.set_ylabel(r\"$u(x,y)$ at $x^2+y^2=1$\", fontsize=18)\n",
    "ax.set_xlabel(r\"$\\phi$\", fontsize=18)\n",
    "ax.set_xlim(-np.pi, np.pi)\n",
    "fig.tight_layout()\n",
    "fig.savefig(\"ch11-fem-2d-ex5.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Mesh refining"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "domain = mshr.Circle(dolfin.Point(.0, .0), 1.0) - mshr.Rectangle(dolfin.Point(0.0, -1.0), dolfin.Point(1.0, 0.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = mshr.generate_mesh(domain, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "refined_mesh.topology().dim()-1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "refined_mesh = mesh\n",
    "for r in [0.5, 0.25]:\n",
    "    cell_markers = dolfin.MeshFunction(\"bool\", refined_mesh, 2)\n",
    "    cell_markers.set_all(False)\n",
    "    for cell in dolfin.cells(refined_mesh):\n",
    "        if cell.distance(dolfin.Point(.0, .0)) < r:\n",
    "            cell_markers[cell] = True\n",
    "    refined_mesh = dolfin.refine(refined_mesh, cell_markers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 4))\n",
    " \n",
    "ax1.triplot(mesh_triangulation(mesh))\n",
    "ax2.triplot(mesh_triangulation(refined_mesh))\n",
    "\n",
    "for ax in [ax1, ax2]:\n",
    "    for side in ['bottom','right','top','left']:\n",
    "        ax.spines[side].set_visible(False)\n",
    "        ax.set_xticks([])\n",
    "        ax.set_yticks([])\n",
    "        ax.xaxis.set_ticks_position('none')\n",
    "        ax.yaxis.set_ticks_position('none')\n",
    "\n",
    "    ax.set_xlabel(r\"$x$\", fontsize=18)\n",
    "    ax.set_ylabel(r\"$y$\", fontsize=18)\n",
    " \n",
    "fig.savefig(\"ch11-fem-2d-mesh-refine.pdf\")\n",
    "fig.savefig(\"ch11-fem-2d-mesh-refine.png\")\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Refined mesh with Dirichlet boundary conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = refined_mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = dolfin.FunctionSpace(mesh, 'Lagrange', 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u = dolfin.TrialFunction(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v = dolfin.TestFunction(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = dolfin.inner(dolfin.nabla_grad(u), dolfin.nabla_grad(v)) * dolfin.dx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = dolfin.Constant(0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "L = f * v * dolfin.dx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def u0_vertical_boundary(x, on_boundary):\n",
    "    x, y = x[0], x[1]\n",
    "    return on_boundary and abs(x) < 1e-2 and y < 0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def u0_horizontal_boundary(x, on_boundary):\n",
    "    x, y = x[0], x[1]\n",
    "    return on_boundary and abs(y) < 1e-2 and x > 0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bc_vertical = dolfin.DirichletBC(V, dolfin.Constant(-5), u0_vertical_boundary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bc_horizontal = dolfin.DirichletBC(V, dolfin.Constant(5), u0_horizontal_boundary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bcs = [bc_vertical, bc_horizontal]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_sol = dolfin.Function(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dolfin.solve(a == L, u_sol, bcs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "triangulation = mesh_triangulation(mesh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))\n",
    " \n",
    "ax1.triplot(triangulation, color='k')\n",
    "ax1.set_xlabel(r\"$x$\", fontsize=18)\n",
    "ax1.set_ylabel(r\"$y$\", fontsize=18)\n",
    " \n",
    "c = ax2.tripcolor(triangulation, np.array(u_sol.vector()), cmap=mpl.cm.get_cmap(\"RdBu_r\"))\n",
    "cb = plt.colorbar(c, ax=ax2)\n",
    "ax2.set_xlabel(r\"$x$\", fontsize=18)\n",
    "ax2.set_ylabel(r\"$y$\", fontsize=18)\n",
    "cb.set_label(r\"$u(x, y)$\", fontsize=18)\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Refined mesh with Dirichlet and von Neumann boundary conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = refined_mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = dolfin.FunctionSpace(mesh, 'Lagrange', 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u = dolfin.TrialFunction(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v = dolfin.TestFunction(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "boundary_parts = dolfin.MeshFunction(\"size_t\", mesh, mesh.topology().dim()-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def v_boundary_func(x, on_boundary):\n",
    "    \"\"\" the vertical edge of the mesh, where x = 0 and y < 0\"\"\"\n",
    "    x, y = x[0], x[1]\n",
    "    return on_boundary and abs(x) < 1e-4 and y < 0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v_boundary = dolfin.AutoSubDomain(v_boundary_func)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v_boundary.mark(boundary_parts, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def h_boundary_func(x, on_boundary):\n",
    "    \"\"\" the horizontal edge of the mesh, where y = 0 and x > 0\"\"\"\n",
    "    x, y = x[0], x[1]\n",
    "    return on_boundary and abs(y) < 1e-4 and x > 0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h_boundary = dolfin.AutoSubDomain(h_boundary_func)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h_boundary.mark(boundary_parts, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def outer_boundary_func(x, on_boundary):\n",
    "    x, y = x[0], x[1]\n",
    "    return on_boundary and abs(x**2 + y**2-1) < 1e-2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "outer_boundary = dolfin.AutoSubDomain(outer_boundary_func)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "outer_boundary.mark(boundary_parts, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bc = dolfin.DirichletBC(V, dolfin.Constant(0.0), boundary_parts, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = dolfin.inner(dolfin.nabla_grad(u), dolfin.nabla_grad(v)) * dolfin.dx(domain=mesh, subdomain_data=boundary_parts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = dolfin.Constant(0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g_v = dolfin.Constant(-2.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g_h = dolfin.Constant(1.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "L = f * v * dolfin.dx(domain=mesh, subdomain_data=boundary_parts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "L += g_v * v * dolfin.ds(0, domain=mesh, subdomain_data=boundary_parts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "L += g_h * v * dolfin.ds(1, domain=mesh, subdomain_data=boundary_parts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_sol = dolfin.Function(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dolfin.solve(a == L, u_sol, bc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "triangulation = mesh_triangulation(mesh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))\n",
    " \n",
    "ax1.triplot(triangulation)\n",
    "ax1.set_xlabel(r\"$x$\", fontsize=18)\n",
    "ax1.set_ylabel(r\"$y$\", fontsize=18)\n",
    " \n",
    "data = np.array(u_sol.vector())\n",
    "norm = mpl.colors.Normalize(-abs(data).max(), abs(data).max())\n",
    "c = ax2.tripcolor(triangulation, data, norm=norm, cmap=mpl.cm.get_cmap(\"RdBu_r\"))\n",
    "cb = plt.colorbar(c, ax=ax2)\n",
    "ax2.set_xlabel(r\"$x$\", fontsize=18)\n",
    "ax2.set_ylabel(r\"$y$\", fontsize=18)\n",
    "cb.set_label(r\"$u(x, y)$\", fontsize=18)\n",
    "cb.set_ticks([-.5, 0, .5])\n",
    " \n",
    "fig.savefig(\"ch11-fem-2d-ex5.pdf\")\n",
    "fig.savefig(\"ch11-fem-2d-ex5.png\")\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Versions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext version_information\n",
    "%version_information numpy, scipy, matplotlib, dolfin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.6 (NPM fenics)",
   "language": "python",
   "name": "py3.6-npm-fenics"
  },
  "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}