spd_learn.functional.bures_wasserstein_geodesic#

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

Compute the geodesic interpolation under the Bures-Wasserstein metric.

The geodesic between two SPD matrices \(A\) and \(B\) under the Bures-Wasserstein metric is given by [Malagò et al., 2018]:

\[\gamma(t) = (1-t)^2 A + t^2 B + t(1-t) \left((AB)^{1/2} + (BA)^{1/2}\right)\]

where \((AB)^{1/2} = A^{1/2} (A^{1/2} B A^{1/2})^{1/2} A^{-1/2}\).

Parameters:

A (torch.Tensor) – Starting point SPD matrices of shape (…, n, n).
B (torch.Tensor) – End point SPD matrices of shape (…, n, n).
t (float or torch.Tensor) – Interpolation parameter(s). When t = 0, returns A. When t = 1, returns B. Can be a scalar or tensor broadcastable with the batch dimensions.

Returns:

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

Return type:

torch.Tensor

Examples

>>> import torch
>>> from spd_learn.functional.bures_wasserstein import bures_wasserstein_geodesic
>>> A = torch.eye(3)
>>> B = 2 * torch.eye(3)
>>> # Midpoint
>>> C = bures_wasserstein_geodesic(A, B, 0.5)
>>> # Endpoints
>>> torch.allclose(bures_wasserstein_geodesic(A, B, 0.0), A)
True
>>> torch.allclose(bures_wasserstein_geodesic(A, B, 1.0), B)
True

Notes

The cross-term \((AB)^{1/2}\) is computed as \(A^{1/2} M A^{-1/2}\) where \(M = (A^{1/2} B A^{1/2})^{1/2}\), since \((A^{1/2} M A^{-1/2})^2 = A^{1/2} M^2 A^{-1/2} = AB\).