/[escript]/trunk/pyvisi/py_src/datacollector.py
ViewVC logotype

Contents of /trunk/pyvisi/py_src/datacollector.py

Parent Directory Parent Directory | Revision Log Revision Log


Revision 1100 - (show annotations)
Tue Apr 17 05:36:24 2007 UTC (12 years ago) by jongui
File MIME type: text/x-python
File size: 12756 byte(s)
Fixed the problem related to the use of point and cell data. Both point and cell data should now be able to work properly together. i.e. plotting a velocity field using 'vector point data' but coloring the arrows using 'scalar cell data'.
1 """
2 @author: John NGUI
3 """
4
5 import vtk
6 import tempfile, os, sys
7 from constant import Source, VizType, ColorMode
8 try:
9 import esys.escript
10 except ImportError:
11 print "Warning: importing esys.escript failed."
12
13 class DataCollector:
14 """
15 Class that defines a data collector. A data collector is used to read
16 data from an XML file or from an escript object directly.
17
18 @attention: A DataCollector instance can only be used to specify one
19 scalar, vector and tensor attribute from a source at any one time. If a
20 second scalar, vector or tensor attribute needs to be specified from the
21 same source, a second DataCollector instance must be created.
22
23 @attention: When a series of XML files or escript objects are read
24 (using 'setFileName' or 'setData' in a for-loop), the 'setActiveScalar' /
25 'setActiveVector' / 'setActiveTensor' have to be called for each new file
26 (provided a specific field needs to be loaded) as all active fields
27 specified from the previous file goes back to the default once a new file
28 is read.
29 """
30
31 def __init__(self, source = Source.XML):
32 """
33 Initialise the data collector.
34
35 @type source: L{Source <constant.Source>} constant
36 @param source: Source type
37 """
38
39 self.__source = source
40 self.__count = 0 # Keeps track of the number of files/sources read.
41
42 if(source == Source.XML): # Source is an XML file.
43 self.__vtk_xml_reader = vtk.vtkXMLUnstructuredGridReader()
44 # Source is a escript data object using a temp file in the background.
45 elif (self.__source == Source.ESCRIPT):
46 self.__vtk_xml_reader = vtk.vtkXMLUnstructuredGridReader()
47 # Create a temporary .xml file and retrieve its path.
48 self.__tmp_file = tempfile.mkstemp(suffix=".xml")[1]
49
50 def __del__(self):
51 """
52 Perform some clean up of the temporary file.
53 """
54
55 if (self.__source == Source.ESCRIPT):
56 if os.access(self.__tmp_file,os.F_OK): os.unlink(self.__tmp_file)
57
58 def setFileName(self, file_name):
59 """
60 Set the XML file name to read.
61
62 @type file_name: String
63 @param file_name: Name of the file to read
64 """
65
66 if(self.__source == Source.XML):
67 # Check whether the specified file exists, otherwise exit.
68 if not(os.access(file_name, os.F_OK)):
69 print "\nERROR: '%s' file does NOT exists.\n" % file_name
70 sys.exit(1)
71
72 self.__vtk_xml_reader.SetFileName(file_name)
73 # Update must be called after SetFileName to make the reader
74 # up-to-date. Otherwise, some output values may be incorrect.
75 self.__vtk_xml_reader.Update()
76 self.__output = self.__vtk_xml_reader.GetOutput()
77 self.__get_attribute_lists()
78
79 # Count has to be larger than zero because when setFileName is
80 # called for the first time, the data set mapper has not yet been
81 # instantiated. Therefore, the range of the mapper can only be
82 # updated after the first file/source has been read.
83 if(self.__count > 0):
84 self._updateRange()
85
86 self.__count+=1
87
88 else:
89 raise ValueError("Source type %s does not support \
90 'setFileName'\n" % self.__source)
91
92 def setData(self,**args):
93 """
94 Create data using the <name>=<data> pairing. Assumption is made
95 that the data will be given in the appropriate format.
96 """
97
98 if self.__source == Source.ESCRIPT:
99 esys.escript.saveVTK(self.__tmp_file,**args)
100 self.__vtk_xml_reader.SetFileName(self.__tmp_file)
101 # Modified must be called for setData but NOT for
102 # setFileName. If Modified is not called, only the first file
103 # will always be displayed. The reason Modified is used is
104 # because the same temporary file name is always used
105 # (previous file is overwritten). Modified MUST NOT be used in
106 # setFileName, it can cause incorrect output such as map.
107 self.__vtk_xml_reader.Modified()
108 # Update must be called after Modified. If Update is called before
109 # Modified, then the first/second image(s) may not be updated
110 # correctly.
111 self.__vtk_xml_reader.Update()
112 self.__output = self.__vtk_xml_reader.GetOutput()
113 self.__get_attribute_lists()
114
115 if(self.__count > 0):
116 self._updateRange()
117
118 self.__count+=1
119 else:
120 raise ValueError("Source type %s does not support 'setData'\n" \
121 % self.__source)
122
123 def setActiveScalar(self, scalar):
124 """
125 Specify the scalar field to load.
126
127 @type scalar: String
128 @param scalar: Scalar field to load from the file.
129 """
130
131 # Check whether the specified scalar is available in either point
132 # or cell data. If not available, program exits.
133
134 # NOTE: This check is similar to the check used in _getScalarRange
135 # but this is used only when a scalar attribute has been specified.
136 if scalar in self.__point_attribute['scalars']:
137 self._getOutput().GetPointData().SetActiveScalars(scalar)
138 elif scalar in self.__cell_attribute['scalars']:
139 self._getOutput().GetCellData().SetActiveScalars(scalar)
140 else:
141 print "\nERROR: No scalar called '%s' is available.\n" % scalar
142 sys.exit(1)
143
144 def setActiveVector(self, vector):
145 """
146 Specify the vector field to load.
147
148 @type vector: String
149 @param vector: Vector field to load from the file.
150 """
151
152 # Check whether the specified vector is available in either point
153 # or cell data. If not available, program exits.
154
155 # NOTE: This check is similar to the check used in _getVectorRange
156 # but this is used only when a vector attribute has been specified.
157 if vector in self.__point_attribute['vectors']:
158 self._getOutput().GetPointData().SetActiveVectors(vector)
159 elif vector in self.__cell_attribute['vectors']:
160 self._getOutput().GetCellData().SetActiveVectors(vector)
161 else:
162 print "\nERROR: No vector called '%s' is available.\n" % vector
163 sys.exit(1)
164
165 def setActiveTensor(self, tensor):
166 """
167 Specify the tensor field to load.
168
169 @type tensor: String
170 @param tensor: Tensor field to load from the file.
171 """
172
173 # Check whether the specified tensor is available in either point
174 # or cell data. If not available, program exits.
175
176 # NOTE: This check is similar to the check used in _getTensorRange
177 # but this is used only when a tensor attribute has been specified.
178 if tensor in self.__point_attribute['tensors']:
179 self._getOutput().GetPointData().SetActiveTensors(tensor)
180 elif tensor in self.__cell_attribute['tensors']:
181 self._getOutput().GetCellData().SetActiveTensors(tensor)
182 else:
183 print "\nERROR: No tensor called '%s' is available.\n" % tensor
184 sys.exit(0)
185
186 # 'object' is set to 'None' because some types of visualization have
187 # two ranges that needs to be updated while others only have one.
188 def _paramForUpdatingMultipleSources(self, viz_type, color_mode, mapper,
189 object = None):
190 """
191 Parameters required to update the necessary data when two or more
192 files or escript objects are read.
193
194 @type viz_type: : L{VizType <constant.VizType>} constant
195 @param viz_type: Type if visualization
196 @type color_mode: L{ColorMode <constant.ColorMode>} constant
197 @param color_mode: Type of color mode
198 @type mapper: vtkDataSetMapper
199 @param mapper: Mapped data
200 @type object: vtkPolyDataAlgorithm (i.e. vtkContourFilter, vtkGlyph3D, \
201 etc)
202 @param object: Polygonal data
203 """
204
205 self.__viz_type = viz_type
206 self.__color_mode = color_mode
207 self.__mapper = mapper
208 self.__object = object
209
210 def _updateRange(self):
211 """
212 Update the necessary range(s) when two or more files or escript objects
213 are read.
214 """
215
216 if self.__viz_type == VizType.MAP or \
217 self.__viz_type == VizType.ELLIPSOID or \
218 self.__viz_type == VizType.CARPET:
219 self.__mapper.SetScalarRange(self._getScalarRange())
220 elif self.__viz_type == VizType.VELOCITY:
221 if self.__color_mode == ColorMode.VECTOR:
222 self.__object.SetRange(self._getVectorRange())
223 self.__mapper.SetScalarRange(self._getVectorRange())
224 elif self.__color_mode == ColorMode.SCALAR:
225 self.__object.SetRange(self._getScalarRange())
226 self.__mapper.SetScalarRange(self._getScalarRange())
227 elif self.__viz_type == VizType.CONTOUR:
228 self.__object.GenerateValues(
229 self.__object.GetNumberOfContours(),
230 self._getScalarRange()[0],
231 self._getScalarRange()[1])
232 self.__mapper.SetScalarRange(self._getScalarRange())
233 elif self.__viz_type == VizType.STREAMLINE:
234 if self.__color_mode == ColorMode.VECTOR:
235 self.__mapper.SetScalarRange(self._getVectorRange())
236 elif self.__color_mode == ColorMode.SCALAR:
237 self.__mapper.SetScalarRange(self._getScalarRange())
238
239 def __get_array_type(self, arr):
240 """
241 Return whether an array type is scalar, vector or tensor by looking
242 at the number of components in the array.
243
244 @type arr: vtkDataArray
245 @param arr: An array from the source.
246 @rtype: String
247 @return: Array type ('scalar', vector' or 'tensor')
248 """
249
250 # Number of components in an array.
251 num_components = arr.GetNumberOfComponents()
252
253 if num_components == 1:
254 return 'scalars'
255 elif num_components == 3:
256 return 'vectors'
257 elif num_components == 9:
258 return 'tensors'
259
260 def __get_attribute_list(self, data):
261 """
262 Return the available scalar, vector and tensor attributes
263 (either point or cell data).
264
265 @type data: vtkPointData or vtkCellData
266 @param data: Available point data or cell data from the source
267 @rtype: Dictionary
268 @return: Dictionary containing the available scalar, vector and \
269 tensor attributes
270 """
271
272 attribute = {'scalars':[], 'vectors':[], 'tensors':[]}
273 if data:
274 num_arrays = data.GetNumberOfArrays() # Number of arrays.
275 for i in range(num_arrays):
276 name = data.GetArrayName(i) # Get an array name.
277 type = self.__get_array_type(data.GetArray(i)) # Get array type.
278 attribute[type].extend([name]) # Add array name to dictionary.
279
280 return attribute
281
282 def __get_attribute_lists(self):
283 """
284 Get all the available point and cell data attributes from the source.
285 """
286
287 # Get all the available point data attributes into a list.
288 self.__point_attribute = \
289 self.__get_attribute_list(self._getOutput().GetPointData())
290 # Get all the available cell data attribute into another list.
291 self.__cell_attribute = \
292 self.__get_attribute_list(self._getOutput().GetCellData())
293
294 def _getScalarRange(self):
295 """
296 Return the scalar range.
297
298 @rtype: Two column tuple containing numbers
299 @return: Scalar range
300 """
301
302 # Check whether any scalar is available in either point or cell data.
303 # If not available, program exits.
304
305 # NOTE: This check is similar to the check used in _setActiveScalar
306 # but this is used only when no scalar attribute has been specified.
307 if(len(self.__point_attribute['scalars']) != 0):
308 return self._getOutput().GetPointData().GetScalars().GetRange(-1)
309 elif(len(self.__cell_attribute['scalars']) != 0):
310 return self._getOutput().GetCellData().GetScalars().GetRange(-1)
311 else:
312 print "\nERROR: No scalar is available.\n"
313 sys.exit(1)
314
315 def _getVectorRange(self):
316 """
317 Return the vector range.
318
319 @rtype: Two column tuple containing numbers
320 @return: Vector range
321 """
322
323 # Check whether any vector is available in either point or cell data.
324 # If not available, program exits.
325
326 # NOTE: This check is similar to the check used in _setActiveVector
327 # but this is used only when no vector attribute has been specified.
328
329 # NOTE: Generally GetRange(-1) returns the correct vector range.
330 # However, there are certain data sets where GetRange(-1) seems
331 # to return incorrect mimimum vector although the maximum vector is
332 # correct. As a result, the mimimum vector has been hard coded to 0.0
333 # to accommodate for the incorrect cases.
334 if(len(self.__point_attribute['vectors']) != 0):
335 vector_range = \
336 self._getOutput().GetPointData().GetVectors().GetRange(-1)
337 return (0.0, vector_range[1])
338 elif(len(self.__cell_attribute['vectors']) != 0):
339 vector_range = \
340 self._getOutput().GetCellData().GetVectors().GetRange(-1)
341 return (0.0, vector_range[1])
342 else:
343 print "\nERROR: No vector is available.\n"
344 sys.exit(0)
345
346 def _getTensorRange(self):
347 """
348 Return the tensor range.
349
350 @rtype: Two column tuple containing numbers
351 @return: Tensor range
352 """
353
354 # Check whether any tensor is available in either point or cell data.
355 # If not available, program exits.
356
357 # NOTE: This check is similar to the check used in _setActiveTensor
358 # but this is used only when no tensor attribute has been specified.
359 if(len(self.__point_attribute['tensors']) != 0):
360 return self._getOutput().GetPointData().GetTensors().GetRange(-1)
361 elif(len(self.__cell_attribute['tensors']) != 0):
362 return self._getOutput().GetCellData().GetTensors().GetRange(-1)
363 else:
364 print "\nERROR: No tensor is available.\n"
365 sys.exit(1)
366
367 def _getOutput(self):
368 """
369 Return the output of the data collector.
370
371 @rtype: vtkUnstructuredGrid
372 @return: Unstructured grid
373 """
374
375 return self.__output
376

  ViewVC Help
Powered by ViewVC 1.1.26