1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135
|
"""Contains pyvista_ndarray a numpy ndarray type used in pyvista."""
from __future__ import annotations
from collections.abc import Iterable
from typing import TYPE_CHECKING
from typing import cast
import numpy as np
from . import _vtk_core as _vtk
from .utilities.arrays import FieldAssociation
from .utilities.arrays import convert_array
from .utilities.misc import _NoNewAttrMixin
if TYPE_CHECKING:
from typing import Any
import numpy.typing as npt
from pyvista import DataSet
from ._typing_core import ArrayLike
from ._typing_core import NumpyArray
class pyvista_ndarray(_NoNewAttrMixin, np.ndarray): # numpydoc ignore=PR02 # noqa: N801
"""A ndarray which references the owning dataset and the underlying vtk array.
This array can be acted upon just like a :class:`numpy.ndarray`.
Parameters
----------
array : ArrayLike or :vtk:`vtkAbstractArray`
Array like.
dataset : DataSet
Input dataset.
association : pyvista.core.utilities.arrays.FieldAssociation
Field association.
Examples
--------
Return the points of a Sphere as a :class:`pyvista.pyvista_ndarray`.
>>> import pyvista as pv
>>> mesh = pv.Sphere()
>>> mesh.points # doctest:+SKIP
pyvista_ndarray([[-5.5511151e-17, 0.0000000e+00, -5.0000000e-01],
[ 5.5511151e-17, 0.0000000e+00, 5.0000000e-01],
[-5.4059509e-02, 0.0000000e+00, -4.9706897e-01],
...,
[-1.5616201e-01, -3.3193260e-02, 4.7382659e-01],
[-1.0513641e-01, -2.2347433e-02, 4.8831028e-01],
[-5.2878179e-02, -1.1239604e-02, 4.9706897e-01]],
dtype=float32)
"""
def __new__( # noqa: PYI034
cls: type[pyvista_ndarray],
array: ArrayLike[float] | _vtk.vtkAbstractArray,
dataset: DataSet | _vtk.vtkDataSet | _vtk.VTKObjectWrapper | None = None,
association: FieldAssociation = FieldAssociation.NONE,
) -> pyvista_ndarray:
"""Allocate the array."""
if isinstance(array, _vtk.vtkAbstractArray):
obj = convert_array(array).view(cls)
obj.VTKObject = array
elif isinstance(array, Iterable):
obj = np.asarray(array).view(cls)
else:
msg = ( # type: ignore[unreachable]
f'pyvista_ndarray got an invalid type {type(array)}. '
'Expected an Iterable or vtk.vtkAbstractArray'
)
raise TypeError(msg)
obj.association = association
obj.dataset = _vtk.vtkWeakReference()
if isinstance(dataset, _vtk.VTKObjectWrapper):
obj.dataset.Set(dataset.VTKObject)
else:
obj.dataset.Set(cast('_vtk.vtkDataSet', dataset))
return obj
def __array_finalize__(self: pyvista_ndarray, obj: npt.NDArray[Any] | None) -> None:
"""Finalize array (associate with parent metadata)."""
# this is necessary to ensure that views/slices of pyvista_ndarray
# objects stay associated with those of their parents.
#
# the VTKArray class uses attributes called `DataSet` and `Association`
# to hold this data. I don't know why this class doesn't use the same
# convention, but here we just map those over to the appropriate
# attributes of this class
_vtk.VTKArray.__array_finalize__(self, obj) # type: ignore[arg-type]
if np.shares_memory(self, obj):
self.dataset = getattr(obj, 'dataset', None)
self.association = getattr(obj, 'association', FieldAssociation.NONE)
self.VTKObject = getattr(obj, 'VTKObject', None)
else:
self.dataset = None
self.association = FieldAssociation.NONE
self.VTKObject = None
def __setitem__(self: pyvista_ndarray, key: int | NumpyArray[int], value: Any) -> None: # type: ignore[override]
"""Implement [] set operator.
When the array is changed it triggers "Modified()" which updates
all upstream objects, including any render windows holding the
object.
"""
super().__setitem__(key, value)
if self.VTKObject is not None:
self.VTKObject.Modified()
# the associated dataset should also be marked as modified
dataset = self.dataset
if dataset is not None and dataset.Get():
dataset.Get().Modified()
def __array_wrap__(self: pyvista_ndarray, out_arr, context=None, return_scalar: bool = False): # noqa: ANN001, ANN204, FBT001, FBT002
"""Return a numpy scalar if array is 0d.
See https://github.com/numpy/numpy/issues/5819
"""
if out_arr.ndim:
return super().__array_wrap__(out_arr, context, return_scalar)
# Match numpy's behavior and return a numpy dtype scalar
return out_arr[()]
__getattr__ = _vtk.VTKObjectWrapperCheckSnakeCase.__getattr__
|