convert dimensionality metric into plugin format

brainscore_vision/benchmarks/islam2021/benchmark.py

@@ -5,7 +5,7 @@
 
 from brainscore_core import Score
 from brainscore_core.benchmarks import BenchmarkBase
-from brainscore_vision import BrainModel, load_stimulus_set
+from brainscore_vision import BrainModel, load_stimulus_set, load_metric
 from brainscore_vision.metrics.dimensionality import Dimensionality
 
 BIBTEX = """@inproceedings{
@@ -29,7 +29,7 @@ def __init__(self,region,factor,deterministic=True):
         self.stimulus_set  = load_stimulus_set("Islam2021")
         self.region = region
         self.deterministic = deterministic
-        self._metric = Dimensionality(factor_idx)
+        self._metric = load_metric('dimensionality', factor=factor_idx)
         self._number_of_trials = 1
         super(_Islam2021Dimensionality, self).__init__(
             identifier=f'Islam2021-{region + "_" + factor + "_dimensionality"}', version=1,

brainscore_vision/metrics/dimensionality/__init__.py

@@ -0,0 +1,13 @@
+from brainscore_vision import metric_registry
+from .metric import Dimensionality
+
+metric_registry['factor_dimensionality'] = Dimensionality
+
+BIBTEX = """@inproceedings{esser2020disentangling,
+  title={A disentangling invertible interpretation network for explaining latent representations},
+  author={Esser, Patrick and Rombach, Robin and Ommer, Bjorn},
+  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
+  pages={9223--9232},
+  year={2020},
+  url={https://arxiv.org/abs/2004.13166}
+}"""

brainscore_vision/metrics/dimensionality.py → brainscore_vision/metrics/dimensionality/metric.py

@@ -1,29 +1,19 @@
 import numpy as np
-from brainscore_core import Score
+from brainscore_core import Metric, Score
 
-"""
-@inproceedings{esser2020disentangling,
-  title={A disentangling invertible interpretation network for explaining latent representations},
-  author={Esser, Patrick and Rombach, Robin and Ommer, Bjorn},
-  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
-  pages={9223--9232},
-  year={2020},
-  url={https://arxiv.org/abs/2004.13166}
-}
-"""
 
-class Dimensionality:
+class Dimensionality(Metric):
     def __init__(self, factor):
         self.factor = factor #factor to return
-        
-    def __call__(self,assembly1, assembly2):
-        #each assembly should have 3 dimensions: (factor, sample_num, neuron_i) 
+
+    def __call__(self, assembly1, assembly2) -> Score:
+        #each assembly should have 3 dimensions: (factor, sample_num, neuron_i)
         dims_percent = self._dim_est(assembly1,assembly2)
         factor_dim = dims_percent[self.factor]
         return Score(factor_dim)
-        
+
     def _dim_est(self, za, zb):
-        score_by_factor = dict() 
+        score_by_factor = dict()
         zall = np.concatenate([za,zb], 0)
         mean = np.mean(zall, 0, keepdims=True) #mean and variance per neuron
         var = np.sum(np.mean((zall-mean)*(zall-mean), 0))
@@ -33,15 +23,15 @@ def _dim_est(self, za, zb):
             mean_f_factor = 0.5*(np.mean(za_factor, 0, keepdims=True)+np.mean(zb_factor, 0, keepdims=True))
             cov_f_factor = np.mean((za_factor-mean_f_factor)*(zb_factor-mean_f_factor), 0)
             raw_score_f_factor = np.sum(cov_f_factor)
-            score_by_factor[f] = raw_score_f_factor/var   
-        
-        assert "residual" not in score_by_factor 
+            score_by_factor[f] = raw_score_f_factor/var
+
+        assert "residual" not in score_by_factor
         score_by_factor["residual"] = 1.0 #by default
         score_names = score_by_factor.keys()
         scores = np.fromiter(score_by_factor.values(),dtype=float)
-        dims_percent = self._softmax_dim(score_names,scores,za.shape[1])  
+        dims_percent = self._softmax_dim(score_names,scores,za.shape[1])
         return dims_percent
-    
+
     def _softmax_dim(self, score_names, scores, N):
         m = np.max(scores)
         e = np.exp(scores-m)