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 cut(inp, origin, normal):
151 """ 152 VTK operation: cut. 153 154 @param inp: input VTK object. 155 @type inp: vtk.vtkobject 156 @param origin: a 3-tuple for cut origin. 157 @type origin: tuple 158 @param normal: a 3-tuple for cut normal. 159 @type normal: tuple 160 @return: output VTK object. 161 @rtype: vtk.vtkobject 162 """ 163 import vtk 164 pne = vtk.vtkPlane() 165 pne.SetOrigin(origin) 166 pne.SetNormal(normal) 167 cut = vtk.vtkCutter() 168 cut.SetInputConnection(inp.GetOutputPort()) 169 cut.SetCutFunction(pne) 170 return cut
171 @staticmethod
172 - def contour_value(inp, num, value):
173 """ 174 VTK operation: contour by a single value. 175 176 @param inp: input VTK object. 177 @type inp: vtk.vtkobject 178 @param num: the index of the contour line. 179 @type num: int 180 @param value: the value of the contour line. 181 @type value: float 182 @return: output VTK object. 183 @rtype: vtk.vtkobject 184 """ 185 import vtk 186 cnr = vtk.vtkContourFilter() 187 cnr.SetInputConnection(inp.GetOutputPort()) 188 cnr.SetValue(num, value) 189 return cnr
190 @staticmethod
191 - def contour_range(inp, num, begin, end):
192 """ 193 VTK operation: contour by a range. 194 195 @param inp: input VTK object. 196 @type inp: vtk.vtkobject 197 @param num: the number of the contour lines. 198 @type num: int 199 @param begin: the start of the range. 200 @type begin: float 201 @param end: the end of the range. 202 @type end: float 203 @return: output VTK object. 204 @rtype: vtk.vtkobject 205 """ 206 import vtk 207 cnr = vtk.vtkContourFilter() 208 cnr.SetInputConnection(inp.GetOutputPort()) 209 cnr.GenerateValues(num, begin, end) 210 return cnr
211 @staticmethod
212 - def lump_poly(*args):
213 """ 214 Lump all passed-in poly together. FIXME: when lumped together, vectors 215 are messed up. 216 217 @return: output VTK object. 218 @rtype: vtk.vtkobject 219 """ 220 import vtk 221 apd = vtk.vtkAppendPolyData() 222 for vbj in args: 223 apd.AddInput(vbj.GetOutput()) 224 return apd
225 @staticmethod
226 - def write_poly(inp, outfn):
227 """ 228 Write VTK polydata to a file. 229 230 @param inp: input VTK object. 231 @type inp: vtk.vtkobject 232 @param outfn: output file name. 233 @type outfn: str 234 @return: nothing 235 """ 236 import vtk 237 wtr = vtk.vtkXMLPolyDataWriter() 238 wtr.EncodeAppendedDataOff() 239 wtr.SetInput(inp.GetOutput()) 240 wtr.SetFileName(outfn) 241 wtr.Write()
242 Vop = VtkOperation 243