statistics.py

"""Processes data recorded during the n back experiment and calculates statistics

This script will read all the json files that it will be handed as input. It will
compute the average for each difficulty level n for every json data file and output
it as a json file and optionally as a csv file.

On top of the it will compute the global average for each difficulty level n across
all subjects. This information will be included in the JSON file but not the csv file

...

Arguments:
    ExperimentData (str):
        pathes to the JSON files generated by JsPsych during the n-back experiment
    output (str):
        path to the output file. Must end in .json. Default: statistics.json
    csv (flag):
        if set a csv file with the same name as given on 'output' is generated on
        top of the JSON file. This will not include the global average
    heatmap-plot(str):
        if set a heatmap of the score count will be plotted and saved to the directory
        specified by this parameter
"""


import argparse
import json
import os

import lib.plot as plot
import lib.statistics as statistic


parser = argparse.ArgumentParser()
parser.add_argument('ExperimentData', nargs='*',
                     help='All the experiment json files that should be included')
parser.add_argument('--output', '-o', default='statistics.json',
                     help='The output JSON file to store the statistics')
parser.add_argument('--csv', action='store_true',
                     help="If set the script will create a csv file as well")
parser.add_argument('--heatmap-plot', '-hmp',
                     help='Where to save the heatmap plot picture if set')
arguments = parser.parse_args()

# get the absolute paths to the csv files
exp_data_files = [os.path.abspath(file) for file in arguments.ExperimentData]
out_file       = os.path.abspath(arguments.output)

assert all([
    os.path.exists(file) and
    os.path.isfile(file) and
    os.path.splitext(file)[1] == '.json'
    for file in exp_data_files
])

assert os.path.splitext(out_file)[1] == '.json'

if arguments.heatmap_plot:
    assert os.path.exists(arguments.heatmap_plot)

# statistic object
stat = statistic.Statistic()

# add all the data from the experiments
for data in exp_data_files:
    stat.add_subject_data(data)

# get the relevant metrics
average_subject_scores = stat.average_scores_all_subjects()
global_average_scores  = stat.global_average_scores()
global_std_deviation   = stat.global_std_deviation()

# create the statistics object
statistics = {}

# average subject scores
statistics['average_scores'] = []
for i in range(average_subject_scores.shape[0]):
    statistics['average_scores'].append(average_subject_scores[i,...].tolist())
# global average
statistics['global_average'] = global_average_scores.tolist()
# standard deviation
statistics['std_deviation'] = global_std_deviation.tolist()
# global score count
statistics['score_count'] = stat.global_score_count().tolist()




# write the output as json
with open(out_file, 'w') as out:
    json.dump(statistics, out)
    print('Saved output to {}'.format(out_file))

if(arguments.csv):
# write the output as csv
    csv_file = '{}.csv'.format(os.path.splitext(out_file)[0])
    with open(csv_file, 'w') as out:
        for i in range(average_subject_scores.shape[0]):
            out.write('p{:02},'.format(i+1))
            out.write(','.join([str(int(value)) for value in (average_subject_scores[i,...]*10).tolist()]))
            out.write('\n')
        print('Saved output to {}'.format(csv_file))

if arguments.heatmap_plot:
    import matplotlib.pyplot as plt

    plot.score_heatmap(stat.global_score_count())
    plt.savefig(os.path.join(arguments.heatmap_plot, 'score_heatmap.png'))