brainscore-vision 2.1__py3-none-any.whl

brainscore_vision/__init__.py

@@ -0,0 +1,105 @@
+import logging
+from typing import Dict, Any, Union, Callable
+
+from brainio.assemblies import DataAssembly
+from brainio.stimuli import StimulusSet
+from brainscore_core.benchmarks import Benchmark
+from brainscore_core.metrics import Metric, Score
+from brainscore_core.plugin_management.conda_score import wrap_score
+from brainscore_core.plugin_management.import_plugin import import_plugin
+from brainscore_vision.metrics import Ceiling
+from brainscore_vision.model_interface import BrainModel
+
+_logger = logging.getLogger(__name__)
+
+data_registry: Dict[str, Callable[[], Union[DataAssembly, Any]]] = {}
+""" Pool of available data """
+
+stimulus_set_registry: Dict[str, Callable[[], StimulusSet]] = {}
+""" Pool of available stimulus sets"""
+
+metric_registry: Dict[str, Callable[[], Metric]] = {}
+""" Pool of available metrics """
+
+benchmark_registry: Dict[str, Callable[[], Benchmark]] = {}
+""" Pool of available benchmarks """
+
+model_registry: Dict[str, Callable[[], BrainModel]] = {}
+""" Pool of available models """
+
+
+def load_dataset(identifier: str) -> Union[DataAssembly, Any]:
+    import_plugin('brainscore_vision', 'data', identifier)
+
+    return data_registry[identifier]()
+
+
+def load_stimulus_set(identifier: str) -> Any:
+    # dataset and stimulus_set are kept in the same directory
+    import_plugin('brainscore_vision', 'data', identifier, registry_prefix='stimulus_set')
+
+    return stimulus_set_registry[identifier]()
+
+
+def load_metric(identifier: str, *args, **kwargs) -> Metric:
+    import_plugin('brainscore_vision', 'metrics', identifier)
+
+    return metric_registry[identifier](*args, **kwargs)
+
+
+def load_ceiling(identifier: str, *args, **kwargs) -> Ceiling:
+    """
+    Convenience method to load ceilings with the correct type -- internally uses :meth:`~brainscore_vision.load_metric`,
+    i.e. the ceiling has to be registered as a metric.
+    """
+    return load_metric(identifier, *args, **kwargs)
+
+
+def load_benchmark(identifier: str) -> Benchmark:
+    import_plugin('brainscore_vision', 'benchmarks', identifier)
+
+    return benchmark_registry[identifier]()
+
+
+def load_model(identifier: str) -> BrainModel:
+    import_plugin('brainscore_vision', 'models', identifier)
+
+    return model_registry[identifier]()
+
+
+def _run_score(model_identifier: str, benchmark_identifier: str) -> Score:
+    """
+    Score the model referenced by the `model_identifier` on the benchmark referenced by the `benchmark_identifier`.
+    """
+    model: BrainModel = load_model(model_identifier)
+    benchmark: Benchmark = load_benchmark(benchmark_identifier)
+    score: Score = benchmark(model)
+    score.attrs['model_identifier'] = model_identifier
+    score.attrs['benchmark_identifier'] = benchmark_identifier
+    try:  # attempt to look up the layer commitment if model uses a standard layer model
+        score.attrs['comment'] = f"layers: {model.layer_model.region_layer_map}"
+    except Exception:
+        pass
+    return score
+
+
+def score(model_identifier: str, benchmark_identifier: str, conda_active: bool = False) -> Score:
+    """
+    Score the model referenced by the `model_identifier` on the benchmark referenced by the `benchmark_identifier`.
+    The model needs to implement the :class:`~brainscore_vision.model_interface.BrainModel` interface
+    so that the benchmark can interact with it.
+    The benchmark will be looked up from the :data:`~brainscore_vision.benchmarks` and evaluates the model
+    (looked up from :data:`~brainscore_vision.models`) on how brain-like it is under that benchmark's
+    experimental paradigm, primate measurements, comparison metric, and ceiling.
+    This results in a quantitative
+    `Score <https://brain-score-core.readthedocs.io/en/latest/modules/metrics.html#brainscore_core.metrics.Score>`_
+    ranging from 0 (least brain-like) to 1 (most brain-like under this benchmark).
+
+    :param model_identifier: the identifier for the model
+    :param benchmark_identifier: the identifier for the benchmark to test the model against
+    :return: a Score of how brain-like the candidate model is under this benchmark. The score is normalized by
+        this benchmark's ceiling such that 1 means the model matches the data to ceiling level.
+    """
+    return wrap_score(__file__,
+                      model_identifier=model_identifier, benchmark_identifier=benchmark_identifier,
+                      score_function=_run_score, conda_active=conda_active)

brainscore_vision/__main__.py

@@ -0,0 +1,20 @@
+from pathlib import Path
+
+import fire
+
+from brainscore_vision import score as _score_function
+
+
+def score(model_identifier: str, benchmark_identifier: str, conda_active: bool=False):
+    result = _score_function(model_identifier, benchmark_identifier, conda_active)
+    print(result)  # print instead of return because fire has issues with xarray objects
+
+
+def test_plugins(*args, **kwargs):
+    from brainscore_core.plugin_management.test_plugins import run_args as _core_test_plugins
+    root_directory = Path(__file__).parent
+    _core_test_plugins(*args, root_directory=root_directory, **kwargs)
+
+
+if __name__ == '__main__':
+    fire.Fire()

brainscore_vision/benchmark_helpers/__init__.py

@@ -0,0 +1,67 @@
+from typing import Union
+
+import numpy as np
+import hashlib
+
+from brainio.assemblies import NeuroidAssembly, DataAssembly
+from brainscore_core import Score
+from brainscore_vision.model_interface import BrainModel
+
+
+class PrecomputedFeatures(BrainModel):
+    def __init__(self, features: Union[DataAssembly, dict], visual_degrees):
+        """
+        :param features: The precomputed features. Either an assembly of features, indexable with `stimulus_id` or
+            a dictionary mapping from stimulus identifier to feature assemblies.
+        :param visual_degrees: Some visual degrees to use for the precomputed features. Since features are precomputed,
+            this should only affect the `place_on_screen` in the benchmark's __call__ method.
+        """
+        self.features = features
+        self._visual_degrees = visual_degrees
+
+    @property
+    def identifier(self) -> str:
+        # serialize the features to a string and create hash
+        features_data = str(self.features)
+        features_hash = hashlib.md5(features_data.encode('utf-8')).hexdigest()
+        return f"precomputed-{features_hash}"
+
+    def visual_degrees(self) -> int:
+        return self._visual_degrees
+
+    def start_task(self, task, fitting_stimuli=None):
+        pass
+
+    def start_recording(self, region, *args, **kwargs):
+        pass
+
+    def look_at(self, stimuli, number_of_trials=1):
+        features = self.features[stimuli.identifier] if isinstance(self.features, dict) else self.features
+        missing_stimulus_ids = set(stimuli['stimulus_id'].values) - set(features['stimulus_id'].values)
+        assert not missing_stimulus_ids, f"stored features do not contain stimulus_ids {missing_stimulus_ids}"
+        image_indices = [np.where(features['stimulus_id'].values == image_id)[0][0]
+                         for image_id in stimuli['stimulus_id'].values]
+        features = features.isel(presentation=image_indices)
+        assert all(features['stimulus_id'].values == stimuli['stimulus_id'].values)
+        return features
+
+
+def check_standard_format(assembly, nans_expected=False):
+    assert isinstance(assembly, NeuroidAssembly)
+    assert set(assembly.dims).issuperset({'presentation', 'neuroid'})
+    assert hasattr(assembly, 'stimulus_id')
+    assert hasattr(assembly, 'neuroid_id')
+    if not nans_expected:
+        assert not np.isnan(assembly).any()
+    assert 'stimulus_set_identifier' in assembly.attrs
+
+
+def bound_score(score: Score):
+    """
+    Force score value to be between 0 and 1.
+    If score is lower than 0, set to 0. If score is greater than 1, set to 1.
+    """
+    if score < 0:
+        score.loc[()] = 0
+    elif score > 1:
+        score.loc[()] = 1

brainscore_vision/benchmark_helpers/neural_common.py

@@ -0,0 +1,70 @@
+import numpy as np
+
+from brainio.assemblies import array_is_element, walk_coords
+from brainscore_core import Score
+from brainscore_vision.benchmarks import BenchmarkBase
+from brainscore_vision.model_interface import BrainModel
+from .screen import place_on_screen
+
+
+class NeuralBenchmark(BenchmarkBase):
+    def __init__(self, identifier, assembly, similarity_metric, visual_degrees, number_of_trials, **kwargs):
+        super(NeuralBenchmark, self).__init__(identifier=identifier, **kwargs)
+        self._assembly = assembly
+        self._similarity_metric = similarity_metric
+        region = np.unique(self._assembly['region'])
+        assert len(region) == 1
+        self.region = region[0]
+        timebins = timebins_from_assembly(self._assembly)
+        self.timebins = timebins
+        self._visual_degrees = visual_degrees
+        self._number_of_trials = number_of_trials
+
+    def __call__(self, candidate: BrainModel):
+        candidate.start_recording(self.region, time_bins=self.timebins)
+        stimulus_set = place_on_screen(self._assembly.stimulus_set, target_visual_degrees=candidate.visual_degrees(),
+                                       source_visual_degrees=self._visual_degrees)
+        source_assembly = candidate.look_at(stimulus_set, number_of_trials=self._number_of_trials)
+        if 'time_bin' in source_assembly.dims and source_assembly.sizes['time_bin'] == 1:
+            source_assembly = source_assembly.squeeze('time_bin')  # static case for these benchmarks
+        raw_score = self._similarity_metric(source_assembly, self._assembly)
+        ceiled_score = explained_variance(raw_score, self.ceiling)
+        return ceiled_score
+
+
+def timebins_from_assembly(assembly):
+    timebins = assembly['time_bin'].values
+    if 'time_bin' not in assembly.dims:
+        timebins = [timebins]  # only single time-bin
+    return timebins
+
+
+def explained_variance(score: Score, ceiling: Score) -> Score:
+    # ro(X, Y)
+    # = (r(X, Y) / sqrt(r(X, X) * r(Y, Y)))^2
+    # = (r(X, Y) / sqrt(r(Y, Y) * r(Y, Y)))^2  # assuming that r(Y, Y) ~ r(X, X) following Yamins 2014
+    # = (r(X, Y) / r(Y, Y))^2
+    r_square = np.power(score.values /
+                        ceiling.values, 2)
+    ceiled_score = Score(r_square)
+    if 'error' in score.attrs:
+        ceiled_score.attrs['error'] = score.attrs['error']
+    ceiled_score.attrs[Score.RAW_VALUES_KEY] = score
+    ceiled_score.attrs['ceiling'] = ceiling
+    return ceiled_score
+
+
+def average_repetition(assembly):
+    def avg_repr(assembly):
+        presentation_coords = [coord for coord, dims, values in walk_coords(assembly)
+                               if array_is_element(dims, 'presentation') and coord != 'repetition']
+        assembly = assembly.multi_groupby(presentation_coords).mean(dim='presentation', skipna=True)
+        return assembly
+
+    return apply_keep_attrs(assembly, avg_repr)
+
+def apply_keep_attrs(assembly, fnc):  # workaround to keeping attrs
+    attrs = assembly.attrs
+    assembly = fnc(assembly)
+    assembly.attrs = attrs
+    return assembly