[WIP] Wrapper for OT barycenter solvers with free support

[cedricvincentcuaz]

Types of changes

• creation of the class ot.utils.BaryResult extending OTResult

• Long term objectives: Wrap all barycenter solvers available in POT and more, ideally establishing links between these wrappers and those for inner OT solvers ot.solve_sample, ot.solve, and ot.solve_gromov - respectively associated with their barycenter solver ot.solve_bary_sample (free support); ot.solve_bary (free masses); ot.solve_bary_gromov (free joint support)

• a) ot.solvers.solve_bary_sample systematically optimizes the support of the barycenter using BCD with closed forms if the parameter metric is in ('sqeuclidean', 'euclidean') calling an hidden function _bary_sample_bcd which works with most OT problems in ot.solve_sample. metric can also be a callable or a list of callable function following the API of lp.free_support_barycenter_generic_costs for balanced OT only.

• b) handle variations in masses across OT problems only for unbalanced problems - keeping masses unchanged for exact and inexact sinkhorn like solvers - via parameter update_masses in ot.solvers._bary_sample_bcd.

• c) Handle given initialization of transport plans & potentials + warmstart strategies : added parameter warmstart to solve_bary_sample, and based on the requested inner OT solver, we pick the adequate strategy warmstart_plan or warmstart_potentials passed to _bary_sample_bcd .

• d) Options for stopping criterion between losses and barycenter variations: added parameter stopping_criterion in ('loss', 'bary') passed to _bary_sample_bcd.

Extensions let for future works:

• a) include lazy inner OT solvers referenced in the list lst_method_lazy (in ot.solvers)

How has this been tested (if it applies)

• a) test.test_solvers.assert_allclose_bary_sol method to compare solutions as BaryResult objects.
• b) test_solve_bary_sample grid across all available barycenter solvers depending on those available in ot.solvers.solve_sample.
• c) test_solve_bary_sample_NotImplemented test error messages for methods not implemented yet.
• d) test_solve_bary_sample_ValueError Test ValueError cases including stopping_criterion and X_b_init shape.
• e) test_solve_bary_sample_callable_metric test barycenter solvers for balanced OT with callable metrics.

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.

[codecov]

Codecov Report

❌ Patch coverage is 96.47303% with 17 lines in your changes missing coverage. Please review.
✅ Project coverage is 96.79%. Comparing base (44fd091) to head (b0f51d2).

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #730      +/-   ##
==========================================
- Coverage   96.82%   96.79%   -0.04%     
==========================================
  Files         124      127       +3     
  Lines       24536    24831     +295     
==========================================
+ Hits        23757    24035     +278     
- Misses        779      796      +17     

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[rflamary]

Some quick comments. Will have a more detailed look when solvers.py is split

[rflamary]

Maybe plot first the barycenter and belwo the barycenter + the OT plan? this is a bit dens and hard to parse visually.

[rflamary]

Thanks @cedricvincentcuaz

I took the time to go over the detailed implementation and it is very nice but found a few things. First this solver should not change the b of the bary (expect for the weird variant returned from generic) so the sum of marginal you update in the algorithm shod not be b but only a normalization term when computing barycenters.

few other comments below

[rflamary]

you should do it in two pats: first sharp vs entropic baryenter (easier to compare) and then a second section where you vary the marginal violation weight so that we see the transformation from exact bary to points on the left and right (as we can see there)