also accept vtkRectilinearGrid etc as type

This commit is contained in:
Martin Diehl 2020-03-19 08:04:15 +01:00
parent 6d0c4b9765
commit fbd92599e8
1 changed files with 19 additions and 16 deletions

View File

@ -30,6 +30,7 @@ class VTK:
""" """
self.geom = geom self.geom = geom
@staticmethod @staticmethod
def from_rectilinearGrid(grid,size,origin=np.zeros(3)): def from_rectilinearGrid(grid,size,origin=np.zeros(3)):
""" """
@ -123,9 +124,9 @@ class VTK:
Parameters Parameters
---------- ----------
fname : str fname : str
Filename for reading. Valid extensions are .vtk, .vtr, .vtu, and .vtp. Filename for reading. Valid extensions are *.vtr, *.vtu, *.vtp, and *.vtk.
dataset_type : str, optional dataset_type : str, optional
Name of the vtk.vtkDataSet subclass when opening an .vtk file. Valid types are vtkRectilinearGrid, Name of the vtk.vtkDataSet subclass when opening an *.vtk file. Valid types are vtkRectilinearGrid,
vtkUnstructuredGrid, and vtkPolyData. vtkUnstructuredGrid, and vtkPolyData.
""" """
@ -134,11 +135,13 @@ class VTK:
reader = vtk.vtkGenericDataObjectReader() reader = vtk.vtkGenericDataObjectReader()
reader.SetFileName(fname) reader.SetFileName(fname)
reader.Update() reader.Update()
if 'rectilineargrid' in dataset_type.lower(): if not dataset_type:
raise TypeError('Dataset type for *.vtk file not given.')
elif dataset_type.lower().endswith('rectilineargrid'):
geom = reader.GetRectilinearGridOutput() geom = reader.GetRectilinearGridOutput()
elif 'unstructuredgrid' in dataset_type.lower(): elif dataset_type.lower().endswith('unstructuredgrid'):
geom = reader.GetUnstructuredGridOutput() geom = reader.GetUnstructuredGridOutput()
elif 'polydata' in dataset_type.lower(): elif dataset_type.lower().endswith('polydata'):
geom = reader.GetPolyDataOutput() geom = reader.GetPolyDataOutput()
else: else:
raise TypeError('Unknown dataset type for vtk file {}'.format(dataset_type)) raise TypeError('Unknown dataset type for vtk file {}'.format(dataset_type))