1 |
""" |
2 |
@author: John NGUI |
3 |
""" |
4 |
|
5 |
import vtk |
6 |
import tempfile, os, sys |
7 |
from constant import Source, 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 a XML file or from an escript object directly. Writing XML |
17 |
files are expensive, but this approach has the advantage given that the |
18 |
results can be analyzed easily after the simulation has completed. |
19 |
|
20 |
@attention: A DataCollector instance can only be used to specify one |
21 |
scalar, vector and tensor attribute from a source at any one time. If a |
22 |
second scalar, vector or tensor attribute needs to be specified from the |
23 |
same source, a second DataCollector instance must be created. |
24 |
""" |
25 |
|
26 |
def __init__(self, source = Source.XML): |
27 |
""" |
28 |
Initialise the data collector. |
29 |
|
30 |
@type source: L{Source <constant.Source>} constant |
31 |
@param source: Source type |
32 |
""" |
33 |
|
34 |
self.__source = source |
35 |
# Keeps track on whether DataCollector have been modified. |
36 |
self.__modified = True |
37 |
# Keeps track on whether any specific scalar, vector or tensor |
38 |
# field have been specified. |
39 |
self.__set_scalar = False |
40 |
self.__set_vector= False |
41 |
self.__set_tensor= False |
42 |
|
43 |
if(source == Source.XML): # Source is an XML file. |
44 |
self.__vtk_xml_reader = vtk.vtkXMLUnstructuredGridReader() |
45 |
# Source is a escript data object using a temp file in the background. |
46 |
elif (self.__source == Source.ESCRIPT): |
47 |
self.__vtk_xml_reader = vtk.vtkXMLUnstructuredGridReader() |
48 |
# Create a temporary .xml file and retrieve its path. |
49 |
self.__tmp_file = tempfile.mkstemp(suffix=".xml")[1] |
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 XML file name to read. |
62 |
|
63 |
@type file_name: String |
64 |
@param file_name: Name of the file to read |
65 |
""" |
66 |
|
67 |
self.__modified = True |
68 |
|
69 |
if(self.__source == Source.XML): |
70 |
# Check whether the specified file exists, otherwise exit. |
71 |
if not(os.access(file_name, os.F_OK)): |
72 |
raise IOError("\nERROR: '%s' file does NOT exists.\n" % \ |
73 |
file_name) |
74 |
|
75 |
self.__vtk_xml_reader.SetFileName(file_name) |
76 |
# Update must be called after SetFileName to make the reader |
77 |
# up-to-date. Otherwise, some output values may be incorrect. |
78 |
self.__vtk_xml_reader.Update() |
79 |
self.__get_attribute_lists() |
80 |
else: |
81 |
raise ValueError("Source type %s does not support \ |
82 |
'setFileName'\n" % self.__source) |
83 |
|
84 |
def setData(self,**args): |
85 |
""" |
86 |
Create data using the <name>=<data> pairing. Assumption is made |
87 |
that the data will be given in the appropriate format. |
88 |
""" |
89 |
|
90 |
self.__modified = True |
91 |
|
92 |
if self.__source == Source.ESCRIPT: |
93 |
esys.escript.saveVTK(self.__tmp_file,**args) |
94 |
self.__vtk_xml_reader.SetFileName(self.__tmp_file) |
95 |
|
96 |
# Modified must be called for setData but NOT for |
97 |
# setFileName. If Modified is not called, only the first file |
98 |
# will always be displayed. The reason Modified is used is |
99 |
# because the same temporary file name is always used |
100 |
# (previous file is overwritten). Modified MUST NOT be used in |
101 |
# setFileName, it can cause incorrect output such as map. |
102 |
self.__vtk_xml_reader.Modified() |
103 |
|
104 |
# Update must be called after Modified. If Update is called before |
105 |
# Modified, then the first/second image(s) may not be updated |
106 |
# correctly. |
107 |
self.__vtk_xml_reader.Update() |
108 |
self.__get_attribute_lists() |
109 |
else: |
110 |
raise ValueError("Source type %s does not support 'setData'\n" \ |
111 |
% self.__source) |
112 |
|
113 |
# This method is used to delay the execution of setting the active scalar |
114 |
# until 'setFileName' or 'setData' have been executed. |
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 |
self.__set_scalar = True |
124 |
self.__active_scalar = scalar |
125 |
|
126 |
def _setActiveScalar(self): |
127 |
""" |
128 |
Load the specified scalar field. |
129 |
""" |
130 |
|
131 |
# Check whether the specified scalar is available in either point |
132 |
# or cell data. If not available, error is raised. |
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 self.__active_scalar in self.__point_attribute['scalars']: |
137 |
self._getDataCollectorOutput().GetPointData().SetActiveScalars( |
138 |
self.__active_scalar) |
139 |
elif self.__active_scalar in self.__cell_attribute['scalars']: |
140 |
self._getDataCollectorOutput().GetCellData().SetActiveScalars( |
141 |
self.__active_scalar) |
142 |
else: |
143 |
raise IOError("ERROR: No scalar called '%s' is available." % \ |
144 |
self.__active_scalar) |
145 |
|
146 |
# This method is used to delay the execution of setting the active vector |
147 |
# until 'setFileName' or 'setData' have been executed. |
148 |
def setActiveVector(self, vector): |
149 |
""" |
150 |
Specify the vector field to load. |
151 |
|
152 |
@type vector: String |
153 |
@param vector: Vector field to load from the file. |
154 |
""" |
155 |
|
156 |
self.__set_vector = True |
157 |
self.__active_vector = vector |
158 |
|
159 |
def _setActiveVector(self): |
160 |
""" |
161 |
Load the specified vector field. |
162 |
""" |
163 |
|
164 |
# Check whether the specified vector is available in either point |
165 |
# or cell data. If not available, error is raised. |
166 |
|
167 |
# NOTE: This check is similar to the check used in _getVectorRange |
168 |
# but this is used only when a vector attribute has been specified. |
169 |
if self.__active_vector in self.__point_attribute['vectors']: |
170 |
self._getDataCollectorOutput().GetPointData().SetActiveVectors( |
171 |
self.__active_vector) |
172 |
elif self.__active_vector in self.__cell_attribute['vectors']: |
173 |
self._getDataCollectorOutput().GetCellData().SetActiveVectors( |
174 |
self.__active_vector) |
175 |
else: |
176 |
raise IOError("ERROR: No vector called '%s' is available." % \ |
177 |
self.__active_vector) |
178 |
|
179 |
# This method is used to delay the execution of setting the active tensor |
180 |
# until 'setFileName' or 'setData' have been executed. |
181 |
def setActiveTensor(self, tensor): |
182 |
""" |
183 |
Specify the tensor field to load. |
184 |
|
185 |
@type tensor: String |
186 |
@param tensor: Tensor field to load from the file. |
187 |
""" |
188 |
|
189 |
self.__set_tensor = True |
190 |
self.__active_tensor = tensor |
191 |
|
192 |
def _setActiveTensor(self): |
193 |
""" |
194 |
Load the the specified tensor field. |
195 |
""" |
196 |
|
197 |
# Check whether the specified tensor is available in either point |
198 |
# or cell data. If not available, error is raised. |
199 |
|
200 |
# NOTE: This check is similar to the check used in _getTensorRange |
201 |
# but this is used only when a tensor attribute has been specified. |
202 |
if self.__active_tensor in self.__point_attribute['tensors']: |
203 |
self._getDataCollectorOutput().GetPointData().SetActiveTensors( |
204 |
self.__active_tensor) |
205 |
elif self.__active_tensor in self.__cell_attribute['tensors']: |
206 |
self._getDataCollectorOutput().GetCellData().SetActiveTensors( |
207 |
self.__active_tensor) |
208 |
else: |
209 |
raise IOError("ERROR: No tensor called '%s' is available." % \ |
210 |
self.__active_tensor) |
211 |
|
212 |
def __get_array_type(self, arr): |
213 |
""" |
214 |
Return whether an array type is scalar, vector or tensor by looking |
215 |
at the number of components in the array. |
216 |
|
217 |
@type arr: vtkDataArray |
218 |
@param arr: An array from the source. |
219 |
@rtype: String |
220 |
@return: Array type ('scalar', vector' or 'tensor') |
221 |
""" |
222 |
|
223 |
# Number of components in an array. |
224 |
num_components = arr.GetNumberOfComponents() |
225 |
|
226 |
if num_components == 1: |
227 |
return 'scalars' |
228 |
elif num_components == 3: |
229 |
return 'vectors' |
230 |
elif num_components == 9: |
231 |
return 'tensors' |
232 |
|
233 |
def __get_attribute_list(self, data): |
234 |
""" |
235 |
Return the available scalar, vector and tensor attributes |
236 |
(either point or cell data). |
237 |
|
238 |
@type data: vtkPointData or vtkCellData |
239 |
@param data: Available point data or cell data from the source |
240 |
@rtype: Dictionary |
241 |
@return: Dictionary containing the available scalar, vector and \ |
242 |
tensor attributes |
243 |
""" |
244 |
|
245 |
attribute = {'scalars':[], 'vectors':[], 'tensors':[]} |
246 |
if data: |
247 |
num_arrays = data.GetNumberOfArrays() # Number of arrays. |
248 |
for i in range(num_arrays): |
249 |
name = data.GetArrayName(i) # Get an array name. |
250 |
type = self.__get_array_type(data.GetArray(i)) # Get array type. |
251 |
attribute[type].extend([name]) # Add array name to dictionary. |
252 |
|
253 |
return attribute |
254 |
|
255 |
def __get_attribute_lists(self): |
256 |
""" |
257 |
Get all the available point and cell data attributes from the source. |
258 |
""" |
259 |
|
260 |
# Get all the available point data attributes into a list. |
261 |
self.__point_attribute = \ |
262 |
self.__get_attribute_list( |
263 |
self._getDataCollectorOutput().GetPointData()) |
264 |
|
265 |
# Get all the available cell data attribute into another list. |
266 |
self.__cell_attribute = \ |
267 |
self.__get_attribute_list( |
268 |
self._getDataCollectorOutput().GetCellData()) |
269 |
|
270 |
def _getScalarRange(self): |
271 |
""" |
272 |
Return the scalar range. |
273 |
|
274 |
@rtype: Two column tuple containing numbers |
275 |
@return: Scalar range |
276 |
""" |
277 |
|
278 |
# Check whether any scalar is available in either point or cell data. |
279 |
# If not available, program exits. |
280 |
|
281 |
# NOTE: This check is similar to the check used in _setActiveScalar |
282 |
# but this is used only when no scalar attribute has been specified. |
283 |
if(len(self.__point_attribute['scalars']) != 0): |
284 |
return self._getDataCollectorOutput().GetPointData().\ |
285 |
GetScalars().GetRange(-1) |
286 |
elif(len(self.__cell_attribute['scalars']) != 0): |
287 |
return self._getDataCollectorOutput().GetCellData().\ |
288 |
GetScalars().GetRange(-1) |
289 |
else: |
290 |
raise IOError("\nERROR: No scalar is available.\n") |
291 |
|
292 |
def _getVectorRange(self): |
293 |
""" |
294 |
Return the vector range. |
295 |
|
296 |
@rtype: Two column tuple containing numbers |
297 |
@return: Vector range |
298 |
""" |
299 |
|
300 |
# Check whether any vector is available in either point or cell data. |
301 |
# If not available, program exits. |
302 |
|
303 |
# NOTE: This check is similar to the check used in _setActiveVector |
304 |
# but this is used only when no vector attribute has been specified. |
305 |
|
306 |
# NOTE: Generally GetRange(-1) returns the correct vector range. |
307 |
# However, there are certain data sets where GetRange(-1) seems |
308 |
# to return incorrect mimimum vector although the maximum vector is |
309 |
# correct. As a result, the mimimum vector has been hard coded to 0.0 |
310 |
# to accommodate for the incorrect cases. |
311 |
if(len(self.__point_attribute['vectors']) != 0): |
312 |
vector_range = \ |
313 |
self._getDataCollectorOutput().GetPointData().\ |
314 |
GetVectors().GetRange(-1) |
315 |
return (0.0, vector_range[1]) |
316 |
elif(len(self.__cell_attribute['vectors']) != 0): |
317 |
vector_range = \ |
318 |
self._getDataCollectorOutput().GetCellData().\ |
319 |
GetVectors().GetRange(-1) |
320 |
return (0.0, vector_range[1]) |
321 |
else: |
322 |
raise IOError("\nERROR: No vector is available.\n") |
323 |
|
324 |
def _getTensorRange(self): |
325 |
""" |
326 |
Return the tensor range. |
327 |
|
328 |
@rtype: Two column tuple containing numbers |
329 |
@return: Tensor range |
330 |
""" |
331 |
|
332 |
# Check whether any tensor is available in either point or cell data. |
333 |
# If not available, program exits. |
334 |
|
335 |
# NOTE: This check is similar to the check used in _setActiveTensor |
336 |
# but this is used only when no tensor attribute has been specified. |
337 |
if(len(self.__point_attribute['tensors']) != 0): |
338 |
return self._getDataCollectorOutput().GetPointData().\ |
339 |
GetTensors().GetRange(-1) |
340 |
elif(len(self.__cell_attribute['tensors']) != 0): |
341 |
return self._getDataCollectorOutput().GetCellData().\ |
342 |
GetTensors().GetRange(-1) |
343 |
else: |
344 |
raise IOError("\nERROR: No tensor is available.\n") |
345 |
|
346 |
def _getDataCollectorOutput(self): |
347 |
""" |
348 |
Return the output of the data collector. |
349 |
|
350 |
@rtype: vtkUnstructuredGrid |
351 |
@return: Unstructured grid |
352 |
""" |
353 |
return self.__vtk_xml_reader.GetOutput() |
354 |
|
355 |
def _isModified(self): |
356 |
""" |
357 |
Return whether the DataCollector has been modified. |
358 |
|
359 |
@rtype: Boolean |
360 |
@return: True or False |
361 |
""" |
362 |
|
363 |
if(self.__modified == True): |
364 |
self.__modified = False |
365 |
return True |
366 |
else: |
367 |
return False |
368 |
|
369 |
def _isScalarSet(self): |
370 |
""" |
371 |
Return whether a specific scalar field has been specified. |
372 |
|
373 |
@rtype: Boolean |
374 |
@return: True or False |
375 |
""" |
376 |
|
377 |
return self.__set_scalar |
378 |
|
379 |
def _isVectorSet(self): |
380 |
""" |
381 |
Return whether a specific vector field has been specified. |
382 |
|
383 |
@rtype: Boolean |
384 |
@return: True or False |
385 |
""" |
386 |
|
387 |
return self.__set_vector |
388 |
|
389 |
def _isTensorSet(self): |
390 |
""" |
391 |
Return whether a specific tensor field has been specified. |
392 |
|
393 |
@rtype: Boolean |
394 |
@return: True or False |
395 |
""" |
396 |
|
397 |
return self.__set_tensor |
398 |
|
399 |
def _getCenter(self): |
400 |
""" |
401 |
Return the center of the rendered object. |
402 |
|
403 |
@rtype: Three column tuple containing numbers |
404 |
@return: Center of the rendered object |
405 |
""" |
406 |
|
407 |
return self._getDataCollectorOutput().GetCenter() |
408 |
|