fix: apply ruff auto-fixes to experanto source files

Author: binary69

experanto/dataloaders.py

@@ -1,6 +1,6 @@
 import logging
 import warnings
-from typing import Dict, List, Optional, Union
+from typing import Optional
 
 from omegaconf import DictConfig
 
@@ -16,8 +16,8 @@
 
 
 def get_multisession_dataloader(
-    paths: List[str],
-    configs: Optional[Union[DictConfig, Dict, List[Union[DictConfig, Dict]]]] = None,
+    paths: list[str],
+    configs: DictConfig | dict | list[DictConfig | dict] | None = None,
     shuffle_keys: bool = False,
     **kwargs,
 ) -> LongCycler:
@@ -90,10 +90,10 @@ def get_multisession_dataloader(
 
 
 def get_multisession_concat_dataloader(
-    paths: List[str],
-    configs: Optional[Union[Dict, List[Dict]]] = None,
-    seed: Optional[int] = 0,
-    dataloader_config: Optional[Dict] = None,
+    paths: list[str],
+    configs: dict | list[dict] | None = None,
+    seed: int | None = 0,
+    dataloader_config: dict | None = None,
     **kwargs,
 ) -> Optional["FastSessionDataLoader"]:
     """Create a concatenated multi-session dataloader.

experanto/datasets.py

@@ -6,7 +6,7 @@
 import os
 from collections.abc import Iterable
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 import numpy as np
 import torch
@@ -212,15 +212,15 @@ class ChunkDataset(Dataset):
     def __init__(
         self,
         root_folder: str,
-        global_sampling_rate: Optional[float] = None,
-        global_chunk_size: Optional[int] = None,
+        global_sampling_rate: float | None = None,
+        global_chunk_size: int | None = None,
         add_behavior_as_channels: bool = False,
         replace_nans_with_means: bool = False,
         cache_data: bool = False,
-        out_keys: Optional[Iterable] = None,
+        out_keys: Iterable | None = None,
         normalize_timestamps: bool = True,
         modality_config: dict = DEFAULT_MODALITY_CONFIG,
-        seed: Optional[int] = None,
+        seed: int | None = None,
         safe_interval_threshold: float = 0.5,
         interpolate_precision: int = 5,
     ) -> None:
@@ -362,8 +362,8 @@ def initialize_statistics(self) -> None:
                         None, ...
                     ]
                 elif mode == "screen_default":
-                    means = np.array((80))
-                    stds = np.array((60))
+                    means = np.array(80)
+                    stds = np.array(60)
 
                 self._statistics[device_name]["mean"] = means.reshape(
                     1, -1
@@ -385,7 +385,7 @@ def initialize_transforms(self):
         for device_name in self.device_names:
             if device_name == "screen":
                 add_channel = Lambda(self.add_channel_function)
-                transform_list: List[Any] = []
+                transform_list: list[Any] = []
 
                 for v in self.modality_config.screen.transforms.values():  # type: ignore[union-attr]
                     if isinstance(v, dict):  # config dict
@@ -395,7 +395,7 @@ def initialize_transforms(self):
 
                 transform_list.insert(0, add_channel)
             else:
-                transform_list: List[Any] = [ToTensor()]
+                transform_list: list[Any] = [ToTensor()]
 
             # Normalization.
             if self.modality_config[device_name].transforms.get("normalization", False):
@@ -468,8 +468,8 @@ def _get_callable_filter(self, filter_config):
 
     def get_valid_intervals_from_filters(
         self, visualize: bool = False
-    ) -> List[TimeInterval]:
-        valid_intervals: Optional[List[TimeInterval]] = None
+    ) -> list[TimeInterval]:
+        valid_intervals: list[TimeInterval] | None = None
         for modality in self.modality_config:
             if "filters" in self.modality_config[modality]:
                 device = self._experiment.devices[modality]
@@ -478,7 +478,7 @@ def get_valid_intervals_from_filters(
                 ].items():
                     # Get the final callable filter function
                     filter_function = self._get_callable_filter(filter_config)
-                    valid_intervals_: List[TimeInterval] = filter_function(device_=device)  # type: ignore[assignment]
+                    valid_intervals_: list[TimeInterval] = filter_function(device_=device)  # type: ignore[assignment]
                     if visualize:
                         logger.info("modality: %s, filter: %s", modality, filter_name)
                         visualization_string = get_stats_for_valid_interval(
@@ -495,7 +495,7 @@ def get_valid_intervals_from_filters(
         return valid_intervals if valid_intervals is not None else []
 
     def get_condition_mask_from_meta_conditions(
-        self, valid_conditions_sum_of_product: List[dict]
+        self, valid_conditions_sum_of_product: list[dict]
     ) -> np.ndarray:
         """Create a boolean mask for trials satisfying given conditions.
 
@@ -517,7 +517,7 @@ def get_condition_mask_from_meta_conditions(
         ``[{'tier': 'train', 'stim_type': 'natural'}, {'tier': 'blank'}]``
         matches trials that are either (train AND natural) OR blank.
         """
-        all_conditions: Optional[np.ndarray] = None
+        all_conditions: np.ndarray | None = None
         for valid_conditions_product in valid_conditions_sum_of_product:
             conditions_of_product = None
             for k, valid_condition in valid_conditions_product.items():
@@ -540,7 +540,7 @@ def get_condition_mask_from_meta_conditions(
     def get_screen_sample_mask_from_meta_conditions(
         self,
         satisfy_for_next: int,
-        valid_conditions_sum_of_product: List[dict],
+        valid_conditions_sum_of_product: list[dict],
         filter_for_valid_intervals: bool = True,
     ) -> np.ndarray:
         """Create a boolean mask for screen samples satisfying given conditions.
@@ -702,7 +702,7 @@ def get_data_key_from_root_folder(self, root_folder):
         # Check if the file exists before trying to open it
         if os.path.isfile(meta_file_path):
             try:
-                with open(meta_file_path, "r") as file:
+                with open(meta_file_path) as file:
                     meta = json.load(file)
 
                 # Get data_key from meta if it exists
@@ -811,14 +811,14 @@ def __getitem__(self, idx: int) -> dict:
 
         return final_out
 
-    def get_state(self) -> Dict[str, Any]:
+    def get_state(self) -> dict[str, Any]:
         """Return the current state of the dataset's RNG."""
         return {
             "rng_state": self._rng.get_state() if self.seed is not None else None,
             "valid_screen_times": self._valid_screen_times.copy(),
         }
 
-    def set_state(self, state: Dict[str, Any]) -> None:
+    def set_state(self, state: dict[str, Any]) -> None:
         """Restore the dataset's RNG state."""
         if state["rng_state"] is not None and self.seed is not None:
             self._rng.set_state(state["rng_state"])

experanto/experiment.py

@@ -3,7 +3,6 @@
 import logging
 import warnings
 from pathlib import Path
-from typing import Union
 
 import numpy as np
 from hydra.utils import instantiate
@@ -130,9 +129,9 @@ def device_names(self):
     def interpolate(
         self,
         times: np.ndarray,
-        device: Union[str, Interpolator, None] = None,
+        device: str | Interpolator | None = None,
         return_valid: bool = False,
-    ) -> Union[tuple[dict, dict], dict, tuple[np.ndarray, np.ndarray], np.ndarray]:
+    ) -> tuple[dict, dict] | dict | tuple[np.ndarray, np.ndarray] | np.ndarray:
         """Interpolate data from one or all devices at specified time points.
 
         Parameters
@@ -200,7 +199,7 @@ def interpolate(
             else:
                 return values
         elif isinstance(device, str):
-            assert device in self.devices, "Unknown device '{}'".format(device)
+            assert device in self.devices, f"Unknown device '{device}'"
             res = self.devices[device].interpolate(times, return_valid=return_valid)
             return res
         else:

experanto/interpolators.py

@@ -4,11 +4,10 @@
 import logging
 import os
 import re
-import typing
 import warnings
 from abc import abstractmethod
 from pathlib import Path
-from typing import Union, cast
+from typing import cast
 
 import cv2
 import numpy as np
@@ -71,7 +70,7 @@ def load_meta(self):
     @abstractmethod
     def interpolate(
         self, times: np.ndarray, return_valid: bool = False
-    ) -> Union[tuple[np.ndarray, np.ndarray], np.ndarray]:
+    ) -> tuple[np.ndarray, np.ndarray] | np.ndarray:
         """Map an array of time points to interpolated data values."""
         ...
 
@@ -85,7 +84,7 @@ def __exit__(self, *exc):
         self.close()
 
     @staticmethod
-    def create(root_folder: str, cache_data: bool = False, **kwargs) -> "Interpolator":
+    def create(root_folder: str, cache_data: bool = False, **kwargs) -> Interpolator:
         """Factory method to create the appropriate interpolator for a modality.
 
         Reads the ``meta.yml`` file in the folder to determine the modality type
@@ -110,7 +109,7 @@ def create(root_folder: str, cache_data: bool = False, **kwargs) -> "Interpolato
         ValueError
             If the modality type is not supported.
         """
-        with open(Path(root_folder) / "meta.yml", "r") as file:
+        with open(Path(root_folder) / "meta.yml") as file:
             meta_data = yaml.safe_load(file)
         modality = meta_data.get("modality")
 
@@ -197,7 +196,7 @@ def __init__(
         interpolation_mode: str = "nearest_neighbor",
         normalize: bool = False,
         normalize_subtract_mean: bool = False,
-        normalize_std_threshold: typing.Optional[float] = None,  # or 0.01
+        normalize_std_threshold: float | None = None,  # or 0.01
         **kwargs,
     ) -> None:
         super().__init__(root_folder)
@@ -262,7 +261,7 @@ def normalize_data(self, data):
 
     def interpolate(
         self, times: np.ndarray, return_valid: bool = False
-    ) -> Union[tuple[np.ndarray, np.ndarray], np.ndarray]:
+    ) -> tuple[np.ndarray, np.ndarray] | np.ndarray:
         valid = self.valid_times(times)
         valid_times = times[valid]
 
@@ -361,7 +360,7 @@ def __init__(
         interpolation_mode: str = "nearest_neighbor",
         normalize: bool = False,
         normalize_subtract_mean: bool = False,
-        normalize_std_threshold: typing.Optional[float] = None,  # or 0.01
+        normalize_std_threshold: float | None = None,  # or 0.01
         **kwargs,
     ) -> None:
         super().__init__(
@@ -385,7 +384,7 @@ def __init__(
 
     def interpolate(
         self, times: np.ndarray, return_valid: bool = False
-    ) -> Union[tuple[np.ndarray, np.ndarray], np.ndarray]:
+    ) -> tuple[np.ndarray, np.ndarray] | np.ndarray:
         valid = self.valid_times(times)
         valid_times = times[valid]
 
@@ -500,7 +499,7 @@ def __init__(
         root_folder: str,
         cache_data: bool = False,  # New parameter
         rescale: bool = False,
-        rescale_size: typing.Optional[tuple[int, int]] = None,
+        rescale_size: tuple[int, int] | None = None,
         normalize: bool = False,
         **kwargs,
     ) -> None:
@@ -565,7 +564,7 @@ def is_numbered_yml(file_name):
 
         # Read each YAML file and store under its filename
         for meta_file in meta_files:
-            with open(meta_file, "r") as file:
+            with open(meta_file) as file:
                 file_base_name = meta_file.stem
                 yaml_content = yaml.safe_load(file)
                 all_data[file_base_name] = yaml_content
@@ -579,7 +578,7 @@ def read_combined_meta(self) -> tuple[list, list]:
             logger.info("Combining metadata files...")
             self._combine_metadatas()
 
-        with open(self.root_folder / "combined_meta.json", "r") as file:
+        with open(self.root_folder / "combined_meta.json") as file:
             self.combined_meta = json.load(file)
 
         metadatas = []
@@ -605,7 +604,7 @@ def _parse_trials(self) -> None:
 
     def interpolate(
         self, times: np.ndarray, return_valid: bool = False
-    ) -> Union[tuple[np.ndarray, np.ndarray], np.ndarray]:
+    ) -> tuple[np.ndarray, np.ndarray] | np.ndarray:
         valid = self.valid_times(times)
         valid_times = times[valid]
         valid_times += 1e-4  # add small offset to avoid numerical issues
@@ -724,7 +723,7 @@ def __init__(self, root_folder: str, cache_data: bool = False, **kwargs):
 
     def interpolate(
         self, times: np.ndarray, return_valid: bool = False
-    ) -> Union[tuple[np.ndarray, np.ndarray], np.ndarray]:
+    ) -> tuple[np.ndarray, np.ndarray] | np.ndarray:
         valid = self.valid_times(times)
         valid_times = times[valid]
 
@@ -791,7 +790,7 @@ class ScreenTrial:
 
     def __init__(
         self,
-        data_file_name: Union[str, Path],
+        data_file_name: str | Path,
         meta_data: dict,
         image_size: tuple,
         first_frame_idx: int,
@@ -811,10 +810,10 @@ def __init__(
 
     @staticmethod
     def create(
-        data_file_name: Union[str, Path],
+        data_file_name: str | Path,
         meta_data: dict,
         cache_data: bool = False,
-    ) -> "ScreenTrial":
+    ) -> ScreenTrial:
         modality = meta_data.get("modality")
         assert modality is not None
         class_name = modality.lower().capitalize() + "Trial"
@@ -1050,7 +1049,7 @@ def __init__(
 
     def interpolate(
         self, times: np.ndarray, return_valid: bool = False
-    ) -> Union[tuple[np.ndarray, np.ndarray], np.ndarray]:
+    ) -> tuple[np.ndarray, np.ndarray] | np.ndarray:
         # 1. Filter for valid times
         valid = self.valid_times(times)
         valid_times = times[valid]