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/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/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/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
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/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 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/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/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/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/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/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/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 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


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
    n_samples, n_channels = traces.shape
    # Take a subset if necessary.
    if interval_samples is not None:
        start, end = interval_samples
    else:
        return traces, 0
    assert 0 <= start < end
    # WARNING: this loads all traces into memory! To fix this properly,
    # we'll have to implement lazy chunking in ConcatenatedTraces.
    size = (end - start) * traces.shape[1] * traces.dtype.itemsize / 1024. ** 3
    if size > 1:
        logger.warn("Loading all traces in memory: this will require %.3f GB "
                    "of RAM! ", size)
        logger.warn("To avoid this, do not specify `--interval`. "
                    "This bug will be fixed later.")
    traces = traces[start:end, ...]
    if start > 0:
        # TODO: add offset to the spike samples...
        raise NotImplementedError("Need to add `start` to the "
                                  "spike samples")
    return traces, start


def _subtract_offsets(samples, offsets):
    """Subtract the recording offsets from spike samples.

    Return the subtracted spike samples and the spike_recordings array.

    """
    if samples is None:
        return None, None
    assert isinstance(samples, (list, np.ndarray))
    samples = np.asarray(samples).copy()
    assert len(offsets) >= 2
    assert offsets[0] == 0
    assert offsets[-1] >= samples[-1]
    n = len(offsets) - 1

    # Find where to insert the offsets in the spike samples.
    ind = np.searchsorted(samples, offsets)
    assert ind[0] == 0
    assert len(ind) == n + 1

    spike_recordings = np.zeros(len(samples), dtype=np.int32)

    # Loop over all recordings.
    for rec in range(n):
        # Start of the current recording.
        start_rec = offsets[rec]
        # Spike indices of the first spikes in every recording.
        i, j = ind[rec], ind[rec + 1]
        # Ensure that the selected spikes belong to the current recording.
        if i < len(samples) - 1:
            assert start_rec <= samples[i]
        # Subtract the current recording offset to the selected spikes.
        samples[i:j] -= start_rec
        # Create the spike_recordings array.
        spike_recordings[i:j] = rec
        assert np.all(samples[i:j] >= 0)
        assert np.all(samples[i:j] <= (offsets[rec + 1] - start_rec))

    return samples, spike_recordings


def _relative_channels(channels, adjacency):
    """Renumber channels from absolute indices to relative indices,
    to match arrays used in the detection code.

    Parameters
    ----------

    channels : dict
        A dict {group: list_of_channels}
    adjacency : dict
        A dict {group: set_of_neighbors}

    """
    ch_out = {}
    adj_out = {}
    mapping = {}
    offset = 0
    for group in channels:
        ch = channels[group]
        n = len(ch)
        ch_out[group] = [i + offset for i in range(n)]
        # From absolute to relative indices.
        mapping.update({c: (i + offset) for i, c in enumerate(ch)})
        offset += n
    # Recreate the adjacency dict by applying the mapping to
    # both the keys and values.
    for c, i in mapping.items():
        adj_out[i] = set(mapping[_] for _ in adjacency.get(c, set())
                         if _ in mapping)
    return ch_out, adj_out


#------------------------------------------------------------------------------
# Spike detection class
#------------------------------------------------------------------------------

class SpikeDetekt(object):
    """Spike detection class.

    Parameters
    ----------

    tempdir : str
        Path to the temporary directory used by the algorithm. It should be
        on a SSD for best performance.
    probe : dict
        The probe dictionary.
    **kwargs : dict
        Spike detection parameters.

    """
    def __init__(self, tempdir=None, probe=None, **kwargs):
        super(SpikeDetekt, self).__init__()
        self._tempdir = tempdir
        # Load a probe.
        if probe is not None:

Answer:             kwargs['probe_channels'] = _channels_per_group(probe)