Source code for pyprocar.core.surface

__author__ = "Pedram Tavadze and Logan Lang"
__maintainer__ = "Pedram Tavadze and Logan Lang"
__email__ = "petavazohi@mail.wvu.edu, lllang@mix.wvu.edu"
__date__ = "March 31, 2020"

from shutil import which

import pyvista
# import trimesh
import numpy as np
from matplotlib import cm
from matplotlib import colors as mpcolors


# TODO add python typing to all of the functions
# TODO add trimesh 
# TODO what is the point of self.test?
# TODO uncomment self.polydata and self.mesh
# TODO add an __str__ method
# TODO change boolean add method to __add__ method and get rid of try and expect

[docs] class Surface(pyvista.PolyData): """ Surface is a class that holds information about a surface. This class inherits from the pyvista.PolyData class. To create a surface the minimum requirements are verts and faces Parameters ---------- verts : list of float (nverts,3) The list of verticies that create the surface. faces : list of integers (nfaces,3) The default is None. The list of connectivity between verts that create the surface. face_normals : list of float (nfaces,3) The list of normal vectors to each face. vert_normals : list of float (nverts,3) The list of normal vectors to each vertex. face_colors : list of tuples floats (nfaces,3) The list of colors of each face. **example**:``face_colors=[(1,0,0),(1,0.5,0),...,(1,0,0)]`` vert_colors : list of tuples floats (nfaces,3) The list of colors of each vertex. vectors : list of floats (nfaces,3) The list of vectors one wants to attach to the surface(glyphs) Only useful in pyvista objects scalars : list of floats (nfaces,) The list of scalars for each face. This can represent the color using a color map """
[docs] def __init__( self, verts:np.ndarray=None, faces:np.ndarray=None, face_normals:np.ndarray=None, vert_normals:np.ndarray=None, face_colors:np.ndarray=None, vert_colors:np.ndarray=None, vectors:np.ndarray=None, scalars:np.ndarray=None, ): super().__init__( var_inp = verts, faces = np.array(faces)) # print(faces) # super().__init__( var_inp = verts, faces = np.array(faces)) # self.face_normals = face_normals # self.vert_normals = vert_normals self.face_colors = face_colors self.vert_colors = vert_colors # self.vectors = vectors self.scalars = scalars self.trimesh_obj = None # if self.verts is not None and self.faces is not None: # self._create_trimesh() # if self.face_normals is None: # if self.pyvista_obj.face_normals is not None: # self.face_normals = self.pyvista_obj.face_normals # if self.vert_normals is None: # self.vert_normals=self.pyvista_obj.point_normals return None
@property def centers(self): """ Centers of faces Returns ------- centers : list of floats (n,3) A list of centers of faces. """ return self.cell_centers().points @property def faces_array(self): """ The faces listed in a list of list which contains the faces. Returns ------- new_faces : list A list of faces """ new_faces = [] face = [] count = 0 for iverts_in_face,verts_in_face in enumerate(self.faces): if iverts_in_face == 0: num_verts = verts_in_face face = [num_verts] else: if count == num_verts: count = 0 new_faces.append(face) num_verts = verts_in_face face = [num_verts] elif iverts_in_face == len(self.faces)-1: face.append(verts_in_face) new_faces.append(face) else: count += 1 face.append(verts_in_face) return new_faces # def _create_trimesh(self): # """ # creates a trimesh object # """ # # if np.any(np.array([len(x) for x in self.faces]) > 3): # # faces = [] # # for i in range(0, len(self.pyvista_obj.triangulate().faces), 4): # # point_1 = self.pyvista_obj.triangulate().faces[i + 1] # # point_2 = self.pyvista_obj.triangulate().faces[i + 2] # # point_3 = self.pyvista_obj.triangulate().faces[i + 3] # # faces.append([point_1, point_2, point_3]) # # self.trimesh_obj = trimesh.Trimesh(vertices=self.verts, faces=faces) # # else: # self.trimesh_obj = trimesh.Trimesh(vertices=self.points, faces=self.faces)
[docs] def set_scalars( self, scalars: np.ndarray, scalar_name: str="scalars" ): """ Sets/Updates the scalars of the surface. Scalars represent a color using a color map. Parameters ---------- scalars : list Scalars should be the same size as the number of faces. scalar_name : str The name of the scalar """ self.scalars = scalars self[scalar_name] = self.scalars # self.cell_data[scalar_name] = self.scalars # self.point_data[scalar_name] = self.scalars self.set_active_scalars(scalar_name,preference='cell')
[docs] def set_vectors(self, vectors_X:np.ndarray, vectors_Y:np.ndarray, vectors_Z:np.ndarray, vectors_name: str="vectors"): """Sets/Updates the vectors of the surface. Parameters ---------- vectors_X : np.ndarray The x values of the vector vectors_Y : np.ndarray The y values of the vector vectors_Z : np.ndarray The z values of the vector vectors_name : str, optional The name of the vector, by default "vectors" """ def mag(vectors): return np.array([(vector[0]**2 + vector[1]**2 + vector[2]**2)**0.5 for vector in vectors]) vectors = np.vstack([vectors_X, vectors_Y, vectors_Z]).T self.point_data[vectors_name] = vectors self.point_data[vectors_name + "_magnitude"] = mag(vectors) # self.set_active_scalars('vectors') return None
[docs] def set_color_with_cmap(self, cmap:str="viridis", vmin:float=None, vmax:float=None): """ Sets colors for the trimesh object using the color map provided Parameters ---------- cmap : string The colormap. The default is 'viridis'. vmin : float, optional The minimum normalizing value. The default is None. vmax : float, optional The maximum normalizing value. The default is None. """ if vmin is None: vmin = min(self.scalars) if vmax is None: vmax = max(self.scalars) norm = mpcolors.Normalize(vmin=vmin, vmax=vmax) cmap = cm.get_cmap(cmap) colors = np.array([cmap(norm(x)) for x in self.scalars]).reshape(-1, 4) self.face_colors = colors # This next line will make all the surfaces double sided if you want # only show one side comment the next line if len(self.trimesh_obj.faces) == self.n_faces: self.trimesh_obj.faces = np.vstack( (self.trimesh_obj.faces, np.fliplr(self.trimesh_obj.faces)) ) if len(self.trimesh_obj.faces) == self.n_faces: self.trimesh_obj.visual.face_colors = colors else: self.trimesh_obj.visual.face_colors = np.append(colors, colors, axis=0) return None
[docs] def export(self, file_obj:str="output.glb", file_type:str="glb"): """ This function uses the export function from trimesh Parameters ---------- file_obj : str, optional The default is 'output.glb'. file_type : str, optional The default is 'glb'. Returns ------- None """ self.trimesh_obj.export(file_obj, file_type) return None
def convert_from_pyvista_faces(pyvista_obj): """ pyvista mesh faces are written in a 1d array, This function returns faces in a conventional way. A list of lists, where each list contains integers numbers of vert conections Parameters ---------- pyvista_obj : PyVista mesh The pyvista mesh. Returns ------- new_faces : list of lists A list of lists, where each list contains integers numbers of vert conections """ new_faces = [] courser = 0 for iface in range(pyvista_obj.n_faces): start = courser + 1 end = start + pyvista_obj.faces[courser] face = pyvista_obj.faces[start:end] courser = end new_faces.append(face) return new_faces def boolean_add(surfaces): """ This functtion uses boolean add from PyVista Parameters ---------- surfaces : list of pyprocar.Surface DESCRIPTION. Returns ------- surf : pyprocar surface The unionized surface from surfaces """ try: ret = surfaces[0].pyvista_obj.copy() for isurface in range(1, len(surfaces)): ret = ret.boolean_add(surfaces[isurface].pyvista_obj, inplace=False) except: try: ret = surfaces[0].copy() for isurface in range(1, len(surfaces)): ret = ret.boolean_add(surfaces[isurface], inplace=False) except: print("Not a valid surface") surf = Surface(verts=ret.points, faces=convert_from_pyvista_faces(ret), face_normals=ret.face_normals, vert_normals=ret.point_normals, scalars=ret.active_scalars) # ret = surfaces[0].pyvista_obj.copy() # for isurface in range(1, len(surfaces)): # ret = ret.boolean_add(surfaces[isurface].pyvista_obj, inplace=False) # surf = Surface( # verts=ret.points, # faces=convert_from_pyvista_faces(ret), # face_normals=ret.face_normals, # vert_normals=ret.point_normals, # scalars=ret.active_scalars, # ) return surf