scripts: plotcsv.py: Move plotting functionality to own function

This way we can call it from random-muscle-activation.py to inspect
generated data directly.

1 files changed, 48 insertions, 18 deletions

diff --git a/scripts/plotcsv.py b/scripts/plotcsv.py
index 931e4e8..fd6d084 100755
--- a/scripts/plotcsv.py
+++ b/scripts/plotcsv.py
@@ -11,35 +11,59 @@ X_MAX = 100
 Y_MAX = 50
 
 def usage():
-    print """usage: %s [OPTION...] CSV-FILE...""" % os.path.basename(sys.argv[0])
+    print("""usage: {} [OPTION...] CSV-FILE...
 
-def read_and_plot_csv(csv_file, x_max=X_MAX, y_max=Y_MAX):
-    reader = csv.reader(open(csv_file, 'r'), delimiter=',')
-    x_name,y_name = reader.next() # header line
+    -c  do cubic interpolation
+    -l  do linear interpolation
+    -h  show this help and exit""".format(os.path.basename(sys.argv[0])))
 
-    X = np.array([[float(_x), float(_y)] for _x,_y in reader ])
-    x = X[:,0]
-    y = X[:,1]
+def plot(x, y, x_max, y_max, x_name, y_name, title, cubic, linear, xnew=None, f=None, f2=None):
+    plt.plot(x, y, 'o')
+    if linear:
+        plt.plot(xnew, f(xnew), '-')
+    if cubic:
+        plt.plot(xnew, f2(xnew), '--')
 
-    # interpolate data points
-    xnew = np.linspace(min(x), max(x), len(x)*5)
-    f = scipy.interpolate.interp1d(x, y)
-    f2 = scipy.interpolate.interp1d(x, y, kind='cubic')
+    leg = [ 'data' ]
+    if linear:
+        leg.append('linear')
+    if cubic:
+        leg.append('cubic')
+    plt.legend(leg, loc='best')
 
-    fig = plt.figure()
-    plt.plot(x, y, 'o', xnew, f(xnew), '-', xnew, f2(xnew), '--')
-    plt.legend(['data', 'linear interp.', 'cubic interp.'], loc='best')
     plt.axis([0, x_max, 0, y_max], 'equal')
     plt.xlabel(x_name)
     plt.ylabel(y_name)
-    plt.title(csv_file)
+    plt.title(title)
     plt.grid(True)
 
     plt.show()
 
+def read_and_plot_csv(csv_file, cubic, linear, x_max=X_MAX, y_max=Y_MAX):
+    print("Reading data...")
+    reader = csv.reader(open(csv_file, 'r'), delimiter=',')
+    x_name,y_name = reader.next() # header line
+
+    X = np.array([[float(_x), float(_y)] for _x,_y in reader ])
+    x = X[:,0]
+    y = X[:,1]
+    xnew = f = f2 = None
+
+    # interpolate data points
+    if cubic or linear:
+        xnew = np.linspace(min(x), max(x), len(x)*5)
+    if linear:
+        print("Calculating linear interpolation...")
+        f = scipy.interpolate.interp1d(x, y)
+    if cubic:
+        print("Calculating cubic interpolation...")
+        f2 = scipy.interpolate.interp1d(x, y, kind='cubic')
+
+    plot(x, y, x_max, y_max, x_name, y_name, csv_file, cubic, linear, xnew, f, f2)
+
 def main():
     try:
-        opts, args = getopt.getopt(sys.argv[1:], "h")
+        opts, args = getopt.getopt(sys.argv[1:], "clh")
     except getopt.GetoptError, err:
         print(str(err))
         usage()
@@ -49,8 +73,14 @@ def main():
         usage()
         sys.exit(-1)
 
+    cubic = False
+    linear = False
     for o, a in opts:
-        if o == '-h':
+        if o == '-c':
+            cubic = True
+        elif o == '-l':
+            linear = True
+        elif o == '-h':
             usage()
             sys.exit(0)
         else:
@@ -60,7 +90,7 @@ def main():
         if not os.path.exists(csv):
             print "Error: File %s not found, skipping" % csv
             continue
-        read_and_plot_csv(csv)
+        read_and_plot_csv(csv, cubic, linear)
 
 if __name__ == '__main__':
     main()