diff --git a/RELEASES.md b/RELEASES.md index d01ff927f..f79266793 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -18,6 +18,7 @@ #### Closed issues +- Fix `ot.solve(..., unbalanced=float("inf"))` (documented as the balanced limit): the exact and L-BFGSB paths returned a NaN cost, and L-BFGSB could return a plan that silently satisfied both marginals. The limit is now routed to the balanced solvers only where they minimize the same objective; `reg_type="L2"`, `"entropy"` and `"KL"` with a custom `c` (when `unbalanced_type` is not `"KL"`) raise `NotImplementedError` rather than silently solving a different problem, and one-sided infinite relaxations are rejected outside the entropic KL path. - Fix duplicate function and class descriptions in generated module API pages while preserving top-level aliases and example backreferences (Issue #527) - Fix internal arrays created on the torch default device instead of the input device in `ot.partial.partial_wasserstein`, `ot.wasserstein_circle` (and `ot.sliced_wasserstein_sphere`), the high-dimensional Bures-Wasserstein functions in `ot.gaussian`, `ot.gmm`, `ot.lowrank_sinkhorn` with `init="deterministic"`, the semi-relaxed (F)GW barycenters and `ot.utils.projection_sparse_simplex`, which made them fail on GPU inputs. Add a `meta_default_device` test fixture to catch such allocations on CPU (PR #883, Issue #852) - 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) diff --git a/ot/solvers/_linear.py b/ot/solvers/_linear.py index 32d21ec4b..9ef436204 100644 --- a/ot/solvers/_linear.py +++ b/ot/solvers/_linear.py @@ -8,7 +8,13 @@ # # License: MIT License -from ..utils import OTResult, dist, split_sample_ratio, apply_scaler +from ..utils import ( + OTResult, + dist, + split_sample_ratio, + apply_scaler, + get_parameter_pair, +) from ..lp import emd2, emd2_lazy, wasserstein_1d from ..backend import get_backend from ..unbalanced import mm_unbalanced, sinkhorn_knopp_unbalanced, lbfgsb_unbalanced @@ -58,6 +64,35 @@ all_valid_methods_solve_sample = lst_valid_methods_solve + lst_method_solve_sample +def _is_pos_inf(value): + r"""Return True when every entry of ``value`` is the ``+inf`` sentinel. + + ``unbalanced=float("inf")`` is a documented way to request hard marginal + constraints in :any:`ot.solve`. The sentinel is detected here, before + dispatch, because the solvers behind the exact (MM) and L-BFGSB paths do + not accept infinite marginal relaxations. Array-like inputs are only + reported as infinite when *all* their entries are ``+inf``; anything else + (including ``-inf`` or ``nan``) returns False and keeps the previous + behaviour. Traced values (e.g. a JAX tracer inside ``jit``) are not + supported and also return False. + + Parameters + ---------- + value : float or array-like + Marginal relaxation term to test. + + Returns + ------- + is_inf : bool + """ + try: + if hasattr(value, "all"): + return bool((value == float("inf")).all()) + return bool(value == float("inf")) + except (TypeError, ValueError, RuntimeError): + return False + + def solve( M, a=None, @@ -113,10 +148,22 @@ def solve( Marginal relaxation term. If it is a scalar or an indexable object of length 1, then the same relaxation is applied to both marginal relaxations. - The balanced OT can be recovered using :math:`unbalanced=float("inf")`. + The balanced OT can be recovered using :math:`unbalanced=float("inf")` + (both marginal relaxations infinite) when the hard-constraint limit of + the objective is the balanced problem itself: exact OT (``reg=0``), the + entropic KL solver, ``reg_type="KL"`` with the default reference measure + ``c = a b^T``, and a user supplied ``(f, df)`` pair. Other regularizers + (``reg_type="L2"``, ``"entropy"``, or ``"KL"`` with a custom ``c`` when + ``unbalanced_type`` is not ``"KL"``) raise :class:`NotImplementedError`, + because their balanced branch minimizes a different objective. A ``c`` + passed explicitly is treated as a custom reference measure even when it + is numerically equal to ``a b^T``. For semi-relaxed case, use either :math:`unbalanced=(float("inf"), scalar)` or - :math:`unbalanced=(scalar, float("inf"))`. + :math:`unbalanced=(scalar, float("inf"))`. Semi-relaxed OT is only + available for the entropic KL solver (``reg>0``, ``reg_type="KL"`` and + ``unbalanced_type="KL"``); other combinations raise + :class:`NotImplementedError`. If unbalanced is an array, it must have the same backend as input arrays `(a, b, M)`. unbalanced_type : str, optional @@ -307,11 +354,73 @@ def solve( if b is None: b = nx.ones(M.shape[1], type_as=M) / M.shape[1] if c is None: + c_is_default = True c = a[:, None] * b[None, :] + else: + c_is_default = False if reg is None: reg = 0 + # ``unbalanced=float("inf")`` is documented as the balanced limit and must + # be normalized *before* dispatch. Only the Sinkhorn KL path implements the + # sentinel itself: the exact (MM) path would build a kernel from + # ``inf / inf`` and the L-BFGSB path would multiply a zero penalty by + # ``inf``, silently returning NaN or solving the balanced problem instead of + # the requested semi-relaxed one. + unbalanced_is_inf = (False, False) + if unbalanced is not None: + rho_1, rho_2 = get_parameter_pair(unbalanced) + unbalanced_is_inf = (_is_pos_inf(rho_1), _is_pos_inf(rho_2)) + # The entropic KL Sinkhorn path implements the sentinel itself (it skips + # the infinite penalty term and keeps the hard constraint), and it is + # the only path that also honours `potentials_init` and `method`. It is + # therefore left untouched, so that already working calls keep their + # exact behaviour. + entropic_kl_path = ( + reg > 0 + and isinstance(reg_type, str) + and reg_type.lower() == "kl" + and isinstance(unbalanced_type, str) + and unbalanced_type.lower() == "kl" + ) + if all(unbalanced_is_inf) and not entropic_kl_path: + # A two-sided infinite relaxation is the hard-constraint limit of + # the *same* objective, which may only be replaced by a balanced + # solver when that solver minimizes the same objective: + # * ``reg == 0``: the regularizer vanishes, so the limit is exact + # balanced OT whatever ``c`` is; + # * ``reg_type="kl"`` with the default reference measure + # ``c = a b^T``: the limit is exactly the balanced KL problem; + # * a user supplied ``(f, df)`` pair: the balanced CG solver + # minimizes the same ``f``. + # The L2 and entropy regularizers use a different coefficient, a + # different reference measure and/or a different value convention in + # the balanced branch, so their hard-constraint limit is *not* the + # balanced problem. Dispatching them would silently solve another + # optimization problem (e.g. returning ``a b^T`` instead of ``c``, + # or a value off by ``reg``), so they are rejected instead. + balanced_matches_limit = isinstance(reg_type, tuple) or ( + isinstance(reg_type, str) and reg_type.lower() == "kl" and c_is_default + ) + if reg == 0 or balanced_matches_limit: + unbalanced = None + else: + reg_name = reg_type if isinstance(reg_type, str) else "custom" + ut_name = ( + unbalanced_type if isinstance(unbalanced_type, str) else "custom" + ) + raise NotImplementedError( + "The hard-constraint limit (unbalanced=float('inf')) of " + f"reg_type={reg_name!r} with unbalanced_type={ut_name!r} is not " + "the balanced problem solved by unbalanced=None: the balanced " + "branch uses a different regularizer coefficient, reference " + "measure or value convention. Returning it would silently solve " + "another optimization problem. Pass a finite `unbalanced` for " + "the unbalanced problem, or `unbalanced=None` to explicitly " + "solve the balanced problem that this reg_type implements." + ) + if method is not None and method not in lst_valid_methods_solve: raise ValueError( "Unknown method={} for solve, must be in {}".format( @@ -348,6 +457,15 @@ def solve( plan = log["G"] status = log["warning"] if log["warning"] is not None else "Converged" + elif any(unbalanced_is_inf): + raise NotImplementedError( + "Semi-relaxed exact unbalanced OT (reg=0 with exactly one " + 'infinite marginal relaxation, e.g. unbalanced=(float("inf"), ' + 'rho)) is not implemented. Use reg>0 with reg_type="KL" and ' + 'unbalanced_type="KL" for semi-relaxed entropic OT, or pass a ' + "finite `unbalanced` for exact unbalanced OT." + ) + elif unbalanced_type.lower() in ["kl", "l2"]: # unbalanced exact OT # default values for exact unbalanced OT if max_iter is None: @@ -574,6 +692,20 @@ def solve( isinstance(reg_type, tuple) or reg_type.lower() in ["kl", "l2", "entropy"] ) and unbalanced_type.lower() in ["kl", "l2", "tv"]: + if any(unbalanced_is_inf): + raise NotImplementedError( + "The L-BFGSB solver path does not support an infinite " + "marginal relaxation (semi-relaxed OT). The infinite " + "penalty cannot be represented, so the solver would " + "ignore it and silently solve the balanced problem " + "(with a NaN cost) instead of the requested " + "semi-relaxed one. The combination is rejected before " + "optimization. Use " + 'reg_type="KL" with unbalanced_type="KL" for ' + "semi-relaxed entropic OT, or pass unbalanced=None for " + "balanced OT." + ) + if isinstance(reg_type, tuple) and nx.__name__ == "tf": raise NotImplementedError( "Custom regularization is not implemented for TensorFlow backend." diff --git a/ot/unbalanced/_lbfgs.py b/ot/unbalanced/_lbfgs.py index 9591b976d..04040c9f5 100644 --- a/ot/unbalanced/_lbfgs.py +++ b/ot/unbalanced/_lbfgs.py @@ -207,6 +207,9 @@ def lbfgsb_unbalanced( regularization term >=0 reg_m: float or indexable object of length 1 or 2 Marginal relaxation term: nonnegative (including 0) but cannot be infinity. + Note that :any:`ot.solve` accepts ``float("inf")`` as a documented + sentinel for hard marginal constraints and resolves or rejects it before calling + this solver; infinite values are not supported here. If :math:`\mathrm{reg_{m}}` is a scalar or an indexable object of length 1, then the same :math:`\mathrm{reg_{m}}` is applied to both marginal relaxations. If :math:`\mathrm{reg_{m}}` is an array, it must be a Numpy array. @@ -396,6 +399,9 @@ def lbfgsb_unbalanced2( regularization term >=0 reg_m: float or indexable object of length 1 or 2 Marginal relaxation term: nonnegative (including 0) but cannot be infinity. + Note that :any:`ot.solve` accepts ``float("inf")`` as a documented + sentinel for hard marginal constraints and resolves or rejects it before calling + this solver; infinite values are not supported here. If :math:`\mathrm{reg_{m}}` is a scalar or an indexable object of length 1, then the same :math:`\mathrm{reg_{m}}` is applied to both marginal relaxations. If :math:`\mathrm{reg_{m}}` is an array, it must be a Numpy array. diff --git a/ot/unbalanced/_mm.py b/ot/unbalanced/_mm.py index 897375259..6f8201d43 100644 --- a/ot/unbalanced/_mm.py +++ b/ot/unbalanced/_mm.py @@ -65,6 +65,9 @@ def mm_unbalanced( loss matrix reg_m: float or indexable object of length 1 or 2 Marginal relaxation term: nonnegative but cannot be infinity. + Note that :any:`ot.solve` accepts ``float("inf")`` as a documented + sentinel for hard marginal constraints and resolves or rejects it before calling + this solver; infinite values are not supported here. If :math:`\mathrm{reg_{m}}` is a scalar or an indexable object of length 1, then the same :math:`\mathrm{reg_{m}}` is applied to both marginal relaxations. If :math:`\mathrm{reg_{m}}` is an array, @@ -263,6 +266,9 @@ def mm_unbalanced2( loss matrix reg_m: float or indexable object of length 1 or 2 Marginal relaxation term: nonnegative but cannot be infinity. + Note that :any:`ot.solve` accepts ``float("inf")`` as a documented + sentinel for hard marginal constraints and resolves or rejects it before calling + this solver; infinite values are not supported here. If :math:`\mathrm{reg_{m}}` is a scalar or an indexable object of length 1, then the same :math:`\mathrm{reg_{m}}` is applied to both marginal relaxations. If :math:`\mathrm{reg_{m}}` is an array, diff --git a/test/unbalanced/test_inf_marginal.py b/test/unbalanced/test_inf_marginal.py new file mode 100644 index 000000000..9cc0599a9 --- /dev/null +++ b/test/unbalanced/test_inf_marginal.py @@ -0,0 +1,348 @@ +"""Contract tests for infinite marginal relaxation in :any:`ot.solve`. + +``unbalanced=float("inf")`` is documented as the balanced limit, and +``unbalanced=(float("inf"), rho)`` / ``(rho, float("inf"))`` as the +semi-relaxed case. + +A two-sided infinite relaxation is the hard-constraint limit of the *same* +objective. That limit is only equal to the problem solved by ``unbalanced=None`` +when the balanced branch minimizes the same objective. These tests therefore +compare against an **independent analytic oracle** (a bounded scalar search on +the feasible set), not against another ``ot.solve`` call, and they require the +combinations whose limit is *not* representable to fail loudly instead of +returning a solution to a different optimization problem. + +License: MIT License +""" + +import itertools + +import numpy as np +import pytest +from scipy.optimize import minimize_scalar + +import ot + +INF = float("inf") + +# 2x2 symmetric fixtures: with a = b = (0.5, 0.5) the hard-constraint feasible +# set is one dimensional, P = [[x, .5-x], [.5-x, x]] with x in [0, .5], which +# makes an exact, solver-independent oracle possible. +A2 = np.array([0.5, 0.5]) +B2 = np.array([0.5, 0.5]) +M2 = np.array([[0.0, 1.0], [1.0, 0.0]]) +M2_ZERO = np.zeros((2, 2)) +C2_DEFAULT = np.outer(A2, B2) +C2_CUSTOM = np.array([[0.45, 0.05], [0.05, 0.45]]) + + +def _data(n=9, m=7, seed=3): + """Return a deterministic, non-uniform OT problem as Numpy arrays.""" + rng = np.random.RandomState(seed) + a = rng.rand(n) + a /= a.sum() + b = rng.rand(m) + b /= b.sum() + M = rng.rand(n, m) + return a, b, M + + +def _gkl(x, y): + """Generalized KL divergence, matching the ``mass=True`` convention.""" + return np.sum(x * np.log(np.maximum(x, 1e-300) / y) - x + y) + + +def _sym_plan(x): + return np.array([[x, 0.5 - x], [0.5 - x, x]]) + + +def _objective(x, reg, reg_type, c, M): + """The documented objective restricted to the hard-constraint feasible set.""" + P = _sym_plan(x) + value = float(np.sum(P * M)) + if reg > 0: + if reg_type == "kl": + value += reg * _gkl(P, c) + elif reg_type == "l2": + value += reg * 0.5 * np.sum((P - c) ** 2) + elif reg_type == "entropy": + value += reg * (np.sum(P * np.log(np.maximum(P, 1e-300))) - P.sum()) + return value + + +def _hard_limit_oracle(reg, reg_type, c, M=M2): + """Minimize the original objective over the hard-constraint feasible set.""" + res = minimize_scalar( + _objective, + bounds=(0.0, 0.5), + args=(reg, reg_type, c, M), + method="bounded", + options={"xatol": 1e-14}, + ) + return float(res.x), float(res.fun) + + +@pytest.mark.parametrize("unbalanced", [INF, (INF, INF), [INF, INF]]) +def test_balanced_limit_exact(nx, unbalanced): + """``unbalanced=inf`` with ``reg=0`` must reproduce exact balanced OT.""" + a, b, M = _data() + a_nx, b_nx, M_nx = nx.from_numpy(a, b, M) + + ref = ot.solve(M_nx, a_nx, b_nx, reg=0) + res = ot.solve(M_nx, a_nx, b_nx, reg=0, unbalanced=unbalanced) + + plan = nx.to_numpy(res.plan) + assert np.isfinite(plan).all() + assert np.isfinite(float(nx.to_numpy(res.value))) + np.testing.assert_allclose( + float(nx.to_numpy(res.value)), float(nx.to_numpy(ref.value)), rtol=1e-12 + ) + np.testing.assert_allclose(plan, nx.to_numpy(ref.plan), atol=1e-12) + # both marginals are hard constraints in the balanced limit + np.testing.assert_allclose(plan.sum(1), a, atol=1e-10) + np.testing.assert_allclose(plan.sum(0), b, atol=1e-10) + + +def test_balanced_limit_exact_independent_reference(nx): + """Check the balanced limit against a brute-force optimum, not against POT. + + With uniform weights on a square problem, the exact OT optimum is attained + at a permutation, so the value can be computed independently of any OT + solver. This keeps the regression test from re-deriving POT's own output. + """ + # shifted so that the optimum is not trivially 0 (a zero optimum would make + # the assertion vacuous against a solver that always returns 0) + M = np.array([[1.0, 2.0, 3.0], [2.0, 1.0, 4.0], [3.0, 4.0, 1.0]]) + a = np.full(3, 1.0 / 3.0) + b = np.full(3, 1.0 / 3.0) + best = min( + sum(M[i, perm[i]] for i in range(3)) / 3.0 + for perm in itertools.permutations(range(3)) + ) + + a_nx, b_nx, M_nx = nx.from_numpy(a, b, M) + res = ot.solve(M_nx, a_nx, b_nx, reg=0, unbalanced=INF) + + np.testing.assert_allclose(float(nx.to_numpy(res.value)), best, atol=1e-12) + np.testing.assert_allclose(np.asarray(nx.to_numpy(res.plan)).sum(), 1.0, atol=1e-12) + + +@pytest.mark.parametrize("unbalanced_type", ["KL", "L2", "TV"]) +def test_balanced_limit_kl_matches_hard_limit_oracle(nx, unbalanced_type): + """KL regularization with the default reference: the limit is the oracle. + + ``reg_type="kl"`` with ``c = a b^T`` makes the balanced problem identical to + the hard-constraint limit, so the dispatched result must match the analytic + oracle exactly - including the value, which a value-only guard would miss. + """ + x_star, v_star = _hard_limit_oracle(1.0, "kl", C2_DEFAULT) + a_nx, b_nx, M_nx = nx.from_numpy(A2, B2, M2) + + res = ot.solve( + M_nx, a_nx, b_nx, reg=1.0, reg_type="kl", + unbalanced_type=unbalanced_type, unbalanced=INF, + ) + plan = nx.to_numpy(res.plan) + + assert np.isfinite(plan).all() + np.testing.assert_allclose(plan[0, 0], x_star, atol=1e-6) + np.testing.assert_allclose(float(nx.to_numpy(res.value)), v_star, rtol=1e-8) + np.testing.assert_allclose(plan.sum(1), A2, atol=1e-6) + np.testing.assert_allclose(plan.sum(0), B2, atol=1e-6) + + +def test_kl_custom_reference_is_honoured(nx): + """A custom ``c`` must survive the hard-constraint limit. + + With ``M = 0`` and ``c`` already feasible, the unique minimizer is ``c`` + itself (KL >= 0 with equality iff P = c). The balanced KL branch ignores + ``c`` and would return ``a b^T`` instead, so this test would fail if the + limit were dispatched there. + """ + x_star, v_star = _hard_limit_oracle(1.0, "kl", C2_CUSTOM, M2_ZERO) + a_nx, b_nx, M_nx = nx.from_numpy(A2, B2, M2_ZERO) + c_nx = nx.from_numpy(C2_CUSTOM) + + res = ot.solve( + M_nx, a_nx, b_nx, reg=1.0, reg_type="kl", + unbalanced_type="KL", c=c_nx, unbalanced=INF, + ) + plan = nx.to_numpy(res.plan) + + np.testing.assert_allclose(plan, C2_CUSTOM, atol=1e-6) + np.testing.assert_allclose(plan[0, 0], x_star, atol=1e-6) + np.testing.assert_allclose(float(nx.to_numpy(res.value)), v_star, atol=1e-8) + + +@pytest.mark.parametrize( + "reg_type,unbalanced_type,custom_c", + [ + # L2, reg=1 on this fixture: the balanced branch *solves* + # reg/2 * ||P||^2 with no reference measure and *reports* + # + reg * ||P||^2, while the UOT hard limit is + # + reg/2 * ||P - c||^2. Oracle: x*=0.5, V*=0.125, but the + # balanced call returns V=0.5. The plans coincide only because a and b + # are uniform here; with a=(0.7,0.3), b=(0.6,0.4), M=0, reg=1 they + # differ (0.42 vs 0.40). At reg=2 the same pair is x*=0.5, 0.25 vs 1.0. + ("l2", "KL", False), + ("l2", "L2", False), + ("l2", "TV", False), + ("l2", "KL", True), + # entropy, reg=1: same argmin, but the balanced value omits -sum(P), so + # the reported value is off by exactly `reg` (oracle V*=-2.0064, the + # balanced call returns -1.0064). + ("entropy", "KL", False), + ("entropy", "L2", False), + ("entropy", "KL", True), + ("l2", "TV", True), + # KL with a custom c: the limit is c, the balanced branch would drop it. + ("kl", "L2", True), + ("kl", "TV", True), + ], +) +def test_hard_limit_rejects_different_objective(nx, reg_type, unbalanced_type, custom_c): + """Combinations whose limit is not the balanced problem must raise. + + Silently returning the balanced solution here would solve a different + optimization problem (wrong reference measure, coefficient or value + convention), which is exactly the bug these tests lock down. + """ + a_nx, b_nx, M_nx = nx.from_numpy(A2, B2, M2) + kwargs = dict( + reg=1.0, reg_type=reg_type, unbalanced_type=unbalanced_type, unbalanced=INF + ) + if custom_c: + kwargs["c"] = nx.from_numpy(C2_CUSTOM) + + with pytest.raises(NotImplementedError, match="hard-constraint limit"): + ot.solve(M_nx, a_nx, b_nx, **kwargs) + + +def test_hard_limit_custom_regularizer_matches_balanced(nx): + """A user supplied (f, df) pair is preserved by the balanced CG solver. + + This is a consistency check between the two paths (both call the same + ``cg``), not an independent oracle: it guards against the limit being routed + somewhere that drops ``f``. + """ + if nx.__name__ != "numpy": + pytest.skip( + "ot.solve with a custom (f, df) reg_type is numpy-only: the CG " + "solver passes backend arrays straight to the user callables." + ) + + def f(P): + return np.sum(P * np.log(np.maximum(P, 1e-300))) + + def df(P): + return np.log(np.maximum(P, 1e-300)) + 1 + + res = ot.solve( + nx.from_numpy(M2), nx.from_numpy(A2), nx.from_numpy(B2), + reg=1.0, reg_type=(f, df), unbalanced=INF, + ) + ref = ot.solve( + nx.from_numpy(M2), nx.from_numpy(A2), nx.from_numpy(B2), + reg=1.0, reg_type=(f, df), + ) + np.testing.assert_allclose( + float(nx.to_numpy(res.value)), float(nx.to_numpy(ref.value)), rtol=1e-10 + ) + np.testing.assert_allclose( + nx.to_numpy(res.plan), nx.to_numpy(ref.plan), atol=1e-10 + ) + + +@pytest.mark.parametrize("hard_side", ["left", "right"]) +def test_semi_relaxed_sinkhorn(nx, hard_side): + """The supported semi-relaxed path must honour the hard side only.""" + a, b, M = _data() + a_nx, b_nx, M_nx = nx.from_numpy(a, b, M) + reg = 0.5 + rho = 1.0 + unbalanced = (INF, rho) if hard_side == "left" else (rho, INF) + + res = ot.solve(M_nx, a_nx, b_nx, reg=reg, unbalanced=unbalanced) + plan = nx.to_numpy(res.plan) + value = float(nx.to_numpy(res.value)) + + assert np.isfinite(plan).all() + assert np.isfinite(value) + + # the hard constraint is enforced, the relaxed side is *not* forced. + # 1e-5 is the order of the Sinkhorn stopping criterion used by ot.solve. + if hard_side == "left": + np.testing.assert_allclose(plan.sum(1), a, atol=1e-5) + assert np.abs(plan.sum(0) - b).max() > 1e-4 + expected = np.sum(plan * M) + reg * _gkl(plan, np.outer(a, b)) + expected += rho * _gkl(plan.sum(0), b) + else: + np.testing.assert_allclose(plan.sum(0), b, atol=1e-5) + assert np.abs(plan.sum(1) - a).max() > 1e-4 + expected = np.sum(plan * M) + reg * _gkl(plan, np.outer(a, b)) + expected += rho * _gkl(plan.sum(1), a) + + # independent recomputation of the entropic semi-relaxed objective + np.testing.assert_allclose(value, expected, rtol=1e-8) + + +@pytest.mark.parametrize( + "kwargs", + [ + # exact semi-relaxed OT has no solver: reject instead of returning NaN + dict(reg=0, unbalanced=(INF, 1.0)), + dict(reg=0, unbalanced=(1.0, INF)), + dict(reg=0, unbalanced=(INF, 1.0), unbalanced_type="L2"), + dict(reg=0, unbalanced=(1.0, INF), unbalanced_type="L2"), + dict(reg=0, unbalanced=(INF, 1.0), unbalanced_type="TV"), + # the L-BFGSB path would silently solve the balanced problem + dict(reg=0.5, reg_type="l2", unbalanced=(INF, 1.0)), + dict(reg=0.5, reg_type="l2", unbalanced=(1.0, INF)), + dict(reg=0.5, unbalanced=(INF, 1.0), unbalanced_type="L2"), + dict(reg=0.5, unbalanced=(INF, 1.0), unbalanced_type="TV"), + dict(reg=0.5, reg_type="entropy", unbalanced=(INF, 1.0)), + ], +) +def test_unsupported_semi_relaxed_raises(nx, kwargs): + """Unsupported semi-relaxed combinations must fail loudly, not return NaN.""" + a, b, M = _data() + a_nx, b_nx, M_nx = nx.from_numpy(a, b, M) + + with pytest.raises(NotImplementedError, match="semi-relaxed|L-BFGSB|hard-constraint"): + ot.solve(M_nx, a_nx, b_nx, **kwargs) + + +def test_lbfgsb_infinite_relaxation_not_silently_balanced(nx): + """Regression: the L-BFGSB path used to return a balanced plan with NaN value.""" + a, b, M = _data() + a_nx, b_nx, M_nx = nx.from_numpy(a, b, M) + + # a one-sided hard constraint must be rejected ... + with pytest.raises(NotImplementedError): + ot.solve(M_nx, a_nx, b_nx, reg=0.5, reg_type="l2", unbalanced=(INF, 1.0)) + + # ... while the two-sided limit stays available for KL with the default c + res = ot.solve(M_nx, a_nx, b_nx, reg=0.5, reg_type="kl", unbalanced=INF) + plan = nx.to_numpy(res.plan) + assert np.isfinite(float(nx.to_numpy(res.value))) + # the balanced limit enforces both marginals + np.testing.assert_allclose(plan.sum(1), a, atol=1e-6) + np.testing.assert_allclose(plan.sum(0), b, atol=1e-6) + + +def test_finite_relaxation_unchanged(nx): + """Finite relaxations must keep using the unbalanced solvers unchanged.""" + a, b, M = _data() + a_nx, b_nx, M_nx = nx.from_numpy(a, b, M) + + res = ot.solve(M_nx, a_nx, b_nx, reg=0.5, unbalanced=1.0) + plan = nx.to_numpy(res.plan) + assert np.isfinite(plan).all() + # unbalanced OT does not satisfy the marginals exactly + assert np.abs(plan.sum(1) - a).max() > 1e-6 + np.testing.assert_allclose( + float(nx.to_numpy(res.value)), + np.sum(plan * M) + 0.5 * _gkl(plan, np.outer(a, b)) + _gkl(plan.sum(1), a) + + _gkl(plan.sum(0), b), + rtol=1e-8, + )