/[escript]/trunk/pyvisi/py_src/renderers/vtk/arrow_plot.py
ViewVC logotype

Contents of /trunk/pyvisi/py_src/renderers/vtk/arrow_plot.py

Parent Directory Parent Directory | Revision Log Revision Log


Revision 720 - (show annotations)
Thu Apr 27 10:16:05 2006 UTC (13 years, 6 months ago) by gross
File MIME type: text/x-python
File size: 15916 byte(s)
formatting errors in epydoc tags fixed
1 """
2 Class and functions associated with a pyvisi ArrowPlot objects
3
4 @var __author__: name of author
5 @var __license__: licence agreement
6 @var __copyright__: copyrights
7 @var __url__: url entry point on documentation
8 @var __version__: version
9 @var __date__: date of the version
10 """
11
12 __copyright__=""" Copyright (c) 2006 by ACcESS MNRF
13 http://www.access.edu.au
14 Primary Business: Queensland, Australia"""
15 __license__="""Licensed under the Open Software License version 3.0
16 http://www.opensource.org/licenses/osl-3.0.php"""
17 __author__="Paul Cochrane"
18 __url__="http://www.iservo.edu.au/esys"
19 __version__="$Revision$"
20 __date__="$Date$"
21
22
23 # generic imports
24 from common import debugMsg
25 import os
26 import copy
27
28 # module specific imports
29 from plot import Plot
30
31 class ArrowPlot(Plot):
32 """
33 Arrow field plot
34 """
35 def __init__(self, scene):
36 """
37 Initialisation of the ArrowPlot class
38
39 @param scene: The Scene to render the plot in
40 @type scene: Scene object
41 """
42 debugMsg("Called ArrowPlot.__init__()")
43 Plot.__init__(self, scene)
44
45 self.renderer = scene.renderer
46
47 # default values for shared info
48 self.fname = None
49 self.format = None
50 self.vectors = None
51 self.escriptData = False
52 self.otherData = False
53
54 # add the plot to the scene
55 scene.add(self)
56
57 def setData(self, *dataList, **options):
58 """
59 Set data to the plot
60
61 @param dataList: List of data to set to the plot
62 @type dataList: tuple
63
64 @param options: Dictionary of extra options
65 @type options: dict
66
67 @keyword fname: the name of the input vtk file
68 @type fname: string
69
70 @keyword format: the format of the input vtk file ('vtk' or 'vtk-xml')
71 @type format: string
72
73 @keyword vectors: the name of the vector data in the vtk file to use
74 @type vectors: string
75 """
76 debugMsg("Called setData() in ArrowPlot()")
77 self.renderer.runString("# ArrowPlot.setData()")
78
79 # get the options, if any
80 ## fname
81 if options.has_key('fname'):
82 fname = options['fname']
83 else:
84 fname = None
85 ## format
86 if options.has_key('format'):
87 format = options['format']
88 else:
89 format = None
90 ## vectors
91 if options.has_key('vectors'):
92 vectors = options['vectors']
93 else:
94 vectors = None
95
96 # we want to pass this info around
97 self.fname = fname
98 self.format = format
99 self.vectors = vectors
100
101 # reset the default values for shared info
102 self.escriptData = False
103 self.otherData = False
104
105 # do some sanity checking on the input args
106 if len(dataList) == 0 and fname is None:
107 raise ValueError, \
108 "You must specify a data list or an input filename"
109
110 if len(dataList) != 0 and fname is not None:
111 raise ValueError, \
112 "You cannot specify a data list as well as an input file"
113
114 if fname is not None and vectors is None:
115 debugMsg("No vectors specified; using default in vtk")
116
117 if fname is not None and format is None:
118 raise ValueError, "You must specify an input file format"
119
120 if fname is None and format is not None:
121 raise ValueError, "Format specified but no input filename"
122
123 # if have just a data list, check the objects passed in to see if
124 # they are escript data objects or not
125 if len(dataList) != 0:
126 for obj in dataList:
127 try:
128 obj.convertToNumArray()
129 # ok, we've got escript data, set the flag
130 self.escriptData = True
131 except AttributeError:
132 self.otherData = True
133
134 # if we have both escript data and other data, barf as can't handle
135 # that yet
136 if self.escriptData and self.otherData:
137 raise TypeError, \
138 "Sorry can't handle both escript and other data yet"
139
140 # now generate the code for the case when we have just escript data
141 # passed into setData()
142 if self.escriptData:
143 # get the relevant bits of data
144 if len(dataList) == 1:
145 # only one data variable, will need to get the domain from it
146 escriptZ = dataList[0]
147 escriptX = escriptZ.getDomain().getX()
148 else:
149 errorString = \
150 "Expecting only 1 element in data list. I got %d" \
151 % len(dataList)
152 raise ValueError, errorString
153
154 domainData = escriptX.convertToNumArray()
155 fieldData = escriptZ.convertToNumArray()
156
157 # now check the shapes
158 if len(domainData.shape) != 2:
159 raise ValueError, \
160 "domainData shape is not 2D. I got %d dims" % \
161 len(domainData.shape)
162
163 if len(fieldData.shape) != 2:
164 raise ValueError, \
165 "fieldData shape is not 2D. I got %d dims" % \
166 len(fieldData.shape)
167
168 # we expect only 2D vectors, so make sure that the second
169 # dimension of the array is equal to 2
170 if fieldData.shape[1] != 2:
171 raise ValueError, \
172 "fieldData vectors not 2D. I got %d dims" % \
173 fieldData.shape[1]
174
175 # make sure the lengths agree
176 if domainData.shape[0] != fieldData.shape[0]:
177 raise ValueError, \
178 "domainData and fieldData lengths don't agree"
179
180 # split the domainData and fieldData up into x and y parts
181 xData = domainData[:,0]
182 yData = domainData[:,1]
183
184 dxData = fieldData[:,0]
185 dyData = fieldData[:,1]
186
187 # now pass the data to the render dictionary so that the render code
188 # knows what it's supposed to plot
189 # x data
190 self.renderer.renderDict['_x'] = copy.deepcopy(xData)
191
192 # y data
193 self.renderer.renderDict['_y'] = copy.deepcopy(yData)
194
195 # dx data
196 self.renderer.renderDict['_dx'] = copy.deepcopy(dxData)
197
198 # dy data
199 self.renderer.renderDict['_dy'] = copy.deepcopy(dyData)
200
201 # keep the number of points for future reference
202 numPoints = len(xData)
203
204 # construct the points data
205 evalString = "_points = vtk.vtkPoints()\n"
206 evalString += "_points.SetNumberOfPoints(%d)\n" % numPoints
207 evalString += "for _j in range(%d):\n" % numPoints
208 evalString += " _points.InsertPoint(_j, _x[_j], _y[_j], 0.0)\n"
209 self.renderer.runString(evalString)
210
211 # construct the vectors
212 evalString = "_vectors = vtk.vtkFloatArray()\n"
213 evalString += "_vectors.SetNumberOfComponents(3)\n"
214 evalString += "_vectors.SetNumberOfTuples(%d)\n" % numPoints
215 evalString += "_vectors.SetName(\"vectors\")\n"
216 evalString += "for _j in range(%d):\n" % numPoints
217 evalString += \
218 " _vectors.InsertTuple3(_j, _dx[_j], _dy[_j], 0.0)\n"
219 self.renderer.runString(evalString)
220
221 # construct the grid
222 evalString = "_grid = vtk.vtkUnstructuredGrid()\n"
223 evalString += "_grid.SetPoints(_points)\n"
224 evalString += "_grid.GetPointData().AddArray(_vectors)\n"
225 evalString += "_grid.GetPointData().SetActiveVectors(\"vectors\")"
226 self.renderer.runString(evalString)
227
228 elif self.otherData:
229
230 # do some sanity checking on the data
231 if len(dataList) != 4:
232 raise ValueError, \
233 "Must have four vectors as input: x, y, dx, dy"
234
235 for i in range(len(dataList)):
236 if len(dataList[i].shape) != len(dataList[0].shape):
237 raise ValueError, "All arrays must be of the same shape"
238
239 for i in range(len(dataList)):
240 if len(dataList[i].shape) != 1 and len(dataList[i].shape) != 2:
241 errorString = \
242 "Can only handle 1D or 2D arrays: dim=%d" % \
243 len(dataList[i].shape)
244 raise ValueError, errorString
245
246 for i in range(len(dataList)):
247 if len(dataList[0]) != len(dataList[i]):
248 raise ValueError, \
249 "Input vectors must all be the same length"
250
251 # if we have 2D arrays as input, we need to flatten them to plot the
252 # data properly
253 if len(dataList[0].shape) == 1:
254 xData = dataList[0]
255 yData = dataList[1]
256 dxData = dataList[2]
257 dyData = dataList[3]
258 elif len(dataList[0].shape) == 2:
259 xData = dataList[0].flat
260 yData = dataList[1].flat
261 dxData = dataList[2].flat
262 dyData = dataList[3].flat
263 else:
264 raise ValueError, "Input vectors can only be 1D or 2D"
265
266 # now pass the data to the render dictionary so that the render code
267 # knows what it's supposed to plot
268 # x data
269 self.renderer.renderDict['_x'] = copy.deepcopy(xData)
270
271 # y data
272 self.renderer.renderDict['_y'] = copy.deepcopy(yData)
273
274 # dx data
275 self.renderer.renderDict['_dx'] = copy.deepcopy(dxData)
276
277 # dy data
278 self.renderer.renderDict['_dy'] = copy.deepcopy(dyData)
279
280 # keep the number of points for future reference
281 numPoints = len(xData)
282
283 # construct the points data
284 evalString = "_points = vtk.vtkPoints()\n"
285 evalString += "_points.SetNumberOfPoints(%d)\n" % numPoints
286 evalString += "for _j in range(%d):\n" % numPoints
287 evalString += " _points.InsertPoint(_j, _x[_j], _y[_j], 0.0)\n"
288 self.renderer.runString(evalString)
289
290 # construct the vectors
291 evalString = "_vectors = vtk.vtkFloatArray()\n"
292 evalString += "_vectors.SetNumberOfComponents(3)\n"
293 evalString += "_vectors.SetNumberOfTuples(%d)\n" % numPoints
294 evalString += "_vectors.SetName(\"vectors\")\n"
295 evalString += "for _j in range(%d):\n" % numPoints
296 evalString += \
297 " _vectors.InsertTuple3(_j, _dx[_j], _dy[_j], 0.0)\n"
298 self.renderer.runString(evalString)
299
300 # construct the grid
301 evalString = "_grid = vtk.vtkUnstructuredGrid()\n"
302 evalString += "_grid.SetPoints(_points)\n"
303 evalString += "_grid.GetPointData().AddArray(_vectors)\n"
304 evalString += "_grid.GetPointData().SetActiveVectors(\"vectors\")"
305 self.renderer.runString(evalString)
306
307 # run the stuff for when we're reading from file
308 if fname is not None:
309
310 # had best make sure it exists
311 if not os.path.exists(fname):
312 raise SystemError, "File: %s not found" % fname
313
314 if format == 'vtk':
315 # read old-style vtk files
316 evalString = "_reader = vtk.vtkUnstructuredGridReader()\n"
317 elif format == 'vtk-xml':
318 # read vtk xml files
319 evalString = "_reader = vtk.vtkXMLUnstructuredGridReader()\n"
320 else:
321 # barf
322 raise ValueError, "Unknown format. I got %s" % format
323
324 evalString += "_reader.SetFileName(\"%s\")\n" % fname
325 evalString += "_reader.Update()"
326 self.renderer.runString(evalString)
327
328 # read the output into an unstructured grid
329 evalString = "_grid = _reader.GetOutput()\n"
330 evalString += \
331 "_grid.GetPointData().SetActiveVectors(\"%s\")" % vectors
332 self.renderer.runString(evalString)
333
334 def render(self):
335 """
336 Does ArrowPlot specific rendering tasks
337 """
338 debugMsg("Called render() in ArrowPlot")
339 self.renderer.runString("# ArrowPlot.render()")
340
341 # set up the lookup table and reverse the order of the colours
342 evalString = "_lut = vtk.vtkLookupTable()\n"
343 evalString += "_lut.Build()\n"
344 evalString += "_refLut = vtk.vtkLookupTable()\n"
345 evalString += "_refLut.Build()\n"
346 evalString += "for _i in range(256):\n"
347 evalString += \
348 " _lut.SetTableValue(_i, _refLut.GetTableValue(255-_i))"
349 self.renderer.runString(evalString)
350
351 # make the arrow source
352 self.renderer.runString("_arrow = vtk.vtkArrowSource()")
353
354 # make the glyph
355 evalString = "_glyph = vtk.vtkGlyph2D()\n"
356 evalString += "_glyph.ScalingOn()\n"
357 evalString += "_glyph.SetScaleModeToScaleByVector()\n"
358 evalString += "_glyph.SetColorModeToColorByVector()\n"
359 evalString += "_glyph.SetScaleFactor(0.5)\n"
360 evalString += "_glyph.SetSource(_arrow.GetOutput())\n"
361 evalString += "_glyph.SetInput(_grid)\n"
362 evalString += "_glyph.ClampingOff()"
363 self.renderer.runString(evalString)
364
365 # set up a stripper for faster rendering
366 evalString = "_stripper = vtk.vtkStripper()\n"
367 evalString += "_stripper.SetInput(_glyph.GetOutput())"
368 self.renderer.runString(evalString)
369
370 # get the maximum norm of the data
371 evalString = "_maxNorm = _grid.GetPointData().GetVectors().GetMaxNorm()"
372 self.renderer.runString(evalString)
373
374 # set up the mapper
375 evalString = "_mapper = vtk.vtkPolyDataMapper()\n"
376 evalString += "_mapper.SetInput(_stripper.GetOutput())\n"
377 evalString += "_mapper.SetScalarRange(0, _maxNorm)"
378 self.renderer.runString(evalString)
379
380 # set up the actor
381 evalString = "_actor = vtk.vtkActor()\n"
382 evalString += "_actor.SetMapper(_mapper)"
383 self.renderer.runString(evalString)
384
385 # add the actor
386 self.renderer.runString("_renderer.AddActor(_actor)")
387
388 # text properties
389 evalString = "_font_size = 14\n" # this will need to be an option!!
390 evalString += "_textProp = vtk.vtkTextProperty()\n"
391 evalString += "_textProp.SetFontSize(_font_size)\n"
392 evalString += "_textProp.SetFontFamilyToArial()\n"
393 evalString += "_textProp.BoldOff()\n"
394 evalString += "_textProp.ItalicOff()\n"
395 evalString += "_textProp.ShadowOff()\n"
396 evalString += "_textProp.SetColor(0,0,0)\n"
397 self.renderer.runString(evalString)
398
399 # set the title if set
400 if self.title is not None:
401 # add a title
402 evalString = "_titleMapper = vtk.vtkTextMapper()\n"
403 evalString += "_titleMapper.SetInput(\"%s\")\n" % self.title
404
405 evalString += "_titleProp = _titleMapper.GetTextProperty()\n"
406 evalString += "_titleProp.ShallowCopy(_textProp)\n"
407 evalString += "_titleProp.SetJustificationToCentered()\n"
408 evalString += "_titleProp.SetVerticalJustificationToTop()\n"
409 evalString += "_titleProp.SetFontSize(18)\n"
410
411 # set up the text actor
412 evalString += "_titleActor = vtk.vtkTextActor()\n"
413 evalString += "_titleActor.SetMapper(_titleMapper)\n"
414 evalString += "_titleActor.GetPositionCoordinate()."
415 evalString += "SetCoordinateSystemToNormalizedDisplay()\n"
416 evalString += "_titleActor.GetPositionCoordinate()."
417 evalString += "SetValue(0.5, 0.95)\n"
418
419 evalString += "_renderer.AddActor(_titleActor)"
420 self.renderer.runString(evalString)
421
422 # set up some axes
423 evalString = "_axes = vtk.vtkCubeAxesActor2D()\n"
424 evalString += "_axes.SetCamera(_renderer.GetActiveCamera())\n"
425 evalString += "_axes.SetFlyModeToOuterEdges()\n"
426 evalString += "_axes.SetBounds(min(_x), max(_x)+_maxNorm, "
427 evalString += "min(_y), max(_y)+_maxNorm, 0, 0)\n"
428
429 if self.xlabel is None:
430 evalString += "_axes.SetXLabel(\"\")\n"
431 else:
432 evalString += "_axes.SetXLabel(\"%s\")\n" % self.xlabel
433
434 if self.ylabel is None:
435 evalString += "_axes.SetYLabel(\"\")\n"
436 else:
437 evalString += "_axes.SetYLabel(\"%s\")\n" % self.ylabel
438
439 evalString += "_axes.SetZLabel(\"\")\n"
440 evalString += "_axes.YAxisVisibilityOff()\n" # but this is the z axis!!
441
442 # set up the axes properties
443 evalString += "_axesProp = _axes.GetProperty()\n"
444 evalString += "_axesProp.SetColor(0,0,0)\n"
445
446 # set up the axes title properties
447 evalString += "_axesTitleProp = _axes.GetAxisTitleTextProperty()\n"
448 evalString += "_axesTitleProp.ShallowCopy(_textProp)\n"
449
450 # set up the axes label properties
451 evalString += "_axesLabelProp = _axes.GetAxisLabelTextProperty()\n"
452 evalString += "_axesLabelProp.ShallowCopy(_textProp)\n"
453 evalString += "_axesLabelProp.SetFontSize(8)\n"
454 self.renderer.runString(evalString)
455
456 # add the axes to the renderer
457 self.renderer.runString("_renderer.AddActor(_axes)")
458
459 # reset the camera, will make things look nicer
460 ### is this the right place to put this???
461 self.renderer.runString("_renderer.ResetCamera()")
462
463 ### this should be somewhere else too...
464 self.renderer.runString("_renderer.SetBackground(1,1,1)")
465
466 return
467
468
469 # vim: expandtab shiftwidth=4:

  ViewVC Help
Powered by ViewVC 1.1.26