Merge branch 'release/2.4.0'

Author: juhuntenburg

brainbox/behavior/dlc.py

@@ -4,7 +4,9 @@
 import numpy as np
 import scipy.interpolate as interpolate
 import logging
+import warnings
 from one.api import ONE
+from ibllib.dsp.smooth import smooth_interpolate_savgol
 
 logger = logging.getLogger('ibllib')
 
@@ -141,3 +143,74 @@ def get_dlc_everything(dlc_cam, camera):
     dlc_cam['aligned'] = aligned
 
     return dlc_cam
+
+
+def get_pupil_diameter(dlc):
+    """
+    Estimates pupil diameter by taking median of different computations.
+
+    The two most straightforward estimates: d1 = top - bottom, d2 = left - right
+    In addition, assume the pupil is a circle and estimate diameter from other pairs of points
+
+    :param dlc: dlc pqt table with pupil estimates, should be likelihood thresholded (e.g. at 0.9)
+    :return: np.array, pupil diameter estimate for each time point, shape (n_frames,)
+    """
+    diameters = []
+    # Get the x,y coordinates of the four pupil points
+    top, bottom, left, right = [np.vstack((dlc[f'pupil_{point}_r_x'], dlc[f'pupil_{point}_r_y']))
+                                for point in ['top', 'bottom', 'left', 'right']]
+    # First compute direct diameters
+    diameters.append(np.linalg.norm(top - bottom, axis=0))
+    diameters.append(np.linalg.norm(left - right, axis=0))
+
+    # For non-crossing edges, estimate diameter via circle assumption
+    for pair in [(top, left), (top, right), (bottom, left), (bottom, right)]:
+        diameters.append(np.linalg.norm(pair[0] - pair[1], axis=0) * 2 ** 0.5)
+
+    # Ignore all nan runtime warning
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=RuntimeWarning)
+        return np.nanmedian(diameters, axis=0)
+
+
+def get_smooth_pupil_diameter(diameter_raw, camera, std_thresh=5, nan_thresh=1):
+    """
+    :param diameter_raw: np.array, raw pupil diameters, calculated from (thresholded) dlc traces
+    :param camera: str ('left', 'right'), which camera to run the smoothing for
+    :param std_thresh: threshold (in standard deviations) beyond which a point is labeled as an outlier
+    :param nan_thresh: threshold (in seconds) above which we will not interpolate nans, but keep them
+                       (for long stretches interpolation may not be appropriate)
+    :return:
+    """
+    # set framerate of camera
+    if camera == 'left':
+        fr = 60  # set by hardware
+        window = 31  # works well empirically
+    elif camera == 'right':
+        fr = 150  # set by hardware
+        window = 75  # works well empirically
+    else:
+        raise NotImplementedError("camera has to be 'left' or 'right")
+
+    # run savitzy-golay filter on non-nan time points to denoise
+    diameter_smoothed = smooth_interpolate_savgol(diameter_raw, window=window, order=3, interp_kind='linear')
+
+    # find outliers and set them to nan
+    difference = diameter_raw - diameter_smoothed
+    outlier_thresh = std_thresh * np.nanstd(difference)
+    without_outliers = np.copy(diameter_raw)
+    without_outliers[(difference < -outlier_thresh) | (difference > outlier_thresh)] = np.nan
+    # run savitzy-golay filter again on (possibly reduced) non-nan timepoints to denoise
+    diameter_smoothed = smooth_interpolate_savgol(without_outliers, window=window, order=3, interp_kind='linear')
+
+    # don't interpolate long strings of nans
+    t = np.diff(np.isnan(without_outliers).astype(int))
+    begs = np.where(t == 1)[0]
+    ends = np.where(t == -1)[0]
+    if begs.shape[0] > ends.shape[0]:
+        begs = begs[:ends.shape[0]]
+    for b, e in zip(begs, ends):
+        if (e - b) > (fr * nan_thresh):
+            diameter_smoothed[(b + 1):(e + 1)] = np.nan  # offset by 1 due to earlier diff
+
+    return diameter_smoothed

brainbox/ephys_plots.py

@@ -404,6 +404,8 @@ def plot_brain_regions(channel_ids, channel_depths=None, brain_regions=None, dis
     if display:
         if ax is None:
             fig, ax = plt.subplots()
+        else:
+            fig = ax.get_figure()
 
         for reg, col in zip(regions, region_colours):
             height = np.abs(reg[1] - reg[0])
@@ -423,7 +425,7 @@ def plot_brain_regions(channel_ids, channel_depths=None, brain_regions=None, dis
 
 
 def plot_cdf(spike_amps, spike_depths, spike_times, n_amp_bins=10, d_bin=40, amp_range=None, d_range=None,
-             display=False, cmap='hot'):
+             display=False, cmap='hot', ax=None):
     """
     Plot cumulative amplitude of spikes across depth
     :param spike_amps:
@@ -466,7 +468,7 @@ def histc(x, bins):
                     ylabel='Distance from probe tip (um)', clabel='Firing Rate (Hz)')
 
     if display:
-        fig, ax = plot_image(data.convert2dict(), fig_kwargs={'figsize': [3, 7]})
+        fig, ax = plot_image(data.convert2dict(), fig_kwargs={'figsize': [3, 7]}, ax=ax)
         return data.convert2dict(), fig, ax
 
     return data

brainbox/io/one.py

@@ -331,6 +331,9 @@ def _load_channel_locations_traj(eid, probe=None, one=None, revision=None, align
 
             channels[probe] = _channels_traj2bunch(chans, brain_atlas)
 
+        channels[probe]['axial_um'] = chn_coords[:, 1]
+        channels[probe]['lateral_um'] = chn_coords[:, 0]
+
     else:
         _logger.warning(f'Histology tracing for {probe} does not exist. '
                         f'No channels for {probe}')

examples/one/histology/docs_find_nearby_trajectories.py

@@ -15,7 +15,6 @@
 import ibllib.atlas as atlas
 from atlaselectrophysiology import rendering
 
-mlab.init_notebook()
 # Instantiate brain atlas and one
 brain_atlas = atlas.AllenAtlas(25)
 one = ONE(base_url='https://openalyx.internationalbrainlab.org')

examples/one/histology/docs_visualize_session_coronal_tilted.py

@@ -29,4 +29,4 @@
 
 cax = ba.plot_tilted_slice(xyz=picks, axis=1, volume='image')
 cax.plot(picks[:, 0] * 1e6, picks[:, 2] * 1e6)
-cax.plot(channels[probe_label].x * 1e6, channels[probe_label].z * 1e6, 'g*')
+cax.plot(channels[probe_label]['x'] * 1e6, channels[probe_label]['z'] * 1e6, 'g*')

ibllib/atlas/cosmos.npy

@@ -1,5 +1,6 @@
 import numpy as np
 import matplotlib.pyplot as plt
+from scipy.interpolate import interp1d
 
 import ibllib.dsp.fourier as ft
 
@@ -77,6 +78,141 @@ def rolling_window(x, window_len=11, window='blackman'):
     return y[round((window_len / 2 - 1)):round(-(window_len / 2))]
 
 
+def non_uniform_savgol(x, y, window, polynom):
+    """Applies a Savitzky-Golay filter to y with non-uniform spacing as defined in x.
+    This is based on
+    https://dsp.stackexchange.com/questions/1676/savitzky-golay-smoothing-filter-for-not-equally-spaced-data
+    The borders are interpolated like scipy.signal.savgol_filter would do
+    https://dsp.stackexchange.com/a/64313
+    Parameters
+    ----------
+    x : array_like
+        List of floats representing the x values of the data
+    y : array_like
+        List of floats representing the y values. Must have same length as x
+    window : int (odd)
+        Window length of datapoints. Must be odd and smaller than x
+    polynom : int
+        The order of polynom used. Must be smaller than the window size
+    Returns
+    -------
+    np.array
+        The smoothed y values
+    """
+
+    if len(x) != len(y):
+        raise ValueError('"x" and "y" must be of the same size')
+    if len(x) < window:
+        raise ValueError('The data size must be larger than the window size')
+    if type(window) is not int:
+        raise TypeError('"window" must be an integer')
+    if window % 2 == 0:
+        raise ValueError('The "window" must be an odd integer')
+    if type(polynom) is not int:
+        raise TypeError('"polynom" must be an integer')
+    if polynom >= window:
+        raise ValueError('"polynom" must be less than "window"')
+
+    half_window = window // 2
+    polynom += 1
+
+    # Initialize variables
+    A = np.empty((window, polynom))  # Matrix
+    tA = np.empty((polynom, window))  # Transposed matrix
+    t = np.empty(window)  # Local x variables
+    y_smoothed = np.full(len(y), np.nan)
+
+    # Start smoothing
+    for i in range(half_window, len(x) - half_window, 1):
+        # Center a window of x values on x[i]
+        for j in range(0, window, 1):
+            t[j] = x[i + j - half_window] - x[i]
+
+        # Create the initial matrix A and its transposed form tA
+        for j in range(0, window, 1):
+            r = 1.0
+            for k in range(0, polynom, 1):
+                A[j, k] = r
+                tA[k, j] = r
+                r *= t[j]
+
+        # Multiply the two matrices
+        tAA = np.matmul(tA, A)
+        # Invert the product of the matrices
+        tAA = np.linalg.inv(tAA)
+        # Calculate the pseudoinverse of the design matrix
+        coeffs = np.matmul(tAA, tA)
+        # Calculate c0 which is also the y value for y[i]
+        y_smoothed[i] = 0
+        for j in range(0, window, 1):
+            y_smoothed[i] += coeffs[0, j] * y[i + j - half_window]
+
+        # If at the end or beginning, store all coefficients for the polynom
+        if i == half_window:
+            first_coeffs = np.zeros(polynom)
+            for j in range(0, window, 1):
+                for k in range(polynom):
+                    first_coeffs[k] += coeffs[k, j] * y[j]
+        elif i == len(x) - half_window - 1:
+            last_coeffs = np.zeros(polynom)
+            for j in range(0, window, 1):
+                for k in range(polynom):
+                    last_coeffs[k] += coeffs[k, j] * y[len(y) - window + j]
+
+    # Interpolate the result at the left border
+    for i in range(0, half_window, 1):
+        y_smoothed[i] = 0
+        x_i = 1
+        for j in range(0, polynom, 1):
+            y_smoothed[i] += first_coeffs[j] * x_i
+            x_i *= x[i] - x[half_window]
+
+    # Interpolate the result at the right border
+    for i in range(len(x) - half_window, len(x), 1):
+        y_smoothed[i] = 0
+        x_i = 1
+        for j in range(0, polynom, 1):
+            y_smoothed[i] += last_coeffs[j] * x_i
+            x_i *= x[i] - x[-half_window - 1]
+
+    return y_smoothed
+
+
+def smooth_interpolate_savgol(signal, window=31, order=3, interp_kind='cubic'):
+    """Run savitzy-golay filter on signal, interpolate through nan points.
+
+    Parameters
+    ----------
+    signal : np.ndarray
+        original noisy signal of shape (t,), may contain nans
+    window : int
+        window of polynomial fit for savitzy-golay filter
+    order : int
+        order of polynomial for savitzy-golay filter
+    interp_kind : str
+        type of interpolation for nans, e.g. 'linear', 'quadratic', 'cubic'
+    Returns
+    -------
+    np.array
+        smoothed, interpolated signal for each time point, shape (t,)
+    """
+
+    signal_noisy_w_nans = np.copy(signal)
+    timestamps = np.arange(signal_noisy_w_nans.shape[0])
+    good_idxs = np.where(~np.isnan(signal_noisy_w_nans))[0]
+    # perform savitzky-golay filtering on non-nan points
+    signal_smooth_nonans = non_uniform_savgol(
+        timestamps[good_idxs], signal_noisy_w_nans[good_idxs], window=window, polynom=order)
+    signal_smooth_w_nans = np.copy(signal_noisy_w_nans)
+    signal_smooth_w_nans[good_idxs] = signal_smooth_nonans
+    # interpolate nan points
+    interpolater = interp1d(
+        timestamps[good_idxs], signal_smooth_nonans, kind=interp_kind, fill_value='extrapolate')
+    signal = interpolater(timestamps)
+
+    return signal
+
+
 def smooth_demo():
 
     t = np.linspace(-4, 4, 100)

ibllib/dsp/smooth.py

@@ -165,10 +165,16 @@ def create_sessions(self, root_data_folder, glob_pattern='**/create_me.flag', dr
             if dry:
                 print(flag_file)
                 continue
-            _logger.info('creating session for ' + str(flag_file.parent))
-            # providing a false flag stops the registration after session creation
-            self.create_session(flag_file.parent)
-            flag_file.unlink()
+            try:
+                _logger.info('creating session for ' + str(flag_file.parent))
+                # providing a false flag stops the registration after session creation
+                self.create_session(flag_file.parent)
+                flag_file.unlink()
+            except BaseException as e:
+                _logger.error(f'Error creating session for {flag_file.parent}\n{e}')
+                _logger.warning(f'Skipping {flag_file.parent}')
+                continue
+
         return [ff.parent for ff in flag_files]
 
     def create_session(self, session_path):