253 lines
9.1 KiB
Python
253 lines
9.1 KiB
Python
"""This module adds support to easily import and export NumPy
|
|
(http://numpy.scipy.org) arrays into/out of VTK arrays. The code is
|
|
loosely based on TVTK (https://svn.enthought.com/enthought/wiki/TVTK).
|
|
|
|
This code depends on an addition to the VTK data arrays made by Berk
|
|
Geveci to make it support Python's buffer protocol (on Feb. 15, 2008).
|
|
|
|
The main functionality of this module is provided by the two functions:
|
|
numpy_to_vtk,
|
|
vtk_to_numpy.
|
|
|
|
|
|
Caveats:
|
|
--------
|
|
|
|
- Bit arrays in general do not have a numpy equivalent and are not
|
|
supported. Char arrays are also not easy to handle and might not
|
|
work as you expect. Patches welcome.
|
|
|
|
- You need to make sure you hold a reference to a Numpy array you want
|
|
to import into VTK. If not you'll get a segfault (in the best case).
|
|
The same holds in reverse when you convert a VTK array to a numpy
|
|
array -- don't delete the VTK array.
|
|
|
|
|
|
Created by Prabhu Ramachandran in Feb. 2008.
|
|
"""
|
|
|
|
from . import vtkConstants
|
|
from vtkmodules.vtkCommonCore import vtkDataArray, vtkIdTypeArray, vtkLongArray
|
|
import numpy
|
|
|
|
# Useful constants for VTK arrays.
|
|
VTK_ID_TYPE_SIZE = vtkIdTypeArray().GetDataTypeSize()
|
|
if VTK_ID_TYPE_SIZE == 4:
|
|
ID_TYPE_CODE = numpy.int32
|
|
elif VTK_ID_TYPE_SIZE == 8:
|
|
ID_TYPE_CODE = numpy.int64
|
|
|
|
VTK_LONG_TYPE_SIZE = vtkLongArray().GetDataTypeSize()
|
|
if VTK_LONG_TYPE_SIZE == 4:
|
|
LONG_TYPE_CODE = numpy.int32
|
|
ULONG_TYPE_CODE = numpy.uint32
|
|
elif VTK_LONG_TYPE_SIZE == 8:
|
|
LONG_TYPE_CODE = numpy.int64
|
|
ULONG_TYPE_CODE = numpy.uint64
|
|
|
|
|
|
def get_vtk_array_type(numpy_array_type):
|
|
"""Returns a VTK typecode given a numpy array."""
|
|
# This is a Mapping from numpy array types to VTK array types.
|
|
_np_vtk = {numpy.uint8:vtkConstants.VTK_UNSIGNED_CHAR,
|
|
numpy.uint16:vtkConstants.VTK_UNSIGNED_SHORT,
|
|
numpy.uint32:vtkConstants.VTK_UNSIGNED_INT,
|
|
numpy.uint64:vtkConstants.VTK_UNSIGNED_LONG_LONG,
|
|
numpy.int8:vtkConstants.VTK_SIGNED_CHAR,
|
|
numpy.int16:vtkConstants.VTK_SHORT,
|
|
numpy.int32:vtkConstants.VTK_INT,
|
|
numpy.int64:vtkConstants.VTK_LONG_LONG,
|
|
numpy.float32:vtkConstants.VTK_FLOAT,
|
|
numpy.float64:vtkConstants.VTK_DOUBLE,
|
|
numpy.complex64:vtkConstants.VTK_FLOAT,
|
|
numpy.complex128:vtkConstants.VTK_DOUBLE}
|
|
for key, vtk_type in _np_vtk.items():
|
|
if numpy_array_type == key or \
|
|
numpy.issubdtype(numpy_array_type, key) or \
|
|
numpy_array_type == numpy.dtype(key):
|
|
return vtk_type
|
|
raise TypeError(
|
|
'Could not find a suitable VTK type for %s' % (str(numpy_array_type)))
|
|
|
|
def get_vtk_to_numpy_typemap():
|
|
"""Returns the VTK array type to numpy array type mapping."""
|
|
_vtk_np = {vtkConstants.VTK_BIT:numpy.uint8,
|
|
vtkConstants.VTK_CHAR:numpy.int8,
|
|
vtkConstants.VTK_SIGNED_CHAR:numpy.int8,
|
|
vtkConstants.VTK_UNSIGNED_CHAR:numpy.uint8,
|
|
vtkConstants.VTK_SHORT:numpy.int16,
|
|
vtkConstants.VTK_UNSIGNED_SHORT:numpy.uint16,
|
|
vtkConstants.VTK_INT:numpy.int32,
|
|
vtkConstants.VTK_UNSIGNED_INT:numpy.uint32,
|
|
vtkConstants.VTK_LONG:LONG_TYPE_CODE,
|
|
vtkConstants.VTK_LONG_LONG:numpy.int64,
|
|
vtkConstants.VTK_UNSIGNED_LONG:ULONG_TYPE_CODE,
|
|
vtkConstants.VTK_UNSIGNED_LONG_LONG:numpy.uint64,
|
|
vtkConstants.VTK_ID_TYPE:ID_TYPE_CODE,
|
|
vtkConstants.VTK_FLOAT:numpy.float32,
|
|
vtkConstants.VTK_DOUBLE:numpy.float64}
|
|
return _vtk_np
|
|
|
|
|
|
def get_numpy_array_type(vtk_array_type):
|
|
"""Returns a numpy array typecode given a VTK array type."""
|
|
return get_vtk_to_numpy_typemap()[vtk_array_type]
|
|
|
|
|
|
def create_vtk_array(vtk_arr_type):
|
|
"""Internal function used to create a VTK data array from another
|
|
VTK array given the VTK array type.
|
|
"""
|
|
return vtkDataArray.CreateDataArray(vtk_arr_type)
|
|
|
|
|
|
def numpy_to_vtk(num_array, deep=0, array_type=None):
|
|
"""Converts a real numpy Array to a VTK array object.
|
|
|
|
This function only works for real arrays.
|
|
Complex arrays are NOT handled. It also works for multi-component
|
|
arrays. However, only 1, and 2 dimensional arrays are supported.
|
|
This function is very efficient, so large arrays should not be a
|
|
problem.
|
|
|
|
If the second argument is set to 1, the array is deep-copied from
|
|
from numpy. This is not as efficient as the default behavior
|
|
(shallow copy) and uses more memory but detaches the two arrays
|
|
such that the numpy array can be released.
|
|
|
|
WARNING: You must maintain a reference to the passed numpy array, if
|
|
the numpy data is gc'd and VTK will point to garbage which will in
|
|
the best case give you a segfault.
|
|
|
|
Parameters:
|
|
|
|
num_array
|
|
a 1D or 2D, real numpy array.
|
|
|
|
"""
|
|
|
|
z = numpy.asarray(num_array)
|
|
if not z.flags.contiguous:
|
|
z = numpy.ascontiguousarray(z)
|
|
|
|
shape = z.shape
|
|
assert z.flags.contiguous, 'Only contiguous arrays are supported.'
|
|
assert len(shape) < 3, \
|
|
"Only arrays of dimensionality 2 or lower are allowed!"
|
|
assert not numpy.issubdtype(z.dtype, numpy.dtype(complex).type), \
|
|
"Complex numpy arrays cannot be converted to vtk arrays."\
|
|
"Use real() or imag() to get a component of the array before"\
|
|
" passing it to vtk."
|
|
|
|
# First create an array of the right type by using the typecode.
|
|
if array_type:
|
|
vtk_typecode = array_type
|
|
else:
|
|
vtk_typecode = get_vtk_array_type(z.dtype)
|
|
result_array = create_vtk_array(vtk_typecode)
|
|
|
|
# Fixup shape in case its empty or scalar.
|
|
try:
|
|
testVar = shape[0]
|
|
except:
|
|
shape = (0,)
|
|
|
|
# Find the shape and set number of components.
|
|
if len(shape) == 1:
|
|
result_array.SetNumberOfComponents(1)
|
|
else:
|
|
result_array.SetNumberOfComponents(shape[1])
|
|
|
|
# We don't need to call result_array.SetNumberOfTuples(shape[0])
|
|
# because we will use result_array.SetVoidPointer
|
|
# which takes care of setting the NumberOfTuples
|
|
# Calling SetNumberOfTuples will result in a memory allocation
|
|
# that will be deleted on SetVoidPointer.
|
|
|
|
# Ravel the array appropriately.
|
|
arr_dtype = get_numpy_array_type(vtk_typecode)
|
|
if numpy.issubdtype(z.dtype, arr_dtype) or \
|
|
z.dtype == numpy.dtype(arr_dtype):
|
|
z_flat = numpy.ravel(z)
|
|
else:
|
|
z_flat = numpy.ravel(z).astype(arr_dtype)
|
|
# z_flat is now a standalone object with no references from the caller.
|
|
# As such, it will drop out of this scope and cause memory issues if we
|
|
# do not deep copy its data.
|
|
deep = 1
|
|
|
|
# Point the VTK array to the numpy data. The last argument (1)
|
|
# tells the array not to deallocate.
|
|
result_array.SetVoidArray(z_flat, len(z_flat), 1)
|
|
if deep:
|
|
copy = result_array.NewInstance()
|
|
copy.DeepCopy(result_array)
|
|
result_array = copy
|
|
else:
|
|
result_array._numpy_reference = z
|
|
return result_array
|
|
|
|
def numpy_to_vtkIdTypeArray(num_array, deep=0):
|
|
isize = vtkIdTypeArray().GetDataTypeSize()
|
|
dtype = num_array.dtype
|
|
if isize == 4:
|
|
if dtype != numpy.int32:
|
|
raise ValueError(
|
|
'Expecting a numpy.int32 array, got %s instead.' % (str(dtype)))
|
|
else:
|
|
if dtype != numpy.int64:
|
|
raise ValueError(
|
|
'Expecting a numpy.int64 array, got %s instead.' % (str(dtype)))
|
|
|
|
return numpy_to_vtk(num_array, deep, vtkConstants.VTK_ID_TYPE)
|
|
|
|
def vtk_to_numpy(vtk_array):
|
|
"""Converts a VTK data array to a numpy array.
|
|
|
|
Given a subclass of vtkDataArray, this function returns an
|
|
appropriate numpy array containing the same data -- it actually
|
|
points to the same data.
|
|
|
|
Parameters
|
|
|
|
vtk_array
|
|
The VTK data array to be converted.
|
|
|
|
"""
|
|
typ = vtk_array.GetDataType()
|
|
assert typ in get_vtk_to_numpy_typemap().keys(), \
|
|
"Unsupported array type %s"%typ
|
|
|
|
shape = vtk_array.GetNumberOfTuples(), \
|
|
vtk_array.GetNumberOfComponents()
|
|
|
|
# Get the data via the buffer interface
|
|
dtype = get_numpy_array_type(typ)
|
|
try:
|
|
if typ != vtkConstants.VTK_BIT:
|
|
result = numpy.frombuffer(vtk_array, dtype=dtype)
|
|
else:
|
|
result = numpy.unpackbits(vtk_array, count=shape[0])
|
|
except ValueError:
|
|
# http://mail.scipy.org/pipermail/numpy-tickets/2011-August/005859.html
|
|
# numpy 1.5.1 (and maybe earlier) has a bug where if frombuffer is
|
|
# called with an empty buffer, it throws ValueError exception. This
|
|
# handles that issue.
|
|
if shape[0] == 0:
|
|
# create an empty array with the given shape.
|
|
result = numpy.empty(shape, dtype=dtype)
|
|
else:
|
|
raise
|
|
if shape[1] == 1:
|
|
shape = (shape[0], )
|
|
try:
|
|
result.shape = shape
|
|
except ValueError:
|
|
if shape[0] == 0:
|
|
# Refer to https://github.com/numpy/numpy/issues/2536 .
|
|
# For empty array, reshape fails. Create the empty array explicitly
|
|
# if that happens.
|
|
result = numpy.empty(shape, dtype=dtype)
|
|
else: raise
|
|
return result
|