improve doc

3 files changed, 28 insertions, 25 deletions

diff --git a/docs/source/quickstart.rst b/docs/source/quickstart.rst
index 9729664..978eaff 100644
--- a/docs/source/quickstart.rst
+++ b/docs/source/quickstart.rst
@@ -577,10 +577,10 @@ Unbalanced optimal transport
 
 Unbalanced OT is a relaxation of the entropy regularized OT problem where the violation of
 the constraint on the marginals is added to the objective of the optimization
-problem. The unbalanced OT metric between two histograms a and b is defined as [25]_ [10]_:
+problem. The unbalanced OT metric between two unbalanced histograms a and b is defined as [25]_ [10]_:
 
 .. math::
-    W_u(a, b) = \min_\gamma \quad \sum_{i,j}\gamma_{i,j}M_{i,j} + reg\cdot\Omega(\gamma) + \alpha KL(\gamma 1, a) + \alpha KL(\gamma^T 1, b)
+    W_u(a, b) = \min_\gamma \quad \sum_{i,j}\gamma_{i,j}M_{i,j} + reg\cdot\Omega(\gamma) + reg_m KL(\gamma 1, a) + reg_m KL(\gamma^T 1, b)
 
     s.t. \quad  \gamma\geq 0
 
@@ -593,13 +593,11 @@ in :any:`ot.unbalanced`. Computing the optimal transport
 plan or the transport cost is similar to the balanced case. The Sinkhorn-Knopp
 algorithm is implemented in :any:`ot.sinkhorn_unbalanced` and :any:`ot.sinkhorn_unbalanced2`
 that return respectively the OT matrix and the value of the
-linear term. Note that the regularization parameter :math:`\alpha` in the
-equation above is given to those functions with the parameter :code:`reg_m`.
-
+linear term.
 
 .. note::
     The main function to solve entropic regularized UOT is :any:`ot.sinkhorn_unbalanced`.
-    This function is a wrapper and the parameter :code:`method` help you select
+    This function is a wrapper and the parameter :code:`method` helps you select
     the actual algorithm used to solve the problem:
 
     + :code:`method='sinkhorn'` calls :any:`ot.unbalanced.sinkhorn_knopp_unbalanced`
diff --git a/ot/__init__.py b/ot/__init__.py
index 99f288e..df0ef8a 100644
--- a/ot/__init__.py
+++ b/ot/__init__.py
@@ -77,4 +77,5 @@ __all__ = ["emd", "emd2", 'emd_1d', "sinkhorn", "sinkhorn2", "utils", 'datasets'
            'bregman', 'lp', 'tic', 'toc', 'toq', 'gromov',
            'emd_1d', 'emd2_1d', 'wasserstein_1d',
            'dist', 'unif', 'barycenter', 'sinkhorn_lpl1_mm', 'da', 'optim',
-           'sinkhorn_unbalanced', "barycenter_unbalanced"]
+           'sinkhorn_unbalanced', 'barycenter_unbalanced',
+           'sinkhorn_unbalanced2']
diff --git a/ot/bregman.py b/ot/bregman.py
index 02aeb6d..2cd832b 100644
--- a/ot/bregman.py
+++ b/ot/bregman.py
@@ -35,7 +35,7 @@ def sinkhorn(a, b, M, reg, method='sinkhorn', numItermax=1000,
 
     - M is the (dim_a, dim_b) metric cost matrix
     - :math:`\Omega` is the entropic regularization term :math:`\Omega(\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})`
-    - a and b are source and target weights (sum to 1)
+    - a and b are source and target weights (histograms, both sum to 1)
 
     The algorithm used for solving the problem is the Sinkhorn-Knopp matrix scaling algorithm as proposed in [2]_
 
@@ -143,7 +143,7 @@ def sinkhorn2(a, b, M, reg, method='sinkhorn', numItermax=1000,
 
     - M is the (dim_a, dim_b) metric cost matrix
     - :math:`\Omega` is the entropic regularization term :math:`\Omega(\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})`
-    - a and b are source and target weights (sum to 1)
+    - a and b are source and target weights (histograms, both sum to 1)
 
     The algorithm used for solving the problem is the Sinkhorn-Knopp matrix scaling algorithm as proposed in [2]_
 
@@ -251,7 +251,7 @@ def sinkhorn_knopp(a, b, M, reg, numItermax=1000,
 
     - M is the (dim_a, dim_b) metric cost matrix
     - :math:`\Omega` is the entropic regularization term :math:`\Omega(\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})`
-    - a and b are source and target weights (sum to 1)
+    - a and b are source and target weights (histograms, both sum to 1)
 
     The algorithm used for solving the problem is the Sinkhorn-Knopp matrix scaling algorithm as proposed in [2]_
 
@@ -432,7 +432,7 @@ def greenkhorn(a, b, M, reg, numItermax=10000, stopThr=1e-9, verbose=False,
 
     - M is the (dim_a, dim_b) metric cost matrix
     - :math:`\Omega` is the entropic regularization term :math:`\Omega(\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})`
-    - a and b are source and target weights (sum to 1)
+    - a and b are source and target weights (histograms, both sum to 1)
 
 
 
@@ -578,7 +578,8 @@ def sinkhorn_stabilized(a, b, M, reg, numItermax=1000, tau=1e3, stopThr=1e-9,
 
     - M is the (dim_a, dim_b) metric cost matrix
     - :math:`\Omega` is the entropic regularization term :math:`\Omega(\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})`
-    - a and b are source and target weights (sum to 1)
+    - a and b are source and target weights (histograms, both sum to 1)
+
 
     The algorithm used for solving the problem is the Sinkhorn-Knopp matrix
     scaling algorithm as proposed in [2]_ but with the log stabilization
@@ -808,7 +809,8 @@ def sinkhorn_epsilon_scaling(a, b, M, reg, numItermax=100, epsilon0=1e4,
 
     - M is the (dim_a, dim_b) metric cost matrix
     - :math:`\Omega` is the entropic regularization term :math:`\Omega(\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})`
-    - a and b are source and target weights (sum to 1)
+    - a and b are source and target weights (histograms, both sum to 1)
+
 
     The algorithm used for solving the problem is the Sinkhorn-Knopp matrix
     scaling algorithm as proposed in [2]_ but with the log stabilization
@@ -1229,7 +1231,6 @@ def barycenter_stabilized(A, M, reg, tau=1e10, weights=None, numItermax=1000,
         absorbing = False
         if (u > tau).any() or (v > tau).any():
             absorbing = True
-            print("YEAH absorbing")
             alpha = alpha + reg * np.log(np.max(u, 1))
             beta = beta + reg * np.log(np.max(v, 1))
             K = np.exp((alpha[:, None] + beta[None, :] -
@@ -1394,9 +1395,12 @@ def unmix(a, D, M, M0, h0, reg, reg0, alpha, numItermax=1000,
     where :
 
     - :math:`W_{M,reg}(\cdot,\cdot)` is the entropic regularized Wasserstein distance with M loss matrix (see ot.bregman.sinkhorn)
-    - :math:`\mathbf{a}` is an observed distribution,  :math:`\mathbf{h}_0` is aprior on unmixing
-    - reg and :math:`\mathbf{M}` are respectively the regularization term and the cost matrix for OT data fitting
-    - reg0 and :math:`\mathbf{M0}` are respectively the regularization term and the cost matrix for regularization
+    - :math: `\mathbf{D}` is a dictionary of `n_atoms` atoms of dimension `dim_a`, its expected shape is `(dim_a, n_atoms)`
+    - :math:`\mathbf{h}` is the estimated unmixing of dimension `n_atoms`
+    - :math:`\mathbf{a}` is an observed distribution of dimension `dim_a`
+    - :math:`\mathbf{h}_0` is a prior on `h` of dimension `dim_prior`
+    - reg and :math:`\mathbf{M}` are respectively the regularization term and the cost matrix (dim_a, dim_a) for OT data fitting
+    - reg0 and :math:`\mathbf{M0}` are respectively the regularization term and the cost matrix (dim_prior, n_atoms) regularization
     - :math:`\\alpha`weight data fitting and regularization
 
     The optimization problem is solved suing the algorithm described in [4]
@@ -1404,16 +1408,16 @@ def unmix(a, D, M, M0, h0, reg, reg0, alpha, numItermax=1000,
 
     Parameters
     ----------
-    a : ndarray, shape (n_observed)
-        observed distribution
-    D : ndarray, shape (dim, dim)
+    a : ndarray, shape (dim_a)
+        observed distribution (histogram, sums to 1)
+    D : ndarray, shape (dim_a, n_atoms)
         dictionary matrix
-    M : ndarray, shape (dim, dim)
+    M : ndarray, shape (dim_a, dim_a)
         loss matrix
-    M0 : ndarray, shape (n_observed, n_observed)
+    M0 : ndarray, shape (n_atoms, dim_prior)
         loss matrix
-    h0 : ndarray, shape (dim,)
-        prior on h
+    h0 : ndarray, shape (n_atoms,)
+        prior on the estimated unmixing h
     reg : float
         Regularization term >0 (Wasserstein data fitting)
     reg0 : float
@@ -1432,7 +1436,7 @@ def unmix(a, D, M, M0, h0, reg, reg0, alpha, numItermax=1000,
 
     Returns
     -------
-    a : ndarray, shape (dim,)
+    h : ndarray, shape (n_atoms,)
         Wasserstein barycenter
     log : dict
         log dictionary return only if log==True in parameters