rocket_powered_landing.py

"""

A rocket powered landing with successive convexification

author: Sven Niederberger
        Atsushi Sakai

Reference:
- Python implementation of 'Successive Convexification for 6-DoF Mars Rocket Powered Landing with Free-Final-Time' paper
by Michael Szmuk and Behcet Acıkmese.

- EmbersArc/SuccessiveConvexificationFreeFinalTime: Implementation of "Successive Convexification for 6-DoF Mars Rocket Powered Landing with Free-Final-Time" https://github.com/EmbersArc/SuccessiveConvexificationFreeFinalTime

"""
import warnings
from time import time
import numpy as np
from scipy.integrate import odeint
import cvxpy
import matplotlib.pyplot as plt

# Trajectory points
K = 50

# Max solver iterations
iterations = 30

# Weight constants
W_SIGMA = 1  # flight time
W_DELTA = 1e-3  # difference in state/input
W_DELTA_SIGMA = 1e-1  # difference in flight time
W_NU = 1e5  # virtual control

print(cvxpy.installed_solvers())
solver = 'ECOS'
verbose_solver = False

show_animation = True


class Rocket_Model_6DoF:
    """
    A 6 degree of freedom rocket landing problem.
    """

    def __init__(self, rng):
        """
        A large r_scale for a small scale problem will
        ead to numerical problems as parameters become excessively small
        and (it seems) precision is lost in the dynamics.
        """
        self.n_x = 14
        self.n_u = 3

        # Mass
        self.m_wet = 3.0  # 30000 kg
        self.m_dry = 2.2  # 22000 kg

        # Flight time guess
        self.t_f_guess = 10.0  # 10 s

        # State constraints
        self.r_I_final = np.array((0., 0., 0.))
        self.v_I_final = np.array((-1e-1, 0., 0.))
        self.q_B_I_final = self.euler_to_quat((0, 0, 0))
        self.w_B_final = np.deg2rad(np.array((0., 0., 0.)))

        self.w_B_max = np.deg2rad(60)

        # Angles
        max_gimbal = 20
        max_angle = 90
        glidelslope_angle = 20

        self.tan_delta_max = np.tan(np.deg2rad(max_gimbal))
        self.cos_theta_max = np.cos(np.deg2rad(max_angle))
        self.tan_gamma_gs = np.tan(np.deg2rad(glidelslope_angle))

        # Thrust limits
        self.T_max = 5.0
        self.T_min = 0.3

        # Angular moment of inertia
        self.J_B = 1e-2 * np.diag([1., 1., 1.])

        # Gravity
        self.g_I = np.array((-1, 0., 0.))

        # Fuel consumption
        self.alpha_m = 0.01

        # Vector from thrust point to CoM
        self.r_T_B = np.array([-1e-2, 0., 0.])

        self.set_random_initial_state(rng)

        self.x_init = np.concatenate(
            ((self.m_wet,), self.r_I_init, self.v_I_init, self.q_B_I_init, self.w_B_init))
        self.x_final = np.concatenate(
            ((self.m_dry,), self.r_I_final, self.v_I_final, self.q_B_I_final, self.w_B_final))

        self.r_scale = np.linalg.norm(self.r_I_init)
        self.m_scale = self.m_wet

    def set_random_initial_state(self, rng):
        if rng is None:
            rng = np.random.default_rng()

        self.r_I_init = np.array((0., 0., 0.))
        self.r_I_init[0] = rng.uniform(3, 4)
        self.r_I_init[1:3] = rng.uniform(-2, 2, size=2)

        self.v_I_init = np.array((0., 0., 0.))
        self.v_I_init[0] = rng.uniform(-1, -0.5)
        self.v_I_init[1:3] = rng.uniform(-0.5, -0.2,
                                         size=2) * self.r_I_init[1:3]

        self.q_B_I_init = self.euler_to_quat((0,
                                              rng.uniform(-30, 30),
                                              rng.uniform(-30, 30)))
        self.w_B_init = np.deg2rad((0,
                                    rng.uniform(-20, 20),
                                    rng.uniform(-20, 20)))

    def f_func(self, x, u):
        m, _, _, _, vx, vy, vz, q0, q1, q2, q3, wx, wy, wz = x[0], x[1], x[
            2], x[3], x[4], x[5], x[6], x[7], x[8], x[9], x[10], x[11], x[12], x[13]
        ux, uy, uz = u[0], u[1], u[2]

        return np.array([
            [-0.01 * np.sqrt(ux**2 + uy**2 + uz**2)],
            [vx],
            [vy],
            [vz],
            [(-1.0 * m - ux * (2 * q2**2 + 2 * q3**2 - 1) - 2 * uy
              * (q0 * q3 - q1 * q2) + 2 * uz * (q0 * q2 + q1 * q3)) / m],
            [(2 * ux * (q0 * q3 + q1 * q2) - uy * (2 * q1**2
                                                   + 2 * q3**2 - 1) - 2 * uz * (q0 * q1 - q2 * q3)) / m],
            [(-2 * ux * (q0 * q2 - q1 * q3) + 2 * uy
              * (q0 * q1 + q2 * q3) - uz * (2 * q1**2 + 2 * q2**2 - 1)) / m],
            [-0.5 * q1 * wx - 0.5 * q2 * wy - 0.5 * q3 * wz],
            [0.5 * q0 * wx + 0.5 * q2 * wz - 0.5 * q3 * wy],
            [0.5 * q0 * wy - 0.5 * q1 * wz + 0.5 * q3 * wx],
            [0.5 * q0 * wz + 0.5 * q1 * wy - 0.5 * q2 * wx],
            [0],
            [1.0 * uz],
            [-1.0 * uy]
        ])

    def A_func(self, x, u):
        m, _, _, _, _, _, _, q0, q1, q2, q3, wx, wy, wz = x[0], x[1], x[
            2], x[3], x[4], x[5], x[6], x[7], x[8], x[9], x[10], x[11], x[12], x[13]
        ux, uy, uz = u[0], u[1], u[2]

        return np.array([
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
            [(ux * (2 * q2**2 + 2 * q3**2 - 1) + 2 * uy * (q0 * q3 - q1 * q2) - 2 * uz * (q0 * q2 + q1 * q3)) / m**2, 0, 0, 0, 0, 0, 0, 2 * (q2 * uz
                                                                                                                                             - q3 * uy) / m, 2 * (q2 * uy + q3 * uz) / m, 2 * (q0 * uz + q1 * uy - 2 * q2 * ux) / m, 2 * (-q0 * uy + q1 * uz - 2 * q3 * ux) / m, 0, 0, 0],
            [(-2 * ux * (q0 * q3 + q1 * q2) + uy * (2 * q1**2 + 2 * q3**2 - 1) + 2 * uz * (q0 * q1 - q2 * q3)) / m**2, 0, 0, 0, 0, 0, 0, 2 * (-q1 * uz
                                                                                                                                              + q3 * ux) / m, 2 * (-q0 * uz - 2 * q1 * uy + q2 * ux) / m, 2 * (q1 * ux + q3 * uz) / m, 2 * (q0 * ux + q2 * uz - 2 * q3 * uy) / m, 0, 0, 0],
            [(2 * ux * (q0 * q2 - q1 * q3) - 2 * uy * (q0 * q1 + q2 * q3) + uz * (2 * q1**2 + 2 * q2**2 - 1)) / m**2, 0, 0, 0, 0, 0, 0, 2 * (q1 * uy
                                                                                                                                             - q2 * ux) / m, 2 * (q0 * uy - 2 * q1 * uz + q3 * ux) / m, 2 * (-q0 * ux - 2 * q2 * uz + q3 * uy) / m, 2 * (q1 * ux + q2 * uy) / m, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, -0.5 * wx, -0.5 * wy,
             - 0.5 * wz, -0.5 * q1, -0.5 * q2, -0.5 * q3],
            [0, 0, 0, 0, 0, 0, 0, 0.5 * wx, 0, 0.5 * wz,
             - 0.5 * wy, 0.5 * q0, -0.5 * q3, 0.5 * q2],
            [0, 0, 0, 0, 0, 0, 0, 0.5 * wy, -0.5 * wz, 0,
             0.5 * wx, 0.5 * q3, 0.5 * q0, -0.5 * q1],
            [0, 0, 0, 0, 0, 0, 0, 0.5 * wz, 0.5 * wy,
             - 0.5 * wx, 0, -0.5 * q2, 0.5 * q1, 0.5 * q0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])

    def B_func(self, x, u):
        m, _, _, _, _, _, _, q0, q1, q2, q3, _, _, _ = x[0], x[1], x[
            2], x[3], x[4], x[5], x[6], x[7], x[8], x[9], x[10], x[11], x[12], x[13]
        ux, uy, uz = u[0], u[1], u[2]

        return np.array([
            [-0.01 * ux / np.sqrt(ux**2 + uy**2 + uz**2),
             -0.01 * uy / np.sqrt(ux ** 2 + uy**2 + uz**2),
             -0.01 * uz / np.sqrt(ux**2 + uy**2 + uz**2)],
            [0, 0, 0],
            [0, 0, 0],
            [0, 0, 0],
            [(-2 * q2**2 - 2 * q3**2 + 1) / m, 2
             * (-q0 * q3 + q1 * q2) / m, 2 * (q0 * q2 + q1 * q3) / m],
            [2 * (q0 * q3 + q1 * q2) / m, (-2 * q1**2 - 2
                                           * q3**2 + 1) / m, 2 * (-q0 * q1 + q2 * q3) / m],
            [2 * (-q0 * q2 + q1 * q3) / m, 2 * (q0 * q1 + q2 * q3)
             / m, (-2 * q1**2 - 2 * q2**2 + 1) / m],
            [0, 0, 0],
            [0, 0, 0],
            [0, 0, 0],
            [0, 0, 0],
            [0, 0, 0],
            [0, 0, 1.0],
            [0, -1.0, 0]
        ])

    def euler_to_quat(self, a):
        a = np.deg2rad(a)

        cy = np.cos(a[1] * 0.5)
        sy = np.sin(a[1] * 0.5)
        cr = np.cos(a[0] * 0.5)
        sr = np.sin(a[0] * 0.5)
        cp = np.cos(a[2] * 0.5)
        sp = np.sin(a[2] * 0.5)

        q = np.zeros(4)

        q[0] = cy * cr * cp + sy * sr * sp
        q[1] = cy * sr * cp - sy * cr * sp
        q[3] = cy * cr * sp + sy * sr * cp
        q[2] = sy * cr * cp - cy * sr * sp

        return q

    def skew(self, v):
        return np.array([
            [0, -v[2], v[1]],
            [v[2], 0, -v[0]],
            [-v[1], v[0], 0]
        ])

    def dir_cosine(self, q):
        return np.array([
            [1 - 2 * (q[2] ** 2 + q[3] ** 2), 2 * (q[1] * q[2]
                                                   + q[0] * q[3]), 2 * (q[1] * q[3] - q[0] * q[2])],
            [2 * (q[1] * q[2] - q[0] * q[3]), 1 - 2
             * (q[1] ** 2 + q[3] ** 2), 2 * (q[2] * q[3] + q[0] * q[1])],
            [2 * (q[1] * q[3] + q[0] * q[2]), 2 * (q[2] * q[3]
                                                   - q[0] * q[1]), 1 - 2 * (q[1] ** 2 + q[2] ** 2)]
        ])

    def omega(self, w):
        return np.array([
            [0, -w[0], -w[1], -w[2]],
            [w[0], 0, w[2], -w[1]],
            [w[1], -w[2], 0, w[0]],
            [w[2], w[1], -w[0], 0],
        ])

    def initialize_trajectory(self, X, U):
        """
        Initialize the trajectory with linear approximation.
        """
        K = X.shape[1]

        for k in range(K):
            alpha1 = (K - k) / K
            alpha2 = k / K

            m_k = (alpha1 * self.x_init[0] + alpha2 * self.x_final[0],)
            r_I_k = alpha1 * self.x_init[1:4] + alpha2 * self.x_final[1:4]
            v_I_k = alpha1 * self.x_init[4:7] + alpha2 * self.x_final[4:7]
            q_B_I_k = np.array([1, 0, 0, 0])
            w_B_k = alpha1 * self.x_init[11:14] + alpha2 * self.x_final[11:14]

            X[:, k] = np.concatenate((m_k, r_I_k, v_I_k, q_B_I_k, w_B_k))
            U[:, k] = m_k * -self.g_I

        return X, U

    def get_constraints(self, X_v, U_v, X_last_p, U_last_p):
        """
        Get model specific constraints.

        :param X_v: cvx variable for current states
        :param U_v: cvx variable for current inputs
        :param X_last_p: cvx parameter for last states
        :param U_last_p: cvx parameter for last inputs
        :return: A list of cvx constraints
        """
        # Boundary conditions:
        constraints = [
            X_v[0, 0] == self.x_init[0],
            X_v[1:4, 0] == self.x_init[1:4],
            X_v[4:7, 0] == self.x_init[4:7],
            # X_v[7:11, 0] == self.x_init[7:11],  # initial orientation is free
            X_v[11:14, 0] == self.x_init[11:14],

            # X_[0, -1] final mass is free
            X_v[1:, -1] == self.x_final[1:],
            U_v[1:3, -1] == 0,
        ]

        constraints += [
            # State constraints:
            X_v[0, :] >= self.m_dry,  # minimum mass
            cvxpy.norm(X_v[2: 4, :], axis=0) <= X_v[1, :] / \
            self.tan_gamma_gs,  # glideslope
            cvxpy.norm(X_v[9:11, :], axis=0) <= np.sqrt(
                (1 - self.cos_theta_max) / 2),  # maximum angle
            # maximum angular velocity
            cvxpy.norm(X_v[11: 14, :], axis=0) <= self.w_B_max,

            # Control constraints:
            cvxpy.norm(U_v[1:3, :], axis=0) <= self.tan_delta_max * \
            U_v[0, :],  # gimbal angle constraint
            cvxpy.norm(U_v, axis=0) <= self.T_max,  # upper thrust constraint
        ]

        # linearized lower thrust constraint
        rhs = [U_last_p[:, k] / cvxpy.norm(U_last_p[:, k]) @ U_v[:, k]
               for k in range(X_v.shape[1])]
        constraints += [
            self.T_min <= cvxpy.vstack(rhs)
        ]

        return constraints


class Integrator:
    def __init__(self, m, K):
        self.K = K
        self.m = m
        self.n_x = m.n_x
        self.n_u = m.n_u

        self.A_bar = np.zeros([m.n_x * m.n_x, K - 1])
        self.B_bar = np.zeros([m.n_x * m.n_u, K - 1])
        self.C_bar = np.zeros([m.n_x * m.n_u, K - 1])
        self.S_bar = np.zeros([m.n_x, K - 1])
        self.z_bar = np.zeros([m.n_x, K - 1])

        # vector indices for flat matrices
        x_end = m.n_x
        A_bar_end = m.n_x * (1 + m.n_x)
        B_bar_end = m.n_x * (1 + m.n_x + m.n_u)
        C_bar_end = m.n_x * (1 + m.n_x + m.n_u + m.n_u)
        S_bar_end = m.n_x * (1 + m.n_x + m.n_u + m.n_u + 1)
        z_bar_end = m.n_x * (1 + m.n_x + m.n_u + m.n_u + 2)
        self.x_ind = slice(0, x_end)
        self.A_bar_ind = slice(x_end, A_bar_end)
        self.B_bar_ind = slice(A_bar_end, B_bar_end)
        self.C_bar_ind = slice(B_bar_end, C_bar_end)
        self.S_bar_ind = slice(C_bar_end, S_bar_end)
        self.z_bar_ind = slice(S_bar_end, z_bar_end)

        self.f, self.A, self.B = m.f_func, m.A_func, m.B_func

        # integration initial condition
        self.V0 = np.zeros((m.n_x * (1 + m.n_x + m.n_u + m.n_u + 2),))
        self.V0[self.A_bar_ind] = np.eye(m.n_x).reshape(-1)

        self.dt = 1. / (K - 1)

    def calculate_discretization(self, X, U, sigma):
        """
        Calculate discretization for given states, inputs and total time.

        :param X: Matrix of states for all time points
        :param U: Matrix of inputs for all time points
        :param sigma: Total time
        :return: The discretization matrices
        """
        for k in range(self.K - 1):
            self.V0[self.x_ind] = X[:, k]
            V = np.array(odeint(self._ode_dVdt, self.V0, (0, self.dt),
                                args=(U[:, k], U[:, k + 1], sigma))[1, :])

            # using \Phi_A(\tau_{k+1},\xi) = \Phi_A(\tau_{k+1},\tau_k)\Phi_A(\xi,\tau_k)^{-1}
            # flatten matrices in column-major (Fortran) order for CVXPY
            Phi = V[self.A_bar_ind].reshape((self.n_x, self.n_x))
            self.A_bar[:, k] = Phi.flatten(order='F')
            self.B_bar[:, k] = np.matmul(Phi, V[self.B_bar_ind].reshape(
                (self.n_x, self.n_u))).flatten(order='F')
            self.C_bar[:, k] = np.matmul(Phi, V[self.C_bar_ind].reshape(
                (self.n_x, self.n_u))).flatten(order='F')
            self.S_bar[:, k] = np.matmul(Phi, V[self.S_bar_ind])
            self.z_bar[:, k] = np.matmul(Phi, V[self.z_bar_ind])

        return self.A_bar, self.B_bar, self.C_bar, self.S_bar, self.z_bar

    def _ode_dVdt(self, V, t, u_t0, u_t1, sigma):
        """
        ODE function to compute dVdt.

        :param V: Evaluation state V = [x, Phi_A, B_bar, C_bar, S_bar, z_bar]
        :param t: Evaluation time
        :param u_t0: Input at start of interval
        :param u_t1: Input at end of interval
        :param sigma: Total time
        :return: Derivative at current time and state dVdt
        """
        alpha = (self.dt - t) / self.dt
        beta = t / self.dt
        x = V[self.x_ind]
        u = u_t0 + beta * (u_t1 - u_t0)

        # using \Phi_A(\tau_{k+1},\xi) = \Phi_A(\tau_{k+1},\tau_k)\Phi_A(\xi,\tau_k)^{-1}
        # and pre-multiplying with \Phi_A(\tau_{k+1},\tau_k) after integration
        Phi_A_xi = np.linalg.inv(
            V[self.A_bar_ind].reshape((self.n_x, self.n_x)))

        A_subs = sigma * self.A(x, u)
        B_subs = sigma * self.B(x, u)
        f_subs = self.f(x, u)

        dVdt = np.zeros_like(V)
        dVdt[self.x_ind] = sigma * f_subs.transpose()
        dVdt[self.A_bar_ind] = np.matmul(
            A_subs, V[self.A_bar_ind].reshape((self.n_x, self.n_x))).reshape(-1)
        dVdt[self.B_bar_ind] = np.matmul(Phi_A_xi, B_subs).reshape(-1) * alpha
        dVdt[self.C_bar_ind] = np.matmul(Phi_A_xi, B_subs).reshape(-1) * beta
        dVdt[self.S_bar_ind] = np.matmul(Phi_A_xi, f_subs).transpose()
        z_t = -np.matmul(A_subs, x) - np.matmul(B_subs, u)
        dVdt[self.z_bar_ind] = np.dot(Phi_A_xi, z_t.T).flatten()

        return dVdt


class SCProblem:
    """
    Defines a standard Successive Convexification problem and
            adds the model specific constraints and objectives.

    :param m: The model object
    :param K: Number of discretization points
    """

    def __init__(self, m, K):
        # Variables:
        self.var = dict()
        self.var['X'] = cvxpy.Variable((m.n_x, K))
        self.var['U'] = cvxpy.Variable((m.n_u, K))
        self.var['sigma'] = cvxpy.Variable(nonneg=True)
        self.var['nu'] = cvxpy.Variable((m.n_x, K - 1))
        self.var['delta_norm'] = cvxpy.Variable(nonneg=True)
        self.var['sigma_norm'] = cvxpy.Variable(nonneg=True)

        # Parameters:
        self.par = dict()
        self.par['A_bar'] = cvxpy.Parameter((m.n_x * m.n_x, K - 1))
        self.par['B_bar'] = cvxpy.Parameter((m.n_x * m.n_u, K - 1))
        self.par['C_bar'] = cvxpy.Parameter((m.n_x * m.n_u, K - 1))
        self.par['S_bar'] = cvxpy.Parameter((m.n_x, K - 1))
        self.par['z_bar'] = cvxpy.Parameter((m.n_x, K - 1))

        self.par['X_last'] = cvxpy.Parameter((m.n_x, K))
        self.par['U_last'] = cvxpy.Parameter((m.n_u, K))
        self.par['sigma_last'] = cvxpy.Parameter(nonneg=True)

        self.par['weight_sigma'] = cvxpy.Parameter(nonneg=True)
        self.par['weight_delta'] = cvxpy.Parameter(nonneg=True)
        self.par['weight_delta_sigma'] = cvxpy.Parameter(nonneg=True)
        self.par['weight_nu'] = cvxpy.Parameter(nonneg=True)

        # Constraints:
        constraints = []

        # Model:
        constraints += m.get_constraints(
            self.var['X'], self.var['U'], self.par['X_last'], self.par['U_last'])

        # Dynamics:
        # x_t+1 = A_*x_t+B_*U_t+C_*U_T+1*S_*sigma+zbar+nu
        constraints += [
            self.var['X'][:, k + 1] ==
            cvxpy.reshape(self.par['A_bar'][:, k], (m.n_x, m.n_x), order='F') @
            self.var['X'][:, k] +
            cvxpy.reshape(self.par['B_bar'][:, k], (m.n_x, m.n_u), order='F') @
            self.var['U'][:, k] +
            cvxpy.reshape(self.par['C_bar'][:, k], (m.n_x, m.n_u), order='F') @
            self.var['U'][:, k + 1] +
            self.par['S_bar'][:, k] * self.var['sigma'] +
            self.par['z_bar'][:, k] +
            self.var['nu'][:, k]
            for k in range(K - 1)
        ]

        # Trust regions:
        dx = cvxpy.sum(cvxpy.square(
            self.var['X'] - self.par['X_last']), axis=0)
        du = cvxpy.sum(cvxpy.square(
            self.var['U'] - self.par['U_last']), axis=0)
        ds = self.var['sigma'] - self.par['sigma_last']
        constraints += [cvxpy.norm(dx + du, 1) <= self.var['delta_norm']]
        constraints += [cvxpy.norm(ds, 'inf') <= self.var['sigma_norm']]

        # Flight time positive:
        constraints += [self.var['sigma'] >= 0.1]

        # Objective:
        sc_objective = cvxpy.Minimize(
            self.par['weight_sigma'] * self.var['sigma'] +
            self.par['weight_nu'] * cvxpy.norm(self.var['nu'], 'inf') +
            self.par['weight_delta'] * self.var['delta_norm'] +
            self.par['weight_delta_sigma'] * self.var['sigma_norm']
        )

        objective = sc_objective

        self.prob = cvxpy.Problem(objective, constraints)

    def set_parameters(self, **kwargs):
        """
        All parameters have to be filled before calling solve().
        Takes the following arguments as keywords:

        A_bar
        B_bar
        C_bar
        S_bar
        z_bar
        X_last
        U_last
        sigma_last
        E
        weight_sigma
        weight_nu
        radius_trust_region
        """

        for key in kwargs:
            if key in self.par:
                self.par[key].value = kwargs[key]
            else:
                print(f'Parameter \'{key}\' does not exist.')

    def get_variable(self, name):
        if name in self.var:
            return self.var[name].value
        else:
            print(f'Variable \'{name}\' does not exist.')
            return None

    def solve(self, **kwargs):
        error = False
        try:
            with warnings.catch_warnings():  # For User warning from solver
                warnings.simplefilter('ignore')
                self.prob.solve(verbose=verbose_solver,
                                solver=solver)
        except cvxpy.SolverError:
            error = True

        stats = self.prob.solver_stats

        info = {
            'setup_time': stats.setup_time,
            'solver_time': stats.solve_time,
            'iterations': stats.num_iters,
            'solver_error': error
        }

        return info


def axis3d_equal(X, Y, Z, ax):

    max_range = np.array([X.max() - X.min(), Y.max()
                          - Y.min(), Z.max() - Z.min()]).max()
    Xb = 0.5 * max_range * np.mgrid[-1:2:2, -1:2:2,
                                    - 1:2:2][0].flatten() + 0.5 * (X.max() + X.min())
    Yb = 0.5 * max_range * np.mgrid[-1:2:2, -1:2:2,
                                    - 1:2:2][1].flatten() + 0.5 * (Y.max() + Y.min())
    Zb = 0.5 * max_range * np.mgrid[-1:2:2, -1:2:2,
                                    - 1:2:2][2].flatten() + 0.5 * (Z.max() + Z.min())
    # Comment or uncomment following both lines to test the fake bounding box:
    for xb, yb, zb in zip(Xb, Yb, Zb):
        ax.plot([xb], [yb], [zb], 'w')


def plot_animation(X, U):  # pragma: no cover

    fig = plt.figure()
    ax = fig.add_subplot(projection='3d')
    # for stopping simulation with the esc key.
    fig.canvas.mpl_connect('key_release_event',
                           lambda event: [exit(0) if event.key == 'escape' else None])

    for k in range(K):
        plt.cla()
        ax.plot(X[2, :], X[3, :], X[1, :])  # trajectory
        ax.scatter3D([0.0], [0.0], [0.0], c="r",
                     marker="x")  # target landing point
        axis3d_equal(X[2, :], X[3, :], X[1, :], ax)

        rx, ry, rz = X[1:4, k]
        # vx, vy, vz = X[4:7, k]
        qw, qx, qy, qz = X[7:11, k]

        CBI = np.array([
            [1 - 2 * (qy ** 2 + qz ** 2), 2 * (qx * qy + qw * qz),
             2 * (qx * qz - qw * qy)],
            [2 * (qx * qy - qw * qz), 1 - 2
             * (qx ** 2 + qz ** 2), 2 * (qy * qz + qw * qx)],
            [2 * (qx * qz + qw * qy), 2 * (qy * qz - qw * qx),
             1 - 2 * (qx ** 2 + qy ** 2)]
        ])

        Fx, Fy, Fz = np.dot(np.transpose(CBI), U[:, k])
        dx, dy, dz = np.dot(np.transpose(CBI), np.array([1., 0., 0.]))

        # attitude vector
        ax.quiver(ry, rz, rx, dy, dz, dx, length=0.5, linewidth=3.0,
                  arrow_length_ratio=0.0, color='black')

        # thrust vector
        ax.quiver(ry, rz, rx, -Fy, -Fz, -Fx, length=0.1,
                  arrow_length_ratio=0.0, color='red')

        ax.set_title("Rocket powered landing")
        plt.pause(0.5)


def main(rng=None):
    print("start!!")
    m = Rocket_Model_6DoF(rng)

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

    # INITIALIZATION
    sigma = m.t_f_guess
    X, U = m.initialize_trajectory(X, U)

    integrator = Integrator(m, K)
    problem = SCProblem(m, K)

    converged = False
    w_delta = W_DELTA
    for it in range(iterations):
        t0_it = time()
        print('-' * 18 + f' Iteration {str(it + 1).zfill(2)} ' + '-' * 18)

        A_bar, B_bar, C_bar, S_bar, z_bar = integrator.calculate_discretization(
            X, U, sigma)

        problem.set_parameters(A_bar=A_bar, B_bar=B_bar, C_bar=C_bar, S_bar=S_bar, z_bar=z_bar,
                               X_last=X, U_last=U, sigma_last=sigma,
                               weight_sigma=W_SIGMA, weight_nu=W_NU,
                               weight_delta=w_delta, weight_delta_sigma=W_DELTA_SIGMA)
        problem.solve()

        X = problem.get_variable('X')
        U = problem.get_variable('U')
        sigma = problem.get_variable('sigma')

        delta_norm = problem.get_variable('delta_norm')
        sigma_norm = problem.get_variable('sigma_norm')
        nu_norm = np.linalg.norm(problem.get_variable('nu'), np.inf)

        print('delta_norm', delta_norm)
        print('sigma_norm', sigma_norm)
        print('nu_norm', nu_norm)

        if delta_norm < 1e-3 and sigma_norm < 1e-3 and nu_norm < 1e-7:
            converged = True

        w_delta *= 1.5

        print('Time for iteration', time() - t0_it, 's')

        if converged:
            print(f'Converged after {it + 1} iterations.')
            break

    if show_animation:  # pragma: no cover
        plot_animation(X, U)

    print("done!!")


if __name__ == '__main__':
    main()