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

  ViewVC Help
Powered by ViewVC 1.1.26