Commit 0f04af0c authored by Lukas Riedel's avatar Lukas Riedel

Add full docstrings for VTKReader and improve properties

parent 4335909f
......@@ -6,8 +6,32 @@ from vtk import vtkXMLGenericDataObjectReader, \
import numpy as np
class VTKReader:
"""Read scalar or vector data arrays from an unstructured VTK file."""
"""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.
See :func:`~set_data_array` for details on the arguments regarding the
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
......@@ -22,15 +46,25 @@ class VTKReader:
self._locator.SetDataSet(self._dataset)
self._locator.BuildLocator()
# set the evaluation data array
self.set_data_array(array_name, data_type)
# cache for cell and ID
self._cell = vtkGenericCell()
self._cell_id = -1
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."""
"""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.
Raises:
RuntimeError: Data array was not found.
"""
self._data_array = None
def retrieve_array (data, name):
......@@ -44,14 +78,14 @@ class VTKReader:
# retrieve cell data array...
self._data_array = retrieve_array(self._dataset.GetCellData(),
array_name)
self._eval = self.__eval_cell
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._eval = self.__eval_vertex
self._eval = self.__eval_vertex__
if self._data_array is None:
if data_type is None:
......@@ -64,15 +98,41 @@ class VTKReader:
self._num_comp = self._data_array.GetNumberOfComponents()
def get_bounds (self):
"""Return the outer bounds of the grid."""
@property
def bounds (self):
"""Return the outer bounds of the grid.
Returns:
np.array: The rectangular bounds in any spatial direction.
Shape: ``(3, 2)``.
"""
bounds = self._dataset.GetBounds()
return np.reshape(bounds, (3, 2))
@property
def number_of_components (self):
"""Number of components of the evaluated quantity."""
return self._num_comp
def evaluate (self, position):
"""Evaluate the selected data array at a certain position"""
"""Evaluate the selected data array at a certain position.
Args:
position (list): 3D coordinate probing the data array.
Returns:
np.array: The value of the data array at the given position.
The output is squeezed: Scalar data arrays return scalar values,
vector arrays return numpy arrays with shape ``(3)``.
Raises:
ValueError: ``position`` argument not a 3D list or array.
ValueError: Position outside grid.
RuntimeError: Numerical error in cell locator algorithm.
"""
if len(position) != 3:
raise ValueError("Position must be a 3D vector/array")
raise ValueError("Position must be a 3D lest/array")
# unused stubs
zeros = [0, 0, 0]
......@@ -108,11 +168,11 @@ class VTKReader:
data = self._eval(weights)
return np.squeeze(data)
def __eval_cell (self, weights):
def __eval_cell__ (self, weights):
"""Evaluation function for a cell-type data array."""
return np.asarray(self._data_array.GetTuple(self._cell_id))
def __eval_vertex (self, weights):
def __eval_vertex__ (self, weights):
"""Evaluation function for a vertex-type data array."""
data = np.zeros(self._num_comp)
for point in range(self._cell.GetNumberOfPoints()):
......
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