MultiTaskBattery 1.3.0__tar.gz

multitaskbattery-1.3.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Joern Diedrichsen
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

multitaskbattery-1.3.0/MultiTaskBattery/__init__.py

@@ -0,0 +1,4 @@
+"""MultiTaskBattery: a toolbox for running multi-domain cognitive and motor
+tasks during fMRI sessions, built on PsychoPy."""
+
+__version__ = "1.3.0"

multitaskbattery-1.3.0/MultiTaskBattery/battery.py

@@ -0,0 +1,223 @@
+import os
+import numpy as np
+import pandas as pd
+import nibabel as nb
+import pooch
+from numba import njit
+from time import time
+
+
+ZENODO_RECORD = "18793343"
+ZENODO_URL = f"https://zenodo.org/api/records/{ZENODO_RECORD}/files"
+CACHE_DIR = pooch.os_cache("MTB_task_library")
+
+
+def fetch_task_library(version='V1', atlas='multiatlasHCP', structures=None):
+    """
+    Fetch task activation library from Zenodo.
+
+    Args:
+        version: Library version (e.g., 'V1')
+        atlas: Atlas name (e.g., 'multiatlasHCP')
+        structures: List of CIFTI structure names to include. If None, loads all.
+
+    Returns:
+        library_data: ndarray (n_conditions, n_measurement_channels)
+        library_info: DataFrame with condition info
+    """
+    data_fname = f"desc-tasklibrary_space-{atlas}_{version}.dscalar.nii"
+    info_fname = f"desc-tasklibrary_{version}_info.tsv"
+
+    data_path = pooch.retrieve(
+        url=f"https://zenodo.org/records/{ZENODO_RECORD}/files/{data_fname}",
+        known_hash=None,
+        path=CACHE_DIR,
+        fname=data_fname,
+    )
+    info_path = pooch.retrieve(
+        url=f"https://zenodo.org/records/{ZENODO_RECORD}/files/{info_fname}",
+        known_hash=None,
+        path=CACHE_DIR,
+        fname=info_fname,
+    )
+
+    return load_library(data_path, info_path, structures=structures)
+
+
+def load_library(data_path, info_path, structures=None):
+    """
+    Load task library_data from local files.
+
+    Args:
+        data_path: Path to library data file (.nii, .nii.gz, .dscalar.nii), must be shape (n_conditions, n_measurement_channels)
+        info_path: Path to library info file (.tsv) with 'full_code' column
+        structures: List of CIFTI structure names to include (e.g., ['CIFTI_STRUCTURE_CORTEX_LEFT',
+                    'CIFTI_STRUCTURE_CORTEX_RIGHT']). If None, loads all data.
+
+    Returns:
+        library_data: ndarray (n_conditions, n_measurement_channels)
+        library_info: DataFrame with condition info
+    """
+    img = nb.load(data_path)
+    library_data = img.get_fdata()
+    library_info = pd.read_csv(info_path, sep='\t')
+
+    if 'full_code' not in library_info.columns:
+        raise ValueError("info file must have 'full_code' column")
+
+    # Filter by structures if specified (CIFTI files only)
+    if structures is not None:
+        brain_models = img.header.get_axis(1)
+
+        # Convert short names to full CIFTI structure names
+        cifti_structures = []
+        for s in structures:
+            name = s.upper()
+            if not name.startswith('CIFTI_STRUCTURE_'):
+                name = f'CIFTI_STRUCTURE_{name}'
+            cifti_structures.append(name)
+
+        indices = []
+        for struct_name, idx, bm in brain_models.iter_structures():
+            if struct_name in cifti_structures:
+                indices.extend(range(idx.start, idx.stop))
+        if not indices:
+            available = [name for name, _, _ in brain_models.iter_structures()]
+            raise ValueError(f"No matching structures found. Available: {available}")
+        library_data = library_data[:, indices]
+    
+    print(f"Loaded {library_data.shape[0]} conditions, {library_data.shape[1]} greyordinates")
+    return library_data, library_info
+
+
+@njit
+def _compute_nit(G):
+    """
+    Computes negative inverse trace of the battery task-by-task covariance matrix
+    """
+    N = G.shape[0]
+
+    # Center the matrix
+    H = np.eye(N) - np.ones((N, N)) / N
+    G_mc = H @ G @ H
+
+    # Eigenvalues
+    l_mc = np.linalg.eigvalsh(G_mc)
+
+    # Negative inverse trace (exclude zero eigenvalue from centering, otherwise it will add an offset to all batteries that makes the numbers huge and ugly)
+    l_mc = l_mc[::-1]
+    l_mc = l_mc[l_mc > 1e-12]
+    nit = -np.sum(1 / l_mc)
+    return nit
+
+
+def evaluate_battery(library_data, library_info, battery_full_codes):
+    """
+    Evaluate a task battery by computing the negative inverse trace (NIT)
+    of the task-by-task covariance matrix.
+
+    Args:
+        library_data: Task activation data (n_conditions, n_measurement_channels)
+        library_info: DataFrame with condition info (must have 'full_code' column)
+        battery_full_codes: List of task condition codes from the 'full_code' column in the info file (e.g., ['task1_cond1', 'task2_cond2'])
+
+    Returns: NIT score (float) - higher values indicate more separable task patterns and more optimal battery
+    """
+    full_codes = library_info['full_code'].str.lower().tolist()
+    comb_names_lower = [name.lower() for name in battery_full_codes]
+
+    # Check for missing tasks
+    missing = [name for name in comb_names_lower if name not in full_codes]
+    if missing:
+        raise ValueError(f"Tasks not found in info: {missing}")
+
+    idx = np.array([full_codes.index(name) for name in comb_names_lower])
+    sub_data = library_data[idx]
+    G = sub_data @ sub_data.T
+    return _compute_nit(G)
+
+
+def get_top_batteries(library_data, library_info, n_samples,
+                    battery_size=8, n_top_batteries=10, forced_tasks=None,
+                    verbose=True):
+    """
+    Random search over task combinations to find highest NIT.
+
+    Args:
+        library_data: Task activation data (n_conditions, n_measurement_channels)
+        library_info: DataFrame with condition info (must have 'full_code' column)
+        n_samples: Total number of random batteries to test
+        battery_size: Number of task conditions in each battery
+        n_top_batteries: Number of top batteries to keep
+        forced_tasks: List of task names to include in every battery
+
+    Returns: DataFrame with columns 'rank', 'evaluation', 'battery'
+    """
+    condition_labels = library_info['full_code'].tolist()
+    n_lib_task = library_data.shape[0]
+
+    # Compute full G matrix once
+    G_full = library_data @ library_data.T
+
+    # Handle forced tasks - convert labels to indices
+    if forced_tasks is None:
+        forced_tasks = []
+    label_to_idx = {label: i for i, label in enumerate(condition_labels)}
+    forced_indices = np.array([label_to_idx[t] for t in forced_tasks], dtype=np.int64)
+    n_forced = len(forced_indices)
+    n_random = battery_size - n_forced
+
+    # Available tasks for random selection (exclude forced)
+    available_tasks = np.array([i for i in range(n_lib_task) if i not in forced_indices])
+
+    # Track top N results
+    top_results = []
+    min_nit = float('-inf')
+
+    if verbose:
+        print(f"Starting computation... (forced: {n_forced}, random: {n_random})")
+    start_time = time()
+
+    for i in range(n_samples):
+        # Generate random combination
+        random_part = np.random.choice(available_tasks, size=n_random, replace=False)
+        idx = np.sort(np.concatenate([forced_indices, random_part]))
+
+        # Compute NIT
+        G = G_full[np.ix_(idx, idx)]
+        nit = _compute_nit(G)
+
+        # Keep top N
+        if len(top_results) < n_top_batteries:
+            top_results.append((nit, tuple(idx)))
+            if len(top_results) == n_top_batteries:
+                top_results.sort()
+                min_nit = top_results[0][0]
+        elif nit > min_nit:
+            top_results[0] = (nit, tuple(idx))
+            top_results.sort()
+            min_nit = top_results[0][0]
+
+        if verbose:
+            # Progress updates every 10%
+            if (i + 1) % (n_samples // 10) == 0:
+                elapsed = time() - start_time
+                print(f"Processed {i + 1:,} / {n_samples:,} ({100 * (i + 1) / n_samples:.0f}%) | "
+                    f"Rate: {(i + 1) / elapsed:.0f}/s")
+
+    # Sort results by NIT descending
+    top_results.sort(reverse=True)
+
+    df = pd.DataFrame({
+        'rank': range(1, len(top_results) + 1),
+        'evaluation': [r[0] for r in top_results],
+        'battery': [[condition_labels[i] for i in r[1]] for r in top_results]
+    })
+    return df
+
+
+if __name__ == "__main__":
+    structures = None
+    data, info = fetch_task_library(version='V1',structures=structures)
+    # batteries = get_top_batteries(data, info, n_samples=1000000, battery_size=8, n_top_batteries=10, forced_tasks=None)
+    # print(batteries)

multitaskbattery-1.3.0/MultiTaskBattery/experiment_block.py

@@ -0,0 +1,340 @@
+# Defines the Experiment as a class
+# March 2021: First version: Ladan Shahshahani  - Maedbh King - Suzanne Witt,
+# Revised 2023: Bassel Arafat, Jorn Diedrichsen, Ince Husain
+
+import pandas as pd
+import sys
+import numpy as np
+from datetime import datetime
+from psychopy import visual, gui, event, core
+import MultiTaskBattery.utils as ut
+import MultiTaskBattery.task_blocks as tasks
+from MultiTaskBattery.ttl_clock import TTLClock
+from MultiTaskBattery.screen import Screen
+# import pylink as pl # to connect to eyelink
+
+
+class Experiment:
+    def __init__(self, const, subj_id, screen=None):
+        """    A general class with attributes common to experiments
+
+               Args:
+                    const (module):
+                        local constants.py module (see example_experiment/constants.py) as example
+                    subj_id (str):
+                        id for the subject
+                    screen (Screen, optional):
+                        an existing Screen to reuse; if None, a new one is opened
+                Returns:
+                    self (object):
+                        an instance of the Experiment class
+        """
+
+        self.exp_name   = const.exp_name
+        self.subj_id    = subj_id
+        self.run_number = 0
+        self.const = const
+        self.ttl_clock = TTLClock()
+        self.ttl_clock.ttl_button = getattr(const, 'ttl_key', 't')  # key the scanner trigger sends; differs by site
+        # open screen and display fixation cross
+        ### set the resolution of the subject screen here:
+        self.screen = screen if screen is not None else Screen(const.screen)
+
+        # quit key: works at any point of the experiment (checked on every flip / key read); set quit_key = None in constants.py to disable
+        quit_key = getattr(const, 'quit_key', 'escape')
+        if quit_key:
+            try:
+                event.globalKeys.remove(quit_key)  # set by an earlier Experiment in this session
+            except KeyError:
+                pass
+            event.globalKeys.add(key=quit_key, func=self.quit, name='quit')
+
+        # connect to the eyetracker already
+        if self.const.eye_tracker:
+            import pylink as pl
+            self.pl = pl  # keep a handle so run() can call pylink (optional import, only when eye_tracker=True)
+            # create an Eyelink class
+            ## the default ip address is 100.1.1.1.
+            ## in the ethernet settings of the laptop,
+            ## set the ip address of the EyeLink ethernet connection
+            ## to 100.1.1.2 and the subnet mask to 255.255.255.0
+            self.tk = pl.EyeLink('100.1.1.1')
+
+    def quit(self):
+        """Closes the window and quits the experiment (bound to const.quit_key)."""
+        self.screen.window.close()
+        core.quit()
+
+    def confirm_run_info(self):
+        """
+        Presents a GUI to confirm the settings for the run:
+
+        The following parameters will be set:
+        run_number      - run number
+        subj_id         - id assigned to the subject
+        ttl_flag        - should the program wait for the ttl pulse or not? For scanning THIS FLAG HAS TO BE SET TO TRUE
+
+        Args:
+        """
+        if not self.const.debug:
+            # Step 1: Create a temporary small window for user feedback
+            self.screen.start_screen(height=getattr(self.const, 'instruction_text_height', None),
+                                     wrap_width=getattr(self.const, 'instruction_wrap_width', None))
+
+            # a dialog box pops up so you can enter info
+            #Set up input box
+            inputDlg = gui.Dlg(title = f"{self.exp_name}")
+            inputDlg.addField('Enter Subject id (str):',initial = self.subj_id)      # run number (int)
+            inputDlg.addField('Enter Run Number (int):',initial = self.run_number+1)      # run number (int)
+            inputDlg.addField('Run File name (str):',initial = self.const.default_run_filename.format(f'{self.run_number + 1:02d}'))      # run number (int)
+            inputDlg.addField('Wait for TTL pulse?', initial = True) # a checkbox for ttl pulse (set it true for scanning)
+
+            inputDlg.show()
+
+            if inputDlg.OK:
+                self.subj_id        = str(inputDlg.data[0])
+                self.run_number     = int(inputDlg.data[1])
+                self.run_filename   = str(inputDlg.data[2])
+                self.wait_ttl       = bool(inputDlg.data[3])
+            else:
+                sys.exit()
+
+        else:
+            print("running in debug mode")
+            # pass on the values for your debugging with the following keywords
+            self.subj_id = 'test00'
+            self.run_number = self.run_number+1
+            self.run_filename = 'run_01.tsv'
+            self.wait_ttl = True
+
+    def init_run(self):
+        """initializing the run:
+            making sure a directory is created for the behavioral results
+            getting run file
+            Initializes all the tasks for a run
+        """
+
+        # 1. get the run file info: creates self.run_info
+        self.run_info = pd.read_csv(self.const.run_dir / self.run_filename,sep='\t')
+
+        # 2. Initialize the all tasks that we need
+        self.task_obj_list = [] # a list containing task objects in the run
+        task_table = ut.get_task_table(self.const.exp_dir, getattr(self.const, 'task_modules', None)) # Get the task table for the experiment: next to the task modules, the experiment's own table (if it exists), then the general task table
+        for t_num, task_info in self.run_info.iterrows():
+            # create a task object for the current task, reads the trial file, and append it to the list
+            t = task_table[task_table['name']== task_info.task_name]
+            task_info['code'] = t.code
+            task_info['descriptive_name'] = t.descriptive_name.iloc[0].capitalize()
+            class_name = t.task_class.iloc[0]
+            # Get the task class from the task modules
+            TaskClass = ut.get_task_class(self.const,class_name)
+            Task_obj  = TaskClass(task_info,
+                                 screen = self.screen,
+                                 ttl_clock = self.ttl_clock,
+                                 const = self.const,
+                                 subj_id = self.subj_id)
+            Task_obj.init_task()
+            self.task_obj_list.append(Task_obj)
+
+        # 3. make subject folder in data/raw/<subj_id>
+        subj_dir = self.const.data_dir / self.subj_id
+        ut.dircheck(subj_dir) # making sure the directory is created!
+        self.run_data_file = self.const.data_dir / self.subj_id / f"{self.subj_id}.tsv"
+
+    def run(self):
+        """
+        run a run of the experiment
+        """
+        print(f"Running the experiment - run {self.run_number}")
+        self.screen.fixation_cross()
+        self.ttl_clock.reset()
+        self.ttl_clock.wait_for_first_ttl(wait = self.wait_ttl)
+        run_data = []
+        run_end_timestamp = None
+
+        # Start the eyetracker
+        if self.const.eye_tracker:
+            self.start_eyetracker()
+
+        for t_num, task in enumerate(self.task_obj_list):
+            # clear keys in buffer
+            event.clearEvents()
+
+            # padded, to overwrite the TR counter line (printed with end='\r')
+            print(f"Starting block {t_num+1}: {task.name}".ljust(60))
+
+            # Take the task data from the run_info dataframe
+            r_data = self.run_info.iloc[t_num].copy()
+
+            # wait till it's time to start the task
+            r_data['real_start_time'],r_data['start_ttl'],r_data['start_ttl_time'] = self.ttl_clock.wait_until(r_data.start_time)
+
+            ## sending a message to the edf file specifying task name
+            if self.const.eye_tracker:
+                self.pl.sendMessageToFile(f"task_name: {task.name} start_track: {self.pl.currentUsec()} real start time {r_data.real_start_time} TR count {self.ttl_clock.ttl_count}")
+
+            # display the instruction text for the task. (instructions are task specific)
+            if r_data.instruction_dur > 0:
+                task.display_instructions()
+
+                # wait for a time period equal to instruction duration
+                self.ttl_clock.wait_until(r_data.start_time + r_data.instruction_dur)
+
+            # Run the task (which saves its data to the target)
+            task.start_time = self.ttl_clock.get_time()
+            r_data['acc'],r_data['rt'] = task.run()
+
+            # Add the end time of the task
+            r_data['real_end_time'] = self.ttl_clock.get_time()
+            if getattr(self.const, 'record_task_end_timestamp', False):
+                r_data['task_end_timestamp'] = datetime.now().isoformat()
+            # Run-level stamp: a single wall-clock time taken after the last task.
+            # Also written onto the last row so it survives the save_per_task path.
+            if getattr(self.const, 'record_run_end_timestamp', False) and t_num == len(self.task_obj_list) - 1:
+                run_end_timestamp = datetime.now().isoformat()
+                r_data['run_end_timestamp'] = run_end_timestamp
+            run_data.append(r_data)
+
+            # Optional: persist this task's data immediately so completed tasks survive a mid-run crash
+            if getattr(self.const, 'save_per_task', False):
+                task.save_data(self.subj_id, self.run_number)
+                one_row = pd.DataFrame([r_data])
+                one_row.insert(0, 'run_num', self.run_number)
+                ut.append_data_to_file(self.run_data_file, one_row)
+
+            self.screen.fixation_cross()
+
+            # If last task, wait until the endtime for the last task, which for imaging could be longer than the task duration
+            # Note that endtime is not used for any task but the last one
+            if t_num == len(self.task_obj_list)-1:
+                self.ttl_clock.wait_until(self.ttl_clock.get_time() + 1.5)
+                self.screen.fixation_cross()
+                self.ttl_clock.wait_until(r_data.end_time)
+
+        # Stop the eyetracker
+        if self.const.eye_tracker:
+            self.stop_eyetracker()
+            self.tk.receiveDataFile(self.tk_filename, self.tk_filename)
+
+        run_data = pd.DataFrame(run_data)
+        if not getattr(self.const, 'save_per_task', False):
+            # save the run data to the run file
+            run_data.insert(0,'run_num',[self.run_number]*len(run_data))
+            if getattr(self.const, 'record_run_end_timestamp', False) and run_end_timestamp is not None:
+                run_data['run_end_timestamp'] = run_end_timestamp
+            ut.append_data_to_file(self.run_data_file, run_data )
+
+            for task in self.task_obj_list:
+                task.save_data(self.subj_id, self.run_number)
+
+        # show the scoreboard (set show_scoreboard = False in constants.py to skip it)
+        if getattr(self.const, 'show_scoreboard', True):
+            self.display_run_feedback(run_data)
+
+
+    def display_run_feedback(self, run_data):
+        """ Displays a score board for the tasks in the task_list
+
+            Args:
+                run_data (pd.DataFrame): a dataframe containing the run data
+        """
+        # Optional: Summarize the behavioural data from different conditions of the same task
+        ignore_conditions = True
+        if ignore_conditions:
+            # Ignore Nans
+            run_data = run_data.groupby('task_name', as_index=False).agg({'acc': lambda x: x.mean(skipna=True), 'rt': lambda x: x.mean(skipna=True)})
+
+        # Define column headers
+        scores = [["Score:", "", ""], ["", "Correct", "Time"]]
+        # Add each task's data to the score text
+        task_object_names = [obj.name for obj in self.task_obj_list]
+        for r, name in enumerate(run_data.task_name.values):
+            task_idx = task_object_names.index(name)
+            if self.task_obj_list[task_idx].feedback_type.lower() != 'none':
+                task_name = f"{self.task_obj_list[task_idx].descriptive_name}"
+                acc = run_data.acc[r]
+                accuracy = f"{int(acc*100)} %" if not np.isnan(acc) else "-"
+                rt = f"{run_data.rt[r]:.2f} s" if not np.isnan(run_data.rt[r]) else "-"
+                scores.append([task_name, accuracy, rt])
+
+        # Display settings
+        # Height of the text elements: allow experiment-specific override via
+        # const.scoreboard_text_height; fall back to original default (1.3 deg).
+        height = getattr(self.const, 'scoreboard_text_height', None) or 1.3
+        gap = [8 * height / 1.3, 1.6 * height / 1.3] # Gap between entries [between_columns, between_rows]; scales with height so the whole board shrinks together
+        cols=len(scores[0])
+        rws=len(scores)
+
+        # flatten scores
+        scores_flat = [item for sublist in scores for item in sublist]
+
+        # Generate position points
+        points = np.array([[(1 + i - (cols + 1) / 2) * gap[0], (1 + j - (rws + 1) / 2) * gap[1]]
+                        for j in range(rws) for i in range(cols)])
+        # Reverse the order of the rows (to run from top to bottom)
+        points = points.reshape(rws, cols, 2)[::-1].reshape(-1, 2)
+
+        # Make the score display elements
+        elements = []
+        for Idx in range(len(scores_flat)):
+            # The task names are the first column. Anchor them on the right so that a
+            # long name grows away from the scores instead of running into them.
+            name_column = Idx % cols == 0
+            element = visual.TextStim(win=self.screen.window,
+                        text = str(scores_flat[Idx]), # Variable text
+                        font = 'Ubuntu Mono',
+                        pos = points[Idx],
+                        height = height,
+                        wrapWidth = 100,
+                        color=[-1, -1, -1],
+                        ori = 0,
+                        anchorHoriz = 'right' if name_column else 'center',
+                        alignText = 'right' if name_column else 'center',
+                        units='deg')
+            elements.append(element)
+        # Draw the scores
+        [el.draw() for el in elements]
+        self.screen.window.flip()
+
+        # Wait for a key press to continue.
+        # By default, *any* key will dismiss the scoreboard.
+        # If `continue_key` is defined in `constants.py`,
+        # only that key will be accepted (useful for tasks where participants press keys accidentally).
+        continue_key = getattr(self.const, "continue_key", None)
+        if continue_key:
+            event.waitKeys(keyList=[continue_key])
+        else:
+            event.waitKeys()
+
+
+    def start_eyetracker(self):
+        """
+        sets up a connection with the eyetracker and start recording eye position
+        """
+
+        # opening an edf file to store eye recordings
+        ## the file name should not have too many characters (<=8?)
+        ### get the run number
+        self.tk_filename = f"s_{self.run_number}.edf"
+        # self.tk_filename = f"{self.subj_id}_r{self.run_number}.edf"
+        self.tk.openDataFile(self.tk_filename)
+        # set the sampling rate for the eyetracker
+        ## you can set it to 500 or 250
+        self.tk.sendCommand("sample_rate  500")
+        # start eyetracking and send a text message to tag the start of the file
+        self.tk.startRecording(1, 1, 1, 1)
+        # pl.sendMessageToFile(f"task_name: {self.task_name} start_track: {pl.currentUsec()}")
+        return
+
+    def stop_eyetracker(self):
+        """
+        Stop recording and close the EDF file.
+
+        - Receiving the EDF file takes time and might be problematic during scanning.
+        - It may be better to copy the EDF files from the EyeLink host computer afterwards.
+        """
+        self.tk.stopRecording()
+        self.tk.closeDataFile()
+        # self.tk.receiveDataFile(self.tk_filename, self.tk_filename)
+        self.tk.close()
+        return

multitaskbattery-1.3.0/MultiTaskBattery/fetch.py

@@ -0,0 +1,77 @@
+"""Download task stimuli from Zenodo into a local cache folder.
+
+    mtb fetch                      # stimuli of every task
+    mtb fetch reading movie        # only these tasks
+
+Run it once while online (e.g. before going to an offline scanner computer).
+Set the MTB_STIMULI environment variable to keep the stimuli somewhere else.
+"""
+import argparse
+import json
+import os
+import shutil
+import urllib.request
+import zipfile
+from pathlib import Path
+import pooch
+
+# MultiTaskBattery stimuli v1.0.0 (doi:10.5281/zenodo.23749589). Each version of
+# the record has its own ID, so this pins the stimuli to the ones MTB was tested with.
+ZENODO_RECORD = '23749589'
+
+# The OS's per-user cache folder (e.g. ~/.cache/MultiTaskBattery on Linux), unless MTB_STIMULI is set
+CACHE_DIR = Path(os.environ.get('MTB_STIMULI') or pooch.os_cache('MultiTaskBattery') / 'stimuli')
+
+
+def record_files():
+    """Return {task: md5} for the zips in the Zenodo record (one zip per task)."""
+    with urllib.request.urlopen(f'https://zenodo.org/api/records/{ZENODO_RECORD}') as r:
+        files = json.load(r)['files']
+    return {f['key'][:-len('.zip')]: f['checksum'].split(':')[1] for f in files if f['key'].endswith('.zip')}
+
+
+def fetch(tasks=None):
+    """Download and unpack the stimuli of the given tasks (default: all) into CACHE_DIR.
+
+    A task already unpacked from the same zip is skipped, so fetch can be rerun
+    safely; a task unpacked from an older zip is replaced.
+
+    Args:
+        tasks (list of str): task names; tasks that use no stimuli are skipped.
+    Returns:
+        Path: the cache folder.
+    """
+    available = record_files()
+    if tasks is None:
+        tasks = list(available)
+    elif unknown := [t for t in tasks if t not in available]:
+        print(f'No stimuli to fetch for: {", ".join(unknown)} (these tasks need none, or are not in MTB)')
+        tasks = [t for t in tasks if t in available]
+    for task in tasks:
+        md5 = available[task]
+        marker = CACHE_DIR / task / '.zenodo_md5'
+        if marker.exists() and marker.read_text() == md5:
+            continue
+        print(f'Fetching {task} stimuli ...')
+        zip_path = pooch.retrieve(
+            f'https://zenodo.org/records/{ZENODO_RECORD}/files/{task}.zip?download=1',
+            known_hash=f'md5:{md5}', fname=f'{task}.zip', path=CACHE_DIR / '_zips')
+        shutil.rmtree(CACHE_DIR / task, ignore_errors=True)
+        with zipfile.ZipFile(zip_path) as zf:
+            zf.extractall(CACHE_DIR)
+        marker.write_text(md5)
+        os.remove(zip_path)
+    return CACHE_DIR
+
+
+def main():
+    parser = argparse.ArgumentParser(prog='mtb', description='MultiTaskBattery command line tools')
+    sub = parser.add_subparsers(dest='command', required=True)
+    p = sub.add_parser('fetch', help='download task stimuli from Zenodo')
+    p.add_argument('tasks', nargs='*', help='tasks to fetch (default: all)')
+    args = parser.parse_args()
+    print(f'Stimuli are in {fetch(args.tasks or None)}')
+
+
+if __name__ == '__main__':
+    main()