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