diff --git a/RELEASES.md b/RELEASES.md index 82bd61421..5874aff79 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -4,6 +4,7 @@ #### New features +- Use `ot.utils.check_marginal` (and shape-tuple support in `ot.utils.unif`) to fill and validate default marginals consistently across solvers (Gromov, low-rank, stochastic, barycenter, factored) (PR #856) - Add stereographic spherical sliced Wasserstein distance in `ot.sliced.stereographic_sliced_wasserstein_sphere`, with its rotationally invariant extension (PR #836) - Add Quasi-Monte Carlo sliced Wasserstein sampling (QSW/RQSW) via generalized spiral points, selectable with `sampling_slices` in `sliced_wasserstein_distance`, diff --git a/ot/bregman/_barycenter.py b/ot/bregman/_barycenter.py index 89732f9df..9f388989c 100644 --- a/ot/bregman/_barycenter.py +++ b/ot/bregman/_barycenter.py @@ -13,7 +13,7 @@ import warnings import numpy as np -from ..utils import dist, list_to_array, unif +from ..utils import dist, list_to_array, unif, check_marginal from ..backend import get_backend from ._utils import geometricBar, geometricMean, projR, projC @@ -352,8 +352,7 @@ def free_support_sinkhorn_barycenter( N = len(measures_locations) k = X_init.shape[0] d = X_init.shape[1] - if b is None: - b = nx.ones((k,), type_as=X_init) / k + b = check_marginal(b, k, type_as=X_init) if weights is None: weights = nx.ones((N,), type_as=X_init) / N diff --git a/ot/bregman/_geomloss.py b/ot/bregman/_geomloss.py index 30129ada4..d94bec0f9 100644 --- a/ot/bregman/_geomloss.py +++ b/ot/bregman/_geomloss.py @@ -14,7 +14,7 @@ from geomloss import SamplesLoss import torch from torch.autograd import grad - from ..utils import get_backend, LazyTensor, dist + from ..utils import get_backend, LazyTensor, dist, check_marginal if geomloss.__version__ < "0.3.1": old_geomloss = True @@ -199,10 +199,8 @@ def empirical_sinkhorn2_geomloss( if nx.__name__ not in ["torch", "numpy"]: raise ValueError("geomloss only support torch or numpy backend") - if a is None: - a = nx.ones(X_s.shape[0], type_as=X_s) / X_s.shape[0] - if b is None: - b = nx.ones(X_t.shape[0], type_as=X_t) / X_t.shape[0] + a = check_marginal(a, X_s.shape[0], type_as=X_s) + b = check_marginal(b, X_t.shape[0], type_as=X_t) if nx.__name__ == "numpy": X_s_torch = torch.tensor(X_s) diff --git a/ot/factored.py b/ot/factored.py index f1b9f28c4..9119c87de 100644 --- a/ot/factored.py +++ b/ot/factored.py @@ -7,7 +7,7 @@ # License: MIT License from .backend import get_backend -from .utils import dist, get_lowrank_lazytensor +from .utils import dist, get_lowrank_lazytensor, check_marginal from .lp import emd from .bregman import sinkhorn @@ -104,10 +104,8 @@ def factored_optimal_transport( n_b = Xb.shape[0] d = Xa.shape[1] - if a is None: - a = nx.ones((n_a), type_as=Xa) / n_a - if b is None: - b = nx.ones((n_b), type_as=Xb) / n_b + a = check_marginal(a, n_a, type_as=Xa) + b = check_marginal(b, n_b, type_as=Xb) if X0 is None: X = nx.randn(r, d, type_as=Xa) diff --git a/ot/gromov/_gw.py b/ot/gromov/_gw.py index 99d5bcf10..87a206b68 100644 --- a/ot/gromov/_gw.py +++ b/ot/gromov/_gw.py @@ -17,7 +17,7 @@ from ..utils import dist, UndefinedParameter, list_to_array from ..optim import cg, line_search_armijo, solve_1d_linesearch_quad -from ..utils import check_random_state, unif +from ..utils import check_marginal, check_random_state, unif from ..backend import get_backend, NumpyBackend from ._utils import init_matrix, gwloss, gwggrad @@ -373,10 +373,8 @@ def gromov_wasserstein2( nx = get_backend(C1, C2) # init marginals if set as None - if p is None: - p = unif(C1.shape[0], type_as=C1) - if q is None: - q = unif(C2.shape[0], type_as=C1) + p = check_marginal(p, C1.shape[0], type_as=C1) + q = check_marginal(q, C2.shape[0], type_as=C1) T, log_gw = gromov_wasserstein( C1, @@ -780,10 +778,8 @@ def fused_gromov_wasserstein2( nx = get_backend(C1, C2, M) # init marginals if set as None - if p is None: - p = unif(C1.shape[0], type_as=M) - if q is None: - q = unif(C2.shape[0], type_as=M) + p = check_marginal(p, C1.shape[0], type_as=M) + q = check_marginal(q, C2.shape[0], type_as=M) T, log_fgw = fused_gromov_wasserstein( M, diff --git a/ot/gromov/_lowrank.py b/ot/gromov/_lowrank.py index ca7e1f72f..497ac928c 100644 --- a/ot/gromov/_lowrank.py +++ b/ot/gromov/_lowrank.py @@ -7,7 +7,7 @@ # License: MIT License import warnings -from ..utils import unif, get_lowrank_lazytensor +from ..utils import get_lowrank_lazytensor, check_marginal from ..backend import get_backend from ..lowrank import compute_lr_sqeuclidean_matrix, _init_lr_sinkhorn, _LR_Dysktra @@ -180,10 +180,8 @@ def lowrank_gromov_wasserstein_samples( ns, nt = X_s.shape[0], X_t.shape[0] # Initialize weights a, b - if a is None: - a = unif(ns, type_as=X_s) - if b is None: - b = unif(nt, type_as=X_t) + a = check_marginal(a, ns, type_as=X_s) + b = check_marginal(b, nt, type_as=X_t) # Compute rank (see Section 3.1, def 1) r = rank diff --git a/ot/gromov/_partial.py b/ot/gromov/_partial.py index afb76eff6..4064193dd 100644 --- a/ot/gromov/_partial.py +++ b/ot/gromov/_partial.py @@ -9,7 +9,7 @@ # # License: MIT License -from ..utils import list_to_array, unif +from ..utils import check_marginal, list_to_array, unif from ..backend import get_backend, NumpyBackend from ..partial import entropic_partial_wasserstein from ._utils import _transform_matrix, gwloss, gwggrad @@ -451,10 +451,8 @@ def partial_gromov_wasserstein2( nx = get_backend(C1, C2) # init marginals if set as None - if p is None: - p = unif(C1.shape[0], type_as=C1) - if q is None: - q = unif(C2.shape[0], type_as=C1) + p = check_marginal(p, C1.shape[0], type_as=C1) + q = check_marginal(q, C2.shape[0], type_as=C1) T, log_pgw = partial_gromov_wasserstein( C1, @@ -925,10 +923,8 @@ def partial_fused_gromov_wasserstein2( nx = get_backend(M, C1, C2) # init marginals if set as None - if p is None: - p = unif(C1.shape[0], type_as=C1) - if q is None: - q = unif(C2.shape[0], type_as=C1) + p = check_marginal(p, C1.shape[0], type_as=C1) + q = check_marginal(q, C2.shape[0], type_as=C1) T, log_pfgw = partial_fused_gromov_wasserstein( M, @@ -1207,10 +1203,8 @@ def entropic_partial_gromov_wasserstein( nx = get_backend(*arr) - if p is None: - p = nx.ones(C1.shape[0], type_as=C1) / C1.shape[0] - if q is None: - q = nx.ones(C2.shape[0], type_as=C2) / C2.shape[0] + p = check_marginal(p, C1.shape[0], type_as=C1) + q = check_marginal(q, C2.shape[0], type_as=C2) if m is None: m = min(nx.sum(p), nx.sum(q)) @@ -1565,10 +1559,8 @@ def entropic_partial_fused_gromov_wasserstein( nx = get_backend(*arr) - if p is None: - p = nx.ones(C1.shape[0], type_as=C1) / C1.shape[0] - if q is None: - q = nx.ones(C2.shape[0], type_as=C2) / C2.shape[0] + p = check_marginal(p, C1.shape[0], type_as=C1) + q = check_marginal(q, C2.shape[0], type_as=C2) if m is None: m = min(nx.sum(p), nx.sum(q)) diff --git a/ot/gromov/_semirelaxed.py b/ot/gromov/_semirelaxed.py index 02fafc874..d28939b51 100644 --- a/ot/gromov/_semirelaxed.py +++ b/ot/gromov/_semirelaxed.py @@ -11,7 +11,7 @@ import numpy as np -from ..utils import list_to_array, unif, dist, UndefinedParameter +from ..utils import list_to_array, unif, dist, UndefinedParameter, check_marginal from ..optim import semirelaxed_cg, solve_1d_linesearch_quad from ..backend import get_backend @@ -336,8 +336,7 @@ def semirelaxed_gromov_wasserstein2( nx = get_backend(C1, C2) # init marginals if set as None - if p is None: - p = unif(C1.shape[0], type_as=C1) + p = check_marginal(p, C1.shape[0], type_as=C1) T, log_srgw = semirelaxed_gromov_wasserstein( C1, @@ -698,8 +697,7 @@ def semirelaxed_fused_gromov_wasserstein2( nx = get_backend(C1, C2) # init marginals if set as None - if p is None: - p = unif(C1.shape[0], type_as=C1) + p = check_marginal(p, C1.shape[0], type_as=C1) T, log_fgw = semirelaxed_fused_gromov_wasserstein( M, diff --git a/ot/lowrank.py b/ot/lowrank.py index 923e7f29b..fef964949 100644 --- a/ot/lowrank.py +++ b/ot/lowrank.py @@ -8,7 +8,7 @@ # License: MIT License import warnings -from .utils import unif, dist, get_lowrank_lazytensor +from .utils import dist, get_lowrank_lazytensor, check_marginal from .backend import get_backend from .bregman import sinkhorn @@ -420,10 +420,8 @@ def lowrank_sinkhorn( ns, nt = X_s.shape[0], X_t.shape[0] # Initialize weights a, b - if a is None: - a = unif(ns, type_as=X_s) - if b is None: - b = unif(nt, type_as=X_t) + a = check_marginal(a, ns, type_as=X_s) + b = check_marginal(b, nt, type_as=X_t) # Compute rank (see Section 3.1, def 1) r = rank diff --git a/ot/lp/_barycenter_solvers.py b/ot/lp/_barycenter_solvers.py index b938cb799..e7207068c 100644 --- a/ot/lp/_barycenter_solvers.py +++ b/ot/lp/_barycenter_solvers.py @@ -9,7 +9,7 @@ # License: MIT License from ..backend import get_backend -from ..utils import dist +from ..utils import dist, check_marginal from ._network_simplex import emd, emd2 import numpy as np @@ -240,8 +240,7 @@ def free_support_barycenter( N = len(measures_locations) k = X_init.shape[0] d = X_init.shape[1] - if b is None: - b = nx.ones((k,), type_as=X_init) / k + b = check_marginal(b, k, type_as=X_init) if weights is None: weights = nx.ones((N,), type_as=X_init) / N @@ -396,8 +395,7 @@ def generalized_free_support_barycenter( if Y_init is None: Y_init = nx.randn(n_samples_bary, d, type_as=X_list[0]) - if b is None: - b = nx.ones(n_samples_bary, type_as=X_list[0]) / n_samples_bary # not optimized + b = check_marginal(b, n_samples_bary, type_as=X_list[0]) # not optimized out = free_support_barycenter( Z_list, diff --git a/ot/stochastic.py b/ot/stochastic.py index 25b08c6ab..6df3ad771 100644 --- a/ot/stochastic.py +++ b/ot/stochastic.py @@ -10,7 +10,7 @@ # License: MIT License import numpy as np -from .utils import dist, check_random_state +from .utils import dist, check_random_state, check_marginal from .backend import get_backend ############################################################################## @@ -653,11 +653,8 @@ def loss_dual_entropic(u, v, xs, xt, reg=1, ws=None, wt=None, metric="sqeuclidea nx = get_backend(u, v, xs, xt) - if ws is None: - ws = nx.ones(xs.shape[0], type_as=xs) / xs.shape[0] - - if wt is None: - wt = nx.ones(xt.shape[0], type_as=xt) / xt.shape[0] + ws = check_marginal(ws, xs.shape[0], type_as=xs) + wt = check_marginal(wt, xt.shape[0], type_as=xt) if callable(metric): M = metric(xs, xt) @@ -711,11 +708,8 @@ def plan_dual_entropic(u, v, xs, xt, reg=1, ws=None, wt=None, metric="sqeuclidea nx = get_backend(u, v, xs, xt) - if ws is None: - ws = nx.ones(xs.shape[0], type_as=xs) / xs.shape[0] - - if wt is None: - wt = nx.ones(xt.shape[0], type_as=xt) / xt.shape[0] + ws = check_marginal(ws, xs.shape[0], type_as=xs) + wt = check_marginal(wt, xt.shape[0], type_as=xt) if callable(metric): M = metric(xs, xt) @@ -769,11 +763,8 @@ def loss_dual_quadratic(u, v, xs, xt, reg=1, ws=None, wt=None, metric="sqeuclide nx = get_backend(u, v, xs, xt) - if ws is None: - ws = nx.ones(xs.shape[0], type_as=xs) / xs.shape[0] - - if wt is None: - wt = nx.ones(xt.shape[0], type_as=xt) / xt.shape[0] + ws = check_marginal(ws, xs.shape[0], type_as=xs) + wt = check_marginal(wt, xt.shape[0], type_as=xt) if callable(metric): M = metric(xs, xt) @@ -827,11 +818,8 @@ def plan_dual_quadratic(u, v, xs, xt, reg=1, ws=None, wt=None, metric="sqeuclide nx = get_backend(u, v, xs, xt) - if ws is None: - ws = nx.ones(xs.shape[0], type_as=xs) / xs.shape[0] - - if wt is None: - wt = nx.ones(xt.shape[0], type_as=xt) / xt.shape[0] + ws = check_marginal(ws, xs.shape[0], type_as=xs) + wt = check_marginal(wt, xt.shape[0], type_as=xt) if callable(metric): M = metric(xs, xt) diff --git a/ot/utils.py b/ot/utils.py index 6336db07a..bb7dc1884 100644 --- a/ot/utils.py +++ b/ot/utils.py @@ -298,7 +298,7 @@ def check_marginal(a, shape, type_as=None, nx=None): raise ValueError(f"marginal has shape {tuple(a.shape)}, expected {size}") if type_as is not None: if nx is None: - nx = get_backend(type_as) + nx = get_backend(type_as, a) nx.assert_same_dtype_device(type_as, a) return a @@ -2139,8 +2139,7 @@ def split_sample_ratio( n_a = X_a.shape[0] - if a is None: - a = nx.ones(n_a, type_as=X_a) / n_a + a = check_marginal(a, n_a, type_as=X_a) if random_split: if random_state is not None: