Commit 2ab991d9 authored by Lukas Riedel's avatar Lukas Riedel

Add generic VTK XML file reader

Capabilities:
* Load any VTK XML file containing a VTK DataSet.
* Evaluate data arrays of cell or point (vertex) data.
* Select data arrays by name.
* Return data set bounds.
parent ebbc64b0
import vtk
import numpy as np
class VTKReader:
"""Read scalar or vector data arrays from an unstructured VTK file."""
def __init__(self, file_name, array_name, data_type=None):
self._eval = None
self._data_array = None
# set up the reader
self._reader = vtk.vtkXMLGenericDataObjectReader()
self._reader.SetFileName(file_name)
self._reader.Update()
self._dataset = vtk.vtkDataSet.SafeDownCast(self._reader.GetOutput())
# build the locator
self._locator = vtk.vtkCellTreeLocator()
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 = vtk.vtkGenericCell()
self._cell_id = -1
def set_data_array (self, array_name, data_type=None):
"""Switch the data array (dataset) to evaluate."""
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 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
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))
self._num_comp = self._data_array.GetNumberOfComponents()
def get_bounds (self):
"""Return the outer bounds of the grid."""
bounds = self._dataset.GetBounds()
return np.reshape(bounds, (3, 2))
def evaluate (self, position):
"""Evaluate the selected data array at a certain position"""
if len(position) != 3:
raise ValueError("Position must be a 3D vector/array")
# unused stubs
zeros = [0, 0, 0]
zero_ref = vtk.reference(0)
# interpolation weights
weights = np.zeros(10)
# check if we cached the correct cell
on_cell = self._cell.EvaluatePosition(position,
zeros,
zero_ref,
zeros,
zero_ref,
weights)
if on_cell == 0 or self._cell_id < 0:
# locate correct cell
self._cell_id = self._locator.FindCell(position,
0,
self._cell,
zeros,
weights)
if self._cell_id < 0:
raise ValueError("Position outside domain: {}".format(position))
elif on_cell < 0:
raise RuntimeError("Error while evaluating if position {} is on "
"cell {}".format(position, self._cell_id))
# evaluate the data
data = self._eval(weights)
return np.squeeze(data)
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):
"""Evaluation function for a vertex-type data array."""
data = np.zeros(self._num_comp)
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
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