[MRG] Preserve label penalties across transport estimator refits

RELEASES.md

@@ -19,6 +19,8 @@
 #### Closed issues
 
 - 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)
+- Keep the configured `limit_max` multiplier unchanged when refitting domain-adaptation transport estimators, so repeated fits do not compound the label penalty or change the transport plan.
+
 - 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)
 - Remove a leftover debug `print` from `ot.utils.projection_sparse_simplex` with `axis=1`, and make the `ot.datasets.make_gauss_hd` docstring a raw string so importing `ot` no longer emits a `SyntaxWarning` (PR #860)
 - Fix the value reported by the entropic unbalanced OT solvers (`ot.unbalanced.sinkhorn_unbalanced`, `sinkhorn_unbalanced2`, `sinkhorn_knopp_unbalanced`, `sinkhorn_stabilized_unbalanced` and `sinkhorn_unbalanced_translation_invariant`, also exposed as `ot.solve(..., reg=..., unbalanced=...).value`): every penalization is now evaluated with the generalized KL divergence, and with `reg_type="entropy"` the constant `dim_a * dim_b` is removed so that the reported value matches the documented negative entropy regularizer. The transport plans are unchanged. `reg_type` is now handled case insensitively by all the methods and an unknown value raises a `ValueError` instead of silently falling back to `"kl"`. The docstrings now state the generalized KL convention, the regularizer `Omega_reg_type(gamma, c)` with its two cases, the keys of the `log` dictionary and the multi-histogram behaviour (PR #874)

ot/da.py

@@ -598,8 +598,9 @@ class label
             )
 
             if (ys is not None) and (yt is not None):
-                if self.limit_max != np.inf:
-                    self.limit_max = self.limit_max * nx.max(self.cost_)
+                limit_max = self.limit_max
+                if limit_max != np.inf:
+                    limit_max = limit_max * nx.max(self.cost_)
 
                 # 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
@@ -618,7 +619,7 @@ class label
                 # (as we explicitly cover potential NANs later)
                 with warnings.catch_warnings():
                     warnings.simplefilter("ignore", category=RuntimeWarning)
-                    cost_correction = label_mismatch * present_labels * self.limit_max
+                    cost_correction = label_mismatch * present_labels * 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)
@@ -1454,6 +1455,8 @@ class EMDTransport(BaseTransport):
         Controls the semi supervised mode. Transport between labeled source
         and target samples of different classes will exhibit an infinite cost
         (10 times the maximum value of the cost matrix)
+        The multiplier is applied separately on each fit without changing
+        the configured value of ``limit_max``.
     max_iter : int, optional (default=100000)
         The maximum number of iterations before stopping the optimization
         algorithm if it has not converged.

test/test_da.py

@@ -264,8 +264,8 @@ def test_sinkhorn_l1l2_transport_class(nx):
 
     # check that the coupling forbids mass transport between labeled source
     # and labeled target samples
-    mass_semi = nx.sum(otda_semi.coupling_[otda_semi.cost_ == otda_semi.limit_max])
-    mass_semi = otda_semi.coupling_[otda_semi.cost_ == otda_semi.limit_max]
+    label_mismatch = (ys[:, None] != yt_semi[None, :]) & (yt_semi[None, :] != -1)
+    mass_semi = otda_semi.coupling_[label_mismatch]
     assert_allclose(
         nx.to_numpy(mass_semi), np.zeros_like(mass_semi), rtol=1e-9, atol=1e-9
     )
@@ -605,15 +605,35 @@ def test_emd_transport_class(nx):
 
     # check that the coupling forbids mass transport between labeled source
     # and labeled target samples
-    mass_semi = nx.sum(otda_semi.coupling_[otda_semi.cost_ == otda_semi.limit_max])
-    mass_semi = otda_semi.coupling_[otda_semi.cost_ == otda_semi.limit_max]
+    mass_semi = otda_semi.coupling_[ys[:, None] != yt[None, :]]
 
     # we need to use a small tolerance here, otherwise the test breaks
     assert_allclose(
         nx.to_numpy(mass_semi), np.zeros(list(mass_semi.shape)), rtol=1e-2, atol=1e-2
     )
 
 
+@pytest.skip_backend("tf")
+@pytest.mark.parametrize("transport", [ot.da.EMDTransport, ot.da.SinkhornTransport])
+@pytest.mark.parametrize("limit_max", [10.0, np.inf])
+def test_transport_refit_label_penalty(nx, transport, limit_max):
+    ys, yt = nx.from_numpy(np.array([0, 1]), np.array([1, 0]))
+    kwargs = {"limit_max": limit_max}
+    if transport is ot.da.SinkhornTransport:
+        kwargs["reg_e"] = 0.01
+    reused = transport(**kwargs)
+
+    # Repeated data and changing scales must behave like a fresh estimator.
+    for scale in [0.1, 0.1, 2.0, 0.1]:
+        Xs, Xt = nx.from_numpy(np.array([[0.0], [scale]]), np.array([[0.0], [scale]]))
+        fresh = transport(**kwargs).fit(Xs=Xs, ys=ys, Xt=Xt, yt=yt)
+        reused.fit(Xs=Xs, ys=ys, Xt=Xt, yt=yt)
+        assert_allclose(nx.to_numpy(reused.cost_), nx.to_numpy(fresh.cost_))
+        assert_allclose(nx.to_numpy(reused.coupling_), nx.to_numpy(fresh.coupling_))
+
+    assert reused.get_params()["limit_max"] == limit_max
+
+
 @pytest.skip_backend("tf")
 def test_semisupervised_cost_correction(nx):
     # gh-664: the label-aware cost correction must push apart labeled source
@@ -631,7 +651,7 @@ def test_semisupervised_cost_correction(nx):
     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))
+    limit = 10 * nx.to_numpy(ot.dist(Xs, Xt)).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"
@@ -642,7 +662,7 @@ def test_semisupervised_cost_correction(nx):
     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))
+    limit2 = 10 * nx.to_numpy(ot.dist(Xs, Xt)).max()
     unlabeled = nx.to_numpy(yt_semi) == -1
     assert np.all(cost2[:, unlabeled] < limit2), "unlabeled targets wrongly forbidden"