processed/covtoken_code/arch/volumetric.py

"""Phase 6 — volumetric two-level economy (slice-level + token-level), inference-time.

The edge-deployment payoff. For a 3D CT volume:
  1. SLICE level: a cheap SHALLOW pass (first L_route blocks) scores every slice by lesion-
     subspace coverage (top-k token membership). Only the top-S fraction of slices -- the
     lesion-bearing ones -- get the full deep pass.
  2. TOKEN level: within kept slices, route tokens by coverage (routed_depth) at fraction f.

Compute model (block-token units):
  dense      = N * L_total
  two_level  = N * L_route                       (shallow scoring, ALL slices)
             + S*N * (L_total - L_route) * f      (deep pass, kept slices, routed tokens)
  reduction  = dense / two_level

Volume-level lesion sensitivity = fraction of total lesion mass (lesion patches summed over
the whole volume) that survives BOTH selections (slice kept AND token in the deep set).
"""
from __future__ import annotations

import numpy as np


def slice_score(token_membership: np.ndarray, topk: int = 8) -> float:
    """Slice lesion-presence score = mean of the top-k token coverage memberships."""
    s = np.sort(token_membership)[::-1]
    return float(s[:topk].mean())


def two_level_reduction(S: float, f: float, L_route: int, L_total: int = 12) -> float:
    dense = L_total
    two = L_route + S * (L_total - L_route) * f
    return float(dense / two)


def select_top_fraction(scores: np.ndarray, frac: float) -> np.ndarray:
    n = len(scores)
    k = max(1, int(round(frac * n)))
    keep = np.zeros(n, bool)
    keep[np.argsort(-scores)[:k]] = True
    return keep