1 |
|
2 |
######################################################## |
3 |
# |
4 |
# Copyright (c) 2003-2008 by University of Queensland |
5 |
# Earth Systems Science Computational Center (ESSCC) |
6 |
# http://www.uq.edu.au/esscc |
7 |
# |
8 |
# Primary Business: Queensland, Australia |
9 |
# Licensed under the Open Software License version 3.0 |
10 |
# http://www.opensource.org/licenses/osl-3.0.php |
11 |
# |
12 |
######################################################## |
13 |
|
14 |
__copyright__="""Copyright (c) 2003-2008 by University of Queensland |
15 |
Earth Systems Science Computational Center (ESSCC) |
16 |
http://www.uq.edu.au/esscc |
17 |
Primary Business: Queensland, Australia""" |
18 |
__license__="""Licensed under the Open Software License version 3.0 |
19 |
http://www.opensource.org/licenses/osl-3.0.php""" |
20 |
__url__="http://www.uq.edu.au/esscc/escript-finley" |
21 |
|
22 |
|
23 |
import vtk |
24 |
import sys, os, re |
25 |
import getopt |
26 |
|
27 |
(opts, args) = getopt.getopt(sys.argv[1:], |
28 |
"d:f:ph", |
29 |
["dirname=", "fnamestem=", "preproc", "help"], |
30 |
) |
31 |
|
32 |
def usage(): |
33 |
print "Usage:" |
34 |
print " velVis.py -d <dirname> -f <fnamestem> -p -i\n" |
35 |
print " Arguments:" |
36 |
print " -d/--dirname: directory of vtk files to process (req)" |
37 |
print " -f/--fnamestem: filename stem of vtk files to process (req)" |
38 |
print " -p/--preproc: preprocess vtk files to find global vector max? (opt)" |
39 |
print " -i/--interactive: interactive use of first frame? (opt)" |
40 |
|
41 |
#dirName = '/data/raid2/matt/results/050516/m3dpData/vis' |
42 |
dirName = None |
43 |
fnameStem = None |
44 |
interactive = False |
45 |
preproc = False # preprocess to find the global max of the data |
46 |
|
47 |
for option, arg in opts: |
48 |
if option in ('-d', '--dirname'): |
49 |
dirName = arg |
50 |
elif option in ('-f', '--fnamestem'): |
51 |
fnameStem = arg |
52 |
elif option in ('-p', '--preproc'): |
53 |
preproc = True |
54 |
elif option in ('-i', '--interactive'): |
55 |
interactive = True |
56 |
elif option in ('-h', '--help'): |
57 |
usage() |
58 |
sys.exit(0) |
59 |
|
60 |
if dirName is None: |
61 |
print "You must supply a directory of vtk files to process\n" |
62 |
usage() |
63 |
sys.exit(1) |
64 |
|
65 |
if fnameStem is None: |
66 |
print "You must supply the filename stem of the files to process" |
67 |
print "The filename stem is the part of the filename without the" |
68 |
print "frame number or extension\n" |
69 |
usage() |
70 |
sys.exit(1) |
71 |
|
72 |
# now need to determine the number of files to process |
73 |
dirList = os.listdir(dirName) |
74 |
r = re.compile( "%s\\d+\\.\\w+" % fnameStem ) |
75 |
count = 0 # counter for the number of files found |
76 |
fnames = [] |
77 |
for fname in dirList: |
78 |
if r.match(fname): |
79 |
fnames.append(fname) |
80 |
count += 1 |
81 |
|
82 |
# determine the first frame number and number of frames |
83 |
fnames.sort() |
84 |
firstFile = fnames[0] |
85 |
|
86 |
r = re.compile("([a-zA-Z-])(\\d+)(\\.)(\\w+)") |
87 |
firstNum = r.findall(firstFile) |
88 |
firstNum = int(firstNum[0][1]) |
89 |
|
90 |
# grab the filename extension as well, just in case |
91 |
extn = r.findall(firstFile) |
92 |
extn = extn[0][3] |
93 |
|
94 |
# the number of frames should be equal to the count |
95 |
if interactive: |
96 |
numFrames = 1 |
97 |
mesa = False |
98 |
else: |
99 |
numFrames = count |
100 |
mesa = True |
101 |
|
102 |
# try and work out the format of the zero-padded frame number |
103 |
if count < 10: |
104 |
formatString = '"%01d"' |
105 |
elif count < 100: |
106 |
formatString = '"%02d"' |
107 |
elif count < 1000: |
108 |
formatString = '"%03d"' |
109 |
elif count < 10000: |
110 |
formatString = '"%04d"' |
111 |
elif count < 100000: |
112 |
formatString = '"%05d"' |
113 |
else: |
114 |
print "Sorry, can't handle greater than 100000 input frames" |
115 |
sys.exit(1) |
116 |
|
117 |
# use mesa if required |
118 |
if mesa: |
119 |
factGraphics = vtk.vtkGraphicsFactory() |
120 |
factGraphics.SetUseMesaClasses(1) |
121 |
|
122 |
factImage = vtk.vtkImagingFactory() |
123 |
factImage.SetUseMesaClasses(1) |
124 |
|
125 |
# perform some preprocessing to find the global max of the vector data norm |
126 |
maxNorm = 0 |
127 |
if not interactive and preproc: |
128 |
print "Preprocessing vector data" |
129 |
sys.stdout.write("Completed: ") |
130 |
for i in range(numFrames): |
131 |
frameNum = eval(formatString) % i |
132 |
reader = vtk.vtkXMLUnstructuredGridReader() |
133 |
reader.SetFileName("%s/%s%s.%s" % (dirName,fnameStem,frameNum,extn)) |
134 |
reader.Update() |
135 |
|
136 |
grid = reader.GetOutput() |
137 |
norm = grid.GetPointData().GetVectors().GetMaxNorm() |
138 |
if norm > maxNorm: |
139 |
maxNorm = norm |
140 |
|
141 |
sys.stdout.write("%3.0f%%\b\b\b\b" % (float(i)/float(numFrames), )) |
142 |
print "" |
143 |
|
144 |
for i in range(firstNum, numFrames+firstNum): |
145 |
|
146 |
frameNum = eval(formatString) % i |
147 |
reader = vtk.vtkXMLUnstructuredGridReader() |
148 |
reader.SetFileName("%s/%s%s.%s" % (dirName,fnameStem,frameNum,extn)) |
149 |
reader.Update() |
150 |
|
151 |
print "Processing %s" % fnames[i] |
152 |
|
153 |
grid = reader.GetOutput() |
154 |
|
155 |
# grab the model centre and bounds |
156 |
centre = grid.GetCenter() |
157 |
bounds = grid.GetBounds() |
158 |
|
159 |
# grab the norm of the vectors |
160 |
norm = vtk.vtkVectorNorm() |
161 |
norm.SetInput(grid) |
162 |
|
163 |
if not preproc: |
164 |
maxNorm = grid.GetPointData().GetVectors().GetMaxNorm() |
165 |
|
166 |
# to make arrow glyphs need an arrow source |
167 |
arrow = vtk.vtkArrowSource() |
168 |
|
169 |
# the arrows are 3D glyphs so set that up now |
170 |
glyph = vtk.vtkGlyph3D() |
171 |
glyph.ScalingOn() |
172 |
glyph.SetScaleModeToScaleByScalar() |
173 |
glyph.SetColorModeToColorByScalar() |
174 |
glyph.SetVectorModeToUseVector() |
175 |
glyph.SetScaleFactor(0.1/maxNorm) |
176 |
glyph.SetInput(norm.GetOutput()) |
177 |
glyph.SetSource(arrow.GetOutput()) |
178 |
glyph.ClampingOff() |
179 |
|
180 |
# set up a stripper to speed up rendening |
181 |
stripper = vtk.vtkStripper() |
182 |
stripper.SetInput(glyph.GetOutput()) |
183 |
|
184 |
# make a lookup table for the colour map and invert it (colours look |
185 |
# better when it's inverted) |
186 |
lut = vtk.vtkLookupTable() |
187 |
refLut = vtk.vtkLookupTable() |
188 |
lut.Build() |
189 |
refLut.Build() |
190 |
for j in range(256): |
191 |
lut.SetTableValue(j, refLut.GetTableValue(255-j)) |
192 |
|
193 |
# set up the mapper |
194 |
mapper = vtk.vtkPolyDataMapper() |
195 |
mapper.SetInput(stripper.GetOutput()) |
196 |
mapper.SetScalarRange(0,maxNorm) |
197 |
mapper.SetLookupTable(lut) |
198 |
|
199 |
# set up the actor |
200 |
actor = vtk.vtkActor() |
201 |
actor.SetMapper(mapper) |
202 |
|
203 |
# set up the text properties for nice text |
204 |
font_size = 20 |
205 |
textProp = vtk.vtkTextProperty() |
206 |
textProp.SetFontSize(font_size) |
207 |
textProp.SetFontFamilyToArial() |
208 |
textProp.BoldOff() |
209 |
textProp.ItalicOff() |
210 |
textProp.ShadowOff() |
211 |
textProp.SetColor(0.0, 0.0, 0.0) |
212 |
|
213 |
# make a title |
214 |
title = vtk.vtkTextMapper() |
215 |
title.SetInput("Velocity flow vectors in mantle convection") |
216 |
|
217 |
# make the title text use the text properties |
218 |
titleProp = title.GetTextProperty() |
219 |
titleProp.ShallowCopy(textProp) |
220 |
titleProp.SetJustificationToCentered() |
221 |
titleProp.SetVerticalJustificationToTop() |
222 |
titleProp.BoldOn() |
223 |
|
224 |
# make the actor for the title |
225 |
titleActor = vtk.vtkTextActor() |
226 |
titleActor.SetMapper(title) |
227 |
titleActor.GetPositionCoordinate().SetCoordinateSystemToNormalizedDisplay() |
228 |
titleActor.GetPositionCoordinate().SetValue(0.5, 0.95) |
229 |
|
230 |
# make a frame counter |
231 |
counter = vtk.vtkTextMapper() |
232 |
counter.SetInput("frame: %d" % i) |
233 |
|
234 |
# make the counter use the text properties |
235 |
counterProp = counter.GetTextProperty() |
236 |
counterProp.ShallowCopy(textProp) |
237 |
counterProp.SetJustificationToLeft() |
238 |
counterProp.SetVerticalJustificationToTop() |
239 |
counterProp.SetFontSize(14) |
240 |
|
241 |
# make the actor for the frame counter |
242 |
counterActor = vtk.vtkTextActor() |
243 |
counterActor.SetMapper(counter) |
244 |
counterActor.GetPositionCoordinate().SetCoordinateSystemToNormalizedDisplay() |
245 |
counterActor.GetPositionCoordinate().SetValue(0.05, 0.8) |
246 |
|
247 |
# use a scalar bar |
248 |
scalarBar = vtk.vtkScalarBarActor() |
249 |
scalarBar.SetLookupTable(lut) |
250 |
scalarBar.SetWidth(0.1) |
251 |
scalarBar.SetHeight(0.7) |
252 |
scalarBar.SetPosition(0.9, 0.2) |
253 |
scalarBar.SetTitle("|v|") |
254 |
|
255 |
# set up its title text properties |
256 |
scalarBarProp = scalarBar.GetTitleTextProperty() |
257 |
scalarBarProp.ShallowCopy(textProp) |
258 |
scalarBarProp.SetFontSize(10) |
259 |
|
260 |
# set up the label text properties |
261 |
scalarBarTextProp = scalarBar.GetLabelTextProperty() |
262 |
scalarBarTextProp.ShallowCopy(textProp) |
263 |
scalarBarTextProp.SetFontSize(10) |
264 |
|
265 |
# put an outline around the data |
266 |
outline = vtk.vtkOutlineSource() |
267 |
outline.SetBounds(bounds) |
268 |
|
269 |
# make its mapper |
270 |
outlineMapper = vtk.vtkPolyDataMapper() |
271 |
outlineMapper.SetInput(outline.GetOutput()) |
272 |
|
273 |
# make its actor |
274 |
outlineActor = vtk.vtkActor() |
275 |
outlineActor.SetMapper(outlineMapper) |
276 |
outlineActor.GetProperty().SetColor(0,0,0) |
277 |
|
278 |
# set up the renderer and render window |
279 |
ren = vtk.vtkRenderer() |
280 |
renWin = vtk.vtkRenderWindow() |
281 |
|
282 |
renWin.SetSize(800,600) |
283 |
renWin.AddRenderer(ren) |
284 |
ren.SetBackground(1,1,1) |
285 |
|
286 |
# add the relevant actors |
287 |
ren.AddActor(actor) |
288 |
ren.AddActor(titleActor) |
289 |
ren.AddActor(counterActor) |
290 |
ren.AddActor(scalarBar) |
291 |
ren.AddActor(outlineActor) |
292 |
|
293 |
cam = ren.GetActiveCamera() |
294 |
cam.Azimuth(0) |
295 |
cam.Elevation(-90) |
296 |
cam.Zoom(1.2) |
297 |
ren.SetActiveCamera(cam) |
298 |
ren.ResetCameraClippingRange() |
299 |
|
300 |
# add some axes |
301 |
axes = vtk.vtkCubeAxesActor2D() |
302 |
axes.SetInput(grid) |
303 |
axes.SetCamera(ren.GetActiveCamera()) |
304 |
axes.SetLabelFormat("%6.4g") |
305 |
axes.SetFlyModeToOuterEdges() |
306 |
axes.SetFontFactor(0.8) |
307 |
axes.SetAxisTitleTextProperty(textProp) |
308 |
axes.SetAxisLabelTextProperty(textProp) |
309 |
axes.SetXLabel("x") |
310 |
axes.SetYLabel("y") |
311 |
axes.SetZLabel("z") |
312 |
axes.SetNumberOfLabels(5) |
313 |
axes.GetProperty().SetColor(0,0,0) |
314 |
ren.AddProp(axes) |
315 |
|
316 |
if interactive: |
317 |
# set up stuff for interactive viewing |
318 |
iren = vtk.vtkRenderWindowInteractor() |
319 |
iren.SetRenderWindow(renWin) |
320 |
|
321 |
iren.Initialize() |
322 |
renWin.Render() |
323 |
iren.Start() |
324 |
else: |
325 |
renWin.OffScreenRenderingOn() |
326 |
renWin.Render() |
327 |
|
328 |
# the WindowToImageFilter is what one uses to save the window to an |
329 |
# image file |
330 |
win2img = vtk.vtkWindowToImageFilter() |
331 |
win2img.SetInput(renWin) |
332 |
|
333 |
# set up the PNMWriter as we're saving to pnm |
334 |
writer = vtk.vtkPNMWriter() |
335 |
writer.SetFileName("%s%04d.pnm" % (fnameStem,i)) |
336 |
writer.SetInput(win2img.GetOutput()) |
337 |
writer.Write() |
338 |
|
339 |
print "Wrote %s%04d.pnm" % (fnameStem,i) |
340 |
|
341 |
# vim: expandtab shiftwidth=4: |