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)]
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]:
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]:
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]:
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]:
In [ ]: