spd_learn.functional.bures_wasserstein_transport#

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

Compute the optimal transport map from A to B applied to X.

Given SPD matrices A and B, the optimal transport map T that pushes the Gaussian \(\mathcal{N}(0, A)\) to \(\mathcal{N}(0, B)\) is:

\[T_{A \to B} = A^{-1/2} (A^{1/2} B A^{1/2})^{1/2} A^{-1/2}\]

When applied to a covariance matrix X (representing a Gaussian with covariance X), the transported covariance is:

\[T_{A \to B}(X) = T_{A \to B} \, X \, T_{A \to B}^T\]

This is useful for domain adaptation and covariance alignment tasks.

Parameters:

A (torch.Tensor) – Source SPD matrices of shape (…, n, n).
B (torch.Tensor) – Target SPD matrices of shape (…, n, n).
X (torch.Tensor) – SPD matrices to transport of shape (…, n, n).

Returns:

Transported SPD matrices of shape (…, n, n).

Return type:

torch.Tensor

Examples

>>> import torch
>>> from spd_learn.functional.bures_wasserstein import bures_wasserstein_transport
>>> # Transport A to B (should give B)
>>> A = torch.eye(3)
>>> B = 2 * torch.eye(3)
>>> transported = bures_wasserstein_transport(A, B, A)
>>> torch.allclose(transported, B, atol=1e-6)
True

Notes

When X = A, the transport map gives exactly B. This property can be used to verify the correctness of the implementation [Givens and Shortt, 1984].