rocket_landing.py
import scvxgen as scvx

#======================#
#*———* Parameters *———*#
#======================#

p = scvx.Parameters('gI l J1:4 g0 Isp m_dry v_max w_max del_max theta_max T_min T_max gamma beta c_ax c_ayz S_a rho l_p')

#==================#
#*———* States *———*#
#==================#

x = scvx.States('m r1:4 v1:4 q1:5 w1:4')

#====================#
#*———* Controls *———*#
#====================#

u = scvx.Controls('T1:4')

#====================#
#*———* Dynamics *———*#
#====================#

gI, l, J1, J2, J3, g0, Isp, m_dry, v_max, w_max, del_max, theta_max, T_min, T_max, gamma, beta, c_ax, c_ayz, S_a, rho, l_p = p
m, r1, r2, r3, v1, v2, v3, q1, q2, q3, q4, w1, w2, w3 = x
T1, T2, T3 = u

v = scvx.Vector(v1, v2, v3)      # velocity
q = scvx.Vector(q1, q2, q3, q4)  # quaternion
w = scvx.Vector(w1, w2, w3)      # angular velocity
T = scvx.Vector(T1, T2, T3)      # thrust vector

CBI = scvx.sympy.Matrix([[q4**2 + q1**2 - q2**2 - q3**2,  2*(q1*q2-q4*q3), 2*(q4*q2 + q1*q3)],
                         [2*(q4*q3 + q1*q2), q4**2 - q1**2 + q2**2 - q3**2, 2*(q2*q3 - q4*q1)],
                         [2*(q1*q3 - q4*q2), 2*(q4*q1 + q2*q3), q4**2 - q1**2 - q2**2 + q3**2]]).T  # direction cosine matrix (DCM)

r_arm = scvx.Vector(-l, 0, 0)                   # lever arm
J = scvx.sympy.diag(J1, J2, J3)                 # inertia tensor
CA = scvx.sympy.diag(c_ax, c_ayz, c_ayz)        # aerodynamic coefficient matrix
A = -0.5 * rho * v.norm() * S_a * CA @ CBI @ v  # aerodynamic force
r_cp = scvx.Vector(l_p, 0, 0)                   # center of pressure

q1_dot =  0.5 * (w1*q4 - w2*q3 + w3*q2)
q2_dot =  0.5 * (w1*q3 - w3*q1 + w2*q4)
q3_dot =  0.5 * (w2*q1 - w1*q2 + w3*q4)
q4_dot = -0.5 * (w1*q1 + w2*q2 + w3*q3)

m_dot = -(1 / (Isp * g0)) * T.norm() - beta                          # mass depletion
r_dot = v                                                            # translation kinematics
v_dot = (1 / m) * CBI.T @ (T + A) + scvx.Vector(-gI, 0, 0)           # translation dynamics
q_dot = scvx.Vector(q1_dot, q2_dot, q3_dot, q4_dot)                  # quaternion kinematics
w_dot = J.inv() @ (r_arm.cross(T) + r_cp.cross(A) - w.cross(J @ w))  # attitude dynamics

f = scvx.Dynamics(m_dot, r_dot, v_dot, q_dot, w_dot)

#==========================================#
#*———* Inequality constraints (`≤ 0`) *———*#
#==========================================#

g = scvx.InequalityConstraints(
    m_dry - m,                                                                        # minimum mass
    scvx.Vector(r2, r3).norm() - scvx.sympy.tan(gamma * scvx.sympy.pi / 180.0) * r1,  # maximum glideslope angle
    v.norm()**2 - v_max**2,                                                           # maximum speed
    scvx.sympy.cos(theta_max * scvx.sympy.pi / 180.0) - 1 + 2 * (q2**2 + q3**2),      # maximum tilt angle
    w.norm()**2 - w_max**2,                                                           # maximum angular speed
    T.norm() - T1 / scvx.sympy.cos(del_max * scvx.sympy.pi / 180.0),                  # maximum thrust pointing angle
    T.norm()**2 - T_max**2,                                                           # maximum thrust magnitude
    T_min**2 - T.norm()**2                                                            # minimum thrust magnitude
)

#============================#
#*———* Numerical values *———*#
#============================#

#*———* Final time guess *———*#

t = scvx.Time(guess=10.0, t_max_factor=2.0, t_min_factor=0.2)

#*———* System/constraint parameters *———*#

p.set_value(1.0, 0.25, 0.168 * 2e-2, 0.168, 0.168, 1.0, 30.0, 1.0, 3.0, 0.3752, 20.0, 75.0, 1.5, 6.5, 75.0, 0.01, 0.5, 1.0, 0.5, 1.0, 0.05)

#*———* State parameters *———*#

x.set_initial(2.0, 7.5, 4.5, 2.0, -0.5, -2.8, 0.0, scvx.Free(guess=0.0), scvx.Free(guess=0.0), scvx.Free(guess=0.0), scvx.Free(), 0.0, 0.0, 0.0)
x.set_final(scvx.Maximize(), 0.0, 0.0, 0.0, -0.1, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0)
x.set_maximum(2.0, 10.0, 10.0, 10.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, p.w_max.value, p.w_max.value, p.w_max.value)
x.set_minimum(2.0, -10.0, -10.0, -10.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -p.w_max.value, -p.w_max.value, -p.w_max.value)

#*———* Control input parameters *———*#

u.set_initial(p.gI.value * x.m.initial, 0.0, 0.0)
u.set_final(p.gI.value * p.m_dry.value, 0.0, 0.0)
u.set_minimum(-p.T_max.value, -p.T_max.value, -p.T_max.value)
u.set_maximum(p.T_max.value, p.T_max.value, p.T_max.value)

#====================#
#*———* Settings *———*#
#====================#

settings = scvx.Settings()

#===============================#
#*———* Instantiate problem *———*#
#===============================#

problem = scvx.Problem(time=t, parameters=p, states=x, controls=u, dynamics=f, inequality_constraints=g, settings=settings)

#=========================#
#*———* Generate code *———*#
#=========================#

problem.generate_code()

#=========================#
#*———* Solve problem *———*#
#=========================#

output = problem.solve()