Commit ff81d71f authored by Lukas Riedel's avatar Lukas Riedel

Split VTKReader into reader and VTKDataArray

This enables convenient iteration over all data arrays in a VTK file,
even without their names given. These can now be retrieved from the
reader directly.

* Make VTKReader a mapping: Use it like dict to access data arrays.
* Wrap data arrays into VTKDataArray for evaluation.
parent 930da619
from collections.abc import Mapping
from vtk import vtkXMLGenericDataObjectReader, \
vtkCellTreeLocator, \
vtkDataSet, \
......@@ -5,114 +6,57 @@ from vtk import vtkXMLGenericDataObjectReader, \
reference
import numpy as np
class VTKReader:
"""A generic VTK reader for any type of VTK XML files.
It is intended to evaluate cell or vertex (point) data arrays on any point
of the grid.
Attributes:
_reader: Generic VTK XML data file reader.
_dataset: The complete data set loaded by the reader.
_locator: Fast locator of grid cells based on global coordinates.
_cell: Generic grid cell implementation.
_cell_id: ID of the cell on the grid.
_eval: Evaluation method based on the selected data array type.
_data_array: The actual data to evaluate.
"""
def __init__(self, file_name, array_name, data_type=None):
"""Open the target VTK file, build a cell locator and open the data array.
class VTKDataArray:
"""Wrapper for evaluating a VTK data array"""
See :func:`~set_data_array` for details on the arguments regarding the
data array.
def __init__(self, dataset, locator, data_array, data_type):
"""Initialize this wrapper from a dataset, a cell locator,
and the underlying VTK data array.
Args:
file_name (str): Path to the VTK file.
array_name (str): Name of the data array to evaluate.
data_type (str): Type of the data array.
"""
self._eval = None
self._data_array = None
# set up the reader
self._reader = vtkXMLGenericDataObjectReader()
self._reader.SetFileName(file_name)
self._reader.Update()
self._dataset = vtkDataSet.SafeDownCast(self._reader.GetOutput())
# build the locator
self._locator = vtkCellTreeLocator()
self._locator.SetDataSet(self._dataset)
self._locator.BuildLocator()
# cache for cell and ID
self._cell = vtkGenericCell()
self._cell_id = -1 # invalid cell
# set the evaluation data array
self.set_data_array(array_name, data_type)
def set_data_array (self, array_name, data_type=None):
"""Switch the data array (dataset) to evaluate.
Args:
array_name (str): The name of the data array to evaluate.
data_type (str, optional): The type of the data array (either
``cell`` or ``vertex``). If this is ``None``, vertex arrays
take precedence.
dataset: VTK dataset object the data array is stored in.
locator: Locator for finding cells based on global coordinates.
data_array: VTK data array containing the data to evaluate.
data_type (str): Type of data stored in the array (``point`` or
``cell``).
Raises:
RuntimeError: Data array was not found.
ValueError: Unsupported data type.
"""
self._data_array = None
def retrieve_array (data, name):
"""Return the array if it exists in the data."""
if data.HasArray(name) == 1:
return data.GetArray(name)
else:
return None
if data_type == "cell" or data_type is None:
# retrieve cell data array...
self._data_array = retrieve_array(self._dataset.GetCellData(),
array_name)
self._eval = self.__eval_cell__
if self._data_array is None \
and (data_type == "vertex" or data_type is None):
# ...or vertex data array
self._data_array = retrieve_array(self._dataset.GetPointData(),
array_name)
self._locator = locator
self._data_array = data_array
if data_type == "point":
self._eval = self.__eval_vertex__
elif data_type == "cell":
self._eval = self.__eval_cell__
else:
raise ValueError("Unsupported VTK array data type: {}".format(
data_type
))
if self._data_array is None:
if data_type is None:
raise RuntimeError("Data array '{}' not found in VTK file"
.format(array_name))
else:
raise RuntimeError("Data array '{}' of type '{}' not found in "
"VTK file. Supported types are: 'cell', "
"'vertex'".format(array_name, data_type))
# store physical dataset bounds
self._bounds = np.reshape(dataset.GetBounds(), (3, 2))
self._num_comp = self._data_array.GetNumberOfComponents()
# cache for cell and ID
self._cell = vtkGenericCell()
self._cell_id = -1 # invalid cell
@property
def bounds (self):
"""Return the outer bounds of the grid.
"""Return the outer bounds of the grid this data array lives on.
Returns:
np.array: The rectangular bounds in any spatial direction.
Shape: ``(3, 2)``.
"""
bounds = self._dataset.GetBounds()
return np.reshape(bounds, (3, 2))
return self._bounds
@property
def number_of_components (self):
"""Number of components of the evaluated quantity."""
return self._num_comp
return self._data_array.GetNumberOfComponents()
def evaluate (self, position):
"""Evaluate the selected data array at a certain position.
......@@ -174,10 +118,106 @@ class VTKReader:
def __eval_vertex__ (self, weights):
"""Evaluation function for a vertex-type data array."""
data = np.zeros(self._num_comp)
data = np.zeros(self.number_of_components)
for point in range(self._cell.GetNumberOfPoints()):
point_id = self._cell.GetPointId(point)
tup = np.asarray(self._data_array.GetTuple(point_id))
data += weights[point] * tup
return data
class VTKReader (Mapping):
"""A generic VTK reader for any type of VTK XML files.
It maps the data array names to data array evaluators, see
:py:class:`VTKDataArray`. Use it like any Python dict.
Warning:
VTK files separately store vertex (point) and cell data and associate
the data arrays by their names. If the loaded file contains one array
of each type with the same name, the *vertex (point)* array takes
precedence. **The cell array cannot be accessed!** The key sequence
will contain the name twice, but the associated value will be the
vertex data set in both cases.
Args:
file_name (str): Name of the VTK file to open. Any file in XML
format is supported.
Attributes:
_reader: Generic VTK XML data file reader.
_dataset: The complete data set loaded by the reader.
_locator: Fast locator of grid cells based on global coordinates.
_point_data: Pointer to the VTK point data arrays.
_cell_data: Pointer to the VTK cell data arrays.
"""
def __init__(self, file_name):
"""Open the target VTK file, and build a cell locator."""
# set up the reader
self._reader = vtkXMLGenericDataObjectReader()
self._reader.SetFileName(file_name)
self._reader.Update()
self._dataset = vtkDataSet.SafeDownCast(self._reader.GetOutput())
# build the locator
self._locator = vtkCellTreeLocator()
self._locator.SetDataSet(self._dataset)
self._locator.BuildLocator()
# ID of the current array
self._point_data = self._dataset.GetPointData()
self._cell_data = self._dataset.GetCellData()
def __len__(self):
return self._point_data.GetNumberOfArrays() \
+ self._cell_data.GetNumberOfArrays()
def __getitem__(self, key):
if self._point_data.HasArray(key) == 1:
return VTKDataArray(self._dataset,
self._locator,
self._point_data.GetArray(key),
"point")
elif self._cell_data.HasArray(key) == 1:
return VTKDataArray(self._dataset,
self._locator,
self._cell_data.GetArray(key),
"cell")
else:
raise KeyError("Data array not found: {}".format(key))
def __iter__(self):
for i in range(len(self)):
yield self._get_array_name(i)
def _get_array_name (self, array_id):
"""Return the name of the array at the given ID.
The VTK data set stores separate IDs for cell and point data arrays.
This method subtracts the number of point data arrays from the given
array ID if it is larger than the number of point data arrays, to
properly access cell data arrays.
Parameters:
array_id (int): Internal ID of the data array.
Returns:
str: The name of the array associated with the given ID.
"""
if array_id >= len(self):
raise ValueError("Array id '{}' out of bounds of dataset (size "
"is: {})".format(array_id, len(self)))
num_point_arrays = self._point_data.GetNumberOfArrays()
if array_id < num_point_arrays:
data = self._point_data
# start counting from zero again for cell arrays
else:
array_id = array_id - num_point_arrays
data = self._cell_data
return data.GetArrayName(array_id)
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment