+mul_scalar, tests restructured and cosmetics

4 files changed, 74 insertions, 0 deletions

diff --git a/pyspike/distances.py b/pyspike/distances.py
index 10b1d3c..f4be625 100644
--- a/pyspike/distances.py
+++ b/pyspike/distances.py
@@ -9,6 +9,10 @@ import numpy as np
 
 from pyspike import PieceWiseConstFunc, PieceWiseLinFunc
 
+
+############################################################
+# add_auxiliary_spikes
+############################################################
 def add_auxiliary_spikes( spike_train, T_end , T_start=0.0):
     """ Adds spikes at the beginning (T_start) and end (T_end) of the 
     observation interval.
@@ -29,6 +33,10 @@ def add_auxiliary_spikes( spike_train, T_end , T_start=0.0):
         spike_train = np.append(spike_train, T_end)
     return spike_train
 
+
+############################################################
+# isi_distance
+############################################################
 def isi_distance(spikes1, spikes2):
     """ Computes the instantaneous isi-distance S_isi (t) of the two given 
     spike trains. The spike trains are expected to have auxiliary spikes at the
@@ -95,6 +103,9 @@ def isi_distance(spikes1, spikes2):
     return PieceWiseConstFunc(spike_events[:index+1], isi_values[:index])
 
 
+############################################################
+# get_min_dist
+############################################################
 def get_min_dist(spike_time, spike_train, start_index=0):
     """ Returns the minimal distance |spike_time - spike_train[i]| 
     with i>=start_index.
@@ -111,6 +122,9 @@ def get_min_dist(spike_time, spike_train, start_index=0):
     return d
 
 
+############################################################
+# spike_distance
+############################################################
 def spike_distance(spikes1, spikes2):
     """ Computes the instantaneous spike-distance S_spike (t) of the two given
     spike trains. The spike trains are expected to have auxiliary spikes at the
diff --git a/pyspike/function.py b/pyspike/function.py
index b705293..3a5a01c 100644
--- a/pyspike/function.py
+++ b/pyspike/function.py
@@ -109,6 +109,13 @@ class PieceWiseConstFunc:
         self.x = x_new[:index+2]
         self.y = y_new[:index+1]
 
+    def mul_scalar(self, fac):
+        """ Multiplies the function with a scalar value
+        Params:
+        - fac: Value to multiply
+        """
+        self.y *= fac
+
 
 ##############################################################
 # PieceWiseLinFunc
@@ -236,3 +243,11 @@ class PieceWiseLinFunc:
         self.x = x_new[:index+2]
         self.y1 = y1_new[:index+1]
         self.y2 = y2_new[:index+1]
+
+    def mul_scalar(self, fac):
+        """ Multiplies the function with a scalar value
+        Params:
+        - fac: Value to multiply
+        """
+        self.y1 *= fac
+        self.y2 *= fac
diff --git a/pyspike/spikes.py b/pyspike/spikes.py
index 6b2eea3..70b48ff 100644
--- a/pyspike/spikes.py
+++ b/pyspike/spikes.py
@@ -7,6 +7,9 @@ Copyright 2014, Mario Mulansky <mario.mulansky@gmx.net>
 
 import numpy as np
 
+############################################################
+# spike_train_from_string
+############################################################
 def spike_train_from_string(s, sep=' '):
     """ Converts a string of times into an array of spike times.
     Params:
@@ -18,6 +21,9 @@ def spike_train_from_string(s, sep=' '):
     return np.fromstring(s, sep=sep)
 
 
+############################################################
+# merge_spike_trains
+############################################################
 def merge_spike_trains(spike_trains):
     """ Merges a number of spike trains into a single spike train.
     Params:
diff --git a/test/test_function.py b/test/test_function.py
index 014ecac..7420011 100644
--- a/test/test_function.py
+++ b/test/test_function.py
@@ -29,6 +29,13 @@ def test_pwc():
     assert_almost_equal(f.abs_avrg(), (1.0+0.5+0.5*1.5+1.5*0.75)/4.0,
                         decimal=16)
 
+
+def test_pwc_add():
+    # some random data
+    x = [0.0, 1.0, 2.0, 2.5, 4.0]
+    y = [1.0, -0.5, 1.5, 0.75]
+    f = spk.PieceWiseConstFunc(x, y)
+
     f1 = copy(f)
     x = [0.0, 0.75, 2.0, 2.5, 2.7, 4.0]
     y = [0.5, 1.0, -0.25, 0.0, 1.5]
@@ -48,6 +55,17 @@ def test_pwc():
     assert_array_almost_equal(f1.x, f2.x, decimal=16)
     assert_array_almost_equal(f1.y, 2*f2.y, decimal=16)
 
+def test_pwc_mul():
+    x = [0.0, 1.0, 2.0, 2.5, 4.0]
+    y = [1.0, -0.5, 1.5, 0.75]
+    f = spk.PieceWiseConstFunc(x, y)
+    
+    f.mul_scalar(1.5)
+    assert_array_almost_equal(f.x, x, decimal=16)
+    assert_array_almost_equal(f.y, 1.5*np.array(y), decimal=16)
+    f.mul_scalar(1.0/5.0)
+    assert_array_almost_equal(f.y, 1.5/5.0*np.array(y), decimal=16)
+
 
 def test_pwl():
     x = [0.0, 1.0, 2.0, 2.5, 4.0]
@@ -67,6 +85,13 @@ def test_pwl():
     abs_avrg_expected = (1.25 + 0.45 + 0.75 + 1.5*0.5) / 4.0
     assert_almost_equal(f.abs_avrg(), abs_avrg_expected, decimal=16)
 
+
+def test_pwl_add():
+    x = [0.0, 1.0, 2.0, 2.5, 4.0]
+    y1 = [1.0, -0.5, 1.5, 0.75]
+    y2 = [1.5, -0.4, 1.5, 0.25]
+    f = spk.PieceWiseLinFunc(x, y1, y2)
+
     f1 = copy(f)
     x = [0.0, 0.75, 2.0, 2.5, 2.7, 4.0]
     y1 = [0.5, 1.0, -0.25, 0.0, 1.5]
@@ -94,5 +119,19 @@ def test_pwl():
     assert_array_almost_equal(f1.y2, 2*f2.y2, decimal=16)
 
 
+def test_pwc_mul():
+    x = [0.0, 1.0, 2.0, 2.5, 4.0]
+    y1 = [1.0, -0.5, 1.5, 0.75]
+    y2 = [1.5, -0.4, 1.5, 0.25]
+    f = spk.PieceWiseLinFunc(x, y1, y2)
+    
+    f.mul_scalar(1.5)
+    assert_array_almost_equal(f.x, x, decimal=16)
+    assert_array_almost_equal(f.y1, 1.5*np.array(y1), decimal=16)
+    assert_array_almost_equal(f.y2, 1.5*np.array(y2), decimal=16)
+    f.mul_scalar(1.0/5.0)
+    assert_array_almost_equal(f.y1, 1.5/5.0*np.array(y1), decimal=16)
+    assert_array_almost_equal(f.y2, 1.5/5.0*np.array(y2), decimal=16)
+
 if __name__ == "__main__":
     test_pwc()