Linearized OT

We have many point clouds, each drawn from a member of a family of shapes. The aims are to find the family back from the clouds alone, and then to say which member a new cloud comes from, in the Wasserstein sense.

Point clouds are not vectors. Draw twice from the same shape, and the $i$-th point of the first cloud has nothing to do with the $i$-th point of the second: a PCA of the coordinates of the points, or of their displacements, means nothing. What the clouds do share is a reference measure, here the Lebesgue measure $\rho$ of the unit square. Each cloud ${\mu }_k=\frac{1}{n}\sum _i{\delta }_{y_i^k}$ is the image of $\rho$ by its optimal transport map $T_k$, and these maps all live in the same space, $L^2(\rho )$:

$${\mu }_k=(T_k{)}_{\mathrm{\#}}\rho ,\parallel T_k-T_l{\parallel }_{L^2(\rho )}\ge W_2({\mu }_k,{\mu }_l).$$

This is linearized optimal transport (Wang et al., 2013). With a semi-discrete transport, $T_k$ is constant on each power cell: it sends the cell of $y_i^k$ to $y_i^k$. Close to a reference measure that is smooth enough, the distance between the maps is close to $W_2$ (Mérigot, Delalande and Chazal, Quantitative stability of optimal transport maps and linearization of the 2-Wasserstein space, 2020). Every linear tool then works on the maps: the mean, the PCA, least squares.

The map of a cloud

sdot.plan gives the weights $w_i$ of the power diagram, and a point $x$ of the square goes to the dirac of its cell, $\arg \underset{i}{min}|x-y_i{|}^2-w_i$. The map is read at 4096 points drawn uniformly in the square, once and for all: the same points for every cloud, so that the maps can be compared. That gives one vector of ${\mathbb{R}}^{4096\times 2}$ per cloud, and the mean over these points stands for the integral against $\rho$.

import numpy as np
import sdot

x = np.random.default_rng( 123 ).random( ( 4096, 2 ) )   # where every map is read

def transport_map( y ):
    """T( x ): the dirac of the power cell of each reading point"""
    w = np.asarray( sdot.plan( sdot.SumOfDiracs( y ), sdot.Box() ).weights )
    return y[ ( ( y * y ).sum( axis = 1 ) - w - 2 * x @ y.T ).argmin( axis = 1 ) ]

The argmin over all the diracs is the power cell of each reading point. For many diracs, the cell complex finds it with a tree: cells, _ = sdot.PowerDiagram( y, w, boundaries = ... ).cell_complex().locate( x ).

The family of the example is a ring, centered in the square. Its inner and outer boundaries each go from a circle ($s=0$) to a square ($s=1$), by interpolating the polar radius: two parameters, $s_{\text{in}}$ and $s_{\text{out}}$ in $[0,1]$. Below are the power cells of four draws (of 1000 diracs), at the four corners of the family. Use the slider at the top of the viewer to go from one draw to the next.

Open in a separate page ↗

The star in the middle is where the map tears. The square has no hole and the ring has one, so no continuous map goes from one to the other. The center of the square is split between all the diracs of the inner boundary. This tear decides most of what follows.

Learning the family

200 draws of ${10}^4$ diracs, with $(s_{\text{in}},s_{\text{out}})$ uniform in $[0,1{]}^2$, give 200 maps. The parameters are not used to learn anything. The PCA of the maps:

maps = np.stack( [ transport_map( y ).ravel() for y in draws ] )   # [ 200, 4096 * 2 ]
mean = maps.mean( axis = 0 )
_, sv, modes = np.linalg.svd( maps - mean, full_matrices = False )
coords = ( maps - mean ) @ modes[ :4 ].T

mode	1	2	3	4	5	6
share of the variance	0.616	0.081	0.066	0.058	0.016	0.012

The first mode is the outer boundary, $s_{\text{out}}$: it moves the most mass, the farthest. The inner boundary moves much less mass, and shares modes 2 to 4 with the tear, which the sampling noise moves from one draw to the next.

A point of the learned space is a measure: $(\overline{T}+\sum _k{c}_k{v}_k{)}_{\mathrm{\#}}\rho$, the image of the reading points by the reconstructed map. Below, the first two coordinates, the others at 0. Far from the training draws, the reconstruction extrapolates and folds onto itself: it is no longer a transport map. The sliders therefore span the range of the training draws, shrunk towards the mean until the reconstruction folds nowhere. The dashed lines are the circles and squares the family moves between. $c_1$ turns the outside from a circle into a square, while $c_2$ mostly pushes the hole aside.

Open in a separate page ↗

Identifying a new draw

A new cloud is identified by its map. Its 4 coordinates $c=V^T(T-\overline{T})$ give the closest member of the learned space, $(\overline{T}+Vc{)}_{\mathrm{\#}}\rho$. The distance from $T$ to its projection is an upper bound of $W_2$ between the cloud and this member. To check the result against the truth, and only for that, an affine least squares fit of the training coordinates onto $(s_{\text{in}},s_{\text{out}})$ reads the coordinates as parameters.

On 10 new draws, the mean error is 0.039 on $s_{\text{in}}$ and 0.007 on $s_{\text{out}}$, and the distance to the learned space ranges from 0.0079 to 0.0104. Each frame shows the diracs of a draw (grey), its reconstruction (red), and its true boundaries:

Open in a separate page ↗

Outside of the family

A ring whose outer square is turned by 45° is not a member of the family. Its distance to the learned space, 0.057, is more than five times that of the test draws (at most 0.0104), while its unturned twin is at 0.0097. Reading its coordinates as parameters would give $(0.68,-0.03)$, a meaningless answer that the distance flags. Use the slider to compare the twin and the turned square:

Open in a separate page ↗

The script

docs/examples/lot_dictionary.py computes everything on this page. It is NumPy only, with one sdot.plan per draw, and takes about three minutes on a laptop.

python docs/examples/lot_dictionary.py           # every page and table
python docs/examples/lot_dictionary.py 1d cells  # the quick ones