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

  ViewVC Help
Powered by ViewVC 1.1.26