Merge pull request #115 from iamamutt/main

Remove ibllib dependency

Author: Thinh Nguyen

.gitignore

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

element_array_ephys/ephys_no_curation.py

@@ -102,7 +102,7 @@ def make(self, key):
 
         from .plotting.unit_level import (
             plot_waveform,
-            plot_correlogram,
+            plot_auto_correlogram,
             plot_depth_waveforms,
         )
 
@@ -119,7 +119,7 @@ def make(self, key):
             waveform=peak_electrode_waveform, sampling_rate=sampling_rate
         )
 
-        correlogram_fig = plot_correlogram(
+        correlogram_fig = plot_auto_correlogram(
             spike_times=spike_times, bin_size=0.001, window_size=1
         )
 

element_array_ephys/ephys_report.py

@@ -0,0 +1,147 @@
+"""Code adapted from International Brain Laboratory, T. (2021). ibllib [Computer software]. https://github.com/int-brain-lab/ibllib
+"""
+
+import numpy as np
+
+
+def _index_of(arr: np.ndarray, lookup: np.ndarray):
+    """Replace scalars in an array by their indices in a lookup table."""
+
+    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: np.ndarray,
+    spike_clusters: np.ndarray,
+    bin_size: float,
+    window_size: int,
+) -> np.ndarray:
+    """Compute all pairwise cross-correlograms among the clusters appearing in `spike_clusters`.
+
+    Args:
+        spike_times (np.ndarray): Spike times in seconds.
+        spike_clusters (np.ndarray): Spike-cluster mapping.
+        bin_size (float): Size of the time bin in seconds.
+        window_size (int): Size of the correlogram window in seconds.
+
+    Returns:
+         np.ndarray: cross-correlogram array
+    """
+    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: np.ndarray, bin_size: float, window_size: int) -> np.ndarray:
+    """Compute the auto-correlogram of a unit.
+
+    Args:
+        spike_times (np.ndarray): Spike times in seconds.
+        bin_size (float, optional): Size of the time bin in seconds.
+        window_size (int, optional): Size of the correlogram window in seconds.
+
+    Returns:
+        np.ndarray: auto-correlogram array (winsize_samples,)
+    """
+    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/corr.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/probe_level.py

@@ -1,12 +1,22 @@
-from .. import probe
 from modulefinder import Module
+
 import numpy as np
 import pandas as pd
 import plotly.graph_objs as go
+from typing import Any
+from .. import probe
 
 
 def plot_waveform(waveform: np.ndarray, sampling_rate: float) -> go.Figure:
+    """Plot unit waveform.
+
+    Args:
+        waveform (np.ndarray): Amplitude of a spike waveform in μV.
+        sampling_rate (float): Sampling rate in kHz.
 
+    Returns:
+        go.Figure: Plotly figure object.
+    """
     waveform_df = pd.DataFrame(data={"waveform": waveform})
     waveform_df["timestamp"] = waveform_df.index / sampling_rate
 
@@ -31,11 +41,20 @@ def plot_waveform(waveform: np.ndarray, sampling_rate: float) -> go.Figure:
     return fig
 
 
-def plot_correlogram(
+def plot_auto_correlogram(
     spike_times: np.ndarray, bin_size: float = 0.001, window_size: int = 1
 ) -> go.Figure:
+    """Plot the auto-correlogram of a unit.
 
-    from brainbox.singlecell import acorr
+    Args:
+        spike_times (np.ndarray): Spike timestamps in seconds
+        bin_size (float, optional): Size of the time bin (lag) in seconds. Defaults to 0.001.
+        window_size (int, optional): Size of the correlogram window in seconds. Defaults to 1.
+
+    Returns:
+        go.Figure: Plotly figure object.
+    """
+    from .corr import acorr
 
     correlogram = acorr(
         spike_times=spike_times, bin_size=bin_size, window_size=window_size
@@ -74,9 +93,19 @@ def plot_correlogram(
 
 def plot_depth_waveforms(
     ephys: Module,
-    unit_key: dict,
+    unit_key: dict[str, Any],
     y_range: float = 60,
 ) -> go.Figure:
+    """Plot waveforms
+
+    Args:
+        ephys (Module): Imported ephys module.
+        unit_key (dict[str, Any]): Key dictionary from ephys.CuratedClustering.Unit table.
+        y_range (float, optional): Vertical range to show waveforms relative to the peak waveform in μm. Defaults to 60.
+
+    Returns:
+        go.Figure: Plotly figure object.
+    """
 
     sampling_rate = (ephys.EphysRecording & unit_key).fetch1(
         "sampling_rate"

element_array_ephys/plotting/unit_level.py

@@ -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
+plotly
+seaborn