spd_learn.functional.frechet_mean#

spd_learn.functional.frechet_mean(X: Tensor, max_iter: int = 1, weights: Tensor | None = None, return_distances: bool = False) → Tensor | Tuple[Tensor, Tensor][source]#

Fréchet mean of SPD matrices under the AIRM via Karcher flow.

Computes the minimizer of the sum of squared geodesic distances:

\[\bar{X} = \arg\min_{G \in \mathcal{S}_{++}^n} \sum_{i=1}^{N} w_i \, d_{\text{AIRM}}^2(G, X_i)\]

using iterative Karcher flow initialized from the (weighted) Euclidean mean.

Parameters:

X (torch.Tensor) – Batch of SPD matrices with shape (batch_size, ..., n, n).
max_iter (int, default=1) – Number of Karcher flow iterations. A single iteration is often sufficient for batch normalization; use more (e.g. 50) when a high-accuracy mean is needed.
weights (torch.Tensor, optional) – Per-sample weights with shape broadcastable to X. When None, uniform weights 1/N are used.
return_distances (bool, default=False) – If True, also returns the geodesic distances from each sample to the mean.

Returns:

mean (torch.Tensor) – Fréchet mean with shape (1, ..., n, n).
distances (torch.Tensor) – Only returned when return_distances=True. Geodesic distances from each sample to the mean, with shape (batch_size, ...).