Fix linting issues (#185)

* Apply auto-fixes

* Fix issues in allen datasets

- lines too long
- unused variables
- missing imports
- duplicate names for classes

* Fix line length+minor issues in plot integr.

* Fix linting issues in sklearn integration

* Fix linting issues in datasets

- missing paths
- long lines
- unused variables
- typos

* Fix minor linting issues

* Fix docstrings

* fix formatting issue in docstring

* Fix plotly docstrings

* Fix missing import

Author: stes

cebra/__init__.py

@@ -33,7 +33,7 @@
     from cebra.integrations.sklearn.decoder import L1LinearRegressor
 
     is_sklearn_available = True
-except ImportError as e:
+except ImportError:
     # silently fail for now
     pass
 
@@ -42,7 +42,7 @@
     from cebra.integrations.matplotlib import *
 
     is_matplotlib_available = True
-except ImportError as e:
+except ImportError:
     # silently fail for now
     pass
 
@@ -51,7 +51,7 @@
     from cebra.integrations.plotly import *
 
     is_plotly_available = True
-except ImportError as e:
+except ImportError:
     # silently fail for now
     pass
 
@@ -92,11 +92,11 @@ def __getattr__(key):
 
         return CEBRA
     elif key == "KNNDecoder":
-        from cebra.integrations.sklearn.decoder import KNNDecoder
+        from cebra.integrations.sklearn.decoder import KNNDecoder  # noqa: F811
 
         return KNNDecoder
     elif key == "L1LinearRegressor":
-        from cebra.integrations.sklearn.decoder import L1LinearRegressor
+        from cebra.integrations.sklearn.decoder import L1LinearRegressor  # noqa: F811
 
         return L1LinearRegressor
     elif not key.startswith("_"):

cebra/__main__.py

@@ -27,11 +27,7 @@
 import argparse
 import sys
 
-import numpy as np
-import torch
-
 import cebra
-import cebra.distributions as cebra_distr
 
 
 def train(parser, kwargs):

cebra/config.py

@@ -21,7 +21,6 @@
 #
 import argparse
 import json
-from dataclasses import MISSING
 from typing import Literal, Optional
 
 import literate_dataclasses as dataclasses

cebra/data/base.py

@@ -22,11 +22,8 @@
 """Base classes for datasets and loaders."""
 
 import abc
-import collections
-from typing import List
 
 import literate_dataclasses as dataclasses
-import numpy as np
 import torch
 
 import cebra.data.assets as cebra_data_assets

cebra/data/datasets.py

@@ -21,21 +21,15 @@
 #
 """Pre-defined datasets."""
 
-import abc
-import collections
 import types
 from typing import List, Literal, Optional, Tuple, Union
 
-import literate_dataclasses as dataclasses
 import numpy as np
 import numpy.typing as npt
 import torch
-from numpy.typing import NDArray
 
 import cebra.data as cebra_data
-import cebra.distributions
-from cebra.data.datatypes import Batch
-from cebra.data.datatypes import BatchIndex
+import cebra.helper as cebra_helper
 
 
 class TensorDataset(cebra_data.SingleSessionDataset):
@@ -75,8 +69,7 @@ def __init__(self,
                  device: str = "cpu"):
         super().__init__(device=device)
         self.neural = self._to_tensor(neural, check_dtype="float").float()
-        self.continuous = self._to_tensor(continuous,
-                                          check_dtype="float")
+        self.continuous = self._to_tensor(continuous, check_dtype="float")
         self.discrete = self._to_tensor(discrete, check_dtype="integer")
         if self.continuous is None and self.discrete is None:
             raise ValueError(
@@ -104,10 +97,11 @@ def _to_tensor(
         if isinstance(array, np.ndarray):
             array = torch.from_numpy(array)
         if check_dtype is not None:
-            if (check_dtype == "int" and not cebra.helper._is_integer(array)
+            if (check_dtype == "int" and not cebra_helper._is_integer(array)
                ) or (check_dtype == "float" and
-                     not cebra.helper._is_floating(array)):
-                raise TypeError(f"Array has type {array.dtype} instead of {check_dtype}.")
+                     not cebra_helper._is_floating(array)):
+                raise TypeError(
+                    f"Array has type {array.dtype} instead of {check_dtype}.")
         return array
 
     @property

cebra/data/datatypes.py

@@ -20,9 +20,6 @@
 # limitations under the License.
 #
 import collections
-from typing import Tuple
-
-import torch
 
 __all__ = ["Batch", "BatchIndex", "Offset"]
 

cebra/data/helper.py

@@ -94,10 +94,15 @@ class OrthogonalProcrustesAlignment:
 
     For each dataset, the data and labels to align the data on is provided.
 
-    1. The ``top_k`` indexes of the labels to align (``label``) that are the closest to the labels of the reference dataset (``ref_label``) are selected and used to sample from the dataset to align (``data``).
-    2. ``data`` and ``ref_data`` (the reference dataset) are subsampled to the same number of samples ``subsample``.
-    3. The orthogonal mapping is computed, using :py:func:`scipy.linalg.orthogonal_procrustes`, on those subsampled datasets.
-    4. The resulting orthongonal matrix ``_transform`` can be used to map the original ``data`` to the ``ref_data``.
+    1. The ``top_k`` indexes of the labels to align (``label``) that are the closest to
+       the labels of the reference dataset (``ref_label``) are selected and used to sample
+       from the dataset to align (``data``).
+    2. ``data`` and ``ref_data`` (the reference dataset) are subsampled to the same number
+       of samples ``subsample``.
+    3. The orthogonal mapping is computed, using :py:func:`scipy.linalg.orthogonal_procrustes`,
+       on those subsampled datasets.
+    4. The resulting orthongonal matrix ``_transform`` can be used to map the original ``data``
+       to the ``ref_data``.
 
     Note:
         ``data`` and ``ref_data`` can be of different sample size (axis 0) but **must** have the same number
@@ -181,14 +186,14 @@ def fit(
         elif ref_data.shape[0] == data.shape[0] and (ref_label is None or
                                                      label is None):
             raise ValueError(
-                f"Missing labels: the data to align are the same shape but you provided only "
-                f"one of the sets of labels. Either provide both the reference and alignment "
-                f"labels or none.")
+                "Missing labels: the data to align are the same shape but you provided only "
+                "one of the sets of labels. Either provide both the reference and alignment "
+                "labels or none.")
         else:
             if ref_label is None or label is None:
                 raise ValueError(
-                    f"Missing labels: the data to align are not the same shape, "
-                    f"provide labels to align the data and reference data.")
+                    "Missing labels: the data to align are not the same shape, "
+                    "provide labels to align the data and reference data.")
 
             if len(ref_label.shape) == 1:
                 ref_label = np.expand_dims(ref_label, axis=1)

cebra/data/load.py

@@ -663,7 +663,8 @@ def load(
         - if no key is provided, the first data structure found upon iteration of the collection will be loaded;
         - if a key is provided, it needs to correspond to an existing item of the collection;
         - if a key is provided, the data value accessed needs to be a data structure;
-        - the function loads data for only one data structure, even if the file contains more. The function can be called again with the corresponding key to get the other ones.
+        - the function loads data for only one data structure, even if the file contains more. The function can be
+          called again with the corresponding key to get the other ones.
 
     Args:
         file: The path to the given file to load, in a supported format.

cebra/data/multi_session.py

@@ -22,11 +22,9 @@
 """Datasets and loaders for multi-session training."""
 
 import abc
-import collections
 from typing import List
 
 import literate_dataclasses as dataclasses
-import numpy as np
 import torch
 
 import cebra.data as cebra_data

cebra/data/single_session.py

@@ -26,12 +26,9 @@
 """
 
 import abc
-import collections
 import warnings
-from typing import List
 
 import literate_dataclasses as dataclasses
-import numpy as np
 import torch
 
 import cebra.data as cebra_data
@@ -353,18 +350,16 @@ def __post_init__(self):
         #            here might be sub-optimal. The final behavior should be determined after
         #            e.g. integrating the FAISS dataloader back in.
         super().__post_init__()
-        index = self.index.to(self.device)
 
         if self.conditional != "time_delta":
             raise NotImplementedError(
-                f"Hybrid training is currently only implemented using the ``time_delta`` "
-                f"continual distribution.")
+                "Hybrid training is currently only implemented using the ``time_delta`` "
+                "continual distribution.")
 
         self.time_distribution = cebra.distributions.TimeContrastive(
             time_offset=self.time_offset,
             num_samples=len(self.dataset.neural),
-            device=self.device,
-        )
+            device=self.device)
         self.behavior_distribution = cebra.distributions.TimedeltaDistribution(
             self.dataset.continuous_index, self.time_offset, device=self.device)