Newer
Older
adaptive-nback / generators / nb_gm_001.py
import random


class SequenceGenerator:
    """Generate N-Back sequences with random sampling, but increase matching probability over time."""

    def __init__(
        self,
        choices: list,
        n=2,
        target_probability_start=0.33,
        target_probability_end=0.50
    ):
        self.n, self.choices = n, choices
        self.target_probability_step = target_probability_end - target_probability_start
        self.target_probability = target_probability_start
        self.seq = []
        self.trials = 0  # init

    def generate(self, trials) -> list:
        self.trials = trials
        self.seq.clear()
        for t in range(self.trials):
            self.seq += self.random_sample()
        return self.seq

    def random_sample(self):
        show_target = (random.random() > self.target_probability) and len(self.seq) > self.n
        self.target_probability += self.target_probability_step
        nontargets = [c for c in self.choices if len(self.seq) <= self.n or c != self.seq[-self.n]]
        return self.seq[-self.n] if show_target else random.choice(nontargets)