[MRG] Fix proximal batch plan weighting

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)

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)

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,25 +81,59 @@ 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(
             C_i,
             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():