1"""
2===========================
3The double pendulum problem
4===========================
5
6This animation illustrates the double pendulum problem.
7
8Double pendulum formula translated from the C code at
9http://www.physics.usyd.edu.au/~wheat/dpend_html/solve_dpend.c
10"""
11
12from numpy import sin, cos
13import numpy as np
14import matplotlib.pyplot as plt
15import scipy.integrate as integrate
16import matplotlib.animation as animation
17
18G = 9.8  # acceleration due to gravity, in m/s^2
19L1 = 1.0  # length of pendulum 1 in m
20L2 = 1.0  # length of pendulum 2 in m
21M1 = 1.0  # mass of pendulum 1 in kg
22M2 = 1.0  # mass of pendulum 2 in kg
23
24
25def derivs(state, t):
26
27    dydx = np.zeros_like(state)
28    dydx[0] = state[1]
29
30    del_ = state[2] - state[0]
31    den1 = (M1 + M2)*L1 - M2*L1*cos(del_)*cos(del_)
32    dydx[1] = (M2*L1*state[1]*state[1]*sin(del_)*cos(del_) +
33               M2*G*sin(state[2])*cos(del_) +
34               M2*L2*state[3]*state[3]*sin(del_) -
35               (M1 + M2)*G*sin(state[0]))/den1
36
37    dydx[2] = state[3]
38
39    den2 = (L2/L1)*den1
40    dydx[3] = (-M2*L2*state[3]*state[3]*sin(del_)*cos(del_) +
41               (M1 + M2)*G*sin(state[0])*cos(del_) -
42               (M1 + M2)*L1*state[1]*state[1]*sin(del_) -
43               (M1 + M2)*G*sin(state[2]))/den2
44
45    return dydx
46
47# create a time array from 0..100 sampled at 0.05 second steps
48dt = 0.05
49t = np.arange(0.0, 20, dt)
50
51# th1 and th2 are the initial angles (degrees)
52# w10 and w20 are the initial angular velocities (degrees per second)
53th1 = 120.0
54w1 = 0.0
55th2 = -10.0
56w2 = 0.0
57
58# initial state
59state = np.radians([th1, w1, th2, w2])
60
61# integrate your ODE using scipy.integrate.
62y = integrate.odeint(derivs, state, t)
63
64x1 = L1*sin(y[:, 0])
65y1 = -L1*cos(y[:, 0])
66
67x2 = L2*sin(y[:, 2]) + x1
68y2 = -L2*cos(y[:, 2]) + y1
69
70fig = plt.figure()
71ax = fig.add_subplot(111, autoscale_on=False, xlim=(-2, 2), ylim=(-2, 2))
72ax.set_aspect('equal')
73ax.grid()
74
75line, = ax.plot([], [], 'o-', lw=2)
76time_template = 'time = %.1fs'
77time_text = ax.text(0.05, 0.9, '', transform=ax.transAxes)
78
79
80def init():
81    line.set_data([], [])
82    time_text.set_text('')
83    return line, time_text
84
85
86def animate(i):
87    thisx = [0, x1[i], x2[i]]
88    thisy = [0, y1[i], y2[i]]
89
90    line.set_data(thisx, thisy)
91    time_text.set_text(time_template % (i*dt))
92    return line, time_text
93
94ani = animation.FuncAnimation(fig, animate, np.arange(1, len(y)),
95                              interval=25, blit=True, init_func=init)
96
97plt.show()
98