csv2sto: update to merge multiple CSV files

Still same time spacing is needed...

1 files changed, 72 insertions, 33 deletions

diff --git a/scripts/csv2sto.py b/scripts/csv2sto.py
index 8e5300b..5515e84 100755
--- a/scripts/csv2sto.py
+++ b/scripts/csv2sto.py
@@ -3,19 +3,22 @@
 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)
-  -h  show this help and exit""" % (DEFAULT_NAME, os.path.basename(sys.argv[0]))
+  -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:], "fh")
+        opts, args = getopt.getopt(sys.argv[1:], "fn:hT:")
     except getopt.GetoptError, err:
         print(str(err))
         usage()
@@ -24,61 +27,97 @@ def main():
     if len(args) < 2:
         usage()
         sys.exit(-1)
-    elif len(args) > 2:
-        print "Merging CSV files currently not supported, only %s considered" % args[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"
 
-    f_csv = args[0]
+    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 os.path.exists(f_csv):
-        print "Error: CSV file %s does not exist" % f_csv
-        sys.exit(-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
-    fd = open(f_csv, 'r')
-    nRows = nColumns = 0
-    for line in fd:
-        if nRows == 0:
-            nCols = len(line.split(','))
-        nRows += 1
-    fd.close()
+    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()
 
     fd_sto = open(f_sto, 'w')
     fd_sto.write(name + "\n")
-    fd_sto.write("nRows=" + str(nRows) + "\n")
-    fd_sto.write("nColumns=" + str(nColumns) + "\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")
 
     print "Writing sto file %s with name %s" % (f_sto, name)
-
-    colname = os.path.splitext(f_csv)[0]
-    print "[+] Adding muscle %s..." % colname
-    fd_sto.write("time\t" + colname + "\n")
-
-    fd_csv = open(f_csv, 'r')
-    cr = csv.reader(fd_csv, delimiter=',')
-
-    for i, row in enumerate(cr):
-        if i == 0:
-            continue    # skip header line
-        fd_sto.write("\t".join(row) + "\n")
-
-    fd_csv.close()
+    if maxtime != DEFAULT_MAXTIME:
+        print "Scaling to maxtime %f" % maxtime
+
+    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__':