/
double_pendulum.py
149 lines (116 loc) · 4.56 KB
/
double_pendulum.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
# Credit for this code is given to author at
# https://jakevdp.github.io/blog/2012/08/18/matplotlib-animation-tutorial/
from numpy import sin, cos
import numpy as np
import matplotlib.pyplot as plt
import scipy.integrate as integrate
import matplotlib.animation as animation
class DoublePendulum:
""" Double Pendulum class
init_state is [theta1, omega1, theta2, omega2] in degrees,
where theta1, omega1 is the angular position and velocity of
the first pendulum arm, and theta2, omega2 is that of the second
pendulum arm
"""
def __init__(self,
init_state = [120, 0, -20, 0],
L1 = 1.0, # length of arm 1
L2 = 1.0, # length of arm 2
M1 = 1.0, # mass 1
M2 = 1.0, # mass 2
G = 9.81, # accel due to gravity
origin=( 0 , 0 )):
self.init_state = np.asarray(init_state, dtype='float')
self.params = (L1, L2, M1, M2, G)
self.origin = origin
self.time_elapsed = 0
self.state = self.init_state * np.pi / 180.0
def position(self):
"""compute the current x,y positions of the pendulum arms"""
(L1, L2, M1, M2, G) = self. params
x = np.cumsum([self.origin[0],
L1 * sin(self.state[0]),
L2 * sin(self.state[2])])
y = np.cumsum([self.origin[1],
-L1 * cos(self.state[0]),
-L2 * cos(self.state[2])])
return (x,y)
def energy(self):
"""compute the energy of the current state"""
(L1, L2, M1, M2, G) = self.params
x = np.cumsum([L1 * sin(self.state[0]),
L2 * sin(self.state[2])])
y = np.cumsum([-L1 * cos(self.state[0]),
-L2 * cos(self.state[2])])
vx = np.cumsum([L1 * self.state[1] * cos(self.state[0]),
L2 * self.state[3] * cos(self.state[2])])
vy = np.cumsum([L1 * self.state[1] * sin(self.state[0]),
L2 * self.state[3] * sin(self.state[2])])
U = G * (M1 * y[0] + M2 * y[1])
K = 0.5 * (M1 * np.dot(vx, vx) + M2 * np.dot(vy, vy))
return U + K
def dstate_dt(self, state, t):
""" compute the derivative of the given state"""
(L1, L2, M1, M2, G) = self.params
dydx = np.zeros_like(state)
dydx[0] = state[1]
dydx[2] = state[3]
cos_delta = cos(state[2] - state[0])
sin_delta = sin(state[2] - state[0])
den1 = (M1 + M2) * L1 - M2 * L1 * cos_delta * cos_delta
dydx[1] = (M2 * L1 * state[1] * state[1] * sin_delta * cos_delta
+ M2 * G * sin(state[2]) * cos_delta
+ M2 * L2 * state[3] * state[3] * sin_delta
- (M1 + M2) * G * sin(state[0])) / den1
den2 = (L2 / L1) * den1
dydx[3] = (-M2 * L2 * state[3] * state[3] * sin_delta * cos_delta
+ (M1 + M2) * G * sin(state[0]) * cos_delta
- (M1 + M2) * L1 * state[1] * state[1] * sin_delta
- (M1 + M2) * G * sin(state[2])) / den2
return dydx
def step(self, dt):
"""execute one time step of length dt and update state"""
self.state = integrate.odeint(self.dstate_dt, self.state, [0, dt])[1]
self.time_elapsed += dt
#------------------------------------------------------------
# set up initial state and global variables
pendulum = DoublePendulum([90., 0.0, 90., 0.0])
dt = 1./30 # 30 fps
#------------------------------------------------------------
# set up figure and animation
fig = plt.figure()
ax = fig.add_subplot(111, aspect='equal', autoscale_on=False,
xlim=(-2, 2), ylim=(-2, 2))
ax.grid()
line, = ax.plot([], [], 'o-', lw=2)
time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes)
energy_text = ax.text(0.02, 0.90, '', transform=ax.transAxes)
def init():
"""initialize animation"""
line.set_data([], [])
time_text.set_text('')
energy_text.set_text('')
return line, time_text, energy_text
def animate(i):
"""perform animation step"""
global pendulum, dt
pendulum.step(dt)
line.set_data(*pendulum.position())
time_text.set_text('time = %.1f' % pendulum.time_elapsed)
energy_text.set_text('energy = %.3f J' % pendulum.energy())
return line, time_text, energy_text
# choose the interval based on dt and the time to animate one step
from time import time
t0 = time()
animate(0)
t1 = time()
interval = 1000 * dt - (t1 - t0)
ani = animation.FuncAnimation(fig, animate, frames=300,
interval=interval, blit=True, init_func=init)
# save the animation as an mp4. This requires ffmpeg or mencoder to be
# installed. The extra_args ensure that the x264 codec is used, so that
# the video can be embedded in html5. You may need to adjust this for
# your system: for more information, see
# http://matplotlib.sourceforge.net/api/animation_api.html
#ani.save('double_pendulum.mp4', fps=30, extra_args=['-vcodec', 'libx264'])
plt.show()