Skip to content

Commit 4af9c8e

Browse files
AtsushiSakaiAtsushiSakai
committed
code cleanup
1 parent d3f03c6 commit 4af9c8e

File tree

1 file changed

+102
-102
lines changed

1 file changed

+102
-102
lines changed

‎AerialNavigation/rocket_powered_landing/rocket_powered_landing.py‎

Lines changed: 102 additions & 102 deletions
Original file line numberDiff line numberDiff line change
@@ -38,107 +38,7 @@
3838
show_animation = True
3939

4040

41-
class Integrator:
42-
def __init__(self, m, K):
43-
self.K = K
44-
self.m = m
45-
self.n_x = m.n_x
46-
self.n_u = m.n_u
47-
48-
self.A_bar = np.zeros([m.n_x * m.n_x, K - 1])
49-
self.B_bar = np.zeros([m.n_x * m.n_u, K - 1])
50-
self.C_bar = np.zeros([m.n_x * m.n_u, K - 1])
51-
self.S_bar = np.zeros([m.n_x, K - 1])
52-
self.z_bar = np.zeros([m.n_x, K - 1])
53-
54-
# vector indices for flat matrices
55-
x_end = m.n_x
56-
A_bar_end = m.n_x * (1 + m.n_x)
57-
B_bar_end = m.n_x * (1 + m.n_x + m.n_u)
58-
C_bar_end = m.n_x * (1 + m.n_x + m.n_u + m.n_u)
59-
S_bar_end = m.n_x * (1 + m.n_x + m.n_u + m.n_u + 1)
60-
z_bar_end = m.n_x * (1 + m.n_x + m.n_u + m.n_u + 2)
61-
self.x_ind = slice(0, x_end)
62-
self.A_bar_ind = slice(x_end, A_bar_end)
63-
self.B_bar_ind = slice(A_bar_end, B_bar_end)
64-
self.C_bar_ind = slice(B_bar_end, C_bar_end)
65-
self.S_bar_ind = slice(C_bar_end, S_bar_end)
66-
self.z_bar_ind = slice(S_bar_end, z_bar_end)
67-
68-
self.f, self.A, self.B = m.f_func, m.A_func, m.B_func
69-
70-
# integration initial condition
71-
self.V0 = np.zeros((m.n_x * (1 + m.n_x + m.n_u + m.n_u + 2),))
72-
self.V0[self.A_bar_ind] = np.eye(m.n_x).reshape(-1)
73-
74-
self.dt = 1. / (K - 1)
75-
76-
def calculate_discretization(self, X, U, sigma):
77-
"""
78-
Calculate discretization for given states, inputs and total time.
79-
80-
:param X: Matrix of states for all time points
81-
:param U: Matrix of inputs for all time points
82-
:param sigma: Total time
83-
:return: The discretization matrices
84-
"""
85-
for k in range(self.K - 1):
86-
self.V0[self.x_ind] = X[:, k]
87-
V = np.array(odeint(self._ode_dVdt, self.V0, (0, self.dt),
88-
args=(U[:, k], U[:, k + 1], sigma))[1, :])
89-
90-
# using \Phi_A(\tau_{k+1},\xi) = \Phi_A(\tau_{k+1},\tau_k)\Phi_A(\xi,\tau_k)^{-1}
91-
# flatten matrices in column-major (Fortran) order for CVXPY
92-
Phi = V[self.A_bar_ind].reshape((self.n_x, self.n_x))
93-
self.A_bar[:, k] = Phi.flatten(order='F')
94-
self.B_bar[:, k] = np.matmul(Phi, V[self.B_bar_ind].reshape(
95-
(self.n_x, self.n_u))).flatten(order='F')
96-
self.C_bar[:, k] = np.matmul(Phi, V[self.C_bar_ind].reshape(
97-
(self.n_x, self.n_u))).flatten(order='F')
98-
self.S_bar[:, k] = np.matmul(Phi, V[self.S_bar_ind])
99-
self.z_bar[:, k] = np.matmul(Phi, V[self.z_bar_ind])
100-
101-
return self.A_bar, self.B_bar, self.C_bar, self.S_bar, self.z_bar
102-
103-
def _ode_dVdt(self, V, t, u_t0, u_t1, sigma):
104-
"""
105-
ODE function to compute dVdt.
106-
107-
:param V: Evaluation state V = [x, Phi_A, B_bar, C_bar, S_bar, z_bar]
108-
:param t: Evaluation time
109-
:param u_t0: Input at start of interval
110-
:param u_t1: Input at end of interval
111-
:param sigma: Total time
112-
:return: Derivative at current time and state dVdt
113-
"""
114-
alpha = (self.dt - t) / self.dt
115-
beta = t / self.dt
116-
x = V[self.x_ind]
117-
u = u_t0 + beta * (u_t1 - u_t0)
118-
119-
# using \Phi_A(\tau_{k+1},\xi) = \Phi_A(\tau_{k+1},\tau_k)\Phi_A(\xi,\tau_k)^{-1}
120-
# and pre-multiplying with \Phi_A(\tau_{k+1},\tau_k) after integration
121-
Phi_A_xi = np.linalg.inv(
122-
V[self.A_bar_ind].reshape((self.n_x, self.n_x)))
123-
124-
A_subs = sigma * self.A(x, u)
125-
B_subs = sigma * self.B(x, u)
126-
f_subs = self.f(x, u)
127-
128-
dVdt = np.zeros_like(V)
129-
dVdt[self.x_ind] = sigma * f_subs.transpose()
130-
dVdt[self.A_bar_ind] = np.matmul(
131-
A_subs, V[self.A_bar_ind].reshape((self.n_x, self.n_x))).reshape(-1)
132-
dVdt[self.B_bar_ind] = np.matmul(Phi_A_xi, B_subs).reshape(-1) * alpha
133-
dVdt[self.C_bar_ind] = np.matmul(Phi_A_xi, B_subs).reshape(-1) * beta
134-
dVdt[self.S_bar_ind] = np.matmul(Phi_A_xi, f_subs).transpose()
135-
z_t = -np.matmul(A_subs, x) - np.matmul(B_subs, u)
136-
dVdt[self.z_bar_ind] = np.dot(Phi_A_xi, z_t.T).flatten()
137-
138-
return dVdt
139-
140-
141-
class Model_6DoF:
41+
class Rocket_Model_6DoF:
14242
"""
14343
A 6 degree of freedom rocket landing problem.
14444
"""
@@ -413,6 +313,106 @@ def get_constraints(self, X_v, U_v, X_last_p, U_last_p):
413313
return constraints
414314

415315

316+
class Integrator:
317+
def __init__(self, m, K):
318+
self.K = K
319+
self.m = m
320+
self.n_x = m.n_x
321+
self.n_u = m.n_u
322+
323+
self.A_bar = np.zeros([m.n_x * m.n_x, K - 1])
324+
self.B_bar = np.zeros([m.n_x * m.n_u, K - 1])
325+
self.C_bar = np.zeros([m.n_x * m.n_u, K - 1])
326+
self.S_bar = np.zeros([m.n_x, K - 1])
327+
self.z_bar = np.zeros([m.n_x, K - 1])
328+
329+
# vector indices for flat matrices
330+
x_end = m.n_x
331+
A_bar_end = m.n_x * (1 + m.n_x)
332+
B_bar_end = m.n_x * (1 + m.n_x + m.n_u)
333+
C_bar_end = m.n_x * (1 + m.n_x + m.n_u + m.n_u)
334+
S_bar_end = m.n_x * (1 + m.n_x + m.n_u + m.n_u + 1)
335+
z_bar_end = m.n_x * (1 + m.n_x + m.n_u + m.n_u + 2)
336+
self.x_ind = slice(0, x_end)
337+
self.A_bar_ind = slice(x_end, A_bar_end)
338+
self.B_bar_ind = slice(A_bar_end, B_bar_end)
339+
self.C_bar_ind = slice(B_bar_end, C_bar_end)
340+
self.S_bar_ind = slice(C_bar_end, S_bar_end)
341+
self.z_bar_ind = slice(S_bar_end, z_bar_end)
342+
343+
self.f, self.A, self.B = m.f_func, m.A_func, m.B_func
344+
345+
# integration initial condition
346+
self.V0 = np.zeros((m.n_x * (1 + m.n_x + m.n_u + m.n_u + 2),))
347+
self.V0[self.A_bar_ind] = np.eye(m.n_x).reshape(-1)
348+
349+
self.dt = 1. / (K - 1)
350+
351+
def calculate_discretization(self, X, U, sigma):
352+
"""
353+
Calculate discretization for given states, inputs and total time.
354+
355+
:param X: Matrix of states for all time points
356+
:param U: Matrix of inputs for all time points
357+
:param sigma: Total time
358+
:return: The discretization matrices
359+
"""
360+
for k in range(self.K - 1):
361+
self.V0[self.x_ind] = X[:, k]
362+
V = np.array(odeint(self._ode_dVdt, self.V0, (0, self.dt),
363+
args=(U[:, k], U[:, k + 1], sigma))[1, :])
364+
365+
# using \Phi_A(\tau_{k+1},\xi) = \Phi_A(\tau_{k+1},\tau_k)\Phi_A(\xi,\tau_k)^{-1}
366+
# flatten matrices in column-major (Fortran) order for CVXPY
367+
Phi = V[self.A_bar_ind].reshape((self.n_x, self.n_x))
368+
self.A_bar[:, k] = Phi.flatten(order='F')
369+
self.B_bar[:, k] = np.matmul(Phi, V[self.B_bar_ind].reshape(
370+
(self.n_x, self.n_u))).flatten(order='F')
371+
self.C_bar[:, k] = np.matmul(Phi, V[self.C_bar_ind].reshape(
372+
(self.n_x, self.n_u))).flatten(order='F')
373+
self.S_bar[:, k] = np.matmul(Phi, V[self.S_bar_ind])
374+
self.z_bar[:, k] = np.matmul(Phi, V[self.z_bar_ind])
375+
376+
return self.A_bar, self.B_bar, self.C_bar, self.S_bar, self.z_bar
377+
378+
def _ode_dVdt(self, V, t, u_t0, u_t1, sigma):
379+
"""
380+
ODE function to compute dVdt.
381+
382+
:param V: Evaluation state V = [x, Phi_A, B_bar, C_bar, S_bar, z_bar]
383+
:param t: Evaluation time
384+
:param u_t0: Input at start of interval
385+
:param u_t1: Input at end of interval
386+
:param sigma: Total time
387+
:return: Derivative at current time and state dVdt
388+
"""
389+
alpha = (self.dt - t) / self.dt
390+
beta = t / self.dt
391+
x = V[self.x_ind]
392+
u = u_t0 + beta * (u_t1 - u_t0)
393+
394+
# using \Phi_A(\tau_{k+1},\xi) = \Phi_A(\tau_{k+1},\tau_k)\Phi_A(\xi,\tau_k)^{-1}
395+
# and pre-multiplying with \Phi_A(\tau_{k+1},\tau_k) after integration
396+
Phi_A_xi = np.linalg.inv(
397+
V[self.A_bar_ind].reshape((self.n_x, self.n_x)))
398+
399+
A_subs = sigma * self.A(x, u)
400+
B_subs = sigma * self.B(x, u)
401+
f_subs = self.f(x, u)
402+
403+
dVdt = np.zeros_like(V)
404+
dVdt[self.x_ind] = sigma * f_subs.transpose()
405+
dVdt[self.A_bar_ind] = np.matmul(
406+
A_subs, V[self.A_bar_ind].reshape((self.n_x, self.n_x))).reshape(-1)
407+
dVdt[self.B_bar_ind] = np.matmul(Phi_A_xi, B_subs).reshape(-1) * alpha
408+
dVdt[self.C_bar_ind] = np.matmul(Phi_A_xi, B_subs).reshape(-1) * beta
409+
dVdt[self.S_bar_ind] = np.matmul(Phi_A_xi, f_subs).transpose()
410+
z_t = -np.matmul(A_subs, x) - np.matmul(B_subs, u)
411+
dVdt[self.z_bar_ind] = np.dot(Phi_A_xi, z_t.T).flatten()
412+
413+
return dVdt
414+
415+
416416
class SCProblem:
417417
"""
418418
Defines a standard Successive Convexification problem and
@@ -605,7 +605,7 @@ def plot_animation(X, U):
605605

606606
def main():
607607
print("start!!")
608-
m = Model_6DoF()
608+
m = Rocket_Model_6DoF()
609609

610610
# state and input list
611611
X = np.empty(shape=[m.n_x, K])

0 commit comments

Comments
 (0)