brainbox io: spike sorting loader can return collection

Author: oliche

brainbox/io/one.py

@@ -276,7 +276,6 @@ def _load_channel_locations_traj(eid, probe=None, one=None, revision=None, align
     collections = one.list_collections(eid, filename='channels*', collection=collection,
                                        revision=revision)
     probe_collection = _get_spike_sorting_collection(collections, probe)
-    print(probe_collection)
     chn_coords = one.load_dataset(eid, 'channels.localCoordinates', collection=probe_collection)
     depths = chn_coords[:, 1]
 
@@ -388,7 +387,7 @@ def load_channel_locations(eid, probe=None, one=None, aligned=False, brain_atlas
 
 
 def load_spike_sorting_fast(eid, one=None, probe=None, dataset_types=None, spike_sorter=None, revision=None,
-                            brain_regions=None, nested=True, collection=None):
+                            brain_regions=None, nested=True, collection=None, return_collection=False):
     """
     From an eid, loads spikes and clusters for all probes
     The following set of dataset types are loaded:
@@ -407,7 +406,8 @@ def load_spike_sorting_fast(eid, one=None, probe=None, dataset_types=None, spike
     :param return_channels: (bool) defaults to False otherwise tries and load channels from disk
     :param brain_regions: ibllib.atlas.regions.BrainRegions object - will label acronyms if provided
     :param nested: if a single probe is required, do not output a dictionary with the probe name as key
-    :return: spikes, clusters (dict of bunch, 1 bunch per probe)
+    :param return_collection: (False) if True, will return the collection used to load
+    :return: spikes, clusters, channels (dict of bunch, 1 bunch per probe)
     """
     if collection is None:
         collection = _collection_filter_from_args(probe, spike_sorter)
@@ -421,7 +421,10 @@ def load_spike_sorting_fast(eid, one=None, probe=None, dataset_types=None, spike
         channels = channels[k]
         clusters = clusters[k]
         spikes = spikes[k]
-    return spikes, clusters, channels
+    if return_collection:
+        return spikes, clusters, channels, collection
+    else:
+        return spikes, clusters, channels
 
 
 def load_spike_sorting(eid, one=None, probe=None, dataset_types=None, spike_sorter=None, revision=None,