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

Contents of /trunk/pyvisi/pyvisi/renderers/vtk/ellipsoid_plot.py

Parent Directory Parent Directory | Revision Log Revision Log


Revision 422 - (show annotations)
Fri Jan 6 03:10:54 2006 UTC (13 years, 5 months ago) by cochrane
File MIME type: text/x-python
File size: 15907 byte(s)
Implemented plotting of escript Data objects.  This only works with the vtk
renderer module at this stage.


1 # Copyright (C) 2004-2005 Paul Cochrane
2 #
3 # This program is free software; you can redistribute it and/or
4 # modify it under the terms of the GNU General Public License
5 # as published by the Free Software Foundation; either version 2
6 # of the License, or (at your option) any later version.
7 #
8 # This program is distributed in the hope that it will be useful,
9 # but WITHOUT ANY WARRANTY; without even the implied warranty of
10 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11 # GNU General Public License for more details.
12 #
13 # You should have received a copy of the GNU General Public License
14 # along with this program; if not, write to the Free Software
15 # Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
16
17 # $Id: ellipsoid_plot.py,v 1.2 2006/01/05 05:05:13 paultcochrane Exp $
18
19 """
20 Class and functions associated with a pyvisi EllipsoidPlot objects
21 """
22
23 # generic imports
24 from pyvisi.renderers.vtk.common import debugMsg
25 import Numeric
26 import os
27 import copy
28 import numarray
29
30 # module specific imports
31 from pyvisi.renderers.vtk.plot import Plot
32
33 __revision__ = '$Revision: 1.2 $'
34
35 class EllipsoidPlot(Plot):
36 """
37 Ellipsoid plot
38 """
39 def __init__(self, scene):
40 """
41 Initialisation of the EllipsoidPlot class
42
43 @param scene: The Scene to render the plot in
44 @type scene: Scene object
45 """
46 debugMsg("Called EllipsoidPlot.__init__()")
47 Plot.__init__(self, scene)
48
49 self.renderer = scene.renderer
50 self.renderer.addToInitStack("# EllipsoidPlot.__init__()")
51
52 # labels and stuff
53 self.title = None
54 self.xlabel = None
55 self.ylabel = None
56 self.zlabel = None
57
58 # default values for fname, format and tensors
59 self.fname = None
60 self.format = None
61 self.tensors = None
62
63 # default values for shared info
64 self.escriptData = False
65 self.otherData = False
66
67 # add the plot to the scene
68 scene.add(self)
69
70 def setData(self, *dataList, **options):
71 """
72 Set data to the plot
73
74 @param dataList: List of data to set to the plot
75 @type dataList: tuple
76
77 @param options: Dictionary of keyword options to the method
78 @type options: dict
79
80 @param fname: the name of the input vtk file
81 @type fname: string
82
83 @param format: the format of the input vtk file ('vtk' or 'vtk-xml')
84 @type format: string
85
86 @param tensors: the name of the tensor data in the vtk file to use
87 @type tensors: string
88 """
89 debugMsg("Called setData() in EllipsoidPlot()")
90
91 self.renderer.runString("# EllipsoidPlot.setData()")
92
93 # process the options, if any
94 ## fname
95 if options.has_key('fname'):
96 fname = options['fname']
97 else:
98 fname = None
99 ## format
100 if options.has_key('format'):
101 format = options['format']
102 else:
103 format = None
104 ## tensors
105 if options.has_key('tensors'):
106 tensors = options['tensors']
107 else:
108 tensors = None
109
110 # we want to pass this info around
111 self.fname = fname
112 self.format = format
113 self.tensors = tensors
114
115 # reset the default values for shared info
116 self.escriptData = False
117 self.otherData = False
118
119 # do some sanity checking on the input args
120 if len(dataList) == 0 and fname is None:
121 raise ValueError, \
122 "You must specify a data list or an input filename"
123
124 if len(dataList) != 0 and fname is not None:
125 raise ValueError, \
126 "You cannot specify a data list as well as an input file"
127
128 if fname is not None and tensors is None:
129 debugMsg("No tensors specified; using default in vtk")
130
131 if fname is not None and format is None:
132 raise ValueError, "You must specify an input file format"
133
134 if fname is None and format is not None:
135 raise ValueError, "Format specified but no input filename"
136
137 # if have just a data list, check the objects passed in to see if
138 # they are escript data objects or not
139 if len(dataList) != 0:
140 for obj in dataList:
141 try:
142 obj.convertToNumArray()
143 # ok, we've got escript data, set the flag
144 self.escriptData = True
145 except AttributeError:
146 self.otherData = True
147
148 # if we have both escript data and other data, barf as can't handle
149 # that yet
150 if self.escriptData and self.otherData:
151 raise TypeError, \
152 "Sorry can't handle both escript and other data yet"
153
154 # now generate the code for the case when we have just escript data
155 # passed into setData()
156 if self.escriptData:
157 # get the relevant bits of data
158 if len(dataList) == 1:
159 # only one data variable, will need to get the domain from it
160 escriptZ = dataList[0]
161 escriptX = escriptZ.getDomain().getX()
162 else:
163 errorString = \
164 "Expecting only 1 element in data list. I got %d" \
165 % len(dataList)
166 raise ValueError, errorString
167
168 # convert the Data objects to numarrays
169 domainData = escriptX.convertToNumArray()
170 fieldData = escriptZ.convertToNumArray()
171
172 # make sure that the shapes are correct
173 if len(domainData.shape) != 2:
174 raise ValueError, \
175 "domainData shape not 2D. I got %d dims" % \
176 len(domainData.shape)
177
178 if len(fieldData.shape) != 3:
179 raise ValueError, \
180 "fieldData shape not 3D. I got %d dims" % \
181 len(fieldData.shape)
182
183 # check the array lengths
184 if domainData.shape[0] != fieldData.shape[0]:
185 raise ValueError, \
186 "domainData and fieldData array lengths don't agree"
187
188 # check the array dimensions
189 if domainData.shape[1] != 2 and domainData.shape[1] != 3:
190 raise ValueError, \
191 "domainData array not 2D or 3D. I got %d dims" % \
192 domainData.shape[1]
193
194 if fieldData.shape[1] != 2 and fieldData.shape[1] != 3:
195 raise ValueError, \
196 "fieldData first array not 2D or 3D. I got %d dims" \
197 % fieldData.shape[1]
198
199 if fieldData.shape[2] != 2 and fieldData.shape[2] != 3:
200 raise ValueError, \
201 "fieldData second array not 2D or 3D. I got %d dims" \
202 % fieldData.shape[2]
203
204 # check that the fieldData arrays are square
205 if fieldData.shape[1] != fieldData.shape[2]:
206 errorString = "fieldData arrays not square. "
207 errorString += "First dim: %d. " % fieldData.shape[1]
208 errorString += "Second dim %d." % fieldData.shape[2]
209 raise ValueError, errorString
210
211 # get the x, y and z data from the domain
212 xData = domainData[:,0]
213 numPoints = len(xData) # handy number to keep hanging around
214 yData = domainData[:,1]
215 # handle the case where no zData is specified
216 if domainData.shape[1] == 2:
217 zData = numarray.zeros(numPoints)
218 else:
219 zData = domainData[:,2]
220
221 ####!!!! now process the data properly so that I can plot it
222 print "domainData.shape = %s" % str(domainData.shape)
223 print "fieldData.shape = %s" % str(fieldData.shape)
224
225 # now pass the data to the render dictionary so that the render
226 # code knows what it's supposed to plot
227
228 # x data
229 self.renderer.renderDict['_x'] = copy.deepcopy(xData)
230
231 # y data
232 self.renderer.renderDict['_y'] = copy.deepcopy(yData)
233
234 # z data
235 self.renderer.renderDict['_z'] = copy.deepcopy(zData)
236
237 # field data
238 self.renderer.renderDict['_data'] = copy.deepcopy(fieldData)
239
240 # make the points for the grid
241 evalString = "_points = vtk.vtkPoints()\n"
242 evalString += "_points.SetNumberOfPoints(%d)\n" % numPoints
243 evalString += "for _j in range(%d):\n" % numPoints
244 evalString += \
245 " _points.InsertPoint(_j, _x[_j], _y[_j], _z[_j])\n"
246 self.renderer.runString(evalString)
247
248 # make the array of data
249 evalString = "_tensors = vtk.vtkFloatArray()\n"
250 evalString += "_tensors.SetNumberOfComponents(9)\n"
251 evalString += "_tensors.SetNumberOfTuples(%d)\n" % numPoints
252 evalString += "_tensors.SetName(\"tensors\")\n"
253 evalString += "for _j in range(%d):\n" % numPoints
254 if fieldData.shape[1] == 2:
255 evalString += " _tensors.InsertTuple9(_j, "
256 evalString += "_data[_j][0][0], _data[_j][0][1], 0.0, "
257 evalString += "_data[_j][1][0], _data[_j][1][1], 0.0, "
258 evalString += "0.0, 0.0, 0.0)\n"
259 elif fieldData.shape[1] == 3:
260 evalString += " _tensors.InsertTuple9(_j, "
261 evalString += \
262 "_data[_j][0][0], _data[_j][0][1], _data[_j][0][2], "
263 evalString += \
264 "_data[_j][1][0], _data[_j][1][1], _data[_j][1][2], "
265 evalString += \
266 "_data[_j][2][0], _data[_j][2][1], _data[_j][2][2])\n"
267 else:
268 raise ValueError, "Incorrect fieldData shape. I got %d" % \
269 fieldData.shape[1]
270 self.renderer.runString(evalString)
271
272 # make the grid
273 evalString = "_grid = vtk.vtkUnstructuredGrid()\n"
274 evalString += "_grid.SetPoints(_points)\n"
275 evalString += "_grid.GetPointData().AddArray(_tensors)\n"
276 evalString += \
277 "_grid.GetPointData().SetActiveTensors(\"tensors\")\n"
278 self.renderer.runString(evalString)
279
280 # now extract the tensor components
281 evalString = "_extract = vtk.vtkExtractTensorComponents()\n"
282 evalString += "_extract.SetInput(_grid)\n"
283 evalString += "_extract.SetScalarModeToEffectiveStress()\n"
284 evalString += "_extract.ExtractScalarsOn()\n"
285 evalString += "_extract.PassTensorsToOutputOn()\n"
286 evalString += "_extract.ScalarIsEffectiveStress()\n"
287
288 evalString += "_extractGrid = _extract.GetOutput()\n"
289 evalString += "_extractGrid.Update()\n"
290 evalString += "_extractScalarRange = "
291 evalString += \
292 "_extractGrid.GetPointData().GetScalars().GetRange()\n"
293 self.renderer.runString(evalString)
294
295 elif self.otherData:
296
297 # do some checks to make sure have the right kind of shape for
298 # the data and then generate the code
299
300 raise ImplementationError, "Can't process plain array data yet"
301
302 if fname is not None:
303
304 # check to see if the file exists
305 if not os.path.exists(fname):
306 raise SystemError, "File %s not found" % fname
307
308 if format == 'vtk':
309 # read old-style vtk files
310 evalString = "_reader = vtk.vtkUnstructuredGridReader()\n"
311 elif format == 'vtk-xml':
312 # read vtk xml files
313 evalString = "_reader = vtk.vtkXMLUnstructuredGridReader()\n"
314 else:
315 # barf
316 raise ValueError, "Unknown format. I got %s" % format
317
318 evalString += "_reader.SetFileName(\"%s\")\n" % fname
319 evalString += "_reader.Update()"
320
321 self.renderer.runString(evalString)
322
323 # grab the grid of the data
324 self.renderer.runString("_grid = _reader.GetOutput()")
325
326 # convert the cell data to point data
327 evalString = "_c2p = vtk.vtkCellDataToPointData()\n"
328 evalString += "_c2p.SetInput(_grid)"
329 self.renderer.runString(evalString)
330
331 # now extract the tensor components
332 evalString = "_extract = vtk.vtkExtractTensorComponents()\n"
333 evalString += "_extract.SetInput(_c2p.GetOutput())\n"
334 evalString += "_extract.SetScalarModeToEffectiveStress()\n"
335 evalString += "_extract.ExtractScalarsOn()\n"
336 evalString += "_extract.PassTensorsToOutputOn()\n"
337 evalString += "_extract.ScalarIsEffectiveStress()\n"
338
339 evalString += "_extractGrid = _extract.GetOutput()\n"
340 evalString += "_extractGrid.Update()\n"
341 evalString += "_extractScalarRange = "
342 evalString += \
343 "_extractGrid.GetPointData().GetScalars().GetRange()\n"
344 self.renderer.runString(evalString)
345
346 return
347
348 def render(self):
349 """
350 Does EllipsoidPlot object specific (pre)rendering stuff
351 """
352 debugMsg("Called EllipsoidPlot.render()")
353
354 self.renderer.runString("# EllipsoidPlot.render()")
355
356 # make a sphere source for the glyphs
357 evalString = "_sphere = vtk.vtkSphereSource()\n"
358 evalString += "_sphere.SetThetaResolution(6)\n"
359 evalString += "_sphere.SetPhiResolution(6)\n"
360 evalString += "_sphere.SetRadius(0.5)"
361 self.renderer.runString(evalString)
362
363 # make tensor glyphs
364 evalString = "_glyph = vtk.vtkTensorGlyph()\n"
365 evalString += "_glyph.SetSource(_sphere.GetOutput())\n"
366 evalString += "_glyph.SetInput(_extractGrid)\n"
367 evalString += "_glyph.SetColorModeToScalars()\n"
368 evalString += "_glyph.ScalingOn()\n"
369 evalString += "_glyph.SetMaxScaleFactor(5.0)\n"
370 evalString += "_glyph.SetScaleFactor(1.0)\n"
371 evalString += "_glyph.ClampScalingOn()"
372 self.renderer.runString(evalString)
373
374 # make a stripper for faster rendering
375 evalString = "_stripper = vtk.vtkStripper()\n"
376 evalString += "_stripper.SetInput(_glyph.GetOutput())"
377 self.renderer.runString(evalString)
378
379 # make the normals of the data
380 evalString = "_normals = vtk.vtkPolyDataNormals()\n"
381 evalString += "_normals.SetInput(_stripper.GetOutput())"
382 self.renderer.runString(evalString)
383
384 # make the mapper for the data
385 evalString = "_mapper = vtk.vtkPolyDataMapper()\n"
386 evalString += "_mapper.SetInput(_normals.GetOutput())\n"
387 evalString += "_mapper.SetLookupTable(_lut)\n"
388 evalString += "_mapper.SetScalarRange(_extractScalarRange)"
389 self.renderer.runString(evalString)
390
391 # make the actor
392 evalString = "_actor = vtk.vtkActor()\n"
393 evalString += "_actor.SetMapper(_mapper)"
394 self.renderer.runString(evalString)
395
396 # add the actor
397 self.renderer.runString("_renderer.AddActor(_actor)")
398
399 # set up the text properties for nice text
400 evalString = "_textProp = vtk.vtkTextProperty()\n"
401 evalString += "_textProp.SetFontFamilyToArial()\n"
402 evalString += "_textProp.BoldOff()\n"
403 evalString += "_textProp.ItalicOff()\n"
404 evalString += "_textProp.ShadowOff()\n"
405 evalString += "_textProp.SetColor(0.0, 0.0, 0.0)"
406 self.renderer.runString(evalString)
407
408 # if a title is set, put it in here
409 if self.title is not None:
410 # make a title
411 evalString = "_title = vtk.vtkTextMapper()\n"
412 evalString += "_title.SetInput(\"%s\")\n" % self.title
413
414 # make the title text use the text properties
415 evalString += "_titleProp = _title.GetTextProperty()\n"
416 evalString += "_titleProp.ShallowCopy(_textProp)\n"
417 evalString += "_titleProp.SetJustificationToCentered()\n"
418 evalString += "_titleProp.SetVerticalJustificationToTop()\n"
419 evalString += "_titleProp.SetFontSize(20)\n"
420 evalString += "_titleProp.BoldOn()\n"
421
422 # make the actor for the title
423 evalString += "_titleActor = vtk.vtkTextActor()\n"
424 evalString += "_titleActor.SetMapper(_title)\n"
425 evalString += "_titleActor.GetPositionCoordinate()."
426 evalString += "SetCoordinateSystemToNormalizedDisplay()\n"
427 evalString += "_titleActor.GetPositionCoordinate()."
428 evalString += "SetValue(0.5, 0.95)"
429 self.renderer.runString(evalString)
430
431 # add to the renderer
432 evalString = "_renderer.AddActor(_titleActor)"
433 self.renderer.runString(evalString)
434
435 # add a scalar bar
436 evalString = "_scalarBar = vtk.vtkScalarBarActor()\n"
437 evalString += "_scalarBar.SetLookupTable(_lut)\n"
438 evalString += "_scalarBar.SetWidth(0.1)\n"
439 evalString += "_scalarBar.SetHeight(0.8)\n"
440 evalString += "_scalarBar.SetPosition(0.9, 0.15)"
441 self.renderer.runString(evalString)
442
443 # set up the label text properties
444 evalString = "_scalarBarTextProp = _scalarBar.GetLabelTextProperty()\n"
445 evalString += "_scalarBarTextProp.ShallowCopy(_textProp)\n"
446 evalString += "_scalarBarTextProp.SetFontSize(10)\n"
447
448 evalString += "_renderer.AddActor(_scalarBar)"
449 self.renderer.runString(evalString)
450
451 return
452
453
454 # vim: expandtab shiftwidth=4:

  ViewVC Help
Powered by ViewVC 1.1.26