[MRG] Fix label-aware cost correction in ot.da (closes #664)

RELEASES.md

@@ -44,6 +44,7 @@ This new release adds support for sparse cost matrices and a new lazy exact OT s
 
 #### Closed issues
 
+- Fix label-aware cost correction in `ot.da` transport classes: the large cost was applied to unlabeled pairs instead of labeled pairs with different labels, so `ys`/`yt` had no effect when all labels were known (PR #833, Issue #664)
 - Mitigate NaN regime of `entropic_partial_wasserstein` at small `reg` via a new log-domain solver, reachable with `entropic_partial_wasserstein(..., method='sinkhorn_log')` (Issue #723; the default `method='sinkhorn'` path is unchanged — callers opt into the log-domain variant)
 - Fix NumPy 2.x compatibility in Brenier potential bounds (PR #788)
 - Fix MSVC Windows build by removing **restrict** keyword (PR #788)

ot/da.py

@@ -48,6 +48,9 @@
     joint_OT_mapping_kernel,
 )
 
+# value used in ys/yt to mark a sample whose label is unknown
+MISSING_LABEL = -1
+
 
 def sinkhorn_lpl1_mm(
     a,
@@ -598,23 +601,24 @@ class label
                 if self.limit_max != np.inf:
                     self.limit_max = self.limit_max * nx.max(self.cost_)
 
-                # missing_labels is a (ns, nt) matrix of {0, 1} such that
-                # the cells (i, j) has 0 iff either ys[i] or yt[j] is masked
-                missing_ys = (ys == -1) + nx.zeros(ys.shape, type_as=ys)
-                missing_yt = (yt == -1) + nx.zeros(yt.shape, type_as=yt)
-                missing_labels = missing_ys[:, None] @ missing_yt[None, :]
-                # labels_match is a (ns, nt) matrix of {True, False} such that
-                # the cells (i, j) has False if ys[i] != yt[i]
-                label_match = (ys[:, None] - yt[None, :]) != 0
+                # present_labels is a (ns, nt) matrix of {0, 1} such that
+                # the cell (i, j) is 1 iff both ys[i] and yt[j] are labeled
+                # (i.e. neither is masked with MISSING_LABEL)
+                present_ys = (ys != MISSING_LABEL) + nx.zeros(ys.shape, type_as=ys)
+                present_yt = (yt != MISSING_LABEL) + nx.zeros(yt.shape, type_as=yt)
+                present_labels = present_ys[:, None] @ present_yt[None, :]
+                # label_mismatch is a (ns, nt) matrix of {True, False} such that
+                # the cell (i, j) is True if ys[i] != yt[j]
+                label_mismatch = (ys[:, None] - yt[None, :]) != 0
                 # cost correction is a (ns, nt) matrix of {-Inf, float, Inf} such
                 # that he cells (i, j) has -Inf where there's no correction necessary
-                # by 'correction' we mean setting cost to a large value when
-                # labels do not match
+                # by 'correction' we mean setting cost to a large value when two
+                # labeled samples have different labels
                 # we suppress potential RuntimeWarning caused by Inf multiplication
                 # (as we explicitly cover potential NANs later)
                 with warnings.catch_warnings():
                     warnings.simplefilter("ignore", category=RuntimeWarning)
-                    cost_correction = label_match * missing_labels * self.limit_max
+                    cost_correction = label_mismatch * present_labels * self.limit_max
                 # this operation is necessary because 0 * Inf = NAN
                 # thus is irrelevant when limit_max is finite
                 cost_correction = nx.nan_to_num(cost_correction, -np.inf)

test/test_da.py

@@ -166,7 +166,7 @@ def test_sinkhorn_lpl1_transport_class(nx):
     n_semisup = nx.sum(otda_semi.cost_)
 
     # check that the cost matrix norms are indeed different
-    assert np.allclose(
+    assert not np.allclose(
         n_unsup, n_semisup, atol=1e-7
     ), "semisupervised mode is not working"
 
@@ -258,7 +258,7 @@ def test_sinkhorn_l1l2_transport_class(nx):
     n_semisup = nx.sum(otda_semi.cost_)
 
     # check that the cost matrix norms are indeed different
-    assert np.allclose(
+    assert not np.allclose(
         n_unsup, n_semisup, atol=1e-7
     ), "semisupervised mode is not working"
 
@@ -356,7 +356,7 @@ def test_sinkhorn_transport_class(nx):
     n_semisup = nx.sum(otda_semi.cost_)
 
     # check that the cost matrix norms are indeed different
-    assert np.allclose(
+    assert not np.allclose(
         n_unsup, n_semisup, atol=1e-7
     ), "semisupervised mode is not working"
 
@@ -472,7 +472,7 @@ def test_unbalanced_sinkhorn_transport_class(nx):
         n_semisup = nx.sum(otda_semi.cost_)
 
         # check that the cost matrix norms are indeed different
-        assert np.allclose(
+        assert not np.allclose(
             n_unsup, n_semisup, atol=1e-7
         ), "semisupervised mode is not working"
 
@@ -571,7 +571,7 @@ def test_emd_transport_class(nx):
     n_semisup = nx.sum(otda_semi.cost_)
 
     # check that the cost matrix norms are indeed different
-    assert np.allclose(
+    assert not np.allclose(
         n_unsup, n_semisup, atol=1e-7
     ), "semisupervised mode is not working"
 
@@ -586,6 +586,39 @@ def test_emd_transport_class(nx):
     )
 
 
+@pytest.skip_backend("tf")
+def test_semisupervised_cost_correction(nx):
+    # gh-664: the label-aware cost correction must push apart labeled source
+    # and target samples with *different* labels, and must leave pairs
+    # involving an unlabeled (-1) sample untouched.
+    rng = np.random.RandomState(0)
+    Xs = rng.randn(6, 2)
+    Xt = rng.randn(6, 2)
+    ys = np.array([0, 0, 0, 1, 1, 1])
+    yt = np.array([0, 1, 0, 1, 0, 1])
+    Xs, ys, Xt, yt = nx.from_numpy(Xs, ys, Xt, yt)
+
+    # all labels known: different-label pairs get the (scaled) limit_max,
+    # same-label pairs keep their original, smaller cost
+    otda = ot.da.EMDTransport(limit_max=10)
+    otda.fit(Xs=Xs, ys=ys, Xt=Xt, yt=yt)
+    cost = nx.to_numpy(otda.cost_)
+    limit = float(nx.to_numpy(otda.limit_max))
+    ys_np, yt_np = nx.to_numpy(ys), nx.to_numpy(yt)
+    mismatch = ys_np[:, None] != yt_np[None, :]
+    assert np.all(cost[mismatch] >= limit - 1e-9), "different labels not penalized"
+    assert np.all(cost[~mismatch] < limit), "same labels wrongly penalized"
+
+    # semi-supervised: unlabeled (-1) targets must never be forbidden
+    yt_semi = nx.from_numpy(np.array([0, -1, 0, -1, 0, -1]))
+    otda2 = ot.da.EMDTransport(limit_max=10)
+    otda2.fit(Xs=Xs, ys=ys, Xt=Xt, yt=yt_semi)
+    cost2 = nx.to_numpy(otda2.cost_)
+    limit2 = float(nx.to_numpy(otda2.limit_max))
+    unlabeled = nx.to_numpy(yt_semi) == -1
+    assert np.all(cost2[:, unlabeled] < limit2), "unlabeled targets wrongly forbidden"
+
+
 @pytest.skip_backend("jax")
 @pytest.skip_backend("tf")
 @pytest.mark.parametrize("kernel", ["linear", "gaussian"])