Implement pseudo-spectral method with RK4 in Python

Question

I am using pseudo-spectral method to solve for KdV equation u_t + u*u_x + u_xxx = 0. After simplification by Fourier Transform, I got two equations with two variables:

1. uhat = vhat * exp(-i*k^3*t)
2. d(vhat)/dt =-0.5i * k * exp(-i*k^3*t)*F[u^2]

where F represents Fourier Transform, uhat = F[u], vhat = F[v]

I'd like to use RK4 to solve for u by è ifft(uhat)è eventually. Now I got two variables uhat and vhat, I have 2 ideas for solving uhat and vhat:

1. To treat them as a system of ODEs to implement RK4.

2. Treat the equation 2 above as the to-be solved ODE to solve chat by RK4, and calculate uhat = vhat * exp(-i*k^2*delta_t) after each time step delta_t.

I got problem to implement both ideas. Here are the codes for the 2nd idea above.

import numpy as np
import math
from matplotlib import pyplot as plt
from matplotlib import animation


#----- Numerical integration of ODE via fixed-step classical Runge-Kutta -----

def RK4Step(yprime,t,y,dt):
    k1 = yprime(t       , y          )
    k2 = yprime(t+0.5*dt, y+0.5*k1*dt)
    k3 = yprime(t+0.5*dt, y+0.5*k2*dt)
    k4 = yprime(t+    dt, y+    k3*dt)
    return y + (k1+2*k2+2*k3+k4)*dt/6. 

def RK4(yprime,times,y0):
    y = np.empty(times.shape+y0.shape,dtype=y0.dtype)
    y[0,:] = y0                         # enter initial conditions in y
    steps = 4
    for i in range(times.size-1):
        dt = (times[i+1]-times[i])/steps
        y[i+1,:] = y[i,:]
        for k in range(steps):
            y[i+1,:] = RK4Step(yprime, times[i]+k*dt, y[i+1,:], dt)
            y[i+1,:] = y[i+1,:] * np.exp(delta_t * 1j * kx**3)
    return y

#====================================================================
#----- Parameters for PDE -----
L       = 100
n       = 512
delta_t = 0.1
tmax    = 20
c1      = 1.5  # amplitude of 1st wave
c2      = 0.5  # amplitude of 2nd wave

#----- Constructing the grid and kernel functions
x2      = np.linspace(-L/2,L/2, n+1)
x       = x2[:n]                         # periodic B.C. #0 = #n

kx1     = np.linspace(0,n/2-1,n/2)
kx2     = np.linspace(1,n/2,  n/2)
kx2     = -1*kx2[::-1]
kx      = (2.* np.pi/L)*np.concatenate((kx1,kx2)); kx[0] = 1e-6

#----- Initial Condition -----
z1      = np.sqrt(c1)/2. * (x-0.1*L)
z2      = np.sqrt(c2)/2. * (x-0.4*L)
soliton = 6*(0.5 * c1 / (np.cosh(z1))**2 + 0.5 * c2/ (np.cosh(z2))**2)
uhat0   = np.fft.fft(soliton)
vhat0   = uhat0

#----- Define ODE -----
def wprime(t,vhat):  
    g = -0.5 * 1j* kx * np.exp(-1j * kx**3 * t)
    return g * np.fft.fft(np.real(np.fft.ifft(np.exp(1j*kx**3*t)*vhat)**2))

#====================================================================
#----- Compute the numerical solution -----
TimeStart = 0.
TimeEnd   = tmax+delta_t
TimeSpan  = np.arange(TimeStart, TimeEnd, delta_t)
w_sol     = RK4(wprime,TimeSpan, vhat0)

#----- Animate the numerical solution -----
fig = plt.figure()
ims = []
for i in TimeSpan:
    w = np.real(np.fft.ifft(w_sol[i,:]))
    im = plt.plot(x,w)
    ims.append([im])

ani = animation.ArtistAnimation(fig, ims, interval=100, blit=False)

plt.show()

The RK4 part was from @LutzL.

Ardent

Answer 1

In general it looks good. Better use a FuncAnimation, reduces the memory footprint and works:

#----- Animate the numerical solution -----
fig = plt.figure()
ax = plt.axes(ylim=(-5,5))
line, = ax.plot(x, soliton, '-')

NT = TimeSpan.size;

def animate(i):
    i = i % (NT + 10)
    if i<NT:
        w = np.real(np.fft.ifft(w_sol[i,:]))
        line.set_ydata(w)
    return line,


anim = animation.FuncAnimation(fig, animate, interval=50, blit=True)

plt.show()

However, the first graph remains standing, I do not know why.

Physiognomy

I changed the 'blit=True' to 'False', then the animation part works. However, the behavior of the two solitons is weird, which is not what I am looking for. What I am looking for is the result like in this video (youtube.com/watch?v=v5MGNcCnuE4). So, there must be something wrong in the RK4 implementation part in this case, I am not sure how to treat the such two variables in Fourier Space. Thank you! – Ardent