brainscore-vision 2.1__py3-none-any.whl

tests/test_model_helpers/brain_transformation/test_temporal.py

@@ -0,0 +1,66 @@
+import numpy as np
+import pytest
+
+from brainio.assemblies import NeuroidAssembly
+from brainscore_vision.model_helpers.brain_transformation import TemporalAligned
+from brainscore_vision.model_interface import BrainModel
+
+
+class LayerMappedModelMock:
+    def __init__(self, output_temporal=False):
+        # attributes that TemporalAligned copies for outside use
+        self.region_layer_map = None
+        self.activations_model = None
+        self.start_task = None
+        self.output_temporal = output_temporal
+
+    def start_recording(self, *args, **kwargs):
+        pass
+
+    def look_at(self, *args, **kwargs):
+        if self.output_temporal:
+            return NeuroidAssembly([[[1, 1], [2, 2], [3, 3]], [[1, 1], [2, 2], [3, 3]]], 
+                                coords={'stimulus_id': ('presentation', ['image1', 'image2']),
+                                                        'object_name': (
+                                                            'presentation', ['number', 'number']),
+                                                        'neuroid_id': ('neuroid', [1, 2, 3]),
+                                                        'region': ('neuroid', ['IT'] * 3), 
+                                                        'time_bin_start': ('time_bin', [0, 150]), 
+                                                        'time_bin_end': ('time_bin', [150, 300]), },
+                                dims=['presentation', 'neuroid', 'time_bin'])
+        else:
+            return NeuroidAssembly([[1, 2, 3], [1, 2, 3]], coords={'stimulus_id': ('presentation', ['image1', 'image2']),
+                                                                'object_name': (
+                                                                    'presentation', ['number', 'number']),
+                                                                'neuroid_id': ('neuroid', [1, 2, 3]),
+                                                                'region': ('neuroid', ['IT'] * 3), },
+                                dims=['presentation', 'neuroid'])
+
+
+class TestTemporalAligned:
+    @pytest.mark.parametrize('output_temporal', [False, True])
+    def test_single_timebin(self, output_temporal):
+        model = TemporalAligned(layer_model=LayerMappedModelMock(output_temporal))
+        model.start_recording(recording_target=BrainModel.RecordingTarget.IT, time_bins=[(70, 170)])
+        recordings = model.look_at('dummy')
+        assert set(recordings.dims) == {'presentation', 'neuroid'}  # squeezed time-bin
+
+    @pytest.mark.parametrize('output_temporal', [False, True])
+    def test_two_timebins(self, output_temporal):
+        layer_model = LayerMappedModelMock(output_temporal)
+        model = TemporalAligned(layer_model=layer_model)
+        model.start_recording(recording_target=BrainModel.RecordingTarget.IT, time_bins=[(70, 170), (170, 270)])
+        recordings = model.look_at('dummy')
+        assert set(recordings.dims) == {'presentation', 'neuroid', 'time_bin'}
+        np.testing.assert_array_equal(recordings['time_bin_start'].values, [70, 170])
+        np.testing.assert_array_equal(recordings['time_bin_end'].values, [170, 270])
+
+    @pytest.mark.parametrize('output_temporal', [False, True])
+    def test_18_timebins(self, output_temporal):
+        model = TemporalAligned(layer_model=LayerMappedModelMock(output_temporal))
+        time_bins = [(70 + i * 10, 80 + i * 10) for i in range(18)]
+        model.start_recording(recording_target=BrainModel.RecordingTarget.IT, time_bins=time_bins)
+        recordings = model.look_at('dummy')
+        assert set(recordings.dims) == {'presentation', 'neuroid', 'time_bin'}
+        np.testing.assert_array_equal(recordings['time_bin_start'].values, [start for start, end in time_bins])
+        np.testing.assert_array_equal(recordings['time_bin_end'].values, [end for start, end in time_bins])

tests/test_model_helpers/temporal/activations/test_extractor.py

@@ -0,0 +1,96 @@
+import numpy as np
+import os
+import pytest
+
+from brainio.stimuli import StimulusSet
+from brainscore_vision.model_helpers.activations.temporal.inputs.base import Stimulus
+from brainscore_vision.model_helpers.activations.temporal.model import ActivationWrapper
+from brainscore_vision.model_helpers.activations.temporal.core import TemporalInferencer, CausalInferencer 
+from collections import OrderedDict
+
+
+"""This module tests model_helpers.activations.temporal.core.extractor using the default settings"""
+
+
+video_paths = [
+    os.path.join(os.path.dirname(__file__), "..", "dots1.mp4"),
+    os.path.join(os.path.dirname(__file__), "..", "dots2.mp4"),
+]
+video_durations = [2000, 6000]
+img_path = os.path.join(os.path.dirname(__file__), "../../activations/rgb.jpg")
+fps = 1
+
+
+def get_fake_models(causal=False, **kwargs):
+    def transform_video(video):
+        frames = video.to_numpy()[:, :12, :12]
+        return frames
+    
+    class FakeActivationWrapper(ActivationWrapper):
+        def __init__(self, **kwargs):
+            super().__init__("dummy", transform_video, **kwargs)
+
+        def get_activations(self, inputs, layers):
+            ret = OrderedDict()
+            for layer in layers:
+                ret[layer] = np.stack(inputs)
+            return ret
+
+    layer_activation_format = {**{f'layer{i}': "THWC" for i in range(1, 3)}}
+
+    inferencer_cls = TemporalInferencer if not causal else CausalInferencer
+    if inferencer_cls is CausalInferencer: 
+        kwargs['duration'] = (0, 3000)
+    wrapper = FakeActivationWrapper(inferencer_cls=inferencer_cls, fps=fps, 
+                             layer_activation_format=layer_activation_format, max_workers=1, batch_size=4, **kwargs)
+    layers = list(layer_activation_format.keys())
+    return wrapper, layers
+
+
+@pytest.mark.memory_intense
+@pytest.mark.parametrize(["causal", "padding", "time_align"], [(False, True, "ignore_time"), (True, False, "per_frame_aligned")])
+def test_from_video_path(causal, padding, time_align):
+    video_names = ["dots1.mp4", "dots2.mp4"]  # 2s vs 6s
+    stimuli_paths = video_paths
+
+    activations_extractor, layers = get_fake_models(causal=causal, 
+                                                           batch_padding=padding, time_alignment=time_align)
+    activations = activations_extractor.from_paths(stimuli_paths=stimuli_paths,
+                                                   layers=layers)
+
+    assert activations is not None
+    assert len(activations['stimulus_path']) == 2
+    assert len(np.unique(activations['layer'])) == len(layers)
+
+    expected_num_time_bins = 6 * fps
+    if causal:
+        assert activations.sizes['time_bin'] == expected_num_time_bins 
+
+    import gc
+    gc.collect()  # free some memory, we're piling up a lot of activations at this point
+
+
+def _build_stimulus_set(video_names):
+    stimulus_set = StimulusSet([{'stimulus_id': video_name, 'some_meta': video_name[::-1]}
+                                for video_name in video_names])
+    stimulus_set.stimulus_paths = {video_name: path
+                                   for video_name, path in zip(video_names, video_paths)}
+    return stimulus_set
+
+
+@pytest.mark.memory_intense
+@pytest.mark.parametrize(["causal", "padding"], [(False, True), (True, False)])
+def test_from_stimulus_set(causal, padding):
+    video_names = ["dots1.mp4", "dots2.mp4"]
+    stimulus_set = _build_stimulus_set(video_names)
+
+    activations_extractor, layers = get_fake_models(causal=causal, batch_padding=padding)
+    activations = activations_extractor(stimulus_set, layers=layers)
+    
+    assert activations is not None
+    assert set(activations['stimulus_id'].values) == set(video_names)
+    assert all(activations['some_meta'].values == [video_name[::-1] for video_name in video_names])
+    assert len(np.unique(activations['layer'])) == len(layers)
+
+    import gc
+    gc.collect()  # free some memory, we're piling up a lot of activations at this point

tests/test_model_helpers/temporal/activations/test_inferencer.py

@@ -0,0 +1,189 @@
+import numpy as np
+import os
+import pytest
+
+from collections import OrderedDict
+
+from brainscore_vision.model_helpers.activations.temporal.core import Inferencer, TemporalInferencer, CausalInferencer, BlockInferencer
+from brainscore_vision.model_helpers.activations.temporal.inputs import Video, Stimulus
+
+
+"""This module tests model_helpers.activations.temporal.core.inferencer
+
+    Different inferencers are tested:
+    - Inferencer: the basic inferencer that does not enforce any temporal context
+    - TemporalInferencer: the inferencer that aligns the activations to the video time
+    - CausalInferencer: the inferencer that ensures the activations are causal
+    - BlockInferencer: the inferencer that divides the video into blocks and infer the activations for each block
+"""
+
+
+video_paths = [
+    os.path.join(os.path.dirname(__file__), "..", "dots1.mp4"),
+    os.path.join(os.path.dirname(__file__), "..", "dots2.mp4"),
+]
+video_durations = [2000, 6000]
+img_path = os.path.join(os.path.dirname(__file__), "../../activations/rgb.jpg")
+
+
+def dummy_get_features(model_inputs, layers):
+    feature = np.stack(model_inputs)
+    B, F, H, W, C = feature.shape
+    feature = feature.reshape(B, F, H//80, 80, W//80, 80, C).mean((3, 5))[..., :2]  # BFHWC=B,F,6,3,2
+    batch_activation = OrderedDict({
+        "layer1": feature,
+        "layer2": feature[:, 0, 0, 0]
+    })
+    return batch_activation
+
+def dummy_preprocess(video):
+    feature = video.to_numpy()[:, 200:680, 200:440, :]
+    return feature
+
+def time_down_sample_preprocess(video):
+    feature = video.to_numpy()[::2, 200:680, 200:440, :]
+    return feature
+
+dummy_layer_activation_format = {
+    "layer1": "THWC",
+    "layer2": "C",
+}
+
+dummy_layers = ["layer1", "layer2"]
+
+
+@pytest.mark.memory_intense
+@pytest.mark.parametrize("max_spatial_size", [None, 2, 4])
+def test_inferencer(max_spatial_size):
+    inferencer = Inferencer(dummy_get_features, dummy_preprocess, dummy_layer_activation_format, 
+                            Video, max_workers=1, max_spatial_size=max_spatial_size, batch_grouper=lambda s: s.duration)
+    model_assembly = inferencer(video_paths, layers=dummy_layers)
+    if max_spatial_size is None:
+        # 6 second video with fps 60 has 360 frames
+        # the model simply return the same number of frames as the temporal size of activations
+        # so the number of channel_temporal should be 360
+        assert model_assembly.sizes["neuroid"] == 360*6*3*2 + 2
+    else:
+        assert model_assembly.sizes["neuroid"] == 360*max_spatial_size*(max_spatial_size//2) * 2 + 2
+    assert model_assembly.sizes["stimulus_path"] == 2 
+
+
+@pytest.mark.parametrize("time_alignment", ["evenly_spaced", "ignore_time"])
+@pytest.mark.parametrize("fps", [10, 30, 45])
+def test_temporal_inferencer(time_alignment, fps):
+    inferencer = TemporalInferencer(dummy_get_features, dummy_preprocess, 
+                                    dummy_layer_activation_format, max_workers=1, 
+                                    fps=fps, time_alignment=time_alignment)
+    model_assembly = inferencer(video_paths, layers=dummy_layers)
+    assert model_assembly['time_bin_start'].values[0] == 0
+    assert model_assembly['time_bin_end'].values[-1] == max(video_durations)
+
+    if time_alignment != "ignore_time":
+        # since the longer video lasts for 6 seconds, and the temporal inferencer align all output assembly to have fps
+        # specified when constructing the inferencer, the number of time bins should be 6*fps
+        assert model_assembly.sizes["time_bin"] == 6 * fps
+        assert np.isclose(model_assembly['time_bin_end'].values[0] - model_assembly['time_bin_start'].values[0], 1000/fps)
+    else:
+        assert model_assembly.sizes["time_bin"] == 1
+        assert model_assembly['time_bin_end'].values[0] - model_assembly['time_bin_start'].values[0] == max(video_durations)
+
+    # manual computation check
+    output_values = model_assembly.sel(stimulus_path=video_paths[1])\
+                                    .isel(neuroid=model_assembly.layer=="layer1")\
+                                    .transpose('time_bin', 'neuroid').values.reshape(-1)
+    
+    manual_compute_values = []
+    video = Video.from_path(video_paths[1]).set_fps(fps)
+    manual_compute_values = dummy_get_features([dummy_preprocess(video)], ["layer1"])["layer1"][0].reshape(-1)
+    manual_compute_values = manual_compute_values.astype(output_values.dtype)
+    assert np.allclose(output_values, manual_compute_values)
+
+
+@pytest.mark.memory_intense
+def test_img_input():
+    fps = 30
+    inferencer = TemporalInferencer(dummy_get_features, dummy_preprocess, 
+                                    dummy_layer_activation_format, max_workers=1, 
+                                    fps=fps, convert_img_to_video=True, img_duration=1000)
+    model_assembly = inferencer([img_path], layers=dummy_layers)
+    assert model_assembly.sizes["time_bin"] == fps
+
+
+def test_compute_temporal_context():
+    fps=10
+    inferencer = CausalInferencer(None, None, None, fps=fps, duration=(200, 1000), temporal_context_strategy="greedy")
+    assert inferencer._compute_temporal_context() == (200, 1000)
+
+    inferencer = CausalInferencer(None, None, None, fps=fps, duration=(200, 1000), temporal_context_strategy="conservative")
+    assert inferencer._compute_temporal_context() == (200, 200)
+
+    inferencer = CausalInferencer(None, None, None, fps=fps, duration=(0, 1000), num_frames=(2, 5), temporal_context_strategy="greedy")
+    assert inferencer._compute_temporal_context() == (200, 500)
+
+    inferencer = CausalInferencer(None, None, None, fps=fps, duration=(0, 1000), num_frames=(2, 15), temporal_context_strategy="greedy")
+    assert inferencer._compute_temporal_context() == (200, 1000)
+
+    inferencer = CausalInferencer(None, None, None, fps=fps, duration=(0, 1000), num_frames=(2, 15), temporal_context_strategy="fix", fixed_temporal_context=500)
+    assert inferencer._compute_temporal_context() == (200, 500)
+
+
+@pytest.mark.memory_intense
+@pytest.mark.parametrize("preprocess", ["normal", "downsample"])
+@pytest.mark.parametrize("fps", [1, 40])
+def test_causal_inferencer(preprocess, fps):
+    if preprocess == "normal":
+        preprocess = dummy_preprocess
+    else:
+        preprocess = time_down_sample_preprocess
+    inferencer = CausalInferencer(dummy_get_features, dummy_preprocess, 
+                                    dummy_layer_activation_format, 
+                                    fps=fps, max_workers=1)
+    model_assembly = inferencer(video_paths, layers=dummy_layers)
+    assert model_assembly.sizes["time_bin"] == 6 * fps
+    assert np.isclose(model_assembly['time_bin_end'].values[0] - model_assembly['time_bin_start'].values[0], 1000/fps)
+    assert inferencer._compute_temporal_context() == (1000/fps, np.inf)
+
+    # manual computation check
+    output_values = model_assembly.sel(stimulus_path=video_paths[1])\
+                                    .isel(neuroid=model_assembly.layer=="layer1")\
+                                    .transpose('time_bin', 'neuroid').values
+    manual_compute_values = []
+    video = Video.from_path(video_paths[1]).set_fps(fps)
+    interval = 1000/fps
+    for time_end in np.arange(interval, 6000+interval, interval):
+        clip = video.set_window(0, time_end)
+        manual_compute_values.append(dummy_get_features([dummy_preprocess(clip)], ["layer1"])["layer1"][0, -1])
+    manual_compute_values = np.stack(manual_compute_values).reshape(len(manual_compute_values), -1).astype(output_values.dtype)
+    assert np.allclose(output_values, manual_compute_values)
+
+
+@pytest.mark.memory_intense
+@pytest.mark.parametrize("preprocess", ["normal", "downsample"])
+@pytest.mark.parametrize("fps", [1, 40])
+def test_block_inferencer(preprocess, fps):
+    if preprocess == "normal":
+        preprocessing = dummy_preprocess
+    else:
+        preprocessing = time_down_sample_preprocess
+    inferencer = BlockInferencer(dummy_get_features, preprocessing, dummy_layer_activation_format, fps=fps, 
+                                 duration=(200, 4000), temporal_context_strategy="greedy", max_workers=1)
+    model_assembly = inferencer(video_paths, layers=dummy_layers)
+    assert model_assembly.sizes["time_bin"] == 8 * fps  # block overflow 2 x 4 seconds
+    assert np.isclose(model_assembly['time_bin_end'].values[0] - model_assembly['time_bin_start'].values[0], 1000/fps)
+
+    # manual computation check
+    output_values = model_assembly.sel(stimulus_path=video_paths[1])\
+                                    .isel(neuroid=model_assembly.layer=="layer1")\
+                                    .transpose('time_bin', 'neuroid').values
+    manual_compute_values = []
+    video = Video.from_path(video_paths[1]).set_fps(fps)
+    interval = 4000
+    for time_end in np.arange(interval, 6000+interval, interval):
+        time_start = time_end - interval
+        clip = video.set_window(time_start, time_end)
+        manual_compute_values.append(dummy_get_features([preprocessing(clip)], ["layer1"])["layer1"][0])
+    manual_compute_values = np.concatenate(manual_compute_values)
+    manual_compute_values = manual_compute_values.reshape(len(manual_compute_values), -1).astype(output_values.dtype)
+    if preprocess == "downsample":
+        output_values = output_values[::2]
+    assert np.allclose(output_values, manual_compute_values)

tests/test_model_helpers/temporal/activations/test_inputs.py

@@ -0,0 +1,103 @@
+import os
+import numpy as np
+
+from brainscore_vision.model_helpers.activations.temporal.inputs import Video, Image
+
+
+"""This module tests model_helpers.activations.temporal.inputs
+
+    Different inputs are tested:
+    - Video
+    - Image
+    Specifically, the different transformations (set fps, set size, etc.) of the inputs are tested.
+"""
+
+video_paths = [
+    os.path.join(os.path.dirname(__file__), "..", "dots1.mp4"),
+    os.path.join(os.path.dirname(__file__), "..", "dots2.mp4"),
+]
+video_durations = [2000, 6000]
+img_path = os.path.join(os.path.dirname(__file__), "../../activations/rgb.jpg")
+
+
+def test_video_load_frames():
+    video1 = Video.from_path(video_paths[0])
+    fps = video1.fps
+    duration = video1.duration
+    all_frames = video1.to_frames()
+    interval = 1000 / fps
+    for t, f in zip(np.arange(interval, duration, interval), all_frames):
+        v = video1.set_window(0, t)
+        this_f = v.to_frames()[-1]
+        assert (this_f == f).all()
+
+def test_video():
+    video1 = Video.from_path(video_paths[0])
+    video2 = Video.from_img_path(img_path, 1000, 30)
+
+    assert video2.duration == 1000
+
+    video3 = video1.set_window(-10, 0, padding="repeat")
+    video4 = video1.set_window(-20, -10, padding="repeat")
+    assert (video3.to_numpy() == video4.to_numpy()).all()
+    assert (video3.to_numpy()[0] == video1.to_numpy()[0]).all()
+
+    assert video2.fps == 30
+    assert video2.set_fps(1).to_numpy().shape[0] == 1
+
+    video5 = video1.set_size((120, 100))
+    assert tuple(video5.to_numpy().shape[1:3]) == (100, 120)
+
+    video6 = video1.set_fps(30)
+    assert (video6.to_numpy()[1] == video1.to_numpy()[2]).all()
+    assert (video6.to_numpy()[2] == video1.to_numpy()[4]).all()
+
+    video6 = video1.set_fps(20)
+    assert (video6.to_numpy()[1] == video1.to_numpy()[3]).all()
+    assert (video6.to_numpy()[2] == video1.to_numpy()[6]).all()
+
+    video7 = video1.set_window(-100, 100).set_window(100, 200)
+    assert video7.duration == 100
+    assert (video7.to_numpy() == video1.set_window(0, 100).to_numpy()).all()
+
+    video8 = video1.set_window(300, 500).set_window(0, 100)
+    assert video8.duration == 100
+    assert (video8.to_numpy() == video1.set_window(300, 400).to_numpy()).all()
+
+    # test copy
+    video9 = video1.set_fps(5).copy().set_fps(30).copy()
+    assert (video9.to_numpy()[1] == video1.to_numpy()[2]).all()
+    assert (video9.to_numpy()[2] == video1.to_numpy()[4]).all()
+
+    for frame in [10, 50, 100]:
+        time_start = 1000 / video1.fps * frame
+        video10 = video1.set_window(time_start, time_start+1000/video1.fps)
+        assert video10.to_numpy().shape[0] == 1
+        assert (video10.to_numpy()[0] == video1.to_numpy()[frame]).all()
+
+        video10 = video1.set_window(0, time_start+1000/video1.fps)
+        assert video10.to_numpy().shape[0] == frame+1
+        assert (video10.to_numpy()[frame] == video1.to_numpy()[frame]).all()
+
+        video10 = video1.set_window(time_start, video1.duration)
+        assert video10.to_numpy().shape[0] == video1.to_numpy().shape[0] - frame
+        assert (video10.to_numpy()[0] == video1.to_numpy()[frame]).all()
+
+    for fps in [7.5, 9, 1, 43, 1000/video1.duration, 1001/video1.duration]:
+        video11 = video1.set_fps(fps)
+        assert video11.to_numpy().shape[0] == np.ceil(video1.duration * fps / 1000)
+
+    for v in [video1, video2]:
+        target_num_frames = 7
+        duration = 1000 / v.fps * target_num_frames
+        common = list(np.arange(0, v.duration, 100))
+        extra1 = 1000 / v.fps * 3 + v.duration
+        extra2 = 1000 / v.fps * 2 + v.duration
+        extra3 = 1000 / v.fps * 1
+        for t in [extra1, extra2, extra3] + common:
+            video = v.set_window(t-duration, t, padding="repeat")
+            assert video.to_numpy().shape[0] == target_num_frames
+
+def test_image():
+    img = Image.from_path(img_path)
+    assert img.set_size((10, 12)).to_numpy().shape[:2] == (12, 10)

tests/test_model_helpers/temporal/brain_transformation/test_temporal_ops.py

@@ -0,0 +1,122 @@
+from brainscore_vision.model_helpers.brain_transformation.temporal import time_align, assembly_time_align
+import numpy as np
+
+# imports
+import numpy as np
+from brainio.assemblies import DataAssembly
+
+
+"""
+    This module tests the time alignment functionalities: 
+    ie., given a set of target time bins, align the neural assembly with a set of source time bins to them.
+"""
+
+
+T = 3
+P = 6
+N = 4
+time_bins = np.array([(0, 10), (10, 20), (20, 30)])
+stimulus_id = ["A", "B", "C", "D", "E", "F"]
+neuroid_id = ["n1", "n2", "n3", "n4"]
+
+# latent variable that generates both "video" and "neural signals"
+latent1 = [1, 3, 2]
+latent2 = [2, 3, 1]
+latent3 = [1, 2, 3]
+latent4 = [3, 2, 1]
+latent5 = [3, 1, 2]
+latent6 = [2, 1, 3]
+latent = np.array([latent1, latent2, latent3, latent4, latent5, latent6])
+
+# videos
+from PIL import Image
+vs = np.random.rand(P, T, 8, 8, 3) + latent[..., None, None, None]  # 0~4
+vs = (vs * 255//4).astype(np.uint8)
+videos = [[Image.fromarray(img) for img in video] for video in vs]
+video_paths = [f"video{i}.mp4"for i in range(P)]
+
+# neural signals at time 1 2 3:
+n1 = latent * 2
+n2 = latent * 1
+n3 = latent * .5
+n4 = latent * .1
+perfect_response = np.array([n1, n2, n3, n4])
+noisy_response = perfect_response + np.random.randn(N, P, T) * 0.1
+
+# assemblies
+def _make_neural_assembly(data):
+    assembly = DataAssembly(
+        data,
+        dims=["neuroid", "presentation", "time_bin"],
+        coords={
+            "neuroid_id": ("neuroid", neuroid_id),
+            "stimulus_id": ("presentation", stimulus_id),
+            "time_bin_start": ("time_bin", time_bins[:, 0]),
+            "time_bin_end": ("time_bin", time_bins[:, 1]),
+        }
+    )
+    return assembly
+
+assembly = _make_neural_assembly(noisy_response)
+
+def _except(func, *args, **kwargs):
+    try:
+        func(*args, **kwargs)
+    except:
+        return
+    else:
+        raise False
+
+
+def test_time_align():
+    # case 1
+    source_time_bins = [(0, 100), (100, 200), (200, 300)]
+    target_time_bins = [(0, 50), (250, 300)]
+
+    ret = time_align(source_time_bins, target_time_bins, mode = "center")
+    belong_to = np.array([[1, 0, 0], [0, 0, 1]])
+    assert (ret == belong_to).all()
+
+    ret = time_align(source_time_bins, target_time_bins, mode = "portion")
+    belong_to = np.array([[0.5, 0, 0], [0, 0, 0.5]])
+    assert (ret == belong_to).all()
+
+    # case 2
+    starts = np.arange(0, 100, 10)
+    ends = starts + 10
+    target_time_bins = np.stack([starts, ends], axis=-1)
+    starts = np.arange(0, 100, 1000/30)
+    ends = starts + 1000/30
+    source_time_bins = np.stack([starts, ends], axis=-1)
+    assert (time_align(source_time_bins, target_time_bins).argmax(1) == np.array([0,0,0,1,1,1,1,2,2,2])).all()
+
+
+def test_assembly_time_align():
+    ret = assembly_time_align(assembly, [(0, 10), (20, 30), (10, 20)])
+    assert (assembly.isel(time_bin=[0, 2, 1]) == ret.isel(time_bin=[0, 1, 2])).all()
+    _except(assembly_time_align, assembly, [(0, 10), (20, 30), (40, 50)])
+    _except(assembly_time_align, assembly, [(30, 20)])
+    ret = assembly_time_align(assembly, [(20, 30), (0, 10), (10, 20)])
+    assert (assembly.isel(time_bin=[2, 0, 1]) == ret.isel(time_bin=[0, 1, 2])).all()
+
+    ret = assembly_time_align(assembly, [(0, 30), (0, 10), (10, 20)], mode="portion")
+    assert (assembly.isel(time_bin=[1, 2]) == ret.isel(time_bin=[2, 1])).all()
+    assert (ret.isel(time_bin=0).values == assembly.mean("time_bin").values).all()
+
+    ret = assembly_time_align(assembly, [(5, 25)], mode="portion")
+    val = (assembly.isel(time_bin=0).values * 0.5 + assembly.isel(time_bin=1).values + assembly.isel(time_bin=2).values * 0.5) / 2
+    assert np.isclose(ret.isel(time_bin=0).values, val).all()
+
+    ret = assembly_time_align(assembly, [(9, 21)], mode="portion")
+    val = (assembly.isel(time_bin=0).values * .1 + assembly.isel(time_bin=1).values + assembly.isel(time_bin=2).values * .1) / 1.2
+    assert np.isclose(ret.isel(time_bin=0).values, val).all()
+
+    ret = assembly_time_align(assembly, [(9, 13)], mode="portion")
+    val = (assembly.isel(time_bin=0).values * 0.25 + assembly.isel(time_bin=1).values * 0.75) 
+    assert np.isclose(ret.isel(time_bin=0).values, val).all()
+
+    _except(assembly_time_align, assembly, [(20, 35)], mode="portion")
+    _except(assembly_time_align, assembly, [(0, 31)], mode="portion")
+    _except(assembly_time_align, assembly, [(0, 30.00001)], mode="portion")
+
+    assert (assembly_time_align(assembly, [(0, 10)], mode="portion").values == assembly_time_align(assembly, [(0, 10)], mode="center").values).all()

tests/test_model_helpers/temporal/test_utils.py

@@ -0,0 +1,61 @@
+import numpy as np
+from brainscore_vision.model_helpers.activations.temporal.utils import (
+    batch_2d_resize
+)
+
+
+# this tests whether the batch_2d_resize with "pool" mode works as expected (ie, average over regions)
+def test_proportional_average_pooling():
+    arr = np.array([
+        [1, 2],
+        [3, 4],
+        [5, 6],
+    ])[...,None].astype(float)
+
+    size_0 = (2, 1)
+    arr_0 = np.array([
+        [3, 4],
+    ])[...,None]
+
+    def proportional_average_pooling(arr, size):
+        return batch_2d_resize(arr[None,:], size, "pool")[0]
+
+    assert np.allclose(proportional_average_pooling(arr, size_0), arr_0)
+
+    size_1 = (2, 2)
+    arr_1 = np.array([
+        [1*2/3+3*1/3, 2*2/3+4*1/3],
+        [3*1/3+5*2/3, 4*1/3+6*2/3],
+    ])[...,None]
+
+    assert np.allclose(proportional_average_pooling(arr, size_1), arr_1)
+
+    size_2 = (1, 2)
+    arr_2 = np.array([
+        [(1*2/3+3*1/3+2*2/3+4*1/3)/2],
+        [(3*1/3+5*2/3+4*1/3+6*2/3)/2],
+    ])[...,None]
+
+    assert np.allclose(proportional_average_pooling(arr, size_2), arr_2)
+
+    size_3 = (3, 3)
+    arr_3 = np.array([
+        [1, 1.5, 2],
+        [3, 3.5, 4],
+        [5, 5.5, 6],
+    ])[...,None]
+
+    assert np.allclose(proportional_average_pooling(arr, size_3), arr_3)
+
+    size_4 = (1, 1)
+    arr_4 = arr.reshape(-1, 1).mean(0)[None, None, :]
+    assert np.allclose(proportional_average_pooling(arr, size_4), arr_4)
+
+    size_5 = (2, 4)
+    arr_5 = np.array([
+        [1, 2],
+        [(1*1+3*2)/3, (2*1+4*2)/3],
+        [(3*2+5*1)/3, (4*2+6*1)/3],
+        [5, 6],
+    ])[...,None].astype(float)
+    assert np.allclose(proportional_average_pooling(arr, size_5), arr_5)