adapted code to POT

Author: Kilian Fatras

examples/plot_stochastic.py

@@ -31,15 +31,17 @@
 
 n_source = 7
 n_target = 4
-eps = 1
-nb_iter = 10000
+reg = 1
+numItermax = 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, :])
+a = ot.utils.unif(n_source)
+b = ot.utils.unif(n_target)
+
+rng = np.random.RandomState(0)
+X_source = rng.randn(n_source, 2)
+Y_target = rng.randn(n_target, 2)
+M = ot.dist(X_source, Y_target)
 
 #############################################################################
 #
@@ -50,9 +52,8 @@
 # results.
 
 method = "SAG"
-sag_pi = ot.stochastic.transportation_matrix_entropic(method, eps, a, b, M,
-                                                      n_source, n_target,
-                                                      nb_iter, lr)
+sag_pi = ot.stochastic.transportation_matrix_entropic(a, b, M, reg, method,
+                                                      numItermax, lr)
 print(sag_pi)
 
 #############################################################################
@@ -66,15 +67,17 @@
 
 n_source = 7
 n_target = 4
-eps = 1
-nb_iter = 10000
-lr = 0.1
+reg = 1
+numItermax = 300000
+lr = 1
+
+a = ot.utils.unif(n_source)
+b = ot.utils.unif(n_target)
 
-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, :])
+rng = np.random.RandomState(0)
+X_source = rng.randn(n_source, 2)
+Y_target = rng.randn(n_target, 2)
+M = ot.dist(X_source, Y_target)
 
 #############################################################################
 #
@@ -86,9 +89,8 @@
 # results.
 
 method = "ASGD"
-asgd_pi = ot.stochastic.transportation_matrix_entropic(method, eps, a, b, M,
-                                                       n_source, n_target,
-                                                       nb_iter, lr)
+asgd_pi = ot.stochastic.transportation_matrix_entropic(a, b, M, reg, method,
+                                                       numItermax, lr)
 print(asgd_pi)
 
 #############################################################################

ot/__init__.py

@@ -18,6 +18,7 @@
 from . import datasets
 from . import da
 from . import gromov
+from . import stochastic
 
 # OT functions
 from .lp import emd, emd2