spd_learn.functional.bures_wasserstein_distance#

spd_learn.functional.bures_wasserstein_distance(A, B)[source]#

Compute the Bures-Wasserstein distance between SPD matrices.

The Bures-Wasserstein (BW) distance, also known as the Wasserstein-2 distance for centered Gaussian distributions, is defined as:

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

This metric is symmetric and satisfies the triangle inequality. It is equivalent to the optimal transport distance between centered Gaussian distributions \(\mathcal{N}(0, A)\) and \(\mathcal{N}(0, B)\) [Bhatia et al., 2019].

Parameters:

A (torch.Tensor) – SPD matrices of shape (…, n, n).
B (torch.Tensor) – SPD matrices of shape (…, n, n). Must be broadcastable with A.

Returns:

Bures-Wasserstein distances of shape (…).

Return type:

torch.Tensor

Examples

>>> import torch
>>> from spd_learn.functional.bures_wasserstein import bures_wasserstein_distance
>>> # Single pair of 3x3 SPD matrices
>>> A = torch.eye(3)
>>> B = 2 * torch.eye(3)
>>> d = bures_wasserstein_distance(A, B)
>>> print(f"Distance: {d.item():.4f}")
Distance: 0.8787

>>> # Batch of matrices
>>> A = torch.eye(3).unsqueeze(0).expand(10, 3, 3)
>>> B = torch.randn(10, 3, 3)
>>> B = B @ B.transpose(-1, -2) + 0.1 * torch.eye(3)  # Make SPD
>>> distances = bures_wasserstein_distance(A, B)
>>> print(f"Shape: {distances.shape}")
Shape: torch.Size([10])