-
Notifications
You must be signed in to change notification settings - Fork 16
/
pendulum.py
152 lines (117 loc) · 4.38 KB
/
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
149
150
151
152
from collections import namedtuple
import numpy as np
np.seterr(over='raise')
g = 9.81
Pendulum = namedtuple('Pendulum', 'm x y u v')
class DoublePendulum:
"""Double pendulum consisting of two massless bodies and a point masses on
their ends.
Attributes:
M (np.array, (4, 4)): mass matrix
M_inv (np.array, (4, 4)): inverse of mass matrix
f (np.array, (4, 1)): vector of gravitational forces
y_0 (np.array, (8,)): initial positions and velocities
z_0 (np.array, (2, 1)): initial Lagrange multipliers
"""
def __init__(self, body_1, body_2):
"""Initializes double pendulum instance.
Args:
body_1, body_2 (namedtuple):
m: point mass on the end of the body
x: x-coordinate of the point mass
y: y-coordinate of the point mass
u: velocity in x-direction of the point mass
v: velocity in y-direction of the point mass
"""
self._check_consistency(body_1, body_2)
self.M = np.diag([body_1.m, body_2.m, body_1.m, body_2.m])
self.M_inv = np.linalg.inv(self.M)
self.f = -g * np.array([0, 0, body_1.m, body_2.m]).reshape((4, 1))
self.y_0 = np.array([body_1.x, body_2.x, body_1.y, body_2.y,
body_1.u, body_2.u, body_1.v, body_2.v])
self.z_0 = self.get_z(self.y_0)
self._b1 = body_1
self._b2 = body_2
def get_z(self, y):
"""Calculate Lagrange multipliers for `y`.
See AP98, p. 248.
Args:
y (np.array, (8,)): <x1, x2, y1, y2, u1, u2, v1, v2>
Returns:
np.array, (2,)
"""
G = self.get_G(y)
g_uu = self.get_g_uu(y)
G_Mi = G.dot(self.M_inv)
A = G_Mi.dot(G.T)
b = G_Mi.dot(self.f) + g_uu
return np.linalg.solve(A, b).reshape(2,)
def get_dy(self, y):
"""Calculate `y' = dy/dt` by solving the system
[[M G.T] [[du] [[ f ]
[G 0 ]] @ [z ]] = [g_uu]]
See HW96, p. 465 for more details.
Args:
y (np.array, (8,)): <x1, x2, y1, y2, u1, u2, v1, v2>
Returns:
np.array, (8,): <dx1, dx2, dy1, dy2, du1, du2, dv1, dv2>
"""
G = self.get_G(y)
g_uu = self.get_g_uu(y)
A = np.zeros((6, 6))
b = np.zeros((6, 1))
A[:4, :4] = self.M
A[:4, -2:] = G.T
A[-2:, :4] = G
b[:4] = self.f
b[-2:] = -1 * g_uu
x = np.linalg.solve(A, b)
return np.concatenate((y[-4:], x[:4, 0]))
@staticmethod
def get_G(y):
"""Calculate matrix `G = grad(g)`, where `g` are algebraic constraints.
Args:
y (np.array, (8,)): <x1, x2, y1, y2, u1, u2, v1, v2>
Returns:
np.array, (2, 4)
"""
x1, x2, y1, y2 = y[:4]
return np.array([[x1, 0, y1, 0],
[-(x2-x1), (x2-x1), -(y2-y1), (y2-y1)]])
@staticmethod
def get_g_uu(y):
"""Calculate vector `g_uu = G' @ u`, needed to solve double pendulum
as index 1 DAE system. (See HW96, p. 465)
Args:
y (np.array, (8,)): <x1, x2, y1, y2, u1, u2, v1, v2>
Returns:
np.array, (2, 1)
"""
u1, u2, v1, v2 = y[-4:]
return np.array([[u1**2 + v1**2],
[(u2 - u1)**2 + (v2 - v1)**2]])
@staticmethod
def _check_consistency(body_1, body_2):
"""Checks if the initial values satisfy an index 2 constraint."""
if body_1.x * body_1.u + body_1.y * body_1.v != 0:
raise InconsistentInitialValues("First pendulum has inconsistent "
"initial values. Try setting velocities to zero.")
if ((body_2.x - body_1.x) * (body_2.u - body_1.u) +
(body_2.y - body_1.y) * (body_2.v - body_1.v)) != 0:
raise InconsistentInitialValues("Second pendulum has inconsistent "
"initial values. Try setting velocities to zero.")
class InconsistentInitialValues(Exception):
pass
def __run_basic_example():
p1 = Pendulum(m=5, x=1.5, y=-2, u=10, v=0)
p2 = Pendulum(m=15, x=5.5, y=-5, u=0, v=20)
dp = DoublePendulum(p1, p2)
help(dp)
print(dp.M)
print(dp.M_inv)
print(dp.f)
print(dp.y_0)
print(dp.z_0)
print(dp.get_dy(dp.y_0))
if __name__ == '__main__':
__run_basic_example()