diff --git a/processing/post/vtk_scalars2vectors.py b/processing/post/vtk_scalars2vectors.py new file mode 100755 index 000000000..a07c7bf92 --- /dev/null +++ b/processing/post/vtk_scalars2vectors.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python + +import os, sys, string, re, shutil +from optparse import OptionParser +from vtk import * + + +# ----------------------------- +# MAIN FUNCTION STARTS HERE +# ----------------------------- + +# --- input parsing + +parser = OptionParser(usage='%prog [options] vtkfile', description = """ +""" + string.replace('$Id: $','\n','\\n') +) + +parser.add_option('-v','--vector', nargs=3, dest='vector', type='string', \ + help='suffices indicating vector components [%default]') +parser.add_option('-s','--separator', dest='separator', type='string', \ + help='separator between label and suffix [%default]') + +parser.set_defaults(vector = ['x','y','z']) +parser.set_defaults(separator = '.') + +(options, filenames) = parser.parse_args() + + +# --- sanity checks + +if filenames == []: + parser.print_help() + parser.error('no file specified...') + +for filename in filenames: + if not os.path.isfile(filename): + parser.print_help() + parser.error('invalid file "%s" specified...'%filename) + + +# --- ITERATE OVER FILES AND PROCESS THEM + +for filename in filenames: + + # Read the source file + + sys.stdout.write('read file "%s" ...'%filename) + sys.stdout.flush() + suffix = os.path.splitext(filename)[1] + if suffix == '.vtk': + reader = vtkUnstructuredGridReader() + reader.ReadAllScalarsOn() + reader.ReadAllVectorsOn() + reader.ReadAllTensorsOn() + elif suffix == '.vtu': + reader = vtkXMLUnstructuredGridReader() + else: + parser.error('filetype "%s" not supported'%suffix) + reader.SetFileName(filename) + reader.Update() + uGrid = reader.GetOutput() + sys.stdout.write(' done\n') + sys.stdout.flush() + + + # Read the scalar data + + scalarData = {} + scalarsToBeRemoved = [] + Nscalars = uGrid.GetCellData().GetNumberOfArrays() + for i in range(Nscalars): + sys.stdout.write("\rread scalar data %d%%" %(100*i/Nscalars)) + sys.stdout.flush() + scalarName = uGrid.GetCellData().GetArrayName(i) + if scalarName.split(options.separator)[-1] in options.vector: + label,suffix = scalarName.split(options.separator) + if label not in scalarData: + scalarData[label] = [[],[],[]] + uGrid.GetCellData().SetActiveScalars(scalarName) + scalarData[label][options.vector.index(suffix)] = uGrid.GetCellData().GetScalars(scalarName) + scalarsToBeRemoved.append(scalarName) + for scalarName in scalarsToBeRemoved: + uGrid.GetCellData().RemoveArray(scalarName) + # uGrid.UpdateData() + sys.stdout.write('\rread scalar data done\n') + sys.stdout.flush() + + + # Convert the scalar data to vector data + + NscalarData = len(scalarData) + for n,label in enumerate(scalarData): + sys.stdout.write("\rconvert to vector data %d%%" %(100*n/NscalarData)) + sys.stdout.flush() + Nvalues = scalarData[label][0].GetNumberOfTuples() + vectorData = vtkDoubleArray() + vectorData.SetName(label) + vectorData.SetNumberOfComponents(3) # set this before NumberOfTuples !!! + vectorData.SetNumberOfTuples(Nvalues) + for i in range(Nvalues): + for j in range(3): + vectorData.SetComponent(i,j,scalarData[label][j].GetValue(i)) + uGrid.GetCellData().AddArray(vectorData) + # uGrid.GetCellData().SetActiveVectors(label) + sys.stdout.write('\rconvert to vector data done\n') + + + # Write to new vtk file + + outfilename = os.path.splitext(filename)[0]+'.vtu' + sys.stdout.write('write to file "%s" ...'%outfilename) + sys.stdout.flush() + writer = vtkXMLUnstructuredGridWriter() + writer.SetFileName(outfilename+'_tmp') + writer.SetDataModeToAscii() + writer.SetInput(uGrid) + writer.Write() + sys.stdout.write(' done\n') + sys.stdout.flush() + shutil.move(outfilename+'_tmp',outfilename) + + + +# --------------------------- DONE --------------------------------