manifold_harmonics.py

import numpy as np
import scipy
import trimesh
import polyscope as ps
from pathlib import Path
from tqdm import tqdm


def double_area(geom):
    i = geom.faces[:, 0]  # [num_faces]
    j = geom.faces[:, 1]  # [num_faces]
    k = geom.faces[:, 2]  # [num_faces]
    V = geom.vertices
    x_kj = V[k] - V[j]  # [num_faces, 3]
    x_ik = V[i] - V[k]  # [num_faces, 3]

    # norm of cross products of two edges = twise area of the triangle
    n = np.cross(x_kj, x_ik)
    dbl_a = np.linalg.norm(n, axis=1).reshape(-1, 1)
    return dbl_a


def calc_gradient(geom):
    # geom: trimesh.Trimesh
    num_verts = geom.vertices.shape[0]
    num_faces = geom.faces.shape[0]

    # vertex indices
    i = geom.faces[:, 0]  # [num_faces]
    j = geom.faces[:, 1]  # [num_faces]
    k = geom.faces[:, 2]  # [num_faces]

    V = geom.vertices  # [num_verts, 3]

    x_kj = V[k] - V[j]  # [num_faces, 3]
    x_ik = V[i] - V[k]  # [num_faces, 3]
    x_ji = V[j] - V[i]  # [num_faces, 3]

    # norm of cross products of two edges = twise area of the triangle
    n = np.cross(x_kj, x_ik)
    dbl_a = double_area(geom)

    u = n / dbl_a  # uniformed normal of the triangle
    eperp_ji = np.cross(u, x_ji) / dbl_a
    eperp_ik = np.cross(u, x_ik) / dbl_a

    Fs = np.tile(np.arange(num_faces, dtype=np.int64), (1, 4))
    mi = np.concatenate([num_faces * 0 + Fs,
                         num_faces * 1 + Fs,
                         num_faces * 2 + Fs]).reshape(-1, 1)
    mj = np.tile(np.array([geom.faces[:, 1],
                           geom.faces[:, 0],
                           geom.faces[:, 2],
                           geom.faces[:, 0]]), (3, 1)).reshape(-1, 1)
    mv = np.concatenate([eperp_ik[:, 0], -eperp_ik[:, 0], eperp_ji[:, 0], -eperp_ji[:, 0],
                         eperp_ik[:, 1], -eperp_ik[:, 1], eperp_ji[:, 1], -eperp_ji[:, 1],
                         eperp_ik[:, 2], -eperp_ik[:, 2], eperp_ji[:, 2], -eperp_ji[:, 2]])
    mv = mv.reshape(-1, 1)

    # Gradient
    G = scipy.sparse.csr_matrix((mv.reshape(-1),
                                 (mi.reshape(-1), mj.reshape(-1))),
                                shape=(num_faces * 3, num_verts))
    return G


def calc_alt_mass_mat(geom):
    d = double_area(geom).reshape(-1)
    T = scipy.sparse.diags((np.hstack([d, d, d]) * 0.5))
    return T


def laplace_beltrami(geom):
    G = calc_gradient(geom)
    T = calc_alt_mass_mat(geom)
    return - G.T @ T @ G


if __name__ == "__main__":
    ROOT_PATH = Path(".").resolve()
    asset_dir = ROOT_PATH / "assets"

    # geom = trimesh.load(str(asset_dir / "garg.off"), process=False)
    geom = trimesh.load(str(asset_dir / "bunny.ply"), process=False)
    print("geom: ", geom.vertices.shape)
    L = laplace_beltrami(geom)

    # lamb, e = scipy.sparse.linalg.eigs(L, k=1000)
    lamb, e = np.linalg.eig(L.toarray())
    out = np.zeros_like(geom.vertices)

    f = geom.vertices
    ps.init()
    ps.register_surface_mesh("origin", geom.vertices, geom.faces)
    for i in tqdm(range(e.shape[1])):
        ei = e[:, i].reshape(-1, 1).real
        out += (ei.T @ f) * ei
        if (i + 1) % 100 == 0:
            ps.register_surface_mesh(f"out_{i + 1}", out, geom.faces, enabled=False)
    ps.show()