add sort_by_depth kwargs to RasterWidget

Author: tayheau

src/spikeinterface/widgets/rasters.py

@@ -363,6 +363,8 @@ class RasterWidget(BaseRasterWidget):
         A sorting object. Deprecated.
     sorting_analyzer : SortingAnalyzer  | None, default: None
         A sorting analyzer object. Deprecated.
+    sort_by_depth: bool = False
+        Wether or not to sort units by depth, default: False
     """
 
     def __init__(
@@ -375,6 +377,7 @@ def __init__(
         backend: str | None = None,
         sorting: BaseSorting | None = None,
         sorting_analyzer: SortingAnalyzer | None = None,
+        sort_by_depth : bool = False,
         **backend_kwargs,
     ):
         if sorting is not None:
@@ -387,13 +390,23 @@ def __init__(
             warn(deprecation_msg, category=DeprecationWarning, stacklevel=2)
             sorting_analyzer_or_sorting = sorting_analyzer
 
+        
+
         sorting = self.ensure_sorting(sorting_analyzer_or_sorting)
 
         segment_indices = validate_segment_indices(segment_indices, sorting)
 
         if unit_ids is None:
             unit_ids = sorting.unit_ids
 
+        if sort_by_depth :
+            # print("hey")
+            if not sorting_analyzer_or_sorting.has_extension("unit_locations"):
+                raise AttributeError(f"'unit_locations' necessary for sort_by_depth is True")
+            depths = sorting_analyzer_or_sorting.get_extension("unit_locations").get_data(outputs="numpy")
+            s_args_depths = np.argsort(depths[:, 1])
+            unit_ids = unit_ids[s_args_depths]
+
         # Create dict of dicts structure
         spike_train_data = {}
         y_axis_data = {}