[WIP] Spectral-Grassmann OT by thibaut-germain · Pull Request #792 · PythonOT/POT

thibaut-germain · 2026-02-09T15:26:58Z

Types of changes

Adding sgot file in the ot folder.

Motivation and context / Related issue

Keep track of SGOT implementation in POT.

How has this been tested (if it applies)

Not tested yet.

PR checklist

I have read the CONTRIBUTING document.
The documentation is up-to-date with the changes I made (check build artifacts).
[] All tests passed, and additional code has been covered with new tests.
I have added the PR and Issue fix to the RELEASES.md file.

codecov · 2026-02-11T15:03:43Z

Codecov Report

❌ Patch coverage is 96.08541% with 11 lines in your changes missing coverage. Please review.
✅ Project coverage is 96.86%. Comparing base (41a4d57) to head (9b3cdc7).

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #792      +/-   ##
==========================================
- Coverage   96.87%   96.86%   -0.02%     
==========================================
  Files         113      115       +2     
  Lines       23062    23339     +277     
==========================================
+ Hits        22342    22608     +266     
- Misses        720      731      +11

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

rflamary

Hello @osheasienna and @thibaut-germain this is a nice first step.

Here are below a few comments that we can discuss together

rflamary · 2026-02-23T07:53:08Z

+    return C
+
+
+def metric(


Suggested change

def metric(

def sgot_metric(

rflamary · 2026-02-23T07:54:35Z

+    return prod ** (q / 2)
+
+
+def ot_plan(C, Ws=None, Wt=None, nx=None):


this function is not needed, this is two lines and the ormalization wrt ws and wt are not oK because it rcan retrun very weird things

rflamary · 2026-02-23T07:55:44Z

+### SPECTRAL-GRASSMANNIAN WASSERSTEIN METRIC ###
+#####################################################################################################################################
+#####################################################################################################################################
+def cost(


Suggested change

def cost(

def sgot_cost_matrix(

rflamary · 2026-02-23T07:56:51Z

+    imag_scale=1.0,
+    nx=None,
+):
+    """Compute the SGOT cost matrix between two spectral decompositions.


recall here the equation with eta and define with math teh different acceptable metrics

rflamary · 2026-02-23T07:57:42Z

+        raise ValueError(f"cost() expects Dt to be 1D (n,), got shape {Dt.shape}")
+    lam2 = Dt
+
+    lam1 = nx.astype(lam1, "complex128")


is that necessary? seems overkill to add a function to the backend for that . When and why does it fails?

rflamary · 2026-02-23T13:03:37Z

+    logits_s = rng.randn(r)
+    logits_t = rng.randn(r)
+
+    Ws = np.exp(logits_s)


simpler and return only positive values

Suggested change

Ws = np.exp(logits_s)

Ws = rng.rand(r)

rflamary · 2026-02-23T13:04:28Z

+    """Create test_cost for each trial: sweep over HPs and run cost()."""
+    grassmann_types = ["geodesic", "chordal", "procrustes", "martin"]
+    n_trials = 10
+    for _ in range(n_trials):


no need for trials ;)

rflamary · 2026-02-23T13:04:43Z

+def test_hyperparameter_sweep():
+    grassmann_types = ["geodesic", "chordal", "procrustes", "martin"]
+
+    for _ in range(10):


rflamary · 2026-02-23T13:05:48Z

 This new release adds support for sparse cost matrices and a new lazy EMD solver that computes distances on-the-fly from coordinates, reducing memory usage from O(n×m) to O(n+m). Both implementations are backend-agnostic and preserve gradient computation for automatic differentiation.

 #### New features 
 - Add lazy EMD solver with on-the-fly distance computation from coordinates (PR #788)


add feature here

rflamary · 2026-02-23T13:06:00Z

+## Upcomming 0.9.7.post1
+
+#### New features
+The next release will add cost functions between linear operators following [A Spectral-Grassmann Wasserstein metric for operator representations of dynamical systems](https://arxiv.org/pdf/2509.24920).


move this text to the new feature of 0.9.7.dev0 this is what we are working on. Also add a line in the Itemize with the PR number

rflamary

A few comments from talking together

rflamary · 2026-02-23T14:17:10Z

+    if grassman_metric == "procrustes":
+        return 2.0 * (1.0 - delta)
+    if grassman_metric == "martin":
+        return -nx.log(nx.clip(delta**2, eps, 1e300))


remove upper threshold

rflamary · 2026-02-23T14:27:26Z

+    C_grass = _grassmann_distance_squared(delta, grassman_metric=grassman_metric, nx=nx)
+
+    C2 = eta * C_lambda + (1.0 - eta) * C_grass
+    C = C2 ** (p / 2.0)


Suggested change

C = C2 ** (p / 2.0)

C = nx.real(C2) ** (p / 2.0)

rflamary · 2026-02-23T14:29:53Z

+    q=1,
+    r=2,
+    grassman_metric="chordal",
+    real_scale=1.0,


lets call this eigen_scaling and set it to None by default

rflamary · 2026-02-23T14:30:23Z

+    nx=None,
+):
+    """Compute the SGOT metric between two spectral decompositions.
+


add equation that illustrate p q and r

rflamary · 2026-02-23T14:32:36Z

+import numpy as np
+import pytest
+
+from ot.backend import get_backend


Suggested change

from ot.backend import get_backend

from ot.backend import get_backend, torch, jax

rflamary · 2026-02-23T14:35:15Z

+rng = np.random.RandomState(0)
+
+
+def rand_complex(shape):


Suggested change

def rand_complex(shape):

def rand_complex(shape,rng):

rflamary · 2026-02-23T14:35:35Z

+    return real + 1j * imag
+
+
+def random_atoms(d=8, r=4):


Suggested change

def random_atoms(d=8, r=4):

def random_atoms(d=8, r=4,seed=42):

rflamary · 2026-02-23T14:39:42Z

+
+
+@pytest.mark.parametrize("backend_name", ["numpy", "torch", "jax"])
+def test_cost_backend_consistency(backend_name):


Suggested change

def test_cost_backend_consistency(backend_name):

def test_cost_backend_consistency(nx):

rflamary · 2026-02-23T14:45:46Z

+# ---------------------------------------------------------------------
+
+
+def test_hyperparameter_sweep_cost(nx):


Suggested change

def test_hyperparameter_sweep_cost(nx):

def test_hyperparameter_sweep_cost(nx,grassmann_types,p,q,r,eta):

rflamary · 2026-02-25T16:33:22Z

+    Ws = Ws / nx.sum(Ws)
+    Wt = Wt / nx.sum(Wt)
+
+    P = ot.emd2(Ws, Wt, nx.real(C))


emd2 retruns directly obj no need to compute it again below

thibaut-germain · 2026-03-02T08:25:00Z

+        else:
+            real_scale, imag_scale = eigen_scaling[0], eigen_scaling[1]
+
+    Dsn = nx.real(Ds) * real_scale + 1j * nx.imag(Ds) * imag_scale


Suggested change

Dsn = nx.real(Ds) * real_scale + 1j * nx.imag(Ds) * imag_scale

C_real = nx.real(Dsn)[:,None] - nx.real(Dtn)[None,:]

C_real = C_real**2

C_imag = nx.imag(Dsn)[:,None] - nx.imag(Dtn)[None,:]

C_imag = C_imag**2

prod = C_real + C_imag

return prod ** (q / 2)

thibaut-germain · 2026-03-02T09:01:02Z

+    A_norm: array-like, shape (d, n)
+        Column-normalized array.
+    """
+    nrm = nx.sqrt(nx.sum(A * nx.conj(A), axis=0, keepdims=True))


Suggested change

nrm = nx.sqrt(nx.sum(A * nx.conj(A), axis=0, keepdims=True))

nrm = nx.norm(A, axis=0, keepdims=True)

You can replace it with the function nx.norm which manages the case of complex number

thibaut-germain · 2026-03-02T09:06:55Z

+    return delta
+
+
+def _grassmann_distance_squared(delta, grassman_metric="chordal", nx=None, eps=1e-300):


Epsilon is too small for the machine precision, you can set it to 1e-12 for instance.

thibaut-germain · 2026-03-02T09:08:49Z

+    if nx is None:
+        nx = get_backend(delta)
+
+    delta = nx.clip(delta, 0.0, 1.0)


If delta is not in [0,1] it should raise an error, this is an issue in the computation of delta outside of this function.

thibaut-germain · 2026-03-02T09:11:34Z

+### SPECTRAL-GRASSMANNIAN WASSERSTEIN METRIC ###
+#####################################################################################################################################
+#####################################################################################################################################
+def sgot_cost_matrix(


You should add eps in the definition of the function as this parameters appears in downstream functions. Keep the same epsilon for all functions.

thibaut-germain

Thibaut's review

thibaut-germain · 2026-03-04T08:14:38Z

+        # information-geometric interpretation in Germain et al. (2025).
+        delta2 = nx.maximum(delta**2, eps)
+        return -nx.log(delta2)
+    raise ValueError(f"Unknown grassman_metric: {grassman_metric}")


In this function the power q should also be a parameter:
for any distance you can set:
result = square_ditance(delta)
then
return nx.real(result)**(q/2)
Set by default q to the same value as for eigenvalue cost

thibaut-germain · 2026-03-04T08:14:49Z

+    C_lambda = eigenvalue_cost_matrix(Ds, Dt, q=q, eigen_scaling=eigen_scaling, nx=nx)
+
+    delta = _delta_matrix_1d(Rs, Ls, Rt, Lt, nx=nx)
+    C_grass = _grassmann_distance_squared(delta, grassman_metric=grassman_metric, nx=nx)


the power parameter q should also affect the Grassmann cost

thibaut-germain · 2026-03-04T08:16:43Z

+    C_grass = _grassmann_distance_squared(delta, grassman_metric=grassman_metric, nx=nx)
+
+    C2 = eta * C_lambda + (1.0 - eta) * C_grass
+    C = nx.real(C2) ** (p / 2.0)


you your cost function already return a real no need for nx.real here

thibaut-germain · 2026-03-04T08:28:18Z

+    return C
+
+
+def _validate_sgot_metric_inputs(Ds, Dt):


You can add verifications you wrote in line 272-290 in this function and also add verifications than source and target have the same shapes.

thibaut-germain · 2026-03-04T08:29:04Z

+        )
+
+
+def sgot_metric(


You will need to add eps also in this function.

thibaut-germain added 2 commits February 9, 2026 14:29

Sienna & Thibaut works

fd7e2b4

add sgot file

9b49c76

rflamary changed the title ~~Sgot~~ [WIP] Spactral-Gromov OT Feb 9, 2026

rflamary changed the title ~~[WIP] Spactral-Gromov OT~~ [WIP] Spectral-Grassman OT Feb 9, 2026

rflamary changed the title ~~[WIP] Spectral-Grassman OT~~ [WIP] Spectral-Grassmann OT Feb 9, 2026

Merge branch 'master' into sgot

5e775a5

osheasienna and others added 5 commits February 15, 2026 21:52

first draft of sgot.py

ba33032

Merge branch 'master' into sgot

2eaf221

Merge branch 'master' into sgot

1c73452

rewrite backend and refactor OT metric

3f10111

refactor sgot, and implement tests for sgot and backend

d5fef5b

github-actions bot added Tests ot.backend labels Feb 22, 2026

rflamary and others added 2 commits February 23, 2026 08:38

Merge branch 'master' into sgot

c61c398

fix astype in backend

630e359

rflamary reviewed Feb 23, 2026

View reviewed changes

Merge branch 'master' into sgot

6b4d17b

rflamary reviewed Feb 23, 2026

View reviewed changes

rflamary and others added 2 commits February 24, 2026 16:43

Merge branch 'master' into sgot

055c5e4

edits as per PR PythonOT#792

57c7a0c

rflamary reviewed Feb 25, 2026

View reviewed changes

osheasienna added 3 commits February 27, 2026 14:03

cost & metric fixed in test_sgot

a0e74be

correct issues on test_sgot

e994303

fixing test failures

e9f8be7

thibaut-germain commented Mar 2, 2026

View reviewed changes

thibaut-germain commented Mar 4, 2026

View reviewed changes

sgot documentation

37aa721

github-actions bot added the documentation label Mar 14, 2026

rflamary and others added 2 commits March 19, 2026 11:27

Merge branch 'master' into sgot

c0c7acf

fix backend error and example notebook

7641f7f

github-actions bot added the Examples label Mar 31, 2026

osheasienna added 7 commits March 31, 2026 10:28

removed notebook, added python version

b8a6cbe

fix SGOT example and relax SGOT test tolerances

0b90ab7

gaussian metric on graphs

9cffb45

updated rotation paragraph across graphs

dae7b75

equation colours

579cce4

fixed sgot distance vs frequency (omega)

2882f32

fix RELEASES

9b3cdc7

		return prod ** (q / 2)


		def ot_plan(C, Ws=None, Wt=None, nx=None):

	from ot.backend import get_backend
	from ot.backend import get_backend, torch, jax

	def random_atoms(d=8, r=4):
	def random_atoms(d=8, r=4,seed=42):



		@pytest.mark.parametrize("backend_name", ["numpy", "torch", "jax"])
		def test_cost_backend_consistency(backend_name):

	def test_cost_backend_consistency(backend_name):
	def test_cost_backend_consistency(nx):

		# ---------------------------------------------------------------------


		def test_hyperparameter_sweep_cost(nx):

	def test_hyperparameter_sweep_cost(nx):
	def test_hyperparameter_sweep_cost(nx,grassmann_types,p,q,r,eta):

-    Dsn = nx.real(Ds) * real_scale + 1j * nx.imag(Ds) * imag_scale
+   C_real = nx.real(Dsn)[:,None] - nx.real(Dtn)[None,:]
+   C_real = C_real**2
+   C_imag = nx.imag(Dsn)[:,None] - nx.imag(Dtn)[None,:]
+   C_imag = C_imag**2
+   prod = C_real + C_imag
+   return prod ** (q / 2)

	nrm = nx.sqrt(nx.sum(A * nx.conj(A), axis=0, keepdims=True))
	nrm = nx.norm(A, axis=0, keepdims=True)

		return delta


		def _grassmann_distance_squared(delta, grassman_metric="chordal", nx=None, eps=1e-300):

Conversation

thibaut-germain commented Feb 9, 2026

Types of changes

Motivation and context / Related issue

How has this been tested (if it applies)

PR checklist

Uh oh!

codecov bot commented Feb 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

rflamary left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rflamary left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

thibaut-germain left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov bot commented Feb 11, 2026 •

edited

Loading