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()