[MRG] Nystrom sinkhorn (#742)

* Create _approx_kernel.py

* first implem of nystrom

* Put everything in low rank + doc

* Add: doc, tests for Nystrom and improve computation low rank sinkhorn

* Delete test_ny.py

* Update RELEASES.md

* better names + docs

* test+bugfix+examples

* Update .gitignore

* fix doc and add test

* remove docs

* doc fix + more coverage

* Update test_lowrank.py

* remove multiple targets from low rank sinkhorn + little more coverage

* Update RELEASES.md

* Add in solver + randperm backened + tests

* number paper

* fix test and plot

* fix tf backened + nx.sqrt + doc

* fix none seed

* improve coverage and fix doc

* pass seed none

* move to lowrank folder

---------

Co-authored-by: Rémi Flamary <remi.flamary@gmail.com>

Author: tvayer

README.md

@@ -20,7 +20,7 @@ Source Code (MIT):
 
 POT has the following main features:
 * A large set of differentiable solvers for optimal transport problems, including:
-  *  Exact linear OT, entropic and quadratic regularized OT, 
+  *  Exact linear OT, entropic and quadratic regularized OT,
   *  Gromov-Wasserstein (GW) distances, Fused GW distances and variants of
      quadratic OT,
   *  Unbalanced and partial OT for different divergences,
@@ -444,3 +444,5 @@ Artificial Intelligence.
 [78] Martin, R. D., Medri, I., Bai, Y., Liu, X., Yan, K., Rohde, G. K., & Kolouri, S. (2024). [LCOT: Linear Circular Optimal Transport](https://openreview.net/forum?id=49z97Y9lMq). International Conference on Learning Representations.
 
 [79] Liu, X., Bai, Y., Martín, R. D., Shi, K., Shahbazi, A., Landman, B. A., Chang, C., & Kolouri, S. (2025). [Linear Spherical Sliced Optimal Transport: A Fast Metric for Comparing Spherical Data](https://openreview.net/forum?id=fgUFZAxywx). International Conference on Learning Representations.
+
+[80] Altschuler, J., Bach, F., Rudi, A., Niles-Weed, J., [Massively scalable Sinkhorn distances via the Nyström method](https://proceedings.neurips.cc/paper_files/paper/2019/file/f55cadb97eaff2ba1980e001b0bd9842-Paper.pdf), Advances in Neural Information Processing Systems, 2019.

RELEASES.md

@@ -20,6 +20,7 @@
 - Backend implementation of `ot.dist` for (PR #701)
 - Updated documentation Quickstart guide and User guide with new API (PR #726)
 - Fix jax version for auto-grad (PR #732)
+- Add Nystrom kernel approximation for Sinkhorn (PR #742)
 - Added `ot.solver_1d.linear_circular_ot` and `ot.sliced.linear_sliced_wasserstein_sphere` (PR #736)
 - Implement 1d solver for partial optimal transport (PR #741)
 - Fix reg_div function compatibility with numpy in `ot.unbalanced.lbfgsb_unbalanced` via new function `ot.utils.fun_to_numpy` (PR #731)
@@ -48,7 +49,7 @@ This new release contains several new features, starting with
 a novel [Gaussian Mixture Model Optimal Transport (GMM-OT)](https://pythonot.github.io/master/gen_modules/ot.gmm.html#examples-using-ot-gmm-gmm-ot-apply-map) solver to compare GMM while enforcing the transport plan to remain a GMM, that benefits from a closed-form solution making it practical for high-dimensional matching problems. We also extended our general unbalanced OT solvers to support any non-negative reference measure in the regularization terms, before adding the novel [translation invariant UOT](https://pythonot.github.io/master/auto_examples/unbalanced-partial/plot_conv_sinkhorn_ti.html) solver showcasing a higher convergence speed. We also implemented several new solvers and enhanced existing ones to perform OT across spaces. These include a [semi-relaxed FGW barycenter](https://pythonot.github.io/master/auto_examples/gromov/plot_semirelaxed_gromov_wasserstein_barycenter.html) solver, coupled with new initialization heuristics for the inner divergence computation, to perform graph partitioning or dictionary learning. Followed by novel [unbalanced FGW and Co-optimal transport](https://pythonot.github.io/master/auto_examples/others/plot_outlier_detection_with_COOT_and_unbalanced_COOT.html) solvers to promote robustness to outliers in such matching problems. And we finally updated the implementation of partial GW now supporting asymmetric structures and the KL divergence, while leveraging a new generic conditional gradient solver for partial transport problems enabling significant speed improvements. These latest updates required some modifications to the line search functions of our generic conditional gradient solver, paving the way for future improvements to other GW-based solvers. Last but not least, we implemented a pre-commit scheme to automatically correct common programming mistakes likely to be made by our future contributors.
 
 This release also contains few bug fixes, concerning the support of any metric in `ot.emd_1d` / `ot.emd2_1d`, and the support of any weights in `ot.gaussian`.
- 
+
 #### Breaking change
 - Custom functions provided as parameter `line_search` to `ot.optim.generic_conditional_gradient` must now have the signature `line_search(cost, G, deltaG, Mi, cost_G, df_G, **kwargs)`, adding as input `df_G` the gradient of the regularizer evaluated at the transport plan `G`. This change aims at improving speed of solvers having quadratic polynomial functions as regularizer such as the Gromov-Wassertein loss (PR #663).
 

examples/lowrank/plot_nystroem_approximation.py

@@ -0,0 +1,169 @@
+# -*- coding: utf-8 -*-
+"""
+============================
+Nyström approximation for OT
+============================
+
+Shows how to use Nyström kernel approximation for approximating the Sinkhorn algorithm in linear time.
+
+
+"""
+
+# Author: Titouan Vayer <titouan.vayer@inria.fr>
+#
+# License: MIT License
+
+# sphinx_gallery_thumbnail_number = 2
+
+import numpy as np
+from ot.lowrank import kernel_nystroem, sinkhorn_low_rank_kernel
+from ot.bregman import empirical_sinkhorn_nystroem
+import math
+import ot
+import matplotlib.pyplot as plt
+from matplotlib.colors import LogNorm
+
+##############################################################################
+# Generate data
+# -------------
+
+# %%
+offset = 1
+n_samples_per_blob = 500  # We use 2D ''blobs'' data
+random_state = 42
+std = 0.2  # standard deviation
+np.random.seed(random_state)
+
+centers = np.array(
+    [
+        [-offset, -offset],  # Class 0 - blob 1
+        [-offset, offset],  # Class 0 - blob 2
+        [offset, -offset],  # Class 1 - blob 1
+        [offset, offset],  # Class 1 - blob 2
+    ]
+)
+
+X_list = []
+y_list = []
+
+for i, center in enumerate(centers):
+    blob_points = np.random.randn(n_samples_per_blob, 2) * std + center
+    label = 0 if i < 2 else 1
+    X_list.append(blob_points)
+    y_list.append(np.full(n_samples_per_blob, label))
+
+X = np.vstack(X_list)
+y = np.concatenate(y_list)
+Xs = X[y == 0]  # source data
+Xt = X[y == 1]  # target data
+
+
+##############################################################################
+# Plot data
+# ---------
+
+# %%
+plt.scatter(Xs[:, 0], Xs[:, 1], label="Source")
+plt.scatter(Xt[:, 0], Xt[:, 1], label="Target")
+plt.legend()
+
+##############################################################################
+# Compute the Nyström approximation of the Gaussian kernel
+# --------------------------------------------------------
+
+# %%
+reg = 5.0  # proportional to the std of the Gaussian kernel
+anchors = 10  # number of anchor points for the Nyström approximation
+ot.tic()
+left_factor, right_factor = kernel_nystroem(
+    Xs, Xt, anchors=anchors, sigma=math.sqrt(reg / 2.0), random_state=random_state
+)
+ot.toc()
+
+##############################################################################
+# Use this approximation in a Sinkhorn algorithm with low rank kernel.
+# Each matrix/vector product in the Sinkhorn is accelerated
+# since :math:`Kv = K_1 (K_2^\top v)` can be computed in :math:`O(nr)` time
+# instead of :math:`O(n^2)`
+
+# %%
+numItermax = 1000
+stopThr = 1e-7
+verbose = True
+a, b = None, None
+warn = True
+warmstart = None
+ot.tic()
+u, v, dict_log = sinkhorn_low_rank_kernel(
+    K1=left_factor,
+    K2=right_factor,
+    a=a,
+    b=b,
+    numItermax=numItermax,
+    stopThr=stopThr,
+    verbose=verbose,
+    log=True,
+    warn=warn,
+    warmstart=warmstart,
+)
+ot.toc()
+##############################################################################
+# Compare with Sinkhorn
+# ---------------------
+
+# %%
+M = ot.dist(Xs, Xt)
+ot.tic()
+G, log_ = ot.sinkhorn(
+    a=[],
+    b=[],
+    M=M,
+    reg=reg,
+    numItermax=numItermax,
+    verbose=verbose,
+    log=True,
+    warn=warn,
+    warmstart=warmstart,
+)
+ot.toc()
+
+##############################################################################
+# Use directly ot.bregman.empirical_sinkhorn_nystroem
+# --------------------------------------------------
+
+# %%
+ot.tic()
+G_nys = empirical_sinkhorn_nystroem(
+    Xs,
+    Xt,
+    anchors=anchors,
+    reg=reg,
+    numItermax=numItermax,
+    verbose=True,
+    random_state=random_state,
+)[:]
+ot.toc()
+# %%
+ot.tic()
+G_sinkh = ot.bregman.empirical_sinkhorn(
+    Xs, Xt, reg=reg, numIterMax=numItermax, verbose=True
+)
+ot.toc()
+
+##############################################################################
+# Compare OT plans
+# ----------------
+
+fig, ax = plt.subplots(1, 2, figsize=(10, 4), constrained_layout=True)
+vmin = min(G_sinkh.min(), G_nys.min())
+vmax = max(G_sinkh.max(), G_nys.max())
+norm = LogNorm(vmin=vmin, vmax=vmax)
+im0 = ax[0].imshow(G_sinkh, norm=norm, cmap="coolwarm")
+im1 = ax[1].imshow(G_nys, norm=norm, cmap="coolwarm")
+cbar = fig.colorbar(im1, ax=ax, orientation="vertical", fraction=0.046, pad=0.04)
+ax[0].set_title("OT plan Sinkhorn")
+ax[1].set_title("OT plan Nyström Sinkhorn")
+for a in ax:
+    a.set_xticks([])
+    a.set_yticks([])
+plt.show()

ot/backend.py

@@ -779,6 +779,16 @@ def randn(self, *size, type_as=None):
         """
         raise NotImplementedError()
 
+    def randperm(self, size, type_as=None):
+        r"""
+        Returns a random permutation of integers from 0 to n-1.
+
+        This function follows the api from :any:`torch.randperm`
+
+        See: https://docs.pytorch.org/docs/stable/generated/torch.randperm.html
+        """
+        raise NotImplementedError()
+
     def coo_matrix(self, data, rows, cols, shape=None, type_as=None):
         r"""
         Creates a sparse tensor in COOrdinate format.
@@ -929,6 +939,16 @@ def inv(self, a):
         """
         raise NotImplementedError()
 
+    def pinv(self, a, hermitian=False):
+        r"""
+        Computes the pseudo inverse of a matrix.
+
+        This function follows the api from :any:`numpy.linalg.pinv`.
+
+        See: https://numpy.org/devdocs/reference/generated/numpy.linalg.pinv.html
+        """
+        raise NotImplementedError()
+
     def sqrtm(self, a):
         r"""
         Computes the matrix square root.
@@ -1283,6 +1303,11 @@ def rand(self, *size, type_as=None):
     def randn(self, *size, type_as=None):
         return self.rng_.randn(*size)
 
+    def randperm(self, size, type_as=None):
+        if not isinstance(size, int):
+            raise ValueError("size must be an integer")
+        return self.rng_.permutation(size)
+
     def coo_matrix(self, data, rows, cols, shape=None, type_as=None):
         if type_as is None:
             return coo_matrix((data, (rows, cols)), shape=shape)
@@ -1368,6 +1393,9 @@ def trace(self, a):
     def inv(self, a):
         return scipy.linalg.inv(a)
 
+    def pinv(self, a, hermitian=False):
+        return np.linalg.pinv(a, hermitian=hermitian)
+
     def sqrtm(self, a):
         L, V = np.linalg.eigh(a)
         L = np.sqrt(L)
@@ -1690,6 +1718,15 @@ def randn(self, *size, type_as=None):
         else:
             return jax.random.normal(subkey, shape=size)
 
+    def randperm(self, size, type_as=None):
+        self.rng_, subkey = jax.random.split(self.rng_)
+        if not isinstance(size, int):
+            raise ValueError("size must be an integer")
+        if type_as is not None:
+            return jax.random.permutation(subkey, size).astype(type_as.dtype)
+        else:
+            return jax.random.permutation(subkey, size)
+
     def coo_matrix(self, data, rows, cols, shape=None, type_as=None):
         # Currently, JAX does not support sparse matrices
         data = self.to_numpy(data)
@@ -1781,6 +1818,9 @@ def trace(self, a):
     def inv(self, a):
         return jnp.linalg.inv(a)
 
+    def pinv(self, a, hermitian=False):
+        return jnp.linalg.pinv(a, hermitian=hermitian)
+
     def sqrtm(self, a):
         L, V = jnp.linalg.eigh(a)
         L = jnp.sqrt(L)
@@ -2161,7 +2201,9 @@ def reshape(self, a, shape):
         return torch.reshape(a, shape)
 
     def seed(self, seed=None):
-        if isinstance(seed, int):
+        if seed is None:
+            pass
+        elif isinstance(seed, int):
             self.rng_.manual_seed(seed)
             self.rng_cuda_.manual_seed(seed)
         elif isinstance(seed, torch.Generator):
@@ -2200,6 +2242,22 @@ def randn(self, *size, type_as=None):
         else:
             return torch.randn(size=size, generator=self.rng_)
 
+    def randperm(self, size, type_as=None):
+        if not isinstance(size, int):
+            raise ValueError("size must be an integer")
+        if type_as is not None:
+            generator = (
+                self.rng_cuda_ if self.device_type(type_as) == "GPU" else self.rng_
+            )
+            return torch.randperm(
+                n=size,
+                dtype=type_as.dtype,
+                generator=generator,
+                device=type_as.device,
+            )
+        else:
+            return torch.randperm(n=size, generator=self.rng_)
+
     def coo_matrix(self, data, rows, cols, shape=None, type_as=None):
         if type_as is None:
             return torch.sparse_coo_tensor(torch.stack([rows, cols]), data, size=shape)
@@ -2314,6 +2372,9 @@ def trace(self, a):
     def inv(self, a):
         return torch.linalg.inv(a)
 
+    def pinv(self, a, hermitian=False):
+        return torch.linalg.pinv(a, hermitian=hermitian)
+
     def sqrtm(self, a):
         L, V = torch.linalg.eigh(a)
         L = torch.sqrt(L)
@@ -2624,6 +2685,15 @@ def randn(self, *size, type_as=None):
             with cp.cuda.Device(type_as.device):
                 return self.rng_.randn(*size, dtype=type_as.dtype)
 
+    def randperm(self, size, type_as=None):
+        if not isinstance(size, int):
+            raise ValueError("size must be an integer")
+        if type_as is None:
+            return self.rng_.permutation(size)
+        else:
+            with cp.cuda.Device(type_as.device):
+                return self.rng_.permutation(size).astype(type_as.dtype)
+
     def coo_matrix(self, data, rows, cols, shape=None, type_as=None):
         data = self.from_numpy(data)
         rows = self.from_numpy(rows)
@@ -2728,6 +2798,9 @@ def trace(self, a):
     def inv(self, a):
         return cp.linalg.inv(a)
 
+    def pinv(self, a, hermitian=False):
+        return cp.linalg.pinv(a)
+
     def sqrtm(self, a):
         L, V = cp.linalg.eigh(a)
         L = cp.sqrt(L)
@@ -3048,6 +3121,19 @@ def randn(self, *size, type_as=None):
         else:
             return self.rng_.normal(size, dtype=type_as.dtype)
 
+    def randperm(self, size, type_as=None):
+        if not isinstance(size, int):
+            raise ValueError("size must be an integer")
+        local_seed = self.rng_.make_seeds(2)[0]
+        if type_as is None:
+            return tf.random.experimental.stateless_shuffle(
+                tf.range(size), seed=local_seed
+            )
+        else:
+            return tf.random.experimental.stateless_shuffle(
+                tf.range(size, dtype=type_as.dtype), seed=local_seed
+            )
+
     def _convert_to_index_for_coo(self, tensor):
         if isinstance(tensor, self.__type__):
             return int(self.max(tensor)) + 1
@@ -3164,6 +3250,9 @@ def trace(self, a):
     def inv(self, a):
         return tf.linalg.inv(a)
 
+    def pinv(self, a, hermitian=False):
+        return tf.linalg.pinv(a)
+
     def sqrtm(self, a):
         L, V = tf.linalg.eigh(a)
         L = tf.sqrt(L)