code/metrics.py

"""Geometric metrics and Jacobian Effective Rank (Appendix I)."""

from __future__ import annotations

from typing import Callable

import numpy as np
import torch
import torch.nn.functional as F


def participation_ratio(eigenvalues: np.ndarray, normalize_by_d: bool = False) -> float:
    lam = np.asarray(eigenvalues, dtype=np.float64)
    lam = lam[lam > 0]
    if lam.size == 0:
        return 0.0
    pr = (lam.sum() ** 2) / (lam ** 2).sum()
    if normalize_by_d:
        pr = pr / float(lam.size)
    return float(pr)


def isotropy_score(eigenvalues: np.ndarray) -> float:
    lam = np.asarray(eigenvalues, dtype=np.float64)
    total = lam.sum()
    if total <= 0:
        return 0.0
    return float(1.0 - lam.max() / total)


def covariance_eigs(embeddings: np.ndarray) -> np.ndarray:
    x = embeddings.astype(np.float64)
    x = x - x.mean(axis=0, keepdims=True)
    n = x.shape[0]
    cov = (x.T @ x) / max(n - 1, 1)
    eigs = np.linalg.eigvalsh(cov)
    return np.sort(eigs)[::-1]


def global_geometry(embeddings: np.ndarray) -> dict[str, float]:
    eigs = covariance_eigs(embeddings)
    d = embeddings.shape[1]
    return {
        "G.PR": participation_ratio(eigs, normalize_by_d=True),
        "G.Iso": isotropy_score(eigs),
        "dim": float(d),
        "n": float(embeddings.shape[0]),
    }


def local_isotropy(embeddings: np.ndarray, k: int = 32, n_anchors: int = 500, seed: int = 42) -> float:
    rng = np.random.default_rng(seed)
    x = embeddings.astype(np.float64)
    n = x.shape[0]
    n_anchors = min(n_anchors, n)
    anchors = rng.choice(n, size=n_anchors, replace=False)
    # Cosine neighborhood in embedding space
    xn = x / (np.linalg.norm(x, axis=1, keepdims=True) + 1e-12)
    scores = []
    for idx in anchors:
        sims = xn @ xn[idx]
        sims[idx] = -np.inf
        nn = np.argpartition(-sims, kth=min(k, n - 2))[:k]
        local = x[nn]
        eigs = covariance_eigs(local)
        scores.append(isotropy_score(eigs))
    return float(np.mean(scores))


def _orthonormal_probes(n_params: int, k: int, device: torch.device, seed: int) -> torch.Tensor:
    g = torch.Generator(device="cpu")
    g.manual_seed(seed)
    m = torch.randn(n_params, k, generator=g, dtype=torch.float32)
    q, _ = torch.linalg.qr(m, mode="reduced")
    return q.to(device)


@torch.no_grad()
def jacobian_effective_rank_noise(
    encode_fn: Callable[[torch.Tensor], torch.Tensor],
    *,
    n_images: int = 100,
    k: int = 32,
    image_size: int = 224,
    device: torch.device | str = "cuda",
    seed: int = 42,
) -> float:
    """JER on Gaussian-noise probe (Appendix E / Table 14)."""
    device = torch.device(device)
    rng = np.random.default_rng(seed)
    jer_vals = []
    n_in = 3 * image_size * image_size
    probes = _orthonormal_probes(n_in, k, device, seed)

    for i in range(n_images):
        noise = rng.normal(0.45, 0.225, size=(3, image_size, image_size)).astype(np.float32)
        noise = np.clip(noise, 0.0, 1.0)
        # ImageNet normalize
        mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)[:, None, None]
        std = np.array([0.229, 0.224, 0.225], dtype=np.float32)[:, None, None]
        x_np = (noise - mean) / std
        x = torch.from_numpy(x_np).unsqueeze(0).to(device)
        x = x.detach().requires_grad_(True)

        # Build Y = J @ V via JVPs. Force MATH SDPA so ViT attention is differentiable.
        cols = []
        for j in range(k):
            v = probes[:, j].reshape(1, 3, image_size, image_size)

            def f(inp: torch.Tensor) -> torch.Tensor:
                return encode_fn(inp)

            with torch.enable_grad():
                try:
                    from torch.nn.attention import SDPBackend, sdpa_kernel

                    with sdpa_kernel(SDPBackend.MATH):
                        _, jv = torch.autograd.functional.jvp(f, (x,), (v,), create_graph=False)
                except Exception:
                    try:
                        with torch.backends.cuda.sdp_kernel(
                            enable_flash=False, enable_math=True, enable_mem_efficient=False
                        ):
                            _, jv = torch.autograd.functional.jvp(f, (x,), (v,), create_graph=False)
                    except Exception:
                        _, jv = torch.autograd.functional.jvp(f, (x,), (v,), create_graph=False)
            cols.append(jv.reshape(-1).detach())
        y = torch.stack(cols, dim=1)  # D x k
        # Singular values of Y approximate projected spectrum of J
        s = torch.linalg.svdvals(y).detach().float().cpu().numpy()
        jer_vals.append(participation_ratio(s, normalize_by_d=False))
    return float(np.mean(jer_vals))


def cosine_sim(a: np.ndarray, b: np.ndarray) -> float:
    a = a.astype(np.float64)
    b = b.astype(np.float64)
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-12))


def evaluate_binding(encode_images: Callable[[list], np.ndarray], trials) -> float:
    correct = 0
    for t in trials:
        imgs = [t.query] + t.candidates
        embs = encode_images(imgs)
        q = embs[0]
        sims = [cosine_sim(q, embs[i + 1]) for i in range(len(t.candidates))]
        pred = int(np.argmax(sims))
        correct += int(pred == t.target_index)
    return correct / len(trials)


def evaluate_samediff(encode_images: Callable[[list], np.ndarray], trials) -> float:
    dists = []
    labels = []
    for t in trials:
        embs = encode_images([t.image_a, t.image_b])
        d = 1.0 - cosine_sim(embs[0], embs[1])
        dists.append(d)
        labels.append(1 if t.same else 0)
    dists = np.asarray(dists)
    labels = np.asarray(labels)
    best = 0.0
    for p in range(0, 101, 5):
        tau = np.percentile(dists, p)
        pred = (dists < tau).astype(int)
        acc = (pred == labels).mean()
        best = max(best, float(acc))
    return best