DAMASK_EICMD/python/damask/_vtk.py

251 lines
7.9 KiB
Python
Raw Normal View History

2020-03-11 12:02:03 +05:30
import os
import pandas as pd
import numpy as np
import vtk
2020-03-19 15:29:53 +05:30
from vtk.util.numpy_support import numpy_to_vtk as np_to_vtk
from vtk.util.numpy_support import numpy_to_vtkIdTypeArray as np_to_vtkIdTypeArray
2020-03-13 00:22:33 +05:30
from . import Table
from . import Environment
from . import version
2020-03-13 00:22:33 +05:30
class VTK:
"""
Spatial visualization (and potentially manipulation).
High-level interface to VTK.
"""
def __init__(self,geom):
"""
Set geometry and topology.
Parameters
----------
geom : subclass of vtk.vtkDataSet
Description of geometry and topology. Valid types are vtk.vtkRectilinearGrid,
vtk.vtkUnstructuredGrid, or vtk.vtkPolyData.
"""
self.geom = geom
@staticmethod
def from_rectilinearGrid(grid,size,origin=np.zeros(3)):
"""
Create VTK of type vtk.vtkRectilinearGrid.
This is the common type for results from the grid solver.
Parameters
----------
grid : numpy.ndarray of shape (3) of np.dtype = int
Number of cells.
size : numpy.ndarray of shape (3)
Physical length.
origin : numpy.ndarray of shape (3), optional
Spatial origin.
"""
geom = vtk.vtkRectilinearGrid()
geom.SetDimensions(*(grid+1))
2020-03-19 15:29:53 +05:30
geom.SetXCoordinates(np_to_vtk(np.linspace(origin[0],origin[0]+size[0],grid[0]+1),deep=True))
geom.SetYCoordinates(np_to_vtk(np.linspace(origin[1],origin[1]+size[1],grid[1]+1),deep=True))
geom.SetZCoordinates(np_to_vtk(np.linspace(origin[2],origin[2]+size[2],grid[2]+1),deep=True))
2020-03-11 12:02:03 +05:30
return VTK(geom)
@staticmethod
def from_unstructuredGrid(nodes,connectivity,cell_type):
"""
Create VTK of type vtk.vtkUnstructuredGrid.
This is the common type for results from FEM solvers.
Parameters
----------
nodes : numpy.ndarray of shape (:,3)
Spatial position of the nodes.
connectivity : numpy.ndarray of np.dtype = int
Cell connectivity (0-based), first dimension determines #Cells, second dimension determines #Nodes/Cell.
cell_type : str
Name of the vtk.vtkCell subclass. Tested for TRIANGLE, QUAD, and HEXAHEDRON.
"""
vtk_nodes = vtk.vtkPoints()
2020-03-12 12:45:12 +05:30
vtk_nodes.SetData(np_to_vtk(nodes))
cells = vtk.vtkCellArray()
cells.SetNumberOfCells(connectivity.shape[0])
T = np.concatenate((np.ones((connectivity.shape[0],1),dtype=np.int64)*connectivity.shape[1],
2020-03-12 12:45:12 +05:30
connectivity),axis=1).ravel()
2020-03-19 15:29:53 +05:30
cells.SetCells(connectivity.shape[0],np_to_vtkIdTypeArray(T,deep=True))
geom = vtk.vtkUnstructuredGrid()
geom.SetPoints(vtk_nodes)
2020-03-12 12:45:12 +05:30
geom.SetCells(eval('vtk.VTK_{}'.format(cell_type.split('_',1)[-1].upper())),cells)
return VTK(geom)
2020-03-11 12:02:03 +05:30
@staticmethod
2020-03-12 03:05:58 +05:30
def from_polyData(points):
"""
Create VTK of type vtk.polyData.
This is the common type for point-wise data.
Parameters
----------
points : numpy.ndarray of shape (:,3)
Spatial position of the points.
"""
2020-03-12 03:05:58 +05:30
vtk_points= vtk.vtkPoints()
2020-03-12 12:45:12 +05:30
vtk_points.SetData(np_to_vtk(points))
2020-03-12 03:05:58 +05:30
geom = vtk.vtkPolyData()
geom.SetPoints(vtk_points)
return VTK(geom)
2020-03-12 12:45:12 +05:30
2020-03-12 04:24:36 +05:30
@staticmethod
2020-03-12 12:45:12 +05:30
def from_file(fname,dataset_type=None):
"""
Create VTK from file.
Parameters
----------
fname : str
Filename for reading. Valid extensions are *.vtr, *.vtu, *.vtp, and *.vtk.
2020-03-12 12:45:12 +05:30
dataset_type : str, optional
Name of the vtk.vtkDataSet subclass when opening an *.vtk file. Valid types are vtkRectilinearGrid,
vtkUnstructuredGrid, and vtkPolyData.
"""
2020-03-12 04:24:36 +05:30
ext = os.path.splitext(fname)[1]
if ext == '.vtk':
reader = vtk.vtkGenericDataObjectReader()
reader.SetFileName(fname)
reader.Update()
2020-03-21 15:37:21 +05:30
if dataset_type is None:
raise TypeError('Dataset type for *.vtk file not given.')
elif dataset_type.lower().endswith('rectilineargrid'):
2020-03-12 04:24:36 +05:30
geom = reader.GetRectilinearGridOutput()
elif dataset_type.lower().endswith('unstructuredgrid'):
2020-03-12 04:24:36 +05:30
geom = reader.GetUnstructuredGridOutput()
elif dataset_type.lower().endswith('polydata'):
2020-03-12 04:24:36 +05:30
geom = reader.GetPolyDataOutput()
else:
2020-03-12 12:45:12 +05:30
raise TypeError('Unknown dataset type for vtk file {}'.format(dataset_type))
2020-03-12 04:24:36 +05:30
else:
if ext == '.vtr':
reader = vtk.vtkXMLRectilinearGridReader()
elif ext == '.vtu':
reader = vtk.vtkXMLUnstructuredGridReader()
elif ext == '.vtp':
reader = vtk.vtkXMLPolyDataReader()
else:
2020-03-12 12:45:12 +05:30
raise TypeError('Unknown file extension {}'.format(ext))
2020-03-12 04:24:36 +05:30
reader.SetFileName(fname)
reader.Update()
geom = reader.GetOutput()
return VTK(geom)
2020-03-12 04:24:36 +05:30
def write(self,fname):
"""
Write to file.
Parameters
----------
fname : str
Filename for writing.
"""
2020-03-11 12:20:11 +05:30
if (isinstance(self.geom,vtk.vtkRectilinearGrid)):
2020-03-11 12:02:03 +05:30
writer = vtk.vtkXMLRectilinearGridWriter()
2020-03-11 12:20:11 +05:30
elif(isinstance(self.geom,vtk.vtkUnstructuredGrid)):
writer = vtk.vtkXMLUnstructuredGridWriter()
2020-03-11 12:20:11 +05:30
elif(isinstance(self.geom,vtk.vtkPolyData)):
2020-03-11 12:02:03 +05:30
writer = vtk.vtkXMLPolyDataWriter()
2020-03-19 22:04:31 +05:30
default_ext = writer.GetDefaultFileExtension()
name, ext = os.path.splitext(fname)
2020-03-19 23:13:49 +05:30
if ext and ext != '.'+default_ext:
raise ValueError('Given extension {} is not .{}'.format(ext,default_ext))
2020-03-19 22:04:31 +05:30
writer.SetFileName('{}.{}'.format(name,default_ext))
2020-03-11 12:02:03 +05:30
writer.SetCompressorTypeToZLib()
writer.SetDataModeToBinary()
writer.SetInputData(self.geom)
writer.Write()
# Check https://blog.kitware.com/ghost-and-blanking-visibility-changes/ for missing data
# Needs support for pd.DataFrame and/or table
2020-03-12 04:24:36 +05:30
def add(self,data,label=None):
"""Add data to either cells or points."""
N_points = self.geom.GetNumberOfPoints()
N_cells = self.geom.GetNumberOfCells()
2020-03-12 04:24:36 +05:30
if isinstance(data,np.ndarray):
2020-03-12 12:45:12 +05:30
d = np_to_vtk(num_array=data.reshape(data.shape[0],-1),deep=True)
2020-03-21 15:37:21 +05:30
if label is None:
2020-03-19 23:13:49 +05:30
raise ValueError('No label defined for numpy.ndarray')
2020-03-12 04:24:36 +05:30
d.SetName(label)
if data.shape[0] == N_cells:
2020-03-12 04:24:36 +05:30
self.geom.GetCellData().AddArray(d)
elif data.shape[0] == N_points:
2020-03-12 04:24:36 +05:30
self.geom.GetPointData().AddArray(d)
elif isinstance(data,pd.DataFrame):
2020-03-19 23:13:49 +05:30
raise NotImplementedError('pd.DataFrame')
2020-03-13 00:22:33 +05:30
elif isinstance(data,Table):
2020-03-19 23:13:49 +05:30
raise NotImplementedError('damask.Table')
2020-03-19 13:33:38 +05:30
else:
raise TypeError
2020-03-11 12:02:03 +05:30
def __repr__(self):
"""ASCII representation of the VTK data."""
writer = vtk.vtkDataSetWriter()
2020-03-12 03:05:58 +05:30
writer.SetHeader('# DAMASK.VTK v{}'.format(version))
2020-03-11 12:02:03 +05:30
writer.WriteToOutputStringOn()
writer.SetInputData(self.geom)
writer.Write()
return writer.GetOutputString()
2020-03-12 18:15:47 +05:30
def show(self):
"""
Render.
See http://compilatrix.com/article/vtk-1 for further ideas.
"""
mapper = vtk.vtkDataSetMapper()
mapper.SetInputData(self.geom)
actor = vtk.vtkActor()
actor.SetMapper(mapper)
ren = vtk.vtkRenderer()
renWin = vtk.vtkRenderWindow()
renWin.AddRenderer(ren)
ren.AddActor(actor)
ren.SetBackground(0.2,0.2,0.2)
2020-03-13 00:22:33 +05:30
renWin.SetSize(Environment().screen_width,Environment().screen_height)
2020-03-12 18:15:47 +05:30
iren = vtk.vtkRenderWindowInteractor()
iren.SetRenderWindow(renWin)
iren.Initialize()
renWin.Render()
iren.Start()