Merge branch 'main' into stes/add-split

Author: MMathisLab

README.md

@@ -35,7 +35,12 @@ To receive updates on code releases, please 👀 watch or ⭐️ star this repos
 It can jointly use behavioral and neural data in a hypothesis- or discovery-driven manner to produce consistent, high-performance latent spaces. While it is not specific to neural and behavioral data, this is the first domain we used the tool in. This application case is to obtain a consistent representation of latent variables driving activity and behavior, improving decoding accuracy of behavioral variables over standard supervised learning, and obtaining embeddings which are robust to domain shifts.
 
 
-# Reference
+# References
+
+- 📄 **Publication April 2025**:
+  [Time-series attribution maps with regularized contrastive learning.](https://arxiv.org/abs/2502.12977)
+Steffen Schneider, Rodrigo González Laiz, Anastasiia Filipova, Markus Frey, Mackenzie Weygandt Mathis. AISTATS 2025.
+
 
 - 📄 **Publication May 2023**:
   [Learnable latent embeddings for joint behavioural and neural analysis.](https://doi.org/10.1038/s41586-023-06031-6)

cebra/helper.py

@@ -32,7 +32,7 @@
 
 import numpy as np
 import numpy.typing as npt
-import pkg_resources
+import packaging.version
 import requests
 import torch
 
@@ -75,8 +75,8 @@ def download_file_from_zip_url(url, *, file):
 
 def _is_mps_availabe(torch):
     available = False
-    if pkg_resources.parse_version(
-            torch.__version__) >= pkg_resources.parse_version("1.12"):
+    if packaging.version.parse(
+            torch.__version__) >= packaging.version.parse("1.12"):
         if torch.backends.mps.is_available():
             if torch.backends.mps.is_built():
                 available = True
@@ -159,17 +159,17 @@ def requires_package_version(module, version: str):
             the required ``version``.
     """
 
-    required_version = pkg_resources.parse_version(version)
+    required_version = packaging.version.parse(version)
 
     def _requires_package_version(function):
 
         @wraps(function)
         def wrapper(*args, patched_version=None, **kwargs):
             if patched_version is not None:
-                installed_version = pkg_resources.parse_version(
+                installed_version = packaging.version.parse(
                     patched_version)  # Use the patched version if provided
             else:
-                installed_version = pkg_resources.parse_version(
+                installed_version = packaging.version.parse(
                     module.__version__)
 
             if installed_version < required_version:

cebra/integrations/matplotlib.py

@@ -1235,7 +1235,7 @@ def compare_models(
 
     # check the color of the traces
     if color is None:
-        cebra_map = plt.get_cmap(cmap)
+        cebra_map = matplotlib.colormaps.get_cmap(cmap)
         colors = matplotlib.colors.ListedColormap(
             cebra_map.resampled(n_models)(np.arange(n_models))).colors
     else:

cebra/integrations/plotly.py

@@ -87,7 +87,7 @@ def _define_colorscale(self, cmap: str):
         Returns:
             colorscale: List of scaled colors to plot the embeddings
         """
-        colorscale = _convert_cmap2colorscale(matplotlib.cm.get_cmap(cmap))
+        colorscale = _convert_cmap2colorscale(matplotlib.colormaps.get_cmap(cmap))
 
         return colorscale
 

cebra/integrations/sklearn/cebra.py

@@ -28,7 +28,7 @@
 import numpy as np
 import numpy.typing as npt
 import packaging.version
-import pkg_resources
+import importlib.metadata
 import sklearn
 import sklearn.utils.validation as sklearn_utils_validation
 import torch
@@ -1397,7 +1397,7 @@ def save(self,
                             'numpy_version':
                                 np.__version__,
                             'sklearn_version':
-                                pkg_resources.get_distribution("scikit-learn"
+                                importlib.metadata.distribution("scikit-learn"
                                                               ).version
                         }
                     }, filename)

setup.cfg

@@ -39,9 +39,7 @@ install_requires =
     scipy
     torch>=2.4.0
     tqdm
-    # NOTE(stes): Remove pin once https://github.com/AdaptiveMotorControlLab/CEBRA/issues/240
-    # is resolved.
-    matplotlib<3.11
+    matplotlib
     requests
 
 [options.extras_require]

tests/test_plot.py

@@ -25,7 +25,7 @@
 import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
-import pkg_resources
+import packaging.version
 import pytest
 import torch
 from sklearn.exceptions import NotFittedError
@@ -190,8 +190,8 @@ def test_compare_models_with_different_versions(matplotlib_version):
     # minimum version of matplotlib
     minimum_version = "3.6"
 
-    if pkg_resources.parse_version(
-            matplotlib_version) < pkg_resources.parse_version(minimum_version):
+    if packaging.version.parse(
+            matplotlib_version) < packaging.version.parse(minimum_version):
         with pytest.raises(ImportError):
             cebra_plot.compare_models(models=fitted_models,
                                       patched_version=matplotlib_version)

tests/test_plotly.py

@@ -31,7 +31,7 @@
 
 @pytest.mark.parametrize("cmap", ["viridis", "plasma", "inferno", "magma"])
 def test_colorscale(cmap):
-    cmap = matplotlib.cm.get_cmap(cmap)
+    cmap = matplotlib.colormaps.get_cmap(cmap)
     colorscale = cebra_plotly._convert_cmap2colorscale(cmap)
     assert isinstance(colorscale, list)
 

tests/test_sklearn.py

@@ -26,7 +26,7 @@
 import _util
 import _utils_deprecated
 import numpy as np
-import pkg_resources
+import packaging.version
 import pytest
 import sklearn.utils.estimator_checks
 import torch
@@ -1320,8 +1320,8 @@ def test_check_device():
     with pytest.raises(ValueError):
         cebra_sklearn_utils.check_device(device)
 
-    if pkg_resources.parse_version(
-            torch.__version__) >= pkg_resources.parse_version("1.12"):
+    if packaging.version.parse(
+            torch.__version__) >= packaging.version.parse("1.12"):
 
         device = "mps"
         torch.backends.mps.is_available = lambda: True