refactor(deps): ➖ remove ibllib deps and add acorr func

Author: Joseph Burling

.gitignore

@@ -122,3 +122,4 @@ docker-compose.y*ml
 
 # vscode settings
 .vscode
+*.code-workspace

element_array_ephys/plotting/acorr.py

@@ -0,0 +1,165 @@
+"""autocorrelation code adapted from ibllib
+https://github.com/int-brain-lab/ibllib
+"""
+
+import numpy as np
+
+
+def _index_of(arr, lookup):
+    """Replace scalars in an array by their indices in a lookup table.
+
+    Implicitly assume that:
+
+    * All elements of arr and lookup are non-negative integers.
+    * All elements or arr belong to lookup.
+
+    This is not checked for performance reasons.
+
+    """
+    # Equivalent of np.digitize(arr, lookup) - 1, but much faster.
+    # TODO: assertions to disable in production for performance reasons.
+    # TODO: np.searchsorted(lookup, arr) is faster on small arrays with large
+    # values
+    lookup = np.asarray(lookup, dtype=np.int32)
+    m = (lookup.max() if len(lookup) else 0) + 1
+    tmp = np.zeros(m + 1, dtype=int)
+    # Ensure that -1 values are kept.
+    tmp[-1] = -1
+    if len(lookup):
+        tmp[lookup] = np.arange(len(lookup))
+    return tmp[arr]
+
+
+def _increment(arr, indices):
+    """Increment some indices in a 1D vector of non-negative integers.
+    Repeated indices are taken into account."""
+    bbins = np.bincount(indices)
+    arr[: len(bbins)] += bbins
+    return arr
+
+
+def _diff_shifted(arr, steps=1):
+    return arr[steps:] - arr[: len(arr) - steps]
+
+
+def _create_correlograms_array(n_clusters, winsize_bins):
+    return np.zeros((n_clusters, n_clusters, winsize_bins // 2 + 1), dtype=np.int32)
+
+
+def _symmetrize_correlograms(correlograms):
+    """Return the symmetrized version of the CCG arrays."""
+
+    n_clusters, _, n_bins = correlograms.shape
+    assert n_clusters == _
+
+    # We symmetrize c[i, j, 0].
+    # This is necessary because the algorithm in correlograms()
+    # is sensitive to the order of identical spikes.
+    correlograms[..., 0] = np.maximum(correlograms[..., 0], correlograms[..., 0].T)
+
+    sym = correlograms[..., 1:][..., ::-1]
+    sym = np.transpose(sym, (1, 0, 2))
+
+    return np.dstack((sym, correlograms))
+
+
+def xcorr(spike_times, spike_clusters, bin_size=None, window_size=None):
+    """Compute all pairwise cross-correlograms among the clusters appearing in `spike_clusters`.
+
+    Parameters
+    ----------
+
+    :param spike_times: Spike times in seconds.
+    :type spike_times: array-like
+    :param spike_clusters: Spike-cluster mapping.
+    :type spike_clusters: array-like
+    :param bin_size: Size of the bin, in seconds.
+    :type bin_size: float
+    :param window_size: Size of the window, in seconds.
+    :type window_size: float
+
+    Returns an `(n_clusters, n_clusters, winsize_samples)` array with all pairwise
+    cross-correlograms.
+
+    """
+    assert np.all(np.diff(spike_times) >= 0), "The spike times must be increasing."
+    assert spike_times.ndim == 1
+    assert spike_times.shape == spike_clusters.shape
+
+    # Find `binsize`.
+    bin_size = np.clip(bin_size, 1e-5, 1e5)  # in seconds
+
+    # Find `winsize_bins`.
+    window_size = np.clip(window_size, 1e-5, 1e5)  # in seconds
+    winsize_bins = 2 * int(0.5 * window_size / bin_size) + 1
+
+    # Take the cluster order into account.
+    clusters = np.unique(spike_clusters)
+    n_clusters = len(clusters)
+
+    # Like spike_clusters, but with 0..n_clusters-1 indices.
+    spike_clusters_i = _index_of(spike_clusters, clusters)
+
+    # Shift between the two copies of the spike trains.
+    shift = 1
+
+    # At a given shift, the mask precises which spikes have matching spikes
+    # within the correlogram time window.
+    mask = np.ones_like(spike_times, dtype=bool)
+
+    correlograms = _create_correlograms_array(n_clusters, winsize_bins)
+
+    # The loop continues as long as there is at least one spike with
+    # a matching spike.
+    while mask[:-shift].any():
+        # Interval between spike i and spike i+shift.
+        spike_diff = _diff_shifted(spike_times, shift)
+
+        # Binarize the delays between spike i and spike i+shift.
+        spike_diff_b = np.round(spike_diff / bin_size).astype(np.int64)
+
+        # Spikes with no matching spikes are masked.
+        mask[:-shift][spike_diff_b > (winsize_bins / 2)] = False
+
+        # Cache the masked spike delays.
+        m = mask[:-shift].copy()
+        d = spike_diff_b[m]
+
+        # Find the indices in the raveled correlograms array that need
+        # to be incremented, taking into account the spike clusters.
+        indices = np.ravel_multi_index(
+            (spike_clusters_i[:-shift][m], spike_clusters_i[+shift:][m], d),
+            correlograms.shape,
+        )
+
+        # Increment the matching spikes in the correlograms array.
+        _increment(correlograms.ravel(), indices)
+
+        shift += 1
+
+    return _symmetrize_correlograms(correlograms)
+
+
+def acorr(spike_times, bin_size=None, window_size=None):
+    """Compute the auto-correlogram of a neuron.
+
+    Parameters
+    ----------
+
+    :param spike_times: Spike times in seconds.
+    :type spike_times: array-like
+    :param bin_size: Size of the bin, in seconds.
+    :type bin_size: float
+    :param window_size: Size of the window, in seconds.
+    :type window_size: float
+
+    Returns an `(winsize_samples,)` array with the auto-correlogram.
+
+    """
+    xc = xcorr(
+        spike_times,
+        np.zeros_like(spike_times, dtype=np.int32),
+        bin_size=bin_size,
+        window_size=window_size,
+    )
+    return xc[0, 0, :]

element_array_ephys/plotting/probe_level.py

@@ -1,6 +1,6 @@
-import numpy as np
 import matplotlib
 import matplotlib.pyplot as plt
+import numpy as np
 import seaborn as sns
 
 

element_array_ephys/plotting/unit_level.py

@@ -1,9 +1,11 @@
-from .. import probe
 from modulefinder import Module
+
 import numpy as np
 import pandas as pd
 import plotly.graph_objs as go
 
+from .. import probe
+
 
 def plot_waveform(waveform: np.ndarray, sampling_rate: float) -> go.Figure:
 
@@ -35,7 +37,7 @@ def plot_correlogram(
     spike_times: np.ndarray, bin_size: float = 0.001, window_size: int = 1
 ) -> go.Figure:
 
-    from brainbox.singlecell import acorr
+    from .acorr import acorr
 
     correlogram = acorr(
         spike_times=spike_times, bin_size=bin_size, window_size=window_size

requirements.txt

@@ -3,5 +3,5 @@ pyopenephys @ git+https://github.com/datajoint-company/pyopenephys.git
 openpyxl
 pynwb>=2.0.0
 element-interface @ git+https://github.com/datajoint/element-interface.git
-ibllib @ git+https://github.com/int-brain-lab/ibllib.git
 plotly==5.9.0
+seaborn