From 2b64a171ac5662ee4a436db7ba3d81ebf73b76b3 Mon Sep 17 00:00:00 2001
From: Ben Lonnqvist <ben.lonnqvist@epfl.ch>
Date: Wed, 29 Nov 2023 11:16:03 +0100
Subject: [PATCH 01/18] Update requirements

---
 model_tools/activations/core.py | 2 ++
 setup.py                        | 1 +
 2 files changed, 3 insertions(+)

diff --git a/model_tools/activations/core.py b/model_tools/activations/core.py
index ceec93d..e584ed7 100644
--- a/model_tools/activations/core.py
+++ b/model_tools/activations/core.py
@@ -1,5 +1,7 @@
 import copy
 import os
+import cv2
+import tempfile
 
 import functools
 import logging
diff --git a/setup.py b/setup.py
index fa23390..f41739c 100644
--- a/setup.py
+++ b/setup.py
@@ -19,6 +19,7 @@
     "keras==2.3.1",
     "protobuf<4",  # keras import fails on newer protobuf http://braintree.mit.edu:8080/job/unittest_model_tools/132/
     "scikit-learn",
+    "opencv-python",
     "result_caching @ git+https://github.com/brain-score/result_caching",
 ]
 

From 60dc0f18eeb4aed9f93b2800fe15c013853a4a68 Mon Sep 17 00:00:00 2001
From: Ben Lonnqvist <ben.lonnqvist@epfl.ch>
Date: Fri, 1 Dec 2023 14:56:34 +0100
Subject: [PATCH 02/18] Add microsaccade implementation

---
 model_tools/activations/core.py            | 177 ++++++++++++++++++---
 model_tools/brain_transformation/neural.py |  13 +-
 2 files changed, 164 insertions(+), 26 deletions(-)

diff --git a/model_tools/activations/core.py b/model_tools/activations/core.py
index e584ed7..e5332cd 100644
--- a/model_tools/activations/core.py
+++ b/model_tools/activations/core.py
@@ -35,16 +35,37 @@ def __init__(self, get_activations, preprocessing, identifier=False, batch_size=
         self._stimulus_set_hooks = {}
         self._batch_activations_hooks = {}
 
-    def __call__(self, stimuli, layers, stimuli_identifier=None):
+    def __call__(self, stimuli, layers, stimuli_identifier=None, model_requirements=None):
         """
         :param stimuli_identifier: a stimuli identifier for the stored results file. False to disable saving.
+        :param model_requirements: a dictionary containing any requirements a benchmark might have for models, e.g.
+            microsaccades for getting variable responses from non-stochastic models to the same stimuli.
+
+            model_requirements['microsaccades']: list of tuples of x and y shifts to apply to each image to model
+            microsaccades. Note that the shifts happen in pixel space of the original input image, not the preprocessed
+            image.
+            Human microsaccade amplitude varies by who you ask, an estimate might be <0.1 deg = 360 arcsec = 6arcmin.
+            The goal of microsaccades is to obtain multiple different neural activities to the same input stimulus
+            from non-stochastic models. This is to improve estimates of e.g. psychophysical functions, but also other
+            things.
+            --> Rolfs 2009 "Microsaccades: Small steps on a long way" Vision Research, Volume 49, Issue 20, 15
+            October 2009, Pages 2415-2441.
+            --> Haddad & Steinmann 1973 "The smallest voluntary saccade: Implications for fixation" Vision
+            Research Volume 13, Issue 6, June 1973, Pages 1075-1086, IN5-IN6.
+            Huge thanks to Johannes Mehrer for the implementation of microsaccades in the Brain-Score core.py and
+            neural.py files.
+
         """
         if isinstance(stimuli, StimulusSet):
-            return self.from_stimulus_set(stimulus_set=stimuli, layers=layers, stimuli_identifier=stimuli_identifier)
+            return self.from_stimulus_set(stimulus_set=stimuli, layers=layers, stimuli_identifier=stimuli_identifier,
+                                          model_requirements=model_requirements)
         else:
-            return self.from_paths(stimuli_paths=stimuli, layers=layers, stimuli_identifier=stimuli_identifier)
+            return self.from_paths(stimuli_paths=stimuli,
+                                   layers=layers,
+                                   stimuli_identifier=stimuli_identifier,
+                                   model_requirements=model_requirements)
 
-    def from_stimulus_set(self, stimulus_set, layers, stimuli_identifier=None):
+    def from_stimulus_set(self, stimulus_set, layers, stimuli_identifier=None, model_requirements=None):
         """
         :param stimuli_identifier: a stimuli identifier for the stored results file.
             False to disable saving. None to use `stimulus_set.identifier`
@@ -55,15 +76,17 @@ def from_stimulus_set(self, stimulus_set, layers, stimuli_identifier=None):
             stimulus_set = hook(stimulus_set)
         stimuli_paths = [str(stimulus_set.get_stimulus(stimulus_id)) for stimulus_id in stimulus_set['stimulus_id']]
         activations = self.from_paths(stimuli_paths=stimuli_paths, layers=layers, stimuli_identifier=stimuli_identifier)
-        activations = attach_stimulus_set_meta(activations, stimulus_set)
+        activations = attach_stimulus_set_meta(activations, stimulus_set, model_requirements=model_requirements)
         return activations
 
-    def from_paths(self, stimuli_paths, layers, stimuli_identifier=None):
+    def from_paths(self, stimuli_paths, layers, stimuli_identifier=None, model_requirements=None):
         if layers is None:
             layers = ['logits']
         if self.identifier and stimuli_identifier:
             fnc = functools.partial(self._from_paths_stored,
-                                    identifier=self.identifier, stimuli_identifier=stimuli_identifier)
+                                    identifier=self.identifier,
+                                    stimuli_identifier=stimuli_identifier,
+                                    model_requirements=model_requirements)
         else:
             self._logger.debug(f"self.identifier `{self.identifier}` or stimuli_identifier {stimuli_identifier} "
                                f"are not set, will not store")
@@ -72,21 +95,22 @@ def from_paths(self, stimuli_paths, layers, stimuli_identifier=None):
         # to be run individually, compute activations for those, and then expand the activations to all paths again.
         # This is done here, before storing, so that we only store the reduced activations.
         reduced_paths = self._reduce_paths(stimuli_paths)
-        activations = fnc(layers=layers, stimuli_paths=reduced_paths)
+        activations = fnc(layers=layers, stimuli_paths=reduced_paths, model_requirements=model_requirements)
         activations = self._expand_paths(activations, original_paths=stimuli_paths)
         return activations
 
     @store_xarray(identifier_ignore=['stimuli_paths', 'layers'], combine_fields={'layers': 'layer'})
-    def _from_paths_stored(self, identifier, layers, stimuli_identifier, stimuli_paths):
+    def _from_paths_stored(self, identifier, layers, stimuli_identifier, stimuli_paths, model_requirements):
         return self._from_paths(layers=layers, stimuli_paths=stimuli_paths)
 
-    def _from_paths(self, layers, stimuli_paths):
+    def _from_paths(self, layers, stimuli_paths, model_requirements=None):
         if len(layers) == 0:
             raise ValueError("No layers passed to retrieve activations from")
         self._logger.info('Running stimuli')
-        layer_activations = self._get_activations_batched(stimuli_paths, layers=layers, batch_size=self._batch_size)
+        layer_activations = self._get_activations_with_model_requirements(stimuli_paths, layers=layers,
+                                                                          model_requirements=model_requirements)
         self._logger.info('Packaging into assembly')
-        return self._package(layer_activations, stimuli_paths)
+        return self._package(layer_activations, stimuli_paths, model_requirements)
 
     def _reduce_paths(self, stimuli_paths):
         return list(set(stimuli_paths))
@@ -127,6 +151,12 @@ def register_stimulus_set_hook(self, hook):
         self._stimulus_set_hooks[handle.id] = hook
         return handle
 
+    def _get_activations_with_model_requirements(self, paths, layers, model_requirements):
+        if 'microsaccades' in model_requirements.keys():
+            return self._get_microsaccade_activations_batched(paths, layers=layers, batch_size=self._batch_size,
+                                                              shifts=model_requirements['microsaccades'])
+        return self._get_activations_batched(paths, layers=layers, batch_size=self._batch_size)
+
     def _get_activations_batched(self, paths, layers, batch_size):
         layer_activations = None
         for batch_start in tqdm(range(0, len(paths), batch_size), unit_scale=batch_size, desc="activations"):
@@ -144,6 +174,30 @@ def _get_activations_batched(self, paths, layers, batch_size):
 
         return layer_activations
 
+    def _get_microsaccade_activations_batched(self, paths, layers, batch_size, shifts):
+        runs_per_image = len(shifts)
+        batch_size = batch_size // runs_per_image
+        layer_activations = None
+        for batch_start in tqdm(range(0, len(paths), batch_size), unit_scale=batch_size, desc="activations"):
+            batch_end = min(batch_start + batch_size, len(paths))
+            batch_inputs = paths[batch_start:batch_end]
+            batch_activations = self._get_microsaccade_batch_activations(batch_inputs,
+                                                                         layer_names=layers,
+                                                                         batch_size=batch_size,
+                                                                         runs_per_image=runs_per_image,
+                                                                         shifts=shifts)
+            # copy to avoid handle re-enabling messing with the loop
+            for hook in self._batch_activations_hooks.copy().values():
+                batch_activations = hook(batch_activations)
+
+            if layer_activations is None:
+                layer_activations = copy.copy(batch_activations)
+            else:
+                for layer_name, layer_output in batch_activations.items():
+                    layer_activations[layer_name] = np.concatenate((layer_activations[layer_name], layer_output))
+
+        return layer_activations
+
     def _get_batch_activations(self, inputs, layer_names, batch_size):
         inputs, num_padding = self._pad(inputs, batch_size)
         preprocessed_inputs = self.preprocess(inputs)
@@ -152,22 +206,53 @@ def _get_batch_activations(self, inputs, layer_names, batch_size):
         activations = self._unpad(activations, num_padding)
         return activations
 
-    def _pad(self, batch_images, batch_size):
+    def _get_microsaccade_batch_activations(self, inputs, layer_names, batch_size, runs_per_image, shifts):
+        inputs, num_padding = self._pad(inputs, batch_size, runs_per_image)
+        preprocessed_inputs = self.translate_and_preprocess(inputs, shifts)
+        activations = self.get_activations(preprocessed_inputs,
+                                           layer_names)  # img0_run0, img1_run0, img1_run1, img1_run1
+        assert isinstance(activations, OrderedDict)
+        activations = self._unpad(activations, num_padding)
+        return activations
+
+    def translate_and_preprocess(self, images, shifts):
+        """
+        Saves translated file to disc temporarily to be able to use the exact same preprocessing
+        as if image was not shifted (i.e. as if microsaccades were not considered).
+        """
+        preprocessed_images = []
+        for image_path in images:
+            temp_files = []
+            try:
+                for shift in shifts:
+                    translated_image = self.translate(cv2.imread(image_path), shift)
+                    fp = tempfile.mkstemp(suffix=".png")[1]
+                    temp_files.append(fp)
+                    cv2.imwrite(fp, translated_image)
+                preprocessed_images.extend(self.preprocess(temp_files))
+            finally:
+                for fp in temp_files:
+                    os.remove(fp)
+
+        return preprocessed_images
+
+    def _pad(self, batch_images, batch_size, runs_per_image=1):
         num_images = len(batch_images)
-        if num_images % batch_size == 0:
+        if (num_images * runs_per_image) % batch_size == 0:
             return batch_images, 0
-        num_padding = batch_size - (num_images % batch_size)
-        padding = np.repeat(batch_images[-1:], repeats=num_padding, axis=0)
+        num_padding = batch_size * runs_per_image - (num_images % (batch_size * runs_per_image))
+        padding = np.repeat(batch_images[-runs_per_image:], repeats=num_padding, axis=0)
         return np.concatenate((batch_images, padding)), num_padding
 
     def _unpad(self, layer_activations, num_padding):
         return change_dict(layer_activations, lambda values: values[:-num_padding or None])
 
-    def _package(self, layer_activations, stimuli_paths):
+    def _package(self, layer_activations, stimuli_paths, model_requirements):
         shapes = [a.shape for a in layer_activations.values()]
         self._logger.debug(f"Activations shapes: {shapes}")
         self._logger.debug("Packaging individual layers")
-        layer_assemblies = [self._package_layer(single_layer_activations, layer=layer, stimuli_paths=stimuli_paths) for
+        layer_assemblies = [self._package_layer(single_layer_activations, layer=layer,
+                                                stimuli_paths=stimuli_paths, model_requirements=model_requirements) for
                             layer, single_layer_activations in tqdm(layer_activations.items(), desc='layer packaging')]
         # merge manually instead of using merge_data_arrays since `xarray.merge` is very slow with these large arrays
         # complication: (non)neuroid_coords are taken from the structure of layer_assemblies[0] i.e. the 1st assembly;
@@ -193,8 +278,13 @@ def _package(self, layer_activations, stimuli_paths):
                                                    dims=layer_assemblies[0].dims)
         return model_assembly
 
-    def _package_layer(self, layer_activations, layer, stimuli_paths):
-        assert layer_activations.shape[0] == len(stimuli_paths)
+    def _package_layer(self, layer_activations, layer, stimuli_paths, model_requirements):
+        if 'microsaccades' in model_requirements.keys():
+            runs_per_image = len(model_requirements['microsaccades'])
+            stimuli_paths = np.repeat(stimuli_paths, runs_per_image)
+        else:
+            runs_per_image = 1
+        assert layer_activations.shape[0] == len(stimuli_paths) * runs_per_image
         activations, flatten_indices = flatten(layer_activations, return_index=True)  # collapse for single neuroid dim
         flatten_coord_names = None
         if flatten_indices.shape[1] == 1:  # fully connected, e.g. classifier
@@ -233,6 +323,16 @@ def insert_attrs(self, wrapper):
         wrapper.register_batch_activations_hook = self.register_batch_activations_hook
         wrapper.register_stimulus_set_hook = self.register_stimulus_set_hook
 
+    @staticmethod
+    def translate(image, shift):
+        rows, cols, _ = image.shape
+        # translation matrix
+        M = np.float32([[1, 0, shift[0]], [0, 1, shift[1]]])
+
+        # Apply translation, filling new line(s) with line(s) closest to it(them).
+        translated_image = cv2.warpAffine(image, M, (cols, rows), borderMode=cv2.BORDER_REPLICATE)
+        return translated_image
+
 
 def change_dict(d, change_function, keep_name=False, multithread=False):
     if not multithread:
@@ -263,19 +363,54 @@ def lstrip_local(path):
     return path
 
 
-def attach_stimulus_set_meta(assembly, stimulus_set):
+def attach_stimulus_set_meta(assembly, stimulus_set, model_requirements):
+    if 'microsaccades' in model_requirements.keys():
+        return attach_stimulus_set_meta_with_microsaccades(assembly, stimulus_set, model_requirements)
     stimulus_paths = [str(stimulus_set.get_stimulus(stimulus_id)) for stimulus_id in stimulus_set['stimulus_id']]
     stimulus_paths = [lstrip_local(path) for path in stimulus_paths]
     assembly_paths = [lstrip_local(path) for path in assembly['stimulus_path'].values]
+
     assert (np.array(assembly_paths) == np.array(stimulus_paths)).all()
     assembly['stimulus_path'] = stimulus_set['stimulus_id'].values
     assembly = assembly.rename({'stimulus_path': 'stimulus_id'})
+
     for column in stimulus_set.columns:
         assembly[column] = 'stimulus_id', stimulus_set[column].values
     assembly = assembly.stack(presentation=('stimulus_id',))
     return assembly
 
 
+def attach_stimulus_set_meta_with_microsaccades(assembly, stimulus_set, model_requirements):
+    stimulus_paths = [str(stimulus_set.get_stimulus(stimulus_id)) for stimulus_id in stimulus_set['stimulus_id']]
+    stimulus_paths = [lstrip_local(path) for path in stimulus_paths]
+    assembly_paths = [lstrip_local(path) for path in assembly['stimulus_path'].values]
+
+    replication_factor = len(model_requirements['microsaccades']) if model_requirements['microsaccades'] else 1
+    repeated_stimulus_paths = np.repeat(stimulus_paths, replication_factor)
+    assert (np.array(assembly_paths) == np.array(repeated_stimulus_paths)).all()
+
+    repeated_stimulus_ids = np.repeat(stimulus_set['stimulus_id'].values, replication_factor)
+    assembly['stimulus_path'] = repeated_stimulus_ids
+    assembly = assembly.rename({'stimulus_path': 'stimulus_id'})
+
+    for column in stimulus_set.columns:
+        assembly[column] = 'stimulus_id', np.repeat(stimulus_set[column].values, replication_factor)
+
+    # Handle shifts
+    repeated_shifts = np.tile(np.array(model_requirements['microsaccades']), (len(stimulus_set['stimulus_id']), 1))
+    assembly.coords['shift_x'] = ('stimulus_id', repeated_shifts[:, 0])
+    assembly.coords['shift_y'] = ('stimulus_id', repeated_shifts[:, 1])
+
+    # Preserve existing coordinates and dimensions
+    for coord in assembly.coords:
+        assembly[coord] = ('stimulus_id', np.repeat(assembly[coord].data, replication_factor))
+
+    # Set MultiIndex
+    assembly = assembly.set_index(presentation=['stimulus_id'] + list(assembly.coords))
+
+    return assembly
+
+
 class HookHandle:
     next_id = 0
 
diff --git a/model_tools/brain_transformation/neural.py b/model_tools/brain_transformation/neural.py
index f7bfa9d..97ac3af 100644
--- a/model_tools/brain_transformation/neural.py
+++ b/model_tools/brain_transformation/neural.py
@@ -1,6 +1,6 @@
 import logging
 from tqdm import tqdm
-from typing import Optional, Union
+from typing import Optional, Union, Dict, List
 
 from brainscore.metrics import Score
 from brainscore.model_interface import BrainModel
@@ -23,7 +23,7 @@ def __init__(self, identifier, activations_model, region_layer_map, visual_degre
     def identifier(self):
         return self._identifier
 
-    def look_at(self, stimuli, number_of_trials=1):
+    def look_at(self, stimuli, number_of_trials=1, model_requirements: Optional[Dict[str, List]] = None):
         layer_regions = {}
         for region in self.recorded_regions:
             layers = self.region_layer_map[region]
@@ -31,12 +31,15 @@ def look_at(self, stimuli, number_of_trials=1):
             for layer in layers:
                 assert layer not in layer_regions, f"layer {layer} has already been assigned for {layer_regions[layer]}"
                 layer_regions[layer] = region
-        activations = self.run_activations(stimuli, layers=list(layer_regions.keys()), number_of_trials=number_of_trials)
+        activations = self.run_activations(stimuli,
+                                           layers=list(layer_regions.keys()),
+                                           number_of_trials=number_of_trials,
+                                           model_requirements=model_requirements)
         activations['region'] = 'neuroid', [layer_regions[layer] for layer in activations['layer'].values]
         return activations
 
-    def run_activations(self, stimuli, layers, number_of_trials=1):
-        activations = self.activations_model(stimuli, layers=layers)
+    def run_activations(self, stimuli, layers, number_of_trials=1, model_requirements=None):
+        activations = self.activations_model(stimuli, layers=layers, model_requirements=model_requirements)
         return activations
 
     def start_task(self, task):

From 0498196b897108f337ba5230fbe2938e43ddf0f8 Mon Sep 17 00:00:00 2001
From: Ben Lonnqvist <ben.lonnqvist@epfl.ch>
Date: Tue, 5 Dec 2023 11:42:07 +0100
Subject: [PATCH 03/18] fix bug with None dict indexing

---
 model_tools/activations/core.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/model_tools/activations/core.py b/model_tools/activations/core.py
index e5332cd..a2ed930 100644
--- a/model_tools/activations/core.py
+++ b/model_tools/activations/core.py
@@ -152,7 +152,7 @@ def register_stimulus_set_hook(self, hook):
         return handle
 
     def _get_activations_with_model_requirements(self, paths, layers, model_requirements):
-        if 'microsaccades' in model_requirements.keys():
+        if model_requirements is not None and 'microsaccades' in model_requirements.keys():
             return self._get_microsaccade_activations_batched(paths, layers=layers, batch_size=self._batch_size,
                                                               shifts=model_requirements['microsaccades'])
         return self._get_activations_batched(paths, layers=layers, batch_size=self._batch_size)
@@ -279,7 +279,7 @@ def _package(self, layer_activations, stimuli_paths, model_requirements):
         return model_assembly
 
     def _package_layer(self, layer_activations, layer, stimuli_paths, model_requirements):
-        if 'microsaccades' in model_requirements.keys():
+        if model_requirements is not None and 'microsaccades' in model_requirements.keys():
             runs_per_image = len(model_requirements['microsaccades'])
             stimuli_paths = np.repeat(stimuli_paths, runs_per_image)
         else:
@@ -364,7 +364,7 @@ def lstrip_local(path):
 
 
 def attach_stimulus_set_meta(assembly, stimulus_set, model_requirements):
-    if 'microsaccades' in model_requirements.keys():
+    if model_requirements is not None and 'microsaccades' in model_requirements.keys():
         return attach_stimulus_set_meta_with_microsaccades(assembly, stimulus_set, model_requirements)
     stimulus_paths = [str(stimulus_set.get_stimulus(stimulus_id)) for stimulus_id in stimulus_set['stimulus_id']]
     stimulus_paths = [lstrip_local(path) for path in stimulus_paths]

From 47c94230f879594cb3bb843c15e750381f0b8805 Mon Sep 17 00:00:00 2001
From: Ben Lonnqvist <ben.lonnqvist@epfl.ch>
Date: Tue, 5 Dec 2023 14:22:47 +0100
Subject: [PATCH 04/18] Add tests for activations and temporary file handling

---
 model_tools/activations/core.py    |  1 +
 tests/activations/test___init__.py | 55 ++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+)

diff --git a/model_tools/activations/core.py b/model_tools/activations/core.py
index a2ed930..f4869f3 100644
--- a/model_tools/activations/core.py
+++ b/model_tools/activations/core.py
@@ -153,6 +153,7 @@ def register_stimulus_set_hook(self, hook):
 
     def _get_activations_with_model_requirements(self, paths, layers, model_requirements):
         if model_requirements is not None and 'microsaccades' in model_requirements.keys():
+            assert len(model_requirements['microsaccades'] > 0)
             return self._get_microsaccade_activations_batched(paths, layers=layers, batch_size=self._batch_size,
                                                               shifts=model_requirements['microsaccades'])
         return self._get_activations_batched(paths, layers=layers, batch_size=self._batch_size)
diff --git a/tests/activations/test___init__.py b/tests/activations/test___init__.py
index f0a8abe..c805aef 100644
--- a/tests/activations/test___init__.py
+++ b/tests/activations/test___init__.py
@@ -5,6 +5,7 @@
 import xarray as xr
 from pathlib import Path
 
+
 from brainio.stimuli import StimulusSet
 from model_tools.activations import KerasWrapper, PytorchWrapper, TensorflowSlimWrapper
 from model_tools.activations.core import flatten
@@ -189,6 +190,60 @@ def test_from_image_path(model_ctr, layers, image_name, pca_components, logits):
     return activations
 
 
+@pytest.mark.parametrize("image_name", ['rgb.jpg', 'grayscale.png', 'grayscale2.jpg', 'grayscale_alpha.png',
+                                        'palletized.png'])
+@pytest.mark.parametrize(["model_ctr", "layers"], models_layers)
+@pytest.mark.parametrize("shifts", [[(2, 2), (0, 0), (1, -1), (0, 1)], [(0, 1), (2, 3)]])
+def test_microsaccades_from_image_path(model_ctr, layers, image_name, shifts):
+    stimulus_paths = [os.path.join(os.path.dirname(__file__), image_name)]  # Replace with a suitable test image
+    activations_extractor = model_ctr()
+
+    model_requirements = {'microsaccades': shifts}
+    activations = activations_extractor(stimuli=stimulus_paths, layers=layers, model_requirements=model_requirements)
+
+    assert activations is not None
+    assert 'shift_x' in activations.coords and 'shift_y' in activations.coords
+    assert len(activations['shift_x']) == len(shifts) * len(stimulus_paths)
+    assert len(activations['shift_y']) == len(shifts) * len(stimulus_paths)
+
+
+@pytest.mark.parametrize("image_name", ['rgb.jpg', 'grayscale.png', 'grayscale2.jpg', 'grayscale_alpha.png',
+                                        'palletized.png'])
+@pytest.mark.parametrize(["model_ctr", "layers"], models_layers)
+@pytest.mark.parametrize("model_requirements", [None, {'microsaccades': [(0, 0), (1, -1)]}])
+def test_model_requirements(model_ctr, layers, image_name, model_requirements):
+    stimulus_paths = [os.path.join(os.path.dirname(__file__), image_name)]  # Replace with a suitable test image
+    activations_extractor = model_ctr()
+
+    activations_with_req = activations_extractor(stimuli=stimulus_paths, layers=layers,
+                                                 model_requirements=model_requirements)
+    activations_without_req = activations_extractor(stimuli=stimulus_paths, layers=layers,
+                                                    model_requirements=None)
+
+    assert activations_with_req is not None
+    assert activations_without_req is not None
+    if model_requirements:
+        assert len(activations_with_req['neuroid']) > len(activations_without_req['neuroid'])
+    else:
+        assert len(activations_with_req['neuroid']) == len(activations_without_req['neuroid'])
+
+
+@pytest.mark.parametrize("image_name", ['rgb.jpg', 'grayscale.png', 'grayscale2.jpg', 'grayscale_alpha.png',
+                                        'palletized.png'])
+@pytest.mark.parametrize(["model_ctr", "layers"], models_layers)
+def test_temporary_file_handling(model_ctr, layers, image_name):
+    import tempfile
+    stimulus_paths = [os.path.join(os.path.dirname(__file__), image_name)]  # Replace with a suitable test image
+    activations_extractor = model_ctr()
+    model_requirements = {'microsaccades': [(2, 2)]}
+
+    activations = activations_extractor(stimuli=stimulus_paths, layers=layers, model_requirements=model_requirements)
+    temp_files = [f for f in os.listdir(tempfile.gettempdir()) if f.startswith('temp') and f.endswith('.png')]
+
+    assert activations is not None
+    assert len(temp_files) == 0  # No temporary files should remain
+
+
 def _build_stimulus_set(image_names):
     stimulus_set = StimulusSet([{'stimulus_id': image_name, 'some_meta': image_name[::-1]}
                                 for image_name in image_names])

From 642889815c4422bee657f8e52a950792865e87e5 Mon Sep 17 00:00:00 2001
From: Ben Lonnqvist <ben.lonnqvist@epfl.ch>
Date: Tue, 5 Dec 2023 14:26:20 +0100
Subject: [PATCH 05/18] Improve description of model_requirements

---
 model_tools/activations/core.py            |  3 +++
 model_tools/brain_transformation/neural.py | 21 +++++++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/model_tools/activations/core.py b/model_tools/activations/core.py
index f4869f3..fea649e 100644
--- a/model_tools/activations/core.py
+++ b/model_tools/activations/core.py
@@ -48,6 +48,9 @@ def __call__(self, stimuli, layers, stimuli_identifier=None, model_requirements=
             The goal of microsaccades is to obtain multiple different neural activities to the same input stimulus
             from non-stochastic models. This is to improve estimates of e.g. psychophysical functions, but also other
             things.
+            Example usage:
+                model_requirements = {'microsaccades': [(0, 0), (0, 1), (1, 0), (1, 1)]}
+            More information:
             --> Rolfs 2009 "Microsaccades: Small steps on a long way" Vision Research, Volume 49, Issue 20, 15
             October 2009, Pages 2415-2441.
             --> Haddad & Steinmann 1973 "The smallest voluntary saccade: Implications for fixation" Vision
diff --git a/model_tools/brain_transformation/neural.py b/model_tools/brain_transformation/neural.py
index 97ac3af..181c77d 100644
--- a/model_tools/brain_transformation/neural.py
+++ b/model_tools/brain_transformation/neural.py
@@ -24,6 +24,27 @@ def identifier(self):
         return self._identifier
 
     def look_at(self, stimuli, number_of_trials=1, model_requirements: Optional[Dict[str, List]] = None):
+        """
+        :param model_requirements: a dictionary containing any requirements a benchmark might have for models, e.g.
+            microsaccades for getting variable responses from non-stochastic models to the same stimuli.
+
+            model_requirements['microsaccades']: list of tuples of x and y shifts to apply to each image to model
+            microsaccades. Note that the shifts happen in pixel space of the original input image, not the preprocessed
+            image.
+            Human microsaccade amplitude varies by who you ask, an estimate might be <0.1 deg = 360 arcsec = 6arcmin.
+            The goal of microsaccades is to obtain multiple different neural activities to the same input stimulus
+            from non-stochastic models. This is to improve estimates of e.g. psychophysical functions, but also other
+            things.
+            Example usage:
+                model_requirements = {'microsaccades': [(0, 0), (0, 1), (1, 0), (1, 1)]}
+            More information:
+            --> Rolfs 2009 "Microsaccades: Small steps on a long way" Vision Research, Volume 49, Issue 20, 15
+            October 2009, Pages 2415-2441.
+            --> Haddad & Steinmann 1973 "The smallest voluntary saccade: Implications for fixation" Vision
+            Research Volume 13, Issue 6, June 1973, Pages 1075-1086, IN5-IN6.
+            Huge thanks to Johannes Mehrer for the implementation of microsaccades in the Brain-Score core.py and
+            neural.py files.
+        """
         layer_regions = {}
         for region in self.recorded_regions:
             layers = self.region_layer_map[region]

From 3b3e9796c6e540d4a0aee382035d3377adc5d5a1 Mon Sep 17 00:00:00 2001
From: Ben Lonnqvist <ben.lonnqvist@epfl.ch>
Date: Tue, 5 Dec 2023 14:37:11 +0100
Subject: [PATCH 06/18] Improve descriptions and minor cleanup

---
 model_tools/activations/core.py            | 6 +++++-
 model_tools/brain_transformation/neural.py | 3 ++-
 tests/activations/test___init__.py         | 8 ++++----
 3 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/model_tools/activations/core.py b/model_tools/activations/core.py
index fea649e..ca42cb9 100644
--- a/model_tools/activations/core.py
+++ b/model_tools/activations/core.py
@@ -47,7 +47,8 @@ def __call__(self, stimuli, layers, stimuli_identifier=None, model_requirements=
             Human microsaccade amplitude varies by who you ask, an estimate might be <0.1 deg = 360 arcsec = 6arcmin.
             The goal of microsaccades is to obtain multiple different neural activities to the same input stimulus
             from non-stochastic models. This is to improve estimates of e.g. psychophysical functions, but also other
-            things.
+            things. Note that microsaccades are also applied to stochastic models to make them comparable within-
+            benchmark to non-stochastic models.
             Example usage:
                 model_requirements = {'microsaccades': [(0, 0), (0, 1), (1, 0), (1, 1)]}
             More information:
@@ -179,6 +180,9 @@ def _get_activations_batched(self, paths, layers, batch_size):
         return layer_activations
 
     def _get_microsaccade_activations_batched(self, paths, layers, batch_size, shifts):
+        """
+        :param shifts: a list of tuples containing the pixel shifts to apply to the paths.
+        """
         runs_per_image = len(shifts)
         batch_size = batch_size // runs_per_image
         layer_activations = None
diff --git a/model_tools/brain_transformation/neural.py b/model_tools/brain_transformation/neural.py
index 181c77d..e17e098 100644
--- a/model_tools/brain_transformation/neural.py
+++ b/model_tools/brain_transformation/neural.py
@@ -34,7 +34,8 @@ def look_at(self, stimuli, number_of_trials=1, model_requirements: Optional[Dict
             Human microsaccade amplitude varies by who you ask, an estimate might be <0.1 deg = 360 arcsec = 6arcmin.
             The goal of microsaccades is to obtain multiple different neural activities to the same input stimulus
             from non-stochastic models. This is to improve estimates of e.g. psychophysical functions, but also other
-            things.
+            things. Note that microsaccades are also applied to stochastic models to make them comparable within-
+            benchmark to non-stochastic models.
             Example usage:
                 model_requirements = {'microsaccades': [(0, 0), (0, 1), (1, 0), (1, 1)]}
             More information:
diff --git a/tests/activations/test___init__.py b/tests/activations/test___init__.py
index c805aef..b7e2a5f 100644
--- a/tests/activations/test___init__.py
+++ b/tests/activations/test___init__.py
@@ -195,7 +195,7 @@ def test_from_image_path(model_ctr, layers, image_name, pca_components, logits):
 @pytest.mark.parametrize(["model_ctr", "layers"], models_layers)
 @pytest.mark.parametrize("shifts", [[(2, 2), (0, 0), (1, -1), (0, 1)], [(0, 1), (2, 3)]])
 def test_microsaccades_from_image_path(model_ctr, layers, image_name, shifts):
-    stimulus_paths = [os.path.join(os.path.dirname(__file__), image_name)]  # Replace with a suitable test image
+    stimulus_paths = [os.path.join(os.path.dirname(__file__), image_name)]
     activations_extractor = model_ctr()
 
     model_requirements = {'microsaccades': shifts}
@@ -212,7 +212,7 @@ def test_microsaccades_from_image_path(model_ctr, layers, image_name, shifts):
 @pytest.mark.parametrize(["model_ctr", "layers"], models_layers)
 @pytest.mark.parametrize("model_requirements", [None, {'microsaccades': [(0, 0), (1, -1)]}])
 def test_model_requirements(model_ctr, layers, image_name, model_requirements):
-    stimulus_paths = [os.path.join(os.path.dirname(__file__), image_name)]  # Replace with a suitable test image
+    stimulus_paths = [os.path.join(os.path.dirname(__file__), image_name)]
     activations_extractor = model_ctr()
 
     activations_with_req = activations_extractor(stimuli=stimulus_paths, layers=layers,
@@ -233,7 +233,7 @@ def test_model_requirements(model_ctr, layers, image_name, model_requirements):
 @pytest.mark.parametrize(["model_ctr", "layers"], models_layers)
 def test_temporary_file_handling(model_ctr, layers, image_name):
     import tempfile
-    stimulus_paths = [os.path.join(os.path.dirname(__file__), image_name)]  # Replace with a suitable test image
+    stimulus_paths = [os.path.join(os.path.dirname(__file__), image_name)]
     activations_extractor = model_ctr()
     model_requirements = {'microsaccades': [(2, 2)]}
 
@@ -241,7 +241,7 @@ def test_temporary_file_handling(model_ctr, layers, image_name):
     temp_files = [f for f in os.listdir(tempfile.gettempdir()) if f.startswith('temp') and f.endswith('.png')]
 
     assert activations is not None
-    assert len(temp_files) == 0  # No temporary files should remain
+    assert len(temp_files) == 0
 
 
 def _build_stimulus_set(image_names):

From fe1d3f95909b6918aef954a4bd6a41a7f59d0877 Mon Sep 17 00:00:00 2001
From: Ben Lonnqvist <ben.lonnqvist@epfl.ch>
Date: Tue, 5 Dec 2023 16:08:42 +0100
Subject: [PATCH 07/18] fix bug with a parenthesis

---
 model_tools/activations/core.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/model_tools/activations/core.py b/model_tools/activations/core.py
index ca42cb9..316d1f4 100644
--- a/model_tools/activations/core.py
+++ b/model_tools/activations/core.py
@@ -157,7 +157,7 @@ def register_stimulus_set_hook(self, hook):
 
     def _get_activations_with_model_requirements(self, paths, layers, model_requirements):
         if model_requirements is not None and 'microsaccades' in model_requirements.keys():
-            assert len(model_requirements['microsaccades'] > 0)
+            assert len(model_requirements['microsaccades']) > 0
             return self._get_microsaccade_activations_batched(paths, layers=layers, batch_size=self._batch_size,
                                                               shifts=model_requirements['microsaccades'])
         return self._get_activations_batched(paths, layers=layers, batch_size=self._batch_size)

From 9173d625c93ccf1b18f1a2d95032b5c0b3fb3be3 Mon Sep 17 00:00:00 2001
From: Ben Lonnqvist <ben.lonnqvist@epfl.ch>
Date: Tue, 5 Dec 2023 16:31:34 +0100
Subject: [PATCH 08/18] Fix issue with batch_size < runs_per_image (thanks
 David Tang)

---
 model_tools/activations/core.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/model_tools/activations/core.py b/model_tools/activations/core.py
index 316d1f4..ccc5066 100644
--- a/model_tools/activations/core.py
+++ b/model_tools/activations/core.py
@@ -184,7 +184,7 @@ def _get_microsaccade_activations_batched(self, paths, layers, batch_size, shift
         :param shifts: a list of tuples containing the pixel shifts to apply to the paths.
         """
         runs_per_image = len(shifts)
-        batch_size = batch_size // runs_per_image
+        batch_size = max(batch_size // runs_per_image, 1)
         layer_activations = None
         for batch_start in tqdm(range(0, len(paths), batch_size), unit_scale=batch_size, desc="activations"):
             batch_end = min(batch_start + batch_size, len(paths))

From bafc04ae266f16d184a9883ffb46ff33dec2d639 Mon Sep 17 00:00:00 2001
From: Ben Lonnqvist <ben.lonnqvist@epfl.ch>
Date: Wed, 6 Dec 2023 11:25:21 +0100
Subject: [PATCH 09/18] Move np.repeat to the correct place

---
 model_tools/activations/core.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/model_tools/activations/core.py b/model_tools/activations/core.py
index ccc5066..fa7f0da 100644
--- a/model_tools/activations/core.py
+++ b/model_tools/activations/core.py
@@ -289,10 +289,10 @@ def _package(self, layer_activations, stimuli_paths, model_requirements):
     def _package_layer(self, layer_activations, layer, stimuli_paths, model_requirements):
         if model_requirements is not None and 'microsaccades' in model_requirements.keys():
             runs_per_image = len(model_requirements['microsaccades'])
-            stimuli_paths = np.repeat(stimuli_paths, runs_per_image)
         else:
             runs_per_image = 1
         assert layer_activations.shape[0] == len(stimuli_paths) * runs_per_image
+        stimuli_paths = np.repeat(stimuli_paths, runs_per_image)
         activations, flatten_indices = flatten(layer_activations, return_index=True)  # collapse for single neuroid dim
         flatten_coord_names = None
         if flatten_indices.shape[1] == 1:  # fully connected, e.g. classifier

From b77e49b3ffc8db46bf9ca0af4d8a4ead088e621d Mon Sep 17 00:00:00 2001
From: Ben Lonnqvist <ben.lonnqvist@epfl.ch>
Date: Fri, 8 Dec 2023 16:13:21 +0100
Subject: [PATCH 10/18] Fix batch size bug

---
 model_tools/activations/core.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/model_tools/activations/core.py b/model_tools/activations/core.py
index fa7f0da..7616dab 100644
--- a/model_tools/activations/core.py
+++ b/model_tools/activations/core.py
@@ -246,7 +246,7 @@ def translate_and_preprocess(self, images, shifts):
 
     def _pad(self, batch_images, batch_size, runs_per_image=1):
         num_images = len(batch_images)
-        if (num_images * runs_per_image) % batch_size == 0:
+        if ((num_images * runs_per_image) % batch_size == 0) or num_images < batch_size:
             return batch_images, 0
         num_padding = batch_size * runs_per_image - (num_images % (batch_size * runs_per_image))
         padding = np.repeat(batch_images[-runs_per_image:], repeats=num_padding, axis=0)

From cc2516bfd75ed5a32267842c1952b85963aa54b1 Mon Sep 17 00:00:00 2001
From: Ben Lonnqvist <ben.lonnqvist@epfl.ch>
Date: Tue, 12 Dec 2023 11:22:29 +0100
Subject: [PATCH 11/18] rewrite logic for batching to avoid convoluted padding
 and improve speed (thanks @YingtianDT)

---
 model_tools/activations/core.py | 88 ++++++++++++++++-----------------
 1 file changed, 42 insertions(+), 46 deletions(-)

diff --git a/model_tools/activations/core.py b/model_tools/activations/core.py
index 7616dab..aaba534 100644
--- a/model_tools/activations/core.py
+++ b/model_tools/activations/core.py
@@ -183,28 +183,36 @@ def _get_microsaccade_activations_batched(self, paths, layers, batch_size, shift
         """
         :param shifts: a list of tuples containing the pixel shifts to apply to the paths.
         """
-        runs_per_image = len(shifts)
-        batch_size = max(batch_size // runs_per_image, 1)
-        layer_activations = None
+        layer_activations = OrderedDict()
         for batch_start in tqdm(range(0, len(paths), batch_size), unit_scale=batch_size, desc="activations"):
             batch_end = min(batch_start + batch_size, len(paths))
             batch_inputs = paths[batch_start:batch_end]
-            batch_activations = self._get_microsaccade_batch_activations(batch_inputs,
-                                                                         layer_names=layers,
-                                                                         batch_size=batch_size,
-                                                                         runs_per_image=runs_per_image,
-                                                                         shifts=shifts)
-            # copy to avoid handle re-enabling messing with the loop
+
+            batch_activations = OrderedDict()
+            # compute activations on the entire batch one shift at a time
+            for shift in shifts:
+                assert type(shift) == tuple
+                shifted_batch = self.translate_and_preprocess(batch_inputs, shift)
+                activations = self.get_activations(shifted_batch, layers)
+                for layer_name, layer_output in activations.items():
+                    batch_activations.setdefault(layer_name, []).append(layer_output)
+
+            # concatenate all shifts into this batch
+            for layer_name, layer_outputs in batch_activations.items():
+                batch_activations[layer_name] = np.concatenate(layer_outputs)
+
             for hook in self._batch_activations_hooks.copy().values():
                 batch_activations = hook(batch_activations)
 
-            if layer_activations is None:
-                layer_activations = copy.copy(batch_activations)
-            else:
-                for layer_name, layer_output in batch_activations.items():
-                    layer_activations[layer_name] = np.concatenate((layer_activations[layer_name], layer_output))
+            # add this batch to layer_activations
+            for layer_name, layer_output in batch_activations.items():
+                layer_activations.setdefault(layer_name, []).append(layer_output)
 
-        return layer_activations
+        # fast concat all batches
+        for layer_name, layer_outputs in layer_activations.items():
+            layer_activations[layer_name] = np.concatenate(layer_outputs)
+
+        return layer_activations  # this is all batches
 
     def _get_batch_activations(self, inputs, layer_names, batch_size):
         inputs, num_padding = self._pad(inputs, batch_size)
@@ -214,42 +222,30 @@ def _get_batch_activations(self, inputs, layer_names, batch_size):
         activations = self._unpad(activations, num_padding)
         return activations
 
-    def _get_microsaccade_batch_activations(self, inputs, layer_names, batch_size, runs_per_image, shifts):
-        inputs, num_padding = self._pad(inputs, batch_size, runs_per_image)
-        preprocessed_inputs = self.translate_and_preprocess(inputs, shifts)
-        activations = self.get_activations(preprocessed_inputs,
-                                           layer_names)  # img0_run0, img1_run0, img1_run1, img1_run1
-        assert isinstance(activations, OrderedDict)
-        activations = self._unpad(activations, num_padding)
-        return activations
-
-    def translate_and_preprocess(self, images, shifts):
-        """
-        Saves translated file to disc temporarily to be able to use the exact same preprocessing
-        as if image was not shifted (i.e. as if microsaccades were not considered).
-        """
-        preprocessed_images = []
+    def translate_and_preprocess(self, images, shift):
+        assert type(images) == list
+        temp_file_paths = []
         for image_path in images:
-            temp_files = []
-            try:
-                for shift in shifts:
-                    translated_image = self.translate(cv2.imread(image_path), shift)
-                    fp = tempfile.mkstemp(suffix=".png")[1]
-                    temp_files.append(fp)
-                    cv2.imwrite(fp, translated_image)
-                preprocessed_images.extend(self.preprocess(temp_files))
-            finally:
-                for fp in temp_files:
-                    os.remove(fp)
-
+            fp = self.translate_image(image_path, shift)
+            temp_file_paths.append(fp)
+        preprocessed_images = self.preprocess(temp_file_paths)
+        for temp_file_path in temp_file_paths:
+            os.remove(temp_file_path)
         return preprocessed_images
 
-    def _pad(self, batch_images, batch_size, runs_per_image=1):
+    def translate_image(self, image_path: str, shift: np.array) -> str:
+        """Translates and saves a temporary image to temporary_fp."""
+        translated_image = self.translate(cv2.imread(image_path), shift)
+        temporary_fp = tempfile.mkstemp(suffix=".png")[1]
+        cv2.imwrite(temporary_fp, translated_image)
+        return temporary_fp
+
+    def _pad(self, batch_images, batch_size):
         num_images = len(batch_images)
-        if ((num_images * runs_per_image) % batch_size == 0) or num_images < batch_size:
+        if num_images % batch_size == 0:
             return batch_images, 0
-        num_padding = batch_size * runs_per_image - (num_images % (batch_size * runs_per_image))
-        padding = np.repeat(batch_images[-runs_per_image:], repeats=num_padding, axis=0)
+        num_padding = batch_size - (num_images % batch_size)
+        padding = np.repeat(batch_images[-1:], repeats=num_padding, axis=0)
         return np.concatenate((batch_images, padding)), num_padding
 
     def _unpad(self, layer_activations, num_padding):

From b9b913b9d6b2dbd8877416585f638bf6f30666bd Mon Sep 17 00:00:00 2001
From: Ben Lonnqvist <ben.lonnqvist@epfl.ch>
Date: Tue, 12 Dec 2023 12:52:57 +0100
Subject: [PATCH 12/18] add shift metadata

---
 model_tools/activations/core.py | 26 ++++++++++++++++++++------
 1 file changed, 20 insertions(+), 6 deletions(-)

diff --git a/model_tools/activations/core.py b/model_tools/activations/core.py
index aaba534..d047b82 100644
--- a/model_tools/activations/core.py
+++ b/model_tools/activations/core.py
@@ -230,7 +230,8 @@ def translate_and_preprocess(self, images, shift):
             temp_file_paths.append(fp)
         preprocessed_images = self.preprocess(temp_file_paths)
         for temp_file_path in temp_file_paths:
-            os.remove(temp_file_path)
+            # os.remove(temp_file_path)
+            pass
         return preprocessed_images
 
     def translate_image(self, image_path: str, shift: np.array) -> str:
@@ -305,11 +306,14 @@ def _package_layer(self, layer_activations, layer, stimuli_paths, model_requirem
             self._logger.debug(f"Unknown layer activations shape {layer_activations.shape}, not inferring channels")
 
         # build assembly
-        coords = {'stimulus_path': stimuli_paths,
-                  'neuroid_num': ('neuroid', list(range(activations.shape[1]))),
-                  'model': ('neuroid', [self.identifier] * activations.shape[1]),
-                  'layer': ('neuroid', [layer] * activations.shape[1]),
-                  }
+        if model_requirements is not None and 'microsaccades' in model_requirements.keys():
+            coords = self.build_microsaccade_coords(activations, layer, stimuli_paths, model_requirements)
+        else:
+            coords = {'stimulus_path': stimuli_paths,
+                      'neuroid_num': ('neuroid', list(range(activations.shape[1]))),
+                      'model': ('neuroid', [self.identifier] * activations.shape[1]),
+                      'layer': ('neuroid', [layer] * activations.shape[1]),
+                      }
         if flatten_coord_names:
             flatten_coords = {flatten_coord_names[i]: [sample_index[i] if i < flatten_indices.shape[1] else np.nan
                                                        for sample_index in flatten_indices]
@@ -321,6 +325,16 @@ def _package_layer(self, layer_activations, layer, stimuli_paths, model_requirem
         layer_assembly['neuroid_id'] = 'neuroid', neuroid_id
         return layer_assembly
 
+    def build_microsaccade_coords(self, activations, layer, stimuli_paths, model_requirements):
+        coords = {'stimulus_path': stimuli_paths,
+                  'neuroid_num': ('neuroid', list(range(activations.shape[1]))),
+                  'model': ('neuroid', [self.identifier] * activations.shape[1]),
+                  'layer': ('neuroid', [layer] * activations.shape[1]),
+                  'shift_x': ('stimulus_path', [element[0] for element in model_requirements['microsaccades']]),
+                  'shift_y': ('stimulus_path', [element[1] for element in model_requirements['microsaccades']])
+                  }
+        return coords
+
     def insert_attrs(self, wrapper):
         wrapper.from_stimulus_set = self.from_stimulus_set
         wrapper.from_paths = self.from_paths

From ec7fbf1f2cb9e389caf1e7544af00c442e89b83a Mon Sep 17 00:00:00 2001
From: Ben Lonnqvist <ben.lonnqvist@epfl.ch>
Date: Tue, 12 Dec 2023 12:54:07 +0100
Subject: [PATCH 13/18] re-add temp file removing

---
 model_tools/activations/core.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/model_tools/activations/core.py b/model_tools/activations/core.py
index d047b82..4eb2762 100644
--- a/model_tools/activations/core.py
+++ b/model_tools/activations/core.py
@@ -230,8 +230,7 @@ def translate_and_preprocess(self, images, shift):
             temp_file_paths.append(fp)
         preprocessed_images = self.preprocess(temp_file_paths)
         for temp_file_path in temp_file_paths:
-            # os.remove(temp_file_path)
-            pass
+            os.remove(temp_file_path)
         return preprocessed_images
 
     def translate_image(self, image_path: str, shift: np.array) -> str:

From f50e3808fd68cf8252fd49fadaff86ad67d362d1 Mon Sep 17 00:00:00 2001
From: Ben Lonnqvist <ben.lonnqvist@epfl.ch>
Date: Wed, 13 Dec 2023 11:12:24 +0100
Subject: [PATCH 14/18] fix issue with NeuroidAssembly coords (thanks
 @mschrimpf)

---
 model_tools/activations/core.py | 40 ++++++++++++++++++++-------------
 1 file changed, 24 insertions(+), 16 deletions(-)

diff --git a/model_tools/activations/core.py b/model_tools/activations/core.py
index 4eb2762..349958b 100644
--- a/model_tools/activations/core.py
+++ b/model_tools/activations/core.py
@@ -98,9 +98,12 @@ def from_paths(self, stimuli_paths, layers, stimuli_identifier=None, model_requi
         # In case stimuli paths are duplicates (e.g. multiple trials), we first reduce them to only the paths that need
         # to be run individually, compute activations for those, and then expand the activations to all paths again.
         # This is done here, before storing, so that we only store the reduced activations.
-        reduced_paths = self._reduce_paths(stimuli_paths)
-        activations = fnc(layers=layers, stimuli_paths=reduced_paths, model_requirements=model_requirements)
-        activations = self._expand_paths(activations, original_paths=stimuli_paths)
+        if model_requirements is not None and 'microsaccades' in model_requirements.keys():
+            activations = fnc(layers=layers, stimuli_paths=stimuli_paths, model_requirements=model_requirements)
+        else:
+            reduced_paths = self._reduce_paths(stimuli_paths)
+            activations = fnc(layers=layers, stimuli_paths=reduced_paths, model_requirements=model_requirements)
+            activations = self._expand_paths(activations, original_paths=stimuli_paths)
         return activations
 
     @store_xarray(identifier_ignore=['stimuli_paths', 'layers'], combine_fields={'layers': 'layer'})
@@ -125,7 +128,7 @@ def _expand_paths(self, activations, original_paths):
         sorted_x = activations_paths[argsort_indices]
         sorted_index = np.searchsorted(sorted_x, original_paths)
         index = [argsort_indices[i] for i in sorted_index]
-        return activations[{'stimulus_path': index}]
+        return activations[{'presentation': index}]
 
     def register_batch_activations_hook(self, hook):
         r"""
@@ -273,9 +276,11 @@ def _package(self, layer_activations, stimuli_paths, model_requirements):
             for coord in neuroid_coords:
                 neuroid_coords[coord][1] = np.concatenate((neuroid_coords[coord][1], layer_assembly[coord].values))
             assert layer_assemblies[0].dims == layer_assembly.dims
-            for dim in set(layer_assembly.dims) - {'neuroid'}:
-                for coord in layer_assembly[dim].coords:
-                    assert (layer_assembly[coord].values == nonneuroid_coords[coord][1]).all()
+            for coord, dims, values in walk_coords(layer_assembly):
+                if set(dims) == {'neuroid'}:
+                    continue
+                assert (values == nonneuroid_coords[coord][1]).all()
+
         neuroid_coords = {coord: (dims_values[0], dims_values[1])  # re-package as tuple instead of list for xarray
                           for coord, dims_values in neuroid_coords.items()}
         model_assembly = type(layer_assemblies[0])(model_assembly, coords={**nonneuroid_coords, **neuroid_coords},
@@ -308,30 +313,33 @@ def _package_layer(self, layer_activations, layer, stimuli_paths, model_requirem
         if model_requirements is not None and 'microsaccades' in model_requirements.keys():
             coords = self.build_microsaccade_coords(activations, layer, stimuli_paths, model_requirements)
         else:
-            coords = {'stimulus_path': stimuli_paths,
+            coords = {'stimulus_path': ('presentation', stimuli_paths),
+                      'stimulus_path2': ('presentation', stimuli_paths),  # to avoid DataAssembly dim collapse
                       'neuroid_num': ('neuroid', list(range(activations.shape[1]))),
                       'model': ('neuroid', [self.identifier] * activations.shape[1]),
                       'layer': ('neuroid', [layer] * activations.shape[1]),
                       }
+
         if flatten_coord_names:
             flatten_coords = {flatten_coord_names[i]: [sample_index[i] if i < flatten_indices.shape[1] else np.nan
                                                        for sample_index in flatten_indices]
                               for i in range(len(flatten_coord_names))}
             coords = {**coords, **{coord: ('neuroid', values) for coord, values in flatten_coords.items()}}
-        layer_assembly = NeuroidAssembly(activations, coords=coords, dims=['stimulus_path', 'neuroid'])
+        layer_assembly = NeuroidAssembly(activations, coords=coords, dims=['presentation', 'neuroid'])
         neuroid_id = [".".join([f"{value}" for value in values]) for values in zip(*[
             layer_assembly[coord].values for coord in ['model', 'layer', 'neuroid_num']])]
         layer_assembly['neuroid_id'] = 'neuroid', neuroid_id
         return layer_assembly
 
     def build_microsaccade_coords(self, activations, layer, stimuli_paths, model_requirements):
-        coords = {'stimulus_path': stimuli_paths,
-                  'neuroid_num': ('neuroid', list(range(activations.shape[1]))),
-                  'model': ('neuroid', [self.identifier] * activations.shape[1]),
-                  'layer': ('neuroid', [layer] * activations.shape[1]),
-                  'shift_x': ('stimulus_path', [element[0] for element in model_requirements['microsaccades']]),
-                  'shift_y': ('stimulus_path', [element[1] for element in model_requirements['microsaccades']])
-                  }
+        coords = {
+            'stimulus_path': ('presentation', stimuli_paths),
+            'shift_x': ('presentation', [shift[0] for shift in model_requirements['microsaccades']]),
+            'shift_y': ('presentation', [shift[1] for shift in model_requirements['microsaccades']]),
+            'neuroid_num': ('neuroid', list(range(activations.shape[1]))),
+            'model': ('neuroid', [self.identifier] * activations.shape[1]),
+            'layer': ('neuroid', [layer] * activations.shape[1]),
+        }
         return coords
 
     def insert_attrs(self, wrapper):

From 0128df3269c157abcc58a9a47a7ac3cb6fc0330f Mon Sep 17 00:00:00 2001
From: Ben Lonnqvist <ben.lonnqvist@epfl.ch>
Date: Wed, 13 Dec 2023 14:12:45 +0100
Subject: [PATCH 15/18] Update tests to conform to new structure

---
 tests/activations/test___init__.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/activations/test___init__.py b/tests/activations/test___init__.py
index b7e2a5f..5c535fd 100644
--- a/tests/activations/test___init__.py
+++ b/tests/activations/test___init__.py
@@ -202,7 +202,8 @@ def test_microsaccades_from_image_path(model_ctr, layers, image_name, shifts):
     activations = activations_extractor(stimuli=stimulus_paths, layers=layers, model_requirements=model_requirements)
 
     assert activations is not None
-    assert 'shift_x' in activations.coords and 'shift_y' in activations.coords
+    assert list(activations['shift_x'].values) == [shift[0] for shift in shifts]
+    assert list(activations['shift_y'].values) == [shift[1] for shift in shifts]
     assert len(activations['shift_x']) == len(shifts) * len(stimulus_paths)
     assert len(activations['shift_y']) == len(shifts) * len(stimulus_paths)
 

From ef7a2b560b9dd2f88e63549ef29aaa1060e3692e Mon Sep 17 00:00:00 2001
From: Ben Lonnqvist <ben.lonnqvist@epfl.ch>
Date: Fri, 15 Dec 2023 10:48:34 +0100
Subject: [PATCH 16/18] add correct coord on test following structure change

---
 tests/activations/test___init__.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tests/activations/test___init__.py b/tests/activations/test___init__.py
index 5c535fd..036a30e 100644
--- a/tests/activations/test___init__.py
+++ b/tests/activations/test___init__.py
@@ -224,9 +224,8 @@ def test_model_requirements(model_ctr, layers, image_name, model_requirements):
     assert activations_with_req is not None
     assert activations_without_req is not None
     if model_requirements:
-        assert len(activations_with_req['neuroid']) > len(activations_without_req['neuroid'])
-    else:
-        assert len(activations_with_req['neuroid']) == len(activations_without_req['neuroid'])
+        assert len(activations_with_req['presentation']) > len(activations_without_req['presentation'])
+    assert len(activations_with_req['neuroid']) == len(activations_without_req['neuroid'])
 
 
 @pytest.mark.parametrize("image_name", ['rgb.jpg', 'grayscale.png', 'grayscale2.jpg', 'grayscale_alpha.png',

From b20777396df406c8d0b178efbc0d1b01999d2091 Mon Sep 17 00:00:00 2001
From: Ben Lonnqvist <ben.lonnqvist@epfl.ch>
Date: Fri, 15 Dec 2023 15:10:34 +0100
Subject: [PATCH 17/18] add metadata to stimulus sets to match new structure

---
 model_tools/activations/core.py | 51 +++++++++++++++++++++------------
 1 file changed, 33 insertions(+), 18 deletions(-)

diff --git a/model_tools/activations/core.py b/model_tools/activations/core.py
index 349958b..fd89331 100644
--- a/model_tools/activations/core.py
+++ b/model_tools/activations/core.py
@@ -10,6 +10,7 @@
 
 import numpy as np
 from tqdm.auto import tqdm
+import xarray as xr
 
 from brainio.assemblies import NeuroidAssembly, walk_coords
 from brainio.stimuli import StimulusSet
@@ -396,12 +397,20 @@ def attach_stimulus_set_meta(assembly, stimulus_set, model_requirements):
     assembly_paths = [lstrip_local(path) for path in assembly['stimulus_path'].values]
 
     assert (np.array(assembly_paths) == np.array(stimulus_paths)).all()
-    assembly['stimulus_path'] = stimulus_set['stimulus_id'].values
+    assembly = assembly.reset_index('presentation')
+    assembly.assign_coords(stimulus_path=('presentation', stimulus_set['stimulus_id'].values))
+
     assembly = assembly.rename({'stimulus_path': 'stimulus_id'})
 
+    all_columns = []
     for column in stimulus_set.columns:
-        assembly[column] = 'stimulus_id', stimulus_set[column].values
-    assembly = assembly.stack(presentation=('stimulus_id',))
+        assembly = assembly.assign_coords({column: ('presentation', stimulus_set[column].values)})
+        all_columns.append(column)
+    if 'stimulus_id' in all_columns:
+        all_columns.remove('stimulus_id')
+    if 'stimulus_path2' in all_columns:
+        all_columns.remove('stimulus_path2')
+    assembly = assembly.set_index(presentation=['stimulus_id', 'stimulus_path2'] + all_columns)
     return assembly
 
 
@@ -410,30 +419,36 @@ def attach_stimulus_set_meta_with_microsaccades(assembly, stimulus_set, model_re
     stimulus_paths = [lstrip_local(path) for path in stimulus_paths]
     assembly_paths = [lstrip_local(path) for path in assembly['stimulus_path'].values]
 
-    replication_factor = len(model_requirements['microsaccades']) if model_requirements['microsaccades'] else 1
+    replication_factor = len(model_requirements['microsaccades'])
     repeated_stimulus_paths = np.repeat(stimulus_paths, replication_factor)
     assert (np.array(assembly_paths) == np.array(repeated_stimulus_paths)).all()
-
     repeated_stimulus_ids = np.repeat(stimulus_set['stimulus_id'].values, replication_factor)
-    assembly['stimulus_path'] = repeated_stimulus_ids
-    assembly = assembly.rename({'stimulus_path': 'stimulus_id'})
 
-    for column in stimulus_set.columns:
-        assembly[column] = 'stimulus_id', np.repeat(stimulus_set[column].values, replication_factor)
+    # repeat over the presentation dimension to accommodate multiple runs per stimulus
+    repeated_assembly = xr.concat([assembly for _ in range(replication_factor)], dim='presentation')
+    repeated_assembly = repeated_assembly.reset_index('presentation')
+    repeated_assembly.assign_coords(stimulus_path=('presentation', repeated_stimulus_ids))
+    repeated_assembly = repeated_assembly.rename({'stimulus_path': 'stimulus_id'})
 
-    # Handle shifts
-    repeated_shifts = np.tile(np.array(model_requirements['microsaccades']), (len(stimulus_set['stimulus_id']), 1))
-    assembly.coords['shift_x'] = ('stimulus_id', repeated_shifts[:, 0])
-    assembly.coords['shift_y'] = ('stimulus_id', repeated_shifts[:, 1])
+    # hack to capture columns
+    all_columns = []
+    for column in stimulus_set.columns:
+        repeated_values = np.repeat(stimulus_set[column].values, replication_factor)
+        repeated_assembly = repeated_assembly.assign_coords({column: ('presentation', repeated_values)})
+        all_columns.append(column)
+    if 'stimulus_id' in all_columns:
+        all_columns.remove('stimulus_id')
+    if 'stimulus_path2' in all_columns:
+        all_columns.remove('stimulus_path2')
 
-    # Preserve existing coordinates and dimensions
-    for coord in assembly.coords:
-        assembly[coord] = ('stimulus_id', np.repeat(assembly[coord].data, replication_factor))
+    repeated_assembly.coords['shift_x'] = ('presentation', [shift[0] for shift in model_requirements['microsaccades']])
+    repeated_assembly.coords['shift_y'] = ('presentation', [shift[1] for shift in model_requirements['microsaccades']])
 
     # Set MultiIndex
-    assembly = assembly.set_index(presentation=['stimulus_id'] + list(assembly.coords))
+    index = ['stimulus_id', 'stimulus_path2'] + all_columns + ['shift_x', 'shift_y']
+    repeated_assembly = repeated_assembly.set_index(presentation=index)
 
-    return assembly
+    return repeated_assembly
 
 
 class HookHandle:

From 155df600df8f61da0b8d627b44b3836a045b4459 Mon Sep 17 00:00:00 2001
From: Ben Lonnqvist <ben.lonnqvist@epfl.ch>
Date: Fri, 15 Dec 2023 16:46:24 +0100
Subject: [PATCH 18/18] add within-shift batch padding to accommodate TF models

---
 model_tools/activations/core.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/model_tools/activations/core.py b/model_tools/activations/core.py
index fd89331..d8ee2ec 100644
--- a/model_tools/activations/core.py
+++ b/model_tools/activations/core.py
@@ -196,8 +196,9 @@ def _get_microsaccade_activations_batched(self, paths, layers, batch_size, shift
             # compute activations on the entire batch one shift at a time
             for shift in shifts:
                 assert type(shift) == tuple
-                shifted_batch = self.translate_and_preprocess(batch_inputs, shift)
+                shifted_batch, num_padding = self.translate_and_preprocess(batch_inputs, shift, batch_size)
                 activations = self.get_activations(shifted_batch, layers)
+                activations = self._unpad(activations, num_padding)
                 for layer_name, layer_output in activations.items():
                     batch_activations.setdefault(layer_name, []).append(layer_output)
 
@@ -226,16 +227,19 @@ def _get_batch_activations(self, inputs, layer_names, batch_size):
         activations = self._unpad(activations, num_padding)
         return activations
 
-    def translate_and_preprocess(self, images, shift):
+    def translate_and_preprocess(self, images, shift, batch_size):
         assert type(images) == list
         temp_file_paths = []
         for image_path in images:
             fp = self.translate_image(image_path, shift)
             temp_file_paths.append(fp)
-        preprocessed_images = self.preprocess(temp_file_paths)
+        # having to pad here causes a considerable slowdown for small datasets (n_images << batch_size),
+        #  but is necessary to accommodate tensorflow models
+        padded_images, num_padding = self._pad(temp_file_paths, batch_size)
+        preprocessed_images = self.preprocess(padded_images)
         for temp_file_path in temp_file_paths:
             os.remove(temp_file_path)
-        return preprocessed_images
+        return preprocessed_images, num_padding
 
     def translate_image(self, image_path: str, shift: np.array) -> str:
         """Translates and saves a temporary image to temporary_fp."""