first_try

app.py

@@ -0,0 +1,272 @@
+"""
+Simple Gradio app for two-mesh initialization and run phases.
+- Upload two meshes (.ply, .obj, .off)
+- (Optional) upload a YAML config to override defaults
+- Adjust a few numeric settings (sane ranges). Defaults pulled from the provided YAML when present.
+- Click **Init** to generate "initialization maps" (here: position/normal-based vertex colors) for both meshes.
+- Click **Run** to simulate an iterative evolution with a progress bar, then output another pair of colored meshes.
+
+Replace the bodies of `make_initialization_maps` and `run_evolution` with your real pipeline as needed.
+
+Tested with: gradio >= 4.0, trimesh, pyyaml, numpy.
+"""
+from __future__ import annotations
+import os
+import io
+import time
+import json
+import tempfile
+from typing import Dict, Tuple, Optional
+from omegaconf import OmegaConf
+import gradio as gr
+import numpy as np
+import trimesh
+import zero_shot
+import yaml
+from utils.surfaces import Surface
+import notebook_helpers as helper
+from utils.meshplot import visu_pts
+from utils.torch_fmap import extract_p2p_torch_fmap, torch_zoomout
+import torch
+import argparse
+# -----------------------------
+# Utils
+# -----------------------------
+SUPPORTED_EXTS = {".ply", ".obj", ".off", ".stl", ".glb", ".gltf"}
+
+def _safe_ext(path: str) -> str:
+    for ext in SUPPORTED_EXTS:
+        if path.lower().endswith(ext):
+            return ext
+    return os.path.splitext(path)[1].lower()
+
+
+def normalize_vertices(vertices: np.ndarray) -> np.ndarray:
+    v = vertices.astype(np.float64)
+    v = v - v.mean(axis=0, keepdims=True)
+    scale = np.linalg.norm(v, axis=1).max()
+    if scale == 0:
+        scale = 1.0
+    v = v / scale
+    return v.astype(np.float32)
+
+
+
+def ensure_vertex_colors(mesh: trimesh.Trimesh, colors: np.ndarray) -> trimesh.Trimesh:
+    out = mesh.copy()
+    if colors.shape[1] == 3:
+        rgba = np.concatenate([colors, 255*np.ones((colors.shape[0],1), dtype=np.uint8)], axis=1)
+    else:
+        rgba = colors
+    out.visual.vertex_colors = rgba
+    return out
+
+
+def export_for_view(surf: Surface, colors: np.ndarray, basename: str, outdir: str) -> Tuple[str, str]:
+    """Export to PLY (with vertex colors) and GLB for Model3D preview."""
+    glb_path = os.path.join(outdir, f"{basename}.glb")
+    mesh = trimesh.Trimesh(surf.vertices, surf.faces, process=False)
+    colored_mesh = ensure_vertex_colors(mesh, colors)
+    colored_mesh.export(glb_path)
+    return glb_path
+
+
+# -----------------------------
+# Algorithm placeholders (replace with your real pipeline)
+# -----------------------------
+DEFAULT_SETTINGS = {
+    "deepfeat_conf.fmap.lambda_": 1,
+    "sds_conf.zoomout": 40.0,
+    "diffusion.time": 1.0,
+    "opt.n_loop": 300,
+    "loss.sds": 1.0,
+    "loss.proper": 1.0,
+}
+
+FLOAT_SLIDERS = {
+    # name: (min, max, step)
+    "deepfeat_conf.fmap.lambda_": (1e-3, 10.0, 1e-3),
+    "sds_conf.zoomout": (1e-3, 10.0, 1e-3),
+    "diffusion.time": (1e-3, 10.0, 1e-3),
+    "loss.sds": (1e-3, 10.0, 1e-3),
+    "loss.proper": (1e-3, 10.0, 1e-3),
+}
+
+INT_SLIDERS = {
+    "opt.n_loop": (1, 5000, 1),
+}
+
+
+def flatten_yaml_floats(d: Dict, prefix: str = "") -> Dict[str, float]:
+    flat = {}
+    for k, v in d.items():
+        key = f"{prefix}.{k}" if prefix else str(k)
+        if isinstance(v, dict):
+            flat.update(flatten_yaml_floats(v, key))
+        elif isinstance(v, (int, float)):
+            flat[key] = float(v)
+    return flat
+
+
+def read_yaml_defaults(yaml_path: Optional[str]) -> Dict[str, float]:
+    if yaml_path and os.path.exists(yaml_path):
+        with open(yaml_path, "r") as f:
+            data = yaml.safe_load(f)
+        flat = flatten_yaml_floats(data)
+        # Only keep known keys we expose as controls
+        defaults = DEFAULT_SETTINGS.copy()
+        for k in list(DEFAULT_SETTINGS.keys()):
+            if k in flat:
+                defaults[k] = float(flat[k])
+        return defaults
+    return DEFAULT_SETTINGS.copy()
+
+
+
+
+class Datadicts:
+    def __init__(self, shape_path, target_path):
+        self.shape_path = shape_path
+        basename_1 = os.path.basename(shape_path)
+        self.shape_dict, _ = helper.load_data(shape_path, "tmp/" + os.path.splitext(basename_1)[0]+".npz", "source", make_cache=True)
+        self.shape_surf = Surface(filename=shape_path)
+        self.target_path = target_path
+        basename_2 = os.path.basename(target_path)
+        self.target_dict, _ = helper.load_data(target_path, "tmp/" + os.path.splitext(basename_2)[0]+".npz", "target", make_cache=True)
+        self.target_surf = Surface(filename=target_path)
+        self.cmap1 = visu_pts(self.shape_surf)
+
+# -----------------------------
+# Gradio UI
+# -----------------------------
+TMP_ROOT = tempfile.mkdtemp(prefix="meshapp_")
+
+def save_array_txt(arr):
+    # Create a temporary file with .txt suffix
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".txt", mode="w") as f:
+        np.savetxt(f, arr.astype(int), fmt="%d")  # save as text
+        return f.name
+
+def build_outputs(surf_a: Surface, surf_b: Surface, cmap_a: np.ndarray, p2p: np.ndarray, tag: str) -> Tuple[str, str, str, str]:
+    outdir = os.path.join(TMP_ROOT, tag)
+    os.makedirs(outdir, exist_ok=True)
+    glb_a = export_for_view(surf_a, cmap_a, f"A_{tag}", outdir)
+    cmap_b = cmap_a[p2p]
+    glb_b = export_for_view(surf_b, cmap_b, f"B_{tag}", outdir)
+    out_file = save_array_txt(p2p)
+    return glb_a, glb_b, out_file
+
+
+def init_clicked(mesh1_path, mesh2_path,
+                 lambda_val, zoomout_val, time_val, nloop_val, sds_val, proper_val):
+    cfg.deepfeat_conf.fmap.lambda_ = lambda_val
+    cfg.sds_conf.zoomout = zoomout_val
+    cfg.deepfeat_conf.fmap.diffusion.time = time_val
+    cfg.opt.n_loop = nloop_val
+    cfg.loss.sds = sds_val
+    cfg.loss.proper = proper_val
+    matcher.reconf(cfg)
+    if not mesh1_path or not mesh2_path:
+        raise gr.Error("Please upload both meshes.")
+    matcher._init()
+    global datadicts
+    datadicts = Datadicts(mesh1_path, mesh2_path)
+
+    C12_pred_init, C21_pred_init, feat1, feat2, evecs_trans1, evecs_trans2 = matcher.fmap_model({"shape1": datadicts.shape_dict, "shape2": datadicts.target_dict}, diff_model=matcher.diffusion_model, scale=matcher.fmap_cfg.diffusion.time)
+    C12_pred, C12_obj, mask_12 = C12_pred_init
+    p2p_init, _ = extract_p2p_torch_fmap(C12_obj, datadicts.shape_dict["evecs"], datadicts.target_dict["evecs"])
+    return build_outputs(datadicts.shape_surf, datadicts.target_surf, datadicts.cmap1, p2p_init, tag="init")
+
+
+def run_clicked(mesh1_path, mesh2_path, yaml_path, lambda_val, zoomout_val, time_val, nloop_val, sds_val, proper_val, progress=gr.Progress(track_tqdm=True)):
+    if not mesh1_path or not mesh2_path:
+        raise gr.Error("Please upload both meshes.")
+
+    cfg.deepfeat_conf.fmap.lambda_ = lambda_val
+    cfg.sds_conf.zoomout = zoomout_val
+    cfg.deepfeat_conf.fmap.diffusion.time = time_val
+    cfg.opt.n_loop = nloop_val
+    cfg.loss.sds = sds_val
+    cfg.loss.proper = proper_val
+    matcher.reconf(cfg)
+    if not mesh1_path or not mesh2_path:
+        raise gr.Error("Please upload both meshes.")
+    matcher._init()
+    global datadicts
+    if datadicts is None:
+        datadicts = Datadicts(mesh1_path, mesh2_path)
+    elif datadicts is not None:
+        if not (datadicts.shape_path == mesh1_path and datadicts.target_path == mesh2_path):
+            datadicts = Datadicts(mesh1_path, mesh2_path)
+
+    target_normals = torch.from_numpy(datadicts.target_surf.surfel/np.linalg.norm(datadicts.target_surf.surfel, axis=-1, keepdims=True)).float().to("cuda")
+
+    C12_new, p2p, p2p_init, _, loss_save = matcher.optimize(datadicts.shape_dict, datadicts.target_dict, target_normals)
+    evecs1, evecs2 = datadicts.shape_dict["evecs"], datadicts.target_dict["evecs"]
+    evecs_2trans = evecs2.t() @ torch.diag(datadicts.target_dict["mass"])
+    C12_end_zo = torch_zoomout(evecs1, evecs2, evecs_2trans, C12_new.squeeze()[:15, :15], 150)# matcher.cfg.sds_conf.zoomout)
+    p2p_zo, _ = extract_p2p_torch_fmap(C12_end_zo, datadicts.shape_dict["evecs"], datadicts.target_dict["evecs"])
+    return build_outputs(datadicts.shape_surf, datadicts.target_surf, datadicts.cmap1, p2p_zo, tag="run")
+
+
+with gr.Blocks(title="DiffuMatch demo") as demo:
+    text_in = "Upload two meshes and try our ICCV zero-shot method DiffuMatch! \n"
+    text_in += "*Init* will give you a rough correspondence, and you can click on *Run* to see if our method is able to match the two shapes! \n"
+    text_in += "*Recommended*: The method requires that the meshes are aligned (rotation-wise) to work well. Also might not work with scans (but try it out!)."
+    gr.Markdown(text_in)
+    with gr.Row():
+        mesh1 = gr.File(label="Mesh A (.ply/.obj/.off)")
+        mesh2 = gr.File(label="Mesh B (.ply/.obj/.off)")
+        yaml_file = gr.File(label="Optional YAML config", file_types=[".yaml", ".yml"], visible=True)
+    # except Exception:
+    with gr.Accordion("Settings", open=True):
+        with gr.Row():
+            lambda_val = gr.Slider(minimum=FLOAT_SLIDERS["deepfeat_conf.fmap.lambda_"][0], maximum=FLOAT_SLIDERS["deepfeat_conf.fmap.lambda_"][1], step=FLOAT_SLIDERS["deepfeat_conf.fmap.lambda_"][2], value=1, label="deepfeat_conf.fmap.lambda_")
+            zoomout_val = gr.Slider(minimum=FLOAT_SLIDERS["sds_conf.zoomout"][0], maximum=FLOAT_SLIDERS["sds_conf.zoomout"][1], step=FLOAT_SLIDERS["sds_conf.zoomout"][2], value=40, label="sds_conf.zoomout")
+            time_val = gr.Slider(minimum=FLOAT_SLIDERS["diffusion.time"][0], maximum=FLOAT_SLIDERS["diffusion.time"][1], step=FLOAT_SLIDERS["diffusion.time"][2], value=1, label="diffusion.time")
+        with gr.Row():
+            nloop_val = gr.Slider(minimum=INT_SLIDERS["opt.n_loop"][0], maximum=INT_SLIDERS["opt.n_loop"][1], step=INT_SLIDERS["opt.n_loop"][2], value=300, label="opt.n_loop")
+            sds_val = gr.Slider(minimum=FLOAT_SLIDERS["loss.sds"][0], maximum=FLOAT_SLIDERS["loss.sds"][1], step=FLOAT_SLIDERS["loss.sds"][2], value=1, label="loss.sds")
+            proper_val = gr.Slider(minimum=FLOAT_SLIDERS["loss.proper"][0], maximum=FLOAT_SLIDERS["loss.proper"][1], step=FLOAT_SLIDERS["loss.proper"][2], value=1, label="loss.proper")
+    
+    with gr.Row():
+        init_btn = gr.Button("Init", variant="primary")
+        run_btn = gr.Button("Run", variant="secondary")
+
+    gr.Markdown("### Outputs\nEach stage exports both **GLB** (preview below) and **PLY** (download links) with per‑vertex colors.")
+    with gr.Tab("Init"):
+        with gr.Row():
+            init_view_a = gr.Model3D(label="Shape")
+            init_view_b = gr.Model3D(label="Target correspondence (init)")
+        with gr.Row():
+            out_file_init = gr.File(label="Download correspondences TXT")
+    with gr.Tab("Run"):
+        with gr.Row():
+            run_view_a = gr.Model3D(label="Shape")
+            run_view_b = gr.Model3D(label="Target correspondence (run)")
+        with gr.Row():
+            out_file = gr.File(label="Download correspondences TXT")
+
+    init_btn.click(
+        fn=init_clicked,
+        inputs=[mesh1, mesh2, lambda_val, zoomout_val, time_val, nloop_val, sds_val, proper_val],
+        outputs=[init_view_a, init_view_b, out_file_init],
+        api_name="init",
+    )
+
+    run_btn.click(
+        fn=run_clicked,
+        inputs=[mesh1, mesh2, yaml_file, lambda_val, zoomout_val, time_val, nloop_val, sds_val, proper_val],
+        outputs=[run_view_a, run_view_b, out_file],
+        api_name="run",
+    )
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Launch the gradio demo")
+    parser.add_argument('--config', type=str, default="config/matching/sds.yaml", help='Config file location')
+    parser.add_argument('--share', action="store_true")
+    args = parser.parse_args()
+    cfg = OmegaConf.load(args.config)
+    matcher = zero_shot.Matcher(cfg)
+    datadicts = None
+    demo.launch(share=args.share)

notebook_helpers.py

@@ -0,0 +1,88 @@
+from utils.mesh import load_mesh
+from utils.geometry import get_operators, load_operators
+import os
+from utils.utils_func import convert_dict
+from utils.surfaces import Surface
+import numpy as np
+
+
+device = "cuda:0"
+
+def load_data(file, cache_path, name, num_evecs=128, make_cache=False, factor=None):
+    if factor is None:
+        verts_shape, faces, vnormals, area_shape, center_shape = load_mesh(file, return_vnormals=True)
+    else:
+        verts_shape, faces, vnormals, area_shape, center_shape = load_mesh(file, return_vnormals=True, scale=False)
+        verts_shape = verts_shape/factor
+        area_shape /= factor**2
+    # print("Cache is: ", cache_path)
+    if not os.path.exists(cache_path) or make_cache:
+        print("Computing operators ...")
+        get_operators(verts_shape, faces, num_evecs, cache_path, vnormals)
+    data_dict = load_operators(cache_path)
+    data_dict['name'] = name
+    data_dict['normals'] = vnormals
+    data_dict['vertices'] = verts_shape
+    data_dict_torch = convert_dict(data_dict, device)
+    #batchify_dict(data_dict_torch)
+    return data_dict_torch, area_shape
+
+def get_map_info(file_1, file_2, dict_1, dict_2, dataset):
+    shape_dict, target_dict = dict_1, dict_2
+    name_1, name_2 = shape_dict["name"], target_dict["name"]
+    if dataset is None:
+        vts_1, vts_2 = np.arange(shape_dict['vertices'].shape[0]), np.arange(target_dict['vertices'].shape[0])
+        map_info = (file_1, file_2, vts_1, vts_2)
+    file_vts_1 = file_1.replace("shapes", "correspondences")[:-4] + ".vts"
+    vts_1 = np.loadtxt(file_vts_1).astype(np.int32) - 1
+
+    file_vts_2 = file_2.replace("shapes", "correspondences")[:-4] + ".vts"
+    vts_2 = np.loadtxt(file_vts_2).astype(np.int32) - 1
+    if "DT4D" in dataset:    
+        file_vts_1 = file_1.replace("shapes", "correspondences")[:-4] + ".vts"
+        vts_1 = np.loadtxt(file_vts_1).astype(np.int32) - 1
+    
+        file_vts_2 = file_2.replace("shapes", "correspondences")[:-4] + ".vts"
+        vts_2 = np.loadtxt(file_vts_2).astype(np.int32) - 1
+        if name_1 == name_2:
+            map_info = (file_1, file_2, vts_1, vts_2)
+        elif ("crypto" in name_1) or ("crypto" in name_2):
+            name_cat_1, name_cat_2 = name_1.split(os.sep)[0], name_2.split(os.sep)[0]
+            data_path = os.path.dirname(os.path.dirname(os.path.dirname(file_1)))
+            map_file = os.path.join(data_path, "correspondences/cross_category_corres", f"{name_cat_1}_{name_cat_2}.vts")
+            if os.path.exists(map_file):
+                map_idx = np.loadtxt(map_file).astype(np.int32) - 1
+                map_info = (file_1, file_2, vts_1, vts_2[map_idx])
+        else:
+            print("NO GROUND TRUTH PAIRS")
+    else:
+        map_info = (file_1,  file_2, vts_1, vts_2)
+    return map_info
+
+
+def load_pair(cache, id_1, id_2, name_1, name_2, dataset):
+    if "SCAPE" in dataset:
+        os.makedirs(cache, exist_ok=True) 
+        cache_file = os.path.join(cache, f"mesh{id_1:03d}_mesh_256k_0n.npz")
+        shape_file = f"data/{dataset}/shapes/mesh{id_1:03d}.ply"
+        shape_surf = Surface(filename=shape_file)
+        shape_dict, _  = load_data(shape_file, cache_file, str(id_1))
+        
+        cache_file = os.path.join(cache, f"mesh{id_2:03d}_mesh_256k_0n.npz")
+        target_file = f"data/{dataset}/shapes/mesh{id_2:03d}.ply"
+        target_surf = Surface(filename=target_file)
+        target_dict, _  = load_data(target_file, cache_file, str(id_2))
+        map_info = get_map_info(shape_file, target_file, shape_dict, target_dict, "SCAPE")
+    elif "DT4D" in dataset:
+        cache_file = os.path.join(cache, f"{name_1}_mesh_256k_0n.npz")
+        shape_file = f"data/DT4D_r_ori/shapes/{name_1}.ply"
+        shape_surf = Surface(filename=shape_file)
+        shape_dict, _  = load_data(shape_file, cache_file, name_1)
+        cache_file = os.path.join(cache, f"{name_2}_mesh_256k_0n.npz")
+        # cache_file = f"../nonrigiddiff/cache/snk/{name_2}.npz"
+        cache_file = f"../nonrigiddiff/cache/attentive/{name_2}_mesh_256k_0n.npz"
+        target_file = f"data/DT4D_r_ori/shapes/{name_2}.ply"
+        target_surf = Surface(filename=target_file)
+        target_dict, _  = load_data(target_file, cache_file, name_2)
+        map_info = get_map_info(shape_file, target_file, shape_dict, target_dict, "DT4D")
+    return shape_surf, target_surf, shape_dict, target_dict, map_info

requirements.txt

@@ -0,0 +1,45 @@
+comm==0.2.1
+configer
+configparser==6.0.0
+contourpy==1.1.1
+cycler==0.12.1
+fonttools==4.47.2
+freetype-py==2.4.0
+imageio==2.33.1
+ipygany==0.5.0
+ipywidgets==8.1.7
+nbformat==5.10.4
+jupyter
+kiwisolver==1.4.5
+lxml==5.1.0
+matplotlib==3.7.4
+networkx==3.1
+open3d>=0.18.0
+pyglet==2.0.10
+pyopengl==3.1.0
+pyparsing==3.1.1
+pyrender==0.1.45
+scipy>=1.10.1
+shapely==2.0.2
+smplx[all]
+tqdm==4.66.1
+trimesh==4.0.10
+vtk==9.3.0
+timm==1.0.11
+potpourri3d 
+pythreejs==2.4.2
+lapy
+roma
+webdataset
+lmdb
+robust_laplacian
+termcolor
+omegaconf
+pykeops==2.2.3
+scikit-image
+pyfmaps
+wandb
+gradio==4.44.1
+pydantic==2.10.6
+traitlets==5.7.1
+https://github.com/skoch9/meshplot/archive/0.4.0.tar.gz

shape_models/__init__.py

@@ -0,0 +1,3 @@
+from .utils import *
+from .geometry import *
+from .layers import *

shape_models/encoder.py

@@ -0,0 +1,24 @@
+import torch 
+import torch.nn as nn
+from .layers import DiffusionNet
+
+class Encoder(nn.Module):
+    def __init__(self, with_grad=True, key_verts="vertices"):
+        super(Encoder, self).__init__()
+        self.diff_net = DiffusionNet(
+             C_in=3,
+             C_out=512,
+             C_width=128,
+             N_block=4,
+             dropout=True,
+             with_gradient_features=with_grad,
+             with_gradient_rotations=with_grad,
+        )
+        self.key_verts = key_verts
+
+
+    def forward(self, shape_dict):
+        feats = self.diff_net(shape_dict[self.key_verts], shape_dict["mass"], shape_dict["L"], evals=shape_dict["evals"], 
+                               evecs=shape_dict["evecs"], gradX=shape_dict["gradX"], gradY=shape_dict["gradY"], faces=shape_dict["faces"])
+        x_out = torch.max(feats, dim=0).values
+        return x_out

shape_models/fmap.py

@@ -0,0 +1,186 @@
+from copy import deepcopy
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# feature extractor
+from .layers import DiffusionNet
+from omegaconf import OmegaConf
+
+
+def get_mask(evals1, evals2, gamma=0.5, device="cpu"):
+    scaling_factor = max(torch.max(evals1), torch.max(evals2))
+    evals1, evals2 = evals1.to(device) / scaling_factor, evals2.to(device) / scaling_factor
+    evals_gamma1, evals_gamma2 = (evals1 ** gamma)[None, :], (evals2 ** gamma)[:, None]
+
+    M_re = evals_gamma2 / (evals_gamma2.square() + 1) - evals_gamma1 / (evals_gamma1.square() + 1)
+    M_im = 1 / (evals_gamma2.square() + 1) - 1 / (evals_gamma1.square() + 1)
+    return M_re.square() + M_im.square()
+
+def get_CXX(feat_x, feat_y, evecs_trans_x, evecs_trans_y):
+    # compute linear operator matrix representation C1 and C2
+
+    F_hat = torch.bmm(evecs_trans_x, feat_x)
+    G_hat = torch.bmm(evecs_trans_y, feat_y)
+    A, B = F_hat, G_hat
+    
+    A_t, B_t = A.transpose(1, 2), B.transpose(1, 2)
+    A_A_t, B_B_t = torch.bmm(A, A_t), torch.bmm(B, B_t)
+    B_A_t, A_B_t = torch.bmm(B, A_t), torch.bmm(A, B_t)
+
+    C12 = torch.bmm(B_A_t, torch.inverse(A_A_t))
+
+    C21 = torch.bmm(A_B_t, torch.inverse(B_B_t))
+    return [C12, C21]
+
+def get_mask_noise(C, diff_model, scale=1, N_est=200, device="cuda", normalize=True, absolute=False):
+    with torch.no_grad():
+        sig = torch.ones([1, 1, 1, 1], device=device) * scale
+        
+        noise = torch.randn((N_est, 1, 30, 30), device=device)
+        if absolute:
+            #noisy_new = torch.abs(C[None, :, :] + noise *scale)
+            noisy_new = torch.abs(torch.abs(C[None, :, :]) + noise * scale)
+        else:
+            noisy_new = C[None, :, :] + noise *scale
+        denoised = diff_model.net(noisy_new, sig)
+        mask_squared = torch.mean(torch.abs(noisy_new - denoised)/(2*scale), dim=0)/torch.mean(torch.abs(noisy_new), dim=0)
+        mask_median = torch.median(mask_squared).item()
+        if normalize:
+            mask_median = torch.median(mask_squared).item()
+            M_denoised = torch.sqrt(torch.clamp(mask_squared-mask_median/2, 0, mask_median*2))
+        else:
+            M_denoised = torch.sqrt(mask_squared)
+    return M_denoised-M_denoised.min()#, mask_median*2)
+
+
+class RegularizedFMNet(nn.Module):
+    """Compute the functional map matrix representation."""
+
+    def __init__(self, lambda_=1e-3, resolvant_gamma=0.5, use_resolvent=False):
+        super().__init__()
+        self.lambda_ = lambda_
+        self.use_resolvent = use_resolvent
+        self.resolvant_gamma = resolvant_gamma
+
+    def forward(self, feat_x, feat_y, evals_x, evals_y, evecs_trans_x, evecs_trans_y, diff_conf=None):
+        # compute linear operator matrix representation C1 and C2
+        evecs_trans_x, evecs_trans_y = evecs_trans_x.unsqueeze(0), evecs_trans_y.unsqueeze(0)
+        evals_x, evals_y = evals_x.unsqueeze(0), evals_y.unsqueeze(0)
+
+        F_hat = torch.bmm(evecs_trans_x, feat_x)
+        G_hat = torch.bmm(evecs_trans_y, feat_y)
+        A, B = F_hat, G_hat
+
+
+
+        if diff_conf is not None:
+            diff_model, scale, normalize, absolute, N_est = diff_conf
+            C12_raw, C21_raw = get_CXX(feat_x, feat_y, evecs_trans_x, evecs_trans_y)
+            D12 = get_mask_noise(C12_raw, diff_model, scale=scale, N_est=N_est, normalize=normalize, absolute=absolute)
+            D21 = get_mask_noise(C21_raw, diff_model, scale=scale, N_est=N_est, normalize=normalize, absolute=absolute)
+        elif self.use_resolvent:
+            D12 = get_mask(evals_x.flatten(), evals_y.flatten(), self.resolvant_gamma, feat_x.device).unsqueeze(0)
+            D21 = get_mask(evals_y.flatten(), evals_x.flatten(), self.resolvant_gamma, feat_x.device).unsqueeze(0)
+        else:
+            D12 = (torch.unsqueeze(evals_y, 2) - torch.unsqueeze(evals_x, 1))**2
+            D21 = (torch.unsqueeze(evals_x, 2) - torch.unsqueeze(evals_y, 1))**2
+        
+        A_t, B_t = A.transpose(1, 2), B.transpose(1, 2)
+        A_A_t, B_B_t = torch.bmm(A, A_t), torch.bmm(B, B_t)
+        B_A_t, A_B_t = torch.bmm(B, A_t), torch.bmm(A, B_t)
+
+        C12_i = []
+        for i in range(evals_x.size(1)):
+            D12_i = torch.cat([torch.diag(D12[bs, i, :].flatten()).unsqueeze(0) for bs in range(evals_x.size(0))], dim=0)
+            C12 = torch.bmm(torch.inverse(A_A_t + self.lambda_ * D12_i), B_A_t[:, i, :].unsqueeze(1).transpose(1, 2))
+            C12_i.append(C12.transpose(1, 2))
+        C12 = torch.cat(C12_i, dim=1)
+
+        C21_i = []
+        for i in range(evals_y.size(1)):
+            D21_i = torch.cat([torch.diag(D21[bs, i, :].flatten()).unsqueeze(0) for bs in range(evals_y.size(0))], dim=0)
+            C21 = torch.bmm(torch.inverse(B_B_t + self.lambda_ * D21_i), A_B_t[:, i, :].unsqueeze(1).transpose(1, 2))
+            C21_i.append(C21.transpose(1, 2))
+        C21 = torch.cat(C21_i, dim=1)
+        if diff_conf is not None:
+            return [C12_raw, C12, D12], [C21_raw, C21, D21]
+        else:
+            return [C12, C21]
+
+
+class DFMNet(nn.Module):
+    """
+    Compilation of the global model :
+    - diffusion net as feature extractor
+    - fmap + q-fmap
+    - unsupervised loss
+    """
+
+    def __init__(self, cfg):
+        super().__init__()
+
+        # feature extractor #
+        with_grad=True
+        self.feat = cfg["feat"]
+
+        self.feature_extractor = DiffusionNet(
+             C_in=cfg["C_in"],
+             C_out=cfg["n_feat"],
+             C_width=128,
+             N_block=4,
+             dropout=True,
+             with_gradient_features=with_grad,
+             with_gradient_rotations=with_grad,
+        )
+
+        # regularized fmap
+        self.fmreg_net = RegularizedFMNet(lambda_=cfg["lambda_"],
+                                          resolvant_gamma=cfg.get("resolvent_gamma", 0.5), use_resolvent=cfg.get("use_resolvent", False))
+        # parameters
+        self.n_fmap = cfg["n_fmap"]
+        if cfg.get("diffusion", None) is not None:
+            self.normalize = cfg["diffusion"]["normalize"]
+            self.abs = cfg["diffusion"]["abs"]
+            self.N_est = cfg.diffusion.get("batch_mask", 100)
+
+    def forward(self, batch, diff_model=None, scale=1):
+        if self.feat == "xyz":
+            feat_1, feat_2 = batch["shape1"]["vertices"], batch["shape2"]["vertices"]
+        elif self.feat == "wks":
+            feat_1, feat_2 = batch["shape1"]["wks"], batch["shape2"]["wks"]
+        elif self.feat == "hks":
+            feat_1, feat_2 = batch["shape1"]["hks"], batch["shape2"]["hks"]
+        else:
+            raise Exception("Unknow Feature")
+
+
+        verts1, faces1, mass1, L1, evals1, evecs1, gradX1, gradY1 = (feat_1, batch["shape1"]["faces"],
+                                                                     batch["shape1"]["mass"], batch["shape1"]["L"],
+                                                                     batch["shape1"]["evals"], batch["shape1"]["evecs"],
+                                                                     batch["shape1"]["gradX"], batch["shape1"]["gradY"])
+        verts2, faces2, mass2, L2, evals2, evecs2, gradX2, gradY2 = (feat_2, batch["shape2"]["faces"],
+                                                                     batch["shape2"]["mass"], batch["shape2"]["L"],
+                                                                     batch["shape2"]["evals"], batch["shape2"]["evecs"],
+                                                                     batch["shape2"]["gradX"], batch["shape2"]["gradY"])
+
+        # set features to vertices
+        features1, features2 = verts1, verts2
+        # print(features1.shape, features2.shape)
+
+        feat1 = self.feature_extractor(features1, mass1, L=L1, evals=evals1, evecs=evecs1,
+                                       gradX=gradX1, gradY=gradY1, faces=faces1).unsqueeze(0)
+        feat2 = self.feature_extractor(features2, mass2, L=L2, evals=evals2, evecs=evecs2,
+                                       gradX=gradX2, gradY=gradY2, faces=faces2).unsqueeze(0)
+
+        evecs_trans1, evecs_trans2 = evecs1.t()[:self.n_fmap] @ torch.diag(mass1.squeeze()), evecs2.squeeze().t()[:self.n_fmap].squeeze() @ torch.diag(mass2.squeeze())
+        evals1, evals2 = evals1[:self.n_fmap], evals2[:self.n_fmap]
+
+        #
+        if diff_model is not None:
+            C12_pred, C21_pred = self.fmreg_net(feat1, feat2, evals1, evals2, evecs_trans1, evecs_trans2,
+                                            diff_conf=[diff_model, scale, self.normalize, self.abs, self.N_est])
+        else:
+            C12_pred, C21_pred = self.fmreg_net(feat1, feat2, evals1, evals2, evecs_trans1, evecs_trans2)
+
+        return C12_pred, C21_pred, feat1, feat2, evecs_trans1, evecs_trans2

shape_models/geometry.py

@@ -0,0 +1,897 @@
+import scipy
+import scipy.sparse.linalg as sla
+# ^^^ we NEED to import scipy before torch, or it crashes :(
+# (observed on Ubuntu 20.04 w/ torch 1.6.0 and scipy 1.5.2 installed via conda)
+
+import os.path
+import sys
+import random
+from multiprocessing import Pool
+
+import numpy as np
+import scipy.spatial
+import torch
+import sklearn.neighbors
+
+import robust_laplacian
+import potpourri3d as pp3d
+
+
+def norm(x, highdim=False):
+    """
+    Computes norm of an array of vectors. Given (shape,d), returns (shape) after norm along last dimension
+    """
+    return torch.norm(x, dim=len(x.shape) - 1)
+
+
+def norm2(x, highdim=False):
+    """
+    Computes norm^2 of an array of vectors. Given (shape,d), returns (shape) after norm along last dimension
+    """
+    return dot(x, x)
+
+
+def normalize(x, divide_eps=1e-6, highdim=False):
+    """
+    Computes norm^2 of an array of vectors. Given (shape,d), returns (shape) after norm along last dimension
+    """
+    if (len(x.shape) == 1):
+        raise ValueError("called normalize() on single vector of dim " + str(x.shape) + " are you sure?")
+    if (not highdim and x.shape[-1] > 4):
+        raise ValueError("called normalize() with large last dimension " + str(x.shape) + " are you sure?")
+    return x / (norm(x, highdim=highdim) + divide_eps).unsqueeze(-1)
+
+
+def face_coords(verts, faces):
+    coords = verts[faces]
+    return coords
+
+
+def cross(vec_A, vec_B):
+    return torch.cross(vec_A, vec_B, dim=-1)
+
+
+def dot(vec_A, vec_B):
+    return torch.sum(vec_A * vec_B, dim=-1)
+
+
+# Given (..., 3) vectors and normals, projects out any components of vecs
+# which lies in the direction of normals. Normals are assumed to be unit.
+
+
+def project_to_tangent(vecs, unit_normals):
+    dots = dot(vecs, unit_normals)
+    return vecs - unit_normals * dots.unsqueeze(-1)
+
+
+def face_area(verts, faces):
+    coords = face_coords(verts, faces)
+    vec_A = coords[:, 1, :] - coords[:, 0, :]
+    vec_B = coords[:, 2, :] - coords[:, 0, :]
+
+    raw_normal = cross(vec_A, vec_B)
+    return 0.5 * norm(raw_normal)
+
+
+def face_normals(verts, faces, normalized=True):
+    coords = face_coords(verts, faces)
+    vec_A = coords[:, 1, :] - coords[:, 0, :]
+    vec_B = coords[:, 2, :] - coords[:, 0, :]
+
+    raw_normal = cross(vec_A, vec_B)
+
+    if normalized:
+        return normalize(raw_normal)
+
+    return raw_normal
+
+
+def neighborhood_normal(points):
+    # points: (N, K, 3) array of neighborhood psoitions
+    # points should be centered at origin
+    # out: (N,3) array of normals
+    # numpy in, numpy out
+    (u, s, vh) = np.linalg.svd(points, full_matrices=False)
+    normal = vh[:, 2, :]
+    return normal / np.linalg.norm(normal, axis=-1, keepdims=True)
+
+
+def mesh_vertex_normals(verts, faces):
+    # numpy in / out
+    face_n = toNP(face_normals(torch.tensor(verts), torch.tensor(faces)))  # ugly torch <---> numpy
+
+    vertex_normals = np.zeros(verts.shape)
+    for i in range(3):
+        np.add.at(vertex_normals, faces[:, i], face_n)
+
+    vertex_normals = vertex_normals / np.linalg.norm(vertex_normals, axis=-1, keepdims=True)
+
+    return vertex_normals
+
+
+def vertex_normals(verts, faces, n_neighbors_cloud=30):
+    verts_np = toNP(verts)
+
+    if faces.numel() == 0:  # point cloud
+
+        _, neigh_inds = find_knn(verts, verts, n_neighbors_cloud, omit_diagonal=True, method='cpu_kd')
+        neigh_points = verts_np[neigh_inds, :]
+        neigh_points = neigh_points - verts_np[:, np.newaxis, :]
+        normals = neighborhood_normal(neigh_points)
+
+    else:  # mesh
+
+        normals = mesh_vertex_normals(verts_np, toNP(faces))
+
+        # if any are NaN, wiggle slightly and recompute
+        bad_normals_mask = np.isnan(normals).any(axis=1, keepdims=True)
+        if bad_normals_mask.any():
+            bbox = np.amax(verts_np, axis=0) - np.amin(verts_np, axis=0)
+            scale = np.linalg.norm(bbox) * 1e-4
+            wiggle = (np.random.RandomState(seed=777).rand(*verts.shape) - 0.5) * scale
+            wiggle_verts = verts_np + bad_normals_mask * wiggle
+            normals = mesh_vertex_normals(wiggle_verts, toNP(faces))
+
+        # if still NaN assign random normals (probably means unreferenced verts in mesh)
+        bad_normals_mask = np.isnan(normals).any(axis=1)
+        if bad_normals_mask.any():
+            normals[bad_normals_mask, :] = (np.random.RandomState(seed=777).rand(*verts.shape) - 0.5)[bad_normals_mask, :]
+            normals = normals / np.linalg.norm(normals, axis=-1)[:, np.newaxis]
+
+    normals = torch.from_numpy(normals).to(device=verts.device, dtype=verts.dtype)
+
+    if torch.any(torch.isnan(normals)):
+        raise ValueError("NaN normals :(")
+
+    return normals
+
+
+def build_tangent_frames(verts, faces, normals=None):
+
+    V = verts.shape[0]
+    dtype = verts.dtype
+    device = verts.device
+
+    if normals == None:
+        vert_normals = vertex_normals(verts, faces)  # (V,3)
+    else:
+        vert_normals = normals
+
+    # = find an orthogonal basis
+
+    basis_cand1 = torch.tensor([1, 0, 0]).to(device=device, dtype=dtype).expand(V, -1)
+    basis_cand2 = torch.tensor([0, 1, 0]).to(device=device, dtype=dtype).expand(V, -1)
+
+    basisX = torch.where((torch.abs(dot(vert_normals, basis_cand1)) < 0.9).unsqueeze(-1), basis_cand1, basis_cand2)
+    basisX = project_to_tangent(basisX, vert_normals)
+    basisX = normalize(basisX)
+    basisY = cross(vert_normals, basisX)
+    frames = torch.stack((basisX, basisY, vert_normals), dim=-2)
+
+    if torch.any(torch.isnan(frames)):
+        raise ValueError("NaN coordinate frame! Must be very degenerate")
+
+    return frames
+
+
+def build_grad_point_cloud(verts, frames, n_neighbors_cloud=30):
+
+    verts_np = toNP(verts)
+    frames_np = toNP(frames)
+
+    _, neigh_inds = find_knn(verts, verts, n_neighbors_cloud, omit_diagonal=True, method='cpu_kd')
+    neigh_points = verts_np[neigh_inds, :]
+    neigh_vecs = neigh_points - verts_np[:, np.newaxis, :]
+
+    # TODO this could easily be way faster. For instance we could avoid the weird edges format and the corresponding pure-python loop via some numpy broadcasting of the same logic. The way it works right now is just to share code with the mesh version. But its low priority since its preprocessing code.
+
+    edge_inds_from = np.repeat(np.arange(verts.shape[0]), n_neighbors_cloud)
+    edges = np.stack((edge_inds_from, neigh_inds.flatten()))
+    edge_tangent_vecs = edge_tangent_vectors(verts, frames, edges)
+
+    return build_grad(verts_np, torch.tensor(edges), edge_tangent_vecs)
+
+
+def edge_tangent_vectors(verts, frames, edges):
+    edge_vecs = verts[edges[1, :], :] - verts[edges[0, :], :]
+    basisX = frames[edges[0, :], 0, :]
+    basisY = frames[edges[0, :], 1, :]
+
+    compX = dot(edge_vecs, basisX)
+    compY = dot(edge_vecs, basisY)
+    edge_tangent = torch.stack((compX, compY), dim=-1)
+
+    return edge_tangent
+
+
+def build_grad(verts, edges, edge_tangent_vectors):
+    """
+    Build a (V, V) complex sparse matrix grad operator. Given real inputs at vertices, produces a complex (vector value) at vertices giving the gradient. All values pointwise.
+    - edges: (2, E)
+    """
+
+    edges_np = toNP(edges)
+    edge_tangent_vectors_np = toNP(edge_tangent_vectors)
+
+    # TODO find a way to do this in pure numpy?
+
+    # Build outgoing neighbor lists
+    N = verts.shape[0]
+    vert_edge_outgoing = [[] for i in range(N)]
+    for iE in range(edges_np.shape[1]):
+        tail_ind = edges_np[0, iE]
+        tip_ind = edges_np[1, iE]
+        if tip_ind != tail_ind:
+            vert_edge_outgoing[tail_ind].append(iE)
+
+    # Build local inversion matrix for each vertex
+    row_inds = []
+    col_inds = []
+    data_vals = []
+    eps_reg = 1e-5
+    for iV in range(N):
+        n_neigh = len(vert_edge_outgoing[iV])
+
+        lhs_mat = np.zeros((n_neigh, 2))
+        rhs_mat = np.zeros((n_neigh, n_neigh + 1))
+        ind_lookup = [iV]
+        for i_neigh in range(n_neigh):
+            iE = vert_edge_outgoing[iV][i_neigh]
+            jV = edges_np[1, iE]
+            ind_lookup.append(jV)
+
+            edge_vec = edge_tangent_vectors[iE][:]
+            w_e = 1.
+
+            lhs_mat[i_neigh][:] = w_e * edge_vec
+            rhs_mat[i_neigh][0] = w_e * (-1)
+            rhs_mat[i_neigh][i_neigh + 1] = w_e * 1
+
+        lhs_T = lhs_mat.T
+        lhs_inv = np.linalg.inv(lhs_T @ lhs_mat + eps_reg * np.identity(2)) @ lhs_T
+
+        sol_mat = lhs_inv @ rhs_mat
+        sol_coefs = (sol_mat[0, :] + 1j * sol_mat[1, :]).T
+
+        for i_neigh in range(n_neigh + 1):
+            i_glob = ind_lookup[i_neigh]
+
+            row_inds.append(iV)
+            col_inds.append(i_glob)
+            data_vals.append(sol_coefs[i_neigh])
+
+    # build the sparse matrix
+    row_inds = np.array(row_inds)
+    col_inds = np.array(col_inds)
+    data_vals = np.array(data_vals)
+    mat = scipy.sparse.coo_matrix((data_vals, (row_inds, col_inds)), shape=(N, N)).tocsc()
+
+    return mat
+
+
+def compute_operators(verts, faces, k_eig, normals=None):
+    """
+    Builds spectral operators for a mesh/point cloud. Constructs mass matrix, eigenvalues/vectors for Laplacian, and gradient matrix.
+
+    See get_operators() for a similar routine that wraps this one with a layer of caching.
+
+    Torch in / torch out.
+
+    Arguments:
+      - vertices: (V,3) vertex positions
+      - faces: (F,3) list of triangular faces. If empty, assumed to be a point cloud.
+      - k_eig: number of eigenvectors to use
+
+    Returns:
+      - frames: (V,3,3) X/Y/Z coordinate frame at each vertex. Z coordinate is normal (e.g. [:,2,:] for normals)
+      - massvec: (V) real diagonal of lumped mass matrix
+      - L: (VxV) real sparse matrix of (weak) Laplacian
+      - evals: (k) list of eigenvalues of the Laplacian
+      - evecs: (V,k) list of eigenvectors of the Laplacian 
+      - gradX: (VxV) sparse matrix which gives X-component of gradient in the local basis at the vertex
+      - gradY: same as gradX but for Y-component of gradient
+
+    PyTorch doesn't seem to like complex sparse matrices, so we store the "real" and "imaginary" (aka X and Y) gradient matrices separately, rather than as one complex sparse matrix.
+
+    Note: for a generalized eigenvalue problem, the mass matrix matters! The eigenvectors are only othrthonormal with respect to the mass matrix, like v^H M v, so the mass (given as the diagonal vector massvec) needs to be used in projections, etc.
+    """
+
+    device = verts.device
+    dtype = verts.dtype
+    V = verts.shape[0]
+    is_cloud = faces.numel() == 0
+
+    eps = 1e-8
+
+    verts_np = toNP(verts).astype(np.float64)
+    faces_np = toNP(faces)
+    frames = build_tangent_frames(verts, faces, normals=normals)
+    frames_np = toNP(frames)
+
+    # Build the scalar Laplacian
+    if is_cloud:
+        L, M = robust_laplacian.point_cloud_laplacian(verts_np)
+        massvec_np = M.diagonal()
+    else:
+        # L, M = robust_laplacian.mesh_laplacian(verts_np, faces_np)
+        # massvec_np = M.diagonal()
+        L = pp3d.cotan_laplacian(verts_np, faces_np, denom_eps=1e-10)
+        massvec_np = pp3d.vertex_areas(verts_np, faces_np)
+        massvec_np += eps * np.mean(massvec_np)
+
+    if (np.isnan(L.data).any()):
+        raise RuntimeError("NaN Laplace matrix")
+    if (np.isnan(massvec_np).any()):
+        raise RuntimeError("NaN mass matrix")
+
+    # Read off neighbors & rotations from the Laplacian
+    L_coo = L.tocoo()
+    inds_row = L_coo.row
+    inds_col = L_coo.col
+
+    # === Compute the eigenbasis
+    if k_eig > 0:
+
+        # Prepare matrices
+        L_eigsh = (L + scipy.sparse.identity(L.shape[0]) * eps).tocsc()
+        massvec_eigsh = massvec_np
+        Mmat = scipy.sparse.diags(massvec_eigsh)
+        eigs_sigma = eps
+
+        failcount = 0
+        while True:
+            try:
+                # We would be happy here to lower tol or maxiter since we don't need these to be super precise, but for some reason those parameters seem to have no effect
+                evals_np, evecs_np = sla.eigsh(L_eigsh, k=k_eig, M=Mmat, sigma=eigs_sigma)
+
+                # Clip off any eigenvalues that end up slightly negative due to numerical weirdness
+                evals_np = np.clip(evals_np, a_min=0., a_max=float('inf'))
+
+                break
+            except Exception as e:
+                print(e)
+                if (failcount > 3):
+                    raise ValueError("failed to compute eigendecomp")
+                failcount += 1
+                print("--- decomp failed; adding eps ===> count: " + str(failcount))
+                L_eigsh = L_eigsh + scipy.sparse.identity(L.shape[0]) * (eps * 10**failcount)
+
+    else:  #k_eig == 0
+        evals_np = np.zeros((0))
+        evecs_np = np.zeros((verts.shape[0], 0))
+
+    # == Build gradient matrices
+
+    # For meshes, we use the same edges as were used to build the Laplacian. For point clouds, use a whole local neighborhood
+    if is_cloud:
+        grad_mat_np = build_grad_point_cloud(verts, frames)
+    else:
+        edges = torch.tensor(np.stack((inds_row, inds_col), axis=0), device=device, dtype=faces.dtype)
+        edge_vecs = edge_tangent_vectors(verts, frames, edges)
+        grad_mat_np = build_grad(verts.cpu(), edges.cpu(), edge_vecs.cpu())
+
+    # Split complex gradient in to two real sparse mats (torch doesn't like complex sparse matrices)
+    gradX_np = np.real(grad_mat_np)
+    gradY_np = np.imag(grad_mat_np)
+
+    # === Convert back to torch
+    massvec = torch.from_numpy(massvec_np).to(device=device, dtype=dtype)
+    L = utils.sparse_np_to_torch(L).to(device=device, dtype=dtype)
+    evals = torch.from_numpy(evals_np).to(device=device, dtype=dtype)
+    evecs = torch.from_numpy(evecs_np).to(device=device, dtype=dtype)
+    gradX = utils.sparse_np_to_torch(gradX_np).to(device=device, dtype=dtype)
+    gradY = utils.sparse_np_to_torch(gradY_np).to(device=device, dtype=dtype)
+
+    return frames, massvec, L, evals, evecs, gradX, gradY
+
+
+def get_all_operators(verts_list, faces_list, k_eig, op_cache_dir=None, normals=None):
+    N = len(verts_list)
+
+    frames = [None] * N
+    massvec = [None] * N
+    L = [None] * N
+    evals = [None] * N
+    evecs = [None] * N
+    gradX = [None] * N
+    gradY = [None] * N
+
+    inds = [i for i in range(N)]
+    # process in random order
+    # random.shuffle(inds)
+
+    for num, i in enumerate(inds):
+        print("get_all_operators() processing {} / {} {:.3f}%".format(num, N, num / N * 100))
+        if normals is None:
+            outputs = get_operators(verts_list[i], faces_list[i], k_eig, op_cache_dir)
+        else:
+            outputs = get_operators(verts_list[i], faces_list[i], k_eig, op_cache_dir, normals=normals[i])
+        frames[i] = outputs[0]
+        massvec[i] = outputs[1]
+        L[i] = outputs[2]
+        evals[i] = outputs[3]
+        evecs[i] = outputs[4]
+        gradX[i] = outputs[5]
+        gradY[i] = outputs[6]
+
+    return frames, massvec, L, evals, evecs, gradX, gradY
+
+
+def get_operators(verts, faces, k_eig=128, op_cache_dir=None, normals=None, overwrite_cache=False):
+    """
+    See documentation for compute_operators(). This essentailly just wraps a call to compute_operators, using a cache if possible.
+    All arrays are always computed using double precision for stability, then truncated to single precision floats to store on disk, and finally returned as a tensor with dtype/device matching the `verts` input.
+    """
+
+    device = verts.device
+    dtype = verts.dtype
+    verts_np = toNP(verts)
+    faces_np = toNP(faces)
+    is_cloud = faces.numel() == 0
+
+    if (np.isnan(verts_np).any()):
+        raise RuntimeError("tried to construct operators from NaN verts")
+
+    # Check the cache directory
+    # Note 1: Collisions here are exceptionally unlikely, so we could probably just use the hash...
+    #         but for good measure we check values nonetheless.
+    # Note 2: There is a small possibility for race conditions to lead to bucket gaps or duplicate
+    #         entries in this cache. The good news is that that is totally fine, and at most slightly
+    #         slows performance with rare extra cache misses.
+    found = False
+    if op_cache_dir is not None:
+        utils.ensure_dir_exists(op_cache_dir)
+        hash_key_str = str(utils.hash_arrays((verts_np, faces_np)))
+        # print("Building operators for input with hash: " + hash_key_str)
+
+        # Search through buckets with matching hashes.  When the loop exits, this
+        # is the bucket index of the file we should write to.
+        i_cache_search = 0
+        while True:
+
+            # Form the name of the file to check
+            search_path = os.path.join(op_cache_dir, hash_key_str + "_" + str(i_cache_search) + ".npz")
+
+            try:
+                # print('loading path: ' + str(search_path))
+                npzfile = np.load(search_path, allow_pickle=True)
+                cache_verts = npzfile["verts"]
+                cache_faces = npzfile["faces"]
+                cache_k_eig = npzfile["k_eig"].item()
+
+                # If the cache doesn't match, keep looking
+                if (not np.array_equal(verts, cache_verts)) or (not np.array_equal(faces, cache_faces)):
+                    i_cache_search += 1
+                    print("hash collision! searching next.")
+                    continue
+
+                # print("  cache hit!")
+
+                # If we're overwriting, or there aren't enough eigenvalues, just delete it; we'll create a new
+                # entry below more eigenvalues
+                if overwrite_cache:
+                    print("  overwriting cache by request")
+                    os.remove(search_path)
+                    break
+
+                if cache_k_eig < k_eig:
+                    print("  overwriting cache --- not enough eigenvalues")
+                    os.remove(search_path)
+                    break
+
+                if "L_data" not in npzfile:
+                    print("  overwriting cache --- entries are absent")
+                    os.remove(search_path)
+                    break
+
+                def read_sp_mat(prefix):
+                    data = npzfile[prefix + "_data"]
+                    indices = npzfile[prefix + "_indices"]
+                    indptr = npzfile[prefix + "_indptr"]
+                    shape = npzfile[prefix + "_shape"]
+                    mat = scipy.sparse.csc_matrix((data, indices, indptr), shape=shape)
+                    return mat
+
+                # This entry matches! Return it.
+                frames = npzfile["frames"]
+                mass = npzfile["mass"]
+                L = read_sp_mat("L")
+                evals = npzfile["evals"][:k_eig]
+                evecs = npzfile["evecs"][:, :k_eig]
+                gradX = read_sp_mat("gradX")
+                gradY = read_sp_mat("gradY")
+
+                frames = torch.from_numpy(frames).to(device=device, dtype=dtype)
+                mass = torch.from_numpy(mass).to(device=device, dtype=dtype)
+                L = utils.sparse_np_to_torch(L).to(device=device, dtype=dtype)
+                evals = torch.from_numpy(evals).to(device=device, dtype=dtype)
+                evecs = torch.from_numpy(evecs).to(device=device, dtype=dtype)
+                gradX = utils.sparse_np_to_torch(gradX).to(device=device, dtype=dtype)
+                gradY = utils.sparse_np_to_torch(gradY).to(device=device, dtype=dtype)
+
+                found = True
+
+                break
+
+            except FileNotFoundError:
+                print("  cache miss -- constructing operators")
+                break
+
+            except Exception as E:
+                print("unexpected error loading file: " + str(E))
+                print("-- constructing operators")
+                break
+
+    if not found:
+
+        # No matching entry found; recompute.
+        frames, mass, L, evals, evecs, gradX, gradY = compute_operators(verts, faces, k_eig, normals=normals)
+
+        dtype_np = np.float32
+
+        # Store it in the cache
+        if op_cache_dir is not None:
+
+            L_np = utils.sparse_torch_to_np(L).astype(dtype_np)
+            gradX_np = utils.sparse_torch_to_np(gradX).astype(dtype_np)
+            gradY_np = utils.sparse_torch_to_np(gradY).astype(dtype_np)
+
+            np.savez(
+                search_path,
+                verts=verts_np.astype(dtype_np),
+                frames=toNP(frames).astype(dtype_np),
+                faces=faces_np,
+                k_eig=k_eig,
+                mass=toNP(mass).astype(dtype_np),
+                L_data=L_np.data.astype(dtype_np),
+                L_indices=L_np.indices,
+                L_indptr=L_np.indptr,
+                L_shape=L_np.shape,
+                evals=toNP(evals).astype(dtype_np),
+                evecs=toNP(evecs).astype(dtype_np),
+                gradX_data=gradX_np.data.astype(dtype_np),
+                gradX_indices=gradX_np.indices,
+                gradX_indptr=gradX_np.indptr,
+                gradX_shape=gradX_np.shape,
+                gradY_data=gradY_np.data.astype(dtype_np),
+                gradY_indices=gradY_np.indices,
+                gradY_indptr=gradY_np.indptr,
+                gradY_shape=gradY_np.shape,
+            )
+
+    return frames, mass, L, evals, evecs, gradX, gradY
+
+
+def to_basis(values, basis, massvec):
+    """
+    Transform data in to an orthonormal basis (where orthonormal is wrt to massvec)
+    Inputs:
+      - values: (B,V,D)
+      - basis: (B,V,K)
+      - massvec: (B,V)
+    Outputs:
+      - (B,K,D) transformed values
+    """
+    basisT = basis.transpose(-2, -1)
+    return torch.matmul(basisT, values * massvec.unsqueeze(-1))
+
+
+def from_basis(values, basis):
+    """
+    Transform data out of an orthonormal basis
+    Inputs:
+      - values: (K,D)
+      - basis: (V,K)
+    Outputs:
+      - (V,D) reconstructed values
+    """
+    if values.is_complex() or basis.is_complex():
+        return utils.cmatmul(utils.ensure_complex(basis), utils.ensure_complex(values))
+    else:
+        return torch.matmul(basis, values)
+
+
+def compute_hks(evals, evecs, scales):
+    """
+    Inputs:
+      - evals: (K) eigenvalues
+      - evecs: (V,K) values
+      - scales: (S) times
+    Outputs:
+      - (V,S) hks values
+    """
+
+    # expand batch
+    if len(evals.shape) == 1:
+        expand_batch = True
+        evals = evals.unsqueeze(0)
+        evecs = evecs.unsqueeze(0)
+        scales = scales.unsqueeze(0)
+    else:
+        expand_batch = False
+
+    # TODO could be a matmul
+    power_coefs = torch.exp(-evals.unsqueeze(1) * scales.unsqueeze(-1)).unsqueeze(1)  # (B,1,S,K)
+    terms = power_coefs * (evecs * evecs).unsqueeze(2)  # (B,V,S,K)
+
+    out = torch.sum(terms, dim=-1)  # (B,V,S)
+
+    if expand_batch:
+        return out.squeeze(0)
+    else:
+        return out
+
+
+def compute_hks_autoscale(evals, evecs, count):
+    # these scales roughly approximate those suggested in the hks paper
+    scales = torch.logspace(-2, 0., steps=count, device=evals.device, dtype=evals.dtype)
+    return compute_hks(evals, evecs, scales)
+
+
+def normalize_positions(pos, faces=None, method='mean', scale_method='max_rad'):
+    # center and unit-scale positions
+
+    if method == 'mean':
+        # center using the average point position
+        pos = (pos - torch.mean(pos, dim=-2, keepdim=True))
+    elif method == 'bbox':
+        # center via the middle of the axis-aligned bounding box
+        bbox_min = torch.min(pos, dim=-2).values
+        bbox_max = torch.max(pos, dim=-2).values
+        center = (bbox_max + bbox_min) / 2.
+        pos -= center.unsqueeze(-2)
+    else:
+        raise ValueError("unrecognized method")
+
+    if scale_method == 'max_rad':
+        scale = torch.max(norm(pos), dim=-1, keepdim=True).values.unsqueeze(-1)
+        pos = pos / scale
+    elif scale_method == 'area':
+        if faces is None:
+            raise ValueError("must pass faces for area normalization")
+        coords = pos[faces]
+        vec_A = coords[:, 1, :] - coords[:, 0, :]
+        vec_B = coords[:, 2, :] - coords[:, 0, :]
+        face_areas = torch.norm(torch.cross(vec_A, vec_B, dim=-1), dim=1) * 0.5
+        total_area = torch.sum(face_areas)
+        scale = (1. / torch.sqrt(total_area))
+        pos = pos * scale
+    else:
+        raise ValueError("unrecognized scale method")
+    return pos
+
+
+# Finds the k nearest neighbors of source on target.
+# Return is two tensors (distances, indices). Returned points will be sorted in increasing order of distance.
+def find_knn(points_source, points_target, k, largest=False, omit_diagonal=False, method='brute'):
+
+    if omit_diagonal and points_source.shape[0] != points_target.shape[0]:
+        raise ValueError("omit_diagonal can only be used when source and target are same shape")
+
+    if method != 'cpu_kd' and points_source.shape[0] * points_target.shape[0] > 1e8:
+        method = 'cpu_kd'
+        print("switching to cpu_kd knn")
+
+    if method == 'brute':
+
+        # Expand so both are NxMx3 tensor
+        points_source_expand = points_source.unsqueeze(1)
+        points_source_expand = points_source_expand.expand(-1, points_target.shape[0], -1)
+        points_target_expand = points_target.unsqueeze(0)
+        points_target_expand = points_target_expand.expand(points_source.shape[0], -1, -1)
+
+        diff_mat = points_source_expand - points_target_expand
+        dist_mat = norm(diff_mat)
+
+        if omit_diagonal:
+            torch.diagonal(dist_mat)[:] = float('inf')
+
+        result = torch.topk(dist_mat, k=k, largest=largest, sorted=True)
+        return result
+
+    elif method == 'cpu_kd':
+
+        if largest:
+            raise ValueError("can't do largest with cpu_kd")
+
+        points_source_np = toNP(points_source)
+        points_target_np = toNP(points_target)
+
+        # Build the tree
+        kd_tree = sklearn.neighbors.KDTree(points_target_np)
+
+        k_search = k + 1 if omit_diagonal else k
+        _, neighbors = kd_tree.query(points_source_np, k=k_search)
+
+        if omit_diagonal:
+            # Mask out self element
+            mask = neighbors != np.arange(neighbors.shape[0])[:, np.newaxis]
+
+            # make sure we mask out exactly one element in each row, in rare case of many duplicate points
+            mask[np.sum(mask, axis=1) == mask.shape[1], -1] = False
+
+            neighbors = neighbors[mask].reshape((neighbors.shape[0], neighbors.shape[1] - 1))
+
+        inds = torch.tensor(neighbors, device=points_source.device, dtype=torch.int64)
+        dists = norm(points_source.unsqueeze(1).expand(-1, k, -1) - points_target[inds])
+
+        return dists, inds
+
+    else:
+        raise ValueError("unrecognized method")
+
+
+def farthest_point_sampling(points, n_sample):
+    # Torch in, torch out. Returns a |V| mask with n_sample elements set to true.
+
+    N = points.shape[0]
+    if (n_sample > N):
+        raise ValueError("not enough points to sample")
+
+    chosen_mask = torch.zeros(N, dtype=torch.bool, device=points.device)
+    min_dists = torch.ones(N, dtype=points.dtype, device=points.device) * float('inf')
+
+    # pick the centermost first point
+    points = normalize_positions(points)
+    i = torch.min(norm2(points), dim=0).indices
+    chosen_mask[i] = True
+
+    for _ in range(n_sample - 1):
+
+        # update distance
+        dists = norm2(points[i, :].unsqueeze(0) - points)
+        min_dists = torch.minimum(dists, min_dists)
+
+        # take the farthest
+        i = torch.max(min_dists, dim=0).indices.item()
+        chosen_mask[i] = True
+
+    return chosen_mask
+
+
+def geodesic_label_errors(target_verts,
+                          target_faces,
+                          pred_labels,
+                          gt_labels,
+                          normalization='diameter',
+                          geodesic_cache_dir=None):
+    """
+    Return a vector of distances between predicted and ground-truth lables (normalized by geodesic diameter or area)
+
+    This method is SLOW when it needs to recompute geodesic distances.
+    """
+
+    # move all to numpy cpu
+    target_verts = toNP(target_verts)
+    target_faces = toNP(target_faces)
+
+    pred_labels = toNP(pred_labels)
+    gt_labels = toNP(gt_labels)
+
+    dists = get_all_pairs_geodesic_distance(target_verts, target_faces, geodesic_cache_dir)
+
+    result_dists = dists[pred_labels, gt_labels]
+
+    if normalization == 'diameter':
+        geodesic_diameter = np.max(dists)
+        normalized_result_dists = result_dists / geodesic_diameter
+    elif normalization == 'area':
+        total_area = torch.sum(face_area(torch.tensor(target_verts), torch.tensor(target_faces)))
+        normalized_result_dists = result_dists / torch.sqrt(total_area)
+    else:
+        raise ValueError('unrecognized normalization')
+
+    return normalized_result_dists
+
+
+# This function and the helper class below are to support parallel computation of all-pairs geodesic distance
+def all_pairs_geodesic_worker(verts, faces, i):
+    import igl
+
+    N = verts.shape[0]
+
+    # TODO: this re-does a ton of work, since it is called independently each time. Some custom C++ code could surely make it faster.
+    sources = np.array([i])[:, np.newaxis]
+    targets = np.arange(N)[:, np.newaxis]
+    dist_vec = igl.exact_geodesic(verts, faces, sources, targets)
+
+    return dist_vec
+
+
+class AllPairsGeodesicEngine(object):
+
+    def __init__(self, verts, faces):
+        self.verts = verts
+        self.faces = faces
+
+    def __call__(self, i):
+        return all_pairs_geodesic_worker(self.verts, self.faces, i)
+
+
+def get_all_pairs_geodesic_distance(verts_np, faces_np, geodesic_cache_dir=None):
+    """
+    Return a gigantic VxV dense matrix containing the all-pairs geodesic distance matrix. Internally caches, recomputing only if necessary.
+
+    (numpy in, numpy out)
+    """
+
+    # need libigl for geodesic call
+    try:
+        import igl
+    except ImportError as e:
+        raise ImportError("Must have python libigl installed for all-pairs geodesics. `conda install -c conda-forge igl`")
+
+    # Check the cache
+    found = False
+    if geodesic_cache_dir is not None:
+        utils.ensure_dir_exists(geodesic_cache_dir)
+        hash_key_str = str(utils.hash_arrays((verts_np, faces_np)))
+        # print("Building operators for input with hash: " + hash_key_str)
+
+        # Search through buckets with matching hashes.  When the loop exits, this
+        # is the bucket index of the file we should write to.
+        i_cache_search = 0
+        while True:
+
+            # Form the name of the file to check
+            search_path = os.path.join(geodesic_cache_dir, hash_key_str + "_" + str(i_cache_search) + ".npz")
+
+            try:
+                npzfile = np.load(search_path, allow_pickle=True)
+                cache_verts = npzfile["verts"]
+                cache_faces = npzfile["faces"]
+
+                # If the cache doesn't match, keep looking
+                if (not np.array_equal(verts_np, cache_verts)) or (not np.array_equal(faces_np, cache_faces)):
+                    i_cache_search += 1
+                    continue
+
+                # This entry matches! Return it.
+                found = True
+                result_dists = npzfile["dist"]
+                break
+
+            except FileNotFoundError:
+                break
+
+    if not found:
+
+        print("Computing all-pairs geodesic distance (warning: SLOW!)")
+
+        # Not found, compute from scratch
+        # warning: slowwwwwww
+
+        N = verts_np.shape[0]
+
+        try:
+            pool = Pool(None)  # on 8 processors
+            engine = AllPairsGeodesicEngine(verts_np, faces_np)
+            outputs = pool.map(engine, range(N))
+        finally:  # To make sure processes are closed in the end, even if errors happen
+            pool.close()
+            pool.join()
+
+        result_dists = np.array(outputs)
+
+        # replace any failed values with nan
+        result_dists = np.nan_to_num(result_dists, nan=np.nan, posinf=np.nan, neginf=np.nan)
+
+        # we expect that this should be a symmetric matrix, but it might not be. Take the min of the symmetric values to make it symmetric
+        result_dists = np.fmin(result_dists, np.transpose(result_dists))
+
+        # on rare occaisions MMP fails, yielding nan/inf; set it to the largest non-failed value if so
+        max_dist = np.nanmax(result_dists)
+        result_dists = np.nan_to_num(result_dists, nan=max_dist, posinf=max_dist, neginf=max_dist)
+
+        print("...finished computing all-pairs geodesic distance")
+
+        # put it in the cache if possible
+        if geodesic_cache_dir is not None:
+
+            print("saving geodesic distances to cache: " + str(geodesic_cache_dir))
+
+            # TODO we're potentially saving a double precision but only using a single
+            # precision here; could save storage by always saving as floats
+            np.savez(search_path, verts=verts_np, faces=faces_np, dist=result_dists)
+
+    return result_dists

shape_models/layers.py

@@ -0,0 +1,453 @@
+import sys
+import os
+import os.path
+import random
+
+import scipy
+import scipy.sparse.linalg as sla
+# ^^^ we NEED to import scipy before torch, or it crashes :(
+# (observed on Ubuntu 20.04 w/ torch 1.6.0 and scipy 1.5.2 installed via conda)
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+
+def to_basis(values, basis, massvec):
+    """
+    Transform data in to an orthonormal basis (where orthonormal is wrt to massvec)
+    Inputs:
+      - values: (B,V,D)
+      - basis: (B,V,K)
+      - massvec: (B,V)
+    Outputs:
+      - (B,K,D) transformed values
+    """
+    basisT = basis.transpose(-2, -1)
+    return torch.matmul(basisT, values * massvec.unsqueeze(-1))
+
+
+def from_basis(values, basis):
+    """
+    Transform data out of an orthonormal basis
+    Inputs:
+      - values: (K,D)
+      - basis: (V,K)
+    Outputs:
+      - (V,D) reconstructed values
+    """
+    if values.is_complex() or basis.is_complex():
+        return utils.cmatmul(utils.ensure_complex(basis), utils.ensure_complex(values))
+    else:
+        return torch.matmul(basis, values)
+
+
+class LearnedTimeDiffusion(nn.Module):
+    """
+    Applies diffusion with learned per-channel t.
+
+    In the spectral domain this becomes 
+        f_out = e ^ (lambda_i t) f_in
+
+    Inputs:
+      - values: (V,C) in the spectral domain
+      - L: (V,V) sparse laplacian
+      - evals: (K) eigenvalues
+      - mass: (V) mass matrix diagonal
+
+      (note: L/evals may be omitted as None depending on method)
+    Outputs:
+      - (V,C) diffused values 
+    """
+
+    def __init__(self, C_inout, method='spectral'):
+        super(LearnedTimeDiffusion, self).__init__()
+        self.C_inout = C_inout
+        self.diffusion_time = nn.Parameter(torch.Tensor(C_inout))  # (C)
+        self.method = method  # one of ['spectral', 'implicit_dense']
+
+        nn.init.constant_(self.diffusion_time, 0.0)
+
+    def forward(self, x, L, mass, evals, evecs):
+
+        # project times to the positive halfspace
+        # (and away from 0 in the incredibly rare chance that they get stuck)
+        with torch.no_grad():
+            self.diffusion_time.data = torch.clamp(self.diffusion_time, min=1e-8)
+
+        if x.shape[-1] != self.C_inout:
+            raise ValueError("Tensor has wrong shape = {}. Last dim shape should have number of channels = {}".format(
+                x.shape, self.C_inout))
+
+        if self.method == 'spectral':
+
+            # Transform to spectral
+            x_spec = to_basis(x, evecs, mass)
+
+            # Diffuse
+            time = self.diffusion_time
+            diffusion_coefs = torch.exp(-evals.unsqueeze(-1) * time.unsqueeze(0))
+            x_diffuse_spec = diffusion_coefs * x_spec
+
+            # Transform back to per-vertex
+            x_diffuse = from_basis(x_diffuse_spec, evecs)
+
+        elif self.method == 'implicit_dense':
+            V = x.shape[-2]
+
+            # Form the dense matrices (M + tL) with dims (B,C,V,V)
+            mat_dense = L.to_dense().unsqueeze(1).expand(-1, self.C_inout, V, V).clone()
+            mat_dense *= self.diffusion_time.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+            mat_dense += torch.diag_embed(mass).unsqueeze(1)
+
+            # Factor the system
+            cholesky_factors = torch.linalg.cholesky(mat_dense)
+
+            # Solve the system
+            rhs = x * mass.unsqueeze(-1)
+            rhsT = torch.transpose(rhs, 1, 2).unsqueeze(-1)
+            sols = torch.cholesky_solve(rhsT, cholesky_factors)
+            x_diffuse = torch.transpose(sols.squeeze(-1), 1, 2)
+
+        else:
+            raise ValueError("unrecognized method")
+
+        return x_diffuse
+
+
+class SpatialGradientFeatures(nn.Module):
+    """
+    Compute dot-products between input vectors. Uses a learned complex-linear layer to keep dimension down.
+    
+    Input:
+        - vectors: (V,C,2)
+    Output:
+        - dots: (V,C) dots 
+    """
+
+    def __init__(self, C_inout, with_gradient_rotations=True):
+        super(SpatialGradientFeatures, self).__init__()
+
+        self.C_inout = C_inout
+        self.with_gradient_rotations = with_gradient_rotations
+
+        if (self.with_gradient_rotations):
+            self.A_re = nn.Linear(self.C_inout, self.C_inout, bias=False)
+            self.A_im = nn.Linear(self.C_inout, self.C_inout, bias=False)
+        else:
+            self.A = nn.Linear(self.C_inout, self.C_inout, bias=False)
+
+        # self.norm = nn.InstanceNorm1d(C_inout)
+
+    def forward(self, vectors):
+
+        vectorsA = vectors  # (V,C)
+
+        if self.with_gradient_rotations:
+            vectorsBreal = self.A_re(vectors[..., 0]) - self.A_im(vectors[..., 1])
+            vectorsBimag = self.A_re(vectors[..., 1]) + self.A_im(vectors[..., 0])
+        else:
+            vectorsBreal = self.A(vectors[..., 0])
+            vectorsBimag = self.A(vectors[..., 1])
+
+        dots = vectorsA[..., 0] * vectorsBreal + vectorsA[..., 1] * vectorsBimag
+
+        return torch.tanh(dots)
+
+
+class MiniMLP(nn.Sequential):
+    '''
+    A simple MLP with configurable hidden layer sizes.
+    '''
+
+    def __init__(self, layer_sizes, dropout=False, activation=nn.ReLU, name="miniMLP"):
+        super(MiniMLP, self).__init__()
+
+        for i in range(len(layer_sizes) - 1):
+            is_last = (i + 2 == len(layer_sizes))
+
+            if dropout and i > 0:
+                self.add_module(name + "_mlp_layer_dropout_{:03d}".format(i), nn.Dropout(p=.5))
+
+            # Affine map
+            self.add_module(
+                name + "_mlp_layer_{:03d}".format(i),
+                nn.Linear(
+                    layer_sizes[i],
+                    layer_sizes[i + 1],
+                ),
+            )
+
+            # Nonlinearity
+            # (but not on the last layer)
+            if not is_last:
+                self.add_module(name + "_mlp_act_{:03d}".format(i), activation())
+
+
+class DiffusionNetBlock(nn.Module):
+    """
+    Inputs and outputs are defined at vertices
+    """
+
+    def __init__(self,
+                 C_width,
+                 mlp_hidden_dims,
+                 dropout=True,
+                 diffusion_method='spectral',
+                 with_gradient_features=True,
+                 with_gradient_rotations=True):
+        super(DiffusionNetBlock, self).__init__()
+
+        # Specified dimensions
+        self.C_width = C_width
+        self.mlp_hidden_dims = mlp_hidden_dims
+
+        self.dropout = dropout
+        self.with_gradient_features = with_gradient_features
+        self.with_gradient_rotations = with_gradient_rotations
+
+        # Diffusion block
+        self.diffusion = LearnedTimeDiffusion(self.C_width, method=diffusion_method)
+
+        self.MLP_C = 2 * self.C_width
+
+        if self.with_gradient_features:
+            self.gradient_features = SpatialGradientFeatures(self.C_width, with_gradient_rotations=self.with_gradient_rotations)
+            self.MLP_C += self.C_width
+
+        # MLPs
+        self.mlp = MiniMLP([self.MLP_C] + self.mlp_hidden_dims + [self.C_width], dropout=self.dropout)
+
+    def forward(self, x_in, mass, L, evals, evecs, gradX, gradY):
+
+        # Manage dimensions
+        B = x_in.shape[0]  # batch dimension
+        if x_in.shape[-1] != self.C_width:
+            raise ValueError("Tensor has wrong shape = {}. Last dim shape should have number of channels = {}".format(
+                x_in.shape, self.C_width))
+
+        # Diffusion block
+        x_diffuse = self.diffusion(x_in, L, mass, evals, evecs)
+
+        # Compute gradient features, if using
+        if self.with_gradient_features:
+
+            # Compute gradients
+            x_grads = [
+            ]  # Manually loop over the batch (if there is a batch dimension) since torch.mm() doesn't support batching
+            for b in range(B):
+                # gradient after diffusion
+                x_gradX = torch.mm(gradX[b, ...], x_diffuse[b, ...])
+                x_gradY = torch.mm(gradY[b, ...], x_diffuse[b, ...])
+
+                x_grads.append(torch.stack((x_gradX, x_gradY), dim=-1))
+            x_grad = torch.stack(x_grads, dim=0)
+
+            # Evaluate gradient features
+            x_grad_features = self.gradient_features(x_grad)
+
+            # Stack inputs to mlp
+            feature_combined = torch.cat((x_in, x_diffuse, x_grad_features), dim=-1)
+        else:
+            # Stack inputs to mlp
+            feature_combined = torch.cat((x_in, x_diffuse), dim=-1)
+
+        # Apply the mlp
+        x0_out = self.mlp(feature_combined)
+
+        # Skip connection
+        x0_out = x0_out + x_in
+
+        return x0_out
+
+
+class DiffusionNet(nn.Module):
+
+    def __init__(self,
+                 C_in,
+                 C_out,
+                 C_width=128,
+                 N_block=4,
+                 last_activation=None,
+                 outputs_at='vertices',
+                 mlp_hidden_dims=None,
+                 dropout=True,
+                 with_gradient_features=True,
+                 with_gradient_rotations=True,
+                 diffusion_method='spectral',
+                 num_eigenbasis=128):
+        """
+        Construct a DiffusionNet.
+
+        Parameters:
+            C_in (int):                     input dimension 
+            C_out (int):                    output dimension 
+            last_activation (func)          a function to apply to the final outputs of the network, such as torch.nn.functional.log_softmax (default: None)
+            outputs_at (string)             produce outputs at various mesh elements by averaging from vertices. One of ['vertices', 'edges', 'faces', 'global_mean']. (default 'vertices', aka points for a point cloud)
+            C_width (int):                  dimension of internal DiffusionNet blocks (default: 128)
+            N_block (int):                  number of DiffusionNet blocks (default: 4)
+            mlp_hidden_dims (list of int):  a list of hidden layer sizes for MLPs (default: [C_width, C_width])
+            dropout (bool):                 if True, internal MLPs use dropout (default: True)
+            diffusion_method (string):      how to evaluate diffusion, one of ['spectral', 'implicit_dense']. If implicit_dense is used, can set k_eig=0, saving precompute.
+            with_gradient_features (bool):  if True, use gradient features (default: True)
+            with_gradient_rotations (bool): if True, use gradient also learn a rotation of each gradient. Set to True if your surface has consistently oriented normals, and False otherwise (default: True)
+            num_eigenbasis (int):           for trunking the eigenvalues eigenvectors
+        """
+
+        super(DiffusionNet, self).__init__()
+
+        ## Store parameters
+
+        # Basic parameters
+        self.C_in = C_in
+        self.C_out = C_out
+        self.C_width = C_width
+        self.N_block = N_block
+
+        # Outputs
+        self.last_activation = last_activation
+        self.outputs_at = outputs_at
+        if outputs_at not in ['vertices', 'edges', 'faces', 'global_mean']:
+            raise ValueError("invalid setting for outputs_at")
+
+        # MLP options
+        if mlp_hidden_dims == None:
+            mlp_hidden_dims = [C_width, C_width]
+        self.mlp_hidden_dims = mlp_hidden_dims
+        self.dropout = dropout
+
+        # Diffusion
+        self.diffusion_method = diffusion_method
+        if diffusion_method not in ['spectral', 'implicit_dense']:
+            raise ValueError("invalid setting for diffusion_method")
+        self.num_eigenbasis = num_eigenbasis
+
+        # Gradient features
+        self.with_gradient_features = with_gradient_features
+        self.with_gradient_rotations = with_gradient_rotations
+
+        ## Set up the network
+
+        # First and last affine layers
+        self.first_lin = nn.Linear(C_in, C_width)
+        self.last_lin = nn.Linear(C_width, C_out)
+
+        # DiffusionNet blocks
+        self.blocks = []
+        for i_block in range(self.N_block):
+            block = DiffusionNetBlock(C_width=C_width,
+                                      mlp_hidden_dims=mlp_hidden_dims,
+                                      dropout=dropout,
+                                      diffusion_method=diffusion_method,
+                                      with_gradient_features=with_gradient_features,
+                                      with_gradient_rotations=with_gradient_rotations)
+
+            self.blocks.append(block)
+            self.add_module("block_" + str(i_block), self.blocks[-1])
+
+    def forward(self, x_in, mass, L=None, evals=None, evecs=None, gradX=None, gradY=None, edges=None, faces=None):
+        """
+        A forward pass on the DiffusionNet.
+
+        In the notation below, dimension are:
+            - C is the input channel dimension (C_in on construction)
+            - C_OUT is the output channel dimension (C_out on construction)
+            - N is the number of vertices/points, which CAN be different for each forward pass
+            - B is an OPTIONAL batch dimension
+            - K_EIG is the number of eigenvalues used for spectral acceleration
+        Generally, our data layout it is [N,C] or [B,N,C].
+
+        Call get_operators() to generate geometric quantities mass/L/evals/evecs/gradX/gradY. Note that depending on the options for the DiffusionNet, not all are strictly necessary.
+
+        Parameters:
+            x_in (tensor):      Input features, dimension [N,C] or [B,N,C]
+            mass (tensor):      Mass vector, dimension [N] or [B,N]
+            L (tensor):         Laplace matrix, sparse tensor with dimension [N,N] or [B,N,N]
+            evals (tensor):     Eigenvalues of Laplace matrix, dimension [K_EIG] or [B,K_EIG]
+            evecs (tensor):     Eigenvectors of Laplace matrix, dimension [N,K_EIG] or [B,N,K_EIG]
+            gradX (tensor):     Half of gradient matrix, sparse real tensor with dimension [N,N] or [B,N,N]
+            gradY (tensor):     Half of gradient matrix, sparse real tensor with dimension [N,N] or [B,N,N]
+
+        Returns:
+            x_out (tensor):    Output with dimension [N,C_out] or [B,N,C_out]
+        """
+
+        ## Check dimensions, and append batch dimension if not given
+        if x_in.shape[-1] != self.C_in:
+            raise ValueError("DiffusionNet was constructed with C_in={}, but x_in has last dim={}".format(
+                self.C_in, x_in.shape[-1]))
+        N = x_in.shape[-2]
+        if len(x_in.shape) == 2:
+            appended_batch_dim = True
+
+            # add a batch dim to all inputs
+            x_in = x_in.unsqueeze(0)
+            mass = mass.unsqueeze(0)
+            if L != None:
+                L = L.unsqueeze(0)
+            if evals != None:
+                evals = evals.unsqueeze(0)
+            if evecs != None:
+                evecs = evecs.unsqueeze(0)
+            if gradX != None:
+                gradX = gradX.unsqueeze(0)
+            if gradY != None:
+                gradY = gradY.unsqueeze(0)
+            if edges != None:
+                edges = edges.unsqueeze(0)
+            if faces != None:
+                faces = faces.unsqueeze(0)
+
+        elif len(x_in.shape) == 3:
+            appended_batch_dim = False
+
+        else:
+            raise ValueError("x_in should be tensor with shape [N,C] or [B,N,C]")
+
+        evals = evals[..., :self.num_eigenbasis]
+        evecs = evecs[..., :self.num_eigenbasis]
+
+        # Apply the first linear layer
+        x = self.first_lin(x_in)
+
+        # Apply each of the blocks
+        for b in self.blocks:
+            x = b(x, mass, L, evals, evecs, gradX, gradY)
+
+        # Apply the last linear layer
+        x = self.last_lin(x)
+
+        # Remap output to faces/edges if requested
+        if self.outputs_at == 'vertices':
+            x_out = x
+
+        elif self.outputs_at == 'edges':
+            # Remap to edges
+            x_gather = x.unsqueeze(-1).expand(-1, -1, -1, 2)
+            edges_gather = edges.unsqueeze(2).expand(-1, -1, x.shape[-1], -1)
+            xe = torch.gather(x_gather, 1, edges_gather)
+            x_out = torch.mean(xe, dim=-1)
+
+        elif self.outputs_at == 'faces':
+            # Remap to faces
+            x_gather = x.unsqueeze(-1).expand(-1, -1, -1, 3)
+            faces_gather = faces.unsqueeze(2).expand(-1, -1, x.shape[-1], -1)
+            xf = torch.gather(x_gather, 1, faces_gather)
+            x_out = torch.mean(xf, dim=-1)
+
+        elif self.outputs_at == 'global_mean':
+            # Produce a single global mean ouput.
+            # Using a weighted mean according to the point mass/area is discretization-invariant.
+            # (A naive mean is not discretization-invariant; it could be affected by sampling a region more densely)
+            x_out = torch.sum(x * mass.unsqueeze(-1), dim=-2) / torch.sum(mass, dim=-1, keepdim=True)
+
+        # Apply last nonlinearity if specified
+        if self.last_activation != None:
+            x_out = self.last_activation(x_out)
+
+        # Remove batch dim if we added it
+        if appended_batch_dim:
+            x_out = x_out.squeeze(0)
+
+        return x_out

shape_models/utils.py

@@ -0,0 +1,123 @@
+import sys
+import os
+import time
+
+import torch
+import hashlib
+import numpy as np
+import scipy
+
+# == Pytorch things
+
+
+def toNP(x):
+    """
+    Really, definitely convert a torch tensor to a numpy array
+    """
+    return x.detach().to(torch.device('cpu')).numpy()
+
+
+def label_smoothing_log_loss(pred, labels, smoothing=0.0):
+    n_class = pred.shape[-1]
+    one_hot = torch.zeros_like(pred)
+    one_hot[labels] = 1.
+    one_hot = one_hot * (1 - smoothing) + (1 - one_hot) * smoothing / (n_class - 1)
+    loss = -(one_hot * pred).sum(dim=-1).mean()
+    return loss
+
+
+# Randomly rotate points.
+# Torch in, torch out
+# Note fornow, builds rotation matrix on CPU.
+def random_rotate_points(pts, randgen=None):
+    R = random_rotation_matrix(randgen)
+    R = torch.from_numpy(R).to(device=pts.device, dtype=pts.dtype)
+    return torch.matmul(pts, R)
+
+
+def random_rotate_points_y(pts):
+    angles = torch.rand(1, device=pts.device, dtype=pts.dtype) * (2. * np.pi)
+    rot_mats = torch.zeros(3, 3, device=pts.device, dtype=pts.dtype)
+    rot_mats[0, 0] = torch.cos(angles)
+    rot_mats[0, 2] = torch.sin(angles)
+    rot_mats[2, 0] = -torch.sin(angles)
+    rot_mats[2, 2] = torch.cos(angles)
+    rot_mats[1, 1] = 1.
+
+    pts = torch.matmul(pts, rot_mats)
+    return pts
+
+
+# Numpy things
+
+
+# Numpy sparse matrix to pytorch
+def sparse_np_to_torch(A):
+    Acoo = A.tocoo()
+    values = Acoo.data
+    indices = np.vstack((Acoo.row, Acoo.col))
+    shape = Acoo.shape
+    return torch.sparse.FloatTensor(torch.LongTensor(indices), torch.FloatTensor(values), torch.Size(shape)).coalesce()
+
+
+# Pytorch sparse to numpy csc matrix
+def sparse_torch_to_np(A):
+    if len(A.shape) != 2:
+        raise RuntimeError("should be a matrix-shaped type; dim is : " + str(A.shape))
+
+    indices = toNP(A.indices())
+    values = toNP(A.values())
+
+    mat = scipy.sparse.coo_matrix((values, indices), shape=A.shape).tocsc()
+
+    return mat
+
+
+# Hash a list of numpy arrays
+def hash_arrays(arrs):
+    running_hash = hashlib.sha1()
+    for arr in arrs:
+        binarr = arr.view(np.uint8)
+        running_hash.update(binarr)
+    return running_hash.hexdigest()
+
+
+def random_rotation_matrix(randgen=None):
+    """
+    Creates a random rotation matrix.
+    randgen: if given, a np.random.RandomState instance used for random numbers (for reproducibility)
+    """
+    # adapted from http://www.realtimerendering.com/resources/GraphicsGems/gemsiii/rand_rotation.c
+
+    if randgen is None:
+        randgen = np.random.RandomState()
+
+    theta, phi, z = tuple(randgen.rand(3).tolist())
+
+    theta = theta * 2.0 * np.pi  # Rotation about the pole (Z).
+    phi = phi * 2.0 * np.pi  # For direction of pole deflection.
+    z = z * 2.0  # For magnitude of pole deflection.
+
+    # Compute a vector V used for distributing points over the sphere
+    # via the reflection I - V Transpose(V).  This formulation of V
+    # will guarantee that if x[1] and x[2] are uniformly distributed,
+    # the reflected points will be uniform on the sphere.  Note that V
+    # has length sqrt(2) to eliminate the 2 in the Householder matrix.
+
+    r = np.sqrt(z)
+    Vx, Vy, Vz = V = (np.sin(phi) * r, np.cos(phi) * r, np.sqrt(2.0 - z))
+
+    st = np.sin(theta)
+    ct = np.cos(theta)
+
+    R = np.array(((ct, st, 0), (-st, ct, 0), (0, 0, 1)))
+    # Construct the rotation matrix  ( V Transpose(V) - I ) R.
+
+    M = (np.outer(V, V) - np.eye(3)).dot(R)
+    return M
+
+
+# Python string/file utilities
+def ensure_dir_exists(d):
+    if not os.path.exists(d):
+        os.makedirs(d)

utils/__init__.py

@@ -0,0 +1,3 @@
+# from .utils_legacy import *
+# from .geometry import *
+# from .layers import *

utils/descriptors.py

@@ -0,0 +1,250 @@
+import numpy as np
+
+# https://github.com/RobinMagnet/SimplifiedFmapsLearning/blob/main/learn_zo/data/utils.py
+# https://github.com/RobinMagnet/pyFM/blob/master/pyFM/signatures/HKS_functions.py
+def HKS(evals, evects, time_list,scaled=False):
+    """
+    Returns the Heat Kernel Signature for num_T different values.
+    The values of the time are interpolated in logscale between the limits
+    given in the HKS paper. These limits only depends on the eigenvalues.
+
+    Parameters
+    ------------------------
+    evals     : (K,) array of the K eigenvalues
+    evecs     : (N,K) array with the K eigenvectors
+    time_list : (num_T,) Time values to use
+    scaled    : (bool) whether to scale for each time value
+
+    Output
+    ------------------------
+    HKS : (N,num_T) array where each line is the HKS for a given t
+    """
+    evals_s = np.asarray(evals).flatten()
+    t_list = np.asarray(time_list).flatten()
+
+    coefs = np.exp(-np.outer(t_list, evals_s))  # (num_T,K)
+    # weighted_evects = evects[None, :, :] * coefs[:, None,:]  # (num_T,N,K)
+    # natural_HKS = np.einsum('tnk,nk->nt', weighted_evects, evects)
+    natural_HKS = np.einsum('tk,nk,nk->nt', coefs, evects, evects)
+
+    if scaled:
+        inv_scaling = coefs.sum(1)  # (num_T)
+        return (1/inv_scaling)[None,:] * natural_HKS
+
+    return natural_HKS
+
+
+def lm_HKS(evals, evects, landmarks, time_list, scaled=False):
+    """
+    Returns the Heat Kernel Signature for some landmarks and time values.
+
+
+    Parameters
+    ------------------------
+    evects      : (N,K) array with the K eigenvectors of the Laplace Beltrami operator
+    evals       : (K,) array of the K corresponding eigenvalues
+    landmarks   : (p,) indices of landmarks to compute
+    time_list   : (num_T,) values of t to use
+
+    Output
+    ------------------------
+    landmarks_HKS : (N,num_E*p) array where each column is the HKS for a given t for some landmark
+    """
+
+    evals_s = np.asarray(evals).flatten()
+    t_list = np.asarray(time_list).flatten()
+
+    coefs = np.exp(-np.outer(t_list, evals_s))  # (num_T,K)
+    weighted_evects = evects[None, landmarks, :] * coefs[:,None,:]  # (num_T,p,K)
+
+    landmarks_HKS = np.einsum('tpk,nk->ptn', weighted_evects, evects)  # (p,num_T,N)
+    landmarks_HKS = np.einsum('tk,pk,nk->ptn', coefs, evects[landmarks, :], evects)  # (p,num_T,N)
+
+    if scaled:
+        inv_scaling = coefs.sum(1)  # (num_T,)
+        landmarks_HKS = (1/inv_scaling)[None,:,None] * landmarks_HKS  # (p,num_T,N)
+
+    return rearrange(landmarks_HKS, 'p T N -> N (p T)')
+
+
+def auto_HKS(evals, evects, num_T, landmarks=None, scaled=True):
+    """
+    Compute HKS with an automatic choice of tile values
+
+    Parameters
+    ------------------------
+    evals       : (K,) array of  K eigenvalues
+    evects      : (N,K) array with K eigenvectors
+    landmarks   : (p,) if not None, indices of landmarks to compute.
+    num_T       : (int) number values of t to use
+    Output
+    ------------------------
+    HKS or lm_HKS : (N,num_E) or (N,p*num_E)  array where each column is the WKS for a given e
+                    for some landmark
+    """
+
+    abs_ev = sorted(np.abs(evals))
+    t_list = np.geomspace(4*np.log(10)/abs_ev[-1], 4*np.log(10)/abs_ev[1], num_T)
+
+    if landmarks is None:
+        return HKS(abs_ev, evects, t_list, scaled=scaled)
+    else:
+        return lm_HKS(abs_ev, evects, landmarks, t_list, scaled=scaled)
+
+
+# https://github.com/RobinMagnet/pyFM/blob/master/pyFM/signatures/WKS_functions.py
+def WKS(evals, evects, energy_list, sigma, scaled=False):
+    """
+    Returns the Wave Kernel Signature for some energy values.
+
+    Parameters
+    ------------------------
+    evects      : (N,K) array with the K eigenvectors of the Laplace Beltrami operator
+    evals       : (K,) array of the K corresponding eigenvalues
+    energy_list : (num_E,) values of e to use
+    sigma       : (float) [positive] standard deviation to use
+    scaled      : (bool) Whether to scale each energy level
+
+    Output
+    ------------------------
+    WKS : (N,num_E) array where each column is the WKS for a given e
+    """
+    assert sigma > 0, f"Sigma should be positive ! Given value : {sigma}"
+
+    evals = np.asarray(evals).flatten()
+    indices = np.where(evals > 1e-5)[0].flatten()
+    evals = evals[indices]
+    evects = evects[:, indices]
+
+    e_list = np.asarray(energy_list)
+    coefs = np.exp(-np.square(e_list[:,None] - np.log(np.abs(evals))[None,:])/(2*sigma**2))  # (num_E,K)
+
+    # weighted_evects = evects[None, :, :] * coefs[:,None, :]  # (num_E,N,K)
+
+    # natural_WKS = np.einsum('tnk,nk->nt', weighted_evects, evects)  # (N,num_E)
+    natural_WKS = np.einsum('tk,nk,nk->nt', coefs, evects, evects)
+
+    if scaled:
+        inv_scaling = coefs.sum(1)  # (num_E)
+        return (1/inv_scaling)[None,:] * natural_WKS
+    
+    return natural_WKS
+
+
+def lm_WKS(evals, evects, landmarks, energy_list, sigma, scaled=False):
+    """
+    Returns the Wave Kernel Signature for some landmarks and energy values.
+
+
+    Parameters
+    ------------------------
+    evects      : (N,K) array with the K eigenvectors of the Laplace Beltrami operator
+    evals       : (K,) array of the K corresponding eigenvalues
+    landmarks   : (p,) indices of landmarks to compute
+    energy_list : (num_E,) values of e to use
+    sigma       : int - standard deviation
+
+    Output
+    ------------------------
+    landmarks_WKS : (N,num_E*p) array where each column is the WKS for a given e for some landmark
+    """
+    assert sigma > 0, f"Sigma should be positive ! Given value : {sigma}"
+
+    evals = np.asarray(evals).flatten()
+    indices = np.where(evals > 1e-2)[0].flatten()
+    evals = evals[indices]
+    evects = evects[:,indices]
+
+    e_list = np.asarray(energy_list)
+    coefs = np.exp(-np.square(e_list[:, None] - np.log(np.abs(evals))[None, :]) / (2*sigma**2))  # (num_E,K)
+    # weighted_evects = evects[None, landmarks, :] * coefs[:,None,:]  # (num_E,p,K)
+
+    # landmarks_WKS = np.einsum('tpk,nk->ptn', weighted_evects, evects)  # (p,num_E,N)
+    landmarks_WKS = np.einsum('tk,pk,nk->ptn', coefs, evects[landmarks, :], evects)  # (p,num_E,N)
+
+    if scaled:
+        inv_scaling = coefs.sum(1)  # (num_E,)
+        landmarks_WKS = (1/inv_scaling)[None,:,None] * landmarks_WKS
+
+    # return landmarks_WKS.reshape(-1,evects.shape[0]).T  # (N,p*num_E)
+    return rearrange(landmarks_WKS, 'p T N -> N (p T)')
+
+
+def auto_WKS(evals, evects, num_E, landmarks=None, scaled=True):
+    """
+    Compute WKS with an automatic choice of scale and energy
+
+    Parameters
+    ------------------------
+    evals       : (K,) array of  K eigenvalues
+    evects      : (N,K) array with K eigenvectors
+    landmarks   : (p,) If not None, indices of landmarks to compute.
+    num_E       : (int) number values of e to use
+    Output
+    ------------------------
+    WKS or lm_WKS : (N,num_E) or (N,p*num_E)  array where each column is the WKS for a given e
+                    and possibly for some landmarks
+    """
+    abs_ev = sorted(np.abs(evals))
+
+    e_min,e_max = np.log(abs_ev[1]),np.log(abs_ev[-1])
+    sigma = 7*(e_max-e_min)/num_E
+
+    e_min += 2*sigma
+    e_max -= 2*sigma
+
+    energy_list = np.linspace(e_min,e_max,num_E)
+
+    if landmarks is None:
+        return WKS(abs_ev, evects, energy_list, sigma, scaled=scaled)
+    else:
+        return lm_WKS(abs_ev, evects, landmarks, energy_list, sigma, scaled=scaled)
+
+
+def compute_hks(evecs, evals, mass, n_descr=100, subsample_step=5, n_eig=35, normalize=True):
+    """
+    Compute Heat Kernel Signature (HKS) descriptors.
+    
+    Args:
+        evecs: (N, K) eigenvectors of the Laplace-Beltrami operator
+        evals: (K,) eigenvalues of the Laplace-Beltrami operator
+        mass: (N,) vertex masses
+        n_descr: (int) number of descriptors
+        subsample_step: (int) subsampling step
+        n_eig: (int) number of eigenvectors to use
+    
+    Returns:
+        feats: (N, n_descr) HKS descriptors
+    """
+    feats = auto_HKS(evals[:n_eig], evecs[:, :n_eig], n_descr, scaled=True)
+    feats = feats[:, np.arange(0, feats.shape[1], subsample_step)]
+    if normalize:
+        feats_norm2 = np.einsum('np,n->p', feats**2, mass).flatten()
+        feats /= np.sqrt(feats_norm2)[None, :]
+    return feats.astype(np.float32)
+
+
+def compute_wks(evecs, evals, mass, n_descr=100, subsample_step=5, n_eig=35, normalize=True):
+    """
+    Compute Wave Kernel Signature (WKS) descriptors.
+
+    Args:
+        evecs: (N, K) eigenvectors of the Laplace-Beltrami operator
+        evals: (K,) eigenvalues of the Laplace-Beltrami operator
+        mass: (N,) vertex masses
+        n_descr: (int) number of descriptors
+        subsample_step: (int) subsampling step
+        n_eig: (int) number of eigenvectors to use
+    
+    Returns:
+        feats: (N, n_descr) WKS descriptors
+    """
+    feats = auto_WKS(evals[:n_eig], evecs[:, :n_eig], n_descr, scaled=True)
+    feats = feats[:, np.arange(0, feats.shape[1], subsample_step)]
+    # print("wks_shape",feats.shape, mass.shape)
+    if normalize:
+        feats_norm2 = np.einsum('np,n->p', feats**2, mass).flatten()
+        feats /= np.sqrt(feats_norm2)[None, :]
+    # feats_norm2 = np.einsum('np,n->p', feats**2, mass).flatten()
+    # feats /= np.sqrt(feats_norm2)[None, :]
+    return feats.astype(np.float32)

utils/eval.py

@@ -0,0 +1,25 @@
+def accuracy(p2p, gt_p2p, D1_geod, return_all=False, sqrt_area=None):
+    """
+    Computes the geodesic accuracy of a vertex to vertex map. The map goes from
+    the target shape to the source shape.
+    Borrowed from Robin.
+    Parameters
+    ----------------------
+    p2p        : (n2,) - vertex to vertex map giving the index of the matched vertex on the source shape
+                 for each vertex on the target shape (from a functional map point of view)
+    gt_p2p     : (n2,) - ground truth mapping between the pairs
+    D1_geod    : (n1,n1) - geodesic distance between pairs of vertices on the source mesh
+    return_all : bool - whether to return all the distances or only the average geodesic distance
+
+    Output
+    -----------------------
+    acc   : float - average accuracy of the vertex to vertex map
+    dists : (n2,) - if return_all is True, returns all the pairwise distances
+    """
+
+    dists = D1_geod[(p2p, gt_p2p)]
+    if sqrt_area is not None:
+        dists /= sqrt_area
+    if return_all:
+        return dists.mean(), dists
+    return dists.mean()

utils/fmap.py

@@ -0,0 +1,121 @@
+import os
+import os.path as osp
+import sys
+import numpy as np
+import scipy.linalg
+from tqdm import tqdm
+
+ROOT_DIR = osp.join(osp.abspath(osp.dirname(__file__)), '../')
+if ROOT_DIR not in sys.path:
+    sys.path.append(ROOT_DIR)
+
+from utils_fmaps.misc import KNNSearch
+
+try:
+    import pynndescent
+    index = pynndescent.NNDescent(np.random.random((100, 3)), n_jobs=2)
+    del index
+    ANN = True
+except ImportError:
+    ANN = False
+
+# https://github.com/RobinMagnet/pyFM
+
+
+def FM_to_p2p(FM, eigvects1, eigvects2, use_ANN=False):
+    if use_ANN and not ANN:
+        raise ValueError('Please install pydescent to achieve Approximate Nearest Neighbor')
+
+    k2, k1 = FM.shape
+
+    assert k1 <= eigvects1.shape[1], \
+        f'At least {k1} should be provided, here only {eigvects1.shape[1]} are given'
+    assert k2 <= eigvects2.shape[1], \
+        f'At least {k2} should be provided, here only {eigvects2.shape[1]} are given'
+
+    if use_ANN:
+        index = pynndescent.NNDescent(eigvects1[:, :k1] @ FM.T, n_jobs=8)
+        matches, _ = index.query(eigvects2[:, :k2], k=1)
+        matches = matches.flatten()
+    else:
+        tree = KNNSearch(eigvects1[:, :k1] @ FM.T)
+        matches = tree.query(eigvects2[:, :k2], k=1).flatten()
+
+    return matches
+
+
+def p2p_to_FM(p2p, eigvects1, eigvects2, A2=None):
+    if A2 is not None:
+        if A2.shape[0] != eigvects2.shape[0]:
+            raise ValueError("Can't compute pseudo inverse with subsampled eigenvectors")
+
+        if len(A2.shape) == 1:
+            return eigvects2.T @ (A2[:, None] * eigvects1[p2p, :])
+
+        return eigvects2.T @ A2 @ eigvects1[p2p, :]
+
+    return scipy.linalg.lstsq(eigvects2, eigvects1[p2p, :])[0]
+
+
+def zoomout_iteration(eigvects1, eigvects2, FM, step=1, A2=None, use_ANN=False):
+    k2, k1 = FM.shape
+    try:
+        step1, step2 = step
+    except TypeError:
+        step1 = step
+        step2 = step
+    new_k1, new_k2 = k1 + step1, k2 + step2
+
+    p2p = FM_to_p2p(FM, eigvects1, eigvects2, use_ANN=use_ANN)
+    FM_zo = p2p_to_FM(p2p, eigvects1[:, :new_k1], eigvects2[:, :new_k2], A2=A2)
+
+    return FM_zo
+
+
+def zoomout_refine(eigvects1,
+                   eigvects2,
+                   FM,
+                   nit=10,
+                   step=1,
+                   A2=None,
+                   subsample=None,
+                   use_ANN=False,
+                   return_p2p=False,
+                   verbose=False):
+    k2_0, k1_0 = FM.shape
+    try:
+        step1, step2 = step
+    except TypeError:
+        step1 = step
+        step2 = step
+
+    assert k1_0 + nit*step1 <= eigvects1.shape[1], \
+        f"Not enough eigenvectors on source : \
+        {k1_0 + nit*step1} are needed when {eigvects1.shape[1]} are provided"
+    assert k2_0 + nit*step2 <= eigvects2.shape[1], \
+        f"Not enough eigenvectors on target : \
+        {k2_0 + nit*step2} are needed when {eigvects2.shape[1]} are provided"
+
+    use_subsample = False
+    if subsample is not None:
+        use_subsample = True
+        sub1, sub2 = subsample
+
+    FM_zo = FM.copy()
+
+    ANN_adventage = False
+    iterable = range(nit) if not verbose else tqdm(range(nit))
+    for it in iterable:
+        ANN_adventage = use_ANN and (FM_zo.shape[0] > 90) and (FM_zo.shape[1] > 90)
+
+        if use_subsample:
+            FM_zo = zoomout_iteration(eigvects1[sub1], eigvects2[sub2], FM_zo, A2=None, step=step, use_ANN=ANN_adventage)
+
+        else:
+            FM_zo = zoomout_iteration(eigvects1, eigvects2, FM_zo, A2=A2, step=step, use_ANN=ANN_adventage)
+
+    if return_p2p:
+        p2p_zo = FM_to_p2p(FM_zo, eigvects1, eigvects2, use_ANN=False)
+        return FM_zo, p2p_zo
+
+    return FM_zo

utils/geometry.py

@@ -0,0 +1,951 @@
+import scipy
+import scipy.sparse.linalg as sla
+# ^^^ we NEED to import scipy before torch, or it crashes :(
+# (observed on Ubuntu 20.04 w/ torch 1.6.0 and scipy 1.5.2 installed via conda)
+
+import os.path
+import sys
+import random
+from multiprocessing import Pool
+
+import numpy as np
+import scipy.spatial
+import torch
+import sklearn.neighbors
+
+import robust_laplacian
+import potpourri3d as pp3d
+
+from utils.utils_legacy import toNP, ensure_dir_exists, sparse_np_to_torch, sparse_torch_to_np
+from utils.descriptors import compute_hks, compute_wks
+
+
+def norm(x, highdim=False):
+    """
+    Computes norm of an array of vectors. Given (shape,d), returns (shape) after norm along last dimension
+    """
+    return torch.norm(x, dim=len(x.shape) - 1)
+
+
+def norm2(x, highdim=False):
+    """
+    Computes norm^2 of an array of vectors. Given (shape,d), returns (shape) after norm along last dimension
+    """
+    return dot(x, x)
+
+
+def normalize(x, divide_eps=1e-6, highdim=False):
+    """
+    Computes norm^2 of an array of vectors. Given (shape,d), returns (shape) after norm along last dimension
+    """
+    if (len(x.shape) == 1):
+        raise ValueError("called normalize() on single vector of dim " + str(x.shape) + " are you sure?")
+    if (not highdim and x.shape[-1] > 4):
+        raise ValueError("called normalize() with large last dimension " + str(x.shape) + " are you sure?")
+    return x / (norm(x, highdim=highdim) + divide_eps).unsqueeze(-1)
+
+
+def face_coords(verts, faces):
+    coords = verts[faces]
+    return coords
+
+
+def cross(vec_A, vec_B):
+    return torch.cross(vec_A, vec_B, dim=-1)
+
+
+def dot(vec_A, vec_B):
+    return torch.sum(vec_A * vec_B, dim=-1)
+
+
+# Given (..., 3) vectors and normals, projects out any components of vecs
+# which lies in the direction of normals. Normals are assumed to be unit.
+
+
+def project_to_tangent(vecs, unit_normals):
+    dots = dot(vecs, unit_normals)
+    return vecs - unit_normals * dots.unsqueeze(-1)
+
+
+def face_area(verts, faces):
+    coords = face_coords(verts, faces)
+    vec_A = coords[:, 1, :] - coords[:, 0, :]
+    vec_B = coords[:, 2, :] - coords[:, 0, :]
+
+    raw_normal = cross(vec_A, vec_B)
+    return 0.5 * norm(raw_normal)
+
+
+def face_normals(verts, faces, normalized=True):
+    coords = face_coords(verts, faces)
+    vec_A = coords[:, 1, :] - coords[:, 0, :]
+    vec_B = coords[:, 2, :] - coords[:, 0, :]
+
+    raw_normal = cross(vec_A, vec_B)
+
+    if normalized:
+        return normalize(raw_normal)
+
+    return raw_normal
+
+
+def neighborhood_normal(points):
+    # points: (N, K, 3) array of neighborhood psoitions
+    # points should be centered at origin
+    # out: (N,3) array of normals
+    # numpy in, numpy out
+    (u, s, vh) = np.linalg.svd(points, full_matrices=False)
+    normal = vh[:, 2, :]
+    return normal / np.linalg.norm(normal, axis=-1, keepdims=True)
+
+
+def mesh_vertex_normals(verts, faces):
+    # numpy in / out
+    face_n = toNP(face_normals(torch.tensor(verts), torch.tensor(faces)))  # ugly torch <---> numpy
+    eps = 1e-3
+    vertex_normals = np.zeros(verts.shape)
+    for i in range(3):
+        np.add.at(vertex_normals, faces[:, i], face_n)
+
+    vertex_normals = vertex_normals / (eps + np.linalg.norm(vertex_normals, axis=-1, keepdims=True))
+
+    return vertex_normals
+
+
+def vertex_normals(verts, faces, n_neighbors_cloud=30):
+    verts_np = toNP(verts)
+
+    if faces.numel() == 0:  # point cloud
+
+        _, neigh_inds = find_knn(verts, verts, n_neighbors_cloud, omit_diagonal=True, method='cpu_kd')
+        neigh_points = verts_np[neigh_inds, :]
+        neigh_points = neigh_points - verts_np[:, np.newaxis, :]
+        normals = neighborhood_normal(neigh_points)
+
+    else:  # mesh
+
+        normals = mesh_vertex_normals(verts_np, toNP(faces))
+
+        # if any are NaN, wiggle slightly and recompute
+        bad_normals_mask = np.isnan(normals).any(axis=1, keepdims=True)
+        if bad_normals_mask.any():
+            bbox = np.amax(verts_np, axis=0) - np.amin(verts_np, axis=0)
+            scale = np.linalg.norm(bbox) * 1e-4
+            wiggle = (np.random.RandomState(seed=777).rand(*verts.shape) - 0.5) * scale
+            wiggle_verts = verts_np + bad_normals_mask * wiggle
+            normals = mesh_vertex_normals(wiggle_verts, toNP(faces))
+
+        # if still NaN assign random normals (probably means unreferenced verts in mesh)
+        bad_normals_mask = np.isnan(normals).any(axis=1)
+        if bad_normals_mask.any():
+            normals[bad_normals_mask, :] = (np.random.RandomState(seed=777).rand(*verts.shape) - 0.5)[bad_normals_mask, :]
+            normals = normals / np.linalg.norm(normals, axis=-1)[:, np.newaxis]
+
+    normals = torch.from_numpy(normals).to(device=verts.device, dtype=verts.dtype)
+
+    if torch.any(torch.isnan(normals)):
+        raise ValueError("NaN normals :(")
+
+    return normals
+
+
+def build_tangent_frames(verts, faces, normals=None):
+
+    V = verts.shape[0]
+    dtype = verts.dtype
+    device = verts.device
+
+    if normals is None:
+        vert_normals = vertex_normals(verts, faces)  # (V,3)
+    elif isinstance(normals, np.ndarray):
+        vert_normals = torch.from_numpy(normals).to(dtype=dtype, device=device)
+    else:
+        vert_normals = normals
+
+    # = find an orthogonal basis
+
+    basis_cand1 = torch.tensor([1, 0, 0]).to(device=device, dtype=dtype).expand(V, -1)
+    basis_cand2 = torch.tensor([0, 1, 0]).to(device=device, dtype=dtype).expand(V, -1)
+
+    basisX = torch.where((torch.abs(dot(vert_normals, basis_cand1)) < 0.9).unsqueeze(-1), basis_cand1, basis_cand2)
+    basisX = project_to_tangent(basisX, vert_normals)
+    basisX = normalize(basisX)
+    basisY = cross(vert_normals, basisX)
+    frames = torch.stack((basisX, basisY, vert_normals), dim=-2)
+
+    if torch.any(torch.isnan(frames)):
+        raise ValueError("NaN coordinate frame! Must be very degenerate")
+
+    return frames
+
+
+def build_grad_point_cloud(verts, frames, n_neighbors_cloud=30):
+
+    verts_np = toNP(verts)
+    frames_np = toNP(frames)
+
+    _, neigh_inds = find_knn(verts, verts, n_neighbors_cloud, omit_diagonal=True, method='cpu_kd')
+    neigh_points = verts_np[neigh_inds, :]
+    neigh_vecs = neigh_points - verts_np[:, np.newaxis, :]
+
+    # TODO this could easily be way faster. For instance we could avoid the weird edges format and the corresponding pure-python loop via some numpy broadcasting of the same logic. The way it works right now is just to share code with the mesh version. But its low priority since its preprocessing code.
+
+    edge_inds_from = np.repeat(np.arange(verts.shape[0]), n_neighbors_cloud)
+    edges = np.stack((edge_inds_from, neigh_inds.flatten()))
+    edge_tangent_vecs = edge_tangent_vectors(verts, frames, edges)
+
+    return build_grad(verts_np, torch.tensor(edges), edge_tangent_vecs)
+
+
+def edge_tangent_vectors(verts, frames, edges):
+    edge_vecs = verts[edges[1, :], :] - verts[edges[0, :], :]
+    basisX = frames[edges[0, :], 0, :]
+    basisY = frames[edges[0, :], 1, :]
+
+    compX = dot(edge_vecs, basisX)
+    compY = dot(edge_vecs, basisY)
+    edge_tangent = torch.stack((compX, compY), dim=-1)
+
+    return edge_tangent
+
+
+def build_grad(verts, edges, edge_tangent_vectors):
+    """
+    Build a (V, V) complex sparse matrix grad operator. Given real inputs at vertices, produces a complex (vector value) at vertices giving the gradient. All values pointwise.
+    - edges: (2, E)
+    """
+
+    edges_np = toNP(edges)
+    edge_tangent_vectors_np = toNP(edge_tangent_vectors)
+
+    # TODO find a way to do this in pure numpy?
+
+    # Build outgoing neighbor lists
+    N = verts.shape[0]
+    vert_edge_outgoing = [[] for i in range(N)]
+    for iE in range(edges_np.shape[1]):
+        tail_ind = edges_np[0, iE]
+        tip_ind = edges_np[1, iE]
+        if tip_ind != tail_ind:
+            vert_edge_outgoing[tail_ind].append(iE)
+
+    # Build local inversion matrix for each vertex
+    row_inds = []
+    col_inds = []
+    data_vals = []
+    eps_reg = 1e-5
+    for iV in range(N):
+        n_neigh = len(vert_edge_outgoing[iV])
+
+        lhs_mat = np.zeros((n_neigh, 2))
+        rhs_mat = np.zeros((n_neigh, n_neigh + 1))
+        ind_lookup = [iV]
+        for i_neigh in range(n_neigh):
+            iE = vert_edge_outgoing[iV][i_neigh]
+            jV = edges_np[1, iE]
+            ind_lookup.append(jV)
+
+            edge_vec = edge_tangent_vectors[iE][:]
+            w_e = 1.
+
+            lhs_mat[i_neigh][:] = w_e * edge_vec
+            rhs_mat[i_neigh][0] = w_e * (-1)
+            rhs_mat[i_neigh][i_neigh + 1] = w_e * 1
+
+        lhs_T = lhs_mat.T
+        lhs_inv = np.linalg.inv(lhs_T @ lhs_mat + eps_reg * np.identity(2)) @ lhs_T
+
+        sol_mat = lhs_inv @ rhs_mat
+        sol_coefs = (sol_mat[0, :] + 1j * sol_mat[1, :]).T
+
+        for i_neigh in range(n_neigh + 1):
+            i_glob = ind_lookup[i_neigh]
+
+            row_inds.append(iV)
+            col_inds.append(i_glob)
+            data_vals.append(sol_coefs[i_neigh])
+
+    # build the sparse matrix
+    row_inds = np.array(row_inds)
+    col_inds = np.array(col_inds)
+    data_vals = np.array(data_vals)
+    mat = scipy.sparse.coo_matrix((data_vals, (row_inds, col_inds)), shape=(N, N)).tocsc()
+
+    return mat
+
+
+def compute_operators(verts, faces, k_eig, normals=None, normalize_desc=True):
+    """
+    Builds spectral operators for a mesh/point cloud. Constructs mass matrix, eigenvalues/vectors for Laplacian, and gradient matrix.
+
+    See get_operators() for a similar routine that wraps this one with a layer of caching.
+
+    Torch in / torch out.
+
+    Arguments:
+      - vertices: (V,3) vertex positions
+      - faces: (F,3) list of triangular faces. If empty, assumed to be a point cloud.
+      - k_eig: number of eigenvectors to use
+
+    Returns:
+      - frames: (V,3,3) X/Y/Z coordinate frame at each vertex. Z coordinate is normal (e.g. [:,2,:] for normals)
+      - massvec: (V) real diagonal of lumped mass matrix
+      - L: (VxV) real sparse matrix of (weak) Laplacian
+      - evals: (k) list of eigenvalues of the Laplacian
+      - evecs: (V,k) list of eigenvectors of the Laplacian 
+      - gradX: (VxV) sparse matrix which gives X-component of gradient in the local basis at the vertex
+      - gradY: same as gradX but for Y-component of gradient
+
+    PyTorch doesn't seem to like complex sparse matrices, so we store the "real" and "imaginary" (aka X and Y) gradient matrices separately, rather than as one complex sparse matrix.
+
+    Note: for a generalized eigenvalue problem, the mass matrix matters! The eigenvectors are only othrthonormal with respect to the mass matrix, like v^H M v, so the mass (given as the diagonal vector massvec) needs to be used in projections, etc.
+    """
+    device = verts.device
+    dtype = verts.dtype
+    V = verts.shape[0]
+    is_cloud = faces.numel() == 0
+
+    eps = 1e-8
+
+    verts_np = toNP(verts).astype(np.float64)
+    faces_np = toNP(faces)
+    frames = build_tangent_frames(verts, faces, normals=normals)
+    frames_np = toNP(frames)
+
+    # Build the scalar Laplacian
+    if is_cloud:
+        L, M = robust_laplacian.point_cloud_laplacian(verts_np)
+        massvec_np = M.diagonal()
+    else:
+        # L, M = robust_laplacian.mesh_laplacian(verts_np, faces_np)
+        # massvec_np = M.diagonal()
+        L = pp3d.cotan_laplacian(verts_np, faces_np, denom_eps=1e-10)
+        massvec_np = pp3d.vertex_areas(verts_np, faces_np)
+        massvec_np += eps * np.mean(massvec_np)
+
+    if (np.isnan(L.data).any()):
+        raise RuntimeError("NaN Laplace matrix")
+    if (np.isnan(massvec_np).any()):
+        raise RuntimeError("NaN mass matrix")
+
+    # Read off neighbors & rotations from the Laplacian
+    L_coo = L.tocoo()
+    inds_row = L_coo.row
+    inds_col = L_coo.col
+
+    # === Compute the eigenbasis
+    if k_eig > 0:
+
+        # Prepare matrices
+        L_eigsh = (L + scipy.sparse.identity(L.shape[0]) * eps).tocsc()
+        massvec_eigsh = massvec_np
+        Mmat = scipy.sparse.diags(massvec_eigsh)
+        eigs_sigma = eps
+
+        failcount = 0
+        while True:
+            try:
+                # We would be happy here to lower tol or maxiter since we don't need these to be super precise, but for some reason those parameters seem to have no effect
+                evals_np, evecs_np = sla.eigsh(L_eigsh, k=k_eig, M=Mmat, sigma=eigs_sigma)
+
+                # Clip off any eigenvalues that end up slightly negative due to numerical weirdness
+                evals_np = np.clip(evals_np, a_min=0., a_max=float('inf'))
+
+                break
+            except Exception as e:
+                print(e)
+                if (failcount > 3):
+                    raise ValueError("failed to compute eigendecomp")
+                failcount += 1
+                print("--- decomp failed; adding eps ===> count: " + str(failcount))
+                L_eigsh = L_eigsh + scipy.sparse.identity(L.shape[0]) * (eps * 10**failcount)
+
+    else:  #k_eig == 0
+        evals_np = np.zeros((0))
+        evecs_np = np.zeros((verts.shape[0], 0))
+
+    # == Build gradient matrices
+
+    # For meshes, we use the same edges as were used to build the Laplacian. For point clouds, use a whole local neighborhood
+    if is_cloud:
+        grad_mat_np = build_grad_point_cloud(verts, frames)
+    else:
+        edges = torch.tensor(np.stack((inds_row, inds_col), axis=0), device=device, dtype=faces.dtype)
+        edge_vecs = edge_tangent_vectors(verts, frames, edges)
+        grad_mat_np = build_grad(verts.cpu(), edges.cpu(), edge_vecs.cpu())
+
+    # Split complex gradient in to two real sparse mats (torch doesn't like complex sparse matrices)
+    gradX_np = np.real(grad_mat_np)
+    gradY_np = np.imag(grad_mat_np)
+
+    # === Convert back to torch
+    massvec = torch.from_numpy(massvec_np).to(device=device, dtype=dtype)
+    L = sparse_np_to_torch(L).to(device=device, dtype=dtype)
+    evals = torch.from_numpy(evals_np).to(device=device, dtype=dtype)
+    evecs = torch.from_numpy(evecs_np).to(device=device, dtype=dtype)
+    gradX = sparse_np_to_torch(gradX_np).to(device=device, dtype=dtype)
+    gradY = sparse_np_to_torch(gradY_np).to(device=device, dtype=dtype)
+
+    hks = torch.from_numpy(compute_hks(evecs_np, evals_np, massvec_np, n_descr=128, subsample_step=1, n_eig=128, normalize=normalize_desc)).to(device=device, dtype=dtype)
+    wks = torch.from_numpy(compute_wks(evecs_np, evals_np, massvec_np, n_descr=128, subsample_step=1, n_eig=128, normalize=normalize_desc)).to(device=device, dtype=dtype)
+
+
+    return frames, massvec, L, evals, evecs, gradX, gradY, hks, wks
+
+
+def get_all_operators(verts_list, faces_list, k_eig, op_cache_dir=None, normals=None):
+    N = len(verts_list)
+
+    frames = [None] * N
+    massvec = [None] * N
+    L = [None] * N
+    evals = [None] * N
+    evecs = [None] * N
+    gradX = [None] * N
+    gradY = [None] * N
+
+    inds = [i for i in range(N)]
+    # process in random order
+    # random.shuffle(inds)
+
+    for num, i in enumerate(inds):
+        print("get_all_operators() processing {} / {} {:.3f}%".format(num, N, num / N * 100))
+        if normals is None:
+            outputs = get_operators(verts_list[i], faces_list[i], k_eig, op_cache_dir)
+        else:
+            outputs = get_operators(verts_list[i], faces_list[i], k_eig, op_cache_dir, normals=normals[i])
+        frames[i] = outputs[0]
+        massvec[i] = outputs[1]
+        L[i] = outputs[2]
+        evals[i] = outputs[3]
+        evecs[i] = outputs[4]
+        gradX[i] = outputs[5]
+        gradY[i] = outputs[6]
+
+    return frames, massvec, L, evals, evecs, gradX, gradY
+
+
+def get_operators(verts, faces, k_eig=128, cache_path=None, normals=None, overwrite_cache=False):
+    """
+    See documentation for compute_operators(). This essentailly just wraps a call to compute_operators, using a cache if possible.
+    All arrays are always computed using double precision for stability, then truncated to single precision floats to store on disk, and finally returned as a tensor with dtype/device matching the `verts` input.
+    """
+    if type(verts) == torch.Tensor:
+        device = verts.device
+        dtype = verts.dtype
+        verts_np = toNP(verts)
+    else:
+        device = "cpu"
+        dtype = torch.float32
+        verts_np = verts.copy()
+        verts = torch.from_numpy(verts).float()
+    if type(faces) == torch.Tensor:
+        faces_np = toNP(faces)
+    else:
+        faces_np = faces.copy()
+        faces = torch.from_numpy(faces).to(device, dtype=torch.int64)
+    is_cloud = faces.numel() == 0
+
+    if (np.isnan(verts_np).any()):
+        raise RuntimeError("tried to construct operators from NaN verts")
+
+    # Check the cache directory
+    # Note 1: Collisions here are exceptionally unlikely, so we could probably just use the hash...
+    #         but for good measure we check values nonetheless.
+    # Note 2: There is a small possibility for race conditions to lead to bucket gaps or duplicate
+    #         entries in this cache. The good news is that that is totally fine, and at most slightly
+    #         slows performance with rare extra cache misses.
+    found = False
+    
+    if cache_path is not None:
+        op_cache_dir = os.path.dirname(cache_path)
+        ensure_dir_exists(op_cache_dir)
+        # print("Building operators for input with hash: " + hash_key_str)
+
+        # Search through buckets with matching hashes.  When the loop exits, this
+        # is the bucket index of the file we should write to.
+        i_cache_search = 0
+        while True:
+            try:
+                # print('loading path: ' + str(search_path))
+                npzfile = np.load(cache_path, allow_pickle=True)
+                cache_verts = npzfile["verts"]
+                cache_faces = npzfile["faces"]
+                cache_k_eig = npzfile["k_eig"].item()
+
+                # If the cache doesn't match, keep looking
+                if (not np.array_equal(verts, cache_verts)) or (not np.array_equal(faces, cache_faces)):
+                    i_cache_search += 1
+                    print("hash collision! searching next.")
+                    continue
+
+                # print("  cache hit!")
+
+                # If we're overwriting, or there aren't enough eigenvalues, just delete it; we'll create a new
+                # entry below more eigenvalues
+                if overwrite_cache:
+                    print("  overwriting cache by request")
+                    os.remove(cache_path)
+                    break
+
+                if cache_k_eig < k_eig:
+                    print("  overwriting cache --- not enough eigenvalues")
+                    os.remove(cache_path)
+                    break
+
+                if "L_data" not in npzfile:
+                    print("  overwriting cache --- entries are absent")
+                    os.remove(cache_path)
+                    break
+
+                def read_sp_mat(prefix):
+                    data = npzfile[prefix + "_data"]
+                    indices = npzfile[prefix + "_indices"]
+                    indptr = npzfile[prefix + "_indptr"]
+                    shape = npzfile[prefix + "_shape"]
+                    mat = scipy.sparse.csc_matrix((data, indices, indptr), shape=shape)
+                    return mat
+
+                # This entry matches! Return it.
+                frames = npzfile["frames"]
+                mass = npzfile["mass"]
+                L = read_sp_mat("L")
+                evals = npzfile["evals"][:k_eig]
+                evecs = npzfile["evecs"][:, :k_eig]
+                gradX = read_sp_mat("gradX")
+                gradY = read_sp_mat("gradY")
+                frames = npzfile["hks"]
+                mass = npzfile["wks"]
+
+                frames = torch.from_numpy(frames).to(device=device, dtype=dtype)
+                mass = torch.from_numpy(mass).to(device=device, dtype=dtype)
+                L = sparse_np_to_torch(L).to(device=device, dtype=dtype)
+                evals = torch.from_numpy(evals).to(device=device, dtype=dtype)
+                evecs = torch.from_numpy(evecs).to(device=device, dtype=dtype)
+                gradX = sparse_np_to_torch(gradX).to(device=device, dtype=dtype)
+                gradY = sparse_np_to_torch(gradY).to(device=device, dtype=dtype)
+                hks = torch.from_numpy(hks).to(device=device, dtype=dtype)
+                wks = torch.from_numpy(wks).to(device=device, dtype=dtype)
+
+                found = True
+
+                break
+
+            except FileNotFoundError:
+                print("  cache miss -- constructing operators")
+                break
+
+            except Exception as E:
+                print("unexpected error loading file: " + str(E))
+                print("-- constructing operators")
+                break
+
+    if not found:
+
+        # No matching entry found; recompute.
+        frames, mass, L, evals, evecs, gradX, gradY, hks, wks = compute_operators(verts, faces, k_eig, normals=normals)
+
+        dtype_np = np.float32
+
+        # Store it in the cache
+        if op_cache_dir is not None:
+
+            L_np = sparse_torch_to_np(L).astype(dtype_np)
+            gradX_np = sparse_torch_to_np(gradX).astype(dtype_np)
+            gradY_np = sparse_torch_to_np(gradY).astype(dtype_np)
+
+            np.savez(
+                cache_path,
+                verts=verts_np.astype(dtype_np),
+                frames=toNP(frames).astype(dtype_np),
+                faces=faces_np,
+                k_eig=k_eig,
+                mass=toNP(mass).astype(dtype_np),
+                L_data=L_np.data.astype(dtype_np),
+                L_indices=L_np.indices,
+                L_indptr=L_np.indptr,
+                L_shape=L_np.shape,
+                evals=toNP(evals).astype(dtype_np),
+                evecs=toNP(evecs).astype(dtype_np),
+                gradX_data=gradX_np.data.astype(dtype_np),
+                gradX_indices=gradX_np.indices,
+                gradX_indptr=gradX_np.indptr,
+                gradX_shape=gradX_np.shape,
+                gradY_data=gradY_np.data.astype(dtype_np),
+                gradY_indices=gradY_np.indices,
+                gradY_indptr=gradY_np.indptr,
+                gradY_shape=gradY_np.shape,
+                hks=hks,
+                wks=wks
+            )
+
+    return frames, mass, L, evals, evecs, gradX, gradY, hks, wks
+
+
+def load_operators(filepath):
+    npzfile = np.load(filepath, allow_pickle=False)
+
+    def read_sp_mat(prefix):
+        data = npzfile[prefix + '_data']
+        indices = npzfile[prefix + '_indices']
+        indptr = npzfile[prefix + '_indptr']
+        shape = npzfile[prefix + '_shape']
+        mat = scipy.sparse.csc_matrix((data, indices, indptr), shape=shape)
+        return mat
+    if 'verts' in npzfile:
+        keyverts = 'verts'
+    else:
+        keyverts = 'vertices'
+    return dict(
+        vertices=npzfile[keyverts],
+        faces=npzfile['faces'],
+        frames=npzfile['frames'],
+        mass=npzfile['mass'],
+        L=read_sp_mat('L'),
+        evals=npzfile['evals'],
+        evecs=npzfile['evecs'],
+        gradX=read_sp_mat('gradX'),
+        gradY=read_sp_mat('gradY'),
+        hks=npzfile['hks'],
+        wks=npzfile['wks'], 
+    )
+
+
+def compute_operators_small(verts, faces, k_eig):
+    """
+    Builds spectral operators for a mesh/point cloud. Constructs mass matrix, eigenvalues/vectors for Laplacian, and gradient matrix.
+
+    See get_operators() for a similar routine that wraps this one with a layer of caching.
+
+    Torch in / torch out.
+
+    Arguments:
+      - vertices: (V,3) vertex positions
+      - faces: (F,3) list of triangular faces. If empty, assumed to be a point cloud.
+      - k_eig: number of eigenvectors to use
+
+    Returns:
+      - massvec: (V) real diagonal of lumped mass matrix
+      - L: (VxV) real sparse matrix of (weak) Laplacian
+      - evals: (k) list of eigenvalues of the Laplacian
+      - evecs: (V,k) list of eigenvectors of the Laplacian 
+      - gradX: (VxV) sparse matrix which gives X-component of gradient in the local basis at the vertex
+      - gradY: same as gradX but for Y-component of gradient
+
+    PyTorch doesn't seem to like complex sparse matrices, so we store the "real" and "imaginary" (aka X and Y) gradient matrices separately, rather than as one complex sparse matrix.
+
+    Note: for a generalized eigenvalue problem, the mass matrix matters! The eigenvectors are only othrthonormal with respect to the mass matrix, like v^H M v, so the mass (given as the diagonal vector massvec) needs to be used in projections, etc.
+    """
+
+    device = verts.device
+    dtype = verts.dtype
+    V = verts.shape[0]
+    is_cloud = faces.numel() == 0
+
+    eps = 1e-8
+
+    verts_np = toNP(verts).astype(np.float64)
+    faces_np = toNP(faces)
+    # Build the scalar Laplacian
+    if is_cloud:
+        L, M = robust_laplacian.point_cloud_laplacian(verts_np)
+        massvec_np = M.diagonal()
+    else:
+        # L, M = robust_laplacian.mesh_laplacian(verts_np, faces_np)
+        # massvec_np = M.diagonal()
+        L = pp3d.cotan_laplacian(verts_np, faces_np, denom_eps=1e-10)
+        massvec_np = pp3d.vertex_areas(verts_np, faces_np)
+        massvec_np += eps * np.mean(massvec_np)
+
+    if (np.isnan(L.data).any()):
+        raise RuntimeError("NaN Laplace matrix")
+    if (np.isnan(massvec_np).any()):
+        raise RuntimeError("NaN mass matrix")
+
+    # Read off neighbors & rotations from the Laplacian
+    L_coo = L.tocoo()
+    inds_row = L_coo.row
+    inds_col = L_coo.col
+
+    # === Compute the eigenbasis
+    if k_eig > 0:
+
+        # Prepare matrices
+        L_eigsh = (L + scipy.sparse.identity(L.shape[0]) * eps).tocsc()
+        massvec_eigsh = massvec_np
+        Mmat = scipy.sparse.diags(massvec_eigsh)
+        eigs_sigma = eps
+
+        failcount = 0
+        while True:
+            try:
+                # We would be happy here to lower tol or maxiter since we don't need these to be super precise, but for some reason those parameters seem to have no effect
+                evals_np, evecs_np = sla.eigsh(L_eigsh, k=k_eig, M=Mmat, sigma=eigs_sigma)
+
+                # Clip off any eigenvalues that end up slightly negative due to numerical weirdness
+                evals_np = np.clip(evals_np, a_min=0., a_max=float('inf'))
+
+                break
+            except Exception as e:
+                print(e)
+                if (failcount > 3):
+                    raise ValueError("failed to compute eigendecomp")
+                failcount += 1
+                print("--- decomp failed; adding eps ===> count: " + str(failcount))
+                L_eigsh = L_eigsh + scipy.sparse.identity(L.shape[0]) * (eps * 10**failcount)
+
+    else:  #k_eig == 0
+        evals_np = np.zeros((0))
+        evecs_np = np.zeros((verts.shape[0], 0))
+    # Split complex gradient in to two real sparse mats (torch doesn't like complex sparse matrices)
+
+    # === Convert back to torch
+    massvec = torch.from_numpy(massvec_np).to(device=device, dtype=dtype)
+    L = sparse_np_to_torch(L).to(device=device, dtype=dtype)
+    evals = torch.from_numpy(evals_np).to(device=device, dtype=dtype)
+    evecs = torch.from_numpy(evecs_np).to(device=device, dtype=dtype)
+
+    return massvec, L, evals, evecs
+
+
+def get_operators_small(verts, faces, k_eig=128, cache_path=None, overwrite_cache=False):
+    """
+    See documentation for compute_operators(). This essentailly just wraps a call to compute_operators, using a cache if possible.
+    All arrays are always computed using double precision for stability, then truncated to single precision floats to store on disk, and finally returned as a tensor with dtype/device matching the `verts` input.
+    """
+    if type(verts) == torch.Tensor:
+        device = verts.device
+        dtype = verts.dtype
+        verts_np = toNP(verts)
+    else:
+        device = "cpu"
+        dtype = torch.float32
+        verts_np = verts.copy()
+        verts = torch.from_numpy(verts).float()
+    if type(faces) == torch.Tensor:
+        faces_np = toNP(faces)
+    else:
+        faces_np = faces.copy()
+        faces = torch.from_numpy(faces).to(device, dtype=torch.int64)
+    is_cloud = faces.numel() == 0
+
+    if (np.isnan(verts_np).any()):
+        raise RuntimeError("tried to construct operators from NaN verts")
+
+    # Check the cache directory
+    # Note 1: Collisions here are exceptionally unlikely, so we could probably just use the hash...
+    #         but for good measure we check values nonetheless.
+    # Note 2: There is a small possibility for race conditions to lead to bucket gaps or duplicate
+    #         entries in this cache. The good news is that that is totally fine, and at most slightly
+    #         slows performance with rare extra cache misses.
+    found = False
+    
+    if cache_path is not None:
+        op_cache_dir = os.path.dirname(cache_path)
+        ensure_dir_exists(op_cache_dir)
+        # print("Building operators for input with hash: " + hash_key_str)
+
+        # Search through buckets with matching hashes.  When the loop exits, this
+        # is the bucket index of the file we should write to.
+        i_cache_search = 0
+        while True:
+            try:
+                # print('loading path: ' + str(search_path))
+                npzfile = np.load(cache_path, allow_pickle=True)
+                cache_verts = npzfile["verts"]
+                cache_faces = npzfile["faces"]
+                cache_k_eig = npzfile["k_eig"].item()
+
+                # If the cache doesn't match, keep looking
+                if (not np.array_equal(verts, cache_verts)) or (not np.array_equal(faces, cache_faces)):
+                    i_cache_search += 1
+                    print("hash collision! searching next.")
+                    continue
+
+                # print("  cache hit!")
+
+                # If we're overwriting, or there aren't enough eigenvalues, just delete it; we'll create a new
+                # entry below more eigenvalues
+                if overwrite_cache:
+                    print("  overwriting cache by request")
+                    os.remove(cache_path)
+                    break
+
+                if cache_k_eig < k_eig:
+                    print("  overwriting cache --- not enough eigenvalues")
+                    os.remove(cache_path)
+                    break
+
+                if "L_data" not in npzfile:
+                    print("  overwriting cache --- entries are absent")
+                    os.remove(cache_path)
+                    break
+
+                def read_sp_mat(prefix):
+                    data = npzfile[prefix + "_data"]
+                    indices = npzfile[prefix + "_indices"]
+                    indptr = npzfile[prefix + "_indptr"]
+                    shape = npzfile[prefix + "_shape"]
+                    mat = scipy.sparse.csc_matrix((data, indices, indptr), shape=shape)
+                    return mat
+
+                # This entry matches! Return it.
+                mass = npzfile["mass"]
+                L = read_sp_mat("L")
+                evals = npzfile["evals"][:k_eig]
+                evecs = npzfile["evecs"][:, :k_eig]
+                mass = torch.from_numpy(mass).to(device=device, dtype=dtype)
+                L = sparse_np_to_torch(L).to(device=device, dtype=dtype)
+                evals = torch.from_numpy(evals).to(device=device, dtype=dtype)
+                evecs = torch.from_numpy(evecs).to(device=device, dtype=dtype)
+
+                found = True
+
+                break
+
+            except FileNotFoundError:
+                print(cache_path)
+                print("  cache miss -- constructing operators")
+                break
+
+            except Exception as E:
+                print("unexpected error loading file: " + str(E))
+                print("-- constructing operators")
+                break
+
+    if not found:
+
+        # No matching entry found; recompute.
+        mass, L, evals, evecs, = compute_operators_small(verts, faces, k_eig)
+
+        dtype_np = np.float32
+
+        # Store it in the cache
+        if op_cache_dir is not None:
+
+            L_np = sparse_torch_to_np(L).astype(dtype_np)
+
+            np.savez(
+                cache_path,
+                verts=verts_np.astype(dtype_np),
+                faces=faces_np,
+                k_eig=k_eig,
+                mass=toNP(mass).astype(dtype_np),
+                L_data=L_np.data.astype(dtype_np),
+                L_indices=L_np.indices,
+                L_indptr=L_np.indptr,
+                L_shape=L_np.shape,
+                evals=toNP(evals).astype(dtype_np),
+                evecs=toNP(evecs).astype(dtype_np),
+            )
+
+    return mass, L, evals, evecs
+
+
+def to_basis(values, basis, massvec):
+    """
+    Transform data in to an orthonormal basis (where orthonormal is wrt to massvec)
+    Inputs:
+      - values: (B,V,D)
+      - basis: (B,V,K)
+      - massvec: (B,V)
+    Outputs:
+      - (B,K,D) transformed values
+    """
+    basisT = basis.transpose(-2, -1)
+    return torch.matmul(basisT, values * massvec.unsqueeze(-1))
+
+
+
+def normalize_positions(pos, faces=None, method='mean', scale_method='max_rad'):
+    # center and unit-scale positions
+
+    if method == 'mean':
+        # center using the average point position
+        pos = (pos - torch.mean(pos, dim=-2, keepdim=True))
+    elif method == 'bbox':
+        # center via the middle of the axis-aligned bounding box
+        bbox_min = torch.min(pos, dim=-2).values
+        bbox_max = torch.max(pos, dim=-2).values
+        center = (bbox_max + bbox_min) / 2.
+        pos -= center.unsqueeze(-2)
+    else:
+        raise ValueError("unrecognized method")
+
+    if scale_method == 'max_rad':
+        scale = torch.max(norm(pos), dim=-1, keepdim=True).values.unsqueeze(-1)
+        pos = pos / scale
+    elif scale_method == 'area':
+        if faces is None:
+            raise ValueError("must pass faces for area normalization")
+        coords = pos[faces]
+        vec_A = coords[:, 1, :] - coords[:, 0, :]
+        vec_B = coords[:, 2, :] - coords[:, 0, :]
+        face_areas = torch.norm(torch.cross(vec_A, vec_B, dim=-1), dim=1) * 0.5
+        total_area = torch.sum(face_areas)
+        scale = (1. / torch.sqrt(total_area))
+        pos = pos * scale
+    else:
+        raise ValueError("unrecognized scale method")
+    return pos
+
+
+# Finds the k nearest neighbors of source on target.
+# Return is two tensors (distances, indices). Returned points will be sorted in increasing order of distance.
+def find_knn(points_source, points_target, k, largest=False, omit_diagonal=False, method='brute'):
+
+    if omit_diagonal and points_source.shape[0] != points_target.shape[0]:
+        raise ValueError("omit_diagonal can only be used when source and target are same shape")
+
+    if method != 'cpu_kd' and points_source.shape[0] * points_target.shape[0] > 1e8:
+        method = 'cpu_kd'
+        print("switching to cpu_kd knn")
+
+    if method == 'brute':
+
+        # Expand so both are NxMx3 tensor
+        points_source_expand = points_source.unsqueeze(1)
+        points_source_expand = points_source_expand.expand(-1, points_target.shape[0], -1)
+        points_target_expand = points_target.unsqueeze(0)
+        points_target_expand = points_target_expand.expand(points_source.shape[0], -1, -1)
+
+        diff_mat = points_source_expand - points_target_expand
+        dist_mat = norm(diff_mat)
+
+        if omit_diagonal:
+            torch.diagonal(dist_mat)[:] = float('inf')
+
+        result = torch.topk(dist_mat, k=k, largest=largest, sorted=True)
+        return result
+
+    elif method == 'cpu_kd':
+
+        if largest:
+            raise ValueError("can't do largest with cpu_kd")
+
+        points_source_np = toNP(points_source)
+        points_target_np = toNP(points_target)
+
+        # Build the tree
+        kd_tree = sklearn.neighbors.KDTree(points_target_np)
+
+        k_search = k + 1 if omit_diagonal else k
+        _, neighbors = kd_tree.query(points_source_np, k=k_search)
+
+        if omit_diagonal:
+            # Mask out self element
+            mask = neighbors != np.arange(neighbors.shape[0])[:, np.newaxis]
+
+            # make sure we mask out exactly one element in each row, in rare case of many duplicate points
+            mask[np.sum(mask, axis=1) == mask.shape[1], -1] = False
+
+            neighbors = neighbors[mask].reshape((neighbors.shape[0], neighbors.shape[1] - 1))
+
+        inds = torch.tensor(neighbors, device=points_source.device, dtype=torch.int64)
+        dists = norm(points_source.unsqueeze(1).expand(-1, k, -1) - points_target[inds])
+
+        return dists, inds
+
+    else:
+        raise ValueError("unrecognized method")
+

utils/io.py

@@ -0,0 +1,78 @@
+from pathlib import Path
+import json
+import os
+import os.path as osp
+import re
+import shutil
+
+
+def is_number(s):
+    try:
+        float(s)
+        return True
+    except ValueError:
+        return False
+
+
+def may_create_folder(folder_path):
+    if not osp.exists(folder_path):
+        oldmask = os.umask(000)
+        os.makedirs(folder_path, mode=0o777)
+        os.umask(oldmask)
+        return True
+    return False
+
+
+def make_clean_folder(folder_path):
+    success = may_create_folder(folder_path)
+    if not success:
+        shutil.rmtree(folder_path)
+        may_create_folder(folder_path)
+
+
+def sorted_alphanum(file_list_ordered):
+    convert = lambda text: int(text) if text.isdigit() else text
+    alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key) if len(c) > 0]
+    return sorted(file_list_ordered, key=alphanum_key)
+
+
+def list_files(folder_path, name_filter, alphanum_sort=False):
+    file_list = [p.name for p in list(Path(folder_path).glob(name_filter))]
+    if alphanum_sort:
+        return sorted_alphanum(file_list)
+    else:
+        return sorted(file_list)
+
+
+def list_folders(folder_path, name_filter=None, alphanum_sort=False):
+    folders = list()
+    for subfolder in Path(folder_path).iterdir():
+        if subfolder.is_dir() and not subfolder.name.startswith('.'):
+            folder_name = subfolder.name
+            if name_filter is not None:
+                if name_filter in folder_name:
+                    folders.append(folder_name)
+            else:
+                folders.append(folder_name)
+    if alphanum_sort:
+        return sorted_alphanum(folders)
+    else:
+        return sorted(folders)
+
+
+def read_lines(file_path):
+    with open(file_path, 'r') as fin:
+        lines = [line.strip() for line in fin.readlines() if len(line.strip()) > 0]
+    return lines
+
+
+def read_strings(file_path):
+    with open(file_path, 'r') as fin:
+        ret = fin.readlines()
+    return ''.join(ret)
+
+
+def read_json(filepath):
+    with open(filepath, 'r') as fh:
+        ret = json.load(fh)
+    return ret

utils/layers.py

@@ -0,0 +1,430 @@
+import sys
+import os
+import os.path
+import random
+
+import scipy
+import scipy.sparse.linalg as sla
+# ^^^ we NEED to import scipy before torch, or it crashes :(
+# (observed on Ubuntu 20.04 w/ torch 1.6.0 and scipy 1.5.2 installed via conda)
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+ROOT_DIR = os.path.join(os.path.abspath(os.path.dirname(__file__)), '../')
+if ROOT_DIR not in sys.path:
+    sys.path.append(ROOT_DIR)
+
+from diffusion_net.utils import toNP
+from diffusion_net.geometry import to_basis, from_basis
+
+
+class LearnedTimeDiffusion(nn.Module):
+    """
+    Applies diffusion with learned per-channel t.
+
+    In the spectral domain this becomes 
+        f_out = e ^ (lambda_i t) f_in
+
+    Inputs:
+      - values: (V,C) in the spectral domain
+      - L: (V,V) sparse laplacian
+      - evals: (K) eigenvalues
+      - mass: (V) mass matrix diagonal
+
+      (note: L/evals may be omitted as None depending on method)
+    Outputs:
+      - (V,C) diffused values 
+    """
+
+    def __init__(self, C_inout, method='spectral'):
+        super(LearnedTimeDiffusion, self).__init__()
+        self.C_inout = C_inout
+        self.diffusion_time = nn.Parameter(torch.Tensor(C_inout))  # (C)
+        self.method = method  # one of ['spectral', 'implicit_dense']
+
+        nn.init.constant_(self.diffusion_time, 0.0)
+
+    def forward(self, x, L, mass, evals, evecs):
+
+        # project times to the positive halfspace
+        # (and away from 0 in the incredibly rare chance that they get stuck)
+        with torch.no_grad():
+            self.diffusion_time.data = torch.clamp(self.diffusion_time, min=1e-8)
+
+        if x.shape[-1] != self.C_inout:
+            raise ValueError("Tensor has wrong shape = {}. Last dim shape should have number of channels = {}".format(
+                x.shape, self.C_inout))
+
+        if self.method == 'spectral':
+
+            # Transform to spectral
+            x_spec = to_basis(x, evecs, mass)
+
+            # Diffuse
+            time = self.diffusion_time
+            diffusion_coefs = torch.exp(-evals.unsqueeze(-1) * time.unsqueeze(0))
+            x_diffuse_spec = diffusion_coefs * x_spec
+
+            # Transform back to per-vertex
+            x_diffuse = from_basis(x_diffuse_spec, evecs)
+
+        elif self.method == 'implicit_dense':
+            V = x.shape[-2]
+
+            # Form the dense matrices (M + tL) with dims (B,C,V,V)
+            mat_dense = L.to_dense().unsqueeze(1).expand(-1, self.C_inout, V, V).clone()
+            mat_dense *= self.diffusion_time.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+            mat_dense += torch.diag_embed(mass).unsqueeze(1)
+
+            # Factor the system
+            cholesky_factors = torch.linalg.cholesky(mat_dense)
+
+            # Solve the system
+            rhs = x * mass.unsqueeze(-1)
+            rhsT = torch.transpose(rhs, 1, 2).unsqueeze(-1)
+            sols = torch.cholesky_solve(rhsT, cholesky_factors)
+            x_diffuse = torch.transpose(sols.squeeze(-1), 1, 2)
+
+        else:
+            raise ValueError("unrecognized method")
+
+        return x_diffuse
+
+
+class SpatialGradientFeatures(nn.Module):
+    """
+    Compute dot-products between input vectors. Uses a learned complex-linear layer to keep dimension down.
+    
+    Input:
+        - vectors: (V,C,2)
+    Output:
+        - dots: (V,C) dots 
+    """
+
+    def __init__(self, C_inout, with_gradient_rotations=True):
+        super(SpatialGradientFeatures, self).__init__()
+
+        self.C_inout = C_inout
+        self.with_gradient_rotations = with_gradient_rotations
+
+        if (self.with_gradient_rotations):
+            self.A_re = nn.Linear(self.C_inout, self.C_inout, bias=False)
+            self.A_im = nn.Linear(self.C_inout, self.C_inout, bias=False)
+        else:
+            self.A = nn.Linear(self.C_inout, self.C_inout, bias=False)
+
+        # self.norm = nn.InstanceNorm1d(C_inout)
+
+    def forward(self, vectors):
+
+        vectorsA = vectors  # (V,C)
+
+        if self.with_gradient_rotations:
+            vectorsBreal = self.A_re(vectors[..., 0]) - self.A_im(vectors[..., 1])
+            vectorsBimag = self.A_re(vectors[..., 1]) + self.A_im(vectors[..., 0])
+        else:
+            vectorsBreal = self.A(vectors[..., 0])
+            vectorsBimag = self.A(vectors[..., 1])
+
+        dots = vectorsA[..., 0] * vectorsBreal + vectorsA[..., 1] * vectorsBimag
+
+        return torch.tanh(dots)
+
+
+class MiniMLP(nn.Sequential):
+    '''
+    A simple MLP with configurable hidden layer sizes.
+    '''
+
+    def __init__(self, layer_sizes, dropout=False, activation=nn.ReLU, name="miniMLP"):
+        super(MiniMLP, self).__init__()
+
+        for i in range(len(layer_sizes) - 1):
+            is_last = (i + 2 == len(layer_sizes))
+
+            if dropout and i > 0:
+                self.add_module(name + "_mlp_layer_dropout_{:03d}".format(i), nn.Dropout(p=.5))
+
+            # Affine map
+            self.add_module(
+                name + "_mlp_layer_{:03d}".format(i),
+                nn.Linear(
+                    layer_sizes[i],
+                    layer_sizes[i + 1],
+                ),
+            )
+
+            # Nonlinearity
+            # (but not on the last layer)
+            if not is_last:
+                self.add_module(name + "_mlp_act_{:03d}".format(i), activation())
+
+
+class DiffusionNetBlock(nn.Module):
+    """
+    Inputs and outputs are defined at vertices
+    """
+
+    def __init__(self,
+                 C_width,
+                 mlp_hidden_dims,
+                 dropout=True,
+                 diffusion_method='spectral',
+                 with_gradient_features=True,
+                 with_gradient_rotations=True):
+        super(DiffusionNetBlock, self).__init__()
+
+        # Specified dimensions
+        self.C_width = C_width
+        self.mlp_hidden_dims = mlp_hidden_dims
+
+        self.dropout = dropout
+        self.with_gradient_features = with_gradient_features
+        self.with_gradient_rotations = with_gradient_rotations
+
+        # Diffusion block
+        self.diffusion = LearnedTimeDiffusion(self.C_width, method=diffusion_method)
+
+        self.MLP_C = 2 * self.C_width
+
+        if self.with_gradient_features:
+            self.gradient_features = SpatialGradientFeatures(self.C_width, with_gradient_rotations=self.with_gradient_rotations)
+            self.MLP_C += self.C_width
+
+        # MLPs
+        self.mlp = MiniMLP([self.MLP_C] + self.mlp_hidden_dims + [self.C_width], dropout=self.dropout)
+
+    def forward(self, x_in, mass, L, evals, evecs, gradX, gradY):
+
+        # Manage dimensions
+        B = x_in.shape[0]  # batch dimension
+        if x_in.shape[-1] != self.C_width:
+            raise ValueError("Tensor has wrong shape = {}. Last dim shape should have number of channels = {}".format(
+                x_in.shape, self.C_width))
+
+        # Diffusion block
+        x_diffuse = self.diffusion(x_in, L, mass, evals, evecs)
+
+        # Compute gradient features, if using
+        if self.with_gradient_features:
+
+            # Compute gradients
+            x_grads = [
+            ]  # Manually loop over the batch (if there is a batch dimension) since torch.mm() doesn't support batching
+            for b in range(B):
+                # gradient after diffusion
+                x_gradX = torch.mm(gradX[b, ...], x_diffuse[b, ...])
+                x_gradY = torch.mm(gradY[b, ...], x_diffuse[b, ...])
+
+                x_grads.append(torch.stack((x_gradX, x_gradY), dim=-1))
+            x_grad = torch.stack(x_grads, dim=0)
+
+            # Evaluate gradient features
+            x_grad_features = self.gradient_features(x_grad)
+
+            # Stack inputs to mlp
+            feature_combined = torch.cat((x_in, x_diffuse, x_grad_features), dim=-1)
+        else:
+            # Stack inputs to mlp
+            feature_combined = torch.cat((x_in, x_diffuse), dim=-1)
+
+        # Apply the mlp
+        x0_out = self.mlp(feature_combined)
+
+        # Skip connection
+        x0_out = x0_out + x_in
+
+        return x0_out
+
+
+class DiffusionNet(nn.Module):
+
+    def __init__(self,
+                 C_in,
+                 C_out,
+                 C_width=128,
+                 N_block=4,
+                 last_activation=None,
+                 outputs_at='vertices',
+                 mlp_hidden_dims=None,
+                 dropout=True,
+                 with_gradient_features=True,
+                 with_gradient_rotations=True,
+                 diffusion_method='spectral',
+                 num_eigenbasis=128):
+        """
+        Construct a DiffusionNet.
+
+        Parameters:
+            C_in (int):                     input dimension 
+            C_out (int):                    output dimension 
+            last_activation (func)          a function to apply to the final outputs of the network, such as torch.nn.functional.log_softmax (default: None)
+            outputs_at (string)             produce outputs at various mesh elements by averaging from vertices. One of ['vertices', 'edges', 'faces', 'global_mean']. (default 'vertices', aka points for a point cloud)
+            C_width (int):                  dimension of internal DiffusionNet blocks (default: 128)
+            N_block (int):                  number of DiffusionNet blocks (default: 4)
+            mlp_hidden_dims (list of int):  a list of hidden layer sizes for MLPs (default: [C_width, C_width])
+            dropout (bool):                 if True, internal MLPs use dropout (default: True)
+            diffusion_method (string):      how to evaluate diffusion, one of ['spectral', 'implicit_dense']. If implicit_dense is used, can set k_eig=0, saving precompute.
+            with_gradient_features (bool):  if True, use gradient features (default: True)
+            with_gradient_rotations (bool): if True, use gradient also learn a rotation of each gradient. Set to True if your surface has consistently oriented normals, and False otherwise (default: True)
+            num_eigenbasis (int):           for trunking the eigenvalues eigenvectors
+        """
+
+        super(DiffusionNet, self).__init__()
+
+        ## Store parameters
+
+        # Basic parameters
+        self.C_in = C_in
+        self.C_out = C_out
+        self.C_width = C_width
+        self.N_block = N_block
+
+        # Outputs
+        self.last_activation = last_activation
+        self.outputs_at = outputs_at
+        if outputs_at not in ['vertices', 'edges', 'faces', 'global_mean']:
+            raise ValueError("invalid setting for outputs_at")
+
+        # MLP options
+        if mlp_hidden_dims == None:
+            mlp_hidden_dims = [C_width, C_width]
+        self.mlp_hidden_dims = mlp_hidden_dims
+        self.dropout = dropout
+
+        # Diffusion
+        self.diffusion_method = diffusion_method
+        if diffusion_method not in ['spectral', 'implicit_dense']:
+            raise ValueError("invalid setting for diffusion_method")
+        self.num_eigenbasis = num_eigenbasis
+
+        # Gradient features
+        self.with_gradient_features = with_gradient_features
+        self.with_gradient_rotations = with_gradient_rotations
+
+        ## Set up the network
+
+        # First and last affine layers
+        self.first_lin = nn.Linear(C_in, C_width)
+        self.last_lin = nn.Linear(C_width, C_out)
+
+        # DiffusionNet blocks
+        self.blocks = []
+        for i_block in range(self.N_block):
+            block = DiffusionNetBlock(C_width=C_width,
+                                      mlp_hidden_dims=mlp_hidden_dims,
+                                      dropout=dropout,
+                                      diffusion_method=diffusion_method,
+                                      with_gradient_features=with_gradient_features,
+                                      with_gradient_rotations=with_gradient_rotations)
+
+            self.blocks.append(block)
+            self.add_module("block_" + str(i_block), self.blocks[-1])
+
+    def forward(self, x_in, mass, L=None, evals=None, evecs=None, gradX=None, gradY=None, edges=None, faces=None):
+        """
+        A forward pass on the DiffusionNet.
+
+        In the notation below, dimension are:
+            - C is the input channel dimension (C_in on construction)
+            - C_OUT is the output channel dimension (C_out on construction)
+            - N is the number of vertices/points, which CAN be different for each forward pass
+            - B is an OPTIONAL batch dimension
+            - K_EIG is the number of eigenvalues used for spectral acceleration
+        Generally, our data layout it is [N,C] or [B,N,C].
+
+        Call get_operators() to generate geometric quantities mass/L/evals/evecs/gradX/gradY. Note that depending on the options for the DiffusionNet, not all are strictly necessary.
+
+        Parameters:
+            x_in (tensor):      Input features, dimension [N,C] or [B,N,C]
+            mass (tensor):      Mass vector, dimension [N] or [B,N]
+            L (tensor):         Laplace matrix, sparse tensor with dimension [N,N] or [B,N,N]
+            evals (tensor):     Eigenvalues of Laplace matrix, dimension [K_EIG] or [B,K_EIG]
+            evecs (tensor):     Eigenvectors of Laplace matrix, dimension [N,K_EIG] or [B,N,K_EIG]
+            gradX (tensor):     Half of gradient matrix, sparse real tensor with dimension [N,N] or [B,N,N]
+            gradY (tensor):     Half of gradient matrix, sparse real tensor with dimension [N,N] or [B,N,N]
+
+        Returns:
+            x_out (tensor):    Output with dimension [N,C_out] or [B,N,C_out]
+        """
+
+        ## Check dimensions, and append batch dimension if not given
+        if x_in.shape[-1] != self.C_in:
+            raise ValueError("DiffusionNet was constructed with C_in={}, but x_in has last dim={}".format(
+                self.C_in, x_in.shape[-1]))
+        N = x_in.shape[-2]
+        if len(x_in.shape) == 2:
+            appended_batch_dim = True
+
+            # add a batch dim to all inputs
+            x_in = x_in.unsqueeze(0)
+            mass = mass.unsqueeze(0)
+            if L != None:
+                L = L.unsqueeze(0)
+            if evals != None:
+                evals = evals.unsqueeze(0)
+            if evecs != None:
+                evecs = evecs.unsqueeze(0)
+            if gradX != None:
+                gradX = gradX.unsqueeze(0)
+            if gradY != None:
+                gradY = gradY.unsqueeze(0)
+            if edges != None:
+                edges = edges.unsqueeze(0)
+            if faces != None:
+                faces = faces.unsqueeze(0)
+
+        elif len(x_in.shape) == 3:
+            appended_batch_dim = False
+
+        else:
+            raise ValueError("x_in should be tensor with shape [N,C] or [B,N,C]")
+
+        evals = evals[..., :self.num_eigenbasis]
+        evecs = evecs[..., :self.num_eigenbasis]
+
+        # Apply the first linear layer
+        x = self.first_lin(x_in)
+
+        # Apply each of the blocks
+        for b in self.blocks:
+            x = b(x, mass, L, evals, evecs, gradX, gradY)
+
+        # Apply the last linear layer
+        x = self.last_lin(x)
+
+        # Remap output to faces/edges if requested
+        if self.outputs_at == 'vertices':
+            x_out = x
+
+        elif self.outputs_at == 'edges':
+            # Remap to edges
+            x_gather = x.unsqueeze(-1).expand(-1, -1, -1, 2)
+            edges_gather = edges.unsqueeze(2).expand(-1, -1, x.shape[-1], -1)
+            xe = torch.gather(x_gather, 1, edges_gather)
+            x_out = torch.mean(xe, dim=-1)
+
+        elif self.outputs_at == 'faces':
+            # Remap to faces
+            x_gather = x.unsqueeze(-1).expand(-1, -1, -1, 3)
+            faces_gather = faces.unsqueeze(2).expand(-1, -1, x.shape[-1], -1)
+            xf = torch.gather(x_gather, 1, faces_gather)
+            x_out = torch.mean(xf, dim=-1)
+
+        elif self.outputs_at == 'global_mean':
+            # Produce a single global mean ouput.
+            # Using a weighted mean according to the point mass/area is discretization-invariant.
+            # (A naive mean is not discretization-invariant; it could be affected by sampling a region more densely)
+            x_out = torch.sum(x * mass.unsqueeze(-1), dim=-2) / torch.sum(mass, dim=-1, keepdim=True)
+
+        # Apply last nonlinearity if specified
+        if self.last_activation != None:
+            x_out = self.last_activation(x_out)
+
+        # Remove batch dim if we added it
+        if appended_batch_dim:
+            x_out = x_out.squeeze(0)
+
+        return x_out

utils/mesh.py

@@ -0,0 +1,214 @@
+import argparse
+from pathlib import Path
+import os 
+from tqdm import tqdm
+import potpourri3d as pp3d
+import open3d as o3d
+import scipy.io as sio
+import numpy as np
+import re
+import sys
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+from shape_data import get_data_dirs
+
+# List of file extensions to consider as "mesh" files.
+# Kudos to chatgpt!
+# Add or remove extensions here as needed.
+MESH_EXTENSIONS = {".ply", ".obj", ".off", ".stl", ".fbx", ".gltf", ".glb"}
+
+def sorted_alphanum(file_list_ordered):
+    def convert(text):
+        return int(text) if text.isdigit() else text
+    def alphanum_key(key):
+        return [convert(c) for c in re.split('([0-9]+)', str(key)) if len(c) > 0]
+    return sorted(file_list_ordered, key=alphanum_key)
+
+
+def list_files(folder_path, name_filter, alphanum_sort=False):
+    file_list = [p.name for p in list(Path(folder_path).glob(name_filter))]
+    if alphanum_sort:
+        return sorted_alphanum(file_list)
+    else:
+        return sorted(file_list)
+
+def find_mesh_files(directory: Path, extensions: set=MESH_EXTENSIONS, alphanum_sort=False) -> list[Path]:
+    """
+    Recursively find all files in 'directory' whose suffix (lowercased) is in 'extensions'.
+    Returns a list of Path objects.
+    """
+    matches = []
+    for path in directory.rglob("*"):
+        if path.is_file() and path.suffix.lower() in extensions:
+            matches.append(path)
+    if alphanum_sort:
+        return sorted_alphanum(matches)
+    else:
+        return sorted(matches)
+
+def save_ply(file_name, V, F, Rho=None, color=None):
+    """Save mesh information either as an ASCII ply file.
+    https://github.com/emmanuel-hartman/BaRe-ESA/blob/main/utils/input_output.py
+    Input:
+        - file_name: specified path for saving mesh [string]
+        - V: vertices of the triangulated surface [nVx3 numpy ndarray]
+        - F: faces of the triangulated surface [nFx3 numpy ndarray]
+        - Rho: weights defined on the vertices of the triangulated surface [nVx1 numpy ndarray, default=None]
+        - color: colormap [nVx3 numpy ndarray of RGB triples]
+
+    Output:
+        - file_name.mat or file_name.ply file containing mesh information
+    """
+
+    # Save as .ply file
+    nV = V.shape[0]
+    nF = F.shape[0]   
+    if not ".ply" in file_name:
+        file_name += ".ply"
+    file = open(file_name, "w")
+    lines = ("ply","\n","format ascii 1.0","\n", "element vertex {}".format(nV),"\n","property float x","\n","property float y","\n","property float z","\n")
+    
+    if color is not None:
+        lines += ("property uchar red","\n","property uchar green","\n","property uchar blue","\n")
+        if Rho is not None:
+            lines += ("property uchar alpha","\n")
+    
+    lines += ("element face {}".format(nF),"\n","property list uchar int vertex_index","\n","end_header","\n")
+
+    file.writelines(lines)
+    lines = []
+    for i in range(0,nV):
+        for j in range(0,3):
+            lines.append(str(V[i][j]))
+            lines.append(" ")
+        if color is not None:
+            for j in range(0,3):
+                lines.append(str(color[i][j]))
+                lines.append(" ")
+            if Rho is not None:
+                lines.append(str(Rho[i]))
+                lines.append(" ")
+                    
+        lines.append("\n")
+    for i in range(0,nF):
+        l = len(F[i,:])
+        lines.append(str(l))
+        lines.append(" ")
+
+        for j in range(0,l):
+            lines.append(str(F[i,j]))
+            lines.append(" ")
+        lines.append("\n")
+
+    file.writelines(lines)
+    file.close()
+
+def numpy_to_open3d_mesh(V, F):
+    # Create an empty TriangleMesh object
+    mesh = o3d.geometry.TriangleMesh()
+    # Set vertices
+    mesh.vertices = o3d.utility.Vector3dVector(V)
+    # Set triangles
+    mesh.triangles = o3d.utility.Vector3iVector(F)
+    return mesh
+
+
+
+def load_mesh(filepath, scale=True, return_vnormals=False):
+    V, F = pp3d.read_mesh(filepath)
+    mesh = numpy_to_open3d_mesh(V, F)
+
+    tmat = np.identity(4, dtype=np.float32)
+    center = mesh.get_center()
+    tmat[:3, 3] = -center
+    area = mesh.get_surface_area()
+    if scale:
+        smat = np.identity(4, dtype=np.float32)
+        
+        smat[:3, :3] = np.identity(3, dtype=np.float32) / np.sqrt(area)
+        tmat = smat @ tmat
+    mesh.transform(tmat)
+
+    vertices = np.asarray(mesh.vertices, dtype=np.float32)
+    faces = np.asarray(mesh.triangles, dtype=np.int32)
+    if return_vnormals:
+        mesh.compute_vertex_normals()
+        vnormals = np.asarray(mesh.vertex_normals, dtype=np.float32)
+        if scale:
+            return vertices, faces, vnormals, area, center
+        return vertices, faces, vnormals, area, center
+    else:
+        return vertices, faces
+
+
+def mesh_geod_matrix(vertices, faces, do_tqdm=False, verbose=False):
+    if verbose:
+        print("Setting Geodesic matrix bw vertices")
+    n_vertices = vertices.shape[0]
+    distmat = np.zeros((n_vertices, n_vertices))
+    solver = pp3d.MeshHeatMethodDistanceSolver(vertices, faces)
+    if do_tqdm:
+        iterable = tqdm(range(n_vertices))
+    else:
+        iterable = range(n_vertices)
+    for vertind in iterable:
+        distmat[vertind] = np.maximum(solver.compute_distance(vertind), 0)
+    geod_mat = distmat
+    return geod_mat
+
+
+def prepare_geod_mats(shapes_folder, out, basename=None):
+    if basename is None:
+        basename = os.path.basename(os.path.dirname(shapes_folder)) #+ "_" + os.path.basename(shapes_folder)
+    case = basename
+    case_folder_out = os.path.join(out, case)
+    os.makedirs(case_folder_out, exist_ok=True)
+    all_shapes = [f for f in os.listdir(shapes_folder) if (".ply" in f) or (".off" in f) or ('.obj' in f)]
+    for shape in tqdm(all_shapes, "Processing " + os.path.basename(shapes_folder)):
+        vertices, faces = pp3d.read_mesh(os.path.join(shapes_folder, shape))
+        areas = pp3d.face_areas(vertices, faces)
+        mat = mesh_geod_matrix(vertices, faces, verbose=False)
+        dict_save = {
+            'geod_dist': mat,
+            'areas_f': areas
+        }
+        sio.savemat(os.path.join(case_folder_out, shape[:-4]+'.mat'), dict_save)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="What to do.")
+    parser.add_argument('--make_geods', required=True, type=int, default=0, help='launch computation of geod matrices')
+    parser.add_argument('--data', required=False, type=str, default=None)
+    parser.add_argument('--datadir', type=str, default="data", help='path where datasets are store')
+    parser.add_argument('--basename', required=False, type=str, default=None)
+    args = parser.parse_args()
+    if args.make_geods:
+        # from config import get_geod_path, get_dataset_path, get_template_path
+        # output = get_geod_path()
+        output = os.path.join(args.datadir, "geomats")
+        if args.data == "humans":
+            all_folders = [get_data_dirs(args.datadir, "faust", 'test')[0], get_data_dirs(args.datadir, "scape", 'test')[0], get_data_dirs(args.datadir, "shrec19", 'test')[0]]
+            for folder in all_folders:
+                prepare_geod_mats(folder, output)
+        elif args.data == "dt4d":
+            data_dir, _, corr_dir = get_data_dirs(args.datadir, args.data, 'test')
+            all_folders = sorted([f for f in os.listdir(data_dir) if "cross" not in f])
+            for folder in all_folders:
+                prepare_geod_mats(os.path.join(data_dir, folder), os.path.join(output, "DT4D"), basename=folder)
+        elif args.data is not None:
+            data_dir, _, corr_dir = get_data_dirs(args.datadir, args.data, 'test')
+            prepare_geod_mats(data_dir, output, args.basename)
+    # parser = argparse.ArgumentParser(description="Find all mesh files in a folder and list their paths.")
+    # parser.add_argument("--folder", type=Path, help="Path to the folder to search (will search recursively).")
+    # args = parser.parse_args()
+
+    # search_folder = args.folder
+    # if not search_folder.is_dir():
+    #     print(f"Error: '{search_folder}' is not a valid directory.")
+    # else:
+    #     # Find all matching files
+    #     mesh_files = find_mesh_files(search_folder)
+
+    #     # Sort the results for consistency
+    #     mesh_files.sort()
+    #     for p in mesh_files:
+    #         print(p.resolve())

utils/meshplot.py

@@ -0,0 +1,67 @@
+import numpy as np
+import potpourri3d as pp3d 
+import torch
+import meshplot as mp
+from ipygany import PolyMesh, Scene
+
+
+
+colors = np.array([[255, 119, 0], #orange
+                           [163, 51, 219], #violet
+                           [0, 246, 205], #bleu clair
+                           [0, 131, 246], #bleu floncé 
+                           [246, 234, 0], #jaune
+                           [143, 188, 143], #rouge
+                    [255, 0, 0]]) 
+    
+
+def double_plot_surf(surf_1, surf_2,cmap1=None,cmap2=None):
+    d = mp.subplot(surf_1.vertices, surf_1.faces, c=cmap1, s=[2, 2, 0])
+    mp.subplot(surf_2.vertices, surf_2.faces, c=cmap2, s=[2, 2, 1], data=d)
+
+def visu_pts(surf, colors=colors, idx_points=None, n_kpts=5):
+    areas = np.linalg.norm(surf.surfel, axis=-1, keepdims=True)
+    area = np.sqrt(areas.sum()/2)
+    if idx_points is None:
+        center = (surf.centers*(areas)).sum(axis=0)/areas.sum()
+        surf.updateVertices((surf.vertices - center)/area)
+        surf.cotanLaplacian()
+        idx_points = surf.get_keypoints(n_points=n_kpts)
+    solver = pp3d.MeshHeatMethodDistanceSolver(surf.vertices, surf.faces)
+    norm_center = solver.compute_distance(idx_points[-1])
+    color_array = np.zeros(surf.vertices.shape)
+    for i in range(len(idx_points)):
+        coeff = 0.1 + 1.5*norm_center[idx_points[i]]
+        i_v = idx_points[i]
+        dist = solver.compute_distance(i_v)*2
+        #color_array += np.clip(1-dist, 0, np.inf)[:, None]*colors[i][None, :]
+        color_array += np.exp(-dist**2/coeff)[:, None]*colors[i][None, :]
+    color_array = np.clip(color_array, 0, 255.)
+    return color_array
+
+def toNP(tens):
+    if isinstance(tens, torch.Tensor):
+        return tens.detach().squeeze().cpu().numpy()
+    return tens
+
+
+def overlay_surf(shape_vertices, shape_faces, target_vertices, target_faces, colors=["tomato", "darksalmon"]):
+    shape_vertices = toNP(shape_vertices)
+    target_vertices = toNP(target_vertices)
+    mesh_1 = PolyMesh(
+        vertices=shape_vertices,
+        triangle_indices=shape_faces
+    )
+    mesh_1.default_color = colors[0]
+    
+    
+    mesh_2 = PolyMesh(
+        vertices=target_vertices,
+        triangle_indices=target_faces
+    )
+    mesh_2.default_color = colors[1]
+    
+    
+    scene = Scene([mesh_1, mesh_2])
+    #scene = Scene([mesh_5])
+    return scene, [mesh_1, mesh_2]

utils/misc.py

@@ -0,0 +1,122 @@
+import os.path as osp
+import random
+import numpy as np
+import torch
+import yaml
+import omegaconf
+from omegaconf import OmegaConf
+from scipy.spatial import cKDTree
+from pathlib import Path
+
+
+class KNNSearch(object):
+    DTYPE = np.float32
+    NJOBS = 4
+
+    def __init__(self, data):
+        self.data = np.asarray(data, dtype=self.DTYPE)
+        self.kdtree = cKDTree(self.data)
+
+    def query(self, kpts, k, return_dists=False):
+        kpts = np.asarray(kpts, dtype=self.DTYPE)
+        nndists, nnindices = self.kdtree.query(kpts, k=k, workers=self.NJOBS)
+        if return_dists:
+            return nnindices, nndists
+        else:
+            return nnindices
+
+    def query_ball(self, kpt, radius):
+        kpt = np.asarray(kpt, dtype=self.DTYPE)
+        assert kpt.ndim == 1
+        nnindices = self.kdtree.query_ball_point(kpt, radius, n_jobs=self.NJOBS)
+        return nnindices
+
+
+def validate_str(x):
+    return x is not None and x != ''
+
+
+def hashing(arr, M):
+    assert isinstance(arr, np.ndarray) and arr.ndim == 2
+    N, D = arr.shape
+
+    hash_vec = np.zeros(N, dtype=np.int64)
+    for d in range(D):
+        hash_vec += arr[:, d] * M**d
+    return hash_vec
+
+
+def omegaconf_to_dotdict(hparams):
+
+    def _to_dot_dict(cfg):
+        res = {}
+        for k, v in cfg.items():
+            if v is None:
+                res[k] = v
+            elif isinstance(v, omegaconf.DictConfig):
+                res.update({k + "." + subk: subv for subk, subv in _to_dot_dict(v).items()})
+            elif isinstance(v, (str, int, float, bool)):
+                res[k] = v
+            elif isinstance(v, omegaconf.ListConfig):
+                res[k] = omegaconf.OmegaConf.to_container(v, resolve=True)
+            else:
+                raise RuntimeError('The type of {} is not supported.'.format(type(v)))
+        return res
+
+    return _to_dot_dict(hparams)
+
+
+def incrange(start, end, step):
+    assert step > 0
+    res = [start]
+    if start <= end:
+        while res[-1] + step <= end:
+            res.append(res[-1] + step)
+    else:
+        while res[-1] - step >= end:
+            res.append(res[-1] - step)
+    return res
+
+
+def seeding(seed=0):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+
+    torch.backends.cudnn.enabled = True
+    torch.backends.cudnn.benchmark = True
+    torch.backends.cudnn.deterministic = True
+
+
+def run_trainer(trainer_cls):
+    cfg_cli = OmegaConf.from_cli()
+
+    assert cfg_cli.run_mode is not None
+    if cfg_cli.run_mode == 'train':
+        assert cfg_cli.run_cfg is not None
+        cfg = OmegaConf.merge(
+            OmegaConf.load(cfg_cli.run_cfg),
+            cfg_cli,
+        )
+        OmegaConf.resolve(cfg)
+        cfg = omegaconf_to_dotdict(cfg)
+        seeding(cfg['seed'])
+        trainer = trainer_cls(cfg)
+        trainer.train()
+        trainer.test()
+    elif cfg_cli.run_mode == 'test':
+        assert cfg_cli.run_ckpt is not None
+        log_dir = str(Path(cfg_cli.run_ckpt).parent)
+        cfg = OmegaConf.merge(
+            OmegaConf.load(osp.join(log_dir, 'config.yml')),
+            cfg_cli,
+        )
+        OmegaConf.resolve(cfg)
+        cfg = omegaconf_to_dotdict(cfg)
+        cfg['test_ckpt'] = cfg_cli.run_ckpt
+        seeding(cfg['seed'])
+        trainer = trainer_cls(cfg)
+        trainer.test()
+    else:
+        raise RuntimeError(f'Mode {cfg_cli.run_mode} is not supported.')

utils/pickle_stuff.py

@@ -0,0 +1,38 @@
+import pickle
+import io
+import importlib
+import sys 
+import os 
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+
+# Mapping of old module names to new module names
+MODULE_RENAME_MAP = {
+    'module': 'diffu_models',
+    'module.model': 'diffu_models.precond',
+    'module.dit_models': 'diffu_models.dit_models',
+    'module.model': 'diffu_models.precond',
+    # add more as needed
+}
+
+class RenameUnpickler(pickle.Unpickler):
+    def find_class(self, module, name):
+        if module in MODULE_RENAME_MAP:
+            module = MODULE_RENAME_MAP[module]
+        try:
+            return super().find_class(module, name)
+        except ModuleNotFoundError as e:
+            raise ModuleNotFoundError(f"Could not find module '{module}'. You may need to update MODULE_RENAME_MAP.") from e
+
+# Usage
+def load_renamed_pickle(file_path):
+    with open(file_path, 'rb') as f:
+        return RenameUnpickler(f).load()
+
+def safe_load_with_fallback(file_path):
+    try:
+        with open(file_path, 'rb') as f:
+            return pickle.load(f)
+    except ModuleNotFoundError:
+        with open(file_path, 'rb') as f:
+            return RenameUnpickler(f).load()

utils/surfaces.py

@@ -0,0 +1,1377 @@
+import numpy as np
+import scipy as sp
+import scipy.linalg as spLA
+import scipy.sparse.linalg as sla
+import os
+import logging
+import copy
+import potpourri3d as pp3d
+from sklearn.cluster import KMeans#, DBSCAN, SpectralClustering
+from scipy.spatial import cKDTree
+import trimesh as tm
+
+
+try:
+    from vtk import *
+    import vtk.util.numpy_support as v2n
+
+    gotVTK = True
+except ImportError:
+    print('could not import VTK functions')
+    gotVTK = False
+
+
+# import kernelFunctions as kfun
+
+class vtkFields:
+    def __init__(self):
+        self.scalars = []
+        self.vectors = []
+        self.normals = []
+        self.tensors = []
+
+
+# General surface class, fibers possible. The fiber is a vector field of norm 1, defined on each vertex.
+# It must be an array of the same size (number of vertices)x3.
+
+class Surface:
+
+    # Fibers: a list of vectors, with i-th element corresponding to the value of the vector field at vertex i
+    # Contains as object:
+    # vertices : all the vertices
+    # centers: the centers of each face
+    # faces: faces along with the id of the faces
+    # surfel: surface element of each face (area*normal)
+    # List of methods:
+    # read : from filename type, call readFILENAME and set all surface attributes
+    # updateVertices: update the whole surface after a modification of the vertices
+    # computeVertexArea and Normals
+    # getEdges
+    # LocalSignedDistance distance function in the neighborhood of a shape
+    # toPolyData: convert the surface to a polydata vtk object
+    # fromPolyDate: guess what
+    # Simplify : simplify the meshes
+    # flipfaces: invert the indices of faces from [a, b, c] to [a, c, b]
+    # smooth: get a smoother surface
+    # Isosurface: compute isosurface
+    # edgeRecove: ensure that orientation is correct
+    # remove isolated: if isolated vertice, remove it
+    # laplacianMatrix: compute Laplacian Matrix of the surface graph
+    # graphLaplacianMatrix: ???
+    # laplacianSegmentation: segment the surface using laplacian properties
+    # surfVolume: compute volume inscribed in the surface (+ inscribed infinitesimal volume for each face)
+    # surfCenter: compute surface Center
+    # surfMoments: compute surface second order moments
+    # surfU: compute the informations for surface rigid alignement
+    # surfEllipsoid: compute ellipsoid representing the surface
+    # savebyu of vitk or vtk2: save the surface in a file
+    # concatenate: concatenate to another surface
+
+    def __init__(self, surf=None, filename=None, FV=None):
+        if surf == None:
+            if FV == None:
+                if filename == None:
+                    self.vertices = np.empty(0)
+                    self.centers = np.empty(0)
+                    self.faces = np.empty(0)
+                    self.surfel = np.empty(0)
+                else:
+                    if type(filename) is list:
+                        fvl = []
+                        for name in filename:
+                            fvl.append(Surface(filename=name))
+                        self.concatenate(fvl)
+                    else:
+                        self.read(filename)
+            else:
+                self.vertices = np.copy(FV[1])
+                self.faces = np.int_(FV[0])
+                self.computeCentersAreas()
+
+        else:
+            self.vertices = np.copy(surf.vertices)
+            self.faces = np.copy(surf.faces)
+            self.surfel = np.copy(surf.surfel)
+            self.centers = np.copy(surf.centers)
+            self.computeCentersAreas()
+        self.volume, self.vols = self.surfVolume()
+        self.center = self.surfCenter()
+        self.cotanLaplacian()
+
+    def read(self, filename):
+        (mainPart, ext) = os.path.splitext(filename)
+        if ext == '.byu':
+            self.readbyu(filename)
+        elif ext == '.off':
+            self.readOFF(filename)
+        elif ext == '.vtk':
+            self.readVTK(filename)
+        elif ext == '.obj':
+            self.readOBJ(filename)
+        elif ext == '.ply':
+            self.readPLY(filename)
+        elif ext == '.tri' or ext == ".ntri":
+            self.readTRI(filename)
+        else:
+            print('Unknown Surface Extension:', ext)
+            self.vertices = np.empty(0)
+            self.centers = np.empty(0)
+            self.faces = np.empty(0)
+            self.surfel = np.empty(0)
+
+    # face centers and area weighted normal
+    def computeCentersAreas(self):
+        xDef1 = self.vertices[self.faces[:, 0], :]
+        xDef2 = self.vertices[self.faces[:, 1], :]
+        xDef3 = self.vertices[self.faces[:, 2], :]
+        self.centers = (xDef1 + xDef2 + xDef3) / 3
+        self.surfel = np.cross(xDef2 - xDef1, xDef3 - xDef1)
+
+    # modify vertices without toplogical change
+    def updateVertices(self, x0):
+        self.vertices = np.copy(x0)
+        xDef1 = self.vertices[self.faces[:, 0], :]
+        xDef2 = self.vertices[self.faces[:, 1], :]
+        xDef3 = self.vertices[self.faces[:, 2], :]
+        self.centers = (xDef1 + xDef2 + xDef3) / 3
+        self.surfel = np.cross(xDef2 - xDef1, xDef3 - xDef1)
+        self.volume, self.vols = self.surfVolume()
+
+    def computeVertexArea(self):
+        # compute areas of faces and vertices
+        V = self.vertices
+        F = self.faces
+        nv = V.shape[0]
+        nf = F.shape[0]
+        AF = np.zeros([nf, 1])
+        AV = np.zeros([nv, 1])
+        for k in range(nf):
+            # determining if face is obtuse
+            x12 = V[F[k, 1], :] - V[F[k, 0], :]
+            x13 = V[F[k, 2], :] - V[F[k, 0], :]
+            n12 = np.sqrt((x12 ** 2).sum())
+            n13 = np.sqrt((x13 ** 2).sum())
+            c1 = (x12 * x13).sum() / (n12 * n13)
+            x23 = V[F[k, 2], :] - V[F[k, 1], :]
+            n23 = np.sqrt((x23 ** 2).sum())
+            # n23 = norm(x23) ;
+            c2 = -(x12 * x23).sum() / (n12 * n23)
+            c3 = (x13 * x23).sum() / (n13 * n23)
+            AF[k] = np.sqrt((np.cross(x12, x13) ** 2).sum()) / 2
+            if (c1 < 0):
+                # face obtuse at vertex 1
+                AV[F[k, 0]] += AF[k] / 2
+                AV[F[k, 1]] += AF[k] / 4
+                AV[F[k, 2]] += AF[k] / 4
+            elif (c2 < 0):
+                # face obuse at vertex 2
+                AV[F[k, 0]] += AF[k] / 4
+                AV[F[k, 1]] += AF[k] / 2
+                AV[F[k, 2]] += AF[k] / 4
+            elif (c3 < 0):
+                # face obtuse at vertex 3
+                AV[F[k, 0]] += AF[k] / 4
+                AV[F[k, 1]] += AF[k] / 4
+                AV[F[k, 2]] += AF[k] / 2
+            else:
+                # non obtuse face
+                cot1 = c1 / np.sqrt(1 - c1 ** 2)
+                cot2 = c2 / np.sqrt(1 - c2 ** 2)
+                cot3 = c3 / np.sqrt(1 - c3 ** 2)
+                AV[F[k, 0]] += ((x12 ** 2).sum() * cot3 + (x13 ** 2).sum() * cot2) / 8
+                AV[F[k, 1]] += ((x12 ** 2).sum() * cot3 + (x23 ** 2).sum() * cot1) / 8
+                AV[F[k, 2]] += ((x13 ** 2).sum() * cot2 + (x23 ** 2).sum() * cot1) / 8
+
+        for k in range(nv):
+            if (np.fabs(AV[k]) < 1e-10):
+                print('Warning: vertex ', k, 'has no face; use removeIsolated')
+        # print('sum check area:', AF.sum(), AV.sum()
+        return AV, AF
+
+    def computeVertexNormals(self):
+        self.computeCentersAreas()
+        normals = np.zeros(self.vertices.shape)
+        F = self.faces
+        for k in range(F.shape[0]):
+            normals[F[k, 0]] += self.surfel[k]
+            normals[F[k, 1]] += self.surfel[k]
+            normals[F[k, 2]] += self.surfel[k]
+        af = np.sqrt((normals ** 2).sum(axis=1))
+        # logging.info('min area = %.4f'%(af.min()))
+        normals /= af.reshape([self.vertices.shape[0], 1])
+
+        return normals
+
+    # Computes edges from vertices/faces
+    def getEdges(self):
+        self.edges = []
+
+        for k in range(self.faces.shape[0]):
+            for kj in (0, 1, 2):
+                u = [self.faces[k, kj], self.faces[k, (kj + 1) % 3]]
+
+                if (u not in self.edges) & (u.reverse() not in self.edges):
+                    self.edges.append(u)
+
+        self.edgeFaces = []
+
+        for u in self.edges:
+            self.edgeFaces.append([])
+
+        for k in range(self.faces.shape[0]):
+            for kj in (0, 1, 2):
+                u = [self.faces[k, kj], self.faces[k, (kj + 1) % 3]]
+
+                if u in self.edges:
+                    kk = self.edges.index(u)
+                else:
+                    u.reverse()
+                    kk = self.edges.index(u)
+
+                self.edgeFaces[kk].append(k)
+
+        self.edges = np.int_(np.array(self.edges))
+        self.bdry = np.int_(np.zeros(self.edges.shape[0]))
+
+        for k in range(self.edges.shape[0]):
+            if len(self.edgeFaces[k]) < 2:
+                self.bdry[k] = 1
+
+    # computes the signed distance function in a small neighborhood of a shape
+    def LocalSignedDistance(self, data, value):
+        d2 = 2 * np.array(data >= value) - 1
+        c2 = np.cumsum(d2, axis=0)
+        for j in range(2):
+            c2 = np.cumsum(c2, axis=j + 1)
+        (n0, n1, n2) = c2.shape
+
+        rad = 3
+        diam = 2 * rad + 1
+        (x, y, z) = np.mgrid[-rad:rad + 1, -rad:rad + 1, -rad:rad + 1]
+        cube = (x ** 2 + y ** 2 + z ** 2)
+        maxval = (diam) ** 3
+        s = 3.0 * rad ** 2
+        res = d2 * s
+        u = maxval * np.ones(c2.shape)
+        u[rad + 1:n0 - rad, rad + 1:n1 - rad, rad + 1:n2 - rad] = (c2[diam:n0, diam:n1, diam:n2]
+                                                                   - c2[0:n0 - diam, diam:n1, diam:n2] - c2[diam:n0,
+                                                                                                         0:n1 - diam,
+                                                                                                         diam:n2] - c2[
+                                                                                                                    diam:n0,
+                                                                                                                    diam:n1,
+                                                                                                                    0:n2 - diam]
+                                                                   + c2[0:n0 - diam, 0:n1 - diam, diam:n2] + c2[diam:n0,
+                                                                                                             0:n1 - diam,
+                                                                                                             0:n2 - diam] + c2[
+                                                                                                                            0:n0 - diam,
+                                                                                                                            diam:n1,
+                                                                                                                            0:n2 - diam]
+                                                                   - c2[0:n0 - diam, 0:n1 - diam, 0:n2 - diam])
+
+        I = np.nonzero(np.fabs(u) < maxval)
+        # print(len(I[0]))
+
+        for k in range(len(I[0])):
+            p = np.array((I[0][k], I[1][k], I[2][k]))
+            bmin = p - rad
+            bmax = p + rad + 1
+            # print(p, bmin, bmax)
+            if (d2[p[0], p[1], p[2]] > 0):
+                # print(u[p[0],p[1], p[2]])
+                # print(d2[bmin[0]:bmax[0], bmin[1]:bmax[1], bmin[2]:bmax[2]].sum())
+                res[p[0], p[1], p[2]] = min(
+                    cube[np.nonzero(d2[bmin[0]:bmax[0], bmin[1]:bmax[1], bmin[2]:bmax[2]] < 0)]) - .25
+            else:
+                res[p[0], p[1], p[2]] = - min(
+                    cube[np.nonzero(d2[bmin[0]:bmax[0], bmin[1]:bmax[1], bmin[2]:bmax[2]] > 0)]) - .25
+
+        return res
+
+    def toPolyData(self):
+        if gotVTK:
+            points = vtkPoints()
+            for k in range(self.vertices.shape[0]):
+                points.InsertNextPoint(self.vertices[k, 0], self.vertices[k, 1], self.vertices[k, 2])
+            polys = vtkCellArray()
+            for k in range(self.faces.shape[0]):
+                polys.InsertNextCell(3)
+                for kk in range(3):
+                    polys.InsertCellPoint(self.faces[k, kk])
+            polydata = vtkPolyData()
+            polydata.SetPoints(points)
+            polydata.SetPolys(polys)
+            return polydata
+        else:
+            raise Exception('Cannot run toPolyData without VTK')
+
+    def fromPolyData(self, g, scales=[1., 1., 1.]):
+        npoints = int(g.GetNumberOfPoints())
+        nfaces = int(g.GetNumberOfPolys())
+        logging.info('Dimensions: %d %d %d' % (npoints, nfaces, g.GetNumberOfCells()))
+        V = np.zeros([npoints, 3])
+        for kk in range(npoints):
+            V[kk, :] = np.array(g.GetPoint(kk))
+            # print(kk, V[kk])
+            # print(kk, np.array(g.GetPoint(kk)))
+        F = np.zeros([nfaces, 3])
+        gf = 0
+        for kk in range(g.GetNumberOfCells()):
+            c = g.GetCell(kk)
+            if (c.GetNumberOfPoints() == 3):
+                for ll in range(3):
+                    F[gf, ll] = c.GetPointId(ll)
+                    # print(kk, gf, F[gf])
+                gf += 1
+
+                # self.vertices = np.multiply(data.shape-V-1, scales)
+        self.vertices = np.multiply(V, scales)
+        self.faces = np.int_(F[0:gf, :])
+        self.computeCentersAreas()
+
+    def Simplify(self, target=1000.0):
+        if gotVTK:
+            polydata = self.toPolyData()
+            dc = vtkQuadricDecimation()
+            red = 1 - min(np.float(target) / polydata.GetNumberOfPoints(), 1)
+            dc.SetTargetReduction(red)
+            dc.SetInput(polydata)
+            dc.Update()
+            g = dc.GetOutput()
+            self.fromPolyData(g)
+            z = self.surfVolume()
+            if (z > 0):
+                self.flipFaces()
+                print('flipping volume', z, self.surfVolume())
+        else:
+            raise Exception('Cannot run Simplify without VTK')
+
+    def flipFaces(self):
+        self.faces = self.faces[:, [0, 2, 1]]
+        self.computeCentersAreas()
+
+    def smooth(self, n=30, smooth=0.1):
+        if gotVTK:
+            g = self.toPolyData()
+            print(g)
+            smoother = vtkWindowedSincPolyDataFilter()
+            smoother.SetInput(g)
+            smoother.SetNumberOfIterations(n)
+            smoother.SetPassBand(smooth)
+            smoother.NonManifoldSmoothingOn()
+            smoother.NormalizeCoordinatesOn()
+            smoother.GenerateErrorScalarsOn()
+            # smoother.GenerateErrorVectorsOn()
+            smoother.Update()
+            g = smoother.GetOutput()
+            self.fromPolyData(g)
+        else:
+            raise Exception('Cannot run smooth without VTK')
+
+    # Computes isosurfaces using vtk
+    def Isosurface(self, data, value=0.5, target=1000.0, scales=[1., 1., 1.], smooth=0.1, fill_holes=1.):
+        if gotVTK:
+            # data = self.LocalSignedDistance(data0, value)
+            if isinstance(data, vtkImageData):
+                img = data
+            else:
+                img = vtkImageData()
+                img.SetDimensions(data.shape)
+                img.SetOrigin(0, 0, 0)
+                if vtkVersion.GetVTKMajorVersion() >= 6:
+                    img.AllocateScalars(VTK_FLOAT, 1)
+                else:
+                    img.SetNumberOfScalarComponents(1)
+                v = vtkDoubleArray()
+                v.SetNumberOfValues(data.size)
+                v.SetNumberOfComponents(1)
+                for ii, tmp in enumerate(np.ravel(data, order='F')):
+                    v.SetValue(ii, tmp)
+                    img.GetPointData().SetScalars(v)
+
+            cf = vtkContourFilter()
+            if vtkVersion.GetVTKMajorVersion() >= 6:
+                cf.SetInputData(img)
+            else:
+                cf.SetInput(img)
+            cf.SetValue(0, value)
+            cf.SetNumberOfContours(1)
+            cf.Update()
+            # print(cf
+            connectivity = vtkPolyDataConnectivityFilter()
+            connectivity.ScalarConnectivityOff()
+            connectivity.SetExtractionModeToLargestRegion()
+            if vtkVersion.GetVTKMajorVersion() >= 6:
+                connectivity.SetInputData(cf.GetOutput())
+            else:
+                connectivity.SetInput(cf.GetOutput())
+            connectivity.Update()
+            g = connectivity.GetOutput()
+
+            if smooth > 0:
+                smoother = vtkWindowedSincPolyDataFilter()
+                if vtkVersion.GetVTKMajorVersion() >= 6:
+                    smoother.SetInputData(g)
+                else:
+                    smoother.SetInput(g)
+                #     else:
+                # smoother.SetInputConnection(contour.GetOutputPort())
+                smoother.SetNumberOfIterations(30)
+                # this has little effect on the error!
+                # smoother.BoundarySmoothingOff()
+                # smoother.FeatureEdgeSmoothingOff()
+                # smoother.SetFeatureAngle(120.0)
+                smoother.SetPassBand(smooth)  # this increases the error a lot!
+                smoother.NonManifoldSmoothingOn()
+                # smoother.NormalizeCoordinatesOn()
+                # smoother.GenerateErrorScalarsOn()
+                # smoother.GenerateErrorVectorsOn()
+                smoother.Update()
+                g = smoother.GetOutput()
+
+            # dc = vtkDecimatePro()
+            red = 1 - min(np.float(target) / g.GetNumberOfPoints(), 1)
+            # print('Reduction: ', red)
+            dc = vtkQuadricDecimation()
+            dc.SetTargetReduction(red)
+            # dc.AttributeErrorMetricOn()
+            # dc.SetDegree(10)
+            # dc.SetSplitting(0)
+            if vtkVersion.GetVTKMajorVersion() >= 6:
+                dc.SetInputData(g)
+            else:
+                dc.SetInput(g)
+                # dc.SetInput(g)
+            # print(dc)
+            dc.Update()
+            g = dc.GetOutput()
+            # print('points:', g.GetNumberOfPoints())
+            cp = vtkCleanPolyData()
+            if vtkVersion.GetVTKMajorVersion() >= 6:
+                cp.SetInputData(dc.GetOutput())
+            else:
+                cp.SetInput(dc.GetOutput())
+                #        cp.SetInput(dc.GetOutput())
+            # cp.SetPointMerging(1)
+            cp.ConvertPolysToLinesOn()
+            cp.SetAbsoluteTolerance(1e-5)
+            cp.Update()
+            g = cp.GetOutput()
+            self.fromPolyData(g, scales)
+            z = self.surfVolume()
+            if (z > 0):
+                self.flipFaces()
+                # print('flipping volume', z, self.surfVolume())
+                logging.info('flipping volume %.2f %.2f' % (z, self.surfVolume()))
+
+            # print(g)
+            # npoints = int(g.GetNumberOfPoints())
+            # nfaces = int(g.GetNumberOfPolys())
+            # print('Dimensions:', npoints, nfaces, g.GetNumberOfCells())
+            # V = np.zeros([npoints, 3])
+            # for kk in range(npoints):
+            #     V[kk, :] = np.array(g.GetPoint(kk))
+            #     #print(kk, V[kk])
+            #     #print(kk, np.array(g.GetPoint(kk)))
+            # F = np.zeros([nfaces, 3])
+            # gf = 0
+            # for kk in range(g.GetNumberOfCells()):
+            #     c = g.GetCell(kk)
+            #     if(c.GetNumberOfPoints() == 3):
+            #         for ll in range(3):
+            #             F[gf,ll] = c.GetPointId(ll)
+            #             #print(kk, gf, F[gf])
+            #         gf += 1
+
+            #         #self.vertices = np.multiply(data.shape-V-1, scales)
+            # self.vertices = np.multiply(V, scales)
+            # self.faces = np.int_(F[0:gf, :])
+            # self.computeCentersAreas()
+        else:
+            raise Exception('Cannot run Isosurface without VTK')
+
+    # Ensures that orientation is correct
+    def edgeRecover(self):
+        v = self.vertices
+        f = self.faces
+        nv = v.shape[0]
+        nf = f.shape[0]
+        # faces containing each oriented edge
+        edg0 = np.int_(np.zeros((nv, nv)))
+        # number of edges between each vertex
+        edg = np.int_(np.zeros((nv, nv)))
+        # contiguous faces
+        edgF = np.int_(np.zeros((nf, nf)))
+        for (kf, c) in enumerate(f):
+            if (edg0[c[0], c[1]] > 0):
+                edg0[c[1], c[0]] = kf + 1
+            else:
+                edg0[c[0], c[1]] = kf + 1
+
+            if (edg0[c[1], c[2]] > 0):
+                edg0[c[2], c[1]] = kf + 1
+            else:
+                edg0[c[1], c[2]] = kf + 1
+
+            if (edg0[c[2], c[0]] > 0):
+                edg0[c[0], c[2]] = kf + 1
+            else:
+                edg0[c[2], c[0]] = kf + 1
+
+            edg[c[0], c[1]] += 1
+            edg[c[1], c[2]] += 1
+            edg[c[2], c[0]] += 1
+
+        for kv in range(nv):
+            I2 = np.nonzero(edg0[kv, :])
+            for kkv in I2[0].tolist():
+                edgF[edg0[kkv, kv] - 1, edg0[kv, kkv] - 1] = kv + 1
+
+        isOriented = np.int_(np.zeros(f.shape[0]))
+        isActive = np.int_(np.zeros(f.shape[0]))
+        I = np.nonzero(np.squeeze(edgF[0, :]))
+        # list of faces to be oriented
+        activeList = [0] + I[0].tolist()
+        lastOriented = 0
+        isOriented[0] = True
+        for k in activeList:
+            isActive[k] = True
+
+        while lastOriented < len(activeList) - 1:
+            i = activeList[lastOriented]
+            j = activeList[lastOriented + 1]
+            I = np.nonzero(edgF[j, :])
+            foundOne = False
+            for kk in I[0].tolist():
+                if (foundOne == False) & (isOriented[kk]):
+                    foundOne = True
+                    u1 = edgF[j, kk] - 1
+                    u2 = edgF[kk, j] - 1
+                    if not ((edg[u1, u2] == 1) & (edg[u2, u1] == 1)):
+                        # reorient face j
+                        edg[f[j, 0], f[j, 1]] -= 1
+                        edg[f[j, 1], f[j, 2]] -= 1
+                        edg[f[j, 2], f[j, 0]] -= 1
+                        a = f[j, 1]
+                        f[j, 1] = f[j, 2]
+                        f[j, 2] = a
+                        edg[f[j, 0], f[j, 1]] += 1
+                        edg[f[j, 1], f[j, 2]] += 1
+                        edg[f[j, 2], f[j, 0]] += 1
+                elif (not isActive[kk]):
+                    activeList.append(kk)
+                    isActive[kk] = True
+            if foundOne:
+                lastOriented = lastOriented + 1
+                isOriented[j] = True
+                # print('oriented face', j, lastOriented,  'out of',  nf,  ';  total active', len(activeList))
+            else:
+                print('Unable to orient face', j)
+                return
+        self.vertices = v;
+        self.faces = f;
+
+        z, _ = self.surfVolume()
+        if (z > 0):
+            self.flipFaces()
+
+    def removeIsolated(self):
+        N = self.vertices.shape[0]
+        inFace = np.int_(np.zeros(N))
+        for k in range(3):
+            inFace[self.faces[:, k]] = 1
+        J = np.nonzero(inFace)
+        self.vertices = self.vertices[J[0], :]
+        logging.info('Found %d isolated vertices' % (J[0].shape[0]))
+        Q = -np.ones(N)
+        for k, j in enumerate(J[0]):
+            Q[j] = k
+        self.faces = np.int_(Q[self.faces])
+
+    def laplacianMatrix(self):
+        F = self.faces
+        V = self.vertices;
+        nf = F.shape[0]
+        nv = V.shape[0]
+
+        AV, AF = self.computeVertexArea()
+
+        # compute edges and detect boundary
+        # edm = sp.lil_matrix((nv,nv))
+        edm = -np.ones([nv, nv]).astype(np.int32)
+        E = np.zeros([3 * nf, 2]).astype(np.int32)
+        j = 0
+        for k in range(nf):
+            if (edm[F[k, 0], F[k, 1]] == -1):
+                edm[F[k, 0], F[k, 1]] = j
+                edm[F[k, 1], F[k, 0]] = j
+                E[j, :] = [F[k, 0], F[k, 1]]
+                j = j + 1
+            if (edm[F[k, 1], F[k, 2]] == -1):
+                edm[F[k, 1], F[k, 2]] = j
+                edm[F[k, 2], F[k, 1]] = j
+                E[j, :] = [F[k, 1], F[k, 2]]
+                j = j + 1
+            if (edm[F[k, 0], F[k, 2]] == -1):
+                edm[F[k, 2], F[k, 0]] = j
+                edm[F[k, 0], F[k, 2]] = j
+                E[j, :] = [F[k, 2], F[k, 0]]
+                j = j + 1
+        E = E[0:j, :]
+
+        edgeFace = np.zeros([j, nf])
+        ne = j
+        # print(E)
+        for k in range(nf):
+            edgeFace[edm[F[k, 0], F[k, 1]], k] = 1
+            edgeFace[edm[F[k, 1], F[k, 2]], k] = 1
+            edgeFace[edm[F[k, 2], F[k, 0]], k] = 1
+
+        bEdge = np.zeros([ne, 1])
+        bVert = np.zeros([nv, 1])
+        edgeAngles = np.zeros([ne, 2])
+        for k in range(ne):
+            I = np.flatnonzero(edgeFace[k, :])
+            # print('I=', I, F[I, :], E.shape)
+            # print('E[k, :]=', k, E[k, :])
+            # print(k, edgeFace[k, :])
+            for u in range(len(I)):
+                f = I[u]
+                i1l = np.flatnonzero(F[f, :] == E[k, 0])
+                i2l = np.flatnonzero(F[f, :] == E[k, 1])
+                # print(f, F[f, :])
+                # print(i1l, i2l)
+                i1 = i1l[0]
+                i2 = i2l[0]
+                s = i1 + i2
+                if s == 1:
+                    i3 = 2
+                elif s == 2:
+                    i3 = 1
+                elif s == 3:
+                    i3 = 0
+                x1 = V[F[f, i1], :] - V[F[f, i3], :]
+                x2 = V[F[f, i2], :] - V[F[f, i3], :]
+                a = (np.cross(x1, x2) * np.cross(V[F[f, 1], :] - V[F[f, 0], :], V[F[f, 2], :] - V[F[f, 0], :])).sum()
+                b = (x1 * x2).sum()
+                if (a > 0):
+                    edgeAngles[k, u] = b / np.sqrt(a)
+                else:
+                    edgeAngles[k, u] = b / np.sqrt(-a)
+            if (len(I) == 1):
+                # boundary edge
+                bEdge[k] = 1
+                bVert[E[k, 0]] = 1
+                bVert[E[k, 1]] = 1
+                edgeAngles[k, 1] = 0
+
+                # Compute Laplacian matrix
+        L = np.zeros([nv, nv])
+
+        for k in range(ne):
+            L[E[k, 0], E[k, 1]] = (edgeAngles[k, 0] + edgeAngles[k, 1]) / 2
+            L[E[k, 1], E[k, 0]] = L[E[k, 0], E[k, 1]]
+
+        for k in range(nv):
+            L[k, k] = - L[k, :].sum()
+
+        A = np.zeros([nv, nv])
+        for k in range(nv):
+            A[k, k] = AV[k]
+
+        return L, A
+
+    def graphLaplacianMatrix(self):
+        F = self.faces
+        V = self.vertices
+        nf = F.shape[0]
+        nv = V.shape[0]
+
+        # compute edges and detect boundary
+        # edm = sp.lil_matrix((nv,nv))
+        edm = -np.ones([nv, nv])
+        E = np.zeros([3 * nf, 2])
+        j = 0
+        for k in range(nf):
+            if (edm[F[k, 0], F[k, 1]] == -1):
+                edm[F[k, 0], F[k, 1]] = j
+                edm[F[k, 1], F[k, 0]] = j
+                E[j, :] = [F[k, 0], F[k, 1]]
+                j = j + 1
+            if (edm[F[k, 1], F[k, 2]] == -1):
+                edm[F[k, 1], F[k, 2]] = j
+                edm[F[k, 2], F[k, 1]] = j
+                E[j, :] = [F[k, 1], F[k, 2]]
+                j = j + 1
+            if (edm[F[k, 0], F[k, 2]] == -1):
+                edm[F[k, 2], F[k, 0]] = j
+                edm[F[k, 0], F[k, 2]] = j
+                E[j, :] = [F[k, 2], F[k, 0]]
+                j = j + 1
+        E = E[0:j, :]
+
+        edgeFace = np.zeros([j, nf])
+        ne = j
+        # print(E)
+
+        # Compute Laplacian matrix
+        L = np.zeros([nv, nv])
+
+        for k in range(ne):
+            L[E[k, 0], E[k, 1]] = 1
+            L[E[k, 1], E[k, 0]] = 1
+
+        for k in range(nv):
+            L[k, k] = - L[k, :].sum()
+
+        return L
+
+    def cotanLaplacian(self, eps=1e-8):
+        L = pp3d.cotan_laplacian(self.vertices, self.faces, denom_eps=1e-10)
+        massvec_np = pp3d.vertex_areas(self.vertices, self.faces)
+        massvec_np += eps * np.mean(massvec_np)
+
+        if (np.isnan(L.data).any()):
+            raise RuntimeError("NaN Laplace matrix")
+        if (np.isnan(massvec_np).any()):
+            raise RuntimeError("NaN mass matrix")
+        self.L = L 
+        self.massvec_np = massvec_np
+        return L, massvec_np
+
+    def lapEigen(self, n_eig=10, eps=1e-8):
+        L_eigsh = (self.L + sp.sparse.identity(self.L.shape[0]) * eps).tocsc()
+        massvec_eigsh = self.massvec_np
+        Mmat = sp.sparse.diags(massvec_eigsh)
+        eigs_sigma = eps#-0.01
+        evals, evecs = sla.eigsh(L_eigsh, k=n_eig, M=Mmat, sigma=eigs_sigma)
+        return evals, evecs
+
+    def laplacianSegmentation(self, k, verbose=False):
+        # (L, AA) = self.laplacianMatrix()
+        # # print((L.shape[0]-k-1, L.shape[0]-2))
+        # (D, y) = spLA.eigh(L, AA, eigvals=(L.shape[0] - k, L.shape[0] - 1))
+        evals, evecs = self.lapEigen(k*2)
+        y = evecs[:, 1:] * np.exp(-2*0.1 * evals[1:]) #**2
+        # N = y.shape[0]
+        # d = y.shape[1]
+        # I = np.argsort(y.sum(axis=1))
+        # I0 = np.floor((N - 1) * sp.linspace(0, 1, num=k)).astype(int)
+        # # print(y.shape, L.shape, N, k, d)
+        # C = y[I0, :].copy()
+        #
+        # eps = 1e-20
+        # Cold = C.copy()
+        # u = ((C.reshape([k, 1, d]) - y.reshape([1, N, d])) ** 2).sum(axis=2)
+        # T = u.min(axis=0).sum() / (N)
+        # # print(T)
+        # j = 0
+        # while j < 5000:
+        #     u0 = u - u.min(axis=0).reshape([1, N])
+        #     w = np.exp(-u0 / T);
+        #     w = w / (eps + w.sum(axis=0).reshape([1, N]))
+        #     # print(w.min(), w.max())
+        #     cost = (u * w).sum() + T * (w * np.log(w + eps)).sum()
+        #     C = np.dot(w, y) / (eps + w.sum(axis=1).reshape([k, 1]))
+        #     # print(j, 'cost0 ', cost)
+        #
+        #     u = ((C.reshape([k, 1, d]) - y.reshape([1, N, d])) ** 2).sum(axis=2)
+        #     cost = (u * w).sum() + T * (w * np.log(w + eps)).sum()
+        #     err = np.sqrt(((C - Cold) ** 2).sum(axis=1)).sum()
+        #     # print(j, 'cost ', cost, err, T)
+        #     if (j > 100) & (err < 1e-4):
+        #         break
+        #     j = j + 1
+        #     Cold = C.copy()
+        #     T = T * 0.99
+        #
+        #     # print(k, d, C.shape)
+        # dst = ((C.reshape([k, 1, d]) - y.reshape([1, N, d])) ** 2).sum(axis=2)
+        # md = dst.min(axis=0)
+        # idx = np.zeros(N).astype(int)
+        # for j in range(N):
+        #     I = np.flatnonzero(dst[:, j] < md[j] + 1e-10)
+        #     idx[j] = I[0]
+        # I = -np.ones(k).astype(int)
+        # kk = 0
+        # for j in range(k):
+        #     if True in (idx == j):
+        #         I[j] = kk
+        #         kk += 1
+        # idx = I[idx]
+        # if idx.max() < (k - 1):
+        #     print('Warning: kmeans convergence with %d clusters instead of %d' % (idx.max(), k))
+        #     # ml = w.sum(axis=1)/N
+        kmeans = KMeans(n_clusters=k, random_state=0, n_init="auto").fit(y)
+        # kmeans = DBSCAN().fit(y)
+        idx = kmeans.labels_
+        nc = idx.max() + 1
+        C = np.zeros([nc, self.vertices.shape[1]])
+        a, foo = self.computeVertexArea()
+        for k in range(nc):
+            I = np.flatnonzero(idx == k)
+            nI = len(I)
+            # print(a.shape, nI)
+            aI = a[I]
+            ak = aI.sum()
+            C[k, :] = (self.vertices[I, :] * aI).sum(axis=0) / ak;
+        mean_eigen = (y*a).sum(axis=0)/a
+        tree = cKDTree(y)
+        _, indices = tree.query(mean_eigen, k=1)
+        index_center = indices[0]
+        _, indices = tree.query(C, k=1)
+        index_C = indices[0]
+
+        mesh = tm.Trimesh(vertices=self.vertices, faces=self.faces)
+        nv = self.computeVertexNormals()
+        rori = self.vertices[index] + 0.001 * (-nv[index, :])
+        rdir = -nv[index, :]
+
+        locations, index_ray, index_tri = mesh.ray.intersects_location(ray_origins=rori[None, :], ray_directions=rdir[None, :])
+        if verbose:
+            print(locations, index_ray, index_tri)
+            print(index_center, index_C)
+        return idx, C, index_center, locations[0]
+
+    def get_keypoints(self, n_eig=30, n_points=5, sym=False):
+        evals, evecs = self.lapEigen(n_eig+1)
+        pts_evecs = evecs[:, 1:] * np.exp(-2*0.1 * evals[1:]) # approximate geod distance
+        tree = cKDTree(pts_evecs)
+        _, center_index = tree.query(np.zeros(pts_evecs.shape[-1]), k=1)
+        indices = fps_gpt_geod(self.vertices, self.faces, 10, center_index)[1:]
+        areas = np.linalg.norm(self.surfel, axis=-1, keepdims=True)
+        area = np.sqrt(areas.sum()/2)
+        print(area)
+        ## Looking for center + provide distances to the center
+        
+        print(center_index)
+        solver = pp3d.MeshHeatMethodDistanceSolver(self.vertices, self.faces)
+        norm_center = solver.compute_distance(center_index)
+    
+        ## Dist matrix just between selected points
+        dist_mat_indices = np.zeros((len(indices), len(indices)))
+        all_ii = np.arange(len(indices)) # just to easy code
+        for ii, index in enumerate(indices):
+            dist_mat_indices[ii, ii!=all_ii] = solver.compute_distance(index)[indices[indices != index]]
+    
+        ## Select points which : farthest from the center, delete their neighbors with dist < 0.5*area
+        keep = []
+        max_norm = norm_center.max()
+        for ii, index in enumerate(indices[:n_points]):
+            # distances_compar = (dist_mat_indices[ii, ii!=all_ii] + norm_center[ii]) / norm_center[indices[ii!=all_ii]]
+            # min_dist = np.amin(np.abs(distances_compar - 1))
+            # print(min_dist/norm_center[ii], min_dist)
+            # if min_dist > 0.3:
+            #     #print(ii, min_dist/norm_center[ii], norm_center[ii], dist_mat_indices[ii, ii!=all_ii], norm_center[indices[indices != index]], indices[7])
+            keep.append(index)
+        ## Add the index of the center
+        ## Add the index of the center 
+        if sym:
+            mesh = tm.base.Trimesh(self.vertices, self.faces, process=False)  # Load your mesh
+            # Given data
+            index = 123  # Your starting point index
+            direction = -mesh.vertex_normals[center_index]
+            
+            # Get the starting vertex position
+            start_point = mesh.vertices[center_index]
+            
+            # Perform ray intersection
+            locations, index_ray, index_tri = mesh.ray.intersects_location(
+                ray_origins=start_point[None, :],  # Starting point
+                ray_directions=direction[None, :]  # Direction of travel
+            )
+            
+            # If intersections exist
+            if len(locations) > 0:
+                # Sort intersections by distance from start_point
+                distances = np.linalg.norm(locations - start_point, axis=1)
+                sorted_indices = np.argsort(distances)
+            
+                # Get the first valid intersection point beyond the start point
+                #next_intersection = locations[sorted_indices]
+                #print(f"Next intersection at: {next_intersection}")
+                intersected_triangle = index_tri[sorted_indices[1]]
+                print(f"Intersected Triangle Index: {mesh.faces[intersected_triangle]}")
+                center_index = mesh.faces[intersected_triangle][0]
+            else:
+                print("No intersection found in the given direction.")
+        keep.append(center_index)
+        return keep
+
+
+    # Computes surface volume
+    def surfVolume(self):
+        f = self.faces
+        v = self.vertices
+        t = v[f, :]
+        vols = np.linalg.det(t) / 6
+        return vols.sum(), vols
+
+    def surfCenter(self):
+        f = self.faces
+        v = self.vertices
+        center_infs = (v[f, :].sum(axis=1) / 4) * self.vols[:, np.newaxis]
+        center = center_infs.sum(axis=0)
+        return center
+
+    def surfMoments(self):
+        f = self.faces
+        v = self.vertices
+        vec_0 = v[f[:, 0], :] + v[f[:, 1], :]
+        s_0 = vec_0[:, :, np.newaxis] * vec_0[:, np.newaxis, :]
+        vec_1 = v[f[:, 0], :] + v[f[:, 2], :]
+        s_1 = vec_1[:, :, np.newaxis] * vec_1[:, np.newaxis, :]
+        vec_2 = v[f[:, 1], :] + v[f[:, 2], :]
+        s_2 = vec_2[:, :, np.newaxis] * vec_2[:, np.newaxis, :]
+        moments_inf = self.vols[:, np.newaxis, np.newaxis] * (1. / 20) * (s_0 + s_1 + s_2)
+        return moments_inf.sum(axis=0)
+
+    def surfF(self):
+        f = self.faces
+        v = self.vertices
+        cent = (v[f, :].sum(axis=1)) / 4.
+        F = self.vols[:, np.newaxis] * np.sign(cent) * (cent ** 2)
+        return F.sum(axis=0)
+
+    def surfU(self):
+        vol = self.volume
+        vertices = self.vertices / pow(vol, 1. / 3)
+        vertices -= self.center
+        self.updateVertices(vertices)
+        moments = self.surfMoments()
+        u, s, vh = np.linalg.svd(moments)
+        F = self.surfF()
+        return np.diag(np.sign(F)) @ u, s
+
+    def surfEllipsoid(self, u, s, moments):
+        coeff = pow(4 * np.pi / 15, 1. / 5) * pow(np.linalg.det(moments), -1. / 10)
+        A = coeff * ((u * np.sqrt(s)) @ u.T)
+        return u, A
+
+    # Reads from .off file
+    def readOFF(self, offfile):
+        with open(offfile, 'r') as f:
+            all_lines = f.readlines()
+    
+    
+        n_vertices = int(all_lines[1].split()[0])
+        n_faces = int(all_lines[1].split()[1])
+
+        vertices_list = []
+        for i in range(n_vertices):
+            vertices_list.append([float(x) for x in all_lines[2+i].split()[:3]])
+
+        faces_list = []
+        for i in range(n_faces):
+            # Be careful to convert to int. Otherwise, you can use np.array(faces_list).astype(np.int32)
+            faces_list.append([int(x) for x in all_lines[2+i+n_vertices].split()[1:4]])
+        self.faces = np.array(faces_list)
+        self.vertices = np.array(vertices_list)
+        self.computeCentersAreas()    # Reads from .byu file
+    def readbyu(self, byufile):
+        with open(byufile, 'r') as fbyu:
+            ln0 = fbyu.readline()
+            ln = ln0.split()
+            # read header
+            ncomponents = int(ln[0])  # number of components
+            npoints = int(ln[1])  # number of vertices
+            nfaces = int(ln[2])  # number of faces
+            # fscanf(fbyu,'%d',1);		% number of edges
+            # %ntest = fscanf(fbyu,'%d',1);		% number of edges
+            for k in range(ncomponents):
+                fbyu.readline()  # components (ignored)
+            # read data
+            self.vertices = np.empty([npoints, 3])
+            k = -1
+            while k < npoints - 1:
+                ln = fbyu.readline().split()
+                k = k + 1;
+                self.vertices[k, 0] = float(ln[0])
+                self.vertices[k, 1] = float(ln[1])
+                self.vertices[k, 2] = float(ln[2])
+                if len(ln) > 3:
+                    k = k + 1;
+                    self.vertices[k, 0] = float(ln[3])
+                    self.vertices[k, 1] = float(ln[4])
+                    self.vertices[k, 2] = float(ln[5])
+
+            self.faces = np.empty([nfaces, 3])
+            ln = fbyu.readline().split()
+            kf = 0
+            j = 0
+            while ln:
+                if kf >= nfaces:
+                    break
+                    # print(nfaces, kf, ln)
+                for s in ln:
+                    self.faces[kf, j] = int(sp.fabs(int(s)))
+                    j = j + 1
+                    if j == 3:
+                        kf = kf + 1
+                        j = 0
+                ln = fbyu.readline().split()
+        self.faces = np.int_(self.faces) - 1
+        xDef1 = self.vertices[self.faces[:, 0], :]
+        xDef2 = self.vertices[self.faces[:, 1], :]
+        xDef3 = self.vertices[self.faces[:, 2], :]
+        self.centers = (xDef1 + xDef2 + xDef3) / 3
+        self.surfel = np.cross(xDef2 - xDef1, xDef3 - xDef1)
+
+    # Saves in .byu format
+    def savebyu(self, byufile):
+        # FV = readbyu(byufile)
+        # reads from a .byu file into matlab's face vertex structure FV
+
+        with open(byufile, 'w') as fbyu:
+            # copy header
+            ncomponents = 1  # number of components
+            npoints = self.vertices.shape[0]  # number of vertices
+            nfaces = self.faces.shape[0]  # number of faces
+            nedges = 3 * nfaces  # number of edges
+
+            str = '{0: d} {1: d} {2: d} {3: d} 0\n'.format(ncomponents, npoints, nfaces, nedges)
+            fbyu.write(str)
+            str = '1 {0: d}\n'.format(nfaces)
+            fbyu.write(str)
+
+            k = -1
+            while k < (npoints - 1):
+                k = k + 1
+                str = '{0: f} {1: f} {2: f} '.format(self.vertices[k, 0], self.vertices[k, 1], self.vertices[k, 2])
+                fbyu.write(str)
+                if k < (npoints - 1):
+                    k = k + 1
+                    str = '{0: f} {1: f} {2: f}\n'.format(self.vertices[k, 0], self.vertices[k, 1], self.vertices[k, 2])
+                    fbyu.write(str)
+                else:
+                    fbyu.write('\n')
+
+            j = 0
+            for k in range(nfaces):
+                for kk in (0, 1):
+                    fbyu.write('{0: d} '.format(self.faces[k, kk] + 1))
+                    j = j + 1
+                    if j == 16:
+                        fbyu.write('\n')
+                        j = 0
+
+                fbyu.write('{0: d} '.format(-self.faces[k, 2] - 1))
+                j = j + 1
+                if j == 16:
+                    fbyu.write('\n')
+                    j = 0
+
+    def saveVTK(self, fileName, scalars=None, normals=None, tensors=None, scal_name='scalars', vectors=None,
+                vect_name='vectors'):
+        vf = vtkFields()
+        # print(scalars)
+        if not (scalars == None):
+            vf.scalars.append(scal_name)
+            vf.scalars.append(scalars)
+        if not (vectors == None):
+            vf.vectors.append(vect_name)
+            vf.vectors.append(vectors)
+        if not (normals == None):
+            vf.normals.append('normals')
+            vf.normals.append(normals)
+        if not (tensors == None):
+            vf.tensors.append('tensors')
+            vf.tensors.append(tensors)
+        self.saveVTK2(fileName, vf)
+
+    # Saves in .vtk format
+    def saveVTK2(self, fileName, vtkFields=None):
+        F = self.faces;
+        V = self.vertices;
+
+        with open(fileName, 'w') as fvtkout:
+            fvtkout.write('# vtk DataFile Version 3.0\nSurface Data\nASCII\nDATASET POLYDATA\n')
+            fvtkout.write('\nPOINTS {0: d} float'.format(V.shape[0]))
+            for ll in range(V.shape[0]):
+                fvtkout.write('\n{0: f} {1: f} {2: f}'.format(V[ll, 0], V[ll, 1], V[ll, 2]))
+            fvtkout.write('\nPOLYGONS {0:d} {1:d}'.format(F.shape[0], 4 * F.shape[0]))
+            for ll in range(F.shape[0]):
+                fvtkout.write('\n3 {0: d} {1: d} {2: d}'.format(F[ll, 0], F[ll, 1], F[ll, 2]))
+            if not (vtkFields == None):
+                wrote_pd_hdr = False
+                if len(vtkFields.scalars) > 0:
+                    if not wrote_pd_hdr:
+                        fvtkout.write(('\nPOINT_DATA {0: d}').format(V.shape[0]))
+                        wrote_pd_hdr = True
+                    nf = len(vtkFields.scalars) / 2
+                    for k in range(nf):
+                        fvtkout.write('\nSCALARS ' + vtkFields.scalars[2 * k] + ' float 1\nLOOKUP_TABLE default')
+                        for ll in range(V.shape[0]):
+                            # print(scalars[ll])
+                            fvtkout.write('\n {0: .5f}'.format(vtkFields.scalars[2 * k + 1][ll]))
+                if len(vtkFields.vectors) > 0:
+                    if not wrote_pd_hdr:
+                        fvtkout.write(('\nPOINT_DATA {0: d}').format(V.shape[0]))
+                        wrote_pd_hdr = True
+                    nf = len(vtkFields.vectors) / 2
+                    for k in range(nf):
+                        fvtkout.write('\nVECTORS ' + vtkFields.vectors[2 * k] + ' float')
+                        vectors = vtkFields.vectors[2 * k + 1]
+                        for ll in range(V.shape[0]):
+                            fvtkout.write(
+                                '\n {0: .5f} {1: .5f} {2: .5f}'.format(vectors[ll, 0], vectors[ll, 1], vectors[ll, 2]))
+                if len(vtkFields.normals) > 0:
+                    if not wrote_pd_hdr:
+                        fvtkout.write(('\nPOINT_DATA {0: d}').format(V.shape[0]))
+                        wrote_pd_hdr = True
+                    nf = len(vtkFields.normals) / 2
+                    for k in range(nf):
+                        fvtkout.write('\nNORMALS ' + vtkFields.normals[2 * k] + ' float')
+                        vectors = vtkFields.normals[2 * k + 1]
+                        for ll in range(V.shape[0]):
+                            fvtkout.write(
+                                '\n {0: .5f} {1: .5f} {2: .5f}'.format(vectors[ll, 0], vectors[ll, 1], vectors[ll, 2]))
+                if len(vtkFields.tensors) > 0:
+                    if not wrote_pd_hdr:
+                        fvtkout.write(('\nPOINT_DATA {0: d}').format(V.shape[0]))
+                        wrote_pd_hdr = True
+                    nf = len(vtkFields.tensors) / 2
+                    for k in range(nf):
+                        fvtkout.write('\nTENSORS ' + vtkFields.tensors[2 * k] + ' float')
+                        tensors = vtkFields.tensors[2 * k + 1]
+                        for ll in range(V.shape[0]):
+                            for kk in range(2):
+                                fvtkout.write(
+                                    '\n {0: .5f} {1: .5f} {2: .5f}'.format(tensors[ll, kk, 0], tensors[ll, kk, 1],
+                                                                           tensors[ll, kk, 2]))
+                fvtkout.write('\n')
+
+    # Reads .vtk file
+    def readVTK(self, fileName):
+        if gotVTK:
+            u = vtkPolyDataReader()
+            u.SetFileName(fileName)
+            u.Update()
+            v = u.GetOutput()
+            # print(v)
+            npoints = int(v.GetNumberOfPoints())
+            nfaces = int(v.GetNumberOfPolys())
+            V = np.zeros([npoints, 3])
+            for kk in range(npoints):
+                V[kk, :] = np.array(v.GetPoint(kk))
+
+            F = np.zeros([nfaces, 3])
+            for kk in range(nfaces):
+                c = v.GetCell(kk)
+                for ll in range(3):
+                    F[kk, ll] = c.GetPointId(ll)
+
+            self.vertices = V
+            self.faces = np.int_(F)
+            xDef1 = self.vertices[self.faces[:, 0], :]
+            xDef2 = self.vertices[self.faces[:, 1], :]
+            xDef3 = self.vertices[self.faces[:, 2], :]
+            self.centers = (xDef1 + xDef2 + xDef3) / 3
+            self.surfel = np.cross(xDef2 - xDef1, xDef3 - xDef1)
+        else:
+            raise Exception('Cannot run readVTK without VTK')
+
+    # Reads .obj file
+    def readOBJ(self, fileName):
+        if gotVTK:
+            u = vtkOBJReader()
+            u.SetFileName(fileName)
+            u.Update()
+            v = u.GetOutput()
+            # print(v)
+            npoints = int(v.GetNumberOfPoints())
+            nfaces = int(v.GetNumberOfPolys())
+            V = np.zeros([npoints, 3])
+            for kk in range(npoints):
+                V[kk, :] = np.array(v.GetPoint(kk))
+
+            F = np.zeros([nfaces, 3])
+            for kk in range(nfaces):
+                c = v.GetCell(kk)
+                for ll in range(3):
+                    F[kk, ll] = c.GetPointId(ll)
+
+            self.vertices = V
+            self.faces = np.int_(F)
+            xDef1 = self.vertices[self.faces[:, 0], :]
+            xDef2 = self.vertices[self.faces[:, 1], :]
+            xDef3 = self.vertices[self.faces[:, 2], :]
+            self.centers = (xDef1 + xDef2 + xDef3) / 3
+            self.surfel = np.cross(xDef2 - xDef1, xDef3 - xDef1)
+        else:
+            raise Exception('Cannot run readOBJ without VTK')
+
+    def readPLY(self, fileName):
+        if gotVTK:
+            u = vtkPLYReader()
+            u.SetFileName(fileName)
+            u.Update()
+            v = u.GetOutput()
+            # print(v)
+            npoints = int(v.GetNumberOfPoints())
+            nfaces = int(v.GetNumberOfPolys())
+            V = np.zeros([npoints, 3])
+            for kk in range(npoints):
+                V[kk, :] = np.array(v.GetPoint(kk))
+
+            F = np.zeros([nfaces, 3])
+            for kk in range(nfaces):
+                c = v.GetCell(kk)
+                for ll in range(3):
+                    F[kk, ll] = c.GetPointId(ll)
+
+            self.vertices = V
+            self.faces = np.int_(F)
+            xDef1 = self.vertices[self.faces[:, 0], :]
+            xDef2 = self.vertices[self.faces[:, 1], :]
+            xDef3 = self.vertices[self.faces[:, 2], :]
+            self.centers = (xDef1 + xDef2 + xDef3) / 3
+            self.surfel = np.cross(xDef2 - xDef1, xDef3 - xDef1)
+        else:
+            raise Exception('Cannot run readPLY without VTK')
+
+    def readTRI(self, fileName):
+        vertices = []
+        faces = []
+        with open(fileName, "r") as f:
+            pithc = f.readlines()
+            line_1 = pithc[0]
+            n_pts, n_tri = line_1.split(' ')
+            for i, line in enumerate(pithc[1:]):
+                pouetage = line.split(' ')
+                if i <= int(n_pts):
+                    vertices.append([float(poupout) for poupout in pouetage[:3]])
+                    # vertices.append([float(poupout) for poupout in pouetage[4:7]])
+                else:
+                    faces.append([int(poupout.split("\n")[0]) for poupout in pouetage[1:]])
+        self.vertices = np.array(vertices)
+        self.faces = np.array(faces)
+        xDef1 = self.vertices[self.faces[:, 0], :]
+        xDef2 = self.vertices[self.faces[:, 1], :]
+        xDef3 = self.vertices[self.faces[:, 2], :]
+        self.centers = (xDef1 + xDef2 + xDef3) / 3
+        self.surfel = np.cross(xDef2 - xDef1, xDef3 - xDef1)
+
+    def concatenate(self, fvl):
+        nv = 0
+        nf = 0
+        for fv in fvl:
+            nv += fv.vertices.shape[0]
+            nf += fv.faces.shape[0]
+        self.vertices = np.zeros([nv, 3])
+        self.faces = np.zeros([nf, 3], dtype='int')
+
+        nv0 = 0
+        nf0 = 0
+        for fv in fvl:
+            nv = nv0 + fv.vertices.shape[0]
+            nf = nf0 + fv.faces.shape[0]
+            self.vertices[nv0:nv, :] = fv.vertices
+            self.faces[nf0:nf, :] = fv.faces + nv0
+            nv0 = nv
+            nf0 = nf
+        self.computeCentersAreas()
+
+def get_surf_area(surf):
+    areas = np.linalg.norm(surf.surfel, axis=-1)
+    return areas.sum()/2
+
+def centroid(surf):
+    areas = np.linalg.norm(surf.surfel, axis=-1, keepdims=True)
+    center = (surf.centers * areas).sum(axis=0) / areas.sum()
+    return center, np.sqrt(areas.sum()/2)
+
+
+def do_bbox_vertices(vertices):
+    new_verts = np.zeros(vertices.shape)
+    new_verts[:, 0] = (vertices[:, 0] - np.amin(vertices[:, 0]))/(np.amax(vertices[:, 0]) - np.amin(vertices[:, 0]))
+    new_verts[:, 1] = (vertices[:, 1] - np.amin(vertices[:, 1]))/(np.amax(vertices[:, 1]) - np.amin(vertices[:, 1]))
+    new_verts[:, 2] = (vertices[:, 2] - np.amin(vertices[:, 2]))/(np.amax(vertices[:, 1]) - np.amin(vertices[:, 2]))*0.5
+    return new_verts
+
+def opt_rot_surf(surf_1, surf_2, areas_c):
+    center_1, _ = centroid(surf_1)
+    pts_c1 = surf_1.centers - center_1
+    center_2, _ = centroid(surf_2)
+    pts_c2 = surf_2.centers - center_2 
+    to_sum = pts_c1[:, :, np.newaxis] * pts_c2[:, np.newaxis, :]
+    A = (to_sum * areas_c[:, :, np.newaxis]).sum(axis=0)
+    #A = to_sum.sum(axis=0)
+    u, _, v = np.linalg.svd(A)
+    a = np.array([[1, 0, 0], [0, 1, 0], [0, 0, np.sign(np.linalg.det(A))]])
+    O = u @ a @ v
+    return O.T
+
+def srnf(surf):
+    areas = np.linalg.norm(surf.surfel, axis=-1, keepdims=True) + 1e-8
+    return surf.surfel/np.sqrt(areas)
+
+def fps_gpt(points, num_samples):
+    """
+    Selects a subset of points using the Farthest Point Sampling algorithm.
+
+    :param points: numpy array of shape (N, D), where N is the number of points and D is the dimension of each point
+    :param num_samples: number of points to select
+    :return: numpy array of indices of the selected points
+    """
+    if num_samples > len(points):
+        raise ValueError("num_samples must be less than or equal to the number of points")
+
+    # Initialize an array to hold the indices of the selected points
+    selected_indices = np.zeros(num_samples, dtype=int)
+
+    center = np.mean(points, axis=0)
+
+    # Array to store the shortest distance of each point to a selected poin
+    shortest_distances = np.full(len(points), np.inf)
+    last_selected_point = center
+    for i in range(1, num_samples):
+        
+
+        # Update the shortest distance of each point to the selected points
+        for j, point in enumerate(points):
+            distance = np.linalg.norm(point - last_selected_point)
+            shortest_distances[j] = min(shortest_distances[j], distance)
+
+        # Select the point that is farthest from all selected points
+        selected_indices[i] = np.argmax(shortest_distances)
+        last_selected_point = points[selected_indices[i]]
+    return selected_indices
+
+
+def fps_gpt_geod(points, faces, num_samples, index_start):
+    """
+    Selects a subset of points using the Farthest Point Sampling algorithm.
+
+    :param points: numpy array of shape (N, D), where N is the number of points and D is the dimension of each point
+    :param num_samples: number of points to select
+    :return: numpy array of indices of the selected points
+    """
+    if num_samples > len(points):
+        raise ValueError("num_samples must be less than or equal to the number of points")
+
+    # Initialize an array to hold the indices of the selected points
+    selected_indices = np.zeros(num_samples, dtype=int)
+
+    # Array to store the shortest distance of each point to a selected poin
+    shortest_distances = np.full(len(points), np.inf)
+    selected_indices[0] = index_start
+    distances = np.zeros((num_samples, len(points)))
+    solver = pp3d.MeshHeatMethodDistanceSolver(points, faces)
+    for i in range(1, num_samples):
+        distances[i] = (solver.compute_distance(selected_indices[i-1]))
+
+        # Update the shortest distance of each point to the selected points
+        for j, point in enumerate(points):
+            distance = distances[i][j]
+            shortest_distances[j] = min(shortest_distances[j], distance)
+
+        # Select the point that is farthest from all selected points
+        selected_indices[i] = np.argmax(shortest_distances)
+    return selected_indices
+
+# class MultiSurface:
+#     def __init__(self, pattern):
+#         self.surf = []
+#         files = glob.glob(pattern)
+#         for f in files:
+#             self.surf.append(Surface(filename=f))

utils/torch_fmap.py

@@ -0,0 +1,77 @@
+import torch
+import torch.nn.functional as F
+from pykeops.torch import LazyTensor
+
+
+def euclidean_dist(x, y):
+    """
+    Args:
+      x: pytorch Variable, with shape [m, d]
+      y: pytorch Variable, with shape [n, d]
+    Returns:
+      dist: pytorch Variable, with shape [m, n]
+    """
+    #bs, m, n = x.size(0), x.size(1), y.size(1)
+    xx = torch.pow(x.squeeze(), 2).sum(1, keepdim=True)
+    yy = torch.pow(y.squeeze(), 2).sum(1, keepdim=True).t()
+    dist = xx + yy - 2 * torch.inner(x.squeeze(), y.squeeze())
+    dist = dist.clamp(min=1e-12).sqrt() 
+    return dist
+
+def knnsearch(x, y, alpha=1./0.07, prod=False):
+    if prod:
+        prods = torch.inner(x.squeeze(), y.squeeze())#/( torch.norm(x.squeeze(), dim=-1)[:, None]*torch.norm(y.squeeze(), dim=-1)[None, :])
+        output = F.softmax(alpha*prods, dim=1)
+    else:
+        distance = euclidean_dist(x, y[None,:])
+        output = F.softmax(-alpha*distance, dim=1)
+    return output.squeeze()
+
+def extract_p2p_torch(reps_shape, reps_template):
+    n_ev = reps_shape.shape[-1]
+    with torch.no_grad():
+        # print((evecs0_dzo @ fmap01_final.squeeze().T).shape)
+        # print(evecs1_dzo.shape)
+        reps_shape_torch = torch.from_numpy(reps_shape).float().cuda()
+        G_i = LazyTensor(reps_shape_torch[:, None, :].contiguous())  # (M**2, 1, 2)
+        reps_template_torch = torch.from_numpy(reps_template).float().cuda()
+        X_j = LazyTensor(reps_template_torch[None, :, :n_ev].contiguous())  # (1, N, 2)
+        D_ij = ((G_i - X_j) ** 2).sum(-1)  # (M**2, N) symbolic matrix of squared distances
+        indKNN = D_ij.argKmin(1, dim=0).squeeze()  # Grid <-> Samples, (M**2, K) integer tensor
+        # pmap10_ref = FM_to_p2p(fmap01_final.detach().squeeze().cpu().numpy(), s_dict['evecs'], template_dict['evecs'])
+        # print(indKNN[:10], pmap10_ref[:10])
+        indKNN_2 = D_ij.argKmin(1, dim=1).squeeze()
+    return indKNN.detach().cpu().numpy(), indKNN_2.detach().cpu().numpy()
+
+def extract_p2p_torch_fmap(fmap_shape_template, evecs_shape, evecs_template):
+    n_ev = fmap_shape_template.shape[-1]
+    with torch.no_grad():
+        # print((evecs0_dzo @ fmap01_final.squeeze().T).shape)
+        # print(evecs1_dzo.shape)
+        G_i = LazyTensor((evecs_shape[:, :n_ev] @ fmap_shape_template.squeeze().T)[:, None, :].contiguous())  # (M**2, 1, 2)
+        X_j = LazyTensor(evecs_template[None, :, :n_ev].contiguous())  # (1, N, 2)
+        D_ij = ((G_i - X_j) ** 2).sum(-1)  # (M**2, N) symbolic matrix of squared distances
+        indKNN = D_ij.argKmin(1, dim=0).squeeze()  # Grid <-> Samples, (M**2, K) integer tensor
+        # pmap10_ref = FM_to_p2p(fmap01_final.detach().squeeze().cpu().numpy(), s_dict['evecs'], template_dict['evecs'])
+        # print(indKNN[:10], pmap10_ref[:10])
+        indKNN_2 = D_ij.argKmin(1, dim=1).squeeze()
+    return indKNN.detach().cpu().numpy(), indKNN_2.detach().cpu().numpy()
+
+def wlstsq(A, B, w):
+    if w is None:
+        return torch.linalg.lstsq(A, B).solution
+    else:
+        assert w.dim() + 1 == A.dim() and w.shape[-1] == A.shape[-2]
+        W = torch.diag_embed(w)
+        return torch.linalg.lstsq(W @ A, W @ B).solution
+
+def torch_zoomout(evecs0, evecs1, evecs_1_trans, fmap01, target_size, step=1):
+    assert fmap01.shape[-2] == fmap01.shape[-1], f"square fmap needed, got {fmap01.shape[-2]} and {fmap01.shape[-1]}"
+    fs = fmap01.shape[0]
+    for i in range(fs, target_size+1, step):
+        indKNN, _ = extract_p2p_torch_fmap(fmap01, evecs0, evecs1)
+        #fmap01 = wlstsq(evecs1[..., :i], evecs0[indKNN, :i], None)
+        fmap01 = evecs_1_trans[:i, :] @ evecs0[indKNN, :i]
+    if fmap01.shape[0] < target_size:
+        fmap01 = evecs_1_trans[:target_size, :] @ evecs0[indKNN, :target_size]
+    return fmap01

utils/utils_func.py

@@ -0,0 +1,123 @@
+import os
+import os.path as osp
+import scipy.sparse as sp
+import shutil
+from pathlib import Path
+import torch
+import numpy as np
+import re
+import os
+import requests
+
+def ensure_pretrained_file(hf_url: str, save_dir: str = "pretrained", filename: str = "pretrained.pkl", token: str = None):
+    """
+    Ensure that a pretrained file exists locally.
+    If the folder is empty, download from Hugging Face.
+
+    Args:
+        hf_url (str): Hugging Face file URL (resolve/main/...).
+        save_dir (str): Directory to store pretrained file.
+        filename (str): Name of file to save.
+        token (str): Optional Hugging Face token (for gated/private repos).
+    """
+    os.makedirs(save_dir, exist_ok=True)
+    save_path = os.path.join(save_dir, filename)
+
+    if os.path.exists(save_path):
+        print(f"✅ Found pretrained file: {save_path}")
+        return save_path
+
+    headers = {"Authorization": f"Bearer {token}"} if token else {}
+
+    print(f"⬇️ Downloading pretrained file from {hf_url} to {save_path} ...")
+    response = requests.get(hf_url, headers=headers, stream=True)
+    response.raise_for_status()
+
+    with open(save_path, "wb") as f:
+        for chunk in response.iter_content(chunk_size=8192):
+            f.write(chunk)
+
+    print("✅ Download complete.")
+    return save_path
+
+
+def may_create_folder(folder_path):
+    if not osp.exists(folder_path):
+        oldmask = os.umask(000)
+        os.makedirs(folder_path, mode=0o777)
+        os.umask(oldmask)
+        return True
+    return False
+
+
+def make_clean_folder(folder_path):
+    success = may_create_folder(folder_path)
+    if not success:
+        shutil.rmtree(folder_path)
+        may_create_folder(folder_path)
+
+def str_delta(delta):
+    s = delta.total_seconds()
+    hours, remainder = divmod(s, 3600)
+    minutes, seconds = divmod(remainder, 60)
+    return '{:02}:{:02}:{:02}'.format(int(hours), int(minutes), int(seconds))
+
+def desc_from_config(config, world_size):
+    cond_str = 'uncond' #'cond' if c.dataset_kwargs.use_labels else 'uncond'
+    dtype_str = 'fp16' if config["perfs"]["fp16"] else 'fp32'
+    name_data = config["data"]["name"]
+    if "abs" in config["data"]:
+        name_data += "abs" if config["data"]["abs"] else ""
+    desc = (f'{name_data:s}-{cond_str:s}-{config["architecture"]["model"]:s}-'
+            f'gpus{world_size:d}-batch{config["hyper_params"]["batch_size"]:d}-{dtype_str:s}')
+    if config["misc"]["precond"]:
+        desc += f'-{config["misc"]["precond"]:s}'
+
+    if "desc" in config["misc"]:
+        if config["misc"]["desc"] is not None:
+            desc += f'-{config["misc"]["desc"]}'
+    return desc
+
+def get_dataset_path(config):
+    name_exp = config["data"]["name"]
+    if "add_name" in config:
+        if config["add_name"]["do"]:
+            name_exp += "_" + config["add_name"]["name"]
+    dataset_path = os.path.join(config["data"]["root_dir"], name_exp)
+    return name_exp, dataset_path
+
+def sparse_np_to_torch(A):
+    Acoo = A.tocoo()
+    values = Acoo.data
+    indices = np.vstack((Acoo.row, Acoo.col))
+    shape = Acoo.shape
+    return torch.sparse_coo_tensor(torch.LongTensor(indices), torch.FloatTensor(values), torch.Size(shape), dtype=torch.float32).coalesce()
+
+def convert_dict(np_dict, device="cpu"):
+    torch_dict = {}
+    for k, value in np_dict.items():
+        if sp.issparse(value):
+            torch_dict[k] = sparse_np_to_torch(value).to(device)
+            if torch_dict[k].dtype == torch.int32:
+                torch_dict[k] = torch_dict[k].long().to(device)
+        elif isinstance(value, np.ndarray):
+            torch_dict[k] = torch.from_numpy(value).to(device)
+            if torch_dict[k].dtype == torch.int32:
+                torch_dict[k] = torch_dict[k].squeeze().long().to(device)
+        else:
+            torch_dict[k] = value
+    return torch_dict
+
+
+def convert_dict_torch(in_dict, device="cpu"):
+    torch_dict = {}
+    for k, value in in_dict.items():
+        if isinstance(value, torch.Tensor):
+            torch_dict[k] = in_dict[k].to(device)
+    return torch_dict
+
+def batchify_dict(torch_dict):
+    for k, value in torch_dict.items():
+        if isinstance(value, torch.Tensor):
+            if torch_dict[k].dtype != torch.int64:
+                torch_dict[k] = torch_dict[k].unsqueeze(0)

utils/utils_legacy.py

@@ -0,0 +1,130 @@
+import sys
+import os
+import time
+
+import torch
+import hashlib
+import numpy as np
+import scipy
+
+def read_lines(file_path):
+    with open(file_path, 'r') as fin:
+        lines = [line.strip() for line in fin.readlines() if len(line.strip()) > 0]
+    return lines
+
+# == Pytorch things
+
+
+def toNP(x):
+    """
+    Really, definitely convert a torch tensor to a numpy array
+    """
+    return x.detach().to(torch.device('cpu')).numpy()
+
+
+def label_smoothing_log_loss(pred, labels, smoothing=0.0):
+    n_class = pred.shape[-1]
+    one_hot = torch.zeros_like(pred)
+    one_hot[labels] = 1.
+    one_hot = one_hot * (1 - smoothing) + (1 - one_hot) * smoothing / (n_class - 1)
+    loss = -(one_hot * pred).sum(dim=-1).mean()
+    return loss
+
+
+# Randomly rotate points.
+# Torch in, torch out
+# Note fornow, builds rotation matrix on CPU.
+def random_rotate_points(pts, randgen=None):
+    R = random_rotation_matrix(randgen)
+    R = torch.from_numpy(R).to(device=pts.device, dtype=pts.dtype)
+    return torch.matmul(pts, R)
+
+
+def random_rotate_points_y(pts):
+    angles = torch.rand(1, device=pts.device, dtype=pts.dtype) * (2. * np.pi)
+    rot_mats = torch.zeros(3, 3, device=pts.device, dtype=pts.dtype)
+    rot_mats[0, 0] = torch.cos(angles)
+    rot_mats[0, 2] = torch.sin(angles)
+    rot_mats[2, 0] = -torch.sin(angles)
+    rot_mats[2, 2] = torch.cos(angles)
+    rot_mats[1, 1] = 1.
+
+    pts = torch.matmul(pts, rot_mats)
+    return pts
+
+
+# Numpy things
+
+
+# Numpy sparse matrix to pytorch
+def sparse_np_to_torch(A):
+    Acoo = A.tocoo()
+    values = Acoo.data
+    indices = np.vstack((Acoo.row, Acoo.col))
+    shape = Acoo.shape
+    return torch.sparse_coo_tensor(torch.LongTensor(indices), torch.FloatTensor(values), torch.Size(shape), dtype=torch.float32).coalesce()
+
+
+# Pytorch sparse to numpy csc matrix
+def sparse_torch_to_np(A):
+    if len(A.shape) != 2:
+        raise RuntimeError("should be a matrix-shaped type; dim is : " + str(A.shape))
+
+    indices = toNP(A.indices())
+    values = toNP(A.values())
+
+    mat = scipy.sparse.coo_matrix((values, indices), shape=A.shape).tocsc()
+
+    return mat
+
+
+# Hash a list of numpy arrays
+def hash_arrays(arrs):
+    running_hash = hashlib.sha1()
+    for arr in arrs:
+        binarr = arr.view(np.uint8)
+        running_hash.update(binarr)
+    return running_hash.hexdigest()
+
+
+def random_rotation_matrix(randgen=None):
+    """
+    Creates a random rotation matrix.
+    randgen: if given, a np.random.RandomState instance used for random numbers (for reproducibility)
+    """
+    # adapted from http://www.realtimerendering.com/resources/GraphicsGems/gemsiii/rand_rotation.c
+
+    if randgen is None:
+        randgen = np.random.RandomState()
+
+    theta, phi, z = tuple(randgen.rand(3).tolist())
+
+    theta = theta * 2.0 * np.pi  # Rotation about the pole (Z).
+    phi = phi * 2.0 * np.pi  # For direction of pole deflection.
+    z = z * 2.0  # For magnitude of pole deflection.
+
+    # Compute a vector V used for distributing points over the sphere
+    # via the reflection I - V Transpose(V).  This formulation of V
+    # will guarantee that if x[1] and x[2] are uniformly distributed,
+    # the reflected points will be uniform on the sphere.  Note that V
+    # has length sqrt(2) to eliminate the 2 in the Householder matrix.
+
+    r = np.sqrt(z)
+    Vx, Vy, Vz = V = (np.sin(phi) * r, np.cos(phi) * r, np.sqrt(2.0 - z))
+
+    st = np.sin(theta)
+    ct = np.cos(theta)
+
+    R = np.array(((ct, st, 0), (-st, ct, 0), (0, 0, 1)))
+    # Construct the rotation matrix  ( V Transpose(V) - I ) R.
+
+    M = (np.outer(V, V) - np.eye(3)).dot(R)
+    return M
+
+
+# Python string/file utilities
+def ensure_dir_exists(d):
+    if not os.path.exists(d):
+        os.makedirs(d)
+
+

zero_shot.py

@@ -0,0 +1,402 @@
+import torch
+import numpy as np
+import scipy
+import os
+import sys
+import random
+import numpy as np
+import torch
+import time
+from datetime import datetime
+import importlib
+import json
+import argparse
+from omegaconf import OmegaConf
+from snk.loss import PrismRegularizationLoss
+from snk.prism_decoder import PrismDecoder
+from shape_models.fmap import DFMNet
+from shape_models.encoder import Encoder
+from diffu_models.losses import VELoss, VPLoss, EDMLoss
+from diffu_models.sds import guidance_grad
+from utils.torch_fmap import torch_zoomout, knnsearch, extract_p2p_torch_fmap
+from utils.utils_func import convert_dict, str_delta, ensure_pretrained_file
+from utils.eval import accuracy
+from utils.mesh import save_ply, load_mesh
+from shape_data import get_data_dirs
+from utils.pickle_stuff import safe_load_with_fallback
+from utils.geometry import compute_operators, load_operators
+from utils.surfaces import Surface
+import sys
+try:
+    import google.colab
+    print("Running Colab")
+    from tqdm import tqdm
+except ImportError:
+    print("Running local")
+    from tqdm.auto import tqdm
+
+
+def seed_everything(seed=42):
+  random.seed(seed)
+  os.environ['PYTHONHASHSEED'] = str(seed)
+  np.random.seed(seed)
+  torch.manual_seed(seed)
+  torch.backends.cudnn.deterministic = True
+  torch.backends.cudnn.benchmark = False
+
+seed_everything()
+
+class Tee:
+    def __init__(self, *outputs):
+        self.outputs = outputs
+
+    def write(self, message):
+        for output in self.outputs:
+            output.write(message)
+            output.flush()  # ensure it's written immediately
+
+    def flush(self):
+        for output in self.outputs:
+            output.flush()
+
+class DiffModel:
+
+    def __init__(self, cfg, device="cuda:0"):
+        if cfg["train_dir"] == "pretrained":
+            url = "https://huggingface.co/daidedou/diffumatch_model/resolve/main/network-snapshot-041216.pkl"
+            network_pkl = ensure_pretrained_file(url, "pretrained")
+            url_json = "https://huggingface.co/daidedou/diffumatch_model/resolve/main/training_options.json"
+            json_filename = ensure_pretrained_file(url_json, "pretrained", filename="training_options.json")
+            train_cfg = json.load(open(json_filename))
+        else:
+            num_exp = cfg["diff_num_exp"]
+            files = os.listdir(cfg["train_dir"])
+            for file in files:
+                if file[:5] == f"{num_exp:05d}":
+                    netdir = os.path.join(cfg["train_dir"], file)
+            train_cfg = json.load(open(os.path.join(netdir, "training_options.json")))
+            pkls = [f for f in os.listdir(netdir) if ".pkl" in f]
+            nice_pkls = sorted(pkls, key=lambda x: int(x.split(".")[0].split("-")[-1]))
+            chosen_pkl = nice_pkls[-1]
+            network_pkl = os.path.join(netdir, chosen_pkl)
+        print(f'Loading network from "{network_pkl}"...')
+        self.net = safe_load_with_fallback(network_pkl)['ema'].to(device)
+        
+        print('Done!')
+        loss_name = train_cfg['hyper_params']['loss_name']
+        self.loss_sde = None
+        if loss_name == "EDMLoss":
+            self.loss_sde = EDMLoss()
+        elif loss_name == "VPLoss":
+            self.loss_sde = VPLoss()
+
+
+class Matcher(object):
+
+    def __init__(self, cfg):
+        self.cfg = cfg
+        self.device = torch.device(f'cuda:{cfg["gpu"]}' if torch.cuda.is_available() else 'cpu')
+        self.diffusion_model = None
+        if self.cfg.get("sds", False):
+            self.diffusion_model = DiffModel(cfg["sds_conf"])
+        self.n_fmap = self.cfg["deepfeat_conf"]["fmap"]["n_fmap"]
+        self.n_loop = 0
+        if self.cfg.get("optimize", False):
+            self.n_loop = self.cfg.opt.get("n_loop", 0)
+        self.snk = self.cfg.get("snk", False)
+        self.fmap_cfg = self.cfg.deepfeat_conf.fmap
+        self.dataloaders = dict()
+
+    def reconf(self, cfg):
+        self.cfg = cfg
+        self.n_fmap = self.cfg["deepfeat_conf"]["fmap"]["n_fmap"]
+        self.n_loop = 0
+        if self.cfg.get("optimize", False):
+            self.n_loop = self.cfg.opt.get("n_loop", 0)
+        self.fmap_cfg = self.cfg.deepfeat_conf.fmap
+        self.dataloaders = dict()
+        
+    def _init(self):
+        cfg = self.cfg
+        self.fmap_model = DFMNet(self.cfg["deepfeat_conf"]["fmap"]).to(self.device)
+        if self.snk:
+            self.encoder = Encoder().to(self.device)
+            self.decoder = PrismDecoder(dim_in=515).to(self.device)
+            self.loss_prism = PrismRegularizationLoss(primo_h=0.02)
+            self.soft_p2p = True
+            params_to_opt = list(self.fmap_model.parameters()) + list(self.encoder.parameters()) + list(self.decoder.parameters())
+        else:
+            params_to_opt = self.fmap_model.parameters()
+        self.optim = torch.optim.Adam(params_to_opt, lr=0.001, betas=(0.9, 0.99))
+        self.eye = torch.eye(self.n_fmap).float().to(self.device)
+        self.eye.requires_grad = False
+
+    def fmap(self, shape_dict, target_dict):
+        if self.fmap_cfg.get("use_diff", False):
+            C12_pred, C21_pred, feat1, feat2, evecs_trans1, evecs_trans2 = self.fmap_model({"shape1": shape_dict, "shape2": target_dict}, diff_model=self.diffusion_model, scale=self.fmap_cfg.diffusion.time)
+            C12_pred, C12_obj, mask_12 = C12_pred
+            C21_pred, C21_obj, mask_21 = C21_pred
+        else:
+            C12_pred, C21_pred, feat1, feat2, evecs_trans1, evecs_trans2 = self.fmap_model({"shape1": shape_dict, "shape2": target_dict})
+            C12_obj, C21_obj = C12_pred, C21_pred
+            mask_12, mask_21 = None, None
+        return C12_pred, C12_obj, C21_pred, C21_obj, feat1, feat2, evecs_trans1, evecs_trans2, mask_12, mask_21
+    
+
+    def zo_shot(self, shape_dict, target_dict):
+        self._init()
+        evecs1, evecs2 = shape_dict["evecs"], target_dict["evecs"]
+        _, C12_mask_init, _, _, _, _, _ , _, _, _ = self.fmap(shape_dict, target_dict)
+        evecs_2trans = evecs2.t() @ torch.diag(target_dict["mass"])
+        new_FM = torch_zoomout(evecs1, evecs2, evecs_2trans, C12_mask_init.squeeze(), self.cfg["zo_shot"])
+        indKNN_new, _ = extract_p2p_torch_fmap(new_FM, evecs1, evecs2)
+        return new_FM, indKNN_new
+
+
+    def optimize(self, shape_dict, target_dict, target_normals):
+        self._init()
+        evecs1, evecs2 = shape_dict["evecs"], target_dict["evecs"]
+        C12_pred_init, _, _, _ , _, _, evecs_trans1, evecs_trans2, _, _ = self.fmap(shape_dict, target_dict)
+        evecs_2trans = evecs2.t() @ torch.diag(target_dict["mass"])
+        evecs_1trans = evecs1.t() @ torch.diag(shape_dict["mass"])
+        n_verts_target = target_dict["vertices"].shape[-2]
+        
+        loss_save = {"cycle": [], "fmap": [], "mse": [], "prism": [], "bij": [], "ortho": [], "sds": [], "lap": [], "proper": []}
+        snk_rec = None
+        for i in tqdm(range(self.n_loop), "Optimizing matching " + shape_dict['name'] + " " + target_dict['name']):
+            C12_pred, C12_obj, C21_pred, C21_obj, feat1, feat2, evecs_trans1, evecs_trans2, _, _ = self.fmap(shape_dict, target_dict)
+            if self.cfg.opt.soft_p2p:
+                ### A la SNK
+                ## P2P 2 -> 1
+                soft_p2p_21 = knnsearch(evecs2[:, :self.n_fmap] @ C12_pred.squeeze(), evecs1[:, :self.n_fmap], prod=True)
+                C12_new = evecs_trans2[:self.n_fmap, :] @ soft_p2p_21 @ evecs1[:, :self.n_fmap]
+                soft_p2p_21 = knnsearch(evecs2[:, :self.n_fmap] @ C12_new.squeeze(), evecs1[:, :self.n_fmap], prod=True)
+
+                ## P2P 1 -> 2 
+                soft_p2p_12 = knnsearch(evecs1[:, :self.n_fmap] @ C21_pred.squeeze(), evecs2[:, :self.n_fmap], prod=True)
+                C21_new = evecs_trans1[:self.n_fmap, :] @ soft_p2p_12 @ evecs2[:, :self.n_fmap]
+                soft_p2p_12 = knnsearch(evecs1[:, :self.n_fmap] @ C21_new.squeeze(), evecs2[:, :self.n_fmap], prod=True)
+
+                l_cycle = ((soft_p2p_12 @ (soft_p2p_21 @ shape_dict["vertices"]) - shape_dict["vertices"])**2).sum(dim=-1).mean()
+            else:
+                C12_new, C21_new = C12_pred, C21_pred
+
+            l_ortho = ((C12_new.squeeze() @ C12_new.squeeze().T - self.eye)**2).mean() + ((C21_new.squeeze() @ C21_new.squeeze().T - self.eye)**2).mean()
+            l_bij = ((C12_new.squeeze() @ C21_new.squeeze() - self.eye)**2).mean() + ((C21_new.squeeze() @ C12_new.squeeze() - self.eye)**2).mean()
+            l_lap = ((C12_new @ torch.diag(shape_dict["evals"][:self.n_fmap]) - torch.diag(target_dict["evals"][:self.n_fmap]) @ C12_new)**2).mean()
+            l_lap += ((C21_new @ torch.diag(target_dict["evals"][:self.n_fmap]) - torch.diag(shape_dict["evals"][:self.n_fmap]) @ C21_new)**2).mean()
+
+
+            l_cycle, l_prism, l_mse = torch.as_tensor(0.).float().to(self.device), torch.as_tensor(0.).float().to(self.device), torch.as_tensor(0.).float().to(self.device)
+            if self.snk:
+                # Latent vector 
+                latents = self.encoder(shape_dict)
+                latents_duplicate = latents[None, :].repeat(n_verts_target, 1)
+
+                # Prism decoder
+                feats_decode = torch.cat((target_dict["vertices"], latents_duplicate), dim=1)
+                snk_rec, prism, rots = self.decoder(target_dict, feats_decode)
+                l_prism = self.loss_prism(prism, rots, target_dict["vertices"], target_dict["faces"], target_normals)
+                l_mse = ((soft_p2p_21 @ shape_dict["vertices"] - snk_rec)**2).sum(dim=-1).mean()
+                l_cycle = ((soft_p2p_12 @ (soft_p2p_21 @ shape_dict["vertices"]) - shape_dict["vertices"])**2).sum(dim=-1).mean()
+            l_sds, l_proper = torch.as_tensor(0.).float().to(self.device), torch.as_tensor(0.).float().to(self.device)
+            if self.fmap_cfg.get("use_diff", False):
+                if self.fmap_cfg.diffusion.get("abs", False):
+                    C12_in, C21_in = torch.abs(C12_pred).squeeze(), torch.abs(C21_pred).squeeze()
+                else:
+                    C12_in, C21_in = C12_pred.squeeze(), C21_pred.squeeze()
+                grad_12, _ = guidance_grad(C12_in, self.diffusion_model.net, grad_scale=1, batch_size=self.fmap_cfg.diffusion.batch_sds, 
+                                           scale_noise=self.fmap_cfg.diffusion.time, device=self.device)
+                with torch.no_grad():
+                    denoised_12 = C12_pred - self.optim.param_groups[0]['lr'] * grad_12
+                targets_12 = torch_zoomout(evecs1, evecs2, evecs_2trans, C12_obj.squeeze(), self.cfg.sds_conf.zoomout)   
+                             
+                l_proper_12 = ((C12_pred.squeeze()[:self.n_fmap, :self.n_fmap] - targets_12.squeeze()[:self.n_fmap, :self.n_fmap])**2).mean()
+
+                grad_21, _ = guidance_grad(C21_in, self.diffusion_model.net, grad_scale=1, batch_size=self.fmap_cfg.diffusion.batch_sds, 
+                                           scale_noise=self.fmap_cfg.diffusion.time, device=self.device)
+                #denoised_21 = C21_pred - self.optim.param_groups[0]['lr'] * grad_21
+                with torch.no_grad():
+                    denoised_21 = C21_pred - self.optim.param_groups[0]['lr'] * grad_21 
+                targets_21 = torch_zoomout(evecs2, evecs1, evecs_1trans, C21_obj.squeeze(), self.cfg.sds_conf.zoomout)#, step=10)
+                l_proper_21 = ((C21_pred.squeeze()[:self.n_fmap, :self.n_fmap] - targets_21.squeeze()[:self.n_fmap, :self.n_fmap])**2).mean()
+                l_proper = l_proper_12 + l_proper_21
+
+                l_sds = ((torch.abs(C12_pred).squeeze()[:self.n_fmap, :self.n_fmap] - denoised_12.squeeze()[:self.n_fmap, :self.n_fmap])**2).mean()
+                l_sds += ((torch.abs(C21_pred).squeeze()[:self.n_fmap, :self.n_fmap] - denoised_21.squeeze()[:self.n_fmap, :self.n_fmap])**2).mean()
+            loss = torch.as_tensor(0.).float().to(self.device)
+            if self.cfg.loss.get("ortho", 0) > 0:
+                loss += self.cfg.loss.get("ortho", 0) *  l_ortho
+            if self.cfg.loss.get("bij", 0) > 0:
+                loss += self.cfg.loss.get("bij", 0) *  l_bij
+            if self.cfg.loss.get("lap", 0) > 0:
+                loss += self.cfg.loss.get("lap", 0) *  l_lap 
+            if self.cfg.loss.get("cycle", 0) > 0:
+                loss += self.cfg.loss.get("cycle", 0) *  l_cycle
+            if self.cfg.loss.get("mse_rec", 0) > 0:
+                loss += self.cfg.loss.get("mse_rec", 0) *  l_mse
+            if self.cfg.loss.get("prism_rec", 0) > 0:
+                loss += self.cfg.loss.get("prism_rec", 0) *  l_prism
+            if self.cfg.loss.get("sds", 0) > 0 and self.fmap_cfg.get("use_diff", False):
+                loss += self.cfg.loss.get("sds", 0) * l_sds
+            if self.cfg.loss.get("proper", 0) > 0 and self.fmap_cfg.get("use_diff", False):
+                loss += self.cfg.loss.get("proper", 0) * l_proper
+        
+            loss.backward()
+            self.optim.step()
+            self.optim.zero_grad()
+
+            loss_save["cycle"].append(l_cycle.item())
+            loss_save["ortho"].append(l_ortho.item())
+            loss_save["bij"].append(l_bij.item())
+            loss_save["sds"].append(l_sds.item())
+            loss_save["proper"].append(l_proper.item())
+            loss_save["mse"].append(l_mse.item())
+            loss_save["prism"].append(l_prism.item())
+        indKNN_new_init, _ = extract_p2p_torch_fmap(C12_pred_init, evecs1, evecs2)
+        indKNN_new, _ = extract_p2p_torch_fmap(C12_new, evecs1, evecs2)
+        return C12_new, indKNN_new, indKNN_new_init, snk_rec, loss_save
+    
+
+
+    def match(self, pair_batch, output_pair, geod_path, refine=True, eval=False):
+        shape_dict, _, target_dict, _, target_normals, mapinfo = pair_batch 
+        shape_dict_device = convert_dict(shape_dict, self.device)
+        target_dict_device = convert_dict(target_dict, self.device)
+        print(shape_dict_device["vertices"].device)
+        os.makedirs(output_pair, exist_ok=True)
+
+
+        if self.cfg["optimize"]:
+            C12_new, p2p, p2p_init, snk_rec, loss_save = self.optimize(shape_dict_device, target_dict_device, target_normals.to(self.device))
+            np.save(os.path.join(output_pair, "p2p_init.npy"), p2p_init)
+            np.save(os.path.join(output_pair, "losses.npy"), loss_save)
+        else:
+            C12_new, p2p = self.zo_shot(shape_dict_device, target_dict_device)
+            snk_rec, loss_save = None, None
+        np.save(os.path.join(output_pair, "fmap.npy"), C12_new.detach().squeeze().cpu().numpy())
+        np.save(os.path.join(output_pair, "p2p.npy"), p2p)
+        if snk_rec is not None:
+            save_ply(os.path.join(output_pair, "rec.ply"), snk_rec.detach().squeeze().cpu().numpy(), target_dict["faces"])
+
+        if refine:
+            evecs1, evecs2 = shape_dict_device["evecs"], target_dict_device["evecs"]
+            evecs_2trans = evecs2.t() @ torch.diag(target_dict_device["mass"])
+            new_FM = torch_zoomout(evecs1, evecs2, evecs_2trans, C12_new.squeeze(), 128)#, step=10)
+            p2p_refined_zo, _ = extract_p2p_torch_fmap(new_FM, evecs1, evecs2)
+            np.save(os.path.join(output_pair, "p2p_zo.npy"), p2p)
+        if eval:
+            file_i, vts_1, vts_2 = mapinfo
+            mat_loaded = scipy.io.loadmat(os.path.join(geod_path, file_i + ".mat"))
+            A_geod, sqrt_area = mat_loaded['geod_dist'], np.sqrt(mat_loaded['areas_f'].sum())
+            _, dist = accuracy(p2p[vts_2], vts_1, A_geod,
+                                            sqrt_area=sqrt_area,
+                                            return_all=True)
+            if refine:
+                _, dist_zo = accuracy(p2p_refined_zo[vts_2], vts_1, A_geod,
+                                            sqrt_area=sqrt_area,
+                                            return_all=True)
+                np.savetxt(os.path.join(output_pair, "dists.txt"), (dist.mean(), dist_zo.mean()))
+                return p2p, p2p_refined_zo, loss_save, dist.mean(), dist_zo.mean()
+            return p2p, loss_save, dist.mean()
+        return p2p, loss_save
+        
+
+
+
+    def _dataset_epoch(self, dataset, name_dataset, save_dir, data_dir):
+        os.makedirs(save_dir, exist_ok=True)
+        # dloader = DataLoader(dataset, collate_fn=collate_default, batch_size=1)
+        num_pairs = len(dataset)
+        id_pair = 0
+        all_accs = []
+        all_accs_zo = []
+        t1 = datetime.now()
+        save_txt = os.path.join(save_dir, "log.txt")
+        # Open a file for writing
+        log_file = open(save_txt, 'w')
+        # Replace sys.stdout with Tee that writes to both console and file
+        sys.stdout = Tee(sys.__stdout__, log_file)
+
+        for batch in dset:
+            shape_dict, _, target_dict, _, _, _ = batch
+            print("Pair: " + shape_dict['name'] + " " + target_dict['name'])
+            name_exp = os.path.join(save_dir, shape_dict['name'], target_dict['name'])
+            if self.cfg.get("refine", False):
+                _, _, _, dist, dist_zo = self.match(batch, name_exp, os.path.join(data_dir, "geomats", name_dataset), eval=True, refine=True)
+            else:
+                _, _, dist = self.match(batch, name_exp, os.path.join(data_dir, "geomats", name_dataset), eval=True, refine=False)
+            delta = datetime.now() - t1
+            fm_delta = str_delta(delta)
+            remains = ((delta/(id_pair+1))*num_pairs) - delta
+            fm_remains = str_delta(remains)
+            all_accs.append(dist)
+            accs_mean = np.mean(all_accs)
+            if self.cfg.get("refine", False):
+                all_accs_zo.append(dist_zo)
+                accs_zo = np.mean(all_accs_zo)
+                print(f"error: {dist}, zo: {dist_zo}, element {id_pair}/{num_pairs}, mean accuracy: {accs_mean}, mean zo: {accs_zo}, full time: {fm_delta}, remains: {fm_remains}")
+            else:
+                print(f"error: {dist}, element {id_pair}/{num_pairs}, mean accuracy: {accs_mean}, full time: {fm_delta}, remains: {fm_remains}")
+            id_pair += 1
+        if self.cfg.get("refine", False):
+            print(f"mean error : {np.mean(all_accs)}, mean error refined: {np.mean(all_accs_zo)}")
+        else:
+            print(f"mean error : {np.mean(all_accs)}")
+        sys.stdout = sys.__stdout__ 
+
+    def load_data(self, file, num_evecs=200, make_cache=False, factor=None):
+        name = os.path.basename(os.path.splitext(file)[0])
+        cache_file = "single_" + name + ".npz"
+        verts_shape, faces, vnormals, area_shape, center_shape = load_mesh(file, return_vnormals=True)
+        cache_path = os.path.join(self.cfg.cache, cache_file)
+        print("Cache is: ", cache_path)
+        if not os.path.exists(cache_path) or make_cache:
+            print("Computing operators ...")
+            compute_operators(verts_shape, faces, vnormals, num_evecs, cache_path, force_save=make_cache)
+        data_dict = load_operators(cache_path)
+        data_dict['name'] = name
+        data_dict_torch = convert_dict(data_dict, self.device)
+        #batchify_dict(data_dict_torch)
+        return data_dict_torch, area_shape
+
+    def match_files(self, file_shape, file_target):
+        batch_shape, _ = self.load_data(file_shape)
+        batch_target, _ = self.load_data(file_target) 
+        target_surf = Surface(filename=file_target)
+        target_normals = torch.from_numpy(target_surf.surfel/np.linalg.norm(target_surf.surfel, axis=-1, keepdims=True)).float().to(self.device)
+        batch = batch_shape, None, batch_target, target_normals, None, None
+        output_folder = os.path.join(self.cfg.output, batch_shape["name"] + "_" + batch_shape["target"])
+        p2p, _ = self.match(batch, output_folder, None)
+        return batch_shape, batch_target, p2p
+
+
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description="Launch the SDS demo over datasets")
+    parser.add_argument('--dataset', type=str, default="SCAPE", help='name of the dataset')
+    parser.add_argument('--config', type=str, default="config/matching/sds.yaml", help='Config file location')    
+    parser.add_argument('--datadir', type=str, default="data", help='path where datasets are store')
+    parser.add_argument('--output', type=str, default="results", help="where to store experience results")
+    args = parser.parse_args()
+
+    arg_cfg = OmegaConf.from_dotlist(
+        [f"{k}={v}" for k, v in vars(args).items() if v is not None]
+    )
+    yaml_cfg = OmegaConf.load(args.config)
+    cfg = OmegaConf.merge(yaml_cfg, arg_cfg)
+    dataset_name = args.dataset.lower()
+    if cfg.get("oriented", False):
+        dataset_name += "_ori"
+    shape_cls = getattr(importlib.import_module(f'shape_data.{args.dataset.lower()}'), 'ShapeDataset')
+    pair_cls = getattr(importlib.import_module(f'shape_data.{args.dataset.lower()}'), 'ShapePairDataset')
+    data_dir, name_data_geo, corr_dir = get_data_dirs(args.datadir, dataset_name, 'test')
+    name_data_geo = "_".join(name_data_geo.split("_")[:2])
+    dset_shape = shape_cls(data_dir, "cache/fmaps", "test", oriented=cfg.get("oriented", False))
+    print("Preprocessing shapes done.")
+    dset = pair_cls(corr_dir, 'test', dset_shape, rotate=cfg.get("rotate", False))
+    exp_time = time.strftime('%y-%m-%d_%H-%M-%S')
+    output_logs = os.path.join(args.output, name_data_geo, exp_time)
+    matcher = Matcher(cfg)
+    matcher._dataset_epoch(dset, name_data_geo, output_logs, args.datadir)