pep8

Author: rflamary

ot/gpu/bregman.py

@@ -8,12 +8,11 @@
 #
 # License: MIT License
 
-import cupy as np # np used for matrix computation
-import cupy as cp # cp used for cupy specific operations
+import cupy as np  # np used for matrix computation
+import cupy as cp  # cp used for cupy specific operations
 from . import utils
 
 
-
 def sinkhorn_knopp(a, b, M, reg, numItermax=1000, stopThr=1e-9,
                    verbose=False, log=False, to_numpy=True, **kwargs):
     """
@@ -159,7 +158,7 @@ def sinkhorn_knopp(a, b, M, reg, numItermax=1000, stopThr=1e-9,
                     np.sum((v - vprev)**2) / np.sum((v)**2)
             else:
                 # compute right marginal tmp2= (diag(u)Kdiag(v))^T1
-                tmp2=np.sum(u[:,None]*K*v[None,:],0)
+                tmp2 = np.sum(u[:, None] * K * v[None, :], 0)
                 #tmp2=np.einsum('i,ij,j->j', u, K, v)
                 err = np.linalg.norm(tmp2 - b)**2  # violation of marginal
             if log:
@@ -177,24 +176,25 @@ def sinkhorn_knopp(a, b, M, reg, numItermax=1000, stopThr=1e-9,
 
     if nbb:  # return only loss
         #res = np.einsum('ik,ij,jk,ij->k', u, K, v, M) (explodes cupy memory)
-        res=np.empty(nbb)
+        res = np.empty(nbb)
         for i in range(nbb):
-            res[i]=np.sum(u[:,None,i]*(K*M)*v[None,:,i])
+            res[i] = np.sum(u[:, None, i] * (K * M) * v[None, :, i])
         if to_numpy:
-            res=utils.to_np(res)            
+            res = utils.to_np(res)
         if log:
             return res, log
         else:
             return res
 
     else:  # return OT matrix
-        res=u.reshape((-1, 1)) * K * v.reshape((1, -1))
+        res = u.reshape((-1, 1)) * K * v.reshape((1, -1))
         if to_numpy:
-            res=utils.to_np(res)
+            res = utils.to_np(res)
         if log:
             return res, log
         else:
             return res
 
+
 # define sinkhorn as sinkhorn_knopp
-sinkhorn=sinkhorn_knopp
+sinkhorn = sinkhorn_knopp

ot/gpu/da.py

@@ -10,15 +10,16 @@
 #
 # License: MIT License
 
-import cupy as np # np used for matrix computation
-import cupy as cp # cp used for cupy specific operations
+import cupy as np  # np used for matrix computation
+import cupy as cp  # cp used for cupy specific operations
 import numpy as npp
 from . import utils
 from .bregman import sinkhorn
 
+
 def sinkhorn_lpl1_mm(a, labels_a, b, M, reg, eta=0.1, numItermax=10,
                      numInnerItermax=200, stopInnerThr=1e-9, verbose=False,
-                     log=False,to_numpy=True):
+                     log=False, to_numpy=True):
     """
     Solve the entropic regularization optimal transport problem with nonconvex
     group lasso regularization
@@ -101,15 +102,14 @@ def sinkhorn_lpl1_mm(a, labels_a, b, M, reg, eta=0.1, numItermax=10,
     ot.optim.cg : General regularized OT
 
     """
-    
+
     a, labels_a, b, M = utils.to_gpu(a, labels_a, b, M)
-    
-    
+
     p = 0.5
     epsilon = 1e-3
 
     indices_labels = []
-    labels_a2=cp.asnumpy(labels_a)
+    labels_a2 = cp.asnumpy(labels_a)
     classes = npp.unique(labels_a2)
     for c in classes:
         idxc, = utils.to_gpu(npp.where(labels_a2 == c))
@@ -120,7 +120,7 @@ def sinkhorn_lpl1_mm(a, labels_a, b, M, reg, eta=0.1, numItermax=10,
     for cpt in range(numItermax):
         Mreg = M + eta * W
         transp = sinkhorn(a, b, Mreg, reg, numItermax=numInnerItermax,
-                          stopThr=stopInnerThr,to_numpy=False)
+                          stopThr=stopInnerThr, to_numpy=False)
         # the transport has been computed. Check if classes are really
         # separated
         W = np.ones(M.shape)

ot/gpu/utils.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 """
-Utility functions for GPU 
+Utility functions for GPU
 """
 
 # Author: Remi Flamary <[email protected]>
@@ -9,9 +9,8 @@
 #
 # License: MIT License
 
-import cupy as np # np used for matrix computation
-import cupy as cp # cp used for cupy specific operations
-
+import cupy as np  # np used for matrix computation
+import cupy as cp  # cp used for cupy specific operations
 
 
 def euclidean_distances(a, b, squared=False, to_numpy=True):
@@ -34,23 +33,24 @@ def euclidean_distances(a, b, squared=False, to_numpy=True):
     c : (n x m) np.ndarray or cupy.ndarray
         pairwise euclidean distance distance matrix
     """
-    
+
     a, b = to_gpu(a, b)
-    
-    a2=np.sum(np.square(a),1)
-    b2=np.sum(np.square(b),1)
-    
-    c=-2*np.dot(a,b.T)
-    c+=a2[:,None]
-    c+=b2[None,:]
-    
+
+    a2 = np.sum(np.square(a), 1)
+    b2 = np.sum(np.square(b), 1)
+
+    c = -2 * np.dot(a, b.T)
+    c += a2[:, None]
+    c += b2[None, :]
+
     if not squared:
         np.sqrt(c, out=c)
     if to_numpy:
         return to_np(c)
     else:
         return c
 
+
 def dist(x1, x2=None, metric='sqeuclidean', to_numpy=True):
     """Compute distance between samples in x1 and x2 on gpu
 
@@ -61,8 +61,8 @@ def dist(x1, x2=None, metric='sqeuclidean', to_numpy=True):
         matrix with n1 samples of size d
     x2 : np.array (n2,d), optional
         matrix with n2 samples of size d (if None then x2=x1)
-    metric : str  
-        Metric from 'sqeuclidean', 'euclidean', 
+    metric : str
+        Metric from 'sqeuclidean', 'euclidean',
 
 
     Returns
@@ -80,7 +80,6 @@ def dist(x1, x2=None, metric='sqeuclidean', to_numpy=True):
         return euclidean_distances(x1, x2, squared=False, to_numpy=to_numpy)
     else:
         raise NotImplementedError
-    
 
 
 def to_gpu(*args):
@@ -91,10 +90,9 @@ def to_gpu(*args):
         return cp.asarray(args[0])
 
 
-
 def to_np(*args):
     """ convert GPU arras to numpy and return them"""
     if len(args) > 1:
         return (cp.asnumpy(x) for x in args)
     else:
-        return cp.asnumpy(args[0])
+        return cp.asnumpy(args[0])