Fixes to MDASortingExtractor

SpikeInterface · Oct 2, 2023 · 5cefdac · 5cefdac
1 parent d44e99c
commit 5cefdac
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/src/spikeinterface/extractors/mdaextractors.py b/src/spikeinterface/extractors/mdaextractors.py
@@ -216,14 +216,14 @@ def write_sorting(sorting, save_path, write_primary_channels=False):
         times_list = []
         labels_list = []
         primary_channels_list = []
-        for unit_id_i, unit_id in enumerate(unit_ids):
+        for unit_index, unit_id in enumerate(unit_ids):
             times = sorting.get_unit_spike_train(unit_id=unit_id)
             times_list.append(times)
             # unit id may not be numeric
-            if unit_id.dtype.kind in "biufc":
-                labels_list.append(np.ones(times.shape) * unit_id)
+            if unit_id.dtype.kind in "iu":
+                labels_list.append(np.ones(times.shape, dtype=unit_id.dtype) * unit_id)
             else:
-                labels_list.append(np.ones(times.shape) * unit_id_i)
+                labels_list.append(np.ones(times.shape, dtype=int) * unit_index)
             if write_primary_channels:
                 if "max_channel" in sorting.get_unit_property_names(unit_id):
                     primary_channels_list.append([sorting.get_unit_property(unit_id, "max_channel")] * times.shape[0])