nulls.py

"""Random-rotation null distribution for pairwise principal angles.

Per prereg §12 E6: the null is "what would the pairwise angles look like
if each probe's leading r-subspace were oriented uniformly at random in
R^d, independent of the others?"

We sample N independent Haar-random orthonormal r-frames (matrices U_i
of shape (d, r) with orthonormal columns) and compute the pairwise max
principal angles among them. Repeat n_draws times to get a null
distribution.

This is equivalent to applying independent Haar-random orthogonal
rotations Q_i to fixed (d, r) bases — only the relative orientation
matters for principal angles.
"""

from __future__ import annotations

import numpy as np
from scipy import stats


def haar_orthonormal_frame(d: int, r: int, rng: np.random.Generator) -> np.ndarray:
    """Sample an orthonormal (d, r) matrix uniformly from the Stiefel manifold V_r(R^d)."""
    if r > d:
        raise ValueError(f"r ({r}) cannot exceed d ({d})")
    A = rng.standard_normal((d, r))
    Q, R = np.linalg.qr(A)
    # Sign correction so that the QR decomposition is uniform on Stiefel
    # (see Mezzadri 2007; without this Q is biased).
    s = np.sign(np.diag(R))
    s[s == 0] = 1.0
    Q = Q * s
    return Q


def _principal_angles_between_frames(U1: np.ndarray, U2: np.ndarray) -> np.ndarray:
    """Stable principal angles between two orthonormal (d, r) frames.

    Same Björck-Golub formulation as misalignment.principal_angles, kept
    here so the null draws use exactly the same computation.
    """
    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
    Y = U2 @ V_r
    Y_perp = Y - U1 @ (U1.T @ Y)
    sin_vec = np.linalg.norm(Y_perp, axis=0)
    cos_vec = np.clip(sigma, -1.0, 1.0)
    return np.arctan2(sin_vec, cos_vec)


def _pair_max_angles_from_frames(frames: list[np.ndarray]) -> np.ndarray:
    """Pairwise max principal angles for a list of orthonormal frames."""
    N = len(frames)
    out = []
    for i in range(N):
        for j in range(i + 1, N):
            ang = _principal_angles_between_frames(frames[i], frames[j])
            out.append(float(ang.max()))
    return np.asarray(out)


def random_rotation_null(
    d: int,
    r: int,
    n_probes: int,
    *,
    n_draws: int = 1000,
    rng: np.random.Generator | None = None,
) -> np.ndarray:
    """Return an array of pairwise max principal angles under the random-orientation null.

    Shape: (n_draws * C(n_probes, 2),). Flattened across all draws.
    """
    if rng is None:
        rng = np.random.default_rng(0)
    samples = []
    for _ in range(n_draws):
        frames = [haar_orthonormal_frame(d, r, rng) for _ in range(n_probes)]
        samples.append(_pair_max_angles_from_frames(frames))
    return np.concatenate(samples)


def compare_to_null(
    actual_angles: np.ndarray,
    d: int,
    r: int,
    n_probes: int,
    *,
    n_draws: int = 1000,
    rng: np.random.Generator | None = None,
    alpha: float = 0.05,
) -> dict:
    """Compare actual pairwise max angles to the random-orientation null via KS test.

    Returns a dict with the KS statistic, p-value, and a `differs_from_null`
    boolean for `alpha`.
    """
    actual_angles = np.asarray(actual_angles, dtype=float)
    actual_angles = actual_angles[np.isfinite(actual_angles)]
    if actual_angles.size == 0:
        return {
            "ks_stat": float("nan"),
            "p_value": float("nan"),
            "differs_from_null": False,
            "alpha": alpha,
            "n_actual": 0,
            "n_null": 0,
            "actual_median_rad": float("nan"),
            "null_median_rad": float("nan"),
        }
    null_samples = random_rotation_null(d, r, n_probes, n_draws=n_draws, rng=rng)
    ks = stats.ks_2samp(actual_angles, null_samples, alternative="two-sided")
    return {
        "ks_stat": float(ks.statistic),
        "p_value": float(ks.pvalue),
        "differs_from_null": bool(ks.pvalue < alpha),
        "alpha": alpha,
        "n_actual": int(actual_angles.size),
        "n_null": int(null_samples.size),
        "actual_median_rad": float(np.median(actual_angles)),
        "null_median_rad": float(np.median(null_samples)),
    }