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

Author: CeliaBenquet

.github/workflows/build.yml

@@ -14,24 +14,15 @@ jobs:
       fail-fast: true
       matrix:
         os: [ubuntu-latest]
-        python-version: ["3.8", "3.10"]
+        python-version: ["3.9", "3.10", "3.12"]
         # 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: ["1.12.1", "2.0.0"]
+        torch-version: ["2.2.2", "2.4.0"]
         include:
-          - os: ubuntu-latest
-            python-version: 3.8
-            torch-version: 1.9.0
           - os: windows-latest
-            torch-version: 2.0.0
+            torch-version: 2.4.0
             python-version: "3.10"
-          - os: ubuntu-latest
-            torch-version: 2.1.1
-            python-version: "3.11"
-          #- os: macos-latest
-          #  torch-version: 2.0.0
-          #  python-version: "3.10"
 
     runs-on: ${{ matrix.os }}
 

CODE_OF_CONDUCT.md

@@ -6,7 +6,7 @@ In the interest of fostering an open and welcoming environment, we as
 contributors and maintainers pledge to making participation in our project and
 our community a harassment-free experience for everyone, regardless of age, body
 size, disability, ethnicity, sex characteristics, gender identity and expression,
-level of experience, education, socio-economic status, nationality, personal
+level of experience, education, socioeconomic status, nationality, personal
 appearance, race, religion, or sexual identity and orientation.
 
 ## Our Standards

Dockerfile

@@ -1,15 +1,13 @@
 ## EXPERIMENT BASE CONTAINER
-FROM nvidia/cuda:11.7.1-cudnn8-runtime-ubuntu20.04 AS cebra-base
+FROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04 AS cebra-base
 
 ENV DEBIAN_FRONTEND=noninteractive
 RUN apt-get update -y \
 	&& apt-get install --no-install-recommends -yy git python3 python3-pip python-is-python3 \
 	&& rm -rf /var/lib/apt/lists/*
 
-RUN pip install --no-cache-dir torch==2.0.0+cu117 \
-	--index-url https://download.pytorch.org/whl/cu117
-RUN pip install --no-cache-dir --pre 'cebra[dev,datasets,integrations]' \
-	&& pip uninstall -y cebra
+RUN pip install --no-cache-dir torch torchvision --index-url https://download.pytorch.org/whl/cu124
+RUN pip install --upgrade pip
 
 
 ## GIT repository

Makefile

@@ -24,6 +24,9 @@ test: clean_test
 doctest: clean_test
 	python -m pytest --ff --doctest-modules -vvv ./docs/source/usage.rst
 
+docker:
+	./tools/build_docker.sh
+
 test_parallel: clean_test
 	python -m pytest -n auto --ff -m "not requires_dataset"  tests
 
@@ -98,4 +101,7 @@ report: check_docker format .coverage .pylint
 	cat .pylint
 	coverage report
 
-.PHONY: dist build archlinux clean_test test doctest test_parallel test_parallel_debug test_all test_fast test_debug test_benchmark interrogate docs docs-touch docs-strict serve_docs serve_page format codespell check_for_binary
+.PHONY: dist build docker archlinux clean_test test doctest test_parallel \
+	test_parallel_debug test_all test_fast test_debug test_benchmark \
+	interrogate docs docs-touch docs-strict serve_docs serve_page \
+	format codespell check_for_binary

cebra/data/assets.py

@@ -21,9 +21,9 @@
 #
 
 import hashlib
-import os
 import re
 import warnings
+from pathlib import Path
 from typing import Optional
 
 import requests
@@ -57,8 +57,10 @@ def download_file_with_progress_bar(url: str,
     """
 
     # Check if the file already exists in the location
-    file_path = os.path.join(location, file_name)
-    if os.path.exists(file_path):
+    location_path = Path(location)
+    file_path = location_path / file_name
+
+    if file_path.exists():
         existing_checksum = calculate_checksum(file_path)
         if existing_checksum == expected_checksum:
             return file_path
@@ -91,10 +93,10 @@ def download_file_with_progress_bar(url: str,
         )
 
     # Create the directory and any necessary parent directories
-    os.makedirs(location, exist_ok=True)
+    location_path.mkdir(exist_ok=True)
 
     filename = filename_match.group(1)
-    file_path = os.path.join(location, filename)
+    file_path = location_path / filename
 
     total_size = int(response.headers.get("Content-Length", 0))
     checksum = hashlib.md5()  # create checksum
@@ -111,7 +113,7 @@ def download_file_with_progress_bar(url: str,
     downloaded_checksum = checksum.hexdigest()  # Get the checksum value
     if downloaded_checksum != expected_checksum:
         warnings.warn(f"Checksum verification failed. Deleting '{file_path}'.")
-        os.remove(file_path)
+        file_path.unlink()
         warnings.warn("File deleted. Retrying download...")
 
         # Retry download using a for loop

cebra/data/datasets.py

@@ -67,14 +67,13 @@ class TensorDataset(cebra_data.SingleSessionDataset):
 
     """
 
-    def __init__(
-        self,
-        neural: Union[torch.Tensor, npt.NDArray],
-        continuous: Union[torch.Tensor, npt.NDArray] = None,
-        discrete: Union[torch.Tensor, npt.NDArray] = None,
-        offset: int = 1,
-    ):
-        super().__init__()
+    def __init__(self,
+                 neural: Union[torch.Tensor, npt.NDArray],
+                 continuous: Union[torch.Tensor, npt.NDArray] = None,
+                 discrete: Union[torch.Tensor, npt.NDArray] = None,
+                 offset: int = 1,
+                 device: str = "cpu"):
+        super().__init__(device=device)
         self.neural = self._to_tensor(neural, torch.FloatTensor).float()
         self.continuous = self._to_tensor(continuous, torch.FloatTensor)
         self.discrete = self._to_tensor(discrete, torch.LongTensor)
@@ -222,9 +221,9 @@ def __init__(
         else:
             self._cindex = None
         if discrete:
-            raise NotImplementedError(
-                "Multisession implementation does not support discrete index yet."
-            )
+            self._dindex = torch.cat(list(
+                self._iter_property("discrete_index")),
+                                     dim=0)
         else:
             self._dindex = None
 

cebra/data/multi_session.py

@@ -164,7 +164,16 @@ def index(self):
 
 @dataclasses.dataclass
 class DiscreteMultiSessionDataLoader(MultiSessionLoader):
-    pass
+    """Contrastive learning conditioned on a discrete behavior variable."""
+
+    # Overwrite sampler with the discrete implementation
+    # Generalize MultisessionSampler to avoid doing this?
+    def __post_init__(self):
+        self.sampler = cebra_distr.DiscreteMultisessionSampler(self.dataset)
+
+    @property
+    def index(self):
+        return self.dataset.discrete_index
 
 
 @dataclasses.dataclass

cebra/datasets/allen/ca_movie_decoding.py

@@ -31,7 +31,6 @@
 
 import glob
 import hashlib
-import os
 import pathlib
 
 import h5py

cebra/datasets/allen/neuropixel_movie.py

@@ -28,7 +28,6 @@
 """
 import glob
 import hashlib
-import os
 import pathlib
 
 import h5py

cebra/datasets/allen/neuropixel_movie_decoding.py

@@ -28,7 +28,6 @@
 """
 import glob
 import hashlib
-import os
 import pathlib
 
 import h5py