clean

1 files changed, 22 insertions, 6 deletions

diff --git a/examples/plot_screenkhorn_1D.py b/examples/plot_screenkhorn_1D.py
index 0eb64b0..103d54c 100644
--- a/examples/plot_screenkhorn_1D.py
+++ b/examples/plot_screenkhorn_1D.py
@@ -14,6 +14,7 @@ Screening Sinkhorn Algorithm for Optimal transport.
 
 import numpy as np
 import matplotlib.pylab as pl
+import time
 import ot.plot
 from ot.datasets import make_1D_gauss as gauss
 from ot.bregman import screenkhorn
@@ -54,17 +55,32 @@ pl.figure(2, figsize=(5, 5))
 ot.plot.plot1D_mat(a, b, M, 'Cost matrix M')
 
 ##############################################################################
-# Solve Screened Sinkhorn
+# Solve Screenkhorn
 # -----------------------
 
 # Screenkhorn
-
-lambd = 1e-2  # entropy parameter
+lambd = 1e-3  # 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
 
-Gsc = screenkhorn(a, b, M, lambd, ns_budget, nt_budget, uniform=False, restricted=True, verbose=True)
+tic = time.time()
+G_screen = screenkhorn(a, b, M, lambd, ns_budget, nt_budget, uniform=False, restricted=True, verbose=True)
+tac_screen = time.time() - tic
+
+# Sinkhorn
+tic = time.time()
+G_sink = ot.sinkhorn(a, b, M, lambd, verbose=False)
+tac_sink = time.time() - tic
+
+
 pl.figure(4, figsize=(5, 5))
-ot.plot.plot1D_mat(a, b, Gsc, 'OT matrix Screenkhorn')
+ot.plot.plot1D_mat(a, b, G_screen, 'OT matrix Screenkhorn')
 
-pl.show()
\ No newline at end of file
+pl.show()
+
+##############################################################################
+# Time complexity
+# -----------------------
+print("Sinkhorn time complexity: %s\n" % tac_sink)
+print("Screenkhorn time complexity: %s\n" % tac_screen)
+print("Time_Sinkhorn / Time_Screenkhorn: %s\n" % (tac_sink / tac_screen))