1D_Heat.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2ab890a2",
   "metadata": {},
   "source": [
    "# 1D Linear FE Classes\n",
    "\n",
    "This notebook contains classes for the solution of steady, 2nd order, linear ODE's in one dimension. The code uses Lagrange interpolation polynomials. It includes an example solution of the steady-state homogeneous heat equation with conduction and advection terms.\n",
    "\n",
    "The code is based on the discussion in a set of presentation [slides](http://www.math.pitt.edu/~sussmanm/3040Summer14/FEM1D.pdf) by M.M. Sussman, University of Pittsburgh, accessed 23/02/2022. The slides are incomplete and were used mainly as a guide to build a class architecture for the problem. The lecture notes in turn are based on code from the [FEniCS project](https://fenicsproject.org). Notebook author: Mattis Voss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f3269a4-2c1e-4ed3-9bd7-75e580025c1d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Install helper libraries\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy import integrate\n",
    "from ipywidgets import interact, interactive, fixed, interact_manual\n",
    "import ipywidgets as widgets\n",
    "from IPython.display import Image\n",
    "\n",
    "\"\"\"Code based on and adapted from presentation slides by M.M. Sussman, \n",
    "University of Pittsburgh.\n",
    "URL: http://www.math.pitt.edu/~sussmanm/3040Summer14/FEM1D.pdf\n",
    "Accessed 23/02/2022. Notebook author: Mattis Voss\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d94e5b5",
   "metadata": {},
   "source": [
    "## Basis Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b5a393d-d9f7-4515-b29c-2c7fa77edf90",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BasisFunctions1D:\n",
    "    ''' Define set of Lagrange basis functions on (local) interval [0,1].\n",
    "    The local spatial coordinate over the element is xi \n",
    "    (as in the Greek letter).\n",
    "    Order can be 1 (for pair of linear functions) or 2 \n",
    "    (for triple of quadratic functions)\n",
    "    \n",
    "    PUBLIC METHODS    \n",
    "        eval(n, xi): evaluate basis function phi[n](xi)\n",
    "        dx(n, xi): evaluate first derivative wrt xi \n",
    "        of basis function dphi[n](xi)\n",
    "        size(): return number of nodes of element\n",
    "        plot(): plot all basis functions and their derivatives wrt xi'''\n",
    "    \n",
    "    def __init__(self, order):\n",
    "        '''\n",
    "        PARAMETERS\n",
    "            order: order of basis functions to use (linear or quadratic)\n",
    "        '''\n",
    "        self.__order = order\n",
    "        if self.__order == 1:\n",
    "            # Vector of basis functions\n",
    "            self.__phi = np.array([lambda xi: 1 - xi, \n",
    "                                   lambda xi: xi])\n",
    "            # Vector of derivates w.r.t. xi of basis functions\n",
    "            self.__dphi_dxi = np.array([lambda xi: -1*np.ones_like(xi), \n",
    "                                       lambda xi: 1*np.ones_like(xi)])\n",
    "        elif self.__order == 2:\n",
    "            # Vector of basis functions\n",
    "            self.__phi = np.array([lambda xi: 2.0 * (xi-0.5) * (xi-1.0),\n",
    "                                    lambda xi: 4.0 * xi * (1.0-xi),\n",
    "                                    lambda xi: 2.0 * xi * (xi-0.5)])\n",
    "            # Vector of derivates w.r.t. xi of basis functions\n",
    "            self.__dphi_dxi = np.array([lambda xi: 2.0 * (xi-0.5) +\n",
    "                                        2.0*(xi - 1.0),\n",
    "                                        lambda xi: -4.0 * xi + 4.0*(1.0 - xi),\n",
    "                                        lambda xi: 2.0 * xi + 2.0*(xi - 0.5)])\n",
    "        else:\n",
    "            raise Exception(\"Polynomial order must be 1 or 2\")\n",
    "        \n",
    "        # Number of nodes (always 2 for 1D linear basis functions, \n",
    "        # 3 for 1D quadratic basis functions)\n",
    "        self.__N = order + 1 \n",
    "        \n",
    "    def eval(self,n,xi):\n",
    "        '''\n",
    "        Evaluate the function phi\n",
    "        PARAMETERS\n",
    "            n := index of the basis function\n",
    "            xi := vector of values of xi at which to evaluate\n",
    "        RETURNS\n",
    "            Vector of values of basis function phi\n",
    "        '''\n",
    "        return self.__phi[n](xi)\n",
    "    \n",
    "    def dx(self,n,xi):\n",
    "        '''\n",
    "        RETURNS: derivate of the function phi w.r.t. x\n",
    "        '''\n",
    "        return self.__dphi_dxi[n](xi)\n",
    "    \n",
    "    def size(self):\n",
    "        '''\n",
    "        RETURNS: number of nodes (1 greater than order of \n",
    "        interpolation polynomial)\n",
    "        '''\n",
    "        return self.__N\n",
    "    \n",
    "    def order(self):\n",
    "        ''' RETURNS: order of interpolation polynomial '''\n",
    "        return self.__order\n",
    "    \n",
    "\n",
    "    def plot(self,xi): \n",
    "        '''\n",
    "        Plot all basis functions of the current instance, and their \n",
    "        first derivatives\n",
    "        '''\n",
    "        fig, (ax1, ax2) = plt.subplots(2,1)\n",
    "        fig.suptitle(f'Local Lagrange polynomial basis functions \\\n",
    "                     of order {self.__order}')       \n",
    "        # Set axis labels\n",
    "        ax1.set_ylabel(r'$\\phi(\\xi)$')\n",
    "        ax1.set_xlabel(r'$\\xi$')\n",
    "        ax2.set_xlabel(r'$\\xi$')\n",
    "        ax2.set_ylabel(r'$\\phi^\\prime(\\xi)$')\n",
    "\n",
    "        for i in range(self.__N):\n",
    "            ax1.plot(xi, self.__phi[i](xi))\n",
    "            ax2.plot(xi, self.__dphi_dxi[i](xi))\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a388557",
   "metadata": {},
   "source": [
    "## Mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46d892fb-5229-40ae-b21e-573cb148f8ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Mesh1D:\n",
    "    '''Define a 1D FE mesh. Note only linear elements available for now.\n",
    "    \n",
    "    PUBLIC METHODS\n",
    "        all_nodes(): return array of all node numbers\n",
    "        cell_nodes(e): return node numbers of cell e\n",
    "        all_coords(): return coordinates (global x-values) of all nodes\n",
    "        cell_coords(e): return coordinates of nodes of cell e\n",
    "        node_coord(n): return coordinate (global x-value) of node n\n",
    "        node_coords(): return array of all node coordinates\n",
    "        size(): return number of elements in mesh\n",
    "        cells(): return ordered array of cell numbers\n",
    "        nnodes(): return total number of global nodes\n",
    "        connectivity(): return element connectivity matrix, left to right\n",
    "        x(): x-vector for evaluation of functions over entire mesh\n",
    "    '''\n",
    "    def __init__(self, C, a, b):\n",
    "        '''\n",
    "        PARAMETERS\n",
    "            C: number of cells in mesh\n",
    "            a, b: left and right ends of mesh\n",
    "        '''\n",
    "        self.__a = a # Left endpoint\n",
    "        self.__b = b # Right endpoint\n",
    "        self.__C = int(C) # Number of cells\n",
    "        self.__N = int(C + 1) # Number of nodes\n",
    "        self.__c = np.abs(b - a) # Domain length \n",
    "        self.__h = self.__c/C # Element length, assume all equal\n",
    "\n",
    "    def all_nodes(self):\n",
    "        '''Return ordered array of all nodes in mesh'''\n",
    "        return np.linspace(0, self.__N - 1, num = self.__N) # List of nodes\n",
    "    \n",
    "    def cell_nodes(self, e):\n",
    "        '''\n",
    "        Return all nodes of cell e. Only linear elements for now.\n",
    "        '''\n",
    "        return (e, e+1)\n",
    "    \n",
    "    def all_coords(self):\n",
    "        '''\n",
    "        RETURNS: ordered array of mesh coordinates. Each tuple represents \n",
    "        coordinates of one cell. Only evenly spaced meshes for now.\n",
    "        '''\n",
    "        return np.array([[self.__a+e*self.__h, self.__a+e*(self.__h)+1]\n",
    "                         for e in range(self.__C) ])\n",
    "    \n",
    "    def cell_coords(self, e):\n",
    "        '''\n",
    "        RETURNS: coordinates of nodes of e^th cell (only linear for now)\n",
    "        '''\n",
    "        return np.array([self.__a+e*self.__h, self.__a+(e+1)*(self.__h)])\n",
    "        \n",
    "    def node_coord(self, n):\n",
    "        '''RETURNS: coordinate of n^th node'''\n",
    "        return self.__a + n*self.__h\n",
    "    \n",
    "    def node_coords(self):\n",
    "        '''RETURNS: coordinates of all global nodes'''\n",
    "        return np.array([self.__a + n*self.__h for n in self.all_nodes()])\n",
    "    \n",
    "    def cells(self):\n",
    "        '''RETURNS: ordered array of cell numbers'''\n",
    "        return np.linspace(0, self.__C-1, num = self.__C)\n",
    "        \n",
    "    def size(self):\n",
    "        '''RETURNS: number of cells in mesh'''\n",
    "        return self.__C\n",
    "    \n",
    "    def nnodes(self):\n",
    "        '''RETURNS: number of nodes'''\n",
    "        return self.__N\n",
    "    \n",
    "    def connectivity(self):\n",
    "        '''RETURNS: element connectivity matrix, left to right'''\n",
    "        return np.hstack((np.atleast_2d(self.all_nodes()[:self.size()]).T, \n",
    "                          np.atleast_2d(self.all_nodes()[:self.size()]+1).T))\n",
    "    \n",
    "    def x(self, n=10):\n",
    "        '''RETURNS: x-vector of length of mesh\n",
    "        PARAMETERS:\n",
    "            n: number of points in vector, defaults to 10\n",
    "        '''\n",
    "        return np.linspace(self.__a, self.__b, self.__C*n)\n",
    " \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f50fe1ef",
   "metadata": {},
   "source": [
    "## Linear Element"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb707e5f-6740-4ffc-ae6d-a944b33a3e2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Element1D:\n",
    "    '''\n",
    "    Defines a single 1D finite element, with either linear or \n",
    "    Quadratic Lagrange polynomial basis functions(only linear for now)\n",
    "        \n",
    "    PUBLIC METHODS\n",
    "        ends(): Return global coordinates of element end points\n",
    "        dof_points(): Return coordinates of locations for DOFs \n",
    "                      (same as end points of element for linear case)\n",
    "        dof_nos(): Return array of DOF numbers\n",
    "        num_dofs(): Return number of degrees of freedom \n",
    "        eval(n, x): Evaluate phi_n(x), i.e. evaluate the n_th \n",
    "                    element basis function \n",
    "                    in terms of the global x-coord (not xi)\n",
    "        dx(n, x): Evaluate first derivative of basis function \n",
    "                  dphi[n](x) (not xi)\n",
    "        integrate(n, f1, f2, include_deriv): return an array of values \n",
    "            for the integral $\\int_e f_1(x)f_2(x)\\phi_n(x)dx$,\n",
    "            one value for each n. The f's are optional. if \n",
    "            include_deriv=True then $\\phi^\\prime$ is put in the integral.\n",
    "    '''\n",
    "    def __init__(self, mesh, bfuncs, e, dofs):\n",
    "        '''\n",
    "        PARAMETERS\n",
    "            mesh: the mesh on which the element is located\n",
    "            bfuncs: the BasisFunctions instance\n",
    "            e: element number (in mesh.cells() array)\n",
    "            ends: a tuple of ints giving the numbers of the \n",
    "                  endpoints of this element in the mesh\n",
    "            dofs: array indexing the numbers of the dofs\n",
    "        '''\n",
    "        self.__e = e\n",
    "        ends = mesh.cell_nodes(e)\n",
    "        self.__endcoords = mesh.cell_coords(e)\n",
    "        self.__numdofs = bfuncs.size()\n",
    "        self.__dofs = dofs\n",
    "        self.__dofcoords = np.linspace(self.__endcoords[0],\n",
    "                                       self.__endcoords[1],self.__numdofs)\n",
    "        self.__bfuncs = bfuncs\n",
    "        self.__mesh = mesh\n",
    "        self.__bfuncs = bfuncs\n",
    "    \n",
    "    # Getter methods for finite element endpoints, dof coordinates etc.\n",
    "    def endcoords(self):\n",
    "        '''Get end points of element'''\n",
    "        return self.__endcoords\n",
    "    \n",
    "    def dofcoords(self):    \n",
    "        '''\n",
    "        Get coordinates of degree of freedom points of element. \n",
    "        These are the same as the node coordinates for linear element\n",
    "        '''\n",
    "        return self.__mesh.cell_coords(self.__e)\n",
    "    \n",
    "    def dofs(self):    \n",
    "        '''Get indices of degrees of freedom. \n",
    "        Note these are same as node numbers for a linear element'''\n",
    "        return self.__dofs\n",
    "    \n",
    "    def numdofs(self):\n",
    "        '''Get number of degrees of freedom, same as \n",
    "        number of basis functions'''\n",
    "        return self.__numdofs \n",
    "    \n",
    "    def length(self):\n",
    "          return self.endcoords()[1] - self.endcoords()[0]\n",
    "    \n",
    "    def e(self):\n",
    "        '''RETURNS: element index number in mesh'''\n",
    "        return self.__e\n",
    "    \n",
    "    # Evaluate shape functions over the element in terms of \n",
    "    # global coordinate x\n",
    "    def phi(self, n, x):\n",
    "        '''Evaluate $\\phi_n$ on the element. First map global to \n",
    "        local coordinates, then use BasisFunctions1D public \n",
    "        eval() method to return resulting array'''\n",
    "        xi = (x - self.__endcoords[0])/self.length()\n",
    "        # Return the evaluation on xi, but only for 0 <= xi <= 1\n",
    "        return self.__bfuncs.eval(n, xi) * (xi >= 0) * (xi <= 1)\n",
    "    \n",
    "    def phi_prime(self, n, x):\n",
    "        '''Evaluate $d\\phi_n/dx$ on the element. First map global \n",
    "        to local coordinates, then use BasisFunctions1D \n",
    "        public dx() method to return resulting array'''\n",
    "        xi = (x - self.__endcoords[0])/self.length()\n",
    "        # Return the derivative on xi, but only for 0 <= xi <= 1\n",
    "        return self.__bfuncs.dx(n, xi) * (xi >= 0) * (xi <= 1)    \n",
    "    \n",
    "    def integrate(self, x, f1=None, f2=None, deriv=False, c=None):\n",
    "        '''Define integration method over the element. \n",
    "        Used when defining the SolutionSpace1D.\n",
    "        PARAMETERS\n",
    "            x: global x-vector\n",
    "            c: a constant to multiply the integral with\n",
    "            f1: optional function, must be mapped to element coordinates \n",
    "                (i.e. must be a vector)\n",
    "            f2: see f1\n",
    "            deriv: boolean to decide whether $\\phi$ or $\\phi^\\prime$ \n",
    "                   is used in the integral\n",
    "        RETURNS\n",
    "            I: array of integral values, one for each basis function\n",
    "        '''\n",
    "        L = self.length()\n",
    "        xi = (x - self.__endcoords[0])/L\n",
    "        \n",
    "        if f1 is None:\n",
    "            f1 = np.ones_like(xi)\n",
    "        if f2 is None:\n",
    "            f2 = np.ones_like(xi)\n",
    "        if c is None:\n",
    "            c = np.ones_like(xi)\n",
    "        if not deriv:\n",
    "            f = np.array([f1*f2*self.__bfuncs.eval(n,xi) \n",
    "                          for n in range(self.numdofs())]) \n",
    "        else:\n",
    "            f = np.array([f1*f2*self.__bfuncs.dx(n,xi)/L \n",
    "                          for n in range(self.numdofs())])\n",
    "        I = np.trapz(f,xi)\n",
    "        return I"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e81c16de",
   "metadata": {},
   "source": [
    "## Solution Space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6193a484",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SolutionSpace1D:\n",
    "    '''Defines a complete trial solution function space \n",
    "    over a mesh of elements. Building blocks are Mesh1D \n",
    "    and BasisFunctions1D objects. \n",
    "    \n",
    "    PUBLIC METHODS\n",
    "        size(): Return number of elements\n",
    "        ndofs(): Return sum of all dofs over all elements\n",
    "        all_dofcoords(): Return array of coordinates \n",
    "                         of all dofs\n",
    "        dof_coord(n): Return coordinate of dof n\n",
    "        int_gal(f, deriv): Return stiffness matrix based on \n",
    "                           integration of Galerkin criterion\n",
    "    '''\n",
    "    def __init__(self, mesh, bfuncs):\n",
    "        '''\n",
    "        PARAMETERS\n",
    "            mesh: a Mesh1D object\n",
    "            bfuncs: a BasisFunctions1D object\n",
    "        '''\n",
    "        self.__mesh = mesh\n",
    "        self.__bfuncs = bfuncs\n",
    "        self.__size = mesh.size()\n",
    "        # List to be populated with elements in a loop below\n",
    "        self.__elems = list([])\n",
    "        # List to be populated with DOF coordinates in the same loop\n",
    "        self.__dofcoords = list([]) \n",
    "        self.__ndofs = 0 # Initialise counter to zero\n",
    "        \n",
    "        # Populate solution space with elements \n",
    "        #(cells and associated basis functions)\n",
    "        for e in range(self.__size):\n",
    "            # Boundary conditions are Dirichlet at \n",
    "            # beginning and end of mesh for now\n",
    "            if e == 0:\n",
    "                self.__ndofs += int(self.__bfuncs.size())\n",
    "                # Note this would be [2*e, 2*e+1, 2*e+2] for \n",
    "                # quadratic element\n",
    "                dofs = [e, e+1] \n",
    "                # This would be range(3) for quadratic elements\n",
    "                newdofs = range(2) \n",
    "            else:\n",
    "                self.__ndofs += int(self.__bfuncs.size() - 1)\n",
    "                dofs = [e, e+1]\n",
    "                # This would be range(1,3) for quadratic elements\n",
    "                newdofs = range(1,2) \n",
    "            # Create new element in mesh\n",
    "            elem = Element1D(self.__mesh, self.__bfuncs, e, dofs) \n",
    "            # Add new element to the array of elements\n",
    "            self.__elems.append(elem) \n",
    "            for i in newdofs:\n",
    "                self.__dofcoords.append(elem.dofcoords()[i])\n",
    "                            \n",
    "    # Getter methods\n",
    "    def bfuncs(self):\n",
    "        return self.__bfuncs\n",
    "    def size(self):\n",
    "        return self.__size\n",
    "    def ndofs(self):\n",
    "        return self.__ndofs\n",
    "    def dofcoords(self):\n",
    "        return self.__dofcoords\n",
    "    def x(self):\n",
    "        return self.__mesh.x()\n",
    "    \n",
    "    def int_gal(self, f=None, deriv=[False, False]):\n",
    "        '''Integration method for Galerkin criterion. \n",
    "        Uses the Element1D `integrate` method.\n",
    "        PARAMETERS\n",
    "            deriv: Pair of booleans, giving the ability to \n",
    "                   mix and match $\\phi$ and $\\phi^\\prime$\n",
    "            f: Optional function that may be included in the derivative\n",
    "        RETURNS\n",
    "            A: Combined stiffness matrix for all elements\n",
    "        '''\n",
    "        # Stiffness matrix to be populated\n",
    "        A = np.zeros([self.__ndofs, self.__ndofs]) \n",
    "        # Loop over all elements and add integrated terms into \n",
    "        # stiffness matrix\n",
    "        for elem in self.__elems:\n",
    "            # List of global degree of freedom indices on current element\n",
    "            d = elem.dofs() \n",
    "            # Number of degrees of freedom on current element, this will \n",
    "            # be same as number of element equations\n",
    "            D = elem.numdofs()  \n",
    "            # Length of element, for conversion from local to \n",
    "            # global coordinates                    \n",
    "            L = elem.length() \n",
    "\n",
    "            for i in range(D):\n",
    "                if deriv[1]:\n",
    "                    phi = elem.phi_prime(i, self.x())/L\n",
    "                    # Because of the chain rule d(xi)/dx = 1/L\n",
    "                else:\n",
    "                    phi = elem.phi(i, self.x())\n",
    "                A[d,d[i]] += elem.integrate(self.x(), \n",
    "                                            phi,f,deriv=deriv[0])\n",
    "                                                                            \n",
    "        return A"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8e36f74",
   "metadata": {},
   "source": [
    "## Galerkin Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1e2b580",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Galerkin:\n",
    "    '''Galerkin solution for a linear, first order homogeneous ODE. \n",
    "    Currently only accepts Dirichlet\n",
    "    boundary conditions at either end of interval.\n",
    "    \n",
    "    PUBLIC METHODS\n",
    "    get_K(): returns the integrated left hand side of the \n",
    "             Galerkin criterion (i.e. the global stiffness matrix)\n",
    "    apply_bc(): applies boundary conditions to stiffness matrix\n",
    "    solve(): uses np.linalg.solve to solve the system for nodal unknowns\n",
    "    plot(): plot solution and exact solution, if available\n",
    "    '''\n",
    "    def __init__(self, solspace, terms, derivs, \n",
    "                 bc=None, exact_soln=None, rhs=None):\n",
    "        '''\n",
    "        PARAMETERS\n",
    "            solspace: solution space in which to evaluate the integral\n",
    "            terms: number of terms in each integral\n",
    "            derivs: 2D array of size terms x 2, each row is a pair of \n",
    "                    booleans, where True gives $\\phi_prime$\n",
    "                    and False gives $\\phi$\n",
    "            bc: array of pairs specifying Dirichlet BCs as [node, value]\n",
    "            rhs: forcing term on rhs of Galerkin criterion\n",
    "            exact_soln: function defining the exact solution \n",
    "                        over the domain\n",
    "        '''\n",
    "        self.__dofcoords = solspace.dofcoords()\n",
    "        self.__exact_soln = exact_soln\n",
    "        self.__bc = bc\n",
    "        self.__rhs = rhs\n",
    "        self.__N = solspace.ndofs()\n",
    "        self.__x = solspace.x()\n",
    "        self.__bfuncs = solspace.bfuncs()\n",
    "        if exact_soln is not None: \n",
    "            self.__y = exact_soln(solspace.x())\n",
    "        \n",
    "        # Find full stiffness matrix \n",
    "        # (before application of boundary conditions)\n",
    "        self.__K_full = 0*solspace.int_gal(deriv=[False, False])\n",
    "        for i in range(terms):\n",
    "            self.__K_full += solspace.int_gal(deriv=derivs[i])        \n",
    "        \n",
    "    def get_K(self):\n",
    "        '''Get stiffness matrix\n",
    "        RETURNS: Numerically integrated lhs of the Galerkin criterion'''\n",
    "        return self.__K_full\n",
    "    \n",
    "    def apply_bc(self):\n",
    "        '''Apply Dirichlet boundary conditions. \n",
    "        RETURNS: final stiffness matrix and right-hand vector'''\n",
    "        # Make stiffness matrix\n",
    "        K = self.get_K()\n",
    "        K[0] = 0\n",
    "        K[0,0] = 1\n",
    "        K[-1] = 0\n",
    "        K[-1, -1] = 1\n",
    "        b = np.zeros(self.__N)\n",
    "        b[0] = self.__bc[0][1]\n",
    "        b[-1] = self.__bc[-1][1]\n",
    "        return (K, b)\n",
    "    \n",
    "    def solve(self):\n",
    "        '''Solve system equations. \n",
    "        RETURNS \n",
    "            Vector of approximate nodal values of field variable'''\n",
    "        (K, b) = self.apply_bc()\n",
    "        return np.round(np.linalg.solve(K, b), 3)       \n",
    "    \n",
    "    def plot(self):\n",
    "        '''Plot approximation over domain, \n",
    "        along with exact solution where available'''\n",
    "        plt.figure(figsize=(12,8))\n",
    "        plt.plot(self.__dofcoords, self.solve())\n",
    "        plt.plot(self.__x, self.__exact_soln(self.__x))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63fa8a79",
   "metadata": {},
   "source": [
    "# Validation Code\n",
    "The following code blocks test the classes defined above on the 1D steady heat equation with no forcing (RHS) terms. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f55458f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Validate solver on 1D steady heat equation\n",
    "def heat1D(elements):\n",
    "    '''Validation function for finite element code.\n",
    "    Uses 1D steady heat equation with \n",
    "    Dirichlet boundary conditions.\n",
    "    PARAMETERS\n",
    "        elements: number of elements in mesh\n",
    "    '''\n",
    "    \n",
    "    # Create BasisFunctions1D object, use linear basis functions\n",
    "    lin_bfuncs = BasisFunctions1D(1)\n",
    "    # Create Mesh1D object, domain [0, 2]\n",
    "    mesh = Mesh1D(elements,0.0,2.0)\n",
    "    # Create SolutionSpace1D object\n",
    "    solspace = SolutionSpace1D(mesh,lin_bfuncs)\n",
    "    # Define exact solution\n",
    "    T = lambda x, k=1: (2*(np.exp(2/k) - 2) + \n",
    "                        2 * np.exp(x/k)) / (np.exp(2/k) - 1)\n",
    "    # Create and plot Galerkin solution\n",
    "    gal = Galerkin(solspace, 2, \n",
    "                   [[False, True], [True, True]],[[0,2],[2,4]], T)\n",
    "    gal.plot()\n",
    "\n",
    "def heat_interact(f):\n",
    "    \"\"\" Streamline Interact function for T_approx \"\"\"\n",
    "    interactive_plot = interactive(f, elements=(1,8,1))\n",
    "    output = interactive_plot.children[1]\n",
    "    output.layout.height = '500px'\n",
    "    return interactive_plot\n",
    "\n",
    "def plot_bf(order):\n",
    "    '''Plot basis functions\n",
    "    PARAMETERS:\n",
    "    order: 1 for linear, 2 for quadratic'''\n",
    "    xi = np.linspace(0, 1, num=10)\n",
    "    lin_bfuncs = BasisFunctions1D(order)\n",
    "    lin_bfuncs.plot(xi)\n",
    "    \n",
    "def bf_interact(f):\n",
    "    '''Streamline Interact function for plot_bf'''\n",
    "    interactive_plot = interactive(f, order=(1,2,1))\n",
    "    output = interactive_plot.children[1]\n",
    "    output.layout.height = '500px'\n",
    "    return interactive_plot  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98ebbd2b",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Interactive plot of solution with \n",
    "# slider for number of elements\n",
    "heat_interact(heat1D)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c4f012b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Interactive plot of basis functions\n",
    "# with slider for order\n",
    "bf_interact(plot_bf)"
   ]
  }
 ],
 "metadata": {
  "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.9.7"
  },
  "toc-autonumbering": true,
  "toc-showcode": true,
  "toc-showmarkdowntxt": false
 },
 "nbformat": 4,
 "nbformat_minor": 5
}