gravview.py

# examples/Python/Basic/mesh.py
import numpy as np
from gravmod import gravity
import open3d as o3d
import trimesh
import matplotlib as plt
from matplotlib import cm
import numba as nb
from scipy.spatial.transform import Rotation

Gcon = np.float64(6.6726e-11)

@nb.njit(fastmath=True,parallel=True)
def cross_product(vec_1,vec_2):
    res=np.empty((vec_1.shape[0],3),dtype=vec_1.dtype)
    for i in nb.prange(vec_1.shape[0]):
        res[i,0]=vec_1[i,1] * vec_2[i,2] - vec_1[i,2] * vec_2[i,1]
        res[i,1]=vec_1[i,2] * vec_2[i,0] - vec_1[i,0] * vec_2[i,2]
        res[i,2]=vec_1[i,0] * vec_2[i,1] - vec_1[i,1] * vec_2[i,0]

    return res

#@nb.njit(fastmath=True,parallel=True)
def unit_vec(vec):
    vecnorms = np.apply_along_axis(np.linalg.norm, 1, vec)
    vechat = (vec.T / vecnorms).T
    return vecnorms, vechat


class smallbod:
    """A model of a self-gravitating small body"""
    def __init__(self):
        self.polyfile = ""
        self.dens = np.float64(1000.0)
        self.rot_cen = np.zeros(3,dtype='float64')
        self.rot_rate = np.zeros(3,dtype='float64')
        #used for Open3d module
        self.o3d_mesh = o3d.geometry.TriangleMesh()
        # used for Trimesh module
        self.mesh = trimesh.Trimesh()
        self.nface = 0
        self.nvert = 0
        self.vert_has_gravity = False
        self.vert_gravity = np.zeros(3,dtype='float64')
        self.vert_potential = np.zeros(1,dtype='float64')
        self.vert_has_slope = False
        self.face_has_gravity = False
        self.face_gravity = np.zeros(3,dtype='float64')
        self.face_potential = np.zeros(1,dtype='float64')
        self.face_has_slope = False
        self.rotP = np.float64(0.0) # Rotation period
        self.angular_momentum = np.float(0.0)
        self.mesh_is_distorted = False
        return

    def io_read_ply(self,filename):
        """Reads in PLY file into Open3d triangular mesh object and computes normals"""
        self.polyfile = filename
        self.mesh = trimesh.load(f"{self.polyfile}")
        trimesh.repair.fix_normals(self.mesh)
        self.mesh.vertex_colors = np.full((self.nvert,4),1.0)
        self.nvert = len(self.mesh.vertices)
        self.nface = len(self.mesh.faces)
        return

    def render(self):
        """Renders output"""
        self.convert_trimesh_to_o3d()
        o3d.visualization.draw_geometries([self.o3d_mesh])
        return

    def convert_trimesh_to_o3d(self):
        """Update Open3D mesh with new vertices and faces, and recompute all normals"""
        self.o3d_mesh.vertices  = o3d.Vector3dVector(self.mesh.vertices)
        self.o3d_mesh.triangles = o3d.Vector3iVector(self.mesh.faces)
        self.o3d_mesh.compute_vertex_normals()
        self.o3d_mesh.compute_triangle_normals()
        self.o3d_mesh.vertex_colors = o3d.Vector3dVector(self.mesh.vertex_colors[:,0:3])
        return

    def paint_roid(self,color_value,cvmin,cvmax):
        #norm = plt.colors.Normalize()
        norm = color_value / cvmax
        self.mesh.vertex_colors  = plt.cm.terrain(norm)
        return

    def compute_vert_gravity(self):
        #Polygrav needs arrays with dimension as the first index, so need to transpose from Trimesh/Open3d
        m = self.mesh.vertices.T
        v = self.mesh.faces.T + 1 # Make vertex index list compatible with Fortran 1-indexed arrays
        mxn = self.mesh.face_normals.T
        self.vert_gravity = np.zeros([self.nvert,3])
        self.vert_potential = np.zeros([self.nvert])
        for i in range(self.nvert):
            x0 = m[:,i]
            self.vert_gravity[i,:], self.vert_potential[i] = gravity.polygrav(x0,m,v,mxn,self.rot_cen,self.rot_rate,self.dens)
        self.vert_has_gravity = True
        return

    def compute_vert_slope(self):
        if not self.vert_has_gravity:
            self.compute_vert_gravity()
        #normalize the gravity vectors
        norms = np.apply_along_axis(np.linalg.norm, 1, self.vert_gravity)
        ghat = (self.vert_gravity.T / norms).T
        self.vert_slope = np.arccos(np.einsum('ij,ij->i', ghat, -self.mesh.vertex_normals))
        self.vert_has_slope = True
        return

    def compute_face_gravity(self):
        #Polygrav needs arrays with dimension as the first index, so need to transpose from Open3D
        m = self.mesh.vertices.T
        v = self.mesh.faces.T + 1 # Make vertex index list compatible with Fortran 1-indexed arrays
        mxn = self.mesh.face_normals.T
        self.face_gravity = np.zeros([self.nface,3])
        self.face_potential = np.zeros([self.nface])
        for i in range(self.nface):
            vert = self.mesh.faces[i,:]
            x0 = (m[:,vert[0]] + m[:,vert[1]] + m[:,vert[2]]) / 3.0
            self.face_gravity[i,:],self.face_potential[i] = gravity.polygrav(x0,m,v,mxn,self.rot_cen,self.rot_rate,self.dens)
        self.face_has_gravity = True
        return

    def compute_face_slope(self):
        if not self.face_has_gravity:
            self.compute_face_gravity()
        # normalize the gravity vectors
        norms = np.apply_along_axis(np.linalg.norm, 1, self.face_gravity)
        ghat = (self.face_gravity.T / norms).T
        self.face_slope = np.arccos(np.einsum('ij,ij->i', ghat, -self.facenorm))
        self.face_has_slope = True
        return

    def distort_mesh_to_spherical_potential(self):
        #Scale to the length of the vertex that it would have if it were in the 1/r gravity potential of an equivalent
        #mass spherical body with the vertex vectors along the same axis as the local gravity+rotation acceleration vector

        if not self.vert_has_gravity:
            self.compute_vert_gravity()

        #Compute rotation matrices to move each vertex to the direction of the gravity unit vector
        vnorms, vhat = unit_vec(self.mesh.vertices)
        gnorms, ghat = unit_vec(self.vert_gravity)
        theta = np.arccos(np.einsum('ij,ij->i', vhat, -ghat))
        rvec = cross_product(vhat,ghat)
        rnorms, rhat = unit_vec(rvec)

        #(vec.T / vecnorms).T
        rotvec = (rhat.T * theta).T
        self.rotmat = Rotation.from_rotvec(rotvec)

        # Scale to the length of the vertex that it would have if it were in the 1/r gravity potential of an equivalent
        self.rscale = Gcon * self.mesh.mass  / self.vert_potential / vnorms

        #Now apply rotations and scale transforemation
        #self.mesh.vertices = self.rotmat.apply(self.mesh.vertices)
        self.mesh.vertices = self.mesh.vertices * self.rscale[:, np.newaxis]
        self.mesh.fix_normals()
        self.mesh_is_distorted = True
        return

    def undistort_mesh(self):
        #Undo the gravity + rotational potential distortion
        if not self.mesh_is_distorted:
            return
        self.mesh.vertices = self.mesh.vertices / self.rscale[:, np.newaxis]
        #self.mesh.vertices = self.rotmat.apply(self.mesh.vertices,inverse=True)
        self.mesh.fix_normals()
        self.mesh_is_distorted = False
        return

#Make sure we are running in parallel
gravity.test_omp()

#Make our sandy mu69
mu69 = smallbod()

mu69.io_read_ply("TestData/mu69_test5h_1_final_oriented.ply")
mu69.mesh.density = 300.00e9
mu69.rotP = np.float64(15.9 * 60 * 60)

# Put in rotation about maximum inertia axis
mu69.mesh.apply_transform(mu69.mesh.principal_inertia_transform)
mu69.rot_cen = np.zeros(3, dtype='float64')
omega = 2 * np.pi / mu69.rotP
mu69.rot_rate[np.argmax(mu69.mesh.principal_inertia_components)] = omega
# Save angular momentum for conservation law
mu69.angular_momentum = np.amax(mu69.mesh.principal_inertia_components) * omega
mu69.mass = mu69.mesh.mass


#Paint the mu69 with vertex slope values
mu69.compute_vert_slope()
gaccel = np.sqrt(np.einsum('ij,ij->i', mu69.vert_gravity,mu69.vert_gravity))
gmin = np.amin(gaccel)
gmax = np.amax(gaccel)
mu69.paint_roid(gaccel,cvmin=gmin,cvmax=gmax)
#mu69.distort_mesh_to_spherical_potential()
mu69.render()