update screenkhorn example

2 files changed, 23 insertions, 51 deletions

diff --git a/examples/plot_screenkhorn_1D.py b/examples/plot_screenkhorn_1D.py
index 22a9ddc..0eb64b0 100644
--- a/examples/plot_screenkhorn_1D.py
+++ b/examples/plot_screenkhorn_1D.py
@@ -1,41 +1,26 @@
-#!/usr/bin/env python
-# coding: utf-8
-
-# In[ ]:
-
-
-from ot.bregman import screenkhorn
-from ot.datasets import make_1D_gauss as gauss
-import ot.plot
-import ot
-import matplotlib.pylab as pl
-import numpy as np
-get_ipython().run_line_magic('matplotlib', 'inline')
-
-
-#
-# # 1D Screened optimal transport
-#
-#
-# This example illustrates the computation of Screenkhorn: Screening Sinkhorn Algorithm for Optimal transport.
-#
-#
-
-# In[13]:
+# -*- 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
 # -------------
-#
-#
-
-# In[14]:
-
 
 #%% parameters
 
@@ -52,14 +37,9 @@ b = gauss(n, m=60, s=10)
 M = ot.dist(x.reshape((n, 1)), x.reshape((n, 1)))
 M /= M.max()
 
-
+##############################################################################
 # Plot distributions and loss matrix
 # ----------------------------------
-#
-#
-
-# In[15]:
-
 
 #%% plot the distributions
 
@@ -73,14 +53,9 @@ pl.legend()
 pl.figure(2, figsize=(5, 5))
 ot.plot.plot1D_mat(a, b, M, 'Cost matrix M')
 
-
+##############################################################################
 # Solve Screened Sinkhorn
-# --------------
-#
-#
-
-# In[21]:
-
+# -----------------------
 
 # Screenkhorn
 
@@ -90,9 +65,6 @@ nt_budget = 30  # budget number of points to be keeped in the target distributio
 
 Gsc = 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, Gs, 'OT matrix Screenkhorn')
-
-pl.show()
-
+ot.plot.plot1D_mat(a, b, Gsc, 'OT matrix Screenkhorn')
 
-# In[ ]:
+pl.show()
\ No newline at end of file
diff --git a/ot/bregman.py b/ot/bregman.py
index 4f24cf4..95b27e4 100644
--- a/ot/bregman.py
+++ b/ot/bregman.py
@@ -1793,13 +1793,13 @@ def empirical_sinkhorn_divergence(X_s, X_t, reg, a=None, b=None, metric='sqeucli
 
 def screenkhorn(a, b, M, reg, ns_budget=None, nt_budget=None, uniform=False, restricted=True,
                 maxiter=10000, maxfun=10000, pgtol=1e-09, verbose=False, log=False):
-    """"
+    r""""
     Screening Sinkhorn Algorithm for Regularized Optimal Transport
 
     The function solves an approximate dual of Sinkhorn divergence [2] which is written as the following optimization problem:
 
     ..math::
-      (u, v) = \argmin_{u, v} 1_{ns}.T B(u,v) 1_{nt} - <\kappa u, a> - <v/\kappa, b> 
+      (u, v) = \argmin_{u, v} 1_{ns}^\top B(u,v) 1_{nt} - <\kappa u, a> - <v/\kappa, b>
 
       where B(u,v) = \diag(e^u) K \diag(e^v), with K = e^{-M/reg} and 
 
@@ -1853,8 +1853,8 @@ def screenkhorn(a, b, M, reg, ns_budget=None, nt_budget=None, uniform=False, res
 
     Dependency
     ----------
-    To gain more efficiency, screenkhorn needs to call the "Bottleneck" package (https://pypi.org/project/Bottleneck/) in the screening pre-processing step. 
-    If Bottleneck isn't installed, the following error message appears: 
+    To gain more efficiency, screenkhorn needs to call the "Bottleneck" package (https://pypi.org/project/Bottleneck/)
+    in the screening pre-processing step. If Bottleneck isn't installed, the following error message appears:
     "Bottleneck module doesn't exist. Install it from https://pypi.org/project/Bottleneck/"