POT icon indicating copy to clipboard operation
POT copied to clipboard

[WIP] Restructure unbalanced OT module

Open 6Ulm opened this issue 7 months ago • 1 comments

Types of changes

This PR aims to improve the current module unbalanced.py on two aspects:

  • The current module contains 3 seperate groups of unbalanced OT solvers: Sinkhorn, Majorization-Minimization and LBFGS-B. It is preferable to break these groups into seperate files/modules for easier maintenance.
  • In the current Sinkhorn solver, the reference measure used in the regulazation is always fixed as $a b^T$ (for reg_type=kl) or one-matrix (for reg_type=entropy). With trivial modification, we can allow for any nonnegative reference measure. This will be useful for the ot.solve, since we consider a very generic unbalanced OT problem.

Major changes:

  • The reg_type argument is removed from Sinkhorn solver and replace by the reference c in the regularization term (c = ab^T by default)
  • Add ot.unbalanced.lbfgsb_unbalanced2, which returns the UOT loss. This is consistent with almost all solvers in POT.
  • Remove option "entropy" for argument reg_div in lbfgsb_unbalanced.

Motivation and context / Related issue

How has this been tested (if it applies)

PR checklist

  • [x] I have read the CONTRIBUTING document.
  • [x] The documentation is up-to-date with the changes I made (check build artifacts).
  • [x] 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.

6Ulm avatar Jul 16 '24 15:07 6Ulm