Package solvcon :: Module visual_vtk
[hide private]
[frames] | no frames]

Source Code for Module solvcon.visual_vtk

  1  # -*- coding: UTF-8 -*- 
  2  # 
  3  # Copyright (C) 2010 Yung-Yu Chen <yyc@solvcon.net>. 
  4  # 
  5  # This program is free software; you can redistribute it and/or modify 
  6  # it under the terms of the GNU General Public License as published by 
  7  # the Free Software Foundation; either version 2 of the License, or 
  8  # (at your option) any later version. 
  9  # 
 10  # This program is distributed in the hope that it will be useful, 
 11  # but WITHOUT ANY WARRANTY; without even the implied warranty of 
 12  # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the 
 13  # GNU General Public License for more details. 
 14  # 
 15  # You should have received a copy of the GNU General Public License along 
 16  # with this program; if not, write to the Free Software Foundation, Inc., 
 17  # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. 
 18   
 19  """ 
 20  Visualizing algorithm by VTK. 
 21  """ 
 22   
 23  VANMAP = dict(float32='vtkFloatArray', float64='vtkDoubleArray') 
24 25 -def make_ust_from_blk(blk):
26 """ 27 Create vtk.vtkUnstructuredGrid object from Block object. 28 29 @param blk: input block. 30 @type blk: solvcon.block.Block 31 @return: the ust object. 32 @rtype: vtk.vtkUnstructuredGrid 33 """ 34 from numpy import array 35 import vtk 36 cltpm = array([1, 3, 9, 5, 12, 10, 13, 14], dtype='int32') 37 nnode = blk.nnode 38 ndcrd = blk.ndcrd 39 ncell = blk.ncell 40 cltpn = blk.cltpn 41 clnds = blk.clnds 42 # create vtkPoints. 43 pts = vtk.vtkPoints() 44 pts.SetNumberOfPoints(nnode) 45 ind = 0 46 while ind < nnode: 47 pts.SetPoint(ind, ndcrd[ind,:]) 48 ind += 1 49 # create vtkUnstructuredGrid. 50 ust = vtk.vtkUnstructuredGrid() 51 ust.Allocate(ncell, ncell) 52 ust.SetPoints(pts) 53 icl = 0 54 while icl < ncell: 55 tpn = cltpm[cltpn[icl]] 56 ids = vtk.vtkIdList() 57 for inl in range(1, clnds[icl,0]+1): 58 ids.InsertNextId(clnds[icl,inl]) 59 ust.InsertNextCell(tpn, ids) 60 icl += 1 61 return ust
62
63 -def valid_vector(arr):
64 """ 65 A valid vector must have 3 compoments. If it has only 2, pad it. If 66 it has more than 3, raise ValueError. 67 68 @param arr: input vector array. 69 @type arr: numpy.ndarray 70 @return: validated array. 71 @rtype: numpy.ndarray 72 """ 73 from numpy import empty 74 if arr.shape[1] < 3: 75 arrn = empty((arr.shape[0], 3), dtype=arr.dtype) 76 arrn[:,2] = 0.0 77 try: 78 arrn[:,:2] = arr[:,:] 79 except ValueError, e: 80 args = e.args[:] 81 args.append('arrn.shape=%s, arr.shape=%s' % ( 82 str(arrn.shape), str(arr.shape))) 83 e.args = args 84 raise 85 arr = arrn 86 elif arr.shape[1] > 3: 87 raise IndexError('arr.shape[1] = %d > 3'%arr.shape[1]) 88 return arr
89
90 -def set_array(arr, name, fpdtype, ust):
91 """ 92 Set the data of a ndarray to vtk array and return the set vtk array. 93 If the array of the specified name existed, use the existing array. 94 95 @param arr: input array. 96 @type arr: numpy.ndarray 97 @param name: array name. 98 @type name: str 99 @param fpdtype: floating point dtype. 100 @type fpdtype: str 101 @param ust: the UnstructuredGrid object to store the array. 102 @type ust: vtk.vtkobject 103 @return: the set VTK array object. 104 """ 105 import vtk 106 if ust.GetCellData().HasArray(name): 107 vaj = ust.GetCellData().GetArray(name) 108 else: 109 vaj = getattr(vtk, VANMAP[fpdtype])() 110 # prepare for vector. 111 if len(arr.shape) > 1: 112 vaj.SetNumberOfComponents(3) 113 # set number of tuples to allocate. 114 vaj.SetNumberOfTuples(arr.shape[0]) 115 # cache. 116 vaj.SetName(name) 117 ust.GetCellData().AddArray(vaj) 118 # set data. 119 nt = arr.shape[0] 120 it = 0 121 if len(arr.shape) > 1: 122 while it < nt: 123 vaj.SetTuple3(it, *arr[it]) 124 it += 1 125 else: 126 while it < nt: 127 vaj.SetValue(it, arr[it]) 128 it += 1 129 return vaj
130
131 -class VtkOperation(object):
132 """ 133 Pre-defined VTK operations. 134 """ 135 @staticmethod
136 - def c2p(inp):
137 """ 138 VTK operation: cell to point. 139 140 @param inp: input VTK object. 141 @type inp: vtk.vtkobject 142 @return: output VTK object. 143 @rtype: vtk.vtkobject 144 """ 145 import vtk 146 usp = vtk.vtkCellDataToPointData() 147 usp.SetInput(inp) 148 return usp
149 @staticmethod
150 - def clip(inp, origin, normal, inside_out=False, take_cell=False):
151 """ 152 VTK operation: clip. A vtkGeometryFilter is used to convert the 153 resulted vtkUnstructuredMesh object into a vtkPolyData object. 154 155 @param inp: input VTK object. 156 @type inp: vtk.vtkobject 157 @param origin: a 3-tuple for cut origin. 158 @type origin: tuple 159 @param normal: a 3-tuple for cut normal. 160 @type normal: tuple 161 @keyword inside_out: make inside out. Default False. 162 @type inside_out: bool 163 @keyword take_cell: treat the input VTK object with values on cells. 164 Default False. 165 @type: take_cell: bool 166 @return: output VTK object. 167 @rtype: vtk.vtkobject 168 """ 169 import vtk 170 pne = vtk.vtkPlane() 171 pne.SetOrigin(origin) 172 pne.SetNormal(normal) 173 clip = vtk.vtkClipDataSet() 174 if take_cell: 175 clip.SetInput(inp) 176 else: 177 clip.SetInputConnection(inp.GetOutputPort()) 178 clip.SetClipFunction(pne) 179 if inside_out: 180 clip.InsideOutOn() 181 parison = vtk.vtkGeometryFilter() 182 parison.SetInputConnection(clip.GetOutputPort()) 183 return parison
184 @staticmethod
185 - def cut(inp, origin, normal):
186 """ 187 VTK operation: cut. 188 189 @param inp: input VTK object. 190 @type inp: vtk.vtkobject 191 @param origin: a 3-tuple for cut origin. 192 @type origin: tuple 193 @param normal: a 3-tuple for cut normal. 194 @type normal: tuple 195 @return: output VTK object. 196 @rtype: vtk.vtkobject 197 """ 198 import vtk 199 pne = vtk.vtkPlane() 200 pne.SetOrigin(origin) 201 pne.SetNormal(normal) 202 cut = vtk.vtkCutter() 203 cut.SetInputConnection(inp.GetOutputPort()) 204 cut.SetCutFunction(pne) 205 return cut
206 @staticmethod
207 - def contour_value(inp, num, value):
208 """ 209 VTK operation: contour by a single value. 210 211 @param inp: input VTK object. 212 @type inp: vtk.vtkobject 213 @param num: the index of the contour line. 214 @type num: int 215 @param value: the value of the contour line. 216 @type value: float 217 @return: output VTK object. 218 @rtype: vtk.vtkobject 219 """ 220 import vtk 221 cnr = vtk.vtkContourFilter() 222 cnr.SetInputConnection(inp.GetOutputPort()) 223 cnr.SetValue(num, value) 224 return cnr
225 @staticmethod
226 - def contour_range(inp, num, begin, end):
227 """ 228 VTK operation: contour by a range. 229 230 @param inp: input VTK object. 231 @type inp: vtk.vtkobject 232 @param num: the number of the contour lines. 233 @type num: int 234 @param begin: the start of the range. 235 @type begin: float 236 @param end: the end of the range. 237 @type end: float 238 @return: output VTK object. 239 @rtype: vtk.vtkobject 240 """ 241 import vtk 242 cnr = vtk.vtkContourFilter() 243 cnr.SetInputConnection(inp.GetOutputPort()) 244 cnr.GenerateValues(num, begin, end) 245 return cnr
246 @staticmethod
247 - def lump_poly(*args):
248 """ 249 Lump all passed-in poly together. FIXME: when lumped together, vectors 250 are messed up. 251 252 @return: output VTK object. 253 @rtype: vtk.vtkobject 254 """ 255 import vtk 256 apd = vtk.vtkAppendPolyData() 257 for vbj in args: 258 apd.AddInput(vbj.GetOutput()) 259 return apd
260 @staticmethod
261 - def write_poly(inp, outfn):
262 """ 263 Write VTK polydata to a file. 264 265 @param inp: input VTK object. 266 @type inp: vtk.vtkobject 267 @param outfn: output file name. 268 @type outfn: str 269 @return: nothing 270 """ 271 import vtk 272 wtr = vtk.vtkXMLPolyDataWriter() 273 wtr.EncodeAppendedDataOff() 274 wtr.SetInput(inp.GetOutput()) 275 wtr.SetFileName(outfn) 276 wtr.Write()
277 Vop = VtkOperation 278