NeuralJacobianFields/MeshProcessor.py

from multiprocessing import process
import warnings
warnings.filterwarnings("ignore")


from scipy.sparse import load_npz, save_npz
from .PoissonSystem import poisson_system_matrices_from_mesh, PoissonSystemMatrices, SparseMat
import os
import trimesh
from easydict import EasyDict
import numpy
import numpy as np
import scipy
import scipy.sparse
import igl
from scipy.sparse import save_npz
from time import time
import torch

NUM_SAMPLES = 1024
WKS_DIM = 100

class MeshProcessor:
    '''
    Extracts all preprocessing-related data (sample points  for pointnet; wave-kernel-signature, etc.)
    '''
    def __init__(self, vertices, faces, ttype, source_dir=None,from_file = False,
                 cpuonly=False, load_wks_samples=False, load_wks_centroids=False,
                compute_splu=True, load_splu=False):
        '''
        :param vertices:
        :param faces:
        :param ttype: the torch data type to use (float, half, double)
        :param source_dir: the directory to load the preprocessed data from; if given, will try to load the data before computing, if not given, always compute
        '''
        
        self.ttype = ttype
        self.num_samples = NUM_SAMPLES
        self.vertices = vertices.squeeze()
        self.faces = faces.squeeze()
        self.normals =  igl.per_vertex_normals(self.vertices, self.faces)
        # self.__use_wks = use_wks
        self.samples = EasyDict()
        self.samples.xyz = None
        self.samples.normals = None
        self.samples.wks = None
        self.centroids = EasyDict()
        self.centroids.points_and_normals = None
        self.centroids.wks = None
        self.diff_ops = EasyDict()
        self.diff_ops.splu = EasyDict()
        self.diff_ops.splu.L = None
        self.diff_ops.splu.U = None
        self.diff_ops.splu.perm_c = None
        self.diff_ops.splu.perm_r = None
        self.diff_ops.frames = None
        self.diff_ops.rhs = None
        self.diff_ops.grad = None
        self.diff_ops.poisson_sys_mat = None
        self.faces_wks = None
        self.vert_wks = None
        self.diff_ops.poisson = None
        self.source_dir = source_dir
        self.from_file = from_file
        self.cpuonly = cpuonly
        self.load_wks_samples = load_wks_samples
        self.load_wks_centroids = load_wks_centroids
        self.compute_splu = compute_splu
        self.load_splu = load_splu

    @staticmethod
    def meshprocessor_from_directory(source_dir, ttype, cpuonly=False, load_wks_samples=False, load_wks_centroids=False):
        try:
            vertices = np.load(os.path.join(source_dir, "vertices.npy"))
            faces = np.load(os.path.join(source_dir, "faces.npy"))
        except:
            print(os.path.join(source_dir, "vertices.npy"))
            import traceback
            traceback.print_exc()
        return MeshProcessor(vertices,faces,ttype,source_dir, cpuonly=cpuonly, load_wks_samples=load_wks_samples, load_wks_centroids=load_wks_centroids, compute_splu=False)

    @staticmethod
    def meshprocessor_from_file(fname, ttype, cpuonly=False, load_wks_samples=False, load_wks_centroids=False):
        if fname[-4:] == '.obj':
            V, _, _, F, _, _ = igl.read_obj(fname)
        elif fname[-4:] == '.off':
            V,F,_ = igl.read_off(fname)
        elif fname[-4:] == '.ply':
            V,F = igl.read_triangle_mesh(fname)
        return MeshProcessor(V,F,ttype,os.path.dirname(fname),True, cpuonly=cpuonly, load_wks_samples=load_wks_samples, load_wks_centroids=load_wks_centroids, compute_splu=False)

    @staticmethod
    def meshprocessor_from_array(vertices, faces, source_dir, ttype, cpuonly=False, load_wks_samples=False, load_wks_centroids=False):
        return MeshProcessor(vertices,faces,ttype,source_dir, cpuonly=cpuonly, load_wks_samples=load_wks_samples, load_wks_centroids=load_wks_centroids, compute_splu=False)



    def get_vertices(self):
        return self.vertices

    def get_faces(self):
        return self.faces

    def load_centroids(self):
        self.centroids.points_and_normals = np.load(os.path.join(self.source_dir, "centroids_and_normals.npy"))
        if self.load_wks_centroids:
            self.centroids.wks = np.load(os.path.join(self.source_dir, "centroids_wks.npy"))

    def get_samples(self):
        if self.samples.xyz is None:
            if True:
                try:
                    self.load_samples()
                except Exception as e:
                    self.compute_samples()
                    self.save_samples()
        return self.samples

    def load_samples(self):
        if self.samples.xyz is None:
            self.samples.xyz = np.load(os.path.join(self.source_dir, 'samples.npy'))
        if self.samples.normals is None:
            self.samples.normals = np.load(os.path.join(self.source_dir, 'samples_normals.npy'))
        if self.load_wks_samples:
            if  self.samples.wks is None:
                self.samples.wks = np.load(os.path.join(self.source_dir, 'samples_wks.npy'))
            if self.centroids.wks is None:
                self.centroids.wks = np.load(os.path.join(self.source_dir, 'centroid_wks.npy'))

    def save_samples(self):
        os.makedirs(self.source_dir, exist_ok=True)
        np.save(os.path.join(self.source_dir, 'samples.npy'), self.samples.xyz)
        np.save(os.path.join(self.source_dir, 'samples_normals.npy'), self.samples.normals)
        if self.load_wks_samples:
            np.save(os.path.join(self.source_dir, 'samples_wks.npy'), self.samples.wks)
            np.save(os.path.join(self.source_dir, 'centroid_wks.npy'), self.centroids.wks)

    def compute_samples(self):
        sstime = time()
        if self.load_wks_centroids or self.load_wks_centroids:
            self.computeWKS()
        # print(f"WKS {time() - sstime}")
        pt_samples, normals_samples, wks_samples, bary = self.sample_points( self.num_samples)
        self.samples.xyz = pt_samples
        self.samples.normals = normals_samples
        self.samples.wks = wks_samples
        self.centroids.wks = self.faces_wks

    def get_centroids(self):
        if self.centroids.points_and_normals is None:
            if True:
                try:
                    self.load_centroids()
                except Exception as e:
                    self.compute_centroids()
                    # self.save_centroids() # centroid WKS and samples WKS are intertwined right now and you cannot really use one without the other. So this is redondont with function save_samples
        return self.centroids

    def compute_centroids(self):
        m = trimesh.Trimesh(vertices=self.vertices, faces=self.faces, process=False)
        self.centroids.points_and_normals = np.hstack((np.mean(m.triangles, axis=1), m.face_normals))
        self.get_samples()# this is to compute WKS for centroids

    def get_differential_operators(self):
        if self.diff_ops.grad is None:
            if True:
                try:
                    self.load_differential_operators()
                except Exception as e:
                    warnings.warn(f'while loading data, got file not exists exception: {e} ')
                    self.compute_differential_operators()
                    self.save_differential_operators()
        if self.load_splu:
            self.get_poisson_system()
        return self.diff_ops


    def load_poisson_system(self):
        try:
            self.diff_ops.splu.L = load_npz(os.path.join(self.source_dir, 'lap_L.npz'))
            self.diff_ops.splu.U = load_npz(os.path.join(self.source_dir, 'lap_U.npz'))
            self.diff_ops.splu.perm_c = np.load(os.path.join(self.source_dir, 'lap_perm_c.npy'))
            self.diff_ops.splu.perm_r = np.load(os.path.join(self.source_dir, 'lap_perm_r.npy'))
        except:
            print(f"FAILED load poisson on: {os.path.join(self.source_dir)}")
            raise Exception("FAILED load poisson on: {os.path.join(self.source_dir)}")


    def load_differential_operators(self): 
        self.diff_ops.rhs = SparseMat.from_coo(load_npz(os.path.join(self.source_dir, 'new_rhs.npz')), ttype=torch.float64)
        self.diff_ops.grad = SparseMat.from_coo(load_npz(os.path.join(self.source_dir, 'new_grad.npz')), ttype=torch.float64)
        self.diff_ops.frames = np.load(os.path.join(self.source_dir, 'w.npy'))
        self.diff_ops.laplacian = SparseMat.from_coo(load_npz(os.path.join(self.source_dir, 'laplacian.npz')), ttype=torch.float64)

    def save_differential_operators(self):  
        save_npz(os.path.join(self.source_dir, 'new_rhs.npz'), self.diff_ops.rhs.to_coo())
        save_npz(os.path.join(self.source_dir, 'new_grad.npz'), self.diff_ops.grad.to_coo())
        np.save(os.path.join(self.source_dir, 'w.npy'), self.diff_ops.frames)
        save_npz(os.path.join(self.source_dir, 'laplacian.npz'), self.diff_ops.laplacian.to_coo())

    def compute_differential_operators(self):
        '''
        process the given mesh
        '''
        poisson_sys_mat = poisson_system_matrices_from_mesh(V= self.vertices, F=self.faces,  cpuonly=self.cpuonly)
        self.diff_ops.grad = poisson_sys_mat.igl_grad
        self.diff_ops.rhs = poisson_sys_mat.rhs
        self.diff_ops.laplacian = poisson_sys_mat.lap
        self.diff_ops.frames = poisson_sys_mat.w
        self.diff_ops.poisson_sys_mat = poisson_sys_mat
    

    def compute_poisson(self):
        poissonsolver = poissonbuilder.compute_poisson_solver_from_laplacian(compute_splu=self.compute_splu)
        # new_grad = poissonbuilder.get_new_grad() # This is now done in poisson_system_matrices_from_mesh
        if self.compute_splu:
            self.diff_ops.splu.L, self.diff_ops.splu.U , self.diff_ops.splu.perm_c , self.diff_ops.splu.perm_r = poissonbuilder.compute_splu()
        self.diff_ops.frames = poissonbuilder.w


    def prepare_differential_operators_for_use(self,ttype):
        diff_ops = self.get_differential_operators() # call 1
        ## WARNING : we commented these two lines because they seemed redundant.
        if self.diff_ops.poisson_sys_mat is None: # not created if loaded from disk the diff ops
            diff_ops.poisson_sys_mat = PoissonSystemMatrices(self.vertices, self.faces, diff_ops.grad, diff_ops.rhs, diff_ops.frames, ttype, lap = diff_ops.laplacian,  cpuonly=self.cpuonly)
        
        self.diff_ops.poisson_solver = diff_ops.poisson_sys_mat.create_poisson_solver() # call 2
        self.diff_ops.MyCuSPLU_solver = diff_ops.poisson_sys_mat.create_poisson_solver() #create_poisson_solver_from_splu_old

    def get_writeable(self):
        '''
        get dictionaries to write numpy and npz
        :return: two args, np, npz, each dicts with field_name --> data to save
        '''
        out_np = {}
        out_npz = {}
        out_np['vertices'] = self.vertices
        out_np['faces'] = self.faces
        if self.samples is not None:
            out_np["samples"] = self.samples.xyz
            out_np["samples_normals"] = self.samples.normals
            out_np["samples_wks"] = self.samples.wks
        if self.centroids is not None:
            out_np["centroids_wks"] = self.centroids.wks
            out_np["centroids_and_normals"] = self.centroids.points_and_normals
        if self.diff_ops is not None:
            out_np['lap_perm_c'] = self.diff_ops.splu.perm_c
            out_np['lap_perm_r'] = self.diff_ops.splu.perm_r
            out_np['w'] =self.diff_ops.frames
            out_npz['new_grad'] = self.diff_ops.grad.to_coo()
            out_npz['new_rhs'] = self.diff_ops.rhs
            out_npz['lap_L'] = self.diff_ops.splu.L
            out_npz['lap_U'] = self.diff_ops.splu.U
            out_npz['lap'] = self.diff_ops.poisson.lap
        return {key: value for key, value in out_np.items() if value is not None}, {key: value for key, value in out_npz.items() if value is not None}
        
    def get_data(self, key,file_type = 'npy'):
        if key == 'samples':
            return self.get_samples().xyz
        elif key == "samples_normals":
            return self.get_samples().normals
        elif key == "samples_wks":
            return self.get_samples().wks
        elif key == 'vertices':
            return self.vertices
        elif key == 'faces':
            return self.faces
        if file_type == 'npy':
            return np.load(os.path.join(self.source_dir, f'{key}.npy'))
        elif file_type == 'npz':
            return load_npz(os.path.join(self.source_dir, f'{key}.npz'))
        else:
            raise RuntimeError("wrong file type")

    def computeWKS(self):
        if self.faces_wks is  None or self.vert_wks is  None:
            st = time()
            w = WaveKernelSignature(self.vertices, self.faces, top_k_eig=50)
            w.compute()
            print(f"Ellapsed {time() - st}")
            wk = w.wks
            faces_wks = np.zeros((self.faces.shape[0], wk.shape[1]))
            for i in range(3):
                faces_wks += wk[self.faces[:, i], :]
            faces_wks /= 3
            self.faces_wks = faces_wks
            self.vert_wks = wk
            assert (self.faces_wks.shape[0] == self.faces.shape[0])
            assert (self.vert_wks.shape[0] == self.vertices.shape[0])


    def sample_points(self, n):
        bary, found_faces = igl.random_points_on_mesh(n, self.vertices, self.faces)
        vert_ind =  self.faces[found_faces]
        point_samples =  self.vertices[vert_ind[:,0]] * bary[:,0:1] + self.vertices[vert_ind[:,1]] * bary[:,1:2] + self.vertices[vert_ind[:,2]] * bary[:,2:3]
        normal_samples = self.normals[vert_ind[:,0]] * bary[:,0:1] + self.normals[vert_ind[:,1]] * bary[:,1:2] + self.normals[vert_ind[:,2]] * bary[:,2:3]
        wks_samples = None
        if self.load_wks_centroids or self.load_wks_samples:
            wks_samples = self.vert_wks[vert_ind[:,0]] * bary[:,0:1] + self.vert_wks[vert_ind[:,1]] * bary[:,1:2] + self.vert_wks[vert_ind[:,2]] * bary[:,2:3]
        return point_samples, normal_samples, wks_samples, bary

# This is insane to me
def sample_points(V, F, n):
    '''
    samples n points on the given mesh, along with normals and wks. Also return WKS of original faces (by averaging wks of 3 vertices of each face)
    :return:
    '''
    newF = F
    newV = V
    for iter in range(n):
        newV, newF = _sample_point(newV, newF)

    w = WaveKernelSignature(newV, newF, top_k_eig=100)
    w.compute()
    wk = w.wks
    sample_ks = wk[len(V):, :]
    org_ks = wk[:len(V), :]
    normals = igl.per_vertex_normals(newV, newF)
    normals = normals[len(V):, :]

    # get per-face wks by averaging its vertices
    faces_wks = np.zeros((F.shape[0], org_ks.shape[1]))
    for i in range(3):
        faces_wks += org_ks[F[:, i], :]
    faces_wks /= 3
    return newV[len(V):, :], normals, sample_ks, faces_wks


def _sample_point(VV, FF):
    while (True):
        bary, found_faces = igl.random_points_on_mesh(1, VV, FF)
        if (found_faces >= FF.shape[0]):
            continue
        # use to be 0.01
        if not numpy.any(bary < 0.05):
            break
    ret = numpy.zeros((1, VV.shape[1]))
    for i in range(VV.shape[1]):
        res = np.multiply(VV[FF[found_faces, :], i], bary)
        ret[:, i] = np.sum(res)
    newF = FF
    new_index = len(VV)
    new_tris = _insert_triangle(FF[found_faces, :], new_index)
    newF = numpy.concatenate((newF, new_tris), axis=0)
    newF = numpy.delete(newF, found_faces, axis=0)
    newV = numpy.concatenate((VV, ret), 0)
    return newV, newF


def _insert_triangle(old_tri, new_index):
    d = new_index
    a, b, c = (old_tri[0], old_tri[1], old_tri[2])
    new_tris = numpy.array([[a, b, d], [b, c, d], [c, a, d]])
    return new_tris





class WaveKernelSignatureError(Exception):
    pass

class WaveKernelSignature:
    '''
    Computes wave kernel signature for a given mesh
    '''
    def __init__(self,
                 vertices,
                 faces,
                 top_k_eig=200,
                 timestamps=WKS_DIM):
        # vertices, faces are both numpy arrays.
        self.vertices = vertices
        self.faces = faces

        # self.vertices_gpu = torch.from_numpy(vertices).cuda()
        # self.faces_gpu = torch.from_numpy(faces).cuda()

        self.top_k_eig = top_k_eig
        self.timestamps = timestamps
        self.max_iter = 10000

    def compute(self):
        '''
        compute the wks. Afterwards WKS stores in self.wks
        '''
        cp = igl.connected_components(igl.adjacency_matrix(self.faces))
        assert(cp[0]==1), f"{cp}"
        L = -igl.cotmatrix(self.vertices, self.faces) # this is fast 0.04 seconds
        M = igl.massmatrix(self.vertices, self.faces, igl.MASSMATRIX_TYPE_VORONOI)
        # assert(not numpy.any(numpy.isinf(L)))
        try:
            try:
                self.eig_vals, self.eig_vecs = scipy.sparse.linalg.eigsh(
                    L, self.top_k_eig, M, sigma=0, which='LM', maxiter=self.max_iter)
            except:
                self.eig_vals, self.eig_vecs = scipy.sparse.linalg.eigsh(
                    L, self.top_k_eig, M, sigma=1e-4, which='LM', maxiter=self.max_iter)                
        except:
            raise WaveKernelSignatureError("Error in computing WKS")

        # print(np.linalg.norm(self.eig_vecs, axis=0, keepdims=True))
        # print(np.max(self.eig_vecs))
        self.eig_vecs /= 200 #np.linalg.norm(self.eig_vecs, axis=0, keepdims=True)
        # np.save("norm_v2.npy", np.max(np.abs(self.eig_vecs), axis=0, keepdims=True))
        # np.save("norm_v2.npy", np.max(np.abs(self.eig_vecs), axis=0, keepdims=True))
        # print(np.linalg.norm(self.eig_vecs, axis=0))
        # print(np.max(self.eig_vecs, axis=0))
        # print(np.min(self.eig_vecs, axis=0))
        # print(self.eig_vals)

        # self.eig_vecs /= np.load('norm_v1.npy')

        # self.eig_vecs = self.eig_vecs / np.max(np.abs(self.eig_vecs), axis=0, keepdims=True)
        # self.eig_vecs = self.eig_vecs * np.load('norm_v2.npy')


        # nn = np.load('norm2.npy')
        # self.eig_vecs /= nn[:,:50]


        # ==== VISUALIZATION CODE ==========
        if False:
            num_mesh_to_viz = 6
            meshes = []
            for i in range(num_mesh_to_viz):
                meshes.append(trimesh.Trimesh(self.vertices + np.array([i*1,0,0]), self.faces, process=False))

            # mesh = meshes[0].union( meshes[1])
            # mesh = mesh.union( meshes[2])
            # mesh = mesh.union( meshes[3])
            meshes = [trimesh.util.concatenate(meshes)]

            from vedo import trimesh2vedo, show, screenshot, Plotter

            vp = Plotter(axes=0, offscreen=True)

            vmeshes = trimesh2vedo(meshes)
            cmaps = ('jet', 'PuOr', 'viridis')
            scals =  self.eig_vecs[:,:num_mesh_to_viz].transpose((1,0)).reshape(-1)
            vmeshes[0].cmap(cmaps[0], scals).lighting('plastic')

            # add a 2D scalar bar to a mesh
            vmeshes[0].addScalarBar(title=f"scalarbar #{0}", c='k') 

            vp.show(vmeshes, axes=1)
            screenshot(f"test_{time()}.png")
            import sys
            sys.exit(0)
        #  ================




        # range between biggest and smallest eigenvalue : 
        # 6 0.09622419080119388
        # 6_bis 0.09651935545457718
        delta = (np.log(self.eig_vals[-1]) - np.log(self.eig_vals[1])) / self.timestamps
        sigma = 7 * delta
        e_min = np.log(self.eig_vals[1]) + 2 * delta
        e_max = np.log(self.eig_vals[-1]) - 2 * delta
        es = np.linspace(e_min, e_max, self.timestamps)  # T
        self.delta = delta

        
        coef = np.expand_dims(es, 0) - np.expand_dims(np.log(self.eig_vals[1:]), 1)  # (K-1)xT
        coef = np.exp(-np.square(coef) / (2 * sigma * sigma))  # (K-1)xT #element wise square
        sum_coef = coef.sum(0)  # T
        K = np.matmul(np.square(self.eig_vecs[:, 1:]), coef)  # VxT. Scaling of the eigen vectors by coef. Coef depends only on the eigen values. Triangulation agnostic.
        self.wks = K / np.expand_dims(sum_coef, 0)  # VxT Scaling of the eigen vectors by sum_coef. Coef depends only on the eigen values. Triangulation agnostic.
        # print(np.linalg.norm(self.wks, axis=0))
        # print(np.linalg.norm(self.wks, axis=1))