Spectral Methods for Linear Advection

As an extremely simple example, we solve the linear advection equation $$\frac{\partial u}{\partial t} = c \frac{\partial u}{\partial x}$$ using spectral methods.

For simplicity, we study the equation on the domain $$\Omega=[-\pi,\pi]$$ with periodic boundary conditions.

We represent our function as a sum over positive and negative frequencies of Fourier modes for a total of $2N+1$ modes. $$u(t,x) = \sum_{k=0}^N\left[ a_k(t) \cos(kx) + b_k(t) \sin(kx) \right]$$ our vector of coefficients, then, is $$\vec{u} = [a_0,b_0,a_1,b_1,...,a_N,b_N]$$

We can solve this numerically using the techniques we've learned so far.

Let's get to it. First, let's import our modules.


In [1]:
# First we import our modules
%matplotlib inline
import numpy 
import sympy
from scipy import integrate
from numpy import linalg
import matplotlib.pyplot as plt
from matplotlib import animation
from JSAnimation.IPython_display import display_animation
from sympy import init_printing
from matplotlib import rcParams
rcParams['font.family'] = 'serif'
rcParams['font.size'] = 16
init_printing()

Define global constants


In [2]:
# We define the ORDER of approximation 
N = 20
# And the speed of wave propogation
c = -1
# and the min and max of the domain
xmin = -numpy.pi
xmax = numpy.pi

For convenience, we define the functional inner product using scipy's integration feature.


In [3]:
def inner(a,b):
    """
    Compute the inner product of functions a and b on the interval
    """
    return (2.0/(xmax-xmin))*integrate.quad(lambda x: a(x)*b(x),xmin,xmax)[0]

Define basis functions. We define them as: $$\phi_k(x) = \begin{cases}\frac{1}{2}&\text{if }k = 0\\\cos((k/2)x)&\text{if }k\text{ even}\\\sin((1+k/2)x)&\text{if }k\text{ odd}\end{cases},$$ where $k/2$ is defined as integer division and always returns the floor of a floating point number.

To calculate differentiation properties, we will need to use sympy. This follows below.


In [19]:
# First we define the function analyticially
y = sympy.Symbol('y',real=True)
def phi_symb(k):
    k=int(k) #typecast. Be aware this kind of behavior can be dangerous.
    if k == 0:
        return numpy.sqrt(1./2)
    elif k % 2 == 0:
        return sympy.cos(int(k/2)*y)
    else:
        return sympy.sin(int(1+k/2)*y)

In [22]:
print "example basis: ",[phi_symb(k) for k in range(2*3+1)]


example basis:  [0.70710678118654757, sin(y), cos(y), sin(2*y), cos(2*y), sin(3*y), cos(3*y)]

We also define an array of basis elements an array of their derivatives:


In [6]:
phi = numpy.array([sympy.lambdify([y],phi_symb(k),numpy) for k in range(N)])
dphi = numpy.array([sympy.lambdify([y],sympy.diff(phi_symb(k),y),numpy)\
                    for k in range(N)])

Note that functions that have been lambdified do not, a-priori, accept vector arguments. To do that, you should use the "vectorize" numpy array. I.e., if you wanted to plot the first few basis functions:


In [7]:
x = numpy.linspace(xmin,xmax,100)
[plt.plot(x,numpy.vectorize(phi[k])(x)) for k in range(3)]


Out[7]:
[[<matplotlib.lines.Line2D at 0x7f3e4753b9d0>],
 [<matplotlib.lines.Line2D at 0x7f3e475d9c50>],
 [<matplotlib.lines.Line2D at 0x7f3e474cb310>]]

Next, we calculate the mass and stiffness matrices and use them to calculate the differentiation matrix.


In [8]:
# mass matrix
M = numpy.empty((N,N),dtype=float)
for i in range(N):
    for j in range(N):
        M[i,j] = inner(phi[i],phi[j])
        
# stiffness matrix
S = numpy.empty((N,N),dtype=float)
for i in range(N):
    for j in range(N):
        S[i,j] = inner(phi[i],dphi[j])

# differentiation matrix.
D = numpy.dot(linalg.inv(M),S)

Now, assume we have a vector $\vec{u}=[u_0,u_1,...,u_n]$ of Fourier coefficients. How do we extract the $u(x)$? Let's define a function to do it.


In [9]:
def solution(u):
    """
    Takes a vector u of coefficients u_i and returns the true function
    u(x) = sum(u_i * e^(inx))
    """
    return lambda x: numpy.sum(numpy.array([u[k]*phi[k](x) for k in range(N)]))

Let's get it over with and define an initial function. We'll use $$u0=u(0,x) = \sin(e^{-(x-\pi)^2}).$$ Let's plot it:


In [10]:
x = numpy.linspace(xmin,xmax,100)
u_initial = lambda x: numpy.sin(numpy.exp(-x**2))
plt.plot(x,u_initial(x))


Out[10]:
[<matplotlib.lines.Line2D at 0x7fae5a31f150>]

How do we convert this into a spectral representation? Like this!


In [11]:
def to_modal(u_func):
    """
    Projects a function u_func onto our discrete 
    Fourier basis and returns a vector u
    of Fourier coefficients
    """
    return numpy.array([inner(u_func,phi[k])for k in range(N)])

u0 = to_modal(u_initial)

Let's see if it looks right.


In [12]:
plt.plot(x,numpy.vectorize(solution(u0))(x))


Out[12]:
[<matplotlib.lines.Line2D at 0x7fae5a1f7550>]

We're ready to define an integration scheme. For technical reasons, we cannot use an Euler step. So we'll instead use RK2.


In [13]:
def rhs(u):
    """
    Given a vector of Fourier coefficeints u, returns du/dt
    """
    return c*numpy.dot(D,u)
    

def RK2_step(u,dt):
    """
    Maps u at time t to u at time t + dt, where u is a vector of Fourier coefficients
    """
    u_temp = u + 0.5*dt*rhs(u)
    return u + dt*rhs(u_temp)

Okay, let's evolve this thing in time.

A word on the CFL condition. Here assume that $\Delta x$ is the wavelength of your smallest mode. In our case, this means that $\Delta x$ should be $\pi/(N)$.


In [14]:
NUM_T = 100
tmax=2*numpy.pi
t = numpy.linspace(0,tmax,NUM_T)
dt=t[1]-t[0]
u = numpy.empty((len(t),len(u0)))
u[0] = u0 # set initial condition
# Integrate!
for i in range(1,len(t)):
    u[i] = RK2_step(u[i-1],dt)

Let's animate it!


In [15]:
fig = plt.figure()
ax = plt.axes(xlim=(xmin,xmax),ylim=(0,1),xlabel=('x'),ylabel=('u'))
line, = ax.plot([],[],color='#003366', lw=2)
def animate(i):
    line.set_data(x,numpy.vectorize(solution(u[i]))(x))
    return line
animation.FuncAnimation(fig, animate, frames=len(t), interval=10)


Out[15]:


Once Loop Reflect

In [ ]: