Merge branch 'master' into doc_modules

8 files changed, 919 insertions, 3 deletions

diff --git a/README.md b/README.md
index 288ba13..caa5759 100644
--- a/README.md
+++ b/README.md
@@ -27,6 +27,7 @@ It provides the following solvers:
 * Gromov-Wasserstein distances and barycenters ([13] and regularized [12])
 * Stochastic Optimization for Large-scale Optimal Transport (semi-dual problem [18] and dual problem [19])
 * Non regularized free support Wasserstein barycenters [20].
+* Unbalanced OT with KL relaxation distance and barycenter [10, 25].
 
 Some demonstrations (both in Python and Jupyter Notebook format) are available in the examples folder.
 
@@ -178,6 +179,7 @@ The contributors to this library are
 * [Kilian Fatras](https://kilianfatras.github.io/)
 * [Alain Rakotomamonjy](https://sites.google.com/site/alainrakotomamonjy/home)
 * [Vayer Titouan](https://tvayer.github.io/)
+* [Hicham Janati](https://hichamjanati.github.io/) (Unbalanced OT)
 
 This toolbox benefit a lot from open source research and we would like to thank the following persons for providing some code (in various languages):
 
@@ -249,3 +251,5 @@ You can also post bug reports and feature requests in Github issues. Make sure t
 [23] Aude, G., Peyré, G., Cuturi, M., [Learning Generative Models with Sinkhorn Divergences](https://arxiv.org/abs/1706.00292), Proceedings of the Twenty-First International Conference on Artficial Intelligence and Statistics, (AISTATS) 21, 2018
 
 [24] Vayer, T., Chapel, L., Flamary, R., Tavenard, R. and Courty, N. (2019). [Optimal Transport for structured data with application on graphs](http://proceedings.mlr.press/v97/titouan19a.html) Proceedings of the 36th International Conference on Machine Learning (ICML).
+
+[25] Frogner C., Zhang C., Mobahi H., Araya-Polo M., Poggio T. (2019). [Learning with a Wasserstein Loss](http://cbcl.mit.edu/wasserstein/)  Advances in Neural Information Processing Systems (NIPS).
diff --git a/docs/source/all.rst b/docs/source/all.rst
index 32930fd..c968aa1 100644
--- a/docs/source/all.rst
+++ b/docs/source/all.rst
@@ -43,7 +43,7 @@ ot.da
 
 .. automodule:: ot.da
   :members:
-  
+
 ot.gpu
 --------
 
@@ -80,3 +80,9 @@ ot.stochastic
 
 .. automodule:: ot.stochastic
    :members:
+
+ot.unbalanced
+-------------
+
+.. automodule:: ot.unbalanced
+  :members:
diff --git a/examples/plot_UOT_1D.py b/examples/plot_UOT_1D.py
new file mode 100644
index 0000000..2ea8b05
--- /dev/null
+++ b/examples/plot_UOT_1D.py
@@ -0,0 +1,76 @@
+# -*- coding: utf-8 -*-
+"""
+===============================
+1D Unbalanced optimal transport
+===============================
+
+This example illustrates the computation of Unbalanced Optimal transport
+using a Kullback-Leibler relaxation.
+"""
+
+# Author: Hicham Janati <hicham.janati@inria.fr>
+#
+# License: MIT License
+
+import numpy as np
+import matplotlib.pylab as pl
+import ot
+import ot.plot
+from ot.datasets import make_1D_gauss as gauss
+
+##############################################################################
+# Generate data
+# -------------
+
+
+#%% parameters
+
+n = 100  # nb bins
+
+# bin positions
+x = np.arange(n, dtype=np.float64)
+
+# Gaussian distributions
+a = gauss(n, m=20, s=5)  # m= mean, s= std
+b = gauss(n, m=60, s=10)
+
+# make distributions unbalanced
+b *= 5.
+
+# loss matrix
+M = ot.dist(x.reshape((n, 1)), x.reshape((n, 1)))
+M /= M.max()
+
+
+##############################################################################
+# Plot distributions and loss matrix
+# ----------------------------------
+
+#%% plot the distributions
+
+pl.figure(1, figsize=(6.4, 3))
+pl.plot(x, a, 'b', label='Source distribution')
+pl.plot(x, b, 'r', label='Target distribution')
+pl.legend()
+
+# plot distributions and loss matrix
+
+pl.figure(2, figsize=(5, 5))
+ot.plot.plot1D_mat(a, b, M, 'Cost matrix M')
+
+
+##############################################################################
+# Solve Unbalanced Sinkhorn
+# --------------
+
+
+# Sinkhorn
+
+epsilon = 0.1  # entropy parameter
+alpha = 1.  # Unbalanced KL relaxation parameter
+Gs = ot.unbalanced.sinkhorn_unbalanced(a, b, M, epsilon, alpha, verbose=True)
+
+pl.figure(4, figsize=(5, 5))
+ot.plot.plot1D_mat(a, b, Gs, 'UOT matrix Sinkhorn')
+
+pl.show()
diff --git a/examples/plot_UOT_barycenter_1D.py b/examples/plot_UOT_barycenter_1D.py
new file mode 100644
index 0000000..c8d9d3b
--- /dev/null
+++ b/examples/plot_UOT_barycenter_1D.py
@@ -0,0 +1,164 @@
+# -*- coding: utf-8 -*-
+"""
+===========================================================
+1D Wasserstein barycenter demo for Unbalanced distributions
+===========================================================
+
+This example illustrates the computation of regularized Wassersyein Barycenter
+as proposed in [10] for Unbalanced inputs.
+
+
+[10] Chizat, L., Peyré, G., Schmitzer, B., & Vialard, F. X. (2016). Scaling algorithms for unbalanced transport problems. arXiv preprint arXiv:1607.05816.
+
+"""
+
+# Author: Hicham Janati <hicham.janati@inria.fr>
+#
+# License: MIT License
+
+import numpy as np
+import matplotlib.pylab as pl
+import ot
+# necessary for 3d plot even if not used
+from mpl_toolkits.mplot3d import Axes3D  # noqa
+from matplotlib.collections import PolyCollection
+
+##############################################################################
+# Generate data
+# -------------
+
+# parameters
+
+n = 100  # nb bins
+
+# bin positions
+x = np.arange(n, dtype=np.float64)
+
+# Gaussian distributions
+a1 = ot.datasets.make_1D_gauss(n, m=20, s=5)  # m= mean, s= std
+a2 = ot.datasets.make_1D_gauss(n, m=60, s=8)
+
+# make unbalanced dists
+a2 *= 3.
+
+# creating matrix A containing all distributions
+A = np.vstack((a1, a2)).T
+n_distributions = A.shape[1]
+
+# loss matrix + normalization
+M = ot.utils.dist0(n)
+M /= M.max()
+
+##############################################################################
+# Plot data
+# ---------
+
+# plot the distributions
+
+pl.figure(1, figsize=(6.4, 3))
+for i in range(n_distributions):
+    pl.plot(x, A[:, i])
+pl.title('Distributions')
+pl.tight_layout()
+
+##############################################################################
+# Barycenter computation
+# ----------------------
+
+# non weighted barycenter computation
+
+weight = 0.5  # 0<=weight<=1
+weights = np.array([1 - weight, weight])
+
+# l2bary
+bary_l2 = A.dot(weights)
+
+# wasserstein
+reg = 1e-3
+alpha = 1.
+
+bary_wass = ot.unbalanced.barycenter_unbalanced(A, M, reg, alpha, weights)
+
+pl.figure(2)
+pl.clf()
+pl.subplot(2, 1, 1)
+for i in range(n_distributions):
+    pl.plot(x, A[:, i])
+pl.title('Distributions')
+
+pl.subplot(2, 1, 2)
+pl.plot(x, bary_l2, 'r', label='l2')
+pl.plot(x, bary_wass, 'g', label='Wasserstein')
+pl.legend()
+pl.title('Barycenters')
+pl.tight_layout()
+
+##############################################################################
+# Barycentric interpolation
+# -------------------------
+
+# barycenter interpolation
+
+n_weight = 11
+weight_list = np.linspace(0, 1, n_weight)
+
+
+B_l2 = np.zeros((n, n_weight))
+
+B_wass = np.copy(B_l2)
+
+for i in range(0, n_weight):
+    weight = weight_list[i]
+    weights = np.array([1 - weight, weight])
+    B_l2[:, i] = A.dot(weights)
+    B_wass[:, i] = ot.unbalanced.barycenter_unbalanced(A, M, reg, alpha, weights)
+
+
+# plot interpolation
+
+pl.figure(3)
+
+cmap = pl.cm.get_cmap('viridis')
+verts = []
+zs = weight_list
+for i, z in enumerate(zs):
+    ys = B_l2[:, i]
+    verts.append(list(zip(x, ys)))
+
+ax = pl.gcf().gca(projection='3d')
+
+poly = PolyCollection(verts, facecolors=[cmap(a) for a in weight_list])
+poly.set_alpha(0.7)
+ax.add_collection3d(poly, zs=zs, zdir='y')
+ax.set_xlabel('x')
+ax.set_xlim3d(0, n)
+ax.set_ylabel(r'$\alpha$')
+ax.set_ylim3d(0, 1)
+ax.set_zlabel('')
+ax.set_zlim3d(0, B_l2.max() * 1.01)
+pl.title('Barycenter interpolation with l2')
+pl.tight_layout()
+
+pl.figure(4)
+cmap = pl.cm.get_cmap('viridis')
+verts = []
+zs = weight_list
+for i, z in enumerate(zs):
+    ys = B_wass[:, i]
+    verts.append(list(zip(x, ys)))
+
+ax = pl.gcf().gca(projection='3d')
+
+poly = PolyCollection(verts, facecolors=[cmap(a) for a in weight_list])
+poly.set_alpha(0.7)
+ax.add_collection3d(poly, zs=zs, zdir='y')
+ax.set_xlabel('x')
+ax.set_xlim3d(0, n)
+ax.set_ylabel(r'$\alpha$')
+ax.set_ylim3d(0, 1)
+ax.set_zlabel('')
+ax.set_zlim3d(0, B_l2.max() * 1.01)
+pl.title('Barycenter interpolation with Wasserstein')
+pl.tight_layout()
+
+pl.show()
diff --git a/ot/__init__.py b/ot/__init__.py
index 16b0dd6..89bdc78 100644
--- a/ot/__init__.py
+++ b/ot/__init__.py
@@ -30,10 +30,12 @@ from . import da
 from . import gromov
 from . import smooth
 from . import stochastic
+from . import unbalanced
 
 # OT functions
 from .lp import emd, emd2
 from .bregman import sinkhorn, sinkhorn2, barycenter
+from .unbalanced import sinkhorn_unbalanced, barycenter_unbalanced
 from .da import sinkhorn_lpl1_mm
 
 # utils functions
@@ -43,4 +45,5 @@ __version__ = "0.5.1"
 
 __all__ = ["emd", "emd2", "sinkhorn", "sinkhorn2", "utils", 'datasets',
            'bregman', 'lp', 'tic', 'toc', 'toq', 'gromov',
-           'dist', 'unif', 'barycenter', 'sinkhorn_lpl1_mm', 'da', 'optim']
+           'dist', 'unif', 'barycenter', 'sinkhorn_lpl1_mm', 'da', 'optim',
+           'sinkhorn_unbalanced', "barycenter_unbalanced"]
diff --git a/ot/bregman.py b/ot/bregman.py
index 321712b..09716e6 100644
--- a/ot/bregman.py
+++ b/ot/bregman.py
@@ -241,7 +241,7 @@ def sinkhorn2(a, b, M, reg, method='sinkhorn', numItermax=1000,
 
     b = np.asarray(b, dtype=np.float64)
     if len(b.shape) < 2:
-        b = b.reshape((-1, 1))
+        b = b[:, None]
 
     return sink()
 
diff --git a/ot/unbalanced.py b/ot/unbalanced.py
new file mode 100644
index 0000000..484ce95
--- /dev/null
+++ b/ot/unbalanced.py
@@ -0,0 +1,517 @@
+# -*- coding: utf-8 -*-
+"""
+Regularized Unbalanced OT
+"""
+
+# Author: Hicham Janati <hicham.janati@inria.fr>
+# License: MIT License
+
+import warnings
+import numpy as np
+# from .utils import unif, dist
+
+
+def sinkhorn_unbalanced(a, b, M, reg, alpha, method='sinkhorn', numItermax=1000,
+                        stopThr=1e-9, verbose=False, log=False, **kwargs):
+    u"""
+    Solve the unbalanced entropic regularization optimal transport problem and return the loss
+
+    The function solves the following optimization problem:
+
+    .. math::
+        W = \min_\gamma <\gamma,M>_F + reg\cdot\Omega(\gamma) + \alpha KL(\gamma 1, a) + \alpha KL(\gamma^T 1, b)
+
+        s.t.
+             \gamma\geq 0
+    where :
+
+    - M is the (ns, nt) 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
+    - KL is the Kullback-Leibler divergence
+
+    The algorithm used for solving the problem is the generalized Sinkhorn-Knopp matrix scaling algorithm as proposed in [10, 23]_
+
+
+    Parameters
+    ----------
+    a : np.ndarray (ns,)
+        samples weights in the source domain
+    b : np.ndarray (nt,) or np.ndarray (nt,n_hists)
+        samples in the target domain, compute sinkhorn with multiple targets
+        and fixed M if b is a matrix (return OT loss + dual variables in log)
+    M : np.ndarray (ns, nt)
+        loss matrix
+    reg : float
+        Entropy regularization term > 0
+    alpha : float
+        Marginal relaxation term > 0
+    method : str
+        method used for the solver either 'sinkhorn',  'sinkhorn_stabilized' or
+        'sinkhorn_epsilon_scaling', see those function for specific parameters
+    numItermax : int, optional
+        Max number of iterations
+    stopThr : float, optional
+        Stop threshol on error (> 0)
+    verbose : bool, optional
+        Print information along iterations
+    log : bool, optional
+        record log if True
+
+
+    Returns
+    -------
+    W : (nt) ndarray or float
+        Optimal transportation matrix for the given parameters
+    log : dict
+        log dictionary return only if log==True in parameters
+
+    Examples
+    --------
+
+    >>> import ot
+    >>> a=[.5, .5]
+    >>> b=[.5, .5]
+    >>> M=[[0., 1.], [1., 0.]]
+    >>> ot.sinkhorn_unbalanced(a, b, M, 1, 1)
+    array([[0.51122823, 0.18807035],
+          [0.18807035, 0.51122823]])
+
+
+    References
+    ----------
+
+    .. [2] M. Cuturi, Sinkhorn Distances : Lightspeed Computation of Optimal Transport, Advances in Neural Information Processing Systems (NIPS) 26, 2013
+
+    .. [9] Schmitzer, B. (2016). Stabilized Sparse Scaling Algorithms for Entropy Regularized Transport Problems. arXiv preprint arXiv:1610.06519.
+
+    .. [10] Chizat, L., Peyré, G., Schmitzer, B., & Vialard, F. X. (2016). Scaling algorithms for unbalanced transport problems. arXiv preprint arXiv:1607.05816.
+
+    .. [25] Frogner C., Zhang C., Mobahi H., Araya-Polo M., Poggio T. : Learning with a Wasserstein Loss,  Advances in Neural Information Processing Systems (NIPS) 2015
+
+
+    See Also
+    --------
+    ot.unbalanced.sinkhorn_knopp_unbalanced : Unbalanced Classic Sinkhorn [10]
+    ot.unbalanced.sinkhorn_stabilized_unbalanced: Unbalanced Stabilized sinkhorn [9][10]
+    ot.unbalanced.sinkhorn_epsilon_scaling_unbalanced: Unbalanced Sinkhorn with epslilon scaling [9][10]
+
+    """
+
+    if method.lower() == 'sinkhorn':
+        def sink():
+            return sinkhorn_knopp_unbalanced(a, b, M, reg, alpha,
+                                             numItermax=numItermax,
+                                             stopThr=stopThr, verbose=verbose,
+                                             log=log, **kwargs)
+
+    elif method.lower() in ['sinkhorn_stabilized', 'sinkhorn_epsilon_scaling']:
+        warnings.warn('Method not implemented yet. Using classic Sinkhorn Knopp')
+
+        def sink():
+            return sinkhorn_knopp_unbalanced(a, b, M, reg, alpha,
+                                             numItermax=numItermax,
+                                             stopThr=stopThr, verbose=verbose,
+                                             log=log, **kwargs)
+    else:
+        raise ValueError('Unknown method. Using classic Sinkhorn Knopp')
+
+    return sink()
+
+
+def sinkhorn_unbalanced2(a, b, M, reg, alpha, method='sinkhorn',
+                         numItermax=1000, stopThr=1e-9, verbose=False,
+                         log=False, **kwargs):
+    u"""
+    Solve the entropic regularization unbalanced optimal transport problem and return the loss
+
+    The function solves the following optimization problem:
+
+    .. math::
+        W = \min_\gamma <\gamma,M>_F + reg\cdot\Omega(\gamma) + \alpha KL(\gamma 1, a) + \alpha KL(\gamma^T 1, b)
+
+        s.t.
+             \gamma\geq 0
+    where :
+
+    - M is the (ns, nt) 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
+    - KL is the Kullback-Leibler divergence
+
+    The algorithm used for solving the problem is the generalized Sinkhorn-Knopp matrix scaling algorithm as proposed in [10, 23]_
+
+
+    Parameters
+    ----------
+    a : np.ndarray (ns,)
+        samples weights in the source domain
+    b : np.ndarray (nt,) or np.ndarray (nt, n_hists)
+        samples in the target domain, compute sinkhorn with multiple targets
+        and fixed M if b is a matrix (return OT loss + dual variables in log)
+    M : np.ndarray (ns,nt)
+        loss matrix
+    reg : float
+        Entropy regularization term > 0
+    alpha : float
+        Marginal relaxation term > 0
+    method : str
+        method used for the solver either 'sinkhorn',  'sinkhorn_stabilized' or
+        'sinkhorn_epsilon_scaling', see those function for specific parameters
+    numItermax : int, optional
+        Max number of iterations
+    stopThr : float, optional
+        Stop threshol on error (>0)
+    verbose : bool, optional
+        Print information along iterations
+    log : bool, optional
+        record log if True
+
+
+    Returns
+    -------
+    W : (nt) ndarray or float
+        Optimal transportation matrix for the given parameters
+    log : dict
+        log dictionary return only if log==True in parameters
+
+    Examples
+    --------
+
+    >>> import ot
+    >>> a=[.5, .10]
+    >>> b=[.5, .5]
+    >>> M=[[0., 1.],[1., 0.]]
+    >>> ot.unbalanced.sinkhorn_unbalanced2(a, b, M, 1., 1.)
+    array([0.31912866])
+
+
+
+    References
+    ----------
+
+    .. [2] M. Cuturi, Sinkhorn Distances : Lightspeed Computation of Optimal Transport, Advances in Neural Information Processing Systems (NIPS) 26, 2013
+
+    .. [9] Schmitzer, B. (2016). Stabilized Sparse Scaling Algorithms for Entropy Regularized Transport Problems. arXiv preprint arXiv:1610.06519.
+
+    .. [10] Chizat, L., Peyré, G., Schmitzer, B., & Vialard, F. X. (2016). Scaling algorithms for unbalanced transport problems. arXiv preprint arXiv:1607.05816.
+
+    .. [25] Frogner C., Zhang C., Mobahi H., Araya-Polo M., Poggio T. : Learning with a Wasserstein Loss,  Advances in Neural Information Processing Systems (NIPS) 2015
+
+    See Also
+    --------
+    ot.unbalanced.sinkhorn_knopp : Unbalanced Classic Sinkhorn [10]
+    ot.unbalanced.sinkhorn_stabilized: Unbalanced Stabilized sinkhorn [9][10]
+    ot.unbalanced.sinkhorn_epsilon_scaling: Unbalanced Sinkhorn with epslilon scaling [9][10]
+
+    """
+
+    if method.lower() == 'sinkhorn':
+        def sink():
+            return sinkhorn_knopp_unbalanced(a, b, M, reg, alpha,
+                                             numItermax=numItermax,
+                                             stopThr=stopThr, verbose=verbose,
+                                             log=log, **kwargs)
+
+    elif method.lower() in ['sinkhorn_stabilized', 'sinkhorn_epsilon_scaling']:
+        warnings.warn('Method not implemented yet. Using classic Sinkhorn Knopp')
+
+        def sink():
+            return sinkhorn_knopp_unbalanced(a, b, M, reg, alpha,
+                                             numItermax=numItermax,
+                                             stopThr=stopThr, verbose=verbose,
+                                             log=log, **kwargs)
+    else:
+        raise ValueError('Unknown method. Using classic Sinkhorn Knopp')
+
+    b = np.asarray(b, dtype=np.float64)
+    if len(b.shape) < 2:
+        b = b[:, None]
+
+    return sink()
+
+
+def sinkhorn_knopp_unbalanced(a, b, M, reg, alpha, numItermax=1000,
+                              stopThr=1e-9, verbose=False, log=False, **kwargs):
+    """
+    Solve the entropic regularization unbalanced optimal transport problem and return the loss
+
+    The function solves the following optimization problem:
+
+    .. math::
+        W = \min_\gamma <\gamma,M>_F + reg\cdot\Omega(\gamma) + \alpha KL(\gamma 1, a) + \alpha KL(\gamma^T 1, b)
+
+        s.t.
+             \gamma\geq 0
+    where :
+
+    - M is the (ns, nt) 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
+    - KL is the Kullback-Leibler divergence
+
+    The algorithm used for solving the problem is the generalized Sinkhorn-Knopp matrix scaling algorithm as proposed in [10, 23]_
+
+
+    Parameters
+    ----------
+    a : np.ndarray (ns,)
+        samples weights in the source domain
+    b : np.ndarray (nt,) or np.ndarray (nt, n_hists)
+        samples in the target domain, compute sinkhorn with multiple targets
+        and fixed M if b is a matrix (return OT loss + dual variables in log)
+    M : np.ndarray (ns,nt)
+        loss matrix
+    reg : float
+        Entropy regularization term > 0
+    alpha : float
+        Marginal relaxation term > 0
+    numItermax : int, optional
+        Max number of iterations
+    stopThr : float, optional
+        Stop threshol on error (>0)
+    verbose : bool, optional
+        Print information along iterations
+    log : bool, optional
+        record log if True
+
+
+    Returns
+    -------
+    gamma : (ns x nt) ndarray
+        Optimal transportation matrix for the given parameters
+    log : dict
+        log dictionary return only if log==True in parameters
+
+    Examples
+    --------
+
+    >>> import ot
+    >>> a=[.5, .15]
+    >>> b=[.5, .5]
+    >>> M=[[0., 1.],[1., 0.]]
+    >>> ot.sinkhorn_knopp_unbalanced(a, b, M, 1., 1.)
+    array([[0.52761554, 0.22392482],
+           [0.10286295, 0.32257641]])
+
+    References
+    ----------
+
+    .. [10] Chizat, L., Peyré, G., Schmitzer, B., & Vialard, F. X. (2016). Scaling algorithms for unbalanced transport problems. arXiv preprint arXiv:1607.05816.
+
+    .. [25] Frogner C., Zhang C., Mobahi H., Araya-Polo M., Poggio T. : Learning with a Wasserstein Loss,  Advances in Neural Information Processing Systems (NIPS) 2015
+
+    See Also
+    --------
+    ot.lp.emd : Unregularized OT
+    ot.optim.cg : General regularized OT
+
+    """
+
+    a = np.asarray(a, dtype=np.float64)
+    b = np.asarray(b, dtype=np.float64)
+    M = np.asarray(M, dtype=np.float64)
+
+    n_a, n_b = M.shape
+
+    if len(a) == 0:
+        a = np.ones(n_a, dtype=np.float64) / n_a
+    if len(b) == 0:
+        b = np.ones(n_b, dtype=np.float64) / n_b
+
+    if len(b.shape) > 1:
+        n_hists = b.shape[1]
+    else:
+        n_hists = 0
+
+    if log:
+        log = {'err': []}
+
+    # we assume that no distances are null except those of the diagonal of
+    # distances
+    if n_hists:
+        u = np.ones((n_a, 1)) / n_a
+        v = np.ones((n_b, n_hists)) / n_b
+        a = a.reshape(n_a, 1)
+    else:
+        u = np.ones(n_a) / n_a
+        v = np.ones(n_b) / n_b
+
+    # print(reg)
+    # Next 3 lines equivalent to K= np.exp(-M/reg), but faster to compute
+    K = np.empty(M.shape, dtype=M.dtype)
+    np.divide(M, -reg, out=K)
+    np.exp(K, out=K)
+
+    # print(np.min(K))
+    fi = alpha / (alpha + reg)
+
+    cpt = 0
+    err = 1.
+
+    while (err > stopThr and cpt < numItermax):
+        uprev = u
+        vprev = v
+
+        Kv = K.dot(v)
+        u = (a / Kv) ** fi
+        Ktu = K.T.dot(u)
+        v = (b / Ktu) ** fi
+
+        if (np.any(Ktu == 0.)
+                or np.any(np.isnan(u)) or np.any(np.isnan(v))
+                or np.any(np.isinf(u)) or np.any(np.isinf(v))):
+            # we have reached the machine precision
+            # come back to previous solution and quit loop
+            warnings.warn('Numerical errors at iteration', cpt)
+            u = uprev
+            v = vprev
+            break
+        if cpt % 10 == 0:
+            # we can speed up the process by checking for the error only all
+            # the 10th iterations
+            err = np.sum((u - uprev)**2) / np.sum((u)**2) + \
+                np.sum((v - vprev)**2) / np.sum((v)**2)
+            if log:
+                log['err'].append(err)
+            if verbose:
+                if cpt % 200 == 0:
+                    print(
+                        '{:5s}|{:12s}'.format('It.', 'Err') + '\n' + '-' * 19)
+                print('{:5d}|{:8e}|'.format(cpt, err))
+        cpt = cpt + 1
+    if log:
+        log['u'] = u
+        log['v'] = v
+
+    if n_hists:  # return only loss
+        res = np.einsum('ik,ij,jk,ij->k', u, K, v, M)
+        if log:
+            return res, log
+        else:
+            return res
+
+    else:  # return OT matrix
+
+        if log:
+            return u[:, None] * K * v[None, :], log
+        else:
+            return u[:, None] * K * v[None, :]
+
+
+def barycenter_unbalanced(A, M, reg, alpha, weights=None, numItermax=1000,
+                          stopThr=1e-4, verbose=False, log=False):
+    """Compute the entropic regularized unbalanced wasserstein barycenter of distributions A
+
+     The function solves the following optimization problem:
+
+    .. math::
+       \mathbf{a} = arg\min_\mathbf{a} \sum_i Wu_{reg}(\mathbf{a},\mathbf{a}_i)
+
+    where :
+
+    - :math:`Wu_{reg}(\cdot,\cdot)` is the unbalanced entropic regularized Wasserstein distance (see ot.unbalanced.sinkhorn_unbalanced)
+    - :math:`\mathbf{a}_i` are training distributions in the columns of matrix :math:`\mathbf{A}`
+    - reg and :math:`\mathbf{M}` are respectively the regularization term and the cost matrix for OT
+    - alpha is the marginal relaxation hyperparameter
+    The algorithm used for solving the problem is the generalized Sinkhorn-Knopp matrix scaling algorithm as proposed in [10]_
+
+    Parameters
+    ----------
+    A : np.ndarray (d,n)
+        n training distributions a_i of size d
+    M : np.ndarray (d,d)
+        loss matrix   for OT
+    reg : float
+        Entropy regularization term > 0
+    alpha : float
+        Marginal relaxation term > 0
+    weights : np.ndarray (n,)
+        Weights of each histogram a_i on the simplex (barycentric coodinates)
+    numItermax : int, optional
+        Max number of iterations
+    stopThr : float, optional
+        Stop threshol on error (>0)
+    verbose : bool, optional
+        Print information along iterations
+    log : bool, optional
+        record log if True
+
+
+    Returns
+    -------
+    a : (d,) ndarray
+        Unbalanced Wasserstein barycenter
+    log : dict
+        log dictionary return only if log==True in parameters
+
+
+    References
+    ----------
+
+    .. [3] Benamou, J. D., Carlier, G., Cuturi, M., Nenna, L., & Peyré, G. (2015). Iterative Bregman projections for regularized transportation problems. SIAM Journal on Scientific Computing, 37(2), A1111-A1138.
+    .. [10] Chizat, L., Peyré, G., Schmitzer, B., & Vialard, F. X. (2016). Scaling algorithms for unbalanced transport problems. arXiv preprint arXiv:1607.05816.
+
+
+    """
+    p, n_hists = A.shape
+    if weights is None:
+        weights = np.ones(n_hists) / n_hists
+    else:
+        assert(len(weights) == A.shape[1])
+
+    if log:
+        log = {'err': []}
+
+    K = np.exp(- M / reg)
+
+    fi = alpha / (alpha + reg)
+
+    v = np.ones((p, n_hists)) / p
+    u = np.ones((p, 1)) / p
+
+    cpt = 0
+    err = 1.
+
+    while (err > stopThr and cpt < numItermax):
+        uprev = u
+        vprev = v
+
+        Kv = K.dot(v)
+        u = (A / Kv) ** fi
+        Ktu = K.T.dot(u)
+        q = ((Ktu ** (1 - fi)).dot(weights))
+        q = q ** (1 / (1 - fi))
+        Q = q[:, None]
+        v = (Q / Ktu) ** fi
+
+        if (np.any(Ktu == 0.)
+                or np.any(np.isnan(u)) or np.any(np.isnan(v))
+                or np.any(np.isinf(u)) or np.any(np.isinf(v))):
+            # we have reached the machine precision
+            # come back to previous solution and quit loop
+            warnings.warn('Numerical errors at iteration %s' % cpt)
+            u = uprev
+            v = vprev
+            break
+        if cpt % 10 == 0:
+            # we can speed up the process by checking for the error only all
+            # the 10th iterations
+            err = np.sum((u - uprev) ** 2) / np.sum((u) ** 2) + \
+                np.sum((v - vprev) ** 2) / np.sum((v) ** 2)
+            if log:
+                log['err'].append(err)
+            if verbose:
+                if cpt % 50 == 0:
+                    print(
+                        '{:5s}|{:12s}'.format('It.', 'Err') + '\n' + '-' * 19)
+                print('{:5d}|{:8e}|'.format(cpt, err))
+
+    cpt += 1
+    if log:
+        log['niter'] = cpt
+        log['u'] = u
+        log['v'] = v
+        return q, log
+    else:
+        return q
diff --git a/test/test_unbalanced.py b/test/test_unbalanced.py
new file mode 100644
index 0000000..1395fe1
--- /dev/null
+++ b/test/test_unbalanced.py
@@ -0,0 +1,146 @@
+"""Tests for module Unbalanced OT with entropy regularization"""
+
+# Author: Hicham Janati <hicham.janati@inria.fr>
+#
+# License: MIT License
+
+import numpy as np
+import ot
+import pytest
+
+
+@pytest.mark.parametrize("method", ["sinkhorn"])
+def test_unbalanced_convergence(method):
+    # test generalized sinkhorn for unbalanced OT
+    n = 100
+    rng = np.random.RandomState(42)
+
+    x = rng.randn(n, 2)
+    a = ot.utils.unif(n)
+
+    # make dists unbalanced
+    b = ot.utils.unif(n) * 1.5
+
+    M = ot.dist(x, x)
+    epsilon = 1.
+    alpha = 1.
+    K = np.exp(- M / epsilon)
+
+    G, log = ot.unbalanced.sinkhorn_unbalanced(a, b, M, reg=epsilon, alpha=alpha,
+                                               stopThr=1e-10, method=method,
+                                               log=True)
+    loss = ot.unbalanced.sinkhorn_unbalanced2(a, b, M, epsilon, alpha,
+                                              method=method)
+    # check fixed point equations
+    fi = alpha / (alpha + epsilon)
+    v_final = (b / K.T.dot(log["u"])) ** fi
+    u_final = (a / K.dot(log["v"])) ** fi
+
+    np.testing.assert_allclose(
+        u_final, log["u"], atol=1e-05)
+    np.testing.assert_allclose(
+        v_final, log["v"], atol=1e-05)
+
+    # check if sinkhorn_unbalanced2 returns the correct loss
+    np.testing.assert_allclose((G * M).sum(), loss, atol=1e-5)
+
+
+@pytest.mark.parametrize("method", ["sinkhorn"])
+def test_unbalanced_multiple_inputs(method):
+    # test generalized sinkhorn for unbalanced OT
+    n = 100
+    rng = np.random.RandomState(42)
+
+    x = rng.randn(n, 2)
+    a = ot.utils.unif(n)
+
+    # make dists unbalanced
+    b = rng.rand(n, 2)
+
+    M = ot.dist(x, x)
+    epsilon = 1.
+    alpha = 1.
+    K = np.exp(- M / epsilon)
+
+    loss, log = ot.unbalanced.sinkhorn_unbalanced(a, b, M, reg=epsilon,
+                                                  alpha=alpha,
+                                                  stopThr=1e-10, method=method,
+                                                  log=True)
+    # check fixed point equations
+    fi = alpha / (alpha + epsilon)
+    v_final = (b / K.T.dot(log["u"])) ** fi
+
+    u_final = (a[:, None] / K.dot(log["v"])) ** fi
+
+    np.testing.assert_allclose(
+        u_final, log["u"], atol=1e-05)
+    np.testing.assert_allclose(
+        v_final, log["v"], atol=1e-05)
+
+    assert len(loss) == b.shape[1]
+
+
+def test_unbalanced_barycenter():
+    # test generalized sinkhorn for unbalanced OT barycenter
+    n = 100
+    rng = np.random.RandomState(42)
+
+    x = rng.randn(n, 2)
+    A = rng.rand(n, 2)
+
+    # make dists unbalanced
+    A = A * np.array([1, 2])[None, :]
+    M = ot.dist(x, x)
+    epsilon = 1.
+    alpha = 1.
+    K = np.exp(- M / epsilon)
+
+    q, log = ot.unbalanced.barycenter_unbalanced(A, M, reg=epsilon, alpha=alpha,
+                                                 stopThr=1e-10,
+                                                 log=True)
+    # check fixed point equations
+    fi = alpha / (alpha + epsilon)
+    v_final = (q[:, None] / K.T.dot(log["u"])) ** fi
+    u_final = (A / K.dot(log["v"])) ** fi
+
+    np.testing.assert_allclose(
+        u_final, log["u"], atol=1e-05)
+    np.testing.assert_allclose(
+        v_final, log["v"], atol=1e-05)
+
+
+def test_implemented_methods():
+    IMPLEMENTED_METHODS = ['sinkhorn']
+    TO_BE_IMPLEMENTED_METHODS = ['sinkhorn_stabilized',
+                                 'sinkhorn_epsilon_scaling']
+    NOT_VALID_TOKENS = ['foo']
+    # test generalized sinkhorn for unbalanced OT barycenter
+    n = 3
+    rng = np.random.RandomState(42)
+
+    x = rng.randn(n, 2)
+    a = ot.utils.unif(n)
+
+    # make dists unbalanced
+    b = ot.utils.unif(n) * 1.5
+
+    M = ot.dist(x, x)
+    epsilon = 1.
+    alpha = 1.
+    for method in IMPLEMENTED_METHODS:
+        ot.unbalanced.sinkhorn_unbalanced(a, b, M, epsilon, alpha,
+                                          method=method)
+        ot.unbalanced.sinkhorn_unbalanced2(a, b, M, epsilon, alpha,
+                                           method=method)
+    with pytest.warns(UserWarning, match='not implemented'):
+        for method in set(TO_BE_IMPLEMENTED_METHODS):
+            ot.unbalanced.sinkhorn_unbalanced(a, b, M, epsilon, alpha,
+                                              method=method)
+            ot.unbalanced.sinkhorn_unbalanced2(a, b, M, epsilon, alpha,
+                                               method=method)
+    with pytest.raises(ValueError):
+        for method in set(NOT_VALID_TOKENS):
+            ot.unbalanced.sinkhorn_unbalanced(a, b, M, epsilon, alpha,
+                                              method=method)
+            ot.unbalanced.sinkhorn_unbalanced2(a, b, M, epsilon, alpha,
+                                               method=method)