evaluate_run_sim.py

""" Evaluation script of simulation experiments generated by sim_exp.py """
import argparse
import gc
import json
import os
import pickle
import time
from typing import Union, List

import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib
from matplotlib import pyplot as plt, patches
from matplotlib.collections import LineCollection
from scipy.stats import norm
from tqdm import tqdm

from ship_ice_planner import PATH_DIR, FULL_SCALE_SIM_EXP_CONFIG
from ship_ice_planner.controller.sim_dynamics import STATE_HISTORY_FILE_NAME
from ship_ice_planner.geometry.utils import Rxy
from ship_ice_planner.ship import FULL_SCALE_PSV_VERTICES
from ship_ice_planner.utils.utils import compute_path_length, tracking_error

pd.set_option('display.max_columns', None)

TRIAL_RESULTS_CSV_FILE_NAME = 'results.csv'
ALL_TRIAL_RESULTS_RAW_CSV_FILE_NAME = 'raw_results.csv'
ALL_TRIAL_RESULTS_MEAN_CSV_FILE_NAME = 'mean.csv'
BOXPLOT_MEAN_IMPULSE = 'boxplot_mean_impulse.pdf'
BOXPLOT_SHIP_KE_LOSS = 'boxplot_ship_ke_loss.pdf'
CUMULATIVE_SHIP_KE_LOSS_PLOT = 'cumulative_total_ship_ke_loss.pdf'
ALL_IMPACT_LOCS_IMPULSE_PLOT = 'all_impact_locs_impulse.png'  # need rasterized plot
MASS_PROB_DENSITY_PLOT = 'floe_mass_prob_density.pdf'


# ---- helpers to process simulation collision data ---- #
def floe_mass_hist_plot(sim_data, floe_masses, save_fig=None):
    collided_obs_mass = get_collided_obs_mass(sim_data, floe_masses)
    f, ax = plt.subplots()
    ax.hist([collided_obs_mass, floe_masses], label=['impact', 'all'])
    ax.set_title('Number of floes {}\nNumber of floes collided with ship {}'
                 .format(len(floe_masses), len(collided_obs_mass)))
    ax.set_xlabel('Mass (kg)')
    plt.legend(loc='upper right')
    if save_fig:
        f.savefig(save_fig, dpi=300)


def ke_impulse_vs_time_plot(sim_data, save_fig=None):
    ship_ke_loss = get_ship_ke_loss(sim_data)
    system_ke_loss = get_system_ke_loss(sim_data)
    delta_ke_ice = get_delta_ke_ice(sim_data)
    total_impulse = get_total_impulse(sim_data)
    total_impulse_agg_max = get_total_impulse(sim_data, aggregate='max')

    t = sim_data['time']
    f, ax = plt.subplots(2, 1, figsize=(7, 8), sharex=True)
    ax[0].plot(t, -ship_ke_loss, label=r'Ship $\Delta$KE')
    ax[0].plot(t, -np.asarray(system_ke_loss), label=r'System $\Delta$KE')
    ax[0].plot(t, delta_ke_ice, label=r'Ice $\Delta$KE')
    ax[0].set_title('Change in kinetic energy from collisions')
    ax[0].set_xlabel('Time (s)')
    ax[0].set_ylabel('J')
    ax[0].text(0.05, 0.95, f'Total Ship: {-np.sum(ship_ke_loss):.3e}\n'
                           f'Total System: {-np.sum(system_ke_loss):.3e}\n'
                           f'Total Ice: {np.sum(delta_ke_ice):.3e}',
               verticalalignment='top', transform=ax[0].transAxes,
               bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
    ax[0].legend(loc='lower left')
    ax[1].plot(t, total_impulse_agg_max)
    ax[1].set_title('Impulse')
    ax[1].set_ylabel('N s')
    ax[1].set_xlabel('Time (s)')
    ax[1].text(0.05, 0.95, f'Mean impulse: {np.mean(total_impulse):.3e}\n'
                           f'Max impulse: {np.max(total_impulse):.3e}',
               verticalalignment='top', transform=ax[1].transAxes,
               bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))

    if save_fig:
        f.savefig(save_fig, dpi=300)


def impact_locs_plot(sim_data, ship_vertices, save_fig=None):
    contact_pts = get_contact_pts(sim_data)
    f, ax = plt.subplots()
    ax.plot(*contact_pts.T, 'b.', alpha=0.1)
    ax.set_aspect('equal')
    ax.set_title('Impact locations on ship')
    ax.add_patch(patches.Polygon(ship_vertices, True, fill=False))

    if save_fig:
        f.savefig(save_fig, dpi=300)


def impact_locs_hist_plot(sim_data, save_fig=None):
    contact_pts = get_contact_pts(sim_data)
    f, ax = plt.subplots()
    ax.hist(contact_pts.T[1])
    ax.set_title('Histogram of lateral distances of impact locations on ship')
    ax.set_xlabel('Lateral distance (m)')

    if save_fig:
        f.savefig(save_fig, dpi=300)


def impact_locs_impulse_plot(sim_data: Union[dict, pd.DataFrame],
                             ship_vertices=FULL_SCALE_PSV_VERTICES,
                             max_impulse: float = None,
                             max_arrow_size=10., save_fig=None,
                             interactive_plot=True,
                             collision_time=None):
    """
    Plot the impulse vectors at the impact locations on the ship
    Can use a lot of memory if the simulation data is large!

    Can plot impact force instead of impulse by passing in the collision_time
    Where impact force = impulse / collision_time
    """
    ship_vertices = np.asarray(ship_vertices)
    current_backend = plt.get_backend()

    if not interactive_plot and current_backend != 'agg':
        matplotlib.use('Agg')  # use rasterized rendering to avoid memory issues
        print('Switched to Agg matplotlib backend!')

    if type(sim_data) == pd.DataFrame:
        fig, ax = plt.subplots()
        ax = [ax]
        sim_data = {'': [sim_data]}  # turn into dict

    else:
        planners = list(sim_data.keys())
        fig, ax = plt.subplots(len(planners), 1, sharex=True, sharey=True, figsize=(8, 6), dpi=300)
        if len(planners) == 1:
            ax = [ax]

    for axes in ax:
        axes.add_patch(patches.Polygon(ship_vertices, True, fill=False, alpha=0.1, lw=2))

    # color map to show the impulse magnitude
    cmap = plt.get_cmap('plasma')

    for axes, planner in zip(ax, sim_data):

        segs = []  # list of line segments to plot
        color_list = []
        num_collisions = 0

        for fp in sim_data[planner]:
            if type(fp) == str:
                with open(fp, 'r') as f:
                    curr_sim_data = pd.DataFrame(
                        [json.loads(item) for item in f.readlines()]
                    )[['total_impulse', 'contact_pts', 'psi']]  # only need these columns
            else:
                curr_sim_data = fp  # already a dataframe

            if max_impulse is None:
                max_impulse = np.max(get_total_impulse(curr_sim_data))

            for idx, (impulse, contact_pts) in enumerate(
                    zip(curr_sim_data['total_impulse'], curr_sim_data['contact_pts'])
            ):
                if impulse == []:
                    continue

                # multiple collisions may have occurred during a timestep
                for i in range(len(impulse)):
                    norm_impulse = np.linalg.norm(impulse[i])
                    if norm_impulse > 0:
                        num_collisions += 1

                        # get the color based on the impulse magnitude
                        color = cmap(int(norm_impulse / max_impulse * 255))
                        # get the impulse vector in the ship body frame (need inverse of rotation matrix)
                        curr_impulse = Rxy(curr_sim_data['psi'].iloc[idx]).T @ impulse[i]
                        segs.append(
                             [[contact_pts[i][0], contact_pts[i][1]],  # x, y
                              [contact_pts[i][0] - min(curr_impulse[0] / max_impulse, 1.) * max_arrow_size,
                               contact_pts[i][1] - min(curr_impulse[1] / max_impulse, 1.) * max_arrow_size]]
                        )
                        color_list.append(color)

            del curr_sim_data
            gc.collect()

        lc = LineCollection(segs, colors=color_list, linestyles='-', linewidths=1, alpha=0.7, rasterized=True)
        axes.add_collection(lc)

        # we need to set the plot limits, they will not autoscale
        axes.set_xlim(np.min(ship_vertices[:, 0]) - max_arrow_size,
                      np.max(ship_vertices[:, 0]) + max_arrow_size)
        axes.set_ylim(np.min(ship_vertices[:, 1]) - max_arrow_size,
                      np.max(ship_vertices[:, 1]) + max_arrow_size)
        axes.set_xlabel('')
        axes.set_xticks([])
        axes.set_ylabel('')
        axes.set_yticks([])
        axes.set_aspect('equal')
        axes.set_title(planner)

        del segs, color_list
        gc.collect()

        print('Number of ship-ice collisions for planner {}: {:.2e}'.format(planner, num_collisions))

    if len(ax) == 1:
        ax[0].set_title('Impulse vectors at impact locations')
    else:
        fig.suptitle('Impulse vectors at impact locations')

    # add color bar to plot
    if collision_time:
        sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin=0, vmax=max_impulse / collision_time))
        fig.colorbar(sm, ax=ax, label='Impact force (N)', shrink=0.95)
    else:
        sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin=0, vmax=max_impulse))
        fig.colorbar(sm, ax=ax, label='Impulse magnitude (N s)', shrink=0.95)
    sm.set_array([])

    if save_fig:
        fig.savefig(save_fig, dpi=300)
        plt.close(fig)


def get_contact_pts(sim_data) -> np.ndarray:
    return np.concatenate([item for item in sim_data['contact_pts'] if item != []])


def get_total_impulse(sim_data, aggregate=None) -> np.ndarray:
    assert aggregate in [None, 'mean', 'max'], 'Invalid aggregate function'

    # each item at ith index in the list corresponds to either no collision
    # or multiple collisions that occurred during ith timestep of the main simulation loop
    total_impulse = []
    for i in range(len(sim_data)):
        if sim_data['total_impulse'].iloc[i] == []:
            total_impulse.append([0])
            continue

        row = sim_data.iloc[i]
        curr_ke_ice = []
        for j in range(len(row['total_impulse'])):  # several collisions may have occurred during timestep i
            norm_impulse = np.linalg.norm(row['total_impulse'][j])
            if norm_impulse > 0:
                curr_ke_ice.append(norm_impulse)

        if len(curr_ke_ice) > 0:
            total_impulse.append(curr_ke_ice)
        else:
            total_impulse.append([0])

    # either apply mean or max for each timestep that may have multiple collisions
    if aggregate == 'mean':
        return np.asarray([np.mean(item) for item in total_impulse])
    elif aggregate == 'max':
        return np.asarray([np.max(item) for item in total_impulse])

    total_impulse = np.concatenate(total_impulse)  # flatten the list, so no longer associated with time steps

    return total_impulse[total_impulse > 0]  # ignore zero impulse collisions


def get_delta_ke_ice(sim_data) -> np.ndarray:
    delta_ke_ice = []
    for i in range(len(sim_data)):
        if sim_data['total_impulse'].iloc[i] == []:
            delta_ke_ice.append(0)
            continue

        row = sim_data.iloc[i]
        curr_ke_ice = []
        for j in range(len(row['total_impulse'])):  # several collisions may have occurred during timestep i
            if np.linalg.norm(row['total_impulse'][j]) > 0:  # ignore this collision if total impulse was zero
                curr_ke_ice.append(row['delta_ke_ice'][j])

        delta_ke_ice.append(sum(curr_ke_ice))

    return np.asarray(delta_ke_ice)


def get_system_ke_loss(sim_data) -> np.ndarray:
    system_ke_loss = []
    for i in range(len(sim_data)):
        if sim_data['total_impulse'].iloc[i] == []:
            system_ke_loss.append(0)
            continue

        row = sim_data.iloc[i]
        curr_ke_loss = []
        for j in range(len(row['total_impulse'])):  # several collisions may have occurred during timestep i
            if np.linalg.norm(row['total_impulse'][j]) > 0:  # ignore this collision if total impulse was zero
                curr_ke_loss.append(row['system_ke_loss'][j])

        system_ke_loss.append(sum(curr_ke_loss))

    return np.asarray(system_ke_loss)


def get_ship_ke_loss(sim_data) -> np.ndarray:
    # get the kinetic energy loss of the ship for each timestep
    ship_ke_loss = []
    for i in range(len(sim_data)):
        if sim_data['total_impulse'].iloc[i] == []:
            ship_ke_loss.append(0)
            continue

        row = sim_data.iloc[i]
        curr_ke_loss = []
        for j in range(len(row['total_impulse'])):  # several collisions may have occurred during timestep i
            if np.linalg.norm(row['total_impulse'][j]) > 0:  # ignore this collision if total impulse was zero
                curr_ke_loss.append(
                    row['system_ke_loss'][j] + row['delta_ke_ice'][j]
                )

        ship_ke_loss.append(sum(curr_ke_loss))

    return np.asarray(ship_ke_loss)


def get_collided_obs_mass(sim_data, floe_masses) -> np.ndarray:
    collided_ob_idx_set = set(np.concatenate([item for item in sim_data['collided_ob_idx'] if item != []]))
    collided_obs_mass = [floe_masses[idx] for idx in collided_ob_idx_set]

    return np.asarray(collided_obs_mass)


# ---- helpers to analyze controller performance ---- #
def tracking_error_plot(sim_data, path, map_shape, save_fig=None):
    f, ax = plt.subplots()
    ax.plot(path.T[0], path.T[1], '.r', label='Target path')
    ax.plot(sim_data['x'], sim_data['y'], '-b', label='Actual path')
    ax.set_aspect('equal')
    ax.legend()
    ax.set_xlabel('x (m)')
    ax.set_ylabel('y (m)')
    ax.set_xlim(0, map_shape[1])
    cross_track_error, heading_error = tracking_error(sim_data[['x', 'y', 'psi']].to_numpy(), path)
    # NOTE: computing tracking error here does not work when planner does replanning
    ax.set_title('Tracking error (cross track {:.2f} m, heading {:.2f} deg)'.format(
        cross_track_error, heading_error * 180 / np.pi)
    )

    if save_fig:
        f.savefig(save_fig, dpi=300)


def state_vs_time_plot(sim_data, save_fig=None):
    f, ax = plt.subplots(9, 1, sharex='all', figsize=(5, 12))
    t = sim_data['time']
    x, y, psi, u, v, r = sim_data[['x', 'y', 'psi', 'u', 'v', 'r']].to_numpy().T
    x_d, y_d, psi_d = np.vstack(sim_data['setpoint'].to_numpy()).T
    ax[0].plot(t, x)
    ax[0].set_title('x (m)')
    ax[1].plot(t, y)
    ax[1].set_title('y (m)')
    ax[2].plot(t, psi)
    ax[2].set_title('yaw (rad)')
    ax[3].plot(t, u)
    ax[3].set_title('surge (m/s)')
    ax[4].plot(t, v)
    ax[4].set_title('sway (m/s)')
    ax[5].plot(t, r * 180 / np.pi)
    ax[5].set_title('yaw rate (deg/s)')
    ax[6].plot(t, x_d)
    ax[6].set_title('x_d (m)')
    ax[7].plot(t, y_d)
    ax[7].set_title('y_d (m)')
    ax[8].plot(t, psi_d)
    ax[8].set_title('yaw_d (rad)')
    ax[-1].set_xlabel('time (s)')
    f.tight_layout(

    )
    if save_fig:
        f.savefig(save_fig, dpi=300)


def control_vs_time_plot(sim_data, dim_U, control_labels, save_fig=None):
    f, ax = plt.subplots(dim_U, 1, figsize=(5, 12), sharex='all')
    t = sim_data['time']
    for i, label in zip(range(dim_U), control_labels):
        ax[i].plot(t, sim_data['u_control'].apply(lambda x: x[i]), 'b', label='command')
        ax[i].set_title(label)
        if len(sim_data['u_actual'].iloc[-1]) > 0:
            ax[i].plot(t, sim_data['u_actual'].apply(lambda x: x[i]), 'r', label='actual')
    ax[-1].legend()
    ax[-1].set_xlabel('time (s)')
    f.tight_layout()

    if save_fig:
        f.savefig(save_fig, dpi=300)
# ------------------------------------------------------- #


def metric_vs_time_plot(planners, metric_scores, ylabel, save_fig):
    fig, ax = plt.subplots()

    for planner, score in zip(planners, metric_scores):
        ax.plot(*score, label=planner)

    ax.legend()
    ax.set_xlabel('time (s)')
    ax.set_ylabel(ylabel)
    fig.tight_layout()
    fig.savefig(save_fig, dpi=300)
    plt.close(fig)


def process_trials(dir_path,
                   obs_dicts=None,
                   collided_obs_results=None,
                   all_sim_data_fps=None,
                   check_existing=False) -> Union[None, pd.DataFrame]:
    """
    Process all trials for a particular ice field
    """
    if not os.path.exists(dir_path):
        return

    if check_existing and os.path.exists(os.path.join(dir_path, TRIAL_RESULTS_CSV_FILE_NAME)):
        planners = [p for p in os.listdir(dir_path) if os.path.isdir(os.path.join(dir_path, p))]
        for p in planners:
            sim_data_fp = os.path.join(dir_path, p, STATE_HISTORY_FILE_NAME)
            if all_sim_data_fps is not None:
                if p not in all_sim_data_fps:
                    all_sim_data_fps[p] = [sim_data_fp]
                else:
                    all_sim_data_fps[p].append(sim_data_fp)

        return pd.read_csv(os.path.join(dir_path, TRIAL_RESULTS_CSV_FILE_NAME), index_col=0)

    # get the planners
    planners = [p.lower() for p in os.listdir(dir_path) if os.path.isdir(os.path.join(dir_path, p))]

    # for storing results
    results = []

    # cumulative sum of total ship ke loss
    cumsum_ship_ke_loss = []

    if obs_dicts is not None:
        # get ice floe masses
        floe_masses = [obs['mass'] for obs in obs_dicts]  # this may slightly differ from get_floe_masses() in sim_utils

    for p in planners:
        planner_results = {}
        fp = os.path.join(dir_path, p)
        sim_data_fp = os.path.join(fp, STATE_HISTORY_FILE_NAME)

        with open(sim_data_fp, 'r') as f:
            sim_data = pd.DataFrame([json.loads(item) for item in f.readlines()])

        if all_sim_data_fps is not None:
            if p not in all_sim_data_fps:
                all_sim_data_fps[p] = [sim_data_fp]
            else:
                all_sim_data_fps[p].append(sim_data_fp)

        planner_results['Total Time (s)'] = sim_data.iloc[-1]['time']

        actual_path = sim_data[['x', 'y', 'psi']].to_numpy()
        planner_results['Actual Path Length (m)'] = compute_path_length(actual_path[:, :2])

        # compute total energy use from ship actuators
        planner_results['Total Energy Use (J)'] = sim_data['energy_use'].sum()

        # compute metrics from collision data
        ship_ke_loss = get_ship_ke_loss(sim_data)
        cumsum_ship_ke_loss.append([sim_data['time'],
                                    np.cumsum(ship_ke_loss)])
        system_ke_loss = get_system_ke_loss(sim_data)
        delta_ke_ice = get_delta_ke_ice(sim_data)
        total_impulse = get_total_impulse(sim_data)

        if obs_dicts is not None:
            collided_obs_mass = get_collided_obs_mass(sim_data, floe_masses)

        # cross_track_error, heading_error = tracking_error(
        #     actual_path,
        #     path=pickle.load(open(os.path.join(fp, PATH_DIR, '0.pkl'), 'rb'))['path'].T
        # )

        planner_results['Total Ship KE Loss (J)'] = ship_ke_loss.sum()
        planner_results['Total System KE Loss (J)'] = system_ke_loss.sum()
        planner_results['Total Delta KE Ice (J)'] = delta_ke_ice.sum()
        planner_results['Mean Collision Impulse (N s)'] = np.mean(total_impulse)
        planner_results['Max Collision Impulse (N s)'] = np.max(total_impulse)
        planner_results['Number of Collisions'] = len(total_impulse)

        if obs_dicts is not None:
            planner_results['Number of Collided Floes'] = len(collided_obs_mass)
            planner_results['Mean Collided Obstacle Mass (kg)'] = collided_obs_mass.mean()

        # planner_results['Tracking error - Cross track (m)'] = cross_track_error
        # planner_results['Tracking error - Heading (deg)'] = heading_error * 180 / np.pi

        # add to results
        results.append(planner_results)

        if collided_obs_results is not None:
            if p not in collided_obs_results:
                collided_obs_results[p] = [collided_obs_mass]
            else:
                collided_obs_results[p].append(collided_obs_mass)

    if 'straight' in planners:
        straight_idx = planners.index('straight')
        # normalize some metrics
        cols_to_normalize = ['Total Ship KE Loss (J)',
                             'Mean Collision Impulse (N s)',
                             'Max Collision Impulse (N s)']
        for p, res in zip(planners, results):
            for col in cols_to_normalize:
                res[col + ' Normalized'] = res[col] / results[straight_idx][col]

    df = pd.DataFrame(results, index=planners)
    df = df.reindex(sorted(df.columns), axis=1)
    df.to_csv(os.path.join(dir_path, TRIAL_RESULTS_CSV_FILE_NAME))

    # plot cumulative total ship ke loss vs time
    metric_vs_time_plot(planners,
                        cumsum_ship_ke_loss,
                        ylabel='Cumulative Total Ship KE Loss (J)',
                        save_fig=os.path.join(dir_path, CUMULATIVE_SHIP_KE_LOSS_PLOT))

    return df


def boxplots(df: pd.DataFrame, metrics: List[str], save_figs: List[str] = None):
    for idx, metric in enumerate(metrics):
        f, ax = plt.subplots(figsize=(5, 5))
        sns.boxplot(x='Concentration', y=metric, hue='Planner', ax=ax,
                    data=df, palette='Set2', showfliers=False, flierprops={'marker': 'o'},)
        ax.legend(loc='upper left')
        plt.tight_layout()

        if save_figs is not None:
            f.savefig(save_figs[idx], dpi=300)


def floe_mass_prob_density_plot(collided_obs: dict, save_fig=None):
    fig, ax = plt.subplots(figsize=(6, 6))
    ax.set_xlabel('Mass (kg)')
    ax.set_ylabel('probability density')

    for planner in collided_obs:
        mass = np.sort(np.concatenate(collided_obs[planner]))
        ax.plot(mass, norm.pdf(mass, loc=np.mean(mass), scale=np.std(mass)), label=planner)

    ax.legend()

    if save_fig:
        fig.savefig(save_fig, dpi=300)


def compute_experiment_statistics(df, dir_path) -> pd.DataFrame:
    df_mean = df.groupby(['Concentration', 'Planner']).mean()
    df_mean.drop(columns=['Ice Field Index'], inplace=True)

    # normalize some metrics
    cols_to_normalize = ['Total Ship KE Loss (J)',
                         'Mean Collision Impulse (N s)',
                         'Max Collision Impulse (N s)']
    for concentration in df['Concentration'].unique():
        for col in cols_to_normalize:
            straight_val = df_mean.loc[(concentration, 'straight'), col]
            df_mean.loc[(concentration, slice(None)), col + ' Normalized'] = \
                df_mean.loc[(concentration, slice(None)), col] / straight_val

    df_mean.to_csv(os.path.join(dir_path, ALL_TRIAL_RESULTS_MEAN_CSV_FILE_NAME))

    return df_mean


def process_experiment(root_dir,
                       exp_dict: Union[str, dict],
                       process_trials_only=False,
                       check_existing=False,
                       science_plot_style=False):
    print('Processing experiment data...')
    t0 = time.time()
    assert os.path.exists(root_dir), 'Directory {} does not exist!'.format(root_dir)

    if science_plot_style:
        import scienceplots; plt.style.use('science')

    # load the experiment configuration
    if type(exp_dict) == str:
        with open(exp_dict, 'rb') as f:
            exp_data = pickle.load(f)['exp']
    else:
        exp_data = exp_dict['exp']

    if process_trials_only:
        process_trials(dir_path=root_dir,
                       obs_dicts=None,)
        return

    # aggregate results
    all_results = []

    # get all the collided obstacles
    collided_obs_results = {}

    # get all the simulation data file paths
    all_sim_data_fps = {}

    # iterate over trials first
    for concentration in tqdm(exp_data):
        for ice_field_idx in tqdm(exp_data[concentration]):
            trial_df = process_trials(dir_path=os.path.join(root_dir, str(concentration), str(ice_field_idx)),
                                      obs_dicts=exp_data[concentration][ice_field_idx]['obstacles'],
                                      collided_obs_results=collided_obs_results,
                                      all_sim_data_fps=all_sim_data_fps,
                                      check_existing=check_existing)

            if trial_df is not None:
                trial_df['Concentration'] = concentration
                trial_df['Ice Field Index'] = ice_field_idx
                all_results.append(trial_df)

    # combine dataframes
    df = pd.concat(all_results)
    df['Planner'] = df.index
    df.index = range(len(df))
    df.reindex(
        columns=['Concentration', 'Ice Field Index', 'Planner',
                 *[col for col in df.columns if col not in ['Concentration', 'Ice Field Index', 'Planner']]]
    ).to_csv(os.path.join(root_dir, ALL_TRIAL_RESULTS_RAW_CSV_FILE_NAME))

    # compute some statistics
    compute_experiment_statistics(df, root_dir)

    # make some plots
    print('Generating plots...')
    boxplots(df, metrics=['Total Ship KE Loss (J)',
                          'Mean Collision Impulse (N s)'],
             save_figs=[os.path.join(root_dir, BOXPLOT_SHIP_KE_LOSS),
                        os.path.join(root_dir, BOXPLOT_MEAN_IMPULSE)])

    floe_mass_prob_density_plot(collided_obs_results,
                                save_fig=os.path.join(root_dir, MASS_PROB_DENSITY_PLOT))

    # this plot may take a while... and may use a lot of memory!
    # would be better to pass in `all_sim_data_fps` but it seems to use up too much memory!
    # so instead make a separate plot for each planner
    for planner in all_sim_data_fps:
        impact_locs_impulse_plot({planner: all_sim_data_fps[planner]},
                                 max_impulse=df['Max Collision Impulse (N s)'].quantile(0.99),
                                 save_fig=os.path.join(root_dir,
                                                       ALL_IMPACT_LOCS_IMPULSE_PLOT.split('.')[0] +
                                                       f'_{planner}.png'),
                                 interactive_plot=False,
                                 # collision_time=1 / 50 / 4  # option to plot impact force instead of impulse
                                 )
    # impact_locs_impulse_plot(all_sim_data_fps,
    #                          max_impulse=df['Max Collision Impulse (N s)'].max(),
    #                          save_fig=os.path.join(root_dir, ALL_IMPACT_LOCS_IMPULSE_PLOT),
    #                          interactive_plot=False)

    print('Done! Time taken: {:.2f} hours'.format((time.time() - t0) / 3600))
    plt.close('all')


if __name__ == '__main__':
    # setup arg parser
    parser = argparse.ArgumentParser(description='Process simulation experiment data')
    parser.add_argument('exp_config_file', nargs='?', type=str, help='File path to experiment config pickle file '
                                                                     'generated by generate_rand_exp.py',
                        default=FULL_SCALE_SIM_EXP_CONFIG)
    parser.add_argument('--output_dir', type=str, help='Root directory of the experiment data')
    parser.add_argument('--process_trials_only', action='store_true', help='Skip aggregating results')
    parser.add_argument('--check_existing', action='store_true', help='Skip trials that have been already processed')
    parser.add_argument('--science_plot_style', action='store_true', help='Option to set the pyplot style to "science"')
    args = parser.parse_args()

    process_experiment(args.output_dir,
                       args.exp_config_file,
                       args.process_trials_only,
                       args.check_existing,
                       args.science_plot_style)