df = pd.DataFrame(perf_metrics)

df['removed_features'] = removed_features

df['param_counts'] = param_counts

df[['RMSEx','RMSEy','RMSEz']] = pd.DataFrame(df['Per Axis RMSE'].values.tolist(),index=df.index)

df[['nRMSEx','nRMSEy','nRMSEz']] = pd.DataFrame(df['Per Axis nRMSE'].values.tolist(),index=df.index)

df = df.drop(labels=['Per Axis RMSE','Per Axis nRMSE'],axis=1)

df_melt = pd.melt(frame=df,

id_vars = ['removed_features','param_counts'],

value_vars= quantity ,

value_name='Error Value',

var_name='Type of Error')

sns.catplot(data=df_melt,x='removed_features',y = 'Error Value',hue='Type of Error',kind='point')

'''set up some persistent tracking variables'''

remove_features = [] # list of features we are removing

metrics_list = [] # list storing dictionary of performance metrics

# feature indexes

full_state_index = np.arange(7,61)

input_state = 54

# create loss function

criterion = nn.MSELoss(reduction='sum')

# define optimization method

optimizer = opt.Adam(model.parameters(),lr=0.01)

param_count = []

param_count.append(count_params(model))

current_feature_list = np.array(qty)

# create the dataloader

dataloaders,dataset_sizes = dat.init_dataset(train_list,val_list,val_list,model_type,config_dict)

print('evaluating full model predictions...')

predictions = mdl.evaluate_model(model, dataloaders['test'], model_type=model_type,no_pbar=True)

# compute the loss statistics

print('computing full model performance metrics...')

metrics = model_eval.compute_loss_metrics(predictions, dataloaders['test'].dataset.label_array[:,1:4])

metrics_list.append(metrics)

print('Performance Summary of Full Model:')

print(metrics)

print('Running ablation study on model type:' + model_type)

for iteration in range(num_ablations):

print('-'*10)

print('Begin ablation run: {}/{}'.format(iteration+1,num_ablations))

print('-'*10)

# compute the backprop values:

gbp_data = model_eval.compute_GBP(model,dataloaders['test'],

num_state_inputs=input_state,

model_type=model_type,

no_pbar=True)

# evaluate means

df_gbp_means = model_eval.compute_and_plot_gbp(gbp_data,current_feature_list,True,suppress_plots=True)

# group by feature type and rank by value

df_gbp_means = df_gbp_means.groupby('feature').mean().sort_values(by='gbp',ascending=False).reset_index()

# get top ranking value and append to removal list

feature_to_remove = df_gbp_means.iloc[0,0]

print("removing " + feature_to_remove + "...")

remove_features.append(feature_to_remove)

# create the mask

mask = np.isin(qty,remove_features,invert=True)

# mask the full state vector in config_dict global variable

config_dict['custom_state'] = full_state_index[mask]

current_feature_list = np.array(qty)[mask] #update the current feature list

# decrease the input dimension of the model by one

input_state = input_state - 1

# redefine the models

print('redefining model with input state dims: {}'.format(input_state))

if model_type == "VS":

model = mdl.StateVisionModel(30, input_state, 3,feature_extract=feat_extract)

elif model_type == "S":

model = mdl.StateModel(input_state, 3)

# recalculate the number of parameters

param_count.append(count_params(model))

# redefine the optimizer

optimizer = opt.Adam(model.parameters(),lr=0.01)

# redefine the dataloader

dataloaders,dataset_sizes = dat.init_dataset(train_list,val_list,val_list,model_type,config_dict)

# retrain the model

model,train_history,val_history = mdl.train_model(model,

criterion, optimizer,

dataloaders, dataset_sizes,

num_epochs=50,

model_type= model_type,

weight_file=weight_file,

no_pbar=True)

print('retraining completed')

# do inference

print('evaluating model predictions...')

predictions = mdl.evaluate_model(model, dataloaders['test'], model_type=model_type,no_pbar=True)

# compute the loss statistics

print('computing performance metrics...')

metrics = model_eval.compute_loss_metrics(predictions, dataloaders['test'].dataset.label_array[:,1:4])

metrics_list.append(metrics)

print('Performance Summary:')

print(metrics)

'q1','q2','q3','q4','q5','q6','q7',

'vq1','vq2','vq3','vq4','vq5','vq6','vq7',

'tq1','tq2','tq3','tq4','tq5','tq6','tq7',

'q1d','q2d','q3d','q4d','q5d','q6d','q7d',

'tq1d','tq2d','tq3d','tq4d','tq5d','tq6d','tq7d',

'include_torque': False,

'custom_state': None,

'batch_size': 16,

'crop_list': crop_list,

'spatial_forces': force_align,