diff options
author | RĂ©mi Flamary <remi.flamary@gmail.com> | 2020-01-24 15:15:13 +0100 |
---|---|---|
committer | GitHub <noreply@github.com> | 2020-01-24 15:15:13 +0100 |
commit | da375139c75117a560546a16031787c7f0d50e74 (patch) | |
tree | 2ddafa7d8553534efa87693603d0f007ba7ccff7 /examples | |
parent | 1b58440457b25aace9dac56aa21144286e60f16e (diff) | |
parent | 6d4ccaca7705646c9b46b1d01a001943b6c778a9 (diff) |
Merge pull request #121 from mzalaya/master
[WIP] add screenkhorn: screening Sinkhorn algorithm
Diffstat (limited to 'examples')
-rw-r--r-- | examples/plot_screenkhorn_1D.py | 68 |
1 files changed, 68 insertions, 0 deletions
diff --git a/examples/plot_screenkhorn_1D.py b/examples/plot_screenkhorn_1D.py new file mode 100644 index 0000000..840ead8 --- /dev/null +++ b/examples/plot_screenkhorn_1D.py @@ -0,0 +1,68 @@ +# -*- coding: utf-8 -*- +""" +=============================== +1D Screened optimal transport +=============================== + +This example illustrates the computation of Screenkhorn: +Screening Sinkhorn Algorithm for Optimal transport. +""" + +# Author: Mokhtar Z. Alaya <mokhtarzahdi.alaya@gmail.com> +# +# License: MIT License + +import numpy as np +import matplotlib.pylab as pl +import ot.plot +from ot.datasets import make_1D_gauss as gauss +from ot.bregman import screenkhorn + +############################################################################## +# Generate data +# ------------- + +#%% parameters + +n = 100 # nb bins + +# bin positions +x = np.arange(n, dtype=np.float64) + +# Gaussian distributions +a = gauss(n, m=20, s=5) # m= mean, s= std +b = gauss(n, m=60, s=10) + +# loss matrix +M = ot.dist(x.reshape((n, 1)), x.reshape((n, 1))) +M /= M.max() + +############################################################################## +# Plot distributions and loss matrix +# ---------------------------------- + +#%% plot the distributions + +pl.figure(1, figsize=(6.4, 3)) +pl.plot(x, a, 'b', label='Source distribution') +pl.plot(x, b, 'r', label='Target distribution') +pl.legend() + +# plot distributions and loss matrix + +pl.figure(2, figsize=(5, 5)) +ot.plot.plot1D_mat(a, b, M, 'Cost matrix M') + +############################################################################## +# Solve Screenkhorn +# ----------------------- + +# Screenkhorn +lambd = 2e-03 # entropy parameter +ns_budget = 30 # budget number of points to be keeped in the source distribution +nt_budget = 30 # budget number of points to be keeped in the target distribution + +G_screen = screenkhorn(a, b, M, lambd, ns_budget, nt_budget, uniform=False, restricted=True, verbose=True) +pl.figure(4, figsize=(5, 5)) +ot.plot.plot1D_mat(a, b, G_screen, 'OT matrix Screenkhorn') +pl.show() |