[MRG] New API for gromov solvers

RELEASES.md

@@ -8,6 +8,7 @@
 + Callbacks for generalized conditional gradient in `ot.da.sinkhorn_l1l2_gl` are now vectorized to improve performance (PR #507)
 + The `linspace` method of the backends now has the `type_as` argument to convert to the same dtype and device. (PR #533)
 + The `convolutional_barycenter2d` and `convolutional_barycenter2d_debiased` functions now work with different devices.. (PR #533)
++ New API for Gromov-Wasserstein solvers with `ot.solve_gromov` function (PR #536)
 
 #### Closed issues
 - Fix line search evaluating cost outside of the interpolation range (Issue #502, PR #504)

ot/__init__.py

@@ -50,7 +50,7 @@
                      gromov_barycenters, fused_gromov_wasserstein, fused_gromov_wasserstein2)
 from .weak import weak_optimal_transport
 from .factored import factored_optimal_transport
-from .solvers import solve
+from .solvers import solve, solve_gromov
 
 # utils functions
 from .utils import dist, unif, tic, toc, toq
@@ -65,7 +65,7 @@
            'sinkhorn_unbalanced2', 'sliced_wasserstein_distance', 'sliced_wasserstein_sphere',
            'gromov_wasserstein', 'gromov_wasserstein2', 'gromov_barycenters', 'fused_gromov_wasserstein',
            'fused_gromov_wasserstein2', 'max_sliced_wasserstein_distance', 'weak_optimal_transport',
-           'factored_optimal_transport', 'solve',
+           'factored_optimal_transport', 'solve', 'solve_gromov',
            'smooth', 'stochastic', 'unbalanced', 'partial', 'regpath', 'solvers',
            'binary_search_circle', 'wasserstein_circle',
            'semidiscrete_wasserstein2_unif_circle', 'sliced_wasserstein_sphere_unif']

ot/da.py

@@ -2274,6 +2274,7 @@ class NearestBrenierPotential(BaseTransport):
     ot.mapping.nearest_brenier_potential_fit : Fitting the SSNB on source and target data
     ot.mapping.nearest_brenier_potential_predict_bounds : Predicting SSNB images on new source data
     """
+
     def __init__(self, strongly_convex_constant=0.6, gradient_lipschitz_constant=1.4, log=False, its=100, seed=None):
         self.strongly_convex_constant = strongly_convex_constant
         self.gradient_lipschitz_constant = gradient_lipschitz_constant

ot/gromov/_bregman.py

@@ -330,6 +330,7 @@ def entropic_gromov_wasserstein2(
         learning for graph matching and node embedding. In International
         Conference on Machine Learning (ICML), 2019.
     """
+
     T, logv = entropic_gromov_wasserstein(
         C1, C2, p, q, loss_fun, epsilon, symmetric, G0, max_iter,
         tol, solver, warmstart, verbose, log=True, **kwargs)
@@ -815,12 +816,21 @@ def entropic_fused_gromov_wasserstein2(
         (ICML). 2019.
 
     """
+
+    nx = get_backend(M, C1, C2)
+
     T, logv = entropic_fused_gromov_wasserstein(
         M, C1, C2, p, q, loss_fun, epsilon, symmetric, alpha, G0, max_iter,
         tol, solver, warmstart, verbose, log=True, **kwargs)
 
     logv['T'] = T
 
+    lin_term = nx.sum(T * M)
+    gw_term = (logv['fgw_dist'] - (1 - alpha) * lin_term) / alpha
+
+    logv['quad_loss'] = gw_term * alpha
+    logv['lin_loss'] = lin_term * (1 - alpha)
+
     if log:
         return logv['fgw_dist'], logv
     else:

ot/gromov/_gw.py

@@ -582,6 +582,13 @@ def fused_gromov_wasserstein2(M, C1, C2, p=None, q=None, loss_fun='square_loss',
     fgw_dist = log_fgw['fgw_dist']
     log_fgw['T'] = T
 
+    # compute separate terms for gradients and log
+    lin_term = nx.sum(T * M)
+    gw_term = (fgw_dist - (1 - alpha) * lin_term) / alpha
+
+    log_fgw['quad_loss'] = gw_term * alpha
+    log_fgw['lin_loss'] = lin_term * (1 - alpha)
+
     if loss_fun == 'square_loss':
         gC1 = 2 * C1 * nx.outer(p, p) - 2 * nx.dot(T, nx.dot(C2, T.T))
         gC2 = 2 * C2 * nx.outer(q, q) - 2 * nx.dot(T.T, nx.dot(C1, T))
@@ -591,8 +598,7 @@ def fused_gromov_wasserstein2(M, C1, C2, p=None, q=None, loss_fun='square_loss',
                                          log_fgw['v'] - nx.mean(log_fgw['v']),
                                          alpha * gC1, alpha * gC2, (1 - alpha) * T))
         else:
-            lin_term = nx.sum(T * M)
-            gw_term = (fgw_dist - (1 - alpha) * lin_term) / alpha
+
             fgw_dist = nx.set_gradients(fgw_dist, (p, q, C1, C2, M, alpha),
                                         (log_fgw['u'] - nx.mean(log_fgw['u']),
                                          log_fgw['v'] - nx.mean(log_fgw['v']),

ot/gromov/_semirelaxed.py

@@ -488,6 +488,8 @@ def semirelaxed_fused_gromov_wasserstein2(M, C1, C2, p=None, loss_fun='square_lo
     q = nx.sum(T, 0)
     srfgw_dist = log_fgw['srfgw_dist']
     log_fgw['T'] = T
+    log_fgw['lin_loss'] = nx.sum(M * T) * (1 - alpha)
+    log_fgw['quad_loss'] = srfgw_dist - log_fgw['lin_loss']
 
     if loss_fun == 'square_loss':
         gC1 = 2 * C1 * nx.outer(p, p) - 2 * nx.dot(T, nx.dot(C2, T.T))
@@ -979,7 +981,9 @@ def df(G):
     if log:
         qG = nx.sum(G, 0)
         marginal_product = nx.outer(ones_p, nx.dot(qG, fC2t))
-        log['srfgw_dist'] = alpha * gwloss(constC + marginal_product, hC1, hC2, G, nx) + (1 - alpha) * nx.sum(M * G)
+        log['lin_loss'] = nx.sum(M * G) * (1 - alpha)
+        log['quad_loss'] = alpha * gwloss(constC + marginal_product, hC1, hC2, G, nx)
+        log['srfgw_dist'] = log['lin_loss'] + log['quad_loss']
         return G, log
     else:
         return G