add test for barycenter

1 files changed, 7 insertions, 5 deletions

diff --git a/ot/lp/cvx.py b/ot/lp/cvx.py
index c62da6a..fe9ac76 100644
--- a/ot/lp/cvx.py
+++ b/ot/lp/cvx.py
@@ -39,7 +39,9 @@ def barycenter(A, M, weights=None, verbose=False, log=False, solver='interior-po
     - :math:`\mathbf{a}_i` are training distributions in the columns of matrix :math:`\mathbf{A}`
 
     The linear program is solved using the interior point solver from scipy.optimize.
-    If cvxopt solver if installed it can use cvxopt.
+    If cvxopt solver if installed it can use cvxopt
+
+    Note that this problem do not scale well (both in memory and computational time).
 
     Parameters
     ----------
@@ -114,14 +116,14 @@ def barycenter(A, M, weights=None, verbose=False, log=False, solver='interior-po
     A_eq = sps.vstack((A_eq1, A_eq2))
     b_eq = np.concatenate((b_eq1, b_eq2))
 
-    if not cvxopt or solver in ['interior-point']:  
+    if not cvxopt or solver in ['interior-point']:
         # cvxopt not installed or interior point
 
         if solver is None:
             solver = 'interior-point'
 
         options = {'sparse': True, 'disp': verbose}
-        sol = sp.optimize.linprog(c, A_eq=A_eq, b_eq=b_eq, method=solver, 
+        sol = sp.optimize.linprog(c, A_eq=A_eq, b_eq=b_eq, method=solver,
                                   options=options)
         x = sol.x
         b = x[-n:]
@@ -131,8 +133,8 @@ def barycenter(A, M, weights=None, verbose=False, log=False, solver='interior-po
         h = np.zeros((n_distributions * n2 + n))
         G = -sps.eye(n_distributions * n2 + n)
 
-        sol = solvers.lp(matrix(c), scipy_sparse_to_spmatrix(G), matrix(h), 
-                         A=scipy_sparse_to_spmatrix(A_eq), b=matrix(b_eq), 
+        sol = solvers.lp(matrix(c), scipy_sparse_to_spmatrix(G), matrix(h),
+                         A=scipy_sparse_to_spmatrix(A_eq), b=matrix(b_eq),
                          solver=solver)
 
         x = np.array(sol['x'])