added possibility to specify formula with array return type

added support for numpy

processing/post/addCalculation.py

@@ -1,8 +1,12 @@
 #!/usr/bin/env python
 
-import os,re,sys,math,string,damask
+import os,re,sys,math,string,damask,numpy
 from optparse import OptionParser, Option
 
+def unravel(item):
+  if hasattr(item,'__contains__'): return ' '.join(map(unravel,item))
+  else: return str(item)
+
 # -----------------------------
 class extendableOption(Option):
 # -----------------------------
@@ -40,7 +44,6 @@ parser.add_option('-l','--label',   dest='labels', action='extend', type='string
                                      help='(list of) new column labels', metavar='<LIST>')
 parser.add_option('-f','--formula',  dest='formulas', action='extend', type='string', \
                                      help='(list of) formulas corresponding to labels', metavar='<LIST>')
-
 parser.set_defaults(labels= [])
 parser.set_defaults(formulas= [])
 
@@ -76,11 +79,9 @@ for file in files:
   table.info_append(string.replace('$Id$','\n','\\n') + \
                     '\t' + ' '.join(sys.argv[1:]))
 
-  column = {}
   evaluator = {}
 
   for label,formula in zip(options.labels,options.formulas):
-    table.labels_append(label)
     interpolator = []
     for position,operand in enumerate(set(re.findall(r'#(.+?)#',formula))):
       formula = formula.replace('#'+operand+'#','{%i}'%position)
@@ -99,18 +100,33 @@ for file in files:
   
     evaluator[label] = "'" + formula + "'.format(" + ','.join(interpolator) + ")"
 
+# ------------------------------------------ calculate one result to get length of labels  ------
+  labelLen = {}
+  table.data_read()
+  for label in options.labels:
+    labelLen[label] = numpy.size(eval(eval(evaluator[label])))
+    print label, labelLen[label]
 
 # ------------------------------------------ assemble header ---------------------------------------  
+  for label,formula in zip(options.labels,options.formulas):
+    if labelLen[label] == 0:
+      print 'label ',label,' has length 0'
+      sys.exit()
+    elif labelLen[label] == 1:
+      table.labels_append(label)
+    else:
+      table.labels_append(['%i_%s'%(i+1,label) for i in xrange(labelLen[label])])
 
   table.head_write()
 
 # ------------------------------------------ process data ---------------------------------------  
 
   outputAlive = True
+  table.data_rewind()
   while outputAlive and table.data_read():                                  # read next data line of ASCII table
 
     specials['_row_'] += 1                                                  # count row
-    for label in options.labels: table.data_append(eval(eval(evaluator[label])))
+    for label in options.labels: table.data_append(unravel(eval(eval(evaluator[label]))))
     outputAlive = table.data_write()                                        # output processed line
 
 # ------------------------------------------ output result ---------------------------------------