spd_learn.functional.bures_wasserstein_mean#

spd_learn.functional.bures_wasserstein_mean(matrices, weights=None, max_iter=50, tol=1e-08, init=None, return_info=False)[source]#

Compute the Bures-Wasserstein barycenter of SPD matrices.

The Bures-Wasserstein barycenter (also known as the Fréchet mean under the BW metric) of a set of SPD matrices \(\{A_1, \ldots, A_K\}\) with weights \(\{w_1, \ldots, w_K\}\) is defined as:

\[\bar{A} = \arg\min_{M \succ 0} \sum_{i=1}^{K} w_i \, d_{BW}^2(M, A_i)\]

The barycenter is computed using the fixed-point iteration:

\[M_{k+1} = \sum_{i=1}^{K} w_i \left(M_k^{1/2} A_i M_k^{1/2}\right)^{1/2}\]

which converges to the unique barycenter.

Parameters:

matrices (torch.Tensor) – SPD matrices of shape (K, …, n, n) where K is the number of matrices to average.
weights (torch.Tensor, optional) – Non-negative weights of shape (K,). If None, uniform weights are used. Weights are automatically normalized to sum to 1.
max_iter (int, default=50) – Maximum number of fixed-point iterations.
tol (float, default=1e-8) – Convergence tolerance. Iteration stops when the relative change in the Frobenius norm is below this threshold.
init (torch.Tensor, optional) – Initial estimate of the barycenter of shape (…, n, n). If None, the arithmetic mean is used as initialization.
return_info (bool, default=False) – If True, return a dictionary with convergence information.

Returns:

barycenter (torch.Tensor) – The Bures-Wasserstein barycenter of shape (…, n, n).
info (dict, optional) – Convergence information (only if return_info=True): - "n_iter": Number of iterations performed - "converged": Whether the algorithm converged - "relative_change": Final relative change in Frobenius norm

Examples

>>> import torch
>>> from spd_learn.functional.bures_wasserstein import bures_wasserstein_mean
>>> # Create 5 random SPD matrices
>>> K, n = 5, 3
>>> matrices = torch.randn(K, n, n)
>>> matrices = matrices @ matrices.transpose(-1, -2) + 0.1 * torch.eye(n)
>>> # Compute barycenter
>>> mean = bures_wasserstein_mean(matrices)
>>> print(f"Barycenter shape: {mean.shape}")
Barycenter shape: torch.Size([3, 3])

>>> # With custom weights
>>> weights = torch.tensor([0.5, 0.2, 0.1, 0.1, 0.1])
>>> mean = bures_wasserstein_mean(matrices, weights=weights)

>>> # With convergence info
>>> mean, info = bures_wasserstein_mean(matrices, return_info=True)
>>> print(f"Converged in {info['n_iter']} iterations")

Notes

The fixed-point iteration [Agueh and Carlier, 2011] is guaranteed to converge for any initialization. However, the convergence rate depends on the condition numbers of the input matrices. For ill-conditioned matrices, more iterations may be needed [Álvarez-Esteban et al., 2016].