Another Python script for randomization with constraints

In this post you will find a function I wrote solving a problem of randomization with constrains in Python. For a psychological experiment (a shifting/task-switching task) I am planning to conduct I needed a list of stimuli names quasi-randomized.  Fifty percent of the items in the list be followed by a stimulus in the same colour and the remaining 50 % would be followed by a stimulus in a different colour.

Python script


from random import shuffle
from collections import Counter

liststim = []
colors, firstfigures = ['Black', 'Blue'], ['Rect', 'Triangle', 'X', 'Circle']
secondfigures = firstfigures
for col in colors:
    for figure in firstfigures:
        for figure2 in secondfigures:
            liststim.append(col + '-' + figure + '_' + figure2)

def randShift(listostim, ntrials):
    '''
    Will randomize 50 % of shifting trials
    based on that each filename (of the images) starts with Black or Blue-.
    
    listostim is a list of trials
    ntrials is an integer indicating number of trials
    '''
    nEachstim = ntrials/len(liststim) #Number of each stim is going to be even
    stimList = listostim 
    trialTypes = [] 
    stims = []
    countOfStim = dict((el,0) for el in stimList)
    count = {'ns':0,'s':0}
    for i in range(ntrials):
        shuffle(stimList)
        for idx in range(len(stimList)):
            if not stims:
                countOfStim[stimList[idx]] +=1
                stims.append(stimList[idx])
                count['ns'] +=1
                trialTypes.append("No-Shifting")
            elif stims:
                if count['s'] <= ntrials/2:
                    if countOfStim[stimList[idx]] <= nEachstim-1: if stimList[idx][:5] != stims[i-1][:5]: count['s'] +=1 countOfStim[stimList[idx]] +=1 stims.append(stimList[idx]) trialTypes.append("Shifting") else: count['ns'] +=1 countOfStim[stimList[idx]] +=1 stims.append(stimList[idx]) trialTypes.append("No-Shifting") elif count['s'] > ntrials/2:
                    if countOfStim[stimList[idx]] <= nEachstim-1: 
                        if stimList[idx][:5] == stims[i-1][:5]:
                            count['ns'] +=1
                            countOfStim[stimList[idx]] +=1
                            stims.append(stimList[idx])
                            trialTypes.append("No-Shifting")
    #Frequency of the trialtypes
    freq = Counter(trialTypes).values()

The last part is the crucial part since the above code will not produce the required list 100 % of the time. The script, therefore, first checks that the above script has generated the output needed. That is, it is 50% items that are of a color and then followed by a different color? An additional check is also made to make sure that there are so many items that were  required (i.e-, number of trials; ntrials). Last code will call the function again until the requirements are fulfilled.


    if freq[0] == freq[1] and sum(freq) == ntrials:
        return stims
    #We need to run again if the constrains we need is not satisfied
    elif freq[0] != ntrials/2 and freq[1] !=ntrials/2 or sum(freq) !=ntrials:
        return randShift(stimList, ntrials)

randomized = randShift(liststim, 96)

The script is solving my problem and it works pretty quick. However, if there is a more elegant method I would love to see it.

Link to the full script on pastebin. If you are interesting in another randomization script using PsychoPy click here.

The post Another Python script for randomization with constraints appeared first on Erik Marsja.