diff --git a/RELEASES.md b/RELEASES.md index 860f2babb..abb8f04c8 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -17,6 +17,7 @@ - Allow `NumpyBackend.seed` to adopt an existing `np.random.RandomState` instance and remove NumPy-specific random sampling paths in sliced utilities (PR #849, Issue #848) - Remove a leftover debug `print` from `ot.utils.projection_sparse_simplex` with `axis=1`, and make the `ot.datasets.make_gauss_hd` docstring a raw string so importing `ot` no longer emits a `SyntaxWarning` (PR #860) +- Fix the value reported by the entropic unbalanced OT solvers (`ot.unbalanced.sinkhorn_unbalanced`, `sinkhorn_unbalanced2`, `sinkhorn_knopp_unbalanced`, `sinkhorn_stabilized_unbalanced` and `sinkhorn_unbalanced_translation_invariant`, also exposed as `ot.solve(..., reg=..., unbalanced=...).value`): every penalization is now evaluated with the generalized KL divergence, and with `reg_type="entropy"` the constant `dim_a * dim_b` is removed so that the reported value matches the documented negative entropy regularizer. The transport plans are unchanged. `reg_type` is now handled case insensitively by all the methods and an unknown value raises a `ValueError` instead of silently falling back to `"kl"`. The docstrings now state the generalized KL convention, the regularizer `Omega_reg_type(gamma, c)` with its two cases, the keys of the `log` dictionary and the multi-histogram behaviour (PR #874) - Fix `ot.dist` ignoring the weights `w` for `metric="cityblock"`, which returned the unweighted distance although the weights are documented for this metric (PR #859) - Fix swapped arguments to `div_to_product` in `ot.gromov.fused_unbalanced_across_spaces_cost`: with `reg_type="independent"` (UCOOT) the entropic terms used the plan marginals as the reference measures and vice versa (PR #855, Issue #854) - Fix device placement in `ot.batch.bregman_projection_batch` so `ot.solve_batch(..., method="sinkhorn")` no longer crashes on GPU when the torch default device is CPU (PR #851) diff --git a/ot/unbalanced/_sinkhorn.py b/ot/unbalanced/_sinkhorn.py index d338e1652..b32507cec 100644 --- a/ot/unbalanced/_sinkhorn.py +++ b/ot/unbalanced/_sinkhorn.py @@ -15,6 +15,32 @@ from ..backend import get_backend from ..utils import list_to_array, get_parameter_pair +REG_TYPES = ("kl", "entropy") + + +def _check_reg_type(reg_type): + r"""Validate `reg_type` and return its lower case form. + + Only ``'kl'`` and ``'entropy'`` are implemented. Any other value used to be + silently handled as ``'kl'``, which silently changed the solved problem, so + it is now rejected explicitly. + + Parameters + ---------- + reg_type : str + Name of the regularizer, either 'kl' or 'entropy' (case insensitive). + + Returns + ------- + reg_type : str + The lower case name of the regularizer. + """ + if not isinstance(reg_type, str) or reg_type.lower() not in REG_TYPES: + raise ValueError( + "Unknown reg_type '{}'. Must be either 'kl' or 'entropy'.".format(reg_type) + ) + return reg_type.lower() + def sinkhorn_unbalanced( a, @@ -40,7 +66,7 @@ def sinkhorn_unbalanced( .. math:: W = \arg \min_\gamma \ \langle \gamma, \mathbf{M} \rangle_F + - \mathrm{reg} \cdot \mathrm{KL}(\gamma, \mathbf{c}) + + \mathrm{reg} \cdot \Omega_{\texttt{reg\_type}}(\gamma, \mathbf{c}) + \mathrm{reg_{m1}} \cdot \mathrm{KL}(\gamma \mathbf{1}, \mathbf{a}) + \mathrm{reg_{m2}} \cdot \mathrm{KL}(\gamma^T \mathbf{1}, \mathbf{b}) @@ -52,7 +78,17 @@ def sinkhorn_unbalanced( - :math:`\mathbf{M}` is the (`dim_a`, `dim_b`) metric cost matrix - :math:`\mathbf{a}` and :math:`\mathbf{b}` are source and target unbalanced distributions - :math:`\mathbf{c}` is a reference distribution for the regularization - - KL is the Kullback-Leibler divergence + - :math:`\mathrm{KL}` is the generalized Kullback-Leibler divergence + :math:`\mathrm{KL}(\mathbf{P}, \mathbf{Q}) = \sum_{i,j} \mathbf{P}_{i,j} \log(\mathbf{P}_{i,j} / \mathbf{Q}_{i,j}) - \mathbf{P}_{i,j} + \mathbf{Q}_{i,j}` + + and :math:`\Omega_{\texttt{reg\_type}}(\gamma, \mathbf{c})` is the regularizer + + .. math:: + \Omega_{\texttt{reg\_type}}(\gamma, \mathbf{c}) = + \begin{cases} + \mathrm{KL}(\gamma, \mathbf{c}) & \text{if } \texttt{reg\_type} = \text{'kl'}, \\ + \sum_{i,j} \gamma_{i,j} \log(\gamma_{i,j}) - \sum_{i,j} \gamma_{i,j} & \text{if } \texttt{reg\_type} = \text{'entropy'}. + \end{cases} The algorithm used for solving the problem is the generalized Sinkhorn-Knopp matrix scaling algorithm as proposed in :ref:`[10, 25] @@ -92,13 +128,19 @@ def sinkhorn_unbalanced( method used for the solver either 'sinkhorn', 'sinkhorn_stabilized', 'sinkhorn_translation_invariant' or 'sinkhorn_reg_scaling', see those function for specific parameters reg_type : string, optional - Regularizer term. Can take two values: + Name of the regularizer :math:`\Omega_{\texttt{reg\_type}}(\gamma, \mathbf{c})`, + case insensitive. Can take two values: - Negative entropy: 'entropy': :math:`\Omega(\gamma) = \sum_{i,j} \gamma_{i,j} \log(\gamma_{i,j}) - \sum_{i,j} \gamma_{i,j}`. - This is equivalent (up to a constant) to :math:`\Omega(\gamma) = \text{KL}(\gamma, 1_{dim_a} 1_{dim_b}^T)`. - - Kullback-Leibler divergence (default): 'kl': - :math:`\Omega(\gamma) = \text{KL}(\gamma, \mathbf{a} \mathbf{b}^T)`. + It is equal to :math:`\mathrm{KL}(\gamma, 1_{dim_a} 1_{dim_b}^T)` minus the + constant :math:`dim_a \times dim_b`, and the reference measure :math:`\mathbf{c}` + is overwritten by the all-ones matrix. + - Generalized Kullback-Leibler divergence (default): 'kl': + :math:`\Omega(\gamma, \mathbf{c}) = \mathrm{KL}(\gamma, \mathbf{c})` with + :math:`\mathbf{c} = \mathbf{a} \mathbf{b}^T` when `c` is None. + + Any other value raises a :class:`ValueError`. c : array-like, shape (dim_a, dim_b), optional (default=None) Reference measure for the regularization. If None, then use :math:`\mathbf{c} = \mathbf{a} \mathbf{b}^T`. @@ -118,7 +160,7 @@ def sinkhorn_unbalanced( Returns ------- - if n_hists == 1: + if n_hists == 0: - gamma : array-like, shape(dim_a, dim_b) Optimal transportation matrix for the given parameters - log : dict @@ -129,6 +171,19 @@ def sinkhorn_unbalanced( - log : dict log dictionary returned only if `log` is `True` + Notes + ----- + When `log=True`, the returned dictionary is the one of the solver selected by + `method`, see :any:`ot.unbalanced.sinkhorn_knopp_unbalanced`, + :any:`ot.unbalanced.sinkhorn_stabilized_unbalanced` and + :any:`ot.unbalanced.sinkhorn_unbalanced_translation_invariant`. + + .. note:: + When `b` is a 2d array of shape (`dim_b`, `n_hists`) -- which includes + :math:`n_{hists} = 1` -- only the negative entropy regularization is + implemented: `reg_type` and `c` are ignored, and only the linear cost is + returned. + Examples -------- @@ -274,7 +329,7 @@ def sinkhorn_unbalanced2( .. math:: \min_\gamma \quad \langle \gamma, \mathbf{M} \rangle_F + - \mathrm{reg} \cdot \mathrm{KL}(\gamma, \mathbf{c}) + + \mathrm{reg} \cdot \Omega_{\texttt{reg\_type}}(\gamma, \mathbf{c}) + \mathrm{reg_{m1}} \cdot \mathrm{KL}(\gamma \mathbf{1}, \mathbf{a}) + \mathrm{reg_{m2}} \cdot \mathrm{KL}(\gamma^T \mathbf{1}, \mathbf{b}) @@ -285,7 +340,17 @@ def sinkhorn_unbalanced2( - :math:`\mathbf{M}` is the (`dim_a`, `dim_b`) metric cost matrix - :math:`\mathbf{a}` and :math:`\mathbf{b}` are source and target unbalanced distributions - :math:`\mathbf{c}` is a reference distribution for the regularization - - KL is the Kullback-Leibler divergence + - :math:`\mathrm{KL}` is the generalized Kullback-Leibler divergence + :math:`\mathrm{KL}(\mathbf{P}, \mathbf{Q}) = \sum_{i,j} \mathbf{P}_{i,j} \log(\mathbf{P}_{i,j} / \mathbf{Q}_{i,j}) - \mathbf{P}_{i,j} + \mathbf{Q}_{i,j}` + + and :math:`\Omega_{\texttt{reg\_type}}(\gamma, \mathbf{c})` is the regularizer + + .. math:: + \Omega_{\texttt{reg\_type}}(\gamma, \mathbf{c}) = + \begin{cases} + \mathrm{KL}(\gamma, \mathbf{c}) & \text{if } \texttt{reg\_type} = \text{'kl'}, \\ + \sum_{i,j} \gamma_{i,j} \log(\gamma_{i,j}) - \sum_{i,j} \gamma_{i,j} & \text{if } \texttt{reg\_type} = \text{'entropy'}. + \end{cases} The algorithm used for solving the problem is the generalized Sinkhorn-Knopp matrix scaling algorithm as proposed in :ref:`[10, 25] @@ -324,13 +389,19 @@ def sinkhorn_unbalanced2( method used for the solver either 'sinkhorn', 'sinkhorn_stabilized', 'sinkhorn_translation_invariant' or 'sinkhorn_reg_scaling', see those function for specific parameters reg_type : string, optional - Regularizer term. Can take two values: + Name of the regularizer :math:`\Omega_{\texttt{reg\_type}}(\gamma, \mathbf{c})`, + case insensitive. Can take two values: - Negative entropy: 'entropy': :math:`\Omega(\gamma) = \sum_{i,j} \gamma_{i,j} \log(\gamma_{i,j}) - \sum_{i,j} \gamma_{i,j}`. - This is equivalent (up to a constant) to :math:`\Omega(\gamma) = \text{KL}(\gamma, 1_{dim_a} 1_{dim_b}^T)`. - - Kullback-Leibler divergence: 'kl': - :math:`\Omega(\gamma) = \text{KL}(\gamma, \mathbf{a} \mathbf{b}^T)`. + It is equal to :math:`\mathrm{KL}(\gamma, 1_{dim_a} 1_{dim_b}^T)` minus the + constant :math:`dim_a \times dim_b`, and the reference measure :math:`\mathbf{c}` + is overwritten by the all-ones matrix. + - Generalized Kullback-Leibler divergence (default): 'kl': + :math:`\Omega(\gamma, \mathbf{c}) = \mathrm{KL}(\gamma, \mathbf{c})` with + :math:`\mathbf{c} = \mathbf{a} \mathbf{b}^T` when `c` is None. + + Any other value raises a :class:`ValueError`. c : array-like, shape (dim_a, dim_b), optional (default=None) Reference measure for the regularization. If None, then use :math:`\mathbf{c} = \mathbf{a} \mathbf{b}^T`. @@ -358,6 +429,19 @@ def sinkhorn_unbalanced2( log : dict log dictionary returned only if `log` is `True` + Notes + ----- + When `log=True`, the returned dictionary is the one of the solver selected by + `method`, see :any:`ot.unbalanced.sinkhorn_knopp_unbalanced`, + :any:`ot.unbalanced.sinkhorn_stabilized_unbalanced` and + :any:`ot.unbalanced.sinkhorn_unbalanced_translation_invariant`. + + .. note:: + When `b` is a 2d array of shape (`dim_b`, `n_hists`) -- which includes + :math:`n_{hists} = 1` -- only the negative entropy regularization is + implemented: `reg_type` and `c` are ignored, and only the linear cost is + returned. + Examples -------- @@ -487,8 +571,13 @@ def sinkhorn_unbalanced2( return cost else: - if reg_type == "kl": - warnings.warn("Reg_type not implemented yet. Use entropy.") + if returnCost not in ("linear", "total"): + raise ValueError("Unknown returnCost = {}".format(returnCost)) + if returnCost != "linear": + warnings.warn( + "returnCost='total' is not available with a 2d b (n_hists={}): " + "the linear cost is returned.".format(b.shape[1]) + ) if method.lower() == "sinkhorn": return sinkhorn_knopp_unbalanced( @@ -585,7 +674,7 @@ def sinkhorn_knopp_unbalanced( .. math:: W = \arg \min_\gamma \quad \langle \gamma, \mathbf{M} \rangle_F + - \mathrm{reg} \cdot \mathrm{KL}(\gamma, \mathbf{c}) + + \mathrm{reg} \cdot \Omega_{\texttt{reg\_type}}(\gamma, \mathbf{c}) + \mathrm{reg_{m1}} \cdot \mathrm{KL}(\gamma \mathbf{1}, \mathbf{a}) + \mathrm{reg_{m2}} \cdot \mathrm{KL}(\gamma^T \mathbf{1}, \mathbf{b}) @@ -597,7 +686,17 @@ def sinkhorn_knopp_unbalanced( - :math:`\mathbf{M}` is the (`dim_a`, `dim_b`) metric cost matrix - :math:`\mathbf{a}` and :math:`\mathbf{b}` are source and target unbalanced distributions - :math:`\mathbf{c}` is a reference distribution for the regularization - - KL is the Kullback-Leibler divergence + - :math:`\mathrm{KL}` is the generalized Kullback-Leibler divergence + :math:`\mathrm{KL}(\mathbf{P}, \mathbf{Q}) = \sum_{i,j} \mathbf{P}_{i,j} \log(\mathbf{P}_{i,j} / \mathbf{Q}_{i,j}) - \mathbf{P}_{i,j} + \mathbf{Q}_{i,j}` + + and :math:`\Omega_{\texttt{reg\_type}}(\gamma, \mathbf{c})` is the regularizer + + .. math:: + \Omega_{\texttt{reg\_type}}(\gamma, \mathbf{c}) = + \begin{cases} + \mathrm{KL}(\gamma, \mathbf{c}) & \text{if } \texttt{reg\_type} = \text{'kl'}, \\ + \sum_{i,j} \gamma_{i,j} \log(\gamma_{i,j}) - \sum_{i,j} \gamma_{i,j} & \text{if } \texttt{reg\_type} = \text{'entropy'}. + \end{cases} The algorithm used for solving the problem is the generalized Sinkhorn-Knopp matrix scaling algorithm as proposed in :ref:`[10, 25] ` @@ -632,13 +731,19 @@ def sinkhorn_knopp_unbalanced( If :math:`\mathrm{reg_{m}}` is an array, it must have the same backend as input arrays `(a, b, M)`. reg_type : string, optional - Regularizer term. Can take two values: + Name of the regularizer :math:`\Omega_{\texttt{reg\_type}}(\gamma, \mathbf{c})`, + case insensitive. Can take two values: - Negative entropy: 'entropy': :math:`\Omega(\gamma) = \sum_{i,j} \gamma_{i,j} \log(\gamma_{i,j}) - \sum_{i,j} \gamma_{i,j}`. - This is equivalent (up to a constant) to :math:`\Omega(\gamma) = \text{KL}(\gamma, 1_{dim_a} 1_{dim_b}^T)`. - - Kullback-Leibler divergence: 'kl': - :math:`\Omega(\gamma) = \text{KL}(\gamma, \mathbf{a} \mathbf{b}^T)`. + It is equal to :math:`\mathrm{KL}(\gamma, 1_{dim_a} 1_{dim_b}^T)` minus the + constant :math:`dim_a \times dim_b`, and the reference measure :math:`\mathbf{c}` + is overwritten by the all-ones matrix. + - Generalized Kullback-Leibler divergence (default): 'kl': + :math:`\Omega(\gamma, \mathbf{c}) = \mathrm{KL}(\gamma, \mathbf{c})` with + :math:`\mathbf{c} = \mathbf{a} \mathbf{b}^T` when `c` is None. + + Any other value raises a :class:`ValueError`. c : array-like, shape (dim_a, dim_b), optional (default=None) Reference measure for the regularization. If None, then use :math:`\mathbf{c} = \mathbf{a} \mathbf{b}^T`. @@ -658,7 +763,7 @@ def sinkhorn_knopp_unbalanced( Returns ------- - if n_hists == 1: + if n_hists == 0: - gamma : array-like, shape (dim_a, dim_b) Optimal transportation matrix for the given parameters - log : dict @@ -669,6 +774,25 @@ def sinkhorn_knopp_unbalanced( - log : dict log dictionary returned only if `log` is `True` + Notes + ----- + The `log` dictionary returned when `log=True` contains: + + - 'err' : list of float, the error at each iteration; + - 'logu', 'logv' : array-like, the log of the scaling vectors; + - 'cost' : float, the linear cost :math:`\langle \gamma, \mathbf{M} \rangle_F`; + - 'total_cost' : float, the value of the optimization problem above. + + 'cost' and 'total_cost' are only computed when `b` is a single histogram. + + .. note:: + When `b` is a 2d array of shape (`dim_b`, `n_hists`) -- which includes + :math:`n_{hists} = 1`, since any 2d `b` selects this branch -- the function + returns the cost of each column instead of a plan, and **only the negative + entropy regularization is implemented**: `reg_type` and `c` are ignored and + the reference measure is the all-ones matrix. In that case `log` only + contains 'err', 'logu' and 'logv'. + Examples -------- @@ -702,6 +826,8 @@ def sinkhorn_knopp_unbalanced( M, a, b = list_to_array(M, a, b) nx = get_backend(M, a, b) + reg_type = _check_reg_type(reg_type) + dim_a, dim_b = M.shape if len(a) == 0: @@ -714,6 +840,13 @@ def sinkhorn_knopp_unbalanced( else: n_hists = 0 + if n_hists and (reg_type != "entropy" or c is not None): + warnings.warn( + "With a 2d b (n_hists={}) only the negative entropy regularization is " + "implemented: reg_type and c are ignored and the reference measure is " + "the all-ones matrix.".format(n_hists) + ) + reg_m1, reg_m2 = get_parameter_pair(reg_m) if log: @@ -732,10 +865,13 @@ def sinkhorn_knopp_unbalanced( else: u, v = nx.exp(warmstart[0]), nx.exp(warmstart[1]) - if reg_type.lower() == "entropy": - warnings.warn( - "If reg_type = entropy, then the matrix c is overwritten by the one matrix." - ) + if reg_type == "entropy": + # in 2d mode the warning above already reports that c is ignored + if c is not None and not n_hists: + warnings.warn( + "reg_type='entropy' ignores the provided c: the reference measure " + "of the regularization is the all-ones matrix." + ) c = nx.ones((dim_a, dim_b), type_as=M) if n_hists: @@ -806,11 +942,25 @@ def sinkhorn_knopp_unbalanced( linear_cost = nx.sum(plan * M) dict_log["cost"] = linear_cost - total_cost = linear_cost + reg * nx.kl_div(plan, c) + # The regularizer is the generalized KL divergence + # KL(plan, c) = sum(plan * log(plan / c) - plan + c). + # With reg_type="entropy" the reference measure is c = 1 and the + # regularizer is the negative entropy + # Omega(plan) = sum(plan * log(plan) - plan) + # = KL(plan, 1) - dim_a * dim_b, + # so the constant dim_a * dim_b must be removed. + reg_cost = nx.kl_div(plan, c, mass=True) + if reg_type == "entropy": + reg_cost = reg_cost - dim_a * dim_b + total_cost = linear_cost + reg * reg_cost if reg_m1 != float("inf"): - total_cost = total_cost + reg_m1 * nx.kl_div(nx.sum(plan, 1), a) + total_cost = total_cost + reg_m1 * nx.kl_div( + nx.sum(plan, 1), a, mass=True + ) if reg_m2 != float("inf"): - total_cost = total_cost + reg_m2 * nx.kl_div(nx.sum(plan, 0), b) + total_cost = total_cost + reg_m2 * nx.kl_div( + nx.sum(plan, 0), b, mass=True + ) dict_log["total_cost"] = total_cost return plan, dict_log @@ -843,7 +993,7 @@ def sinkhorn_stabilized_unbalanced( .. math:: W = \arg \min_\gamma \quad \langle \gamma, \mathbf{M} \rangle_F + - \mathrm{reg} \cdot \mathrm{KL}(\gamma, \mathbf{c}) + + \mathrm{reg} \cdot \Omega_{\texttt{reg\_type}}(\gamma, \mathbf{c}) + \mathrm{reg_{m1}} \cdot \mathrm{KL}(\gamma \mathbf{1}, \mathbf{a}) + \mathrm{reg_{m2}} \cdot \mathrm{KL}(\gamma^T \mathbf{1}, \mathbf{b}) @@ -855,7 +1005,17 @@ def sinkhorn_stabilized_unbalanced( - :math:`\mathbf{M}` is the (`dim_a`, `dim_b`) metric cost matrix - :math:`\mathbf{a}` and :math:`\mathbf{b}` are source and target unbalanced distributions - :math:`\mathbf{c}` is a reference distribution for the regularization - - KL is the Kullback-Leibler divergence + - :math:`\mathrm{KL}` is the generalized Kullback-Leibler divergence + :math:`\mathrm{KL}(\mathbf{P}, \mathbf{Q}) = \sum_{i,j} \mathbf{P}_{i,j} \log(\mathbf{P}_{i,j} / \mathbf{Q}_{i,j}) - \mathbf{P}_{i,j} + \mathbf{Q}_{i,j}` + + and :math:`\Omega_{\texttt{reg\_type}}(\gamma, \mathbf{c})` is the regularizer + + .. math:: + \Omega_{\texttt{reg\_type}}(\gamma, \mathbf{c}) = + \begin{cases} + \mathrm{KL}(\gamma, \mathbf{c}) & \text{if } \texttt{reg\_type} = \text{'kl'}, \\ + \sum_{i,j} \gamma_{i,j} \log(\gamma_{i,j}) - \sum_{i,j} \gamma_{i,j} & \text{if } \texttt{reg\_type} = \text{'entropy'}. + \end{cases} The algorithm used for solving the problem is the generalized Sinkhorn-Knopp matrix scaling algorithm as proposed in :ref:`[10, 25] ` @@ -890,13 +1050,19 @@ def sinkhorn_stabilized_unbalanced( method used for the solver either 'sinkhorn', 'sinkhorn_stabilized' or 'sinkhorn_reg_scaling', see those function for specific parameters reg_type : string, optional - Regularizer term. Can take two values: + Name of the regularizer :math:`\Omega_{\texttt{reg\_type}}(\gamma, \mathbf{c})`, + case insensitive. Can take two values: - Negative entropy: 'entropy': :math:`\Omega(\gamma) = \sum_{i,j} \gamma_{i,j} \log(\gamma_{i,j}) - \sum_{i,j} \gamma_{i,j}`. - This is equivalent (up to a constant) to :math:`\Omega(\gamma) = \text{KL}(\gamma, 1_{dim_a} 1_{dim_b}^T)`. - - Kullback-Leibler divergence: 'kl': - :math:`\Omega(\gamma) = \text{KL}(\gamma, \mathbf{a} \mathbf{b}^T)`. + It is equal to :math:`\mathrm{KL}(\gamma, 1_{dim_a} 1_{dim_b}^T)` minus the + constant :math:`dim_a \times dim_b`, and the reference measure :math:`\mathbf{c}` + is overwritten by the all-ones matrix. + - Generalized Kullback-Leibler divergence (default): 'kl': + :math:`\Omega(\gamma, \mathbf{c}) = \mathrm{KL}(\gamma, \mathbf{c})` with + :math:`\mathbf{c} = \mathbf{a} \mathbf{b}^T` when `c` is None. + + Any other value raises a :class:`ValueError`. c : array-like, shape (dim_a, dim_b), optional (default=None) Reference measure for the regularization. If None, then use :math:`\mathbf{c} = \mathbf{a} \mathbf{b}^T`. @@ -922,7 +1088,7 @@ def sinkhorn_stabilized_unbalanced( Returns ------- - if n_hists == 1: + if n_hists == 0: - gamma : array-like, shape (dim_a, dim_b) Optimal transportation matrix for the given parameters - log : dict @@ -932,6 +1098,25 @@ def sinkhorn_stabilized_unbalanced( the OT cost between :math:`\mathbf{a}` and each of the histograms :math:`\mathbf{b}_i` - log : dict log dictionary returned only if `log` is `True` + Notes + ----- + The `log` dictionary returned when `log=True` contains: + + - 'err' : list of float, the error at each iteration; + - 'logu', 'logv' : array-like, the log of the scaling vectors; + - 'cost' : float, the linear cost :math:`\langle \gamma, \mathbf{M} \rangle_F`; + - 'total_cost' : float, the value of the optimization problem above. + + 'cost' and 'total_cost' are only computed when `b` is a single histogram. + + .. note:: + When `b` is a 2d array of shape (`dim_b`, `n_hists`) -- which includes + :math:`n_{hists} = 1`, since any 2d `b` selects this branch -- the function + returns the cost of each column instead of a plan, and **only the negative + entropy regularization is implemented**: `reg_type` and `c` are ignored and + the reference measure is the all-ones matrix. In that case `log` only + contains 'err', 'logu' and 'logv'. + Examples -------- @@ -964,6 +1149,8 @@ def sinkhorn_stabilized_unbalanced( a, b, M = list_to_array(a, b, M) nx = get_backend(M, a, b) + reg_type = _check_reg_type(reg_type) + dim_a, dim_b = M.shape if len(a) == 0: @@ -976,6 +1163,13 @@ def sinkhorn_stabilized_unbalanced( else: n_hists = 0 + if n_hists and (reg_type != "entropy" or c is not None): + warnings.warn( + "With a 2d b (n_hists={}) only the negative entropy regularization is " + "implemented: reg_type and c are ignored and the reference measure is " + "the all-ones matrix.".format(n_hists) + ) + reg_m1, reg_m2 = get_parameter_pair(reg_m) if log: @@ -995,9 +1189,12 @@ def sinkhorn_stabilized_unbalanced( u, v = nx.exp(warmstart[0]), nx.exp(warmstart[1]) if reg_type == "entropy": - warnings.warn( - "If reg_type = entropy, then the matrix c is overwritten by the one matrix." - ) + # in 2d mode the warning above already reports that c is ignored + if c is not None and not n_hists: + warnings.warn( + "reg_type='entropy' ignores the provided c: the reference measure " + "of the regularization is the all-ones matrix." + ) c = nx.ones((dim_a, dim_b), type_as=M) if n_hists: @@ -1106,11 +1303,25 @@ def sinkhorn_stabilized_unbalanced( linear_cost = nx.sum(plan * M) dict_log["cost"] = linear_cost - total_cost = linear_cost + reg * nx.kl_div(plan, c) + # The regularizer is the generalized KL divergence + # KL(plan, c) = sum(plan * log(plan / c) - plan + c). + # With reg_type="entropy" the reference measure is c = 1 and the + # regularizer is the negative entropy + # Omega(plan) = sum(plan * log(plan) - plan) + # = KL(plan, 1) - dim_a * dim_b, + # so the constant dim_a * dim_b must be removed. + reg_cost = nx.kl_div(plan, c, mass=True) + if reg_type == "entropy": + reg_cost = reg_cost - dim_a * dim_b + total_cost = linear_cost + reg * reg_cost if reg_m1 != float("inf"): - total_cost = total_cost + reg_m1 * nx.kl_div(nx.sum(plan, 1), a) + total_cost = total_cost + reg_m1 * nx.kl_div( + nx.sum(plan, 1), a, mass=True + ) if reg_m2 != float("inf"): - total_cost = total_cost + reg_m2 * nx.kl_div(nx.sum(plan, 0), b) + total_cost = total_cost + reg_m2 * nx.kl_div( + nx.sum(plan, 0), b, mass=True + ) dict_log["total_cost"] = total_cost return plan, dict_log @@ -1141,7 +1352,7 @@ def sinkhorn_unbalanced_translation_invariant( .. math:: W = \arg \min_\gamma \ \langle \gamma, \mathbf{M} \rangle_F + - \mathrm{reg} \cdot \mathrm{KL}(\gamma, \mathbf{c}) + + \mathrm{reg} \cdot \Omega_{\texttt{reg\_type}}(\gamma, \mathbf{c}) + \mathrm{reg_{m1}} \cdot \mathrm{KL}(\gamma \mathbf{1}, \mathbf{a}) + \mathrm{reg_{m2}} \cdot \mathrm{KL}(\gamma^T \mathbf{1}, \mathbf{b}) @@ -1153,7 +1364,17 @@ def sinkhorn_unbalanced_translation_invariant( - :math:`\mathbf{M}` is the (`dim_a`, `dim_b`) metric cost matrix - :math:`\Omega` is the entropic regularization term,KL divergence - :math:`\mathbf{a}` and :math:`\mathbf{b}` are source and target unbalanced distributions - - KL is the Kullback-Leibler divergence + - :math:`\mathrm{KL}` is the generalized Kullback-Leibler divergence + :math:`\mathrm{KL}(\mathbf{P}, \mathbf{Q}) = \sum_{i,j} \mathbf{P}_{i,j} \log(\mathbf{P}_{i,j} / \mathbf{Q}_{i,j}) - \mathbf{P}_{i,j} + \mathbf{Q}_{i,j}` + + and :math:`\Omega_{\texttt{reg\_type}}(\gamma, \mathbf{c})` is the regularizer + + .. math:: + \Omega_{\texttt{reg\_type}}(\gamma, \mathbf{c}) = + \begin{cases} + \mathrm{KL}(\gamma, \mathbf{c}) & \text{if } \texttt{reg\_type} = \text{'kl'}, \\ + \sum_{i,j} \gamma_{i,j} \log(\gamma_{i,j}) - \sum_{i,j} \gamma_{i,j} & \text{if } \texttt{reg\_type} = \text{'entropy'}. + \end{cases} The algorithm used for solving the problem is the translation invariant Sinkhorn algorithm as proposed in :ref:`[73] ` @@ -1177,11 +1398,19 @@ def sinkhorn_unbalanced_translation_invariant( `reg_m=(float("inf"), scalar)` or `reg_m=(scalar, float("inf"))`. If reg_m is an array, it must have the same backend as input arrays (a, b, M). reg_type : string, optional - Regularizer term. Can take two values: - 'entropy' (negative entropy) - :math:`\Omega(\gamma) = \sum_{i,j} \gamma_{i,j} \log(\gamma_{i,j}) - \sum_{i,j} \gamma_{i,j}`, or - 'kl' (Kullback-Leibler) - :math:`\Omega(\gamma) = \text{KL}(\gamma, \mathbf{a} \mathbf{b}^T)`. + Name of the regularizer :math:`\Omega_{\texttt{reg\_type}}(\gamma, \mathbf{c})`, + case insensitive. Can take two values: + + - Negative entropy: 'entropy': + :math:`\Omega(\gamma) = \sum_{i,j} \gamma_{i,j} \log(\gamma_{i,j}) - \sum_{i,j} \gamma_{i,j}`. + It is equal to :math:`\mathrm{KL}(\gamma, 1_{dim_a} 1_{dim_b}^T)` minus the + constant :math:`dim_a \times dim_b`, and the reference measure :math:`\mathbf{c}` + is overwritten by the all-ones matrix. + - Generalized Kullback-Leibler divergence (default): 'kl': + :math:`\Omega(\gamma, \mathbf{c}) = \mathrm{KL}(\gamma, \mathbf{c})` with + :math:`\mathbf{c} = \mathbf{a} \mathbf{b}^T` when `c` is None. + + Any other value raises a :class:`ValueError`. c : array-like, shape (dim_a, dim_b), optional (default=None) Reference measure for the regularization. If None, then use :math:`\mathbf{c} = \mathbf{a} \mathbf{b}^T`. @@ -1201,7 +1430,7 @@ def sinkhorn_unbalanced_translation_invariant( Returns ------- - if n_hists == 1: + if n_hists == 0: - gamma : array-like, shape (dim_a, dim_b) Optimal transportation matrix for the given parameters - log : dict @@ -1212,6 +1441,25 @@ def sinkhorn_unbalanced_translation_invariant( - log : dict log dictionary returned only if `log` is `True` + Notes + ----- + The `log` dictionary returned when `log=True` contains: + + - 'err' : list of float, the error at each iteration; + - 'logu', 'logv' : array-like, the log of the scaling vectors; + - 'cost' : float, the linear cost :math:`\langle \gamma, \mathbf{M} \rangle_F`; + - 'total_cost' : float, the value of the optimization problem above. + + 'cost' and 'total_cost' are only computed when `b` is a single histogram. + + .. note:: + When `b` is a 2d array of shape (`dim_b`, `n_hists`) -- which includes + :math:`n_{hists} = 1`, since any 2d `b` selects this branch -- the function + returns the cost of each column instead of a plan, and **only the negative + entropy regularization is implemented**: `reg_type` and `c` are ignored and + the reference measure is the all-ones matrix. In that case `log` only + contains 'err', 'logu' and 'logv'. + Examples -------- @@ -1235,6 +1483,8 @@ def sinkhorn_unbalanced_translation_invariant( M, a, b = list_to_array(M, a, b) nx = get_backend(M, a, b) + reg_type = _check_reg_type(reg_type) + dim_a, dim_b = M.shape if len(a) == 0: @@ -1247,6 +1497,13 @@ def sinkhorn_unbalanced_translation_invariant( else: n_hists = 0 + if n_hists and (reg_type != "entropy" or c is not None): + warnings.warn( + "With a 2d b (n_hists={}) only the negative entropy regularization is " + "implemented: reg_type and c are ignored and the reference measure is " + "the all-ones matrix.".format(n_hists) + ) + reg_m1, reg_m2 = get_parameter_pair(reg_m) if log: @@ -1268,9 +1525,12 @@ def sinkhorn_unbalanced_translation_invariant( u_, v_ = u, v if reg_type == "entropy": - warnings.warn( - "If reg_type = entropy, then the matrix c is overwritten by the one matrix." - ) + # in 2d mode the warning above already reports that c is ignored + if c is not None and not n_hists: + warnings.warn( + "reg_type='entropy' ignores the provided c: the reference measure " + "of the regularization is the all-ones matrix." + ) c = nx.ones((dim_a, dim_b), type_as=M) if n_hists: @@ -1389,11 +1649,25 @@ def sinkhorn_unbalanced_translation_invariant( linear_cost = nx.sum(plan * M) dict_log["cost"] = linear_cost - total_cost = linear_cost + reg * nx.kl_div(plan, c) + # The regularizer is the generalized KL divergence + # KL(plan, c) = sum(plan * log(plan / c) - plan + c). + # With reg_type="entropy" the reference measure is c = 1 and the + # regularizer is the negative entropy + # Omega(plan) = sum(plan * log(plan) - plan) + # = KL(plan, 1) - dim_a * dim_b, + # so the constant dim_a * dim_b must be removed. + reg_cost = nx.kl_div(plan, c, mass=True) + if reg_type == "entropy": + reg_cost = reg_cost - dim_a * dim_b + total_cost = linear_cost + reg * reg_cost if reg_m1 != float("inf"): - total_cost = total_cost + reg_m1 * nx.kl_div(nx.sum(plan, 1), a) + total_cost = total_cost + reg_m1 * nx.kl_div( + nx.sum(plan, 1), a, mass=True + ) if reg_m2 != float("inf"): - total_cost = total_cost + reg_m2 * nx.kl_div(nx.sum(plan, 0), b) + total_cost = total_cost + reg_m2 * nx.kl_div( + nx.sum(plan, 0), b, mass=True + ) dict_log["total_cost"] = total_cost return plan, dict_log diff --git a/test/unbalanced/test_sinkhorn.py b/test/unbalanced/test_sinkhorn.py index be7694309..9369d8ad5 100644 --- a/test/unbalanced/test_sinkhorn.py +++ b/test/unbalanced/test_sinkhorn.py @@ -7,6 +7,8 @@ # License: MIT License import itertools +import warnings + import numpy as np import ot import pytest @@ -809,3 +811,293 @@ def test_implemented_methods(nx): ot.unbalanced.sinkhorn_unbalanced(a, b, M, epsilon, reg_m, method=method) ot.unbalanced.sinkhorn_unbalanced2(a, b, M, epsilon, reg_m, method=method) barycenter_unbalanced(A, M, reg=epsilon, reg_m=reg_m, method=method) + + +METHODS_UNBALANCED = [ + "sinkhorn", + "sinkhorn_stabilized", + "sinkhorn_translation_invariant", +] + + +def _unbalanced_problem(nx, n=6, m=7, seed=0): + rng = np.random.RandomState(seed) + a = rng.rand(n) + a /= a.sum() + b = rng.rand(m) + b /= b.sum() + M = rng.rand(n, m) + M /= M.max() + return nx.from_numpy(a, b, M) + + +def _gen_kl(P, Q): + r"""Generalized KL divergence, computed with plain numpy. + + Deliberately independent from ``nx.kl_div`` so that the reference value does + not reuse the implementation under test. + """ + P = np.asarray(P, dtype=np.float64) + Q = np.asarray(Q, dtype=np.float64) + return float(np.sum(P * np.log(P / Q) - P + Q)) + + +def _negative_entropy(P): + r"""Negative entropy :math:`\sum_{ij} P_{ij} \log(P_{ij}) - P_{ij}`.""" + P = np.asarray(P, dtype=np.float64) + return float(np.sum(P * np.log(P) - P)) + + +@pytest.mark.parametrize("method", METHODS_UNBALANCED) +def test_unbalanced_total_cost_kl(nx, method): + """Non-regression test for the generalized KL objective (PR #874). + + The penalizations of the unbalanced OT problem are generalized KL + divergences, so the reported total cost must be at least the linear cost and + must match the objective recomputed from the returned plan. + """ + a, b, M = _unbalanced_problem(nx) + reg, reg_m = 0.3, 1.0 + + G, log = ot.unbalanced.sinkhorn_unbalanced( + a, + b, + M, + reg=reg, + reg_m=reg_m, + method=method, + reg_type="kl", + log=True, + numItermax=5000, + stopThr=1e-13, + ) + + G_np = nx.to_numpy(G) + a_np = nx.to_numpy(a) + b_np = nx.to_numpy(b) + M_np = nx.to_numpy(M) + c_np = a_np[:, None] * b_np[None, :] + + expected = ( + float(np.sum(G_np * M_np)) + + reg * _gen_kl(G_np, c_np) + + reg_m * _gen_kl(G_np.sum(1), a_np) + + reg_m * _gen_kl(G_np.sum(0), b_np) + ) + + total_cost = float(nx.to_numpy(log["total_cost"])) + np.testing.assert_allclose(total_cost, expected, atol=1e-6) + # with reg_type="kl" every penalization is a divergence, hence non negative + np.testing.assert_array_less(float(nx.to_numpy(log["cost"])) - 1e-5, total_cost) + + +@pytest.mark.parametrize("method", METHODS_UNBALANCED) +def test_unbalanced_total_cost_entropy(nx, method): + """The entropy regularizer must not carry the constant ``dim_a * dim_b``. + + The reported value must be the objective built on + :math:`\\Omega(\\gamma) = \\sum \\gamma \\log \\gamma - \\gamma`, which is + ``KL(gamma, 1) - dim_a * dim_b``. It may be negative, so no comparison with + the linear cost is made here. + """ + a, b, M = _unbalanced_problem(nx) + reg, reg_m = 0.3, 1.0 + + G, log = ot.unbalanced.sinkhorn_unbalanced( + a, + b, + M, + reg=reg, + reg_m=reg_m, + method=method, + reg_type="entropy", + log=True, + numItermax=5000, + stopThr=1e-13, + ) + + G_np = nx.to_numpy(G) + a_np = nx.to_numpy(a) + b_np = nx.to_numpy(b) + M_np = nx.to_numpy(M) + + expected = ( + float(np.sum(G_np * M_np)) + + reg * _negative_entropy(G_np) + + reg_m * _gen_kl(G_np.sum(1), a_np) + + reg_m * _gen_kl(G_np.sum(0), b_np) + ) + total_cost = float(nx.to_numpy(log["total_cost"])) + np.testing.assert_allclose(total_cost, expected, atol=1e-6) + + # the constant reg * dim_a * dim_b must have been removed + without_constant_removal = expected + reg * G_np.size + assert abs(total_cost - without_constant_removal) > 1e-3 + + +@pytest.mark.parametrize("method", METHODS_UNBALANCED) +def test_unbalanced_entropy_constant(nx, method): + """The entropy constant is exactly ``reg * dim_a * dim_b``. + + With ``c`` set to the all-ones matrix, ``reg_type="kl"`` and + ``reg_type="entropy"`` minimize the same objective up to that constant, so + the plans must coincide and the two reported total costs must differ by + exactly ``reg * dim_a * dim_b``. + """ + a, b, M = _unbalanced_problem(nx) + reg, reg_m = 0.3, 1.0 + n, m = nx.to_numpy(M).shape + c_ones = nx.from_numpy(np.ones((n, m))) + + G_kl, log_kl = ot.unbalanced.sinkhorn_unbalanced( + a, + b, + M, + reg=reg, + reg_m=reg_m, + method=method, + reg_type="kl", + c=c_ones, + log=True, + numItermax=5000, + stopThr=1e-13, + ) + G_en, log_en = ot.unbalanced.sinkhorn_unbalanced( + a, + b, + M, + reg=reg, + reg_m=reg_m, + method=method, + reg_type="entropy", + log=True, + numItermax=5000, + stopThr=1e-13, + ) + + # same objective up to a constant -> same plan + np.testing.assert_allclose(nx.to_numpy(G_kl), nx.to_numpy(G_en), atol=1e-8) + diff = float(nx.to_numpy(log_kl["total_cost"])) - float( + nx.to_numpy(log_en["total_cost"]) + ) + np.testing.assert_allclose(diff, reg * n * m, rtol=1e-6) + + +@pytest.mark.parametrize("method", METHODS_UNBALANCED) +@pytest.mark.parametrize("reg_type", ["kl", "entropy"]) +def test_unbalanced_reg_type_case_insensitive(nx, method, reg_type): + """`reg_type` must be handled case insensitively by every method.""" + a, b, M = _unbalanced_problem(nx) + reg, reg_m = 0.3, 1.0 + + ref_plan, ref_cost = None, None + for variant in [reg_type, reg_type.capitalize(), reg_type.upper()]: + G, log = ot.unbalanced.sinkhorn_unbalanced( + a, + b, + M, + reg=reg, + reg_m=reg_m, + method=method, + reg_type=variant, + log=True, + numItermax=5000, + stopThr=1e-13, + ) + plan = nx.to_numpy(G) + cost = float(nx.to_numpy(log["total_cost"])) + if ref_plan is None: + ref_plan, ref_cost = plan, cost + else: + np.testing.assert_allclose(plan, ref_plan, atol=1e-10) + np.testing.assert_allclose(cost, ref_cost, atol=1e-10) + + +@pytest.mark.parametrize("method", METHODS_UNBALANCED) +def test_unbalanced_unknown_reg_type_raises(nx, method): + """An unknown `reg_type` must be rejected instead of silently using 'kl'.""" + a, b, M = _unbalanced_problem(nx) + with pytest.raises(ValueError, match="Unknown reg_type"): + ot.unbalanced.sinkhorn_unbalanced( + a, b, M, reg=0.3, reg_m=1.0, method=method, reg_type="cryptic divergence" + ) + + +def test_unbalanced_multiple_inputs_reg_type_warns(nx): + """Multi histogram mode only implements the negative entropy regularization. + + `reg_type` and `c` are ignored in that mode, which must be advertised with a + single warning. Asking for 'entropy' (the implemented regularizer) with the + default reference measure must stay silent. + """ + a, b, M = _unbalanced_problem(nx) + n_hists = 3 + m = nx.to_numpy(M).shape[1] + B = nx.from_numpy(np.random.RandomState(1).rand(m, n_hists)) + c_custom = nx.from_numpy(np.full((nx.to_numpy(M).shape[0], m), 0.7)) + + def n_warnings(**kwargs): + with warnings.catch_warnings(record=True) as rec: + warnings.simplefilter("always") + ot.unbalanced.sinkhorn_unbalanced( + a, B, M, reg=0.3, reg_m=1.0, method="sinkhorn", **kwargs + ) + return [str(w.message) for w in rec if issubclass(w.category, UserWarning)] + + # the default reg_type is "kl", which is not implemented in that branch + msgs = n_warnings(reg_type="kl") + assert len(msgs) == 1, msgs + assert "2d b" in msgs[0] and "n_hists=3" in msgs[0] + + # "entropy" is the implemented regularizer and c is not given -> silent + assert n_warnings(reg_type="entropy") == [] + + # an explicit c is ignored too, but must not be reported twice + msgs = n_warnings(reg_type="entropy", c=c_custom) + assert len(msgs) == 1, msgs + + # a 2d b with a single column also takes that branch: the message must be exact + B1 = nx.from_numpy(np.random.RandomState(2).rand(m, 1)) + with warnings.catch_warnings(record=True) as rec: + warnings.simplefilter("always") + ot.unbalanced.sinkhorn_unbalanced( + a, B1, M, reg=0.3, reg_m=1.0, method="sinkhorn", reg_type="kl" + ) + msgs = [str(w.message) for w in rec if issubclass(w.category, UserWarning)] + assert len(msgs) == 1, msgs + assert "n_hists=1" in msgs[0] + + +def test_unbalanced_multiple_inputs_returnCost(nx): + """`returnCost` must be honoured or reported, never silently ignored. + + With a 2d `b` only the linear cost is available. `sinkhorn_unbalanced2` used to + ignore `returnCost` entirely in that branch. + """ + a, b, M = _unbalanced_problem(nx) + m = nx.to_numpy(M).shape[1] + B = nx.from_numpy(np.random.RandomState(1).rand(m, 3)) + + # "linear" is the only supported value and must not warn + with warnings.catch_warnings(): + warnings.simplefilter("error") + loss = ot.unbalanced.sinkhorn_unbalanced2( + a, B, M, reg=0.3, reg_m=1.0, method="sinkhorn", reg_type="entropy", + returnCost="linear", + ) + assert nx.to_numpy(loss).shape == (3,) + + # "total" is not available there: warn and still return the linear cost + with pytest.warns(UserWarning, match="returnCost='total' is not available"): + loss = ot.unbalanced.sinkhorn_unbalanced2( + a, B, M, reg=0.3, reg_m=1.0, method="sinkhorn", reg_type="entropy", + returnCost="total", + ) + assert nx.to_numpy(loss).shape == (3,) + + # an invalid value raises, as in the single histogram branch + with pytest.raises(ValueError, match="Unknown returnCost"): + ot.unbalanced.sinkhorn_unbalanced2( + a, B, M, reg=0.3, reg_m=1.0, method="sinkhorn", reg_type="entropy", + returnCost="invalid", + )