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) = a_0 + \sum_{k=1}^N\left[ a_k(t) \cos(kx) + b_k(t) \sin(kx) \right]$$ our vector of coefficients, then, is $$\vec{u} = [a_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
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 matplotlib import rcParams
rcParams['font.family'] = 'serif'
rcParams['font.size'] = 16
Define global constants
In [2]:
# We define the ORDER of approximation
N = 7
# 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}\sqrt{\frac{1}{2}}&\text{if }k = 0\\\cos((k/2)x)&\text{if }k\text{ even}\\\sin((k+1)x/2)&\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.
In [4]:
# First we define the function analyticially
def phi(k,x):
k=int(k) #typecast. Be aware this kind of behavior can be dangerous.
if k == 0:
# why the ones_like call?
return numpy.sqrt(1./2)*numpy.ones_like(x)
elif k % 2 == 0:
return numpy.cos(int(k/2)*x)
else:
return numpy.sin(int(1+k/2)*x)
Let's plot a few of these:
In [5]:
x = numpy.linspace(xmin,xmax,100)
[plt.plot(x,phi(k,x)) for k in range(3)]
Out[5]:
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 [6]:
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^2}).$$ Let's plot it:
In [7]:
x = numpy.linspace(xmin,xmax,100)
u_initial = lambda x: numpy.sin(numpy.exp(-x**2))
plt.plot(x,u_initial(x))
Out[7]:
This has the (rather nasty) analytic solution $$u(t,x) = \sin(e^{-\text{mod}(t-x,2\pi)^2})+\sin(e^{-\text{mod}(t-x,-2\pi)^2}).$$
How do we convert this into a spectral representation? Like this!
In [8]:
def to_modal(u_func):
"""
Projects a function u_func onto our discrete
Fourier basis and returns a vector u
of Fourier coefficients
"""
# why the lambda call here?
return numpy.array([inner(u_func,lambda x: phi(k,x))for k in range(N)])
u0 = to_modal(u_initial)
Let's see if it looks right.
In [9]:
plt.plot(x,numpy.vectorize(solution(u0))(x))
Out[9]:
Let's define a derivative operator. What's the best way to do that? We have: $$\frac{d\phi_k}{dx} = \begin{cases}0&\text{if }k=0\\-\frac{k}{2}\sin((k/2)x)&\text{if }k\text{ even}\\\frac{k+1}{2}\cos((1+k)x/2)&\text{if }k\text{ odd}\end{cases}$$ And this is: $$\frac{d\phi_k}{dx} = \begin{cases}0&\text{if }k=0\\-\frac{k}{2}\phi_{k-1}&\text{if }k\text{ even}\\\frac{k+1}{2}\phi_{k+1}&\text{if }k\text{ odd}\end{cases}$$
In [10]:
# Differentiation matrix for the phis
Dphi = numpy.zeros((N,N))
for k in range(1,N-1):
if k % 2 == 0:
Dphi[k,k-1] = -k/2
else:
Dphi[k,k+1] = 1+k/2
# Differentiation matrix for the u's
D = Dphi.transpose()
In [11]:
D[1]
Out[11]:
In [12]:
def rhs(u):
"""
Given a vector of Fourier coefficeints u, returns du/dt
"""
return c*numpy.dot(D,u)
# For technical reasons, don't use Euler step.
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 [13]:
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 [14]:
fig = plt.figure()
ax = plt.axes(xlim=(xmin,xmax),ylim=(0,1),xlabel=('x'),ylabel=('u'))
line1, = ax.plot([],[],color='#003366', lw=2)
line2, = ax.plot([],[],'r--',lw=2)
def animate(i):
line1.set_data(x,numpy.vectorize(solution(u[i]))(x))
line2.set_data(x,numpy.sin(numpy.exp(-((t[i]-x)%(-2*numpy.pi))**2))+numpy.sin(numpy.exp(-((t[i]-x)%(2*numpy.pi))**2)))
return line1,line2
animation.FuncAnimation(fig, animate, frames=len(t), interval=10)
Out[14]:
In [14]: