[MRG] Restructure and Augment Partial Gromov-Wasserstein solvers (#663)

* init commit srgw bary

* init commit - restructure partial GW

* remove overlap with adjacent srGW PR

* update tests

* update old exemple

* up

* correct line-search + augment generic cg + complete tests

* fix issues from change in gcg

* trying to fix bugs with kwargs and args

* update optim file

* fix pep8

* updates

* completing tests partial gw

* complete partial tests

* fix pep8

* up

* fix prints in docs

* up

* fix precision tests

* fixing tests

* up

* update tests

* update tests

* up

* fasten partial cg

* fasten partial cg

* put back old partial gw functions

* fix partial

* improve doc of optim.py

* merge with master

* Update RELEASES.md

* improving doc for optim.py and adding breaking change to release.md

* improving doc for optim.py and adding breaking change to release.md

* tipos in doc

---------

Co-authored-by: Rémi Flamary <[email protected]>

Author: cedricvincentcuaz

Diff for: RELEASES.md

@@ -2,6 +2,9 @@
 
 ## 0.9.5dev
 
+#### Breaking change
+- Custom functions provided as parameter `line_search` to `ot.optim.generic_conditional_gradient` must now have the signature `line_search(cost, G, deltaG, Mi, cost_G, df_G, **kwargs)`, adding as input `df_G` the gradient of the regularizer evaluated at the transport plan `G`. This change aims at improving speed of solvers having quadratic polynomial functions as regularizer such as the Gromov-Wassertein loss (PR #663).
+
 #### New features
 - Added feature `mass=True` for `nx.kl_div` (PR #654)
 - Implemented Gaussian Mixture Model OT `ot.gmm` (PR #649)
@@ -13,6 +16,8 @@
 - Restructured `ot.unbalanced` module (PR #658)
 - Added `ot.unbalanced.lbfgsb_unbalanced2` and add flexible reference measure `c` in all unbalanced solvers (PR #658)
 - Implemented Fused unbalanced Gromov-Wasserstein and unbalanced Co-Optimal Transport (PR #677)
+- Notes before depreciating partial Gromov-Wasserstein function in `ot.partial` moved to ot.gromov  (PR #663)
+- Create `ot.gromov._partial` add new features `loss_fun = "kl_loss"` and `symmetry=False` to all solvers while increasing speed + updating adequatly `ot.solvers` (PR #663)
 - Added `ot.unbalanced.sinkhorn_unbalanced_translation_invariant` (PR #676)
 
 #### Closed issues

Diff for: examples/unbalanced-partial/plot_partial_wass_and_gromov.py

@@ -125,10 +125,10 @@
 # transport 100% of the mass
 print('------m = 1')
 m = 1
-res0, log0 = ot.partial.partial_gromov_wasserstein(C1, C2, p, q, m=m, log=True)
-res, log = ot.partial.entropic_partial_gromov_wasserstein(C1, C2, p, q, 10,
-                                                          m=m, log=True,
-                                                          verbose=True)
+res0, log0 = ot.gromov.partial_gromov_wasserstein(C1, C2, p, q, m=m, log=True)
+res, log = ot.gromov.entropic_partial_gromov_wasserstein(C1, C2, p, q, 10,
+                                                         m=m, log=True,
+                                                         verbose=True)
 
 print('Wasserstein distance (m = 1): ' + str(log0['partial_gw_dist']))
 print('Entropic Wasserstein distance (m = 1): ' + str(log['partial_gw_dist']))
@@ -146,11 +146,11 @@
 # transport 2/3 of the mass
 print('------m = 2/3')
 m = 2 / 3
-res0, log0 = ot.partial.partial_gromov_wasserstein(C1, C2, p, q, m=m, log=True,
-                                                   verbose=True)
-res, log = ot.partial.entropic_partial_gromov_wasserstein(C1, C2, p, q, 10,
-                                                          m=m, log=True,
-                                                          verbose=True)
+res0, log0 = ot.gromov.partial_gromov_wasserstein(C1, C2, p, q, m=m, log=True,
+                                                  verbose=True)
+res, log = ot.gromov.entropic_partial_gromov_wasserstein(C1, C2, p, q, 10,
+                                                         m=m, log=True,
+                                                         verbose=True)
 
 print('Partial Wasserstein distance (m = 2/3): ' +
       str(log0['partial_gw_dist']))

Diff for: ot/gromov/__init__.py

@@ -72,33 +72,55 @@
                           unbalanced_co_optimal_transport2,
                           fused_unbalanced_across_spaces_divergence)
 
+from ._partial import (partial_gromov_wasserstein,
+                       partial_gromov_wasserstein2,
+                       solve_partial_gromov_linesearch,
+                       entropic_partial_gromov_wasserstein,
+                       entropic_partial_gromov_wasserstein2)
+
+
 __all__ = ['init_matrix', 'tensor_product', 'gwloss', 'gwggrad',
            'init_matrix_semirelaxed', 'semirelaxed_init_plan',
            'update_barycenter_structure', 'update_barycenter_feature',
            'div_between_product', 'div_to_product', 'fused_unbalanced_across_spaces_cost',
            'uot_cost_matrix', 'uot_parameters_and_measures',
-           'gromov_wasserstein', 'gromov_wasserstein2', 'fused_gromov_wasserstein',
-           'fused_gromov_wasserstein2', 'solve_gromov_linesearch', 'gromov_barycenters',
-           'fgw_barycenters', 'entropic_gromov_wasserstein', 'entropic_gromov_wasserstein2',
-           'BAPG_gromov_wasserstein', 'BAPG_gromov_wasserstein2',
-           'entropic_gromov_barycenters', 'entropic_fused_gromov_wasserstein',
+           'gromov_wasserstein', 'gromov_wasserstein2',
+           'fused_gromov_wasserstein', 'fused_gromov_wasserstein2',
+           'solve_gromov_linesearch', 'gromov_barycenters',
+           'fgw_barycenters', 'entropic_gromov_wasserstein',
+           'entropic_gromov_wasserstein2', 'BAPG_gromov_wasserstein',
+           'BAPG_gromov_wasserstein2', 'entropic_gromov_barycenters',
+           'entropic_fused_gromov_wasserstein',
            'entropic_fused_gromov_wasserstein2', 'BAPG_fused_gromov_wasserstein',
            'BAPG_fused_gromov_wasserstein2', 'entropic_fused_gromov_barycenters',
-           'GW_distance_estimation', 'pointwise_gromov_wasserstein', 'sampled_gromov_wasserstein',
+           'GW_distance_estimation', 'pointwise_gromov_wasserstein',
+           'sampled_gromov_wasserstein',
            'semirelaxed_gromov_wasserstein', 'semirelaxed_gromov_wasserstein2',
-           'semirelaxed_fused_gromov_wasserstein', 'semirelaxed_fused_gromov_wasserstein2',
-           'solve_semirelaxed_gromov_linesearch', 'entropic_semirelaxed_gromov_wasserstein',
-           'entropic_semirelaxed_gromov_wasserstein2', 'entropic_semirelaxed_fused_gromov_wasserstein',
+           'semirelaxed_fused_gromov_wasserstein',
+           'semirelaxed_fused_gromov_wasserstein2',
+           'solve_semirelaxed_gromov_linesearch',
+           'entropic_semirelaxed_gromov_wasserstein',
+           'entropic_semirelaxed_gromov_wasserstein2',
+           'entropic_semirelaxed_fused_gromov_wasserstein',
            'entropic_semirelaxed_fused_gromov_wasserstein2',
            'semirelaxed_fgw_barycenters', 'semirelaxed_gromov_barycenters',
            'gromov_wasserstein_dictionary_learning',
-           'gromov_wasserstein_linear_unmixing', 'fused_gromov_wasserstein_dictionary_learning',
-           'fused_gromov_wasserstein_linear_unmixing', 'lowrank_gromov_wasserstein_samples',
-           'quantized_fused_gromov_wasserstein_partitioned', 'get_graph_partition',
-           'get_graph_representants', 'format_partitioned_graph',
-           'quantized_fused_gromov_wasserstein', 'get_partition_and_representants_samples',
-           'format_partitioned_samples', 'quantized_fused_gromov_wasserstein_samples',
-           'fused_unbalanced_gromov_wasserstein', 'fused_unbalanced_gromov_wasserstein2',
-           'unbalanced_co_optimal_transport', 'unbalanced_co_optimal_transport2',
-           'fused_unbalanced_across_spaces_divergence'
+           'gromov_wasserstein_linear_unmixing',
+           'fused_gromov_wasserstein_dictionary_learning',
+           'fused_gromov_wasserstein_linear_unmixing',
+           'lowrank_gromov_wasserstein_samples',
+           'quantized_fused_gromov_wasserstein_partitioned',
+           'get_graph_partition', 'get_graph_representants',
+           'format_partitioned_graph', 'quantized_fused_gromov_wasserstein',
+           'get_partition_and_representants_samples', 'format_partitioned_samples',
+           'quantized_fused_gromov_wasserstein_samples',
+           'fused_unbalanced_gromov_wasserstein',
+           'fused_unbalanced_gromov_wasserstein2',
+           'unbalanced_co_optimal_transport',
+           'unbalanced_co_optimal_transport2',
+           'fused_unbalanced_across_spaces_divergence',
+           'partial_gromov_wasserstein', 'partial_gromov_wasserstein2',
+           'solve_partial_gromov_linesearch',
+           'entropic_partial_gromov_wasserstein',
+           'entropic_partial_gromov_wasserstein2'
            ]

Diff for: ot/gromov/_gw.py

@@ -167,10 +167,10 @@ def df(G):
             return 0.5 * (gwggrad(constC, hC1, hC2, G, np_) + gwggrad(constCt, hC1t, hC2t, G, np_))
 
     if armijo:
-        def line_search(cost, G, deltaG, Mi, cost_G, **kwargs):
+        def line_search(cost, G, deltaG, Mi, cost_G, df_G, **kwargs):
             return line_search_armijo(cost, G, deltaG, Mi, cost_G, nx=np_, **kwargs)
     else:
-        def line_search(cost, G, deltaG, Mi, cost_G, **kwargs):
+        def line_search(cost, G, deltaG, Mi, cost_G, df_G, **kwargs):
             return solve_gromov_linesearch(G, deltaG, cost_G, hC1, hC2, M=0., reg=1., nx=np_, symmetric=symmetric, **kwargs)
 
     if not nx.is_floating_point(C10):
@@ -475,11 +475,12 @@ def df(G):
             return 0.5 * (gwggrad(constC, hC1, hC2, G, np_) + gwggrad(constCt, hC1t, hC2t, G, np_))
 
     if armijo:
-        def line_search(cost, G, deltaG, Mi, cost_G, **kwargs):
+        def line_search(cost, G, deltaG, Mi, cost_G, df_G, **kwargs):
             return line_search_armijo(cost, G, deltaG, Mi, cost_G, nx=np_, **kwargs)
     else:
-        def line_search(cost, G, deltaG, Mi, cost_G, **kwargs):
+        def line_search(cost, G, deltaG, Mi, cost_G, df_G, **kwargs):
             return solve_gromov_linesearch(G, deltaG, cost_G, hC1, hC2, M=(1 - alpha) * M, reg=alpha, nx=np_, symmetric=symmetric, **kwargs)
+
     if not nx.is_floating_point(M0):
         warnings.warn(
             "Input feature matrix consists of integer. The transport plan will be "
@@ -625,9 +626,11 @@ def fused_gromov_wasserstein2(M, C1, C2, p=None, q=None, loss_fun='square_loss',
     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))
+
     elif loss_fun == 'kl_loss':
         gC1 = nx.log(C1 + 1e-15) * nx.outer(p, p) - nx.dot(T, nx.dot(nx.log(C2 + 1e-15), T.T))
         gC2 = - nx.dot(T.T, nx.dot(C1, T)) / (C2 + 1e-15) + nx.outer(q, q)
+
     if isinstance(alpha, int) or isinstance(alpha, float):
         fgw_dist = nx.set_gradients(fgw_dist, (p, q, C1, C2, M),
                                     (log_fgw['u'] - nx.mean(log_fgw['u']),