pep8 normalization

2 files changed, 70 insertions, 68 deletions

diff --git a/examples/plot_convolutional_barycenter.py b/examples/plot_convolutional_barycenter.py
index d231da9..7ccdbe3 100644
--- a/examples/plot_convolutional_barycenter.py
+++ b/examples/plot_convolutional_barycenter.py
@@ -1,4 +1,4 @@
-    
+
 #%%
 # -*- coding: utf-8 -*-
 """
@@ -32,24 +32,24 @@ f3 = 1 - pl.imread('../data/heart.png')[:, :, 2]
 f4 = 1 - pl.imread('../data/tooth.png')[:, :, 2]
 
 A = []
-f1=f1/np.sum(f1)
-f2=f2/np.sum(f2)
-f3=f3/np.sum(f3)
-f4=f4/np.sum(f4)
+f1 = f1 / np.sum(f1)
+f2 = f2 / np.sum(f2)
+f3 = f3 / np.sum(f3)
+f4 = f4 / np.sum(f4)
 A.append(f1)
 A.append(f2)
 A.append(f3)
 A.append(f4)
-A=np.array(A)
+A = np.array(A)
 
 nb_images = 5
 
 # those are the four corners coordinates that will be interpolated by bilinear
 # interpolation
-v1=np.array((1,0,0,0))
-v2=np.array((0,1,0,0))
-v3=np.array((0,0,1,0))
-v4=np.array((0,0,0,1))
+v1 = np.array((1, 0, 0, 0))
+v2 = np.array((0, 1, 0, 0))
+v3 = np.array((0, 0, 1, 0))
+v4 = np.array((0, 0, 0, 1))
 
 
 ##############################################################################
@@ -57,36 +57,36 @@ v4=np.array((0,0,0,1))
 # ----------------------------------------
 #
 
-pl.figure(figsize=(10,10))
+pl.figure(figsize=(10, 10))
 pl.title('Convolutional Wasserstein Barycenters in POT')
-cm='Blues'
+cm = 'Blues'
 # regularization parameter
-reg=0.004
+reg = 0.004
 for i in range(nb_images):
     for j in range(nb_images):
-        pl.subplot(nb_images,nb_images,i*nb_images+j+1)
-        tx=float(i)/(nb_images-1)
-        ty=float(j)/(nb_images-1)
-        
+        pl.subplot(nb_images, nb_images, i * nb_images + j + 1)
+        tx = float(i) / (nb_images - 1)
+        ty = float(j) / (nb_images - 1)
+
         # weights are constructed by bilinear interpolation
-        tmp1=(1-tx)*v1+tx*v2
-        tmp2=(1-tx)*v3+tx*v4
-        weights=(1-ty)*tmp1+ty*tmp2
-        
-        if i==0 and j==0:
-           pl.imshow(f1,cmap=cm)
-           pl.axis('off')
-        elif i==0 and j==(nb_images-1):
-           pl.imshow(f3,cmap=cm)
-           pl.axis('off')
-        elif i==(nb_images-1) and j==0:
-           pl.imshow(f2,cmap=cm)
-           pl.axis('off')
-        elif i==(nb_images-1) and j==(nb_images-1):
-           pl.imshow(f4,cmap=cm)
-           pl.axis('off')
+        tmp1 = (1 - tx) * v1 + tx * v2
+        tmp2 = (1 - tx) * v3 + tx * v4
+        weights = (1 - ty) * tmp1 + ty * tmp2
+
+        if i == 0 and j == 0:
+            pl.imshow(f1, cmap=cm)
+            pl.axis('off')
+        elif i == 0 and j == (nb_images - 1):
+            pl.imshow(f3, cmap=cm)
+            pl.axis('off')
+        elif i == (nb_images - 1) and j == 0:
+            pl.imshow(f2, cmap=cm)
+            pl.axis('off')
+        elif i == (nb_images - 1) and j == (nb_images - 1):
+            pl.imshow(f4, cmap=cm)
+            pl.axis('off')
         else:
             # call to barycenter computation
-            pl.imshow(ot.convolutional_barycenter2d(A,reg,weights),cmap=cm)
+            pl.imshow(ot.convolutional_barycenter2d(A, reg, weights), cmap=cm)
             pl.axis('off')
-pl.show()
\ No newline at end of file
+pl.show()
diff --git a/ot/bregman.py b/ot/bregman.py
index 05f4d9d..f844f03 100644
--- a/ot/bregman.py
+++ b/ot/bregman.py
@@ -918,7 +918,8 @@ def barycenter(A, M, reg, weights=None, numItermax=1000,
     else:
         return geometricBar(weights, UKv)
 
-def convolutional_barycenter2d(A,reg,weights=None,numItermax = 10000, stopThr=1e-9, verbose=False, log=False):
+
+def convolutional_barycenter2d(A, reg, weights=None, numItermax=10000, stopThr=1e-9, verbose=False, log=False):
     """Compute the entropic regularized wasserstein barycenter of distributions A
        where A is a collection of 2D images. 
 
@@ -979,51 +980,52 @@ def convolutional_barycenter2d(A,reg,weights=None,numItermax = 10000, stopThr=1e
     if log:
         log = {'err': []}
 
-    b=np.zeros_like(A[0,:,:])
-    U=np.ones_like(A)
-    KV=np.ones_like(A)
-    threshold = 1e-30 # in order to avoids numerical precision issues
+    b = np.zeros_like(A[0, :, :])
+    U = np.ones_like(A)
+    KV = np.ones_like(A)
+    threshold = 1e-30  # in order to avoids numerical precision issues
 
     cpt = 0
-    err=1
-    
-    # build the convolution operator 
-    t = np.linspace(0,1,A.shape[1])
-    [Y,X] = np.meshgrid(t,t)
-    xi1 = np.exp(-(X-Y)**2/reg)
-    K = lambda x: np.dot(np.dot(xi1,x),xi1)
-
-    while (err>stopThr and cpt<numItermax):
-        
-        bold=b
-        cpt = cpt +1
-        
-        b=np.zeros_like(A[0,:,:])
+    err = 1
+
+    # build the convolution operator
+    t = np.linspace(0, 1, A.shape[1])
+    [Y, X] = np.meshgrid(t, t)
+    xi1 = np.exp(-(X - Y)**2 / reg)
+
+    def K(x): return np.dot(np.dot(xi1, x), xi1)
+
+    while (err > stopThr and cpt < numItermax):
+
+        bold = b
+        cpt = cpt + 1
+
+        b = np.zeros_like(A[0, :, :])
         for r in range(A.shape[0]):
-            KV[r,:,:]=K(A[r,:,:]/np.maximum(threshold,K(U[r,:,:])))
-            b += weights[r] * np.log(np.maximum(threshold, U[r,:,:]*KV[r,:,:]))
+            KV[r, :, :] = K(A[r, :, :] / np.maximum(threshold, K(U[r, :, :])))
+            b += weights[r] * np.log(np.maximum(threshold, U[r, :, :] * KV[r, :, :]))
         b = np.exp(b)
         for r in range(A.shape[0]):
-            U[r,:,:]=b/np.maximum(threshold,KV[r,:,:])
-    
-        if cpt%10==1:
-            err=np.sum(np.abs(bold-b))
+            U[r, :, :] = b / np.maximum(threshold, KV[r, :, :])
+
+        if cpt % 10 == 1:
+            err = np.sum(np.abs(bold - b))
             # log and verbose print
             if log:
                 log['err'].append(err)
 
             if verbose:
-                if cpt%200 ==0:
-                    print('{:5s}|{:12s}'.format('It.','Err')+'\n'+'-'*19)
-                print('{:5d}|{:8e}|'.format(cpt,err))
+                if cpt % 200 == 0:
+                    print('{:5s}|{:12s}'.format('It.', 'Err') + '\n' + '-' * 19)
+                print('{:5d}|{:8e}|'.format(cpt, err))
 
     if log:
-        log['niter']=cpt
-        log['U']=U
-        return b,log
+        log['niter'] = cpt
+        log['U'] = U
+        return b, log
     else:
-        return b   
-    
+        return b
+
 
 def unmix(a, D, M, M0, h0, reg, reg0, alpha, numItermax=1000,
           stopThr=1e-3, verbose=False, log=False):