DAMASK_EICMD/python/damask/_vtk.py

610 lines
20 KiB
Python
Raw Permalink 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
from typing import Union, Literal, List, Sequence
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
from vtk.util.numpy_support import vtk_to_numpy as vtk_to_np
2022-01-22 12:35:49 +05:30
from ._typehints import FloatSequence, IntSequence
from . import util
2020-03-13 00:22:33 +05:30
from . import Table
from . import Colormap
2020-03-13 00:22:33 +05:30
class VTK:
"""
Spatial visualization (and potentially manipulation).
High-level interface to VTK.
"""
def __init__(self,
vtk_data: vtk.vtkDataSet):
"""
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.
2022-01-13 03:43:38 +05:30
Valid types are vtk.vtkImageData, vtk.vtkUnstructuredGrid,
vtk.vtkPolyData, and vtk.vtkRectilinearGrid.
"""
2020-08-25 13:26:24 +05:30
self.vtk_data = vtk_data
2022-02-23 19:07:23 +05:30
def __repr__(self) -> str:
"""
Return repr(self).
Give short human-readable summary.
"""
2022-02-23 19:07:23 +05:30
info = [self.vtk_data.__vtkname__]
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 += [f' - {l}' for l in self.labels[data]]
return util.srepr(info)
def __eq__(self,
other: object) -> bool:
"""
Return self==other.
Test equality of other.
2022-02-23 19:07:23 +05:30
Parameters
----------
other : damask.VTK
2022-02-23 19:42:55 +05:30
VTK to check for equality.
2022-02-23 19:07:23 +05:30
"""
2022-02-23 19:42:55 +05:30
if not isinstance(other, VTK):
return NotImplemented
return self.as_ASCII() == other.as_ASCII()
2022-02-23 19:07:23 +05:30
def copy(self):
if isinstance(self.vtk_data,vtk.vtkImageData):
dup = vtk.vtkImageData()
elif isinstance(self.vtk_data,vtk.vtkUnstructuredGrid):
2022-02-21 16:47:00 +05:30
dup = vtk.vtkUnstructuredGrid()
elif isinstance(self.vtk_data,vtk.vtkPolyData):
2022-02-21 16:47:00 +05:30
dup = vtk.vtkPolyData()
elif isinstance(self.vtk_data,vtk.vtkRectilinearGrid):
2022-02-21 16:47:00 +05:30
dup = vtk.vtkRectilinearGrid()
else:
raise TypeError
dup.DeepCopy(self.vtk_data)
return VTK(dup)
2022-02-23 19:07:23 +05:30
@property
def comments(self) -> List[str]:
"""Return the comments."""
2022-03-07 15:08:56 +05:30
field_data = self.vtk_data.GetFieldData()
for a in range(field_data.GetNumberOfArrays()):
if field_data.GetArrayName(a) == 'comments':
comments = field_data.GetAbstractArray(a)
2022-02-23 19:07:23 +05:30
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 sequence of str
2022-02-23 19:07:23 +05:30
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)
2022-02-16 03:08:02 +05:30
@property
def N_points(self) -> int:
"""Number of points in vtkdata."""
return self.vtk_data.GetNumberOfPoints()
2022-02-23 19:07:23 +05:30
2022-02-16 03:08:02 +05:30
@property
def N_cells(self) -> int:
"""Number of cells in vtkdata."""
return self.vtk_data.GetNumberOfCells()
2022-02-23 19:07:23 +05:30
2022-02-22 12:38:41 +05:30
@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
2022-02-16 03:08:02 +05:30
2022-02-23 19:07:23 +05:30
@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 : sequence of int, len (3)
Number of cells along each dimension.
size : sequence of float, len (3)
Physical length along each dimension.
origin : sequence 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)))
2022-01-22 12:35:49 +05:30
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.
2021-04-24 18:17:52 +05:30
This is the common type for mesh solver results.
Parameters
----------
2022-02-18 05:19:45 +05:30
nodes : numpy.ndarray, shape (:,3)
Spatial position of the nodes.
connectivity : numpy.ndarray of np.dtype = np.int64
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.
2021-04-24 18:17:52 +05:30
Returns
-------
new : damask.VTK
VTK-based geometry without nodal or cell data.
"""
vtk_nodes = vtk.vtkPoints()
2021-09-02 11:34:09 +05:30
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],
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
2022-01-22 12:35:49 +05:30
def from_poly_data(points: np.ndarray) -> 'VTK':
"""
Create VTK of type vtk.polyData.
This is the common type for point-wise data.
Parameters
----------
2022-02-18 05:19:45 +05:30
points : numpy.ndarray, shape (:,3)
Spatial position of the points.
2021-04-24 18:17:52 +05:30
Returns
-------
new : damask.VTK
VTK-based geometry without nodal or cell data.
"""
N = points.shape[0]
2020-03-22 21:33:28 +05:30
vtk_points = vtk.vtkPoints()
2021-09-02 11:34:09 +05:30
vtk_points.SetData(np_to_vtk(np.ascontiguousarray(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
2022-02-18 22:15:21 +05:30
@staticmethod
def from_rectilinear_grid(grid: FloatSequence) -> 'VTK':
"""
Create VTK of type vtk.vtkRectilinearGrid.
Parameters
----------
grid : sequence of sequences of floats, len (3)
2022-02-18 22:15:21 +05:30
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)
2020-03-12 04:24:36 +05:30
@staticmethod
2022-01-17 19:30:25 +05:30
def load(fname: Union[str, Path],
2022-02-18 22:15:21 +05:30
dataset_type: Literal['ImageData', 'UnstructuredGrid', 'PolyData', 'RectilinearGrid'] = None) -> 'VTK':
"""
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
2022-01-22 12:35:49 +05:30
Filename for reading.
2022-02-18 22:15:21 +05:30
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.
2021-04-24 18:17:52 +05:30
Returns
-------
loaded : damask.VTK
VTK-based geometry from file.
"""
2022-02-21 16:47:00 +05:30
if not Path(fname).expanduser().is_file(): # vtk has a strange error handling
2022-02-22 21:26:12 +05:30
raise FileNotFoundError(f'file "{fname}" not found')
if (ext := Path(fname).suffix) == '.vtk' or dataset_type is not None:
2020-03-12 04:24:36 +05:30
reader = vtk.vtkGenericDataObjectReader()
reader.SetFileName(str(Path(fname).expanduser()))
2020-03-21 15:37:21 +05:30
if dataset_type is None:
2022-02-22 21:26:12 +05:30
raise TypeError('dataset type for *.vtk file not given')
elif dataset_type.lower().endswith(('imagedata','image_data')):
reader.Update()
2021-06-16 00:53:10 +05:30
vtk_data = reader.GetStructuredPointsOutput()
elif dataset_type.lower().endswith(('unstructuredgrid','unstructured_grid')):
reader.Update()
2020-08-25 13:26:24 +05:30
vtk_data = reader.GetUnstructuredGridOutput()
elif dataset_type.lower().endswith(('polydata','poly_data')):
reader.Update()
2020-08-25 13:26:24 +05:30
vtk_data = reader.GetPolyDataOutput()
2022-02-18 22:15:21 +05:30
elif dataset_type.lower().endswith(('rectilineargrid','rectilinear_grid')):
reader.Update()
vtk_data = reader.GetRectilinearGridOutput()
2020-03-12 04:24:36 +05:30
else:
2022-02-22 21:26: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 == '.vti':
reader = vtk.vtkXMLImageDataReader()
2020-03-12 04:24:36 +05:30
elif ext == '.vtu':
reader = vtk.vtkXMLUnstructuredGridReader()
elif ext == '.vtp':
reader = vtk.vtkXMLPolyDataReader()
2022-02-18 22:15:21 +05:30
elif ext == '.vtr':
reader = vtk.vtkXMLRectilinearGridReader()
2020-03-12 04:24:36 +05:30
else:
2022-02-22 21:26:12 +05:30
raise TypeError(f'unknown file extension "{ext}"')
2020-03-12 04:24:36 +05:30
reader.SetFileName(str(Path(fname).expanduser()))
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()
2022-02-23 19:07:23 +05:30
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 save(self,
fname: Union[str, Path],
parallel: bool = True,
compress: bool = 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.
2022-01-13 03:43:38 +05:30
parallel : bool, 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.
"""
if isinstance(self.vtk_data,vtk.vtkImageData):
writer = vtk.vtkXMLImageDataWriter()
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()
2022-02-18 22:15:21 +05:30
elif isinstance(self.vtk_data,vtk.vtkRectilinearGrid):
writer = vtk.vtkXMLRectilinearGridWriter()
2020-03-11 12:02:03 +05:30
2020-11-04 22:38:04 +05:30
default_ext = '.'+writer.GetDefaultFileExtension()
2020-06-03 14:13:07 +05:30
ext = Path(fname).suffix
writer.SetFileName(str(Path(fname).expanduser())+(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
def set(self,
label: str = None,
data: Union[np.ndarray, np.ma.MaskedArray] = None,
2022-05-13 09:32:22 +05:30
info: str = None,
*,
table: 'Table' = None):
"""
2022-05-13 09:32:22 +05:30
Add new or replace existing point or cell data.
Data can either be a numpy.array, which requires a corresponding label,
or a damask.Table.
Parameters
----------
label : str, optional
Label of data array.
data : numpy.ndarray or numpy.ma.MaskedArray, optional
Data to add or replace. First array dimension needs to match either
number of cells or number of points.
2022-05-13 09:32:22 +05:30
info : str, optional
Human-readable information about the data.
table: damask.Table, optional
Data to add or replace. Each table label is individually considered.
Number of rows needs to match either number of cells or number of points.
Notes
-----
If the number of cells equals the number of points, the data is added to both.
"""
2020-03-12 04:24:36 +05:30
2022-02-21 16:47:00 +05:30
def _add_array(vtk_data,
label: str,
data: np.ndarray):
2022-02-16 03:08:02 +05:30
N_p,N_c = vtk_data.GetNumberOfPoints(),vtk_data.GetNumberOfCells()
if (N_data := data.shape[0]) not in [N_p,N_c]:
raise ValueError(f'data count mismatch ({N_data}{N_p} & {N_c})')
2022-02-23 19:07:23 +05:30
data_ = data.reshape(N_data,-1) \
2022-02-16 03:08:02 +05:30
.astype(np.single if data.dtype in [np.double,np.longdouble] else data.dtype)
2022-02-16 03:08:02 +05:30
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)
2020-03-12 04:24:36 +05:30
d.SetName(label)
if N_data == N_p:
2022-02-21 16:47:00 +05:30
vtk_data.GetPointData().AddArray(d)
if N_data == N_c:
2022-02-21 16:47:00 +05:30
vtk_data.GetCellData().AddArray(d)
2022-02-16 03:08:02 +05:30
if data is None and table is None:
raise KeyError('no data given')
if data is not None and table is not None:
raise KeyError('cannot use both, data and table')
2022-02-21 16:47:00 +05:30
dup = self.copy()
2022-02-16 03:08:02 +05:30
if isinstance(data,np.ndarray):
if label is not None:
2022-02-21 16:47:00 +05:30
_add_array(dup.vtk_data,
label,
np.where(data.mask,data.fill_value,data) if isinstance(data,np.ma.MaskedArray) else data)
2022-05-18 01:42:00 +05:30
if info is not None: dup.comments += f'{label}: {info}'
2022-02-16 03:08:02 +05:30
else:
raise ValueError('no label defined for data')
elif isinstance(table,Table):
for l in table.labels:
_add_array(dup.vtk_data,l,table.get(l))
2022-05-18 01:42:00 +05:30
if info is not None: dup.comments += f'{l}: {info}'
2020-03-19 13:33:38 +05:30
else:
raise TypeError
2022-05-13 09:32:22 +05:30
2022-02-21 16:47:00 +05:30
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.
2021-04-24 18:17:52 +05:30
Returns
-------
data : numpy.ndarray
Data stored under the given 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:
try:
return vtk_to_np(cell_data.GetArray(a))
except AttributeError:
vtk_array = cell_data.GetAbstractArray(a) # string array
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:
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:
2022-03-07 15:58:17 +05:30
raise KeyError(f'array "{label}" not found')
2022-02-21 15:49:53 +05:30
def show(self,
label: str = None,
colormap: Union[Colormap, str] = 'cividis'):
2020-03-12 18:15:47 +05:30
"""
Render.
Parameters
----------
label : str, optional
Label of the dataset to show.
colormap : damask.Colormap or str, optional
Colormap for visualization of dataset. Defaults to 'cividis'.
Notes
-----
See http://compilatrix.com/article/vtk-1 for further ideas.
2020-03-12 18:15:47 +05:30
"""
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()
2022-01-29 23:40:12 +05:30
except Exception:
2021-02-12 02:26:53 +05:30
width = 1024
height = 768
lut = vtk.vtkLookupTable()
colormap_ = Colormap.from_predefined(colormap) if isinstance(colormap,str) else \
colormap
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)
2020-03-12 18:15:47 +05:30
mapper = vtk.vtkDataSetMapper()
2020-08-25 13:26:24 +05:30
mapper.SetInputData(self.vtk_data)
mapper.SetLookupTable(lut)
mapper.SetScalarRange(self.vtk_data.GetScalarRange())
2020-03-12 18:15:47 +05:30
actor = vtk.vtkActor()
actor.SetMapper(mapper)
actor.GetProperty().SetColor(230/255,150/255,68/255)
2020-03-12 18:15:47 +05:30
ren = vtk.vtkRenderer()
ren.AddActor(actor)
2022-02-21 15:49:53 +05:30
if label is None:
ren.SetBackground(67/255,128/255,208/255)
else:
colormap_vtk = vtk.vtkScalarBarActor()
colormap_vtk.SetLookupTable(lut)
colormap_vtk.SetTitle(label)
colormap_vtk.SetMaximumWidthInPixels(width//100)
ren.AddActor2D(colormap_vtk)
2022-02-21 15:49:53 +05:30
ren.SetBackground(0.3,0.3,0.3)
2020-03-12 18:15:47 +05:30
2020-03-22 21:33:28 +05:30
window = vtk.vtkRenderWindow()
window.AddRenderer(ren)
2021-02-12 02:26:53 +05:30
window.SetSize(width,height)
window.SetWindowName(util.execution_stamp('VTK','show'))
2020-03-12 18:15:47 +05:30
iren = vtk.vtkRenderWindowInteractor()
2020-03-22 21:33:28 +05:30
iren.SetRenderWindow(window)
if os.name == 'posix' and 'DISPLAY' not in os.environ:
print('Found no rendering device')
else:
window.Render()
iren.Start()