Merge pull request #382 from int-brain-lab/pykilosort

Pykilosort

Author: oliche

ibllib/dsp/voltage.py

@@ -210,13 +210,19 @@ def destripe(x, fs, tr_sel=None, neuropixel_version=1, butter_kwargs=None, k_kwa
     return x_
 
 
-def decompress_destripe_cbin(sr, output_file=None, h=None):
+def decompress_destripe_cbin(sr, output_file=None, h=None, wrot=None, append=False, nc_out=None, ns2add=0):
     """
-    From a spikeglx Reader object, decompresses and apply ADC
+    From a spikeglx Reader object, decompresses and apply ADC.
+    Saves output as a flat binary file in int16
     Production version with optimized FFTs - requires pyfftw
     :param sr: seismic reader object (spikeglx.Reader)
     :param output_file: (optional, defaults to .bin extension of the compressed bin file)
     :param h: (optional)
+    :param wrot: (optional) whitening matrix [nc x nc] or amplitude scalar to apply to the output
+    :param append: (optional, False) for chronic recordings, append to end of file
+    :param nc_out: (optional, True) saves non selected channels (synchronisation trace) in output
+    :param ns2add: (optional) for kilosort, adds padding samples at the end of the file so the total
+    number of samples is a multiple of the batchsize
     :return:
     """
     import pyfftw
@@ -235,6 +241,7 @@ def decompress_destripe_cbin(sr, output_file=None, h=None):
     taper = np.r_[0, scipy.signal.windows.cosine((SAMPLES_TAPER - 1) * 2), 0]
     # create the FFT stencils
     ncv = h['x'].size  # number of channels
+    nc_out = nc_out or sr.nc
     # compute LP filter coefficients
     sos = scipy.signal.butter(**butter_kwargs, output='sos')
     # compute fft stencil for batchsize
@@ -247,7 +254,7 @@ def decompress_destripe_cbin(sr, output_file=None, h=None):
     DEPHAS = np.exp(1j * np.angle(fft_object(dephas)) * h['sample_shift'][:, np.newaxis])
 
     pbar = tqdm(total=sr.ns / sr.fs)
-    with open(output_file, 'wb') as fid:
+    with open(output_file, 'ab' if append else 'wb') as fid:
         first_s = 0
         while True:
             last_s = np.minimum(NBATCH + first_s, sr.ns)
@@ -273,10 +280,14 @@ def decompress_destripe_cbin(sr, output_file=None, h=None):
             chunk = kfilt(chunk, **k_kwargs)
             # add back sync trace and save
             chunk = np.r_[chunk, sr[first_s:last_s, ncv:].T].T
-            (chunk[slice(*ind2save), :] / sr.channel_conversion_sample2v['ap']
-             ).astype(np.int16).tofile(fid)
+            chunk = chunk[slice(*ind2save), :] / sr.channel_conversion_sample2v['ap']
+            if wrot is not None:
+                chunk[:, :ncv] = np.dot(chunk[:, :ncv], wrot)
+            chunk[:, :nc_out].astype(np.int16).tofile(fid)
             first_s += NBATCH - SAMPLES_TAPER * 2
             pbar.update(NBATCH / sr.fs)
             if last_s == sr.ns:
+                if ns2add > 0:
+                    np.tile(chunk[-1, :nc_out].astype(np.int16), (ns2add, 1)).tofile(fid)
                 break
     pbar.close()

ibllib/ephys/spikes.py

@@ -140,8 +140,8 @@ def _sr(ap_file):
     np.save(st_file, interp_times)
     # get the list of output files
     out_files.extend([f for f in out_path.glob("*.*") if
-                      f.name.startswith(('channels.', 'clusters.', 'spikes.', 'templates.',
-                                         '_kilosort_', '_phy_spikes_subset'))])
+                      f.name.startswith(('channels.', 'drift', 'clusters.', 'spikes.', 'templates.',
+                                         '_kilosort_', '_phy_spikes_subset', '_ibl_log.info'))])
     return out_files, 0
 
 

ibllib/io/extractors/signatures.py

@@ -0,0 +1,200 @@
+# RAW DATA
+RAW_BEHAVIOUR = [
+    ('_iblrig_ambientSensorData.raw', 'raw_behavior_data', True),
+    ('_iblrig_encoderEvents.raw', 'raw_behavior_data', True),
+    ('_iblrig_encoderPositions.raw', 'raw_behavior_data', True),
+    ('_iblrig_encoderTrialInfo.raw', 'raw_behavior_data', True),
+    ('_iblrig_micData.raw', 'raw_behavior_data', False),
+    ('_iblrig_stimPositionScreen.raw', 'raw_behavior_data', False),
+    ('_iblrig_syncSquareUpdate.raw', 'raw_behavior_data', False),
+    ('_iblrig_taskData.raw', 'raw_behavior_data', True),
+    ('_iblrig_taskSettings.raw', 'raw_behavior_data', True),
+]
+
+RAW_PASSIVE = [
+    ('_iblrig_RFMapStim.raw', 'raw_passive_data', True),
+    ('_iblrig_stimPositionScreen.raw', 'raw_passive_data', False),
+    ('_iblrig_syncSquareUpdate.raw', 'raw_passive_data', False),
+    ('_iblrig_taskSettings.raw', 'raw_passive_data', False),
+    ('_iblrig_encoderEvents.raw', 'raw_passive_data', False),
+    ('_iblrig_encoderPositions.raw', 'raw_passive_data', False),
+    ('_iblrig_encoderTrialInfo.raw', 'raw_passive_data', False),
+]
+
+# Data common to both 3A and 3B
+RAW_EPHYS = [
+    ['ephysData.raw.ap', 'raw_ephys_data/XX', True],
+    ['ephysData.raw.ch', 'raw_ephys_data/XX', True, ['ap', 'lf']],
+    ['ephysData.raw.lf', 'raw_ephys_data/XX', True],
+    ['ephysData.raw.meta', 'raw_ephys_data/XX', True, ['ap', 'lf']],
+]
+
+# These are for both 3B probes but only main probe for 3A
+RAW_EPHYS_EXTRA = [
+    ['ephysData.raw.sync', 'raw_ephys_data/XX', True],
+    ['ephysData.raw.timestamps', 'raw_ephys_data/XX', True],
+    ['ephysData.raw.wiring', 'raw_ephys_data/XX', False],
+]
+
+# Only for 3B
+RAW_EPHYS_NIDAQ = [
+    ('ephysData.raw.nidq', 'raw_ephys_data', True),
+    ('ephysData.raw.meta', 'raw_ephys_data', True, ['nidq']),
+    ('ephysData.raw.ch', 'raw_ephys_data', True, ['nidq']),
+]
+
+
+RAW_VIDEO = [
+    ('_iblrig_Camera.raw', 'raw_video_data', True, ['left', 'right', 'body']),
+    ('_iblrig_Camera.timestamps', 'raw_video_data', True, ['left', 'right', 'body']),
+    ('_iblrig_Camera.GPIO', 'raw_video_data', False, ['left', 'right', 'body']),
+    ('_iblrig_Camera.frame_counter', 'raw_video_data', False, ['left', 'right', 'body']),
+]
+
+# PROCESSED DATA
+TRIALS = [
+    ('trials.choice', 'alf', True),
+    ('trials.contrastLeft', 'alf', True),
+    ('trials.contrastRight', 'alf', True),
+    ('trials.feedbackType', 'alf', True),
+    ('trials.feedback_times', 'alf', True),
+    ('trials.firstMovement_times', 'alf', True),
+    ('trials.goCueTrigger_times', 'alf', True),
+    ('trials.goCue_times', 'alf', True),
+    ('trials.intervals', 'alf', True),
+    ('trials.intervals', 'alf', True),
+    ('trials.probabilityLeft', 'alf', True),
+    ('trials.response_times', 'alf', True),
+    ('trials.rewardVolume', 'alf', True),
+    ('trials.stimOff_times', 'alf', True),
+    ('trials.stimOn_times', 'alf', True),
+]
+
+WHEEL = [
+    ('wheel.position', 'alf', True),
+    ('wheel.timestamps', 'alf', True),
+    ('wheelMoves.intervals', 'alf', True),
+    ('wheelMoves.peakAmplitude', 'alf', True),
+]
+
+PASSIVE = [
+    ('_ibl_passiveGabor.table', 'alf', True),
+    ('_ibl_passivePeriods.intervalsTable', 'alf', True),
+    ('_ibl_passiveRFM.times', 'alf', True),
+    ('_ibl_passiveStims.table', 'alf', True),
+]
+
+
+DLC = [
+    ('camera.dlc', 'alf', True, ['left', 'right', 'body']),
+    ('camera.times', 'alf', True, ['left', 'right', 'body']),
+    ('camera.ROIMotionEnergy', 'alf', False, ['left', 'right', 'body']),
+    ('ROIMotionEnergy.position', 'alf', False, ['left', 'right', 'body']),
+]
+
+VIDEO = [
+    ('camera.times', 'alf', True, ['left', 'right', 'body']),
+]
+
+# Data common to both 3A and 3B
+EPHYS = [
+    ['_spikeglx_sync.channels', 'raw_ephys_data/XX', True],
+    ['_spikeglx_sync.polarities', 'raw_ephys_data/XX', True],
+    ['_spikeglx_sync.times', 'raw_ephys_data/XX', True],
+    ['_iblqc_ephysSpectralDensity.freqs', 'raw_ephys_data/XX', True, ['ap', 'lf']],
+    ['_iblqc_ephysSpectralDensity.power', 'raw_ephys_data/XX', True, ['ap', 'lf']],
+    ['_iblqc_ephysTimeRms.rms', 'raw_ephys_data/XX', True, ['ap', 'lf']],
+    ['_iblqc_ephysTimeRms.timestamps', 'raw_ephys_data/XX', True, ['ap', 'lf']],
+]
+
+# Only for 3B
+EPHYS_NIDAQ = [
+    ('_spikeglx_sync.channels', 'raw_ephys_data', True),
+    ('_spikeglx_sync.polarities', 'raw_ephys_data', True),
+    ('_spikeglx_sync.times', 'raw_ephys_data', True),
+]
+
+SPIKE_SORTING = [
+    ['_phy_spikes_subset.channels', 'alf/XX', False],
+    ['_phy_spikes_subset.spikes', 'alf/XX', False],
+    ['_phy_spikes_subset.waveforms', 'alf/XX', False],
+    ['channels.brainLocationIds_ccf_2017', 'alf/XX', False],
+    ['channels.mlapdv', 'alf/XX', False],
+    ['channels.localCoordinates', 'alf/XX', True],
+    ['channels.rawInd', 'alfa/XX', True],
+    ['clusters.amps', 'alf/XX', True],
+    ['clusters.brainLocationAcronyms_ccf_2017', 'alf/XX', False],
+    ['clusters.brainLocationIds_ccf_2017', 'alf/XX', False],
+    ['clusters.channels', 'alf/XX', True],
+    ['clusters.depths', 'alf/XX', True],
+    ['clusters.metrics', 'alf/XX', False],
+    ['clusters.mlapdv', 'alf/XX', False],
+    ['clusters.peakToTrough', 'alf/XX', True],
+    ['clusters.uuids', 'alf/XX', True],
+    ['clusters.waveforms', 'alf/XX', True],
+    ['clusters.waveformsChannels', 'alf/XX', True],
+    ['spikes.amps', 'alf/XX', True],
+    ['spikes.clusters', 'alf/XX', True],
+    ['spikes.depths', 'alf/XX', True],
+    ['spikes.samples', 'alf/XX', True],
+    ['spikes.templates', 'alf/XX', True],
+    ['spikes.times', 'alf/XX', True],
+    ['templates.amps', 'alf/XX', True],
+    ['templates.waveforms', 'alf/XX', True],
+    ['templates.waveformsChannels', 'alf/XX', True],
+]
+
+
+# RAW DATA TASKS
+RAW_BEHAVIOUR_TASKS = ['TrainingRegisterRaw', 'EphysAudio']
+RAW_PASSIVE_TASKS = ['TrainingRegisterRaw']
+RAW_EPHYS_TASKS = ['EphysMtscomp']
+RAW_VIDEO_TASKS = ['TrainingRegisterRaw', 'EphysVideoCompress']
+
+# PROCESSED DATA TASKS
+PASSIVE_TASKS = ['EphysPassive']
+EPHYS_TASKS = ['EphysPulses', 'RawEphysQC']
+VIDEO_TASKS = ['EphysVideoCompress']
+TRIAL_TASKS = ['EphysTrials']
+WHEEL_TASKS = ['EphysTrials']
+SPIKE_SORTING_TASKS = ['SpikeSorting', 'EphysCellsQc']
+DLC_TASKS = ['EphysDLC']
+
+# 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),
+               ('_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)]
+
+EPHYSPASSIVE = [('_iblrig_taskSettings.raw*', 'raw_behavior_data', True),
+                ('_spikeglx_sync.channels.*', 'raw_ephys_data*', True),
+                ('_spikeglx_sync.polarities.*', 'raw_ephys_data*', True),
+                ('_spikeglx_sync.times.*', 'raw_ephys_data*', True),
+                ('*.meta', 'raw_ephys_data*', True),
+                ('_iblrig_RFMapStim.raw*', 'raw_passive_data', True)]
+
+RAWEPHYSQC = [('*.meta', 'raw_ephys_data/probe*', True),
+              ('*.ch', 'raw_ephys_data/probe*', True),
+              ('*.cbin', 'raw_ephys_data/probe*', True)]
+
+
+def spike_sorting_signature(pname=None):
+    pname = pname if pname is not None else "probe*"
+    signature = [('*ap.meta', f'raw_ephys_data/{pname}', True),
+                 ('*ap.ch', f'raw_ephys_data/{pname}', True),
+                 ('*ap.cbin', f'raw_ephys_data/{pname}', True),
+                 ('_spikeglx_sync.channels.*', 'raw_ephys_data*', True),
+                 ('_spikeglx_sync.polarities.*', 'raw_ephys_data*', True),
+                 ('_spikeglx_sync.times.*', 'raw_ephys_data*', True),
+                 ('_iblrig_taskData.raw.*', 'raw_behavior_data', True),
+                 ('_iblrig_taskSettings.raw.*', 'raw_behavior_data', True)
+                 ]
+    return signature
+
+
+SPIKESORTING = [spike_sorting_signature()]

ibllib/oneibl/aws.py

@@ -0,0 +1,77 @@
+import logging
+from pathlib import Path
+import sys
+from time import time
+
+import boto3
+import numpy as np
+
+from one.api import ONE
+from one.alf.files import add_uuid_string
+from iblutil.io.parquet import np2str
+
+
+_logger = logging.getLogger('ibllib')
+
+AWS_ROOT_PATH = Path('/mnt/ibl')
+BUCKET_NAME = 'ibl-brain-wide-map-private'
+
+# To get aws credentials follow
+# https://docs.aws.amazon.com/cli/latest/userguide/install-cliv2-linux.html to install aws cli
+# https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-quickstart.html to set up
+# credentials
+
+
+class AWS:
+    def __init__(self, s3_bucket_name=None, one=None):
+        # TODO some initialisation routine to set up credentials for the first time
+
+        s3 = boto3.resource('s3')
+        self.bucket_name = s3_bucket_name or BUCKET_NAME
+        self.bucket = s3.Bucket(self.bucket_name)
+        self.one = one or ONE()
+
+    def _download_datasets(self, datasets):
+
+        for _, d in datasets.iterrows():
+            rel_file_path = Path(d['session_path']).joinpath(d['rel_path'])
+            file_path = Path(self.one.cache_dir).joinpath(rel_file_path)
+            file_path.parent.mkdir(exist_ok=True, parents=True)
+
+            if file_path.exists():
+                # already downloaded, need to have some options for overwrite, clobber, look
+                # for file mismatch like in ONE
+                _logger.info(f'{file_path} already exists wont redownload')
+                continue
+
+            aws_path = AWS_ROOT_PATH.joinpath(add_uuid_string(rel_file_path,
+                                                              np2str(np.r_[d.name[0], d.name[1]])))
+            aws_path = as_aws_path(aws_path)
+            file_path = as_aws_path(file_path)
+            # maybe should avoid this and do a try catch instead?, see here
+            # https://boto3.amazonaws.com/v1/documentation/api/latest/guide/collections.html#filtering
+            # probably better to do filter on collection ? Not for today
+            objects = list(self.bucket.objects.filter(Prefix=aws_path))
+            if len(objects) == 1:
+                ts = time()
+                _logger.info(f'Downloading {aws_path} to {file_path}')
+                self.bucket.download_file(aws_path, file_path)
+                _logger.debug(f'Complete. Time elapsed {time() - ts} for {file_path}')
+            else:
+                _logger.warning(f'{aws_path} not found on s3 bucket: {self.bucket.name}')
+
+
+def as_aws_path(path):
+    """
+    Convert a path into one suitable for the aws. Mainly for windows to convert // to \
+
+    :param path: A Path instance
+    :return: A formatted path string
+
+    """
+    if sys.platform == 'win32':
+        path = '/'.join(str(path).split('\\'))
+    else:
+        path = str(path)
+
+    return path