Fix imports with matplotlib

Author: CeliaBenquet

cebra_lens/__init__.py

@@ -1,6 +1,5 @@
 # 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 .matplotlib import *
 from .quantification.cka_metric import *
 from .quantification.decoder import *
 from .quantification.distance import *
@@ -9,6 +8,7 @@
 from .utils import *
 from .utils_allen import *
 from .utils_hpc import *
+from .utils_plot import *
 
 # selects what files can be imported when doing from CEBRA_Lens import * --> keep env clean
 # __all__ = ['get_layer_activations']

cebra_lens/activations.py

@@ -9,7 +9,7 @@
 import torch
 import torch.nn as nn
 
-from .matplotlib import plot_activations
+from .utils_plot import plot_activations
 
 
 def _cut_array(array: npt.NDArray,

cebra_lens/quantification/cka_metric.py

@@ -12,7 +12,7 @@
 import numpy.typing as npt
 from tqdm import tqdm
 
-import cebra_lens.matplotlib as cebra_lens_matplotlib
+from cebra_lens import utils_plot
 
 from .base import _BaseMetric
 
@@ -300,7 +300,7 @@ def plot(
         matplotlib.axes.Axes
             The axes on which the heatmap is plotted.
         """
-        return cebra_lens_matplotlib.plot_cka_heatmaps(
+        return utils_plot.plot_cka_heatmaps(
             cka_matrices,
             annot,
             show_cbar,

cebra_lens/quantification/decoder.py

@@ -9,7 +9,7 @@
 import torch as pt
 import torch.nn as nn
 
-import cebra_lens.matplotlib as cebra_lens_matplotlib
+from cebra_lens import utils_plot
 
 from ..activations import get_activations_model
 from ..utils_allen import decoding_frames
@@ -373,10 +373,10 @@ def plot(
                 )
 
         if self.output_only:
-            return cebra_lens_matplotlib.plot_decoding(results_dict, palette,
-                                                       self.dataset_label,
-                                                       label, plot_error, ax)
+            return utils_plot.plot_decoding(results_dict, palette,
+                                            self.dataset_label, label,
+                                            plot_error, ax)
         else:
-            return cebra_lens_matplotlib.plot_layer_decoding(
-                results_dict, title, self.dataset_label, label, plot_error,
-                figsize)
+            return utils_plot.plot_layer_decoding(results_dict, title,
+                                                  self.dataset_label, label,
+                                                  plot_error, figsize)

cebra_lens/quantification/distance.py

@@ -7,8 +7,8 @@
 from scipy.spatial.distance import cdist, pdist
 from sklearn.preprocessing import StandardScaler
 
-from ..matplotlib import *
 from ..utils import extract_label
+from ..utils_plot import *
 from .base import _BaseMetric
 from .misc import continuous_binning, discrete_binning, repetition_binning
 

cebra_lens/quantification/rdm_metric.py

@@ -9,7 +9,7 @@
 from scipy.linalg import block_diag
 from scipy.spatial.distance import correlation, pdist, squareform
 
-import cebra_lens.matplotlib as cebra_lens_matplotlib
+from cebra_lens import utils_plot
 
 from .base import _BaseMetric
 from .misc import continuous_binning, discrete_binning
@@ -258,9 +258,9 @@ def plot(
             The figure containing the plotted RDMs.
         """
         if self.bool_oracle:
-            return cebra_lens_matplotlib.plot_rdm_correlation(rdms)
+            return utils_plot.plot_rdm_correlation(rdms)
         else:
-            return cebra_lens_matplotlib.plot_rdm_all(
+            return utils_plot.plot_rdm_all(
                 rdms=rdms,
                 labels=self.label,
                 num_bins=self.num_bins,

cebra_lens/quantification/tsne.py

@@ -2,11 +2,12 @@
 
 from typing import List, Optional, Union
 
+import matplotlib
 import numpy as np
 import numpy.typing as npt
 from sklearn.manifold import TSNE
 
-from ..matplotlib import *
+from ..utils_plot import *
 from .base import _BaseMetric
 
 

cebra_lens/matplotlib.py cebra_lens/utils_plot.pycebra_lens/matplotlib.py renamed to cebra_lens/utils_plot.py

@@ -1,9 +1,9 @@
 import numpy as np
-from ..cebra_lens.quantification.base import _BaseMetric
-from ..cebra_lens.matplotlib import *
+from cebra_lens.quantification.base import _BaseMetric
+from cebra_lens.utils_plot import *
 from typing import List, Optional, Union
 import numpy.typing as npt
-
+import matplotlib
 
 class NewMetric(_BaseMetric):
     """
@@ -86,7 +86,7 @@ def plot(
             The figure containing the NewMetric plot.
         """
 
-        #Need to define the plot_newMetric function in the matplotlib.py
+        #Need to define the plot_newMetric function in the utils_plot.py
         return plot_newMetric(
             embeddings,
             labels,

demos/metric_template.py

@@ -101,7 +101,7 @@ def test_compute_intra_label(dummy_cka, dummy_activations):
     assert isinstance(result, np.ndarray)
 
 
-@patch("cebra_lens.matplotlib.plot_cka_heatmaps")
+@patch("cebra_lens.utils_plot.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)