quick fix for spikeglx probe

spikeinterface/extractors/neoextractors/spikeglx.py

@@ -7,13 +7,13 @@
 
 from packaging import version
 
-HAS_NEO_10_2 = version.parse(neo.__version__) >= version.parse('0.10.2')
+HAS_NEO_10_2 = version.parse(neo.__version__) >= version.parse("0.10.2")
 
 
 class SpikeGLXRecordingExtractor(NeoBaseRecordingExtractor):
     """
     Class for reading data from a SpikeGLX system  (NI-DAQ for neuropixel probe)
-    See https://billkarsh.github.io/SpikeGLX/    
+    See https://billkarsh.github.io/SpikeGLX/
 
     Based on neo.rawio.SpikeGLXRawIO
 
@@ -28,22 +28,32 @@ class SpikeGLXRecordingExtractor(NeoBaseRecordingExtractor):
     stream_id: str or None
         stream for instance : 'imec0.ap' 'nidq' or 'imec0.lf'
     """
-    mode = 'folder'
-    NeoRawIOClass = 'SpikeGLXRawIO'
+    mode = "folder"
+    NeoRawIOClass = "SpikeGLXRawIO"
 
     def __init__(self, folder_path, stream_id=None):
-        neo_kwargs = {'dirname': str(folder_path)}
+        neo_kwargs = {"dirname": str(folder_path)}
         if HAS_NEO_10_2:
-            neo_kwargs['load_sync_channel'] = False
+            neo_kwargs["load_sync_channel"] = False
         NeoBaseRecordingExtractor.__init__(self, stream_id=stream_id, **neo_kwargs)
 
-        #~ # open the corresponding stream probe
-        if HAS_NEO_10_2 and '.ap' in self.stream_id:
-            signals_info_dict = {e['stream_name'] : e for e  in self.neo_reader.signals_info_list}
-            meta_filename = signals_info_dict[self.stream_id]['meta_file']
+        # ~ # open the corresponding stream probe
+        if HAS_NEO_10_2 and "nidq" not in self.stream_id:
+            signals_info_dict = {
+                e["stream_name"]: e for e in self.neo_reader.signals_info_list
+            }
+            meta_filename = signals_info_dict[self.stream_id]["meta_file"]
+            # Load probe geometry if available
+            if "lf" in self.stream_id:
+                meta_filename = meta_filename.replace(".lf", ".ap")
             probe = pi.read_spikeglx(meta_filename)
-            self.set_probe(probe, in_place=True)
 
+            if probe.shank_ids is not None:
+                self.set_probe(probe, in_place=True, group_mode="by_shank")
+            else:
+                self.set_probe(probe, in_place=True)
+            self.set_probe(probe, in_place=True)
+
             # load sample shifts
             imDatPrb_type = probe.annotations["imDatPrb_type"]
             sample_shifts = _get_sample_shifts(self.get_num_channels(), imDatPrb_type)
@@ -60,6 +70,7 @@ def read_spikeglx(*args, **kwargs):
 read_spikeglx.__doc__ = SpikeGLXRecordingExtractor.__doc__
 
 
+# TODO check sample shifts for different configurations!!!
 def _get_sample_shifts(num_contact, imDatPrb_type):
     # calculate sample_shift
     # See adc_shift: https://github.com/int-brain-lab/ibllib/blob/master/ibllib/ephys/neuropixel.py
@@ -72,4 +83,4 @@ def _get_sample_shifts(num_contact, imDatPrb_type):
     sample_shift = np.zeros_like(adc)
     for a in adc:
         sample_shift[adc == a] = np.arange(adc_channels) / adc_channels
-    return sample_shift
+    return sample_shift