Merge branch 'main' into neuron-selection

Author: Vrittigyl

experanto/interpolators.py

@@ -164,16 +164,28 @@ def create(root_folder: str, cache_data: bool = False, **kwargs) -> Interpolator
         if modality == "sequence":
             if meta_data.get("phase_shift_per_signal", False):
                 return PhaseShiftedSequenceInterpolator(
-                    root_folder, cache_data, **kwargs
+                    root_folder, cache_data=cache_data, **kwargs
                 )
             else:
-                return SequenceInterpolator(root_folder, cache_data, **kwargs)
+                return SequenceInterpolator(
+                    root_folder, cache_data=cache_data, **kwargs
+                )
         elif modality == "screen":
-            return ScreenInterpolator(root_folder, cache_data, **kwargs)
+            use_stimuli_names = kwargs.pop(
+                "use_stimuli_names", meta_data.get("use_stimuli_names", False)
+            )
+            return ScreenInterpolator(
+                root_folder,
+                cache_data=cache_data,
+                use_stimuli_names=use_stimuli_names,
+                **kwargs,
+            )
         elif modality == "time_interval":
-            return TimeIntervalInterpolator(root_folder, cache_data, **kwargs)
+            return TimeIntervalInterpolator(
+                root_folder, cache_data=cache_data, **kwargs
+            )
         elif modality == "spikes":
-            return SpikeInterpolator(root_folder, cache_data, **kwargs)
+            return SpikeInterpolator(root_folder, cache_data=cache_data, **kwargs)
         else:
             raise ValueError(
                 f"There is no interpolator for {modality}. Please use 'sequence', 'screen', 'time_interval' as modality or provide a custom interpolator."
@@ -601,6 +613,8 @@ class ScreenInterpolator(Interpolator):
         native image size from metadata.
     normalize : bool, default=False
         If True, normalizes frames using stored mean/std statistics.
+    use_stimuli_names : bool, default=False
+        If True, uses ``stimulus_name`` from metadata to locate data files instead of trial keys.
     **kwargs
         Additional keyword arguments (ignored).
 
@@ -625,6 +639,7 @@ def __init__(
         rescale: bool = False,
         rescale_size: tuple[int, int] | None = None,
         normalize: bool = False,
+        use_stimuli_names: bool = False,
         **kwargs,
     ) -> None:
         super().__init__(root_folder)
@@ -634,6 +649,7 @@ def __init__(
         self.valid_interval = TimeInterval(self.start_time, self.end_time)
         self.rescale = rescale
         self.cache_trials = cache_data  # Store the cache preference
+        self.use_stimuli_names = use_stimuli_names
         self._parse_trials()
 
         # create mapping from image index to file index
@@ -718,8 +734,14 @@ def _parse_trials(self) -> None:
         metadatas, keys = self.read_combined_meta()
 
         for key, metadata in zip(keys, metadatas, strict=True):
-            data_file_name = self.root_folder / "data" / f"{key}.npy"
-            # Pass the cache_trials parameter when creating trials
+            if self.use_stimuli_names:
+                stimulus_name = metadata.get("stimulus_name")
+                assert (
+                    stimulus_name is not None
+                ), f"stimulus_name is required in metadata when use_stimuli_names is True, but not found for key: {key}"
+                data_file_name = self.root_folder / "data" / f"{stimulus_name}.npy"
+            else:
+                data_file_name = self.root_folder / "data" / f"{key}.npy"
             self.trials.append(
                 ScreenTrial.create(
                     data_file_name, metadata, cache_data=self.cache_trials