Improve segment validation to list only. Add unitls function for validation

Author: jakeswann1

src/spikeinterface/widgets/amplitudes.py

@@ -5,7 +5,7 @@
 
 from .rasters import BaseRasterWidget
 from .base import BaseWidget, to_attr
-from .utils import get_some_colors
+from .utils import get_some_colors, validate_segment_indices
 
 from spikeinterface.core.sortinganalyzer import SortingAnalyzer
 
@@ -25,7 +25,7 @@ class AmplitudesWidget(BaseRasterWidget):
     unit_colors : dict | None, default: None
         Dict of colors with unit ids as keys and colors as values. Colors can be any type accepted
         by matplotlib. If None, default colors are chosen using the `get_some_colors` function.
-    segment_index : int or list of int or None, default: None
+    segment_indices : list of int or None, default: None
         Segment index or indices to plot. If None and there are multiple segments, defaults to 0.
         If list, spike trains and amplitudes are concatenated across the specified segments.
     max_spikes_per_unit : int or None, default: None
@@ -52,7 +52,7 @@ def __init__(
         sorting_analyzer: SortingAnalyzer,
         unit_ids=None,
         unit_colors=None,
-        segment_index=None,
+        segment_indices=None,
         max_spikes_per_unit=None,
         y_lim=None,
         scatter_decimate=1,
@@ -74,38 +74,16 @@ def __init__(
         if unit_ids is None:
             unit_ids = sorting.unit_ids
 
-        num_segments = sorting.get_num_segments()
-
         # Handle segment_index input
-        if num_segments > 1:
-            if segment_index is None:
-                warn("More than one segment available! Using `segment_index = 0`.")
-                segment_index = 0
-        else:
-            segment_index = 0
+        segment_indices = validate_segment_indices(segment_indices, sorting)
 
         # Check for SortingView backend
         is_sortingview = backend == "sortingview"
 
         # For SortingView, ensure we're only using a single segment
-        if is_sortingview and isinstance(segment_index, list) and len(segment_index) > 1:
+        if is_sortingview and len(segment_indices) > 1:
             warn("SortingView backend currently supports only single segment. Using first segment.")
-            segment_index = segment_index[0]
-
-        # Convert segment_index to list for consistent processing
-        if isinstance(segment_index, int):
-            segment_indices = [segment_index]
-        elif isinstance(segment_index, list):
-            segment_indices = segment_index
-        else:
-            raise ValueError("segment_index must be an int or a list of ints")
-
-        # Validate segment indices
-        for idx in segment_indices:
-            if not isinstance(idx, int):
-                raise ValueError(f"Each segment index must be an integer, got {type(idx)}")
-            if idx < 0 or idx >= num_segments:
-                raise ValueError(f"segment_index {idx} out of range (0 to {num_segments - 1})")
+            segment_indices = segment_indices[0]
 
         # Create multi-segment data structure (dict of dicts)
         spiketrains_by_segment = {}

src/spikeinterface/widgets/motion.py

@@ -121,7 +121,7 @@ class DriftRasterMapWidget(BaseRasterWidget):
         The recording extractor object (only used to get "real" times).
     sampling_frequency : float, default: None
         The sampling frequency (needed if recording is None).
-    segment_index : int or list of int or None, default: None
+    segment_indices : list of int or None, default: None
         The segment index or indices to display. If None and there's only one segment, it's used.
         If None and there are multiple segments, you must specify which to use.
         If a list of indices is provided, peaks and locations are concatenated across the segments.
@@ -149,7 +149,7 @@ def __init__(
         direction: str = "y",
         recording: BaseRecording | None = None,
         sampling_frequency: float | None = None,
-        segment_index: int | list | None = None,
+        segment_indices: list[int] | None = None,
         depth_lim: tuple[float, float] | None = None,
         color_amplitude: bool = True,
         scatter_decimate: int | None = None,
@@ -197,16 +197,13 @@ def __init__(
 
         unique_segments = np.unique(peaks["segment_index"])
 
-        if segment_index is None:
+        if segment_indices is None:
             if len(unique_segments) == 1:
                 segment_indices = [int(unique_segments[0])]
             else:
                 raise ValueError("segment_index must be specified if there are multiple segments")
-        elif isinstance(segment_index, int):
-            segment_indices = [segment_index]
-        elif isinstance(segment_index, list):
-            segment_indices = segment_index
-        else:
+        
+        if not isinstance(segment_indices, list):
             raise ValueError("segment_index must be an int or a list of ints")
 
         # Validate all segment indices exist in the data

src/spikeinterface/widgets/rasters.py

@@ -4,7 +4,7 @@
 from warnings import warn
 
 from .base import BaseWidget, to_attr, default_backend_kwargs
-from .utils import get_some_colors
+from .utils import get_some_colors, validate_segment_indices
 
 
 class BaseRasterWidget(BaseWidget):
@@ -23,7 +23,7 @@ class BaseRasterWidget(BaseWidget):
         converted to a dict of dicts with segment 0.
     unit_ids : array-like | None, default: None
         List of unit_ids to plot
-    segment_index : int | list | None, default: None
+    segment_indices : list | None, default: None
         For multi-segment data, specifies which segment(s) to plot. If None, uses all available segments.
         For single-segment data, this parameter is ignored.
     total_duration : int | None, default: None
@@ -65,7 +65,7 @@ def __init__(
         spike_train_data: dict,
         y_axis_data: dict,
         unit_ids: list | None = None,
-        segment_index: int | list | None = None,
+        segment_indices: list | None = None,
         total_duration: int | None = None,
         plot_histograms: bool = False,
         bins: int | None = None,
@@ -93,22 +93,17 @@ def __init__(
         available_segments.sort()  # Ensure consistent ordering
 
         # Determine which segments to use
-        if segment_index is None:
+        if segment_indices is None:
             # Use all segments by default
             segments_to_use = available_segments
-        elif isinstance(segment_index, int):
-            # Single segment specified
-            if segment_index not in available_segments:
-                raise ValueError(f"segment_index {segment_index} not found in data")
-            segments_to_use = [segment_index]
-        elif isinstance(segment_index, list):
+        elif isinstance(segment_indices, list):
             # Multiple segments specified
-            for idx in segment_index:
+            for idx in segment_indices:
                 if idx not in available_segments:
-                    raise ValueError(f"segment_index {idx} not found in data")
-            segments_to_use = segment_index
+                    raise ValueError(f"segment_index {idx} not found in avialable segments {available_segments}")
+            segments_to_use = segment_indices
         else:
-            raise ValueError("segment_index must be int, list, or None")
+            raise ValueError("segment_index must be `list` or `None`")
 
         # Get all unit IDs present in any segment if not specified
         if unit_ids is None:
@@ -391,7 +386,7 @@ class RasterWidget(BaseRasterWidget):
         A sorting object
     sorting_analyzer : SortingAnalyzer  | None, default: None
         A sorting analyzer object
-    segment_index : int or list of int or None, default: None
+    segment_indices : list of int or None, default: None
         The segment index or indices to use. If None and there are multiple segments, defaults to 0.
         If a list of indices is provided, spike trains are concatenated across the specified segments.
     unit_ids : list
@@ -406,7 +401,7 @@ def __init__(
         self,
         sorting=None,
         sorting_analyzer=None,
-        segment_index=None,
+        segment_indices=None,
         unit_ids=None,
         time_range=None,
         color="k",
@@ -424,30 +419,7 @@ def __init__(
 
         sorting = self.ensure_sorting(sorting)
 
-        num_segments = sorting.get_num_segments()
-
-        # Handle segment_index input
-        if num_segments > 1:
-            if segment_index is None:
-                warn("More than one segment available! Using `segment_index = 0`.")
-                segment_index = 0
-        else:
-            segment_index = 0
-
-        # Convert segment_index to list for consistent processing
-        if isinstance(segment_index, int):
-            segment_indices = [segment_index]
-        elif isinstance(segment_index, list):
-            segment_indices = segment_index
-        else:
-            raise ValueError("segment_index must be an int or a list of ints")
-
-        # Validate segment indices
-        for idx in segment_indices:
-            if not isinstance(idx, int):
-                raise ValueError(f"Each segment index must be an integer, got {type(idx)}")
-            if idx < 0 or idx >= num_segments:
-                raise ValueError(f"segment_index {idx} out of range (0 to {num_segments - 1})")
+        segment_indices = validate_segment_indices(sorting, segment_indices)
 
         if unit_ids is None:
             unit_ids = sorting.unit_ids

src/spikeinterface/widgets/utils.py

@@ -3,6 +3,7 @@
 from warnings import warn
 import numpy as np
 
+from spikeinterface.core import BaseSorting
 
 def get_some_colors(
     keys,
@@ -349,3 +350,46 @@ def make_units_table_from_analyzer(
                 )
 
     return units_table
+
+def validate_segment_indices(segment_indices: list[int] |  None, sorting: BaseSorting):
+    """
+    Validate a list of segment indices for a sorting object.
+
+    Parameters
+    ----------
+    segment_indices : list of int
+        The segment index or indices to validate.
+    sorting : BaseSorting
+        The sorting object to validate against.
+
+    Returns
+    -------
+    list of int
+        A list of valid segment indices.
+
+    Raises
+    ------
+    ValueError
+        If the segment indices are not valid.
+    """
+    num_segments = sorting.get_num_segments()
+
+    # Handle segment_indices input
+    if segment_indices is None:
+        if num_segments > 1:
+            warn("Segment indices not specified. Using first available segment only.")
+        return [0]
+
+    # Convert segment_index to list for consistent processing
+    if not isinstance(segment_indices, list):
+        raise ValueError("segment_indices must be a list of ints - available segments are: " + list(range(num_segments)))
+
+    # Validate segment indices
+    for idx in segment_indices:
+        if not isinstance(idx, int):
+            raise ValueError(f"Each segment index must be an integer, got {type(idx)}")
+        if idx < 0 or idx >= num_segments:
+            raise ValueError(f"segment_index {idx} out of range (0 to {num_segments - 1})")
+
+
+    return segment_indices