diff --git a/RELEASES.md b/RELEASES.md index 9a5ddef50..8ea7012d3 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -8,6 +8,7 @@ #### Closed issues +- Fix the sign of the previous transport plan in `ot.batch.proximal_bregman_log_plan_batch`, so proximal kernels weight rather than invert the preceding plan (Issue #842) - Load triton before TensorFlow in `ot.backend` so that building a torch optimizer no longer segfaults the interpreter, and remove the `torch<2.12` pin from the doctest and documentation requirements (PR #839, Issue #816) - Fix mean centering in `ot.dr.fda` and `ot.dr.wda`: `np.mean(X)` returned a scalar instead of the per-feature mean, so `proj` did not center the data as documented. In `ot.dr.fda` the same pattern in the class means made the between-class scatter matrix independent of which features separate the classes, and FDA returned a non-discriminant direction (PR #840) - `ot.dr.fda` and `ot.dr.wda` no longer modify the input array `X` in place (PR #840) diff --git a/ot/batch/_utils.py b/ot/batch/_utils.py index 8bc96c3b9..d7b028f13 100644 --- a/ot/batch/_utils.py +++ b/ot/batch/_utils.py @@ -331,7 +331,7 @@ def proximal_bregman_log_plan_batch( .. math:: \mathbf{T}^{(k+1)} = \mathop{\arg \min}_\mathbf{T} \quad \langle \mathbf{C} - \textit{inner\_reg} \cdot \log \mathbf{T}^{(k)}, \mathbf{T} \rangle + (\textit{reg} + \textit{inner\_reg}) \cdot \sum_{i,j} \mathbf{T}_{i,j} \log \mathbf{T}_{i,j} - Denoting :math:`\mathbf{K}^{(k)} = - (\mathbf{C} + \textit{inner\_reg} \cdot \log \mathbf{T}^{(k)})/(\textit{reg} + \textit{inner\_reg})`, the affinity matrix at iteration :math:`k`, the Bregman projection problem is solved in the log-domain with a finite number of inner iterations :math:`\text{inner\_iter}`, i.e., the dual variables :math:`\mathbf{u}` and :math:`\mathbf{v}` are updated as follows: + Denoting :math:`\mathbf{K}^{(k)} = - (\mathbf{C} - \textit{inner\_reg} \cdot \log \mathbf{T}^{(k)})/(\textit{reg} + \textit{inner\_reg})`, the affinity matrix at iteration :math:`k`, the Bregman projection problem is solved in the log-domain with a finite number of inner iterations :math:`\text{inner\_iter}`, i.e., the dual variables :math:`\mathbf{u}` and :math:`\mathbf{v}` are updated as follows: .. math:: \mathbf{u}^{(i+1)} = \log(\mathbf{a}) - \text{LSE}(\mathbf{K}^{(k)} + \mathbf{v}^{(i)}) @@ -418,7 +418,7 @@ def proximal_bregman_log_plan_batch( log_T = nx.zeros(C.shape, type_as=C) for n_iters in range(max_iter): - K_proj = -(C + inner_reg * log_T) / (reg + inner_reg) + K_proj = -(C - inner_reg * log_T) / (reg + inner_reg) for _ in range(inner_iter): u = loga - nx.logsumexp(K_proj + v[:, None, :], axis=2) v = logb - nx.logsumexp(K_proj + u[:, :, None], axis=1) @@ -433,7 +433,7 @@ def proximal_bregman_log_plan_batch( break if grad == "last_step": - K_proj = -(C_ + inner_reg * log_T) / (reg + inner_reg) + K_proj = -(C_ - inner_reg * log_T) / (reg + inner_reg) for _ in range(inner_iter): u = loga - nx.logsumexp(K_proj + v[:, None, :], axis=2) v = logb - nx.logsumexp(K_proj + u[:, :, None], axis=1) diff --git a/test/batch/test_solve_batch.py b/test/batch/test_solve_batch.py index 08e5ff81e..07143de76 100644 --- a/test/batch/test_solve_batch.py +++ b/test/batch/test_solve_batch.py @@ -51,6 +51,7 @@ def test_solve_batch_vs_solve(reg, method, reg_type): base_plan[i] = res_i.plan base_value[i] = res_i.value_linear + inner_reg = 1e-1 if reg is None or reg == 0 else 1e-3 res = solve_batch( C, max_iter=10000, @@ -59,11 +60,15 @@ def test_solve_batch_vs_solve(reg, method, reg_type): reg=reg, method=method, reg_type=reg_type, - inner_reg=1e-3, + inner_reg=inner_reg, ) plan = res.plan value = res.value_linear - np.testing.assert_allclose(plan, base_plan, atol=tol * 10) + if reg is None or reg == 0: + np.testing.assert_allclose(plan.sum(axis=2), 1 / n, atol=tol * 10) + np.testing.assert_allclose(plan.sum(axis=1), 1 / d, atol=tol * 10) + else: + np.testing.assert_allclose(plan, base_plan, atol=tol * 10) np.testing.assert_allclose(value, base_value, atol=tol * 10) @@ -76,16 +81,19 @@ def test_backend_proximal_bregman_log_plan_batch(nx, reg, inner_iter): d = 7 rng = np.random.RandomState(0) C = rng.rand(batchsize, n, d) + inner_reg = 1e-1 if reg is None or reg == 0 else 1e-3 res = proximal_bregman_log_plan_batch( nx.from_numpy(C), reg=reg, - inner_reg=1e-3, + inner_reg=inner_reg, max_iter=10000, tol=tol, inner_iter=inner_iter, grad="detach", ) plan = nx.to_numpy(res["T"]) + np.testing.assert_allclose(plan.sum(axis=2), 1 / n, atol=tol * 10) + np.testing.assert_allclose(plan.sum(axis=1), 1 / d, atol=tol * 10) for i in range(batchsize): C_i = C[i] res_i = solve( @@ -93,8 +101,39 @@ def test_backend_proximal_bregman_log_plan_batch(nx, reg, inner_iter): reg=reg, tol=tol, ) - plan_i = res_i.plan - np.testing.assert_allclose(plan_i, plan[i], atol=tol * 10) + if reg is None or reg == 0: + value = np.sum(C_i * plan[i]) + np.testing.assert_allclose(res_i.value_linear, value, atol=tol * 10) + else: + np.testing.assert_allclose(res_i.plan, plan[i], atol=tol * 10) + + +def test_proximal_bregman_uses_previous_plan(): + """Check that each proximal kernel is weighted by the previous plan.""" + cost = np.array([[[0.0, 1.0], [2.0, 0.0]]]) + result = proximal_bregman_log_plan_batch( + cost, inner_reg=1.0, max_iter=2, inner_iter=10, tol=-1.0 + ) + log_plan = result["log_T"][0] + cross_ratio = log_plan[0, 0] + log_plan[1, 1] - log_plan[0, 1] - log_plan[1, 0] + np.testing.assert_allclose(cross_ratio, 6.0) + + +@pytest.mark.skipif(not torch, reason="torch not installed") +def test_proximal_bregman_last_step_uses_previous_plan(): + """Check that the differentiable final step keeps the proximal sign.""" + cost = torch.tensor([[[0.0, 1.0], [2.0, 0.0]]], requires_grad=True) + result = proximal_bregman_log_plan_batch( + cost, + inner_reg=1.0, + max_iter=2, + inner_iter=10, + tol=-1.0, + grad="last_step", + ) + log_plan = result["log_T"][0] + cross_ratio = log_plan[0, 0] + log_plan[1, 1] - log_plan[0, 1] - log_plan[1, 0] + torch.testing.assert_close(cross_ratio, torch.tensor(9.0)) def test_bregman_batch():