Use SpikeSortingLoader instead of load_spike_sorting_fast

Author: k1o0

brainbox/io/one.py

@@ -599,7 +599,7 @@ class SpikeSortingLoader:
     This class can be instantiated in several manners
     - With Alyx database probe id:
             SpikeSortingLoader(pid=pid, one=one)
-    - With Alyx database eic and probe name:
+    - With Alyx database eid and probe name:
             SpikeSortingLoader(eid=eid, pname='probe00', one=one)
     - From a local session and probe name:
             SpikeSortingLoader(session_path=session_path, pname='probe00')

ibllib/plots/figures.py

@@ -26,7 +26,7 @@
 from brainbox.behavior.dlc import SAMPLING, plot_trace_on_frame, plot_wheel_position, plot_lick_hist, \
     plot_lick_raster, plot_motion_energy_hist, plot_speed_hist, plot_pupil_diameter_hist
 from brainbox.ephys_plots import image_lfp_spectrum_plot, image_rms_plot, plot_brain_regions
-from brainbox.io.one import load_spike_sorting_fast
+from brainbox.io.one import SpikeSortingLoader
 from brainbox.behavior import training
 from iblutil.numerical import ismember
 from ibllib.plots.misc import Density
@@ -360,11 +360,12 @@ def plot_driftmap(self, spikes, clusters, channels, collection):
             if all_here and len(output_files) == len(spike_sorting_runs):
                 return output_files
             logger.info(self.output_directory)
-            for run in spike_sorting_runs:
-                collection = str(Path(run).parent.as_posix())
-                spikes, clusters, channels = load_spike_sorting_fast(
-                    eid=self.eid, probe=self.pname, one=self.one, nested=False, collection=collection,
-                    dataset_types=['spikes.depths'], brain_regions=self.brain_regions)
+            ss = SpikeSortingLoader(
+                one=self.one, pid=self.pid, eid=self.eid, pname=self.pname, session_path=self.session_path)
+            for run in map(self.session_path.joinpath, spike_sorting_runs):
+                sorter = run.without_revision().relative_to(self.session_path / f'alf/{self.pname}')
+                collection = run.relative_to(self.session_path).parent.as_posix()
+                spikes, clusters, channels = ss.load_spike_sorting(spike_sorter=sorter, revision=run.revision or None)
 
                 if 'atlas_id' not in channels.keys():
                     channels = self.get_channels('channels', collection)

ibllib/plots/snapshot.py

@@ -7,6 +7,7 @@
 
 from one.api import ONE
 from one.alf.spec import is_uuid
+from one.alf.path import ensure_alf_path
 from ibllib.pipes import tasks
 from one.alf.exceptions import ALFObjectNotFound
 from neuropixel import trace_header, TIP_SIZE_UM
@@ -70,7 +71,7 @@ def __init__(self, pid, session_path=None, one=None, brain_regions=None, brain_a
         self.content_type = 'probeinsertion'
         self.pid = pid
         self.eid, self.pname = self.one.pid2eid(self.pid)
-        self.session_path = session_path or self.one.eid2path(self.eid)
+        self.session_path = ensure_alf_path(session_path or self.one.eid2path(self.eid))
         self.output_directory = self.session_path.joinpath('snapshot', self.pname)
         self.output_directory.mkdir(exist_ok=True, parents=True)
         self.histology_status = None