/[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 1147 - (show annotations)
Wed May 16 06:39:11 2007 UTC (12 years, 3 months ago) by ksteube
File MIME type: text/x-python
File size: 12913 byte(s)
Added back in some files that were accidentally deleted.

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 try:
14 PYVISI_WORKDIR=os.environ['PYVISI_WORKDIR']
15 except KeyError:
16 PYVISI_WORKDIR='.'
17 try:
18 PYVISI_TEST_DATA_ROOT=os.environ['PYVISI_TEST_DATA_ROOT']
19 except KeyError:
20 PYVISI_TEST_DATA_ROOT='.'
21
22 class DataCollector:
23 """
24 Class that defines a data collector. A data collector is used to read
25 data from an XML file or from an escript object directly.
26
27 @attention: A DataCollector instance can only be used to specify one
28 scalar, vector and tensor attribute from a source at any one time. If a
29 second scalar, vector or tensor attribute needs to be specified from the
30 same source, a second DataCollector instance must be created.
31
32 @attention: When a series of XML files or escript objects are read
33 (using 'setFileName' or 'setData' in a for-loop), the 'setActiveScalar' /
34 'setActiveVector' / 'setActiveTensor' have to be called for each new file
35 (provided a specific field needs to be loaded) as all active fields
36 specified from the previous file goes back to the default once a new file
37 is read.
38 """
39
40 def __init__(self, source = Source.XML):
41 """
42 Initialise the data collector.
43
44 @type source: L{Source <constant.Source>} constant
45 @param source: Source type
46 """
47
48 self.__source = source
49 self.__count = 0 # Keeps track of the number of files/sources read.
50
51 if(source == Source.XML): # Source is an XML file.
52 self.__vtk_xml_reader = vtk.vtkXMLUnstructuredGridReader()
53 # Source is a escript data object using a temp file in the background.
54 elif (self.__source == Source.ESCRIPT):
55 self.__vtk_xml_reader = vtk.vtkXMLUnstructuredGridReader()
56 # Create a temporary .xml file and retrieve its path.
57 self.__tmp_file = tempfile.mkstemp(suffix=".xml")[1]
58
59 def __del__(self):
60 """
61 Perform some clean up of the temporary file.
62 """
63
64 if (self.__source == Source.ESCRIPT):
65 if os.access(self.__tmp_file,os.F_OK): os.unlink(self.__tmp_file)
66
67 def setFileName(self, file_name):
68 """
69 Set the XML file name to read.
70
71 @type file_name: String
72 @param file_name: Name of the file to read
73 """
74
75 if(self.__source == Source.XML):
76 # Check whether the specified file exists, otherwise exit.
77 if not(os.access(file_name, os.F_OK)):
78 raise IOError("ERROR: '%s' file does NOT exists." % file_name)
79
80 self.__vtk_xml_reader.SetFileName(file_name)
81 # Update must be called after SetFileName to make the reader
82 # up-to-date. Otherwise, some output values may be incorrect.
83 self.__vtk_xml_reader.Update()
84 self.__output = self.__vtk_xml_reader.GetOutput()
85 self.__get_attribute_lists()
86
87 # Count has to be larger than zero because when setFileName is
88 # called for the first time, the data set mapper has not yet been
89 # instantiated. Therefore, the range of the mapper can only be
90 # updated after the first file/source has been read.
91 if(self.__count > 0):
92 self._updateRange()
93
94 self.__count+=1
95
96 else:
97 raise ValueError("Source type %s does not support \
98 'setFileName'\n" % self.__source)
99
100 def setData(self,**args):
101 """
102 Create data using the <name>=<data> pairing. Assumption is made
103 that the data will be given in the appropriate format.
104 """
105
106 if self.__source == Source.ESCRIPT:
107 esys.escript.saveVTK(self.__tmp_file,**args)
108 self.__vtk_xml_reader.SetFileName(self.__tmp_file)
109 # Modified must be called for setData but NOT for
110 # setFileName. If Modified is not called, only the first file
111 # will always be displayed. The reason Modified is used is
112 # because the same temporary file name is always used
113 # (previous file is overwritten). Modified MUST NOT be used in
114 # setFileName, it can cause incorrect output such as map.
115 self.__vtk_xml_reader.Modified()
116 # Update must be called after Modified. If Update is called before
117 # Modified, then the first/second image(s) may not be updated
118 # correctly.
119 self.__vtk_xml_reader.Update()
120 self.__output = self.__vtk_xml_reader.GetOutput()
121 self.__get_attribute_lists()
122
123 if(self.__count > 0):
124 self._updateRange()
125
126 self.__count+=1
127 else:
128 raise ValueError("Source type %s does not support 'setData'\n" \
129 % self.__source)
130
131 def setActiveScalar(self, scalar):
132 """
133 Specify the scalar field to load.
134
135 @type scalar: String
136 @param scalar: Scalar field to load from the file.
137 """
138
139 # Check whether the specified scalar is available in either point
140 # or cell data. If not available, program exits.
141
142 # NOTE: This check is similar to the check used in _getScalarRange
143 # but this is used only when a scalar attribute has been specified.
144 if scalar in self.__point_attribute['scalars']:
145 self._getOutput().GetPointData().SetActiveScalars(scalar)
146 elif scalar in self.__cell_attribute['scalars']:
147 self._getOutput().GetCellData().SetActiveScalars(scalar)
148 else:
149 raise IOError("ERROR: No scalar called '%s' is available." % scalar)
150
151 def setActiveVector(self, vector):
152 """
153 Specify the vector field to load.
154
155 @type vector: String
156 @param vector: Vector field to load from the file.
157 """
158
159 # Check whether the specified vector is available in either point
160 # or cell data. If not available, program exits.
161
162 # NOTE: This check is similar to the check used in _getVectorRange
163 # but this is used only when a vector attribute has been specified.
164 if vector in self.__point_attribute['vectors']:
165 self._getOutput().GetPointData().SetActiveVectors(vector)
166 elif vector in self.__cell_attribute['vectors']:
167 self._getOutput().GetCellData().SetActiveVectors(vector)
168 else:
169 raise IOError("ERROR: No vector called '%s' is available." % vector)
170
171 def setActiveTensor(self, tensor):
172 """
173 Specify the tensor field to load.
174
175 @type tensor: String
176 @param tensor: Tensor field to load from the file.
177 """
178
179 # Check whether the specified tensor is available in either point
180 # or cell data. If not available, program exits.
181
182 # NOTE: This check is similar to the check used in _getTensorRange
183 # but this is used only when a tensor attribute has been specified.
184 if tensor in self.__point_attribute['tensors']:
185 self._getOutput().GetPointData().SetActiveTensors(tensor)
186 elif tensor in self.__cell_attribute['tensors']:
187 self._getOutput().GetCellData().SetActiveTensors(tensor)
188 else:
189 raise IOError("ERROR: No tensor called '%s' is available." % tensor)
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 raise IOError("ERROR: No scalar is available.")
318
319 def _getVectorRange(self):
320 """
321 Return the vector range.
322
323 @rtype: Two column tuple containing numbers
324 @return: Vector range
325 """
326
327 # Check whether any vector is available in either point or cell data.
328 # If not available, program exits.
329
330 # NOTE: This check is similar to the check used in _setActiveVector
331 # but this is used only when no vector attribute has been specified.
332
333 # NOTE: Generally GetRange(-1) returns the correct vector range.
334 # However, there are certain data sets where GetRange(-1) seems
335 # to return incorrect mimimum vector although the maximum vector is
336 # correct. As a result, the mimimum vector has been hard coded to 0.0
337 # to accommodate for the incorrect cases.
338 if(len(self.__point_attribute['vectors']) != 0):
339 vector_range = \
340 self._getOutput().GetPointData().GetVectors().GetRange(-1)
341 return (0.0, vector_range[1])
342 elif(len(self.__cell_attribute['vectors']) != 0):
343 vector_range = \
344 self._getOutput().GetCellData().GetVectors().GetRange(-1)
345 return (0.0, vector_range[1])
346 else:
347 print "\nERROR: No vector is available.\n"
348 sys.exit(0)
349
350 def _getTensorRange(self):
351 """
352 Return the tensor range.
353
354 @rtype: Two column tuple containing numbers
355 @return: Tensor range
356 """
357
358 # Check whether any tensor is available in either point or cell data.
359 # If not available, program exits.
360
361 # NOTE: This check is similar to the check used in _setActiveTensor
362 # but this is used only when no tensor attribute has been specified.
363 if(len(self.__point_attribute['tensors']) != 0):
364 return self._getOutput().GetPointData().GetTensors().GetRange(-1)
365 elif(len(self.__cell_attribute['tensors']) != 0):
366 return self._getOutput().GetCellData().GetTensors().GetRange(-1)
367 else:
368 print "\nERROR: No tensor is available.\n"
369 sys.exit(1)
370
371 def _getOutput(self):
372 """
373 Return the output of the data collector.
374
375 @rtype: vtkUnstructuredGrid
376 @return: Unstructured grid
377 """
378
379 return self.__output
380

  ViewVC Help
Powered by ViewVC 1.1.26