Add unified encoder pytorch implementation #251

CeliaBenquet · 2025-05-01T14:36:20Z

This PR adds a PyTorch implementation of a unified CEBRA encoder, which is composed of:

A new sampling scheme that samples across all sessions so that they can be aligned on the neuron axis to train a single encoder.
A unified Dataset and Loader, adapted to the new sampling scheme.
A unified Solver that considers multiple sessions to be aligned at inference.
A new masked modeling training option, with different types of masking.

🚧 A preprint is pending "A Unified Encoder for Modeling Neural Dynamics with Contrastive Learning" by Célia Benquet, Hossein Mirzaei, Steffen Schneider, Mackenzie W. Mathis.

* apply auto-fixes * Fix linting errors in tests/ * Fix version check

…rControlLab#195)

* Add support for new __sklearn_tags__ * fix inheritance order * Add more tests * fix added test

…veMotorControlLab#212)

…earn>=1.6 (AdaptiveMotorControlLab#206)

) * Started implementing improved goodness of fit implementation * add tests and improve implementation * Fix examples * Fix docstring error * Handle batch size = None for goodness of fit computation * adapt GoF implementation * Fix docstring tests * Update docstring for goodness_of_fit_score Co-authored-by: Célia Benquet <[email protected]> * add annotations to goodness_of_fit_history Co-authored-by: Célia Benquet <[email protected]> * fix typo Co-authored-by: Célia Benquet <[email protected]> * improve err message Co-authored-by: Célia Benquet <[email protected]> * make numerical test less conversative * Add tests for exception handling * fix tests --------- Co-authored-by: Célia Benquet <[email protected]>

…trolLab#221) * Drop numpy constraint * Implement workaround for pytables * better error message * pin numpy only for python 3.9 * update dependencies * Upgrade torch version * Fix based on python version * Add support for torch.load with weights_only=True * Implement safe loading for torch models starting in torch 2.6 * Fix windows specs * fix docstring * Revert changes to loading logic

* Make bump_version script runnable on MacOS * Bump version to 0.5.0rc1 * fix minor formatting issues * remove commented code --------- Co-authored-by: Mackenzie Mathis <[email protected]>

* force packaging upgrade to 24.2 for twine * Bump version to 0.5.0rc2 * remove universal compatibility option * revert tag * adapt files to new wheel name due to py3

This is a lazy solution to AdaptiveMotorControlLab#223

…olLab#227) * Change text consistency max from 99 to 100 * Update cebra/integrations/matplotlib.py --------- Co-authored-by: Mackenzie Mathis <[email protected]> Co-authored-by: Steffen Schneider <[email protected]>

…Lab#230) Update assets.py - mkdir was failing in 0.5.0rc1; attempt to fix

@stes

* user note added to usage.rst - link added * Update usage.rst - more detailed note on the effect of temp. * Update usage.rst - add in temp to demo model - testout put thanks @stes * Update docs/source/usage.rst Co-authored-by: Steffen Schneider <[email protected]> * Update docs/source/usage.rst Co-authored-by: Steffen Schneider <[email protected]> * Update docs/source/usage.rst Co-authored-by: Steffen Schneider <[email protected]> --------- Co-authored-by: Steffen Schneider <[email protected]>

* Update installation.rst - python 3.9+ * Update index.rst * Update figures.rst * Update index.rst -typo fix * Update usage.rst - update suggestion on data split * Update docs/source/usage.rst Co-authored-by: Steffen Schneider <[email protected]> * Update usage.rst - indent error fixed * Update usage.rst - changed infoNCE to new GoF * Update usage.rst - finx numpy() doctest * Update usage.rst - small typo fix (label) * Update usage.rst --------- Co-authored-by: Steffen Schneider <[email protected]>

…lLab#231)

- adding link to new notebook icon

* website refresh

* Improve build setup for docs * update pydata theme options * Add README for docs folder * Fix demo notebook build * Finish build setup * update git workflow * add timeout to workflow * add timeout also to docs build * switch build back to sphinx for gh actions * attempt to fix build workflow * update to sphinx-build * fix build workflow * fix indent error * fix build system * revert demos to main * increase timeout to 30

* Allow indexing of the cebra docs * Fix docs workflow

* Add multiobjective solver and regularized training (#783) * Add multiobjective solver and regularized training * Add example for multiobjective training * Add jacobian regularizer and SAM * update license headers * add api draft for multiobjective training * add all necessary modules to run the complete xcebra pipeline * add notebooks to reproduce xcebra pipeline * add first working notebook * add notebook with hybrid learning * add notebook with creation of synthetic data * add notebook with hybrid training * add plot with R2 for different parts of the embedding * add new API * update api wrapper with more checks and messages * add tests and notebook with new api * merge xcebra into attribution * separate xcebra dataset from cebra * some minor refactoring of cebra dataset * separate xcebra loader from cebra * remove xcebra distributions from cebra * minor refactoring with distributions * separate xcebra criterions from cebra * minor refactoring on criterion * separate xcebra models/criterions/layers from cebra * refactoring multiobjective * more refactoring... * separate xcebra solvers from cebra * more refactoring * move xcebra to its own package * move more files into xcebra package * more files and remove changes with the registry * remove unncessary import * add folder structure * move back distributions * add missing init * remove wrong init * make loader and dataset run with new imports * making it run! * make attribution run * Run pre-commit * move xcebra repo one level up * update gitignore and add __init__ from data * add init to distributions * add correct init for attribution pacakge * add correct init for model package * fix remaining imports * fix tests * add examples back to xcebra repo * update imports from graphs_xcebra * add setup.py to create a package * update imports of graph_xcebra * update notebooks * Formatting code for submission Co-authored-by: Rodrigo Gonzalez <[email protected]> * move test into xcebra * Add README * move distributions back to main package * clean up examples * adapt tests * Add LICENSE * add train/eval notebook again * add notebook with clean results * rm synthetic data * change name from xcebra to regcl * change names of modules and adapt imports * change name from graphs_xcebra to synthetic_data * Integrate into CEBRA * Fix remaining imports and make notebook runnable * Add dependencies, add version flag * Remove synthetic data files * reset dockerfile, move vmf * apply pre-commit * Update notice * add some docstrings * Apply license headers * add new scd notebook * add notebook with scd --------- Co-authored-by: Steffen Schneider <[email protected]> * Fix tests * bump version * update dockerfile * fix progress bar * remove outdated test * rename models * Apply fixes to pass ruff tests * Fix typos * Update license headers, fix additional ruff errors * remove unused comment * rename regcl in codebase * change regcl name in dockerfile * Improve attribution module * Fix imports name naming * add basic integration test * temp disable of binary check * Add legacy multiobjective model for backward compat * add synth import back in * Fix docstrings and type annot in cebra/models/jacobian_regularizer.py * add xcebra to tests * add missing cvxpy dep * fix docstrings * more docstrings to fix attr error * Improve build setup for docs * update pydata theme options * Add README for docs folder * Fix demo notebook build * Finish build setup * update git workflow * Move demo notebooks to CEBRA-demos repo See AdaptiveMotorControlLab/CEBRA-demos#28 * revert unneeded changes in solver * formatting in solver * further minimize solver diff * Revert unneeded updates to the solver * fix citation * fix docs build, missing refs * remove file dependency from xcebra int test * remove unneeded change in registry * update gitignore * update docs * exclude some assets * include binary file check again * add timeout to workflow * add timeout also to docs build * switch build back to sphinx for gh actions * pin sphinx version in setup.cfg * attempt workflow fix * attempt to fix build workflow * update to sphinx-build * fix build workflow * fix indent error * fix build system * revert demos to main * adapt workflow for testing * bump version to 0.6.0rc1 * format imports * docs writing * enable build on dev branch * fix some review comments * extend multiobjective docs * Set version to alpha * make tempdir platform independent * Remove ratinabox and ephysiopy as deps * Apply review comments * Update Makefile - setting coverage threshold to 80% to not delay good code being made public. In the near future this can be fixed and raised again to 90%. --------- Co-authored-by: Steffen Schneider <[email protected]> Co-authored-by: Steffen Schneider <[email protected]> Co-authored-by: Mackenzie Mathis <[email protected]>

* apply auto-fixes * Fix linting errors in tests/ * Fix version check

…rControlLab#195)

* Add support for new __sklearn_tags__ * fix inheritance order * Add more tests * fix added test

…veMotorControlLab#212)

…earn>=1.6 (AdaptiveMotorControlLab#206)

…nd-padding

* Update conf.py - adding link to xCEBRA link for demo notebook * Update conf.py * Update conf.py * Update conf.py

Update docs.yml

…rolLab#247)

Co-authored-by: Steffen Schneider <[email protected]>

stes and others added 30 commits October 27, 2024 19:23

Fix linting errors in tests (AdaptiveMotorControlLab#188)

9898850

* apply auto-fixes * Fix linting errors in tests/ * Fix version check

Fix scikit-learn reference in conda environment files (AdaptiveMoto…

36a91c7

…rControlLab#195)

Add support for new __sklearn_tags__ (AdaptiveMotorControlLab#205)

5f46c32

* Add support for new __sklearn_tags__ * fix inheritance order * Add more tests * fix added test

Update workflows to actions/setup-python@v5, actions/cache@v4 (Adapti…

7a4d3fc

…veMotorControlLab#212)

Fix deprecation warning force_all_finite -> ensure_all_finite for skl…

a79c2de

…earn>=1.6 (AdaptiveMotorControlLab#206)

Add tests to check legacy model loading (AdaptiveMotorControlLab#214)

7e74eda

Release 0.5.0rc1 (AdaptiveMotorControlLab#189)

bea2c04

* Make bump_version script runnable on MacOS * Bump version to 0.5.0rc1 * fix minor formatting issues * remove commented code --------- Co-authored-by: Mackenzie Mathis <[email protected]>

Fix pypi action (AdaptiveMotorControlLab#222)

c32ed67

* force packaging upgrade to 24.2 for twine * Bump version to 0.5.0rc2 * remove universal compatibility option * revert tag * adapt files to new wheel name due to py3

Update base.py (AdaptiveMotorControlLab#224)

f99530c

This is a lazy solution to AdaptiveMotorControlLab#223

Update assets.py --> force check for parent dir (AdaptiveMotorControl…

b713387

…Lab#230) Update assets.py - mkdir was failing in 0.5.0rc1; attempt to fix

render plotly in our docs, show code/doc version (AdaptiveMotorContro…

b677e67

…lLab#231)

Update layout.html (AdaptiveMotorControlLab#233)

37ed6f5

Update conf.py (AdaptiveMotorControlLab#234)

09b8974

- adding link to new notebook icon

Refactoring setup.cfg (AdaptiveMotorControlLab#228)

aa0db43

Home page landing update (AdaptiveMotorControlLab#235)

4901966

* website refresh

v0.5.0 (AdaptiveMotorControlLab#238)

b2357fd

Allow indexing of the cebra docs (AdaptiveMotorControlLab#242)

1127432

* Allow indexing of the cebra docs * Fix docs workflow

Fix broken docs coverage workflows (AdaptiveMotorControlLab#246)

d86ccf0

Fix linting errors in tests (AdaptiveMotorControlLab#188)

a09d123

* apply auto-fixes * Fix linting errors in tests/ * Fix version check

Fix scikit-learn reference in conda environment files (AdaptiveMoto…

521f003

…rControlLab#195)

Add support for new __sklearn_tags__ (AdaptiveMotorControlLab#205)

46610e3

* Add support for new __sklearn_tags__ * fix inheritance order * Add more tests * fix added test

Update workflows to actions/setup-python@v5, actions/cache@v4 (Adapti…

e8004ba

…veMotorControlLab#212)

Fix deprecation warning force_all_finite -> ensure_all_finite for skl…

ddc00f4

…earn>=1.6 (AdaptiveMotorControlLab#206)

stes and others added 29 commits April 23, 2025 16:09

fix linting errors

3acbdf4

Run isort, ruff, yapf

5745449

Merge remote-tracking branch 'upstream/main' into batched-inference-a…

fa3cd3e

…nd-padding

Update conf.py (AdaptiveMotorControlLab#237)

f082a1c

* Update conf.py - adding link to xCEBRA link for demo notebook * Update conf.py * Update conf.py * Update conf.py

Update docs.yml to build from main (AdaptiveMotorControlLab#248)

d303e50

Update docs.yml

Update installation.rst --> add link to docker hub (AdaptiveMotorCont…

73f90ee

…rolLab#247)

Merge branch 'main' into batched-inference-and-padding

1453885

Fix gaussian mixture dataset import

acd2111

Fix all tests but xcebra tests

217a8a7

Fix pytorch API usage example

a1218aa

Make xCEBRA compatible with the batched inference & padding in solver

64d1db8

Add some tests on transform() with xCEBRA

9875a38

Add some docstrings and typings and clean unnecessary changes

65fc455

Implement review comments

1d0c498

Fix sklearn test

4a25899

Initial pass at integrating unifiedCEBRA

b8945ae

Add name in NOTE

0d56e44

Co-authored-by: Steffen Schneider <[email protected]>

Implement reviews on tests and typing

c5dc011

Fix import errors

c9fa5c8

Merge branch 'batched-inference-and-padding' into unified-cebra

9ba22bc

Add select_model to aux solvers

4632c04

Merge branch 'batched-inference-and-padding' into unified-cebra

a52f502

Fix tests

c22e40e

Add mask tests

e8a1877

Fix docs error

22e3c47

Merge branch 'batched-inference-and-padding' into unified-cebra

464f4aa

Remove masking init()

57c9494

Remove shuffled neurons in unified dataset

0d953fc

Remove extra datasets

eba09b6

cla-bot bot added the CLA signed label May 1, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add unified encoder pytorch implementation #251

Add unified encoder pytorch implementation #251

CeliaBenquet commented May 1, 2025 •

edited by MMathisLab

Loading

Add unified encoder pytorch implementation #251

Are you sure you want to change the base?

Add unified encoder pytorch implementation #251

Conversation

CeliaBenquet commented May 1, 2025 • edited by MMathisLab Loading

CeliaBenquet commented May 1, 2025 •

edited by MMathisLab

Loading