added test for functions, bug fixes in add

1 files changed, 29 insertions, 14 deletions

diff --git a/pyspike/function.py b/pyspike/function.py
index 7177a3d..ef6b0f1 100644
--- a/pyspike/function.py
+++ b/pyspike/function.py
@@ -20,8 +20,8 @@ class PieceWiseConstFunc:
         function.
         - y: array of length N defining the function values at the intervals.
         """
-        self.x = x
-        self.y = y
+        self.x = np.array(x)
+        self.y = np.array(y)
 
     def get_plottable_data(self):
         """ Returns two arrays containing x- and y-coordinates for immeditate
@@ -66,29 +66,44 @@ class PieceWiseConstFunc:
         assert self.x[0] == f.x[0], "The functions have different intervals"
         assert self.x[-1] == f.x[-1], "The functions have different intervals"
         x_new = np.empty(len(self.x) + len(f.x))
-        y_new = np.empty_like(x_new)
+        y_new = np.empty(len(x_new)-1)
         x_new[0] = self.x[0]
         y_new[0] = self.y[0] + f.y[0]
-        index1 = 1
-        index2 = 1
-        index = 1
-        while (index1+1 < len(self.x)) and (index2+1 < len(f.x)):
+        index1 = 0
+        index2 = 0
+        index = 0
+        while (index1+1 < len(self.y)) and (index2+1 < len(f.y)):
+            index += 1
+            # print(index1+1, self.x[index1+1], self.y[index1+1], x_new[index])
             if self.x[index1+1] < f.x[index2+1]:
-                x_new[index] = self.x[index1]
                 index1 += 1
+                x_new[index] = self.x[index1]
             elif self.x[index1+1] > f.x[index2+1]:
-                x_new[index] = f.x[index2+1]
                 index2 += 1
+                x_new[index] = f.x[index2]
             else: # self.x[index1+1] == f.x[index2+1]:
-                x_new[index] = self.x[index1]
                 index1 += 1
                 index2 += 1
-            index += 1
+                x_new[index] = self.x[index1]
             y_new[index] = self.y[index1] + f.y[index2]
-        # both indices should have reached the maximum simultaneously
-        assert (index1+1 == len(self.x)) and (index2+1 == len(f.x))
+        # one array reached the end -> copy the contents of the other to the end
+        if index1+1 < len(self.y):
+            x_new[index+1:index+1+len(self.x)-index1-1] = self.x[index1+1:]
+            y_new[index+1:index+1+len(self.y)-index1-1] = self.y[index1+1:] + \
+                                                          f.y[-1]
+            index += len(self.x)-index1-2
+        elif index2+1 < len(f.y):
+            x_new[index+1:index+1+len(f.x)-index2-1] = f.x[index2+1:]
+            y_new[index+1:index+1+len(f.y)-index2-1] = f.y[index2+1:] + \
+                                                       self.y[-1]
+            index += len(f.x)-index2-2
+        else: # both arrays reached the end simultaneously
+            # only the last x-value missing
+            x_new[index+1] = self.x[-1]
+        # the last value is again the end of the interval
+        # x_new[index+1] = self.x[-1]
         # only use the data that was actually filled
-        self.x = x_new[:index+1]
+        self.x = x_new[:index+2]
         self.y = y_new[:index+1]
 
 class PieceWiseLinFunc: