Commit 4335909f authored by Lukas Riedel's avatar Lukas Riedel

Only load required components of VTK Python lib

Also fix a bug where cell data arrays would be ignored.
parent 2b3ae9a0
import vtk
from vtk import vtkXMLGenericDataObjectReader, \
vtkCellTreeLocator, \
vtkDataSet, \
vtkGenericCell, \
reference
import numpy as np
class VTKReader:
......@@ -8,13 +12,13 @@ class VTKReader:
self._data_array = None
# set up the reader
self._reader = vtk.vtkXMLGenericDataObjectReader()
self._reader = vtkXMLGenericDataObjectReader()
self._reader.SetFileName(file_name)
self._reader.Update()
self._dataset = vtk.vtkDataSet.SafeDownCast(self._reader.GetOutput())
self._dataset = vtkDataSet.SafeDownCast(self._reader.GetOutput())
# build the locator
self._locator = vtk.vtkCellTreeLocator()
self._locator = vtkCellTreeLocator()
self._locator.SetDataSet(self._dataset)
self._locator.BuildLocator()
......@@ -22,7 +26,7 @@ class VTKReader:
self.set_data_array(array_name, data_type)
# cache for cell and ID
self._cell = vtk.vtkGenericCell()
self._cell = vtkGenericCell()
self._cell_id = -1
def set_data_array (self, array_name, data_type=None):
......@@ -39,13 +43,14 @@ class VTKReader:
if data_type == "cell" or data_type is None:
# retrieve cell data array...
self._data_array = retrieve_array(self._dataset.GetCellData(),
array_name)
array_name)
self._eval = self.__eval_cell
if data_type == "vertex" or data_type is None:
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)
array_name)
self._eval = self.__eval_vertex
if self._data_array is None:
......@@ -71,7 +76,7 @@ class VTKReader:
# unused stubs
zeros = [0, 0, 0]
zero_ref = vtk.reference(0)
zero_ref = reference(0)
# interpolation weights
weights = np.zeros(10)
......
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