kdv_plain.py

kdv_plain.py

from __future__ import print_function

import numpy
import math
from ode import solve_ode
from differential import make_differential_ops
import matplotlib.pyplot as plt
import matplotlib.animation as animation


# KdV equation
# u_{t} = -6 u u_{x} - u_{xxx}
def f_kdv(u, t, df1, df3):
    u_x = df1.dot(u)
    u_xxx = df3.dot(u)
    return -6.0 * u * u_x - u_xxx


if __name__ == '__main__':

    nx = 1000
    x_max = 100.0

    nt = 10001
    t_max = 10.0

    # "runge_kutta" (default), "euler", or "odeint"
    algo = "runge_kutta"

    # x mesh
    x = numpy.linspace(0, x_max, nx, endpoint=False)
    dx = x_max / float(nx)
    print("dx =", dx)

    # initial condition
    u0 = numpy.sin(x * (2.0 * math.pi / x_max))

    # t mesh
    t_array = numpy.linspace(0, t_max, nt)
    dt = t_max / float(nt-1)
    print("dt =", dt)
    print("dt / dx^3 =", dt/dx/dx/dx)

    # differential operators
    op_df1, _, op_df3 = make_differential_ops(nx, dx)

    print("Solving equation...")
    args = (op_df1, op_df3)
    u_tx = solve_ode(f_kdv, u0, t_array, args=args, algo=algo)

    print("Making animation...")
    t_skip = 100
    u_tx_skip = u_tx[::t_skip, :]

    fig = plt.figure()
    def plot(i):
        plt.cla()
        plt.title("t = %f" % float(i*t_skip*dt))
        plt.xlabel("$x$")
        plt.ylabel("$u(x)$")
        plt.ylim([-1.5, 3.0])
        plt.plot(x, u_tx_skip[i, :])

    anim = animation.FuncAnimation(fig, plot, frames=u_tx_skip.shape[0], interval=100, repeat=False)
    anim.save("kdv_plain.gif", writer="imagemagick")
    # plt.show()