klusta/utils.py
def _concatenate(arrs):
    if arrs is None:
        return
    arrs = [_as_array(arr) for arr in arrs if arr is not None]
    if not arrs:
        return
    return np.concatenate(arrs, axis=0)
klusta/traces/pca.py
class PCA(object):
    """Apply PCA to waveforms."""
    def __init__(self, n_pcs=None):
        self._n_pcs = n_pcs
        self._pcs = None

    def fit(self, waveforms, masks=None):
        """Compute the PCs of waveforms.

        Parameters
        ----------

        waveforms : ndarray
            Shape: `(n_spikes, n_samples, n_channels)`
        masks : ndarray
            Shape: `(n_spikes, n_channels)`

        """
        self._pcs = _compute_pcs(waveforms, n_pcs=self._n_pcs, masks=masks)
        return self._pcs

    def transform(self, waveforms, pcs=None):
        """Project waveforms on the PCs.

        Parameters
        ----------

        waveforms : ndarray
            Shape: `(n_spikes, n_samples, n_channels)`

        """
        if pcs is None:
            pcs = self._pcs
        # Need to call fit() if the pcs are None here.
        if pcs is not None:
            return _project_pcs(waveforms, pcs)
klusta/utils.py
def data_chunk(data, chunk, with_overlap=False):
    """Get a data chunk."""
    assert isinstance(chunk, tuple)
    if len(chunk) == 2:
        i, j = chunk
    elif len(chunk) == 4:
        if with_overlap:
            i, j = chunk[:2]
        else:
            i, j = chunk[2:]
    else:
        raise ValueError("'chunk' should have 2 or 4 elements, "
                         "not {0:d}".format(len(chunk)))
    return data[i:j, ...]
klusta/traces/detect.py
class Thresholder(object):
    """Threshold traces to detect spikes.

    Parameters
    ----------

    mode : str
        `'positive'`, `'negative'`, or `'both'`.
    thresholds : dict
        A `{str: float}` mapping for multiple thresholds (e.g. `weak`
        and `strong`).

    Example
    -------

    ```python
    thres = Thresholder('positive', thresholds=(.1, .2))
    crossings = thres(traces)
    ```

    """
    def __init__(self,
                 mode=None,
                 thresholds=None,
                 ):
        assert mode in ('positive', 'negative', 'both')
        if isinstance(thresholds, (float, int, np.ndarray)):
            thresholds = {'default': thresholds}
        if thresholds is None:
            thresholds = {}
        assert isinstance(thresholds, dict)
        self._mode = mode
        self._thresholds = thresholds

    def transform(self, data):
        """Return `data`, `-data`, or `abs(data)` depending on the mode."""
        if self._mode == 'positive':
            return data
        elif self._mode == 'negative':
            return -data
        elif self._mode == 'both':
            return np.abs(data)

    def detect(self, data_t, threshold=None):
        """Perform the thresholding operation."""
        # Accept dictionary of thresholds.
        if isinstance(threshold, (list, tuple)):
            return {name: self(data_t, threshold=name)
                    for name in threshold}
        # Use the only threshold by default (if there is only one).
        if threshold is None:
            assert len(self._thresholds) == 1
            threshold = list(self._thresholds.keys())[0]
        # Fetch the threshold from its name.
        if isinstance(threshold, string_types):
            assert threshold in self._thresholds
            threshold = self._thresholds[threshold]
        # threshold = float(threshold)
        # Threshold the data.
        return data_t > threshold

    def __call__(self, data, threshold=None):
        # Transform the data according to the mode.
        data_t = self.transform(data)
        return self.detect(data_t, threshold=threshold)
klusta/utils.py
class Bunch(dict):
    """A dict with additional dot syntax."""
    def __init__(self, *args, **kwargs):
        super(Bunch, self).__init__(*args, **kwargs)
        self.__dict__ = self

    def copy(self):
        return Bunch(super(Bunch, self).copy())
klusta/traces/store.py
class SpikeDetektStore(ArrayStore):
    """Store the following items:

    * filtered
    * components
    * spike_samples
    * features
    * masks

    """
    def __init__(self, root_dir, groups=None, chunk_keys=None):
        super(SpikeDetektStore, self).__init__(root_dir)
        self._groups = groups
        self._chunk_keys = chunk_keys
        self._spike_counts = SpikeCounts(groups=groups, chunk_keys=chunk_keys)

    def _rel_path(self, name=None, chunk_key=None, group=None):
        assert chunk_key >= 0
        assert group is None or group >= 0
        assert isinstance(name, string_types)
        group = group if group is not None else 'all'
        return 'group_{group}/{name}/chunk_{chunk:d}.npy'.format(
            chunk=chunk_key, name=name, group=group)

    @property
    def groups(self):
        return self._groups

    @property
    def chunk_keys(self):
        return self._chunk_keys

    def _iter(self, group=None, name=None):
        for chunk_key in self.chunk_keys:
            yield self.load(group=group, chunk_key=chunk_key, name=name)

    def spike_samples(self, group=None):
        if group is None:
            return {group: self.spike_samples(group) for group in self._groups}
        return self.concatenate(self._iter(group=group, name='spike_samples'))

    def features(self, group=None):
        """Yield chunk features."""
        if group is None:
            return {group: self.features(group) for group in self._groups}
        return self._iter(group=group, name='features')

    def masks(self, group=None):
        """Yield chunk masks."""
        if group is None:
            return {group: self.masks(group) for group in self._groups}
        return self._iter(group=group, name='masks')

    @property
    def spike_counts(self):
        return self._spike_counts

    def append(self, group=None, chunk_key=None,
               spike_samples=None, features=None, masks=None,
               spike_offset=0):
        if spike_samples is None or len(spike_samples) == 0:
            return
        n = len(spike_samples)
        assert features.shape[0] == n
        assert masks.shape[0] == n
        spike_samples = spike_samples + spike_offset
        assert np.all(np.diff(spike_samples) >= 0)

        self.store(group=group, chunk_key=chunk_key,
                   name='features', data=features)
        self.store(group=group, chunk_key=chunk_key,
                   name='masks', data=masks)
        self.store(group=group, chunk_key=chunk_key,
                   name='spike_samples', data=spike_samples)
        self._spike_counts.append(group=group, chunk_key=chunk_key, count=n)

    def concatenate(self, arrays):
        return _concatenate(arrays)

    def delete_all(self, name):
        """Delete all files for a given data name."""
        for group in self._groups:
            for chunk_key in self._chunk_keys:
                super(SpikeDetektStore, self).delete(name=name, group=group,
                                                     chunk_key=chunk_key)
klusta/traces/filter.py
class Filter(object):
    """Multichannel bandpass filter.

    The filter is applied on every column of a 2D array.

    Example
    -------

    ```python
    fil = Filter(rate=20000., low=5000., high=15000., order=4)
    traces_f = fil(traces)
    ```

    """
    def __init__(self, rate=None, low=None, high=None, order=None):
        self._filter = bandpass_filter(rate=rate,
                                       low=low,
                                       high=high,
                                       order=order,
                                       )

    def __call__(self, data):
        return apply_filter(data, filter=self._filter)
klusta/kwik/mea.py
def _probe_adjacency_list(probe):
    """Return an adjacency list of a whole probe."""
    cgs = probe['channel_groups'].values()
    graphs = [cg['graph'] for cg in cgs]
    edges = list(itertools.chain(*graphs))
    adjacency_list = _edges_to_adjacency_list(edges)
    return adjacency_list
klusta/utils.py
def get_excerpts(data, n_excerpts=None, excerpt_size=None):
    assert n_excerpts is not None
    assert excerpt_size is not None
    if len(data) < n_excerpts * excerpt_size:
        return data
    elif n_excerpts == 0:
        return data[:0]
    elif n_excerpts == 1:
        return data[:excerpt_size]
    out = np.concatenate([data_chunk(data, chunk)
                          for chunk in excerpts(len(data),
                                                n_excerpts=n_excerpts,
                                                excerpt_size=excerpt_size)])
    assert len(out) <= n_excerpts * excerpt_size
    return out
klusta/kwik/mea.py
def _channels_per_group(probe):
    groups = probe['channel_groups'].keys()
    return {group: probe['channel_groups'][group]['channels']
            for group in groups}
klusta/utils.py
def chunk_bounds(n_samples, chunk_size, overlap=0):
    """Return chunk bounds.
    Chunks have the form:
        [ overlap/2 | chunk_size-overlap | overlap/2 ]
        s_start   keep_start           keep_end     s_end
    Except for the first and last chunks which do not have a left/right
    overlap.
    This generator yields (s_start, s_end, keep_start, keep_end).
    """
    s_start = 0
    s_end = chunk_size
    keep_start = s_start
    keep_end = s_end - overlap // 2
    yield s_start, s_end, keep_start, keep_end

    while s_end - overlap + chunk_size < n_samples:
        s_start = s_end - overlap
        s_end = s_start + chunk_size
        keep_start = keep_end
        keep_end = s_end - overlap // 2
        if s_start < s_end:
            yield s_start, s_end, keep_start, keep_end

    s_start = s_end - overlap
    s_end = n_samples
    keep_start = keep_end
    keep_end = s_end
    if s_start < s_end:
        yield s_start, s_end, keep_start, keep_end
klusta/traces/detect.py
def compute_threshold(arr, single_threshold=True, std_factor=None):
    """Compute the threshold(s) of filtered traces.

    Parameters
    ----------

    arr : ndarray
        Filtered traces, shape `(n_samples, n_channels)`.
    single_threshold : bool
        Whether there should be a unique threshold for all channels, or
        one threshold per channel.
    std_factor : float or 2-tuple
        The threshold in unit of signal std. Two values can be specified
        for multiple thresholds (weak and strong).

    Returns
    -------

    thresholds : ndarray
        A `(2,)` or `(2, n_channels)` array with the thresholds.

    """
    assert arr.ndim == 2
    ns, nc = arr.shape

    assert std_factor is not None
    if isinstance(std_factor, (int, float)):
        std_factor = (std_factor, std_factor)
    assert isinstance(std_factor, (tuple, list))
    assert len(std_factor) == 2
    std_factor = np.array(std_factor)

    if not single_threshold:
        std_factor = std_factor[:, None]

    # Get the median of all samples in all excerpts, on all channels.
    if single_threshold:
        median = np.median(np.abs(arr))
    # Or independently for each channel.
    else:
        median = np.median(np.abs(arr), axis=0)

    # Compute the threshold from the median.
    std = median / .6745
    threshold = std_factor * std
    assert isinstance(threshold, np.ndarray)

    if single_threshold:
        assert threshold.ndim == 1
        assert len(threshold) == 2
    else:
        assert threshold.ndim == 2
        assert threshold.shape == (2, nc)
    return threshold
klusta/traces/detect.py
class FloodFillDetector(object):
    """Detect spikes in weak and strong threshold crossings.

    Parameters
    ----------

    probe_adjacency_list : dict
        A dict `{channel: [neighbors]}`.
    join_size : int
        The number of samples defining the tolerance in time for
        finding connected components

    Example
    -------

    ```python
    det = FloodFillDetector(probe_adjacency_list=...,
                            join_size=...)
    components = det(weak_crossings, strong_crossings)
    ```

    `components` is a list of `(n, 2)` int arrays with the sample and channel
    for every sample in the component.

    """
    def __init__(self, probe_adjacency_list=None, join_size=None,
                 channels_per_group=None):
        self._adjacency_list = probe_adjacency_list
        self._join_size = join_size
        self._channels_per_group = channels_per_group

    def __call__(self, weak_crossings=None, strong_crossings=None):
        weak_crossings = np.asarray(weak_crossings, np.bool)
        strong_crossings = np.asarray(strong_crossings, np.bool)
        all_channels = sorted([item for sublist
                              in self._channels_per_group.values()
                              for item in sublist])

        cc = connected_components(weak_crossings=weak_crossings,
                                  strong_crossings=strong_crossings,
                                  probe_adjacency_list=self._adjacency_list,
                                  channels=all_channels,
                                  join_size=self._join_size,
                                  )
        # cc is a list of list of pairs (sample, channel)
        return [np.array(c) for c in cc]
klusta/traces/waveform.py
class WaveformExtractor(object):
    """Extract waveforms after data filtering and spike detection."""
    def __init__(self,
                 extract_before=None,
                 extract_after=None,
                 weight_power=None,
                 thresholds=None,
                 channels_per_group=None,
                 ):
        self._extract_before = extract_before
        self._extract_after = extract_after
        self._weight_power = weight_power if weight_power is not None else 1.
        self._thresholds = thresholds or {}
        self._channels_per_group = channels_per_group
        # mapping channel => channels in the shank
        self._dep_channels = {i: channels
                              for channels in channels_per_group.values()
                              for i in channels}
        self._channel_groups = {i: g
                                for g, channels in channels_per_group.items()
                                for i in channels}

    def _component(self, component, data=None, n_samples=None):
        comp_s = component[:, 0]  # shape: (component_size,)
        comp_ch = component[:, 1]  # shape: (component_size,)
        channel = comp_ch[0]
        if channel not in self._dep_channels:
            raise RuntimeError("Channel `{}` appears to be dead and should "
                               "have been excluded from the threshold "
                               "crossings.".format(channel))
        channels = self._dep_channels[channel]
        group = self._channel_groups[comp_ch[0]]

        # Get the temporal window around the waveform.
        s_min, s_max = (comp_s.min() - 3), (comp_s.max() + 4)
        s_min = max(s_min, 0)
        s_max = min(s_max, n_samples)
        assert s_min < s_max

        return Bunch(comp_s=comp_s,
                     comp_ch=comp_ch,
                     s_min=s_min,
                     s_max=s_max,
                     channels=channels,
                     group=group,
                     )

    def _normalize(self, x):
        x = np.asarray(x)
        tw = self._thresholds['weak']
        ts = self._thresholds['strong']
        return np.clip((x - tw) / (ts - tw), 0, 1)

    def _comp_wave(self, data_t, comp):
        comp_s, comp_ch = comp.comp_s, comp.comp_ch
        s_min, s_max = comp.s_min, comp.s_max
        nc = data_t.shape[1]
        # Data on weak threshold crossings. shape: (some_length, nc)
        wave = np.zeros((s_max - s_min, nc), dtype=data_t.dtype)
        # The sample where the peak is reached, on each channel.
        wave[comp_s - s_min, comp_ch] = data_t[comp_s, comp_ch]
        return wave

    def masks(self, data_t, wave, comp):
        nc = data_t.shape[1]
        channels = comp.channels
        comp_ch = comp.comp_ch
        s_min = comp.s_min

        # Binary mask. shape: (nc,)
        masks_bin = np.zeros(nc, dtype=np.bool)
        masks_bin[np.unique(comp_ch)] = 1

        # Find the peaks (relative to the start of the chunk). shape: (nc,)
        peaks = np.argmax(wave, axis=0) + s_min
        # Peak values on each channel. shape: (nc,)
        peaks_values = data_t[peaks, np.arange(0, nc)] * masks_bin

        # Compute the float masks.
        masks_float = self._normalize(peaks_values)
        # Keep shank channels.
        masks_float = masks_float[channels]
        return masks_float

    def spike_sample_aligned(self, wave, comp):
        s_min, s_max = comp.s_min, comp.s_max
        # Compute the fractional peak.
        wave_n = self._normalize(wave)
        wave_n_p = np.power(wave_n, self._weight_power)
        u = np.arange(s_max - s_min)[:, np.newaxis]
        # Spike aligned time relative to the start of the chunk.
        s_aligned = np.sum(wave_n_p * u) / np.sum(wave_n_p) + s_min
        return s_aligned

    def extract(self, data, s_aligned, channels=None):
        s = int(s_aligned)
        # Get block of given size around peak sample.
        waveform = _get_padded(data,
                               s - self._extract_before - 1,
                               s + self._extract_after + 2)
        return waveform[:, channels]  # Keep shank channels.

    def align(self, waveform, s_aligned):
        s = int(s_aligned)
        sb, sa = self._extract_before, self._extract_after
        # Perform interpolation around the fractional peak.
        old_s = np.arange(s - sb - 1, s + sa + 2)
        new_s = np.arange(s - sb + 0, s + sa + 0) + (s_aligned - s)
        try:
            f = interp1d(old_s, waveform, bounds_error=True,
                         kind='cubic', axis=0)
        except ValueError:
            logger.warn("Interpolation error at time {0:d}".format(s))
            return waveform
        return f(new_s)

    def set_thresholds(self, **kwargs):
        self._thresholds.update(kwargs)

    def __call__(self, component=None, data=None, data_t=None):
        assert data.shape == data_t.shape
        comp = self._component(component,
                               data=data,
                               n_samples=data_t.shape[0],
                               )
        channels = comp.channels

        wave = self._comp_wave(data_t, comp)
        masks = self.masks(data_t, wave, comp)
        s_aligned = self.spike_sample_aligned(wave, comp)

        waveform_unaligned = self.extract(data, s_aligned, channels=channels)
        waveform_aligned = self.align(waveform_unaligned, s_aligned)

        assert waveform_aligned.ndim == 2
        assert masks.ndim == 1
        assert waveform_aligned.shape[1] == masks.shape[0]

        return comp.group, s_aligned, waveform_aligned, masks


Question: from collections import defaultdict
from tqdm import tqdm
from ..utils import (Bunch,
                     get_excerpts,
                     chunk_bounds,
                     data_chunk,
                     _concatenate,
                     )
from ..kwik.mea import (_channels_per_group,
                        _probe_adjacency_list,
                        )
from .detect import Thresholder, compute_threshold, FloodFillDetector
from .filter import Filter
from .pca import PCA
from .store import SpikeDetektStore
from .waveform import WaveformExtractor
import logging
import numpy as np
        Returns a split object.

        Parameters
        ----------
        components : list
            List of connected components.
        traces_f : array
            Filtered data.
        thresholds : dict
            The weak and strong thresholds.
        keep_bounds : tuple
            (keep_start, keep_end).
        s_start : 0
            Start of the chunk.

        """
        n_spikes = len(components)
        if n_spikes == 0:
            return {}

        # Transform the filtered data according to the detection mode.
        thresholder = self._create_thresholder()
        traces_t = thresholder.transform(traces_f)
        # Extract all waveforms.
        extractor = self._create_extractor(thresholds)
        groups, samples, waveforms, masks = zip(*[extractor(component,
                                                            data=traces_f,
                                                            data_t=traces_t,
                                                            )
                                                  for component in components])

        # Create the return arrays.
        groups = np.array(groups, dtype=np.int32)
        assert groups.shape == (n_spikes,)
        assert groups.dtype == np.int32

        samples = np.array(samples, dtype=np.float64)
        assert samples.shape == (n_spikes,)
        assert samples.dtype == np.float64

        # These are lists of arrays of various shapes (because of various
        # groups).
        waveforms = _array_list(waveforms)
        assert waveforms.shape == (n_spikes,)
        assert waveforms.dtype == np.object

        masks = _array_list(masks)
        assert masks.dtype == np.object
        assert masks.shape == (n_spikes,)

        # Reorder the spikes.
        idx = np.argsort(samples)
        groups = groups[idx]
        samples = samples[idx]
        waveforms = waveforms[idx]
        masks = masks[idx]

        # Remove spikes in the overlapping bands.
        # WARNING: add s_start to spike_samples, because spike_samples
        # is relative to the start of the chunk.
        # It is important to add s_start and not keep_start, because of
        # edge effects between overlapping chunks.
        s_start = s_start or 0
        (keep_start, keep_end) = keep_bounds
        idx = _keep_spikes(samples + s_start, (keep_start, keep_end))

        # Split the data according to the channel groups.
        split = _split_spikes(groups, idx=idx, spike_samples=samples,
                              waveforms=waveforms, masks=masks)
        # split: {group: {'spike_samples': ..., 'waveforms':, 'masks':}}

        # Assert that spike samples are increasing.
        for group in split:
            samples = split[group]['spike_samples']
            if samples is not None:
                assert np.all(np.diff(samples) >= 0)

        return split

    def waveform_pcs(self, waveforms, masks):
        """Compute waveform principal components.

        Returns
        -------

        pcs : array
            An `(n_features, n_samples, n_channels)` array.

        """
        pca = self._create_pca()
        if waveforms is None or not len(waveforms):
            return
        assert (waveforms.shape[0], waveforms.shape[2]) == masks.shape
        return pca.fit(waveforms, masks)

    def features(self, waveforms, pcs):
        """Extract features from waveforms.

        Returns
        -------

        features : array
            An `(n_spikes, n_channels, n_features)` array.

        """
        pca = self._create_pca()
        out = pca.transform(waveforms, pcs=pcs)
        assert out.dtype == np.float32
        return out

    # Chunking
    # -------------------------------------------------------------------------

    def iter_chunks(self, n_samples):
        """Iterate over chunks."""
        rate = self._kwargs['sample_rate']
        chunk_size = int(self._kwargs['chunk_size_seconds'] * rate)
        overlap = int(self._kwargs['chunk_overlap_seconds'] * rate)
        for chunk_idx, bounds in enumerate(chunk_bounds(n_samples, chunk_size,
                                                        overlap=overlap)):

Answer:             yield Bunch(bounds=bounds,