Fix tests and add more

Author: CeliaBenquet

cebra_lens/__init__.py

@@ -1,7 +1,6 @@
 # example of structure so that you can directly use the functions get_layer_activations instead of having to do CEBRA_Lens.activations.get_layer_activations
 from .activations import *
-from .quantification import *
-from .quantification.decoding import *
+from .quantification.decoder import *
 from .quantification.distance import *
 from .quantification.cka_metric import *
 from .quantification.rdm_metric import *

cebra_lens/activations.py

@@ -80,7 +80,8 @@ def get_cut_indices(
         cut_indices.append((0, 0))
     elif layer_type == None:
         raise NotImplementedError(
-            "Padding handling not implemented for 'all'.")
+            "Padding handling not implemented to handle activations for all layer types.",
+            "Set layer_type to nn.Conv1d to use the default padding handling.")
     else:
         # need to analyze the padding from the last output of Conv1 and apply the same cut
         raise NotImplementedError(
@@ -94,7 +95,7 @@ def get_activations_model(
     session_id: int = -1,
     name: str = "single",
     instance: int = 0,
-    layer_type: Type[nn.Module] = None,
+    layer_type: Type[nn.Module] = nn.Conv1d,
 ) -> Dict[str, npt.NDArray]:
     """
     Extracts activations from a single model layer.
@@ -112,7 +113,8 @@ def get_activations_model(
     instance : int
         The instance number for the model, used to differentiate between models from the same model category.
     layer_type : Type[nn.Module]
-        The type of layer to extract activations from. Defaults to None, meaning extracts activations from all layers.
+        The type of layer to extract activations from. None means it extracts activations from all layers.
+        Default is nn.Conv1d, which is the most common layer type used in CEBRA models.
 
     Returns:
     --------

cebra_lens/quantification/__init__.py

@@ -2,6 +2,6 @@
 from .rdm_metric import *
 from .misc import *
 from .distance import *
-from .decoding import *
+from .decoder import *
 from .base import *
 from .tsne import *

cebra_lens/quantification/cka_metric.py

@@ -8,9 +8,10 @@
 from tqdm import tqdm
 import numpy as np
 from .base import _BaseMetric
-from ..matplotlib import *
+import cebra_lens.matplotlib as cebra_lens_matplotlib
 from typing import Optional, List, Dict, Tuple
 import numpy.typing as npt
+import matplotlib
 
 
 class CKA(_BaseMetric):
@@ -188,7 +189,8 @@ def _compute_per_layer(
         cka_matrix = np.zeros((len(embeddings_1), len(embeddings_1[0])))
         for j in tqdm(range(len(embeddings_1))):
             if flag:
-                # the situation when there multiple models inside model labels and the same number of models inside each label
+                # the situation when there multiple models inside model labels and the same number of 
+                # models inside each label
                 cka_matrix[j, :] = self._compute_cka(embeddings_1[j],
                                                      embeddings_2[j])
             else:
@@ -207,7 +209,8 @@ def compute(self, activations: Dict[str, npt.NDArray],
         Parameters:
         -----------
         activations : Dict[str, npt.NDArray]
-            A dictionary where keys are strings which represent the model label and values are 2d lists with the corresponding activations per layer.
+            A dictionary where keys are strings which represent the model label and values are 2d lists 
+            with the corresponding activations per layer.
 
         comparison : Tuple[str, str]
             A tuple containing the model labels to compare.
@@ -227,7 +230,8 @@ def compute(self, activations: Dict[str, npt.NDArray],
 
         if len(activations_1) != len(activations_2):
             # if the number of models in a label is different from the other model label
-            # choose embeddings_1 for the one with more models, and then embeddings_2 just compare with the first model
+            # choose embeddings_1 for the one with more models, and then embeddings_2 just compare with 
+            # the first model
             if len(activations_1) > len(activations_2):
                 embeddings_1 = activations_1
                 embeddings_2 = activations_2[0]
@@ -293,7 +297,7 @@ def plot(
         matplotlib.axes.Axes
             The axes on which the heatmap is plotted.
         """
-        return plot_cka_heatmaps(
+        return cebra_lens_matplotlib.plot_cka_heatmaps(
             cka_matrices,
             annot,
             show_cbar,

cebra_lens/quantification/decoding.py cebra_lens/quantification/decoder.pycebra_lens/quantification/decoding.py renamed to cebra_lens/quantification/decoder.py

@@ -5,12 +5,13 @@
 from ..utils_hpc import decoding_pos_dir
 from ..activations import get_activations_model
 from .base import _BaseMetric
-from ..matplotlib import *
+import cebra_lens.matplotlib as cebra_lens_matplotlib
 import numpy.typing as npt
-from typing import Dict, Type, Tuple
+from typing import Dict, Type, Tuple, Optional
 import torch.nn as nn
 import sklearn.metrics
 import torch as pt
+import matplotlib
 
 
 def decoding(
@@ -118,7 +119,7 @@ def __init__(
         test_label: npt.NDArray,
         session_id: int = 0,
         dataset_label: str = None,
-        layer_type: Optional[Type[nn.Module]] = None,
+        layer_type: Optional[Type[nn.Module]] = nn.Conv1d,
         output_only: bool = True,
     ):
 
@@ -315,7 +316,7 @@ def compute(
     def __name__(self):
         return "decode_by_layer"
 
-    def set_output_only(self, output_only):
+    def set_output_only(self, output_only: bool) -> None:
         """
         Set the output_only parameter to True or False. If True, it will compute the decoding scores for the output embeddings of the model, otherwise it will compute the decoding scores for the activations of the model.
 
@@ -369,8 +370,8 @@ def plot(
                 )
 
         if self.output_only:
-            return plot_decoding(results_dict, palette, self.dataset_label,
+            return cebra_lens_matplotlib.plot_decoding(results_dict, palette, self.dataset_label,
                                  label, plot_error, ax)
         else:
-            return plot_layer_decoding(results_dict, title, self.dataset_label,
+            return cebra_lens_matplotlib.plot_layer_decoding(results_dict, title, self.dataset_label,
                                        label, plot_error, figsize)

cebra_lens/quantification/rdm_metric.py

@@ -1,13 +1,15 @@
 """All the functions relative to the Representation Dissimilarity Matrix (RDM) calculation"""
 
+from typing import Dict, List, Optional
 import numpy as np
 from scipy.linalg import block_diag
 from typing import List, Optional, Tuple, Union
 from scipy.spatial.distance import correlation, pdist, squareform
 from .misc import discrete_binning, continuous_binning
 import torch
+import matplotlib
 from .base import _BaseMetric
-from ..matplotlib import *
+import cebra_lens.matplotlib as cebra_lens_matplotlib
 import numpy.typing as npt
 
 
@@ -21,8 +23,8 @@ class RDM(_BaseMetric):
         The data array of shape (num_samples, num_features).
     label : torch.Tensor
         The array of labels corresponding to the data.
-    discrete : bool, optional
-        Whether the labels are discrete or continuous. If None, it will be determined based on the dataset_label.
+    is_discrete_labels : bool, optional
+        Whether the labels are discrete or continuous. By default, it is False, meaning the labels are continuous.
     dataset_label : str, optional
         The dataset type, either 'visual' or 'HPC'. Default is 'visual'.
     metric : str, optional
@@ -37,7 +39,7 @@ def __init__(
         self,
         data: torch.Tensor,
         label: torch.Tensor,
-        is_discrete_labels: bool = None,
+        is_discrete_labels: bool = False,
         dataset_label: str = None,
         metric: str = "correlation",
         bool_oracle: bool = True,
@@ -254,9 +256,9 @@ def plot(
             The figure containing the plotted RDMs.
         """
         if self.bool_oracle:
-            return plot_rdm_correlation(rdms)
+            return cebra_lens_matplotlib.plot_rdm_correlation(rdms)
         else:
-            return plot_rdm_all(
+            return cebra_lens_matplotlib.plot_rdm_all(
                 rdms=rdms,
                 labels=self.label,
                 num_bins=self.num_bins,

cebra_lens/utils.py

@@ -6,7 +6,7 @@
 import numpy.typing as npt
 from tqdm import tqdm
 from torch import nn
-from .quantification.decoding import Decoding
+from .quantification.decoder import Decoding
 from .quantification.rdm_metric import RDM
 from .quantification.cka_metric import CKA
 from .quantification.tsne import Tsne

tests/test_activations.py

@@ -1,5 +1,5 @@
-import pytest
 import torch
+import pytest
 import numpy as np
 from collections import namedtuple
 from unittest.mock import MagicMock
@@ -22,7 +22,7 @@ def test_cut_array_with_cut():
     np.testing.assert_array_equal(result, np.array([[2, 3, 4]]))
 
 
-def test_get_cut_indices_conv1d():
+def test_get_cut_indices():
     Offset = namedtuple("Offset", ["left", "right"])
 
     # Mock the model's get_offset behavior
@@ -32,6 +32,9 @@ def test_get_cut_indices_conv1d():
     result = get_cut_indices(model_mock, torch.nn.Conv1d, [3, 3])
     assert isinstance(result, list)
     assert all(isinstance(x, tuple) and len(x) == 2 for x in result)
+    
+    with pytest.raises(NotImplementedError, match="Padding handling not implemented*"):
+        get_cut_indices(model_mock, None, [3, 3])
 
 
 def make_mock_cebra_model():

tests/test_cka.py

@@ -0,0 +1,86 @@
+import pytest
+import numpy as np
+import torch
+from unittest.mock import patch, MagicMock
+import cebra_lens
+
+@pytest.fixture
+def dummy_comparisons():
+    return [("A", "B")]
+
+@pytest.fixture
+def dummy_cka(dummy_comparisons):
+    return cebra_lens.quantification.cka_metric.CKA(comparisons=dummy_comparisons)
+
+@pytest.fixture
+def dummy_activations():
+    # Simulate Conv1D and Linear layer activations as 2D arrays (samples, features)
+    batch_size = 10
+    conv_channels = 4
+    conv_length = 8
+    linear_features = 5
+
+    # Conv1D output: (batch_size, conv_channels, conv_length) -> flatten to (batch_size, conv_channels * conv_length)
+    conv1d_activations_A = np.random.rand(batch_size, conv_channels, conv_length).reshape(batch_size, -1)
+    linear_activations_A = np.random.rand(batch_size, linear_features)
+    conv1d_activations_B = np.random.rand(batch_size, conv_channels, conv_length).reshape(batch_size, -1)
+    linear_activations_B = np.random.rand(batch_size, linear_features)
+
+    # Each group has a list of 2D arrays (one per layer)
+    return {
+        "A": np.array([[np.random.rand(5, 10), np.random.rand(5, 10)], 
+              [np.random.rand(5, 10), np.random.rand(5, 10)]]),
+        "B": np.array([[np.random.rand(5, 10), np.random.rand(5, 10)], 
+              [np.random.rand(5, 10), np.random.rand(5, 10)]]),
+    }
+
+def test_center_gram_symmetry(dummy_cka):
+    mat = np.eye(5)
+    centered = dummy_cka.center_gram(mat)
+    assert np.allclose(centered, centered.T)
+
+def test_center_gram_unbiased(dummy_cka):
+    mat = np.eye(5)
+    centered = dummy_cka.center_gram(mat, unbiased=True)
+    assert np.allclose(centered, centered.T)
+
+def test_gram_linear(dummy_cka):
+    x = np.random.rand(10, 5)
+    gram = dummy_cka.gram_linear(x)
+    assert gram.shape == (10, 10)
+    assert np.allclose(gram, gram.T)
+
+def test_cka_value(dummy_cka):
+    x = np.random.rand(10, 5)
+    y = np.random.rand(10, 5)
+    gram_x = dummy_cka.gram_linear(x)
+    gram_y = dummy_cka.gram_linear(y)
+    val = dummy_cka.cka(gram_x, gram_y)
+    assert isinstance(val, float) or isinstance(val, np.floating)
+
+def test_compute_cka_shape(dummy_cka):
+    emb1 = [np.random.rand(5, 10), np.random.rand(5, 10)]
+    emb2 = [np.random.rand(5, 10), np.random.rand(5, 10)]
+    result = dummy_cka._compute_cka(emb1, emb2)
+    assert result.shape == (1, 2)
+
+def test_compute_per_layer_shape(dummy_cka):
+    emb1 = [ [np.random.rand(5, 10), np.random.rand(5, 10)] for _ in range(3) ]
+    emb2 = [ [np.random.rand(5, 10), np.random.rand(5, 10)] for _ in range(3) ]
+    result = dummy_cka._compute_per_layer(emb1, emb2, flag=True)
+    assert result.shape == (3, 2)
+
+def test_compute(dummy_cka, dummy_activations):
+    result = dummy_cka.compute(dummy_activations, ("A", "B"))
+    assert isinstance(result, np.ndarray)
+    
+    
+def test_compute_intra_label(dummy_cka, dummy_activations):
+    result = dummy_cka.compute(dummy_activations, ("A", "A"))
+    assert isinstance(result, np.ndarray)
+
+@patch("cebra_lens.matplotlib.plot_cka_heatmaps")
+def test_plot_calls_heatmap(mock_plot, dummy_cka):
+    cka_matrices = {"A": np.random.rand(2, 2)}
+    dummy_cka.plot(cka_matrices, annot=True)
+    assert mock_plot.called