Add unified encoder pytorch implementation

.dockerignore

@@ -5,3 +5,5 @@ tests/
 third_party/
 tools/
 PKGBUILD
+
+!docs/requirements.txt

.github/workflows/build.yml

@@ -10,6 +10,7 @@ on:
 
 jobs:
   build:
+    timeout-minutes: 30
     strategy:
       fail-fast: true
       matrix:
@@ -18,27 +19,33 @@ jobs:
         # We aim to support the versions on pytorch.org
         # as well as selected previous versions on
         # https://pytorch.org/get-started/previous-versions/
-        torch-version: ["2.2.2", "2.4.0"]
+        torch-version: ["2.4.0", "2.6.0"]
+        sklearn-version: ["latest"]
         include:
           - os: windows-latest
             torch-version: 2.4.0
             python-version: "3.10"
+            sklearn-version: "latest"
+          - os: ubuntu-latest
+            torch-version: 2.4.0
+            python-version: "3.10"
+            sklearn-version: "legacy"
 
     runs-on: ${{ matrix.os }}
 
     steps:
       - name: Cache dependencies
         id: pip-cache
-        uses: actions/cache@v3
+        uses: actions/cache@v4
         with:
           path: ~/.cache/pip
-          key: pip-os_${{ runner.os }}-python_${{ matrix.python-version }}-torch_${{ matrix.torch-version }}
+          key: pip-os_${{ runner.os }}-python_${{ matrix.python-version }}-torch_${{ matrix.torch-version }}-sklearn_${{ matrix.sklearn-version }}
 
       - name: Checkout code
         uses: actions/checkout@v2
 
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
 
@@ -48,6 +55,11 @@ jobs:
           python -m pip install torch==${{ matrix.torch-version }} --extra-index-url https://download.pytorch.org/whl/cpu
           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]'
+
       - name: Run the formatter
         run: |
           make format
@@ -56,6 +68,10 @@ jobs:
         run: |
           make codespell
 
+      - name: Check the documentation coverage
+        run: |
+          make interrogate
+
       - name: Check CITATION.cff validity
         run: |
           cffconvert --validate

.github/workflows/docs.yml

@@ -7,8 +7,6 @@ on:
   pull_request:
     branches:
       - main
-      - public
-      - dev
 
 jobs:
   build:
@@ -17,7 +15,7 @@ jobs:
     steps:
       - name: Cache dependencies
         id: pip-cache
-        uses: actions/cache@v3
+        uses: actions/cache@v4
         with:
           path: ~/.cache/pip
           key: ${{ runner.os }}-pip
@@ -51,28 +49,33 @@ jobs:
           path: docs/source/demo_notebooks
           ref: main
 
-      - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v2
+      - name: Set up Python 3.10
+        uses: actions/setup-python@v5
         with:
-          python-version: ${{ matrix.python-version }}
+          python-version: "3.10"
 
       - name: Install package
         run: |
           python -m pip install --upgrade pip setuptools wheel
           # NOTE(stes) Pandoc version must be at least (2.14.2) but less than (4.0.0).
-          # as of 29/10/23. Ubuntu 22.04 which is used for ubuntu-latest only has an 
+          # as of 29/10/23. Ubuntu 22.04 which is used for ubuntu-latest only has an
           # old pandoc version (2.9.). We will hence install the latest version manually.
           # previou: sudo apt-get install -y pandoc
-          wget https://github.com/jgm/pandoc/releases/download/3.1.9/pandoc-3.1.9-1-amd64.deb 
-          sudo dpkg -i pandoc-3.1.9-1-amd64.deb 
-          rm pandoc-3.1.9-1-amd64.deb 
-          pip install torch --extra-index-url https://download.pytorch.org/whl/cpu
-          pip install '.[docs]'
+          # NOTE(stes): Updated to latest version as of 17/04/2025, v3.6.4.
+          wget -q https://github.com/jgm/pandoc/releases/download/3.6.4/pandoc-3.6.4-1-amd64.deb
+          sudo dpkg -i pandoc-3.6.4-1-amd64.deb
+          rm pandoc-3.6.4-1-amd64.deb
+          pip install -r docs/requirements.txt
+
+      - name: Check software versions
+        run: |
+          sphinx-build --version
+          pandoc --version
 
       - name: Build docs
         run: |
           ls docs/source/cebra-figures
-          # later also add the -n option to check for broken links
+          export SPHINXBUILD="sphinx-build"
           export SPHINXOPTS="-W --keep-going -n"
           make docs
 

.github/workflows/release-pypi.yml

@@ -23,11 +23,18 @@ jobs:
     steps:
       - name: Cache dependencies
         id: pip-cache
-        uses: actions/cache@v3
+        uses: actions/cache@v4
         with:
           path: ~/.cache/pip
           key: ${{ runner.os }}-pip
 
+      - name: Install dependencies
+        run: |
+          pip install --upgrade pip
+          pip install wheel
+          # NOTE(stes) see https://github.com/pypa/twine/issues/1216#issuecomment-2629069669
+          pip install "packaging>=24.2"
+
       - name: Checkout code
         uses: actions/checkout@v3
 

.gitignore

@@ -7,6 +7,25 @@ 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/
+auxiliary_behavior_data.h5
+cebra_model.pt
+data.npz
+grid_search_models/
+neural_data.npz
+saved_models/
 
 # Binary files
 *.png

AUTHORS.md

@@ -0,0 +1,26 @@
+<img src="https://github.com/user-attachments/assets/1f327e57-8ee1-4a2f-afd3-2bbce885c2f8" width="200"/>
+
+
+
+CEBRA was initially developed by **Mackenzie Mathis** and **Steffen Schneider** (2021+), who are co-inventors on the patent application [WO2023143843](https://infoscience.epfl.ch/entities/patent/0d9debed-4d22-47b7-bad1-f211e7010323).
+**Jin Hwa Lee** contributed significantly to our first paper:
+
+> **Schneider, S., Lee, J.H., & Mathis, M.W.**
+> [*Learnable latent embeddings for joint behavioural and neural analysis.*](https://doi.org/10.1038/s41586-023-06031-6)
+> Nature 617, 360–368 (2023)
+
+CEBRA is actively developed by [**Mackenzie Mathis**](https://www.mackenziemathislab.org/) and [**Steffen Schneider**](https://dynamical-inference.ai/) and their labs.
+
+It is a publicly available tool that has benefited from contributions and suggestions from many individuals: [CEBRA/graphs/contributors](https://github.com/AdaptiveMotorControlLab/CEBRA/graphs/contributors).
+
+## CEBRA Extensions
+
+### 2023
+- **Steffen Schneider, Rodrigo González Laiz, Markus Frey, Mackenzie W. Mathis**
+  [*Identifiable attribution maps using regularized contrastive learning.*](https://sslneurips23.github.io/paper_pdfs/paper_80.pdf)
+  NeurIPS 4th Workshop on Self-Supervised Learning: Theory and Practice (2023)
+
+### 2025
+- **Steffen Schneider, Rodrigo González Laiz, Anastasiia Filippova, Markus Frey, Mackenzie W. Mathis**
+  [*Time-series attribution maps with regularized contrastive learning.*](https://openreview.net/forum?id=aGrCXoTB4P)
+  AISTATS (2025)

Dockerfile

@@ -40,7 +40,7 @@ RUN make dist
 FROM cebra-base
 
 # install the cebra wheel
-ENV WHEEL=cebra-0.4.0-py2.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.4.0
+CEBRA_VERSION := 0.6.0a1
 
 dist:
 	python3 -m pip install virtualenv
@@ -55,7 +55,7 @@ interrogate:
 		--ignore-private \
 		--ignore-magic \
 		--omit-covered-files \
-		-f 90 \
+		-f 80 \
 		cebra
 
 # Build documentation using sphinx