Introduction to Manifold Diffusion Geometry¶

This notebook follows the computation in Manifold Diffusion Geometry: Curvature, Tangent Spaces, and Dimension. The aim is to take a point cloud and compute familiar manifold geometry: tangent spaces, dimension, and scalar curvature.

We assume the data behave like points sampled from a low-dimensional manifold embedded in Euclidean space. We use diffusion geometry to estimate the intrinsic geometry from the heat flow on the data. The heat flow gives a Laplacian, and the Laplacian gives the carré du champ

$$ \Gamma(f,h)=g(\nabla f,\nabla h). $$

Once $\Gamma$ is known, Riemannian geometry becomes linear algebra.

The plots in this notebook are interactive

In [1]:
from pathlib import Path
import site
import sys

import numpy as np
from opt_einsum import contract
import plotly.io as pio

# Keep imports on this repo/conda env when running the notebook directly.
sys.path = [path for path in sys.path if not path.startswith(site.getusersitepackages())]
if Path.cwd().name == "intro_notebooks":
    sys.path.insert(0, "..")

from diffusion_geometry import DiffusionGeometry
from diffusion_geometry.visualisation import plot_scatter_3d, plot_2form_3d

pio.renderers.default = "notebook_connected"
In [2]:
def torus(minor_radius=1.0, major_radius=2.0, n=1000, noise=0.0, seed=0):
    rng = np.random.default_rng(seed)
    theta = rng.uniform(0, 2*np.pi, n)
    phi = rng.uniform(0, 2*np.pi, n)
    data = np.column_stack(((major_radius + minor_radius*np.cos(phi))*np.cos(theta),
                            (major_radius + minor_radius*np.cos(phi))*np.sin(theta),
                             minor_radius*np.sin(phi)))
    data += noise*rng.standard_normal(data.shape)
    return data, phi


def torus_scalar_curvature(phi, minor_radius=1.0, major_radius=2.0):
    return 2*np.cos(phi)/(minor_radius*(major_radius + minor_radius*np.cos(phi)))


def two_forms_from_normals(normals):
    """Ambient 2-form matrices whose Hodge duals are the given normal vectors."""
    forms = np.zeros((len(normals), 3, 3))
    # In 3D, plot_2form_3d draws disks orthogonal to the Hodge-dual vector.
    forms[:, 1, 2], forms[:, 2, 1] = -normals[:, 0], normals[:, 0]
    forms[:, 2, 0], forms[:, 0, 2] = -normals[:, 1], normals[:, 1]
    forms[:, 0, 1], forms[:, 1, 0] = -normals[:, 2], normals[:, 2]
    return forms


camera = dict(eye=dict(x=0.9, y=-1.15, z=0.85),
              center=dict(x=0, y=0, z=-0.05),
              up=dict(x=0, y=0, z=1))

torus_camera = dict(eye=dict(x=1.6, y=-1.7, z=1.1),
                    center=dict(x=0, y=0, z=0),
                    up=dict(x=0, y=0, z=1))

Sample some data¶

We start with a Swiss roll: a two-dimensional manifold embedded in three dimensions.

In [3]:
rng = np.random.default_rng(1)
n_samples = 1200
swiss_roll_parameter = 1.5*np.pi*(1 + 2*rng.random(n_samples))
height = 21*rng.random(n_samples)
data = np.column_stack((swiss_roll_parameter*np.cos(swiss_roll_parameter),
                        height,
                        swiss_roll_parameter*np.sin(swiss_roll_parameter)))
data = (data - data.mean(axis=0))/data.std(axis=0).max()

plot_scatter_3d(data, color=swiss_roll_parameter - swiss_roll_parameter.mean(), size=2, camera=camera).show()

First fundamental form¶

The gradients of the ambient coordinate functions span the tangent space. We use these to form the pointwise Gram matrix

$$ \Gamma_p(x_i,x_j). $$

For a clean $d$-dimensional manifold embedded isometrically, this matrix has $d$ eigenvalues equal to $1$ and the remaining eigenvalues equal to $0$. In data, the same pattern appears approximately. The large eigenvectors give a tangent frame and the small eigenvectors give a normal frame.

Diffusion maps recover the Laplacian only up to a constant scale. We correct that scale using the observation that the top intrinsic eigenvalues should average to one.

In [4]:
dg = DiffusionGeometry.from_point_cloud(data, n_function_basis=100)

# First fundamental form in ambient coordinates: gamma[p, i, j] = Gamma_p(x_i, x_j).
gamma = dg.cache.gamma_coords
eigenvalues, ambient_frame = np.linalg.eigh(gamma)
eigenvalues = eigenvalues[:, ::-1]
ambient_frame = ambient_frame[:, :, ::-1]

# The leading tangent eigenvalues should average to one.
metric_scale = np.median(eigenvalues[:, :2].mean(axis=1))
eigenvalues = eigenvalues/metric_scale

# The dimension is the position of the largest eigenvalue drop.
eigenvalue_gaps = -np.diff(eigenvalues, prepend=1.0, append=0.0, axis=1)
pointwise_dimension = dg._regularise(np.argmax(eigenvalue_gaps, axis=1).astype(float))
intrinsic_dimension = int(np.round(np.median(pointwise_dimension)))

# Split the ambient eigenframe into tangent and normal directions.
tangent_frame = ambient_frame[:, :, :intrinsic_dimension]
normal_frame = ambient_frame[:, :, intrinsic_dimension:]

print("global dimension:", intrinsic_dimension)
print("gamma shape:", gamma.shape)
print("median eigenvalues:", np.median(eigenvalues, axis=0))

plot_scatter_3d(data, color=pointwise_dimension.round(0) - 1, size=5, amax=2, camera=camera).show()

global dimension: 2
gamma shape: (1200, 3, 3)
median eigenvalues: [1.23463209 0.76275985 0.00861364]

Tangent planes¶

A tangent plane in three dimensions can be shown as an oriented 2-form: the Hodge dual of that 2-form is the normal vector. Here the normal frame from the previous cell becomes a small disk at each point, so the picture shows the estimated tangent bundle directly.

In [5]:
# The tangent plane is represented by its normal vector, converted to an ambient 2-form.
fig = plot_2form_3d(data,
                    two_forms_from_normals(normal_frame[:, :, 0]),
                    radius=0.1,
                    n_circle=18,
                    magnitude_scaling=False)
fig.show()

Second fundamental form and curvature¶

We can write the Hessian in terms of the iterated carré du champ:

$$ H(f)(\nabla h_1,\nabla h_2)=\frac12\big(\Gamma(h_1,\Gamma(h_2,f)) + \Gamma(h_2,\Gamma(h_1,f)) - \Gamma(f,\Gamma(h_1,h_2))\big). $$

The coordinate Hessians are already cached as dg.cache.hessian_coords. Projecting them onto the normal and tangent frames gives the second fundamental form

$$ \alpha^\ell_{ij}=H(n^\ell)(e_i,e_j). $$

The Gauss equation then gives the Riemann curvature tensor, and tracing gives Ricci and scalar curvature.

In [6]:
data_torus, phi = torus(n=1200, seed=0)
true_scalar = torus_scalar_curvature(phi)

dg = DiffusionGeometry.from_point_cloud(data_torus)

# Tangent and normal frames from the first fundamental form.
gamma = dg.cache.gamma_coords
eigenvalues, ambient_frame = np.linalg.eigh(gamma)
eigenvalues = eigenvalues[:, ::-1]
ambient_frame = ambient_frame[:, :, ::-1]

metric_scale = np.median(eigenvalues[:, :2].mean(axis=1))
tangent_frame = ambient_frame[:, :, :2]
normal_frame = ambient_frame[:, :, 2:]

# hessian[p, l, i, j] = Hess(x_l)(grad x_i, grad x_j), after metric-scale correction.
hessian = dg.cache.hessian_coords.transpose(0, 3, 1, 2)/metric_scale**2

# Project coordinate Hessians onto one normal direction and two tangent directions:
# alpha[p, normal_l, tangent_i, tangent_j] is the second fundamental form.
second_fundamental_form = contract("pstu,psl,pti,puj->plij",
                                   hessian,
                                   normal_frame,
                                   tangent_frame,
                                   tangent_frame)

# Gauss equation: R_ijkl = <alpha_ik, alpha_jl> - <alpha_jk, alpha_il>.
riemann = contract("pLik,pLjl->pijkl", second_fundamental_form, second_fundamental_form)
riemann -= contract("pLjk,pLil->pijkl", second_fundamental_form, second_fundamental_form)

# Trace the Riemann tensor to get Ricci, then trace Ricci to get scalar curvature.
ricci = contract("pkikj->pij", riemann)
scalar = dg._regularise(contract("pii->p", ricci))

print("riemann tensor shape:", riemann.shape)
print("ricci tensor shape:", ricci.shape)
print("scalar curvature shape:", scalar.shape)

print("\nTrue scalar curvature:")
plot_scatter_3d(data_torus, color=true_scalar, size=2, camera=torus_camera).show()
print("Estimated scalar curvature:")
plot_scatter_3d(data_torus, color=scalar, size=2, camera=torus_camera).show()

riemann tensor shape: (1200, 2, 2, 2, 2)
ricci tensor shape: (1200, 2, 2)
scalar curvature shape: (1200,)

True scalar curvature:

Estimated scalar curvature: