kdv.py

kdv.py

from __future__ import print_function

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


class KDV(object):
    def __init__(self, nx, xmin, xmax, nt, tmax):

        # x mesh
        self.x = numpy.linspace(xmin, xmax, nx, endpoint=False)
        dx = xmax / float(nx)
        print("dx =", dx)

        # t mesh
        self.t = numpy.linspace(0, tmax, nt)
        self.dt = tmax / float(nt-1)
        print("dt =", self.dt)

        print("dt / dx^3 =", self.dt/dx**3)

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

        self.u0 = None  # set in set_initial(9
        self.u = None  # set in solve()

    def solve(self, algo="runge_kutta"):
        print("Solving equation...")
        if self.u0 is None:
            print("Set u0 in advance")
            exit(1)

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

        args = (self.op_df1, self.op_df3)
        self.u = solve_ode(func, self.u0, self.t, args=args, algo=algo)

    def save_animation(self, skip, xlim=None, ylim=None, filename=None):
        print("Making animation...")
        if self.u is None:
            print("Call solve() in advance")
            exit(1)

        u_skip = self.u[::skip, :]

        def plot(i):
            plt.cla()
            plt.title("t = %f" % float(i * self.dt * skip))
            plt.xlabel("$x$")
            plt.ylabel("$u(x)$")
            if xlim is not None:
                plt.xlim(xlim)
            if ylim is not None:
                plt.ylim(ylim)
            plt.plot(self.x, u_skip[i, :])

        fig = plt.figure()
        anim = animation.FuncAnimation(fig, plot, frames=u_skip.shape[0], interval=100, repeat=False)
        if filename is None:
            plt.show()
        else:
            anim.save(filename, writer="imagemagick")