$search
00001 """ 00002 This module contains all the array handling code for TVTK. 00003 00004 The most important functions provided by this module involve the 00005 conversion of numpy arrays/Python lists to different VTK data arrays 00006 and vice-versa. 00007 00008 Warning: Numpy Character arrays will not work properly since there 00009 seems no unique one-to-one VTK data array type to map it to. 00010 00011 """ 00012 # Author: Prabhu Ramachandran <prabhu_r@users.sf.net> 00013 # Copyright (c) 2004-2008, Enthought, Inc. 00014 # License: BSD Style. 00015 00016 import types 00017 import sys 00018 import weakref 00019 import atexit 00020 00021 import vtk 00022 from vtk.util import vtkConstants 00023 try: 00024 from vtk.util import numpy_support 00025 except ImportError: 00026 numpy_support = None 00027 00028 import numpy 00029 00030 # Enthought library imports. 00031 #from enthought.tvtk.array_ext import set_id_type_array 00032 00033 # Useful constants for VTK arrays. 00034 VTK_ID_TYPE_SIZE = vtk.vtkIdTypeArray().GetDataTypeSize() 00035 if VTK_ID_TYPE_SIZE == 4: 00036 ID_TYPE_CODE = numpy.int32 00037 elif VTK_ID_TYPE_SIZE == 8: 00038 ID_TYPE_CODE = numpy.int64 00039 00040 VTK_LONG_TYPE_SIZE = vtk.vtkLongArray().GetDataTypeSize() 00041 if VTK_LONG_TYPE_SIZE == 4: 00042 LONG_TYPE_CODE = numpy.int32 00043 ULONG_TYPE_CODE = numpy.uint32 00044 elif VTK_LONG_TYPE_SIZE == 8: 00045 LONG_TYPE_CODE = numpy.int64 00046 ULONG_TYPE_CODE = numpy.uint64 00047 00048 00049 ###################################################################### 00050 # The array cache. 00051 ###################################################################### 00052 class ArrayCache(object): 00053 00054 """Caches references to numpy arrays that are not copied but views 00055 of which are converted to VTK arrays. The caching prevents the user 00056 from deleting or resizing the numpy array after it has been sent 00057 down to VTK. """ 00058 00059 ###################################################################### 00060 # `object` interface. 00061 ###################################################################### 00062 def __init__(self): 00063 # The cache. 00064 self._cache = {} 00065 00066 def __len__(self): 00067 return len(self._cache) 00068 00069 def __contains__(self, vtk_arr): 00070 key = vtk_arr.__this__ 00071 return self._cache.has_key(key) 00072 00073 ###################################################################### 00074 # `ArrayCache` interface. 00075 ###################################################################### 00076 def add(self, vtk_arr, np_arr): 00077 """Add numpy array corresponding to the vtk array to the 00078 cache.""" 00079 key = vtk_arr.__this__ 00080 cache = self._cache 00081 00082 # Setup a callback so this cached array reference is removed 00083 # when the VTK array is destroyed. Passing the key to the 00084 # `lambda` function is necessary because the callback will not 00085 # receive the object (it will receive `None`) and thus there 00086 # is no way to know which array reference one has to remove. 00087 ob_id = vtk_arr.AddObserver('DeleteEvent', lambda o, e, key=key: \ 00088 self._remove_array(key)) 00089 # Try creating a weakref to the array. 00090 try: 00091 ref = weakref.ref(vtk_arr) 00092 except TypeError: 00093 # Old versions of VTK (<5.0) did not support weakrefs. 00094 ref = None 00095 # Cache the array, weakref and the observer id. 00096 cache[key] = (np_arr, ref, ob_id) 00097 00098 def get(self, vtk_arr): 00099 """Return the cached numpy array given a VTK array.""" 00100 key = vtk_arr.__this__ 00101 return self._cache[key][0] 00102 00103 def on_exit(self): 00104 """Should be registered with atexit.register so the observers are 00105 removed before the VTK arrays destruct. If not done, this 00106 causes segfaults. Note that this is not done by default so the 00107 user must explicitly register this function with atexit. 00108 """ 00109 cache = self._cache 00110 for key in cache: 00111 np_arr, ref, ob_id = cache.get(key) 00112 if ref is not None: 00113 vtk_arr = ref() 00114 if vtk_arr is not None: 00115 vtk_arr.RemoveObserver(ob_id) 00116 00117 ###################################################################### 00118 # Non-public interface. 00119 ###################################################################### 00120 def _remove_array(self, key): 00121 """Private function that removes the cached array. Do not 00122 call this unless you know what you are doing.""" 00123 try: 00124 del self._cache[key] 00125 except KeyError: 00126 pass 00127 00128 00129 ###################################################################### 00130 # Setup a global `_array_cache`. The array object cache caches all the 00131 # converted numpy arrays that are not copied. This prevents the user 00132 # from deleting or resizing the numpy array after it has been sent down 00133 # to VTK. 00134 ###################################################################### 00135 00136 _dummy = None 00137 # This makes the cache work even when the module is reloaded. 00138 for name in ['array_handler', 'enthought.tvtk.array_handler']: 00139 if sys.modules.has_key(name): 00140 mod = sys.modules[name] 00141 if hasattr(mod, '_array_cache'): 00142 _dummy = mod._array_cache 00143 del mod 00144 break 00145 00146 if _dummy: 00147 _array_cache = _dummy 00148 else: 00149 _array_cache = ArrayCache() 00150 # Register function to call when Python exits. 00151 atexit.register(_array_cache.on_exit) 00152 del _dummy 00153 00154 00155 ###################################################################### 00156 # Array conversion functions. 00157 ###################################################################### 00158 def get_vtk_array_type(numeric_array_type): 00159 """Returns a VTK typecode given a numpy array.""" 00160 # This is a Mapping from numpy array types to VTK array types. 00161 _arr_vtk = {numpy.dtype(numpy.character):vtkConstants.VTK_UNSIGNED_CHAR, 00162 numpy.dtype(numpy.uint8):vtkConstants.VTK_UNSIGNED_CHAR, 00163 numpy.dtype(numpy.uint16):vtkConstants.VTK_UNSIGNED_SHORT, 00164 numpy.dtype(ULONG_TYPE_CODE):vtkConstants.VTK_UNSIGNED_LONG, 00165 numpy.dtype(numpy.int8):vtkConstants.VTK_CHAR, 00166 numpy.dtype(numpy.int16):vtkConstants.VTK_SHORT, 00167 numpy.dtype(numpy.int32):vtkConstants.VTK_INT, 00168 numpy.dtype(LONG_TYPE_CODE):vtkConstants.VTK_LONG, 00169 numpy.dtype(numpy.float32):vtkConstants.VTK_FLOAT, 00170 numpy.dtype(numpy.float64):vtkConstants.VTK_DOUBLE, 00171 numpy.dtype(numpy.complex64):vtkConstants.VTK_FLOAT, 00172 numpy.dtype(numpy.complex128):vtkConstants.VTK_DOUBLE} 00173 try: 00174 return _arr_vtk[numeric_array_type] 00175 except KeyError: 00176 for key in _arr_vtk: 00177 if numpy.issubdtype(numeric_array_type, key): 00178 return _arr_vtk[key] 00179 raise TypeError, "Couldn't translate array's type to VTK" 00180 00181 def get_vtk_to_numeric_typemap(): 00182 """Returns the VTK array type to numpy array type mapping.""" 00183 _vtk_arr = {vtkConstants.VTK_BIT:numpy.bool, 00184 vtkConstants.VTK_CHAR:numpy.int8, 00185 vtkConstants.VTK_UNSIGNED_CHAR:numpy.uint8, 00186 vtkConstants.VTK_SHORT:numpy.int16, 00187 vtkConstants.VTK_UNSIGNED_SHORT:numpy.uint16, 00188 vtkConstants.VTK_INT:numpy.int32, 00189 vtkConstants.VTK_UNSIGNED_INT:numpy.uint32, 00190 vtkConstants.VTK_LONG:LONG_TYPE_CODE, 00191 vtkConstants.VTK_UNSIGNED_LONG:ULONG_TYPE_CODE, 00192 vtkConstants.VTK_ID_TYPE:ID_TYPE_CODE, 00193 vtkConstants.VTK_FLOAT:numpy.float32, 00194 vtkConstants.VTK_DOUBLE:numpy.float64} 00195 return _vtk_arr 00196 00197 00198 def get_numeric_array_type(vtk_array_type): 00199 """Returns a numpy array typecode given a VTK array type.""" 00200 return get_vtk_to_numeric_typemap()[vtk_array_type] 00201 00202 00203 def get_sizeof_vtk_array(vtk_array_type): 00204 """Returns the size of a VTK array type.""" 00205 _size_dict = {vtkConstants.VTK_BIT : 1, 00206 vtkConstants.VTK_CHAR : 1, 00207 vtkConstants.VTK_UNSIGNED_CHAR : 1, 00208 vtkConstants.VTK_SHORT : 2, 00209 vtkConstants.VTK_UNSIGNED_SHORT : 2, 00210 vtkConstants.VTK_INT : 4, 00211 vtkConstants.VTK_UNSIGNED_INT : 4, 00212 vtkConstants.VTK_LONG : VTK_LONG_TYPE_SIZE, 00213 vtkConstants.VTK_UNSIGNED_LONG : VTK_LONG_TYPE_SIZE, 00214 vtkConstants.VTK_ID_TYPE : VTK_ID_TYPE_SIZE, 00215 vtkConstants.VTK_FLOAT : 4, 00216 vtkConstants.VTK_DOUBLE : 8 } 00217 return _size_dict[vtk_array_type] 00218 00219 00220 def create_vtk_array(vtk_arr_type): 00221 """Internal function used to create a VTK data array from another 00222 VTK array given the VTK array type. 00223 """ 00224 tmp = vtk.vtkDataArray.CreateDataArray(vtk_arr_type) 00225 # CreateDataArray sets the refcount to 3 and this causes a severe 00226 # memory leak. 00227 tmp.SetReferenceCount(2) 00228 return tmp 00229 00230 00231 def array2vtk(num_array, vtk_array=None): 00232 """Converts a real numpy Array (or a Python list) to a VTK array 00233 object. 00234 00235 This function only works for real arrays. Complex arrays are NOT 00236 handled. It also works for multi-component arrays. However, only 00237 1, and 2 dimensional arrays are supported. This function is very 00238 efficient, so large arrays should not be a problem. 00239 00240 Even in cases when no copy of the numpy array data is performed, 00241 a reference to the array is cached. The passed array can 00242 therefore be deleted safely in all circumstances. 00243 00244 Parameters 00245 ---------- 00246 00247 - num_array : numpy array or Python list/tuple 00248 00249 The input array must be 1 or 2D. A copy of the numeric array 00250 data passed is made in the following circumstances: 00251 00252 1. A Python list/tuple was passed. 00253 2. A non-contiguous numpy array was passed. 00254 3. A `vtkBitArray` instance was passed as the second argument. 00255 4. The types of the `vtk_array` and the `num_array` are not 00256 equivalent to each other. For example if one is an integer 00257 array and the other a float. 00258 00259 00260 - vtk_array : `vtkDataArray` (default: `None`) 00261 00262 If an optional `vtkDataArray` instance, is passed as an argument 00263 then a new array is not created and returned. The passed array 00264 is itself returned. 00265 00266 """ 00267 00268 z = numpy.asarray(num_array) 00269 00270 shape = z.shape 00271 assert len(shape) < 3, \ 00272 "Only arrays of dimensionality 2 or lower are allowed!" 00273 assert not numpy.issubdtype(z.dtype, complex), \ 00274 "Complex numpy arrays cannot be converted to vtk arrays."\ 00275 "Use real() or imag() to get a component of the array before"\ 00276 " passing it to vtk." 00277 00278 # First create an array of the right type by using the typecode. 00279 # Bit arrays need special casing. 00280 bit_array = False 00281 if vtk_array is None: 00282 vtk_typecode = get_vtk_array_type(z.dtype) 00283 result_array = create_vtk_array(vtk_typecode) 00284 elif vtk_array.GetDataType() == vtkConstants.VTK_BIT: 00285 vtk_typecode = vtkConstants.VTK_CHAR 00286 result_array = create_vtk_array(vtkConstants.VTK_CHAR) 00287 bit_array = True 00288 else: 00289 vtk_typecode = vtk_array.GetDataType() 00290 result_array = vtk_array 00291 00292 # Find the shape and set number of components. 00293 if len(shape) == 1: 00294 result_array.SetNumberOfComponents(1) 00295 else: 00296 result_array.SetNumberOfComponents(shape[1]) 00297 00298 result_array.SetNumberOfTuples(shape[0]) 00299 00300 # Ravel the array appropriately. 00301 arr_dtype = get_numeric_array_type(vtk_typecode) 00302 if numpy.issubdtype(z.dtype, arr_dtype): 00303 z_flat = numpy.ravel(z) 00304 else: 00305 z_flat = numpy.ravel(z).astype(arr_dtype) 00306 00307 # Point the VTK array to the numpy data. The last argument (1) 00308 # tells the array not to deallocate. 00309 result_array.SetVoidArray(z_flat, len(z_flat), 1) 00310 00311 if bit_array: 00312 # Handle bit arrays -- they have to be copied. Note that bit 00313 # arrays are used ONLY when the user has passed one as an 00314 # argument to this function. 00315 vtk_array.SetNumberOfTuples(result_array.GetNumberOfTuples()) 00316 vtk_array.SetNumberOfComponents(result_array.GetNumberOfComponents()) 00317 for i in range(result_array.GetNumberOfComponents()): 00318 vtk_array.CopyComponent(i, result_array, i) 00319 result_array = vtk_array 00320 else: 00321 # Save a reference to the flatted array in the array cache. 00322 # This prevents the user from deleting or resizing the array 00323 # and getting into serious trouble. This is only done for 00324 # non-bit array cases where the data is not copied. 00325 global _array_cache 00326 _array_cache.add(result_array, z_flat) 00327 00328 return result_array 00329 00330 00331 def vtk2array(vtk_array): 00332 """Converts a VTK data array to a numpy array. 00333 00334 Given a subclass of vtkDataArray, this function returns an 00335 appropriate numpy array containing the same data. The function 00336 is very efficient since it uses the VTK imaging pipeline to 00337 convert the data. If a sufficiently new version of VTK (5.2) is 00338 installed then it actually uses the buffer interface to return a 00339 view of the VTK array in the returned numpy array. 00340 00341 Parameters 00342 ---------- 00343 00344 - vtk_array : `vtkDataArray` 00345 00346 The VTK data array to be converted. 00347 00348 """ 00349 typ = vtk_array.GetDataType() 00350 assert typ in get_vtk_to_numeric_typemap().keys(), \ 00351 "Unsupported array type %s"%typ 00352 00353 shape = vtk_array.GetNumberOfTuples(), \ 00354 vtk_array.GetNumberOfComponents() 00355 00356 # First check if this array already has a numpy array cached, if 00357 # it does, reshape that and return it. 00358 if vtk_array in _array_cache: 00359 arr = _array_cache.get(vtk_array) 00360 if shape[1] == 1: 00361 shape = (shape[0], ) 00362 arr = numpy.reshape(arr, shape) 00363 return arr 00364 00365 # If VTK's new numpy support is available, use the buffer interface. 00366 if numpy_support is not None and typ != vtkConstants.VTK_BIT: 00367 dtype = get_numeric_array_type(typ) 00368 result = numpy.frombuffer(vtk_array, dtype=dtype) 00369 if shape[1] == 1: 00370 shape = (shape[0], ) 00371 result.shape = shape 00372 return result 00373 00374 # Setup an imaging pipeline to export the array. 00375 img_data = vtk.vtkImageData() 00376 img_data.SetDimensions(shape[0], 1, 1) 00377 if typ == vtkConstants.VTK_BIT: 00378 iarr = vtk.vtkCharArray() 00379 iarr.DeepCopy(vtk_array) 00380 img_data.GetPointData().SetScalars(iarr) 00381 elif typ == vtkConstants.VTK_ID_TYPE: 00382 # Needed since VTK_ID_TYPE does not work with VTK 4.5. 00383 iarr = vtk.vtkLongArray() 00384 iarr.SetNumberOfTuples(vtk_array.GetNumberOfTuples()) 00385 nc = vtk_array.GetNumberOfComponents() 00386 iarr.SetNumberOfComponents(nc) 00387 for i in range(nc): 00388 iarr.CopyComponent(i, vtk_array, i) 00389 img_data.GetPointData().SetScalars(iarr) 00390 else: 00391 img_data.GetPointData().SetScalars(vtk_array) 00392 00393 img_data.SetNumberOfScalarComponents(shape[1]) 00394 if typ == vtkConstants.VTK_ID_TYPE: 00395 # Hack necessary because vtkImageData can't handle VTK_ID_TYPE. 00396 img_data.SetScalarType(vtkConstants.VTK_LONG) 00397 elif typ == vtkConstants.VTK_BIT: 00398 img_data.SetScalarType(vtkConstants.VTK_CHAR) 00399 else: 00400 img_data.SetScalarType(typ) 00401 img_data.Update() 00402 00403 exp = vtk.vtkImageExport() 00404 exp.SetInput(img_data) 00405 00406 # Create an array of the right size and export the image into it. 00407 im_arr = numpy.empty((shape[0]*shape[1],), get_numeric_array_type(typ)) 00408 exp.Export(im_arr) 00409 00410 # Now reshape it. 00411 if shape[1] == 1: 00412 shape = (shape[0], ) 00413 im_arr = numpy.reshape(im_arr, shape) 00414 return im_arr 00415 00416 00417 def array2vtkCellArray(num_array, vtk_array=None): 00418 """Given a nested Python list or a numpy array, this method 00419 creates a vtkCellArray instance and returns it. 00420 00421 A variety of input arguments are supported as described in the 00422 Parameter documentation. If numpy arrays are given, this method 00423 is highly efficient. This function is most efficient if the 00424 passed numpy arrays have a typecode `ID_TYPE_CODE`. Otherwise a 00425 typecast is necessary and this involves an extra copy. This 00426 method *always copies* the input data. 00427 00428 An alternative and more efficient way to build the connectivity 00429 list is to create a vtkIdTypeArray having data of the form 00430 (npts,p0,p1,...p(npts-1), repeated for each cell) and then call 00431 <vtkCellArray_instance>.SetCells(n_cell, id_list). 00432 00433 Parameters 00434 ---------- 00435 00436 - num_array : numpy array or Python list/tuple 00437 00438 Valid values are: 00439 00440 1. A Python list of 1D lists. Each 1D list can contain one 00441 cell connectivity list. This is very slow and is to be 00442 used only when efficiency is of no consequence. 00443 00444 2. A 2D numpy array with the cell connectivity list. 00445 00446 3. A Python list of 2D numpy arrays. Each numeric array can 00447 have a different shape. This makes it easy to generate a 00448 cell array having cells of different kinds. 00449 00450 - vtk_array : `vtkCellArray` (default: `None`) 00451 00452 If an optional `vtkCellArray` instance, is passed as an argument 00453 then a new array is not created and returned. The passed array 00454 is itself modified and returned. 00455 00456 Example 00457 ------- 00458 00459 >>> a = [[0], [1, 2], [3, 4, 5], [6, 7, 8, 9]] 00460 >>> cells = array_handler.array2vtkCellArray(a) 00461 >>> a = numpy.array([[0,1,2], [3,4,5], [6,7,8]], 'l') 00462 >>> cells = array_handler.array2vtkCellArray(a) 00463 >>> l_a = [a[:,:1], a[:2,:2], a] 00464 >>> cells = array_handler.array2vtkCellArray(l_a) 00465 00466 """ 00467 if vtk_array: 00468 cells = vtk_array 00469 else: 00470 cells = vtk.vtkCellArray() 00471 assert cells.GetClassName() == 'vtkCellArray', \ 00472 'Second argument must be a `vtkCellArray` instance.' 00473 00474 if len(num_array) == 0: 00475 return cells 00476 00477 ######################################## 00478 # Internal functions. 00479 def _slow_array2cells(z, cells): 00480 cells.Reset() 00481 vtk_ids = vtk.vtkIdList() 00482 for i in z: 00483 vtk_ids.Reset() 00484 for j in i: 00485 vtk_ids.InsertNextId(j) 00486 cells.InsertNextCell(vtk_ids) 00487 00488 def _get_tmp_array(arr): 00489 try: 00490 tmp_arr = numpy.asarray(arr, ID_TYPE_CODE) 00491 except TypeError: 00492 tmp_arr = arr.astype(ID_TYPE_CODE) 00493 return tmp_arr 00494 00495 def _set_cells(cells, n_cells, id_typ_arr): 00496 vtk_arr = vtk.vtkIdTypeArray() 00497 array2vtk(id_typ_arr, vtk_arr) 00498 cells.SetCells(n_cells, vtk_arr) 00499 ######################################## 00500 00501 msg = "Invalid argument. Valid types are a Python list of lists,"\ 00502 " a Python list of numpy arrays, or a numpy array." 00503 00504 if issubclass(type(num_array), (types.ListType, types.TupleType)): 00505 assert len(num_array[0]) > 0, "Input array must be 2D." 00506 tp = type(num_array[0]) 00507 if issubclass(tp, types.ListType): # Pure Python list. 00508 _slow_array2cells(num_array, cells) 00509 return cells 00510 elif issubclass(tp, numpy.ndarray): # List of arrays. 00511 # Check shape of array and find total size. 00512 tot_size = 0 00513 n_cells = 0 00514 for arr in num_array: 00515 assert len(arr.shape) == 2, "Each array must be 2D" 00516 shp = arr.shape 00517 tot_size += shp[0]*(shp[1] + 1) 00518 n_cells += shp[0] 00519 # Create an empty array. 00520 id_typ_arr = numpy.empty((tot_size,), ID_TYPE_CODE) 00521 # Now populate it with the ids. 00522 count = 0 00523 for arr in num_array: 00524 tmp_arr = _get_tmp_array(arr) 00525 shp = arr.shape 00526 sz = shp[0]*(shp[1] + 1) 00527 set_id_type_array(tmp_arr, id_typ_arr[count:count+sz]) 00528 count += sz 00529 # Now set them cells. 00530 _set_cells(cells, n_cells, id_typ_arr) 00531 return cells 00532 else: 00533 raise TypeError, msg 00534 elif issubclass(type(num_array), numpy.ndarray): 00535 assert len(num_array.shape) == 2, "Input array must be 2D." 00536 tmp_arr = _get_tmp_array(num_array) 00537 shp = tmp_arr.shape 00538 id_typ_arr = numpy.empty((shp[0]*(shp[1] + 1),), ID_TYPE_CODE) 00539 set_id_type_array(tmp_arr, id_typ_arr) 00540 _set_cells(cells, shp[0], id_typ_arr) 00541 return cells 00542 else: 00543 raise TypeError, msg 00544 00545 00546 def array2vtkPoints(num_array, vtk_points=None): 00547 """Converts a numpy array/Python list to a vtkPoints object. 00548 00549 Unless a Python list/tuple or a non-contiguous array is given, no 00550 copy of the data is made. Thus the function is very efficient. 00551 00552 Parameters 00553 ---------- 00554 00555 - num_array : numpy array or Python list/tuple 00556 00557 The input array must be 2D with `shape[1] == 3`. 00558 00559 - vtk_points : `vtkPoints` (default: `None`) 00560 00561 If an optional `vtkPoints` instance, is passed as an argument 00562 then a new array is not created and returned. The passed array 00563 is itself modified and returned. 00564 00565 """ 00566 if vtk_points: 00567 points = vtk_points 00568 else: 00569 points = vtk.vtkPoints() 00570 00571 arr = numpy.asarray(num_array) 00572 assert len(arr.shape) == 2, "Points array must be 2 dimensional." 00573 assert arr.shape[1] == 3, "Incorrect shape: shape[1] must be 3." 00574 vtk_array = array2vtk(arr) 00575 points.SetData(vtk_array) 00576 return points 00577 00578 00579 def array2vtkIdList(num_array, vtk_idlist=None): 00580 """Converts a numpy array/Python list to a vtkIdList object. 00581 00582 Parameters 00583 ---------- 00584 00585 - num_array : numpy array or Python list/tuple 00586 00587 The input array must be 2D with `shape[1] == 3`. 00588 00589 - vtk_idlist : `vtkIdList` (default: `None`) 00590 00591 If an optional `vtkIdList` instance, is passed as an argument 00592 then a new array is not created and returned. The passed array 00593 is itself modified and returned. 00594 00595 """ 00596 if vtk_idlist: 00597 ids = vtk_idlist 00598 else: 00599 ids = vtk.vtkIdList() 00600 00601 arr = numpy.asarray(num_array) 00602 assert len(arr.shape) == 1, "Array for vtkIdList must be 1D" 00603 ids.SetNumberOfIds(len(arr)) 00604 for i, j in enumerate(arr): 00605 ids.SetId(i, j) 00606 return ids 00607 00608 00609 ###################################################################### 00610 # Array argument handling functions. 00611 ###################################################################### 00612 00613 def is_array(arr): 00614 """Returns True if the passed `arr` is a numpy array or a List.""" 00615 if issubclass(type(arr), (numpy.ndarray, types.ListType)): 00616 return True 00617 return False 00618 00619 00620 def convert_array(arr, vtk_typ=None): 00621 """Convert the given array to the optional type specified by 00622 `vtk_typ`. 00623 00624 Parameters 00625 ---------- 00626 00627 - arr : numpy array/list. 00628 - vtk_typ : `string` or `None` 00629 represents the type the array is to be converted to. 00630 00631 """ 00632 if vtk_typ: 00633 conv = {'vtkCellArray': array2vtkCellArray, 00634 'vtkPoints': array2vtkPoints, 00635 'vtkIdList': array2vtkIdList} 00636 if vtk_typ in conv.keys(): 00637 vtk_arr = getattr(vtk, vtk_typ)() 00638 return conv[vtk_typ](arr, vtk_arr) 00639 elif vtk_typ.find('Array') > -1: 00640 try: 00641 vtk_arr = getattr(vtk, vtk_typ)() 00642 except TypeError: # vtk_typ == 'vtkDataArray' 00643 return array2vtk(arr) 00644 else: 00645 return array2vtk(arr, vtk_arr) 00646 else: 00647 return arr 00648 else: 00649 return array2vtk(arr) 00650 00651 00652 def is_array_sig(s): 00653 """Given a signature, return if the signature has an array.""" 00654 if not isinstance(s, basestring): 00655 return False 00656 arr_types = ['Array', 'vtkPoints', 'vtkIdList'] 00657 for i in arr_types: 00658 if s.find(i) > -1: 00659 return True 00660 return False 00661 00662 00663 def is_array_or_vtkarray(arg): 00664 """Returns True if the argument is an array/Python list or if it 00665 is a vtk array.""" 00666 00667 if is_array(arg): 00668 return True 00669 else: 00670 if hasattr(arg, '_vtk_obj'): 00671 if is_array_sig(arg._vtk_obj.__class__.__name__): 00672 return True 00673 return False 00674 00675 00676 def get_correct_sig(args, sigs): 00677 """Given a list of args and a collection of possible signatures, 00678 this function returns the most appropriate signature. This 00679 function is only called by deref_array. This implies that one of 00680 the signatures has an array type. 00681 00682 """ 00683 # First do the trivial cases. 00684 if sigs is None: 00685 return None 00686 if len(sigs) == 1: 00687 return sigs[0] 00688 else: 00689 # Non-trivial cases. 00690 la = len(args) 00691 candidate_sigs = [s for s in sigs if len(s) == la] 00692 count = len(candidate_sigs) 00693 if count == 0: 00694 # No sig has the right number of args. 00695 msg = "Insufficient number of arguments to method."\ 00696 "Valid arguments are:\n%s"%sigs 00697 raise TypeError, msg 00698 elif count == 1: 00699 # If only one of the sigs has the right number of args, 00700 # return it. 00701 return candidate_sigs[0] 00702 else: 00703 # More than one sig has the same number of args. 00704 # Check if args need conversion at all. 00705 array_idx = [i for i, a in enumerate(args) \ 00706 if is_array_or_vtkarray(a)] 00707 n_arr = len(array_idx) 00708 if n_arr == 0: 00709 # No conversion necessary so signature info is 00710 # useless. 00711 return None 00712 else: 00713 # Need to find the right sig. This is done by finding 00714 # the first signature that matches all the arrays in 00715 # the argument. 00716 for sig in candidate_sigs: 00717 array_in_sig = [is_array_sig(s) for s in sig] 00718 if array_in_sig.count(True) != len(array_idx): 00719 continue 00720 bad = False 00721 for i in array_idx: 00722 if not array_in_sig[i]: 00723 bad = True 00724 if not bad: 00725 return sig 00726 # Could not find any valid signature, so give up. 00727 return None 00728 00729 00730 def deref_vtk(obj): 00731 """Dereferences the VTK object from the object if possible. This 00732 is duplicated from `tvtk_base.py` because I'd like to keep this 00733 module independent of `tvtk_base.py`. 00734 """ 00735 if hasattr(obj, '_vtk_obj'): 00736 return obj._vtk_obj 00737 else: 00738 return obj 00739 00740 00741 def deref_array(args, sigs=None): 00742 """Given a bunch of arguments and optional signature information, 00743 this converts the arguments suitably. If the argument is either a 00744 Python list or a numpy array it is converted to a suitable type 00745 based on the signature information. If it is not an array, but a 00746 TVTK object the VTK object is dereferenced. Otherwise nothing is 00747 done. If no signature information is provided the arrays are 00748 automatically converted (this can sometimes go wrong). The 00749 signature information is provided in the form of a list of lists. 00750 00751 """ 00752 ret = [] 00753 sig = get_correct_sig(args, sigs) 00754 if sig: 00755 for a, s in zip(args, sig): 00756 if is_array(a) and is_array_sig(s): 00757 ret.append(convert_array(a, s)) 00758 else: 00759 ret.append(deref_vtk(a)) 00760 else: 00761 for a in args: 00762 if is_array(a): 00763 ret.append(convert_array(a)) 00764 else: 00765 ret.append(deref_vtk(a)) 00766 return ret 00767