more xCEBRA docstrings

.github/workflows/docs.yml

@@ -47,7 +47,11 @@ jobs:
  with:
  repository: AdaptiveMotorControlLab/cebra-demos
  path: docs/source/demo_notebooks
- ref: main
+ # NOTE(stes): This is a temporary branch to add the xCEBRA demo notebooks
+ # to the docs. Once the notebooks are merged into main, we can remove this
+ # branch and change the ref to main.
+ # ref: main
+ ref: stes/add-xcebra
 
  - name: Set up Python 3.10
  uses: actions/setup-python@v5

.gitignore

@@ -7,6 +7,16 @@ experiments/sweeps
 exports/
 demo_notebooks/
 assets/
+.remove
+
+# demo run
+.vscode/
+auxiliary_behavior_data.h5
+cebra_model.pt
+data.npz
+grid_search_models/
+neural_data.npz
+saved_models/
 
 # demo run
 .vscode/

Dockerfile

@@ -40,7 +40,7 @@ RUN make dist
 FROM cebra-base
 
 # install the cebra wheel
-ENV WHEEL=cebra-0.5.0-py3-none-any.whl
+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]'

Makefile

@@ -1,4 +1,4 @@
-CEBRA_VERSION := 0.5.0
+CEBRA_VERSION := 0.6.0a1
 
 dist:
 python3 -m pip install virtualenv

NOTICE.yml

@@ -35,3 +35,83 @@
  - 'tests/**/*.py'
  - 'docs/**/*.py'
  - 'conda/**/*.yml'
+
+- header: |
+ CEBRA: Consistent EmBeddings of high-dimensional Recordings using Auxiliary variables
+ © Mackenzie W. Mathis & Steffen Schneider (v0.4.0+)
+ Source code:
+ https://github.com/AdaptiveMotorControlLab/CEBRA
+
+ Please see LICENSE.md for the full license document:
+ https://github.com/AdaptiveMotorControlLab/CEBRA/blob/main/LICENSE.md
+
+ Adapted from https://github.com/rpatrik96/nl-causal-representations/blob/master/care_nl_ica/dep_mat.py,
+ licensed under the following MIT License:
+
+ MIT License
+
+ Copyright (c) 2022 Patrik Reizinger
+
+ Permission is hereby granted, free of charge, to any person obtaining a copy
+ of this software and associated documentation files (the "Software"), to deal
+ in the Software without restriction, including without limitation the rights
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ copies of the Software, and to permit persons to whom the Software is
+ furnished to do so, subject to the following conditions:
+
+ The above copyright notice and this permission notice shall be included in all
+ copies or substantial portions of the Software.
+
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ SOFTWARE.
+
+ include:
+ - 'cebra/attribution/jacobian.py'
+
+
+- header: |
+ CEBRA: Consistent EmBeddings of high-dimensional Recordings using Auxiliary variables
+ © Mackenzie W. Mathis & Steffen Schneider (v0.4.0+)
+ Source code:
+ https://github.com/AdaptiveMotorControlLab/CEBRA
+
+ Please see LICENSE.md for the full license document:
+ https://github.com/AdaptiveMotorControlLab/CEBRA/blob/main/LICENSE.md
+
+ This file contains the PyTorch implementation of Jacobian regularization described in [1].
+ Judy Hoffman, Daniel A. Roberts, and Sho Yaida,
+ "Robust Learning with Jacobian Regularization," 2019.
+ [arxiv:1908.02729](https://arxiv.org/abs/1908.02729)
+
+ Adapted from https://github.com/facebookresearch/jacobian_regularizer/blob/main/jacobian/jacobian.py
+ licensed under the following MIT License:
+
+ MIT License
+
+ Copyright (c) Facebook, Inc. and its affiliates.
+
+ Permission is hereby granted, free of charge, to any person obtaining a copy
+ of this software and associated documentation files (the "Software"), to deal
+ in the Software without restriction, including without limitation the rights
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ copies of the Software, and to permit persons to whom the Software is
+ furnished to do so, subject to the following conditions:
+
+ The above copyright notice and this permission notice shall be included in all
+ copies or substantial portions of the Software.
+
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ SOFTWARE.
+
+ include:
+ - 'cebra/models/jacobian_regularizer.py'

PKGBUILD

@@ -1,7 +1,7 @@
 # Maintainer: Steffen Schneider <stes@hey.com>
 pkgname=python-cebra
 _pkgname=cebra
-pkgver=0.5.0
+pkgver=0.6.0a1
 pkgrel=1
 pkgdesc="Consistent Embeddings of high-dimensional Recordings using Auxiliary variables"
 url="https://cebra.ai"

cebra/__init__.py

@@ -66,7 +66,7 @@
 
 import cebra.integrations.sklearn as sklearn
 
-__version__ = "0.5.0"
+__version__ = "0.6.0a1"
 __all__ = ["CEBRA"]
 __allow_lazy_imports = False
 __lazy_imports = {}

cebra/attribution/__init__.py

@@ -0,0 +1,53 @@
+#
+# CEBRA: Consistent EmBeddings of high-dimensional Recordings using Auxiliary variables
+# © Mackenzie W. Mathis & Steffen Schneider (v0.4.0+)
+# Source code:
+# https://github.com/AdaptiveMotorControlLab/CEBRA
+#
+# Please see LICENSE.md for the full license document:
+# https://github.com/AdaptiveMotorControlLab/CEBRA/blob/main/LICENSE.md
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""Attribution methods for CEBRA.
+
+This module provides tools for computing and analyzing attribution maps in CEBRA models.
+Attribution maps help understand how input features contribute to the model's output
+by analyzing the relationship between input variables and model predictions.
+
+The module includes several attribution methods:
+- Jacobian-based attribution
+- Neuron gradient-based attribution
+- Feature ablation-based attribution
+- Integrated gradients-based attribution
+- Neuron gradient SHAP-based attribution
+
+Each method provides different insights into the model's behavior and can be used
+to analyze different aspects of the feature importance.
+
+This module was added in v0.6.0 and contains attribution methods described and benchmarked
+in [Schneider2025]_.
+
+References
+----------
+.. [Schneider2025] Schneider, S., González Laiz, R., Filippova, A., Frey, M., & Mathis, M. W. (2025).
+ Time-series attribution maps with regularized contrastive learning.
+ The 28th International Conference on Artificial Intelligence and Statistics.
+ https://openreview.net/forum?id=aGrCXoTB4P
+"""
+import cebra.registry
+
+cebra.registry.add_helper_functions(__name__)
+
+from cebra.attribution.attribution_models import *
+from cebra.attribution.jacobian_attribution import *