misalignment.py

"""Principal angles and leading-subspace utilities.

Conventions:
- Fisher matrices are real, symmetric, positive semi-definite (PSD).
- `top_r_eigenspace` returns top-r eigenvectors as columns of an
  orthonormal basis, ordered by descending eigenvalue.
- Principal angles use the SVD of U^T V definition; angles are returned
  in radians, ascending order (smallest angle first).

Edge cases (see prereg §12):
- If a Fisher has effective rank < r, top-r is degenerate. We return
  the available eigenvectors plus an explicit `effective_rank` so the
  caller can decide how to handle it.
"""

from __future__ import annotations

import numpy as np


PD_TOL = 1e-10
RANK_TOL = 1e-8  # eigenvalues below this (relative to max) are zero


def top_r_eigenspace(F: np.ndarray, r: int, *, rank_tol: float = RANK_TOL):
    """Top-r eigenvectors of symmetric matrix F, ordered by descending eigenvalue.

    Returns
    -------
    basis : (d, k) array, orthonormal columns
        k = min(r, effective_rank(F))
    eigenvalues : (k,) array
        Eigenvalues corresponding to the columns of basis, descending.
    effective_rank : int
        Number of eigenvalues that exceed rank_tol * max_eigenvalue.
    """
    F = np.asarray(F, dtype=float)
    if F.ndim != 2 or F.shape[0] != F.shape[1]:
        raise ValueError(f"F must be square 2D, got shape {F.shape}")
    # Symmetrize to remove tiny numerical asymmetry
    F_sym = 0.5 * (F + F.T)
    eigvals, eigvecs = np.linalg.eigh(F_sym)  # ascending
    # Reverse to descending
    eigvals = eigvals[::-1]
    eigvecs = eigvecs[:, ::-1]
    max_eig = eigvals[0] if eigvals.size > 0 else 0.0
    if max_eig <= 0:
        effective_rank = 0
    else:
        effective_rank = int(np.sum(eigvals > rank_tol * max_eig))
    k = min(r, effective_rank)
    basis = eigvecs[:, :k]
    eigs_kept = eigvals[:k]
    return basis, eigs_kept, effective_rank


def principal_angles(F1: np.ndarray, F2: np.ndarray, r: int):
    """Principal angles between top-r subspaces of two PSD matrices.

    Uses the Björck-Golub stable formulation: angle = arctan2(sin, cos)
    where cos comes from the SVD of U1^T U2 and sin comes from the
    norm of the residual (I - U1 U1^T) U2 V_r. This is numerically stable
    at both small and large angles (arccos of σ near 1 amplifies round-off
    by 1/sqrt(2(1-σ)), which arctan2 avoids).

    Reference: Björck & Golub, "Numerical methods for computing angles
    between linear subspaces", Math. Comp. 27 (1973).

    Returns
    -------
    angles : (r_eff,) ndarray of radians, ascending
        r_eff = min(r, effective_rank(F1), effective_rank(F2)).
        If r_eff < r the caller should treat the comparison as low-rank.
    info : dict
        Contains 'r_eff', 'effective_rank_1', 'effective_rank_2'.

    Raises
    ------
    ValueError if either matrix is rank 0 (no comparison possible).
    """
    U1, _, rank1 = top_r_eigenspace(F1, r)
    U2, _, rank2 = top_r_eigenspace(F2, r)
    r_eff = min(r, rank1, rank2)
    if r_eff == 0:
        raise ValueError("At least one input has effective rank 0")
    U1 = U1[:, :r_eff]
    U2 = U2[:, :r_eff]

    # Step 1: SVD of U1^T U2 -> cosines and right singular vectors V_r
    cos_M = U1.T @ U2
    _U_left, sigma, V_r_T = np.linalg.svd(cos_M, full_matrices=False)
    V_r = V_r_T.T  # columns are right singular vectors of cos_M

    # Step 2: residual sines via Y_perp = (I - U1 U1^T) (U2 V_r)
    Y = U2 @ V_r
    Y_perp = Y - U1 @ (U1.T @ Y)
    sin_vec = np.linalg.norm(Y_perp, axis=0)

    # Numerical safety: sigma can be slightly > 1 due to floating point
    cos_vec = np.clip(sigma, -1.0, 1.0)

    # Stable angle = arctan2(sin, cos). Always in [0, pi/2] given
    # sin >= 0 and cos in [-1, 1] (cos is non-negative here since the
    # SVD of U1^T U2 returns non-negative singular values).
    angles = np.arctan2(sin_vec, cos_vec)
    angles = np.sort(angles)

    info = {
        "r_eff": r_eff,
        "effective_rank_1": rank1,
        "effective_rank_2": rank2,
    }
    return angles, info


def max_principal_angle(F1: np.ndarray, F2: np.ndarray, r: int) -> tuple[float, dict]:
    """Convenience: maximum principal angle (the worst-case alignment loss)."""
    angles, info = principal_angles(F1, F2, r)
    return float(angles.max()), info


def is_positive_definite(F: np.ndarray, tol: float = PD_TOL) -> bool:
    """True if smallest eigenvalue of symmetric F exceeds tol."""
    F_sym = 0.5 * (F + F.T)
    eigvals = np.linalg.eigvalsh(F_sym)
    return bool(eigvals.min() > tol)