LDDMM: parallel transport

Imagine we are given three meshes A, B, C.

In this notebook we will:

• register B to A, and B to C

• parallel transport BC along BA

import shutil
import string

import numpy as np
import pyvista as pv

import polpo.utils as putils
from polpo.plot.pyvista import RegisteredMeshesGifPlotter
from polpo.preprocessing.mesh.registration import RigidAlignment
from polpo.surface_mesh.deformetrica import (
    LddmmMetric,
    Point,
    RegistrationDir,
    ShootDir,
    TransportDir,
)
from polpo.surface_mesh.generation.blob import create_blob

[KeOps] Warning : CUDA was detected, but driver API could not be initialized. Switching to CPU only.

RECOMPUTE = False

# NB: fix seed before setting it to False
np.random.seed(42)

STATIC_VIZ = True

if STATIC_VIZ:
    pv.set_jupyter_backend("static")

OUTPUTS_DIR = putils.get_results_path() / "transport_abc_blob_example"

if OUTPUTS_DIR.exists() and RECOMPUTE:
    shutil.rmtree(OUTPUTS_DIR)

Generate meshes

n_meshes = 3
bump_amp = 0.2

raw_meshes = [
    create_blob(resolution=10, bump_amp=bump_amp, n_bumps=5, smoothing_iter=10)
    for _ in range(n_meshes)
]

raw_meshes[0].points.shape

(82, 3)

pl = pv.Plotter(border=False)

for mesh in raw_meshes:
    pl.add_mesh(mesh, show_edges=True, opacity=0.5)

pl.show()

prep_pipe = RigidAlignment(known_correspondences=True)

meshes = prep_pipe(raw_meshes)

pl = pv.Plotter(border=False)

for mesh in meshes:
    pl.add_mesh(mesh, show_edges=True, opacity=0.5)

pl.show()

Registrations

kernel_width = 2 * bump_amp
registration_kwargs = dict(
    kernel_width=kernel_width,
    regularisation=1.0,
    max_iter=2000,
    freeze_control_points=False,
    metric="varifold",
    tol=1e-16,
    attachment_kernel_width=bump_amp,
)

metric = LddmmMetric(OUTPUTS_DIR, use_pole_ladder=True, **registration_kwargs)

point_a, point_b, point_c = [
    Point(
        id_=string.ascii_uppercase[index],
        pv_surface=mesh,
        dirname=metric.dir_config.meshes_dir,
    )
    for index, mesh in enumerate(meshes)
]

Closely following LDDMM: how to register a mesh to a template?, we register B to A and B to C.

vec_ba = metric.log(point_a, point_b)
vec_bc = metric.log(point_c, point_b)

for vec in [vec_ba, vec_bc]:
    reg_dir = RegistrationDir.from_dirname(vec.dirname)

    pl = RegisteredMeshesGifPlotter()
    pl.add_title(f"{reg_dir.base_point.id} -> {reg_dir.point.id}")

    flow_meshes = [pt.as_pv() for pt in reg_dir.flow()]
    pl.add_meshes(flow_meshes)
    pl.close()

    pl.show()

Parallel transport (pole ladder)

# NB: way slower?
trans_vec_bc_pole = metric.parallel_transport(vec_bc, point_b, direction=vec_ba)

This is how the deformation B->C looks from A.

trans_point_c_pole = metric.exp(trans_vec_bc_pole, point_a)

shoot_dir_pole = ShootDir.from_dirname(trans_point_c_pole.vtk_path.parent)

pl = RegisteredMeshesGifPlotter()

pl.add_title("A + BC")
pl.add_meshes([pt.as_pv() for pt in shoot_dir_pole.flow()])

pl.close()

pl.show()

Parallel transport (fanning)

metric.use_pole_ladder = False

trans_vec_bc_fan = metric.parallel_transport(vec_bc, point_b, direction=vec_ba)

This is how the deformation B->C looks from A.

trans_point_c_fan = metric.exp(trans_vec_bc_fan, point_a)

shoot_dir_fan = ShootDir.from_dirname(trans_point_c_fan.vtk_path.parent)

pl = RegisteredMeshesGifPlotter()

pl.add_title("A + BC")
flow_meshes = [pt.as_pv() for pt in shoot_dir_fan.flow()]
pl.add_meshes(flow_meshes)

pl.close()

pl.show()

Some comparisons

Extra information is computed with the fanning scheme:

• an approximation of the end point of direction

• a shooted mesh from that approximation using parallel transport

trans_dir_bc_fan = TransportDir.from_dirname(trans_vec_bc_fan.dirname)

pl = pv.Plotter()

pl.add_mesh(
    point_a.as_pv(),
    color="red",
    opacity=0.5,
    label="truth",
)
pl.add_mesh(
    trans_dir_bc_fan.reconstructed().as_pv(),
    color="green",
    opacity=0.5,
    label="fan-rec",
)
pl.add_mesh(
    RegistrationDir.from_dirname(vec_ba.dirname).reconstructed().as_pv(),
    color="yellow",
    opacity=0.5,
    label="registration-rec",
)


pl.add_legend()

pl.show()

pl = pv.Plotter()

pl.add_mesh(
    trans_dir_bc_fan.reconstructed_shooted().as_pv(),
    color="red",
    opacity=0.5,
    label="fan-rec-shoot",
)
pl.add_mesh(
    trans_point_c_fan.as_pv(),
    color="green",
    opacity=0.5,
    label="fan-shoot",
)
pl.add_mesh(
    trans_point_c_pole.as_pv(),
    color="yellow",
    opacity=0.5,
    label="pole-shoot",
)

pl.add_legend()

pl.show()