scripts/plotsto.py: Add script to plot .sto file

1 files changed, 144 insertions, 0 deletions

diff --git a/scripts/plotsto.py b/scripts/plotsto.py
new file mode 100755
index 0000000..01c76f4
--- /dev/null
+++ b/scripts/plotsto.py
@@ -0,0 +1,144 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# plotsto.py -- Plot all muscle activation signals in a .sto file
+#
+# Copyright (C) 2013 Tobias Klauser <tklauser@distanz.ch>
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License version 2 as
+# published by the Free Software Foundation.
+
+import getopt
+import os, sys
+import numpy as np
+import matplotlib.pyplot as plt
+
+def usage():
+    print("""usage: {} [OPTION...] STO-FILE...
+
+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):
+    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
+
+    print("name: " + name)
+    print("version: " + str(version))
+    print("nRows: " + str(nRows))
+    print("nColumns: " + str(nColumns))
+
+    cols = f.readline().strip().split()
+    if len(cols) != nColumns:
+        f.close()
+        print("Error: nColumns does not match number of columns in file")
+        return None
+
+    i = 0
+    for line in f:
+        i += 1
+
+    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()
+    i = 0
+    for line in f:
+        data = line.split()
+        if len(data) != nColumns:
+            f.close()
+            print("Error: row {} has less than {} columns".format(i+7, nColumns))
+            return None
+        for j, val in enumerate(data):
+            act[i,j] = float(val)
+        i += 1
+    f.close()
+
+    return act, cols
+
+def main():
+    try:
+        opts, args = getopt.getopt(sys.argv[1:], "h")
+    except getopt.GetoptError, err:
+        print(str(err))
+        usage()
+        sys.exit(-1)
+
+    if len(args) < 1:
+        usage()
+        sys.exit(-1)
+
+    for o, a in opts:
+        if o == '-h':
+            usage()
+            sys.exit(0)
+        else:
+            assert False, "unhandled option"
+
+    for sto in args:
+        if not os.path.exists(sto):
+            print "Error: File %s not found, skipping" % sto
+            continue
+
+        act, cols = read_sto(sto)
+        # transpose so we can access it more conveniently
+        act = act.transpose()
+        t, ys = act[0], act[1:]
+        l = []
+        # see matplotlib.org/examples/pylab_examples/subplots.demo.html
+        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)
+            l.append(cols[i+1])
+
+        plt.axis([0, t.max(), 0, 1.0], 'equal')
+    #    plt.legend(l, bbox_to_anchor=(1.05, 1.), loc=2, borderaxespad=0.)
+        plt.show()
+
+if __name__ == '__main__':
+    main()