latin_square.py

import itertools as it
import numpy as np

ivs = [
    ('hit', [0, 1, 2]),
    ('layout', 'ABCDEF'),
    ('cue', 'TF'),
]

# At present, the first iv defines the number and orderof trials. Thus,
# if later ivs have fewer possible values, those values will be repeated.
# Making the first iv have the fewest values results in more participants,
# but no repetitions. As long as all iv_lens are multiples of the largest
# iv_len, all possible combinations will be included. Undefined otherwise.

def latin_hyper_square(ivs):
    iv_lens = [len(vals) for _, vals in ivs]
    first_rows = map(np.array, it.product(
        [0], *(range(iv_len) 
               for iv_len in iv_lens[1:]))
    )
    num_trials = iv_lens[0]
    matrices = []
    for row in first_rows:
        matrix = np.array([row + i for i in range(num_trials)])
        for iv, iv_len in enumerate(iv_lens):
            matrix[:, iv] %= iv_len
        matrices.append(matrix)

    return np.array(matrices)

design_matrix = latin_hyper_square(ivs)
print(design_matrix)