Written by Yair Mau. Check out my webpage for more tutorials: http://www.yairmau.com/

In [1]:
# comment these lines if you want interactive mode,
# i.e., if you want to see the animation in real time.
import matplotlib
matplotlib.use('Agg')

In [2]:
import matplotlib.pyplot as plt
import numpy as np
import os
from scipy.integrate import ode

In [3]:
def equations(t, y, args):
""" the equations for the double pendulum """
x1 = y[0] # x1 = theta1, angle
x2 = y[1] # x2 = theta2, angle
p1 = y[2] # p1 = omega1, angular velocity
p2 = y[3] # p2 = omega2, angular velocity
l1,l2,m1,m2,g = args
x1_eq = p1
x2_eq = p2
p1_eq = -((g*(2*m1+m2)*np.sin(x1)+m2*(g*np.sin(x1-2*x2)+2*(l2*p2**2+l1*p1**2*np.cos(x1-x2))*np.sin(x1-x2)))/(2*l1*(m1+m2-m2*(np.cos(x1-x2))**2)))
p2_eq = ((l1*(m1+m2)*p1**2+g*(m1+m2)*np.cos(x1)+l2*m2*p2**2*np.cos(x1-x2))*np.sin(x1-x2))/(l2*(m1+m2-m2*(np.cos(x1-x2))**2))
return [x1_eq, x2_eq, p1_eq, p2_eq]

def calculate_trajectory(args,time,y0):
""" uses scipy's ode itegrator to simulate the equations """
t0,t1,dt = time
r = ode(equations).set_integrator('dopri5')
r.set_initial_value(y0, t0).set_f_params(args)
data=[[t0, y0[0], y0[1], y0[2], y0[3] ]]
while r.successful() and r.t < t1:
r.integrate(r.t+dt)
data.append([r.t, r.y[0], r.y[1], r.y[2], r.y[3] ])
return np.array(data)

def from_angle_to_xy(args,angles):
""" converts angles into xy positions """
l1,l2,m1,m2,g = args
time,theta1,theta2 = angles.T
x1 =  l1*np.sin(theta1)
y1 = -l1*np.cos(theta1)
x2 =  l2*np.sin(theta2) + x1
y2 = -l2*np.cos(theta2) + y1
return np.array([time,x1,y1,x2,y2]).T

In [4]:
l1 = 0.5 # length of arms
l2 = 0.5
m1 = 1.0 # mass of the pendulum
m2 = 1.0
g  = 10.0 # acceleration of gravity
args = [l1,l2,m1,m2,g]
fps = 80
total_time = 5 # seconds
time = [0.0,total_time,1.0/fps] # start, finish, dt
ic   = [np.pi*0.65, np.pi*1.1, 0.0, 0.0]

# here the magic happens
d = calculate_trajectory(args,time,ic)
data_TXY = from_angle_to_xy(args,d[:,:3])

In [5]:
# let's plot stuff, and make a nice movie
make_movie=True
params = {'backend': 'ps',
'font.size': 20,
'font.family':'serif',
'font.serif':['Computer Modern Roman'], # Times, Palatino, New Century Schoolbook, Bookman, Computer Modern Roman
'ps.usedistiller': 'xpdf',
'text.usetex': True,
}
plt.rcParams.update(params)
plt.ion()
fig = plt.figure(figsize=(9.6,5.4),dpi=100) # 1920x1080
ax.axis('off') # no frame

def plot_last_seconds(data,index):
""" Plots a line with the trajectory of the tip of pendulum 2 (x2,y2)
"""
how_long = 1.0 # seconds
n = int(how_long/time[2])
to_plot = data[:index,:]
if index < n:
prepend =  np.tile(data[0],(n-index,1))
to_plot = np.vstack([prepend,to_plot])
index = n
colormap = plt.cm.Greys_r
colors = [colormap(i) for i in np.linspace(0.0, 1.0, n-1)]
plots = []
for j in np.arange(n-1):
p, = ax.plot(to_plot[index-j-1:index-j+1,3],to_plot[index-j-1:index-j+1,4],
color=colors[j], zorder=-1)
plots.append(p)
return plots

# "plot" returns a tuple of line objects, thus the comma
t,x1,y1,x2,y2 = data_TXY[0]
line1, = ax.plot([0.0,x1], [0.0,y1], 'r-')
line2, = ax.plot([x1,x2], [y1,y2], 'r-')
circ1, = ax.plot([x1], [y1], 'ro',markersize=10)
circ2, = ax.plot([x2], [y2], 'ro',markersize=10)
sizeY = 1.2
ax.axis([-sizeY*16/9,sizeY*16/9,-sizeY,sizeY])

frame_names = []
tex=ax.text(0.0,0.85,'',ha="center")

for i,v in enumerate(data_TXY):
t,x1,y1,x2,y2 = v
# print("t={:.2f}".format(t)) # you might want to know how things are going...
line1.set_data([0.0,x1],[0.0,y1])
line2.set_data([x1,x2],[y1,y2])
circ1.set_data([x1],[y1])
circ2.set_data([x2],[y2])

# plot_last_seconds considerably slows down the simulation,
# but makes it much prettier...
pls = plot_last_seconds(data_TXY,i+1)
tex.set_text(r"$t={:.3f}$ s".format(t))
fig.canvas.draw()
if make_movie:
fname = "_tmp{:05d}.png".format(i)
frame_names.append(fname)
fig.savefig(fname,bbox_inches='tight')
for k in pls:
k.remove()

if make_movie:
frames = "_tmp%5d.png"
frames = "_tmp%5d.png"
movie_command = "ffmpeg -y -r {:} -i {:} double.mp4".format(fps,frames)
os.system(movie_command)

for fname in frame_names:
os.remove(fname)