Add klab.Zhang2018search

tests/test_assemblies.py

@@ -28,6 +28,9 @@
         'dicarlo.Kar2019',
         'dicarlo.Kar2018hvm',
         'dicarlo.Kar2018cocogray',
+        'klab.Zhang2018search_obj_array',
+        'klab.Zhang2018search_naturaldesign',
+        'klab.Zhang2018search_waldo',
 ))
 def test_list_assembly(assembly):
     l = brainio_collection.list_assemblies()
@@ -52,6 +55,9 @@ def test_list_assembly(assembly):
     pytest.param('dicarlo.Kar2019', marks=[pytest.mark.private_access]),
     pytest.param('dicarlo.Kar2018hvm', marks=[pytest.mark.private_access]),
     pytest.param('dicarlo.Kar2018cocogray', marks=[pytest.mark.private_access]),
+    pytest.param('klab.Zhang2018search_obj_array', marks=[pytest.mark.private_access]),
+    pytest.param('klab.Zhang2018search_naturaldesign', marks=[pytest.mark.private_access]),
+    pytest.param('klab.Zhang2018search_waldo', marks=[pytest.mark.private_access]),
 ])
 def test_existence(assembly_identifier):
     assert brainio_collection.get_assembly(assembly_identifier) is not None
@@ -148,3 +154,36 @@ def test_stimulus_set_from_assembly():
     for image_id in stimulus_set['image_id']:
         image_path = stimulus_set.get_image(image_id)
         assert os.path.exists(image_path)
+
+
+@pytest.mark.private_access
+def test_klab_Zhang2018search():
+    # object array
+    assembly = brainio_collection.get_assembly('klab.Zhang2018search_obj_array')
+    assert set(assembly.dims) == {'presentation', 'fixation', 'position'}
+    assert len(assembly['presentation']) == 4500
+    assert len(set(assembly['image_id'].values)) == 300
+    assert len(set(assembly['subjects'].values)) == 15
+    assert len(assembly['fixation']) == 8
+    assert len(assembly['position']) == 2
+    assert assembly.stimulus_set is not None
+
+    # naturaldesign
+    assembly = brainio_collection.get_assembly('klab.Zhang2018search_naturaldesign')
+    assert set(assembly.dims) == {'presentation', 'fixation', 'position'}
+    assert len(assembly['presentation']) == 3600
+    assert len(set(assembly['image_id'].values)) == 240
+    assert len(set(assembly['subjects'].values)) == 15
+    assert len(assembly['fixation']) == 66
+    assert len(assembly['position']) == 2
+    assert assembly.stimulus_set is not None
+
+    # waldo
+    assembly = brainio_collection.get_assembly('klab.Zhang2018search_waldo')
+    assert set(assembly.dims) == {'presentation', 'fixation', 'position'}
+    assert len(assembly['presentation']) == 1005
+    assert len(set(assembly['image_id'].values)) == 67
+    assert len(set(assembly['subjects'].values)) == 15
+    assert len(assembly['fixation']) == 81
+    assert len(assembly['position']) == 2
+    assert assembly.stimulus_set is not None

tests/test_stimuli.py

@@ -43,7 +43,33 @@ def test_dicarlohvm(self):
         'dicarlo.objectome.public',
         'dicarlo.objectome.private',
         'dicarlo.Kar2018cocogray',
+        'klab.Zhang2018.search_obj_array',
+        'klab.Zhang2018.search_naturaldesign',
+        'klab.Zhang2018.search_waldo',
 ))
 def test_list_stimulus_set(stimulus_set):
     l = brainio_collection.list_stimulus_sets()
     assert stimulus_set in l
+
+
+@pytest.mark.private_access
+def test_klab_Zhang2018search():
+    # object array
+    stimulus_set = brainio_collection.get_stimulus_set('klab.Zhang2018.search_obj_array')
+    # There are 300 presentation images in the assembly but 606 in the StimulusSet (explanation from @shashikg follows).
+    # For each of the visual search task out of total 300, you need two images (one - the target image,
+    # second - the search space image) plus there are 6 different mask images to mask objects
+    # present at 6 different locations in a specified search image.
+    # Therefore, a total of 300 * 2 + 6 images are there in the stimulus set.
+    assert len(stimulus_set) == 606
+    assert len(set(stimulus_set['image_id'])) == 606
+
+    # naturaldesign
+    stimulus_set = brainio_collection.get_stimulus_set('klab.Zhang2018.search_naturaldesign')
+    assert len(stimulus_set) == 240*3
+    assert len(set(stimulus_set['image_id'])) == 240*3
+
+    # waldo
+    stimulus_set = brainio_collection.get_stimulus_set('klab.Zhang2018.search_waldo')
+    assert len(stimulus_set) == 67*3
+    assert len(set(stimulus_set['image_id'])) == 67*3