NTUT Course Attitue Control

ntut-lab3.py

"""
Full Lab: Spacecraft Motion and Attitude Visualization in Multiple Frames
with Exportable Interactive 3D Visualizations for both Final Attitude and ECI Orbit

This lab simulates a spacecraft orbiting Earth with an attitude control system.
It produces:
  - Static plots (attitude error, control torque, rate error, orientation projection)
  - An animation showing the spacecraft moving in the ECI frame along its orbit,
    with its body frame and Hill (movement) frame axes.
  - Two interactive Plotly 3D figures exported as HTML files:
      1. Final Attitude Visualization (Body vs. Hill frame)
      2. Full ECI Orbit Visualization
These interactive files allow manual rotation and inspection in a web browser.
"""

import os
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # For 3D plotting
import matplotlib.animation as animation

# Import Plotly for interactive 3D export
import plotly.graph_objects as go

# Basilisk module imports
from Basilisk import __path__
from Basilisk.architecture import messaging
from Basilisk.fswAlgorithms import attTrackingError, hillPoint, mrpFeedback
from Basilisk.simulation import extForceTorque, simpleNav, spacecraft
from Basilisk.utilities import (RigidBodyKinematics, SimulationBaseClass, macros,
                                orbitalMotion, simIncludeGravBody, unitTestSupport, vizSupport)

# Determine file name for figure labeling
bskPath = __path__[0]
fileName = os.path.basename(os.path.splitext(__file__)[0])


# ------------------------------ Static Plotting Functions ------------------------------ #
def plot_attitude_error(timeLineSet, dataSigmaBR):
    """Plot the attitude error norm (log scale) vs. time."""
    plt.figure()
    sNorm = np.array([np.linalg.norm(v) for v in dataSigmaBR])
    plt.plot(timeLineSet, sNorm, color=unitTestSupport.getLineColor(1, 3))
    plt.xlabel('Time [min]')
    plt.ylabel(r'Attitude Error Norm $|\sigma_{B/R}|$')
    plt.title("Attitude Error Convergence")
    plt.yscale('log')
    plt.grid()

def plot_control_torque(timeLineSet, dataLr):
    """Plot the control torque components vs. time."""
    plt.figure()
    for idx in range(3):
        plt.plot(timeLineSet, dataLr[:, idx],
                 color=unitTestSupport.getLineColor(idx, 3),
                 label=f'$L_{{r,{idx}}}$')
    plt.legend(loc='lower right')
    plt.xlabel('Time [min]')
    plt.ylabel('Control Torque [Nm]')
    plt.title("Control Torque Components")
    plt.grid()

def plot_rate_error(timeLineSet, dataOmegaBR):
    """Plot the angular rate tracking error vs. time."""
    plt.figure()
    for idx in range(3):
        plt.plot(timeLineSet, dataOmegaBR[:, idx],
                 color=unitTestSupport.getLineColor(idx, 3),
                 label=r'$\omega_{BR,' + str(idx) + '}$')
    plt.legend(loc='lower right')
    plt.xlabel('Time [min]')
    plt.ylabel('Rate Error [rad/s]')
    plt.title("Angular Rate Tracking Error")
    plt.grid()

def plot_orientation(timeLineSet, dataPos, dataVel, dataSigmaBN):
    """
    Plot the projection of the spacecraft body frame axes onto the Hill frame axes.
    These projections should converge to +1 as the frames align.
    """
    numPoints = len(dataPos)
    projData = np.empty([numPoints, 3])
    for idx in range(numPoints):
        ir = dataPos[idx] / np.linalg.norm(dataPos[idx])
        hv = np.cross(dataPos[idx], dataVel[idx])
        ih = hv / np.linalg.norm(hv)
        itheta = np.cross(ih, ir)
        dcmBN = RigidBodyKinematics.MRP2C(dataSigmaBN[idx])
        projData[idx, 0] = np.dot(ir, dcmBN[0])
        projData[idx, 1] = np.dot(itheta, dcmBN[1])
        projData[idx, 2] = np.dot(ih, dcmBN[2])
    plt.figure()
    labelStrings = (r'$\hat\imath_r\cdot \hat b_1$',
                    r'${\hat\imath}_{\theta}\cdot \hat b_2$',
                    r'$\hat\imath_h\cdot \hat b_3$')
    for idx in range(3):
        plt.plot(timeLineSet, projData[:, idx],
                 color=unitTestSupport.getLineColor(idx, 3),
                 label=labelStrings[idx])
    plt.legend(loc='lower right')
    plt.xlabel('Time [min]')
    plt.ylabel('Projection Value')
    plt.title("Orientation Alignment Projection")
    plt.grid()

def plot_final_attitude_3d(navAttLog, navTransLog):
    """
    Create a static 3D plot of the final spacecraft attitude showing:
      - Body-fixed frame axes (b1, b2, b3)
      - Hill frame axes (r̂, θ̂, ĥ)
    """
    sigma_final = navAttLog.sigma_BN[-1]
    r_final = navTransLog.r_BN_N[-1]
    v_final = navTransLog.v_BN_N[-1]
    dcmBN = RigidBodyKinematics.MRP2C(sigma_final)
    b1, b2, b3 = dcmBN[0], dcmBN[1], dcmBN[2]
    ir = r_final / np.linalg.norm(r_final)
    hv = np.cross(r_final, v_final)
    ih = hv / np.linalg.norm(hv)
    itheta = np.cross(ih, ir)
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    scale = 1e6  # Adjust scale for visualization
    ax.quiver(0, 0, 0, b1[0], b1[1], b1[2],
              color='r', length=scale, normalize=True, label='Body Axis b1')
    ax.quiver(0, 0, 0, b2[0], b2[1], b2[2],
              color='g', length=scale, normalize=True, label='Body Axis b2')
    ax.quiver(0, 0, 0, b3[0], b3[1], b3[2],
              color='b', length=scale, normalize=True, label='Body Axis b3')
    ax.quiver(0, 0, 0, ir[0], ir[1], ir[2],
              color='m', length=scale, normalize=True, label='Hill Axis r̂')
    ax.quiver(0, 0, 0, itheta[0], itheta[1], itheta[2],
              color='c', length=scale, normalize=True, label='Hill Axis θ̂')
    ax.quiver(0, 0, 0, ih[0], ih[1], ih[2],
              color='y', length=scale, normalize=True, label='Hill Axis ĥ')
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    ax.set_title("Final Attitude Visualization (Body vs. Hill Frame)")
    ax.legend(loc='upper left')
    plt.tight_layout()
    return fig


# ------------------------------ Animation: Spacecraft Motion in Multiple Frames ------------------------------ #
def animate_spacecraft_motion(navAttLog, navTransLog, interval=50):
    """
    Animate the spacecraft's motion in the ECI frame, showing:
      - The orbit path in the inertial (ECI) frame.
      - The spacecraft's current position.
      - The body frame axes (red, green, blue) drawn at the spacecraft's position.
      - The movement (Hill) frame axes (magenta, cyan, yellow) derived from its position and velocity.
      - Static ECI axes drawn at the origin.

    Args:
        navAttLog: Navigation attitude log (with sigma_BN).
        navTransLog: Navigation translation log (with r_BN_N and v_BN_N).
        interval: Delay between frames in milliseconds.
    """
    sigmaBN = navAttLog.sigma_BN  # Attitude log (MRP values)
    rN = navTransLog.r_BN_N       # Position vectors (ECI) over time
    vN = navTransLog.v_BN_N       # Velocity vectors (ECI) over time
    timeLine = navAttLog.times() * macros.NANO2MIN  # Time in minutes

    # Set up the 3D plot for the ECI frame
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    # Determine limits based on orbit positions
    r_min = np.min(rN, axis=0)
    r_max = np.max(rN, axis=0)
    margin = 0.1 * np.linalg.norm(r_max - r_min)
    ax.set_xlim([r_min[0]-margin, r_max[0]+margin])
    ax.set_ylim([r_min[1]-margin, r_max[1]+margin])
    ax.set_zlim([r_min[2]-margin, r_max[2]+margin])
    ax.set_xlabel('ECI X [m]')
    ax.set_ylabel('ECI Y [m]')
    ax.set_zlabel('ECI Z [m]')
    ax.set_title("Spacecraft Motion in ECI with Body & Hill Frames")

    # Draw static ECI axes at the origin
    scaleECI = 1e7  # Adjust scaling as needed
    ax.quiver(0, 0, 0, scaleECI, 0, 0, color='k', linestyle='--', label='ECI X')
    ax.quiver(0, 0, 0, 0, scaleECI, 0, color='k', linestyle='--', label='ECI Y')
    ax.quiver(0, 0, 0, 0, 0, scaleECI, color='k', linestyle='--', label='ECI Z')

    # Initialize an empty line to show the orbit path
    orbit_line, = ax.plot([], [], [], 'gray', linewidth=1, label='Orbit Path')

    def update(frame):
        # Remove previous quiver objects
        for coll in list(ax.collections):
            coll.remove()

        # Update the orbit path (all points up to current frame)
        orbit_line.set_data(rN[:frame, 0], rN[:frame, 1])
        orbit_line.set_3d_properties(rN[:frame, 2])

        # Current spacecraft position (in ECI)
        pos = rN[frame]

        # Compute body frame axes from current attitude (MRP)
        sigma = sigmaBN[frame]
        dcmBN = RigidBodyKinematics.MRP2C(sigma)
        b1, b2, b3 = dcmBN[0], dcmBN[1], dcmBN[2]

        # Compute movement (Hill) frame axes from current position and velocity
        r = rN[frame]
        v = vN[frame]
        r_norm = r / np.linalg.norm(r)
        hv = np.cross(r, v)
        ih = hv / np.linalg.norm(hv)
        itheta = np.cross(ih, r_norm)

        # Set arrow length (adjust as needed for visualization)
        axis_len = 1e6

        # Draw body frame axes (red, green, blue)
        ax.quiver(pos[0], pos[1], pos[2],
                  b1[0], b1[1], b1[2], color='r', length=axis_len, normalize=True)
        ax.quiver(pos[0], pos[1], pos[2],
                  b2[0], b2[1], b2[2], color='g', length=axis_len, normalize=True)
        ax.quiver(pos[0], pos[1], pos[2],
                  b3[0], b3[1], b3[2], color='b', length=axis_len, normalize=True)

        # Draw movement (Hill) frame axes (magenta, cyan, yellow)
        ax.quiver(pos[0], pos[1], pos[2],
                  r_norm[0], r_norm[1], r_norm[2], color='m', length=axis_len, normalize=True)
        ax.quiver(pos[0], pos[1], pos[2],
                  itheta[0], itheta[1], itheta[2], color='c', length=axis_len, normalize=True)
        ax.quiver(pos[0], pos[1], pos[2],
                  ih[0], ih[1], ih[2], color='y', length=axis_len, normalize=True)

        ax.set_title(f"Time: {timeLine[frame]:.2f} min")

    ani = animation.FuncAnimation(fig, update, frames=len(sigmaBN), interval=interval, repeat=True)
    plt.legend()
    plt.show()
    return ani


# ------------------------------ Interactive 3D Export Functions (Plotly) ------------------------------ #
def save_interactive_3d_figure(navAttLog, navTransLog, filename="final_attitude.html"):
    """
    Create an interactive 3D plot (using Plotly) of the final spacecraft attitude
    showing the body frame axes and the Hill frame axes. The figure is saved as an HTML file.

    Args:
        navAttLog: Navigation attitude log (with sigma_BN).
        navTransLog: Navigation translation log (with r_BN_N and v_BN_N).
        filename: The output HTML file name.
    """
    sigma_final = navAttLog.sigma_BN[-1]
    r_final = navTransLog.r_BN_N[-1]
    v_final = navTransLog.v_BN_N[-1]
    dcmBN = RigidBodyKinematics.MRP2C(sigma_final)
    b1, b2, b3 = dcmBN[0], dcmBN[1], dcmBN[2]
    # Compute Hill frame axes
    ir = r_final / np.linalg.norm(r_final)
    hv = np.cross(r_final, v_final)
    ih = hv / np.linalg.norm(hv)
    itheta = np.cross(ih, ir)
    scale = 1e6  # Scale factor for arrow length

    # Create body frame line segments
    body_traces = []
    for vec, color, name in zip([b1, b2, b3], ['red', 'green', 'blue'], ['Body Axis b1', 'Body Axis b2', 'Body Axis b3']):
        body_traces.append(go.Scatter3d(
            x=[0, vec[0]*scale],
            y=[0, vec[1]*scale],
            z=[0, vec[2]*scale],
            mode='lines',
            line=dict(color=color, width=5),
            name=name))

    # Create Hill frame line segments
    hill_traces = []
    for vec, color, name in zip([ir, itheta, ih], ['magenta', 'cyan', 'yellow'], ['Hill Axis r̂', 'Hill Axis θ̂', 'Hill Axis ĥ']):
        hill_traces.append(go.Scatter3d(
            x=[0, vec[0]*scale],
            y=[0, vec[1]*scale],
            z=[0, vec[2]*scale],
            mode='lines',
            line=dict(color=color, width=5),
            name=name))

    fig = go.Figure(data=body_traces + hill_traces)
    fig.update_layout(scene=dict(
        xaxis_title='X',
        yaxis_title='Y',
        zaxis_title='Z'),
        title="Interactive 3D Final Attitude (Body vs. Hill Frame)")

    fig.write_html(filename)
    print(f"Interactive 3D final attitude figure saved to {filename}")
    return fig

def save_interactive_eci_figure(navTransLog, filename="eci_orbit.html"):
    """
    Create an interactive 3D plot (using Plotly) of the spacecraft's orbit in the ECI frame.
    The orbit path is plotted as a 3D line, and the figure is saved as an HTML file.

    Args:
        navTransLog: Navigation translation log (with r_BN_N).
        filename: The output HTML file name.
    """
    r_data = np.array(navTransLog.r_BN_N)  # Convert logged positions to a NumPy array
    trace = go.Scatter3d(
        x=r_data[:, 0].flatten(),
        y=r_data[:, 1].flatten(),
        z=r_data[:, 2].flatten(),
        mode='lines',
        name='ECI Orbit',
        line=dict(color='blue', width=2)
    )
    fig = go.Figure(data=[trace])
    fig.update_layout(title="Interactive 3D ECI Orbit",
                      scene=dict(
                          xaxis_title="ECI X (m)",
                          yaxis_title="ECI Y (m)",
                          zaxis_title="ECI Z (m)"))
    fig.write_html(filename)
    print(f"Interactive ECI orbit figure saved to {filename}")
    return fig


# ------------------------------ Simulation Setup Functions ------------------------------ #
def create_simulation(sim_task_name="simTask", sim_process_name="simProcess", time_step=macros.sec2nano(0.1)):
    """
    Create a Basilisk simulation with a specified process and task.
    """
    sim = SimulationBaseClass.SimBaseClass()
    process = sim.CreateNewProcess(sim_process_name)
    process.addTask(sim.CreateNewTask(sim_task_name, time_step))
    return sim, sim_task_name

def setup_spacecraft(sim, task_name):
    """
    Configure the spacecraft with inertial properties and an initial orbit.
    """
    sc = spacecraft.Spacecraft()
    sc.ModelTag = "bsk-Sat"
    I = [900., 0., 0.,
         0., 800., 0.,
         0., 0., 600.]
    sc.hub.mHub = 750.0
    sc.hub.r_BcB_B = [[0.0], [0.0], [0.0]]
    sc.hub.IHubPntBc_B = unitTestSupport.np2EigenMatrix3d(I)
    oe = orbitalMotion.ClassicElements()
    oe.a = 10000000.0
    oe.e = 0.1
    oe.i = 33.3 * macros.D2R
    oe.Omega = 48.2 * macros.D2R
    oe.omega = 347.8 * macros.D2R
    oe.f = 85.3 * macros.D2R
    gravFactory = simIncludeGravBody.gravBodyFactory()
    earth = gravFactory.createEarth()
    mu = earth.mu
    rN, vN = orbitalMotion.elem2rv(mu, oe)
    sc.hub.r_CN_NInit = rN
    sc.hub.v_CN_NInit = vN
    sc.hub.sigma_BNInit = [[0.1], [0.2], [-0.3]]
    sc.hub.omega_BN_BInit = [[0.001], [-0.01], [0.03]]
    sim.AddModelToTask(task_name, sc)
    return sc

def setup_gravity(sc):
    """
    Attach Earth's gravity model to the spacecraft.
    """
    gravFactory = simIncludeGravBody.gravBodyFactory()
    earth = gravFactory.createEarth()
    earth.isCentralBody = True
    gravFactory.addBodiesTo(sc)

def setup_external_force_torque(sim, task_name, sc):
    """
    Setup the external force/torque module.
    """
    extFT = extForceTorque.ExtForceTorque()
    extFT.ModelTag = "externalDisturbance"
    sc.addDynamicEffector(extFT)
    sim.AddModelToTask(task_name, extFT)
    return extFT

def setup_navigation(sim, task_name, sc):
    """
    Setup the simple navigation module.
    """
    nav = simpleNav.SimpleNav()
    nav.ModelTag = "SimpleNavigation"
    sim.AddModelToTask(task_name, nav)
    nav.scStateInMsg.subscribeTo(sc.scStateOutMsg)
    return nav

def setup_guidance(sim, task_name, nav, use_alt_body_frame=False):
    """
    Setup the hillPoint guidance module.
    """
    guidance = hillPoint.hillPoint()
    guidance.ModelTag = "hillPoint"
    guidance.transNavInMsg.subscribeTo(nav.transOutMsg)
    celBodyData = messaging.EphemerisMsgPayload()
    celBodyMsg = messaging.EphemerisMsg().write(celBodyData)
    guidance.celBodyInMsg.subscribeTo(celBodyMsg)
    sim.AddModelToTask(task_name, guidance)
    return guidance

def setup_attitude_error(sim, task_name, nav, guidance, use_alt_body_frame=False):
    """
    Setup the attitude tracking error module.
    """
    attError = attTrackingError.attTrackingError()
    attError.ModelTag = "attErrorInertial3D"
    if use_alt_body_frame:
        attError.sigma_R0R = [0, 1, 0]
    attError.attRefInMsg.subscribeTo(guidance.attRefOutMsg)
    attError.attNavInMsg.subscribeTo(nav.attOutMsg)
    sim.AddModelToTask(task_name, attError)
    return attError

def setup_controller(sim, task_name, attError):
    """
    Setup the MRP feedback controller.
    """
    controller = mrpFeedback.mrpFeedback()
    controller.ModelTag = "mrpFeedback"
    controller.guidInMsg.subscribeTo(attError.attGuidOutMsg)
    controller.K = 3.5
    controller.Ki = -1.0  # Negative to disable integral feedback
    controller.P = 30.0
    controller.integralLimit = 2. / controller.Ki * 0.1
    sim.AddModelToTask(task_name, controller)
    return controller

def setup_vehicle_config(controller, I):
    """
    Setup and publish the vehicle configuration message.
    """
    vehicleConfig = messaging.VehicleConfigMsgPayload()
    vehicleConfig.ISCPntB_B = I
    configMsg = messaging.VehicleConfigMsg().write(vehicleConfig)
    controller.vehConfigInMsg.subscribeTo(configMsg)

def setup_logging(sim, task_name, simulation_time, simulation_time_step, controller, attError, nav):
    """
    Setup logging for control, attitude error, and navigation.
    """
    numDataPoints = 100
    samplingTime = unitTestSupport.samplingTime(simulation_time, simulation_time_step, numDataPoints)
    mrpLog = controller.cmdTorqueOutMsg.recorder(samplingTime)
    attErrLog = attError.attGuidOutMsg.recorder(samplingTime)
    navAttLog = nav.attOutMsg.recorder(samplingTime)
    navTransLog = nav.transOutMsg.recorder(samplingTime)
    sim.AddModelToTask(task_name, mrpLog)
    sim.AddModelToTask(task_name, attErrLog)
    sim.AddModelToTask(task_name, navAttLog)
    sim.AddModelToTask(task_name, navTransLog)
    return mrpLog, attErrLog, navAttLog, navTransLog

def run_simulation(sim, stop_time):
    """
    Initialize and execute the simulation.
    """
    sim.InitializeSimulation()
    sim.ConfigureStopTime(stop_time)
    sim.ExecuteSimulation()

def plot_results(mrpLog, attErrLog, navAttLog, navTransLog):
    """
    Generate static plots.
    """
    dataLr = mrpLog.torqueRequestBody
    dataSigmaBR = attErrLog.sigma_BR
    dataOmegaBR = attErrLog.omega_BR_B
    dataPos = navTransLog.r_BN_N
    dataVel = navTransLog.v_BN_N
    dataSigmaBN = navAttLog.sigma_BN
    timeLine = attErrLog.times() * macros.NANO2MIN
    plt.close("all")
    plot_attitude_error(timeLine, dataSigmaBR)
    plot_control_torque(timeLine, dataLr)
    plot_rate_error(timeLine, dataOmegaBR)
    plot_orientation(timeLine, dataPos, dataVel, dataSigmaBN)
    fig3d = plot_final_attitude_3d(navAttLog, navTransLog)
    plt.show()
    return {"3D_attitude": fig3d}

def run_attitude_guidance(show_plots=True, use_alt_body_frame=False):
    """
    Run the complete attitude guidance simulation.

    Args:
        show_plots (bool): If True, display static plots, animation, and export interactive 3D figures.
        use_alt_body_frame (bool): If True, use an alternate corrected body frame.

    Returns:
        dict: Dictionary of static figure objects.
    """
    sim_task_name = "simTask"
    sim_process_name = "simProcess"
    simulation_time = macros.min2nano(10.0)
    simulation_time_step = macros.sec2nano(0.1)
    sim, task_name = create_simulation(sim_task_name, sim_process_name, simulation_time_step)
    sc = setup_spacecraft(sim, task_name)
    setup_gravity(sc)
    extFT = setup_external_force_torque(sim, task_name, sc)
    nav = setup_navigation(sim, task_name, sc)
    guidance = setup_guidance(sim, task_name, nav, use_alt_body_frame)
    attError = setup_attitude_error(sim, task_name, nav, guidance, use_alt_body_frame)
    controller = setup_controller(sim, task_name, attError)
    I = [900., 0., 0., 0., 800., 0., 0., 0., 600.]
    setup_vehicle_config(controller, I)
    extFT.cmdTorqueInMsg.subscribeTo(controller.cmdTorqueOutMsg)
    vizSupport.enableUnityVisualization(sim, task_name, sc)
    mrpLog, attErrLog, navAttLog, navTransLog = setup_logging(sim, task_name,
                                                               simulation_time,
                                                               simulation_time_step,
                                                               controller, attError, nav)
    run_simulation(sim, simulation_time)
    figures = {}
    if show_plots:
        figures = plot_results(mrpLog, attErrLog, navAttLog, navTransLog)
        animate_spacecraft_motion(navAttLog, navTransLog)
        # Export interactive 3D figures for manual inspection:
        save_interactive_3d_figure(navAttLog, navTransLog, filename="final_attitude.html")
        save_interactive_eci_figure(navTransLog, filename="eci_orbit.html")
    return figures

if __name__ == "__main__":
    run_attitude_guidance(show_plots=True, use_alt_body_frame=True)