From fb1e0d5fd4dfc332298668a225415c5b795c0192 Mon Sep 17 00:00:00 2001 From: Jonathan Jouty Date: Sun, 4 Feb 2018 11:06:06 +0000 Subject: Make merge_spike_trains work with empty spike trains, and faster 1. Fixes https://github.com/mariomulansky/PySpike/issues/30 2. Code is faster 3. Add test case --- pyspike/spikes.py | 22 ++++------------------ test/test_spikes.py | 10 ++++++++++ 2 files changed, 14 insertions(+), 18 deletions(-) diff --git a/pyspike/spikes.py b/pyspike/spikes.py index 486a4a0..cf47043 100644 --- a/pyspike/spikes.py +++ b/pyspike/spikes.py @@ -116,24 +116,10 @@ def merge_spike_trains(spike_trains): :param spike_trains: list of :class:`.SpikeTrain` :returns: spike train with the merged spike times """ - # get the lengths of the spike trains - lens = np.array([len(st.spikes) for st in spike_trains]) - merged_spikes = np.empty(np.sum(lens)) - index = 0 # the index for merged_spikes - indices = np.zeros_like(lens) # indices of the spike trains - index_list = np.arange(len(indices)) # indices of indices of spike trains - # that have not yet reached the end - # list of the possible events in the spike trains - vals = [spike_trains[i].spikes[indices[i]] for i in index_list] - while len(index_list) > 0: - i = np.argmin(vals) # the next spike is the minimum - merged_spikes[index] = vals[i] # put it to the merged spike train - i = index_list[i] - index += 1 # next index of merged spike train - indices[i] += 1 # next index for the chosen spike train - if indices[i] >= lens[i]: # remove spike train index if ended - index_list = index_list[index_list != i] - vals = [spike_trains[n].spikes[indices[n]] for n in index_list] + # concatenating and sorting with numpy is fast, it also means we can handle + # empty spike trains + merged_spikes = np.concatenate([st.spikes for st in spike_trains]) + merged_spikes.sort() return SpikeTrain(merged_spikes, [spike_trains[0].t_start, spike_trains[0].t_end]) diff --git a/test/test_spikes.py b/test/test_spikes.py index bcface2..ee505b5 100644 --- a/test/test_spikes.py +++ b/test/test_spikes.py @@ -85,6 +85,16 @@ def test_merge_spike_trains(): check_merged_spikes(merged_spikes.spikes, [st.spikes for st in spike_trains]) +def test_merge_empty_spike_trains(): + # first load the data + spike_trains = spk.load_spike_trains_from_txt(TEST_DATA, edges=(0, 4000)) + # take two non-empty trains, and one empty one + empty = spk.SpikeTrain([],[spike_trains[0].t_start,spike_trains[0].t_end]) + merged_spikes = spk.merge_spike_trains([spike_trains[0], empty, spike_trains[1]]) + # test if result is sorted + assert((merged_spikes.spikes == np.sort(merged_spikes.spikes)).all()) + # we don't need to check more, that's done by test_merge_spike_trains + if __name__ == "main": test_load_from_txt() -- cgit v1.2.3