[MRG] Restructure and Augment Partial Gromov-Wasserstein solvers by cedricvincentcuaz · Pull Request #663 · PythonOT/POT

Types of changes

  • create duplicates of partial GW related functions from ot/partial.py to ot/gromov/_partial.py while adding a depreciation notes and warnings within ot/partial.py to maintain the API for now. This includes the following functions: partial_gromov_wasserstein, partial_gromov_wasserstein2, entropic_partial_gromov_wasserstein, entropic_partial_gromov_wasserstein2.

  • Add generic (hidden function) _transform_matrix returning transformed structure matrices for GW in ot/gromov/_utils.py (used to compute only once these transformations in pGW even though there is not constant terms a priori as with the GW problem) and use it to factor code in init_matrix and semirelaxed_init_matrix.

Then adapt solvers to mimic other GW related solvers in order to ease future integration of new solvers (e.g pFGW ones):

  • Make marginals p and q optional and set by default to uniform ones.
  • Add loss_fun parameter within ['square_loss', 'kl_loss'] -> now solvers support kl_loss too.
  • Add symmetry parameter within [None, True, False] -> now solvers support asymmetric structure matrices.
  • restructure implementation of solvers :
    • gradient and loss computations now only rely on the generic corresponding functions for the gw loss (gwloss, gwggrad), therefore functions gwgrad_partial and gwloss_partial in ot/partial.py will also be depreciated (kept for now with a de.
    • Implement generic partial CG solver partial_cg in ot/optim.py.
    • Improve docs for functions in ot/optim.py.
    • Call to solve_partial_gromov_linesearch for the exact line-search of pGW. Note that the latter requires the computation of the gradient of the regularizer at G (like other cg functions) but also at the cg direction Gc. This allows to deduce the next step gradient as a convex combinations of these previous gradients (should reduce the computation time of about 33 %), within the generic_cg solver. This trick will be implemented for all (F)GW based regularizer in a next PR.

Motivation and context / Related issue

How has this been tested (if it applies)

PR checklist

  • I have read the CONTRIBUTING document.
  • The documentation is up-to-date with the changes I made (check build artifacts).
  • All tests passed, and additional code has been covered with new tests.
  • I have added the PR and Issue fix to the RELEASES.md file.