/[escript]/trunk/pyvisi/py_src/datacollector.py
ViewVC logotype

Annotation of /trunk/pyvisi/py_src/datacollector.py

Parent Directory Parent Directory | Revision Log Revision Log


Revision 1147 - (hide annotations)
Wed May 16 06:39:11 2007 UTC (12 years, 5 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 ksteube 1147 """
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