# Simulate the Duffing Oscillator
# Name:
# Date:
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint
import matplotlib.animation as animation

alpha = -1.0
beta = 1.0

trun = 20.0  # Total time to run
dt = 0.10    # Time step for display

x0 = 1.1   # Initial position
v0 = 0.0   # Initial velocity
s0 = [x0, v0]  # Initial condition state vector

def osc(s, t, alpha, beta):
    """Given the current state vector s=[x, v] at time t,
       compute the velocity and acceleration.  The return
       value is a list with two elements, [v, a]."""
    x, v = s
    dsdt = [v, -alpha * x - beta * x**3]
    return dsdt

# List of times at which to evaluate the solution
ts = np.arange(0, trun+dt, dt)

result  = odeint(osc, s0, ts, args=(alpha, beta))

# Convenient shorthand views of the x and v parts of the result array.
xs = result[:, 0]
vs = result[:, 1]

# For animation, we need to set the limits for the plot ahead of time,
# so extract them now.
tmin, tmax = [np.min(ts), np.max(ts)]
xmin, xmax = [np.min(xs), np.max(xs)]
vmin, vmax = [np.min(vs), np.max(vs)]

fig = plt.figure()
ax = plt.axes(xlim=(tmin, tmax), ylim=(xmin, xmax))
line, = ax.plot([], [], 'b', label='Duffing')
plt.xlabel('t')
plt.ylabel('x')
plt.grid()

def init():
    line.set_data([], [])
    return line,

def animate(i):
    line.set_data(ts[:i], xs[:i])
    return line,

ani = animation.FuncAnimation(fig, animate, np.arange(0, len(xs)),
                      repeat=False, interval=5, blit=True, init_func=init)

plt.legend(loc='best')
plt.show()