Merge branch 'develop' into docs

Author: mayofaulkner

brainbox/behavior/wheel.py

@@ -1,10 +1,6 @@
 """
 Set of functions to handle wheel data.
 """
-import logging
-import warnings
-import traceback
-
 import numpy as np
 from numpy import pi
 from iblutil.numerical import between_sorted
@@ -68,42 +64,6 @@ def interpolate_position(re_ts, re_pos, freq=1000, kind='linear', fill_gaps=None
     return yinterp, t
 
 
-def velocity(re_ts, re_pos):
-    """
-    (DEPRECATED) Compute wheel velocity from non-uniformly sampled wheel data. Returns the velocity
-    at the same samples locations as the position through interpolation.
-
-    Parameters
-    ----------
-    re_ts : array_like
-        Array of timestamps
-    re_pos: array_like
-        Array of unwrapped wheel positions
-
-    Returns
-    -------
-    np.ndarray
-        numpy array of velocities
-    """
-    for line in traceback.format_stack():
-        print(line.strip())
-
-    msg = 'brainbox.behavior.wheel.velocity will soon be removed. Use velocity_filtered instead.'
-    warnings.warn(msg, FutureWarning)
-    logging.getLogger(__name__).warning(msg)
-
-    dp = np.diff(re_pos)
-    dt = np.diff(re_ts)
-    # Compute raw velocity
-    vel = dp / dt
-    # Compute velocity time scale
-    tv = re_ts[:-1] + dt / 2
-    # interpolate over original time scale
-    if tv.size > 1:
-        ifcn = interpolate.interp1d(tv, vel, fill_value="extrapolate")
-        return ifcn(re_ts)
-
-
 def velocity_filtered(pos, fs, corner_frequency=20, order=8):
     """
     Compute wheel velocity from uniformly sampled wheel data.
@@ -130,83 +90,6 @@ def velocity_filtered(pos, fs, corner_frequency=20, order=8):
     return vel, acc
 
 
-def velocity_smoothed(pos, freq, smooth_size=0.03):
-    """
-    (DEPRECATED) Compute wheel velocity from uniformly sampled wheel data.
-
-    Parameters
-    ----------
-    pos : array_like
-        Array of wheel positions
-    smooth_size : float
-        Size of Gaussian smoothing window in seconds
-    freq : float
-        Sampling frequency of the data
-
-    Returns
-    -------
-    vel : np.ndarray
-        Array of velocity values
-    acc : np.ndarray
-        Array of acceleration values
-    """
-    for line in traceback.format_stack():
-        print(line.strip())
-
-    msg = 'brainbox.behavior.wheel.velocity_smoothed will be removed. Use velocity_filtered instead.'
-    warnings.warn(msg, FutureWarning)
-    logging.getLogger(__name__).warning(msg)
-
-    # Define our smoothing window with an area of 1 so the units won't be changed
-    std_samps = np.round(smooth_size * freq)  # Standard deviation relative to sampling frequency
-    N = std_samps * 6  # Number of points in the Gaussian covering +/-3 standard deviations
-    gauss_std = (N - 1) / 6
-    win = scipy.signal.windows.gaussian(N, gauss_std)
-    win = win / win.sum()  # Normalize amplitude
-
-    # Convolve and multiply by sampling frequency to restore original units
-    vel = np.insert(scipy.signal.convolve(np.diff(pos), win, mode='same'), 0, 0) * freq
-    acc = np.insert(scipy.signal.convolve(np.diff(vel), win, mode='same'), 0, 0) * freq
-
-    return vel, acc
-
-
-def last_movement_onset(t, vel, event_time):
-    """
-    (DEPRECATED) Find the time at which movement started, given an event timestamp that occurred during the
-    movement.
-
-    Movement start is defined as the first sample after the velocity has been zero for at least 50ms.
-    Wheel inputs should be evenly sampled.
-
-    :param t: numpy array of wheel timestamps in seconds
-    :param vel: numpy array of wheel velocities
-    :param event_time: timestamp anywhere during movement of interest, e.g. peak velocity
-    :return: timestamp of movement onset
-    """
-    for line in traceback.format_stack():
-        print(line.strip())
-
-    msg = 'brainbox.behavior.wheel.last_movement_onset has been deprecated. Use get_movement_onset instead.'
-    warnings.warn(msg, FutureWarning)
-    logging.getLogger(__name__).warning(msg)
-
-    # Look back from timestamp
-    threshold = 50e-3
-    mask = t < event_time
-    times = t[mask]
-    vel = vel[mask]
-    t = None  # Initialize
-    for i, t in enumerate(times[::-1]):
-        i = times.size - i
-        idx = np.min(np.where((t - times) < threshold))
-        if np.max(np.abs(vel[idx:i])) < 0.5:
-            break
-
-    # Return timestamp
-    return t
-
-
 def get_movement_onset(intervals, event_times):
     """
     Find the time at which movement started, given an event timestamp that occurred during the

brainbox/io/one.py

@@ -20,6 +20,7 @@
 import spikeglx
 
 import ibldsp.voltage
+from ibldsp.waveform_extraction import WaveformsLoader
 from iblutil.util import Bunch
 from iblatlas.atlas import AllenAtlas, BrainRegions
 from iblatlas import atlas
@@ -975,6 +976,21 @@ def raw_electrophysiology(self, stream=True, band='ap', **kwargs):
             if cbin_file is not None:
                 return spikeglx.Reader(cbin_file)
 
+    def download_raw_waveforms(self, **kwargs):
+        """
+        Downloads raw waveforms extracted from sorting to local disk.
+        """
+        _logger.debug(f"loading waveforms from {self.collection}")
+        return self.one.load_object(
+            self.eid, "waveforms",
+            attribute=["traces", "templates", "table", "channels"],
+            collection=self._get_spike_sorting_collection("pykilosort"), download_only=True, **kwargs
+        )
+
+    def raw_waveforms(self, **kwargs):
+        wf_paths = self.download_raw_waveforms(**kwargs)
+        return WaveformsLoader(wf_paths[0].parent, wfs_dtype=np.float16)
+
     def load_channels(self, **kwargs):
         """
         Loads channels
@@ -1318,6 +1334,7 @@ class SessionLoader:
     one: One = None
     session_path: Path = ''
     eid: str = ''
+    revision: str = ''
     data_info: pd.DataFrame = field(default_factory=pd.DataFrame, repr=False)
     trials: pd.DataFrame = field(default_factory=pd.DataFrame, repr=False)
     wheel: pd.DataFrame = field(default_factory=pd.DataFrame, repr=False)
@@ -1445,7 +1462,7 @@ def load_trials(self, collection=None):
         # itiDuration frequently has a mismatched dimension, and we don't need it, exclude using regex
         self.one.wildcards = False
         self.trials = self.one.load_object(
-            self.eid, 'trials', collection=collection, attribute=r'(?!itiDuration).*').to_df()
+            self.eid, 'trials', collection=collection, attribute=r'(?!itiDuration).*', revision=self.revision or None).to_df()
         self.one.wildcards = True
         self.data_info.loc[self.data_info['name'] == 'trials', 'is_loaded'] = True
 
@@ -1468,7 +1485,7 @@ def load_wheel(self, fs=1000, corner_frequency=20, order=8, collection=None):
         """
         if not collection:
             collection = self._find_behaviour_collection('wheel')
-        wheel_raw = self.one.load_object(self.eid, 'wheel', collection=collection)
+        wheel_raw = self.one.load_object(self.eid, 'wheel', collection=collection, revision=self.revision or None)
         if wheel_raw['position'].shape[0] != wheel_raw['timestamps'].shape[0]:
             raise ValueError("Length mismatch between 'wheel.position' and 'wheel.timestamps")
         # resample the wheel position and compute velocity, acceleration
@@ -1498,7 +1515,7 @@ def load_pose(self, likelihood_thr=0.9, views=['left', 'right', 'body']):
         # empty the dictionary so that if one loads only one view, after having loaded several, the others don't linger
         self.pose = {}
         for view in views:
-            pose_raw = self.one.load_object(self.eid, f'{view}Camera', attribute=['dlc', 'times'])
+            pose_raw = self.one.load_object(self.eid, f'{view}Camera', attribute=['dlc', 'times'], revision=self.revision or None)
             # Double check if video timestamps are correct length or can be fixed
             times_fixed, dlc = self._check_video_timestamps(view, pose_raw['times'], pose_raw['dlc'])
             self.pose[f'{view}Camera'] = likelihood_threshold(dlc, likelihood_thr)
@@ -1525,7 +1542,8 @@ def load_motion_energy(self, views=['left', 'right', 'body']):
         # empty the dictionary so that if one loads only one view, after having loaded several, the others don't linger
         self.motion_energy = {}
         for view in views:
-            me_raw = self.one.load_object(self.eid, f'{view}Camera', attribute=['ROIMotionEnergy', 'times'])
+            me_raw = self.one.load_object(
+                self.eid, f'{view}Camera', attribute=['ROIMotionEnergy', 'times'], revision=self.revision or None)
             # Double check if video timestamps are correct length or can be fixed
             times_fixed, motion_energy = self._check_video_timestamps(
                 view, me_raw['times'], me_raw['ROIMotionEnergy'])
@@ -1550,7 +1568,7 @@ def load_pupil(self, snr_thresh=5.):
             will be considered unusable and will be discarded.
         """
         # Try to load from features
-        feat_raw = self.one.load_object(self.eid, 'leftCamera', attribute=['times', 'features'])
+        feat_raw = self.one.load_object(self.eid, 'leftCamera', attribute=['times', 'features'], revision=self.revision or None)
         if 'features' in feat_raw.keys():
             times_fixed, feats = self._check_video_timestamps('left', feat_raw['times'], feat_raw['features'])
             self.pupil = feats.copy()

brainbox/tests/test_behavior.py

@@ -124,12 +124,6 @@ def test_get_movement_onset(self):
         with self.assertRaises(ValueError):
             wheel.get_movement_onset(intervals, np.random.permutation(self.trials['feedback_times']))
 
-    def test_velocity_deprecation(self):
-        """Ensure brainbox.behavior.wheel.velocity is removed."""
-        from datetime import datetime
-        self.assertTrue(datetime.today() < datetime(2024, 8, 1),
-                        'remove brainbox.behavior.wheel.velocity, velocity_smoothed and last_movement_onset')
-
 
 class TestTraining(unittest.TestCase):
     def setUp(self):

ibllib/ephys/ephysqc.py

@@ -205,57 +205,65 @@ def run(self, update: bool = False, overwrite: bool = True, stream: bool = None,
         return qc_files
 
 
-def rmsmap(sglx):
+def rmsmap(sglx, spectra=True, nmod=1):
     """
     Computes RMS map in time domain and spectra for each channel of Neuropixel probe
 
     :param sglx: Open spikeglx reader
+    :param spectra: Whether to compute the spectra
+    :param nmod: take every nmod windows, in cases where we don't want to compute over the whole signal
     :return: a dictionary with amplitudes in channeltime space, channelfrequency space, time
      and frequency scales
     """
     rms_win_length_samples = 2 ** np.ceil(np.log2(sglx.fs * RMS_WIN_LENGTH_SECS))
     # the window generator will generates window indices
     wingen = utils.WindowGenerator(ns=sglx.ns, nswin=rms_win_length_samples, overlap=0)
+    nwin = np.ceil(wingen.nwin / nmod).astype(int)
     # pre-allocate output dictionary of numpy arrays
-    win = {'TRMS': np.zeros((wingen.nwin, sglx.nc)),
-           'nsamples': np.zeros((wingen.nwin,)),
+    win = {'TRMS': np.zeros((nwin, sglx.nc)),
+           'nsamples': np.zeros((nwin,)),
            'fscale': fourier.fscale(WELCH_WIN_LENGTH_SAMPLES, 1 / sglx.fs, one_sided=True),
-           'tscale': wingen.tscale(fs=sglx.fs)}
+           'tscale': wingen.tscale(fs=sglx.fs)[::nmod]}
     win['spectral_density'] = np.zeros((len(win['fscale']), sglx.nc))
     # loop through the whole session
     with tqdm(total=wingen.nwin) as pbar:
-        for first, last in wingen.firstlast:
+        for iwindow, (first, last) in enumerate(wingen.firstlast):
+            if np.mod(iwindow, nmod) != 0:
+                continue
+
             D = sglx.read_samples(first_sample=first, last_sample=last)[0].transpose()
             # remove low frequency noise below 1 Hz
             D = fourier.hp(D, 1 / sglx.fs, [0, 1])
-            iw = wingen.iw
+            iw = np.floor(wingen.iw / nmod).astype(int)
             win['TRMS'][iw, :] = utils.rms(D)
             win['nsamples'][iw] = D.shape[1]
-            # the last window may be smaller than what is needed for welch
-            if last - first < WELCH_WIN_LENGTH_SAMPLES:
-                continue
-            # compute a smoothed spectrum using welch method
-            _, w = signal.welch(
-                D, fs=sglx.fs, window='hann', nperseg=WELCH_WIN_LENGTH_SAMPLES,
-                detrend='constant', return_onesided=True, scaling='density', axis=-1
-            )
-            win['spectral_density'] += w.T
+            if spectra:
+                # the last window may be smaller than what is needed for welch
+                if last - first < WELCH_WIN_LENGTH_SAMPLES:
+                    continue
+                # compute a smoothed spectrum using welch method
+                _, w = signal.welch(
+                    D, fs=sglx.fs, window='hann', nperseg=WELCH_WIN_LENGTH_SAMPLES,
+                    detrend='constant', return_onesided=True, scaling='density', axis=-1
+                )
+                win['spectral_density'] += w.T
             # print at least every 20 windows
             if (iw % min(20, max(int(np.floor(wingen.nwin / 75)), 1))) == 0:
                 pbar.update(iw)
     sglx.close()
     return win
 
 
-def extract_rmsmap(sglx, out_folder=None, overwrite=False):
+def extract_rmsmap(sglx, out_folder=None, overwrite=False, spectra=True, nmod=1):
     """
     Wrapper for rmsmap that outputs _ibl_ephysRmsMap and _ibl_ephysSpectra ALF files
 
     :param sglx: Open spikeglx Reader with data for which to compute rmsmap
     :param out_folder: folder in which to store output ALF files. Default uses the folder in which
      the `fbin` file lives.
     :param overwrite: do not re-extract if all ALF files already exist
-    :param label: string or list of strings that will be appended to the filename before extension
+    :param spectra: Whether to compute the spectral density across the signal
+    :param nmod: take every nmod windows, in cases where we don't want to compute over the whole signal
     :return: None
     """
     if out_folder is None:
@@ -271,18 +279,19 @@ def extract_rmsmap(sglx, out_folder=None, overwrite=False):
         _logger.warning(f'RMS map already exists for .{sglx.type} data in {out_folder}, skipping. Use overwrite option.')
         return files_time + files_freq
     # crunch numbers
-    rms = rmsmap(sglx)
+    rms = rmsmap(sglx, spectra=spectra, nmod=nmod)
     # output ALF files, single precision with the optional label as suffix before extension
     if not out_folder.exists():
         out_folder.mkdir()
     tdict = {'rms': rms['TRMS'].astype(np.single), 'timestamps': rms['tscale'].astype(np.single)}
-    fdict = {'power': rms['spectral_density'].astype(np.single),
-             'freqs': rms['fscale'].astype(np.single)}
     out_time = alfio.save_object_npy(
         out_folder, object=alf_object_time, dico=tdict, namespace='iblqc')
-    out_freq = alfio.save_object_npy(
-        out_folder, object=alf_object_freq, dico=fdict, namespace='iblqc')
-    return out_time + out_freq
+    if spectra:
+        fdict = {'power': rms['spectral_density'].astype(np.single),
+                 'freqs': rms['fscale'].astype(np.single)}
+        out_freq = alfio.save_object_npy(
+            out_folder, object=alf_object_freq, dico=fdict, namespace='iblqc')
+    return out_time + out_freq if spectra else out_time
 
 
 def raw_qc_session(session_path, overwrite=False):