Plans and Settings

Three functions, from the most common to the most specific:

	returns	when
sdot.w2( f, g )	$W_2^2$, a scalar	distance (a loss to differentiate, ...)
sdot.plan( f, g )	the transport plan	distance plus barycenters, cell masses, weights, picture
sdot.ot_solve( f, g, ... )	the transport plan	the same, with the settings of the solver

f is the discrete source ( a SumOfDiracs ), g the continuous target ( an Image, a SumOfGaussians... see Distributions ). The cost is the squared Euclidean distance.

w2: the distance

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

The result is an array of the framework of the inputs ( numpy here, see Frameworks and devices ), and it is differentiable with respect to everything that defines f and g.

plan: what the transport does

import numpy as np
import sdot

rng = np.random.default_rng( 0 )
f = sdot.SumOfDiracs( rng.random( ( 100, 2 ) ) )
g = sdot.Image( rng.random( ( 10, 10 ) ) )

p = sdot.plan( f, g )

p.cost                # W2^2, the same value as `sdot.w2( f, g )`
p.barycenters         # [ 100, 2 ]: the barycenter of each cell
p.cell_masses         # [ 100 ]: the mass of each cell ( the mass of its dirac )
p.weights             # [ 100 ]: the weights of the power diagram
p.converged           # True

p.write_html( "plan.html" )    # the cells, the seeds and the density, in a self-contained page

The plan sends the cell of each dirac, in the power diagram of weights, onto that dirac. The barycenters are what the Lloyd algorithm moves the diracs to ( see the Lloyd quantization example ).

p.transport() gives ( cost, barycenters, cell_masses ) at once, and p.power_diagram() the diagram itself, to inspect its cells.

ot_solve: the settings of the solver

ot_solve computes the same plan, and takes the settings of the solver as keywords:

import numpy as np
import sdot

rng = np.random.default_rng( 0 )
positions = rng.random( ( 100, 2 ) )
f = sdot.SumOfDiracs( positions )
g = sdot.Image( rng.random( ( 10, 10 ) ) )

p = sdot.ot_solve( f, g, tol = 1e-12, max_iter = 200 )   # the same plan, with the settings of the solver
print( p.stats[ "nb_iter" ], p.stats[ "residual" ] )

# a nearby problem restarts from a previous plan: a few Newton steps instead of a few dozen
f2 = sdot.SumOfDiracs( positions + 0.01 )
p2 = sdot.ot_solve( f2, g, ot_plan = p )
print( p2.stats[ "nb_iter" ] )

In 2D and more, the solver is a damped Newton method ( Kitagawa, Mérigot, Thibert ), and its settings are:

setting	default
tol	1e-8	stop when every cell mass is within tol of its target ( the masses are normalized, the target of a dirac is 1 / n )
max_iter	100	Newton steps, at most
ot_plan	None	a previous plan to start from
precision	"auto"	"fp64", "fp32", or "mixed" on a GPU ( "auto": fp64 on a CPU, mixed on a GPU )
continuation	"auto"	solve first for a smoothed density: "auto" ( when a cell starts empty ), "always", "never"
on_failure	"raise"	"raise" an OtNotConverged ( which carries the plan reached ), or "ignore" and read p.converged

p.stats tells how it went: nb_iter, residual, status...

In 1D the plan is exact ( a sort ): there is neither a tolerance nor iterations, and the only setting is with_barycenters:

import numpy as np
import sdot

rng = np.random.default_rng( 0 )
f = sdot.SumOfDiracs( rng.random( ( 1000, 1 ) ) )            # diracs on a line
g = sdot.Image( rng.random( 100 ) )                         # 100 pixels on [ 0, 1 ]

p = sdot.ot_solve( f, g, with_barycenters = True )            # 1D: a sort, exact, no tolerance

OtProblem: solving again and again

In a loop ( a descent, a reconstruction, Lloyd ), keep the problem and change its inputs. Each solve starts from the previous solution, which is what makes a nearby transport cost a few Newton steps:

import numpy as np
import sdot

rng = np.random.default_rng( 0 )
positions = rng.random( ( 100, 2 ) )

pb = sdot.OtProblem( sdot.SumOfDiracs( positions ), sdot.Image( rng.random( ( 10, 10 ) ) ) )
for _ in range( 10 ):
    p = pb.solve()                                       # restarts from the previous solution by itself
    positions = p.barycenters
    pb.source = sdot.SumOfDiracs( positions )

pb.solve( sdot.Iterative( tol = ... ) ) takes the settings of ot_solve ( sdot.Direct( ... ) in 1D ).