enthought / mayavi
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-2020,  Enthought, Inc.
14
# License: BSD Style.
15

16 4
import sys
17

18 4
import vtk
19 4
from vtk.util import vtkConstants
20 4
try:
21 4
    from vtk.util import numpy_support
22 0
except ImportError:
23 0
    numpy_support = None
24

25 4
import numpy
26

27
# Enthought library imports.
28 4
try:
29 4
    from tvtk.array_ext import set_id_type_array
30 4
    HAS_ARRAY_EXT = True
31 0
except ImportError:
32 0
    HAS_ARRAY_EXT = False
33

34 4
from tvtk.common import is_old_pipeline
35

36
# Useful constants for VTK arrays.
37 4
VTK_ID_TYPE_SIZE = vtk.vtkIdTypeArray().GetDataTypeSize()
38 4
if VTK_ID_TYPE_SIZE == 4:
39 0
    ID_TYPE_CODE = numpy.int32
40 4
elif VTK_ID_TYPE_SIZE == 8:
41 4
    ID_TYPE_CODE = numpy.int64
42

43 4
VTK_LONG_TYPE_SIZE = vtk.vtkLongArray().GetDataTypeSize()
44 4
if VTK_LONG_TYPE_SIZE == 4:
45 0
    LONG_TYPE_CODE = numpy.int32
46 0
    ULONG_TYPE_CODE = numpy.uint32
47 4
elif VTK_LONG_TYPE_SIZE == 8:
48 4
    LONG_TYPE_CODE = numpy.int64
49 4
    ULONG_TYPE_CODE = numpy.uint64
50

51 4
BASE_REFERENCE_COUNT = vtk.vtkObject().GetReferenceCount()
52

53 4
if sys.version_info[0] > 2:
54 4
    unicode = str
55

56

57 4
def getbuffer(array):
58 4
    return getattr(numpy, 'getbuffer', memoryview)(array)
59

60

61 4
def set_id_type_array_py(id_array, out_array):
62
    """Given a 2D Int array (`id_array`), and a contiguous 1D numarray array
63
    (`out_array`) having the correct size, this function sets the data from
64
    `id_array` into `out_array` so that it can be used in place of a
65
    `vtkIdTypeArray` in order to set the cells of a `vtkCellArray`.
66

67
    Note that if `shape = id_array.shape` then `size(out_array) ==
68
    shape[0]*(shape[1] + 1)` should be true. If not you'll get an
69
    `AssertionError`.
70

71
    `id_array` need not be contiguous but `out_array` must be.
72

73
    """
74 4
    assert numpy.issubdtype(id_array.dtype, numpy.signedinteger)
75 4
    assert out_array.flags.contiguous == 1, \
76
        "out_array must be contiguous."
77 4
    shp = id_array.shape
78 4
    assert len(shp) == 2, "id_array must be a two dimensional array."
79 4
    sz = out_array.size
80 4
    e_sz = shp[0]*(shp[1]+1)
81 4
    assert sz == e_sz, \
82
        "out_array size is incorrect, expected: %s, given: %s" % (e_sz, sz)
83

84
    # we are guaranteed contiguous, so these just change the view (no copy)
85 4
    out_shp = out_array.shape
86 4
    out_array.shape = (shp[0], shp[1] + 1)
87 4
    out_array[:, 0] = shp[1]
88 4
    out_array[:, 1:] = id_array
89 4
    out_array.shape = out_shp
90

91

92 4
if not HAS_ARRAY_EXT:
93 0
    set_id_type_array = set_id_type_array_py
94

95

96
######################################################################
97
# The array cache.
98
######################################################################
99 4
class ArrayCache(object):
100

101
    """Caches references to numpy arrays that are not copied but views
102
    of which are converted to VTK arrays.  The caching prevents the user
103
    from deleting or resizing the numpy array after it has been sent
104
    down to VTK.  The cached arrays are automatically removed when the
105
    VTK array destructs."""
106

107
    ######################################################################
108
    # `object` interface.
109
    ######################################################################
110 4
    def __init__(self):
111
        # The cache.
112 4
        self._cache = {}
113

114 4
    def __len__(self):
115 4
        return len(self._cache)
116

117 4
    def __contains__(self, vtk_arr):
118 4
        key = vtk_arr.__this__
119 4
        return key in self._cache
120

121
    ######################################################################
122
    # `ArrayCache` interface.
123
    ######################################################################
124 4
    def add(self, vtk_arr, np_arr):
125
        """Add numpy array corresponding to the vtk array to the
126
        cache."""
127 4
        key = vtk_arr.__this__
128 4
        cache = self._cache
129

130
        # Setup a callback so this cached array reference is removed
131
        # when the VTK array is destroyed.  Passing the key to the
132
        # `lambda` function is necessary because the callback will not
133
        # receive the object (it will receive `None`) and thus there
134
        # is no way to know which array reference one has to remove.
135 4
        vtk_arr.AddObserver(
136
            'DeleteEvent', lambda o, e, key=key: self._remove_array(key)
137
        )
138

139
        # Cache the array
140 4
        cache[key] = np_arr
141

142 4
    def get(self, vtk_arr):
143
        """Return the cached numpy array given a VTK array."""
144 4
        key = vtk_arr.__this__
145 4
        return self._cache[key]
146

147
    ######################################################################
148
    # Non-public interface.
149
    ######################################################################
150 4
    def _remove_array(self, key):
151
        """Private function that removes the cached array.  Do not
152
        call this unless you know what you are doing."""
153 4
        try:
154 4
            del self._cache[key]
155 4
        except KeyError:
156 4
            pass
157

158

159
######################################################################
160
# Setup a global `_array_cache`.  The array object cache caches all the
161
# converted numpy arrays that are not copied.  This prevents the user
162
# from deleting or resizing the numpy array after it has been sent down
163
# to VTK.
164
######################################################################
165

166 4
_dummy = None
167
# This makes the cache work even when the module is reloaded.
168 4
for name in ['array_handler', 'tvtk.array_handler']:
169 4
    if name in sys.modules:
170 4
        mod = sys.modules[name]
171 4
        if hasattr(mod, '_array_cache'):
172 0
            _dummy = mod._array_cache
173 4
        del mod
174 4
        break
175

176 4
if _dummy:
177 0
    _array_cache = _dummy
178
else:
179 4
    _array_cache = ArrayCache()
180 4
del _dummy
181

182

183 4
def get_vtk_array_type(numeric_array_type):
184
    """Returns a VTK typecode given a numpy array."""
185
    # This is a Mapping from numpy array types to VTK array types.
186 4
    _arr_vtk = {
187
        numpy.dtype('S'): vtkConstants.VTK_UNSIGNED_CHAR,  # numpy.character
188
        numpy.dtype(numpy.uint8): vtkConstants.VTK_UNSIGNED_CHAR,
189
        numpy.dtype(numpy.uint16): vtkConstants.VTK_UNSIGNED_SHORT,
190
        numpy.dtype(numpy.int8): vtkConstants.VTK_CHAR,
191
        numpy.dtype(numpy.int16): vtkConstants.VTK_SHORT,
192
        numpy.dtype(numpy.int32): vtkConstants.VTK_INT,
193
        numpy.dtype(numpy.uint32): vtkConstants.VTK_UNSIGNED_INT,
194
        numpy.dtype(numpy.uint64): vtkConstants.VTK_UNSIGNED_LONG,
195
        numpy.dtype(numpy.float32): vtkConstants.VTK_FLOAT,
196
        numpy.dtype(numpy.float64): vtkConstants.VTK_DOUBLE,
197
        numpy.dtype(numpy.complex64): vtkConstants.VTK_FLOAT,
198
        numpy.dtype(numpy.complex128): vtkConstants.VTK_DOUBLE,
199
    }
200 4
    _extra = {
201
        numpy.dtype(ID_TYPE_CODE): vtkConstants.VTK_ID_TYPE,
202
        numpy.dtype(ULONG_TYPE_CODE): vtkConstants.VTK_UNSIGNED_LONG,
203
        numpy.dtype(LONG_TYPE_CODE): vtkConstants.VTK_LONG,
204
    }
205 4
    for t in _extra:
206 4
        if t not in _arr_vtk:
207 4
            _arr_vtk[t] = _extra[t]
208

209 4
    try:
210 4
        return _arr_vtk[numeric_array_type]
211 0
    except KeyError:
212 4
        for key in _arr_vtk:
213 4
            if numpy.issubdtype(numeric_array_type, key):
214 0
                return _arr_vtk[key]
215 0
    raise TypeError(
216
        "Couldn't translate array's type to VTK %s" % numeric_array_type
217
    )
218

219

220 4
def get_vtk_to_numeric_typemap():
221
    """Returns the VTK array type to numpy array type mapping."""
222 4
    _vtk_arr = {
223
        vtkConstants.VTK_BIT: numpy.bool_,
224
        vtkConstants.VTK_CHAR: numpy.int8,
225
        vtkConstants.VTK_UNSIGNED_CHAR: numpy.uint8,
226
        vtkConstants.VTK_SHORT: numpy.int16,
227
        vtkConstants.VTK_UNSIGNED_SHORT: numpy.uint16,
228
        vtkConstants.VTK_INT: numpy.int32,
229
        vtkConstants.VTK_UNSIGNED_INT: numpy.uint32,
230
        vtkConstants.VTK_LONG: LONG_TYPE_CODE,
231
        vtkConstants.VTK_UNSIGNED_LONG: ULONG_TYPE_CODE,
232
        vtkConstants.VTK_ID_TYPE: ID_TYPE_CODE,
233
        vtkConstants.VTK_FLOAT: numpy.float32,
234
        vtkConstants.VTK_DOUBLE: numpy.float64
235
    }
236 4
    return _vtk_arr
237

238

239 4
def get_numeric_array_type(vtk_array_type):
240
    """Returns a numpy array typecode given a VTK array type."""
241 4
    return get_vtk_to_numeric_typemap()[vtk_array_type]
242

243

244 4
def get_sizeof_vtk_array(vtk_array_type):
245
    """Returns the size of a VTK array type."""
246 0
    _size_dict = {
247
        vtkConstants.VTK_BIT: 1,
248
        vtkConstants.VTK_CHAR: 1,
249
        vtkConstants.VTK_UNSIGNED_CHAR: 1,
250
        vtkConstants.VTK_SHORT: 2,
251
        vtkConstants.VTK_UNSIGNED_SHORT: 2,
252
        vtkConstants.VTK_INT: 4,
253
        vtkConstants.VTK_UNSIGNED_INT: 4,
254
        vtkConstants.VTK_LONG: VTK_LONG_TYPE_SIZE,
255
        vtkConstants.VTK_UNSIGNED_LONG: VTK_LONG_TYPE_SIZE,
256
        vtkConstants.VTK_ID_TYPE: VTK_ID_TYPE_SIZE,
257
        vtkConstants.VTK_FLOAT: 4,
258
        vtkConstants.VTK_DOUBLE: 8
259
    }
260 0
    return _size_dict[vtk_array_type]
261

262

263 4
def create_vtk_array(vtk_arr_type):
264
    """Internal function used to create a VTK data array from another
265
    VTK array given the VTK array type.
266
    """
267 4
    tmp = vtk.vtkDataArray.CreateDataArray(vtk_arr_type)
268
    # CreateDataArray sets the refcount to 3 and this causes a severe
269
    # memory leak.
270 4
    tmp.SetReferenceCount(BASE_REFERENCE_COUNT)
271 4
    return tmp
272

273

274 4
def array2vtk(num_array, vtk_array=None):
275
    """Converts a real numpy Array (or a Python list) to a VTK array
276
    object.
277

278
    This function only works for real arrays.  Complex arrays are NOT
279
    handled.  It also works for multi-component arrays.  However, only
280
    1, and 2 dimensional arrays are supported.  This function is very
281
    efficient, so large arrays should not be a problem.
282

283
    Even in cases when no copy of the numpy array data is performed,
284
    a reference to the array is cached.  The passed array can
285
    therefore be deleted safely in all circumstances.
286

287
    Parameters
288
    ----------
289

290
    - num_array : numpy array or Python list/tuple
291

292
      The input array must be 1 or 2D.  A copy of the numeric array
293
      data passed is made in the following circumstances:
294

295
       1. A Python list/tuple was passed.
296
       2. A non-contiguous numpy array was passed.
297
       3. A `vtkBitArray` instance was passed as the second argument.
298
       4. The types of the `vtk_array` and the `num_array` are not
299
          equivalent to each other.  For example if one is an integer
300
          array and the other a float.
301

302
    - vtk_array : `vtkDataArray` (default: `None`)
303

304
      If an optional `vtkDataArray` instance, is passed as an argument
305
      then a new array is not created and returned.  The passed array
306
      is itself returned.
307

308
    """
309

310 4
    z = numpy.asarray(num_array)
311

312 4
    shape = z.shape
313 4
    assert len(shape) < 3, \
314
        "Only arrays of dimensionality 2 or lower are allowed!"
315 4
    assert not numpy.issubdtype(z.dtype, numpy.complexfloating), \
316
        "Complex numpy arrays cannot be converted to vtk arrays."\
317
        "Use real() or imag() to get a component of the array before"\
318
        " passing it to vtk."
319

320
    # First create an array of the right type by using the typecode.
321
    # Bit arrays need special casing.
322 4
    bit_array = False
323 4
    if vtk_array is None:
324 4
        vtk_typecode = get_vtk_array_type(z.dtype)
325 4
        result_array = create_vtk_array(vtk_typecode)
326 4
    elif vtk_array.GetDataType() == vtkConstants.VTK_BIT:
327 4
        vtk_typecode = vtkConstants.VTK_CHAR
328 4
        result_array = create_vtk_array(vtkConstants.VTK_CHAR)
329 4
        bit_array = True
330
    else:
331 4
        vtk_typecode = vtk_array.GetDataType()
332 4
        result_array = vtk_array
333

334
    # Find the shape and set number of components.
335 4
    if len(shape) == 1:
336 4
        result_array.SetNumberOfComponents(1)
337
    else:
338 4
        result_array.SetNumberOfComponents(shape[1])
339

340 4
    result_array.SetNumberOfTuples(shape[0])
341

342
    # Ravel the array appropriately.
343 4
    arr_dtype = get_numeric_array_type(vtk_typecode)
344 4
    if numpy.issubdtype(z.dtype, arr_dtype):
345 4
        z_flat = numpy.ravel(z)
346
    else:
347 4
        z_flat = numpy.ravel(z).astype(arr_dtype)
348

349
    # Point the VTK array to the numpy data.  The last argument (1)
350
    # tells the array not to deallocate.
351 4
    result_array.SetVoidArray(getbuffer(z_flat), len(z_flat), 1)
352

353 4
    if bit_array:
354
        # Handle bit arrays -- they have to be copied.  Note that bit
355
        # arrays are used ONLY when the user has passed one as an
356
        # argument to this function.
357 4
        vtk_array.SetNumberOfTuples(result_array.GetNumberOfTuples())
358 4
        vtk_array.SetNumberOfComponents(result_array.GetNumberOfComponents())
359 4
        for i in range(result_array.GetNumberOfComponents()):
360 4
            vtk_array.CopyComponent(i, result_array, i)
361 4
        result_array = vtk_array
362
    else:
363
        # Save a reference to the flatted array in the array cache.
364
        # This prevents the user from deleting or resizing the array
365
        # and getting into serious trouble.  This is only done for
366
        # non-bit array cases where the data is not copied.
367
        global _array_cache
368 4
        _array_cache.add(result_array, z_flat)
369

370 4
    return result_array
371

372

373 4
def vtk2array(vtk_array):
374
    """Converts a VTK data array to a numpy array.
375

376
    Given a subclass of vtkDataArray, this function returns an
377
    appropriate numpy array containing the same data.  The function
378
    is very efficient since it uses the VTK imaging pipeline to
379
    convert the data.  If a sufficiently new version of VTK (5.2) is
380
    installed then it actually uses the buffer interface to return a
381
    view of the VTK array in the returned numpy array.
382

383
    Parameters
384
    ----------
385

386
    - vtk_array : `vtkDataArray`
387

388
      The VTK data array to be converted.
389

390
    """
391 4
    typ = vtk_array.GetDataType()
392 4
    assert typ in get_vtk_to_numeric_typemap().keys(), \
393
        "Unsupported array type %s" % typ
394

395 4
    shape = (vtk_array.GetNumberOfTuples(),
396
             vtk_array.GetNumberOfComponents())
397 4
    if shape[0] == 0:
398 0
        dtype = get_numeric_array_type(typ)
399 0
        return numpy.array([], dtype)
400

401
    # First check if this array already has a numpy array cached,
402
    # if it does and the array size has not been changed, reshape
403
    # that and return it.
404 4
    if vtk_array in _array_cache:
405 4
        arr = _array_cache.get(vtk_array)
406 4
        if shape[1] == 1:
407 4
            shape = (shape[0], )
408 4
        if arr.size == numpy.prod(shape):
409 4
            arr = numpy.reshape(arr, shape)
410 4
            return arr
411

412
    # If VTK's new numpy support is available, use the buffer interface.
413 4
    if numpy_support is not None and typ != vtkConstants.VTK_BIT:
414 4
        dtype = get_numeric_array_type(typ)
415 4
        result = numpy.frombuffer(vtk_array, dtype=dtype)
416 4
        if shape[1] == 1:
417 4
            shape = (shape[0], )
418 4
        result.shape = shape
419 4
        return result
420

421
    # Setup an imaging pipeline to export the array.
422 4
    img_data = vtk.vtkImageData()
423 4
    img_data.SetDimensions(shape[0], 1, 1)
424 4
    if typ == vtkConstants.VTK_BIT:
425 4
        iarr = vtk.vtkCharArray()
426 4
        iarr.DeepCopy(vtk_array)
427 4
        img_data.GetPointData().SetScalars(iarr)
428 4
    elif typ == vtkConstants.VTK_ID_TYPE:
429
        # Needed since VTK_ID_TYPE does not work with VTK 4.5.
430 0
        iarr = vtk.vtkLongArray()
431 0
        iarr.SetNumberOfTuples(vtk_array.GetNumberOfTuples())
432 0
        nc = vtk_array.GetNumberOfComponents()
433 0
        iarr.SetNumberOfComponents(nc)
434 4
        for i in range(nc):
435 0
            iarr.CopyComponent(i, vtk_array, i)
436 0
        img_data.GetPointData().SetScalars(iarr)
437
    else:
438 0
        img_data.GetPointData().SetScalars(vtk_array)
439

440 4
    if is_old_pipeline():
441 0
        img_data.SetNumberOfScalarComponents(shape[1])
442 4
        if typ == vtkConstants.VTK_ID_TYPE:
443
            # Hack necessary because vtkImageData can't handle VTK_ID_TYPE.
444 0
            img_data.SetScalarType(vtkConstants.VTK_LONG)
445 0
            r_dtype = get_numeric_array_type(vtkConstants.VTK_LONG)
446 4
        elif typ == vtkConstants.VTK_BIT:
447 0
            img_data.SetScalarType(vtkConstants.VTK_CHAR)
448 0
            r_dtype = get_numeric_array_type(vtkConstants.VTK_CHAR)
449
        else:
450 0
            img_data.SetScalarType(typ)
451 0
            r_dtype = get_numeric_array_type(typ)
452 0
        img_data.Update()
453
    else:
454 4
        if typ == vtkConstants.VTK_ID_TYPE:
455 0
            r_dtype = get_numeric_array_type(vtkConstants.VTK_LONG)
456 4
        elif typ == vtkConstants.VTK_BIT:
457 4
            r_dtype = get_numeric_array_type(vtkConstants.VTK_CHAR)
458
        else:
459 0
            r_dtype = get_numeric_array_type(typ)
460 4
        img_data.Modified()
461

462 4
    exp = vtk.vtkImageExport()
463 4
    if is_old_pipeline():
464 0
        exp.SetInput(img_data)
465
    else:
466 4
        exp.SetInputData(img_data)
467

468
    # Create an array of the right size and export the image into it.
469 4
    im_arr = numpy.empty((shape[0]*shape[1],), r_dtype)
470 4
    exp.Export(im_arr)
471

472
    # Now reshape it.
473 4
    if shape[1] == 1:
474 4
        shape = (shape[0], )
475 4
    im_arr = numpy.reshape(im_arr, shape)
476 4
    return im_arr
477

478

479 4
def array2vtkCellArray(num_array, vtk_array=None):
480
    """Given a nested Python list or a numpy array, this method
481
    creates a vtkCellArray instance and returns it.
482

483
    A variety of input arguments are supported as described in the
484
    Parameter documentation.  If numpy arrays are given, this method
485
    is highly efficient.  This function is most efficient if the
486
    passed numpy arrays have a typecode `ID_TYPE_CODE`.  Otherwise a
487
    typecast is necessary and this involves an extra copy.  This
488
    method *always copies* the input data.
489

490
    An alternative and more efficient way to build the connectivity
491
    list is to create a vtkIdTypeArray having data of the form
492
    (npts,p0,p1,...p(npts-1), repeated for each cell) and then call
493
    <vtkCellArray_instance>.SetCells(n_cell, id_list).
494

495
    Parameters
496
    ----------
497

498
    - num_array : numpy array or Python list/tuple
499

500
      Valid values are:
501

502
        1. A Python list of 1D lists.  Each 1D list can contain one
503
           cell connectivity list.  This is very slow and is to be
504
           used only when efficiency is of no consequence.
505

506
        2. A 2D numpy array with the cell connectivity list.
507

508
        3. A Python list of 2D numpy arrays.  Each numeric array can
509
           have a different shape.  This makes it easy to generate a
510
           cell array having cells of different kinds.
511

512
    - vtk_array : `vtkCellArray` (default: `None`)
513

514
      If an optional `vtkCellArray` instance, is passed as an argument
515
      then a new array is not created and returned.  The passed array
516
      is itself modified and returned.
517

518
    Example
519
    -------
520

521
       >>> a = [[0], [1, 2], [3, 4, 5], [6, 7, 8, 9]]
522
       >>> cells = array_handler.array2vtkCellArray(a)
523
       >>> a = numpy.array([[0,1,2], [3,4,5], [6,7,8]], 'l')
524
       >>> cells = array_handler.array2vtkCellArray(a)
525
       >>> l_a = [a[:,:1], a[:2,:2], a]
526
       >>> cells = array_handler.array2vtkCellArray(l_a)
527

528
    """
529 4
    if vtk_array:
530 4
        cells = vtk_array
531
    else:
532 4
        cells = vtk.vtkCellArray()
533 4
    assert cells.GetClassName() == 'vtkCellArray', \
534
        'Second argument must be a `vtkCellArray` instance.'
535

536 4
    if len(num_array) == 0:
537 4
        return cells
538

539
    ########################################
540
    # Internal functions.
541 4
    def _slow_array2cells(z, cells):
542 4
        cells.Reset()
543 4
        vtk_ids = vtk.vtkIdList()
544 4
        for i in z:
545 4
            vtk_ids.Reset()
546 4
            for j in i:
547 4
                vtk_ids.InsertNextId(j)
548 4
            cells.InsertNextCell(vtk_ids)
549

550 4
    def _get_tmp_array(arr):
551 4
        try:
552 4
            tmp_arr = numpy.asarray(arr, ID_TYPE_CODE)
553 0
        except TypeError:
554 0
            tmp_arr = arr.astype(ID_TYPE_CODE)
555 4
        return tmp_arr
556

557 4
    def _set_cells(cells, n_cells, id_typ_arr):
558 4
        vtk_arr = vtk.vtkIdTypeArray()
559 4
        array2vtk(id_typ_arr, vtk_arr)
560 4
        cells.SetCells(n_cells, vtk_arr)
561
    ########################################
562

563 4
    msg = "Invalid argument.  Valid types are a Python list of lists,"\
564
          " a Python list of numpy arrays, or a numpy array."
565

566 4
    if issubclass(type(num_array), (list, tuple)):
567 4
        assert len(num_array[0]) > 0, "Input array must be 2D."
568 4
        tp = type(num_array[0])
569 4
        if issubclass(tp, list):  # Pure Python list.
570 4
            _slow_array2cells(num_array, cells)
571 4
            return cells
572 4
        elif issubclass(tp, numpy.ndarray):  # List of arrays.
573
            # Check shape of array and find total size.
574 4
            tot_size = 0
575 4
            n_cells = 0
576 4
            for arr in num_array:
577 4
                assert len(arr.shape) == 2, "Each array must be 2D"
578 4
                shp = arr.shape
579 4
                tot_size += shp[0]*(shp[1] + 1)
580 4
                n_cells += shp[0]
581
            # Create an empty array.
582 4
            id_typ_arr = numpy.empty((tot_size,), ID_TYPE_CODE)
583
            # Now populate it with the ids.
584 4
            count = 0
585 4
            for arr in num_array:
586 4
                tmp_arr = _get_tmp_array(arr)
587 4
                shp = arr.shape
588 4
                sz = shp[0]*(shp[1] + 1)
589 4
                set_id_type_array(tmp_arr, id_typ_arr[count:count+sz])
590 4
                count += sz
591
            # Now set them cells.
592 4
            _set_cells(cells, n_cells, id_typ_arr)
593 4
            return cells
594
        else:
595 0
            raise TypeError(msg)
596 4
    elif issubclass(type(num_array), numpy.ndarray):
597 4
        assert len(num_array.shape) == 2, "Input array must be 2D."
598 4
        tmp_arr = _get_tmp_array(num_array)
599 4
        shp = tmp_arr.shape
600 4
        id_typ_arr = numpy.empty((shp[0]*(shp[1] + 1),), ID_TYPE_CODE)
601 4
        set_id_type_array(tmp_arr, id_typ_arr)
602 4
        _set_cells(cells, shp[0], id_typ_arr)
603 4
        return cells
604
    else:
605 0
        raise TypeError(msg)
606

607

608 4
def array2vtkPoints(num_array, vtk_points=None):
609
    """Converts a numpy array/Python list to a vtkPoints object.
610

611
    Unless a Python list/tuple or a non-contiguous array is given, no
612
    copy of the data is made.  Thus the function is very efficient.
613

614
    Parameters
615
    ----------
616

617
    - num_array : numpy array or Python list/tuple
618

619
      The input array must be 2D with `shape[1] == 3`.
620

621
    - vtk_points : `vtkPoints` (default: `None`)
622

623
      If an optional `vtkPoints` instance, is passed as an argument
624
      then a new array is not created and returned.  The passed array
625
      is itself modified and returned.
626

627
    """
628 4
    if vtk_points:
629 4
        points = vtk_points
630
    else:
631 4
        points = vtk.vtkPoints()
632

633 4
    arr = numpy.asarray(num_array)
634 4
    assert len(arr.shape) == 2, "Points array must be 2 dimensional."
635 4
    assert arr.shape[1] == 3, "Incorrect shape: shape[1] must be 3."
636 4
    vtk_array = array2vtk(arr)
637 4
    points.SetData(vtk_array)
638 4
    return points
639

640

641 4
def array2vtkIdList(num_array, vtk_idlist=None):
642
    """Converts a numpy array/Python list to a vtkIdList object.
643

644
    Parameters
645
    ----------
646

647
    - num_array : numpy array or Python list/tuple
648

649
      The input array must be 2D with `shape[1] == 3`.
650

651
    - vtk_idlist : `vtkIdList` (default: `None`)
652

653
      If an optional `vtkIdList` instance, is passed as an argument
654
      then a new array is not created and returned.  The passed array
655
      is itself modified and returned.
656

657
    """
658 4
    if vtk_idlist:
659 4
        ids = vtk_idlist
660
    else:
661 4
        ids = vtk.vtkIdList()
662

663 4
    arr = numpy.asarray(num_array)
664 4
    assert len(arr.shape) == 1, "Array for vtkIdList must be 1D"
665 4
    ids.SetNumberOfIds(len(arr))
666 4
    for i, j in enumerate(arr):
667 4
        ids.SetId(i, j)
668 4
    return ids
669

670

671
######################################################################
672
# Array argument handling functions.
673
######################################################################
674

675 4
def is_array(arr):
676
    """Returns True if the passed `arr` is a numpy array or a List."""
677 4
    if issubclass(type(arr), (numpy.ndarray, list)):
678 4
        return True
679 4
    return False
680

681

682 4
def convert_array(arr, vtk_typ=None):
683
    """Convert the given array to the optional type specified by
684
    `vtk_typ`.
685

686
    Parameters
687
    ----------
688

689
    - arr : numpy array/list.
690
    - vtk_typ : `string` or `None`
691
      represents the type the array is to be converted to.
692

693
    """
694 4
    if vtk_typ:
695 4
        conv = {'vtkCellArray': array2vtkCellArray,
696
                'vtkPoints': array2vtkPoints,
697
                'vtkIdList': array2vtkIdList}
698 4
        if vtk_typ in conv.keys():
699 4
            vtk_arr = getattr(vtk, vtk_typ)()
700 4
            return conv[vtk_typ](arr, vtk_arr)
701 4
        elif vtk_typ.find('Array') > -1:
702 4
            try:
703 4
                vtk_arr = getattr(vtk, vtk_typ)()
704 4
            except TypeError:  # vtk_typ == 'vtkDataArray'
705 4
                return array2vtk(arr)
706
            else:
707 4
                return array2vtk(arr, vtk_arr)
708
        else:
709 0
            return arr
710
    else:
711 0
        return array2vtk(arr)
712

713

714 4
def is_array_sig(s):
715
    """Given a signature, return if the signature has an array."""
716 4
    if not isinstance(s, (unicode, str)):
717 4
        return False
718 4
    arr_types = ['Array', 'vtkPoints', 'vtkIdList']
719 4
    for i in arr_types:
720 4
        if s.find(i) > -1:
721 4
            return True
722 4
    return False
723

724

725 4
def is_array_or_vtkarray(arg):
726
    """Returns True if the argument is an array/Python list or if it
727
    is a vtk array."""
728

729 4
    if is_array(arg):
730 4
        return True
731
    else:
732 4
        if hasattr(arg, '_vtk_obj'):
733 4
            if is_array_sig(arg._vtk_obj.__class__.__name__):
734 4
                return True
735 4
    return False
736

737

738 4
def get_correct_sig(args, sigs):
739
    """Given a list of args and a collection of possible signatures,
740
    this function returns the most appropriate signature.  This
741
    function is only called by deref_array.  This implies that one of
742
    the signatures has an array type.
743

744
    """
745
    # First do the trivial cases.
746 4
    if sigs is None:
747 4
        return None
748 4
    if len(sigs) == 1:
749 4
        return sigs[0]
750
    else:
751
        # Non-trivial cases.
752 4
        la = len(args)
753 4
        candidate_sigs = [s for s in sigs if len(s) == la]
754 4
        count = len(candidate_sigs)
755 4
        if count == 0:
756
            # No sig has the right number of args.
757 4
            msg = "Insufficient number of arguments to method."\
758
                  "Valid arguments are:\n%s" % sigs
759 4
            raise TypeError(msg)
760 4
        elif count == 1:
761
            # If only one of the sigs has the right number of args,
762
            # return it.
763 4
            return candidate_sigs[0]
764
        else:
765
            # More than one sig has the same number of args.
766
            # Check if args need conversion at all.
767 4
            array_idx = [i for i, a in enumerate(args)
768
                         if is_array_or_vtkarray(a)]
769 4
            n_arr = len(array_idx)
770 4
            if n_arr == 0:
771
                # No conversion necessary so signature info is
772
                # useless.
773 4
                return None
774
            else:
775
                # Need to find the right sig.  This is done by finding
776
                # the first signature that matches all the arrays in
777
                # the argument.
778 4
                for sig in candidate_sigs:
779 4
                    array_in_sig = [is_array_sig(s) for s in sig]
780 4
                    if array_in_sig.count(True) != len(array_idx):
781 4
                        continue
782 4
                    bad = False
783 4
                    for i in array_idx:
784 4
                        if not array_in_sig[i]:
785 0
                            bad = True
786 4
                    if not bad:
787 4
                        return sig
788
                # Could not find any valid signature, so give up.
789 0
                return None
790

791

792 4
def deref_vtk(obj):
793
    """Dereferences the VTK object from the object if possible.  This
794
    is duplicated from `tvtk_base.py` because I'd like to keep this
795
    module independent of `tvtk_base.py`.
796
    """
797 4
    if hasattr(obj, '_vtk_obj'):
798 4
        return obj._vtk_obj
799
    else:
800 4
        return obj
801

802

803 4
def deref_array(args, sigs=None):
804
    """Given a bunch of arguments and optional signature information,
805
    this converts the arguments suitably.  If the argument is either a
806
    Python list or a numpy array it is converted to a suitable type
807
    based on the signature information.  If it is not an array, but a
808
    TVTK object the VTK object is dereferenced.  Otherwise nothing is
809
    done.  If no signature information is provided the arrays are
810
    automatically converted (this can sometimes go wrong).  The
811
    signature information is provided in the form of a list of lists.
812

813
    """
814 4
    ret = []
815 4
    sig = get_correct_sig(args, sigs)
816 4
    if sig:
817 4
        for a, s in zip(args, sig):
818 4
            if is_array(a) and is_array_sig(s):
819 4
                ret.append(convert_array(a, s))
820
            else:
821 4
                ret.append(deref_vtk(a))
822
    else:
823 4
        for a in args:
824 4
            if is_array(a):
825 0
                ret.append(convert_array(a))
826
            else:
827 4
                ret.append(deref_vtk(a))
828 4
    return ret

Read our documentation on viewing source code .

Loading