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")