pyvisi/py_src/datacollector.py

"""
@author: John NGUI
"""

import vtk
import tempfile, os, sys
from constant import Source, VizType, ColorMode
try:
    import esys.escript 
except ImportError:
    print "Warning: importing esys.escript failed."

class DataCollector:
    """
    Class that defines a data collector which deals with the source 
    of data for the visualisation.
    """

    def __init__(self, source = Source.XML):
        """
        Initialise the data collector.

        @type source: L{Source <constant.Source>} constant
        @param source: Source type
        """

        self.__source = source
        self.__count = 0 # Keeps track of the number of files/sources read.

        if(source == Source.XML): # Source is an XML file.
            self.__vtk_xml_reader = vtk.vtkXMLUnstructuredGridReader()
        # Source is a escript data object using a temp file in the background. 
        elif (self.__source == Source.ESCRIPT):
            self.__vtk_xml_reader = vtk.vtkXMLUnstructuredGridReader()
            # Create a temporary .xml file and retrieves its path.
            self.__tmp_file = tempfile.mkstemp(suffix=".xml")[1]
            #self.__vtk_xml_reader.SetFileName(self.__tmp_file)

    def __del__(self):
        """
        Perform some clean up of the temporary file.
        """

        if (self.__source == Source.ESCRIPT):
            if os.access(self.__tmp_file,os.F_OK): os.unlink(self.__tmp_file)

    def setFileName(self, file_name):
        """
        Set the source file name to read. 

        @type file_name: String
        @param file_name: Name of the file to read
        """

        if(self.__source == Source.XML):
            self.__vtk_xml_reader.SetFileName(file_name)

            # Update must be called after SetFileName to make the reader 
            # up to date. Otherwise, some output values may be incorrect.
            self.__vtk_xml_reader.Update() 
            self.__output = self.__vtk_xml_reader.GetOutput()
            self.__get_attribute_lists()
            
            # Count has to be larger than zero because when setFileName is 
            # called for the first time, the data set mapper has not yet been
            # instantiated. Therefore, the range of the mapper can only be
            # updated after the first file/source has been read.
            if(self.__count > 0):
                self._updateRange(None)

            self.__count+=1

        else:
            raise ValueError("Source type %s does not support 'setFileName'\n" \
                    % self.__source)

    def setData(self,**args):
        """
        Create data using the <name>=<data> pairing. Assumption is made
        that the data will be given in the appropriate format.
        """

        if self.__source == Source.ESCRIPT:
            esys.escript.saveVTK(self.__tmp_file,**args)
            self.__vtk_xml_reader.SetFileName(self.__tmp_file)
            # Modified must be called for setData but NOT for
            # setFileName. If Modified is not called, only the first file 
            # will always be displayed. The reason Modified is used is 
            # because the same temporary file name is always used 
            # (old file is overwritten). Modified MUST NOT be used in 
            # setFileName, it can cause incorrect output such as map.
            self.__vtk_xml_reader.Modified()
            # Update must be called after Modified. If Update is called before
            # Modified, then the first/second image(s) may not be updated
            # correctly.
            self.__vtk_xml_reader.Update()
            self.__output = self.__vtk_xml_reader.GetOutput()
            self.__get_attribute_lists()
        else:
            raise ValueError("Source type %s does not support 'setData'\n" \
                    % self.__source)

    def setActiveScalar(self, scalar):
        """
        Specify the scalar field to load. 

        @type scalar: String
        @param scalar: Scalar field to load from the file. 
        """
        
        # Check whether the specified scalar is available in either point
        # or cell data. If not available, program exits.

        # NOTE: This check is similar to the check used in _getScalarRange 
        # but is used only when a scalar attribute has been specified.
        if scalar in self.__point_attribute['scalars']:
            self._getOutput().GetPointData().SetActiveScalars(scalar)
            print "SCALAR POINT DATA"
        elif scalar in self.__cell_attribute['scalars']:
            self._getOutput().GetCellData().SetActiveScalars(scalar)
            print "SCALAR CELL DATA"
        else:
            print "\nERROR: No scalar called '%s' is available.\n" % scalar
            sys.exit(1) 

    def setActiveVector(self, vector):
        """
        Specify the vector field to load.

        @type vector: String
        @param vector: Vector field to load from the file. 
        """
        
        # Check whether the specified vector is available in either point
        # or cell data. If not available, program exits.

        # NOTE: This check is similar to the check used in _getVectorRange 
        # but is used only when a vector attribute has been specified.
        if vector in self.__point_attribute['vectors']:
            self._getOutput().GetPointData().SetActiveVectors(vector)
            print "VECTOR POINT DATA" 
        elif vector in self.__cell_attribute['vectors']:
            self._getOutput().GetCellData().SetActiveVectors(vector)
            print "VECTOR CELL DATA"
        else:
            print "\nERROR: No vector called '%s' is available.\n" % vector
            sys.exit(1) 
            
    def setActiveTensor(self, tensor):
        """
        Specify the tensor field to load.

        @type tensor: String
        @param tensor: Tensor field to load from the file. 
        """

        # Check whether the specified tensor is available in either point
        # or cell data. If not available, program exits.

        # NOTE: This check is similar to the check used in _getTensorRange 
        # but is used only when a tensor attribute has been specified.
        if tensor in self.__point_attribute['tensors']:
            self._getOutput().GetPointData().SetActiveTensors(tensor)
        elif tensor in self.__cell_attribute['tensors']:
            self._getOutput().GetCellData().SetActiveTensors(tensor)
        else:
            print "\nERROR: No tensor called '%s' is available.\n" % tensor
            sys.exit(0) 

    # 'object' is set to 'None' because some types of visualization have
    # two different ranges that need to be updated while others only have one.
    def _paramForUpdatingMultipleSources(self, viz_type, color_mode, mapper,
            object = None):
        """
        Parameters required to update the necessary range when two or more 
        files are read.

        @type viz_type: : L{VizType <constant.VizType>} constant
        @param viz_type: Type if visualization (i.e. Map, Velocity, etc)
        @type color_mode: L{ColorMode <constant.ColorMode>} constant
        @param color_mode: Type of color mode
        @type mapper: vtkDataSetMapper
        @param mapper: Mapped data
        @type object: vtkPolyDataAlgorith (i.e. vtkContourFilter, vtkGlyph3D, \
                etc)
        @param object: Poly data
        """

        self.__viz_type = viz_type
        self.__color_mode = color_mode
        self.__mapper = mapper
        self.__object = object

    def _updateRange(self, range):
        """
        Update the necessary range when two or more sources are read in.
        """

        if self.__viz_type == VizType.MAP or \
                self.__viz_type == VizType.ELLIPSOID or \
                self.__viz_type == VizType.CARPET:
            self.__mapper.SetScalarRange(self._getScalarRange())
        elif self.__viz_type == VizType.VELOCITY:
            if self.__color_mode == ColorMode.VECTOR:
                self.__object.SetRange(self._getVectorRange())
                self.__mapper.SetScalarRange(self._getVectorRange())
            elif self.__color_mode == ColorMode.SCALAR:
                self.__object.SetRange(self._getScalarRange())
                self.__mapper.SetScalarRange(self._getScalarRange())
        elif self.__viz_type == VizType.CONTOUR:
            self.__object.GenerateValues(
                    self.__object.GetNumberOfContours(),
                    self._getScalarRange()[0],
                    self._getScalarRange()[1])
            self.__mapper.SetScalarRange(self._getScalarRange())
        elif self.__viz_type == VizType.STREAMLINE:
            if self.__color_mode == ColorMode.VECTOR:
                self.__mapper.SetScalarRange(self._getVectorRange())
            elif self.__color_mode == ColorMode.SCALAR:
                self.__mapper.SetScalarRange(self._getScalarRange())

    def __get_array_type(self, arr):
        """
        Return if the array type is a scalar, vector or tensor by looking 
        at the number of components in the array.

        @type arr: vtkDataArray
        @param arr: An array from the source.
        @rtype: String
        @return: Array type ('scalar', vector', 'tensor')
        """

        # Number of components in an array.
        num_components = arr.GetNumberOfComponents() 

        if num_components == 1:
            return 'scalars'
        elif num_components == 3:
            return 'vectors'
        elif num_components == 9:
            return 'tensors'

    def __get_attribute_list(self, data): 
        """
        Return the available scalar, vector and tensor attributes 
        (either point or cell data).

        @type data: vtkPointData or vtkCellData
        @param data: Available point data or cell data from the source
        @rtype: Dictionary
        @return: Dictionary containing the available scalar, vector and \
                tensor attributes
        """

        attribute = {'scalars':[], 'vectors':[], 'tensors':[]} 
        if data:
            num_arrays = data.GetNumberOfArrays() # Number of arrays.
            for i in range(num_arrays):
                name = data.GetArrayName(i) # Get an array name.
                type = self.__get_array_type(data.GetArray(i)) # Get array type.
                attribute[type].extend([name]) # Add array name to dictionary.

        return attribute

    def __get_attribute_lists(self):
        """
        Get all the available point and cell data attributes from the source.
        """

        # Get all the available point data attributes into a list.
        self.__point_attribute = \
                self.__get_attribute_list(self._getOutput().GetPointData())
        # Get all the available cell data attribute into another list.  
        self.__cell_attribute = \
                self.__get_attribute_list(self._getOutput().GetCellData())

    def _getScalarRange(self):
        """
        Return the scalar range.

        @rtype: Two column tuple containing numbers
        @return: Scalar range
        """

        # Check whether any tensor is available in either point or cell data. 
        # If not available, program exits.

        # NOTE: This check is similar to the check used in _setActiveScalar 
        # but is used only when no scalar attribute has been specified.
        if(len(self.__point_attribute['scalars']) != 0):
            return self._getOutput().GetPointData().GetScalars().GetRange(-1)
        elif(len(self.__cell_attribute['scalars']) != 0):
            return self._getOutput().GetCellData().GetScalars().GetRange(-1)
        else:
            print "\nERROR: No scalar is available.\n"  
            sys.exit(1)

    def _getVectorRange(self):
        """
        Return the vector range.
    
        @rtype: Two column tuple containing numbers
        @return: Vector range
        """
        
        # Check whether any vector is available in either point or cell data. 
        # If not available, program exits.

        # NOTE: This check is similar to the check used in _setActiveVector 
        # but is used only when no vector attribute has been specified.

        # NOTE: Generally GetRange(-1) returns the correct vector range. 
        # However, there are certain data sets where GetRange(-1) seems 
        # to return incorrect mimimum vector although the maximum vector is 
        # correct. As a result, the mimimum vector has been hard coded to 0.0
        # to accommodate the incorrect cases.
        if(len(self.__point_attribute['vectors']) != 0):
            vector_range = \
                    self._getOutput().GetPointData().GetVectors().GetRange(-1)
            return (0.0, vector_range[1])
        elif(len(self.__cell_attribute['vectors']) != 0): 
            vector_range = \
                    self._getOutput().GetCellData().GetVectors().GetRange(-1)
            return (0.0, vector_range[1])
        else:
            print "\nERROR: No vector is available.\n"  
            sys.exit(0)

    def _getTensorRange(self):
        """
        Return the tensor range.

        @rtype: Two column table
        @return: Tensor range
        """

        # Check whether any tensor is available in either point or cell data. 
        # If not available, program exits.

        # NOTE: This check is similar to the check used in _setActiveTensor 
        # but is used only when no tensor attribute has been specified.
        if(len(self.__point_attribute['tensors']) != 0):
            return self._getOutput().GetPointData().GetTensors().GetRange(-1)
        elif(len(self.__cell_attribute['tensors']) != 0): 
            return self._getOutput().GetCellData().GetTensors().GetRange(-1)
        else:
            print "\nERROR: No tensor is available.\n"  
            sys.exit(1)

    def _getOutput(self):
        """
        Return the output of the data collector.

        @rtype: vtkUnstructuredGrid
        @return: Unstructured grid
        """

        return self.__output

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