diff --git a/RELEASES.md b/RELEASES.md index a83c8ba2d..ad764b1e4 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -4,6 +4,9 @@ #### New features + +- `ot.dr.wda` gains `solver='torch'`, a PyTorch autodiff solver with Riemannian gradient descent, usable on installations without autograd or pymanopt, and accepting torch tensors directly (PR #858, Issue #806) +- `ot.dr` dependencies (autograd, pymanopt, scikit-learn, torch) are now imported optionally, so importing `ot.dr` no longer requires all of them; each function raises an explicit `ImportError` naming what it needs (PR #858) - 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 diff --git a/ot/dr.py b/ot/dr.py index 9cfabff5f..e55f92115 100644 --- a/ot/dr.py +++ b/ot/dr.py @@ -18,22 +18,54 @@ from scipy import linalg +# ot.dr offers solvers with different dependencies. Each is imported optionally +# so that, for instance, the PyTorch WDA solver works on an installation with no +# autograd or pymanopt. Functions raise an ImportError naming what they need. try: import autograd.numpy as np - from sklearn.decomposition import PCA + HAS_AUTOGRAD = True +except ImportError: # pragma: no cover - depends on the installation + import numpy as np + + HAS_AUTOGRAD = False + +try: import pymanopt import pymanopt.manifolds import pymanopt.optimizers -except ImportError: - raise ImportError( - "Missing dependency for ot.dr. Requires autograd, pymanopt, scikit-learn. You can install with install with 'pip install POT[dr]', or 'conda install autograd pymanopt scikit-learn'" - ) + + HAS_PYMANOPT = True +except ImportError: # pragma: no cover - depends on the installation + HAS_PYMANOPT = False + +try: + import torch + + HAS_TORCH = True +except ImportError: # pragma: no cover - depends on the installation + HAS_TORCH = False + +try: + from sklearn.decomposition import PCA + + HAS_SKLEARN = True +except ImportError: # pragma: no cover - depends on the installation + HAS_SKLEARN = False from .bregman import sinkhorn as sinkhorn_bregman from .utils import dist as dist_utils, check_random_state +def _require(condition, function, dependencies): + if not condition: + raise ImportError( + f"Missing dependency for ot.dr.{function}. Requires {dependencies}. " + "You can install with 'pip install POT[dr]', or " + "'conda install autograd pymanopt scikit-learn'" + ) + + def dist(x1, x2): r"""Compute squared euclidean distance between samples (autograd)""" x1p2 = np.sum(np.square(x1), 1) @@ -79,6 +111,194 @@ def split_classes(X, y): return [X[y == i, :].astype(np.float32) for i in lstsclass] +def _dist_torch(x1, x2): + r"""Squared euclidean distance between samples (torch).""" + return ( + torch.sum(x1**2, 1).reshape((-1, 1)) + + torch.sum(x2**2, 1).reshape((1, -1)) + - 2 * (x1 @ x2.T) + ) + + +def _sinkhorn_torch(w1, w2, M, reg, k): + r"""Sinkhorn algorithm with fixed number of iterations (torch).""" + K = torch.exp(-M / reg) + ui = torch.ones(M.shape[0], dtype=M.dtype, device=M.device) + vi = torch.ones(M.shape[1], dtype=M.dtype, device=M.device) + for _ in range(k): + vi = w2 / (K.T @ ui + 1e-50) + ui = w1 / (K @ vi + 1e-50) + return ui.reshape((-1, 1)) * K * vi.reshape((1, -1)) + + +def _sinkhorn_log_torch(w1, w2, M, reg, k): + r"""Sinkhorn algorithm in log-domain with fixed iterations (torch).""" + Mr = -M / reg + ui = torch.zeros(M.shape[0], dtype=M.dtype, device=M.device) + vi = torch.zeros(M.shape[1], dtype=M.dtype, device=M.device) + log_w1, log_w2 = torch.log(w1), torch.log(w2) + for _ in range(k): + vi = log_w2 - torch.logsumexp(Mr + ui[:, None], 0) + ui = log_w1 - torch.logsumexp(Mr + vi[None, :], 1) + return torch.exp(ui[:, None] + Mr + vi[None, :]) + + +def _stiefel_retract(P, X): + r"""QR retraction onto the Stiefel manifold, with a sign convention.""" + Q, R = torch.linalg.qr(P + X) + return Q * torch.sign(torch.sign(torch.diagonal(R)) + 0.5) + + +def _stiefel_project(P, G): + r"""Project a euclidean gradient onto the tangent space of Stiefel.""" + W = P.T @ G + return G - P @ (0.5 * (W + W.T)) + + +def _wda_cost_torch(P, xc, wc, regmean, reg, k, sinkhorn_solver): + r"""WDA objective: within-class transport cost over between-class.""" + loss_b, loss_w = 0.0, 0.0 + for i, xi in enumerate(xc): + xi = xi @ P + for j, xj in enumerate(xc[i:]): + xj = xj @ P + M = _dist_torch(xi, xj) + G = sinkhorn_solver(wc[i], wc[j + i], M, reg * regmean[i, j], k) + term = torch.sum(G * M) + if j == 0: + loss_w = loss_w + term + else: + loss_b = loss_b + term + if float(loss_b.detach() if torch.is_tensor(loss_b) else loss_b) == 0.0: + raise ValueError( + "The between-class transport cost underflowed to zero, so the WDA " + "objective is undefined. reg is too small for the scale of the " + "data: exp(-M/reg) underflows. Increase reg, or use " + "sinkhorn_method='sinkhorn_log'." + ) + return loss_w / loss_b + + +def _wda_torch(X, y, p, reg, k, sinkhorn_method, maxiter, verbose, P0, normalize): + r"""WDA solved with PyTorch autodiff and Riemannian gradient descent. + + Mirrors the pymanopt ``SteepestDescent`` path: projected gradient, QR + retraction and backtracking, so both solvers target the same optimum. + """ + dtype = X.dtype + device = X.device + labels = torch.unique(y) + xc = [X[y == c] for c in labels] + wc = [ + torch.full((x.shape[0],), 1.0 / x.shape[0], dtype=dtype, device=device) + for x in xc + ] + d = X.shape[1] + nc = len(xc) + + if not 1 <= p <= d: + raise ValueError(f"Need d >= p >= 1. Values supplied were d = {d} and p = {p}") + + if P0 is None: + P = torch.linalg.qr(torch.randn(d, p, dtype=dtype, device=device))[0] + else: + P = P0.clone().to(dtype) + + regmean = torch.ones((nc, nc), dtype=dtype, device=device) + if P0 is not None and normalize: + with torch.no_grad(): + for i, xi in enumerate(xc): + xi = xi @ P + for j, xj in enumerate(xc[i:]): + xj = xj @ P + regmean[i, j] = torch.sum(_dist_torch(xi, xj)) / ( + xi.shape[0] * xj.shape[0] + ) + + if sinkhorn_method.lower() == "sinkhorn": + solver_fn = _sinkhorn_torch + elif sinkhorn_method.lower() == "sinkhorn_log": + solver_fn = _sinkhorn_log_torch + else: + raise ValueError("Unknown Sinkhorn method '%s'." % sinkhorn_method) + + def value(Q): + with torch.no_grad(): + return _wda_cost_torch(Q, xc, wc, regmean, reg, k, solver_fn) + + f = value(P) + step = 1.0 + for it in range(maxiter): + Q = P.detach().requires_grad_(True) + v = _wda_cost_torch(Q, xc, wc, regmean, reg, k, solver_fn) + (g,) = torch.autograd.grad(v, Q) + direction = -_stiefel_project(P, g) + gnorm = float(torch.linalg.norm(direction)) + if verbose: + print(f"{it + 1:<6d} {float(v.detach()):+.16e} {gnorm:.8e}") + if gnorm <= 1e-12: + break + # start from twice the last accepted step, as pymanopt's backtracking + # line search does, so progress is not throttled by a fixed unit step + step = min(2.0 * step, 1e4 / (gnorm + 1e-12)) + improved = False + for _ in range(40): + Pn = _stiefel_retract(P, step * direction) + fn = value(Pn) + if fn < f: + improved = True + break + step *= 0.5 + if not improved: + break + P, f = Pn, fn + return P.detach() + + +def _wda_torch_entry(X, y, p, reg, k, sinkhorn_method, maxiter, verbose, P0, normalize): + r"""Convert inputs, centre, run the torch solver, return ``(P, proj)``. + + numpy in gives numpy out; a torch tensor in keeps its device and dtype. + """ + was_numpy = not torch.is_tensor(X) + if was_numpy: + # match the autograd path, which promotes to float64 via P + Xt = torch.as_tensor(np.asarray(X, dtype=np.float64)) + else: + Xt = X if torch.is_floating_point(X) else X.to(torch.float64) + # labels may be strings or any hashable, which torch cannot hold, so index + # them by position the way numpy's split_classes does + if torch.is_tensor(y): + yt = y + else: + yt = torch.as_tensor(np.unique(np.asarray(y), return_inverse=True)[1]) + if P0 is None: + P0t = None + else: + P0t = (P0 if torch.is_tensor(P0) else torch.as_tensor(P0)).to(Xt.dtype) + + mx = Xt.mean(dim=0) + Xc = Xt - mx.reshape((1, -1)) + + Popt = _wda_torch( + Xc, yt, p, reg, k, sinkhorn_method, maxiter, verbose, P0t, normalize + ) + + if was_numpy: + Pn = Popt.detach().cpu().numpy() + mxn = mx.detach().cpu().numpy() + + def proj(Z): + return (Z - mxn.reshape((1, -1))).dot(Pn) + + return Pn, proj + + def proj(Z): + return (Z - mx.reshape((1, -1))) @ Popt + + return Popt, proj + + def fda(X, y, p=2, reg=1e-16): r"""Fisher Discriminant Analysis @@ -184,9 +404,17 @@ def wda( Size of dimensionality reduction. reg : float, optional Regularization term >0 (entropic regularization) - solver : None | str, optional - None for steepest descent or 'TrustRegions' for trust regions algorithm - else should be a pymanopt.solvers + solver : None | str | pymanopt.optimizers, optional + Chooses both the autodiff framework and the optimizer. + + - `None` or `'autograd'` (default): autograd and pymanopt + `SteepestDescent`. + - `'TrustRegions'` (or `'tr'`): autograd and pymanopt `TrustRegions`. + - a `pymanopt.optimizers` instance: autograd with that optimizer. + - `'torch'`: PyTorch autodiff with Riemannian gradient descent and a QR + retraction. Requires only `torch`, so it works on installations + without autograd or pymanopt, and accepts torch tensors directly, + keeping their device and dtype. sinkhorn_method : str method used for the Sinkhorn solver, either 'sinkhorn' or 'sinkhorn_log' P0 : ndarray, shape (d, p) @@ -211,6 +439,20 @@ def wda( Wasserstein Discriminant Analysis. arXiv preprint arXiv:1608.08063. """ # noqa + if solver == "torch": + _require(HAS_TORCH, "wda(solver='torch')", "torch") + return _wda_torch_entry( + X, y, p, reg, k, sinkhorn_method, maxiter, verbose, P0, normalize + ) + + _require( + HAS_AUTOGRAD and HAS_PYMANOPT, + "wda(solver='autograd')", + "autograd and pymanopt", + ) + if solver == "autograd": + solver = None + if sinkhorn_method.lower() == "sinkhorn": sinkhorn_solver = sinkhorn elif sinkhorn_method.lower() == "sinkhorn_log": @@ -495,6 +737,7 @@ def ewca( X = X - X.mean(0) if U0 is None: + _require(HAS_SKLEARN, "ewca", "scikit-learn") pca_fitted = PCA(n_components=k).fit(X) U = pca_fitted.components_.T if method == "MM": diff --git a/test/test_dr.py b/test/test_dr.py index dcb477717..86ae9ce32 100644 --- a/test/test_dr.py +++ b/test/test_dr.py @@ -13,10 +13,17 @@ try: # test if autograd and pymanopt are installed import ot.dr - nogo = False + nogo = not (ot.dr.HAS_AUTOGRAD and ot.dr.HAS_PYMANOPT) except ImportError: nogo = True +try: + import torch + + notorch = False +except ImportError: + notorch = True + @pytest.mark.skipif(nogo, reason="Missing modules (autograd or pymanopt)") def test_fda(): @@ -251,3 +258,173 @@ def test_ewca(): U_last_eigvec = np.linalg.svd(X.T, full_matrices=False)[0][:, -k:] _, cos, _ = np.linalg.svd(U.T @ U_last_eigvec, full_matrices=False) assert np.allclose(cos, np.ones(k), atol=1e-3) + + +@pytest.mark.skipif(notorch, reason="Missing module (torch)") +def test_wda_torch_solver(): + rng = np.random.RandomState(0) + xs, ys = ot.datasets.make_data_classif("gaussrot", 90, random_state=rng) + xs = np.hstack((xs, rng.randn(90, 4))) + p = 2 + + P, proj = ot.dr.wda(xs, ys, p, maxiter=10, solver="torch") + + np.testing.assert_allclose(np.sum(P**2, 0), np.ones(p), rtol=1e-6) + assert proj(xs).shape == (90, p) + + +@pytest.mark.skipif(notorch, reason="Missing module (torch)") +def test_wda_torch_accepts_torch_tensors(): + rng = np.random.RandomState(0) + xs, ys = ot.datasets.make_data_classif("gaussrot", 90, random_state=rng) + xt = torch.tensor(xs, dtype=torch.float64) + yt = torch.tensor(ys) + + P, proj = ot.dr.wda(xt, yt, 2, maxiter=5, solver="torch") + + assert torch.is_tensor(P) + assert P.dtype == torch.float64 + assert proj(xt).shape == (90, 2) + + +@pytest.mark.skipif(notorch, reason="Missing module (torch)") +def test_wda_torch_does_not_modify_input(): + rng = np.random.RandomState(0) + xs, ys = ot.datasets.make_data_classif("gaussrot", 90, random_state=rng) + xs = xs + 10.0 + xs_copy = xs.copy() + + ot.dr.wda(xs, ys, 2, maxiter=5, solver="torch") + + np.testing.assert_allclose(xs, xs_copy) + + +@pytest.mark.skipif(notorch, reason="Missing module (torch)") +def test_wda_torch_sinkhorn_log(): + rng = np.random.RandomState(0) + xs, ys = ot.datasets.make_data_classif("gaussrot", 90, random_state=rng) + p = 2 + + P, _ = ot.dr.wda( + xs, ys, p, maxiter=10, solver="torch", sinkhorn_method="sinkhorn_log" + ) + + np.testing.assert_allclose(np.sum(P**2, 0), np.ones(p), rtol=1e-6) + + +@pytest.mark.skipif(nogo or notorch, reason="Missing modules") +def test_wda_backends_agree_on_cost_and_gradient(): + """The torch objective and its gradient must match the autograd ones.""" + import autograd + import autograd.numpy as anp + + rng = np.random.RandomState(0) + n, d, C, reg, k = 180, 6, 3, 1.0, 10 + X = np.vstack([rng.randn(n // C, d) + 3 * rng.randn(1, d) for _ in range(C)]) + y = np.repeat(np.arange(C), n // C) + P0 = np.linalg.qr(rng.randn(d, 2))[0] + Xc = X - X.mean(0) + + xc = [np.ascontiguousarray(Xc[y == c]) for c in range(C)] + wc = [np.ones(x.shape[0]) / x.shape[0] for x in xc] + + def cost_autograd(P): + loss_b, loss_w = 0.0, 0.0 + for i, xi in enumerate(xc): + xi = anp.dot(xi, P) + for j, xj in enumerate(xc[i:]): + xj = anp.dot(xj, P) + M = ot.dr.dist(xi, xj) + G = ot.dr.sinkhorn(wc[i], wc[j + i], M, reg, k) + term = anp.sum(G * M) + if j == 0: + loss_w = loss_w + term + else: + loss_b = loss_b + term + return loss_w / loss_b + + xct = [torch.tensor(x, dtype=torch.float64) for x in xc] + wct = [torch.tensor(w, dtype=torch.float64) for w in wc] + rmt = torch.ones((C, C), dtype=torch.float64) + Pt = torch.tensor(P0, dtype=torch.float64, requires_grad=True) + v = ot.dr._wda_cost_torch(Pt, xct, wct, rmt, reg, k, ot.dr._sinkhorn_torch) + (g,) = torch.autograd.grad(v, Pt) + + np.testing.assert_allclose(float(v.detach()), float(cost_autograd(P0)), rtol=1e-10) + np.testing.assert_allclose( + g.numpy(), autograd.grad(cost_autograd)(P0), rtol=1e-8, atol=1e-10 + ) + + +@pytest.mark.skipif(nogo or notorch, reason="Missing modules") +def test_wda_solvers_reach_comparable_objective(): + """Both solvers minimise the same objective, so neither should be much worse.""" + rng = np.random.RandomState(0) + xs, ys = ot.datasets.make_data_classif("gaussrot", 120, random_state=rng) + xs = np.hstack((xs, rng.randn(120, 3))) + P0 = np.linalg.qr(rng.randn(xs.shape[1], 2))[0] + + Pa, _ = ot.dr.wda(xs, ys, 2, maxiter=40, P0=P0) + Pt, _ = ot.dr.wda(xs, ys, 2, maxiter=40, P0=P0, solver="torch") + + Xc = torch.tensor(xs - xs.mean(0), dtype=torch.float64) + yt = torch.tensor(ys) + xc = [Xc[yt == c] for c in torch.unique(yt)] + wc = [torch.full((x.shape[0],), 1.0 / x.shape[0], dtype=torch.float64) for x in xc] + rm = torch.ones((len(xc), len(xc)), dtype=torch.float64) + + def objective(P): + return float( + ot.dr._wda_cost_torch( + torch.tensor(P, dtype=torch.float64), + xc, + wc, + rm, + 1, + 10, + ot.dr._sinkhorn_torch, + ) + ) + + assert objective(Pt) < 1.15 * objective(Pa) + + +@pytest.mark.skipif(notorch, reason="Missing module (torch)") +def test_wda_torch_unknown_sinkhorn_method(): + rng = np.random.RandomState(0) + xs, ys = ot.datasets.make_data_classif("gaussrot", 60, random_state=rng) + + with pytest.raises(ValueError): + ot.dr.wda(xs, ys, 2, maxiter=2, solver="torch", sinkhorn_method="nope") + + +@pytest.mark.skipif(notorch, reason="Missing module (torch)") +def test_wda_torch_non_numeric_labels(): + """Labels need not be numeric: numpy indexes by value, torch cannot.""" + rng = np.random.RandomState(0) + xs = np.vstack([rng.randn(40, 5) + 3 * rng.randn(1, 5) for _ in range(2)]) + ys = np.array(["cat"] * 40 + ["dog"] * 40) + + P, _ = ot.dr.wda(xs, ys, 2, maxiter=3, solver="torch") + + assert P.shape == (5, 2) + + +@pytest.mark.skipif(notorch, reason="Missing module (torch)") +def test_wda_torch_rejects_p_larger_than_d(): + rng = np.random.RandomState(0) + xs, ys = ot.datasets.make_data_classif("gaussrot", 60, random_state=rng) + + with pytest.raises(ValueError): + ot.dr.wda(xs, ys, xs.shape[1] + 1, maxiter=3, solver="torch") + + +@pytest.mark.skipif(notorch, reason="Missing module (torch)") +def test_wda_torch_numpy_input_returns_float64(): + """numpy in gives float64 out, matching the autograd solver.""" + rng = np.random.RandomState(0) + xs, ys = ot.datasets.make_data_classif("gaussrot", 60, random_state=rng) + + P, _ = ot.dr.wda(xs.astype(np.float32), ys, 2, maxiter=3, solver="torch") + + assert P.dtype == np.float64