Merge pull request #495 from int-brain-lab/atlas

Atlas

Author: oliche

ibllib/atlas/atlas.py

@@ -300,16 +300,34 @@ def _lookup(self, xyz):
         """
         return self._lookup_inds(self.bc.xyz2i(xyz))
 
-    def get_labels(self, xyz, mapping='Allen'):
+    def get_labels(self, xyz, mapping='Allen', radius_um=None):
         """
         Performs a 3D lookup from real world coordinates to the volume labels
         and return the regions ids according to the mapping
         :param xyz: [n, 3] array of coordinates
         :param mapping: brain region mapping (defaults to original Allen mapping)
+        :param radius_um: if not null, returns a regions ids array and an array of proportion
+         of regions in a sphere of size radius around the coordinates.
         :return: n array of region ids
         """
-        regions_indices = self._get_mapping(mapping=mapping)[self.label.flat[self._lookup(xyz)]]
-        return self.regions.id[regions_indices]
+        if radius_um:
+            nrx = int(np.ceil(radius_um / abs(self.bc.dx) / 1e6))
+            nry = int(np.ceil(radius_um / abs(self.bc.dy) / 1e6))
+            nrz = int(np.ceil(radius_um / abs(self.bc.dz) / 1e6))
+            nr = [nrx, nry, nrz]
+            iii = self.bc.xyz2i(xyz)
+            # computing the cube radius and indices is more complicated as volume indices are not
+            # necessariy in ml, ap, dv order so the indices order is dynamic
+            rcube = np.meshgrid(*tuple((np.arange(
+                -nr[i], nr[i] + 1) * self.bc.dxyz[i]) ** 2 for i in self.xyz2dims))
+            rcube = np.sqrt(rcube[0] + rcube[1], rcube[2]) * 1e6
+            icube = tuple(slice(-nr[i] + iii[i], nr[i] + iii[i] + 1) for i in self.xyz2dims)
+            cube = self.regions.mappings[mapping][self.label[icube]]
+            ilabs, counts = np.unique(cube[rcube <= radius_um], return_counts=True)
+            return self.regions.id[ilabs], counts / np.sum(counts)
+        else:
+            regions_indices = self._get_mapping(mapping=mapping)[self.label.flat[self._lookup(xyz)]]
+            return self.regions.id[regions_indices]
 
     def _get_mapping(self, mapping='Allen'):
         """

ibllib/atlas/beryl.npy

@@ -4,14 +4,17 @@
 import logging
 
 import numpy as np
-from scipy import signal
+import scipy.signal
+import scipy.ndimage
 from scipy.io import wavfile
 
+
 from neurodsp.utils import WindowGenerator
 from neurodsp import fourier
 import ibllib.io.raw_data_loaders as ioraw
 from ibllib.io.extractors.training_trials import GoCueTimes
 
+
 logger_ = logging.getLogger('ibllib')
 
 NS_WIN = 2 ** 18  # 2 ** np.ceil(np.log2(1 * fs))
@@ -22,16 +25,19 @@
 READY_TONE_SPL = 85
 
 
-def _running_mean(x, N):
-    cumsum = np.cumsum(np.insert(x, 0, 0))
-    return (cumsum[N:] - cumsum[:-N]) / N
-
-
-def _detect_ready_tone(w, fs):
+def detect_ready_tone(w, fs, ftone=FTONE, threshold=0.8):
+    """
+    Detects a transient sinusoid signal in a time-serie
+    :param w: audio time seried
+    :param fs: sampling frequency (Hz)
+    :param ftone: frequency of the tone to detect
+    :param threshold: ratio of the Hilbert to the signal, between 0 and 1 (set to 0.8)
+    :return:
+    """
     # get envelope of DC free signal and envelope of BP signal around freq of interest
-    h = np.abs(signal.hilbert(w - np.median(w)))
-    fh = np.abs(signal.hilbert(fourier.bp(w, si=1 / fs, b=FTONE * np.array([0.9, 0.95, 1.15, 1.1]))))
-    dtect = _running_mean(fh / (h + 1e-3), int(fs * 0.1)) > 0.8
+    h = np.abs(scipy.signal.hilbert(w - np.median(w)))
+    fh = np.abs(scipy.signal.hilbert(fourier.bp(w, si=1 / fs, b=ftone * np.array([0.9, 0.95, 1.15, 1.1]))))
+    dtect = scipy.ndimage.uniform_filter1d(fh / (h + 1e-3), int(fs * 0.1)) > threshold
     return np.where(np.diff(dtect.astype(int)) == 1)[0]
     # tone = np.sin(2 * np.pi * FTONE * np.arange(0, fs * 0.1) / fs)
     # tone = tone / np.sum(tone ** 2)
@@ -56,7 +62,7 @@ def _get_conversion_factor(unit=UNIT, ready_tone_spl=READY_TONE_SPL):
     return fac
 
 
-def welchogram(fs, wav, nswin=NS_WIN, overlap=OVERLAP, nperseg=NS_WELCH):
+def welchogram(fs, wav, nswin=NS_WIN, overlap=OVERLAP, nperseg=NS_WELCH, detect_kwargs=None):
     """
     Computes a spectrogram on a very large audio file.
 
@@ -65,6 +71,7 @@ def welchogram(fs, wav, nswin=NS_WIN, overlap=OVERLAP, nperseg=NS_WELCH):
     :param nswin: n samples of the sliding window
     :param overlap: n samples of the overlap between windows
     :param nperseg: n samples for the computation of the spectrogram
+    :param detect_kwargs: specified paramaters for detection
     :return: tscale, fscale, downsampled_spectrogram
     """
     ns = wav.shape[0]
@@ -78,16 +85,17 @@ def welchogram(fs, wav, nswin=NS_WIN, overlap=OVERLAP, nperseg=NS_WELCH):
         # load the current window into memory
         w = np.float64(wav[first:last]) * _get_conversion_factor()
         # detection of ready tones
-        a = [d + first for d in _detect_ready_tone(w, fs)]
+        detect_kwargs = detect_kwargs or {}
+        a = [d + first for d in detect_ready_tone(w, fs, **detect_kwargs)]
         if len(a):
             detect += a
         # the last window may not allow a pwelch
         if (last - first) < nperseg:
             continue
         # compute PSD estimate for the current window
         iw = window_generator.iw
-        _, W[iw, :] = signal.welch(w, fs=fs, window='hann', nperseg=nperseg, axis=-1,
-                                   detrend='constant', return_onesided=True, scaling='density')
+        _, W[iw, :] = scipy.signal.welch(w, fs=fs, window='hann', nperseg=nperseg, axis=-1,
+                                         detrend='constant', return_onesided=True, scaling='density')
     # the onset detection may have duplicates with sliding window, average them and remove
     detect = np.sort(np.array(detect)) / fs
     ind = np.where(np.diff(detect) < 0.1)[0]

ibllib/io/extractors/training_audio.py

@@ -347,6 +347,14 @@ def test_lookups(self):
         assert self.ba.get_labels([0, 0, self.ba.bc.i2z(103)], mapping='Cosmos') == 997
         assert self.ba.get_labels([0, 0, 0], mapping='Cosmos') == 0
 
+    def test_lookups_probabilistic(self):
+        xyz = np.array([0, -.0058, -.0038])
+        aid = self.ba.get_labels(xyz, mapping='Beryl')
+        aids, proportions = self.ba.get_labels(xyz, mapping='Beryl', radius_um=250)
+        self.assertEqual(aid, 0)
+        assert (np.all(aids == np.array([0])))
+        assert (np.isclose(proportions, np.array([1.])))
+
     def test_slice(self):
         ba = self.ba
         nx, ny, nz = ba.bc.nxyz