scripts/plotsto.py: make more robust and allow to select muscles

1 files changed, 80 insertions, 53 deletions

diff --git a/scripts/plotsto.py b/scripts/plotsto.py
index 01c76f4..259df30 100755
--- a/scripts/plotsto.py
+++ b/scripts/plotsto.py
@@ -21,47 +21,51 @@ Plot all muscle activation signals in a .sto file.
 
 options:
 
-    -h  show this help and exit""".format(os.path.basename(sys.argv[0])))
-
-def read_sto(sto):
+  -m NAMES  plot only columns with names from comma-separated list NAMES
+  -h        show this help and exit""".format(os.path.basename(sys.argv[0])))
+
+# .sto header format:
+# control
+# version=N
+# nRows=N
+# nColumns=N
+# ... (optional header lines)
+# endheader
+def parse_sto_header(f):
+    endheader = False
+    name = ''
+    nRows = nColumns = 0
+
+    name = f.readline()
+
+    while True:
+        line = f.readline().strip()
+        if line is None:
+            break
+
+        if line == 'endheader':
+            endheader = True
+            break
+
+        k, v = line.split('=')
+        if k == 'nRows':
+            nRows = int(v)
+        elif k == 'nColumns':
+            nColumns = int(v)
+
+    if not endheader:
+        print("Error: header not properly terminated by 'endheader' in {}".format(f.name))
+        return -1, -1
+
+    return nRows, nColumns
+
+def read_sto(sto, muscles=None):
     f = open(sto, 'r')
 
-    # .sto header format:
-    # control
-    # version=N
-    # nRows=N
-    # nColumns=N
-    # endheader
-    name = f.readline().strip()
-    line = f.readline().strip().split('=')
-    if line[0] != 'version' or not line[1].isdigit():
-        f.close()
-        print("Error: invalid version header line in {}: {}".format(sto, "=".join(line)))
-        return None
-    version = int(line[1])
-
-    line = f.readline().strip().split('=')
-    if line [0] != 'nRows' or not line[1].isdigit():
-        f.close()
-        print("Error: invalid nRows header line in {}: {}".format(sto, "=".join(line)))
-        return None
-    nRows = int(line[1])
-
-    line = f.readline().strip().split('=')
-    if line[0] != 'nColumns' or not line[1].isdigit():
-        f.close()
-        print("Error: invalid nColumns header line in {}: {}".format(sto, "=".join(line)))
-        return None
-    nColumns = int(line[1])
-
-    endheader = f.readline().strip()
-    if endheader != 'endheader':
-        f.close()
-        print("Error: header not properly terminated by 'endheader' in {}".format(sto))
-        return None
+    nRows, nColumns = parse_sto_header(f)
+    if nRows <= 0 or nColumns <= 0:
+        return
 
-    print("name: " + name)
-    print("version: " + str(version))
     print("nRows: " + str(nRows))
     print("nColumns: " + str(nColumns))
 
@@ -71,31 +75,50 @@ def read_sto(sto):
         print("Error: nColumns does not match number of columns in file")
         return None
 
-    i = 0
-    for line in f:
-        i += 1
-
+    i = len(f.readlines())
     f.close()
 
     if i != nRows:
         print("Error: nRows does not match number of rows in file")
         return None
 
-    act = np.zeros((nRows, nColumns))
-
-    f = open(sto, 'r')
-    # skip first 6 lines (header)
-    for i in range(6):
-        _ = f.readline()
+    # by default use all columns
+    idx = list(range(nColumns))
+    _cols = []
+    # filter out selected muscles
+    if not muscles is None:
+        for i, m in enumerate(cols):
+            if m in muscles or m == 'time':
+                _cols.append(m)
+            else:
+                idx[i] = None
+        cols = _cols
+
+    act = np.zeros((nRows, len(cols)))
     i = 0
+    endheader = 0
+    f = open(sto, 'r')
     for line in f:
+        # skip header and column title lines (we know there's an endheader line)
+        if line.strip() == 'endheader':
+            endheader = 1
+            continue
+        if endheader == 0:
+            continue
+        elif endheader == 1:
+            endheader += 1
+            continue
+
         data = line.split()
         if len(data) != nColumns:
             f.close()
             print("Error: row {} has less than {} columns".format(i+7, nColumns))
             return None
+        k = 0
         for j, val in enumerate(data):
-            act[i,j] = float(val)
+            if j in idx:
+                act[i,k] = float(val.strip())
+                k += 1
         i += 1
     f.close()
 
@@ -103,7 +126,7 @@ def read_sto(sto):
 
 def main():
     try:
-        opts, args = getopt.getopt(sys.argv[1:], "h")
+        opts, args = getopt.getopt(sys.argv[1:], "m:h")
     except getopt.GetoptError, err:
         print(str(err))
         usage()
@@ -113,8 +136,11 @@ def main():
         usage()
         sys.exit(-1)
 
+    muscles = None
     for o, a in opts:
-        if o == '-h':
+        if o == '-m':
+            muscles = [ x.strip() for x in a.split(',') ]
+        elif o == '-h':
             usage()
             sys.exit(0)
         else:
@@ -122,10 +148,10 @@ def main():
 
     for sto in args:
         if not os.path.exists(sto):
-            print "Error: File %s not found, skipping" % sto
+            print("Error: File %s not found, skipping".format(sto))
             continue
 
-        act, cols = read_sto(sto)
+        act, cols = read_sto(sto, muscles)
         # transpose so we can access it more conveniently
         act = act.transpose()
         t, ys = act[0], act[1:]
@@ -134,6 +160,7 @@ def main():
         f, ax = plt.subplots(nrows=(len(cols) - 1) / 2, ncols=2, sharex='col', sharey='row')
         for i, y in enumerate(ys):
             ax[i/2,i%2].plot(t, y)
+            ax[i/2,i%2].set_title(cols[i+1])
             l.append(cols[i+1])
 
         plt.axis([0, t.max(), 0, 1.0], 'equal')