Solve Burger's equation with spectral method

Burgers.py

# Burger's equation

from numpy import *
from scipy.special import erf
import matplotlib.pyplot as plt

N = 512
x = 2*pi*arange(N)/N
k = fft.rfftfreq(N, 1./N)

u = exp(-(x-pi)**2/0.1)
u_hat = fft.rfft(u)

u0 = zeros_like(u)
du0 = zeros_like(u)
nonlin = zeros_like(u_hat)
u1 = zeros_like(u_hat)
u2 = zeros_like(u_hat)
rhs = zeros_like(u_hat)
a = array([1./6., 1./3., 1./3., 1./6.], dtype=float)
b = array([0.5, 0.5, 1.], dtype=float)

plt.figure()
plt.plot(x, u)
ax = plt.gca()

t = 0
dt = 0.002
T = 10
nu = 0.01
while t < T:
    t += dt
    
    # Integrate using explicit fourth order Runge Kutta
    u2[:] = u1[:] = u_hat
    for rk in range(4):
        # Compute nonlinear convection
        u0[:] = fft.irfft(u_hat)       # No dealiasing (yet). Would go here
        du0[:] = fft.irfft(1j*k*u_hat) # du/dx
        nonlin[:] = fft.rfft(u0*du0)
        # Compute right hand side of equation
        rhs[:] = -nu*k**2*u_hat - nonlin
        if rk < 3:
            u_hat[:] = u1 + b[rk]*dt*rhs
        u2 += a[rk]*dt*rhs
    u_hat[:] = u2

    ax.clear()
    u[:] = fft.irfft(u_hat) 
    plt.plot(x, u, 'b')
    plt.draw()
    plt.pause(1e-6)