add icp refinement from p2p

pyFM/refine/__init__.py

@@ -1,2 +1,6 @@
-from .icp import icp_refine, mesh_icp_refine
-from .zoomout import zoomout_refine, mesh_zoomout_refine, mesh_zoomout_refine_p2p
+from .icp import icp_refine, mesh_icp_refine, mesh_icp_refine_p2p  # noqa: F401
+from .zoomout import (
+    zoomout_refine,  # noqa: F401
+    mesh_zoomout_refine,  # noqa: F401
+    mesh_zoomout_refine_p2p,  # noqa: F401
+)

pyFM/refine/icp.py

@@ -40,8 +40,17 @@ def icp_iteration(FM_12, evects1, evects2, use_adj=False, n_jobs=1):
     return U @ np.eye(k2, k1) @ VT
 
 
-def icp_refine(FM_12, evects1, evects2, nit=10, tol=1e-10, use_adj=False,
-               return_p2p=False, n_jobs=1, verbose=False):
+def icp_refine(
+    FM_12,
+    evects1,
+    evects2,
+    nit=10,
+    tol=1e-10,
+    use_adj=False,
+    return_p2p=False,
+    n_jobs=1,
+    verbose=False,
+):
     """
     Refine a functional map using the standard ICP algorithm.
     One can use the adjoint instead of the functional map for pointwise map computation.
@@ -85,33 +94,47 @@ def icp_refine(FM_12, evects1, evects2, nit=10, tol=1e-10, use_adj=False,
         myrange = range(10000)
 
     for i in myrange:
-        FM_12_icp = icp_iteration(FM_12_curr, evects1, evects2, use_adj=use_adj, n_jobs=n_jobs)
+        FM_12_icp = icp_iteration(
+            FM_12_curr, evects1, evects2, use_adj=use_adj, n_jobs=n_jobs
+        )
 
         if nit is None or nit == 0:
             if verbose:
-                print(f'iteration : {1+i} - mean : {np.square(FM_12_curr - FM_12_icp).mean():.2e}'
-                      f' - max : {np.max(np.abs(FM_12_curr - FM_12_icp)):.2e}')
+                print(
+                    f"iteration : {1+i} - mean : {np.square(FM_12_curr - FM_12_icp).mean():.2e}"
+                    f" - max : {np.max(np.abs(FM_12_curr - FM_12_icp)):.2e}"
+                )
             if np.max(np.abs(FM_12_curr - FM_12_icp)) <= tol:
                 break
 
         FM_12_curr = FM_12_icp.copy()
 
     if nit is None or nit == 0 and verbose:
         run_time = time.time() - start_time
-        print(f'ICP done with {iteration:d} iterations - {run_time:.2f} s')
+        print(f"ICP done with {iteration:d} iterations - {run_time:.2f} s")
 
     if return_p2p:
-        p2p_21_icp = spectral.FM_to_p2p(FM_12_icp, evects1, evects2, use_adj=use_adj, n_jobs=n_jobs)  # (n2,)
+        p2p_21_icp = spectral.FM_to_p2p(
+            FM_12_icp, evects1, evects2, use_adj=use_adj, n_jobs=n_jobs
+        )  # (n2,)
         return FM_12_icp, p2p_21_icp
 
     return FM_12_icp
 
 
-def mesh_icp_refine(FM_12, mesh1, mesh2, nit=10, tol=1e-10, use_adj=False,
-                    return_p2p=False, n_jobs=1, verbose=False):
+def mesh_icp_refine(
+    FM_12,
+    mesh1,
+    mesh2,
+    nit=10,
+    tol=1e-10,
+    use_adj=False,
+    return_p2p=False,
+    n_jobs=1,
+    verbose=False,
+):
     """
-    Refine a functional map using the auxiliar ICP algorithm (different conversion
-    from functional map to vertex-to-vertex)
+    Refine a functional map using the ICP algorithm.
 
     Parameters
     --------------------------
@@ -143,8 +166,83 @@ def mesh_icp_refine(FM_12, mesh1, mesh2, nit=10, tol=1e-10, use_adj=False,
     """
     k2, k1 = FM_12.shape
 
-    result = icp_refine(FM_12, mesh1.eigenvectors[:, :k1], mesh2.eigenvectors[:, :k2],
-                        nit=nit, tol=tol, use_adj=use_adj, return_p2p=return_p2p, n_jobs=n_jobs,
-                        verbose=verbose)
+    result = icp_refine(
+        FM_12,
+        mesh1.eigenvectors[:, :k1],
+        mesh2.eigenvectors[:, :k2],
+        nit=nit,
+        tol=tol,
+        use_adj=use_adj,
+        return_p2p=return_p2p,
+        n_jobs=n_jobs,
+        verbose=verbose,
+    )
+
+    return result
+
+
+def mesh_icp_refine_p2p(
+    p2p_21,
+    mesh1,
+    mesh2,
+    k_init,
+    nit=10,
+    tol=1e-10,
+    use_adj=False,
+    return_p2p=False,
+    n_jobs=1,
+    verbose=False,
+):
+    """
+    Refine a functional map using the auxiliar ICP algorithm (different conversion
+    from functional map to vertex-to-vertex).
+    This algorithm starts from an initial pointwise map instead of a functional map.
+
+    Parameters
+    --------------------------
+    p2p_21      : np.ndarray
+        (n2,) initial pointwise map from mesh2 to mesh1
+    mesh1      : TriMesh
+        Source mesh
+    mesh2      : TriMesh
+        Target mesh
+    k_init : int
+        Initial number of eigenvectors to use
+    nit        : int
+        Number of iterations to perform. If not specified, uses the tol parameter
+    tol        : float
+        Maximum change in a functional map to stop refinement
+                 (only used if nit is not specified)
+    use_adj    :
+        use the adjoint method
+    n_jobs     :
+        number of parallel jobs. Use -1 to use all processes
+    return_p2p : bool
+        if True returns the vertex to vertex map from 2 to 1
+
+    Returns
+    ---------------------------
+    FM_12_icp  : np.ndarray
+        ICP-refined functional map
+    p2p_21_icp : np.ndarray
+        only if return_p2p is set to True - the refined pointwise map
+                 from basis 2 to basis 1
+    """
+
+    FM_12_init = spectral.mesh_p2p_to_FM(
+        p2p_21, mesh1, mesh2, dims=k_init, subsample=None
+    )
+
+    result = mesh_icp_refine(
+        FM_12_init,
+        mesh1,
+        mesh2,
+        nit=nit,
+        tol=tol,
+        use_adj=use_adj,
+        return_p2p=return_p2p,
+        n_jobs=n_jobs,
+        verbose=verbose,
+    )
 
     return result