diff --git a/elephant/statistics.py b/elephant/statistics.py index bb456ac02..80e061853 100644 --- a/elephant/statistics.py +++ b/elephant/statistics.py @@ -1113,9 +1113,6 @@ def optimal_kernel(st): sigma=str(kernel.sigma), invert=kernel.invert) - if pool_spike_trains: - rate = np.mean(rate, axis=1) - rate = neo.AnalogSignal(signal=rate, sampling_period=sampling_period, units=pq.Hz, t_start=t_start, @@ -1139,6 +1136,21 @@ def optimal_kernel(st): rate[:, i] *= len(spiketrain) /\ (np.mean(rate[:, i]).magnitude * duration) + # Pooling happens after the border correction, because the correction is + # defined per spike train: it rescales each column so that the integral + # over that column returns the spike count of the corresponding spike + # train. Pooling first collapses the column axis to a single column and + # leaves the loop above indexing columns that no longer exist. Correcting + # first also keeps the pooled rate equal to the mean over the columns of + # the corresponding `pool_spike_trains=False` result, which is the order + # already used for `elephant.trials` input. + if pool_spike_trains: + rate = neo.AnalogSignal( + signal=np.mean(rate.magnitude, axis=1, keepdims=True), + sampling_period=rate.sampling_period, + units=rate.units, t_start=rate.t_start, + kernel=kernel_annotation) + return rate diff --git a/elephant/test/test_statistics.py b/elephant/test/test_statistics.py index fde6d54e8..a541c8c8a 100644 --- a/elephant/test/test_statistics.py +++ b/elephant/test/test_statistics.py @@ -1107,6 +1107,61 @@ def test_instantaneous_rate_border_correction(self): self.assertLess(np.min(average_estimated_rate), (1. - rtol) * rate.item()) + def test_instantaneous_rate_border_correction_pool_spike_trains(self): + # The border correction rescales every column of the rate estimate + # with the spike count of the corresponding spike train, so it has to + # run before the spike trains are pooled. Pooling first collapses the + # column axis to a single column and the correction loop then indexes + # columns that no longer exist. + np.random.seed(0) + n_spiketrains = 5 + t_start = 0. * pq.ms + t_stop = 1000. * pq.ms + sampling_period = 1. * pq.ms + kernel = kernels.GaussianKernel(sigma=50. * pq.ms) + spiketrains = StationaryPoissonProcess( + rate=30. * pq.Hz, t_start=t_start, t_stop=t_stop + ).generate_n_spiketrains(n_spiketrains) + + kwargs = dict(sampling_period=sampling_period, kernel=kernel, + border_correction=True) + rate_pooled = statistics.instantaneous_rate( + spiketrains, pool_spike_trains=True, **kwargs) + rate_per_spiketrain = statistics.instantaneous_rate( + spiketrains, pool_spike_trains=False, **kwargs) + + n_bins = int(((t_stop - t_start) / sampling_period).simplified) + self.assertIsInstance(rate_pooled, neo.AnalogSignal) + self.assertEqual(rate_pooled.shape, (n_bins, 1)) + self.assertEqual(rate_per_spiketrain.shape, (n_bins, n_spiketrains)) + + # Pooling is documented as an average over spike trains, so the + # pooled rate has to equal the mean over the columns of the + # un-pooled rate. + assert_array_almost_equal( + rate_pooled.magnitude[:, 0], + np.mean(rate_per_spiketrain.magnitude, axis=1)) + + # The border correction makes the integral over each un-pooled rate + # equal to the spike count of the corresponding spike train, hence + # the integral over the pooled rate equals the mean spike count. + mean_spike_count = np.mean([len(st) for st in spiketrains]) + area_under_curve = spint.cumulative_trapezoid( + y=rate_pooled.magnitude[:, 0], + x=rate_pooled.times.rescale('s').magnitude)[-1] + self.assertAlmostEqual(mean_spike_count, area_under_curve, + delta=0.01 * mean_spike_count) + + # The same spike trains wrapped in a Trials object already take this + # order, the trials branch estimates the rates per spike train and + # averages the corrected result afterwards. Both routes have to give + # the same answer. + rate_from_trials = statistics.instantaneous_rate( + TrialsFromLists([spiketrains]), pool_trials=False, + pool_spike_trains=True, **kwargs)[0] + assert_array_almost_equal(rate_from_trials.magnitude, + rate_pooled.magnitude) + def test_instantaneous_rate_trials_pool_trials(self): # Input: # Trials object with self.n_trials, self.n_spiketrains