code review1

Author: tvayer

examples/plot_barycenter_fgw.py

@@ -125,7 +125,11 @@ def graph_colors(nx_graph, vmin=0, vmax=7):
         colors.append(val_map[node])
     return colors
 
-#%% create dataset
+##############################################################################
+# Generate data
+# -------------
+
+#%% circular dataset
 # We build a dataset of noisy circular graphs.
 # Noise is added on the structures by random connections and on the features by gaussian noise.
 
@@ -135,7 +139,11 @@ def graph_colors(nx_graph, vmin=0, vmax=7):
 for k in range(9):
     X0.append(build_noisy_circular_graph(np.random.randint(15, 25), with_noise=True, structure_noise=True, p=3))
 
-#%% Plot dataset
+##############################################################################
+# Plot data
+# ---------
+
+#%% Plot graphs
 
 plt.figure(figsize=(8, 10))
 for i in range(len(X0)):
@@ -146,24 +154,28 @@ def graph_colors(nx_graph, vmin=0, vmax=7):
 plt.suptitle('Dataset of noisy graphs. Color indicates the label', fontsize=20)
 plt.show()
 
+##############################################################################
+# Barycenter computation
+# ----------------------
 
-#%%
-# We compute the barycenter using FGW. Structure matrices are computed using the shortest_path distance in the graph
+#%% We compute the barycenter using FGW. Structure matrices are computed using the shortest_path distance in the graph
 # Features distances are the euclidean distances
 Cs = [shortest_path(nx.adjacency_matrix(x)) for x in X0]
 ps = [np.ones(len(x.nodes())) / len(x.nodes()) for x in X0]
 Ys = [np.array([v for (k, v) in nx.get_node_attributes(x, 'attr_name').items()]).reshape(-1, 1) for x in X0]
 lambdas = np.array([np.ones(len(Ys)) / len(Ys)]).ravel()
 sizebary = 15  # we choose a barycenter with 15 nodes
 
-#%%
-
 A, C, log = fgw_barycenters(sizebary, Ys, Cs, ps, lambdas, alpha=0.95)
 
-#%%
+##############################################################################
+# Plot Barycenter
+# -------------------------
+
+#%% Create the barycenter
 bary = nx.from_numpy_matrix(sp_to_adjency(C, threshinf=0, threshsup=find_thresh(C, sup=100, step=100)[0]))
-for i in range(len(A.ravel())):
-    bary.add_node(i, attr_name=float(A.ravel()[i]))
+for i, v in enumerate(A.ravel()):
+    bary.add_node(i, attr_name=v)
 
 #%%
 pos = nx.kamada_kawai_layout(bary)

examples/plot_fgw.py

@@ -22,12 +22,16 @@
 import ot
 from ot.gromov import gromov_wasserstein, fused_gromov_wasserstein
 
+##############################################################################
+# Generate data
+# ---------
+
 #%% parameters
 # We create two 1D random measures
-n = 20
-n2 = 30
-sig = 1
-sig2 = 0.1
+n = 20  # number of points in the first distribution
+n2 = 30  # number of points in the second distribution
+sig = 1  # std of first distribution
+sig2 = 0.1  # std of second distribution
 
 np.random.seed(0)
 
@@ -43,6 +47,10 @@
 p = ot.unif(n)
 q = ot.unif(n2)
 
+##############################################################################
+# Plot data
+# ---------
+
 #%% plot the distributions
 
 pl.close(10)
@@ -64,15 +72,22 @@
 pl.tight_layout()
 pl.show()
 
+##############################################################################
+# Create structure matrices and across-feature distance matrix
+# ---------
 
 #%% Structure matrices and across-features distance matrix
 C1 = ot.dist(xs)
-C2 = ot.dist(xt).T
+C2 = ot.dist(xt)
 M = ot.dist(ys, yt)
 w1 = ot.unif(C1.shape[0])
 w2 = ot.unif(C2.shape[0])
 Got = ot.emd([], [], M)
 
+##############################################################################
+# Plot matrices
+# ---------
+
 #%%
 cmap = 'Reds'
 pl.close(10)
@@ -112,6 +127,9 @@
 ax3.set_aspect('auto')
 pl.show()
 
+##############################################################################
+# Compute FGW/GW
+# ---------
 
 #%% Computing FGW and GW
 alpha = 1e-3
@@ -123,6 +141,10 @@
 #%reload_ext WGW
 Gg, log = gromov_wasserstein(C1, C2, p, q, loss_fun='square_loss', verbose=True, log=True)
 
+##############################################################################
+# Visualize transport matrices
+# ---------
+
 #%% visu OT matrix
 cmap = 'Blues'
 fs = 15

ot/gromov.py

@@ -10,6 +10,7 @@
 #         Nicolas Courty <[email protected]>
 #         Rémi Flamary <[email protected]>
 #         Titouan Vayer <[email protected]>
+#
 # License: MIT License
 
 import numpy as np
@@ -351,9 +352,9 @@ def df(G):
         return cg(p, q, 0, 1, f, df, G0, armijo=armijo, C1=C1, C2=C2, constC=constC, **kwargs)
 
 
-def fused_gromov_wasserstein(M, C1, C2, p, q, loss_fun='square_loss', alpha=0.5, armijo=False, **kwargs):
+def fused_gromov_wasserstein(M, C1, C2, p, q, loss_fun='square_loss', alpha=0.5, armijo=False, log=False, **kwargs):
     """
-    Computes the FGW distance between two graphs see [3]
+    Computes the FGW transport between two graphs see [24]
     .. math::
         \gamma = arg\min_\gamma (1-\alpha)*<\gamma,M>_F + alpha* \sum_{i,j,k,l} L(C1_{i,k},C2_{j,l})*T_{i,j}*T_{k,l}
         s.t. \gamma 1 = p
@@ -377,7 +378,7 @@ def fused_gromov_wasserstein(M, C1, C2, p, q, loss_fun='square_loss', alpha=0.5,
          distribution in the source space
     q :  ndarray, shape (nt,)
          distribution in the target space
-    loss_fun :  string,optionnal
+    loss_fun :  string,optional
         loss function used for the solver
     max_iter : int, optional
         Max number of iterations
@@ -416,7 +417,86 @@ def f(G):
     def df(G):
         return gwggrad(constC, hC1, hC2, G)
 
-    return cg(p, q, M, alpha, f, df, G0, armijo=armijo, C1=C1, C2=C2, constC=constC, **kwargs)
+    if log:
+        res, log = cg(p, q, M, alpha, f, df, G0, armijo=armijo, C1=C1, C2=C2, constC=constC, log=True, **kwargs)
+        log['fgw_dist'] = log['loss'][::-1][0]
+        return res, log
+    else:
+        return cg(p, q, M, alpha, f, df, G0, armijo=armijo, C1=C1, C2=C2, constC=constC, **kwargs)
+
+
+def fused_gromov_wasserstein2(M, C1, C2, p, q, loss_fun='square_loss', alpha=0.5, armijo=False, log=False, **kwargs):
+    """
+    Computes the FGW distance between two graphs see [24]
+    .. math::
+        \gamma = arg\min_\gamma (1-\alpha)*<\gamma,M>_F + alpha* \sum_{i,j,k,l} L(C1_{i,k},C2_{j,l})*T_{i,j}*T_{k,l}
+        s.t. \gamma 1 = p
+             \gamma^T 1= q
+             \gamma\geq 0
+    where :
+    - M is the (ns,nt) metric cost matrix
+    - :math:`f` is the regularization term ( and df is its gradient)
+    - a and b are source and target weights (sum to 1)
+    - L is a loss function to account for the misfit between the similarity matrices
+    The algorithm used for solving the problem is conditional gradient as discussed in  [1]_
+    Parameters
+    ----------
+    M  : ndarray, shape (ns, nt)
+         Metric cost matrix between features across domains
+    C1 : ndarray, shape (ns, ns)
+         Metric cost matrix respresentative of the structure in the source space
+    C2 : ndarray, shape (nt, nt)
+         Metric cost matrix espresentative of the structure in the target space
+    p :  ndarray, shape (ns,)
+         distribution in the source space
+    q :  ndarray, shape (nt,)
+         distribution in the target space
+    loss_fun :  string,optional
+        loss function used for the solver
+    max_iter : int, optional
+        Max number of iterations
+    tol : float, optional
+        Stop threshold on error (>0)
+    verbose : bool, optional
+        Print information along iterations
+    log : bool, optional
+        record log if True
+    armijo : bool, optional
+        If True the steps of the line-search is found via an armijo research. Else closed form is used.
+        If there is convergence issues use False.
+    **kwargs : dict
+        parameters can be directly pased to the ot.optim.cg solver
+    Returns
+    -------
+    gamma : (ns x nt) ndarray
+        Optimal transportation matrix for the given parameters
+    log : dict
+        log dictionary return only if log==True in parameters
+    References
+    ----------
+    .. [24] Vayer Titouan, Chapel Laetitia, Flamary R{\'e}mi, Tavenard Romain
+          and Courty Nicolas
+        "Optimal Transport for structured data with application on graphs"
+        International Conference on Machine Learning (ICML). 2019.
+    """
+
+    constC, hC1, hC2 = init_matrix(C1, C2, p, q, loss_fun)
+
+    G0 = p[:, None] * q[None, :]
+
+    def f(G):
+        return gwloss(constC, hC1, hC2, G)
+
+    def df(G):
+        return gwggrad(constC, hC1, hC2, G)
+
+    res, log = cg(p, q, M, alpha, f, df, G0, armijo=armijo, C1=C1, C2=C2, constC=constC, log=True, **kwargs)
+    if log:
+        log['fgw_dist'] = log['loss'][::-1][0]
+        log['T'] = res
+        return log['fgw_dist'], log
+    else:
+        return log['fgw_dist']
 
 
 def gromov_wasserstein2(C1, C2, p, q, loss_fun, log=False, armijo=False, **kwargs):
@@ -889,7 +969,7 @@ def gromov_barycenters(N, Cs, ps, p, lambdas, loss_fun,
 
 def fgw_barycenters(N, Ys, Cs, ps, lambdas, alpha, fixed_structure=False, fixed_features=False,
                     p=None, loss_fun='square_loss', max_iter=100, tol=1e-9,
-                    verbose=False, log=True, init_C=None, init_X=None):
+                    verbose=False, log=False, init_C=None, init_X=None):
     """
     Compute the fgw barycenter as presented eq (5) in [24].
     ----------
@@ -919,7 +999,8 @@ def fgw_barycenters(N, Ys, Cs, ps, lambdas, alpha, fixed_structure=False, fixed_
         Barycenters' features
     C : ndarray, shape (N,N)
         Barycenters' structure matrix
-    log_:
+    log_: dictionary
+        Only returned when log=True
         T : list of (N,ns) transport matrices
         Ms : all distance matrices between the feature of the barycenter and the other features dist(X,Ys) shape (N,ns)
     References
@@ -1015,14 +1096,13 @@ class UndefinedParameter(Exception):
         T = [fused_gromov_wasserstein((1 - alpha) * Ms[s], C, Cs[s], p, ps[s], loss_fun, alpha, numItermax=max_iter, stopThr=1e-5, verbose=verbose) for s in range(S)]
 
         # T is N,ns
-
-        log_['Ts_iter'].append(T)
         err_feature = np.linalg.norm(X - Xprev.reshape(N, d))
         err_structure = np.linalg.norm(C - Cprev)
 
         if log:
             log_['err_feature'].append(err_feature)
             log_['err_structure'].append(err_structure)
+            log_['Ts_iter'].append(T)
 
         if verbose:
             if cpt % 200 == 0:
@@ -1032,11 +1112,15 @@ class UndefinedParameter(Exception):
             print('{:5d}|{:8e}|'.format(cpt, err_feature))
 
         cpt += 1
-    log_['T'] = T  # from target to Ys
-    log_['p'] = p
-    log_['Ms'] = Ms  # Ms are N,ns
+    if log:
+        log_['T'] = T  # from target to Ys
+        log_['p'] = p
+        log_['Ms'] = Ms  # Ms are N,ns
 
-    return X, C, log_
+    if log:
+        return X, C, log_
+    else:
+        return X, C
 
 
 def update_sructure_matrix(p, lambdas, T, Cs):

ot/optim.py

@@ -5,6 +5,7 @@
 
 # Author: Remi Flamary <[email protected]>
 #         Titouan Vayer <[email protected]>
+#
 # License: MIT License
 
 import numpy as np
@@ -88,20 +89,20 @@ def do_linesearch(cost, G, deltaG, Mi, f_val,
         Cost matrix of the linearized transport problem. Corresponds to the gradient of the cost
     f_val :  float
         Value of the cost at G
-    armijo : bool, optionnal
+    armijo : bool, optional
             If True the steps of the line-search is found via an armijo research. Else closed form is used.
             If there is convergence issues use False.
-    C1 : ndarray (ns,ns), optionnal
+    C1 : ndarray (ns,ns), optional
         Structure matrix in the source domain. Only used when armijo=False
-    C2 : ndarray (nt,nt), optionnal
+    C2 : ndarray (nt,nt), optional
         Structure matrix in the target domain. Only used when armijo=False
-    reg : float, optionnal
+    reg : float, optional
           Regularization parameter. Only used when armijo=False
     Gc : ndarray (ns,nt)
         Optimal map found by linearization in the FW algorithm. Only used when armijo=False
     constC : ndarray (ns,nt)
              Constant for the gromov cost. See [24]. Only used when armijo=False
-    M : ndarray (ns,nt), optionnal
+    M : ndarray (ns,nt), optional
         Cost matrix between the features. Only used when armijo=False
     Returns
     -------
@@ -223,9 +224,9 @@ def cost(G):
     it = 0
 
     if verbose:
-        print('{:5s}|{:12s}|{:8s}'.format(
-            'It.', 'Loss', 'Delta loss') + '\n' + '-' * 32)
-        print('{:5d}|{:8e}|{:8e}'.format(it, f_val, 0))
+        print('{:5s}|{:12s}|{:8s}|{:8s}'.format(
+            'It.', 'Loss', 'Relative loss', 'Absolute loss') + '\n' + '-' * 48)
+        print('{:5d}|{:8e}|{:8e}|{:8e}'.format(it, f_val, 0, 0))
 
     while loop:
 
@@ -261,8 +262,8 @@ def cost(G):
 
         if verbose:
             if it % 20 == 0:
-                print('{:5s}|{:12s}|{:8s}'.format(
-                    'It.', 'Loss', 'Relative variation loss', 'Absolute variation loss') + '\n' + '-' * 32)
+                print('{:5s}|{:12s}|{:8s}|{:8s}'.format(
+                    'It.', 'Loss', 'Relative loss', 'Absolute loss') + '\n' + '-' * 48)
             print('{:5d}|{:8e}|{:8e}|{:8e}'.format(it, f_val, relative_delta_fval, abs_delta_fval))
 
     if log:
@@ -363,9 +364,9 @@ def cost(G):
     it = 0
 
     if verbose:
-        print('{:5s}|{:12s}|{:8s}'.format(
-            'It.', 'Loss', 'Delta loss') + '\n' + '-' * 32)
-        print('{:5d}|{:8e}|{:8e}'.format(it, f_val, 0))
+        print('{:5s}|{:12s}|{:8s}|{:8s}'.format(
+            'It.', 'Loss', 'Relative loss', 'Absolute loss') + '\n' + '-' * 48)
+        print('{:5d}|{:8e}|{:8e}|{:8e}'.format(it, f_val, 0, 0))
 
     while loop:
 
@@ -402,8 +403,8 @@ def cost(G):
 
         if verbose:
             if it % 20 == 0:
-                print('{:5s}|{:12s}|{:8s}'.format(
-                    'It.', 'Loss', 'Relative variation loss', 'Absolute variation loss') + '\n' + '-' * 32)
+                print('{:5s}|{:12s}|{:8s}|{:8s}'.format(
+                    'It.', 'Loss', 'Relative loss', 'Absolute loss') + '\n' + '-' * 48)
             print('{:5d}|{:8e}|{:8e}|{:8e}'.format(it, f_val, relative_delta_fval, abs_delta_fval))
 
     if log: