572 lines
19 KiB
Python
572 lines
19 KiB
Python
import os
|
|
import multiprocessing as mp
|
|
from pathlib import Path
|
|
from typing import Union, Literal, List, Sequence
|
|
|
|
import numpy as np
|
|
import vtk
|
|
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 ._typehints import FloatSequence, IntSequence
|
|
from . import util
|
|
from . import Table
|
|
from . import Colormap
|
|
|
|
|
|
class VTK:
|
|
"""
|
|
Spatial visualization (and potentially manipulation).
|
|
|
|
High-level interface to VTK.
|
|
"""
|
|
|
|
def __init__(self,
|
|
vtk_data: vtk.vtkDataSet):
|
|
"""
|
|
New spatial visualization.
|
|
|
|
Parameters
|
|
----------
|
|
vtk_data : subclass of vtk.vtkDataSet
|
|
Description of geometry and topology, optionally with attached data.
|
|
Valid types are vtk.vtkImageData, vtk.vtkUnstructuredGrid,
|
|
vtk.vtkPolyData, and vtk.vtkRectilinearGrid.
|
|
|
|
"""
|
|
self.vtk_data = vtk_data
|
|
|
|
|
|
def copy(self):
|
|
if isinstance(self.vtk_data,vtk.vtkImageData):
|
|
dup = vtk.vtkImageData()
|
|
elif isinstance(self.vtk_data,vtk.vtkUnstructuredGrid):
|
|
dup = vtk.vtkUnstructuredGrid()
|
|
elif isinstance(self.vtk_data,vtk.vtkPolyData):
|
|
dup = vtk.vtkPolyData()
|
|
elif isinstance(self.vtk_data,vtk.vtkRectilinearGrid):
|
|
dup = vtk.vtkRectilinearGrid()
|
|
else:
|
|
raise TypeError
|
|
|
|
dup.DeepCopy(self.vtk_data)
|
|
|
|
return VTK(dup)
|
|
|
|
|
|
@property
|
|
def N_points(self) -> int:
|
|
"""Number of points in vtkdata."""
|
|
return self.vtk_data.GetNumberOfPoints()
|
|
|
|
@property
|
|
def N_cells(self) -> int:
|
|
"""Number of cells in vtkdata."""
|
|
return self.vtk_data.GetNumberOfCells()
|
|
|
|
@property
|
|
def labels(self):
|
|
"""Labels of datasets."""
|
|
labels = {}
|
|
|
|
cell_data = self.vtk_data.GetCellData()
|
|
if c := [cell_data.GetArrayName(a) for a in range(cell_data.GetNumberOfArrays())]:
|
|
labels['Cell Data'] = c
|
|
|
|
point_data = self.vtk_data.GetPointData()
|
|
if p := [point_data.GetArrayName(a) for a in range(point_data.GetNumberOfArrays())]:
|
|
labels['Point Data'] = p
|
|
|
|
return labels
|
|
|
|
@staticmethod
|
|
def from_image_data(cells: IntSequence,
|
|
size: FloatSequence,
|
|
origin: FloatSequence = np.zeros(3)) -> 'VTK':
|
|
"""
|
|
Create VTK of type vtk.vtkImageData.
|
|
|
|
This is the common type for grid solver results.
|
|
|
|
Parameters
|
|
----------
|
|
cells : iterable of int, len (3)
|
|
Number of cells along each dimension.
|
|
size : iterable of float, len (3)
|
|
Physical length along each dimension.
|
|
origin : iterable of float, len (3), optional
|
|
Coordinates of grid origin.
|
|
|
|
Returns
|
|
-------
|
|
new : damask.VTK
|
|
VTK-based geometry without nodal or cell data.
|
|
|
|
"""
|
|
vtk_data = vtk.vtkImageData()
|
|
vtk_data.SetDimensions(*(np.array(cells)+1))
|
|
vtk_data.SetOrigin(*(np.array(origin)))
|
|
vtk_data.SetSpacing(*(np.array(size)/np.array(cells)))
|
|
|
|
return VTK(vtk_data)
|
|
|
|
|
|
@staticmethod
|
|
def from_unstructured_grid(nodes: np.ndarray,
|
|
connectivity: np.ndarray,
|
|
cell_type: str) -> 'VTK':
|
|
"""
|
|
Create VTK of type vtk.vtkUnstructuredGrid.
|
|
|
|
This is the common type for mesh solver results.
|
|
|
|
Parameters
|
|
----------
|
|
nodes : numpy.ndarray, 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, TETRA, and HEXAHEDRON.
|
|
|
|
Returns
|
|
-------
|
|
new : damask.VTK
|
|
VTK-based geometry without nodal or cell data.
|
|
|
|
"""
|
|
vtk_nodes = vtk.vtkPoints()
|
|
vtk_nodes.SetData(np_to_vtk(np.ascontiguousarray(nodes)))
|
|
cells = vtk.vtkCellArray()
|
|
cells.SetNumberOfCells(connectivity.shape[0])
|
|
T = np.concatenate((np.ones((connectivity.shape[0],1),dtype=np.int64)*connectivity.shape[1],
|
|
connectivity),axis=1).ravel()
|
|
cells.SetCells(connectivity.shape[0],np_to_vtkIdTypeArray(T,deep=True))
|
|
|
|
vtk_data = vtk.vtkUnstructuredGrid()
|
|
vtk_data.SetPoints(vtk_nodes)
|
|
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)
|
|
|
|
return VTK(vtk_data)
|
|
|
|
|
|
@staticmethod
|
|
def from_poly_data(points: np.ndarray) -> 'VTK':
|
|
"""
|
|
Create VTK of type vtk.polyData.
|
|
|
|
This is the common type for point-wise data.
|
|
|
|
Parameters
|
|
----------
|
|
points : numpy.ndarray, shape (:,3)
|
|
Spatial position of the points.
|
|
|
|
Returns
|
|
-------
|
|
new : damask.VTK
|
|
VTK-based geometry without nodal or cell data.
|
|
|
|
"""
|
|
N = points.shape[0]
|
|
vtk_points = vtk.vtkPoints()
|
|
vtk_points.SetData(np_to_vtk(np.ascontiguousarray(points)))
|
|
|
|
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))
|
|
|
|
vtk_data = vtk.vtkPolyData()
|
|
vtk_data.SetPoints(vtk_points)
|
|
vtk_data.SetVerts(vtk_cells)
|
|
|
|
return VTK(vtk_data)
|
|
|
|
|
|
@staticmethod
|
|
def from_rectilinear_grid(grid: FloatSequence) -> 'VTK':
|
|
"""
|
|
Create VTK of type vtk.vtkRectilinearGrid.
|
|
|
|
Parameters
|
|
----------
|
|
grid : iterables of floats, len (3)
|
|
Grid coordinates along x, y, and z directions.
|
|
|
|
Returns
|
|
-------
|
|
new : damask.VTK
|
|
VTK-based geometry without nodal or cell data.
|
|
|
|
"""
|
|
vtk_data = vtk.vtkRectilinearGrid()
|
|
vtk_data.SetDimensions(*map(len,grid))
|
|
coord = [np_to_vtk(np.array(grid[i]),deep=True) for i in [0,1,2]]
|
|
[coord[i].SetName(n) for i,n in enumerate(['x','y','z'])]
|
|
vtk_data.SetXCoordinates(coord[0])
|
|
vtk_data.SetYCoordinates(coord[1])
|
|
vtk_data.SetZCoordinates(coord[2])
|
|
|
|
return VTK(vtk_data)
|
|
|
|
|
|
@staticmethod
|
|
def load(fname: Union[str, Path],
|
|
dataset_type: Literal['ImageData', 'UnstructuredGrid', 'PolyData', 'RectilinearGrid'] = None) -> 'VTK':
|
|
"""
|
|
Load from VTK file.
|
|
|
|
Parameters
|
|
----------
|
|
fname : str or pathlib.Path
|
|
Filename for reading.
|
|
Valid extensions are .vti, .vtu, .vtp, .vtr, and .vtk.
|
|
dataset_type : {'ImageData', 'UnstructuredGrid', 'PolyData', 'RectilinearGrid'}, optional
|
|
Name of the vtk.vtkDataSet subclass when opening a .vtk file.
|
|
|
|
Returns
|
|
-------
|
|
loaded : damask.VTK
|
|
VTK-based geometry from file.
|
|
|
|
"""
|
|
if not Path(fname).expanduser().is_file(): # vtk has a strange error handling
|
|
raise FileNotFoundError(f'file "{fname}" not found')
|
|
if (ext := Path(fname).suffix) == '.vtk' or dataset_type is not None:
|
|
reader = vtk.vtkGenericDataObjectReader()
|
|
reader.SetFileName(str(Path(fname).expanduser()))
|
|
if dataset_type is None:
|
|
raise TypeError('dataset type for *.vtk file not given')
|
|
elif dataset_type.lower().endswith(('imagedata','image_data')):
|
|
reader.Update()
|
|
vtk_data = reader.GetStructuredPointsOutput()
|
|
elif dataset_type.lower().endswith(('unstructuredgrid','unstructured_grid')):
|
|
reader.Update()
|
|
vtk_data = reader.GetUnstructuredGridOutput()
|
|
elif dataset_type.lower().endswith(('polydata','poly_data')):
|
|
reader.Update()
|
|
vtk_data = reader.GetPolyDataOutput()
|
|
elif dataset_type.lower().endswith(('rectilineargrid','rectilinear_grid')):
|
|
reader.Update()
|
|
vtk_data = reader.GetRectilinearGridOutput()
|
|
else:
|
|
raise TypeError(f'unknown dataset type "{dataset_type}" for vtk file')
|
|
else:
|
|
if ext == '.vti':
|
|
reader = vtk.vtkXMLImageDataReader()
|
|
elif ext == '.vtu':
|
|
reader = vtk.vtkXMLUnstructuredGridReader()
|
|
elif ext == '.vtp':
|
|
reader = vtk.vtkXMLPolyDataReader()
|
|
elif ext == '.vtr':
|
|
reader = vtk.vtkXMLRectilinearGridReader()
|
|
else:
|
|
raise TypeError(f'unknown file extension "{ext}"')
|
|
|
|
reader.SetFileName(str(Path(fname).expanduser()))
|
|
reader.Update()
|
|
vtk_data = reader.GetOutput()
|
|
|
|
return VTK(vtk_data)
|
|
|
|
|
|
@staticmethod
|
|
def _write(writer):
|
|
"""Wrapper for parallel writing."""
|
|
writer.Write()
|
|
|
|
def save(self,
|
|
fname: Union[str, Path],
|
|
parallel: bool = True,
|
|
compress: bool = True):
|
|
"""
|
|
Save as VTK file.
|
|
|
|
Parameters
|
|
----------
|
|
fname : str or pathlib.Path
|
|
Filename for writing.
|
|
parallel : bool, optional
|
|
Write data in parallel background process. Defaults to True.
|
|
compress : bool, optional
|
|
Compress with zlib algorithm. Defaults to True.
|
|
|
|
"""
|
|
if isinstance(self.vtk_data,vtk.vtkImageData):
|
|
writer = vtk.vtkXMLImageDataWriter()
|
|
elif isinstance(self.vtk_data,vtk.vtkUnstructuredGrid):
|
|
writer = vtk.vtkXMLUnstructuredGridWriter()
|
|
elif isinstance(self.vtk_data,vtk.vtkPolyData):
|
|
writer = vtk.vtkXMLPolyDataWriter()
|
|
elif isinstance(self.vtk_data,vtk.vtkRectilinearGrid):
|
|
writer = vtk.vtkXMLRectilinearGridWriter()
|
|
|
|
default_ext = '.'+writer.GetDefaultFileExtension()
|
|
ext = Path(fname).suffix
|
|
writer.SetFileName(str(Path(fname).expanduser())+(default_ext if default_ext != ext else ''))
|
|
|
|
if compress:
|
|
writer.SetCompressorTypeToZLib()
|
|
else:
|
|
writer.SetCompressorTypeToNone()
|
|
writer.SetDataModeToBinary()
|
|
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()
|
|
|
|
|
|
# Check https://blog.kitware.com/ghost-and-blanking-visibility-changes/ for missing data
|
|
def add(self,
|
|
data: Union[np.ndarray, np.ma.MaskedArray, 'Table'],
|
|
label: str = None):
|
|
"""
|
|
Add data to either cells or points.
|
|
|
|
Parameters
|
|
----------
|
|
data : numpy.ndarray, numpy.ma.MaskedArray, or damask.Table
|
|
Data to add. First dimension needs to match either
|
|
number of cells or number of points.
|
|
label : str, optional if data is damask.Table
|
|
Data label.
|
|
|
|
"""
|
|
|
|
def _add_array(vtk_data,
|
|
data: np.ndarray,
|
|
label: str):
|
|
|
|
N_data = data.shape[0]
|
|
|
|
data_ = data.reshape(N_data,-1)\
|
|
.astype(np.single if data.dtype in [np.double,np.longdouble] else data.dtype)
|
|
|
|
if data.dtype.type is np.str_:
|
|
d = vtk.vtkStringArray()
|
|
for s in np.squeeze(data_):
|
|
d.InsertNextValue(s)
|
|
else:
|
|
d = np_to_vtk(data_,deep=True)
|
|
|
|
d.SetName(label)
|
|
|
|
if N_data == vtk_data.GetNumberOfPoints():
|
|
vtk_data.GetPointData().AddArray(d)
|
|
elif N_data == vtk_data.GetNumberOfCells():
|
|
vtk_data.GetCellData().AddArray(d)
|
|
else:
|
|
raise ValueError(f'data count mismatch ({N_data} ≠ {self.N_points} & {self.N_cells})')
|
|
|
|
dup = self.copy()
|
|
if isinstance(data,np.ndarray):
|
|
if label is not None:
|
|
_add_array(dup.vtk_data,
|
|
np.where(data.mask,data.fill_value,data) if isinstance(data,np.ma.MaskedArray) else data,
|
|
label)
|
|
else:
|
|
raise ValueError('no label defined for numpy.ndarray')
|
|
elif isinstance(data,Table):
|
|
for l in data.labels:
|
|
_add_array(dup.vtk_data,data.get(l),l)
|
|
else:
|
|
raise TypeError
|
|
|
|
return dup
|
|
|
|
|
|
def get(self,
|
|
label: str) -> np.ndarray:
|
|
"""
|
|
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.
|
|
|
|
Returns
|
|
-------
|
|
data : numpy.ndarray
|
|
Data stored under the given label.
|
|
|
|
"""
|
|
cell_data = self.vtk_data.GetCellData()
|
|
for a in range(cell_data.GetNumberOfArrays()):
|
|
if cell_data.GetArrayName(a) == label:
|
|
try:
|
|
return vtk_to_np(cell_data.GetArray(a))
|
|
except AttributeError:
|
|
vtk_array = cell_data.GetAbstractArray(a) # string array
|
|
|
|
point_data = self.vtk_data.GetPointData()
|
|
for a in range(point_data.GetNumberOfArrays()):
|
|
if point_data.GetArrayName(a) == label:
|
|
try:
|
|
return vtk_to_np(point_data.GetArray(a))
|
|
except AttributeError:
|
|
vtk_array = point_data.GetAbstractArray(a) # string array
|
|
|
|
try:
|
|
# string array
|
|
return np.array([vtk_array.GetValue(i) for i in range(vtk_array.GetNumberOfValues())]).astype(str)
|
|
except UnboundLocalError:
|
|
raise ValueError(f'array "{label}" not found')
|
|
|
|
|
|
@property
|
|
def comments(self) -> List[str]:
|
|
"""Return the comments."""
|
|
fielddata = self.vtk_data.GetFieldData()
|
|
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 []
|
|
|
|
@comments.setter
|
|
def comments(self,
|
|
comments: Union[str, Sequence[str]]):
|
|
"""
|
|
Set comments.
|
|
|
|
Parameters
|
|
----------
|
|
comments : str or list of str
|
|
Comments.
|
|
|
|
"""
|
|
s = vtk.vtkStringArray()
|
|
s.SetName('comments')
|
|
for c in util.tail_repack(comments,self.comments):
|
|
s.InsertNextValue(c)
|
|
self.vtk_data.GetFieldData().AddArray(s)
|
|
|
|
|
|
def __repr__(self) -> str:
|
|
"""Give short human-readable summary."""
|
|
info = []
|
|
if isinstance(self.vtk_data,vtk.vtkImageData):
|
|
info.append('vtkImageData')
|
|
elif isinstance(self.vtk_data,vtk.vtkUnstructuredGrid):
|
|
info.append('vtkUnstructuredGrid')
|
|
elif isinstance(self.vtk_data,vtk.vtkPolyData):
|
|
info.append('vtkPolyData')
|
|
elif isinstance(self.vtk_data,vtk.vtkRectilinearGrid):
|
|
info.append('vtkRectilinearGrid')
|
|
|
|
for data in ['Cell Data', 'Point Data']:
|
|
if data == 'Cell Data': info.append(f'\n# cells: {self.N_cells}')
|
|
if data == 'Point Data': info.append(f'\n# points: {self.N_points}')
|
|
if data in self.labels:
|
|
info.append(f'with {data}:')
|
|
info += self.labels[data]
|
|
|
|
return util.srepr(info)
|
|
|
|
|
|
def __eq__(self,
|
|
other: object) -> bool:
|
|
"""
|
|
Test equality of other.
|
|
|
|
Parameters
|
|
----------
|
|
other : damask.VTK
|
|
VTK to compare self against.
|
|
|
|
"""
|
|
if not isinstance(other, VTK):
|
|
return NotImplemented
|
|
return self.as_ASCII() == other.as_ASCII()
|
|
|
|
|
|
def as_ASCII(self) -> str:
|
|
"""ASCII representation of the VTK data."""
|
|
writer = vtk.vtkDataSetWriter()
|
|
writer.SetHeader(f'# {util.execution_stamp("VTK")}')
|
|
writer.WriteToOutputStringOn()
|
|
writer.SetInputData(self.vtk_data)
|
|
writer.Write()
|
|
return writer.GetOutputString()
|
|
|
|
|
|
def show(self,
|
|
label: str = None,
|
|
colormap: Colormap = Colormap.from_predefined('cividis')):
|
|
"""
|
|
Render.
|
|
|
|
See http://compilatrix.com/article/vtk-1 for further ideas.
|
|
"""
|
|
try:
|
|
import wx
|
|
_ = wx.App(False) # noqa
|
|
width, height = wx.GetDisplaySize()
|
|
except ImportError:
|
|
try:
|
|
import tkinter
|
|
tk = tkinter.Tk()
|
|
width = tk.winfo_screenwidth()
|
|
height = tk.winfo_screenheight()
|
|
tk.destroy()
|
|
except Exception:
|
|
width = 1024
|
|
height = 768
|
|
|
|
lut = vtk.vtkLookupTable()
|
|
lut.SetNumberOfTableValues(len(colormap.colors))
|
|
for i,c in enumerate(colormap.colors):
|
|
lut.SetTableValue(i,c if len(c)==4 else np.append(c,1.0))
|
|
lut.Build()
|
|
|
|
self.vtk_data.GetCellData().SetActiveScalars(label)
|
|
mapper = vtk.vtkDataSetMapper()
|
|
mapper.SetInputData(self.vtk_data)
|
|
mapper.SetLookupTable(lut)
|
|
mapper.SetScalarRange(self.vtk_data.GetScalarRange())
|
|
|
|
actor = vtk.vtkActor()
|
|
actor.SetMapper(mapper)
|
|
actor.GetProperty().SetColor(230/255,150/255,68/255)
|
|
|
|
ren = vtk.vtkRenderer()
|
|
ren.AddActor(actor)
|
|
if label is None:
|
|
ren.SetBackground(67/255,128/255,208/255)
|
|
else:
|
|
colormap = vtk.vtkScalarBarActor()
|
|
colormap.SetLookupTable(lut)
|
|
colormap.SetTitle(label)
|
|
colormap.SetMaximumWidthInPixels(width//100)
|
|
ren.AddActor2D(colormap)
|
|
ren.SetBackground(0.3,0.3,0.3)
|
|
|
|
window = vtk.vtkRenderWindow()
|
|
window.AddRenderer(ren)
|
|
window.SetSize(width,height)
|
|
window.SetWindowName(util.execution_stamp('VTK','show'))
|
|
|
|
iren = vtk.vtkRenderWindowInteractor()
|
|
iren.SetRenderWindow(window)
|
|
if os.name == 'posix' and 'DISPLAY' not in os.environ:
|
|
print('Found no rendering device')
|
|
else:
|
|
window.Render()
|
|
iren.Start()
|