Apply review comments

Author: stes

.github/workflows/build.yml

@@ -7,7 +7,6 @@ on:
   pull_request:
     branches:
       - main
-      - stes/upgrade-docs-rebased
 
 jobs:
   build:
@@ -54,12 +53,12 @@ jobs:
         run: |
           python -m pip install --upgrade pip setuptools wheel
           python -m pip install torch==${{ matrix.torch-version }} --extra-index-url https://download.pytorch.org/whl/cpu
-          pip install '.[dev,datasets,integrations,xcebra]'
+          pip install '.[dev,datasets,integrations]'
 
       - name: Check sklearn legacy version
         if: matrix.sklearn-version == 'legacy'
         run: |
-          pip install scikit-learn==1.4.2 '.[dev,datasets,integrations,xcebra]'
+          pip install scikit-learn==1.4.2 '.[dev,datasets,integrations]'
 
       - name: Run the formatter
         run: |

.github/workflows/docs.yml

@@ -9,7 +9,6 @@ on:
       - main
       - public
       - dev
-      - stes/upgrade-docs-rebased
     paths:
       - '**.py'
       - '**.ipynb'

Dockerfile

@@ -43,7 +43,7 @@ FROM cebra-base
 ENV WHEEL=cebra-0.6.0a1-py3-none-any.whl
 WORKDIR /build
 COPY --from=wheel /build/dist/${WHEEL} .
-RUN pip install --no-cache-dir ${WHEEL}'[dev,integrations,datasets,xcebra]'
+RUN pip install --no-cache-dir ${WHEEL}'[dev,integrations,datasets]'
 RUN rm -rf /build
 
 # add the repository

cebra/data/single_session.py

@@ -359,8 +359,10 @@ def __post_init__(self):
         #            e.g. integrating the FAISS dataloader back in.
         super().__post_init__()
 
-        # BEHAVIOR DISTRIBUTION
+        self._init_behavior_distribution()
+        self._init_time_distribution()
 
+    def _init_behavior_distribution(self):
         if self.conditional == "time":
             self.behavior_distribution = cebra.distributions.TimeContrastive(
                 time_offset=self.time_offset,
@@ -385,7 +387,8 @@ def __post_init__(self):
                 device=self.device,
             )
 
-        # TIME DISTRIBUTION
+    def _init_time_distribution(self):
+
         if self.time_distribution == "time":
             self.time_distribution = cebra.distributions.TimeContrastive(
                 time_offset=self.time_offset,
@@ -403,9 +406,10 @@ def __post_init__(self):
             self.time_distribution = cebra.distributions.DeltaNormalDistribution(
                 self.dataset.continuous_index, self.delta, device=self.device)
 
-        elif self.time_distribution == "delta_vmf":
-            self.time_distribution = cebra.distributions.DeltaVMFDistribution(
-                self.dataset.continuous_index, self.delta, device=self.device)
+        # TODO(stes): Add this distribution from internal xCEBRA codebase at a later point
+        #elif self.time_distribution == "delta_vmf":
+        #    self.time_distribution = cebra.distributions.DeltaVMFDistribution(
+        #        self.dataset.continuous_index, self.delta, device=self.device)
         else:
             raise ValueError
 

cebra/models/multiobjective.py

@@ -29,12 +29,13 @@
 import cebra.models.model as cebra_models_base
 
 
-def create_multiobjective_model(module, **kwargs) -> "MultiobjectiveModel":
+def create_multiobjective_model(module,
+                                **kwargs) -> "SubspaceMultiobjectiveModel":
     assert isinstance(module, cebra_models_base.Model)
     if isinstance(module, cebra.models.ConvolutionalModelMixin):
-        return MultiobjectiveConvolutionalModel(module=module, **kwargs)
+        return SubspaceMultiobjectiveConvolutionalModel(module=module, **kwargs)
     else:
-        return MultiobjectiveModel(module=module, **kwargs)
+        return SubspaceMultiobjectiveModel(module=module, **kwargs)
 
 
 def check_slices_for_gaps(slice_list):
@@ -106,7 +107,7 @@ def forward(self, inp):
         return inp / torch.norm(inp, dim=1, keepdim=True)
 
 
-class LegacyMultiobjectiveModel(nn.Module):
+class MultiobjectiveModel(nn.Module):
     """Wrapper around contrastive learning models to all training with multiple objectives
 
     Multi-objective training splits the last layer's feature representation into multiple
@@ -128,6 +129,13 @@ class LegacyMultiobjectiveModel(nn.Module):
 
     TODO:
         - Update nn.Module type annotation for ``module`` to cebra.models.Model
+
+    Note:
+        This model will be deprecated in a future version. Please use the functionality in
+        :py:mod:`cebra.models.multiobjective` instead, which provides more versatile
+        multi-objective training capabilities. Instantiation of this model will raise a
+        deprecation warning. The new model is :py:class:`cebra.models.multiobjective.SubspaceMultiobjectiveModel`
+        which allows for unlimited subspaces and better configuration of the feature ranges.
     """
 
     class Mode:
@@ -240,7 +248,7 @@ def forward(self, inputs):
         return tuple(outputs)
 
 
-class MultiobjectiveModel(nn.Module):
+class SubspaceMultiobjectiveModel(nn.Module):
     """Wrapper around contrastive learning models to all training with multiple objectives
 
     Multi-objective training splits the last layer's feature representation into multiple
@@ -354,7 +362,6 @@ def forward(self, inputs):
             return output
 
 
-class MultiobjectiveConvolutionalModel(MultiobjectiveModel,
-                                       cebra_models_base.ConvolutionalModelMixin
-                                      ):
+class SubspaceMultiobjectiveConvolutionalModel(
+        SubspaceMultiobjectiveModel, cebra_models_base.ConvolutionalModelMixin):
     pass

cebra/registry.py

@@ -287,10 +287,6 @@ def get_options(pattern: str = None,
         raise RuntimeError(
             f"Registry could not be successfully registered: {module}.")
 
-    # NOTE(stes): Used in xCEBRA initially. If you see this note past 0.6.0, please remove it
-    # as the functionality is no longer needed.
-    #return register, parametrize, init, get_options
-
 
 def add_docstring(module: Union[types.ModuleType, str]):
     """Apply additional information about configuration options to registry modules.

cebra/solver/base.py

@@ -32,6 +32,7 @@
 
 import abc
 import os
+import warnings
 from typing import Callable, Dict, List, Literal, Optional
 
 import literate_dataclasses as dataclasses
@@ -367,11 +368,19 @@ class MultiobjectiveSolver(Solver):
             for time contrastive learning.
         renormalize_features: If ``True``, normalize the behavior and time
             contrastive features individually before computing similarity scores.
+        ignore_deprecation_warning: If ``True``, suppress the deprecation warning.
+
+    Note:
+        This solver will be deprecated in a future version. Please use the functionality in
+        :py:mod:`cebra.solver.multiobjective` instead, which provides more versatile
+        multi-objective training capabilities. Instantiation of this solver will raise a
+        deprecation warning.
     """
 
     num_behavior_features: int = 3
     renormalize_features: bool = False
     output_mode: Literal["overlapping", "separate"] = "overlapping"
+    ignore_deprecation_warning: bool = False
 
     @property
     def num_time_features(self):
@@ -383,8 +392,15 @@ def num_total_features(self):
 
     def __post_init__(self):
         super().__post_init__()
+        if not self.ignore_deprecation_warning:
+            warnings.warn(
+                "MultiobjectiveSolver is deprecated since CEBRA 0.6.0 and will be removed in a future version. "
+                "Use the new functionality in cebra.solver.multiobjective instead, which is more versatile. "
+                "If you see this warning when using the scikit-learn interface, no action is required.",
+                DeprecationWarning,
+                stacklevel=2)
         self._check_dimensions()
-        self.model = cebra.models.LegacyMultiobjectiveModel(
+        self.model = cebra.models.MultiobjectiveModel(
             self.model,
             dimensions=(self.num_behavior_features, self.model.num_output),
             renormalize=self.renormalize_features,

cebra/solver/single_session.py

@@ -22,7 +22,6 @@
 """Single session solvers embed a single pair of time series."""
 
 import copy
-from typing import Dict
 
 import literate_dataclasses as dataclasses
 import torch
@@ -131,16 +130,6 @@ def _inference(self, batch):
 class SingleSessionHybridSolver(abc_.MultiobjectiveSolver):
     """Single session training, contrasting neural data against behavior."""
 
-    log: Dict = dataclasses.field(default_factory=lambda: ({
-        "behavior_pos": [],
-        "behavior_neg": [],
-        "behavior_total": [],
-        "time_pos": [],
-        "time_neg": [],
-        "time_total": [],
-        "temperature": []
-    }))
-
     _variant_name = "single-session-hybrid"
 
     def _inference(self, batch: cebra.data.Batch) -> cebra.data.Batch:

docs/source/api.rst

@@ -38,9 +38,9 @@ these components in other contexts and research code bases.
    api/pytorch/distributions
    api/pytorch/models
    api/pytorch/helpers
-   api/xcebra/multiobjective
-   api/xcebra/regularized
-   api/xcebra/attribution
+   api/pytorch/multiobjective
+   api/pytorch/regularized
+   api/pytorch/attribution
 
 .. toctree::
    :hidden: