Source code for febid.libraries.vtk_rendering.VTK_Rendering

"""
Core visualization module
"""
# Default packages
import os, sys, time
import timeit
import datetime
import copy

# Core packages
import numpy as np
import pyvista as pv
from vtk import vtkDataSetAttributes

# Axillary packeges
from tqdm import tqdm

# Local packages
from febid.Structure import Structure



#### Some colormap(cmap) names: viridis, inferno, plasma, coolwarm, cool, Spectral

filename = ''

[docs]class Render: """ Class implementing rendering utilities for visualizing of Numpy data using Pyvista """ def __init__(self, cell_dim:int, font=12, button_size=25): """ :param cell_dim: cell data spacing for VTK objects :param font: button caption font size :param button_size: size of the show on/off button """ self.p = pv.Plotter() # main object that keeps the plot self.cell_dim = cell_dim self.font = font # button caption font size self.size = button_size # button size self.y_pos = 5 # y-position of a button self.x_pos = self.size + 5 # x-position of a button self.meshes_count = 0 self.arrow = None # serves to indicate beam position
[docs] class SetVisibilityCallback: """ Helper callback to keep a reference to the actor being modified. This helps button show and hide plot elements """ def __init__(self, actor): self.actor = actor
[docs] def __call__(self, state): self.actor.SetVisibility(state)
[docs] def show_full_structure(self, structure:Structure, struct=True, deposit=True, precursor=True, surface=True, semi_surface=True, temperature=True, ghosts=True, t=None, sim_time=None, beam=None, cam_pos=None): """ Render and plot all the structure components :param structure: data object :param struct: if True, plot solid structure :param deposit: if True, plot deposit on the surface :param precursor: if True, plot precursor surface density :param surface: if True, color all surface cells :param semi_surface: if True, color all semi_surface cells :param ghosts: if True, color ghost cells :return: """ if struct: self._add_3Darray(structure.deposit, structure.deposit.min(), -0.01, False, opacity=1, clim=[-2,-1], below_color='red', show_edges=False, scalar_name='Structure', button_name='Structure', cmap='binary', n_colors=1, show_scalar_bar=False) if deposit: self._add_3Darray(structure.deposit, 0.00001, 1, False, opacity=1, clim=[0.00001,1], below_color='red', above_color='red', show_edges=True, scalar_name='Surface deposit', button_name='Deposit', cmap='viridis') if precursor: self._add_3Darray(structure.precursor, 0.00001, 1, False, opacity=1, show_edges=True, scalar_name="Surface precursor density", button_name='Precursor', cmap='plasma') if surface: self._add_3Darray(structure.surface_bool, 1, 1, False, opacity=0.7, show_edges=True, scalar_name="Semi surface prec. density", button_name='Surface', color='red', show_scalar_bar=False) if semi_surface: self._add_3Darray(structure.semi_surface_bool, 1, 1, False, opacity=0.7, show_edges=True, scalar_name="Semi surface prec. density", button_name='Semi-surface', color='green', show_scalar_bar=False) if ghosts: self._add_3Darray(structure.ghosts_bool, 1, 1, False, opacity = 0.7, show_edges=True, scalar_name='ghosts', button_name="Ghosts", color='brown', show_scalar_bar=False) if temperature: self._add_3Darray(structure.temperature, 1, opacity=1, scalar_name='temperature', button_name='Max. temperature', cmap='inferno') init_layer = np.count_nonzero(structure.deposit == -2) # substrate layer total_dep_cells = np.count_nonzero(structure.deposit[structure.deposit < 0]) - init_layer # total number of fully deposited cells self.p.add_text(f'Cells: {total_dep_cells} \n' # showing total number of deposited cells f'Height: {int(np.nonzero(structure.deposit)[0].max() * structure.cell_dimension)} nm \n' # showing current height of the structure f'Deposited volume: {int(total_dep_cells + structure.deposit[structure.deposit>0].sum()) * structure.cell_dimension**3} nm^3\n', position='upper_right', font_size=self.font) text = '' if t: text += f'Time: {t} \n' if sim_time: text += f'Simulation time: {sim_time:.7f} s \n' self.p.add_text(text, position='upper_left', font_size=self.font) if beam is not None: x_pos, y_pos = beam x, y = int(x_pos/self.cell_dim), int(y_pos/self.cell_dim) max_z = structure.deposit[:, y, x].nonzero()[0].max() start = np.array([0, 0, 100]).reshape(1, 3) # position of the center of the arrow end = np.array([0, 0, -100]).reshape(1, 3) # direction and resulting size self.arrow = self.p.add_arrows(start, end, color='tomato') self.arrow.SetPosition(x_pos, y_pos, max_z * self.cell_dim + 30) # relative to the initial position if cam_pos is None: cam_pos = [(463.14450307610286, 271.1171723376318, 156.56895424388603), (225.90027381807235, 164.9577775224395, 71.42188811921902), (-0.27787912231751677, -0.1411181984824172, 0.950194110399093)] return self.show(cam_pos=cam_pos)
def show_mc_result(self, grid=None, pe_traj=None, surface_flux=None, se_traj=None, heat_t=None, heat_pe=None, heat_se=None, t=None, sim_time=None, beam=None, cam_pos=None, interactive=True): pe_traj = copy.deepcopy(pe_traj) se_traj = copy.deepcopy(se_traj) if grid is not None: self._add_3Darray(grid, -2, -0.01, opacity=0.9, show_edges=True, scalar_name='Structure', button_name='Structure', color='white') if pe_traj is not None: self._add_trajectory(pe_traj[:,0], pe_traj[:,1], 0.2, step=1, scalar_name='PE Energy, keV', button_name='PEs', cmap='viridis') if surface_flux is not None: self._add_3Darray(surface_flux, 1, opacity=1, show_edges=False, scalar_name='SE Flux, 1/(nm^2*s)', button_name="SE surface flux", cmap='plasma', log_scale=True) if se_traj is not None: max_trajes = 4000 step = int(se_traj.shape[0]/max_trajes)+1 self._add_trajectory(se_traj, radius=0.1, step=step, button_name='SEs', cmap=None, color='red') if heat_t is not None: self._add_3Darray(heat_t, 1, opacity=0.7, scalar_name='Total energy transfered to heat, eV', button_name='Total heat', cmap='coolwarm', log_scale=True) if heat_pe is not None: self._add_3Darray(heat_pe, 1, opacity=0.7, scalar_name='Total PE energy transferred to heat, eV', button_name='PE heat', cmap='coolwarm', log_scale=True) if heat_se is not None: self._add_3Darray(heat_se, 1, opacity=0.7, scalar_name='Total SE energy transferred to heat, eV', button_name='SE heat', cmap='coolwarm', log_scale=True) text = '' if t: text += f'Time: {t} \n' if sim_time: text += f'Simulation time: {sim_time:.7f} s \n' if beam is not None: text += f'Beam position: {beam[0], beam[1]}' self.p.add_text(text, position='upper_left', font_size=self.font) if cam_pos is None: cam_pos = [(463.14450307610286, 271.1171723376318, 156.56895424388603), (225.90027381807235, 164.9577775224395, 71.42188811921902), (-0.27787912231751677, -0.1411181984824172, 0.950194110399093)] cam_pos = self.show(cam_pos=cam_pos, interactive_update=interactive) return cam_pos def _add_trajectory(self, traj, energies=None, radius=0.7, step=1, scalar_name='scalars_t', button_name='1', color='', cmap='plasma'): """ Adds trajectories to the Pyvista plot :param traj: collection of trajectories, that are represented as a sequence of points :param energies: collection of corresponding energies :param radius: line thickness :param step: a number of trajectories to skip :param scalar_name: name of the scalar bar :param button_name: button caption :param color: color of the trajectories :param cmap: colormap for the trajectories :return: adds PolyData() to Plotter() """ if energies is None: energies = [] obj = self._render_trajectories(traj=traj, energies=energies, radius=radius, step=step, name=scalar_name) self.__prepare_obj(obj, button_name, cmap, color) def _add_3Darray(self, arr, lower_t=None, upper_t=None, exclude_zeros=False, opacity=0.5, clim=None, below_color=None, above_color=None, show_edges=None, nan_opacity=None, scalar_name='scalars_s', button_name='NoName', color=None, show_scalar_bar=True, cmap=None, n_colors=256, log_scale=False, invert=False, texture=None): """ Adds 3D structure from a Numpy array to the Pyvista plot :param arr: numpy array :param lower_t: lower threshold of values :param upper_t: upper threshold of values :param scalar_name: name of the scalar bar :param button_name: button caption :param color: color of the trajectories :param cmap: colormap for the trajectories :return: adds PolyData() to Plotter() """ if nan_opacity == None: nan_opacity = opacity self.obj = self._render_3Darray(arr=arr, lower_t=lower_t, upper_t=upper_t, exclude_zeros=exclude_zeros, name=scalar_name, invert=invert) self.__prepare_obj(self.obj, button_name, cmap, color, show_scalar_bar, clim, below_color, above_color, n_colors, log_scale=log_scale, opacity=opacity, nan_opacity=nan_opacity, show_edges=show_edges, texture=texture) def __prepare_obj(self, obj, name, cmap, color, show_scalar_bar=True, clim=None, below_color=None, above_color=None, n_colors=256, log_scale=False, opacity=0.5, nan_opacity=0.5, show_edges=None, texture=None): while True: try: if not cmap: obj_a = self.p.add_mesh(obj, style='surface', opacity=opacity, nan_opacity=nan_opacity, clim=clim, below_color=below_color, above_color=above_color, name=name, label='Structure', log_scale=log_scale, show_scalar_bar=show_scalar_bar, n_colors=n_colors, color=color, lighting=True, show_edges=show_edges, texture=texture, render=False) # adding data to the plot break else: obj_a = self.p.add_mesh(obj, style='surface', opacity=opacity, use_transparency=False, nan_opacity=nan_opacity, clim=clim, below_color=below_color, above_color=above_color, name=name, label='Structure', log_scale=log_scale, show_scalar_bar=show_scalar_bar, cmap=cmap, n_colors=n_colors, lighting=True, show_edges=show_edges, texture=texture, render=False) break except Exception as e: print(f'Error:{e.args}') return self.p.add_text(name, font_size=self.font, position=(self.x_pos + 5, self.y_pos), name=name+'_caption') # captioning button obj_aa = self.SetVisibilityCallback(obj_a) self.p.add_checkbox_button_widget(obj_aa, value=True, position=(5, self.y_pos), size=self.size, color_on='blue') # adding button self.y_pos += self.size self.meshes_count += 1 def _render_3Darray(self, arr, lower_t=None, upper_t=None, exclude_zeros=False, name='scalars_s', invert=False ): """ Renders a 3D numpy array and trimms values :param arr: array :param lower_t: lower cutoff threshold :param upper_t: upper cutoff threshold :return: pyvista.PolyData object """ # if upper_t is None: upper_t = arr.max() # if lower_t is None: lower_t = arr.min() grid = numpy_to_vtk(arr, self.cell_dim, data_name=name, grid=None) if exclude_zeros: grid.remove_cells((arr==0).flatten()) if upper_t is not None or lower_t is not None: if upper_t is None: upper_t = arr.max() if lower_t is None: lower_t = arr.min() grid = grid.threshold([lower_t,upper_t], continuous=True, invert=invert) # trimming return grid def _render_trajectories(self, traj, energies, radius=0.7, step=1, name='scalars_t'): """ Renders mapped trajectories as splines with the given thickness :param traj: collection of trajectories :param energies: collection of energies :param radius: line width :return: pyvista.PolyData object """ mesh = pv.PolyData() # If energies are provided, they are gonna be used as scalars to color trajectories start = timeit.default_timer() if len(energies) != 0: print('Rendering PEs...', end='') for i in tqdm(range(0, len(traj), step)): # # mesh = mesh + self.__render_trajectory(traj[i], energies[i], radius, name) # mesh[name] = energies mesh_d = self.__render_trajectories(np.asarray(traj[i]), 'line') mesh_d[name] = np.asarray(energies[i]) mesh += mesh_d print(f'took {timeit.default_timer()-start}') else: print('Rendering SEs...', end='') # for i in tqdm(range(0, len(traj), step)): # mesh = mesh + self.__render_trajectory(traj[i], 0, radius, name) traj = traj.reshape(traj.shape[0] * 2, 3) mesh = self.__render_trajectories(traj, 'seg') print(f'took {timeit.default_timer() - start}') return mesh.tube(radius=radius) # it is important for color mapping to create tubes after all trajectories are added def __render_trajectories(self, traj, kind='line'): """ Turn a collection of points to line/lines. kind: 'line' to create a line connecting all the points 'seg' to create separate segments (2 points each) :param traj: collection of points :param kind: type of connection :return: """ if kind not in ['line', 'seg']: raise RuntimeWarning('Wrong \'type\' argument in Render.__render_trajectories. Method accepts \'line\' or \'seg\'') traj[:, 0], traj[:,2] = traj[:, 2], traj[:,0].copy() # swapping x and z coordinates if kind == 'line': mesh = pv.lines_from_points(np.asfortranarray(traj)) if kind == 'seg': mesh = pv.line_segments_from_points(traj) return mesh def __render_trajectory(self, traj, energies=0, radius=0.7, name='scalars'): """ Renders a single trajectory with the given thickness :param traj: collection of points :param energies: energies for every point :param radius: line width :return: pyvista.PolyData object """ points = np.asarray([[t[2], t[1], t[0]] for t in traj]) # coordinates are provided in a numpy array manner [z,y,x], but vista implements [x,y,z] mesh = pv.PolyData() mesh.points = points # assigning points between segments line = np.arange(0, len(points), dtype=np.int_) line = np.insert(line, 0, len(points)) mesh.lines = line # assigning lines that connect the points if energies: mesh[name] = np.asarray(energies) # assigning energies for every point return mesh #.tube(radius=radius) # making line thicker def update_mask(self, mask): index = np.zeros_like(mask, dtype=np.uint8) index[mask == 0] = vtkDataSetAttributes.HIDDENCELL last_scalars = self.p.mesh.array_names[0] self.p.mesh.cell_data[vtkDataSetAttributes.GhostArrayName()] = index.ravel() self.p.mesh.set_active_scalars(last_scalars)
[docs] def save_3Darray(self, filename, arr, data_name='scalar'): """ Dump a Numpy array to a vtk file with a specified name and creation date :param filename: distinct name of the file :param arr: array to save :param data_name: name of the data to include in the vtk dataset :return: """ grid = numpy_to_vtk(arr, self.cell_dim, data_name) print("File is saved in the same directory with current python script. Current time is appended") grid.save(f'{sys.path[0]}{os.sep}{filename}{time.strftime("%H:%M:%S", time.localtime())}.vtk')
[docs] def show(self, screenshot=False, show_grid=True, keep_plot=False, interactive_update=False, cam_pos=None): """ Shows plotting scene :param screenshot: if True, a screenshot of the scene will be saved upon showing :param show_grid: indicates axes and scales :param keep_plot: if True, creates a copy of current Plotter before showing :param interactive_update: if True, code execution does not stop while scene window is opened :param cam_pos: camera view :return: current camera view """ if show_grid: self.p.show_grid() if keep_plot: p1 = copy.deepcopy(self.p) camera_pos = self.p.show(screenshot=screenshot, interactive_update=interactive_update, cpos=cam_pos, return_cpos=True) if keep_plot: self.p = copy.deepcopy(p1) self.y_pos = 5 return camera_pos
[docs] def update(self, time=1, force_redraw=False): """ Update the plot :param time: minimum time before each subsequent update :param force_redraw: redraw the plot immediately :return: """ self.p.update(stime=time, force_redraw=force_redraw) self.y_pos -= self.size*self.meshes_count self.meshes_count = 0
# self.p.clear()
[docs]def read_field_data(vtk_obj): """ Read run time, simulation time and beam position from vtk-file. :param vtk_obj: VTK-object (UniformGrid) :return: """ t = vtk_obj.field_data.get('time', None) sim_time = vtk_obj.field_data.get('simulation_time', None) beam_position = vtk_obj.field_data.get('beam_position', None) if t: t = t[0] if sim_time: sim_time = sim_time[0] if beam_position is not None: beam_position = beam_position[0] return t, sim_time, beam_position
[docs]def numpy_to_vtk(arr, cell_dim, data_name='scalar', grid=None, unstructured=False): """ Convert numpy array to a VTK-datastructure (UniformGrid or UnstructuredGrid). If grid is provided, add new dataset to that grid. :param arr: numpy array :param cell_dim: array cell (cubic) edge length :param data_name: name of data :param grid: existing UniformGrid :param unstructured: if True, return an UnstructuredGrid :return: """ if not grid: grid = pv.UniformGrid() grid.dimensions = np.asarray([arr.shape[2], arr.shape[1], arr.shape[0]]) + 1 # creating a grid with the size of the array grid.spacing = (cell_dim, cell_dim, cell_dim) # assigning dimensions of a cell grid_given = False else: grid_given = True grid.cell_data[data_name] = arr.ravel() # writing values if unstructured and not grid_given: grid = grid.cast_to_unstructured_grid() return grid
[docs]def save_deposited_structure(structure, sim_t=None, t=None, beam_position=None, filename=None): """ Save current deposition result to a vtk file. If filename does not contain path, saves to the current directory. :param structure: an instance of the current state of the process :param sim_t: simulation time, s :param t: run time :param beam_position: (x,y) current position of the beam :param filename: full file name :return: """ cell_dim = structure.cell_dimension # Accumulating data from the array in a VTK datastructure vtk_obj = numpy_to_vtk(structure.deposit, cell_dim, 'deposit', unstructured=False) vtk_obj = numpy_to_vtk(structure.precursor, cell_dim, data_name='precursor_density', grid=vtk_obj) vtk_obj = numpy_to_vtk(structure.surface_bool, cell_dim, data_name='surface_bool', grid=vtk_obj) vtk_obj = numpy_to_vtk(structure.semi_surface_bool, cell_dim, data_name='semi_surface_bool', grid=vtk_obj) vtk_obj = numpy_to_vtk(structure.ghosts_bool, cell_dim, data_name='ghosts_bool', grid=vtk_obj) vtk_obj = numpy_to_vtk(structure.temperature, cell_dim, data_name='temperature', grid=vtk_obj) # Attaching times and beam position vtk_obj.field_data['date'] = [datetime.datetime.now()] vtk_obj.field_data['time'] = [str(datetime.timedelta(seconds=int(t)))] vtk_obj.field_data['simulation_time'] = [sim_t] vtk_obj.field_data['beam_position'] = [beam_position] if filename == None: filename = "Structure" vtk_obj.save(f'{filename}_{time.strftime("%H:%M:%S", time.localtime())}.vtk')
[docs]def export_obj(structure, filename=None): """ Export deposited structure as an .obj file :param structure: Structure class instance, must have 'deposit' array and 'cell_dimension' value :param filename: full path with file name :return: """ grid = numpy_to_vtk(structure.deposit, structure.cell_dimension, 'Deposit', None, True) grid = grid.threshold([-2,-0.001], continuous=True) p = pv.Plotter() p.add_mesh(grid) p.export_obj(filename) return 1
if __name__ == '__main__': raise NotImplementedError