Timescale function option to return nan if there are few spikes (<=2) (…

…NeuralEnsemble#328)

* include option to return np.nan if spiketrains are too short and change raised error to more meaningful ones

elephant/spike_train_correlation.py

@@ -924,12 +924,14 @@ def spike_train_timescale(binned_st, tau_max):
     binned_st : elephant.conversion.BinnedSpikeTrain
         A binned spike train containing the spike train to be evaluated.
     tau_max : pq.Quantity
-        Maximal integration time of the auto-correlation function.
+        Maximal integration time of the auto-correlation function. It needs to
+        be a multiple of the binsize of `binned_st`.
 
     Returns
     -------
     timescale : pq.Quantity
-        The auto-correlation time of the binned spiketrain.
+        The auto-correlation time of the binned spiketrain. If `binned_st` has
+        less than 2 spikes `np.nan` is returned.
 
     Notes
     -----
@@ -949,15 +951,20 @@ def spike_train_timescale(binned_st, tau_max):
         Slow fluctuations in recurrent networks of spiking neurons.
         Physical Review E, 92(4), 040901.
     """
+    if binned_st.get_num_of_spikes() < 2:
+        warnings.warn("Spike train contains less than 2 spikes! "
+                      "np.nan will be returned.")
+        return np.nan
+
     binsize = binned_st.binsize
     if not (tau_max / binsize).simplified.units == pq.dimensionless:
-        raise AssertionError("tau_max needs units of time")
+        raise ValueError("tau_max needs units of time")
 
     # safe casting of tau_max/binsize to integer
     tau_max_bins = int(np.round((tau_max / binsize).simplified.magnitude))
     if not np.isclose(tau_max.simplified.magnitude,
                       (tau_max_bins * binsize).simplified.magnitude):
-        raise AssertionError("tau_max has to be a multiple of the binsize")
+        raise ValueError("tau_max has to be a multiple of the binsize")
 
     cch_window = [-tau_max_bins, tau_max_bins]
     corrfct, bin_ids = cross_correlation_histogram(

elephant/test/test_spike_train_correlation.py

@@ -18,6 +18,7 @@
 import elephant.spike_train_correlation as sc
 from elephant.spike_train_generation import homogeneous_poisson_process,\
     homogeneous_gamma_process
+import math
 
 python_version_major = sys.version_info.major
 
@@ -785,6 +786,44 @@ def test_timescale_calculation(self):
         target = np.allclose(timescale.magnitude, timescale_num, rtol=2e-1)
         self.assertTrue(target)
 
+    def test_timescale_errors(self):
+        spikes = neo.SpikeTrain([1, 5, 7, 8]*pq.ms, t_stop=10*pq.ms)
+        binsize = 1 * pq.ms
+        spikes_bin = conv.BinnedSpikeTrain(spikes, binsize)
+
+        # Tau max with no units
+        tau_max = 1
+        self.assertRaises(ValueError,
+                          sc.spike_train_timescale, spikes_bin, tau_max)
+
+        # Tau max that is not a multiple of the binsize
+        tau_max = 1.1*pq.ms
+        self.assertRaises(ValueError,
+                          sc.spike_train_timescale, spikes_bin, tau_max)
+
+    @unittest.skipUnless(python_version_major == 3,
+                         "assertWarns requires python 3.2")
+    def test_timescale_nan(self):
+        st0 = neo.SpikeTrain([]*pq.ms, t_stop=10*pq.ms)
+        st1 = neo.SpikeTrain([1]*pq.ms, t_stop=10*pq.ms)
+        st2 = neo.SpikeTrain([1, 5]*pq.ms, t_stop=10*pq.ms)
+        st3 = neo.SpikeTrain([1, 5, 6]*pq.ms, t_stop=10*pq.ms)
+        st4 = neo.SpikeTrain([1, 5, 6, 9]*pq.ms, t_stop=10*pq.ms)
+
+        binsize = 1 * pq.ms
+        tau_max = 1 * pq.ms
+
+        for st in [st0, st1]:
+            bst = conv.BinnedSpikeTrain(st, binsize)
+            with self.assertWarns(UserWarning):
+                timescale = sc.spike_train_timescale(bst, tau_max)
+            self.assertTrue(math.isnan(timescale))
+
+        for st in [st2, st3, st4]:
+            bst = conv.BinnedSpikeTrain(st, binsize)
+            timescale = sc.spike_train_timescale(bst, tau_max)
+            self.assertFalse(math.isnan(timescale))
+
 
 if __name__ == '__main__':
     unittest.main()