+ addition for piece wise linear function with tests

1 files changed, 103 insertions, 3 deletions

diff --git a/pyspike/function.py b/pyspike/function.py
index ef6b0f1..b705293 100644
--- a/pyspike/function.py
+++ b/pyspike/function.py
@@ -10,6 +10,9 @@ from __future__ import print_function
 
 import numpy as np
 
+##############################################################
+# PieceWiseConstFunc
+##############################################################
 class PieceWiseConstFunc:
     """ A class representing a piece-wise constant function. """
     
@@ -106,6 +109,10 @@ class PieceWiseConstFunc:
         self.x = x_new[:index+2]
         self.y = y_new[:index+1]
 
+
+##############################################################
+# PieceWiseLinFunc
+##############################################################
 class PieceWiseLinFunc:
     """ A class representing a piece-wise linear function. """
     
@@ -119,9 +126,9 @@ class PieceWiseLinFunc:
         - y2: array of length N defining the function values at the right of the 
         intervals.
         """
-        self.x = x
-        self.y1 = y1
-        self.y2 = y2
+        self.x = np.array(x)
+        self.y1 = np.array(y1)
+        self.y2 = np.array(y2)
 
     def get_plottable_data(self):
         """ Returns two arrays containing x- and y-coordinates for immeditate
@@ -136,3 +143,96 @@ class PieceWiseLinFunc:
         y_plot[1::2] = self.y2
         return x_plot, y_plot
 
+    def avrg(self):
+        """ Computes the average of the piece-wise linear function:
+        a = 1/T int f(x) dx where T is the length of the interval.
+        Returns:
+        - the average a. 
+        """
+        return np.sum((self.x[1:]-self.x[:-1]) * 0.5*(self.y1+self.y2)) / \
+            (self.x[-1]-self.x[0])
+
+    def abs_avrg(self):
+        """ Computes the absolute average of the piece-wise linear function:
+        a = 1/T int |f(x)| dx where T is the length of the interval.
+        Returns:
+        - the average a. 
+        """
+        return np.sum((self.x[1:]-self.x[:-1]) * 0.5 *
+                      (np.abs(self.y1)+np.abs(self.y2)))/(self.x[-1]-self.x[0])
+
+    def add(self, f):
+        """ Adds another PieceWiseLin function to this function. 
+        Note: only functions defined on the same interval can be summed.
+        Params:
+        - f: PieceWiseLin function to be added.
+        """
+        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))
+        y1_new = np.empty(len(x_new)-1)
+        y2_new = np.empty_like(y1_new)
+        x_new[0] = self.x[0]
+        y1_new[0] = self.y1[0] + f.y1[0]
+        index1 = 0 # index for self
+        index2 = 0 # index for f
+        index = 0  # index for new
+        while (index1+1 < len(self.y1)) and (index2+1 < len(f.y1)):
+            # print(index1+1, self.x[index1+1], self.y[index1+1], x_new[index])
+            if self.x[index1+1] < f.x[index2+1]:
+                # first compute the end value of the previous interval
+                # linear interpolation of the interval
+                y = f.y1[index2] + (f.y2[index2]-f.y1[index2]) * \
+                    (self.x[index1+1]-f.x[index2]) / (f.x[index2+1]-f.x[index2])
+                y2_new[index] = self.y2[index1] + y
+                index1 += 1
+                index += 1
+                x_new[index] = self.x[index1]
+                # and the starting value for the next interval
+                y1_new[index] = self.y1[index1] + y
+            elif self.x[index1+1] > f.x[index2+1]:
+                # first compute the end value of the previous interval
+                # linear interpolation of the interval
+                y = self.y1[index1] + (self.y2[index1]-self.y1[index1]) * \
+                    (f.x[index2+1]-self.x[index1]) / \
+                    (self.x[index1+1]-self.x[index1])
+                y2_new[index] = f.y2[index2] + y
+                index2 += 1
+                index += 1
+                x_new[index] = f.x[index2]
+                # and the starting value for the next interval
+                y1_new[index] = f.y1[index2] + y
+            else: # self.x[index1+1] == f.x[index2+1]:
+                y2_new[index] = self.y2[index1] + f.y2[index2]
+                index1 += 1
+                index2 += 1
+                index += 1
+                x_new[index] = self.x[index1]
+                y1_new[index] = self.y1[index1] + f.y1[index2]
+        # one array reached the end -> copy the contents of the other to the end
+        if index1+1 < len(self.y1):
+            # compute the linear interpolations values
+            y = f.y1[index2] + (f.y2[index2]-f.y1[index2]) * \
+                (self.x[index1+1:-1]-f.x[index2]) / (f.x[index2+1]-f.x[index2])
+            x_new[index+1:index+1+len(self.x)-index1-1] = self.x[index1+1:]
+            y1_new[index+1:index+1+len(self.y1)-index1-1] = self.y1[index1+1:]+y
+            y2_new[index:index+len(self.y2)-index1-1] = self.y2[index1:-1] + y
+            index += len(self.x)-index1-2
+        elif index2+1 < len(f.y1):
+            # compute the linear interpolations values
+            y = self.y1[index1] + (self.y2[index1]-self.y1[index1]) * \
+                (f.x[index2+1:-1]-self.x[index1]) / \
+                (self.x[index1+1]-self.x[index1])
+            x_new[index+1:index+1+len(f.x)-index2-1] = f.x[index2+1:]
+            y1_new[index+1:index+1+len(f.y1)-index2-1] = f.y1[index2+1:] + y
+            y2_new[index:index+len(f.y2)-index2-1] = f.y2[index2:-1] + y
+            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]
+        # finally, the end value for the last interval
+        y2_new[index] = self.y2[-1]+f.y2[-1]
+        # only use the data that was actually filled
+        self.x = x_new[:index+2]
+        self.y1 = y1_new[:index+1]
+        self.y2 = y2_new[:index+1]