From 450719dc5a1c21aa1307736531b73a3bc730b2cc Mon Sep 17 00:00:00 2001 From: thibaut-germain Date: Fri, 5 Jun 2026 13:39:39 +0200 Subject: [PATCH 01/13] correction SGOT cost matrix, added to contributor list, move sgot example to other --- CITATION.cff | 6 ++++++ CONTRIBUTORS.md | 2 ++ examples/{ => others}/plot_sgot.py | 0 ot/sgot.py | 2 +- 4 files changed, 9 insertions(+), 1 deletion(-) rename examples/{ => others}/plot_sgot.py (100%) diff --git a/CITATION.cff b/CITATION.cff index 144ec668c..f4bdd166a 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -81,6 +81,12 @@ authors: - given-names: David family-names: Coeurjolly affiliation: CNRS, LIRIS + - given-names: Thibaut + family-names: Germain + affiliation: Ecole Polytechnique + - given-names: Sienna + family-names: O'Shea + affiliation: Ecole Polytechnique identifiers: - type: url diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index aa28dd7ab..08f481b33 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -59,6 +59,8 @@ The contributors to this library are: * [Julie Delon](https://judelo.github.io/) (GMM OT) * [Samuel Boïté](https://samuelbx.github.io/) (GMM OT) * [Nathan Neike](https://github.com/nathanneike) (Sparse EMD solver) +* [Thibaut Germain](https://thibaut-germain.github.io) (SGOT) +* Sienna O'Shea (SGOT) ## Acknowledgments diff --git a/examples/plot_sgot.py b/examples/others/plot_sgot.py similarity index 100% rename from examples/plot_sgot.py rename to examples/others/plot_sgot.py diff --git a/ot/sgot.py b/ot/sgot.py index ff67192c5..22ce5798d 100644 --- a/ot/sgot.py +++ b/ot/sgot.py @@ -124,7 +124,7 @@ def _delta_matrix_1d(Rs, Ls, Rt, Lt, nx=None, eps=1e-12): Ltn = _normalize_columns(Lt, nx=nx, eps=eps) Cr = nx.dot(nx.conj(Rsn).T, Rtn) - Cl = nx.dot(nx.conj(Lsn).T, Ltn) + Cl = nx.dot(Lsn.T, nx.conj(Ltn)) delta = nx.abs(Cr * Cl) delta = nx.clip(delta, 0.0, 1.0) From 6e0bc73cc3adbac2c0007aa5e12f7a303add34e7 Mon Sep 17 00:00:00 2001 From: Sienna O'Shea Date: Fri, 19 Jun 2026 17:00:24 +0200 Subject: [PATCH 02/13] updated graphs --- examples/plot_sgot.py | 576 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 576 insertions(+) create mode 100644 examples/plot_sgot.py diff --git a/examples/plot_sgot.py b/examples/plot_sgot.py new file mode 100644 index 000000000..3d66d5bbf --- /dev/null +++ b/examples/plot_sgot.py @@ -0,0 +1,576 @@ +#!/usr/bin/env python +# coding: utf-8 + +""" +========================================= +Spectral-Grassmann OT on dynamical systems operators +========================================= + +This example presents a synthetic example of Spectral Grassmannian-Wasserstein +Optimal Transport (SGOT) on linear dynamical systems. + +We consider a signal formed by the sum of two damped oscillatory modes evolving +along a rotated direction in the plane. The signal is then associated with an +underlying continuous linear dynamical system, and we study how its spectral +representation varies under rotation. The SGOT cost and metric are used to +compare the reference and rotated systems. + +.. [83] T. Germain; R. Flamary; V. R. Kostic; K. Lounici, A Spectral-Grassmann + Wasserstein Metric for Operator Representations of Dynamical Systems, + arXiv preprint arXiv:2509.24920, 2025. +""" + +# Authors: Sienna O'Shea +# Thibaut Germain +# +# License: MIT License + +import numpy as np +import matplotlib.pyplot as plt + +from ot.sgot import sgot_metric, sgot_cost_matrix + +from scipy.linalg import eig + + +# sampling parameters and time grid +fs = 50 +max_t = 5 +time = np.linspace(0, max_t, fs * max_t) +dt = 1 / fs + + +# %% +# Example: rotating a linear dynamical system in 3D +# ------------------------------------------------- +# +# 1. Build a simple observed signal +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# We begin by assuming that the observed signal is made of two oscillatory +# components: +# +# .. math:: +# +# x_{\text{ref}}(t)=e^{-\tau_1 t}\cos(2\pi\omega_1 t)\,\vec e(\theta) +# \;+\; +# e^{-\tau_2 t}\cos(2\pi\omega_2 t)\,\vec e(\theta), +# +# where :math:`\vec e(\theta)\in\mathbb{R}^2` is a fixed real vector. Thus, +# :math:`x(t)` evolves along the one-dimensional subspace spanned by +# :math:`\vec e(\theta)`, while its time dependence exhibits oscillatory and +# dissipative behaviour. + +tau_0 = np.array([0.08, 0.18]) +freq_0 = np.array([1.0, 2.0]) +theta_0 = np.pi / 4 + + +def generate_data(time, tau, freq, theta): + t_ = np.sin(2 * np.pi * freq[None, :] * time[:, None]) * np.exp( + -tau[None, :] * time[:, None] + ) + t_ = t_.sum(axis=1) + traj_0 = np.zeros((t_.shape[0], 2)) + traj_0[:, 0] = t_ + rotation_matrix = np.array( + [[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]] + ) + traj_0 = traj_0 @ rotation_matrix.T + return traj_0 + + +traj_0 = generate_data(time, tau_0, freq_0, theta_0) + + +# plot the observed signal components and their sum +plt.figure(figsize=(10, 4)) +plt.plot(time, traj_0, label="base trajectory", linewidth=2) +plt.xlabel("time") +plt.ylabel("amplitude") +plt.legend() +plt.title(r"Observed scalar signal along $\vec{e}(\theta)$") +plt.show() + + +# %% +# 2. Interpret the signal as coming from a continuous linear dynamical system +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# We assume that :math:`x(t)` is generated by an underlying continuous linear +# dynamical system. Since the observed signal is a superposition of two +# sinusoidal modes, the corresponding linear dynamics are naturally described +# by a fourth-order model. We therefore introduce the state vector +# +# .. math:: +# +# z(t)= +# \begin{pmatrix} +# x_1(t)\\ +# x_2(t)\\ +# \vdots\\ +# x_1^{(3)}(t)\\ +# x_2^{(3)}(t) +# \end{pmatrix} +# \in\mathbb{R}^8. +# +# where :math:`x^{(n)}(t)` denotes the n-th derivative of :math:`x(t)`. +# +# This allows us to rewrite the dynamics as a first-order linear system: +# +# .. math:: +# +# \dot{z}(t)=Az(t), +# +# where :math:`A\in\mathbb{R}^{8\times 8}`. Its solution is then given by +# +# .. math:: +# +# z(t)=e^{tA}z_0. + +fig = plt.figure(figsize=(9, 6)) +ax = fig.add_subplot(projection="3d") + +ax.plot(time, traj_0[:, 0], traj_0[:, 1]) +ax.set_xlabel("time") +ax.set_ylabel("x₁(t)") +ax.set_title("Observed trajectory in time") + +ax.text2D(1.08, 0.5, "x₂(t)", transform=ax.transAxes, rotation=90, va="center") + +plt.show() + + +# %% +# 3. Sampling and preprocessing discrete trajectories of the dynamical system +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# We now have a bridge between the continuous system and the operator we later +# aim to infer from sampled data. Since in practice we do not observe the full +# continuous trajectory, we work instead with discrete samples of the signal. +# We take snapshots at uniform time intervals :math:`\Delta t`, and write the +# sampled signal as +# +# .. math:: +# +# S= +# \begin{pmatrix} +# x_1(0) & x_2(0)\\ +# x_1(\Delta t) & x_2(\Delta t)\\ +# \vdots\\ +# x_1(N\Delta t) & x_2(N\Delta t)\\ +# \end{pmatrix} +# +# The goal is now to use these observations to recover the operator governing +# the evolution. To do this, we augment the signal :math:`s` using a sliding +# window of length :math:`w`. For each :math:`k`, define +# +# .. math:: +# +# z_k = +# \begin{pmatrix} +# s_k\\ +# s_{k+1}\\ +# \vdots\\ +# s_{k+w-1} +# \end{pmatrix} +# +# We then form the data matrices +# +# .. math:: +# +# X= +# \begin{pmatrix} +# z_1\\ +# z_2\\ +# \vdots\\ +# z_{N-w} +# \end{pmatrix}, +# \qquad +# Y= +# \begin{pmatrix} +# z_2\\ +# z_3\\ +# \vdots\\ +# z_{N-w+1} +# \end{pmatrix}, +# +# so that :math:`X` contains the present windowed states and :math:`Y` the +# corresponding shifted future states. + + +# build a 4-dimensional state using delay embedding +def augment(traj, window_length=2): + Z = np.lib.stride_tricks.sliding_window_view(traj, (window_length, 1)) + Z = Z.reshape(Z.shape[0], -1) + return Z + + +# create the embedded state matrix Z +Z = augment(traj_0, 4) +Z.shape + +# inspect one embedded state vector +Z[0] + +# create X and Y for the SGOT metric +X = Z[:-1] +Y = Z[1:] + +# inspect shapes of X and Y +print("X shape:", X.shape) +print("Y shape:", Y.shape) + + +# %% +# 4. Estimate the discrete-time operator +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# We now identify the operator that maps :math:`X` to :math:`Y`. From +# +# .. math:: +# +# \dot{z}=Az, +# +# we have +# +# .. math:: +# +# z(t+\Delta t)=e^{\Delta tA}z(t). +# +# Setting +# +# .. math:: +# +# B=e^{\Delta tA}, +# +# the corresponding discrete-time evolution is governed by :math:`B`, and we +# seek the best linear map satisfying +# +# .. math:: +# +# Y\approx X B^T. +# +# Equivalently, we solve the optimisation problem +# +# .. math:: +# +# \min_B \|Y-XB\|^2. +# +# We want to recover the best rank-:math:`r` operator, whose estimator is +# defined as follows: +# +# .. math:: +# +# B = C_{xx}^{-\frac{1}{2}}[C_{xx}^{-\frac{1}{2}}C_{xy}]_r +# \quad \text{s.t} \quad C_{xx} = X^T X \quad \text{and} \quad C_{xy} = X^TY. +# +# Here :math:`[\cdot]_r` denotes the best rank-:math:`r` estimator obtained via +# SVD decomposition. [2] +# +# [2] Kostic, V., Novelli, P., Maurer, A., Ciliberto, C., Rosasco, L. and +# Pontil, M., 2022. Learning dynamical systems via Koopman operator regression +# in reproducing kernel Hilbert spaces. Advances in Neural Information +# Processing Systems, 35, pp.4017-4031. + + +def estimator(X, Y, rank=4): + # X: (n_samples, n_features) + # Y: (n_samples, n_features) + + # estimate operator + cxx = X.T @ X + U, S, Vt = np.linalg.svd(cxx) + S_inv = np.divide(1, S, out=np.zeros_like(S), where=S != 0) + cxx_inv_half = Vt.T @ np.diag(np.sqrt(S_inv)) @ U.T + cxy = X.T @ Y + T = cxx_inv_half @ cxy + U, S, Vt = np.linalg.svd(T) + S[rank:] = 0 + T_rank = U @ np.diag(S) @ Vt + T = cxx_inv_half @ T_rank + + # estimate spectral decomposition + val, vl, vr = eig(T, left=True, right=True) + sort_idx = np.argsort(np.abs(val))[::-1] + val = val[sort_idx][:rank] + vl = vl[:, sort_idx][:, :rank] + vr = vr[:, sort_idx][:, :rank] + + return T, {"eig_val": val, "eig_vec_left": vl, "eig_vec_right": vr} + + +B_0, B_0_spec = estimator(X, Y, rank=4) +Y_pred = X @ B_0 + +plt.figure(figsize=(10, 4)) +plt.plot(Y[:, 0], label="true") +plt.plot(Y_pred[:, 0], "--", label="predicted") +plt.xlabel("sample index") +plt.ylabel("first state coordinate") +plt.title("True Signal vs Predicted Signal") +plt.legend() +plt.show() + + +# %% +# The predicted signal is nearly indistinguishable from the true signal, +# indicating that the estimated operator accurately captures the observed +# dynamics. + +# %% +# 6. Recover continuous-time spectral information from the discrete operator +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# To recover the continuous generator :math:`A`, we study the spectral +# structure of :math:`B`. We diagonalise :math:`B` as +# +# .. math:: +# +# B=PDP^{-1}, +# +# where +# +# .. math:: +# +# D=\operatorname{diag}(\mu_1,\dots,\mu_n). +# +# The continuous-time eigenvalues are of the form +# +# .. math:: +# +# \lambda_k=-\tau_k+2\pi i\,\omega_k, +# \qquad k\in\{1,2\}, +# +# and the corresponding eigenvalues of :math:`B` are +# +# .. math:: +# +# \mu_k=e^{\Delta t\lambda_k} +# =e^{\Delta t(-\tau_k+2\pi i\omega_k)}. +# +# Since :math:`B=e^{\Delta tA}`, we recover :math:`A` by taking the logarithm: +# +# .. math:: +# +# A=P\,\frac{\log(D)}{\Delta t}\,P^{-1}. + +D_0 = np.log(B_0_spec["eig_val"]) * fs +L_0 = B_0_spec["eig_vec_left"] +R_0 = B_0_spec["eig_vec_right"] + +recovered_freqs = D_0.imag / (2 * np.pi) +mask = recovered_freqs > 0 +recovered_freqs = recovered_freqs[mask] +decay = -D_0.real[mask] +print(f"First mode: frequency: {recovered_freqs[0]:.2f} Hz -- decay: {decay[0]:.2f}") +print(f"Second mode: frequency: {recovered_freqs[1]:.2f} Hz -- decay: {decay[1]:.2f}") + + +# %% +# Introduction to SGOT for linear operators +# ----------------------------------------- +# +# To compare two linear operators through their spectral structure, we use the +# SGOT framework introduced in Theorem 1 of [1]. For a non-defective +# finite-rank operator :math:`T \in S_r(\mathcal H)`, the theorem associates a +# discrete spectral measure +# +# .. math:: +# +# \mu(T) \triangleq \sum_{j\in[\ell]} +# \frac{m_j}{m_{\mathrm{tot}}}\,\delta_{(\lambda_j,\mathcal V_j)}, +# +# where :math:`\lambda_j` are the eigenvalues of :math:`T`, :math:`m_j` their +# algebraic multiplicities, and :math:`\mathcal V_j` the corresponding +# eigenspaces. Thus, each spectral component of the operator is represented by +# an atom of the form +# +# .. math:: +# +# (\lambda_j,\mathcal V_j), +# +# combining one eigenvalue with its associated invariant subspace. +# +# Theorem 1 then defines a ground cost between two such atoms by combining a +# spectral discrepancy and a geometric discrepancy: +# +# .. math:: +# +# d_\eta\big((\lambda,\mathcal V),(\lambda',\mathcal V')\big) +# \triangleq +# \eta\,|\lambda-\lambda'| + (1-\eta)\, d_{\mathcal G}(\mathcal V,\mathcal V'), +# +# where :math:`d_{\mathcal G}` denotes the grassmann distance between +# eigenspaces and :math:`\eta\in(0,1)` balances the contribution of eigenvalues +# and eigenspaces. +# +# The SGOT distance between two operators :math:`T` and :math:`T'` is then the +# Wasserstein distance between their associated spectral measures: +# +# .. math:: +# +# d_S(T,T') = W_{d_\eta,p}\big(\mu(T),\mu(T')\big). +# +# In this way, SGOT compares linear operators by optimally matching their +# spectral atoms, taking into account both the location of eigenvalues and the +# relative geometry of their eigenspaces. + +# %% +# SGOT distance versus rotation angle +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# We compare the reference signal with a rotated version obtained by changing +# only the observation direction. The shifted signal is +# +# .. math:: +# +# x_{\mathrm{shift}}^{\mathrm{rot}}(t;\theta) +# = +# \sum_{i=1}^{2} +# e^{-\tau_i t}\cos(2\pi\omega_i t)\,\vec e({\color{red}\theta}), +# +# while the reference signal is recovered at :math:`\theta=\theta_0`. Thus, +# this experiment isolates the effect of rotating the underlying one-dimensional +# subspace in the observation plane. + +thetas = np.linspace(0, np.pi / 2, 51) +rotation_scores = [] + +for theta in thetas: + Z = augment(generate_data(time, tau_0, freq_0, theta), 4) + B, B_spec = estimator(Z[:-1], Z[1:]) + D = np.log(B_spec["eig_val"]) * fs + L = B_spec["eig_vec_left"] + R = B_spec["eig_vec_right"] + rotation_scores.append( + sgot_metric(D_0, R_0, L_0, D, R, L, eta=0.9, grassmann_metric="chordal") + ) + +fig, ax = plt.subplots(figsize=(7, 4)) +ax.plot(thetas, rotation_scores, linewidth=1.8) +ax.axvline(theta_0, color="gray", linestyle="--", linewidth=0.8) +ax.set_xlabel(r"Rotation angle $\theta$ (rad)") +ax.set_ylabel(r"$d_S$") +ax.set_title("SGOT distance vs. rotation angle") +fig.tight_layout() +plt.show() + +# %% +# Comparison across Grassmannian metrics for SGOT distance versus rotation angle +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +metrics = ["chordal", "geodesic", "procrustes", "martin"] +styles = {"chordal": "-", "geodesic": "--", "procrustes": "-.", "martin": ":"} +rotation_results = {m: [] for m in metrics} + +for theta in thetas: + Z = augment(generate_data(time, tau_0, freq_0, theta), 4) + B, B_spec = estimator(Z[:-1], Z[1:]) + D = np.log(B_spec["eig_val"]) * fs + L = B_spec["eig_vec_left"] + R = B_spec["eig_vec_right"] + for m in metrics: + rotation_results[m].append( + sgot_metric(D_0, R_0, L_0, D, R, L, eta=0.9, grassmann_metric=m) + ) + +fig, ax = plt.subplots(figsize=(7, 4)) +for m in metrics: + ax.plot(thetas, rotation_results[m], styles[m], label=m, linewidth=1.8) +ax.axvline( + theta_0, color="gray", linestyle="--", linewidth=0.8, label=r"$\theta_0 = \pi/4$" +) +ax.set_xlabel(r"Rotation angle $\theta$ (rad)") +ax.set_ylabel(r"$d_S$") +ax.set_title("SGOT distance vs. rotation angle across Grassmannian metrics") +ax.legend() +fig.tight_layout() +plt.show() + +# %% +# SGOT distance versus frequency +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# In this experiment, we keep the reference direction fixed and perturb one of +# the oscillatory modes. The shifted signal is +# +# .. math:: +# +# x_{\mathrm{shift}}^{\omega}(t) +# = +# e^{-\tau_1 t}\cos(2\pi\omega_1 t)\,\vec e(\theta_0) +# \;+\; +# e^{-\tau_2 t}\cos(2\pi{\color{red}\omega_2'} t)\,\vec e(\theta_0), +# +# where only the second frequency is modified. We then study how the SGOT +# distance changes as a function of the perturbed frequency :math:`\omega_2'`. + +omegas = np.linspace(0.5, 3.0, 21) +frequency_scores = {m: [] for m in metrics} + +for omega in omegas: + Z = augment(generate_data(time, tau_0, np.array([freq_0[0], omega]), theta_0), 4) + B, B_spec = estimator(Z[:-1], Z[1:]) + D = np.log(B_spec["eig_val"]) * fs + L = B_spec["eig_vec_left"] + R = B_spec["eig_vec_right"] + for m in metrics: + frequency_scores[m].append( + sgot_metric(D_0, R_0, L_0, D, R, L, eta=0.9, grassmann_metric=m) + ) + +fig, ax = plt.subplots(figsize=(7, 4)) +for m in metrics: + ax.plot(omegas, frequency_scores[m], styles[m], label=m, linewidth=1.8) +ax.axvline( + freq_0[1], color="gray", linestyle="--", linewidth=0.8, label=r"$\omega_2 = 2.0$ Hz" +) +ax.set_xlabel(r"Frequency $\omega_2'$ (Hz)") +ax.set_ylabel(r"$d_S$") +ax.set_title("SGOT distance vs. frequency across Grassmannian metrics") +ax.legend() +fig.tight_layout() +plt.show() + +# %% +# SGOT distance versus decay +# ~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# We now study the effect of changing the decay rate while keeping the +# observation direction fixed. The shifted signal is +# +# .. math:: +# +# x_{\mathrm{shift}}^{\mathrm{decay}}(t;\tau) +# = +# e^{-{\color{red}\tau} t}\cos(2\pi\omega_1 t)\,\vec e(\theta_0) +# \;+\; +# e^{-{\color{red}\tau} t}\cos(2\pi\omega_2' t)\,\vec e(\theta_0). +# +# In this way, both modes share the same modified decay parameter +# :math:`\tau`, allowing us to isolate the influence of dissipation on the SGOT +# distance. +taus = np.linspace(0.1, 3.0, 21) +decay_scores = {m: [] for m in metrics} + +for tau in taus: + Z = augment(generate_data(time, np.array([tau, tau]), freq_0, theta_0), 4) + B, B_spec = estimator(Z[:-1], Z[1:]) + D = np.log(B_spec["eig_val"]) * fs + L = B_spec["eig_vec_left"] + R = B_spec["eig_vec_right"] + for m in metrics: + decay_scores[m].append( + sgot_metric(D_0, R_0, L_0, D, R, L, eta=0.9, grassmann_metric=m) + ) + +fig, ax = plt.subplots(figsize=(7, 4)) +for m in metrics: + ax.plot(taus, decay_scores[m], styles[m], label=m, linewidth=1.8) +ax.set_xlabel(r"Decay rate $\tau$") +ax.set_ylabel(r"$d_S$") +ax.set_title("SGOT distance vs. decay across Grassmannian metrics") +ax.legend() +fig.tight_layout() +plt.show() From 8d0ddbdd73e5b2e6a4853b696ade8bac3dd781d7 Mon Sep 17 00:00:00 2001 From: Sienna O'Shea Date: Thu, 2 Jul 2026 09:32:42 +0200 Subject: [PATCH 03/13] fix plots --- examples/plot_sgot.py | 64 ++++++++++++++++++++++++++++++++----------- 1 file changed, 48 insertions(+), 16 deletions(-) diff --git a/examples/plot_sgot.py b/examples/plot_sgot.py index 3d66d5bbf..87f7ff3c8 100644 --- a/examples/plot_sgot.py +++ b/examples/plot_sgot.py @@ -274,12 +274,8 @@ def augment(traj, window_length=2): # Processing Systems, 35, pp.4017-4031. -def estimator(X, Y, rank=4): - # X: (n_samples, n_features) - # Y: (n_samples, n_features) - - # estimate operator - cxx = X.T @ X +def estimator(X, Y, rank=4, eps=1e-8): + cxx = X.T @ X + eps * np.eye(X.shape[1]) U, S, Vt = np.linalg.svd(cxx) S_inv = np.divide(1, S, out=np.zeros_like(S), where=S != 0) cxx_inv_half = Vt.T @ np.diag(np.sqrt(S_inv)) @ U.T @@ -416,6 +412,24 @@ def estimator(X, Y, rank=4): # spectral atoms, taking into account both the location of eigenvalues and the # relative geometry of their eigenspaces. +# %% +# A wider delay window for the SGOT experiments below +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# The window of length 4 used above is enough to identify a single reference +# operator, but the experiments below also probe signals whose two modes +# nearly coincide in frequency (e.g. :math:`\omega_2'\to\omega_1`). Telling +# such near-degenerate modes apart requires the delay embedding to span +# enough time to "see" their differing decay, so we re-embed the reference +# signal with a longer window before running the sweeps. + +sgot_window = 10 +Z = augment(traj_0, sgot_window) +_, B_0_spec_sgot = estimator(Z[:-1], Z[1:]) +D_0_sgot = np.log(B_0_spec_sgot["eig_val"]) * fs +L_0_sgot = B_0_spec_sgot["eig_vec_left"] +R_0_sgot = B_0_spec_sgot["eig_vec_right"] + # %% # SGOT distance versus rotation angle # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -438,13 +452,15 @@ def estimator(X, Y, rank=4): rotation_scores = [] for theta in thetas: - Z = augment(generate_data(time, tau_0, freq_0, theta), 4) + Z = augment(generate_data(time, tau_0, freq_0, theta), sgot_window) B, B_spec = estimator(Z[:-1], Z[1:]) D = np.log(B_spec["eig_val"]) * fs L = B_spec["eig_vec_left"] R = B_spec["eig_vec_right"] rotation_scores.append( - sgot_metric(D_0, R_0, L_0, D, R, L, eta=0.9, grassmann_metric="chordal") + sgot_metric( + D_0_sgot, R_0_sgot, L_0_sgot, D, R, L, eta=0.9, grassmann_metric="chordal" + ) ) fig, ax = plt.subplots(figsize=(7, 4)) @@ -465,21 +481,27 @@ def estimator(X, Y, rank=4): rotation_results = {m: [] for m in metrics} for theta in thetas: - Z = augment(generate_data(time, tau_0, freq_0, theta), 4) + Z = augment(generate_data(time, tau_0, freq_0, theta), sgot_window) B, B_spec = estimator(Z[:-1], Z[1:]) D = np.log(B_spec["eig_val"]) * fs L = B_spec["eig_vec_left"] R = B_spec["eig_vec_right"] for m in metrics: rotation_results[m].append( - sgot_metric(D_0, R_0, L_0, D, R, L, eta=0.9, grassmann_metric=m) + sgot_metric( + D_0_sgot, R_0_sgot, L_0_sgot, D, R, L, eta=0.9, grassmann_metric=m + ) ) fig, ax = plt.subplots(figsize=(7, 4)) for m in metrics: ax.plot(thetas, rotation_results[m], styles[m], label=m, linewidth=1.8) ax.axvline( - theta_0, color="gray", linestyle="--", linewidth=0.8, label=r"$\theta_0 = \pi/4$" + theta_0, + color="gray", + linestyle="--", + linewidth=0.8, + label=r"$\theta_0 = \pi/4$ (reference)", ) ax.set_xlabel(r"Rotation angle $\theta$ (rad)") ax.set_ylabel(r"$d_S$") @@ -510,21 +532,29 @@ def estimator(X, Y, rank=4): frequency_scores = {m: [] for m in metrics} for omega in omegas: - Z = augment(generate_data(time, tau_0, np.array([freq_0[0], omega]), theta_0), 4) + Z = augment( + generate_data(time, tau_0, np.array([freq_0[0], omega]), theta_0), sgot_window + ) B, B_spec = estimator(Z[:-1], Z[1:]) D = np.log(B_spec["eig_val"]) * fs L = B_spec["eig_vec_left"] R = B_spec["eig_vec_right"] for m in metrics: frequency_scores[m].append( - sgot_metric(D_0, R_0, L_0, D, R, L, eta=0.9, grassmann_metric=m) + sgot_metric( + D_0_sgot, R_0_sgot, L_0_sgot, D, R, L, eta=0.9, grassmann_metric=m + ) ) fig, ax = plt.subplots(figsize=(7, 4)) for m in metrics: ax.plot(omegas, frequency_scores[m], styles[m], label=m, linewidth=1.8) ax.axvline( - freq_0[1], color="gray", linestyle="--", linewidth=0.8, label=r"$\omega_2 = 2.0$ Hz" + freq_0[1], + color="gray", + linestyle="--", + linewidth=0.8, + label=r"$\omega_2 = 2.0$ Hz (reference)", ) ax.set_xlabel(r"Frequency $\omega_2'$ (Hz)") ax.set_ylabel(r"$d_S$") @@ -555,14 +585,16 @@ def estimator(X, Y, rank=4): decay_scores = {m: [] for m in metrics} for tau in taus: - Z = augment(generate_data(time, np.array([tau, tau]), freq_0, theta_0), 4) + Z = augment(generate_data(time, np.array([tau, tau]), freq_0, theta_0), sgot_window) B, B_spec = estimator(Z[:-1], Z[1:]) D = np.log(B_spec["eig_val"]) * fs L = B_spec["eig_vec_left"] R = B_spec["eig_vec_right"] for m in metrics: decay_scores[m].append( - sgot_metric(D_0, R_0, L_0, D, R, L, eta=0.9, grassmann_metric=m) + sgot_metric( + D_0_sgot, R_0_sgot, L_0_sgot, D, R, L, eta=0.9, grassmann_metric=m + ) ) fig, ax = plt.subplots(figsize=(7, 4)) From f48d6cea11f357c2c12d0a28430b541025882854 Mon Sep 17 00:00:00 2001 From: thibaut-germain Date: Thu, 2 Jul 2026 13:21:01 +0200 Subject: [PATCH 04/13] update example --- examples/others/plot_sgot.py | 249 +++++++------- examples/plot_sgot.py | 608 ----------------------------------- 2 files changed, 130 insertions(+), 727 deletions(-) delete mode 100644 examples/plot_sgot.py diff --git a/examples/others/plot_sgot.py b/examples/others/plot_sgot.py index 0f4ee1dee..4dc3b6a6e 100644 --- a/examples/others/plot_sgot.py +++ b/examples/others/plot_sgot.py @@ -66,6 +66,10 @@ theta_0 = np.pi / 4 +def rotation_matrix(theta): + return np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]) + + def generate_data(time, tau, freq, theta): t_ = np.sin(2 * np.pi * freq[None, :] * time[:, None]) * np.exp( -tau[None, :] * time[:, None] @@ -73,26 +77,24 @@ def generate_data(time, tau, freq, theta): t_ = t_.sum(axis=1) traj_0 = np.zeros((t_.shape[0], 2)) traj_0[:, 0] = t_ - rotation_matrix = np.array( - [[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]] - ) - traj_0 = traj_0 @ rotation_matrix.T + R_ = rotation_matrix(theta) + traj_0 = traj_0 @ R_.T return traj_0 traj_0 = generate_data(time, tau_0, freq_0, theta_0) +traj_0_proj = traj_0 @ rotation_matrix(theta_0)[:, 0] # plot the observed signal components and their sum plt.figure(figsize=(10, 4)) -plt.plot(time, traj_0, label="base trajectory", linewidth=2) +plt.plot(time, traj_0_proj, label="projected trajectory", linewidth=2) plt.xlabel("time") plt.ylabel("amplitude") plt.legend() plt.title(r"Observed scalar signal along $\vec{e}(\theta)$") plt.show() - # %% # 2. Interpret the signal as coming from a continuous linear dynamical system # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -274,12 +276,8 @@ def augment(traj, window_length=2): # Processing Systems, 35, pp.4017-4031. -def estimator(X, Y, rank=4): - # X: (n_samples, n_features) - # Y: (n_samples, n_features) - - # estimate operator - cxx = X.T @ X +def estimator(X, Y, rank=4, eps=1e-8): + cxx = X.T @ X + eps * np.eye(X.shape[1]) U, S, Vt = np.linalg.svd(cxx) S_inv = np.divide(1, S, out=np.zeros_like(S), where=S != 0) cxx_inv_half = Vt.T @ np.diag(np.sqrt(S_inv)) @ U.T @@ -416,6 +414,24 @@ def estimator(X, Y, rank=4): # spectral atoms, taking into account both the location of eigenvalues and the # relative geometry of their eigenspaces. +# %% +# A wider delay window for the SGOT experiments below +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# The window of length 4 used above is enough to identify a single reference +# operator, but the experiments below also probe signals whose two modes +# nearly coincide in frequency (e.g. :math:`\omega_2'\to\omega_1`). Telling +# such near-degenerate modes apart requires the delay embedding to span +# enough time to "see" their differing decay, so we re-embed the reference +# signal with a longer window before running the sweeps. + +sgot_window = 10 +Z = augment(traj_0, sgot_window) +_, B_0_spec_sgot = estimator(Z[:-1], Z[1:]) +D_0_sgot = np.log(B_0_spec_sgot["eig_val"]) * fs +L_0_sgot = B_0_spec_sgot["eig_vec_left"] +R_0_sgot = B_0_spec_sgot["eig_vec_right"] + # %% # SGOT distance versus rotation angle # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -434,52 +450,66 @@ def estimator(X, Y, rank=4): # this experiment isolates the effect of rotating the underlying one-dimensional # subspace in the observation plane. -thetas = np.linspace(0, np.pi / 2, 50) -lst = [] -for i, theta in enumerate(thetas): - traj = generate_data(time, tau_0, freq_0, theta) - Z = augment(traj, 4) - X = Z[:-1] - Y = Z[1:] - B, B_spec = estimator(X, Y, rank=4) - D, R, L = B_spec["eig_val"], B_spec["eig_vec_right"], B_spec["eig_vec_left"] - D = np.log(D) * fs - lst.append(sgot_metric(D_0, R_0, L_0, D, R, L, eta=0.01)) - -plt.figure(figsize=(8, 5)) -plt.plot(thetas, lst) -plt.xlabel("theta") -plt.ylabel("SGOT distance") -plt.title("SGOT distance vs rotation angle") +thetas = np.linspace(0, np.pi / 2, 51) +rotation_scores = [] + +for theta in thetas: + Z = augment(generate_data(time, tau_0, freq_0, theta), sgot_window) + B, B_spec = estimator(Z[:-1], Z[1:]) + D = np.log(B_spec["eig_val"]) * fs + L = B_spec["eig_vec_left"] + R = B_spec["eig_vec_right"] + rotation_scores.append( + sgot_metric( + D_0_sgot, R_0_sgot, L_0_sgot, D, R, L, eta=0.9, grassmann_metric="chordal" + ) + ) + +fig, ax = plt.subplots(figsize=(7, 4)) +ax.plot(thetas, rotation_scores, linewidth=1.8) +ax.axvline(theta_0, color="gray", linestyle="--", linewidth=0.8) +ax.set_xlabel(r"Rotation angle $\theta$ (rad)") +ax.set_ylabel(r"$d_S$") +ax.set_title("SGOT distance vs. rotation angle") +fig.tight_layout() plt.show() # %% # Comparison across Grassmannian metrics for SGOT distance versus rotation angle # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -thetas = np.linspace(0, np.pi / 2, 50) -lst = [] -for i, theta in enumerate(thetas): - traj = generate_data(time, tau_0, freq_0, theta) - Z = augment(traj, 4) - X = Z[:-1] - Y = Z[1:] - B, B_spec = estimator(X, Y, rank=4) - D, R, L = B_spec["eig_val"], B_spec["eig_vec_right"], B_spec["eig_vec_left"] - D = np.log(D) * fs - lst1 = [] - for name in ["chordal", "martin", "geodesic", "procrustes"]: - lst1.append(sgot_metric(D_0, R_0, L_0, D, R, L, eta=0.9, grassmann_metric=name)) - lst.append(lst1) -lst2 = np.array(lst) -plt.figure(figsize=(8, 5)) -for i, name in enumerate(["chordal", "martin", "geodesic", "procrustes"]): - plt.plot(thetas, lst2[:, i], label=name) - -plt.xlabel("theta") -plt.ylabel("SGOT distance") -plt.title("SGOT distance vs rotation angle") -plt.legend() +metrics = ["chordal", "geodesic", "procrustes", "martin"] +styles = {"chordal": "-", "geodesic": "--", "procrustes": "-.", "martin": ":"} +rotation_results = {m: [] for m in metrics} + +for theta in thetas: + Z = augment(generate_data(time, tau_0, freq_0, theta), sgot_window) + B, B_spec = estimator(Z[:-1], Z[1:]) + D = np.log(B_spec["eig_val"]) * fs + L = B_spec["eig_vec_left"] + R = B_spec["eig_vec_right"] + for m in metrics: + rotation_results[m].append( + sgot_metric( + D_0_sgot, R_0_sgot, L_0_sgot, D, R, L, eta=0.9, grassmann_metric=m + ) + ) + +fig, ax = plt.subplots(figsize=(7, 4)) +for m in metrics: + ax.plot(thetas, rotation_results[m], styles[m], label=m, linewidth=1.8) +ax.axvline( + theta_0, + color="gray", + linestyle="--", + linewidth=0.8, + label=r"$\theta_0 = \pi/4$ (reference)", +) +ax.set_xlabel(r"Rotation angle $\theta$ (rad)") +ax.set_ylabel(r"$d_S$") +ax.set_title("SGOT distance vs. rotation angle across Grassmannian metrics") +ax.legend() +fig.tight_layout() plt.show() # %% @@ -501,38 +531,38 @@ def estimator(X, Y, rank=4): # distance changes as a function of the perturbed frequency :math:`\omega_2'`. omegas = np.linspace(0.5, 3.0, 21) -methods = ["chordal", "martin", "geodesic", "procrustes"] -scores_omega = [] -theta = theta_0 +frequency_scores = {m: [] for m in metrics} -eta_fixed = 0.9 for omega in omegas: - freq_1 = np.array([freq_0[0], omega]) - traj = generate_data(time, tau_0, freq_1, theta) - Z = augment(traj, 4) - X = Z[:-1] - Y = Z[1:] - - B, B_spec = estimator(X, Y, rank=4) - D, R, L = B_spec["eig_val"], B_spec["eig_vec_right"], B_spec["eig_vec_left"] - D = np.log(D) * fs - - row = [] - for name in methods: - row.append( - sgot_metric(D_0, R_0, L_0, D, R, L, eta=eta_fixed, grassmann_metric=name) + Z = augment( + generate_data(time, tau_0, np.array([freq_0[0], omega]), theta_0), sgot_window + ) + B, B_spec = estimator(Z[:-1], Z[1:]) + D = np.log(B_spec["eig_val"]) * fs + L = B_spec["eig_vec_left"] + R = B_spec["eig_vec_right"] + for m in metrics: + frequency_scores[m].append( + sgot_metric( + D_0_sgot, R_0_sgot, L_0_sgot, D, R, L, eta=0.9, grassmann_metric=m + ) ) - scores_omega.append(row) - -scores_omega = np.array(scores_omega) -plt.figure(figsize=(8, 5)) -for i, name in enumerate(methods): - plt.plot(omegas, scores_omega[:, i], label=name) -plt.xlabel("omega") -plt.ylabel("SGOT distance") -plt.title("SGOT distance vs omega") -plt.legend() +fig, ax = plt.subplots(figsize=(7, 4)) +for m in metrics: + ax.plot(omegas, frequency_scores[m], styles[m], label=m, linewidth=1.8) +ax.axvline( + freq_0[1], + color="gray", + linestyle="--", + linewidth=0.8, + label=r"$\omega_2 = 2.0$ Hz (reference)", +) +ax.set_xlabel(r"Frequency $\omega_2'$ (Hz)") +ax.set_ylabel(r"$d_S$") +ax.set_title("SGOT distance vs. frequency across Grassmannian metrics") +ax.legend() +fig.tight_layout() plt.show() # %% @@ -553,47 +583,28 @@ def estimator(X, Y, rank=4): # In this way, both modes share the same modified decay parameter # :math:`\tau`, allowing us to isolate the influence of dissipation on the SGOT # distance. -decays = np.linspace(0.1, 3.0, 20) # adjust range as needed -methods = ["chordal", "martin", "geodesic", "procrustes"] -scores_decay = [] -theta = theta_0 - -for tau in decays: - freq_1 = np.array([freq_0[0], recovered_freqs[1]]) - tau_1 = np.array([tau, tau]) # or whatever structure your generator expects - - traj = generate_data(time, tau_1, freq_1, theta) - Z = augment(traj, 4) - X = Z[:-1] - Y = Z[1:] - - B, B_spec = estimator(X, Y, rank=4) - D, R, L = B_spec["eig_val"], B_spec["eig_vec_right"], B_spec["eig_vec_left"] - D = np.log(D) * fs - - row = [] - for name in methods: - row.append( +taus = np.linspace(0.1, 3.0, 21) +decay_scores = {m: [] for m in metrics} + +for tau in taus: + Z = augment(generate_data(time, np.array([tau, tau]), freq_0, theta_0), sgot_window) + B, B_spec = estimator(Z[:-1], Z[1:]) + D = np.log(B_spec["eig_val"]) * fs + L = B_spec["eig_vec_left"] + R = B_spec["eig_vec_right"] + for m in metrics: + decay_scores[m].append( sgot_metric( - D_0, - R_0, - L_0, - D, - R, - L, - eta=0.9, # keep eta fixed here - grassmann_metric=name, + D_0_sgot, R_0_sgot, L_0_sgot, D, R, L, eta=0.9, grassmann_metric=m ) ) - scores_decay.append(row) -scores_decay = np.array(scores_decay) -plt.figure(figsize=(8, 5)) -for i, name in enumerate(methods): - plt.plot(decays, scores_decay[:, i], label=name) - -plt.xlabel("decay") -plt.ylabel("SGOT distance") -plt.title("SGOT distance vs decay") -plt.legend() +fig, ax = plt.subplots(figsize=(7, 4)) +for m in metrics: + ax.plot(taus, decay_scores[m], styles[m], label=m, linewidth=1.8) +ax.set_xlabel(r"Decay rate $\tau$") +ax.set_ylabel(r"$d_S$") +ax.set_title("SGOT distance vs. decay across Grassmannian metrics") +ax.legend() +fig.tight_layout() plt.show() diff --git a/examples/plot_sgot.py b/examples/plot_sgot.py deleted file mode 100644 index 87f7ff3c8..000000000 --- a/examples/plot_sgot.py +++ /dev/null @@ -1,608 +0,0 @@ -#!/usr/bin/env python -# coding: utf-8 - -""" -========================================= -Spectral-Grassmann OT on dynamical systems operators -========================================= - -This example presents a synthetic example of Spectral Grassmannian-Wasserstein -Optimal Transport (SGOT) on linear dynamical systems. - -We consider a signal formed by the sum of two damped oscillatory modes evolving -along a rotated direction in the plane. The signal is then associated with an -underlying continuous linear dynamical system, and we study how its spectral -representation varies under rotation. The SGOT cost and metric are used to -compare the reference and rotated systems. - -.. [83] T. Germain; R. Flamary; V. R. Kostic; K. Lounici, A Spectral-Grassmann - Wasserstein Metric for Operator Representations of Dynamical Systems, - arXiv preprint arXiv:2509.24920, 2025. -""" - -# Authors: Sienna O'Shea -# Thibaut Germain -# -# License: MIT License - -import numpy as np -import matplotlib.pyplot as plt - -from ot.sgot import sgot_metric, sgot_cost_matrix - -from scipy.linalg import eig - - -# sampling parameters and time grid -fs = 50 -max_t = 5 -time = np.linspace(0, max_t, fs * max_t) -dt = 1 / fs - - -# %% -# Example: rotating a linear dynamical system in 3D -# ------------------------------------------------- -# -# 1. Build a simple observed signal -# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -# -# We begin by assuming that the observed signal is made of two oscillatory -# components: -# -# .. math:: -# -# x_{\text{ref}}(t)=e^{-\tau_1 t}\cos(2\pi\omega_1 t)\,\vec e(\theta) -# \;+\; -# e^{-\tau_2 t}\cos(2\pi\omega_2 t)\,\vec e(\theta), -# -# where :math:`\vec e(\theta)\in\mathbb{R}^2` is a fixed real vector. Thus, -# :math:`x(t)` evolves along the one-dimensional subspace spanned by -# :math:`\vec e(\theta)`, while its time dependence exhibits oscillatory and -# dissipative behaviour. - -tau_0 = np.array([0.08, 0.18]) -freq_0 = np.array([1.0, 2.0]) -theta_0 = np.pi / 4 - - -def generate_data(time, tau, freq, theta): - t_ = np.sin(2 * np.pi * freq[None, :] * time[:, None]) * np.exp( - -tau[None, :] * time[:, None] - ) - t_ = t_.sum(axis=1) - traj_0 = np.zeros((t_.shape[0], 2)) - traj_0[:, 0] = t_ - rotation_matrix = np.array( - [[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]] - ) - traj_0 = traj_0 @ rotation_matrix.T - return traj_0 - - -traj_0 = generate_data(time, tau_0, freq_0, theta_0) - - -# plot the observed signal components and their sum -plt.figure(figsize=(10, 4)) -plt.plot(time, traj_0, label="base trajectory", linewidth=2) -plt.xlabel("time") -plt.ylabel("amplitude") -plt.legend() -plt.title(r"Observed scalar signal along $\vec{e}(\theta)$") -plt.show() - - -# %% -# 2. Interpret the signal as coming from a continuous linear dynamical system -# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -# -# We assume that :math:`x(t)` is generated by an underlying continuous linear -# dynamical system. Since the observed signal is a superposition of two -# sinusoidal modes, the corresponding linear dynamics are naturally described -# by a fourth-order model. We therefore introduce the state vector -# -# .. math:: -# -# z(t)= -# \begin{pmatrix} -# x_1(t)\\ -# x_2(t)\\ -# \vdots\\ -# x_1^{(3)}(t)\\ -# x_2^{(3)}(t) -# \end{pmatrix} -# \in\mathbb{R}^8. -# -# where :math:`x^{(n)}(t)` denotes the n-th derivative of :math:`x(t)`. -# -# This allows us to rewrite the dynamics as a first-order linear system: -# -# .. math:: -# -# \dot{z}(t)=Az(t), -# -# where :math:`A\in\mathbb{R}^{8\times 8}`. Its solution is then given by -# -# .. math:: -# -# z(t)=e^{tA}z_0. - -fig = plt.figure(figsize=(9, 6)) -ax = fig.add_subplot(projection="3d") - -ax.plot(time, traj_0[:, 0], traj_0[:, 1]) -ax.set_xlabel("time") -ax.set_ylabel("x₁(t)") -ax.set_title("Observed trajectory in time") - -ax.text2D(1.08, 0.5, "x₂(t)", transform=ax.transAxes, rotation=90, va="center") - -plt.show() - - -# %% -# 3. Sampling and preprocessing discrete trajectories of the dynamical system -# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -# -# We now have a bridge between the continuous system and the operator we later -# aim to infer from sampled data. Since in practice we do not observe the full -# continuous trajectory, we work instead with discrete samples of the signal. -# We take snapshots at uniform time intervals :math:`\Delta t`, and write the -# sampled signal as -# -# .. math:: -# -# S= -# \begin{pmatrix} -# x_1(0) & x_2(0)\\ -# x_1(\Delta t) & x_2(\Delta t)\\ -# \vdots\\ -# x_1(N\Delta t) & x_2(N\Delta t)\\ -# \end{pmatrix} -# -# The goal is now to use these observations to recover the operator governing -# the evolution. To do this, we augment the signal :math:`s` using a sliding -# window of length :math:`w`. For each :math:`k`, define -# -# .. math:: -# -# z_k = -# \begin{pmatrix} -# s_k\\ -# s_{k+1}\\ -# \vdots\\ -# s_{k+w-1} -# \end{pmatrix} -# -# We then form the data matrices -# -# .. math:: -# -# X= -# \begin{pmatrix} -# z_1\\ -# z_2\\ -# \vdots\\ -# z_{N-w} -# \end{pmatrix}, -# \qquad -# Y= -# \begin{pmatrix} -# z_2\\ -# z_3\\ -# \vdots\\ -# z_{N-w+1} -# \end{pmatrix}, -# -# so that :math:`X` contains the present windowed states and :math:`Y` the -# corresponding shifted future states. - - -# build a 4-dimensional state using delay embedding -def augment(traj, window_length=2): - Z = np.lib.stride_tricks.sliding_window_view(traj, (window_length, 1)) - Z = Z.reshape(Z.shape[0], -1) - return Z - - -# create the embedded state matrix Z -Z = augment(traj_0, 4) -Z.shape - -# inspect one embedded state vector -Z[0] - -# create X and Y for the SGOT metric -X = Z[:-1] -Y = Z[1:] - -# inspect shapes of X and Y -print("X shape:", X.shape) -print("Y shape:", Y.shape) - - -# %% -# 4. Estimate the discrete-time operator -# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -# -# We now identify the operator that maps :math:`X` to :math:`Y`. From -# -# .. math:: -# -# \dot{z}=Az, -# -# we have -# -# .. math:: -# -# z(t+\Delta t)=e^{\Delta tA}z(t). -# -# Setting -# -# .. math:: -# -# B=e^{\Delta tA}, -# -# the corresponding discrete-time evolution is governed by :math:`B`, and we -# seek the best linear map satisfying -# -# .. math:: -# -# Y\approx X B^T. -# -# Equivalently, we solve the optimisation problem -# -# .. math:: -# -# \min_B \|Y-XB\|^2. -# -# We want to recover the best rank-:math:`r` operator, whose estimator is -# defined as follows: -# -# .. math:: -# -# B = C_{xx}^{-\frac{1}{2}}[C_{xx}^{-\frac{1}{2}}C_{xy}]_r -# \quad \text{s.t} \quad C_{xx} = X^T X \quad \text{and} \quad C_{xy} = X^TY. -# -# Here :math:`[\cdot]_r` denotes the best rank-:math:`r` estimator obtained via -# SVD decomposition. [2] -# -# [2] Kostic, V., Novelli, P., Maurer, A., Ciliberto, C., Rosasco, L. and -# Pontil, M., 2022. Learning dynamical systems via Koopman operator regression -# in reproducing kernel Hilbert spaces. Advances in Neural Information -# Processing Systems, 35, pp.4017-4031. - - -def estimator(X, Y, rank=4, eps=1e-8): - cxx = X.T @ X + eps * np.eye(X.shape[1]) - U, S, Vt = np.linalg.svd(cxx) - S_inv = np.divide(1, S, out=np.zeros_like(S), where=S != 0) - cxx_inv_half = Vt.T @ np.diag(np.sqrt(S_inv)) @ U.T - cxy = X.T @ Y - T = cxx_inv_half @ cxy - U, S, Vt = np.linalg.svd(T) - S[rank:] = 0 - T_rank = U @ np.diag(S) @ Vt - T = cxx_inv_half @ T_rank - - # estimate spectral decomposition - val, vl, vr = eig(T, left=True, right=True) - sort_idx = np.argsort(np.abs(val))[::-1] - val = val[sort_idx][:rank] - vl = vl[:, sort_idx][:, :rank] - vr = vr[:, sort_idx][:, :rank] - - return T, {"eig_val": val, "eig_vec_left": vl, "eig_vec_right": vr} - - -B_0, B_0_spec = estimator(X, Y, rank=4) -Y_pred = X @ B_0 - -plt.figure(figsize=(10, 4)) -plt.plot(Y[:, 0], label="true") -plt.plot(Y_pred[:, 0], "--", label="predicted") -plt.xlabel("sample index") -plt.ylabel("first state coordinate") -plt.title("True Signal vs Predicted Signal") -plt.legend() -plt.show() - - -# %% -# The predicted signal is nearly indistinguishable from the true signal, -# indicating that the estimated operator accurately captures the observed -# dynamics. - -# %% -# 6. Recover continuous-time spectral information from the discrete operator -# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -# -# To recover the continuous generator :math:`A`, we study the spectral -# structure of :math:`B`. We diagonalise :math:`B` as -# -# .. math:: -# -# B=PDP^{-1}, -# -# where -# -# .. math:: -# -# D=\operatorname{diag}(\mu_1,\dots,\mu_n). -# -# The continuous-time eigenvalues are of the form -# -# .. math:: -# -# \lambda_k=-\tau_k+2\pi i\,\omega_k, -# \qquad k\in\{1,2\}, -# -# and the corresponding eigenvalues of :math:`B` are -# -# .. math:: -# -# \mu_k=e^{\Delta t\lambda_k} -# =e^{\Delta t(-\tau_k+2\pi i\omega_k)}. -# -# Since :math:`B=e^{\Delta tA}`, we recover :math:`A` by taking the logarithm: -# -# .. math:: -# -# A=P\,\frac{\log(D)}{\Delta t}\,P^{-1}. - -D_0 = np.log(B_0_spec["eig_val"]) * fs -L_0 = B_0_spec["eig_vec_left"] -R_0 = B_0_spec["eig_vec_right"] - -recovered_freqs = D_0.imag / (2 * np.pi) -mask = recovered_freqs > 0 -recovered_freqs = recovered_freqs[mask] -decay = -D_0.real[mask] -print(f"First mode: frequency: {recovered_freqs[0]:.2f} Hz -- decay: {decay[0]:.2f}") -print(f"Second mode: frequency: {recovered_freqs[1]:.2f} Hz -- decay: {decay[1]:.2f}") - - -# %% -# Introduction to SGOT for linear operators -# ----------------------------------------- -# -# To compare two linear operators through their spectral structure, we use the -# SGOT framework introduced in Theorem 1 of [1]. For a non-defective -# finite-rank operator :math:`T \in S_r(\mathcal H)`, the theorem associates a -# discrete spectral measure -# -# .. math:: -# -# \mu(T) \triangleq \sum_{j\in[\ell]} -# \frac{m_j}{m_{\mathrm{tot}}}\,\delta_{(\lambda_j,\mathcal V_j)}, -# -# where :math:`\lambda_j` are the eigenvalues of :math:`T`, :math:`m_j` their -# algebraic multiplicities, and :math:`\mathcal V_j` the corresponding -# eigenspaces. Thus, each spectral component of the operator is represented by -# an atom of the form -# -# .. math:: -# -# (\lambda_j,\mathcal V_j), -# -# combining one eigenvalue with its associated invariant subspace. -# -# Theorem 1 then defines a ground cost between two such atoms by combining a -# spectral discrepancy and a geometric discrepancy: -# -# .. math:: -# -# d_\eta\big((\lambda,\mathcal V),(\lambda',\mathcal V')\big) -# \triangleq -# \eta\,|\lambda-\lambda'| + (1-\eta)\, d_{\mathcal G}(\mathcal V,\mathcal V'), -# -# where :math:`d_{\mathcal G}` denotes the grassmann distance between -# eigenspaces and :math:`\eta\in(0,1)` balances the contribution of eigenvalues -# and eigenspaces. -# -# The SGOT distance between two operators :math:`T` and :math:`T'` is then the -# Wasserstein distance between their associated spectral measures: -# -# .. math:: -# -# d_S(T,T') = W_{d_\eta,p}\big(\mu(T),\mu(T')\big). -# -# In this way, SGOT compares linear operators by optimally matching their -# spectral atoms, taking into account both the location of eigenvalues and the -# relative geometry of their eigenspaces. - -# %% -# A wider delay window for the SGOT experiments below -# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -# -# The window of length 4 used above is enough to identify a single reference -# operator, but the experiments below also probe signals whose two modes -# nearly coincide in frequency (e.g. :math:`\omega_2'\to\omega_1`). Telling -# such near-degenerate modes apart requires the delay embedding to span -# enough time to "see" their differing decay, so we re-embed the reference -# signal with a longer window before running the sweeps. - -sgot_window = 10 -Z = augment(traj_0, sgot_window) -_, B_0_spec_sgot = estimator(Z[:-1], Z[1:]) -D_0_sgot = np.log(B_0_spec_sgot["eig_val"]) * fs -L_0_sgot = B_0_spec_sgot["eig_vec_left"] -R_0_sgot = B_0_spec_sgot["eig_vec_right"] - -# %% -# SGOT distance versus rotation angle -# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -# -# We compare the reference signal with a rotated version obtained by changing -# only the observation direction. The shifted signal is -# -# .. math:: -# -# x_{\mathrm{shift}}^{\mathrm{rot}}(t;\theta) -# = -# \sum_{i=1}^{2} -# e^{-\tau_i t}\cos(2\pi\omega_i t)\,\vec e({\color{red}\theta}), -# -# while the reference signal is recovered at :math:`\theta=\theta_0`. Thus, -# this experiment isolates the effect of rotating the underlying one-dimensional -# subspace in the observation plane. - -thetas = np.linspace(0, np.pi / 2, 51) -rotation_scores = [] - -for theta in thetas: - Z = augment(generate_data(time, tau_0, freq_0, theta), sgot_window) - B, B_spec = estimator(Z[:-1], Z[1:]) - D = np.log(B_spec["eig_val"]) * fs - L = B_spec["eig_vec_left"] - R = B_spec["eig_vec_right"] - rotation_scores.append( - sgot_metric( - D_0_sgot, R_0_sgot, L_0_sgot, D, R, L, eta=0.9, grassmann_metric="chordal" - ) - ) - -fig, ax = plt.subplots(figsize=(7, 4)) -ax.plot(thetas, rotation_scores, linewidth=1.8) -ax.axvline(theta_0, color="gray", linestyle="--", linewidth=0.8) -ax.set_xlabel(r"Rotation angle $\theta$ (rad)") -ax.set_ylabel(r"$d_S$") -ax.set_title("SGOT distance vs. rotation angle") -fig.tight_layout() -plt.show() - -# %% -# Comparison across Grassmannian metrics for SGOT distance versus rotation angle -# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -metrics = ["chordal", "geodesic", "procrustes", "martin"] -styles = {"chordal": "-", "geodesic": "--", "procrustes": "-.", "martin": ":"} -rotation_results = {m: [] for m in metrics} - -for theta in thetas: - Z = augment(generate_data(time, tau_0, freq_0, theta), sgot_window) - B, B_spec = estimator(Z[:-1], Z[1:]) - D = np.log(B_spec["eig_val"]) * fs - L = B_spec["eig_vec_left"] - R = B_spec["eig_vec_right"] - for m in metrics: - rotation_results[m].append( - sgot_metric( - D_0_sgot, R_0_sgot, L_0_sgot, D, R, L, eta=0.9, grassmann_metric=m - ) - ) - -fig, ax = plt.subplots(figsize=(7, 4)) -for m in metrics: - ax.plot(thetas, rotation_results[m], styles[m], label=m, linewidth=1.8) -ax.axvline( - theta_0, - color="gray", - linestyle="--", - linewidth=0.8, - label=r"$\theta_0 = \pi/4$ (reference)", -) -ax.set_xlabel(r"Rotation angle $\theta$ (rad)") -ax.set_ylabel(r"$d_S$") -ax.set_title("SGOT distance vs. rotation angle across Grassmannian metrics") -ax.legend() -fig.tight_layout() -plt.show() - -# %% -# SGOT distance versus frequency -# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -# -# In this experiment, we keep the reference direction fixed and perturb one of -# the oscillatory modes. The shifted signal is -# -# .. math:: -# -# x_{\mathrm{shift}}^{\omega}(t) -# = -# e^{-\tau_1 t}\cos(2\pi\omega_1 t)\,\vec e(\theta_0) -# \;+\; -# e^{-\tau_2 t}\cos(2\pi{\color{red}\omega_2'} t)\,\vec e(\theta_0), -# -# where only the second frequency is modified. We then study how the SGOT -# distance changes as a function of the perturbed frequency :math:`\omega_2'`. - -omegas = np.linspace(0.5, 3.0, 21) -frequency_scores = {m: [] for m in metrics} - -for omega in omegas: - Z = augment( - generate_data(time, tau_0, np.array([freq_0[0], omega]), theta_0), sgot_window - ) - B, B_spec = estimator(Z[:-1], Z[1:]) - D = np.log(B_spec["eig_val"]) * fs - L = B_spec["eig_vec_left"] - R = B_spec["eig_vec_right"] - for m in metrics: - frequency_scores[m].append( - sgot_metric( - D_0_sgot, R_0_sgot, L_0_sgot, D, R, L, eta=0.9, grassmann_metric=m - ) - ) - -fig, ax = plt.subplots(figsize=(7, 4)) -for m in metrics: - ax.plot(omegas, frequency_scores[m], styles[m], label=m, linewidth=1.8) -ax.axvline( - freq_0[1], - color="gray", - linestyle="--", - linewidth=0.8, - label=r"$\omega_2 = 2.0$ Hz (reference)", -) -ax.set_xlabel(r"Frequency $\omega_2'$ (Hz)") -ax.set_ylabel(r"$d_S$") -ax.set_title("SGOT distance vs. frequency across Grassmannian metrics") -ax.legend() -fig.tight_layout() -plt.show() - -# %% -# SGOT distance versus decay -# ~~~~~~~~~~~~~~~~~~~~~~~~~~ -# -# We now study the effect of changing the decay rate while keeping the -# observation direction fixed. The shifted signal is -# -# .. math:: -# -# x_{\mathrm{shift}}^{\mathrm{decay}}(t;\tau) -# = -# e^{-{\color{red}\tau} t}\cos(2\pi\omega_1 t)\,\vec e(\theta_0) -# \;+\; -# e^{-{\color{red}\tau} t}\cos(2\pi\omega_2' t)\,\vec e(\theta_0). -# -# In this way, both modes share the same modified decay parameter -# :math:`\tau`, allowing us to isolate the influence of dissipation on the SGOT -# distance. -taus = np.linspace(0.1, 3.0, 21) -decay_scores = {m: [] for m in metrics} - -for tau in taus: - Z = augment(generate_data(time, np.array([tau, tau]), freq_0, theta_0), sgot_window) - B, B_spec = estimator(Z[:-1], Z[1:]) - D = np.log(B_spec["eig_val"]) * fs - L = B_spec["eig_vec_left"] - R = B_spec["eig_vec_right"] - for m in metrics: - decay_scores[m].append( - sgot_metric( - D_0_sgot, R_0_sgot, L_0_sgot, D, R, L, eta=0.9, grassmann_metric=m - ) - ) - -fig, ax = plt.subplots(figsize=(7, 4)) -for m in metrics: - ax.plot(taus, decay_scores[m], styles[m], label=m, linewidth=1.8) -ax.set_xlabel(r"Decay rate $\tau$") -ax.set_ylabel(r"$d_S$") -ax.set_title("SGOT distance vs. decay across Grassmannian metrics") -ax.legend() -fig.tight_layout() -plt.show() From 267674ddf64d09abf3b886e640158bf3bc006ff7 Mon Sep 17 00:00:00 2001 From: thibaut-germain Date: Thu, 2 Jul 2026 13:34:37 +0200 Subject: [PATCH 05/13] fix releases.md --- RELEASES.md | 1 + 1 file changed, 1 insertion(+) diff --git a/RELEASES.md b/RELEASES.md index f10d7b62b..eea94981d 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -38,6 +38,7 @@ This new release adds support for sparse cost matrices and a new lazy exact OT s - Update the geomloss wrapper to the new version and API (PR #826) - Fix docstrings for `lowrank_gromov_wasserstein_samples` and `lowrank_sinkhorn` (PR #823) - Reorganize all tests per backend (PR #828) +- Update sgot cost function and example (PR #830) #### Closed issues From bed3cf1d77b9f4e4fb2eb27f3e2f3d86684dc60b Mon Sep 17 00:00:00 2001 From: thibaut-germain Date: Thu, 2 Jul 2026 16:47:07 +0200 Subject: [PATCH 06/13] reformating --- ot/batch/_utils.py | 138 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 138 insertions(+) diff --git a/ot/batch/_utils.py b/ot/batch/_utils.py index ce5db4664..20c9b900c 100644 --- a/ot/batch/_utils.py +++ b/ot/batch/_utils.py @@ -288,3 +288,141 @@ def bregman_log_projection_batch( log_potentials = (u, v) return {"T": T, "log_T": log_T, "n_iters": n_iters, "potentials": log_potentials} + + +def proximal_bregman_log_plan_batch( + C, + a=None, + b=None, + nx=None, + reg=1e-1, + max_iter=10000, + tol=1e-5, + inner_iter=1, + grad="detach", +): + r""" + Returns optimal transport plans for a batch of cost matrices :math:`\mathbf{C}` using a Bregman divergence based proximal point method :ref:`[1] `. + + This function solves the following optimization problem: + + .. math:: + \begin{aligned} + \mathbf{T} = \mathop{\arg \min}_\mathbf{T} \quad & \langle \mathbf{T}, \mathbf{C} \rangle_F \\ + \text{s.t.} \quad & \mathbf{T} \mathbf{1} = \mathbf{a} \\ + & \mathbf{T}^T \mathbf{1} = \mathbf{b} \\ + & \mathbf{T} \geq 0 + \end{aligned} + + The optimal transport plans are computed iteratively with a proximal point method based on a Bregman divergence where each iteration involves solving a Bregman projection problem: + + .. math:: + \mathbf{T}^{(k+1)} = \mathop{\arg \min}_\mathbf{T} \quad \langle \mathbf{C} - \textit{reg} * \log \mathbf{T}^{(k)}, \mathbf{T} \rangle + \textit{reg} * \sum_{i,j} \mathbf{T}_{i,j} \log \mathbf{T}_{i,j} + + Denoting :math:`\mathbf{K}^{(k)} = - \mathbf{C} / \textit{reg} + \log \mathbf{T}^{(k)}`, the affinity matrix at iteration :math:`k`, the Bregman projection problem is solved in the log-domain with a finite number of inner iterations :math:`\text{inner\_iter}`, i.e., the dual variables :math:`\mathbf{u}` and :math:`\mathbf{v}` are updated as follows: + + .. math:: + \mathbf{u}^{(i+1)} = \log(\mathbf{a}) - \text{LSE}(\mathbf{K}^{(k)} + \mathbf{v}^{(i)}) + + \mathbf{v}^{(i+1)} = \log(\mathbf{b}) - \text{LSE}((\mathbf{K}^{(k)})^T + \mathbf{u}^{(i)}) + + where LSE denotes the log-sum-exp operation. + + Parameters + ---------- + C : array-like, shape (B, n, m) + Cost matrix for each problem in the batch. + a : array-like, shape (B, n), optional + Source distribution for each problem. If None, uniform distribution is used. + b : array-like, shape (B, m), optional + Target distribution for each problem. If None, uniform distribution is used. + nx : backend object, optional + Numerical backend to use for computations. If None, the default backend is used. + max_iter : int, optional + Maximum number of iterations. + inner_iter : int, optional + Number of inner iterations for updating the dual variables u and v. + tol : float, optional + Tolerance for convergence. The solver stops when the maximum change in + the dual variables is below this value. + grad : str, optional + Gradient computation mode: 'detach', 'autodiff', or 'last_step'. + + Returns + ------- + dict + Dictionary containing: + - 'T' : array-like, shape (B, n, m) + Optimal transport plan for each problem. + - 'log_T' : array-like, shape (B, n, m) + Logarithm of the optimal transport plan for each problem. + - 'potentials' : tuple of array-like, shapes ((B, n), (B, m)) + Log-scaling factors (u, v). + - 'n_iters' : int + Number of iterations performed. + + Example + -------- + >>> import numpy as np + >>> from ot.batch import proximal_bregman_log_plan_batch + >>> # Create batch of affinity matrices + >>> C = np.random.rand(5, 10, 15) # 5 problems, 10x15 cost matrices + >>> result = proximal_bregman_log_plan_batch(C) + >>> T = result['T'] # Shape (5, 10, 15) + + .. _references-OT-bregman-proximal-point: + Reference + ---------- + .. [1] Xie, Y., Wang, X., Wang, R., & Zha, H. (2020, August). + A fast proximal point method for computing exact wasserstein distance. + In Uncertainty in artificial intelligence (pp. 433-453). PMLR. + """ + + if nx is None: + nx = get_backend(a, b, C) + + B, n, m = C.shape + + if a is None: + a = nx.ones((B, n)) / n + if b is None: + b = nx.ones((B, m)) / m + + u = nx.zeros((B, n), type_as=C) + v = nx.zeros((B, m), type_as=C) + + K = -C / reg + loga = nx.log(a) + logb = nx.log(b) + + if grad == "detach": + K = nx.detach(K) + elif grad == "last_step": + K_, K = K.clone(), nx.detach(K) + + log_T = nx.zeros(C.shape, type_as=C) + for n_iters in range(max_iter): + K_proj = log_T + K + for _ in range(inner_iter): + u = loga - nx.logsumexp(K_proj + v[:, None, :], axis=2) + v = logb - nx.logsumexp(K_proj + u[:, :, None], axis=1) + log_T = K_proj + u[:, :, None] + v[:, None, :] + + # Check convergence once every 10 iterations + if n_iters % 10 == 0: + T = nx.exp(log_T) + marginal = nx.sum(T, axis=2) + err = nx.max(norm_batch(marginal - a)) + if err < tol: + break + + if grad == "last_step": + K_proj = log_T + K_ + for _ in range(inner_iter): + u = loga - nx.logsumexp(K_proj + v[:, None, :], axis=2) + v = logb - nx.logsumexp(K_proj + u[:, :, None], axis=1) + log_T = K_proj + u[:, :, None] + v[:, None, :] + T = nx.exp(log_T) + log_potentials = (u, v) + + return {"T": T, "log_T": log_T, "n_iters": n_iters, "potentials": log_potentials} From 3f3157d50b34c9cfcd724557d1a1f6d63cc72586 Mon Sep 17 00:00:00 2001 From: thibaut-germain Date: Fri, 3 Jul 2026 11:47:24 +0200 Subject: [PATCH 07/13] updating format --- ot/batch/_linear.py | 162 ++++++++++++++++++++++++++++++-------------- ot/batch/_utils.py | 4 +- 2 files changed, 115 insertions(+), 51 deletions(-) diff --git a/ot/batch/_linear.py b/ot/batch/_linear.py index 1a9ec1955..f5596d747 100644 --- a/ot/batch/_linear.py +++ b/ot/batch/_linear.py @@ -14,6 +14,7 @@ bregman_log_projection_batch, bregman_projection_batch, entropy_batch, + proximal_bregman_log_plan_batch, ) @@ -235,23 +236,38 @@ def dist_batch( def solve_batch( M, - reg, + reg=None, a=None, b=None, max_iter=1000, tol=1e-5, - solver="log_sinkhorn", + solver="proximal", + inner_iter=1, + inner_reg=1e-2, reg_type="entropy", grad="envelope", ): - r"""Batched version of ot.solve, use it to solve many entropic OT problems in parallel. + r""" + Return solutions of a batch of discrete optimal transport problems in a :any:`OTResult` object. + + The function solves in parallel a batch of optimal transport problems: + + .. math:: + \begin{aligned} + \mathbf{T} = \mathop{\arg \min}_\mathbf{T} \quad & \langle \mathbf{T}, \mathbf{M} \rangle_F + \textit{reg} \cdot R(\mathbf{T}) \\ + \text{s.t.} \quad & \mathbf{T} \mathbf{1} = \mathbf{a} \\ + & \mathbf{T}^T \mathbf{1} = \mathbf{b} \\ + & \mathbf{T} \geq 0 + \end{aligned} + + The problem is solved with either a proximal point method :ref:`[1] ` or a Sinkhorn algorithm :ref:`[2] `. Unlike the Sinkhorn algorithm, which assumes a regularization term, the proximal point method aims to solve the exact optimal transport problem leading to sparser transport plans. The choice of solver depends on the value of `reg` and `solver`. By default, the proximal point method is used. The Sinkhorn algorithm is only used when `reg` is greater than 0 and `solver='sinkhorn' or `log_sinkhorn'`. In any other case, the proximal point method is used, i.e. whenever `reg` is None or 0 or `solver='proximal'`. Parameters ---------- M : array-like, shape (B, ns, nt) Cost matrix reg : float - Regularization parameter for entropic regularization + Regularization parameter for entropic regularization. If None or 0, no regularization is applied, proximal solver is used. Default is None. a : array-like, shape (B, ns) Source distribution (optional). If None, uniform distribution is used. b : array-like, shape (B, nt) @@ -261,17 +277,21 @@ def solve_batch( tol : float Tolerance for convergence solver: str - Solver to use, either 'log_sinkhorn' or 'sinkhorn'. Default is "log_sinkhorn" which is more stable. + Solver to use, either 'proximal', 'log_sinkhorn' or 'sinkhorn'. Default is 'proximal'. + inner_iter : int + Number of inner Bregman iterations for the proximal solver. Default is 1. + inner_reg : float + Regularization parameter for the inner Bregman iterations in the proximal solver. Default is 1e-2. reg_type : str, optional - Type of regularization :math:`R` either "KL", or "entropy". Default is "entropy". + Type of regularization :math:`R` either "KL", or "entropy". Only used with Sinkhorn solver. Default is "entropy". grad : str, optional - Type of gradient computation, either or 'autodiff', 'envelope' or 'last_step' used only for - Sinkhorn solver. By default 'autodiff' provides gradients wrt all - outputs (`plan, value, value_linear`) but with important memory cost. - 'envelope' provides gradients only for `value` and and other outputs are - detached. This is useful for memory saving when only the value is needed. 'last_step' provides - gradients only for the last iteration of the Sinkhorn solver, but provides gradient for both the OT plan and the objective values. - 'detach' does not compute the gradients for the Sinkhorn solver. + Type of gradient computation, either 'detach', 'autodiff', 'last_step' or 'envelope'. + 'detach' does not compute the gradients for the Sinkhorn solver. + 'autodiff' provides gradients of all outputs (`plan, value, value_linear`) but with important memory cost. + 'last_step' provides gradients of all outputs (`plan, value, value_linear`) only for the last solver iteration, useful for memory saving. + 'envelope' provides gradients only for `value`. + Default is 'envelope'. + Returns ------- @@ -292,19 +312,39 @@ def solve_batch( >>> X = np.random.randn(5, 10, 3) # 5 batches of 10 samples in 3D >>> Y = np.random.randn(5, 15, 3) # 5 batches of 15 samples in 3D >>> M = dist_batch(X, Y, metric="euclidean") # Compute cost matrices + >>> p_result = solve_batch(M) # Uses proximal solver >>> reg = 0.1 - >>> result = solve_batch(M, reg) - >>> result.plan.shape # Optimal transport plans for each batch + >>> s_result = solve_batch(M, reg, solver="log_sinkhorn") # Uses Sinkhorn solver + >>> s_result.plan.shape # Optimal transport plans for each batch (5, 10, 15) - >>> result.value.shape # Optimal transport values for each batch + >>> s_result.value.shape # Optimal transport values for each batch (5,) See Also -------- ot.batch.dist_batch : batched cost matrix computation for computing M. ot.solve : non-batched version of the OT solver. + + .. _references-batch-solver: + Reference + ---------- + .. [1] Xie, Y., Wang, X., Wang, R., & Zha, H. (2020, August). + A fast proximal point method for computing exact wasserstein distance. + In Uncertainty in artificial intelligence (pp. 433-453). PMLR. + + .. [2] M. Cuturi, Sinkhorn Distances : Lightspeed Computation + of Optimal Transport, Advances in Neural Information Processing + Systems (NIPS) 26, 2013 """ + assert solver in [ + "proximal", + "log_sinkhorn", + "sinkhorn", + ], f"Unknown solver: {solver}" + + use_sinkhorn = solver in ["log_sinkhorn", "sinkhorn"] and reg > 0 + nx = get_backend(a, b, M) B, n, m = M.shape @@ -314,18 +354,29 @@ def solve_batch( if b is None: b = nx.ones((B, m), type_as=M) / m - if solver == "log_sinkhorn": - K = -M / reg - out = bregman_log_projection_batch( - K, a, b, nx=nx, max_iter=max_iter, tol=tol, grad=grad - ) - elif solver == "sinkhorn": - K = nx.exp(-M / reg) - out = bregman_projection_batch( - K, a, b, nx=nx, max_iter=max_iter, tol=tol, grad=grad - ) + if use_sinkhorn: + if solver == "log_sinkhorn": + K = -M / reg + out = bregman_log_projection_batch( + K, a, b, nx=nx, max_iter=max_iter, tol=tol, grad=grad + ) + if solver == "sinkhorn": + K = nx.exp(-M / reg) + out = bregman_projection_batch( + K, a, b, nx=nx, max_iter=max_iter, tol=tol, grad=grad + ) else: - raise ValueError(f"Unknown solver: {solver}") + out = proximal_bregman_log_plan_batch( + M, + a, + b, + nx=nx, + reg=inner_reg, + max_iter=max_iter, + tol=tol, + inner_iter=inner_iter, + grad=grad, + ) T = out["T"] @@ -336,13 +387,14 @@ def solve_batch( T = nx.detach(T) value_linear = loss_linear_batch(M, T) - if reg_type.lower() == "entropy": - entr = -entropy_batch(T, nx=nx) - value = value_linear + reg * entr - elif reg_type.lower() == "kl": - ref = nx.einsum("bi,bj->bij", a, b) - kl = nx.sum(T * nx.log(T / ref + 1e-16), axis=(1, 2)) - value = value_linear + reg * kl + if use_sinkhorn: + if reg_type.lower() == "entropy": + entr = -entropy_batch(T, nx=nx) + value = value_linear + reg * entr + elif reg_type.lower() == "kl": + ref = nx.einsum("bi,bj->bij", a, b) + kl = nx.sum(T * nx.log(T / ref + 1e-16), axis=(1, 2)) + value = value_linear + reg * kl log = {"n_iter": out["n_iters"]} res = OTResult( @@ -360,29 +412,36 @@ def solve_batch( def solve_sample_batch( X_a, X_b, - reg, + reg=None, a=None, b=None, metric="sqeuclidean", p=2, max_iter=1000, tol=1e-5, - solver="log_sinkhorn", + solver="proximal", + inner_iter=1, + inner_reg=1e-2, reg_type="entropy", grad="envelope", ): - r"""Batched version of ot.solve, use it to solve many entropic OT problems in parallel. + r""" + Return solutions of a batch of discrete optimal transport problems in a :any:`OTResult` object computed for source and target samples. + + The problem is solved with either a proximal point method :ref:`[1] ` or a Sinkhorn algorithm :ref:`[2] `. Unlike the Sinkhorn algorithm, which assumes a regularization term, the proximal point method aims to solve the exact optimal transport problem leading to sparser transport plans. The choice of solver depends on the value of `reg` and `solver`. By default, the proximal point method is used. The Sinkhorn algorithm is only used when `reg` is greater than 0 and `solver='sinkhorn' or `log_sinkhorn'`. In any other case, the proximal point method is used, i.e. whenever `reg` is None or 0 or `solver='proximal'`. Parameters ---------- - M : array-like, shape (B, ns, nt) - Cost matrix - reg : float - Regularization parameter for entropic regularization + Xa : array-like, shape (B, ns, d) + Samples from source distribution + Xb : array-like, shape (B, nt, d) + Samples from target distribution metric : str, optional 'sqeuclidean', 'euclidean', 'minkowski' or 'kl' p : float, optional p-norm for the Minkowski metrics. Default value is 2. + reg : float + Regularization parameter for entropic regularization. If None or 0, no regularization is applied, proximal solver is used. Default is None. a : array-like, shape (B, ns) Source distribution (optional). If None, uniform distribution is used. b : array-like, shape (B, nt) @@ -392,17 +451,20 @@ def solve_sample_batch( tol : float Tolerance for convergence solver: str - Solver to use, either 'log_sinkhorn' or 'sinkhorn'. Default is "log_sinkhorn" which is more stable. + Solver to use, either 'proximal', 'log_sinkhorn' or 'sinkhorn'. Default is 'proximal'. + inner_iter : int + Number of inner Bregman iterations for the proximal solver. Default is 1. + inner_reg : float + Regularization parameter for the inner Bregman iterations in the proximal solver. Default is 1e-2. reg_type : str, optional - Type of regularization :math:`R` either "KL", or "entropy". Default is "entropy". + Type of regularization :math:`R` either "KL", or "entropy". Only used with Sinkhorn solver. Default is "entropy". grad : str, optional - Type of gradient computation, either or 'autodiff', 'envelope' or 'last_step' used only for - Sinkhorn solver. By default 'autodiff' provides gradients wrt all - outputs (`plan, value, value_linear`) but with important memory cost. - 'envelope' provides gradients only for `value` and and other outputs are - detached. This is useful for memory saving when only the value is needed. 'last_step' provides - gradients only for the last iteration of the Sinkhorn solver, but provides gradient for both the OT plan and the objective values. + Type of gradient computation, either 'detach', 'autodiff', 'last_step' or 'envelope'. 'detach' does not compute the gradients for the Sinkhorn solver. + 'autodiff' provides gradients of all outputs (`plan, value, value_linear`) but with important memory cost. + 'last_step' provides gradients of all outputs (`plan, value, value_linear`) only for the last solver iteration, useful for memory saving. + 'envelope' provides gradients only for `value`. + Default is 'envelope'. Returns ------- @@ -430,6 +492,8 @@ def solve_sample_batch( max_iter=max_iter, tol=tol, solver=solver, + inner_iter=inner_iter, + inner_reg=inner_reg, reg_type=reg_type, grad=grad, ) diff --git a/ot/batch/_utils.py b/ot/batch/_utils.py index 20c9b900c..18d15b9b4 100644 --- a/ot/batch/_utils.py +++ b/ot/batch/_utils.py @@ -295,7 +295,7 @@ def proximal_bregman_log_plan_batch( a=None, b=None, nx=None, - reg=1e-1, + reg=1e-2, max_iter=10000, tol=1e-5, inner_iter=1, @@ -317,7 +317,7 @@ def proximal_bregman_log_plan_batch( The optimal transport plans are computed iteratively with a proximal point method based on a Bregman divergence where each iteration involves solving a Bregman projection problem: .. math:: - \mathbf{T}^{(k+1)} = \mathop{\arg \min}_\mathbf{T} \quad \langle \mathbf{C} - \textit{reg} * \log \mathbf{T}^{(k)}, \mathbf{T} \rangle + \textit{reg} * \sum_{i,j} \mathbf{T}_{i,j} \log \mathbf{T}_{i,j} + \mathbf{T}^{(k+1)} = \mathop{\arg \min}_\mathbf{T} \quad \langle \mathbf{C} - \textit{reg} \cdot \log \mathbf{T}^{(k)}, \mathbf{T} \rangle + \textit{reg} \cdot \sum_{i,j} \mathbf{T}_{i,j} \log \mathbf{T}_{i,j} Denoting :math:`\mathbf{K}^{(k)} = - \mathbf{C} / \textit{reg} + \log \mathbf{T}^{(k)}`, the affinity matrix at iteration :math:`k`, the Bregman projection problem is solved in the log-domain with a finite number of inner iterations :math:`\text{inner\_iter}`, i.e., the dual variables :math:`\mathbf{u}` and :math:`\mathbf{v}` are updated as follows: From 9fcc185a2eb8d9ac9a01d0ac3a229b0c8b99caed Mon Sep 17 00:00:00 2001 From: thibaut-germain Date: Fri, 3 Jul 2026 13:41:34 +0200 Subject: [PATCH 08/13] updated format --- ot/batch/_linear.py | 1 + ot/batch/_utils.py | 11 ++++++ test/batch/test_solve_batch.py | 71 +++++++++++++++++++++++++++++----- 3 files changed, 74 insertions(+), 9 deletions(-) diff --git a/ot/batch/_linear.py b/ot/batch/_linear.py index f5596d747..450738366 100644 --- a/ot/batch/_linear.py +++ b/ot/batch/_linear.py @@ -5,6 +5,7 @@ # Author: Remi Flamary # Paul Krzakala +# Thibaut Germain # # License: MIT License diff --git a/ot/batch/_utils.py b/ot/batch/_utils.py index 18d15b9b4..984cafd4a 100644 --- a/ot/batch/_utils.py +++ b/ot/batch/_utils.py @@ -1,3 +1,14 @@ +# -*- coding: utf-8 -*- +""" +Utility functions for batch operations in optimal transport. +""" + +# Author: Remi Flamary +# Paul Krzakala +# Thibaut Germain +# +# License: MIT License + from ot.backend import get_backend diff --git a/test/batch/test_solve_batch.py b/test/batch/test_solve_batch.py index 17d459a43..900cff9d3 100644 --- a/test/batch/test_solve_batch.py +++ b/test/batch/test_solve_batch.py @@ -3,6 +3,7 @@ # Author: Remi Flamary # Paul Krzakala # Sonia Mazelet +# Thibaut Germain # # License: MIT License @@ -15,6 +16,7 @@ loss_linear_batch, bregman_projection_batch, bregman_log_projection_batch, + proximal_bregman_log_plan_batch, ) from ot import solve @@ -24,7 +26,7 @@ @pytest.mark.parametrize("solver", ["sinkhorn", "log_sinkhorn"]) @pytest.mark.parametrize("reg_type", ["kl", "entropy"]) -def test_solve_batch(solver, reg_type): +def test_sinkhorn_solve_batch(solver, reg_type): """Check that solve_batch gives the same results as solve for each instance in the batch.""" batchsize = 4 n = 16 @@ -61,6 +63,56 @@ def test_solve_batch(solver, reg_type): np.testing.assert_allclose(value_i, values_batch[i], atol=1e-4) +def test_proximal_solve_batch(): + """Check that proximal_bregman_log_plan_batch gives the same results as solve for each instance in the batch.""" + batchsize = 3 + n = 5 + d = 7 + rng = np.random.RandomState(0) + C = rng.rand(batchsize, n, d) + + exact_plan = np.zeros((batchsize, n, d)) + exact_value = np.zeros(batchsize) + for i in range(batchsize): + C_i = C[i] + res_i = solve(C_i, reg=None, tol=1e-5) + exact_plan[i] = res_i.plan + exact_value[i] = res_i.value_linear + + configs = [ + {"reg": None, "solver": "proximal"}, + {"reg": None, "solver": "log_sinkhorn"}, + {"reg": None, "solver": "sinkhorn"}, + {"reg": 0, "solver": "proximal"}, + {"reg": 1e-2, "solver": "proximal"}, + ] + + for config in configs: + res = solve_batch(C, max_iter=10000, tol=1e-5, grad="detach", **config) + plan = res["T"] + value = res["value_linear"] + np.testing.assert_allclose(plan, exact_plan, atol=1e-5) + np.testing.assert_allclose(value, exact_value, atol=1e-4) + + +@pytest.mark.parametrize("inner_iter", [1, 5, 10]) +def test_proximal_bregman_log_plan_batch(inner_iter): + batchsize = 3 + n = 5 + d = 7 + rng = np.random.RandomState(0) + C = rng.rand(batchsize, n, d) + res = proximal_bregman_log_plan_batch( + C, reg=1e-2, max_iter=10000, tol=1e-5, inner_iter=inner_iter, grad="detach" + ) + plan = res["T"] + for i in range(batchsize): + C_i = C[i] + res_i = solve(C_i, reg=None, tol=1e-5) + plan_i = res_i.plan + np.testing.assert_allclose(plan_i, plan[i], atol=1e-5) + + def test_bregman_batch(): batchsize = 4 d = 2 @@ -78,7 +130,8 @@ def test_bregman_batch(): @pytest.mark.parametrize("metric", ["sqeuclidean", "euclidean", "minkowski", "kl"]) -def test_metrics(metric): +@pytest.mark.parametrize("solver", ["proximal", "sinkhorn", "log_sinkhorn"]) +def test_sample_solve_batch(metric, solver): """Check that all functions run without error.""" batchsize = 2 @@ -93,11 +146,10 @@ def test_metrics(metric): is_positive = M >= 0 np.testing.assert_equal(is_positive.all(), True) - # Solve batch - res = solve_batch(M, reg=0.1, max_iter=10, tol=1e-5) - # Solve sample batch - res = solve_sample_batch(X, X, reg=0.1, max_iter=10, tol=1e-5, metric=metric) + res = solve_sample_batch( + X, X, reg=0.1, max_iter=10, tol=1e-5, metric=metric, solver=solver + ) # Compute loss loss = res.value_linear # loss given by solver @@ -116,7 +168,7 @@ def test_gradients_torch(grad): batchsize = 2 n = 4 d = 2 - for solver in ["sinkhorn", "log_sinkhorn"]: + for solver in ["proximal", "sinkhorn", "log_sinkhorn"]: X = torch.randn((batchsize, n, d), requires_grad=True) M = dist_batch(X, X) res = solve_batch(M, reg=0.1, max_iter=10, tol=1e-5, grad=grad, solver=solver) @@ -140,8 +192,9 @@ def test_backend(nx): X = np.random.randn(batchsize, n, d) X = nx.from_numpy(X) M = dist_batch(X, X) - solve_batch(M, reg=0.1, max_iter=10, tol=1e-5) - solve_sample_batch(X, X, reg=0.1, max_iter=10, tol=1e-5) + for solver in ["proximal", "sinkhorn", "log_sinkhorn"]: + solve_batch(M, reg=0.1, max_iter=10, tol=1e-5, solver=solver) + solve_sample_batch(X, X, reg=0.1, max_iter=10, tol=1e-5, solver=solver) def test_metric_default_parameters(): From 4eca18b592241ed14a50e115d49296d1467566b3 Mon Sep 17 00:00:00 2001 From: thibaut-germain Date: Fri, 3 Jul 2026 14:26:42 +0200 Subject: [PATCH 09/13] updated format --- examples/backends/plot_ot_batch.py | 25 ++++++++++++++++++++++--- ot/batch/__init__.py | 3 +++ ot/batch/_linear.py | 12 ++++++++---- ot/batch/_utils.py | 12 +++++++----- test/batch/test_solve_batch.py | 15 ++++++++++----- 5 files changed, 50 insertions(+), 17 deletions(-) diff --git a/examples/backends/plot_ot_batch.py b/examples/backends/plot_ot_batch.py index dea2657e1..5270c21f1 100644 --- a/examples/backends/plot_ot_batch.py +++ b/examples/backends/plot_ot_batch.py @@ -13,6 +13,7 @@ """ # Author: Paul Krzakala +# Thibaut Germain # License: MIT License # sphinx_gallery_thumbnail_number = 1 @@ -75,19 +76,37 @@ # This is simple but inefficient for large batches. # # Instead, you can use :func:`ot.batch.solve_batch`, which solves all -# problems in parallel. +# problems in parallel. Several solvers are available: ["sinkhorn", "log_sinkhorn"] +# which solve the entropic regularized OT problem, and ["proximal"] which +# solves the classical OT problem using a proximal point method. By default, +# the proximal solver is used, but you can change it with the `solver` argument. reg = 1.0 max_iter = 100 tol = 1e-3 -# Naive approach +# Classical OT problem +## Naive approach +results_values_list = [] +for i in range(n_problems): + res = ot.solve(M_list[i], reg=None, max_iter=max_iter, tol=tol) + results_values_list.append(res.value_linear) + +## Batched approach +results_batch = ot.solve_batch(M=M_batch, reg=None, max_iter=max_iter, tol=tol) +results_values_batch = results_batch.value_linear + +assert np.allclose(np.array(results_values_list), results_values_batch, atol=tol * 10) + + +# Entropic regularized OT problem +## Naive approach results_values_list = [] for i in range(n_problems): res = ot.solve(M_list[i], reg=reg, max_iter=max_iter, tol=tol, reg_type="entropy") results_values_list.append(res.value_linear) -# Batched approach +## Batched approach results_batch = ot.solve_batch( M=M_batch, reg=reg, max_iter=max_iter, tol=tol, reg_type="entropy" ) diff --git a/ot/batch/__init__.py b/ot/batch/__init__.py index d21019b98..93b998cff 100644 --- a/ot/batch/__init__.py +++ b/ot/batch/__init__.py @@ -5,6 +5,7 @@ # Author: Remi Flamary # Paul Krzakala +# Thibaut Germain # # License: MIT License @@ -25,6 +26,7 @@ bregman_log_projection_batch, bregman_projection_batch, entropy_batch, + proximal_bregman_log_plan_batch, ) __all__ = [ @@ -40,4 +42,5 @@ "loss_quadratic_batch", "loss_quadratic_samples_batch", "tensor_batch", + "proximal_bregman_log_plan_batch", ] diff --git a/ot/batch/_linear.py b/ot/batch/_linear.py index 450738366..de0301987 100644 --- a/ot/batch/_linear.py +++ b/ot/batch/_linear.py @@ -244,7 +244,7 @@ def solve_batch( tol=1e-5, solver="proximal", inner_iter=1, - inner_reg=1e-2, + inner_reg=1e-1, reg_type="entropy", grad="envelope", ): @@ -282,7 +282,7 @@ def solve_batch( inner_iter : int Number of inner Bregman iterations for the proximal solver. Default is 1. inner_reg : float - Regularization parameter for the inner Bregman iterations in the proximal solver. Default is 1e-2. + Regularization parameter for the inner Bregman iterations in the proximal solver. Default is 1e-1. reg_type : str, optional Type of regularization :math:`R` either "KL", or "entropy". Only used with Sinkhorn solver. Default is "entropy". grad : str, optional @@ -344,7 +344,9 @@ def solve_batch( "sinkhorn", ], f"Unknown solver: {solver}" - use_sinkhorn = solver in ["log_sinkhorn", "sinkhorn"] and reg > 0 + use_sinkhorn = ( + solver in ["log_sinkhorn", "sinkhorn"] and reg is not None and reg > 0 + ) nx = get_backend(a, b, M) @@ -372,7 +374,7 @@ def solve_batch( a, b, nx=nx, - reg=inner_reg, + inner_reg=inner_reg, max_iter=max_iter, tol=tol, inner_iter=inner_iter, @@ -396,6 +398,8 @@ def solve_batch( ref = nx.einsum("bi,bj->bij", a, b) kl = nx.sum(T * nx.log(T / ref + 1e-16), axis=(1, 2)) value = value_linear + reg * kl + else: + value = value_linear log = {"n_iter": out["n_iters"]} res = OTResult( diff --git a/ot/batch/_utils.py b/ot/batch/_utils.py index 984cafd4a..2aaabd596 100644 --- a/ot/batch/_utils.py +++ b/ot/batch/_utils.py @@ -306,7 +306,7 @@ def proximal_bregman_log_plan_batch( a=None, b=None, nx=None, - reg=1e-2, + inner_reg=1e-1, max_iter=10000, tol=1e-5, inner_iter=1, @@ -328,9 +328,9 @@ def proximal_bregman_log_plan_batch( The optimal transport plans are computed iteratively with a proximal point method based on a Bregman divergence where each iteration involves solving a Bregman projection problem: .. math:: - \mathbf{T}^{(k+1)} = \mathop{\arg \min}_\mathbf{T} \quad \langle \mathbf{C} - \textit{reg} \cdot \log \mathbf{T}^{(k)}, \mathbf{T} \rangle + \textit{reg} \cdot \sum_{i,j} \mathbf{T}_{i,j} \log \mathbf{T}_{i,j} + \mathbf{T}^{(k+1)} = \mathop{\arg \min}_\mathbf{T} \quad \langle \mathbf{C} - \textit{inner\_reg} \cdot \log \mathbf{T}^{(k)}, \mathbf{T} \rangle + \textit{inner\_reg} \cdot \sum_{i,j} \mathbf{T}_{i,j} \log \mathbf{T}_{i,j} - Denoting :math:`\mathbf{K}^{(k)} = - \mathbf{C} / \textit{reg} + \log \mathbf{T}^{(k)}`, the affinity matrix at iteration :math:`k`, the Bregman projection problem is solved in the log-domain with a finite number of inner iterations :math:`\text{inner\_iter}`, i.e., the dual variables :math:`\mathbf{u}` and :math:`\mathbf{v}` are updated as follows: + Denoting :math:`\mathbf{K}^{(k)} = - \mathbf{C} / \textit{inner\_reg} + \log \mathbf{T}^{(k)}`, the affinity matrix at iteration :math:`k`, the Bregman projection problem is solved in the log-domain with a finite number of inner iterations :math:`\text{inner\_iter}`, i.e., the dual variables :math:`\mathbf{u}` and :math:`\mathbf{v}` are updated as follows: .. math:: \mathbf{u}^{(i+1)} = \log(\mathbf{a}) - \text{LSE}(\mathbf{K}^{(k)} + \mathbf{v}^{(i)}) @@ -352,7 +352,9 @@ def proximal_bregman_log_plan_batch( max_iter : int, optional Maximum number of iterations. inner_iter : int, optional - Number of inner iterations for updating the dual variables u and v. + Number of inner iterations for updating the dual variables u and v. Default is 1. + inner_reg : float, optional + Regularization parameter for the Bregman divergence. Default is 1e-1. tol : float, optional Tolerance for convergence. The solver stops when the maximum change in the dual variables is below this value. @@ -402,7 +404,7 @@ def proximal_bregman_log_plan_batch( u = nx.zeros((B, n), type_as=C) v = nx.zeros((B, m), type_as=C) - K = -C / reg + K = -C / inner_reg loga = nx.log(a) logb = nx.log(b) diff --git a/test/batch/test_solve_batch.py b/test/batch/test_solve_batch.py index 900cff9d3..7ec648f33 100644 --- a/test/batch/test_solve_batch.py +++ b/test/batch/test_solve_batch.py @@ -89,9 +89,9 @@ def test_proximal_solve_batch(): for config in configs: res = solve_batch(C, max_iter=10000, tol=1e-5, grad="detach", **config) - plan = res["T"] - value = res["value_linear"] - np.testing.assert_allclose(plan, exact_plan, atol=1e-5) + plan = res.plan + value = res.value_linear + np.testing.assert_allclose(plan, exact_plan, atol=1e-4) np.testing.assert_allclose(value, exact_value, atol=1e-4) @@ -103,14 +103,19 @@ def test_proximal_bregman_log_plan_batch(inner_iter): rng = np.random.RandomState(0) C = rng.rand(batchsize, n, d) res = proximal_bregman_log_plan_batch( - C, reg=1e-2, max_iter=10000, tol=1e-5, inner_iter=inner_iter, grad="detach" + C, + inner_reg=1e-1, + max_iter=10000, + tol=1e-5, + inner_iter=inner_iter, + grad="detach", ) plan = res["T"] for i in range(batchsize): C_i = C[i] res_i = solve(C_i, reg=None, tol=1e-5) plan_i = res_i.plan - np.testing.assert_allclose(plan_i, plan[i], atol=1e-5) + np.testing.assert_allclose(plan_i, plan[i], atol=1e-4) def test_bregman_batch(): From 4d0efca839947d8c2b8bac7d0871f6ac92530722 Mon Sep 17 00:00:00 2001 From: thibaut-germain Date: Fri, 3 Jul 2026 14:42:16 +0200 Subject: [PATCH 10/13] added PR and references --- README.md | 3 +++ RELEASES.md | 1 + ot/batch/_linear.py | 4 ++-- ot/batch/_utils.py | 4 ++-- 4 files changed, 8 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 4ce4082ab..4f5c54c74 100644 --- a/README.md +++ b/README.md @@ -470,3 +470,6 @@ Artificial Intelligence. [89] Tipping, M.E. & Bishop, C.M. (1999). [Probabilistic principal component analysis]. Journal of the Royal Statistical Society Series B: Statistical Methodology 61.3 (1999): 611-622. [90] Genans, F., Godichon-Baggioni, A., Vialard, F. X., & Wintenberger, O. (2025). [Decreasing Entropic Regularization Averaged Gradient for Semi-Discrete Optimal Transport](https://proceedings.neurips.cc/paper_files/paper/2025/file/d7efa12e98f5e0dd8b4f48cd60b4e3aa-Paper-Conference.pdf). Advances in Neural Information Processing Systems, 38, 146913-146949. + +[91] Xie, Y., Wang, X., Wang, R., & Zha, H. (2020, August). +[A fast proximal point method for computing exact wasserstein distance.](https://proceedings.mlr.press/v115/xie20b/xie20b.pdf) In Uncertainty in artificial intelligence (pp. 433-453). PMLR. \ No newline at end of file diff --git a/RELEASES.md b/RELEASES.md index eea94981d..184389646 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -39,6 +39,7 @@ This new release adds support for sparse cost matrices and a new lazy exact OT s - Fix docstrings for `lowrank_gromov_wasserstein_samples` and `lowrank_sinkhorn` (PR #823) - Reorganize all tests per backend (PR #828) - Update sgot cost function and example (PR #830) +- Implemented batch proximal point solver for OT problems `ot.batch.proximal_bregman_log_plan_batch` function and updated wrapper functions `ot.batch.solve_batch` and `ot.batch.solve_sample_batch` (PR #832) #### Closed issues diff --git a/ot/batch/_linear.py b/ot/batch/_linear.py index de0301987..4e764ffcb 100644 --- a/ot/batch/_linear.py +++ b/ot/batch/_linear.py @@ -261,7 +261,7 @@ def solve_batch( & \mathbf{T} \geq 0 \end{aligned} - The problem is solved with either a proximal point method :ref:`[1] ` or a Sinkhorn algorithm :ref:`[2] `. Unlike the Sinkhorn algorithm, which assumes a regularization term, the proximal point method aims to solve the exact optimal transport problem leading to sparser transport plans. The choice of solver depends on the value of `reg` and `solver`. By default, the proximal point method is used. The Sinkhorn algorithm is only used when `reg` is greater than 0 and `solver='sinkhorn' or `log_sinkhorn'`. In any other case, the proximal point method is used, i.e. whenever `reg` is None or 0 or `solver='proximal'`. + The problem is solved with either a proximal point method :ref:`[91] ` or a Sinkhorn algorithm :ref:`[2] `. Unlike the Sinkhorn algorithm, which assumes a regularization term, the proximal point method aims to solve the exact optimal transport problem leading to sparser transport plans. The choice of solver depends on the value of `reg` and `solver`. By default, the proximal point method is used. The Sinkhorn algorithm is only used when `reg` is greater than 0 and `solver='sinkhorn' or `log_sinkhorn'`. In any other case, the proximal point method is used, i.e. whenever `reg` is None or 0 or `solver='proximal'`. Parameters ---------- @@ -329,7 +329,7 @@ def solve_batch( .. _references-batch-solver: Reference ---------- - .. [1] Xie, Y., Wang, X., Wang, R., & Zha, H. (2020, August). + .. [91] Xie, Y., Wang, X., Wang, R., & Zha, H. (2020, August). A fast proximal point method for computing exact wasserstein distance. In Uncertainty in artificial intelligence (pp. 433-453). PMLR. diff --git a/ot/batch/_utils.py b/ot/batch/_utils.py index 2aaabd596..a0582f1a5 100644 --- a/ot/batch/_utils.py +++ b/ot/batch/_utils.py @@ -313,7 +313,7 @@ def proximal_bregman_log_plan_batch( grad="detach", ): r""" - Returns optimal transport plans for a batch of cost matrices :math:`\mathbf{C}` using a Bregman divergence based proximal point method :ref:`[1] `. + Returns optimal transport plans for a batch of cost matrices :math:`\mathbf{C}` using a Bregman divergence based proximal point method :ref:`[91] `. This function solves the following optimization problem: @@ -386,7 +386,7 @@ def proximal_bregman_log_plan_batch( .. _references-OT-bregman-proximal-point: Reference ---------- - .. [1] Xie, Y., Wang, X., Wang, R., & Zha, H. (2020, August). + .. [91] Xie, Y., Wang, X., Wang, R., & Zha, H. (2020, August). A fast proximal point method for computing exact wasserstein distance. In Uncertainty in artificial intelligence (pp. 433-453). PMLR. """ From 3151702b7ce0ad6f0069a41263444a1da1197c3b Mon Sep 17 00:00:00 2001 From: thibaut-germain Date: Mon, 6 Jul 2026 13:32:09 +0200 Subject: [PATCH 11/13] aligned usage of solve and solve_batch + added batch tests --- examples/backends/plot_ot_batch.py | 12 +- ot/batch/_linear.py | 135 ++++++++++++-------- ot/batch/_utils.py | 21 +-- test/batch/test_solve_batch.py | 198 +++++++++++------------------ 4 files changed, 174 insertions(+), 192 deletions(-) diff --git a/examples/backends/plot_ot_batch.py b/examples/backends/plot_ot_batch.py index 5270c21f1..d628fd6b2 100644 --- a/examples/backends/plot_ot_batch.py +++ b/examples/backends/plot_ot_batch.py @@ -76,12 +76,13 @@ # This is simple but inefficient for large batches. # # Instead, you can use :func:`ot.batch.solve_batch`, which solves all -# problems in parallel. Several solvers are available: ["sinkhorn", "log_sinkhorn"] -# which solve the entropic regularized OT problem, and ["proximal"] which -# solves the classical OT problem using a proximal point method. By default, -# the proximal solver is used, but you can change it with the `solver` argument. +# problems in parallel. Several methods are available: ["sinkhorn", "log_sinkhorn"] +# which solve regularized OT problems, and ["proximal"] which +# solves regularized and unregularized OT problem using a proximal point scheme. +# By default, the method is set to "auto" which automatically selects the appropriate +# method based on the value of `reg`. If `reg` is None or 0, the proximal point method +# is used. If `reg` is greater than 0, the Sinkhorn algorithm is used. -reg = 1.0 max_iter = 100 tol = 1e-3 @@ -101,6 +102,7 @@ # Entropic regularized OT problem ## Naive approach +reg = 1.0 results_values_list = [] for i in range(n_problems): res = ot.solve(M_list[i], reg=reg, max_iter=max_iter, tol=tol, reg_type="entropy") diff --git a/ot/batch/_linear.py b/ot/batch/_linear.py index 4e764ffcb..16347faa9 100644 --- a/ot/batch/_linear.py +++ b/ot/batch/_linear.py @@ -19,6 +19,13 @@ ) +solve_batch_method_lst = ["auto", "proximal", "log_sinkhorn", "sinkhorn"] + +solve_batch_reg_type_lst = ["kl", "entropy"] + +solve_batch_grad_lst = ["detach", "autodiff", "last_step", "envelope"] + + def dist_lp_batch(X, Y, p=2, q=1, nx=None): r"""Computes the cost matrix for a batch of samples using the Lp norm. @@ -242,7 +249,7 @@ def solve_batch( b=None, max_iter=1000, tol=1e-5, - solver="proximal", + method="auto", inner_iter=1, inner_reg=1e-1, reg_type="entropy", @@ -261,14 +268,14 @@ def solve_batch( & \mathbf{T} \geq 0 \end{aligned} - The problem is solved with either a proximal point method :ref:`[91] ` or a Sinkhorn algorithm :ref:`[2] `. Unlike the Sinkhorn algorithm, which assumes a regularization term, the proximal point method aims to solve the exact optimal transport problem leading to sparser transport plans. The choice of solver depends on the value of `reg` and `solver`. By default, the proximal point method is used. The Sinkhorn algorithm is only used when `reg` is greater than 0 and `solver='sinkhorn' or `log_sinkhorn'`. In any other case, the proximal point method is used, i.e. whenever `reg` is None or 0 or `solver='proximal'`. + The problem is solved with either a proximal point method :ref:`[91] ` or a Sinkhorn algorithm :ref:`[2] `. Unlike the Sinkhorn algorithm, which assumes a regularization term, the proximal point method can solve both regularized and unregularized optimal transport problems. When `method` is set to 'auto', the function automatically selects the appropriate method based on the value of `reg`. if `reg` is None or 0, the proximal point method is used. If `reg` is greater than 0, the Sinkhorn algorithm is used. Parameters ---------- M : array-like, shape (B, ns, nt) Cost matrix reg : float - Regularization parameter for entropic regularization. If None or 0, no regularization is applied, proximal solver is used. Default is None. + Regularization parameter. Default is None. a : array-like, shape (B, ns) Source distribution (optional). If None, uniform distribution is used. b : array-like, shape (B, nt) @@ -277,19 +284,19 @@ def solve_batch( Maximum number of iterations tol : float Tolerance for convergence - solver: str - Solver to use, either 'proximal', 'log_sinkhorn' or 'sinkhorn'. Default is 'proximal'. + method: str + Method to use, either 'auto', 'proximal', 'log_sinkhorn' or 'sinkhorn'. Default is 'auto'. inner_iter : int - Number of inner Bregman iterations for the proximal solver. Default is 1. + Number of inner Bregman iterations for the proximal method. Default is 1. inner_reg : float - Regularization parameter for the inner Bregman iterations in the proximal solver. Default is 1e-1. + Regularization parameter for the inner Bregman iterations in the proximal method. Default is 1e-1. reg_type : str, optional - Type of regularization :math:`R` either "KL", or "entropy". Only used with Sinkhorn solver. Default is "entropy". + Type of regularization :math:`R` either "KL", or "entropy". Default is "entropy". grad : str, optional Type of gradient computation, either 'detach', 'autodiff', 'last_step' or 'envelope'. - 'detach' does not compute the gradients for the Sinkhorn solver. + 'detach' does not compute the gradients. 'autodiff' provides gradients of all outputs (`plan, value, value_linear`) but with important memory cost. - 'last_step' provides gradients of all outputs (`plan, value, value_linear`) only for the last solver iteration, useful for memory saving. + 'last_step' provides gradients of all outputs (`plan, value, value_linear`) only for the last method iteration, useful for memory saving. 'envelope' provides gradients only for `value`. Default is 'envelope'. @@ -313,9 +320,9 @@ def solve_batch( >>> X = np.random.randn(5, 10, 3) # 5 batches of 10 samples in 3D >>> Y = np.random.randn(5, 15, 3) # 5 batches of 15 samples in 3D >>> M = dist_batch(X, Y, metric="euclidean") # Compute cost matrices - >>> p_result = solve_batch(M) # Uses proximal solver + >>> p_result = solve_batch(M) # Uses proximal method >>> reg = 0.1 - >>> s_result = solve_batch(M, reg, solver="log_sinkhorn") # Uses Sinkhorn solver + >>> s_result = solve_batch(M, reg, method="log_sinkhorn") # Uses Sinkhorn method >>> s_result.plan.shape # Optimal transport plans for each batch (5, 10, 15) >>> s_result.value.shape # Optimal transport values for each batch @@ -324,7 +331,7 @@ def solve_batch( See Also -------- ot.batch.dist_batch : batched cost matrix computation for computing M. - ot.solve : non-batched version of the OT solver. + ot.solve : non-batched version of the solve_batch function. .. _references-batch-solver: Reference @@ -338,15 +345,31 @@ def solve_batch( Systems (NIPS) 26, 2013 """ - assert solver in [ - "proximal", - "log_sinkhorn", - "sinkhorn", - ], f"Unknown solver: {solver}" + if method not in solve_batch_method_lst: + raise ValueError( + f"Unknown method: {method}. Must be one of {solve_batch_method_lst}." + ) - use_sinkhorn = ( - solver in ["log_sinkhorn", "sinkhorn"] and reg is not None and reg > 0 - ) + if reg_type not in solve_batch_reg_type_lst: + raise ValueError( + f"Unknown reg_type: {reg_type}. Must be one of {solve_batch_reg_type_lst}." + ) + + if grad not in solve_batch_grad_lst: + raise ValueError( + f"Unknown grad: {grad}. Must be one of {solve_batch_grad_lst}." + ) + + if method in ["sinkhorn", "log_sinkhorn"] and (reg is None or reg <= 0): + raise ValueError( + "Sinkhorn methods require a strictly positive reg parameter. Please provide a valid reg value." + ) + + if method == "auto": + if reg is None or reg == 0: + method = "proximal" + else: + method = "log_sinkhorn" nx = get_backend(a, b, M) @@ -357,23 +380,23 @@ def solve_batch( if b is None: b = nx.ones((B, m), type_as=M) / m - if use_sinkhorn: - if solver == "log_sinkhorn": - K = -M / reg - out = bregman_log_projection_batch( - K, a, b, nx=nx, max_iter=max_iter, tol=tol, grad=grad - ) - if solver == "sinkhorn": - K = nx.exp(-M / reg) - out = bregman_projection_batch( - K, a, b, nx=nx, max_iter=max_iter, tol=tol, grad=grad - ) - else: + if method == "log_sinkhorn": + K = -M / reg + out = bregman_log_projection_batch( + K, a, b, nx=nx, max_iter=max_iter, tol=tol, grad=grad + ) + if method == "sinkhorn": + K = nx.exp(-M / reg) + out = bregman_projection_batch( + K, a, b, nx=nx, max_iter=max_iter, tol=tol, grad=grad + ) + if method == "proximal": out = proximal_bregman_log_plan_batch( M, a, b, nx=nx, + reg=reg, inner_reg=inner_reg, max_iter=max_iter, tol=tol, @@ -390,14 +413,13 @@ def solve_batch( T = nx.detach(T) value_linear = loss_linear_batch(M, T) - if use_sinkhorn: - if reg_type.lower() == "entropy": - entr = -entropy_batch(T, nx=nx) - value = value_linear + reg * entr - elif reg_type.lower() == "kl": - ref = nx.einsum("bi,bj->bij", a, b) - kl = nx.sum(T * nx.log(T / ref + 1e-16), axis=(1, 2)) - value = value_linear + reg * kl + if reg_type == "entropy" and reg is not None: + entr = -entropy_batch(T, nx=nx) + value = value_linear + reg * entr + elif reg_type == "kl" and reg is not None: + ref = nx.einsum("bi,bj->bij", a, b) + kl = nx.sum(T * nx.log(T / ref + 1e-16), axis=(1, 2)) + value = value_linear + reg * kl else: value = value_linear log = {"n_iter": out["n_iters"]} @@ -424,29 +446,29 @@ def solve_sample_batch( p=2, max_iter=1000, tol=1e-5, - solver="proximal", + method="auto", inner_iter=1, inner_reg=1e-2, reg_type="entropy", grad="envelope", ): r""" - Return solutions of a batch of discrete optimal transport problems in a :any:`OTResult` object computed for source and target samples. + Return solutions of a batch of discrete optimal transport problems in a :any:`OTResult` object computed from batches of source and target samples. - The problem is solved with either a proximal point method :ref:`[1] ` or a Sinkhorn algorithm :ref:`[2] `. Unlike the Sinkhorn algorithm, which assumes a regularization term, the proximal point method aims to solve the exact optimal transport problem leading to sparser transport plans. The choice of solver depends on the value of `reg` and `solver`. By default, the proximal point method is used. The Sinkhorn algorithm is only used when `reg` is greater than 0 and `solver='sinkhorn' or `log_sinkhorn'`. In any other case, the proximal point method is used, i.e. whenever `reg` is None or 0 or `solver='proximal'`. + The problem is solved with either a proximal point method :ref:`[91] ` or a Sinkhorn algorithm :ref:`[2] `. Unlike the Sinkhorn algorithm, which assumes a regularization term, the proximal point method can solve both regularized and unregularized optimal transport problems. When `method` is set to 'auto', the function automatically selects the appropriate method based on the value of `reg`. if `reg` is None or 0, the proximal point method is used. If `reg` is greater than 0, the Sinkhorn algorithm is used. Parameters ---------- - Xa : array-like, shape (B, ns, d) + X_a : array-like, shape (B, ns, d) Samples from source distribution - Xb : array-like, shape (B, nt, d) + X_b : array-like, shape (B, nt, d) Samples from target distribution metric : str, optional 'sqeuclidean', 'euclidean', 'minkowski' or 'kl' p : float, optional p-norm for the Minkowski metrics. Default value is 2. reg : float - Regularization parameter for entropic regularization. If None or 0, no regularization is applied, proximal solver is used. Default is None. + Regularization parameter. Default is None. a : array-like, shape (B, ns) Source distribution (optional). If None, uniform distribution is used. b : array-like, shape (B, nt) @@ -455,19 +477,19 @@ def solve_sample_batch( Maximum number of iterations tol : float Tolerance for convergence - solver: str - Solver to use, either 'proximal', 'log_sinkhorn' or 'sinkhorn'. Default is 'proximal'. + method: str + Method to use, either 'auto', 'proximal', 'log_sinkhorn' or 'sinkhorn'. Default is 'auto'. inner_iter : int - Number of inner Bregman iterations for the proximal solver. Default is 1. + Number of inner Bregman iterations for the proximal method. Default is 1. inner_reg : float - Regularization parameter for the inner Bregman iterations in the proximal solver. Default is 1e-2. + Regularization parameter for the inner Bregman iterations in the proximal method. Default is 1e-2. reg_type : str, optional - Type of regularization :math:`R` either "KL", or "entropy". Only used with Sinkhorn solver. Default is "entropy". + Type of regularization :math:`R` either "KL", or "entropy". Default is "entropy". grad : str, optional Type of gradient computation, either 'detach', 'autodiff', 'last_step' or 'envelope'. - 'detach' does not compute the gradients for the Sinkhorn solver. + 'detach' does not compute the gradients. 'autodiff' provides gradients of all outputs (`plan, value, value_linear`) but with important memory cost. - 'last_step' provides gradients of all outputs (`plan, value, value_linear`) only for the last solver iteration, useful for memory saving. + 'last_step' provides gradients of all outputs (`plan, value, value_linear`) only for the last method iteration, useful for memory saving. 'envelope' provides gradients only for `value`. Default is 'envelope'. @@ -485,10 +507,11 @@ def solve_sample_batch( See Also -------- - ot.batch.solve_batch : solver for computing the optimal T from arbitrary cost matrix M. + ot.batch.solve_batch : function for computing the optimal T from arbitrary cost matrix M. """ M = dist_batch(X_a, X_b, metric=metric, p=p) + return solve_batch( M, reg, @@ -496,7 +519,7 @@ def solve_sample_batch( b=b, max_iter=max_iter, tol=tol, - solver=solver, + method=method, inner_iter=inner_iter, inner_reg=inner_reg, reg_type=reg_type, diff --git a/ot/batch/_utils.py b/ot/batch/_utils.py index a0582f1a5..54f035dcc 100644 --- a/ot/batch/_utils.py +++ b/ot/batch/_utils.py @@ -306,7 +306,8 @@ def proximal_bregman_log_plan_batch( a=None, b=None, nx=None, - inner_reg=1e-1, + reg=None, + inner_reg=1e-2, max_iter=10000, tol=1e-5, inner_iter=1, @@ -319,7 +320,7 @@ def proximal_bregman_log_plan_batch( .. math:: \begin{aligned} - \mathbf{T} = \mathop{\arg \min}_\mathbf{T} \quad & \langle \mathbf{T}, \mathbf{C} \rangle_F \\ + \mathbf{T} = \mathop{\arg \min}_\mathbf{T} \quad & \langle \mathbf{T}, \mathbf{C} \rangle_F + \textit{reg} \cdot \sum_{i,j} \mathbf{T}_{i,j} \log \mathbf{T}_{i,j} \\ \text{s.t.} \quad & \mathbf{T} \mathbf{1} = \mathbf{a} \\ & \mathbf{T}^T \mathbf{1} = \mathbf{b} \\ & \mathbf{T} \geq 0 @@ -328,9 +329,9 @@ def proximal_bregman_log_plan_batch( The optimal transport plans are computed iteratively with a proximal point method based on a Bregman divergence where each iteration involves solving a Bregman projection problem: .. math:: - \mathbf{T}^{(k+1)} = \mathop{\arg \min}_\mathbf{T} \quad \langle \mathbf{C} - \textit{inner\_reg} \cdot \log \mathbf{T}^{(k)}, \mathbf{T} \rangle + \textit{inner\_reg} \cdot \sum_{i,j} \mathbf{T}_{i,j} \log \mathbf{T}_{i,j} + \mathbf{T}^{(k+1)} = \mathop{\arg \min}_\mathbf{T} \quad \langle \mathbf{C} - \textit{inner\_reg} \cdot \log \mathbf{T}^{(k)}, \mathbf{T} \rangle + (\textit{reg} + \textit{inner\_reg}) \cdot \sum_{i,j} \mathbf{T}_{i,j} \log \mathbf{T}_{i,j} - Denoting :math:`\mathbf{K}^{(k)} = - \mathbf{C} / \textit{inner\_reg} + \log \mathbf{T}^{(k)}`, the affinity matrix at iteration :math:`k`, the Bregman projection problem is solved in the log-domain with a finite number of inner iterations :math:`\text{inner\_iter}`, i.e., the dual variables :math:`\mathbf{u}` and :math:`\mathbf{v}` are updated as follows: + Denoting :math:`\mathbf{K}^{(k)} = - (\mathbf{C} + \textit{inner\_reg} \cdot \log \mathbf{T}^{(k)})/(\textit{reg} + \textit{inner\_reg})`, the affinity matrix at iteration :math:`k`, the Bregman projection problem is solved in the log-domain with a finite number of inner iterations :math:`\text{inner\_iter}`, i.e., the dual variables :math:`\mathbf{u}` and :math:`\mathbf{v}` are updated as follows: .. math:: \mathbf{u}^{(i+1)} = \log(\mathbf{a}) - \text{LSE}(\mathbf{K}^{(k)} + \mathbf{v}^{(i)}) @@ -401,21 +402,23 @@ def proximal_bregman_log_plan_batch( if b is None: b = nx.ones((B, m)) / m + if reg is None: + reg = 0.0 + u = nx.zeros((B, n), type_as=C) v = nx.zeros((B, m), type_as=C) - K = -C / inner_reg loga = nx.log(a) logb = nx.log(b) if grad == "detach": - K = nx.detach(K) + C = nx.detach(C) elif grad == "last_step": - K_, K = K.clone(), nx.detach(K) + C_, C = C.clone(), nx.detach(C) log_T = nx.zeros(C.shape, type_as=C) for n_iters in range(max_iter): - K_proj = log_T + K + K_proj = -(C + inner_reg * log_T) / (reg + inner_reg) for _ in range(inner_iter): u = loga - nx.logsumexp(K_proj + v[:, None, :], axis=2) v = logb - nx.logsumexp(K_proj + u[:, :, None], axis=1) @@ -430,7 +433,7 @@ def proximal_bregman_log_plan_batch( break if grad == "last_step": - K_proj = log_T + K_ + K_proj = -(C_ + inner_reg * log_T) / (reg + inner_reg) for _ in range(inner_iter): u = loga - nx.logsumexp(K_proj + v[:, None, :], axis=2) v = logb - nx.logsumexp(K_proj + u[:, :, None], axis=1) diff --git a/test/batch/test_solve_batch.py b/test/batch/test_solve_batch.py index 7ec648f33..08e5ff81e 100644 --- a/test/batch/test_solve_batch.py +++ b/test/batch/test_solve_batch.py @@ -20,102 +20,81 @@ ) from ot import solve +from contextlib import nullcontext import pytest from ot.backend import torch -@pytest.mark.parametrize("solver", ["sinkhorn", "log_sinkhorn"]) +@pytest.mark.parametrize("reg", [None, 0, 1e-0]) +@pytest.mark.parametrize("method", ["auto", "proximal", "sinkhorn", "log_sinkhorn"]) @pytest.mark.parametrize("reg_type", ["kl", "entropy"]) -def test_sinkhorn_solve_batch(solver, reg_type): +def test_solve_batch_vs_solve(reg, method, reg_type): """Check that solve_batch gives the same results as solve for each instance in the batch.""" - batchsize = 4 - n = 16 - rng = np.random.RandomState(0) - - M = rng.rand(batchsize, n, n) - - reg = 0.1 - max_iter = 10000 - tol = 1e-5 - - res = solve_batch( - M, - a=None, - b=None, - reg=reg, - max_iter=max_iter, - tol=tol, - solver=solver, - reg_type=reg_type, - grad="detach", - ) - plan_batch = res.plan - values_batch = res.value_linear - for i in range(batchsize): - M_i = M[i] - res_i = solve( - M_i, a=None, b=None, reg=reg, max_iter=max_iter, tol=tol, reg_type=reg_type + should_fail = method in ["sinkhorn", "log_sinkhorn"] and (reg is None or reg <= 0) + + ctx = pytest.raises(Exception) if should_fail else nullcontext() + + with ctx: + tol = 1e-4 + batchsize = 3 + n = 5 + d = 7 + rng = np.random.RandomState(0) + C = rng.rand(batchsize, n, d) + + base_plan = np.zeros((batchsize, n, d)) + base_value = np.zeros(batchsize) + for i in range(batchsize): + C_i = C[i] + res_i = solve(C_i, reg=reg, tol=tol, reg_type=reg_type) + base_plan[i] = res_i.plan + base_value[i] = res_i.value_linear + + res = solve_batch( + C, + max_iter=10000, + tol=tol, + grad="detach", + reg=reg, + method=method, + reg_type=reg_type, + inner_reg=1e-3, ) - plan_i = res_i.plan - value_i = res_i.value_linear - np.testing.assert_allclose(plan_i, plan_batch[i], atol=1e-05) - np.testing.assert_allclose(value_i, values_batch[i], atol=1e-4) - - -def test_proximal_solve_batch(): - """Check that proximal_bregman_log_plan_batch gives the same results as solve for each instance in the batch.""" - batchsize = 3 - n = 5 - d = 7 - rng = np.random.RandomState(0) - C = rng.rand(batchsize, n, d) - - exact_plan = np.zeros((batchsize, n, d)) - exact_value = np.zeros(batchsize) - for i in range(batchsize): - C_i = C[i] - res_i = solve(C_i, reg=None, tol=1e-5) - exact_plan[i] = res_i.plan - exact_value[i] = res_i.value_linear - - configs = [ - {"reg": None, "solver": "proximal"}, - {"reg": None, "solver": "log_sinkhorn"}, - {"reg": None, "solver": "sinkhorn"}, - {"reg": 0, "solver": "proximal"}, - {"reg": 1e-2, "solver": "proximal"}, - ] - - for config in configs: - res = solve_batch(C, max_iter=10000, tol=1e-5, grad="detach", **config) plan = res.plan value = res.value_linear - np.testing.assert_allclose(plan, exact_plan, atol=1e-4) - np.testing.assert_allclose(value, exact_value, atol=1e-4) + np.testing.assert_allclose(plan, base_plan, atol=tol * 10) + np.testing.assert_allclose(value, base_value, atol=tol * 10) +@pytest.mark.parametrize("reg", [None, 0, 1e-0]) @pytest.mark.parametrize("inner_iter", [1, 5, 10]) -def test_proximal_bregman_log_plan_batch(inner_iter): +def test_backend_proximal_bregman_log_plan_batch(nx, reg, inner_iter): + tol = 1e-4 batchsize = 3 n = 5 d = 7 rng = np.random.RandomState(0) C = rng.rand(batchsize, n, d) res = proximal_bregman_log_plan_batch( - C, - inner_reg=1e-1, + nx.from_numpy(C), + reg=reg, + inner_reg=1e-3, max_iter=10000, - tol=1e-5, + tol=tol, inner_iter=inner_iter, grad="detach", ) - plan = res["T"] + plan = nx.to_numpy(res["T"]) for i in range(batchsize): C_i = C[i] - res_i = solve(C_i, reg=None, tol=1e-5) + res_i = solve( + C_i, + reg=reg, + tol=tol, + ) plan_i = res_i.plan - np.testing.assert_allclose(plan_i, plan[i], atol=1e-4) + np.testing.assert_allclose(plan_i, plan[i], atol=tol * 10) def test_bregman_batch(): @@ -135,10 +114,10 @@ def test_bregman_batch(): @pytest.mark.parametrize("metric", ["sqeuclidean", "euclidean", "minkowski", "kl"]) -@pytest.mark.parametrize("solver", ["proximal", "sinkhorn", "log_sinkhorn"]) -def test_sample_solve_batch(metric, solver): +@pytest.mark.parametrize("method", ["proximal", "sinkhorn", "log_sinkhorn"]) +def test_sample_solve_batch_vs_solve_batch(metric, method): """Check that all functions run without error.""" - + tol = 1e-5 batchsize = 2 n = 4 d = 2 @@ -153,7 +132,7 @@ def test_sample_solve_batch(metric, solver): # Solve sample batch res = solve_sample_batch( - X, X, reg=0.1, max_iter=10, tol=1e-5, metric=metric, solver=solver + X, X, reg=0.1, max_iter=10, tol=tol, metric=metric, method=method ) # Compute loss @@ -162,34 +141,35 @@ def test_sample_solve_batch(metric, solver): loss3 = loss_linear_samples_batch( X, X, res.plan, metric=metric ) # recompute loss from plan and samples - np.testing.assert_allclose(loss, loss2, atol=1e-5) - np.testing.assert_allclose(loss, loss3, atol=1e-5) + np.testing.assert_allclose(loss, loss2, atol=tol * 10) + np.testing.assert_allclose(loss, loss3, atol=tol * 10) @pytest.mark.skipif(not torch, reason="torch not installed") @pytest.mark.parametrize("grad", ["detach", "envelope", "autodiff", "last_step"]) -def test_gradients_torch(grad): +@pytest.mark.parametrize("method", ["proximal", "sinkhorn", "log_sinkhorn"]) +def test_gradients_torch(grad, method): """Check that all gradient methods run without error.""" batchsize = 2 n = 4 d = 2 - for solver in ["proximal", "sinkhorn", "log_sinkhorn"]: - X = torch.randn((batchsize, n, d), requires_grad=True) - M = dist_batch(X, X) - res = solve_batch(M, reg=0.1, max_iter=10, tol=1e-5, grad=grad, solver=solver) - loss = res.value_linear.sum() - loss_plan = res.plan.sum() - if grad == "detach": - assert loss.grad == None - elif grad == "envelope": - loss.backward() - assert X.grad is not None - elif grad in ["autodiff", "last_step"]: - loss_plan.backward() - assert X.grad is not None - - -def test_backend(nx): + X = torch.randn((batchsize, n, d), requires_grad=True) + M = dist_batch(X, X) + res = solve_batch(M, reg=0.1, max_iter=10, tol=1e-5, grad=grad, method=method) + loss = res.value_linear.sum() + loss_plan = res.plan.sum() + if grad == "detach": + assert loss.grad == None + elif grad == "envelope": + loss.backward() + assert X.grad is not None + elif grad in ["autodiff", "last_step"]: + loss_plan.backward() + assert X.grad is not None + + +@pytest.mark.parametrize("method", ["proximal", "sinkhorn", "log_sinkhorn"]) +def test_backend(nx, method): """Check that all gradient methods run without error.""" batchsize = 2 n = 4 @@ -197,31 +177,5 @@ def test_backend(nx): X = np.random.randn(batchsize, n, d) X = nx.from_numpy(X) M = dist_batch(X, X) - for solver in ["proximal", "sinkhorn", "log_sinkhorn"]: - solve_batch(M, reg=0.1, max_iter=10, tol=1e-5, solver=solver) - solve_sample_batch(X, X, reg=0.1, max_iter=10, tol=1e-5, solver=solver) - - -def test_metric_default_parameters(): - """Check that all functions with default parameters run without error.""" - - batchsize = 2 - n = 4 - d = 2 - rng = np.random.RandomState(0) - X = rng.rand(batchsize, n, d) - M = dist_batch(X, X) - is_positive = M >= 0 - np.testing.assert_equal(is_positive.all(), True) - - # Solve batch - res = solve_batch(M, reg=0.1, max_iter=10, tol=1e-5) - - # Solve sample batch - res = solve_sample_batch(X, X, reg=0.1) - - # Compute loss - loss_linear_batch(M, res.plan) # recompute loss from plan - loss_linear_samples_batch(X, X, res.plan) # recompute loss from plan and samples - assert np.isfinite(loss_linear_batch(M, res.plan)).all() - assert np.isfinite(loss_linear_samples_batch(X, X, res.plan)).all() + solve_batch(M, reg=0.1, max_iter=10, tol=1e-5, method=method) + solve_sample_batch(X, X, reg=0.1, max_iter=10, tol=1e-5, method=method) From f6a79b520ce6ef6241e8bf4b34979a58c02d10e0 Mon Sep 17 00:00:00 2001 From: thibaut-germain Date: Mon, 6 Jul 2026 17:18:52 +0200 Subject: [PATCH 12/13] fixed example --- README.md | 2 +- examples/backends/plot_ot_batch.py | 17 ++++++++++++----- ot/batch/_linear.py | 12 ++++++------ ot/batch/_utils.py | 8 ++++---- 4 files changed, 23 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index 4f5c54c74..7eb9b8841 100644 --- a/README.md +++ b/README.md @@ -471,5 +471,5 @@ Artificial Intelligence. [90] Genans, F., Godichon-Baggioni, A., Vialard, F. X., & Wintenberger, O. (2025). [Decreasing Entropic Regularization Averaged Gradient for Semi-Discrete Optimal Transport](https://proceedings.neurips.cc/paper_files/paper/2025/file/d7efa12e98f5e0dd8b4f48cd60b4e3aa-Paper-Conference.pdf). Advances in Neural Information Processing Systems, 38, 146913-146949. -[91] Xie, Y., Wang, X., Wang, R., & Zha, H. (2020, August). +[92] Xie, Y., Wang, X., Wang, R., & Zha, H. (2020, August). [A fast proximal point method for computing exact wasserstein distance.](https://proceedings.mlr.press/v115/xie20b/xie20b.pdf) In Uncertainty in artificial intelligence (pp. 433-453). PMLR. \ No newline at end of file diff --git a/examples/backends/plot_ot_batch.py b/examples/backends/plot_ot_batch.py index d628fd6b2..b530ae9cf 100644 --- a/examples/backends/plot_ot_batch.py +++ b/examples/backends/plot_ot_batch.py @@ -83,8 +83,8 @@ # method based on the value of `reg`. If `reg` is None or 0, the proximal point method # is used. If `reg` is greater than 0, the Sinkhorn algorithm is used. -max_iter = 100 -tol = 1e-3 +max_iter = 10000 +tol = 1e-4 # Classical OT problem ## Naive approach @@ -97,8 +97,9 @@ results_batch = ot.solve_batch(M=M_batch, reg=None, max_iter=max_iter, tol=tol) results_values_batch = results_batch.value_linear -assert np.allclose(np.array(results_values_list), results_values_batch, atol=tol * 10) - +exact_validated = np.allclose( + np.array(results_values_list), results_values_batch, atol=tol * 10 +) # Entropic regularized OT problem ## Naive approach @@ -114,7 +115,13 @@ ) results_values_batch = results_batch.value_linear -assert np.allclose(np.array(results_values_list), results_values_batch, atol=tol * 10) +entropic_validated = np.allclose( + np.array(results_values_list), results_values_batch, atol=tol * 10 +) + +print( + f"Exact solve vs proximal batch close: {exact_validated} \nSinkhorn solve vs Sinkhorn solve_batch close: {entropic_validated}" +) ############################################################################# # diff --git a/ot/batch/_linear.py b/ot/batch/_linear.py index 16347faa9..3c07ae639 100644 --- a/ot/batch/_linear.py +++ b/ot/batch/_linear.py @@ -251,7 +251,7 @@ def solve_batch( tol=1e-5, method="auto", inner_iter=1, - inner_reg=1e-1, + inner_reg=1e-3, reg_type="entropy", grad="envelope", ): @@ -268,7 +268,7 @@ def solve_batch( & \mathbf{T} \geq 0 \end{aligned} - The problem is solved with either a proximal point method :ref:`[91] ` or a Sinkhorn algorithm :ref:`[2] `. Unlike the Sinkhorn algorithm, which assumes a regularization term, the proximal point method can solve both regularized and unregularized optimal transport problems. When `method` is set to 'auto', the function automatically selects the appropriate method based on the value of `reg`. if `reg` is None or 0, the proximal point method is used. If `reg` is greater than 0, the Sinkhorn algorithm is used. + The problem is solved with either a proximal point method :ref:`[92] ` or a Sinkhorn algorithm :ref:`[2] `. Unlike the Sinkhorn algorithm, which assumes a regularization term, the proximal point method can solve both regularized and unregularized optimal transport problems. When `method` is set to 'auto', the function automatically selects the appropriate method based on the value of `reg`. if `reg` is None or 0, the proximal point method is used. If `reg` is greater than 0, the Sinkhorn algorithm is used. Parameters ---------- @@ -289,7 +289,7 @@ def solve_batch( inner_iter : int Number of inner Bregman iterations for the proximal method. Default is 1. inner_reg : float - Regularization parameter for the inner Bregman iterations in the proximal method. Default is 1e-1. + Regularization parameter for the inner Bregman iterations in the proximal method. Default is 1e-3. reg_type : str, optional Type of regularization :math:`R` either "KL", or "entropy". Default is "entropy". grad : str, optional @@ -336,7 +336,7 @@ def solve_batch( .. _references-batch-solver: Reference ---------- - .. [91] Xie, Y., Wang, X., Wang, R., & Zha, H. (2020, August). + .. [92] Xie, Y., Wang, X., Wang, R., & Zha, H. (2020, August). A fast proximal point method for computing exact wasserstein distance. In Uncertainty in artificial intelligence (pp. 433-453). PMLR. @@ -448,7 +448,7 @@ def solve_sample_batch( tol=1e-5, method="auto", inner_iter=1, - inner_reg=1e-2, + inner_reg=1e-3, reg_type="entropy", grad="envelope", ): @@ -482,7 +482,7 @@ def solve_sample_batch( inner_iter : int Number of inner Bregman iterations for the proximal method. Default is 1. inner_reg : float - Regularization parameter for the inner Bregman iterations in the proximal method. Default is 1e-2. + Regularization parameter for the inner Bregman iterations in the proximal method. Default is 1e-3. reg_type : str, optional Type of regularization :math:`R` either "KL", or "entropy". Default is "entropy". grad : str, optional diff --git a/ot/batch/_utils.py b/ot/batch/_utils.py index 54f035dcc..8bc96c3b9 100644 --- a/ot/batch/_utils.py +++ b/ot/batch/_utils.py @@ -307,14 +307,14 @@ def proximal_bregman_log_plan_batch( b=None, nx=None, reg=None, - inner_reg=1e-2, + inner_reg=1e-3, max_iter=10000, tol=1e-5, inner_iter=1, grad="detach", ): r""" - Returns optimal transport plans for a batch of cost matrices :math:`\mathbf{C}` using a Bregman divergence based proximal point method :ref:`[91] `. + Returns optimal transport plans for a batch of cost matrices :math:`\mathbf{C}` using a Bregman divergence based proximal point method :ref:`[92] `. This function solves the following optimization problem: @@ -355,7 +355,7 @@ def proximal_bregman_log_plan_batch( inner_iter : int, optional Number of inner iterations for updating the dual variables u and v. Default is 1. inner_reg : float, optional - Regularization parameter for the Bregman divergence. Default is 1e-1. + Regularization parameter for the Bregman divergence. Default is 1e-3. tol : float, optional Tolerance for convergence. The solver stops when the maximum change in the dual variables is below this value. @@ -387,7 +387,7 @@ def proximal_bregman_log_plan_batch( .. _references-OT-bregman-proximal-point: Reference ---------- - .. [91] Xie, Y., Wang, X., Wang, R., & Zha, H. (2020, August). + .. [92] Xie, Y., Wang, X., Wang, R., & Zha, H. (2020, August). A fast proximal point method for computing exact wasserstein distance. In Uncertainty in artificial intelligence (pp. 433-453). PMLR. """ From f7e82ac801dfddcf880f3ac8f7fe88b2d948bc4c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Flamary?= Date: Mon, 6 Jul 2026 17:38:03 +0200 Subject: [PATCH 13/13] Apply suggestions from code review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Rémi Flamary --- examples/backends/plot_ot_batch.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/backends/plot_ot_batch.py b/examples/backends/plot_ot_batch.py index b530ae9cf..54573c5a4 100644 --- a/examples/backends/plot_ot_batch.py +++ b/examples/backends/plot_ot_batch.py @@ -90,11 +90,11 @@ ## Naive approach results_values_list = [] for i in range(n_problems): - res = ot.solve(M_list[i], reg=None, max_iter=max_iter, tol=tol) + res = ot.solve(M_list[i], max_iter=max_iter, tol=tol) results_values_list.append(res.value_linear) ## Batched approach -results_batch = ot.solve_batch(M=M_batch, reg=None, max_iter=max_iter, tol=tol) +results_batch = ot.solve_batch(M=M_batch, max_iter=max_iter, tol=tol) results_values_batch = results_batch.value_linear exact_validated = np.allclose(