/[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 988 - (show annotations)
Tue Feb 20 07:21:12 2007 UTC (12 years, 2 months ago) by jongui
File MIME type: text/x-python
File size: 7823 byte(s)
Both point and cell data should now work properly with Pyvisi. Some testing have been done but a more extensive test will be done at a later stage.
1 """
2 @author: John NGUI
3 """
4
5 import vtk
6 import tempfile, os, sys
7 from constant import Source
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 dealrs 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 data type
25 """
26
27 self.__source=source
28
29 if(source == Source.XML): # Source is an XML file.
30 self.__vtk_xml_reader = vtk.vtkXMLUnstructuredGridReader()
31 # This is a temporary solution using a file:
32 elif (self.__source == Source.ESCRIPT):
33 self.__vtk_xml_reader = vtk.vtkXMLUnstructuredGridReader()
34 self.__tmp_file=tempfile.mkstemp(suffix=".xml")[1]
35 self.__vtk_xml_reader.SetFileName(self.__tmp_file)
36 self.__output = self.__vtk_xml_reader.GetOutput()
37
38 def __del__(self):
39 """
40 Perform some clean up of the temporary file.
41 """
42
43 if (self.__source == Source.ESCRIPT):
44 if os.access(self.__tmp_file,os.F_OK): os.unlink(self.__tmp_file)
45
46 def setFileName(self, file_name):
47 """
48 Set the source file name to read.
49
50 @type file_name: String
51 @param file_name: Name of the file to read
52 """
53
54 if (self.__source == Source.XML):
55 self.__vtk_xml_reader.SetFileName(file_name)
56
57 # NOTE: Update must be called after SetFileName to make the reader
58 # up to date. Otherwise, some output values may be incorrect.
59 self.__vtk_xml_reader.Update()
60 self.__output = self.__vtk_xml_reader.GetOutput()
61 self.__get_attribute_lists()
62 else:
63 raise ValueError("source type %s does not support setFileName"%self.__source)
64
65 def setData(self,**args):
66 """
67 sets the data using. <name>=<data> sets the data tagged by <name>
68 to the object <data>. It is expected that the data are given in
69 an appropriate source type.
70 """
71
72 if self.__source == Source.ESCRIPT:
73 esys.escript.saveVTK(self.__tmp_file,**args)
74 self.__vtk_xml_reader.Update()
75 else:
76 raise ValueError("source type %s does not support setData"%self.__source)
77
78 def _setActiveScalar(self, scalar):
79 """
80 Specify the scalar field to load from the source file.
81
82 @type scalar: String
83 @param scalar: Scalar field to load from the file.
84 """
85
86 if scalar in self.__point_attribute['scalars']:
87 self._getOutput().GetPointData().SetActiveScalars(scalar)
88 elif scalar in self.__cell_attribute['scalars']:
89 self._getOutput().GetCellData().SetActiveScalars(scalar)
90 else:
91 print "\nSorry, no scalar called '%s' is available.\n" % scalar
92 sys.exit(0)
93
94
95 def _setActiveVector(self, vector):
96 """
97 Specify the vector field to load from the source file.
98
99 @type vector: String
100 @param vector: Vector field to load from the file.
101 """
102
103 if vector in self.__point_attribute['vectors']:
104 self._getOutput().GetPointData().SetActiveVectors(vector)
105 elif vector in self.__cell_attribute['vectors']:
106 self._getOutput().GetCellData().SetActiveVectors(vector)
107 else:
108 print "\nSorry, no vector called '%s' is available.\n" % vector
109 sys.exit(0)
110
111
112 def _setActiveTensor(self, tensor):
113 """
114 Specify the tensor field to load from the source file.
115
116 @type tensor: String
117 @param tensor: Tensor field to load from the file.
118 """
119
120 if tensor in self.__point_attribute['tensors']:
121 self._getOutput().GetPointData().SetActiveTensors(tensor)
122 elif tensor in self.__cell_attribute['tensors']:
123 self._getOutput().GetCellData().SetActiveTensors(tensor)
124 else:
125 print "\nSorry, no tensor called '%s' is available.\n" % tensor
126 sys.exit(0)
127
128
129 def __get_array_type(self, arr):
130 num_components = arr.GetNumberOfComponents()
131
132 if num_components == 1:
133 return 'scalars'
134 elif num_components == 3:
135 return 'vectors'
136 elif num_components == 9:
137 return 'tensors'
138
139 def __get_attribute_list(self, data):
140 attribute = {'scalars':[], 'vectors':[], 'tensors':[]}
141 if data:
142 num_arrays = data.GetNumberOfArrays()
143 for i in range(num_arrays):
144 name = data.GetArrayName(i)
145 type = self.__get_array_type(data.GetArray(i))
146 attribute[type].extend([name])
147
148 return attribute
149
150 def __get_attribute_lists(self):
151 self.__point_attribute = \
152 self.__get_attribute_list(self._getOutput().GetPointData())
153
154 self.__cell_attribute = \
155 self.__get_attribute_list(self._getOutput().GetCellData())
156
157 def _getScalarRange(self):
158 """
159 Return the scalar range.
160
161 @rtype: Two column tuple containing numbers
162 @return: Scalar range
163 """
164
165 if(len(self.__point_attribute['scalars']) != 0):
166 return self._getOutput().GetPointData().GetScalars().GetRange(-1)
167 elif(len(self.__cell_attribute['scalars']) != 0):
168 return self._getOutput().GetCellData().GetScalars().GetRange(-1)
169 else:
170 print "\nSorry, no scalar is available.\n"
171 sys.exit(0)
172
173 def _getVectorRange(self):
174 """
175 Return the vector range.
176
177 @rtype: Two column tuple containing numbers
178 @return: Vector range
179 """
180
181
182 # NOTE: Generally GetRange(-1) returns the correct vector range.
183 # However, there are certain data sets where GetRange(-1) seems
184 # to return incorrect mimimum vector although the maximum vector is
185 # correct. As a result, the mimimum vector has been hard coded to 0.0
186 # to accommodate those incorrect cases.
187 if(len(self.__point_attribute['vectors']) != 0):
188 vector_range = \
189 self._getOutput().GetPointData().GetVectors().GetRange(-1)
190 return (0.0, vector_range[1])
191 elif(len(self.__cell_attribute['vectors']) != 0):
192 vector_range = \
193 self._getOutput().GetCellData().GetVectors().GetRange(-1)
194 return (0.0, vector_range[1])
195 else:
196 print "\nSorry, no vector is available.\n"
197 sys.exit(0)
198
199 def _getTensorRange(self):
200 """
201 Return the tensor range.
202
203 @rtype: Two column table
204 @return: Tensor range
205 """
206
207 if(len(self.__point_attribute['tensors']) != 0):
208 return self._getOutput().GetPointData().GetTensors().GetRange(-1)
209 elif(len(self.__cell_attribute['tensors']) != 0):
210 return self._getOutput().GetCellData().GetTensors().GetRange(-1)
211 else:
212 print "\nSorry, no tensor is available.\n"
213 sys.exit(0)
214
215 def _getOutput(self):
216 """
217 Return the output of the data collector.
218
219 @rtype: vtkUnstructuredGrid
220 @return: Unstructured grid
221 """
222 #cell_to_point_data = vtk.vtkCellDataToPointData()
223 #cell_to_point_data.SetInput(self.__output)
224 #cell_to_point_data.Update()
225 #return (cell_to_point_data.GetOutput())
226 return self.__output
227
228
229 ###############################################################################
230
231
232 from constant import ImageFormat
233
234 class ImageReader:
235 """
236 Class that defines an image reader.
237 """
238
239 def __init__(self, format):
240 """
241 Initialise the image reader.
242
243 @type format: String
244 @param format: Format of the image
245 """
246
247 self.__format = format
248 self.__vtk_image_reader = self.getImageReader()
249
250 def getImageReader(self):
251 """
252 Return the appropriate image reader based on the supplied image
253 format.
254
255 @rtype: vtkImageReader2 (i.e. vtkJPEGReader, etc)
256 @return: Image reader
257 """
258
259 if(self.__format == ImageFormat.JPG):
260 return vtk.vtkJPEGReader()
261 elif(self.__format == ImageFormat.BMP):
262 return vtk.vtkBMPReader()
263 elif(self.__format == ImageFormat.PNM):
264 return vtk.vtkPNMReader()
265 elif(self.__format == ImageFormat.PNG):
266 return vtk.vtkPNGReader()
267 elif(self.__format == ImageFormat.TIF):
268 return vtk.vtkTIFFReader()
269
270 def setFileName(self, file_name):
271 """
272 Set the image file name.
273
274 @type file_name: String
275 @param file_name: Image file name to be read
276 """
277
278 self.__vtk_image_reader.SetFileName(file_name)
279
280 def _getOutput(self):
281 """
282 Return the output of the image reader.
283
284 @rtype: vtkImageData
285 @return Image data
286 """
287
288 return self.__vtk_image_reader.GetOutput()
289

  ViewVC Help
Powered by ViewVC 1.1.26