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 |
|