Home Original page

[MRG] New API ot.solve_sample by rflamary · Pull Request #563 · PythonOT/POT

Types of changes

Implement the first shot as the ot.solve_sample API.

For standard parameters (the same as ot.solve) it is a simple wrapper that pr-compute the cost matrix using the metric parameter and return the same solutions. But this function also provides other approximated OT solvers that can be selected with method and large scale lazy solvers that avoid computation of the full cost matrix with lazy=True

Some examples of use below:

import numpy as np
import ot

n = 100
rng = np.random.RandomState(0)

x = rng.randn(n, 2)
x2 = rng.randn(n//2, 2)+5

#%% ot.solve_sample is a wrapper for ot.solve when lazy=False (and method==None)

M = ot.dist(x, x2, metric='sqeuclidean')
sol0  = ot.solve(M)

sol = ot.solve_sample(x,x2)
print(sol.value)

# use anothe metric
sol2 = ot.solve_sample(x,x2, metric='cityblock')
print(sol2.value)
# sol == sol0 for all parameters in ot.solve (juset a wrapper)

#%% other methods

# solve 1D wasserstein in paralel for all dimensions
sol = ot.solve_sample(x,x2, method='1d')
print(sol)
# sol.value return the wassretsein for each dimensions

# Compute the empirical squared Bures wassrestein distance
sol = ot.solve_sample(x,x2, method='gaussian')
print(sol)

# compute factored OT
sol = ot.solve_sample(x,x2, method='factored', rank=10)
print(sol) # sol.plan is returned if lazy=False
print(sol.lazy_plan) # the low rank lazy tensor is always available for factored OT

#%%  Sinkhorn and Lazy Sinkhorn

# Sinkhorn solution
sol0 = ot.solve_sample(x,x2, reg=1)
print(sol0)

# Lazy sinkhorn (O(n) in memory)
sol = ot.solve_sample(x,x2, reg=1, lazy=True)
print(sol)
print(sol.lazy_plan) # the low rank lazy tensor is always available for lazy sinkhorn

Motivation and context / Related issue

How has this been tested (if it applies)

PR checklist

  • I have read the CONTRIBUTING document.
  • The documentation is up-to-date with the changes I made (check build artifacts).
  • All tests passed, and additional code has been covered with new tests.
  • I have added the PR and Issue fix to the RELEASES.md file.