scripts/csv2sto.py: Fix bugs regarding # of rows/columns in .sto

We wrote to few columns when combining multiple CSV files and were
always wrongly counting the header row as a data row for nRows in the
.sto file.

1 files changed, 14 insertions, 10 deletions

diff --git a/scripts/csv2sto.py b/scripts/csv2sto.py
index be8438e..4275fb1 100755
--- a/scripts/csv2sto.py
+++ b/scripts/csv2sto.py
@@ -72,28 +72,32 @@ def main():
         sys.exit(-1)
 
     # determine number of columns and rows
-    tot_nRows = tot_nCols = 0
+    tot_nRows = 0
+    tot_nCols = 1   # initial time column (as we omit it when taking data from each CSV)
     for f_csv in fs_csv:
         fd = open(f_csv, 'r')
-        nRows = nColumns = 0
+
+        nRows = nCols = 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 do not match")
             sys.exit(-1)
-        if tot_nCols == 0:
-            tot_nCols = nCols
-        elif tot_nCols != nCols:
-            print("Error: Number of columns in CSV files do not match")
-            sys.exit(-1)
+
+        # add up to the total number of colums but omit time column
+        tot_nCols += nCols - 1
 
         fd.close()
 
-    print("Writing sto file {} with name {}".format(f_sto, name))
+    # don't count the header line
+    tot_nRows -= 1
+
+    print("Writing sto file {} with name {} ({} muscles)".format(f_sto, name, len(fs_csv)))
     if maxtime != DEFAULT_MAXTIME:
         print("Scaling to maxtime {}".format(maxtime))
 
@@ -119,8 +123,8 @@ def main():
     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)))
+    times = np.zeros((tot_nRows, len(fds_csv)))
+    vals = np.zeros((tot_nRows, len(fds_csv)))
     for i, cr in enumerate(crs):
         cr.next()   # skip header line
         for j, row in enumerate(cr):