[WIP] Restructure unbalanced OT module

RELEASES.md

@@ -10,6 +10,8 @@
 - Improved `ot.plot.plot1D_mat` (PR #649)
 - Added `nx.det` (PR #649)
 - `nx.sqrtm` is now broadcastable (takes ..., d, d) inputs (PR #649)
+- restructure `ot.unbalanced` module (PR #658)
+- add `ot.unbalanced.lbfgsb_unbalanced2` and add flexible reference measure `c` in all unbalanced solvers (PR #658)
 
 #### Closed issues
 - Fixed `ot.gaussian` ignoring weights when computing means (PR #649, Issue #648)

ot/__init__.py

@@ -73,5 +73,5 @@
            'factored_optimal_transport', 'solve', 'solve_gromov', 'solve_sample',
            'smooth', 'stochastic', 'unbalanced', 'partial', 'regpath', 'solvers',
            'binary_search_circle', 'wasserstein_circle',
-           'semidiscrete_wasserstein2_unif_circle', 'sliced_wasserstein_sphere_unif', 'lowrank_sinkhorn', 
+           'semidiscrete_wasserstein2_unif_circle', 'sliced_wasserstein_sphere_unif', 'lowrank_sinkhorn',
            'lowrank_gromov_wasserstein_samples']

ot/solvers.py

@@ -403,7 +403,7 @@ def solve(M, a=None, b=None, reg=None, reg_type="KL", unbalanced=None,
 
                 value_linear = nx.sum(M * plan)
 
-                value = log['loss']
+                value = log['cost']
 
             else:
                 raise (NotImplementedError('Not implemented reg_type="{}" and unbalanced_type="{}"'.format(reg_type, unbalanced_type)))

ot/unbalanced/__init__.py

@@ -0,0 +1,27 @@
+# -*- coding: utf-8 -*-
+"""
+Solvers related to Unbalanced Optimal Transport problems.
+
+"""
+
+# Author: Quang Huy Tran <[email protected]>
+#
+# License: MIT License
+
+# All submodules and packages
+from ._sinkhorn import (sinkhorn_knopp_unbalanced,
+                        sinkhorn_unbalanced,
+                        sinkhorn_stabilized_unbalanced,
+                        sinkhorn_unbalanced2,
+                        barycenter_unbalanced_sinkhorn,
+                        barycenter_unbalanced_stabilized,
+                        barycenter_unbalanced)
+
+from ._mm import (mm_unbalanced, mm_unbalanced2)
+
+from ._lbfgs import (lbfgsb_unbalanced, lbfgsb_unbalanced2)
+
+__all__ = ['sinkhorn_knopp_unbalanced', 'sinkhorn_unbalanced', 'sinkhorn_stabilized_unbalanced',
+           'sinkhorn_unbalanced2', 'barycenter_unbalanced_sinkhorn', 'barycenter_unbalanced_stabilized',
+           'barycenter_unbalanced', 'mm_unbalanced', 'mm_unbalanced2', '_get_loss_unbalanced',
+           'lbfgsb_unbalanced', 'lbfgsb_unbalanced2']