Compute Procrustes#

procrustes #

FUNCTION DESCRIPTION
procrustes

Align X to reference using Procrustes analysis.

Functions#

procrustes(X: torch.Tensor, reference: torch.Tensor, scaling: bool = True, allow_reflection: bool = True, eps: float = torch.finfo(torch.float32).eps) -> torch.Tensor #

Align X to reference using Procrustes analysis.

Finds optimal rotation (and optionally scaling) to align X to reference embedding via orthogonal Procrustes problem.

PARAMETER DESCRIPTION
X

X embedding to align, shape (n_samples, n_components).

TYPE: Tensor

reference

reference embedding, shape (n_samples, n_components).

TYPE: Tensor

scaling

Whether to apply optimal scaling after rotation.

TYPE: bool, by default True DEFAULT: True

allow_reflection

Whether to allow reflections (det = -1). If False, forces proper rotations only (det = +1) by flipping the sign of one singular vector. Setting to True gives the optimal alignment but may flip axes. Setting to False preserves chirality/handedness but may have worse alignment.

TYPE: bool, by default True DEFAULT: True

eps

Regularization term.

TYPE: float, by default float = torch.finfo(torch.float32).eps DEFAULT: eps

RETURNS DESCRIPTION
Tensor

Aligned X embedding, shape (n_samples, n_components).
Notes

The Procrustes transformation minimizes the Frobenius norm: \(\|reference - scale \cdot X \cdot R\|_F\) where \(R\) is an orthogonal rotation matrix and scale is a scalar.

The algorithm:

  1. Center both embeddings (subtract means)
  2. Compute optimal rotation R via SVD: R = U @ V^T where reference^T @ X = U @ S @ V^T
  3. Optionally compute optimal scaling
  4. Apply transformation and re-center to reference mean

Examples:

>>> import torch
>>> from spectre.utils.align import procrustes
>>>
>>> # Two embeddings from different methods
>>> X = torch.randn(100, 3)
>>> reference = torch.randn(100, 3)
>>>
>>> # Align X to reference
>>> aligned = procrustes(X, reference, scaling=True)
>>> aligned.shape
torch.Size([100, 3])
Source code in spectre/compute/procrustes.py 
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
def procrustes(
    X: torch.Tensor,
    reference: torch.Tensor,
    scaling: bool = True,
    allow_reflection: bool = True,
    eps: float = torch.finfo(torch.float32).eps,
) -> torch.Tensor:
    """
    Align X to reference using Procrustes analysis.

    Finds optimal rotation (and optionally scaling) to align X to reference
    embedding via orthogonal Procrustes problem.

    Parameters
    ----------
    X : torch.Tensor
        X embedding to align, shape (n_samples, n_components).

    reference : torch.Tensor
        reference embedding, shape (n_samples, n_components).

    scaling : bool, by default True
        Whether to apply optimal scaling after rotation.

    allow_reflection : bool, by default True
        Whether to allow reflections (det = -1). If `False`, forces proper
        rotations only (det = +1) by flipping the sign of one singular vector.
        Setting to `True` gives the optimal alignment but may flip axes.
        Setting to `False` preserves chirality/handedness but may have worse
        alignment.

    eps : float, by default float = torch.finfo(torch.float32).eps
        Regularization term.

    Returns
    -------
    torch.Tensor
        Aligned X embedding, shape (n_samples, n_components).

    Notes
    -----
    The Procrustes transformation minimizes the Frobenius norm:
    $\\|reference - scale \\cdot X \\cdot R\\|_F$
    where $R$ is an orthogonal rotation matrix and scale is a scalar.

    The algorithm:

    1. Center both embeddings (subtract means)
    2. Compute optimal rotation R via SVD: R = U @ V^T where
       reference^T @ X = U @ S @ V^T
    3. Optionally compute optimal scaling
    4. Apply transformation and re-center to reference mean

    Examples
    --------
    >>> import torch
    >>> from spectre.utils.align import procrustes
    >>>
    >>> # Two embeddings from different methods
    >>> X = torch.randn(100, 3)
    >>> reference = torch.randn(100, 3)
    >>>
    >>> # Align X to reference
    >>> aligned = procrustes(X, reference, scaling=True)
    >>> aligned.shape
    torch.Size([100, 3])
    """
    check_same_shape(X, reference)
    check_same_dtype(X, reference)

    # Center both embeddings
    X_centered = X - X.mean(dim=0, keepdim=True)
    reference_centered = reference - reference.mean(dim=0, keepdim=True)

    # Compute optimal rotation via SVD
    # H = X.T @ reference gives us the cross-covariance matrix
    try:
        U, S, Vt = torch.linalg.svd(
            X_centered.T @ reference_centered,
            full_matrices=False,
        )
        R = U @ Vt

        # Ensure proper rotation (det = +1) by correcting reflections (det = -1)
        if not allow_reflection and torch.det(R) < 0:
            U[:, -1] *= -1
            R = U @ Vt
    except RuntimeError as e:
        import warnings

        warnings.warn(
            f"SVD failed during Procrustes alignment: {e}. Returning unaligned X.",
            RuntimeWarning,
        )
        return X

    # Apply rotation
    X_aligned = X_centered @ R

    # Compute optimal scaling if requested
    if scaling:
        numerator = torch.sum(reference_centered * X_aligned)
        denominator = torch.sum(X_aligned * X_aligned)

        if denominator > 1e-12:
            scale = numerator / denominator
        else:
            scale = 1.0

        X_aligned = scale * X_aligned

    # Add back reference mean
    X_aligned = X_aligned + reference.mean(dim=0, keepdim=True)

    return X_aligned