analyze.py

# Import CSV in the format with the following columns:
# 1. Timestamp (ms) from start of burn
# 2. Fx (N) - force in x direction
# 3. Fy (N) - force in y direction
# 4. Fz (N) - force in z direction
# 5. Mx (Nm) - moment about x axis
# 6. My (Nm) - moment about y axis
# 7. Mz (Nm) - moment about z axis
# Where Z is the thrust axis, X and Y are the lateral and longitudinal axes, and Mx, My, and Mz are moments about the X, Y, and Z axes, respectively.

START_WINDOW_BUFFER = 150
END_WINDOW_BUFFER = 450

import csv
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

def generate_report(csv_file):
    data = pd.read_csv(csv_file)
    print(data.describe())

def find_start_end(csv_file, window_size=50, num_std=1, duration=3):
    data = pd.read_csv(csv_file)
    
    # Calculate rolling mean and standard deviation
    roll_mean = data['Fz'].rolling(window_size).mean()
    roll_std = data['Fz'].rolling(window_size).std()

    # Define start of burn as the point where 'Fz' exceeds the mean by 'num_std' standard deviations for 'duration' samples
    start = ((data['Fz'] - roll_mean) > num_std * roll_std).rolling(window_size).sum().idxmax() - START_WINDOW_BUFFER

    # Define end of burn as the point where 'Fz' drops below the mean by 'num_std' standard deviations for 'duration' samples, after the start
    end = ((roll_mean - data['Fz']) > num_std * roll_std).iloc[start:].rolling(window_size).sum().idxmax() + END_WINDOW_BUFFER

    return start, end

def generate_pruned_normalized_csv(csv_file, window_size=50, num_std=1, duration=3):
    data = pd.read_csv(csv_file)
    
    # Find the approximate start and end of the burn
    start, end = find_start_end(csv_file, window_size, num_std, duration)

    # "zero" the force by subtracting the mean of the samples before the start
    data['Fz'] = data['Fz'] - data['Fz'].iloc[:start].mean()
    data['Fx'] = data['Fx'] - data['Fx'].iloc[:start].mean()
    data['Fy'] = data['Fy'] - data['Fy'].iloc[:start].mean()

    # "zero" the moments by subtracting the mean of the samples before the start
    data['Mx'] = data['Mx'] - data['Mx'].iloc[:start].mean()
    data['My'] = data['My'] - data['My'].iloc[:start].mean()
    data['Mz'] = data['Mz'] - data['Mz'].iloc[:start].mean()

    # # filter out noise using a sinc filter, scale to the same range as the Fz data
    # filtZ = data['Fz'].rolling(window_size).apply(lambda x: np.sum(x * np.sinc(np.arange(-window_size/2, window_size/2))))

    # # Compare filtered data to original data
    # plt.plot(data['Fz'], label='Original')
    # plt.plot(filtZ, label='Filtered')
    # plt.legend()
    # plt.show()

    pruned = data[start:end]
    output_csv = f'pruned_normalized_{csv_file.split(".")[0]}.csv'
    pruned.to_csv(output_csv, index=False)


def generate_pruned_csv(csv_file, window_size=50, num_std=1, duration=3):
    data = pd.read_csv(csv_file)
    
    # Find the approximate start and end of the burn
    start, end = find_start_end(csv_file, window_size, num_std, duration)

    pruned = data[start:end]
    output_csv = f'pruned_{csv_file.split(".")[0]}.csv'
    pruned.to_csv(output_csv, index=False)

def plot_csv_onegraph_separate(csv_file):
    data = pd.read_csv(csv_file)
    
    # Plot forces Fx, Fy, and Fz in one graph
    plt.figure(figsize=(15, 5))
    plt.plot(data['Fx'], label='Fx')
    plt.plot(data['Fy'], label='Fy')
    plt.plot(data['Fz'], label='Fz', color='red')

    plt.title('Force Data')
    plt.xlabel('Relative Time [ms]')
    plt.ylabel('Force [N]')
    plt.legend()

    plt.tight_layout()
    plt.show()

    # Plot moments Mx, My, and Mz in one graph
    plt.figure(figsize=(15, 5))
    plt.plot(data['Mx'], label='Mx')
    plt.plot(data['My'], label='My')
    plt.plot(data['Mz'], label='Mz', color='red')

    plt.title('Moment Data')
    plt.xlabel('Relative Time [ms]')
    plt.ylabel('Moment [Nm]')
    plt.legend()

    plt.tight_layout()
    plt.show()

def plot_csv_all(csv_file):
    # Plot the Fx, Fy, Fz, Mx, My, and Mz data as 6 graphs in one figure
    data = pd.read_csv(csv_file)

    fig, axs = plt.subplots(2, 3, figsize=(15, 10))

    axs[0, 1].plot(data['Fx'])
    axs[0, 1].set_title('Fx')
    axs[0, 1].set_xlabel('Relative Time [ms]')
    axs[0, 1].set_ylabel('Force [N]')

    axs[0, 2].plot(data['Fy'])
    axs[0, 2].set_title('Fy')
    axs[0, 2].set_xlabel('Relative Time [ms]')
    axs[0, 2].set_ylabel('Force [N]')

    axs[0, 0].plot(data['Fz'], color='red')
    axs[0, 0].set_title('Fz')
    axs[0, 0].set_xlabel('Relative Time [ms]')
    axs[0, 0].set_ylabel('Force [N]')

    axs[1, 1].plot(data['Mx'])
    axs[1, 1].set_title('Mx')
    axs[1, 1].set_xlabel('Relative Time [ms]')
    axs[1, 1].set_ylabel('Moment [Nm]')

    axs[1, 2].plot(data['My'])
    axs[1, 2].set_title('My')
    axs[1, 2].set_xlabel('Relative Time [ms]')
    axs[1, 2].set_ylabel('Moment [Nm]')

    axs[1, 0].plot(data['Mz'])
    axs[1, 0].set_title('Mz')
    axs[1, 0].set_xlabel('Relative Time [ms]')
    axs[1, 0].set_ylabel('Moment [Nm]')

    plt.tight_layout()
    plt.show()

def plot_csv_separate(csv_file, scale_to_z = False):
    '''
    Scale to z means to keep the Fx and Fy scale the same as the Fz scale
    '''
    # Plot the Fx, Fy, Fz, Mx, My, and Mz data as two different plots, one with Forces and one with Moments
    data = pd.read_csv(csv_file)

    # Plot the forces, Fx, Fy, and Fz
    fig, axs = plt.subplots(1, 3, figsize=(15, 5))

    axs[0].plot(data['Fz'], color='red')
    axs[0].set_title('Fz Data')
    axs[0].set_xlabel('Relative Time [ms]')
    axs[0].set_ylabel('Force [N]')

    axs[1].plot(data['Fx'])
    axs[1].set_title('Fx Data')
    axs[1].set_xlabel('Relative Time [ms]')
    axs[1].set_ylabel('Force [N]')

    axs[2].plot(data['Fy'])
    axs[2].set_title('Fy Data')
    axs[2].set_xlabel('Relative Time [ms]')
    axs[2].set_ylabel('Force [N]')

    if scale_to_z:
        axs[1].set_ylim(axs[0].get_ylim())
        axs[2].set_ylim(axs[0].get_ylim())

    plt.tight_layout()
    plt.show()

def plot_csv_onegraph(csv_file):
    '''
    Same as plot_csv_onegraph_separate but with both plots in 1 figure
    '''
    data = pd.read_csv(csv_file)

    fig, axs = plt.subplots(2, 1, figsize=(15, 10))
    axs[0].plot(data['Fx'], label='Fx')
    axs[0].plot(data['Fy'], label='Fy')
    axs[0].plot(data['Fz'], label='Fz', color='red')

    axs[0].set_title('Force Data')
    axs[0].set_xlabel('Relative Time [ms]')
    axs[0].set_ylabel('Force [N]')
    axs[0].legend()

    axs[0].yaxis.set_major_locator(plt.MultipleLocator(50))
    axs[0].grid(color='gray', linestyle='--', linewidth=0.5)

    axs[1].plot(data['Mx'], label='Mx')
    axs[1].plot(data['My'], label='My')
    axs[1].plot(data['Mz'], label='Mz', color='red')

    axs[1].set_title('Moment Data')
    axs[1].set_xlabel('Relative Time [ms]')
    axs[1].set_ylabel('Moment [Nm]')
    axs[1].legend()

    axs[1].grid(color='gray', linestyle='--', linewidth=0.5)

    plt.tight_layout()
    plt.show()


#TODO: Work in progress
def motor_statistics(csv_file):
    data = pd.read_csv(csv_file)

    # Aside from the burn, remove any outliers
    start, end = find_start_end(csv_file)
    before = data[:start]
    after = data[end:]

    # Remove outliers from before and after the burn
    #before = before[(before['Fz'] > before['Fz'].mean() - 3 * before['Fz'].std()) & (before['Fz'] < before['Fz'].mean() + 3 * before['Fz'].std())]
    #after = after[(after['Fz'] > after['Fz'].mean() - 3 * after['Fz'].std()) & (after['Fz'] < after['Fz'].mean() + 3 * after['Fz'].std())]

    # Calculate the mean before the start of the burn and after the end of the burn
    mean_before = before['Fz'].mean()
    mean_after = after['Fz'].mean()
    print('-- Motor Statistics --')
    print('Approximate Propellant Mass:', mean_before - mean_after, 'N')

    # Calculate mean and max thrust during the burn
    # Tighten the data to remove any points before and after the burn by rolling a stdev window
    roll_std = data['Fz'].rolling(10).std()
    start = (roll_std > 0.5).idxmax()
    end = (roll_std > 0.5).iloc[start:].idxmax()

    # Plot this data
    plt.plot(data['Fz'])
    
    burn = data[start:end]
    print('Mean Thrust:', burn['Fz'].mean(), 'N')
    print('Max Thrust:', burn['Fz'].max(), 'N')

    print('Burn time:', burn['timestamp'].iloc[-1] - burn['timestamp'].iloc[0], 'ms')

    # Note: We should be able to compensate for linear thrust decay of the burnt propellant... actually didn't we have to do something like this in phys or chem?

    
    # Calculate the mean before the start of the burn
    mean_before = data['Fz'].iloc[:START_WINDOW_BUFFER].mean()

#TODO: Random work in progress stuff
def visualize_rocket(csv_file):
    # Visualize the rocket's trajectory

    # Read the data
    data = pd.read_csv(csv_file)

    # From the forces, integrate over time to get the velocity and position
    # Assume the rocket is at rest at the start
    data['vx'] = data['Fx'].cumsum()
    data['vy'] = data['Fy'].cumsum()
    data['vz'] = data['Fz'].cumsum()

    data['x'] = data['vx'].cumsum()
    data['y'] = data['vy'].cumsum()
    data['z'] = data['vz'].cumsum()

    # Plot the trajectory
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.plot(data['x'], data['y'], data['z'])
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    # Keep the axes equal according to Z axis
    max_range = np.array([data['x'].max()-data['x'].min(), data['y'].max()-data['y'].min(), data['z'].max()-data['z'].min()]).max() / 2.0
    mid_x = (data['x'].max()+data['x'].min()) / 2.0
    mid_y = (data['y'].max()+data['y'].min()) / 2.0
    mid_z = (data['z'].max()+data['z'].min()) / 2.0
    ax.set_xlim(mid_x - max_range, mid_x + max_range)
    ax.set_ylim(mid_y - max_range, mid_y + max_range)
    ax.set_zlim(mid_z - max_range, mid_z + max_range)

    plt.show()

import argparse
if __name__ == '__main__':
    # Argparse
    parser = argparse.ArgumentParser(description='Analyze test stand data')
    # Optional csv_file argument
    parser.add_argument('-f', '--csv_file', type=str, help='CSV file to analyze', default='data.csv')

    args = parser.parse_args()
    csv_file = args.csv_file
    csv_file_no_ext = csv_file.split('.')[0]

    #generate_report(csv_file)
    #plot_csv_onegraph(csv_file)
    generate_pruned_normalized_csv(csv_file)
    generate_report(f'pruned_normalized_{csv_file_no_ext}.csv')
    plot_csv_onegraph(f'pruned_normalized_{csv_file_no_ext}.csv')
    #plot_csv_all('pruned_normalized.csv')
    #plot_csv_separate('pruned_normalized.csv', scale_to_z=False)
    #motor_statistics('pruned_normalized.csv')
    #visualize_rocket('pruned_normalized.csv')