feat: update ephys for nienborglab pipeline

element_array_ephys/ephys_no_curation.py

@@ -4,13 +4,11 @@
 import pathlib
 import re
 from decimal import Decimal
-
 import datajoint as dj
 import numpy as np
 import pandas as pd
-import spikeinterface as si
+
 from element_interface.utils import dict_to_uuid, find_full_path, find_root_directory
-from spikeinterface import exporters, postprocessing, qualitymetrics, sorters
 
 from . import ephys_report, probe
 from .readers import kilosort, openephys, spikeglx
@@ -818,7 +816,7 @@ def infer_output_dir(cls, key, relative: bool = False, mkdir: bool = False):
         output_dir = (
             processed_dir
             / session_dir.relative_to(root_dir)
-            / f'probe_{key["insertion_number"]}'
+            / f'insertion_{key["insertion_number"]}'
             / f'{method}_{key["paramset_idx"]}'
         )
 
@@ -1041,6 +1039,9 @@ def make(self, key):
         si_sorting_dir = output_dir / sorter_name / "spike_sorting"
 
         if si_waveform_dir.exists():  # Read from spikeinterface outputs
+            import spikeinterface as si
+            from spikeinterface import sorters
+
             we: si.WaveformExtractor = si.load_waveforms(
                 si_waveform_dir, with_recording=False
             )
@@ -1274,6 +1275,7 @@ def make(self, key):
 
         si_waveform_dir = output_dir / sorter_name / "waveform"
         if si_waveform_dir.exists():  # read from spikeinterface outputs
+            import spikeinterface as si
             we: si.WaveformExtractor = si.load_waveforms(
                 si_waveform_dir, with_recording=False
             )

element_array_ephys/spike_sorting/si_preprocessing.py

@@ -35,3 +35,14 @@ def IBLdestriping_modified(recording):
         recording, operator="median", reference="global"
     )
     return recording
+
+
+def NienborgLab_preproc(recording):
+    """Preprocessing pipeline for 32chn ephys data from Trellis."""
+    recording = si.preprocessing.bandpass_filter(
+        recording=recording, freq_min=300, freq_max=6000
+    )
+    recording = si.preprocessing.common_reference(
+        recording=recording, operator="median"
+    )
+    return recording

element_array_ephys/spike_sorting/si_spike_sorting.py

@@ -90,6 +90,18 @@ def make(self, key):
                     f"{required_key} must be defined in ClusteringParamSet for SpikeInterface execution"
                 )
 
+        # Get probe information to recording object
+        electrodes_df = (
+            (
+                ephys.EphysRecording.Channel
+                * probe.ElectrodeConfig.Electrode
+                * probe.ProbeType.Electrode
+                & key
+            )
+            .fetch(format="frame")
+            .reset_index()
+        )
+
         # Set directory to store recording file.
         if not output_dir:
             output_dir = ephys.ClusteringTask.infer_output_dir(
@@ -112,37 +124,52 @@ def make(self, key):
             )
             spikeglx_recording.validate_file("ap")
             data_dir = spikeglx_meta_filepath.parent
+
+            si_extractor: si.extractors.neoextractors = (
+                si.extractors.extractorlist.recording_extractor_full_dict["spikeglx"]
+            )
+            stream_names, stream_ids = si.extractors.get_neo_streams(
+                acq_software, folder_path=data_dir
+            )
+            si_recording: si.BaseRecording = si_extractor(
+                folder_path=data_dir, stream_name=stream_names[0]
+            )
         elif acq_software == "Open Ephys":
             oe_probe = ephys.get_openephys_probe_data(key)
             assert len(oe_probe.recording_info["recording_files"]) == 1
             data_dir = oe_probe.recording_info["recording_files"][0]
+            si_extractor: si.extractors.neoextractors = (
+                si.extractors.extractorlist.recording_extractor_full_dict["openephys"]
+            )
+            stream_names, stream_ids = si.extractors.get_neo_streams(
+                acq_software, folder_path=data_dir
+            )
+            si_recording: si.BaseRecording = si_extractor(
+                folder_path=data_dir, stream_name=stream_names[0]
+            )
+        elif acq_software == "Trellis":
+            si_extractor: si.extractors.neoextractors = (
+                si.extractors.extractorlist.recording_extractor_full_dict["blackrock"]
+            )
+
+            nsx5_relpaths = (ephys.EphysRecording.EphysFile & key).fetch("file_path")
+            nsx5_fullpaths = [
+                find_full_path(ephys.get_ephys_root_data_dir(), f + ".ns5")
+                for f in nsx5_relpaths
+            ]
+            si_recs = []
+            for f in nsx5_fullpaths:
+                si_rec = si_extractor(file_path=f, stream_name="nsx5")
+                # find & remove non-ephys channels
+                non_ephys_chns = set(si_rec.channel_ids) - set(
+                    (electrodes_df.channel_idx.values + 1).astype(str)
+                )
+                si_recs.append(si_rec.remove_channels(list(non_ephys_chns)))
+            si_recording = si.concatenate_recordings(si_recs)
         else:
             raise NotImplementedError(
                 f"SpikeInterface processing for {acq_software} not yet implemented."
             )
-        acq_software = acq_software.replace(" ", "").lower()
-        si_extractor: si.extractors.neoextractors = (
-            si.extractors.extractorlist.recording_extractor_full_dict[acq_software]
-        )  # data extractor object
-
-        stream_names, stream_ids = si.extractors.get_neo_streams(
-            acq_software, folder_path=data_dir
-        )
-        si_recording: si.BaseRecording = si_extractor(
-            folder_path=data_dir, stream_name=stream_names[0]
-        )
-
-        # Add probe information to recording object
-        electrodes_df = (
-            (
-                ephys.EphysRecording.Channel
-                * probe.ElectrodeConfig.Electrode
-                * probe.ProbeType.Electrode
-                & key
-            )
-            .fetch(format="frame")
-            .reset_index()
-        )
 
         # Create SI probe object
         si_probe = readers.probe_geometry.to_probeinterface(