-
Notifications
You must be signed in to change notification settings - Fork 0
/
rod.py
162 lines (141 loc) · 5.46 KB
/
rod.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
153
154
155
156
157
158
159
160
161
162
import numpy as np
import scipy.integrate as integrate
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from visualization import Animator, AnimatableLinePlot, AnimatableSegmentData
class Rod:
def __init__(self,
center_pos=[10, 10], # pos: x z
center_theta=0,
d_center_pos=[0, 3],
d_center_theta=0,
h_floor=0,
L=1.0,
mass=1.0,
inertia=1.0,
lamda=0.8,
g=9.8):
self.state = np.array([*center_pos, center_theta, *d_center_pos, d_center_theta])
self.L = L
self.mass = mass
self.inertia = inertia
self.h_floor = h_floor
self.g = g
self.lamda = lamda
def dstate_dt(self, t, state):
"""derivative of state at time t"""
return np.array([
*state[3:5], # d_center_pos
state[5], # d_center_theta
0, -self.g, # dd_center_pos
0 # dd_center_theta
])
def jacobian(self, state):
"""d [pos1, pos2] / d state"""
theta = state[2]
return np.array([
[1, 0, +self.L / 2 * np.sin(theta)], # pos1 x
[0, 1, -self.L / 2 * np.cos(theta)], # pos1 z
[1, 0, -self.L / 2 * np.sin(theta)], # pos2 x
[0, 1, +self.L / 2 * np.cos(theta)] # pos2 z
])
def get_pos1(self):
center = self.state[0:2]
theta = self.state[2]
return center - self.L / 2 * np.array([np.cos(theta), np.sin(theta)])
def get_pos2(self):
center = self.state[0:2]
theta = self.state[2]
return center + self.L / 2 * np.array([np.cos(theta), np.sin(theta)])
def get_theta(self):
return self.state[2]
def get_d_theta(self):
return self.state[5]
def get_center_speed(self):
return self.state[3:5]
def event(self, t, state):
"""a continuous function that its change in sign signals ode to stop progressing"""
centerz = state[1]
theta = state[2]
z1 = centerz - self.L / 2 * np.sin(theta)
z2 = centerz + self.L / 2 * np.sin(theta)
return min(z1, z2) - self.h_floor
def apply_ground_impulse(self):
# switch
# governing equations:
# M ddq = J^T F_tip -> M dq_post - J^T Impulse = M dq
# dp_post = -lambda * dp -> J dq_post = -lambda J dq
# [[M -J^T] [dq_post = [ M dq
# [J 0]] Impulse] = -lambda J dq ]
########
M = np.array([[self.mass, 0, 0],
[0, self.mass, 0],
[0, 0, self.inertia]])
J = self.jacobian(self.state)
if self.get_pos1()[1] < self.get_pos2()[1]: ## pos1 lower
J = np.array([J[1]])
else: ## pos2 lower
J = np.array([J[3]])
dq = self.state[3:].reshape(3, 1)
A = np.block([
[M, -J.T],
[J, np.zeros((J.shape[0], J.shape[0]))]])
b = np.block([
[M @ dq],
[-self.lamda * J @ dq]
])
inversed = np.linalg.inv(A) @ b
dq_post = inversed[0:3]
self.state = np.array([*self.state[0:3], *dq_post.reshape((3,))])
def step(self, dt):
if dt <= 1e-9:
return
new_state = integrate.odeint(self.dstate_dt, self.state, [0, dt], tfirst=True)[1]
if (self.event(dt, new_state) > 0) ^ (self.event(dt, self.state) > 0): # change in even
# for now just ignore finding the exact time of the hit
self.apply_ground_impulse()
else:
self.state = new_state
if __name__ == "__main__":
fig, (ax_rod, ax_speed_log, ax_theta_log, ax_d_theta_log) = plt.subplots(4, 1)
rod = Rod(
center_pos=[0, 5], # pos: x z
center_theta=1,
d_center_pos=[4, 0],
d_center_theta=0,
h_floor=0,
L=3.0,
mass=1.0,
inertia=0.1,
lamda=0.9,
g=9.81)
p1_data = []
p2_data = []
v_data = []
theta_data = []
dtheta_data = []
t_data = []
dt = 0.001
total_time = 10
cur_time = 0
while cur_time < total_time:
p1_data.append(rod.get_pos1())
p2_data.append(rod.get_pos2())
v_data.append(np.linalg.norm(rod.get_center_speed()))
theta_data.append(rod.get_theta())
dtheta_data.append(rod.get_d_theta())
t_data.append(cur_time)
cur_time += dt
rod.step(dt)
animPoint = AnimatableSegmentData(ax_rod, p1_data=p1_data, p2_data=p2_data, name="rod", fix_scale=True)
animSpeedLog = AnimatableLinePlot(ax_speed_log, x_data=t_data, y_data=v_data, name="speed center", fix_scale=False)
animThetaLog = AnimatableLinePlot(ax_theta_log, x_data=t_data, y_data=v_data, name="theta", fix_scale=False)
animDThetaLog = AnimatableLinePlot(ax_d_theta_log, x_data=t_data, y_data=v_data, name="speed theta", fix_scale=False)
animator = Animator(fig=fig, interval=30, total_time=total_time, animatable_datas=[animPoint, animSpeedLog,
animThetaLog, animDThetaLog],
speed=1.0)
anim = FuncAnimation(animator.fig,
animator.animate,
frames=animator.num_frames,
interval=animator.interval)
plt.show()