Files
geMoldInsight/venv/lib/python3.11/site-packages/vtkmodules/util/data_model.py
T

862 lines
26 KiB
Python
Raw Normal View History

2026-02-12 23:22:11 +08:00
"""This module provides classes that allow numpy style access
to VTK datasets. See examples at bottom.
"""
from contextlib import suppress
from vtkmodules.vtkCommonCore import vtkPoints, vtkAbstractArray, vtkDataArray
from vtkmodules.vtkCommonDataModel import (
vtkCellArray,
vtkDataObject,
vtkFieldData,
vtkDataSetAttributes,
vtkPointData,
vtkCellData,
vtkDataObject,
vtkImageData,
vtkMultiBlockDataSet,
vtkPolyData,
vtkStructuredGrid,
vtkRectilinearGrid,
vtkUnstructuredGrid,
vtkOverlappingAMR,
vtkPartitionedDataSet,
vtkPartitionedDataSetCollection,
)
import weakref
NUMPY_AVAILABLE = False
with suppress(ImportError):
import numpy
from vtkmodules.numpy_interface import dataset_adapter as dsa
NUMPY_AVAILABLE = True
class FieldDataBase(object):
def __init__(self):
self.association = None
self.dataset = None
def __getitem__(self, idx):
"""Implements the [] operator. Accepts an array name or index."""
return self.get_array(idx)
def __setitem__(self, name, value):
"""Implements the [] operator. Accepts an array name or index."""
return self.set_array(name, value)
def get_array(self, idx):
"Given an index or name, returns a VTKArray."
if isinstance(idx, int) and idx >= self.GetNumberOfArrays():
raise IndexError("array index out of range")
vtkarray = super().GetArray(idx)
if not NUMPY_AVAILABLE:
return vtkarray if vtkarray else self.GetAbstractArray(idx)
if not vtkarray:
vtkarray = self.GetAbstractArray(idx)
if vtkarray:
return vtkarray
return dsa.NoneArray
array = dsa.vtkDataArrayToVTKArray(vtkarray, self.dataset)
array.Association = self.association
return array
def __contains__(self, aname):
"""Returns true if the container contains arrays
with the given name, false otherwise"""
return self.HasArray(aname)
def keys(self):
"""Returns the names of the arrays as a list."""
kys = []
narrays = self.GetNumberOfArrays()
for i in range(narrays):
name = self.GetAbstractArray(i).GetName()
if name:
kys.append(name)
return tuple(kys)
def values(self):
"""Returns the arrays as a tuple."""
vals = []
narrays = self.GetNumberOfArrays()
for i in range(narrays):
a = self.get_array(i)
if a.GetName():
vals.append(a)
return tuple(vals)
def items(self):
"""Returns a tuple of pairs (name, array)"""
pairs = []
narrays = self.GetNumberOfArrays()
for i in range(narrays):
arr = self.get_array(i)
name = arr.GetName()
if name:
pairs.append((name, arr))
return tuple(pairs)
def set_array(self, name, narray):
"""Appends a new array to the dataset attributes."""
if not NUMPY_AVAILABLE:
if isinstance(narray, vtkAbstractArray):
narray.SetName(name)
self.AddArray(narray)
return
if narray is dsa.NoneArray:
# if NoneArray, nothing to do.
return
if self.association == vtkDataObject.POINT:
arrLength = self.dataset.GetNumberOfPoints()
elif self.association == vtkDataObject.CELL:
arrLength = self.dataset.GetNumberOfCells()
elif (
self.association == vtkDataObject.ROW
and self.dataset.GetNumberOfColumns() > 0
):
arrLength = self.dataset.GetNumberOfRows()
else:
if not isinstance(narray, numpy.ndarray):
arrLength = 1
else:
arrLength = narray.shape[0]
# Fixup input array length:
if (
not isinstance(narray, numpy.ndarray) or numpy.ndim(narray) == 0
): # Scalar input
dtype = narray.dtype if isinstance(narray, numpy.ndarray) else type(narray)
tmparray = numpy.empty(arrLength, dtype=dtype)
tmparray.fill(narray)
narray = tmparray
elif narray.shape[0] != arrLength: # Vector input
components = 1
for l in narray.shape:
components *= l
tmparray = numpy.empty((arrLength, components), dtype=narray.dtype)
tmparray[:] = narray.flatten()
narray = tmparray
shape = narray.shape
if len(shape) == 3:
# Array of matrices. We need to make sure the order in memory is right.
# If column order (c order), transpose. VTK wants row order (fortran
# order). The deep copy later will make sure that the array is contiguous.
# If row order but not contiguous, transpose so that the deep copy below
# does not happen.
size = narray.dtype.itemsize
if (narray.strides[1] / size == 3 and narray.strides[2] / size == 1) or (
narray.strides[1] / size == 1
and narray.strides[2] / size == 3
and not narray.flags.contiguous
):
narray = narray.transpose(0, 2, 1)
# If array is not contiguous, make a deep copy that is contiguous
if not narray.flags.contiguous:
narray = numpy.ascontiguousarray(narray)
# Flatten array of matrices to array of vectors
if len(shape) == 3:
narray = narray.reshape(shape[0], shape[1] * shape[2])
# this handle the case when an input array is directly appended on the
# output. We want to make sure that the array added to the output is not
# referring to the input dataset.
copy = dsa.VTKArray(narray)
try:
copy.VTKObject = narray.VTKObject
except AttributeError:
pass
arr = dsa.numpyTovtkDataArray(copy, name)
self.AddArray(arr)
def __eq__(self, other: object) -> bool:
"""Test dict-like equivalency."""
# here we check if other is the same class or a subclass of self.
if not isinstance(other, type(self)):
return False
if self is other:
return True
"""
If numpy is not available, only check for identity without comparing contents of the data arrays
"""
if not NUMPY_AVAILABLE:
return False
if set(self.keys()) != set(other.keys()):
return False
# verify the value of the arrays
for key, value in other.items():
if not numpy.array_equal(value, self[key]):
return False
return True
def __iter__(self):
return iter(self.keys())
def __len__(self):
return self.GetNumberOfArrays()
@vtkFieldData.override
class FieldData(FieldDataBase, vtkFieldData):
pass
class DataSetAttributesBase(FieldDataBase):
pass
@vtkDataSetAttributes.override
class DataSetAttributes(DataSetAttributesBase, vtkDataSetAttributes):
def __eq__(self, other: object) -> bool:
"""Test dict-like equivalency."""
if not super().__eq__(other):
return False
for attr in [
"GetScalars",
"GetVectors",
"GetNormals",
"GetTangents",
"GetTCoords",
"GetTensors",
"GetGlobalIds",
"GetPedigreeIds",
"GetRationalWeights",
"GetHigherOrderDegrees",
"GetProcessIds",
]:
self_attr = getattr(self, attr)()
other_attr = getattr(other, attr)()
if self_attr and other_attr:
if self_attr.GetName() != other_attr.GetName():
return False
elif self_attr != other_attr:
return False
return True
@vtkPointData.override
class PointData(DataSetAttributesBase, vtkPointData):
pass
@vtkCellData.override
class CellData(DataSetAttributesBase, vtkCellData):
pass
class CompositeDataSetAttributesIterator(object):
def __init__(self, cdsa):
self._cdsa = cdsa
if cdsa:
self._itr = iter(cdsa.keys())
else:
self._itr = None
def __iter__(self):
return self
def __next__(self):
if not self._cdsa:
raise StopIteration
name = next(self._itr)
return self._cdsa[name]
def next(self):
return self.__next__()
class CompositeDataSetAttributes(object):
"""This is a python friendly wrapper for vtkDataSetAttributes for composite
datasets. Since composite datasets themselves don't have attribute data, but
the attribute data is associated with the leaf nodes in the composite
dataset, this class simulates a DataSetAttributes interface by taking a
union of DataSetAttributes associated with all leaf nodes."""
def __init__(self, dataset, association):
self.DataSet = dataset
self.Association = association
self.ArrayNames = []
self.Arrays = {}
# build the set of arrays available in the composite dataset. Since
# composite datasets can have partial arrays, we need to iterate over
# all non-null blocks in the dataset.
self.__determine_arraynames()
def __determine_arraynames(self):
array_set = set()
array_list = []
for dataset in self.DataSet:
dsa = dataset.GetAttributesAsFieldData(self.Association)
for array_name in dsa.keys():
if array_name not in array_set:
array_set.add(array_name)
array_list.append(array_name)
self.ArrayNames = array_list
def modified(self):
"""Rescans the contained dataset to update the
internal list of arrays."""
self.__determine_arraynames()
def __contains__(self, aname):
"""Returns true if the container contains arrays
with the given name, false otherwise"""
return aname in self.ArrayNames
def keys(self):
"""Returns the names of the arrays as a tuple."""
return tuple(self.ArrayNames)
def values(self):
"""Returns all the arrays as a tuple."""
arrays = []
for array in self:
arrays.append(array)
return tuple(arrays)
def items(self):
"""Returns (name, array) pairs as a tuple."""
items = []
for name in self.keys():
items.append((name, self[name]))
return tuple(items)
def __getitem__(self, idx):
"""Implements the [] operator. Accepts an array name."""
return self.get_array(idx)
def __setitem__(self, name, narray):
"""Implements the [] operator. Accepts an array name."""
return self.set_array(name, narray)
def set_array(self, name, narray):
"""Appends a new array to the composite dataset attributes."""
if not NUMPY_AVAILABLE:
# don't know how to handle composite dataset attribute when numpy not around
raise NotImplementedError("Only available with numpy")
if narray is dsa.NoneArray:
# if NoneArray, nothing to do.
return
added = False
if not isinstance(narray, dsa.VTKCompositeDataArray): # Scalar input
for ds in self.DataSet:
ds.GetAttributesAsFieldData(self.Association).set_array(name, narray)
added = True
if added:
self.ArrayNames.append(name)
# don't add the narray since it's a scalar. GetArray() will create a
# VTKCompositeArray on-demand.
else:
for ds, array in zip(self.DataSet, narray.Arrays):
if array is not None:
ds.GetAttributesAsFieldData(self.Association).set_array(name, array)
added = True
if added:
self.ArrayNames.append(name)
self.Arrays[name] = weakref.ref(narray)
def get_array(self, idx):
"""Given a name, returns a VTKCompositeArray."""
arrayname = idx
if not NUMPY_AVAILABLE:
# don't know how to handle composite dataset attribute when numpy not around
raise NotImplementedError("Only available with numpy")
if arrayname not in self.ArrayNames:
return dsa.NoneArray
if arrayname not in self.Arrays or self.Arrays[arrayname]() is None:
array = dsa.VTKCompositeDataArray(
dataset=self.DataSet, name=arrayname, association=self.Association
)
self.Arrays[arrayname] = weakref.ref(array)
else:
array = self.Arrays[arrayname]()
return array
def __iter__(self):
"""Iterators on keys"""
return iter(self.ArrayNames)
def __len__(self):
return len(self.ArrayNames)
# class DataSet(DataObjectBase):
class DataSet(object):
def __init__(self, **kwargs) -> None:
self._numpy_attrs = []
@property
def point_data(self):
pd = super().GetPointData()
pd.dataset = self
pd.association = self.POINT
return pd
@property
def cell_data(self):
cd = super().GetCellData()
cd.dataset = self
cd.association = self.CELL
return cd
@property
def field_data(self):
fd = super().GetFieldData()
if fd:
fd.dataset = self
fd.association = self.FIELD
return fd
def __eq__(self, other: object) -> bool:
"""Test equivalency between data objects."""
if not isinstance(self, type(other)):
return False
if self is other:
return True
"""
If numpy is not available, only check for identity without comparing contents of the data arrays
"""
if not NUMPY_AVAILABLE:
return False
for attr in self._numpy_attrs:
if hasattr(self, attr):
if not numpy.array_equal(getattr(self, attr), getattr(other, attr)):
return False
for attr in ["field_data", "point_data", "cell_data"]:
if getattr(self, attr) != getattr(other, attr):
return False
return True
def convert_to_unstructured_grid(self):
from vtkmodules.vtkFiltersCore import vtkExtractCells
ecells = vtkExtractCells()
ecells.SetInputData(self)
ecells.ExtractAllCellsOn()
ecells.Update()
return ecells.GetOutput()
class PointSet(DataSet):
def __init__(self, **kwargs) -> None:
DataSet.__init__(self, **kwargs)
self._numpy_attrs.append("points")
@property
def points(self):
pts = self.GetPoints()
if not NUMPY_AVAILABLE:
return pts
if not pts or not pts.GetData():
return None
return dsa.vtkDataArrayToVTKArray(pts.GetData())
@points.setter
def points(self, points):
if isinstance(points, vtkPoints):
self.SetPoints(points)
return
if not NUMPY_AVAILABLE:
raise ValueError("Expect vtkPoints")
pts = dsa.numpyTovtkDataArray(points, "points")
vtkpts = vtkPoints()
vtkpts.SetData(pts)
self.SetPoints(vtkpts)
@vtkUnstructuredGrid.override
class UnstructuredGrid(PointSet, vtkUnstructuredGrid):
def __init__(self, **kwargs):
PointSet.__init__(self, **kwargs)
vtkUnstructuredGrid.__init__(self, **kwargs)
@property
def cells(self):
ca = self.GetCells()
conn_vtk = ca.GetConnectivityArray()
offsets_vtk = ca.GetOffsetsArray()
ct_vtk = self.GetCellTypesArray()
if not NUMPY_AVAILABLE:
return {
"connectivity": conn_vtk,
"offsets": offsets_vtk,
"cell_types": ct_vtk,
}
conn = dsa.vtkDataArrayToVTKArray(conn_vtk)
offsets = dsa.vtkDataArrayToVTKArray(offsets_vtk)
ct = dsa.vtkDataArrayToVTKArray(ct_vtk)
return {"connectivity": conn, "offsets": offsets, "cell_types": ct}
@cells.setter
def cells(self, cells):
ca = vtkCellArray()
if not NUMPY_AVAILABLE:
ca.SetData(cells["offsets"], cells["connectivity"])
self.SetCells(cells["cell_types"], ca)
return
conn_vtk = dsa.numpyTovtkDataArray(cells["connectivity"])
offsets_vtk = dsa.numpyTovtkDataArray(cells["offsets"])
cell_types_vtk = dsa.numpyTovtkDataArray(cells["cell_types"])
ca.SetData(offsets_vtk, conn_vtk)
self.SetCells(cell_types_vtk, ca)
@vtkImageData.override
class ImageData(DataSet, vtkImageData):
def __init__(self, **kwargs):
DataSet.__init__(self, **kwargs)
vtkImageData.__init__(self, **kwargs)
@vtkPolyData.override
class PolyData(PointSet, vtkPolyData):
def __init__(self, **kwargs) -> None:
PointSet.__init__(self, **kwargs)
vtkPolyData.__init__(self, **kwargs)
self._numpy_attrs.extend(["verts", "lines", "strips", "polys"])
@property
def verts_arrays(self):
ca = self.GetVerts()
conn_vtk = ca.GetConnectivityArray()
offsets_vtk = ca.GetOffsetsArray()
if not NUMPY_AVAILABLE:
return {
"connectivity": conn_vtk,
"offsets": offsets_vtk,
}
conn = dsa.vtkDataArrayToVTKArray(conn_vtk)
offsets = dsa.vtkDataArrayToVTKArray(offsets_vtk)
return {"connectivity": conn, "offsets": offsets}
@property
def lines_arrays(self):
ca = self.GetLines()
conn_vtk = ca.GetConnectivityArray()
offsets_vtk = ca.GetOffsetsArray()
if not NUMPY_AVAILABLE:
return {
"connectivity": conn_vtk,
"offsets": offsets_vtk,
}
conn = dsa.vtkDataArrayToVTKArray(conn_vtk)
offsets = dsa.vtkDataArrayToVTKArray(offsets_vtk)
return {"connectivity": conn, "offsets": offsets}
@property
def strips_arrays(self):
ca = self.GetStrips()
conn_vtk = ca.GetConnectivityArray()
offsets_vtk = ca.GetOffsetsArray()
if not NUMPY_AVAILABLE:
return {
"connectivity": conn_vtk,
"offsets": offsets_vtk,
}
conn = dsa.vtkDataArrayToVTKArray(conn_vtk)
offsets = dsa.vtkDataArrayToVTKArray(offsets_vtk)
return {"connectivity": conn, "offsets": offsets}
@property
def polys_arrays(self):
ca = self.GetPolys()
conn_vtk = ca.GetConnectivityArray()
offsets_vtk = ca.GetOffsetsArray()
if not NUMPY_AVAILABLE:
return {
"connectivity": conn_vtk,
"offsets": offsets_vtk,
}
conn = dsa.vtkDataArrayToVTKArray(conn_vtk)
offsets = dsa.vtkDataArrayToVTKArray(offsets_vtk)
return {"connectivity": conn, "offsets": offsets}
@vtkRectilinearGrid.override
class RectilinearGrid(DataSet, vtkRectilinearGrid):
def __init__(self, **kwargs) -> None:
DataSet.__init__(self, **kwargs)
vtkRectilinearGrid.__init__(self, **kwargs)
self._numpy_attrs.extend(["x_coordinates", "y_coordinates", "z_coordinates"])
@property
def x_coordinates(self):
pts = self.GetXCoordinates()
if not NUMPY_AVAILABLE:
return pts
if not pts:
return None
return dsa.vtkDataArrayToVTKArray(pts)
@x_coordinates.setter
def x_coordinates(self, points):
if isinstance(points, vtkDataArray):
self.SetXCoordinates(points)
return
if not NUMPY_AVAILABLE:
raise ValueError("Expect vtkDataArray")
pts = dsa.numpyTovtkDataArray(points, "x_coords")
self.SetXCoordinates(pts)
@property
def y_coordinates(self):
pts = self.GetYCoordinates()
if not NUMPY_AVAILABLE:
return pts
if not pts:
return None
return dsa.vtkDataArrayToVTKArray(pts)
@y_coordinates.setter
def y_coordinates(self, points):
if isinstance(points, vtkDataArray):
self.SetYCoordinates(points)
return
if not NUMPY_AVAILABLE:
raise ValueError("Expect vtkDataArray")
pts = dsa.numpyTovtkDataArray(points, "y_coords")
self.SetYCoordinates(pts)
@property
def z_coordinates(self):
pts = self.GetZCoordinates()
if not NUMPY_AVAILABLE:
return pts
if not pts:
return None
return dsa.vtkDataArrayToVTKArray(pts)
@z_coordinates.setter
def z_coordinates(self, points):
if isinstance(points, vtkDataArray):
self.SetZCoordinates(points)
return
if not NUMPY_AVAILABLE:
raise ValueError("Expect vtkDataArray")
pts = dsa.numpyTovtkDataArray(points, "z_coords")
self.SetZCoordinates(pts)
class CompositeDataIterator(object):
"""Wrapper for a vtkCompositeDataIterator class to satisfy
the python iterator protocol. This iterator iterates
over non-empty leaf nodes. To iterate over empty or
non-leaf nodes, use the vtkCompositeDataIterator directly.
"""
def __init__(self, cds):
self.Iterator = cds.NewIterator()
if self.Iterator:
self.Iterator.UnRegister(None)
self.Iterator.GoToFirstItem()
def __iter__(self):
return self
def __next__(self):
if not self.Iterator:
raise StopIteration
if self.Iterator.IsDoneWithTraversal():
raise StopIteration
retVal = self.Iterator.GetCurrentDataObject()
self.Iterator.GoToNextItem()
return retVal
def next(self):
return self.__next__()
def __getattr__(self, name):
"""Returns attributes from the vtkCompositeDataIterator."""
return getattr(self.Iterator, name)
class CompositeDataSetBase(object):
"""A wrapper for vtkCompositeData and subclasses that makes it easier
to access Point/Cell/Field data as VTKCompositeDataArrays. It also
provides a Python type iterator."""
def __init__(self, **kwargs):
self._PointData = None
self._CellData = None
self._FieldData = None
self._Points = None
def __iter__(self):
"Creates an iterator for the contained datasets."
return CompositeDataIterator(self)
def get_attributes(self, type):
"""Returns the attributes specified by the type as a
CompositeDataSetAttributes instance."""
return CompositeDataSetAttributes(self, type)
@property
def point_data(self):
"Returns the point data as a DataSetAttributes instance."
if self._PointData is None or self._PointData() is None:
pdata = self.get_attributes(vtkDataObject.POINT)
self._PointData = weakref.ref(pdata)
return self._PointData()
@property
def cell_data(self):
"Returns the cell data as a DataSetAttributes instance."
if self._CellData is None or self._CellData() is None:
cdata = self.get_attributes(vtkDataObject.CELL)
self._CellData = weakref.ref(cdata)
return self._CellData()
@property
def field_data(self):
"Returns the field data as a DataSetAttributes instance."
if self._FieldData is None or self._FieldData() is None:
fdata = self.get_attributes(vtkDataObject.FIELD)
self._FieldData = weakref.ref(fdata)
return self._FieldData()
@property
def points(self):
"Returns the points as a VTKCompositeDataArray instance."
if not NUMPY_AVAILABLE:
# don't know how to handle composite dataset when numpy not around
raise NotImplementedError("Only available with numpy")
if self._Points is None or self._Points() is None:
pts = []
for ds in self:
try:
_pts = ds.Points
except AttributeError:
_pts = None
if _pts is None:
pts.append(dsa.NoneArray)
else:
pts.append(_pts)
if len(pts) == 0 or all([a is dsa.NoneArray for a in pts]):
cpts = dsa.NoneArray
else:
cpts = dsa.VTKCompositeDataArray(pts, dataset=self)
self._Points = weakref.ref(cpts)
return self._Points()
@vtkPartitionedDataSet.override
class PartitionedDataSet(CompositeDataSetBase, vtkPartitionedDataSet):
def append(self, dataset):
self.SetPartition(self.GetNumberOfPartitions(), dataset)
@vtkPartitionedDataSetCollection.override
class PartitionedDataSetCollection(CompositeDataSetBase, vtkPartitionedDataSetCollection):
def append(self, dataset):
self.SetPartitionedDataSet(self.GetNumberOfPartitionedDataSets(), dataset)
@vtkOverlappingAMR.override
class OverlappingAMR(CompositeDataSetBase, vtkOverlappingAMR):
pass
@vtkMultiBlockDataSet.override
class MultiBlockDataSet(CompositeDataSetBase, vtkMultiBlockDataSet):
pass
@vtkStructuredGrid.override
class StructuredGrid(PointSet, vtkStructuredGrid):
def __init__(self, **kwargs):
PointSet.__init__(self, **kwargs)
vtkStructuredGrid.__init__(self, **kwargs)
@property
def x_coordinates(self):
if not NUMPY_AVAILABLE:
raise NotImplementedError("Only available with numpy")
dims = [0,0,0]
self.GetDimensions(dims)
return self.points[:, 0].reshape(dims, order="F")
@property
def y_coordinates(self):
if not NUMPY_AVAILABLE:
raise NotImplementedError("Only available with numpy")
dims = [0,0,0]
self.GetDimensions(dims)
return self.points[:, 1].reshape(dims, order="F")
@property
def z_coordinates(self):
if not NUMPY_AVAILABLE:
raise NotImplementedError("Only available with numpy")
dims = [0,0,0]
self.GetDimensions(dims)
return self.points[:, 2].reshape(dims, order="F")
# -----------------------------------------------------------------------------
# Handle pickle registration
# -----------------------------------------------------------------------------
with suppress(ImportError):
import copyreg
from vtkmodules.util.pickle_support import serialize_VTK_data_object
copyreg.pickle(PolyData, serialize_VTK_data_object)
copyreg.pickle(UnstructuredGrid, serialize_VTK_data_object)
copyreg.pickle(ImageData, serialize_VTK_data_object)
copyreg.pickle(PartitionedDataSet, serialize_VTK_data_object)
copyreg.pickle(StructuredGrid, serialize_VTK_data_object)