'''Starship 2D trajectory optimization with OpenMDAO'''
import numpy as np
import time
import openmdao.api as om
from modopt import SLSQP, OpenMDAOProblem
g = 9.80665 # gravity (m/s^2)
m = 100000. # mass (kg)
L = 50. # length (m)
W = 10. # width (m)
I = (1/12) * m * L**2 # moment of inertia (kg*m^2)
duration = 16. # duration (s)
min_gimbal = -20 * np.pi / 180 # (rad)
max_gimbal = 20 * np.pi / 180 # (rad)
min_thrust = 884 * 1000. # (N)
max_thrust = 2210 * 1000. # (N)
x_init = np.array([0, 0, 1000, -80, np.pi/2, 0])
x_final = np.array([0., 0., 0., 0., 0., 0.])
class CostComp(om.ExplicitComponent): # Total of ~170 statements excluding comments.
def initialize(self):
self.options.declare('nt', types=int)
def setup(self):
nt = self.options['nt']
self.add_input('x', shape=(6,nt))
self.add_input('u', shape=(2,nt))
self.add_output('cost')
self.declare_partials('cost', 'x', rows=np.zeros(nt), cols=np.arange(5*nt, 6*nt))
self.declare_partials('cost', 'u')
def compute(self, inputs, outputs):
x = inputs['x']
u = inputs['u']
outputs['cost'] = np.sum(u[0,:]**2) + np.sum(u[1,:]**2) + 2 * np.sum(x[5,:]**2)
def compute_partials(self, inputs, partials):
x = inputs['x']
u = inputs['u']
partials['cost', 'x'] = 4 * x[5,:]
partials['cost', 'u'] = 2 * u
class DynamicsComp(om.ExplicitComponent):
def initialize(self):
self.options.declare('nt', types=int)
self.options.declare('max_thrust', types=float)
self.options.declare('g', types=float)
self.options.declare('m', types=float)
self.options.declare('L', types=float)
self.options.declare('I', types=float)
def setup(self):
nt = self.options['nt']
self.add_input('x', shape=(6,nt))
self.add_input('u', shape=(2,nt))
self.add_output('f', shape=(6,nt-1))
rows = np.arange(5*(nt-1))
cols = np.concatenate([np.arange(nt, 2*nt-1), np.arange(4*nt,5*nt-1), np.arange(3*nt,4*nt-1), np.arange(4*nt,5*nt-1), np.arange(5*nt,6*nt-1)])
val = np.ones(5*(nt-1))
self.declare_partials('f', 'x', rows=rows, cols=cols, val=val)
rows = np.concatenate([np.arange(nt-1, 2*(nt-1)), np.arange(3*(nt-1), 4*(nt-1)), np.arange(5*(nt-1), 6*(nt-1))])
rows = np.concatenate([rows, rows])
cols1 = np.concatenate([np.arange(0, nt-1), np.arange(0,nt-1), np.arange(0,nt-1)])
cols2 = np.concatenate([np.arange(nt, 2*nt-1), np.arange(nt, 2*nt-1), np.arange(nt, 2*nt-1)])
cols = np.concatenate([cols1, cols2])
self.declare_partials('f', 'u', rows=rows, cols=cols)
def compute(self, inputs, outputs):
nt = self.options['nt']
max_thrust = self.options['max_thrust']
g = self.options['g']
m = self.options['m']
L = self.options['L']
I = self.options['I']
x = inputs['x']
u = inputs['u']
f = np.zeros((6, nt-1))
thrust = max_thrust * u[0, :-1] # thrust magnitude (N)
theta = x[4, :-1] # rocket angle (rad)
beta = u[1, :-1] # thrust angle / gimbal (rad)
# Dynamics: xdot = f(x,u) = [xdot, xdotdot, ydot, ydotdot, thetadot, thetadotdot]
f[0, :] = x[1, :-1]
f[1, :] = -thrust * np.sin(beta + theta) / m
f[2, :] = x[3, :-1]
f[3, :] = thrust * np.cos(beta + theta) / m - g
f[4, :] = x[5, :-1]
f[5, :] = -0.5 * L * thrust * np.sin(beta) / I
outputs['f'] = f
def compute_partials(self, inputs, partials):
nt = self.options['nt']
max_thrust = self.options['max_thrust']
m = self.options['m']
L = self.options['L']
I = self.options['I']
x = inputs['x']
u = inputs['u']
thrust = max_thrust * u[0, :-1] # thrust magnitude (N)
theta = x[4, :-1] # rocket angle (rad)
beta = u[1, :-1] # thrust angle / gimbal (rad)
partials['f', 'x'][nt-1 : 2*(nt-1)] = -thrust * np.cos(beta + theta) / m
partials['f', 'x'][3*(nt-1) : 4*(nt-1)] = -thrust * np.sin(beta + theta) / m
partials['f', 'u'][0: (nt-1)] = -max_thrust * np.sin(beta + theta) / m
partials['f', 'u'][1*(nt-1) : 2*(nt-1)] = max_thrust * np.cos(beta + theta) / m
partials['f', 'u'][2*(nt-1) : 3*(nt-1)] = -0.5 * L * max_thrust * np.sin(beta) / I
partials['f', 'u'][3*(nt-1) : 4*(nt-1)] = -thrust * np.cos(beta + theta) / m
partials['f', 'u'][4*(nt-1) : 5*(nt-1)] = -thrust * np.sin(beta + theta) / m
partials['f', 'u'][5*(nt-1) : 6*(nt-1)] = -0.5 * L * thrust * np.cos(beta) / I
class ConstraintsComp(om.ExplicitComponent):
def initialize(self):
self.options.declare('nt', types=int)
self.options.declare('duration', types=float)
def setup(self):
nt = self.options['nt']
dt = self.options['duration'] / nt
self.add_input('x', shape=(6,nt))
self.add_input('f', shape=(6,nt-1))
self.add_output('constraints', shape=(6,nt-1))
rows = np.arange(6*(nt-1))
rows = np.concatenate([rows, rows])
x_cols = np.arange(6*nt).reshape((6, nt))
cols = np.concatenate([x_cols[:, 1:].flatten(), x_cols[:, :-1].flatten()])
val = np.concatenate([np.ones(6*(nt-1)), -np.ones(6*(nt-1))])
self.declare_partials('constraints', 'x', rows=rows, cols=cols, val=val)
rows = np.arange(6*(nt-1))
cols = np.arange(6*(nt-1))
val = -dt * np.ones(6*(nt-1))
self.declare_partials('constraints', 'f', rows=rows, cols=cols, val=val)
def compute(self, inputs, outputs):
nt = self.options['nt']
duration = self.options['duration']
dt = duration / nt
x = inputs['x']
f = inputs['f']
outputs['constraints'] = x[:, 1:] - x[:, :-1] - f * dt
class StarshipGroup(om.Group):
def initialize(self):
self.options.declare('g', default=9.80665, types=float)
self.options.declare('m', default=100000., types=float)
self.options.declare('L', default=50., types=float)
self.options.declare('W', default=10., types=float)
self.options.declare('min_gimbal', default=-20 * np.pi / 180, types=float)
self.options.declare('max_gimbal', default= 20 * np.pi / 180, types=float)
self.options.declare('min_thrust', default= 884 * 1000., types=float)
self.options.declare('max_thrust', default=2210 * 1000., types=float)
self.options.declare('duration', default=16., types=float)
self.options.declare('nt', default=20, types=int)
self.options.declare('x_init', types=np.ndarray) # Initial state
self.options.declare('x_final', types=np.ndarray) # Final state
def setup(self):
g = self.options['g']
m = self.options['m']
L = self.options['L']
I = (1/12) * m * L**2
min_gimbal = self.options['min_gimbal']
max_gimbal = self.options['max_gimbal']
min_thrust = self.options['min_thrust']
max_thrust = self.options['max_thrust']
nt = self.options['nt']
xl = np.full((6, nt), -np.inf)
xu = np.full((6, nt), np.inf)
xl[:, 0] = self.options['x_init']
xu[:, 0] = self.options['x_init']
xl[:, -1] = self.options['x_final']
xu[:, -1] = self.options['x_final']
ul = np.full((2, nt), -np.inf)
uu = np.full((2, nt), np.inf)
ul[0, :] = min_thrust / max_thrust
uu[0, :] = 1.0
ul[1, :] = min_gimbal
uu[1, :] = max_gimbal
comp = DynamicsComp(nt=nt, max_thrust=max_thrust, g=g, m=m, L=L, I=I)
self.add_subsystem('dynamics_comp', comp, promotes_inputs=['x', 'u'])
comp = ConstraintsComp(nt=nt, duration=duration)
self.add_subsystem('constraints_comp', comp, promotes_inputs=['x'])
comp = CostComp(nt=nt,)
self.add_subsystem('cost_comp', comp, promotes_inputs=['x', 'u'])
self.connect('dynamics_comp.f', 'constraints_comp.f')
self.add_design_var('x', lower=xl, upper=xu)
self.add_design_var('u', lower=ul, upper=uu)
self.add_objective('cost_comp.cost')
self.add_constraint('constraints_comp.constraints', equals=0.)
# Method 1 - Using OpenMDAO-modOpt interface
def get_problem(nt):
om_prob = om.Problem(model=StarshipGroup(g=g, m=m, L=L, W=W, nt=nt, duration=duration,
min_gimbal=min_gimbal, max_gimbal=max_gimbal,
min_thrust=min_thrust, max_thrust=max_thrust,
x_init=x_init, x_final=x_final))
om_prob.setup()
prob = OpenMDAOProblem(problem_name=f'starship_{nt}_om', om_problem=om_prob)
prob.x0 = np.ones(nt*8)
return prob
if __name__ == '__main__':
# # Test to see if the problem is correctly defined
# nt = 4
# om_prob = om.Problem(model=StarshipGroup(g=g, m=m, L=L, W=W, nt=nt, duration=duration,
# min_gimbal=min_gimbal, max_gimbal=max_gimbal,
# min_thrust=min_thrust, max_thrust=max_thrust,
# x_init=x_init, x_final=x_final))
# om_prob.setup()
# om_prob.check_partials(compact_print=True)
# # om_prob.set_val('x', np.arange(6*nt).reshape((6, nt)))
# # om_prob.set_val('u', np.arange(6*nt, 8*nt).reshape((2, nt)))
# # print(prob.x0)
# prob = OpenMDAOProblem(problem_name=f'starship_{nt}_om', om_problem=om_prob)
# print(prob._compute_objective(np.arange(nt*8))) # 9800.0
# print(prob._compute_constraints(np.arange(nt*8))) # [ -15. -19. -23. 38.5564043 1993.9522483 -1764.75651359
# # -47. -51. -55. -2081.04096363 995.3398582 1511.53434984
# # -79. -83. -87. 69.97044646 -174.99570609 -271.50702618]
# print(np.linalg.norm(prob._compute_objective_gradient(np.arange(nt*8)))) # 232.44784361228218
# print(np.linalg.norm(prob._compute_constraint_jacobian(np.arange(nt*8)))) # 5427.787420199876
# exit()
# SLSQP
print('\tSLSQP \n\t-----')
nt =20
optimizer = SLSQP(get_problem(nt), solver_options={'maxiter': 1000, 'ftol': 1e-9})
start_time = time.time()
optimizer.solve()
opt_time = time.time() - start_time
success = optimizer.results['success']
print('\tTime:', opt_time)
print('\tSuccess:', success)
print('\tOptimized vars:', optimizer.results['x'])
print('\tOptimized obj:', optimizer.results['fun'])
optimizer.print_results()
v = optimizer.results['x']
x = v[:nt*6].reshape((6, nt))
u = v[nt*6:].reshape((2, nt))
import matplotlib.pyplot as plt
plt.figure()
plt.plot(x[0], label='x')
plt.plot(x[1], label='xdot')
plt.plot(x[2], label='y')
plt.plot(x[3], label='ydot')
plt.plot(x[4], label='theta')
plt.plot(x[5], label='thetadot')
plt.legend()
plt.show()
plt.figure()
plt.plot(u[0], label='thrust (percent)')
plt.plot(u[1], label='gimbal (rad)')
plt.legend()
plt.show()
assert np.allclose(optimizer.results['x'],
[0.00000000e+00, 2.03759954e-16, -6.95543787e+00, -1.92272807e+01, -3.70171262e+01,
-5.73251315e+01, -7.70321155e+01, -9.20093998e+01, -1.00655176e+02, -1.02640831e+02,
-9.83525602e+01, -8.89790015e+01, -7.59615262e+01, -6.07761859e+01, -4.48738711e+01,
-2.96598578e+01, -1.64810765e+01, -6.58179904e+00, -9.71561432e-01, 0.00000000e+00,
0.00000000e+00, -8.69429734e+00, -1.53398036e+01, -2.22373069e+01, -2.53850066e+01,
-2.46337300e+01, -1.87216053e+01, -1.08072201e+01, -2.48206896e+00, 5.36033849e+00,
1.17169484e+01, 1.62718441e+01, 1.89816753e+01, 1.98778935e+01, 1.90175166e+01,
1.64734767e+01, 1.23740968e+01, 7.01279701e+00, 1.21445179e+00, 0.00000000e+00,
1.00000000e+03, 9.36000000e+02, 8.63192172e+02, 7.82173074e+02, 6.96126754e+02,
6.08870469e+02, 5.26122516e+02, 4.50428073e+02, 3.81105091e+02, 3.17983654e+02,
2.61262341e+02, 2.10606910e+02, 1.65634822e+02, 1.26062551e+02, 9.17065742e+01,
6.24965119e+01, 3.84847192e+01, 1.98487678e+01, 6.88900882e+00, 0.00000000e+00,
-8.00000000e+01, -9.10097854e+01, -1.01273872e+02, -1.07557900e+02, -1.09070356e+02,
-1.03434942e+02, -9.46180535e+01, -8.66537274e+01, -7.89017958e+01, -7.09016419e+01,
-6.33192886e+01, -5.62151096e+01, -4.94653393e+01, -4.29449709e+01, -3.65125779e+01,
-3.00147408e+01, -2.32949392e+01, -1.61996988e+01, -8.61126103e+00, 0.00000000e+00,
1.57079633e+00, 1.57079633e+00, 1.26700764e+00, 7.31017383e-01, 1.39970913e-01,
-2.70159195e-01, -4.28319588e-01, -4.66403998e-01, -4.43718800e-01, -3.80500367e-01,
-2.90355810e-01, -1.83663691e-01, -6.88789737e-02, 4.57820842e-02, 1.51359763e-01,
2.37157574e-01, 2.89138727e-01, 2.87100964e-01, 2.00698707e-01, 0.00000000e+00,
0.00000000e+00, -3.79735853e-01, -6.69987827e-01, -7.38808088e-01, -5.12662636e-01,
-1.97700490e-01, -4.76055127e-02, 2.83564967e-02, 7.90230420e-02, 1.12680696e-01,
1.33365149e-01, 1.43480896e-01, 1.43326322e-01, 1.31972098e-01, 1.07247264e-01,
6.49764412e-02, -2.54720412e-03, -1.08002821e-01, -2.50873384e-01, 0.00000000e+00,
5.23318890e-01, 4.00000000e-01, 4.00000000e-01, 4.00000000e-01, 7.63668116e-01,
1.00000000e+00, 1.00000000e+00, 1.00000000e+00, 1.00000000e+00, 9.43773192e-01,
8.83937030e-01, 8.39622310e-01, 8.14118607e-01, 8.09027871e-01, 8.23926904e-01,
8.55826765e-01, 8.97815057e-01, 9.32523511e-01, 9.33333248e-01, 4.00000000e-01,
3.49065850e-01, 3.49065850e-01, 8.11839128e-02, -2.69738955e-01, -1.95642976e-01,
-7.08052720e-02, -3.58117687e-02, -2.38835608e-02, -1.58649445e-02, -1.03304826e-02,
-5.39405342e-03, 8.67737450e-05, 6.57369288e-03, 1.44052702e-02, 2.41841512e-02,
3.71968927e-02, 5.53912837e-02, 7.22765885e-02, -1.27034906e-01, 2.61331412e-10],
rtol=0, atol=1e-6)
# # Method 2 - Using OpenMDAO directly and its ScipyOptimizeDriver
# prob = om.Problem(model=StarshipGroup(g=g, m=m, L=L, W=W, nt=nt, duration=duration,
# min_gimbal=min_gimbal, max_gimbal=max_gimbal,
# min_thrust=min_thrust, max_thrust=max_thrust,
# x_init=x_init, x_final=x_final))
# prob.setup()
# prob.driver = om.ScipyOptimizeDriver()
# prob.driver.options['optimizer'] = 'SLSQP'
# prob.driver.options['tol'] = 1e-9
# prob.driver.options['disp'] = True
# prob.setup()
# start = time.time()
# prob.run_driver()
# print('Time:', time.time() - start)
# print('\tOptimized states:', prob['x'])
# print('\tOptimized controls:', prob['u'])
# print('\tOptimized obj:', prob['cost_comp.cost'])
# solution = np.concatenate((prob['x'].flatten(), prob['u'].flatten()))
# assert np.allclose(optimizer.results['x'],
# [0.00000000e+00, 2.03759954e-16, -6.95543787e+00, -1.92272807e+01, -3.70171262e+01,
# -5.73251315e+01, -7.70321155e+01, -9.20093998e+01, -1.00655176e+02, -1.02640831e+02,
# -9.83525602e+01, -8.89790015e+01, -7.59615262e+01, -6.07761859e+01, -4.48738711e+01,
# -2.96598578e+01, -1.64810765e+01, -6.58179904e+00, -9.71561432e-01, 0.00000000e+00,
# 0.00000000e+00, -8.69429734e+00, -1.53398036e+01, -2.22373069e+01, -2.53850066e+01,
# -2.46337300e+01, -1.87216053e+01, -1.08072201e+01, -2.48206896e+00, 5.36033849e+00,
# 1.17169484e+01, 1.62718441e+01, 1.89816753e+01, 1.98778935e+01, 1.90175166e+01,
# 1.64734767e+01, 1.23740968e+01, 7.01279701e+00, 1.21445179e+00, 0.00000000e+00,
# 1.00000000e+03, 9.36000000e+02, 8.63192172e+02, 7.82173074e+02, 6.96126754e+02,
# 6.08870469e+02, 5.26122516e+02, 4.50428073e+02, 3.81105091e+02, 3.17983654e+02,
# 2.61262341e+02, 2.10606910e+02, 1.65634822e+02, 1.26062551e+02, 9.17065742e+01,
# 6.24965119e+01, 3.84847192e+01, 1.98487678e+01, 6.88900882e+00, 0.00000000e+00,
# -8.00000000e+01, -9.10097854e+01, -1.01273872e+02, -1.07557900e+02, -1.09070356e+02,
# -1.03434942e+02, -9.46180535e+01, -8.66537274e+01, -7.89017958e+01, -7.09016419e+01,
# -6.33192886e+01, -5.62151096e+01, -4.94653393e+01, -4.29449709e+01, -3.65125779e+01,
# -3.00147408e+01, -2.32949392e+01, -1.61996988e+01, -8.61126103e+00, 0.00000000e+00,
# 1.57079633e+00, 1.57079633e+00, 1.26700764e+00, 7.31017383e-01, 1.39970913e-01,
# -2.70159195e-01, -4.28319588e-01, -4.66403998e-01, -4.43718800e-01, -3.80500367e-01,
# -2.90355810e-01, -1.83663691e-01, -6.88789737e-02, 4.57820842e-02, 1.51359763e-01,
# 2.37157574e-01, 2.89138727e-01, 2.87100964e-01, 2.00698707e-01, 0.00000000e+00,
# 0.00000000e+00, -3.79735853e-01, -6.69987827e-01, -7.38808088e-01, -5.12662636e-01,
# -1.97700490e-01, -4.76055127e-02, 2.83564967e-02, 7.90230420e-02, 1.12680696e-01,
# 1.33365149e-01, 1.43480896e-01, 1.43326322e-01, 1.31972098e-01, 1.07247264e-01,
# 6.49764412e-02, -2.54720412e-03, -1.08002821e-01, -2.50873384e-01, 0.00000000e+00,
# 5.23318890e-01, 4.00000000e-01, 4.00000000e-01, 4.00000000e-01, 7.63668116e-01,
# 1.00000000e+00, 1.00000000e+00, 1.00000000e+00, 1.00000000e+00, 9.43773192e-01,
# 8.83937030e-01, 8.39622310e-01, 8.14118607e-01, 8.09027871e-01, 8.23926904e-01,
# 8.55826765e-01, 8.97815057e-01, 9.32523511e-01, 9.33333248e-01, 4.00000000e-01,
# 3.49065850e-01, 3.49065850e-01, 8.11839128e-02, -2.69738955e-01, -1.95642976e-01,
# -7.08052720e-02, -3.58117687e-02, -2.38835608e-02, -1.58649445e-02, -1.03304826e-02,
# -5.39405342e-03, 8.67737450e-05, 6.57369288e-03, 1.44052702e-02, 2.41841512e-02,
# 3.71968927e-02, 5.53912837e-02, 7.22765885e-02, -1.27034906e-01, 2.61331412e-10],
# rtol=0, atol=1e-6)