AERO4200-Ass3/main.py

# Aero4200-Ass3
#   Cal Wing, Sem 2 2023

# Import System / Common Libs 
import os, pickle, time
import numpy as np
import pandas as pd
from tqdm import tqdm

from numpy import pi, sin, cos, tan

#Import makeGraph, a custom graphing wrapper I developed, refer to it for documentation  
from makeGraph import makeGraph, pltKeyClose, UQ_COLOURS as UQC # Custom Graphing Lib

# Make sure the relevant folders folder exists
folders = ["./images", "./tmp"]
for folder in folders:
    if not os.path.isdir(folder): os.mkdir(folder)

# This is a cacheing function, it checks if a cache file exists and loads it or it calcs & saves it.
# TL;DR This takes a file path & either a single variable (defArg) or some sort of callable and either 
# loads it or dumps it to disk. Note: It does not validate the data inside!  
def cacheData(dataFilePath: str, calcFunction: callable = lambda x: x, args: tuple = (), kargs: dict = {}, defArg=None):   
    if len(args) == 0 and defArg is not None:
            args = (defArg, )
    data = None
    dataFileExt = dataFilePath.rsplit(".")[-1]
    
    # Check if file exists
    if os.path.isfile(dataFilePath):
        print(f"Found data file \"{dataFilePath}\", loading data.")
        # Check if file is a compressed numpy file
        if dataFileExt == "npz":
            data = np.load(dataFilePath, allow_pickle=True)["arr_0"]
        elif dataFileExt == "pkl":
            #If its not then just pickle it normally
            with open(dataFilePath, 'rb') as handle:
                data = pickle.load(handle)
        elif dataFileExt == "dfz":
                data = pd.read_pickle(dataFilePath, compression="gzip")
        else:
            raise TypeError(f"Cannot determine file type of: {dataFilePath}")
    else:
        print(f"Could not find data file \"{dataFilePath}\", generating & saving data.")
        # Calculate Value
        data = calcFunction(*args, **kargs)
        
        # Check if file is a compressed numpy file
        if dataFileExt == "npz":
            np.savez_compressed(dataFilePath, data, allow_pickle=True)
        
        elif dataFileExt == "dfz":
                data.to_pickle(dataFilePath, compression="gzip") 
        
        elif dataFileExt == "pkl":
            #If its not then just pickle it normally
            with open(dataFilePath, 'wb') as handle:
                pickle.dump(data, handle)
        else:
            raise TypeError(f"Cannot determine file type of: {dataFilePath}")

    if data is None: raise ValueError("Could not import or generate data requested")
    return data



# IMU Data Loading
# I map in to a heading to add units / make things make more sense
## The gyroscopic body angular rates from the IMU are given:
#  - WBE_1 (in rad/s) - the roll rate about the body-fixed x-axis 
#  - WBE_2 (in rad/s) - the pitch rate about the body-fixed y-axis  
#  - WBE_3 (in rad/s) - the yaw rate about the body-fixed z-axis
## Specific forces:
#  - FSP_X (in m/s^2) - the specific force in the body-fixed x-direction
#  - FSP_Y (in m/s^2) - the specific force in the body-fixed y-direction
#  - FSP_Z (in m/s^2) - the specific force in the body-fixed z-direction 
IMU_TIME_HEADER = ["Time [s]"]
IMU_WBE_HEADERS = ["WBE_1 [rad/s]", "WBE_2 [rad/s]", "WBE_3 [rad/s]"] # Roll, Pitch, Yaw - Rate
IMU_FSP_HEADERS = ["FSP_X [m/s^2]", "FSP_Y [m/s^2]", "FSP_Z [m/s^2]"] # Specific Force X, Y , Z
IMU_DATA_HEADER = IMU_TIME_HEADER + IMU_WBE_HEADERS + IMU_FSP_HEADERS
def importIMUData(mission, imu):
    # If IMU is not a string, then convert based on bool eval where "H" == True
    if type(imu) != str: imu = "H" if imu else "L"  

    data = pd.read_csv(
           f"./data/IMU_M{str(mission) + str(imu)}.txt", 
            header=None, skiprows=1,
            names=IMU_DATA_HEADER,
        )

    return data

# Load the Mission Data
m1_IMUData = importIMUData(1, 0), importIMUData(1, 1) #(L, H) Data
m2_IMUData = importIMUData(2, 0), importIMUData(2, 1)

# NED Translation & Force Functions
INIT_EULER_ANGLES = (0, 0, 0)
def attidude_rate_2NED(angles, euler_angles):
    phi, theta, psi = euler_angles
    p, q, r = angles
    
    transMat = np.array([ [1, sin(phi)*tan(theta), cos(phi)*tan(theta) ],
                          [0, cos(phi),            -sin(phi)           ], 
                          [0, sin(phi)/cos(theta), cos(phi)/cos(theta) ]
                        ])
    
    angleMat = np.array([ [p], 
                          [q], 
                          [r]
                        ])

    return np.matmul(transMat, angleMat)

def getNEDForces(NEDPos):
    phi, theta, psi = NEDPos # Roll, Pitch, Yaw

    R41 = np.array([ [cos(psi)*cos(theta), cos(psi)*sin(theta)*sin(phi)-sin(psi)*cos(phi), cos(psi)*sin(theta)*cos(phi)+sin(psi)*sin(phi)],
                          [sin(psi)*cos(theta), sin(psi)*sin(theta)*sin(phi)+cos(psi)*cos(phi), sin(psi)*sin(theta)*cos(phi)-cos(psi)*sin(phi)],
                          [-sin(theta),         cos(theta)*sin(phi),                            cos(theta)*cos(phi)                           ]
                        ])
    
    return R41

# Calculate Mission Translation & Force Data
TRANS_DATA_HEADER = IMU_TIME_HEADER + IMU_WBE_HEADERS + ["Forces", "ForceSum", "ForceCumSum"]
def calculateTranslatedData(missionWBEData) -> pd.DataFrame:
    print("Translating Motion & Calculating Resulting Forces")
    translatedData = pd.DataFrame(columns=TRANS_DATA_HEADER)
    for i in tqdm(range(len(missionWBEData))):
        dataPoint = missionWBEData[IMU_WBE_HEADERS].iloc[i] # Get the time point data 
        trans = attidude_rate_2NED(dataPoint.values, INIT_EULER_ANGLES).flatten() # Translate to Net
        forces = getNEDForces(trans)
        forceSum = np.array([np.sum(forces[:,0].flatten()), np.sum(forces[:,1].flatten()), np.sum(forces[:,2].flatten())])
        privForce = translatedData[TRANS_DATA_HEADER[5]].iloc[i-1] if i > 0 else [0, 0, 0]
        forceCumSum = np.array([ forceSum[ii] + value for ii, value in enumerate(privForce) ])
        translatedData.loc[i] = {
            TRANS_DATA_HEADER[0]: missionWBEData[IMU_TIME_HEADER].iloc[i],
            TRANS_DATA_HEADER[1]: trans[0],
            TRANS_DATA_HEADER[2]: trans[1],
            TRANS_DATA_HEADER[3]: trans[2],
            TRANS_DATA_HEADER[4]: forces,
            TRANS_DATA_HEADER[5]: forceSum,
            TRANS_DATA_HEADER[6]: forceCumSum,
        }
        
    return translatedData

if __name__ == '__main__':
    missionWBEData = m1_IMUData[0][IMU_TIME_HEADER + IMU_WBE_HEADERS]
    translatedData = cacheData("./tmp/m1_transData.dfz", calculateTranslatedData, (missionWBEData,))
    
    x = np.cumsum(translatedData[IMU_WBE_HEADERS[0]])
    y = np.cumsum(translatedData[IMU_WBE_HEADERS[1]])
    z = np.cumsum(translatedData[IMU_WBE_HEADERS[2]])

    import matplotlib.pyplot as plt
    fig = plt.figure()
    ax = plt.axes(projection='3d')
    plots = translatedData[TRANS_DATA_HEADER[6]].values
    #x = [p[0] for p in plots]
    #y = [p[1] for p in plots]
    #z = [p[2] for p in plots]
    ax.plot3D(x, y, z)
    ax.plot3D(x+5, y-2, z+1)
    
    plt.show()
    

    print("Complete")