Example_Poisson_Equation_Solver_by_FEM.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Homework 5 \"Putting them all together\" Solution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solve the following boundary value problem using the \"standard\" (i.e., Ritz-Galerkin-Courant) finite element method.\n",
    "\n",
    "\\begin{align*}\n",
    "\\nabla \\cdot (a(\\mathbf{x})\\nabla u) &= 0 \\quad \\text{in} \\quad \\Omega, \\\\\n",
    "u &= 10 \\quad \\text{on} \\quad \\Gamma_{g1} \\\\\n",
    "u &= 0 \\quad \\text{on} \\quad \\Gamma_{g2}.\n",
    "\\end{align*}\n",
    "\n",
    "<img src=\"./Figures/domain.png\" width=\"480\">\n",
    "\n",
    "The problem domain ($\\Omega$) is the gray-shaded area in the above figure and the two Dirichlet boundaries ($\\Gamma_{g1}$ and $\\Gamma_{g2}$) are marked by bold lines. $a(\\mathbf{x})$ in the above equation is a spatially variable coefficient given as 2 in the dark gray region and 1 in the light gray region. On the rest of the domain boundary, it is assumed that $\\partial u/\\partial n = 0$, where $n$ is the distance along the outward normal direction on the boundary.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load basic modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import numpy.linalg as la\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.tri as mtri"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Discretization Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nnode = 42\n",
    "nelem = 40\n",
    "ndim = 2\n",
    "NODES_PER_ELEMENT = 3 # P1 element used."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Arrays"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Coordinates and connectivity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "coord = np.zeros((nnode,ndim))\n",
    "connectivity = np.zeros((nelem,3),dtype=int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for j in range(7):\n",
    "    for i in range(6):\n",
    "        node = i + j * 6\n",
    "        coord[node,0] = i*1.0\n",
    "        coord[node,1] = j*1.0\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Bottom 10 elements, numbered from right to left\n",
    "connectivity[9,:] = [0,1,6]\n",
    "connectivity[8,:] = [1,7,6]\n",
    "connectivity[7,:] = [1,2,7]\n",
    "connectivity[6,:] = [2,8,7]\n",
    "connectivity[5,:] = [2,3,8]\n",
    "connectivity[4,:] = [3,9,8]\n",
    "connectivity[3,:] = [3,4,9]\n",
    "connectivity[2,:] = [4,10,9]\n",
    "connectivity[1,:] = [4,5,10]\n",
    "connectivity[0,:] = [5,11,10]\n",
    "# 8 elements in the left column, numbered bottom up\n",
    "connectivity[10,:] = [6,7,12]\n",
    "connectivity[11,:] = [7,13,12]\n",
    "connectivity[12,:] = [12,13,18]\n",
    "connectivity[13,:] = [13,19,18]\n",
    "connectivity[14,:] = [18,19,24]\n",
    "connectivity[15,:] = [19,25,24]\n",
    "connectivity[16,:] = [24,25,30]\n",
    "connectivity[17,:] = [25,31,30]\n",
    "# Top 10 elements, numbered from left to right\n",
    "connectivity[18,:] = [30,31,36]\n",
    "connectivity[19,:] = [31,37,36]\n",
    "connectivity[20,:] = [31,32,37]\n",
    "connectivity[21,:] = [32,38,37]\n",
    "connectivity[22,:] = [32,33,38]\n",
    "connectivity[23,:] = [33,39,38]\n",
    "connectivity[24,:] = [33,34,39]\n",
    "connectivity[25,:] = [34,40,39]\n",
    "connectivity[26,:] = [34,35,40]\n",
    "connectivity[27,:] = [35,41,40]\n",
    "# 6 elements in the right column, numbered top down\n",
    "connectivity[28,:] = [29,35,34]\n",
    "connectivity[29,:] = [28,29,34]\n",
    "connectivity[30,:] = [23,29,28]\n",
    "connectivity[31,:] = [22,23,28]\n",
    "connectivity[32,:] = [17,23,22]\n",
    "connectivity[33,:] = [16,17,22]\n",
    "# 4 elements in the central horizontal row, numbered right to left\n",
    "connectivity[34,:] = [16,22,21]\n",
    "connectivity[35,:] = [15,16,21]\n",
    "connectivity[36,:] = [15,21,20]\n",
    "connectivity[37,:] = [14,15,20]\n",
    "# 2 elements in the cnetral end\n",
    "connectivity[38,:] = [20,21,26]\n",
    "connectivity[39,:] = [21,27,26]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gradients of basis functions\n",
    "Since two-dimensional P1 elements are used, the basis functions on three nodes of the reference triangle are given as\n",
    "\\begin{align}\n",
    "\\hat{\\phi}_0 &= 1 - \\xi - \\eta \\\\\n",
    "\\hat{\\phi}_1 &= \\xi \\\\\n",
    "\\hat{\\phi}_2 &= \\eta\n",
    "\\end{align}\n",
    "\n",
    "The gradient of $\\hat{\\phi}_{i}$ is given as\n",
    "$\\nabla \\hat{\\phi}_{i} = (\\partial \\hat{\\phi}_{i}/\\partial \\xi, \\partial \\hat{\\phi}_{i}/\\partial \\eta)$ and their values are constant in the element. So, they are stored as static arrays rather than a function to be evaluated at quadrature nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dphihat = np.zeros((NODES_PER_ELEMENT,ndim))\n",
    "dphihat[0,:] = np.array([-1.0,-1.0])\n",
    "dphihat[1,:] = np.array([1.0,0.0])\n",
    "dphihat[2,:] = np.array([0.0,1.0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Other helper arrays\n",
    "`bcflag` array stores integer flags for all the nodes in the domain such that 0 for non-boundary nodes and 1 for boundary nodes.\n",
    "\n",
    "`bcval` array stores Dirichlet boundary conditions for all the nodes such that 0.0 for non-Dirichlet nodes and given values (as the problem specifies) for nodes on the Dirichlet boundaries.\n",
    "\n",
    "`a_coeff` array stores coefficient $a(\\mathbf{x})$ for all the elements. In principle, they need to provide the coefficient values for all the quadrature nodes but P1 elements used for this problem has a single quadrature node per element (i.e., centroid). Note that the way elements are numbered made it easy to assign different values of the coefficient as the problem specifies: 1.0 for element number 0 to 19; 2.0 for element number >= 20."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bcflag = np.zeros((nnode,1),dtype=int)\n",
    "bcflag[5]  = 1\n",
    "bcflag[11] = 1\n",
    "bcflag[26] = 1\n",
    "bcflag[27] = 1\n",
    "\n",
    "bcval = np.zeros((nnode,1))\n",
    "bcval[5]  =  0.0\n",
    "bcval[11] =  0.0\n",
    "bcval[26] = 10.0\n",
    "bcval[27] = 10.0\n",
    "\n",
    "a_coeff = np.ones(nelem)\n",
    "a_coeff[np.arange(nelem)>=20] = 2.0\n",
    "#print(a_coeff)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Partial verification of the arrays created"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.triplot(coord[:,0], coord[:,1], connectivity, 'k-o')\n",
    "plt.gca().set_aspect('equal')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constructing Stiffness Matrix and Force Vector"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Element-wise calculation and assembly\n",
    "\n",
    "An element stiffness matrix is given as\n",
    "\\begin{equation}\n",
    "k_{ij} = \\int_{\\hat{T}} \\hat{a} \\left( \\nabla \\hat{\\phi}_{i} \\cdot \\nabla \\hat{\\phi}_{j} \\right) \\,J\\,d\\hat{\\Omega}\n",
    "\\end{equation}\n",
    "\n",
    "The integration is exactly evaluated by 1-point quadrature at $(\\xi,\\eta)=(\\frac{1}{3},\\frac{1}{3})$ with the weight of $\\frac{1}{2}$:\n",
    "\\begin{align}\n",
    "k_{ij} &= \\frac{1}{2}\\, \\hat{a}\\, \\left( \\nabla \\hat{\\phi}_{i} \\cdot \\nabla \\hat{\\phi}_{j} \\,J \\right)_{(\\xi,\\eta)=(1/3,1/3)} \\\\\n",
    "       &= \\frac{1}{2}\\, \\hat{a}\\, \\left( \\frac{\\partial \\hat{\\phi}_{i}}{\\partial \\xi} \\frac{\\partial \\hat{\\phi}_{j}}{\\partial \\xi} + \\frac{\\partial \\hat{\\phi}_{i}}{\\partial \\eta} \\frac{\\partial \\hat{\\phi}_{j}}{\\partial \\eta} \\right) J\n",
    "\\end{align}\n",
    "\n",
    "Note that the partial derivatives of the basis functions are all constants. \n",
    "\n",
    "Jacobian for P1 is equal to $\\det(B_{T})$ or the ratio of an element's area to the reference element's: $A/A_{\\text{ref}}$. However, since all the elements in the chosen discretization have the same area with the reference triangle, $J = 1.0$ for all of them.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "K = np.zeros((nnode,nnode))\n",
    "F = np.zeros((nnode,1))\n",
    "for e in range(nelem):\n",
    "    ## Although not needed for this problem, \n",
    "    ## the Jacobian calculation is shown here\n",
    "    x0 = coord[connectivity[e,0],:]\n",
    "    x1 = coord[connectivity[e,1],:]\n",
    "    x2 = coord[connectivity[e,2],:]\n",
    "    B_T = np.matrix([[x1[0]-x0[0],x2[0]-x0[0]],[x1[1]-x0[1],x2[1]-x0[1]]])\n",
    "    J = np.linalg.det(B_T)\n",
    "    #print(x0,B_T,J)\n",
    "\n",
    "    ## element arrays are created\n",
    "    k_matrix = np.zeros((NODES_PER_ELEMENT,NODES_PER_ELEMENT))\n",
    "    f_vector = np.zeros(NODES_PER_ELEMENT)\n",
    "\n",
    "    ## Compute element k. \n",
    "    # We don't do anything yet for element force vector\n",
    "    # because no source term was given.\n",
    "    for i in range(NODES_PER_ELEMENT):\n",
    "        for j in range(NODES_PER_ELEMENT):\n",
    "            for d in range(ndim):\n",
    "                k_matrix[i,j] += dphihat[i,d]*dphihat[j,d]\n",
    "            k_matrix[i,j] *= 0.5*J*a_coeff[e]\n",
    "\n",
    "    ## Deal with Dirichlet BC\n",
    "    # If i-th node is Dirichlet, overwrite the i-th row of k matrix\n",
    "    # with the i-th row of the identity matrix; and i-th element of\n",
    "    # f vector with the corresponding bcval.\n",
    "    for i in range(NODES_PER_ELEMENT):\n",
    "        gI = connectivity[e,i]\n",
    "        if bcflag[gI] == 1:\n",
    "            #print(e,i,gI,bcval[gI])\n",
    "            k_matrix[i,:] = 0.0\n",
    "            k_matrix[i,i] = 1.0\n",
    "            f_vector[i] = bcval[gI]\n",
    "            #print(k_matrix,f_vector)\n",
    "\n",
    "    ## Assemble the global matrices by mapping the local node numbers \n",
    "    ## to the global counterparts.\n",
    "    for i in range(NODES_PER_ELEMENT):\n",
    "        gI = connectivity[e,i]\n",
    "        for j in range(NODES_PER_ELEMENT):\n",
    "            gJ = connectivity[e,j]\n",
    "            K[gI,gJ] += k_matrix[i,j]\n",
    "        F[gI] += f_vector[i]\n",
    "    #print(k_matrix, f_vector)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Partial verification of K and F\n",
    "\n",
    "Check for node 11 and 27, the non-shared node at the lower right and the central end of the domain, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 12-th element of K[11,:] should be 1 and all the others should be 0.\n",
    "# 12-th element of F should be 0, the prescribed value by the Dirichlet BC.\n",
    "print(K[11,:], F[11])\n",
    "\n",
    "# 28-th element of K[27,:] should be 1 and all the others should be 0.\n",
    "# 28-th element of F should be 10, the prescribed value by the Dirichlet BC.\n",
    "print(K[27,:], F[27])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solve the linear system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "solution = la.solve(K,F)\n",
    "#print(solution)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize the solution\n",
    "\n",
    "If the given problem is viewed as a static heat conduction problem, it should be straightforward to predict the solution. From the high temperature end to the low-temperature end, a piecewise linear temperature profile will be set up. \n",
    "\n",
    "In the high conductivity ($a=2.0$) domain, a low temperature gradient is set up while the low-conductivity ($a=1.0$) domain has a high temperature gradient. This is the result from the requirement that the heat flow should be constant when the coefficient is constant and the heat flow ($a\\nabla T$) should be continuous at the interface between the domains of different conductivity.\n",
    "\n",
    "**Is there anything else that is consistent or inconsistent with your physical intuition?**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot in a meshgrid (Visualization second method)\n",
    "fig, ax = plt.subplots(figsize=(6,5))\n",
    "ax.set_aspect('equal')\n",
    "plt.tricontourf(coord[:,0], coord[:,1], connectivity, solution[:,0], levels=41, cmap='inferno', vmin=0, vmax=10)\n",
    "plt.triplot(coord[:,0], coord[:,1], connectivity, 'ko-')\n",
    "plt.title('Solution')\n",
    "plt.colorbar(ticks=np.linspace(0,10,11))\n",
    "plt.show ()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.12 ('base')",
   "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.9.12"
  },
  "vscode": {
   "interpreter": {
    "hash": "55380dfba70b56de74ea1cf14c727166cd7c84acbca533adda14e26e17deaa5c"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}