Getting Started

Installation

pip install sdot

Alternatively, to get the latest version (or to participate), one can use the sources:

git clone git@github.com:sdot-team/sdot.git
cd sdot
pip install -e .

Framework

SDOT can work with a lot of different frameworks.

If you have not installed any (independently of SDOT), here are the links for the official installation instructions for some of them (one is enough):

• JAX
• PyTorch
• CuPy
• NumPy

GPU support — follow the GPU installation instructions of the chosen framework for your platform (CUDA, …). SDOT will automatically use the devices your tensors are on, but some frameworks need specific installation instructions.

Compilers

The kernels for the common cases come precompiled with the package. Less common cases are compiled on first use. For such a case, you will need ninja, a C++ compiler, and nvcc for NVIDIA GPUs.

Hello, World

Here is a simple program to test if the installation works: the Wasserstein 2 distance between 100 diracs and an image.

NumPy

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

JAX

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

PyTorch

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

CuPy

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

sdot.plan( f, g ) gives the whole transport plan, and p.show() displays it ( inline in a notebook, in the browser otherwise ):

Open in a separate page ↗

What's Next

• Performance → what it costs, CPU and GPU
• Background → the mathematics
• Examples gallery → some use cases
• Plans and settings → barycenters, moments, ...
• Distributions → the sources and the targets
• Frameworks and devices → how sdot handles the frameworks