Label normalization performs copy only when necessary

Author: kachayev

ot/da.py

@@ -644,7 +644,7 @@ def transform_labels(self, ys=None):
             transp = nx.nan_to_num(transp, nan=0, posinf=0, neginf=0)
 
             # compute propagated labels
-            labels = label_normalization(nx.copy(ys))
+            labels = label_normalization(ys)
             masks = labels_to_masks(labels, nx=nx, type_as=transp)
             transp_ys = nx.dot(masks.T, transp)
 
@@ -739,7 +739,7 @@ def inverse_transform_labels(self, yt=None):
             transp = nx.nan_to_num(transp, nan=0, posinf=0, neginf=0)
 
             # compute propagated samples
-            labels = label_normalization(nx.copy(yt))
+            labels = label_normalization(yt)
             masks = labels_to_masks(labels, nx=nx, type_as=transp)
             transp_ys = nx.dot(masks.T, transp.T)
 
@@ -2126,7 +2126,7 @@ def transform_labels(self, ys=None):
                 type_as=ys[0]
             )
             for i in range(len(ys)):
-                ysTemp = label_normalization(nx.copy(ys[i]))
+                ysTemp = label_normalization(ys[i])
                 classes = nx.unique(ysTemp)
                 n = len(classes)
                 ns = len(ysTemp)
@@ -2169,7 +2169,7 @@ def inverse_transform_labels(self, yt=None):
         # check the necessary inputs parameters are here
         if check_params(yt=yt):
             transp_ys = []
-            ytTemp = label_normalization(nx.copy(yt))
+            ytTemp = label_normalization(yt)
             classes = nx.unique(ytTemp)
             n = len(classes)
             D1 = nx.zeros((n, len(ytTemp)), type_as=self.coupling_[0])

ot/utils.py

@@ -390,7 +390,7 @@ def is_all_finite(*args):
     return all(not nx.any(~nx.isfinite(arg)) for arg in args)
 
 
-def label_normalization(y, start=0):
+def label_normalization(y, start=0, nx=None):
     r""" Transform labels to start at a given value
 
     Parameters
@@ -399,31 +399,31 @@ def label_normalization(y, start=0):
         The vector of labels to be normalized.
     start : int
         Desired value for the smallest label in :math:`\mathbf{y}` (default=0)
+    nx : Backend, optional
+        Backend to perform computations on. If omitted, the backend defaults to that of `y`.
 
     Returns
     -------
     y : array-like, shape (`n1`, )
         The input vector of labels normalized according to given start value.
     """
-    nx = get_backend(y)
-
+    if nx is None:
+        nx = get_backend(y)
     diff = nx.min(nx.unique(y)) - start
-    if diff != 0:
-        y -= diff
-    return y
+    return y if diff == 0 else (y - diff)
 
 
-def labels_to_masks(y, nx=None, type_as=None):
-    r"""Transforms (n_samples,) vector of labels into a (n_samples, n_labels) matrix of masks. 
+def labels_to_masks(y, type_as=None, nx=None):
+    r"""Transforms (n_samples,) vector of labels into a (n_samples, n_labels) matrix of masks.
 
     Parameters
     ----------
     y : array-like, shape (n_samples, )
         The vector of labels.
-    nx : Backend, optional
-        Backend to perform computations on. If omitted, the backend defaults to that of `y`.
     type_as : array_like
         Array of the same type of the expected output.
+    nx : Backend, optional
+        Backend to perform computations on. If omitted, the backend defaults to that of `y`.
 
     Returns
     -------

test/test_utils.py

@@ -585,7 +585,7 @@ def test_lowrank_LazyTensor(nx):
     np.testing.assert_allclose(nx.to_numpy(T[:]), nx.to_numpy(T0))
 
 
-def test_label_to_mask_helper(nx):
+def test_labels_to_mask_helper(nx):
     y = np.array([1, 0, 2, 2, 1])
     out = np.array([
         [0, 1, 0],
@@ -597,3 +597,14 @@ def test_label_to_mask_helper(nx):
     y = nx.from_numpy(y)
     masks = ot.utils.labels_to_masks(y)
     np.testing.assert_array_equal(out, masks)
+
+
+def test_label_normalization(nx):
+    y = nx.from_numpy(np.arange(5) + 1)
+    out = np.arange(5)
+    # labels are shifted
+    y_normalized = ot.utils.label_normalization(y)
+    np.testing.assert_array_equal(out, y_normalized)
+    # labels are shifted but the shift if expected
+    y_normalized_start = ot.utils.label_normalization(y, start=1)
+    np.testing.assert_array_equal(y, y_normalized_start)