processed/covtoken_code/coverage/energy.py

"""Energy-based lesion-subspace coverage (additive alternative to effective rank).

Motivation (Gate 2 / Gate 4 root cause): the RankMe / coding-rate coverage is an AGGREGATE
over all tokens whose value barely moves when a few small-lesion tokens are added or removed.
An ENERGY coverage is ADDITIVE in tokens, so high-lesion-energy tokens contribute in
proportion to their lesion content:

    C_E(S; x) = sum_{i in S} || P_L z_i ||^2      (total lesion-subspace energy retained)

Removing a lesion token (high ||P_L z||) drops C_E a lot, so the coverage DROP tracks lesion
loss directly. Label-free, differentiable, no SVD. C*_E(x) = C_E({1..n}; x).
"""
from __future__ import annotations

import torch


def energy_coverage(Z_retained: torch.Tensor, P_L: torch.Tensor | None = None) -> torch.Tensor:
    """C_E(S;x): total lesion-subspace energy of retained tokens (scalar)."""
    Z = Z_retained
    if Z.ndim != 2 or Z.shape[0] == 0:
        return torch.zeros((), dtype=Z.dtype, device=Z.device)
    PZ = Z @ P_L.T if P_L is not None else Z
    return PZ.pow(2).sum()


def energy_coverage_drop(Z_full: torch.Tensor, Z_retained: torch.Tensor,
                         P_L: torch.Tensor | None = None) -> torch.Tensor:
    return energy_coverage(Z_full, P_L) - energy_coverage(Z_retained, P_L)