summaryrefslogtreecommitdiff
path: root/scripts/csv2sto.py
blob: 9609a99335b573dae1365baecabd1ba7817292d8 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
#!/usr/bin/env python

import os, sys
import getopt
import csv
import numpy as np

DEFAULT_NAME = 'control'
DEFAULT_MAXTIME = 100.0

def usage():
    print """usage: %s [OPTION...] CSV-FILE STO-FILE

  -f    force overwrite of existing files
  -n    name of the sto file (default: %s)
  -T N  scale time such that it goes from 0.0 to N
  -h    show this help and exit""" % (DEFAULT_NAME, os.path.basename(sys.argv[0]))

def main():
    try:
        opts, args = getopt.getopt(sys.argv[1:], "fn:hT:")
    except getopt.GetoptError, err:
        print(str(err))
        usage()
        sys.exit(-1)

    if len(args) < 2:
        usage()
        sys.exit(-1)

    overwrite = False
    name = DEFAULT_NAME
    maxtime = DEFAULT_MAXTIME
    for o, a in opts:
        if o == '-f':
            overwrite = True
        elif o == '-n':
            name = a
        elif o == '-T':
            maxtime = float(a)
        elif o == '-h':
            usage()
            sys.exit(0)
        else:
            assert False, "unhandled option"

    fs_csv = args[0:-1]
    for f_csv in fs_csv:
        if not os.path.exists(f_csv):
            print "Error: CSV file %s does not exist" % f_csv
            sys.exit(-1)
    f_sto = args[-1]
    if not overwrite and os.path.exists(f_sto):
        print "Error: STO file %s already exists" % f_sto
        sys.exit(-1)

    # determine number of columns and rows
    tot_nRows = tot_nCols = 0
    for f_csv in fs_csv:
        fd = open(f_csv, 'r')
        nRows = nColumns = 0
        for line in fd:
            if nRows == 0:
                nCols = len(line.split(','))
            nRows += 1
        if tot_nRows == 0:
            tot_nRows = nRows
        elif tot_nRows != nRows:
            print "Error: Number of rows in CSV files does not match"
            sys.exit(-1)
        if tot_nCols == 0:
            tot_nCols = nCols
        elif tot_nCols != nCols:
            print "Error: Number of columns in CSV files does not match"
            sys.exit(-1)

        fd.close()

    print "Writing sto file %s with name %s" % (f_sto, name)
    if maxtime != DEFAULT_MAXTIME:
        print "Scaling to maxtime %f" % maxtime

    fd_sto = open(f_sto, 'w')
    fd_sto.write(name + "\n")
    fd_sto.write("version=1\n")
    fd_sto.write("nRows=" + str(tot_nRows) + "\n")
    fd_sto.write("nColumns=" + str(tot_nCols) + "\n")
    fd_sto.write("endheader\n")

    fd_sto.write("time")
    for f_csv in fs_csv:
        colname = os.path.splitext(os.path.basename(f_csv))[0]
        print "[+] Adding muscle %s..." % colname
        fd_sto.write("\t" + colname)
    fd_sto.write("\n")

    fds_csv = [ open(f_csv, 'r') for f_csv in fs_csv ]
    crs = [ csv.reader(fd_csv, delimiter=',') for fd_csv in fds_csv ]

    times = np.zeros((tot_nRows - 1, len(fds_csv)))
    vals = np.zeros((tot_nRows - 1, len(fds_csv)))
    for i, cr in enumerate(crs):
        cr.next()   # skip header line
        for j, row in enumerate(cr):
            t,y = map(float, row)
            if maxtime != DEFAULT_MAXTIME:
                # Assume default 0-100 time in all input files for now
                t = (t / 100.0) * maxtime
            times[j,i] = t
            vals[j,i] = y

    # XXX: For now assume same time spacing in all CSV files
    for i, row in enumerate(times):
        fd_sto.write(str(row[0]))
        for val in vals[i]:
            fd_sto.write("\t" + str(val))
        fd_sto.write("\n");

    for fd_csv in fds_csv:
        fd_csv.close()
    fd_sto.close()

if __name__ == '__main__':
    main()