Merge branch 'main' into stes/better-goodness-of-fit

Author: MMathisLab

tests/_build_legacy_model/.gitignore

@@ -0,0 +1 @@
+*.pt

tests/_build_legacy_model/Dockerfile

@@ -0,0 +1,39 @@
+FROM python:3.12-slim AS base
+RUN pip install torch --index-url https://download.pytorch.org/whl/cpu
+RUN apt-get update && \
+ apt-get install -y --no-install-recommends git && \
+ rm -rf /var/lib/apt/lists/*
+
+FROM base AS cebra-0.4.0-scikit-learn-1.4
+RUN pip install cebra==0.4.0 "scikit-learn<1.5"
+WORKDIR /app
+COPY create_model.py .
+RUN python create_model.py
+
+FROM base AS cebra-0.4.0-scikit-learn-1.6
+RUN pip install cebra==0.4.0 "scikit-learn>=1.6"
+WORKDIR /app
+COPY create_model.py .
+RUN python create_model.py
+
+FROM base AS cebra-rc-scikit-learn-1.4
+# NOTE(stes): Commit where new scikit-learn tag logic was added to the CEBRA class.
+# https://github.com/AdaptiveMotorControlLab/CEBRA/commit/5f46c3257952a08dfa9f9e1b149a85f7f12c1053
+RUN pip install git+https://github.com/AdaptiveMotorControlLab/CEBRA.git@5f46c3257952a08dfa9f9e1b149a85f7f12c1053 "scikit-learn<1.5"
+WORKDIR /app
+COPY create_model.py .
+RUN python create_model.py
+
+FROM base AS cebra-rc-scikit-learn-1.6
+# NOTE(stes): Commit where new scikit-learn tag logic was added to the CEBRA class.
+# https://github.com/AdaptiveMotorControlLab/CEBRA/commit/5f46c3257952a08dfa9f9e1b149a85f7f12c1053
+RUN pip install git+https://github.com/AdaptiveMotorControlLab/CEBRA.git@5f46c3257952a08dfa9f9e1b149a85f7f12c1053 "scikit-learn>=1.6"
+WORKDIR /app
+COPY create_model.py .
+RUN python create_model.py
+
+FROM scratch
+COPY --from=cebra-0.4.0-scikit-learn-1.4 /app/cebra_model.pt /cebra_model_cebra-0.4.0-scikit-learn-1.4.pt
+COPY --from=cebra-0.4.0-scikit-learn-1.6 /app/cebra_model.pt /cebra_model_cebra-0.4.0-scikit-learn-1.6.pt
+COPY --from=cebra-rc-scikit-learn-1.4 /app/cebra_model.pt /cebra_model_cebra-rc-scikit-learn-1.4.pt
+COPY --from=cebra-rc-scikit-learn-1.6 /app/cebra_model.pt /cebra_model_cebra-rc-scikit-learn-1.6.pt

tests/_build_legacy_model/README.md

@@ -0,0 +1,13 @@
+# Helper script to build CEBRA checkpoints
+
+This script builds CEBRA checkpoints for different versions of scikit-learn and CEBRA.
+To build all models, run:
+
+```bash
+./generate.sh
+```
+
+The models are currently also stored in git directly due to their small size.
+
+Related issue: https://github.com/AdaptiveMotorControlLab/CEBRA/issues/207
+Related test: tests/test_sklearn_legacy.py

tests/_build_legacy_model/create_model.py

@@ -0,0 +1,15 @@
+import numpy as np
+
+import cebra
+
+neural_data = np.random.normal(0, 1, (1000, 30)) # 1000 samples, 30 features
+cebra_model = cebra.CEBRA(model_architecture="offset10-model",
+ batch_size=512,
+ learning_rate=1e-4,
+ max_iterations=10,
+ time_offsets=10,
+ num_hidden_units=16,
+ output_dimension=8,
+ verbose=True)
+cebra_model.fit(neural_data)
+cebra_model.save("cebra_model.pt")

tests/_build_legacy_model/generate.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+DOCKER_BUILDKIT=1 docker build --output type=local,dest=. .

tests/test_sklearn_legacy.py

@@ -0,0 +1,41 @@
+import pathlib
+import urllib.request
+
+import numpy as np
+import pytest
+
+from cebra.integrations.sklearn.cebra import CEBRA
+
+MODEL_VARIANTS = [
+ "cebra-0.4.0-scikit-learn-1.4", "cebra-0.4.0-scikit-learn-1.6",
+ "cebra-rc-scikit-learn-1.4", "cebra-rc-scikit-learn-1.6"
+]
+
+
+@pytest.mark.parametrize("model_variant", MODEL_VARIANTS)
+def test_load_legacy_model(model_variant):
+ """Test loading a legacy CEBRA model."""
+
+ X = np.random.normal(0, 1, (1000, 30))
+
+ model_path = pathlib.Path(
+ __file__
+ ).parent / "_build_legacy_model" / f"cebra_model_{model_variant}.pt"
+
+ if not model_path.exists():
+ url = f"https://cebra.fra1.digitaloceanspaces.com/cebra_model_{model_variant}.pt"
+ urllib.request.urlretrieve(url, model_path)
+
+ loaded_model = CEBRA.load(model_path)
+
+ assert loaded_model.model_architecture == "offset10-model"
+ assert loaded_model.output_dimension == 8
+ assert loaded_model.num_hidden_units == 16
+ assert loaded_model.time_offsets == 10
+
+ output = loaded_model.transform(X)
+ assert isinstance(output, np.ndarray)
+ assert output.shape[1] == loaded_model.output_dimension
+
+ assert hasattr(loaded_model, "state_dict_")
+ assert hasattr(loaded_model, "n_features_")