spd_learn.functional.distance_wasserstein#

spd_learn.functional.distance_wasserstein(A, B, squared=False)[source]#

Wasserstein distance between SPSD/HPSD matrices.

The Wasserstein distance between two SPSD/HPSD matrices \(\mathbf{A}\) and \(\mathbf{B}\) is [1] [2]:

\[d(\mathbf{A},\mathbf{B}) = \sqrt{ \text{tr} \left(\mathbf{A} + \mathbf{B} - 2(\mathbf{B}^{1/2} \mathbf{A} \mathbf{B}^{1/2})^{1/2} \right) }\]

Parameters:

A (ndarray, shape (..., n, n)) – First SPSD/HPSD matrices.
B (ndarray, shape (..., n, n)) – Second SPSD/HPSD matrices, same dimensions as A.
squared (bool, default=False) –
Return squared distance.

Added in version 0.5.

Returns:

d – Wasserstein distance between A and B.

Return type:

float or ndarray, shape (…,)

Notes

Changed in version 0.12: Add support for NumPy and PyTorch.