def test_loader_filter():
    traces = np.c_[np.arange(20), np.arange(20, 40)].astype(np.int32)
    n_samples_trace, n_channels = traces.shape
    h = 3

    def my_filter(x, axis=0):
        return x * x

    loader = WaveformLoader(traces,
                            n_samples_waveforms=(h, h),
                            filter=my_filter,
                            filter_margin=2)

    t = 10
    waveform_filtered = loader[t]
    traces_filtered = my_filter(traces)
    assert np.allclose(waveform_filtered, traces_filtered[t - h:t + h, :])
dataset/ETHPy150Open kwikteam/phy/phy/traces/tests/test_waveform.py/test_loader_filter/original