summarize.py

#!/usr/bin/env python

"""
summarize.py
Python 2.7

Summarizes all participant_X.csv files into a single CSV with the following columns:

participant_number   'Participant ID Number
mean_accuracy        'Mean Accuracy (fraction of trials correct)
mean_response_time   'Mean Reaction Time
trial_count          'Total Number of Trials Completed

* Ignores trials with empty 'correct' field.

Usage: ./summarize.py [output_filename] [directory]

"""

import csv
import os
import re
import sys

PARTICIPANT_FILE_RE = re.compile(r'^participant_(\d*)\.csv$', re.I)
DEFAULT_OUTPUT_FILE = 'summary.csv'

class Trial:
  def __init__(self, trial_number, item_stamp, correct, response_time):
    self.trial_number = trial_number
    self.item_stamp = item_stamp
    self.correct = (correct == 'True')
    self.response_time = float(response_time)

class Participant:
  def __init__(self, participant_number, trials=[]):
    self.participant_number = participant_number
    self.trials = trials

def trials(file):
  """ Returns a list of *valid* Trials found in csv file. """
  with open(file, 'rb') as csvfile:
    reader = csv.DictReader(csvfile)

    # filter out trials with an empty 'correct' field
    trials = [t for t in reader if t['correct'] != '']

    trials = [Trial(t['trial_number'], t['item_stamp'], t['correct'], t['response_time']) for t in trials]
    return trials

def participants(dir):
  """ Returns an iterator that traverses all Participants found in dir. """
  for f in os.listdir(dir):
    if os.path.isfile(f) == False:
      continue

    r = re.match(PARTICIPANT_FILE_RE, f)
    if r == None:
      continue

    participant_number = r.group(1)

    p = Participant(participant_number)
    p.trials = trials(os.path.join(dir, f))
    yield p

def participant_summary(p):
  """ Returns tuple (partcipant_number, mean_accuracy, mean_response_time, trial_count). """
  participant_number = p.participant_number
  trial_count = len(p.trials)
  if trial_count == 0:
    mean_accuracy = 'NaN'
    mean_response_time = 'NaN'
  else:
    mean_accuracy = reduce(lambda a, t: a + 1 if t.correct else a, p.trials, 0.0) / trial_count
    mean_response_time = reduce(lambda sum, t: sum + t.response_time, p.trials, 0.0) / trial_count

  return (participant_number, mean_accuracy, mean_response_time, trial_count)

def main(ofile, dir):
  """ Output a summary for each Participant in dir as a row in ofile. """ 
  with open(ofile, 'w') as csvfile:
    fieldnames = ['participant_number', 'mean_accuracy', 'mean_response_time', 'trial_count']
    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

    writer.writeheader()
    for p in participants(dir):
      summary = participant_summary(p)
      row = dict(zip(fieldnames, summary))

      writer.writerow(row)

if __name__ == "__main__":
  ofile = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_OUTPUT_FILE
  dir = sys.argv[2] if len(sys.argv) > 2 else '.'

  main(ofile, dir)