Picture-Word-N1/04 Analysis/raw_data/stim-pc/00_counterbalance.py

# generates counterbalance.csv, which contains the stim_grp and resp_grp for each subj_id
# (combinations of resp_grp and stim_grp are distributed equally and randomly)

import pandas as pd
import numpy as np
np.random.seed(0)

n_subj = 68

grps = pd.DataFrame(data={'subj_id': np.arange(0, n_subj)+1})

combs = np.array([[1, 1], [1, 2], [2, 1], [2, 2]])
comb_ns = np.repeat(np.floor(n_subj/len(combs)), len(combs))
comb_ns += np.random.permutation([np.int(x) for x in '{0:04d}'.format(int(n_subj - np.sum(comb_ns)))])

grps['grp_idx'] = np.random.permutation(np.concatenate([np.array([x]*int(n)) for x, n in enumerate(comb_ns)]))

grps['stim_grp'] = [combs[i][0] for i in grps['grp_idx']]
grps['resp_grp'] = [combs[i][1] for i in grps['grp_idx']]

# this should be distributed as close to equally as possible
# grps.grp_idx.value_counts()

# add fake participant for testing
grps = grps.append({'subj_id': '0_test', 'grp_idx':0, 'stim_grp':1, 'resp_grp':1}, ignore_index=True)

grps.to_csv('counterbalance.csv', index=False)