[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 7.17131% with 233 lines in your changes missing coverage. Please review.
✅ Project coverage is 95.77%. Comparing base (e164e78) to head (3f10111).
⚠️ Report is 2 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #792      +/-   ##
==========================================
- Coverage   96.77%   95.77%   -1.00%     
==========================================
  Files         107      108       +1     
  Lines       22342    22621     +279     
==========================================
+ Hits        21622    21666      +44     
- Misses        720      955     +235

🚀 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

ot/sgot.py

+    return C
+
+
+def metric(


Suggested change

def metric(

def sgot_metric(

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

ot/sgot.py

+    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

ot/sgot.py

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


Suggested change

def cost(

def sgot_cost_matrix(

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

ot/sgot.py

+    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

ot/sgot.py

+        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

test/test_sgot.py

+    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

test/test_sgot.py

+    """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

test/test_sgot.py

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


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

RELEASES.md

 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

RELEASES.md

+## 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

ot/sgot.py

+    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

ot/sgot.py

+    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

ot/sgot.py

+    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

ot/sgot.py

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


add equation that illustrate p q and r

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

test/test_sgot.py

+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

test/test_sgot.py

+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

test/test_sgot.py

+    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

test/test_sgot.py

+
+
+@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

test/test_sgot.py

+# ---------------------------------------------------------------------
+
+
+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

ot/sgot.py

+    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

ot/sgot.py

+        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

ot/sgot.py

+    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

ot/sgot.py

+    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

ot/sgot.py

+    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

ot/sgot.py

+### 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

ot/sgot.py

+        # 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

ot/sgot.py

+    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

ot/sgot.py

+    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

ot/sgot.py

+    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

ot/sgot.py

+        )
+
+
+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

		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