Merge pull request #428 from int-brain-lab/olivier

Olivier

Author: mayofaulkner

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):
+                            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:
@@ -403,12 +402,15 @@ def load_spike_sorting_fast(eid, one=None, probe=None, dataset_types=None, spike
     :param probe: name of probe to load in, if not given all probes for session will be loaded
     :param dataset_types: additional spikes/clusters objects to add to the standard default list
     :param spike_sorter: name of the spike sorting you want to load (None for default)
+    :param collection: name of the spike sorting collection to load - exclusive with spike sorter name ex: "alf/probe00"
     :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)
     """
-    collection = _collection_filter_from_args(probe, spike_sorter)
+    if collection is None:
+        collection = _collection_filter_from_args(probe, spike_sorter)
     _logger.debug(f"load spike sorting with collection filter {collection}")
     kwargs = dict(eid=eid, one=one, collection=collection, revision=revision, dataset_types=dataset_types,
                   brain_regions=brain_regions)
@@ -419,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,

examples/one/histology/register_lasagna_tracks_alyx.py

@@ -24,7 +24,7 @@
 """
 # Author: Olivier, Gaelle
 from pathlib import Path
-
+from ibllib.atlas import AllenAtlas
 from one.api import ONE
 
 from ibllib.pipes import histology
@@ -42,11 +42,12 @@
 
 # ======== DO NOT EDIT BELOW ====
 one = ONE(base_url=ALYX_URL)
+ba = AllenAtlas()
 
 if EXAMPLE_OVERWRITE:
     # TODO Olivier : Function to download examples folder
     cachepath = Path(one.alyx.cache_dir)
     path_tracks = cachepath.joinpath('examples', 'histology', 'tracks_to_add')
 
-histology.register_track_files(path_tracks=path_tracks, one=one, overwrite=True)
-histology.detect_missing_histology_tracks(path_tracks=path_tracks, one=one)
+histology.register_track_files(path_tracks=path_tracks, one=one, overwrite=True, brain_atlas=ba)
+histology.detect_missing_histology_tracks(path_tracks=path_tracks, one=one, brain_atlas=ba)

ibllib/atlas/atlas.py

@@ -2,7 +2,7 @@
 import logging
 import matplotlib.pyplot as plt
 from pathlib import Path, PurePosixPath
-
+from functools import lru_cache
 import numpy as np
 import nrrd
 
@@ -710,6 +710,7 @@ def get_brain_entry(traj, brain_atlas):
         return Insertion._get_surface_intersection(traj, brain_atlas, surface='top')
 
 
+@lru_cache(maxsize=1)
 class AllenAtlas(BrainAtlas):
     """
     Instantiates an atlas.BrainAtlas corresponding to the Allen CCF at the given resolution

ibllib/dsp/voltage.py

@@ -109,7 +109,7 @@ def fk(x, si=.002, dx=1, vbounds=None, btype='highpass', ntr_pad=0, ntr_tap=None
     return xf / gain
 
 
-def car(x, collection=None, lagc=300, butter_kwargs=None):
+def car(x, collection=None, lagc=300, butter_kwargs=None, **kwargs):
     """
     Applies common average referencing with optional automatic gain control
     :param x: the input array to be filtered. dimension, the filtering is considering
@@ -223,7 +223,22 @@ def interpolate_bad_channels(data, channel_labels=None, h=None, p=1.3, kriging_d
     return data
 
 
-def destripe(x, fs, neuropixel_version=1, butter_kwargs=None, k_kwargs=None, channel_labels=None):
+def _get_destripe_parameters(fs, butter_kwargs, k_kwargs, k_filter):
+    """gets the default params for destripe. This is used for both the destripe fcn on a
+    numpy array and the function that actuates on a cbin file"""
+    if butter_kwargs is None:
+        butter_kwargs = {'N': 3, 'Wn': 300 / fs * 2, 'btype': 'highpass'}
+    if k_kwargs is None:
+        k_kwargs = {'ntr_pad': 60, 'ntr_tap': 0, 'lagc': 3000,
+                    'butter_kwargs': {'N': 3, 'Wn': 0.01, 'btype': 'highpass'}}
+    if k_filter:
+        spatial_fcn = lambda dat: kfilt(dat, **k_kwargs)  # noqa
+    else:
+        spatial_fcn = lambda dat: car(dat, **k_kwargs)  # noqa
+    return butter_kwargs, k_kwargs, spatial_fcn
+
+
+def destripe(x, fs, neuropixel_version=1, butter_kwargs=None, k_kwargs=None, channel_labels=None, k_filter=True):
     """Super Car (super slow also...) - far from being set in stone but a good workflow example
     :param x: demultiplexed array (nc, ns)
     :param fs: sampling frequency
@@ -241,18 +256,10 @@ def destripe(x, fs, neuropixel_version=1, butter_kwargs=None, k_kwargs=None, cha
     :param butter_kwargs: (optional, None) butterworth params, see the code for the defaults dict
     :param k_kwargs: (optional, None) K-filter params, see the code for the defaults dict
         can also be set to 'car', in which case the median accross channels will be subtracted
+    :param k_filter (True): applies k-filter by default, otherwise, apply CAR.
     :return: x, filtered array
     """
-    if butter_kwargs is None:
-        butter_kwargs = {'N': 3, 'Wn': 300 / fs * 2, 'btype': 'highpass'}
-    if k_kwargs is None:
-        k_kwargs = {'ntr_pad': 60, 'ntr_tap': 0, 'lagc': 3000,
-                    'butter_kwargs': {'N': 3, 'Wn': 0.01, 'btype': 'highpass'}}
-        spatial_fcn = lambda dat: kfilt(dat, **k_kwargs)  # noqa
-    elif isinstance(k_kwargs, dict):
-        spatial_fcn = lambda dat: kfilt(dat, **k_kwargs)  # noqa
-    else:
-        spatial_fcn = lambda dat: car(dat, lagc=int(0.1 * fs))  # noqa
+    butter_kwargs, k_kwargs, spatial_fcn = _get_destripe_parameters(fs, butter_kwargs, k_kwargs, k_filter)
     h = neuropixel.trace_header(version=neuropixel_version)
     if channel_labels is True:
         channel_labels, _ = detect_bad_channels(x, fs)
@@ -262,8 +269,7 @@ def destripe(x, fs, neuropixel_version=1, butter_kwargs=None, k_kwargs=None, cha
     # channel interpolation
     # apply ADC shift
     if neuropixel_version is not None:
-        sample_shift = h['sample_shift'] if (30000 / fs) < 10 else h['sample_shift'] * fs / 30000
-        x = fshift(x, sample_shift, axis=1)
+        x = fshift(x, h['sample_shift'], axis=1)
     # apply spatial filter only on channels that are inside of the brain
     if channel_labels is not None:
         x = interpolate_bad_channels(x, channel_labels, h)
@@ -275,7 +281,8 @@ def destripe(x, fs, neuropixel_version=1, butter_kwargs=None, k_kwargs=None, cha
 
 
 def decompress_destripe_cbin(sr_file, output_file=None, h=None, wrot=None, append=False, nc_out=None, butter_kwargs=None,
-                             dtype=np.int16, ns2add=0, nbatch=None, nprocesses=None, compute_rms=True, reject_channels=True):
+                             dtype=np.int16, ns2add=0, nbatch=None, nprocesses=None, compute_rms=True, reject_channels=True,
+                             k_kwargs=None, k_filter=True):
     """
     From a spikeglx Reader object, decompresses and apply ADC.
     Saves output as a flat binary file in int16
@@ -295,6 +302,8 @@ def decompress_destripe_cbin(sr_file, output_file=None, h=None, wrot=None, appen
      interp 3:outside of brain and discard
     :param reject_channels: (True) detects noisy or bad channels and interpolate them. Channels outside of the brain are left
      untouched
+    :param k_kwargs: (None) arguments for the kfilter function
+    :param k_filter: (True) Performs a k-filter - if False will do median common average referencing
     :return:
     """
     import pyfftw
@@ -306,9 +315,7 @@ def decompress_destripe_cbin(sr_file, output_file=None, h=None, wrot=None, appen
     if reject_channels:  # get bad channels if option is on
         channel_labels = detect_bad_channels_cbin(sr)
     assert isinstance(sr_file, str) or isinstance(sr_file, Path)
-    butter_kwargs = butter_kwargs or {'N': 3, 'Wn': 300 / sr.fs * 2, 'btype': 'highpass'}
-    k_kwargs = {'ntr_pad': 60, 'ntr_tap': 0, 'lagc': 3000,
-                'butter_kwargs': {'N': 3, 'Wn': 0.01, 'btype': 'highpass'}}
+    butter_kwargs, k_kwargs, spatial_fcn = _get_destripe_parameters(sr.fs, butter_kwargs, k_kwargs, k_filter)
     h = neuropixel.trace_header(version=1) if h is None else h
     ncv = h['sample_shift'].size  # number of channels
     output_file = sr.file_bin.with_suffix('.bin') if output_file is None else output_file
@@ -412,9 +419,9 @@ def my_function(i_chunk, n_chunk):
             if reject_channels:
                 chunk = interpolate_bad_channels(chunk, channel_labels, h=h)
                 inside_brain = np.where(channel_labels != 3)[0]
-                chunk[inside_brain, :] = kfilt(chunk[inside_brain, :], **k_kwargs)  # apply the k-filter
+                chunk[inside_brain, :] = spatial_fcn(chunk[inside_brain, :])  # apply the k-filter / CAR
             else:
-                chunk = kfilt(chunk, **k_kwargs)  # apply the k-filter
+                chunk = spatial_fcn(chunk)  # apply the k-filter / CAR
             # add back sync trace and save
             chunk = np.r_[chunk, _sr[first_s:last_s, ncv:].T].T
             intnorm = 1 / _sr.channel_conversion_sample2v['ap'] if dtype == np.int16 else 1.
@@ -476,7 +483,7 @@ def normalize(z):
     return rcor
 
 
-def detect_bad_channels(raw, fs, similarity_threshold=(-0.5, 1), psd_hf_threshold=0.02):
+def detect_bad_channels(raw, fs, similarity_threshold=(-0.5, 1), psd_hf_threshold=None):
     """
     Bad channels detection for Neuropixel probes
     Labels channels
@@ -540,7 +547,9 @@ def nxcor(x, ref):
     raw = raw - np.mean(raw, axis=-1)[:, np.newaxis]  # removes DC offset
     xcor = channels_similarity(raw)
     fscale, psd = scipy.signal.welch(raw * 1e6, fs=fs)  # units; uV ** 2 / Hz
-
+    if psd_hf_threshold is None:
+        # the LFP band data is obviously much stronger so auto-adjust the default threshold
+        psd_hf_threshold = 1.4 if fs < 5000 else 0.02
     sos_hp = scipy.signal.butter(**{'N': 3, 'Wn': 300 / fs * 2, 'btype': 'highpass'}, output='sos')
     hf = scipy.signal.sosfiltfilt(sos_hp, raw)
     xcorf = channels_similarity(hf)

ibllib/pipes/ephys_preprocessing.py

@@ -20,6 +20,7 @@
 from ibllib.pipes import tasks
 from ibllib.pipes.training_preprocessing import TrainingRegisterRaw as EphysRegisterRaw
 from ibllib.pipes.misc import create_alyx_probe_insertions
+from ibllib.qc.alignment_qc import get_aligned_channels
 from ibllib.qc.task_extractors import TaskQCExtractor
 from ibllib.qc.task_metrics import TaskQC
 from ibllib.qc.camera import run_all_qc as run_camera_qc
@@ -406,6 +407,18 @@ def _run(self, probes=None):
                 tar_dir.mkdir(parents=True, exist_ok=True)
                 out = spikes.ks2_to_tar(ks2_dir, tar_dir, force=self.FORCE_RERUN)
                 out_files.extend(out)
+
+                if self.one:
+                    eid = self.one.path2eid(self.session_path, query_type='remote')
+                    ins = self.one.alyx.rest('insertions', 'list', session=eid, name=label)
+                    if len(ins) != 0:
+                        resolved = ins[0].get('json', {'temp': 0}).get('extended_qc', {'temp': 0}). \
+                            get('alignment_resolved', False)
+                        if resolved:
+                            chns = np.load(probe_out_path.joinpath('channels.localCoordinates.npy'))
+                            out = get_aligned_channels(ins[0], chns, one=self.one, save_dir=probe_out_path)
+                            out_files.extend(out)
+
             except BaseException:
                 _logger.error(traceback.format_exc())
                 self.status = -1