SDOTSemi-Discrete Optimal Transport

Hello, World ​

The Wasserstein distance between 100 diracs and an image:

import numpy as np
import sdot

rng = np.random.default_rng( 0 )
f = sdot.SumOfDiracs( rng.random( ( 100, 2 ) ) )  # 100 diracs of equal mass
g = sdot.Image( rng.random( ( 10, 10 ) ) )  # a 10 x 10 image, on the unit square

print( sdot.w2( f, g ) )  # W2^2, as a numpy array

import jax
import sdot

k1, k2 = jax.random.split( jax.random.key( 0 ) )
f = sdot.SumOfDiracs( jax.random.uniform( k1, ( 100, 2 ) ) )  # 100 diracs of equal mass
g = sdot.Image( jax.random.uniform( k2, ( 10, 10 ) ) )  # a 10 x 10 image, on the unit square

print( sdot.w2( f, g ) )  # W2^2, as a jax array

import torch
import sdot

f = sdot.SumOfDiracs( torch.rand( 100, 2 ) )  # 100 diracs of equal mass
g = sdot.Image( torch.rand( 10, 10 ) )  # a 10 x 10 image, on the unit square

print( sdot.w2( f, g ) )  # W2^2, as a torch tensor

import cupy as cp
import sdot

rng = cp.random.default_rng( 0 )
f = sdot.SumOfDiracs( rng.random( ( 100, 2 ) ) )  # 100 diracs of equal mass
g = sdot.Image( rng.random( ( 10, 10 ) ) )  # a 10 x 10 image, on the unit square

print( sdot.w2( f, g ) )  # W2^2, as a cupy array, computed on the GPU

How it works ​

Given a discrete measure f ( a weighted sum of Dirac masses ) and a continuous density g, semi-discrete OT finds the power diagram ( Laguerre tessellation ) whose cells have the masses of the diracs:

The solution is unique, and a damped Newton method finds it in a handful of iterations, each one a power diagram plus a linear solve, both almost linear in the number of diracs. On an 8-core CPU, 10⁵ diracs in 2D take about 1 s; on a GPU, 0.07 s ( benchmarks ).

Reference: Kitagawa, Mérigot, Thibert — Convergence of a Newton algorithm for semi-discrete optimal transport, JEMS 2019.

For the ground metric, SDOT focuses essentially for now on variants of the euclidean norm.

What's next ​

• Getting started → installation, ...
• Performance → benchmarks against CGAL and geogram, CPU and GPU
• Examples gallery →
• Background → the mathematics
• Plans and settings → barycenters, moments, ...
• Distributions →: images, sum of gaussians, meshes, polytopes...
• Roadmap → what we are working on