136 lines
5.0 KiB
Python
Executable File
136 lines
5.0 KiB
Python
Executable File
#!/usr/bin/env python
|
|
# -*- coding: UTF-8 no BOM -*-
|
|
|
|
import os,sys,vtk
|
|
import damask
|
|
from collections import defaultdict
|
|
from optparse import OptionParser
|
|
|
|
scriptName = os.path.splitext(os.path.basename(__file__))[0]
|
|
scriptID = ' '.join([scriptName,damask.version])
|
|
|
|
# --------------------------------------------------------------------
|
|
# MAIN
|
|
# --------------------------------------------------------------------
|
|
|
|
parser = OptionParser(option_class=damask.extendableOption, usage='%prog options [file[s]]', description = """
|
|
Add scalar and RGB tuples from ASCIItable to existing VTK voxel cloud (.vtu/.vtk).
|
|
|
|
""", version = scriptID)
|
|
|
|
parser.add_option('-v', '--vtk', dest='vtk', \
|
|
help = 'VTK file name')
|
|
parser.add_option('-s', '--scalar', dest='scalar', action='extend', \
|
|
help = 'scalar values')
|
|
parser.add_option('-c', '--color', dest='color', action='extend', \
|
|
help = 'RGB color tuples')
|
|
|
|
parser.set_defaults(scalar = [],
|
|
color = [],
|
|
render = False,
|
|
)
|
|
|
|
(options, filenames) = parser.parse_args()
|
|
|
|
if options.vtk == None or not os.path.exists(options.vtk):
|
|
parser.error('VTK file does not exist')
|
|
|
|
if os.path.splitext(options.vtk)[1] == '.vtu':
|
|
reader = vtk.vtkXMLUnstructuredGridReader()
|
|
reader.SetFileName(options.vtk)
|
|
reader.Update()
|
|
uGrid = reader.GetOutput()
|
|
elif os.path.splitext(options.vtk)[1] == '.vtk':
|
|
reader = vtk.vtkGenericDataObjectReader()
|
|
reader.SetFileName(options.vtk)
|
|
reader.Update()
|
|
uGrid = reader.GetUnstructuredGridOutput()
|
|
else:
|
|
parser.error('unsupported VTK file type extension')
|
|
|
|
Npoints = uGrid.GetNumberOfPoints()
|
|
Ncells = uGrid.GetNumberOfCells()
|
|
|
|
sys.stderr.write('{}: {} points and {} cells...\n'.format(damask.util.emph(options.vtk),Npoints,Ncells))
|
|
|
|
# --- loop over input files -------------------------------------------------------------------------
|
|
|
|
if filenames == []: filenames = [None]
|
|
|
|
for name in filenames:
|
|
try:
|
|
table = damask.ASCIItable(name = name,
|
|
buffered = False, readonly = True)
|
|
except: continue
|
|
damask.util.croak(damask.util.emph(scriptName)+(': '+name if name else ''))
|
|
|
|
# --- interpret header ----------------------------------------------------------------------------
|
|
|
|
table.head_read()
|
|
|
|
remarks = []
|
|
errors = []
|
|
VTKarray = {}
|
|
active = defaultdict(list)
|
|
|
|
for datatype,dimension,label in [['scalar',1,options.scalar],
|
|
['color',3,options.color],
|
|
]:
|
|
for i,dim in enumerate(table.label_dimension(label)):
|
|
me = label[i]
|
|
if dim == -1: remarks.append('{} "{}" not found...'.format(datatype,me))
|
|
elif dim > dimension: remarks.append('"{}" not of dimension{}...'.format(me,dimension))
|
|
else:
|
|
damask.util.croak('adding {} {}'.format(datatype,me))
|
|
active[datatype].append(me)
|
|
|
|
if datatype in ['scalar']:
|
|
VTKarray[me] = vtk.vtkDoubleArray()
|
|
elif datatype == 'color':
|
|
VTKarray[me] = vtk.vtkUnsignedCharArray()
|
|
|
|
VTKarray[me].SetNumberOfComponents(dimension)
|
|
VTKarray[me].SetName(label[i])
|
|
|
|
if remarks != []: damask.util.croak(remarks)
|
|
if errors != []:
|
|
damask.util.croak(errors)
|
|
table.close(dismiss=True)
|
|
continue
|
|
|
|
# ------------------------------------------ process data ---------------------------------------
|
|
|
|
while table.data_read(): # read next data line of ASCII table
|
|
|
|
for datatype,labels in active.items(): # loop over scalar,color
|
|
for me in labels: # loop over all requested items
|
|
theData = [table.data[i] for i in table.label_indexrange(me)] # read strings
|
|
if datatype == 'color':
|
|
VTKarray[me].InsertNextTuple3(*map(lambda x: int(255.*float(x)),theData))
|
|
elif datatype == 'scalar':
|
|
VTKarray[me].InsertNextValue(float(theData[0]))
|
|
|
|
# ------------------------------------------ add data ---------------------------------------
|
|
|
|
for datatype,labels in active.items(): # loop over scalar,color
|
|
if datatype == 'color':
|
|
uGrid.GetCellData().SetScalars(VTKarray[active['color'][0]])
|
|
for label in labels: # loop over all requested items
|
|
uGrid.GetCellData().AddArray(VTKarray[me])
|
|
|
|
uGrid.Modified()
|
|
if vtk.VTK_MAJOR_VERSION <= 5:
|
|
uGrid.Update()
|
|
|
|
# ------------------------------------------ output result ---------------------------------------
|
|
|
|
writer = vtk.vtkXMLUnstructuredGridWriter()
|
|
writer.SetDataModeToBinary()
|
|
writer.SetCompressorTypeToZLib()
|
|
writer.SetFileName(os.path.splitext(options.vtk)[0]+'_added.vtu')
|
|
if vtk.VTK_MAJOR_VERSION <= 5:
|
|
writer.SetInput(uGrid)
|
|
else:
|
|
writer.SetInputData(uGrid)
|
|
writer.Write()
|