diff options
Diffstat (limited to 'ot')
-rw-r--r-- | ot/backend.py | 23 | ||||
-rw-r--r-- | ot/sliced.py | 33 |
2 files changed, 43 insertions, 13 deletions
diff --git a/ot/backend.py b/ot/backend.py index eecf9dd..d661c74 100644 --- a/ot/backend.py +++ b/ot/backend.py @@ -959,6 +959,14 @@ class Backend(): """ raise NotImplementedError() + def matmul(self, a, b): + r""" + Matrix product of two arrays. + + See: https://numpy.org/doc/stable/reference/generated/numpy.matmul.html#numpy.matmul + """ + raise NotImplementedError() + class NumpyBackend(Backend): """ @@ -1293,6 +1301,9 @@ class NumpyBackend(Backend): return args[0] return args + def matmul(self, a, b): + return np.matmul(a, b) + class JaxBackend(Backend): """ @@ -1645,6 +1656,9 @@ class JaxBackend(Backend): return jax.lax.stop_gradient((args[0],))[0] return [jax.lax.stop_gradient((a,))[0] for a in args] + def matmul(self, a, b): + return jnp.matmul(a, b) + class TorchBackend(Backend): """ @@ -2098,6 +2112,9 @@ class TorchBackend(Backend): return args[0].detach() return [a.detach() for a in args] + def matmul(self, a, b): + return torch.matmul(a, b) + class CupyBackend(Backend): # pragma: no cover """ @@ -2474,6 +2491,9 @@ class CupyBackend(Backend): # pragma: no cover return args[0] return args + def matmul(self, a, b): + return cp.matmul(a, b) + class TensorflowBackend(Backend): @@ -2865,3 +2885,6 @@ class TensorflowBackend(Backend): if len(args) == 1: return tf.stop_gradient(args[0]) return [tf.stop_gradient(a) for a in args] + + def matmul(self, a, b): + return tnp.matmul(a, b) diff --git a/ot/sliced.py b/ot/sliced.py index 3a1644d..fd86df9 100644 --- a/ot/sliced.py +++ b/ot/sliced.py @@ -260,7 +260,7 @@ def max_sliced_wasserstein_distance(X_s, X_t, a=None, b=None, n_projections=50, def sliced_wasserstein_sphere(X_s, X_t, a=None, b=None, n_projections=50, - p=2, seed=None, log=False): + p=2, projections=None, seed=None, log=False): r""" Compute the spherical sliced-Wasserstein discrepancy. @@ -287,6 +287,8 @@ def sliced_wasserstein_sphere(X_s, X_t, a=None, b=None, n_projections=50, Number of projections used for the Monte-Carlo approximation p: float, optional (default=2) Power p used for computing the spherical sliced Wasserstein + projections: shape (n_projections, dim, 2), optional + Projection matrix (n_projections and seed are not used in this case) seed: int or RandomState or None, optional Seed used for random number generator log: bool, optional @@ -326,22 +328,25 @@ def sliced_wasserstein_sphere(X_s, X_t, a=None, b=None, n_projections=50, if nx.any(nx.abs(nx.sum(X_s**2, axis=-1) - 1) > 10**(-4)): raise ValueError("X_s is not on the sphere.") if nx.any(nx.abs(nx.sum(X_t**2, axis=-1) - 1) > 10**(-4)): - raise ValueError("Xt is not on the sphere.") + raise ValueError("X_t is not on the sphere.") - # Uniforms and independent samples on the Stiefel manifold V_{d,2} - if isinstance(seed, np.random.RandomState) and str(nx) == 'numpy': - Z = seed.randn(n_projections, d, 2) + if projections is None: + # Uniforms and independent samples on the Stiefel manifold V_{d,2} + if isinstance(seed, np.random.RandomState) and str(nx) == 'numpy': + Z = seed.randn(n_projections, d, 2) + else: + if seed is not None: + nx.seed(seed) + Z = nx.randn(n_projections, d, 2, type_as=X_s) + + projections, _ = nx.qr(Z) else: - if seed is not None: - nx.seed(seed) - Z = nx.randn(n_projections, d, 2, type_as=X_s) - - projections, _ = nx.qr(Z) + n_projections = projections.shape[0] # Projection on S^1 # Projection on plane - Xps = nx.transpose(nx.reshape(nx.dot(nx.transpose(projections, (0, 2, 1))[:, None], X_s[:, :, None]), (n_projections, 2, n)), (0, 2, 1)) - Xpt = nx.transpose(nx.reshape(nx.dot(nx.transpose(projections, (0, 2, 1))[:, None], X_t[:, :, None]), (n_projections, 2, m)), (0, 2, 1)) + Xps = nx.einsum("ikj, lk -> ilj", projections, X_s) + Xpt = nx.einsum("ikj, lk -> ilj", projections, X_t) # Projection on sphere Xps = Xps / nx.sqrt(nx.sum(Xps**2, -1, keepdims=True)) @@ -425,9 +430,11 @@ def sliced_wasserstein_sphere_unif(X_s, a=None, n_projections=50, seed=None, log # Projection on S^1 # Projection on plane - Xps = nx.transpose(nx.reshape(nx.dot(nx.transpose(projections, (0, 2, 1))[:, None], X_s[:, :, None]), (n_projections, 2, n)), (0, 2, 1)) + Xps = nx.einsum("ikj, lk -> ilj", projections, X_s) + # Projection on sphere Xps = Xps / nx.sqrt(nx.sum(Xps**2, -1, keepdims=True)) + # Get coordinates on [0,1[ Xps_coords = nx.reshape(get_coordinate_circle(nx.reshape(Xps, (-1, 2))), (n_projections, n)) |