From f03e6b0c81c15a87e7cb6993c661ee0f80a51590 Mon Sep 17 00:00:00 2001 From: CodingSelim Date: Mon, 6 Jul 2026 15:41:51 +0300 Subject: [PATCH 1/4] fix label-aware cost correction in ot.da the cost correction that pushes apart differently-labeled samples was computed backwards. it used missing_ys = (ys == -1) and the product of the two missing masks, so the large cost was applied only where both labels are missing, and never where two labeled samples have different labels. with all labels known the correction was a no-op, so ys/yt had no effect on the transport (see #664). switched to present masks (ys != -1) so the correction applies exactly to labeled source/target pairs whose labels differ, matching the original pre-vectorized loop. restored the semisupervised tests that had been flipped to assert the buggy no-op (n_unsup == n_semisup) back to asserting the cost actually changes, and added a regression test. closes #664 --- RELEASES.md | 1 + ot/da.py | 23 ++++++++++++----------- test/test_da.py | 42 +++++++++++++++++++++++++++++++++++++----- 3 files changed, 50 insertions(+), 16 deletions(-) diff --git a/RELEASES.md b/RELEASES.md index eea94981d..ea00e4b60 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -43,6 +43,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 (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) diff --git a/ot/da.py b/ot/da.py index 527999cba..62b75e0bc 100644 --- a/ot/da.py +++ b/ot/da.py @@ -598,23 +598,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 -1) + present_ys = (ys != -1) + nx.zeros(ys.shape, type_as=ys) + present_yt = (yt != -1) + 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) diff --git a/test/test_da.py b/test/test_da.py index bb548e27f..434d01ad5 100644 --- a/test/test_da.py +++ b/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,38 @@ def test_emd_transport_class(nx): ) +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"]) From c05fd79bdd60532fc3f9802160704c0e4934a2eb Mon Sep 17 00:00:00 2001 From: CodingSelim Date: Mon, 6 Jul 2026 15:42:46 +0300 Subject: [PATCH 2/4] add PR number to releases entry --- RELEASES.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/RELEASES.md b/RELEASES.md index ea00e4b60..a9ed7bd72 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -43,7 +43,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 (Issue #664) +- 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) From d80515ae72043a6b514ddb7b87cc1ceb0af0f0f8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Mon, 6 Jul 2026 16:16:13 +0200 Subject: [PATCH 3/4] Apply suggestion from @rflamary --- test/test_da.py | 1 + 1 file changed, 1 insertion(+) diff --git a/test/test_da.py b/test/test_da.py index 434d01ad5..df224a247 100644 --- a/test/test_da.py +++ b/test/test_da.py @@ -586,6 +586,7 @@ 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 From 1b40683f176acf578a5a6f737155433bc4e911cf Mon Sep 17 00:00:00 2001 From: CodingSelim Date: Tue, 7 Jul 2026 04:40:39 +0300 Subject: [PATCH 4/4] use MISSING_LABEL constant instead of bare -1 --- ot/da.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/ot/da.py b/ot/da.py index 62b75e0bc..7b4ed7ba2 100644 --- a/ot/da.py +++ b/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, @@ -600,9 +603,9 @@ class label # 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 -1) - present_ys = (ys != -1) + nx.zeros(ys.shape, type_as=ys) - present_yt = (yt != -1) + nx.zeros(yt.shape, type_as=yt) + # (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]