import sys
import os
sys.path.append(os.path.abspath('..'))  # Add parent directory to path

import numpy as np
from diffusion_geometry.visualisation import *
from plotly.subplots import make_subplots
from figures.generate_data import load_image_point_cloud
from diffusion_geometry import DiffusionGeometry

camera = dict(eye=dict(x=0.8, y=-0.9, z=1.),
                center=dict(x=0, y=0, z=-0.2),
                up=dict(x=0, y=0, z=1))
# rescale eye:
camera['eye'] = {k: v * 1.3 for k, v in camera['eye'].items()}

Connection Laplacian

The Levi-Civita connection defines an operator $\nabla: \mathfrak{X}(M) \to \Omega^{0,2}(M)$ from a vector field $X$ to its derivative $\nabla X$, which is a $(0,2)$-tensor.

The composition of this map with its adjoint defines an operator $\nabla^*\nabla: \mathfrak{X}(M) \to \mathfrak{X}(M)$ called the connection Laplacian (or the 'Bochner' or 'rough' Laplacian). Like the Hodge Laplacian, we can diagonalise it to find 'smooth' vector fields. This is the problem addressed by 'Vector Diffusion Maps' (Singer and Wu, 2012). We can straightforwardly implement it with diffusion geometry.

from diffusion_geometry import DiffusionGeometry
from figures.generate_data import gen_3d_data

# Generate torus data:
n = 2000
data = gen_3d_data(kind='torus', minor_radius=1.0, major_radius=2.0, n=n, noise=0.0, seed=0)[0]

# Vector diffusion maps:
dg = DiffusionGeometry.from_point_cloud(data)
connection = dg.levi_civita.adjoint @ dg.levi_civita
eigenvalues, eigenvectors = connection.spectrum()

# Plot the first 3 eigenvectors in a single row:
fig = make_subplots(
    rows=1,
    cols=3,
    specs=[[{"type": "scene"}, {"type": "scene"}, {"type": "scene"}]],
    horizontal_spacing=0.02,
)

for k in range(3):
    fig_k = plot_quiver_3d(data, eigenvectors[k].to_ambient(), scale=0.1, camera=camera, arrow_scale=0.6)
    fig.add_traces(list(fig_k.data), rows=[1] * len(fig_k.data), cols=[k + 1] * len(fig_k.data))

scene_layout = dict(camera=camera, xaxis=dict(visible=False), yaxis=dict(visible=False), zaxis=dict(visible=False))
fig.update_layout(
    scene=scene_layout,
    scene2=scene_layout,
    scene3=scene_layout,
    margin=dict(l=0, r=0, t=0, b=0),
)

fig.show()