DAMASK_EICMD/python/damask/_vtk.py

386 lines
12 KiB
Python
Raw Normal View History

2020-11-18 20:48:57 +05:30
import os
import multiprocessing as mp
2020-06-03 14:13:07 +05:30
from pathlib import Path
import numpy as np
import numpy.ma as ma
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
from vtk.util.numpy_support import vtk_to_numpy as vtk_to_np
from . import util
2020-03-13 00:22:33 +05:30
from . import Table
2020-03-13 00:22:33 +05:30
class VTK:
"""
Spatial visualization (and potentially manipulation).
High-level interface to VTK.
"""
2020-08-25 13:26:24 +05:30
def __init__(self,vtk_data):
"""
New spatial visualization.
Parameters
----------
2020-08-25 13:26:24 +05:30
vtk_data : subclass of vtk.vtkDataSet
Description of geometry and topology, optionally with attached data.
Valid types are vtk.vtkRectilinearGrid, vtk.vtkUnstructuredGrid,
or vtk.vtkPolyData.
"""
2020-08-25 13:26:24 +05:30
self.vtk_data = vtk_data
@staticmethod
2020-10-27 18:12:49 +05:30
def from_rectilinear_grid(grid,size,origin=np.zeros(3)):
"""
Create VTK of type vtk.vtkRectilinearGrid.
This is the common type for grid solver results.
Parameters
----------
grid : iterable of int, len (3)
Number of cells along each dimension.
size : iterable of float, len (3)
Physical lengths along each dimension.
origin : iterable of float, len (3), optional
Spatial origin coordinates.
"""
2020-08-25 13:26:24 +05:30
vtk_data = vtk.vtkRectilinearGrid()
vtk_data.SetDimensions(*(np.array(grid)+1))
coord = [np_to_vtk(np.linspace(origin[i],origin[i]+size[i],grid[i]+1),deep=True) for i in [0,1,2]]
[coord[i].SetName(n) for i,n in enumerate(['x','y','z'])]
2020-08-25 13:26:24 +05:30
vtk_data.SetXCoordinates(coord[0])
vtk_data.SetYCoordinates(coord[1])
vtk_data.SetZCoordinates(coord[2])
2020-03-11 12:02:03 +05:30
2020-08-25 13:26:24 +05:30
return VTK(vtk_data)
@staticmethod
2020-10-27 18:12:49 +05:30
def from_unstructured_grid(nodes,connectivity,cell_type):
"""
Create VTK of type vtk.vtkUnstructuredGrid.
This is the common type for FEM solver results.
Parameters
----------
nodes : numpy.ndarray of shape (:,3)
Spatial position of the nodes.
connectivity : numpy.ndarray of np.dtype = int
2020-12-02 19:07:44 +05:30
Cell connectivity (0-based), first dimension determines #Cells,
second dimension determines #Nodes/Cell.
cell_type : str
2020-03-31 14:34:06 +05:30
Name of the vtk.vtkCell subclass. Tested for TRIANGLE, QUAD, TETRA, 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))
2020-08-25 13:26:24 +05:30
vtk_data = vtk.vtkUnstructuredGrid()
vtk_data.SetPoints(vtk_nodes)
2020-12-02 19:07:44 +05:30
cell_types = {'TRIANGLE':vtk.VTK_TRIANGLE, 'QUAD':vtk.VTK_QUAD,
'TETRA' :vtk.VTK_TETRA, 'HEXAHEDRON':vtk.VTK_HEXAHEDRON}
vtk_data.SetCells(cell_types[cell_type.split("_",1)[-1].upper()],cells)
2020-08-25 13:26:24 +05:30
return VTK(vtk_data)
2020-03-11 12:02:03 +05:30
@staticmethod
2020-10-27 18:12:49 +05:30
def from_poly_data(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.
"""
N = points.shape[0]
2020-03-22 21:33:28 +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
vtk_cells = vtk.vtkCellArray()
vtk_cells.SetNumberOfCells(N)
vtk_cells.SetCells(N,np_to_vtkIdTypeArray(np.stack((np.ones (N,dtype=np.int64),
np.arange(N,dtype=np.int64)),axis=1).ravel(),deep=True))
2020-08-25 13:26:24 +05:30
vtk_data = vtk.vtkPolyData()
vtk_data.SetPoints(vtk_points)
vtk_data.SetVerts(vtk_cells)
2020-03-12 03:05:58 +05:30
2020-08-25 13:26:24 +05:30
return VTK(vtk_data)
2020-03-12 12:45:12 +05:30
2020-03-12 04:24:36 +05:30
@staticmethod
def load(fname,dataset_type=None):
"""
2020-12-04 02:28:24 +05:30
Load from VTK file.
Parameters
----------
2020-06-28 10:47:51 +05:30
fname : str or pathlib.Path
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 a .vtk file. Valid types are vtkRectilinearGrid,
vtkUnstructuredGrid, and vtkPolyData.
"""
if not os.path.isfile(fname): # vtk has a strange error handling
2020-11-24 00:36:34 +05:30
raise FileNotFoundError(f'no such file: {fname}')
2020-06-03 14:13:07 +05:30
ext = Path(fname).suffix
if ext == '.vtk' or dataset_type is not None:
2020-03-12 04:24:36 +05:30
reader = vtk.vtkGenericDataObjectReader()
2020-06-28 10:47:51 +05:30
reader.SetFileName(str(fname))
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'):
reader.Update()
2020-08-25 13:26:24 +05:30
vtk_data = reader.GetRectilinearGridOutput()
elif dataset_type.lower().endswith('unstructuredgrid'):
reader.Update()
2020-08-25 13:26:24 +05:30
vtk_data = reader.GetUnstructuredGridOutput()
elif dataset_type.lower().endswith('polydata'):
reader.Update()
2020-08-25 13:26:24 +05:30
vtk_data = reader.GetPolyDataOutput()
2020-03-12 04:24:36 +05:30
else:
2020-06-24 21:35:12 +05:30
raise TypeError(f'Unknown dataset type {dataset_type} for vtk file')
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-06-24 21:35:12 +05:30
raise TypeError(f'Unknown file extension {ext}')
2020-03-12 04:24:36 +05:30
2020-06-28 10:47:51 +05:30
reader.SetFileName(str(fname))
2020-03-12 04:24:36 +05:30
reader.Update()
2020-08-25 13:26:24 +05:30
vtk_data = reader.GetOutput()
2020-03-12 04:24:36 +05:30
2020-08-25 13:26:24 +05:30
return VTK(vtk_data)
2020-03-12 04:24:36 +05:30
@staticmethod
def _write(writer):
"""Wrapper for parallel writing."""
writer.Write()
def save(self,fname,parallel=True,compress=True):
"""
2020-12-04 02:28:24 +05:30
Save as VTK file.
Parameters
----------
2020-06-28 10:47:51 +05:30
fname : str or pathlib.Path
Filename for writing.
parallel : boolean, optional
Write data in parallel background process. Defaults to True.
2020-09-18 20:02:08 +05:30
compress : bool, optional
Compress with zlib algorithm. Defaults to True.
"""
2020-08-25 13:26:24 +05:30
if isinstance(self.vtk_data,vtk.vtkRectilinearGrid):
2020-03-11 12:02:03 +05:30
writer = vtk.vtkXMLRectilinearGridWriter()
2020-08-25 13:26:24 +05:30
elif isinstance(self.vtk_data,vtk.vtkUnstructuredGrid):
writer = vtk.vtkXMLUnstructuredGridWriter()
2020-08-25 13:26:24 +05:30
elif isinstance(self.vtk_data,vtk.vtkPolyData):
2020-03-11 12:02:03 +05:30
writer = vtk.vtkXMLPolyDataWriter()
2020-11-04 22:38:04 +05:30
default_ext = '.'+writer.GetDefaultFileExtension()
2020-06-03 14:13:07 +05:30
ext = Path(fname).suffix
2020-11-06 02:08:00 +05:30
writer.SetFileName(str(fname)+(default_ext if default_ext != ext else ''))
2020-11-04 22:38:04 +05:30
2020-09-12 17:16:55 +05:30
if compress:
writer.SetCompressorTypeToZLib()
else:
writer.SetCompressorTypeToNone()
2020-03-11 12:02:03 +05:30
writer.SetDataModeToBinary()
2020-08-25 13:26:24 +05:30
writer.SetInputData(self.vtk_data)
if parallel:
try:
mp_writer = mp.Process(target=self._write,args=(writer,))
mp_writer.start()
except TypeError:
writer.Write()
else:
writer.Write()
2020-03-11 12:02:03 +05:30
# Check https://blog.kitware.com/ghost-and-blanking-visibility-changes/ for missing data
# Needs support for damask.Table
2020-03-12 04:24:36 +05:30
def add(self,data,label=None):
"""
Add data to either cells or points.
Parameters
----------
data : numpy.ndarray or numpy.ma.MaskedArray
Data to add. First dimension needs to match either
number of cells or number of points.
label : str
Data label.
"""
2020-08-25 13:26:24 +05:30
N_points = self.vtk_data.GetNumberOfPoints()
N_cells = self.vtk_data.GetNumberOfCells()
2020-03-12 04:24:36 +05:30
if isinstance(data,np.ndarray):
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')
N_data = data.shape[0]
data_ = np.where(data.mask,data.fill_value,data) if isinstance(data,ma.MaskedArray) else\
data
d = np_to_vtk((data_.astype(np.single) if data_.dtype in [np.double, np.longdouble] else
data_).reshape(N_data,-1),deep=True) # avoid large files
2020-03-12 04:24:36 +05:30
d.SetName(label)
if N_data == N_points:
2020-08-25 13:26:24 +05:30
self.vtk_data.GetPointData().AddArray(d)
elif N_data == N_cells:
self.vtk_data.GetCellData().AddArray(d)
else:
raise ValueError(f'Cell / point count ({N_cells} / {N_points}) differs from data ({N_data}).')
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
def get(self,label):
"""
Get either cell or point data.
Cell data takes precedence over point data, i.e. this
function assumes that labels are unique among cell and
point data.
Parameters
----------
label : str
Data label.
"""
2020-08-25 13:26:24 +05:30
cell_data = self.vtk_data.GetCellData()
for a in range(cell_data.GetNumberOfArrays()):
if cell_data.GetArrayName(a) == label:
return vtk_to_np(cell_data.GetArray(a))
2020-08-25 13:26:24 +05:30
point_data = self.vtk_data.GetPointData()
for a in range(point_data.GetNumberOfArrays()):
if point_data.GetArrayName(a) == label:
return vtk_to_np(point_data.GetArray(a))
raise ValueError(f'Array "{label}" not found.')
2020-08-24 02:52:53 +05:30
def get_comments(self):
"""Return the comments."""
2020-08-25 13:26:24 +05:30
fielddata = self.vtk_data.GetFieldData()
2020-08-24 02:52:53 +05:30
for a in range(fielddata.GetNumberOfArrays()):
if fielddata.GetArrayName(a) == 'comments':
comments = fielddata.GetAbstractArray(a)
return [comments.GetValue(i) for i in range(comments.GetNumberOfValues())]
return []
def set_comments(self,comments):
"""
2020-08-25 04:10:14 +05:30
Set comments.
2020-08-24 02:52:53 +05:30
Parameters
----------
comments : str or list of str
Comments.
2020-08-24 02:52:53 +05:30
"""
s = vtk.vtkStringArray()
s.SetName('comments')
for c in [comments] if isinstance(comments,str) else comments:
s.InsertNextValue(c)
2020-08-25 13:26:24 +05:30
self.vtk_data.GetFieldData().AddArray(s)
2020-08-24 02:52:53 +05:30
def add_comments(self,comments):
"""
2020-08-25 04:10:14 +05:30
Add comments.
Parameters
----------
comments : str or list of str
Comments to add.
"""
2020-08-25 04:10:14 +05:30
self.set_comments(self.get_comments() + ([comments] if isinstance(comments,str) else comments))
2020-03-11 12:02:03 +05:30
def __repr__(self):
"""ASCII representation of the VTK data."""
writer = vtk.vtkDataSetWriter()
writer.SetHeader(f'# {util.execution_stamp("VTK")}')
2020-03-11 12:02:03 +05:30
writer.WriteToOutputStringOn()
2020-08-25 13:26:24 +05:30
writer.SetInputData(self.vtk_data)
2020-03-11 12:02:03 +05:30
writer.Write()
return writer.GetOutputString()
2020-08-24 04:04:07 +05:30
2020-03-12 18:15:47 +05:30
def show(self):
"""
Render.
See http://compilatrix.com/article/vtk-1 for further ideas.
"""
2021-02-12 02:26:53 +05:30
try:
import wx
_ = wx.App(False) # noqa
width, height = wx.GetDisplaySize()
except ImportError:
try:
2021-02-12 02:26:53 +05:30
import tkinter
tk = tkinter.Tk()
width = tk.winfo_screenwidth()
height = tk.winfo_screenheight()
tk.destroy()
except Exception as e:
width = 1024
height = 768
2020-03-12 18:15:47 +05:30
mapper = vtk.vtkDataSetMapper()
2020-08-25 13:26:24 +05:30
mapper.SetInputData(self.vtk_data)
2020-03-12 18:15:47 +05:30
actor = vtk.vtkActor()
actor.SetMapper(mapper)
ren = vtk.vtkRenderer()
2020-03-22 21:33:28 +05:30
window = vtk.vtkRenderWindow()
window.AddRenderer(ren)
2020-03-12 18:15:47 +05:30
ren.AddActor(actor)
ren.SetBackground(0.2,0.2,0.2)
2020-03-13 00:22:33 +05:30
2021-02-12 02:26:53 +05:30
window.SetSize(width,height)
2020-03-12 18:15:47 +05:30
iren = vtk.vtkRenderWindowInteractor()
2020-03-22 21:33:28 +05:30
iren.SetRenderWindow(window)
2020-03-12 18:15:47 +05:30
iren.Initialize()
2020-03-22 21:33:28 +05:30
window.Render()
2020-03-12 18:15:47 +05:30
iren.Start()