added spike generation function

3 files changed, 56 insertions, 7 deletions

diff --git a/pyspike/__init__.py b/pyspike/__init__.py
index d2d5b57..d700e7a 100644
--- a/pyspike/__init__.py
+++ b/pyspike/__init__.py
@@ -12,4 +12,4 @@ from distances import isi_profile, isi_distance, \
     isi_profile_multi, isi_distance_multi, isi_distance_matrix, \
     spike_profile_multi, spike_distance_multi, spike_distance_matrix
 from spikes import add_auxiliary_spikes, load_spike_trains_from_txt, \
-    spike_train_from_string, merge_spike_trains
+    spike_train_from_string, merge_spike_trains, generate_poisson_spikes
diff --git a/pyspike/cython_distance.pyx b/pyspike/cython_distance.pyx
index 178fcba..779ff94 100644
--- a/pyspike/cython_distance.pyx
+++ b/pyspike/cython_distance.pyx
@@ -123,6 +123,15 @@ cdef inline double get_min_dist_cython(double spike_time,
 
 
 ############################################################
+# isi_avrg_cython
+############################################################
+cdef inline double isi_avrg_cython(double isi1, double isi2) nogil:
+    return 0.5*(isi1+isi2)*(isi1+isi2)
+    # alternative definition to obtain <S> ~ 0.5 for Poisson spikes
+    # return 0.5*(isi1*isi1+isi2*isi2)
+
+
+############################################################
 # spike_distance_cython
 ############################################################
 def spike_distance_cython(double[:] t1, 
@@ -155,7 +164,7 @@ def spike_distance_cython(double[:] t1,
         isi2 = max(t2[1]-t2[0], t2[2]-t2[1])
         s1 = dt_f1*(t1[1]-t1[0])/isi1
         s2 = dt_f2*(t2[1]-t2[0])/isi2
-        y_starts[0] = (s1*isi2 + s2*isi1) / ((isi1+isi2)**2/2)
+        y_starts[0] = (s1*isi2 + s2*isi1) / isi_avrg_cython(isi1, isi2)
         while True:
             # print(index, index1, index2)
             if t1[index1+1] < t2[index2+1]:
@@ -169,12 +178,12 @@ def spike_distance_cython(double[:] t1,
                 s1 = dt_p1
                 s2 = (dt_p2*(t2[index2+1]-t1[index1]) + 
                       dt_f2*(t1[index1]-t2[index2])) / isi2
-                y_ends[index-1] = (s1*isi2 + s2*isi1)/(0.5*(isi1+isi2)*(isi1+isi2))
+                y_ends[index-1] = (s1*isi2 + s2*isi1)/isi_avrg_cython(isi1, isi2)
                 # now the next interval start value
                 dt_f1 = get_min_dist_cython(t1[index1+1], t2, N2, index2)
                 isi1 = t1[index1+1]-t1[index1]
                 # s2 is the same as above, thus we can compute y2 immediately
-                y_starts[index] = (s1*isi2 + s2*isi1)/(0.5*(isi1+isi2)*(isi1+isi2))
+                y_starts[index] = (s1*isi2 + s2*isi1)/isi_avrg_cython(isi1, isi2)
             elif t1[index1+1] > t2[index2+1]:
                 index2 += 1
                 if index2+1 >= N2:
@@ -185,13 +194,13 @@ def spike_distance_cython(double[:] t1,
                 s1 = (dt_p1*(t1[index1+1]-t2[index2]) + 
                       dt_f1*(t2[index2]-t1[index1])) / isi1
                 s2 = dt_p2
-                y_ends[index-1] = (s1*isi2 + s2*isi1) / (0.5*(isi1+isi2)*(isi1+isi2))
+                y_ends[index-1] = (s1*isi2 + s2*isi1) / isi_avrg_cython(isi1, isi2)
                 # now the next interval start value
                 dt_f2 = get_min_dist_cython(t2[index2+1], t1, N1, index1)
                 #s2 = dt_f2
                 isi2 = t2[index2+1]-t2[index2]
                 # s2 is the same as above, thus we can compute y2 immediately
-                y_starts[index] = (s1*isi2 + s2*isi1)/(0.5*(isi1+isi2)*(isi1+isi2))
+                y_starts[index] = (s1*isi2 + s2*isi1)/isi_avrg_cython(isi1, isi2)
             else: # t1[index1+1] == t2[index2+1] - generate only one event
                 index1 += 1
                 index2 += 1
@@ -214,7 +223,7 @@ def spike_distance_cython(double[:] t1,
         isi2 = max(t2[N2-1]-t2[N2-2], t2[N2-2]-t2[N2-3])
         s1 = dt_p1*(t1[N1-1]-t1[N1-2])/isi1
         s2 = dt_p2*(t2[N2-1]-t2[N2-2])/isi2
-        y_ends[index-1] = (s1*isi2 + s2*isi1) / (0.5*(isi1+isi2)*(isi1+isi2))
+        y_ends[index-1] = (s1*isi2 + s2*isi1) / isi_avrg_cython(isi1, isi2)
     # end nogil
 
     # use only the data added above 
diff --git a/pyspike/spikes.py b/pyspike/spikes.py
index f7172c9..aa25c48 100644
--- a/pyspike/spikes.py
+++ b/pyspike/spikes.py
@@ -129,3 +129,43 @@ def merge_spike_trains(spike_trains):
             index_list = index_list[index_list != i]
         vals = [spike_trains[n][indices[n]] for n in index_list]
     return merged_spikes
+
+
+############################################################
+# generate_poisson_spikes
+############################################################
+def generate_poisson_spikes(rate, time_interval, add_aux_spikes=True):
+    """ Generates a Poisson spike train with the given rate in the given time
+    interval
+
+    :param rate: The rate of the spike trains
+    :param time_interval: A pair (T_start, T_end) of values representing the
+                          start and end time of the spike train measurement or
+                          a single value representing the end time, the T_start
+                          is then assuemd as 0. Auxiliary spikes will be added
+                          to the spike train at the beginning and end of this
+                          interval, if they are not yet present.
+    :type time_interval: pair of doubles or double
+    :returns: Poisson spike train
+    """
+    try:
+        T_start = time_interval[0]
+        T_end = time_interval[1]
+    except:
+        T_start = 0
+        T_end = time_interval
+    # roughly how many spikes are required to fill the interval
+    N = max(1, int(1.2 * rate * (T_end-T_start)))
+    N_append = max(1, int(0.1 * rate * (T_end-T_start)))
+    intervals = np.random.exponential(1.0/rate, N)
+    # make sure we have enough spikes
+    while T_start + sum(intervals) < T_end:
+        print T_start + sum(intervals)
+        intervals = np.append(intervals,
+                              np.random.exponential(1.0/rate, N_append))
+    spikes = T_start + np.cumsum(intervals)
+    spikes = spikes[spikes < T_end]
+    if add_aux_spikes:
+        return add_auxiliary_spikes(spikes, time_interval)
+    else:
+        return spikes