<commit_before># -*- coding: utf-8 -*-

"""Tests of waveform loader."""

#------------------------------------------------------------------------------
# Imports
#------------------------------------------------------------------------------

import numpy as np
from numpy.testing import assert_array_equal as ae
import numpy.random as npr
from pytest import raises

from ...io.mock import artificial_traces
from ..waveform import _slice, WaveformLoader, WaveformExtracter
from ..filter import bandpass_filter, apply_filter


#------------------------------------------------------------------------------
# Tests extracter
#------------------------------------------------------------------------------

def test_extract():
    weak = 1.
    strong = 2.
    nc = 4
    ns = 20
    channels = list(range(nc))
    cpg = {0: channels}
    # graph = {0: [1, 2], 1: [0, 2], 2: [0, 1], 3: []}

    data = np.random.uniform(size=(ns, nc), low=0., high=1.)

    data[10, 0] = 0.5
    data[11, 0] = 1.5
    data[12, 0] = 1.0

    data[10, 1] = 1.5
    data[11, 1] = 2.5
    data[12, 1] = 2.0

    component = np.array([[10, 0],
                          [10, 1],
                          [11, 0],
                          [11, 1],
                          [12, 0],
                          [12, 1],
                          ])

    we = WaveformExtracter(extract_before=3,
                           extract_after=5,
                           thresholds={'weak': weak,
                                       'strong': strong},
                           channels_per_group=cpg,
                           )

    # _component()
    comp = we._component(component, n_samples=ns)
    ae(comp.comp_s, [10, 10, 11, 11, 12, 12])
    ae(comp.comp_ch, [0, 1, 0, 1, 0, 1])
    assert (comp.s_min, comp.s_max) == (10 - 3, 12 + 4)
    ae(comp.channels, range(nc))

    # _normalize()
    assert we._normalize(weak) == 0
    assert we._normalize(strong) == 1
    ae(we._normalize([(weak + strong) / 2.]), [.5])

    # _comp_wave()
    wave = we._comp_wave(data, comp)
    assert wave.shape == (3 + 5 + 1, nc)
    ae(wave[3:6, :], [[0.5, 1.5, 0., 0.],
                      [1.5, 2.5, 0., 0.],
                      [1.0, 2.0, 0., 0.]])

    # masks()
    masks = we.masks(data, wave, comp)
    ae(masks, [.5, 1., 0, 0])

    # spike_sample_aligned()
    s = we.spike_sample_aligned(wave, comp)
    assert 11 <= s < 12

    # extract()
    wave_e = we.extract(data, s, channels=channels)
    assert wave_e.shape[1] == wave.shape[1]
    ae(wave[3:6, :2], wave_e[3:6, :2])

    # align()
    wave_a = we.align(wave_e, s)
    assert wave_a.shape == (3 + 5, nc)

    # Test final call.
    s_f, wave_f, masks_f = we(component, data=data, data_t=data)
    assert s_f == s
    ae(masks_f, masks)
    ae(wave_f, wave_a)


#------------------------------------------------------------------------------
# Tests loader
#------------------------------------------------------------------------------

def test_slice():
    assert _slice(0, (20, 20)) == slice(0, 20, None)


def test_loader():
    n_samples_trace, n_channels = 10000, 100
    n_samples = 40
    n_spikes = n_samples_trace // (2 * n_samples)

    traces = artificial_traces(n_samples_trace, n_channels)
    spike_samples = np.cumsum(npr.randint(low=0, high=2 * n_samples,
                                          size=n_spikes))

    with raises(ValueError):
        WaveformLoader(traces)

    # Create a loader.
    loader = WaveformLoader(traces, n_samples=n_samples)
    assert id(loader.traces) == id(traces)
    loader.traces = traces
<commit_msg>Test waveform extractor with channel order.<commit_after># -*- coding: utf-8 -*-

"""Tests of waveform loader."""

#------------------------------------------------------------------------------
# Imports
#------------------------------------------------------------------------------

import numpy as np
from numpy.testing import assert_array_equal as ae
import numpy.random as npr
from pytest import raises

from ...io.mock import artificial_traces
from ..waveform import _slice, WaveformLoader, WaveformExtracter
from ..filter import bandpass_filter, apply_filter


#------------------------------------------------------------------------------
# Tests extracter
#------------------------------------------------------------------------------

def test_extract_simple():
    weak = 1.
    strong = 2.
    nc = 4
    ns = 20
    channels = list(range(nc))
    cpg = {0: channels}
    # graph = {0: [1, 2], 1: [0, 2], 2: [0, 1], 3: []}

    data = np.random.uniform(size=(ns, nc), low=0., high=1.)

    data[10, 0] = 0.5
    data[11, 0] = 1.5
    data[12, 0] = 1.0

    data[10, 1] = 1.5
    data[11, 1] = 2.5
    data[12, 1] = 2.0

    component = np.array([[10, 0],
                          [10, 1],
                          [11, 0],
                          [11, 1],
                          [12, 0],
                          [12, 1],
                          ])

    we = WaveformExtracter(extract_before=3,
                           extract_after=5,
                           thresholds={'weak': weak,
                                       'strong': strong},
                           channels_per_group=cpg,
                           )

    # _component()
    comp = we._component(component, n_samples=ns)
    ae(comp.comp_s, [10, 10, 11, 11, 12, 12])
    ae(comp.comp_ch, [0, 1, 0, 1, 0, 1])
    assert (comp.s_min, comp.s_max) == (10 - 3, 12 + 4)
    ae(comp.channels, range(nc))

    # _normalize()
    assert we._normalize(weak) == 0
    assert we._normalize(strong) == 1
    ae(we._normalize([(weak + strong) / 2.]), [.5])

    # _comp_wave()
    wave = we._comp_wave(data, comp)
    assert wave.shape == (3 + 5 + 1, nc)
    ae(wave[3:6, :], [[0.5, 1.5, 0., 0.],
                      [1.5, 2.5, 0., 0.],
                      [1.0, 2.0, 0., 0.]])

    # masks()
    masks = we.masks(data, wave, comp)
    ae(masks, [.5, 1., 0, 0])

    # spike_sample_aligned()
    s = we.spike_sample_aligned(wave, comp)
    assert 11 <= s < 12

    # extract()
    wave_e = we.extract(data, s, channels=channels)
    assert wave_e.shape[1] == wave.shape[1]
    ae(wave[3:6, :2], wave_e[3:6, :2])

    # align()
    wave_a = we.align(wave_e, s)
    assert wave_a.shape == (3 + 5, nc)

    # Test final call.
    s_f, wave_f, masks_f = we(component, data=data, data_t=data)
    assert s_f == s
    ae(masks_f, masks)
    ae(wave_f, wave_a)

    # Tests with a different order.
    we = WaveformExtracter(extract_before=3,
                           extract_after=5,
                           thresholds={'weak': weak,
                                       'strong': strong},
                           channels_per_group={0: [1, 0, 3]},
                           )
    s_f_o, wave_f_o, masks_f_o = we(component, data=data, data_t=data)
    assert s_f == s_f_o
    assert np.allclose(wave_f[:, [1, 0, 3]], wave_f_o)
    ae(masks_f_o, [1., 0.5, 0.])


#------------------------------------------------------------------------------
# Tests loader
#------------------------------------------------------------------------------

def test_slice():
    assert _slice(0, (20, 20)) == slice(0, 20, None)


def test_loader():
    n_samples_trace, n_channels = 10000, 100
    n_samples = 40
    n_spikes = n_samples_trace // (2 * n_samples)

    traces = artificial_traces(n_samples_trace, n_channels)
    spike_samples = np.cumsum(npr.randint(low=0, high=2 * n_samples,
                                          size=n_spikes))

    with raises(ValueError):
        WaveformLoader(traces)

    # Create a loader.
    loader = WaveformLoader(traces, n_samples=n_samples)
    assert id(loader.traces) == id(traces)
    loader.traces = traces

{'fraction_non_alphanumeric': 0.250581545619023, 'fraction_numerical': 0.04277591108813647, 'mean_word_length': 1.3254206730769231, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 9, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2994297', 'n_tokens_mistral': 2717, 'n_tokens_neox': 2442, 'n_words': 588}