diff options
Diffstat (limited to 'scripts')
| -rwxr-xr-x | scripts/plotsto.py | 133 |
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') |