Error in RK4 algorithm in Python

Two different errors were obvious in the first parsing.

  1. (was not an error, ifft is indeed the complete inverse of fft. This might not be the case in other libraries.)

  2. In the RK4 step, you have to decide on one place for the factor h. Either (as example, the others analogously)

     k2 = f(t+0.5*h, y+0.5*h*k1)
    

    or

     k2 = h*f(t+0.5*h, y+0.5*k1)
    

However, correcting these points only delays the blowup. That there is the possibility for dynamical blow-up is no wonder, it is to be expected from the cubic term. In general one can only expect “slow” exponential growth if all terms are linear or sub-linear.

To avoid “unphysical” singularities, one has to scale the step size inversely proportional to the Lipschitz constant. Since the Lipschitz constant here is of size u^2, one has to dynamically adapt. I found that using 1000 steps in the interval [0,1], i.e., h=0.001, proceeds without singularity. This still holds true for 10 000 steps on the interval [0,10].


Update The part of the original equation without the time derivative is self-adjoint, which means that the norm square of the function (integral over the square of the absolute value) is preserved in the exact solution. The general picture is thus a high-dimensional “rotation” (cf. discussions of solid-body kinematics for pattern that evolve already in 3 dimensions).

The problem now is that parts of the function might “rotate” with such a small radius or such a high velocity that a time step represents a large part of a rotation or even multiple rotations. This is hard to capture with numerical methods, which requires thus a reduction in the time step. The general name for this phenomenon is “stiff differential equation”: Explicit Runge-Kutta methods are ill-suited for stiff problems.


Update2: Employing methods used before, one can solve the linear part in the decoupled frequency domain using (note that these are all component-wise array operations)

vhat = exp( 0.5j * kx**2 * t) * uhat

which allows for a stable solution with larger step sizes. As in the treatment of the KdV equation, the linear part i*u_t+0.5*u_xx=0 decouples under the DFT to

i*uhat_t-0.5*kx**2*uhat=0 

and can thus be easily solved into the corresponding exponential functions

exp( -0.5j * kx**2 * t).

The full equation is then tackled using variation of constants by setting

uhat = exp( -0.5j * kx**2 * t)*vhat. 

This lifts some of the burden of the stiffness for the larger components of kx, but still, the third power remains. So if the step size gets to large, the numerical solution explodes in very few steps.

Working code below


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(odefunc, t,w,h):
    k1 = odefunc(t,w)
    k2 = odefunc(t+0.5*h, w+0.5*k1*h)
    k3 = odefunc(t+0.5*h, w+0.5*k2*h)
    k4 = odefunc(t+h,     w+k3*h)
    return w + (k1+2*k2+2*k3+k4)*(h/6.)

def RK4Stream(odefunc,TimeSpan,uhat0,nt):
    h = float(TimeSpan[1]-TimeSpan[0])/nt
    print(f"step size {h}") 
    w = uhat0
    t = TimeSpan[0]
    while True:
        w = RK4Step(odefunc, t, w, h)
        t = t+h
        yield t,w

#----- Constructing the grid and kernel functions -----
L   = 40
nx  = 512
x, dx = np.linspace(-L/2,L/2, nx+1, retstep=True)
x   = x[:-1]  # periodic boundary, last same as first

kx  = 2*np.pi*np.fft.fftfreq(nx, dx) # angular frequencies for the fft bins

def uhat2vhat(t,uhat):
    return np.exp( 0.5j * (kx**2) *t) * uhat

def vhat2uhat(t,vhat):
    return np.exp(- 0.5j * (kx**2) *t) * vhat

#----- Define RHS -----
def uhatprime(t, uhat):
    u = np.fft.ifft(uhat)
    return - 0.5j * (kx**2) * uhat + 1j * np.fft.fft((abs(u)**2) * u)

def vhatprime(t, vhat):
    u = np.fft.ifft(vhat2uhat(t,vhat))
    return  1j * uhat2vhat(t, np.fft.fft((abs(u)**2) * u) )

#------ Initial Conditions -----
u0      = 1./np.cosh(x) #+ 1./np.cosh(x+0.4*L)+1./np.cosh(x-0.4*L) #symmetric or remove jump at wrap-around
uhat0   = np.fft.fft(u0)

#------ Solving for ODE -----
t0 = 0; tf = 10.0;
TimeSpan = [t0, tf]
# nt       = 500 # limit case, barely stable, visible spurious bumps in phase
nt       = 1000 # boring  but stable. smaller step sizes give same picture
vhat0 = uhat2vhat(t0,uhat0)

fig = plt.figure()
fig = plt.figure()
gs = fig.add_gridspec(3, 2)
ax1 = fig.add_subplot(gs[0, :]) 
ax2 = fig.add_subplot(gs[1:, :])
ax1.set_ylim(-0.2,2.5); ax1.set_ylabel("$u$ amplitude")
ax2.set_ylim(-6.4,6.4); ax2.set_ylabel("$u$ angle"); ax2.set_xlabel("$x$")

line1, = ax1.plot(x,u0)
line2, = ax2.plot(x,u0*0)

vhatstream = RK4Stream(vhatprime,[t0,tf],vhat0,nt)

def animate(i):
    t,vhat = vhatstream.next()
    print(f"time {t}")
    u = np.fft.ifft(vhat2uhat(t,vhat))
    line1.set_ydata(np.real(np.abs(u)))
    angles = np.real(np.angle(u))
    # connect the angles over multiple periods
    offset = 0;
    tau = 2*np.pi
    if angles[0] > 1.5: offset = -tau
    if angles[0] < -1.5: offset = tau
    for i,a in enumerate(angles[:-1]):
        diff_a = a-angles[i+1]
        angles[i] += offset
        if diff_a > 2 : 
            offset += tau
            if offset > 9: offset = tau-offset
        if diff_a < -2 : 
            offset -= tau
            if offset < -9: offset = -tau-offset
    angles[-1] += offset
    line2.set_ydata(angles)
    return line1,line2

anim = animation.FuncAnimation(fig, animate, interval=15000/nt+10, blit=False)

plt.show()

One can speed up the animation by computing multiple RK4 steps per frame, increasing the visible step size.

If one wants to use the ODE solvers from scipy.integrate, one has to implement some wrappers as these are not hardened against the use of complex-valued data.

# the stepper functions can not handle complex valued data
def RK45Stream(odefunc,TimeSpan,uhat0,nt):
    def odefuncreal(t,ureal):
        u = ureal.reshape([2,-1])
        deriv = odefunc(t,u[0]+1j*u[1])
        return  np.concatenate([deriv.real, deriv.imag])
    t0,tf = TimeSpan
    h = float(tf-t0)/nt
    print("step size ", h) 
    w = np.concatenate([uhat0.real, uhat0.imag])
    t = t0
    stepper = RK45(odefuncreal, t0, w, tf, atol=1e-9, rtol=1e-12)
    out_t = t0
    while True:
        t = t+h
        while t > stepper.t: stepper.step()
        if t>out_t: out_t, sol = stepper.t, stepper.dense_output()
        w = sol(t); w=w.reshape([2,-1])
        yield t,w[0]+1j*w[1]

Leave a Comment