diff --git a/.github/workflows/ruff.yml b/.github/workflows/ruff.yml new file mode 100644 index 00000000..37b0c585 --- /dev/null +++ b/.github/workflows/ruff.yml @@ -0,0 +1,14 @@ +name: Ruff + +on: [push, pull_request] + +jobs: + ruff: + if: github.event_name == 'push' || (github.event_name == 'pull_request' && github.event.pull_request.head.repo.fork == true) + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: astral-sh/ruff-action@v3 + with: + version: "0.15.6" + args: "check --exclude examples/" \ No newline at end of file diff --git a/.gitignore b/.gitignore index 2bc61f88..22b40f7c 100644 --- a/.gitignore +++ b/.gitignore @@ -165,3 +165,4 @@ cython_debug/ *.bak docs/source/generated/ +gsoc_env/ \ No newline at end of file diff --git a/examples/sensorium/data.ipynb b/examples/sensorium/data.ipynb index 5324dea4..6db7901d 100644 --- a/examples/sensorium/data.ipynb +++ b/examples/sensorium/data.ipynb @@ -11,9 +11,7 @@ "%load_ext autoreload\n", "%autoreload 2\n", "\n", - "import numpy as np\n", - "import matplotlib.pyplot as plt\n", - "import torch" + "import matplotlib.pyplot as plt" ] }, { diff --git a/examples/sensorium/experiment.ipynb b/examples/sensorium/experiment.ipynb index 54dc7c75..21423733 100644 --- a/examples/sensorium/experiment.ipynb +++ b/examples/sensorium/experiment.ipynb @@ -14,7 +14,6 @@ "import numpy as np\n", "import matplotlib.pyplot as plt\n", "import matplotlib.animation as animation\n", - "from pathlib import Path\n", "from IPython.display import HTML" ] }, @@ -33,9 +32,7 @@ "if p not in sys.path:\n", " sys.path.append(p)\n", "\n", - "from experanto.experiment import Experiment\n", - "from experanto.interpolators import Interpolator\n", - "from experanto.interpolators import ScreenInterpolator" + "from experanto.experiment import Experiment" ] }, { diff --git a/examples/sensorium/sensorium_min_example.ipynb b/examples/sensorium/sensorium_min_example.ipynb index 0b1d1ee2..82ca25cf 100644 --- a/examples/sensorium/sensorium_min_example.ipynb +++ b/examples/sensorium/sensorium_min_example.ipynb @@ -32,7 +32,6 @@ "import numpy as np\n", "from tqdm import tqdm\n", "import matplotlib.pyplot as plt\n", - "from torch.utils.data import DataLoader\n", "from collections import OrderedDict" ] }, @@ -444,8 +443,6 @@ "\n", "set_random_seed(seed)\n", "\n", - "from sensorium.datasets.mouse_video_loaders import mouse_video_loader\n", - "from sensorium.utility.scores import get_correlations\n", "from nnfabrik.builder import get_trainer\n", "from sensorium.models.make_model import make_video_model" ] @@ -744,7 +741,6 @@ "source": [ "import yaml\n", "import os\n", - "from tqdm import tqdm\n", "from collections import Counter" ] }, diff --git a/experanto/configs.py b/experanto/configs.py index 5f3a0beb..de7293ba 100644 --- a/experanto/configs.py +++ b/experanto/configs.py @@ -37,8 +37,7 @@ from pathlib import Path -from hydra import compose, initialize, initialize_config_dir -from omegaconf import OmegaConf, open_dict +from omegaconf import OmegaConf # get config relative to this file script_dir = Path(__file__).parent diff --git a/experanto/dataloaders.py b/experanto/dataloaders.py index f891bd57..316a4828 100644 --- a/experanto/dataloaders.py +++ b/experanto/dataloaders.py @@ -1,13 +1,8 @@ import logging -import os -import time import warnings -from pathlib import Path -from typing import Any, Dict, List, Optional, Type, Union +from typing import Dict, List, Optional, Union -import numpy as np from omegaconf import DictConfig -from torch.utils.data import DataLoader from .datasets import ChunkDataset from .utils import ( @@ -158,7 +153,6 @@ def get_multisession_concat_dataloader( datasets = [] session_names = [] - start_time = time.time() for i, (path, cfg) in enumerate(zip(paths, configs)): # Create dataset with deterministic seed path_hash = hash(path) % 10000 diff --git a/experanto/datasets.py b/experanto/datasets.py index 31682cda..c2454a26 100644 --- a/experanto/datasets.py +++ b/experanto/datasets.py @@ -1,35 +1,31 @@ from __future__ import annotations -import functools import importlib import json import logging import os from collections.abc import Iterable from pathlib import Path -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Optional import numpy as np import torch import torchvision from hydra.utils import instantiate -from omegaconf import DictConfig, ListConfig, OmegaConf +from omegaconf import DictConfig, ListConfig from torch.utils.data import Dataset -from torchvision.transforms import v2 from torchvision.transforms.v2 import Compose, Lambda, ToTensor from .configs import DEFAULT_MODALITY_CONFIG from .experiment import Experiment -from .interpolators import ImageTrial, VideoTrial from .intervals import ( TimeInterval, find_intersection_between_two_interval_arrays, get_stats_for_valid_interval, ) -from .utils import add_behavior_as_channels, replace_nan_with_batch_mean +from .utils import add_behavior_as_channels # see .configs.py for the definition of DEFAULT_MODALITY_CONFIG -DEFAULT_MODALITY_CONFIG = dict() logger = logging.getLogger(__name__) @@ -768,7 +764,6 @@ def __getitem__(self, idx: int) -> dict: for device_name in self.device_names: sampling_rate = self.sampling_rates[device_name] chunk_size = self.chunk_sizes[device_name] - chunk_s = chunk_size / sampling_rate # convert everything to int to avoid numerical issues start_time = int(round(s * self.scale_precision)) diff --git a/experanto/experiment.py b/experanto/experiment.py index faa2fed1..b5c4850f 100644 --- a/experanto/experiment.py +++ b/experanto/experiment.py @@ -1,11 +1,9 @@ from __future__ import annotations import logging -import re import warnings -from collections.abc import Sequence from pathlib import Path -from typing import Optional, Union +from typing import Union import numpy as np from hydra.utils import instantiate diff --git a/experanto/interpolators.py b/experanto/interpolators.py index 4928fd85..027e87c8 100644 --- a/experanto/interpolators.py +++ b/experanto/interpolators.py @@ -12,7 +12,6 @@ import cv2 import numpy as np -import numpy.lib.format as fmt import yaml from numba import njit, prange from scipy.ndimage import gaussian_filter1d @@ -326,7 +325,7 @@ def interpolate( else: raise NotImplementedError( - f"interpolation_mode should be linear or nearest_neighbor" + "interpolation_mode should be linear or nearest_neighbor" ) def close(self) -> None: @@ -454,7 +453,7 @@ def interpolate( else: raise NotImplementedError( - f"interpolation_mode should be linear or nearest_neighbor" + "interpolation_mode should be linear or nearest_neighbor" ) diff --git a/experanto/intervals.py b/experanto/intervals.py index 16079857..9a1d6c02 100644 --- a/experanto/intervals.py +++ b/experanto/intervals.py @@ -236,7 +236,7 @@ def get_stats_for_valid_interval( """ total_duration = end_time - start_time if total_duration <= 0: - return f"Error: Invalid time range (end_time <= start_time). Total duration must be positive." + return "Error: Invalid time range (end_time <= start_time). Total duration must be positive." # Ensure intervals are unique and sorted, then clamp them to the analysis window unique_intervals = uniquefy_interval_array(intervals) diff --git a/experanto/utils.py b/experanto/utils.py index bd0c7f85..f40c862b 100644 --- a/experanto/utils.py +++ b/experanto/utils.py @@ -1,28 +1,15 @@ import bisect import logging -import math -import multiprocessing # inbuilt libraries -import os -import queue -import random -import threading -import time -import warnings from collections import defaultdict -from copy import deepcopy -from functools import partial -from typing import Any, Dict, Iterator, List, Optional, Tuple, Union # third-party libraries import numpy as np import torch -from omegaconf import DictConfig -from torch.utils.data import ConcatDataset, DataLoader, Dataset, Sampler +from torch.utils.data import DataLoader, Dataset, Sampler # local libraries -from .intervals import TimeInterval logger = logging.getLogger(__name__) @@ -58,7 +45,7 @@ def add_behavior_as_channels(data: dict[str, torch.Tensor]) -> dict: # Process eye_tracker if len(eye_tracker.shape) == 2: # (t, c_eye) - c_eye = eye_tracker.shape[1] + # Reshape to (c_eye, t, h, w) eye_tracker = eye_tracker.transpose(0, 1) # (c_eye, t) eye_tracker = eye_tracker.unsqueeze(-1).unsqueeze(-1) # (c_eye, t, 1, 1) @@ -69,7 +56,7 @@ def add_behavior_as_channels(data: dict[str, torch.Tensor]) -> dict: # Process treadmill if len(treadmill.shape) == 2: # (t, c_tread) - c_tread = treadmill.shape[1] + # Reshape to (c_tread, t, h, w) treadmill = treadmill.transpose(0, 1) # (c_tread, t) treadmill = treadmill.unsqueeze(-1).unsqueeze(-1) # (c_tread, t, 1, 1) diff --git a/tests/test_experiment.py b/tests/test_experiment.py new file mode 100644 index 00000000..e7e846fc --- /dev/null +++ b/tests/test_experiment.py @@ -0,0 +1,163 @@ +import shutil +from contextlib import contextmanager +from pathlib import Path + +import numpy as np +import pytest +import yaml + +from experanto.experiment import Experiment + +EXPERIMENT_ROOT = Path("tests/experiment_data") + + +@contextmanager +def create_two_device_experiment( + device0_start=0.0, + device0_end=10.0, + device1_start=1.0, + device1_end=8.0, + sampling_rate=10.0, +): + """Create a temporary experiment with two sequence devices with different time ranges.""" + try: + for device_name, start, end in [ + ("device_0", device0_start, device0_end), + ("device_1", device1_start, device1_end), + ]: + device_root = EXPERIMENT_ROOT / device_name + (device_root / "meta").mkdir(parents=True, exist_ok=True) + + n_samples = int((end - start) * sampling_rate) + 1 + timestamps = np.linspace(start, end, n_samples) + data = np.random.rand(n_samples, 5) + + np.save(device_root / "timestamps.npy", timestamps) + np.save(device_root / "data.npy", data) + + meta = { + "start_time": start, + "end_time": end, + "modality": "sequence", + "sampling_rate": sampling_rate, + "phase_shift_per_signal": False, + "is_mem_mapped": False, + "n_signals": 5, + "n_timestamps": n_samples, + "dtype": "float64", + } + with open(device_root / "meta.yml", "w") as f: + yaml.safe_dump(meta, f) + + yield EXPERIMENT_ROOT + + finally: + shutil.rmtree(EXPERIMENT_ROOT) + + +@contextmanager +def create_three_device_experiment(): + """Create a temporary experiment with three sequence devices with different time ranges.""" + try: + for device_name, start, end in [ + ("device_0", 0.0, 10.0), + ("device_1", 1.0, 8.0), + ("device_2", 2.0, 9.0), + ]: + device_root = EXPERIMENT_ROOT / device_name + (device_root / "meta").mkdir(parents=True, exist_ok=True) + + sampling_rate = 10.0 + n_samples = int((end - start) * sampling_rate) + 1 + timestamps = np.linspace(start, end, n_samples) + data = np.random.rand(n_samples, 5) + + np.save(device_root / "timestamps.npy", timestamps) + np.save(device_root / "data.npy", data) + + meta = { + "start_time": start, + "end_time": end, + "modality": "sequence", + "sampling_rate": sampling_rate, + "phase_shift_per_signal": False, + "is_mem_mapped": False, + "n_signals": 5, + "n_timestamps": n_samples, + "dtype": "float64", + } + with open(device_root / "meta.yml", "w") as f: + yaml.safe_dump(meta, f) + + yield EXPERIMENT_ROOT + + finally: + shutil.rmtree(EXPERIMENT_ROOT) + + +def get_two_device_config(): + return { + "device_0": {"interpolation": {"sampling_rate": 10.0, "offset": 0.0}}, + "device_1": {"interpolation": {"sampling_rate": 10.0, "offset": 0.0}}, + } + + +def test_experiment_start_end_time_reflects_union(): + """ + Experiment.start_time and end_time should reflect the union + of all device time ranges — the earliest start and latest end + across all devices, not just the last loaded device. + """ + with create_two_device_experiment( + device0_start=1.0, + device0_end=8.0, + device1_start=0.0, + device1_end=10.0, + ) as experiment_path: + experiment = Experiment( + root_folder=experiment_path, + modality_config=get_two_device_config(), + ) + + # Union: start = min(1.0, 0.0) = 0.0, end = max(8.0, 10.0) = 10.0 + assert experiment.start_time == pytest.approx( + 0.0 + ), f"Expected start_time=0.0, got {experiment.start_time}" + assert experiment.end_time == pytest.approx( + 10.0 + ), f"Expected end_time=10.0, got {experiment.end_time}" + + +def test_experiment_single_device_time_range(): + """With a single device, start_time and end_time should match that device's range.""" + with create_two_device_experiment( + device0_start=2.0, + device0_end=9.0, + ) as experiment_path: + config = {"device_0": {"interpolation": {"sampling_rate": 10.0, "offset": 0.0}}} + experiment = Experiment( + root_folder=experiment_path, + modality_config=config, + ) + + assert experiment.start_time == pytest.approx(2.0) + assert experiment.end_time == pytest.approx(9.0) + + +def test_experiment_start_end_time_three_devices(): + """With three devices, start_time and end_time should reflect the union of all three — + earliest start and latest end across all devices.""" + with create_three_device_experiment() as experiment_path: + config = { + "device_0": {"interpolation": {"sampling_rate": 10.0, "offset": 0.0}}, + "device_1": {"interpolation": {"sampling_rate": 10.0, "offset": 0.0}}, + "device_2": {"interpolation": {"sampling_rate": 10.0, "offset": 0.0}}, + } + experiment = Experiment( + root_folder=experiment_path, + modality_config=config, + ) + + # Union: start = min(0.0, 1.0, 2.0) = 0.0, end = max(10.0, 8.0, 9.0) = 10.0 + assert experiment.start_time == pytest.approx(0.0) + assert experiment.end_time == pytest.approx(10.0)