Advertisement
Guest User

array_handler.py

a guest
Oct 10th, 2012
95
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 27.12 KB | None | 0 0
  1. """
  2. This module contains all the array handling code for TVTK.
  3.  
  4. The most important functions provided by this module involve the
  5. conversion of numpy arrays/Python lists to different VTK data arrays
  6. and vice-versa.
  7.  
  8. Warning: Numpy Character arrays will not work properly since there
  9. seems no unique one-to-one VTK data array type to map it to.
  10.  
  11. """
  12. # Author: Prabhu Ramachandran <prabhu_r@users.sf.net>
  13. # Copyright (c) 2004-2008,  Enthought, Inc.
  14. # License: BSD Style.
  15.  
  16. import types
  17. import sys
  18.  
  19. import vtk
  20. from vtk.util import vtkConstants
  21. try:
  22.     from vtk.util import numpy_support
  23. except ImportError:
  24.     numpy_support = None
  25.  
  26. import numpy
  27.  
  28. # Enthought library imports.
  29. from enthought.tvtk.array_ext import set_id_type_array
  30.  
  31. # Useful constants for VTK arrays.
  32. VTK_ID_TYPE_SIZE = vtk.vtkIdTypeArray().GetDataTypeSize()
  33. if VTK_ID_TYPE_SIZE == 4:
  34.     ID_TYPE_CODE = numpy.int32
  35. elif VTK_ID_TYPE_SIZE == 8:
  36.     ID_TYPE_CODE = numpy.int64
  37.  
  38. VTK_LONG_TYPE_SIZE = vtk.vtkLongArray().GetDataTypeSize()
  39. if VTK_LONG_TYPE_SIZE == 4:
  40.     LONG_TYPE_CODE = numpy.int32
  41.     ULONG_TYPE_CODE = numpy.uint32
  42. elif VTK_LONG_TYPE_SIZE == 8:
  43.     LONG_TYPE_CODE = numpy.int64
  44.     ULONG_TYPE_CODE = numpy.uint64
  45.  
  46. BASE_REFERENCE_COUNT = vtk.vtkObject().GetReferenceCount()
  47.  
  48.  
  49. ######################################################################
  50. # The array cache.
  51. ######################################################################
  52. class ArrayCache(object):
  53.  
  54.     """Caches references to numpy arrays that are not copied but views
  55.    of which are converted to VTK arrays.  The caching prevents the user
  56.    from deleting or resizing the numpy array after it has been sent
  57.    down to VTK.  The cached arrays are automatically removed when the
  58.    VTK array destructs."""
  59.  
  60.     ######################################################################
  61.     # `object` interface.
  62.     ######################################################################
  63.     def __init__(self):
  64.         # The cache.
  65.         self._cache = {}
  66.  
  67.     def __len__(self):
  68.         return len(self._cache)
  69.  
  70.     def __contains__(self, vtk_arr):
  71.         key = vtk_arr.__this__
  72.         return self._cache.has_key(key)
  73.  
  74.     ######################################################################
  75.     # `ArrayCache` interface.
  76.     ######################################################################
  77.     def add(self, vtk_arr, np_arr):
  78.         """Add numpy array corresponding to the vtk array to the
  79.        cache."""
  80.         key = vtk_arr.__this__
  81.         cache = self._cache
  82.  
  83.         # Setup a callback so this cached array reference is removed
  84.         # when the VTK array is destroyed.  Passing the key to the
  85.         # `lambda` function is necessary because the callback will not
  86.         # receive the object (it will receive `None`) and thus there
  87.         # is no way to know which array reference one has to remove.
  88.         vtk_arr.AddObserver('DeleteEvent', lambda o, e, key=key: \
  89.                             self._remove_array(key))
  90.  
  91.         # Cache the array
  92.         cache[key] = np_arr
  93.  
  94.     def get(self, vtk_arr):
  95.         """Return the cached numpy array given a VTK array."""
  96.         key = vtk_arr.__this__
  97.         return self._cache[key]
  98.  
  99.     ######################################################################
  100.     # Non-public interface.
  101.     ######################################################################
  102.     def _remove_array(self, key):
  103.         """Private function that removes the cached array.  Do not
  104.        call this unless you know what you are doing."""
  105.         try:
  106.             del self._cache[key]
  107.         except KeyError:
  108.             pass
  109.  
  110.  
  111. ######################################################################
  112. # Setup a global `_array_cache`.  The array object cache caches all the
  113. # converted numpy arrays that are not copied.  This prevents the user
  114. # from deleting or resizing the numpy array after it has been sent down
  115. # to VTK.
  116. ######################################################################
  117.  
  118. _dummy = None
  119. # This makes the cache work even when the module is reloaded.
  120. for name in ['array_handler', 'enthought.tvtk.array_handler']:
  121.     if sys.modules.has_key(name):
  122.         mod = sys.modules[name]
  123.         if hasattr(mod, '_array_cache'):
  124.             _dummy = mod._array_cache
  125.         del mod
  126.         break
  127.  
  128. if _dummy:
  129.     _array_cache = _dummy
  130. else:
  131.     _array_cache = ArrayCache()
  132. del _dummy
  133.  
  134.  
  135.  
  136. ######################################################################
  137. # Array conversion functions.
  138. ######################################################################
  139. def get_vtk_array_type(numeric_array_type):
  140.     """Returns a VTK typecode given a numpy array."""
  141.     # This is a Mapping from numpy array types to VTK array types.
  142.     _arr_vtk = {numpy.dtype(numpy.character):vtkConstants.VTK_UNSIGNED_CHAR,
  143.                 numpy.dtype(numpy.uint8):vtkConstants.VTK_UNSIGNED_CHAR,
  144.                 numpy.dtype(numpy.uint16):vtkConstants.VTK_UNSIGNED_SHORT,
  145.                 numpy.dtype(numpy.int8):vtkConstants.VTK_CHAR,
  146.                 numpy.dtype(numpy.int16):vtkConstants.VTK_SHORT,
  147.                 numpy.dtype(numpy.int32):vtkConstants.VTK_INT,
  148.                 numpy.dtype(numpy.uint32):vtkConstants.VTK_UNSIGNED_INT,
  149.                 numpy.dtype(numpy.float32):vtkConstants.VTK_FLOAT,
  150.                 numpy.dtype(numpy.float64):vtkConstants.VTK_DOUBLE,
  151.                 numpy.dtype(numpy.complex64):vtkConstants.VTK_FLOAT,
  152.                 numpy.dtype(numpy.complex128):vtkConstants.VTK_DOUBLE,
  153.                 }
  154.     _extra = {numpy.dtype(ID_TYPE_CODE):vtkConstants.VTK_ID_TYPE,
  155.               numpy.dtype(ULONG_TYPE_CODE):vtkConstants.VTK_UNSIGNED_LONG,
  156.               numpy.dtype(LONG_TYPE_CODE):vtkConstants.VTK_LONG,
  157.              }
  158.     for t in _extra:
  159.         if t not in _arr_vtk:
  160.             _arr_vtk[t] = _extra[t]
  161.  
  162.     try:
  163.         return _arr_vtk[numeric_array_type]
  164.     except KeyError:
  165.         for key in _arr_vtk:
  166.             if numpy.issubdtype(numeric_array_type, key):
  167.                 return _arr_vtk[key]
  168.     raise TypeError, "Couldn't translate array's type to VTK"
  169.  
  170. def get_vtk_to_numeric_typemap():
  171.     """Returns the VTK array type to numpy array type mapping."""
  172.     _vtk_arr = {vtkConstants.VTK_BIT:numpy.bool,
  173.                 vtkConstants.VTK_CHAR:numpy.int8,
  174.                 vtkConstants.VTK_UNSIGNED_CHAR:numpy.uint8,
  175.                 vtkConstants.VTK_SHORT:numpy.int16,
  176.                 vtkConstants.VTK_UNSIGNED_SHORT:numpy.uint16,
  177.                 vtkConstants.VTK_INT:numpy.int32,
  178.                 vtkConstants.VTK_UNSIGNED_INT:numpy.uint32,
  179.                 vtkConstants.VTK_LONG:LONG_TYPE_CODE,
  180.                 vtkConstants.VTK_UNSIGNED_LONG:ULONG_TYPE_CODE,
  181.                 vtkConstants.VTK_ID_TYPE:ID_TYPE_CODE,
  182.                 vtkConstants.VTK_FLOAT:numpy.float32,
  183.                 vtkConstants.VTK_DOUBLE:numpy.float64}
  184.     return _vtk_arr
  185.  
  186.  
  187. def get_numeric_array_type(vtk_array_type):
  188.     """Returns a numpy array typecode given a VTK array type."""
  189.     return get_vtk_to_numeric_typemap()[vtk_array_type]
  190.  
  191.  
  192. def get_sizeof_vtk_array(vtk_array_type):
  193.     """Returns the size of a VTK array type."""
  194.     _size_dict = {vtkConstants.VTK_BIT : 1,
  195.                   vtkConstants.VTK_CHAR : 1,
  196.                   vtkConstants.VTK_UNSIGNED_CHAR : 1,
  197.                   vtkConstants.VTK_SHORT : 2,
  198.                   vtkConstants.VTK_UNSIGNED_SHORT : 2,
  199.                   vtkConstants.VTK_INT : 4,
  200.                   vtkConstants.VTK_UNSIGNED_INT : 4,
  201.                   vtkConstants.VTK_LONG : VTK_LONG_TYPE_SIZE,
  202.                   vtkConstants.VTK_UNSIGNED_LONG : VTK_LONG_TYPE_SIZE,
  203.                   vtkConstants.VTK_ID_TYPE : VTK_ID_TYPE_SIZE,
  204.                   vtkConstants.VTK_FLOAT : 4,
  205.                   vtkConstants.VTK_DOUBLE : 8 }
  206.     return _size_dict[vtk_array_type]
  207.  
  208.  
  209. def create_vtk_array(vtk_arr_type):
  210.     """Internal function used to create a VTK data array from another
  211.    VTK array given the VTK array type.
  212.    """
  213.     tmp = vtk.vtkDataArray.CreateDataArray(vtk_arr_type)
  214.     # CreateDataArray sets the refcount to 3 and this causes a severe
  215.     # memory leak.
  216.     tmp.SetReferenceCount(BASE_REFERENCE_COUNT)
  217.     return tmp
  218.  
  219.  
  220. def array2vtk(num_array, vtk_array=None):
  221.     """Converts a real numpy Array (or a Python list) to a VTK array
  222.    object.
  223.  
  224.    This function only works for real arrays.  Complex arrays are NOT
  225.    handled.  It also works for multi-component arrays.  However, only
  226.    1, and 2 dimensional arrays are supported.  This function is very
  227.    efficient, so large arrays should not be a problem.
  228.  
  229.    Even in cases when no copy of the numpy array data is performed,
  230.    a reference to the array is cached.  The passed array can
  231.    therefore be deleted safely in all circumstances.
  232.  
  233.    Parameters
  234.    ----------
  235.  
  236.    - num_array : numpy array or Python list/tuple
  237.  
  238.      The input array must be 1 or 2D.  A copy of the numeric array
  239.      data passed is made in the following circumstances:
  240.  
  241.       1. A Python list/tuple was passed.
  242.       2. A non-contiguous numpy array was passed.
  243.       3. A `vtkBitArray` instance was passed as the second argument.
  244.       4. The types of the `vtk_array` and the `num_array` are not
  245.          equivalent to each other.  For example if one is an integer
  246.          array and the other a float.
  247.  
  248.  
  249.    - vtk_array : `vtkDataArray` (default: `None`)
  250.  
  251.      If an optional `vtkDataArray` instance, is passed as an argument
  252.      then a new array is not created and returned.  The passed array
  253.      is itself returned.
  254.  
  255.    """
  256.  
  257.     z = numpy.asarray(num_array)
  258.  
  259.     shape = z.shape
  260.     assert len(shape) < 3, \
  261.            "Only arrays of dimensionality 2 or lower are allowed!"
  262.     assert not numpy.issubdtype(z.dtype, complex), \
  263.            "Complex numpy arrays cannot be converted to vtk arrays."\
  264.            "Use real() or imag() to get a component of the array before"\
  265.            " passing it to vtk."
  266.  
  267.     # First create an array of the right type by using the typecode.
  268.     # Bit arrays need special casing.
  269.     bit_array = False
  270.     if vtk_array is None:
  271.         vtk_typecode = get_vtk_array_type(z.dtype)
  272.         result_array = create_vtk_array(vtk_typecode)
  273.     elif vtk_array.GetDataType() == vtkConstants.VTK_BIT:
  274.         vtk_typecode = vtkConstants.VTK_CHAR
  275.         result_array = create_vtk_array(vtkConstants.VTK_CHAR)
  276.         bit_array = True
  277.     else:
  278.         vtk_typecode = vtk_array.GetDataType()
  279.         result_array = vtk_array
  280.  
  281.     # Find the shape and set number of components.
  282.     if len(shape) == 1:
  283.         result_array.SetNumberOfComponents(1)
  284.     else:
  285.         result_array.SetNumberOfComponents(shape[1])
  286.  
  287.     result_array.SetNumberOfTuples(shape[0])
  288.  
  289.     # Ravel the array appropriately.
  290.     arr_dtype = get_numeric_array_type(vtk_typecode)
  291.     if numpy.issubdtype(z.dtype, arr_dtype):
  292.         z_flat = numpy.ravel(z)
  293.     else:
  294.         z_flat = numpy.ravel(z).astype(arr_dtype)
  295.  
  296.     # Point the VTK array to the numpy data.  The last argument (1)
  297.     # tells the array not to deallocate.
  298.     result_array.SetVoidArray(numpy.getbuffer(z_flat), len(z_flat), 1)
  299.  
  300.     if bit_array:
  301.         # Handle bit arrays -- they have to be copied.  Note that bit
  302.         # arrays are used ONLY when the user has passed one as an
  303.         # argument to this function.
  304.         vtk_array.SetNumberOfTuples(result_array.GetNumberOfTuples())
  305.         vtk_array.SetNumberOfComponents(result_array.GetNumberOfComponents())
  306.         for i in range(result_array.GetNumberOfComponents()):
  307.             vtk_array.CopyComponent(i, result_array, i)
  308.         result_array = vtk_array
  309.     else:
  310.         # Save a reference to the flatted array in the array cache.
  311.         # This prevents the user from deleting or resizing the array
  312.         # and getting into serious trouble.  This is only done for
  313.         # non-bit array cases where the data is not copied.
  314.         global _array_cache
  315.         _array_cache.add(result_array, z_flat)
  316.  
  317.     return result_array
  318.  
  319.  
  320. def vtk2array(vtk_array):
  321.     """Converts a VTK data array to a numpy array.
  322.  
  323.    Given a subclass of vtkDataArray, this function returns an
  324.    appropriate numpy array containing the same data.  The function
  325.    is very efficient since it uses the VTK imaging pipeline to
  326.    convert the data.  If a sufficiently new version of VTK (5.2) is
  327.    installed then it actually uses the buffer interface to return a
  328.    view of the VTK array in the returned numpy array.
  329.  
  330.    Parameters
  331.    ----------
  332.  
  333.    - vtk_array : `vtkDataArray`
  334.  
  335.      The VTK data array to be converted.
  336.  
  337.    """
  338.     typ = vtk_array.GetDataType()
  339.     assert typ in get_vtk_to_numeric_typemap().keys(), \
  340.            "Unsupported array type %s"%typ
  341.  
  342.     shape = vtk_array.GetNumberOfTuples(), \
  343.             vtk_array.GetNumberOfComponents()
  344.     if shape[0] == 0:
  345.         dtype = get_numeric_array_type(typ)
  346.         return numpy.array([], dtype)
  347.  
  348.     # First check if this array already has a numpy array cached, if
  349.     # it does, reshape that and return it.
  350.     if vtk_array in _array_cache:
  351.         arr = _array_cache.get(vtk_array)
  352.         if shape[1] == 1:
  353.             shape = (shape[0], )
  354.         arr = numpy.reshape(arr, shape)
  355.         return arr
  356.  
  357.     # If VTK's new numpy support is available, use the buffer interface.
  358.     if numpy_support is not None and typ != vtkConstants.VTK_BIT:
  359.         dtype = get_numeric_array_type(typ)
  360.         result = numpy.frombuffer(vtk_array, dtype=dtype)
  361.         if shape[1] == 1:
  362.             shape = (shape[0], )
  363.         result.shape = shape
  364.         return result
  365.  
  366.     # Setup an imaging pipeline to export the array.
  367.     img_data = vtk.vtkImageData()
  368.     img_data.SetDimensions(shape[0], 1, 1)
  369.     if typ == vtkConstants.VTK_BIT:
  370.         iarr = vtk.vtkCharArray()
  371.         iarr.DeepCopy(vtk_array)
  372.         img_data.GetPointData().SetScalars(iarr)
  373.     elif typ == vtkConstants.VTK_ID_TYPE:
  374.         # Needed since VTK_ID_TYPE does not work with VTK 4.5.
  375.         iarr = vtk.vtkLongArray()
  376.         iarr.SetNumberOfTuples(vtk_array.GetNumberOfTuples())
  377.         nc = vtk_array.GetNumberOfComponents()
  378.         iarr.SetNumberOfComponents(nc)
  379.         for i in range(nc):
  380.             iarr.CopyComponent(i, vtk_array, i)
  381.         img_data.GetPointData().SetScalars(iarr)
  382.     else:
  383.         img_data.GetPointData().SetScalars(vtk_array)
  384.  
  385.     img_data.SetNumberOfScalarComponents(shape[1])
  386.     if typ == vtkConstants.VTK_ID_TYPE:
  387.         # Hack necessary because vtkImageData can't handle VTK_ID_TYPE.
  388.         img_data.SetScalarType(vtkConstants.VTK_LONG)
  389.         r_dtype = get_numeric_array_type(vtkConstants.VTK_LONG)
  390.     elif typ == vtkConstants.VTK_BIT:
  391.         img_data.SetScalarType(vtkConstants.VTK_CHAR)
  392.         r_dtype = get_numeric_array_type(vtkConstants.VTK_CHAR)
  393.     else:
  394.         img_data.SetScalarType(typ)
  395.         r_dtype = get_numeric_array_type(typ)
  396.     img_data.Update()
  397.  
  398.     exp = vtk.vtkImageExport()
  399.     exp.SetInput(img_data)
  400.  
  401.     # Create an array of the right size and export the image into it.
  402.     im_arr = numpy.empty((shape[0]*shape[1],), r_dtype)
  403.     exp.Export(im_arr)
  404.  
  405.     # Now reshape it.
  406.     if shape[1] == 1:
  407.         shape = (shape[0], )
  408.     im_arr = numpy.reshape(im_arr, shape)
  409.     return im_arr
  410.  
  411.  
  412. def array2vtkCellArray(num_array, vtk_array=None):
  413.     """Given a nested Python list or a numpy array, this method
  414.    creates a vtkCellArray instance and returns it.
  415.  
  416.    A variety of input arguments are supported as described in the
  417.    Parameter documentation.  If numpy arrays are given, this method
  418.    is highly efficient.  This function is most efficient if the
  419.    passed numpy arrays have a typecode `ID_TYPE_CODE`.  Otherwise a
  420.    typecast is necessary and this involves an extra copy.  This
  421.    method *always copies* the input data.
  422.  
  423.    An alternative and more efficient way to build the connectivity
  424.    list is to create a vtkIdTypeArray having data of the form
  425.    (npts,p0,p1,...p(npts-1), repeated for each cell) and then call
  426.    <vtkCellArray_instance>.SetCells(n_cell, id_list).
  427.  
  428.    Parameters
  429.    ----------
  430.  
  431.    - num_array : numpy array or Python list/tuple
  432.  
  433.      Valid values are:
  434.  
  435.        1. A Python list of 1D lists.  Each 1D list can contain one
  436.           cell connectivity list.  This is very slow and is to be
  437.           used only when efficiency is of no consequence.
  438.  
  439.        2. A 2D numpy array with the cell connectivity list.
  440.  
  441.        3. A Python list of 2D numpy arrays.  Each numeric array can
  442.           have a different shape.  This makes it easy to generate a
  443.           cell array having cells of different kinds.
  444.  
  445.    - vtk_array : `vtkCellArray` (default: `None`)
  446.  
  447.      If an optional `vtkCellArray` instance, is passed as an argument
  448.      then a new array is not created and returned.  The passed array
  449.      is itself modified and returned.
  450.  
  451.    Example
  452.    -------
  453.  
  454.       >>> a = [[0], [1, 2], [3, 4, 5], [6, 7, 8, 9]]
  455.       >>> cells = array_handler.array2vtkCellArray(a)
  456.       >>> a = numpy.array([[0,1,2], [3,4,5], [6,7,8]], 'l')
  457.       >>> cells = array_handler.array2vtkCellArray(a)
  458.       >>> l_a = [a[:,:1], a[:2,:2], a]
  459.       >>> cells = array_handler.array2vtkCellArray(l_a)
  460.  
  461.    """
  462.     if vtk_array:
  463.         cells = vtk_array
  464.     else:
  465.         cells = vtk.vtkCellArray()
  466.     assert cells.GetClassName() == 'vtkCellArray', \
  467.            'Second argument must be a `vtkCellArray` instance.'
  468.  
  469.     if len(num_array) == 0:
  470.         return cells
  471.  
  472.     ########################################
  473.     # Internal functions.
  474.     def _slow_array2cells(z, cells):
  475.         cells.Reset()
  476.         vtk_ids = vtk.vtkIdList()
  477.         for i in z:
  478.             vtk_ids.Reset()
  479.             for j in i:
  480.                 vtk_ids.InsertNextId(j)
  481.             cells.InsertNextCell(vtk_ids)
  482.  
  483.     def _get_tmp_array(arr):
  484.         try:
  485.             tmp_arr = numpy.asarray(arr, ID_TYPE_CODE)
  486.         except TypeError:
  487.             tmp_arr = arr.astype(ID_TYPE_CODE)
  488.         return tmp_arr
  489.  
  490.     def _set_cells(cells, n_cells, id_typ_arr):
  491.         vtk_arr = vtk.vtkIdTypeArray()
  492.         array2vtk(id_typ_arr, vtk_arr)
  493.         cells.SetCells(n_cells, vtk_arr)
  494.     ########################################
  495.  
  496.     msg = "Invalid argument.  Valid types are a Python list of lists,"\
  497.           " a Python list of numpy arrays, or a numpy array."
  498.  
  499.     if issubclass(type(num_array), (types.ListType, types.TupleType)):
  500.         assert len(num_array[0]) > 0, "Input array must be 2D."
  501.         tp = type(num_array[0])
  502.         if issubclass(tp, types.ListType): # Pure Python list.
  503.             _slow_array2cells(num_array, cells)
  504.             return cells
  505.         elif issubclass(tp, numpy.ndarray):  # List of arrays.
  506.             # Check shape of array and find total size.
  507.             tot_size = 0
  508.             n_cells = 0
  509.             for arr in num_array:
  510.                 assert len(arr.shape) == 2, "Each array must be 2D"
  511.                 shp = arr.shape
  512.                 tot_size += shp[0]*(shp[1] + 1)
  513.                 n_cells += shp[0]
  514.             # Create an empty array.
  515.             id_typ_arr = numpy.empty((tot_size,), ID_TYPE_CODE)
  516.             # Now populate it with the ids.
  517.             count = 0
  518.             for arr in num_array:
  519.                 tmp_arr = _get_tmp_array(arr)
  520.                 shp = arr.shape
  521.                 sz = shp[0]*(shp[1] + 1)
  522.                 set_id_type_array(tmp_arr, id_typ_arr[count:count+sz])
  523.                 count += sz
  524.             # Now set them cells.
  525.             _set_cells(cells, n_cells, id_typ_arr)
  526.             return cells
  527.         else:
  528.             raise TypeError, msg
  529.     elif issubclass(type(num_array), numpy.ndarray):
  530.         assert len(num_array.shape) == 2, "Input array must be 2D."
  531.         tmp_arr = _get_tmp_array(num_array)
  532.         shp = tmp_arr.shape
  533.         id_typ_arr = numpy.empty((shp[0]*(shp[1] + 1),), ID_TYPE_CODE)
  534.         set_id_type_array(tmp_arr, id_typ_arr)
  535.         _set_cells(cells, shp[0], id_typ_arr)
  536.         return cells
  537.     else:
  538.         raise TypeError, msg
  539.  
  540.  
  541. def array2vtkPoints(num_array, vtk_points=None):
  542.     """Converts a numpy array/Python list to a vtkPoints object.
  543.  
  544.    Unless a Python list/tuple or a non-contiguous array is given, no
  545.    copy of the data is made.  Thus the function is very efficient.
  546.  
  547.    Parameters
  548.    ----------
  549.  
  550.    - num_array : numpy array or Python list/tuple
  551.  
  552.      The input array must be 2D with `shape[1] == 3`.
  553.  
  554.    - vtk_points : `vtkPoints` (default: `None`)
  555.  
  556.      If an optional `vtkPoints` instance, is passed as an argument
  557.      then a new array is not created and returned.  The passed array
  558.      is itself modified and returned.
  559.  
  560.    """
  561.     if vtk_points:
  562.         points  = vtk_points
  563.     else:
  564.         points = vtk.vtkPoints()
  565.  
  566.     arr = numpy.asarray(num_array)
  567.     assert len(arr.shape) == 2, "Points array must be 2 dimensional."
  568.     assert arr.shape[1] == 3, "Incorrect shape: shape[1] must be 3."
  569.     vtk_array = array2vtk(arr)
  570.     points.SetData(vtk_array)
  571.     return points
  572.  
  573.  
  574. def array2vtkIdList(num_array, vtk_idlist=None):
  575.     """Converts a numpy array/Python list to a vtkIdList object.
  576.  
  577.    Parameters
  578.    ----------
  579.  
  580.    - num_array : numpy array or Python list/tuple
  581.  
  582.      The input array must be 2D with `shape[1] == 3`.
  583.  
  584.    - vtk_idlist : `vtkIdList` (default: `None`)
  585.  
  586.      If an optional `vtkIdList` instance, is passed as an argument
  587.      then a new array is not created and returned.  The passed array
  588.      is itself modified and returned.
  589.  
  590.    """
  591.     if vtk_idlist:
  592.         ids = vtk_idlist
  593.     else:
  594.         ids = vtk.vtkIdList()
  595.  
  596.     arr = numpy.asarray(num_array)
  597.     assert len(arr.shape) == 1, "Array for vtkIdList must be 1D"
  598.     ids.SetNumberOfIds(len(arr))
  599.     for i, j in enumerate(arr):
  600.         ids.SetId(i, j)
  601.     return ids
  602.  
  603.  
  604. ######################################################################
  605. # Array argument handling functions.
  606. ######################################################################
  607.  
  608. def is_array(arr):
  609.     """Returns True if the passed `arr` is a numpy array or a List."""
  610.     if issubclass(type(arr), (numpy.ndarray, types.ListType)):
  611.         return True
  612.     return False
  613.  
  614.  
  615. def convert_array(arr, vtk_typ=None):
  616.     """Convert the given array to the optional type specified by
  617.    `vtk_typ`.
  618.  
  619.    Parameters
  620.    ----------
  621.  
  622.    - arr : numpy array/list.
  623.    - vtk_typ : `string` or `None`
  624.      represents the type the array is to be converted to.
  625.  
  626.    """
  627.     if vtk_typ:
  628.         conv = {'vtkCellArray': array2vtkCellArray,
  629.                 'vtkPoints': array2vtkPoints,
  630.                 'vtkIdList': array2vtkIdList}
  631.         if vtk_typ in conv.keys():
  632.             vtk_arr = getattr(vtk, vtk_typ)()
  633.             return conv[vtk_typ](arr, vtk_arr)
  634.         elif vtk_typ.find('Array') > -1:
  635.             try:
  636.                 vtk_arr = getattr(vtk, vtk_typ)()
  637.             except TypeError: # vtk_typ == 'vtkDataArray'
  638.                 return array2vtk(arr)
  639.             else:
  640.                 return array2vtk(arr, vtk_arr)
  641.         else:
  642.             return arr
  643.     else:
  644.         return array2vtk(arr)
  645.  
  646.  
  647. def is_array_sig(s):
  648.     """Given a signature, return if the signature has an array."""
  649.     if not isinstance(s, basestring):
  650.         return False
  651.     arr_types = ['Array', 'vtkPoints', 'vtkIdList']
  652.     for i in arr_types:
  653.         if s.find(i) > -1:
  654.             return True
  655.     return False
  656.  
  657.  
  658. def is_array_or_vtkarray(arg):
  659.     """Returns True if the argument is an array/Python list or if it
  660.    is a vtk array."""
  661.  
  662.     if is_array(arg):
  663.         return True
  664.     else:
  665.         if hasattr(arg, '_vtk_obj'):
  666.             if is_array_sig(arg._vtk_obj.__class__.__name__):
  667.                 return True
  668.     return False
  669.  
  670.  
  671. def get_correct_sig(args, sigs):
  672.     """Given a list of args and a collection of possible signatures,
  673.    this function returns the most appropriate signature.  This
  674.    function is only called by deref_array.  This implies that one of
  675.    the signatures has an array type.
  676.  
  677.    """
  678.     # First do the trivial cases.
  679.     if sigs is None:
  680.         return None
  681.     if len(sigs) == 1:
  682.         return sigs[0]
  683.     else:
  684.         # Non-trivial cases.
  685.         la = len(args)
  686.         candidate_sigs = [s for s in sigs if len(s) == la]
  687.         count = len(candidate_sigs)
  688.         if count == 0:
  689.             # No sig has the right number of args.
  690.             msg = "Insufficient number of arguments to method."\
  691.                   "Valid arguments are:\n%s"%sigs
  692.             raise TypeError, msg
  693.         elif count == 1:
  694.             # If only one of the sigs has the right number of args,
  695.             # return it.
  696.             return candidate_sigs[0]
  697.         else:
  698.             # More than one sig has the same number of args.
  699.             # Check if args need conversion at all.
  700.             array_idx = [i for i, a in enumerate(args) \
  701.                          if is_array_or_vtkarray(a)]
  702.             n_arr = len(array_idx)
  703.             if n_arr == 0:
  704.                 # No conversion necessary so signature info is
  705.                 # useless.
  706.                 return None
  707.             else:
  708.                 # Need to find the right sig.  This is done by finding
  709.                 # the first signature that matches all the arrays in
  710.                 # the argument.
  711.                 for sig in candidate_sigs:
  712.                     array_in_sig = [is_array_sig(s) for s in sig]
  713.                     if array_in_sig.count(True) != len(array_idx):
  714.                         continue
  715.                     bad = False
  716.                     for i in array_idx:
  717.                         if not array_in_sig[i]:
  718.                             bad = True
  719.                     if not bad:
  720.                         return sig
  721.                 # Could not find any valid signature, so give up.
  722.                 return None
  723.  
  724.  
  725. def deref_vtk(obj):
  726.     """Dereferences the VTK object from the object if possible.  This
  727.    is duplicated from `tvtk_base.py` because I'd like to keep this
  728.    module independent of `tvtk_base.py`.
  729.    """
  730.     if hasattr(obj, '_vtk_obj'):
  731.         return obj._vtk_obj
  732.     else:
  733.         return obj
  734.  
  735.  
  736. def deref_array(args, sigs=None):
  737.     """Given a bunch of arguments and optional signature information,
  738.    this converts the arguments suitably.  If the argument is either a
  739.    Python list or a numpy array it is converted to a suitable type
  740.    based on the signature information.  If it is not an array, but a
  741.    TVTK object the VTK object is dereferenced.  Otherwise nothing is
  742.    done.  If no signature information is provided the arrays are
  743.    automatically converted (this can sometimes go wrong).  The
  744.    signature information is provided in the form of a list of lists.
  745.  
  746.    """
  747.     ret = []
  748.     sig = get_correct_sig(args, sigs)
  749.     if sig:
  750.         for a, s in zip(args, sig):
  751.             if is_array(a) and is_array_sig(s):
  752.                 ret.append(convert_array(a, s))
  753.             else:
  754.                 ret.append(deref_vtk(a))
  755.     else:
  756.         for a in args:
  757.             if is_array(a):
  758.                 ret.append(convert_array(a))
  759.             else:
  760.                 ret.append(deref_vtk(a))
  761.     return ret
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement