Linear advection equation using FFT and RK4

#%config InlineBackend.figure_format='svg'
from pylab import *

# Evaluate f:[a,b] -> R at given x by extending periodically
def periodic(a,b,f,x):
    y = a + mod(x-a, b-a)
    fv = vectorize(f)
    return fv(y)

# Set up grid and differentiation matrix:
N = 128; h = 2*pi/N; x = h*arange(1,N+1);
t = 0.0; dt = h/4.0
tmax = 2 * pi; tplot = 0.10;
plotgap = int(round(tplot/dt)); dt = tplot/plotgap;
nplots = int(round(tmax/tplot))

c = 1.0
f = lambda x: sin(20*x) * exp(-5*(x-pi)**2)
v = f(x)

# wave numbers
ik = 1j*zeros(N)
ik[0:N//2+1] = 1j*arange(0,N//2+1)
ik[N//2+1:]  = 1j*arange(-N//2+1,0,1)

def diff(u):
    u_hat = fft(u)
    w_hat = ik * u_hat
    w_hat[N//2] = 0.0
    return real(ifft(w_hat))

# Time-stepping by leap-frog formula
data = []; data.append(list(zip(x, v)))
tdata = []; tdata.append(0.0)
for i in range(1,nplots):
    for n in range(plotgap):
        t = t + dt
        k1 = -c * dt * diff(v)
        k2 = -c * dt * diff(v+0.5*k1)
        k3 = -c * dt * diff(v+0.5*k2)
        k4 = -c * dt * diff(v+k3)
        v += (1.0/6.0)*(k1 + 2.0*k2 + 2.0*k3 + k4)
    data.append(list(zip(x, v)))
    tdata.append(t);

from matplotlib import animation
from matplotlib import rc
rc('animation', html='jshtml')

# First set up the figure, the axis, and the plot element we want to animate
fig = figure(figsize=(10,6))
ax = axes(xlim=(0, 2*pi), ylim=(-1.1, 1.1))
line, = ax.plot([], [], 'r-', lw=2, label='u')
exact, = ax.plot([], [], 'b-', lw=1, label='exact')
ax.legend(loc='upper left')
ax.grid(True); ax.set_xlabel('x'); ax.set_ylabel('u')
close();

# initialization function: plot the background of each frame
def init():
    line.set_data([], [])
    return line,

# animation function.  This is called sequentially
def animate(i):
    x, v = zip(*data[i])
    line.set_data(x, v)
    vexact = periodic(0, 2*pi, f, array(x)-c*tdata[i])
    exact.set_data(x, vexact)
    return line,exact

# call the animator.  blit=True means only re-draw the parts that have changed.
anim = animation.FuncAnimation(fig, animate, init_func=init, repeat=False,
                               frames=len(data), interval=50, blit=True)
# Save to file
try:
    anim.save('linadv1d.mp4', fps=20, extra_args=['-vcodec', 'libx264'])
except:
    print("Cannot save mp4 file")

# Use this for inline display with controls
anim

# Use this for inline display of movie
#from IPython.display import HTML
#HTML(anim.to_html5_video())