Merge pull request #402 from int-brain-lab/globus_download_localserver

Globus download localserver

Author: mayofaulkner

brainbox/io/one.py

@@ -95,14 +95,24 @@ def _channels_traj2bunch(xyz_chans, brain_atlas):
     return channels
 
 
+def _channels_bunch2alf(channels):
+    channels_ = {
+        'mlapdv': np.c_[channels['x'], channels['y'], channels['z']] * 1e6,
+        'brainLocationIds_ccf_2017': channels['atlas_id'],
+        'localCoordinates': np.c_[channels['lateral_um'], channels['axial_um']]}
+    return channels_
+
+
 def _channels_alf2bunch(channels, brain_regions=None):
     # reformat the dictionary according to the standard that comes out of Alyx
     channels_ = {
         'x': channels['mlapdv'][:, 0].astype(np.float64) / 1e6,
         'y': channels['mlapdv'][:, 1].astype(np.float64) / 1e6,
         'z': channels['mlapdv'][:, 2].astype(np.float64) / 1e6,
         'acronym': None,
-        'atlas_id': channels['brainLocationIds_ccf_2017']
+        'atlas_id': channels['brainLocationIds_ccf_2017'],
+        'axial_um': channels['localCoordinates'][:, 1],
+        'lateral_um': channels['localCoordinates'][:, 0],
     }
     if brain_regions:
         channels_['acronym'] = brain_regions.get(channels_['atlas_id'])['acronym']
@@ -207,24 +217,33 @@ def _load_channels_locations_from_disk(eid, collection=None, one=None, revision=
             channels_aligned = one.load_object(eid, 'channels', collection=ac_collection)
             channels[probe] = channel_locations_interpolation(channels_aligned, channels[probe])
             # only have to reformat channels if we were able to load coordinates from disk
-            channels[probe] = _channels_alf2bunch(channels[probe], brain_regions=brain_regions)
+        channels[probe] = _channels_alf2bunch(channels[probe], brain_regions=brain_regions)
     return channels
 
 
-def channel_locations_interpolation(channels_aligned, channels):
+def channel_locations_interpolation(channels_aligned, channels, brain_regions=None):
     """
     oftentimes the channel map for different spike sorters may be different so interpolate the alignment onto
     if there is no spike sorting in the base folder, the alignment doesn't have the localCoordinates field
     so we reconstruct from the Neuropixel map. This only happens for early pykilosort sorts
     :param channels_aligned: Bunch or dictionary of aligned channels containing at least keys
-     'mlapdv' and 'brainLocationIds_ccf_2017' - those are the guide for the interpolation
+     'localCoordinates', 'mlapdv' and 'brainLocationIds_ccf_2017'
+     OR
+      'x', 'y', 'z', 'acronym', 'axial_um'
+      those are the guide for the interpolation
     :param channels: Bunch or dictionary of aligned channels containing at least keys 'localCoordinates'
-    :return: Bunch or dictionary of channels with extra keys 'mlapdv' and 'brainLocationIds_ccf_2017'
+    :param brain_regions: None (default) or ibllib.atlas.BrainRegions object
+     if None will return a dict with keys 'localCoordinates', 'mlapdv', 'brainLocationIds_ccf_2017
+     if a brain region object is provided, outputts a dict with keys
+      'x', 'y', 'z', 'acronym', 'atlas_id', 'axial_um', 'lateral_um'
+    :return: Bunch or dictionary of channels with brain coordinates keys
     """
     nch = channels['localCoordinates'].shape[0]
+    if set(['x', 'y', 'z']).issubset(set(channels_aligned.keys())):
+        channels_aligned = _channels_bunch2alf(channels_aligned)
     if 'localCoordinates' in channels_aligned.keys():
         aligned_depths = channels_aligned['localCoordinates'][:, 1]
-    else:
+    else:  # this is a edge case for a few spike sorting sessions
         assert channels_aligned['mlapdv'].shape[0] == 384
         NEUROPIXEL_VERSION = 1
         from ibllib.ephys.neuropixel import trace_header
@@ -238,7 +257,10 @@ def channel_locations_interpolation(channels_aligned, channels):
     # the brain locations have to be interpolated by nearest neighbour
     fcn_interp = interp1d(depth_aligned, channels_aligned['brainLocationIds_ccf_2017'][ind_aligned], kind='nearest')
     channels['brainLocationIds_ccf_2017'] = fcn_interp(depths)[iinv].astype(np.int32)
-    return channels
+    if brain_regions is not None:
+        return _channels_alf2bunch(channels, brain_regions=brain_regions)
+    else:
+        return channels
 
 
 def _load_channel_locations_traj(eid, probe=None, one=None, revision=None, aligned=False,
@@ -531,7 +553,7 @@ def merge_clusters_channels(dic_clus, channels, keys_to_add_extra=None):
     dic_clus : dict of one.alf.io.AlfBunch
         1 bunch per probe, containing cluster information
     channels : dict of one.alf.io.AlfBunch
-        1 bunch per probe, containing channels bunch with keys ('acronym', 'atlas_id')
+        1 bunch per probe, containing channels bunch with keys ('acronym', 'atlas_id', 'x', 'y', z', 'localCoordinates')
     keys_to_add_extra : list of str
         Any extra keys to load into channels bunches
 
@@ -541,7 +563,7 @@ def merge_clusters_channels(dic_clus, channels, keys_to_add_extra=None):
         clusters (1 bunch per probe) with new keys values.
     """
     probe_labels = list(channels.keys())  # Convert dict_keys into list
-    keys_to_add_default = ['acronym', 'atlas_id', 'x', 'y', 'z']
+    keys_to_add_default = ['acronym', 'atlas_id', 'x', 'y', 'z', 'axial_um', 'lateral_um']
 
     if keys_to_add_extra is None:
         keys_to_add = keys_to_add_default
@@ -550,10 +572,9 @@ def merge_clusters_channels(dic_clus, channels, keys_to_add_extra=None):
         keys_to_add = list(set(keys_to_add_extra + keys_to_add_default))
 
     for label in probe_labels:
-        try:
-            clu_ch = dic_clus[label]['channels']
-
-            for key in keys_to_add:
+        clu_ch = dic_clus[label]['channels']
+        for key in keys_to_add:
+            try:
                 assert key in channels[label].keys()  # Check key is in channels
                 ch_key = channels[label][key]
                 nch_key = len(ch_key) if ch_key is not None else 0
@@ -564,11 +585,9 @@ def merge_clusters_channels(dic_clus, channels, keys_to_add_extra=None):
                         f'Probe {label}: merging channels and clusters for key "{key}" has {nch_key} on channels'
                         f' but expected {max(clu_ch)}. Data in new cluster key "{key}" is returned empty.')
                     dic_clus[label][key] = []
-        except AssertionError:
-            _logger.warning(
-                f'Either clusters or channels does not have key {label}, could not'
-                f' merge')
-            continue
+            except AssertionError:
+                _logger.warning(f'Either clusters or channels does not have key {key}, could not merge')
+                continue
 
     return dic_clus
 

brainbox/task/passive.py

@@ -200,6 +200,8 @@ def get_stim_aligned_activity(stim_events, spike_times, spike_depths, z_score_fl
     stim_activity = {}
     for stim_type, stim_times in stim_events.items():
 
+        # Get rid of any nan values
+        stim_times = stim_times[~np.isnan(stim_times)]
         stim_intervals = np.c_[stim_times - pre_stim, stim_times + post_stim]
         base_intervals = np.c_[stim_times - base_stim, stim_times - pre_stim]
         out_intervals = stim_intervals[:, 1] > times[-1]

ibllib/ephys/np2_converter.py

@@ -4,6 +4,7 @@
 import numpy as np
 from pathlib import Path
 import copy
+import shutil
 import logging
 _logger = logging.getLogger('ibllib')
 
@@ -14,17 +15,19 @@ class NP2Converter:
     individual shanks
     """
 
-    def __init__(self, ap_file, post_check=True, delete_original=False):
+    def __init__(self, ap_file, post_check=True, delete_original=False, compress=True):
         """
         :param ap_file: ap.bin spikeglx file to process
         :param post_check: whether to apply post-check integrity test to ensure split content is
         identical to original content (only applicable to NP2.4)
-        :param delete_original: whether to delete the original ap file after data has been
+        :param delete_original: whether to delete the original ap file after data has been split
+        :param compress: whether to apply mtscomp to extracted .bin files
         split into shanks (only applicable to NP2.4)
         """
         self.ap_file = Path(ap_file)
         self.sr = spikeglx.Reader(ap_file)
         self.post_check = post_check
+        self.compress = compress
         self.delete_original = delete_original
         self.np_version = spikeglx._get_neuropixel_version_from_meta(self.sr.meta)
         self.check_metadata()
@@ -142,8 +145,10 @@ def _process_NP24(self, overwrite=False):
 
         if self.post_check:
             self.check_NP24()
+        if self.compress:
+            self.compress_NP24(overwrite=overwrite)
         if self.delete_original:
-            self.delete()
+            self.delete_NP24()
 
         return 1
 
@@ -174,11 +179,15 @@ def _prepare_files_NP24(self, overwrite=False):
             probe_path = self.ap_file.parent.parent.joinpath(label + chr(97 + int(sh)) + self.extra)
 
             if not probe_path.exists() or overwrite:
+                if self.sr.is_mtscomp:
+                    ap_file_bin = self.ap_file.with_suffix('.bin').name
+                else:
+                    ap_file_bin = self.ap_file.name
                 probe_path.mkdir(parents=True, exist_ok=True)
-                _shank_info['ap_file'] = probe_path.joinpath(self.ap_file.name)
+                _shank_info['ap_file'] = probe_path.joinpath(ap_file_bin)
                 _shank_info['ap_open_file'] = open(_shank_info['ap_file'], 'wb')
                 _shank_info['lf_file'] = probe_path.joinpath(
-                    self.ap_file.name.replace('ap', 'lf'))
+                    ap_file_bin.replace('ap', 'lf'))
                 _shank_info['lf_open_file'] = open(_shank_info['lf_file'], 'wb')
             else:
                 self.already_exists = True
@@ -210,6 +219,7 @@ def _process_NP21(self, overwrite=False):
         wg = WindowGenerator(self.nsamples, self.samples_window, self.samples_overlap)
 
         for first, last in wg.firstlast:
+
             chunk_lf = self.extract_lfp(self.sr[first:last, :self.napch].T)
             chunk_lf_sync = self.extract_lfp_sync(self.sr[first:last, self.idxsyncch:].T)
 
@@ -224,6 +234,9 @@ def _process_NP21(self, overwrite=False):
 
         self._writemetadata_lf()
 
+        if self.compress:
+            self.compress_NP21(overwrite=overwrite)
+
         return 1
 
     def _prepare_files_NP21(self, overwrite=False):
@@ -242,8 +255,9 @@ def _prepare_files_NP21(self, overwrite=False):
         shank_info = {}
         self.already_exists = False
 
-        lf_file = self.ap_file.parent.joinpath(self.ap_file.name.replace('ap', 'lf'))
-        if not lf_file.exists() or overwrite:
+        lf_file = self.ap_file.parent.joinpath(self.ap_file.name.replace('ap', 'lf')).with_suffix('.bin')
+        lf_cbin_file = lf_file.with_suffix('.cbin')
+        if not (lf_file.exists() or lf_cbin_file.exists()) or overwrite:
             for sh in n_shanks:
                 _shank_info = {}
                 # channels for individual shank + sync channel
@@ -293,15 +307,66 @@ def check_NP24(self):
 
         self.check_completed = True
 
+    def compress_NP24(self, overwrite=False, **kwargs):
+        """
+        Compress spikeglx files
+        :return:
+        """
+        for sh in self.shank_info.keys():
+            bin_file = self.shank_info[sh]['ap_file']
+            if overwrite:
+                cbin_file = bin_file.with_suffix('.cbin')
+                cbin_file.unlink()
+
+            sr_ap = spikeglx.Reader(bin_file)
+            cbin_file = sr_ap.compress_file(**kwargs)
+            sr_ap.close()
+            bin_file.unlink()
+            self.shank_info[sh]['ap_file'] = cbin_file
+
+            bin_file = self.shank_info[sh]['lf_file']
+            if overwrite:
+                cbin_file = bin_file.with_suffix('.cbin')
+                cbin_file.unlink()
+            sr_lf = spikeglx.Reader(bin_file)
+            cbin_file = sr_lf.compress_file(**kwargs)
+            sr_lf.close()
+            bin_file.unlink()
+            self.shank_info[sh]['lf_file'] = cbin_file
+
+    def compress_NP21(self, overwrite=False):
+        """
+        Compress spikeglx files
+        :return:
+        """
+        for sh in self.shank_info.keys():
+            if not self.sr.is_mtscomp:
+                cbin_file = self.sr.compress_file()
+                self.sr.close()
+                self.ap_file.unlink()
+                self.ap_file = cbin_file
+                self.sr = spikeglx.Reader(self.ap_file)
+
+            bin_file = self.shank_info[sh]['lf_file']
+            if overwrite:
+                cbin_file = bin_file.with_suffix('.cbin')
+                cbin_file.unlink()
+            sr_lf = spikeglx.Reader(bin_file)
+            cbin_file = sr_lf.compress_file()
+            sr_lf.close()
+            bin_file.unlink()
+            self.shank_info[sh]['lf_file'] = cbin_file
+
     def delete_NP24(self):
         """
         Delete the original ap file that doesn't has all shanks in one file
 
         :return:
         """
         if self.check_completed and self.delete_original:
-            # TODO need to delete the original wahhhhh
-            pass
+            _logger.info(f'Removing original files in folder {self.ap_file.parent}')
+            self.sr.close()
+            shutil.rmtree(self.ap_file.parent)
 
     def _split2shanks(self, chunk, etype='ap'):
         """

ibllib/io/extractors/signatures.py

@@ -163,13 +163,13 @@
 # DATA REQUIRED FOR TASKS
 EPHYSTRIALS = [('_iblrig_taskData.raw.*', 'raw_behavior_data', True),
                ('_iblrig_taskSettings.raw.*', 'raw_behavior_data', True),
-               ('_spikeglx_sync.npy', 'raw_ephys_data/**', True),
-               ('_spikeglx_sync.polarities.npy', 'raw_ephys_data/**', True),
-               ('_spikeglx_sync.times.npy', 'raw_ephys_data/**', True),
+               ('_spikeglx_sync.npy', 'raw_ephys_data*', True),
+               ('_spikeglx_sync.polarities.npy', 'raw_ephys_data*', True),
+               ('_spikeglx_sync.times.npy', 'raw_ephys_data*', True),
                ('_iblrig_encoderEvents.raw', 'raw_behavior_data', True),
                ('_iblrig_encoderPositions.raw', 'raw_behavior_data', True),
-               ('*wiring.json', 'raw_ephys_data/**', False),
-               ('*.meta', 'raw_ephys_data/**', True)]
+               ('*wiring.json', 'raw_ephys_data*', False),
+               ('*.meta', 'raw_ephys_data*', True)]
 
 EPHYSPASSIVE = [('_iblrig_taskSettings.raw*', 'raw_behavior_data', True),
                 ('_spikeglx_sync.channels.*', 'raw_ephys_data*', True),

ibllib/io/spikeglx.py

@@ -589,6 +589,18 @@ def get_neuropixel_version_from_files(ephys_files):
         return '3A'
 
 
+def get_probes_from_folder(session_path):
+    # should glob the ephys files and get out the labels
+    # This assumes the meta files exist on the server (this is the case for now but should it be?)
+    ephys_files = glob_ephys_files(session_path, ext='meta')
+    probes = []
+    for files in ephys_files:
+        if files['label']:
+            probes.append(files['label'])
+
+    return probes
+
+
 def glob_ephys_files(session_path, suffix='.meta', ext='bin', recursive=True, bin_exists=True):
     """
     From an arbitrary folder (usually session folder) gets the ap and lf files and labels