{
 "metadata": {
  "name": ""
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%pylab inline\n",
      "import numpy as np\n",
      "import pylab as plt"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Before we look at the spectral collocation method, let's see just why it's important not to use an equispaced grid.  We'll  try interpolating Runge's function:  \n",
      "$$f(x)=\\frac{1}{1+25x^2}$$  \n",
      "on equispaced and Chebyshev grids. We can do this very easily using the NumPy functions polyfit() and poly1d(). Of course, under the hood polyfit() is just solving a certain linear system (a Vandermonde matrix, in fact) associated with the interpolation problem.  \n",
      "The script below plots $f(x)$ and these two interpolants. It also plots the Chebyshev point locations. Try changing $m$ and see what happens. Is the Chebyshev interpolant more accurate at every point?  "
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "m=10 # This is the number of (interior) points to be used.\n",
      "# Runge's function:\n",
      "f = lambda x: 1./(1+25*x**2)\n",
      "# The grids:\n",
      "x_eq = np.linspace(-1,1,m+2)\n",
      "x_cheb= np.cos(np.pi*np.arange(m+2)/(m+1.))\n",
      "# The polynomial interpolants:\n",
      "p1=np.poly1d(np.polyfit(x_eq,f(x_eq),m+1))\n",
      "p2=np.poly1d(np.polyfit(x_cheb,f(x_cheb),m+1))\n",
      "# The rest is just plotting:\n",
      "x_fine=np.linspace(-1,1,1000)\n",
      "plt.clf()\n",
      "plt.plot(x_fine,f(x_fine),x_fine,p1(x_fine),x_fine,p2(x_fine))\n",
      "plt.legend(['f(x)','Equispaced interpolant','Chebyshev interpolant'],loc='best')\n",
      "plt.plot(x_cheb,f(x_cheb),'ok')\n",
      "plt.title('N='+str(m+2)+' points')\n",
      "plt.axis([-1,1,-0.5,1.5])"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Now we'll look at implementation of a Chebyshev spectral collocation method.  The function below computes a matrix that gives a spectral approximation to the first derivative over the interval $[a,b]$. This is adapted from MATLAB code in Trefethen's text on spectral methods (see Chapter 6 of that text for an explanation). The function also returns the Chebyshev grid points over the interval.  "
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def cheb(m,a,b):\n",
      "    x = np.cos(np.pi*np.arange(m+2)/(m+1))       # Chebyshev points\n",
      "    c = np.ones(m+2); c[0]=2; c[-1]=2;\n",
      "    c = c*np.power(-1,np.arange(m+2))\n",
      "    X = np.tile(x,[m+2,1]).T\n",
      "    dX = X-X.T\n",
      "    D = (np.outer(c,1./c))/(dX+np.eye(m+2))\n",
      "    D = D - np.diag(np.sum(D,axis=1))\n",
      "    D = -D*2./(b-a)                              # Rescale from [-1,1] to [a,b]\n",
      "    x=a + 0.5*(b-a)*(1.+np.cos(np.pi*(1.-np.arange(m+2)/(m+1.))))    #Shifted Chebyshev points\n",
      "    return D,x"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "D,x=cheb(1,-1,1)\n",
      "print D\n",
      "print x"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "The Chebyshev differentiation matrix $D$ approximates the first derivative.  Let's try using it to approximately compute the first derivative of Runge's function:  "
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "m=100\n",
      "D,x_cheb=cheb(m,-1,1)\n",
      "F=f(x_cheb)\n",
      "dF=np.dot(D,F)\n",
      "plt.clf()\n",
      "plt.plot(x_cheb,F,x_cheb,dF)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Now let's use $D$ to solve the BVP:  \n",
      "$$\\begin{align*} u''(x) = \\exp(4x) \\quad \\quad -1\\le x \\le1 \\\\ u(-1)=u(1)=0.\\end{align*}$$  \n",
      "To approximate the second derivative, we just take $D^2$. To implement the Dirichlet boundary conditions, we change the first and last rows of the resulting system of equations:  "
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def solvecheb(m,f,alpha,beta):\n",
      "    D,x=cheb(m,-1.,1.);\n",
      "    D2=np.dot(D,D); \n",
      "    D2=D2[1:-1,:]; \n",
      "    A=np.zeros([m+2,m+2]); A[1:-1,:]=D2; A[0,0]=1; A[-1,-1]=1;\n",
      "    F=np.zeros(m+2)\n",
      "    F[0]=alpha; F[m+1]=beta\n",
      "    F[1:-1]=f(x[1:-1])\n",
      "    return np.linalg.solve(A,F),x\n",
      "f = lambda x: np.exp(4.*x)\n",
      "U,x=solvecheb(15,f,0,0)\n",
      "plt.clf()\n",
      "plt.plot(x,U)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    }
   ],
   "metadata": {}
  }
 ]
}