add tests and examples

Author: Kilian Fatras

examples/plot_stochastic.py

@@ -0,0 +1,133 @@
+"""
+==========================
+Stochastic examples
+==========================
+
+This example is designed to show how to use the stochatic optimization
+algorithms for descrete and semicontinous measures from the POT library.
+
+"""
+
+# Author: Kilian Fatras <kilian.fatras@gmail.com>
+#
+# License: MIT License
+
+import matplotlib.pylab as pl
+import numpy as np
+import ot
+
+
+#############################################################################
+# COMPUTE TRANSPORTATION MATRIX
+#############################################################################
+
+#############################################################################
+# DISCRETE CASE
+# Sample two discrete measures for the discrete case
+# ---------------------------------------------
+#
+# Define 2 discrete measures a and b, the points where are defined the source
+# and the target measures and finally the cost matrix c.
+
+n_source = 7
+n_target = 4
+eps = 1
+nb_iter = 10000
+lr = 0.1
+
+a = (1./n_source) * np.ones(n_source)
+b = (1./n_target) * np.ones(n_target)
+X_source = np.arange(n_source)
+Y_target = np.arange(0, 2 * n_target, 2)
+M = np.abs(X_source[:, None] - Y_target[None, :])
+
+#############################################################################
+#
+# Call the "SAG" method to find the transportation matrix in the discrete case
+# ---------------------------------------------
+#
+# Define the method "SAG", call ot.transportation_matrix_entropic and plot the
+# results.
+
+method = "SAG"
+sag_pi = ot.stochastic.transportation_matrix_entropic(method, eps, a, b, M,
+                                                      n_source, n_target,
+                                                      nb_iter, lr)
+print(sag_pi)
+
+#############################################################################
+# SEMICONTINOUS CASE
+# Sample one general measure a, one discrete measures b for the semicontinous
+# case
+# ---------------------------------------------
+#
+# Define one general measure a, one discrete measures b, the points where
+# are defined the source and the target measures and finally the cost matrix c.
+
+n_source = 7
+n_target = 4
+eps = 1
+nb_iter = 10000
+lr = 0.1
+
+a = (1./n_source) * np.ones(n_source)
+b = (1./n_target) * np.ones(n_target)
+X_source = np.arange(n_source)
+Y_target = np.arange(0, 2 * n_target, 2)
+M = np.abs(X_source[:, None] - Y_target[None, :])
+
+#############################################################################
+#
+# Call the "ASGD" method to find the transportation matrix in the semicontinous
+# case
+# ---------------------------------------------
+#
+# Define the method "ASGD", call ot.transportation_matrix_entropic and plot the
+# results.
+
+method = "ASGD"
+asgd_pi = ot.stochastic.transportation_matrix_entropic(method, eps, a, b, M,
+                                                       n_source, n_target,
+                                                       nb_iter, lr)
+print(asgd_pi)
+
+#############################################################################
+#
+# Compare the results with the Sinkhorn algorithm
+# ---------------------------------------------
+#
+# Call the Sinkhorn algorithm from POT
+
+sinkhorn_pi = ot.sinkhorn(a, b, M, 1)
+print(sinkhorn_pi)
+
+
+##############################################################################
+# PLOT TRANSPORTATION MATRIX
+##############################################################################
+
+##############################################################################
+# Plot SAG results
+# ----------------
+
+pl.figure(4, figsize=(5, 5))
+ot.plot.plot1D_mat(a, b, sag_pi, 'OT matrix SAG')
+pl.show()
+
+
+##############################################################################
+# Plot ASGD results
+# -----------------
+
+pl.figure(4, figsize=(5, 5))
+ot.plot.plot1D_mat(a, b, asgd_pi, 'OT matrix ASGD')
+pl.show()
+
+
+##############################################################################
+# Plot Sinkhorn results
+# ---------------------
+
+pl.figure(4, figsize=(5, 5))
+ot.plot.plot1D_mat(a, b, sinkhorn_pi, 'OT matrix Sinkhorn')
+pl.show()