[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 (6711f1f).

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

A few comments from talking together

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

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.

@osheasienna can you handle this?

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.

rflamary

Looking good, just a few more things to do

rflamary · 2026-05-13T09:51:00Z

+representation varies under rotation. The SGOT cost and metric are used to
+compare the reference and rotated systems.
+
+[1] 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.


Add the reference at the end of the readme file and use the same reference number here

rflamary · 2026-05-13T11:55:05Z

+# Comparison across Grassmannian metrics for SGOT distance versus rotation angle
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+thetas = np.linspace(0, np.pi / 2, 10)


can we have something more detailed such as

Suggested change

thetas = np.linspace(0, np.pi / 2, 10)

thetas = np.linspace(0, np.pi / 2, 50)

rflamary · 2026-05-13T11:55:56Z

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


@osheasienna can you handle this?

rflamary · 2026-05-13T11:58:34Z

+    - :math:`p` is the exponent used in the OT ground cost and the inner
+      Wasserstein root,
+    - :math:`r` is an additional outer root applied to the Wasserstein objective.
+    """


please ad the references with all authors and number form readme here

rflamary · 2026-05-13T11:58:59Z

+
+    References
+    ----------
+    Germain et al., *Spectral-Grassmann Optimal Transport* (SGOT).


detail reference here (with all informations

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

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

Comment thread ot/sgot.py Outdated

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

Comment thread ot/sgot.py

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 8 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

fix typos

6711f1f

rflamary reviewed May 13, 2026

View reviewed changes

osheasienna and others added 4 commits May 14, 2026 18:35

Apply suggestion from @rflamary

782096f

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

Apply suggestion from @rflamary

39f4dde

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

SGOT review comments

1f5e4cc

ci: uninstall triton before pytest

87de0e7

github-actions Bot added the CI label May 15, 2026



		@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):

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

	thetas = np.linspace(0, np.pi / 2, 10)
	thetas = np.linspace(0, np.pi / 2, 50)

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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

rflamary left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

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!

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!

Uh oh!

rflamary left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

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!

Reviewers

Assignees

Labels

codecov Bot commented Feb 11, 2026 •

edited

Loading