Compute Cca#

cca #

FUNCTION DESCRIPTION
cca

Compute CCA canonical correlations between X and Y, canonical directions,

Functions#

cca(X: torch.Tensor, Y: torch.Tensor, n_eigen: int | None = None, correlations_only: bool = False, eps: float = torch.finfo(torch.float32).eps) -> tuple[torch.Tensor, ...] #

Compute CCA canonical correlations between X and Y, canonical directions, and the projections of X and Y onto their canonical directions.

PARAMETER DESCRIPTION
X

Input tensor of shape (n_samples, x_features).

TYPE: Tensor

Y

Input tensor of shape (n_samples, y_features).

TYPE: Tensor

n_eigen

Number of components. If None, uses min(x_features, y_features).

TYPE: int | None DEFAULT: None

correlations_only

Whether to compute canonical correlation between X and Y only. If False performs full CCA.

TYPE: bool, by default False DEFAULT: False

eps

Regularization term.

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

RETURNS DESCRIPTION
tuple[Tensor, ...]
  • canonical_correlations: canonical correlation coefficients
  • canonical_directions_x: canonical directions for X
  • canonical_projections_x: X projected onto canonical directions
  • canonical_directions_y: canonical directions for Y (for correlations_only=False)
  • canonical_projections_y: Y projected onto canonical directions (for correlations_only=False)
Source code in spectre/compute/cca.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
def cca(
    X: torch.Tensor,
    Y: torch.Tensor,
    n_eigen: int | None = None,
    correlations_only: bool = False,
    eps: float = torch.finfo(torch.float32).eps,
) -> tuple[torch.Tensor, ...]:
    """
    Compute CCA canonical correlations between X and Y, canonical directions,
    and the projections of X and Y onto their canonical directions.

    Parameters
    ----------
    X : torch.Tensor
        Input tensor of shape (n_samples, x_features).

    Y : torch.Tensor
        Input tensor of shape (n_samples, y_features).

    n_eigen : int | None
        Number of components. If None, uses min(x_features, y_features).

    correlations_only : bool, by default False
        Whether to compute canonical correlation between X and Y only. If False
        performs full CCA.

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

    Returns
    -------
    tuple[torch.Tensor, ...]

        - canonical_correlations: canonical correlation coefficients
        - canonical_directions_x: canonical directions for X
        - canonical_projections_x: X projected onto canonical directions
        - canonical_directions_y: canonical directions for Y
          (for `correlations_only=False`)
        - canonical_projections_y: Y projected onto canonical directions
          (for `correlations_only=False`)
    """
    n_samples = X.shape[0]
    m_features = min(X.shape[1], Y.shape[1])

    if n_eigen is None:
        n_eigen = m_features
    elif n_eigen > m_features:
        n_eigen = m_features
    elif n_eigen < 1:
        raise ValueError(f"n_eigen must be at least 1, got {n_eigen}.")

    X_centered = X - X.mean(dim=0, keepdim=True)
    Y_centered = Y - Y.mean(dim=0, keepdim=True)

    # Compute covariance matrices with regularization
    Cxx = (X_centered.T @ X_centered) / n_samples
    Cyy = (Y_centered.T @ Y_centered) / n_samples
    Cxy = (X_centered.T @ Y_centered) / n_samples

    Cxx = Cxx + eps * torch.eye(X.shape[1], device=X.device, dtype=X.dtype)
    Cyy = Cyy + eps * torch.eye(Y.shape[1], device=Y.device, dtype=Y.dtype)

    # Solve generalized eigenvalue problem
    eigenvalues, canonical_directions_x = gen_eig_chol(
        Cxy @ torch.linalg.solve(Cyy, Cxy.T),
        Cxx,
        n_eigen=n_eigen,
        sort_descending=True,
    )

    canonical_correlations = torch.sqrt(torch.clamp(eigenvalues, 0.0, 1.0))

    # Canonical projections only for X
    canonical_projections_x = X_centered @ canonical_directions_x

    if correlations_only:
        return (
            canonical_correlations,
            canonical_directions_x,
            canonical_projections_x,
        )
    else:
        # Y's canonical directions (from the dual problem)
        canonical_directions_y = torch.linalg.solve(Cyy, Cxy.T @ canonical_directions_x)

        # Normalize each canonical direction: v_i^T Cyy v_i = 1
        norms = torch.sqrt(
            (canonical_directions_y * (Cyy @ canonical_directions_y)).sum(dim=0) + eps
        )
        canonical_directions_y = canonical_directions_y / norms.unsqueeze(0)

        canonical_projections_y = Y_centered @ canonical_directions_y

        return (
            canonical_correlations,
            canonical_directions_x,
            canonical_projections_x,
            canonical_directions_y,
            canonical_projections_y,
        )