diff --git a/RELEASES.md b/RELEASES.md index 526f6409e..6ead2ce34 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -8,6 +8,7 @@ #### Closed issues +- Preserve input dtype and device for expected sliced plans, avoid materializing dense distance matrices for sparse plans, and fix weighted sparse-distance ordering (PR #846, Issue #845) - Fix the sign issue in updates of the previous transport plan in `ot.batch.proximal_bregman_log_plan_batch` (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) @@ -909,4 +910,4 @@ It provides the following solvers: * Optimal transport for domain adaptation with group lasso regularization * Conditional gradient and Generalized conditional gradient for regularized OT. -Some demonstrations (both in Python and Jupyter Notebook format) are available in the examples folder. \ No newline at end of file +Some demonstrations (both in Python and Jupyter Notebook format) are available in the examples folder. diff --git a/ot/sliced/_sliced_plans.py b/ot/sliced/_sliced_plans.py index d192a24e8..a0b7bca94 100644 --- a/ot/sliced/_sliced_plans.py +++ b/ot/sliced/_sliced_plans.py @@ -12,7 +12,7 @@ import warnings from ..backend import get_backend -from ..utils import list_to_array, sparse_ot_dist, dist +from ..utils import list_to_array, sparse_ot_dist from ._utils import get_random_projections from ..lp import wasserstein_1d from collections import namedtuple @@ -103,9 +103,9 @@ def sliced_plans( m = X_t.shape[0] if a is None: - a = nx.ones(n) / n + a = nx.ones(n, type_as=X_s) / n if b is None: - b = nx.ones(m) / m + b = nx.ones(m, type_as=X_t) / m is_perm = (n == m) and (a == a.sum() / n).all() and (b == a.sum() / n).all() @@ -493,7 +493,7 @@ def expected_sliced_plan( else: # uniform weights if n_projections is None: n_projections = projections.shape[1] - weights = nx.ones(n_projections) / n_projections + weights = nx.ones(n_projections, type_as=X_s) / n_projections weights_e = nx.concatenate([plans[i].data * weights[i] for i in range(len(plans))]) Xs_idx = nx.concatenate([plans[i].rows for i in range(len(plans))]) @@ -505,10 +505,16 @@ def expected_sliced_plan( plan = nx.todense(plan) if beta == 0.0: - if dense: - cost = nx.sum(plan * dist(X_s, X_t, metric=metric, p=p)) - else: - cost = plan.multiply(dist(X_s, X_t, metric=metric, p=p)).sum() + cost = sparse_ot_dist( + X_s, + X_t, + Xs_idx, + Xt_idx, + weights_e, + metric=metric, + p=p, + batch_size=batch_size, + ) if log: log_dict = { "projections": log_dict_plans["projections"], diff --git a/ot/utils.py b/ot/utils.py index fcd823415..5fee6109c 100644 --- a/ot/utils.py +++ b/ot/utils.py @@ -391,7 +391,7 @@ def dist_idxs(idx_x1, idx_x2): for b in range(0, len(i), batch_size): d += nx.sum( dist_idxs(i[b : b + batch_size], j[b : b + batch_size]) - * w[i[b : b + batch_size]] + * w[b : b + batch_size] ) return d diff --git a/test/sliced/test_sliced_plans.py b/test/sliced/test_sliced_plans.py index 1df3522cb..eb12e89ff 100644 --- a/test/sliced/test_sliced_plans.py +++ b/test/sliced/test_sliced_plans.py @@ -223,6 +223,40 @@ def test_expected_sliced(): ot.sliced.expected_sliced_plan(x, y, a, b, projections=projections, batch_size=2) +@pytest.mark.skipif(not torch, reason="PyTorch not installed") +@pytest.mark.parametrize("dense", [True, False]) +@pytest.mark.parametrize("use_weights", [False, True]) +def test_expected_sliced_torch_preserves_device(dense, use_weights): + device = "cuda" if torch.cuda.is_available() else "cpu" + dtype = torch.float64 + x = torch.tensor( + [[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0]], + dtype=dtype, + device=device, + ) + y = torch.tensor( + [[0.5, 0.0], [1.5, 0.0], [0.5, 1.0], [1.5, 1.0]], + dtype=dtype, + device=device, + ) + projections = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=dtype, device=device) + if use_weights: + a = torch.tensor([0.1, 0.2, 0.3, 0.4], dtype=dtype, device=device) + b = torch.tensor([0.4, 0.3, 0.2, 0.1], dtype=dtype, device=device) + else: + a = b = None + + plan, cost = ot.sliced.expected_sliced_plan( + x, y, a, b, projections=projections, dense=dense + ) + + assert plan.device == x.device + assert cost.device == x.device + dense_plan = plan if dense else plan.to_dense() + expected_cost = torch.sum(dense_plan * ot.dist(x, y)) + torch.testing.assert_close(cost, expected_cost) + + def test_sliced_plans_backends(nx): n = 10 m = 24 diff --git a/test/test_utils.py b/test/test_utils.py index b2d6b4508..dce79af6a 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -281,6 +281,20 @@ def test_dist(): ot.dist(x, x, metric="fakeone") +def test_sparse_ot_dist_uses_pair_weights(): + x1 = np.array([[0.0], [1.0]]) + x2 = np.array([[0.0], [5.0]]) + rows = np.array([1, 0]) + cols = np.array([0, 1]) + weights = np.array([0.25, 0.75]) + + cost = ot.utils.sparse_ot_dist( + x1, x2, rows, cols, weights, metric="sqeuclidean", batch_size=1 + ) + + np.testing.assert_allclose(cost, 0.25 * 1.0 + 0.75 * 25.0) + + @pytest.mark.parametrize("metric", lst_metrics) def test_dist_backends(nx, metric): n = 100