In [1]:
from IPython.core.display import HTML
css_file = './custom.css'
HTML(open(css_file, "r").read())
Out[1]:
Welcome to PseudoSpectralPython, a short course that will teach you how to solve wave equations using pseudospectral collocation methods. This notebook is the first lesson, on solving linear problems. Pseudospectral methods are great for wave problems where:
In [ ]:
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
from matplotlib import animation
from JSAnimation import IPython_display
To begin, let's consider the partial differential equation (PDE)
$$u_t + a u_x = \epsilon u_{xx}$$referred to as the advection-diffusion equation, for reasons we'll soon discover. Here we wish to find $u(x,t)$, which might be the density or concentration of some substance. The subscripts denote partial differentiation; e.g. $u_t$ is the partial derivative of $u$ with respect to $t$. The coefficients $a$ and $\epsilon$ are constants that determine the strength of the advective and diffusive effects.
Let's solve this equation on a periodic domain $[-\pi,\pi]$, with some initial data
$$u(x,0) = u_0(x).$$If we suppose for a moment that our solution is composed of a single Fourier mode with wavenumber $\xi$ and time-dependent amplitude $\hat{u}$:
$$u(x,t; \xi) = \hat{u}(t) e^{i\xi x},$$Then we obtain a simple ordinary differential equation (ODE) for $\hat{u}$:
$$\hat{u}'(t; \xi) + i\xi a \hat{u} = -\xi^2 \epsilon \hat{u}$$We can solve this scalar ODE exactly:
$$\hat{u}'(t; \xi) = e^{(-i \xi a - \epsilon \xi^2)t} \hat{u}(0).$$We've transformed the original PDE into a simple ODE, but you may wonder whether this is useful, since we assumed a very simple form for the solution. The marvelous fact is that every solution of our advection-diffusion equation can be written as a linear combination (a superposition) of simple solutions of the form above, with different wavenumbers $\xi$. We can construct the general solution as follows.
First, we take a Fourier transform of the initial data:
$$\hat{u}(t=0;\xi) = \frac{1}{\sqrt{2\pi}} \int_{-\infty}^\infty u_0(x) e^{-i\xi x}dx.$$Then each mode evolves according to the solution of the ODE above:
$$\hat{u}'(t; \xi) = e^{(-i \xi a - \epsilon \xi^2)t} \hat{u}(0;\xi).$$Finally, we construct the solution again by taking the inverse Fourier transform. This just means summing up all the Fourier modes:
$$u(x,t) = \frac{1}{\sqrt{2\pi}} \int_{-\infty}^\infty u_0(x) e^{i\xi x}d\xi.$$If you haven't seen Fourier analysis at all before, now is a good time to go read up a bit and then come back.
We can't evaluate the integrals above exactly on the computer (at least, not for arbitrary initial data $u_0$). Instead, we need to discretize. To do so, we introduce a grid with a finite set of points in space and time:
In [ ]:
# Spatial grid
m=64 # Number of grid points in space
L = 2 * np.pi # Width of spatial domain
x = np.arange(-m/2,m/2)*(L/m) # Grid points
dx = x[1]-x[0] # Grid spacing
# Temporal grid
tmax=4.0 # Final time
N = 25 # number grid points in time
k = tmax/N # interval between output times
and a corresponding set of discrete wavenumber values $\xi$:
In [ ]:
xi = np.fft.fftfreq(m)*m*2*np.pi/L # Wavenumber "grid"
# (this is the order in which numpy's FFT gives the frequencies)
The functions $u, \hat{u}$ discussed above are replaced by finite-dimensional vectors. These vectors are related through the discrete version of the Fourier transform, aptly called the discrete Fourier transform (DFT). We'll look at the DFT in more detail in the next lesson. For now, let's set the initial condition to
$$u_0(x) = \begin{cases} \sin^2(2x) & -\pi \le x < -\pi/2 \\ 0 & x>-\pi/2 \end{cases}$$and compute its DFT:
In [ ]:
# Initial data
u = np.sin(2*x)**2 * (x<-L/4)
uhat0 = np.fft.fft(u)
Next, we set a value for epsilon and compute the solution:
In [ ]:
epsilon=0.01 # Diffusion coefficient
a = 1.0 # Advection coefficient
# Store solutions in a list for plotting later
frames = [u.copy()]
# Now we solve the problem
for n in range(1,N+1):
t = n*k
uhat = np.exp(-(1.j*xi*a + epsilon*xi**2)*t) * uhat0
u = np.real(np.fft.ifft(uhat))
frames.append(u.copy())
We have computed and stored the solution. The code below plots it as an animation.
In [ ]:
# Set up plotting
fig = plt.figure(figsize=(9,4)); axes = fig.add_subplot(111)
line, = axes.plot([],[],lw=3)
axes.set_xlim((x[0],x[-1])); axes.set_ylim((0.,1.))
def plot_frame(i):
#fig = plt.figure()
#plt.plot(x,frames[i])
line.set_data(x,frames[i])
axes.set_title('t='+str(i*k))
fig.canvas.draw()
return fig
# Animate the solution
matplotlib.animation.FuncAnimation(fig, plot_frame,
frames=len(frames),
interval=200,
repeat=False)
What did we just do? We solved a partial differential equation computationally. It's time to think about how accurate the solution is and what approximations we made.
The first approximation we made was to take the initial data and approximate it by just the first terms in its Fourier series. How many terms did we include? The vector $\hat{u}$ we computed contains just the first 64 Fourier modes (because we chose to use 64 points in our spatial grid vector $x$).
What about the evolution in time? In fact, our time evolution of the solution is exact for the initial data vector, since it just uses the exact solution formula for the ODE that we derived above.
In plotting the solution, note that we only used the values at the 64 spatial grid points. The plot() function merely connects these values by straight lines. We could plot a better representation of the solution by evaluating the Fourier series on a finer grid, but for now we won't worry about that.
The approach we just described isn't particular to the advection-diffusion equation. In fact, it can be used to solve any linear evolution PDE (including systems of PDEs, but here we'll stick to scalar problems):
$$u_t = \sum_j \alpha_j \frac{\partial^j u}{\partial x^j}.$$details to be added here