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

  ViewVC Help
Powered by ViewVC 1.1.26