Stimulus randomization and partial interleaving

randomization.py

from collections.abc import Sequence
from random import shuffle
from itertools import chain, pairwise
from warnings import warn


def has_consecutives(x: Sequence) -> bool:
    """Check if a sequence contains at least a pair of consecutive identical items.

    Args:
        x (Sequence): Sequence to check.

    Returns:
        list: True of at least one pair of consecutive identical items has been found, False otherwise.
    """
    return any([a == b for a, b in pairwise(x)])


def shuffle_sequence(
    x: Sequence, allow_consec: bool = False, max_attempts: int = 1e4
) -> list:
    """Shuffle the items in a sequence.

    Args:
        x (Sequence): Sequence of items to shuffle.
        allow_consec (bool, optional): Are consecutive identical items be allowed. If False (default), elements are reshuffled until no consecutive identical elements are found.
        max_attempts (int, optional): Number of maximum times the sequence will be reshuffled if consecutive identical elements are found. Defaults to 10_000, after which a warning is raised and the latest reshuffled sequence is returned. Only applies if ``allow_consec`` is False.

    Returns:
        list: Reshuffled sequence.
    """
    shuffle(x)
    i = 0
    while not allow_consec and has_consecutives(x):
        shuffle(x)
        i += 1
        if i >= max_attempts:
            warn("Maximum attempts reached, possible consecutive duplicates")
            break
    return x


def chunk_sequence(x: Sequence, seq_len: Sequence[tuple[int, int]]) -> list:
    """Chunk a sequence into groups of variable length.

    The length of each chunk is determined by ``seq_len``. If ``seq_len`` is an integer, chunks have constant length ``seq_len``. If a Sequence is provided, chunk lengths are randomly chosen from a list of digits comprising the ``seq_len`` range, with the constrain that the sum of the sizes of all chunks equals the length of ``x``.

    Args:
        x (Sequence): Sequence to chunk.
        seq_len (Seuquence[tuple[int, int]]): A Sequence of tuples of length 2, where the first element indicates the size of the chunk, and the second element indicates the number of chunks with said said will be created. The resulting number of elements must be equal to ``len(x)``.

    Returns:
        list: List of lists, where each embedded list is a chunk of items from ``x``.
    """
    assert sum(a * b for a, b in seq_len) == len(x)
    sizes = []
    for le, t in seq_len:
        sizes.extend([le] * t)
    shuffle(sizes)
    seq, i = [], 0
    for s in sizes:
        seq.append(x[i : i + s])
        i += s
    return seq


def interleave_stim(
    stim: dict[str, list],
    inter_cond: str = "F",
    n_reps: Sequence = None,
    random: bool = True,
    seq_len: int = (1, 3),
) -> list[str]:
    """Make interleaved sequence of stimuli.

    Args:
        stim (dict[str, list]): A dictionary with each item indicating the stimuli group name (key) and a list of stimuli in the group (values).
        inter_cond (str, optional): Condition to interleave between the other conditions. Must be a group in ``stim``. Defaults to "F".
        n_reps (Sequence, optional): Number of repetitions of each stimulus in each stimulus group. If an integer is provided, all stimuli are repeated a ``n_reps`` amount of times. If a Sequence (e.g., list) of same length as the number of stimulus groups is provided, each element indicates the number of repetitions of each stimulus in each particular stimulus group. Defaults to 4.
        random (bool, optional): Should the stimuli in each stimulus group be shuffled before interleaving? Defaults to True.
        seq_len (Sequence[tuple[int, int]], optional): Passed to ``chunk_sequence``.

    Returns:
        list[str]: Sequence of interleaved stimuli.
    """
    if n_reps is None:
        n_reps = [4] * len(stim)
    if isinstance(n_reps, int):
        n_reps = [n_reps] * len(stim)
    else:
        assert len(n_reps) == len(stim)

    x = {k: v * ni for (k, v), ni in zip(stim.items(), n_reps)}
    intl_stim = x.pop(inter_cond)
    main_stim = list(x.values())

    if random:
        shuffle_sequence(intl_stim)
        shuffle(main_stim)
        for m in main_stim:
            shuffle_sequence(m)

    seq = list(chain.from_iterable(zip(*main_stim)))
    seq = [[s] for s in seq]
    intl_seq = chunk_sequence(intl_stim, seq_len=seq_len)
    intl_seq.append([None])

    out = [val for pair in zip(seq, intl_seq) for val in pair]
    print(len(out))
    out = [x for xs in out for x in xs]
    return out[:-1]


if __name__ == "__main__":
    stim = {
        "TF": ["tf-1", "tf-2", "tf-3", "tf-4"],
        "TN": ["tn-1", "tn-2", "tn-3", "tn-4"],
        "F": ["f-5", "f-6", "f-7", "f-8"],
    }
    seq_len = [(3, 6), (2, 5), (1, 4)]
    seq = interleave_stim(stim, inter_cond="F", n_reps=[2, 2, 8], seq_len=seq_len)
    print(seq)

    # check that each stimulus appears the correct amount of times
    for s in sorted(seq):
        print(f"{s}: {seq.count(s)}")