Audit API

The audit module implements the attribute-aware shortcut-potential audit used in the paper: a conditional probe as the single primary test, five secondary detectors, and a conditional permutation null (the attribute shuffled within task-label strata, between patients when patient ids are given). See The paper's audit for the method description.

Running an audit

from shortcut_detect.audit import AuditConfig, conditional_audit

cfg = AuditConfig(n_permutations=200, n_permutations_probe=500, seed=42)
result = conditional_audit(E, A, Y, config=cfg, groups=patient_ids)
probe = result.results["probe"]
print(probe.statistic, probe.p_value, probe.p_bh, probe.flag)

Class Reference

AuditConfig

AuditConfig dataclass

AuditConfig(
    n_permutations: int = 200,
    n_permutations_probe: int | None = None,
    alpha: float = 0.05,
    n_folds: int = 5,
    C: float = 1.0,
    top_k_fraction: float = 0.05,
    n_clusters: int = 8,
    n_bootstrap: int = 2000,
    max_iter: int = 100,
    seed: int = 0,
    permutation: str = "within_y",
    include_y: bool = True,
    detectors: tuple[str, ...] = DETECTORS,
    chunk_problems: int = 1200,
    cv_repeats: int = 5,
    cv_seed: int | None = None,
    attribute_level: str = "patient",
)

Settings fixed by the revision method decisions (sections 2 to 4).

conditional_audit

conditional_audit

conditional_audit(
    E: ndarray,
    A: ndarray,
    Y: ndarray,
    config: AuditConfig | None = None,
    groups: ndarray | None = None,
) -> AuditResult

Run the attribute-aware audit on one (embedding, attribute, task) cell.

E is any 2-D float matrix (n, d): any width d and any modality (image, text, audio, tabular, multimodal encoders), the audit only sees the numbers. A and Y are 1-D arrays of length n.

groups (e.g. patient ids) makes every cross-validation split patient-disjoint, every bootstrap resample whole patients, and the permutation patient-level within Y-pattern strata (see conditional_permutations).

Source code in shortcut_detect/audit/conditional.py

def conditional_audit(
    E: np.ndarray,
    A: np.ndarray,
    Y: np.ndarray,
    config: AuditConfig | None = None,
    groups: np.ndarray | None = None,
) -> AuditResult:
    """Run the attribute-aware audit on one (embedding, attribute, task) cell.

    ``E`` is any 2-D float matrix ``(n, d)``: any width ``d`` and any modality
    (image, text, audio, tabular, multimodal encoders), the audit only sees the
    numbers. ``A`` and ``Y`` are 1-D arrays of length ``n``.

    ``groups`` (e.g. patient ids) makes every cross-validation split
    patient-disjoint, every bootstrap resample whole patients, and the
    permutation patient-level within Y-pattern strata (see
    ``conditional_permutations``).
    """
    cfg = config or AuditConfig()
    # independent streams, so the observed statistic does not depend on K
    perm_ss, null_fold_ss, boot_ss, km_ss, cv_ss = np.random.SeedSequence(cfg.seed).spawn(5)
    perm_rng, null_fold_rng = np.random.default_rng(perm_ss), np.random.default_rng(null_fold_ss)
    boot_rng, km_rng = np.random.default_rng(boot_ss), np.random.default_rng(km_ss)
    cv_rng = np.random.default_rng(cv_ss if cfg.cv_seed is None else cfg.cv_seed)
    E = np.asarray(E)
    Y = np.asarray(Y).astype(np.int64)
    a, names = _encode(A)
    G = len(names)
    if G < 2:
        raise ValueError("the attribute needs at least two groups")
    if E.shape[0] != len(a) or len(a) != len(Y):
        raise ValueError("E, A and Y must have the same number of rows")
    groups_idx = None
    n_mixed = 0
    level = "image"
    if groups is not None:
        groups = np.asarray(groups)
        if len(groups) != len(a):
            raise ValueError("groups must have one entry per row")
        level = cfg.attribute_level
        if level == "patient":
            a, n_mixed = patient_majority(a, groups)
            if len(np.unique(a)) < 2:
                raise ValueError("the patient-level attribute needs at least two groups")
        _, ginv = np.unique(groups.astype(str), return_inverse=True)
        groups_idx = [np.flatnonzero(ginv == g) for g in range(ginv.max() + 1)]

    needs = set(cfg.detectors)
    ctx = _prepare(E, Y, cfg, needs, km_rng)
    ctx["G"] = G
    K = cfg.n_permutations
    Kp = cfg.n_permutations_probe or K
    k_max = max(K, Kp) if "probe" in needs else K
    a_rows = conditional_permutations(a, Y, k_max, perm_rng, cfg.permutation, groups, level)
    R = max(1, int(cfg.cv_repeats))
    cv_needed = bool(needs & {"probe", "frequency_topk"})
    if cv_needed:
        obs_folds, fold_strata = _fold_ids(np.repeat(a[None, :], R, axis=0), Y, cfg.n_folds, cv_rng, groups)
        null_folds, _ = _fold_ids(a_rows[1:], Y, cfg.n_folds, null_fold_rng, groups)
    else:
        fold_strata = "none"
    sizes = {}
    for yv in np.unique(Y):
        for g, nm in enumerate(names):
            sizes[f"Y={yv},A={nm}"] = int(((Y == yv) & (a == g)).sum())

    res = AuditResult(
        n=len(a),
        n_groups=G,
        groups=names,
        group_sizes=sizes,
        fold_strata=fold_strata,
        config=asdict(cfg),
        n_inconsistent_groups=n_mixed,
        attribute_level=level,
    )
    B = cfg.n_bootstrap
    for det in cfg.detectors:
        t0 = time.time()
        ci = (float("nan"), float("nan"))
        rep_sd, n_rep, note = float("nan"), 1, ""
        if det in {"probe", "frequency_topk"}:
            # observed: mean over R repeats of grouped CV; null: one repeat each (conservative)
            k_used = Kp if det == "probe" else K
            fn_cv = _stat_probe if det == "probe" else _stat_frequency
            # observed repeats are fitted in their own batch: the batched optimiser
            # stops on the mean objective, so sharing a batch with the null rows
            # would make the observed value depend on K
            reps, sc, tg = fn_cv(ctx, np.repeat(a[None, :], R, axis=0), obs_folds, cfg)
            null_stats, _, _ = fn_cv(ctx, a_rows[1 : k_used + 1], null_folds[:k_used], cfg)
            reps = np.asarray(reps, dtype=float)
            stats = np.r_[reps.mean(), null_stats]
            n_rep = R
            rep_sd = float(reps.std(ddof=1)) if R > 1 else float("nan")
            if B:
                per = sc.shape[1] // R
                b_r = int(math.ceil(B / R))
                boots = [
                    _bootstrap_auc(sc[:, r * per : (r + 1) * per], tg[:, r * per : (r + 1) * per], b_r, boot_rng, groups_idx)
                    for r in range(R)
                ]
                ci = bc_interval(np.concatenate(boots), float(stats[0]))
                note = "CI: resampling of out-of-fold predictions pooled over CV repeats, no refit; see repeat_sd"
        elif det == "statistical":
            stats = _stat_statistical(ctx, a_rows[: K + 1], cfg)
        else:
            fn = {
                "geometric": _stat_geometric,
                "bias_direction": _stat_bias_direction,
                "hbac_nmi": _stat_hbac_nmi,
            }[det]
            rows = a_rows[: K + 1]
            stats = fn(ctx, rows, cfg)
            if B:
                oh0 = _onehot(a[None, :], G)
                vals = []
                step = max(1, int(4e6 // max(len(a) * max(G, 2), 1)))
                for s in range(0, B, step):
                    m = min(step, B - s)
                    if groups_idx is None:
                        w = boot_rng.multinomial(len(a), np.full(len(a), 1.0 / len(a)), size=m)
                    else:
                        gw = boot_rng.multinomial(
                            len(groups_idx), np.full(len(groups_idx), 1.0 / len(groups_idx)), size=m
                        )
                        w = gw[:, ginv]
                    w = w.astype(np.float32)
                    vals.append(fn(ctx, rows[:1], cfg, w=w, oh_fixed=oh0))
                ci = bc_interval(np.concatenate(vals), float(stats[0]))
        obs, null = float(stats[0]), np.asarray(stats[1:], dtype=float)
        p = (1.0 + np.sum(null >= obs - 1e-12)) / (len(null) + 1.0)
        sdn = float(null.std(ddof=1)) if len(null) > 1 else float("nan")
        res.results[det] = DetectorResult(
            detector=det,
            statistic=obs,
            null_mean=float(null.mean()),
            null_sd=sdn,
            z=(obs - float(null.mean())) / sdn if sdn and sdn > 0 else float("nan"),
            p_value=float(p),
            ci_low=float(ci[0]),
            ci_high=float(ci[1]),
            n_permutations=int(len(null)),
            seconds=round(time.time() - t0, 3),
            note=note,
            repeat_sd=rep_sd,
            n_repeats=n_rep,
        )
    # the primary probe is judged at its own level; BH runs across the secondary detectors only
    if PRIMARY in res.results:
        pr = res.results[PRIMARY]
        pr.p_bh = float("nan")
        pr.flag = bool(pr.p_value < cfg.alpha)
    names_d = [d for d in res.results if d != PRIMARY]
    if names_d:
        q = benjamini_hochberg(np.array([res.results[d].p_value for d in names_d]))
        for d_, qv in zip(names_d, q, strict=True):
            res.results[d_].p_bh = float(qv)
            res.results[d_].flag = bool(qv < cfg.alpha)
    return res

AuditResult

AuditResult dataclass

AuditResult(
    n: int,
    n_groups: int,
    groups: list[str],
    group_sizes: dict[str, int],
    fold_strata: str,
    config: dict[str, Any],
    results: dict[str, DetectorResult] = dict(),
    n_inconsistent_groups: int = 0,
    attribute_level: str = "image",
)

DetectorResult

DetectorResult dataclass

DetectorResult(
    detector: str,
    statistic: float,
    null_mean: float,
    null_sd: float,
    z: float,
    p_value: float,
    p_bh: float = float("nan"),
    flag: bool = False,
    ci_low: float = float("nan"),
    ci_high: float = float("nan"),
    n_permutations: int = 0,
    seconds: float = 0.0,
    note: str = "",
    repeat_sd: float = float("nan"),
    n_repeats: int = 1,
)

conditional_permutations

conditional_permutations

conditional_permutations(
    a: ndarray,
    y: ndarray,
    k: int,
    rng: Generator,
    mode: str = "within_y",
    groups: ndarray | None = None,
    level: str = "patient",
) -> np.ndarray

Return a (k + 1, n) array whose row 0 is the observed attribute and rows 1..k are permuted copies: within each Y stratum, or globally when mode='global'.

With groups (patient ids) the permutation unit is the patient, because images of one patient are correlated and an image-level shuffle gives anti-conservative p-values.

level='patient': the attribute is a patient property. Each patient's value (its majority, which is also row 0) is reassigned among patients with the same Y pattern (all negative, all positive, or mixed).

level='image': the attribute varies between a patient's images (e.g. a marker). Whole attribute vectors (images ordered by Y) are reassigned among patients with the same sorted Y vector, hence the same image count, so the within-patient structure and the joint (A, Y) image counts are kept.

Source code in shortcut_detect/audit/conditional.py

def conditional_permutations(
    a: np.ndarray,
    y: np.ndarray,
    k: int,
    rng: np.random.Generator,
    mode: str = "within_y",
    groups: np.ndarray | None = None,
    level: str = "patient",
) -> np.ndarray:
    """Return a (k + 1, n) array whose row 0 is the observed attribute and rows
    1..k are permuted copies: within each Y stratum, or globally when
    ``mode='global'``.

    With ``groups`` (patient ids) the permutation unit is the patient, because
    images of one patient are correlated and an image-level shuffle gives
    anti-conservative p-values.

    ``level='patient'``: the attribute is a patient property. Each patient's
    value (its majority, which is also row 0) is reassigned among patients with
    the same Y pattern (all negative, all positive, or mixed).

    ``level='image'``: the attribute varies between a patient's images (e.g. a
    marker). Whole attribute vectors (images ordered by Y) are reassigned among
    patients with the same sorted Y vector, hence the same image count, so the
    within-patient structure and the joint (A, Y) image counts are kept.
    """
    n = len(a)
    out = np.empty((k + 1, n), dtype=a.dtype)
    out[0] = a
    if groups is not None:
        _, ginv = np.unique(np.asarray(groups).astype(str), return_inverse=True)
        n_g = ginv.max() + 1
        if level == "image":
            buckets: dict[tuple, list[np.ndarray]] = {}
            for g in range(n_g):
                idx = np.flatnonzero(ginv == g)
                idx = idx[np.argsort(y[idx], kind="stable")]
                key = (len(idx),) if mode == "global" else tuple(int(v) for v in y[idx])
                buckets.setdefault(key, []).append(idx)
            mats = [np.stack(v) for v in buckets.values()]  # (patients, images)
            for i in range(1, k + 1):
                row = a.copy()
                for mat in mats:
                    row[mat] = a[mat[rng.permutation(len(mat))]]
                out[i] = row
            return out
        if level != "patient":
            raise ValueError("level must be 'patient' or 'image'")
        pa = np.empty(n_g, dtype=a.dtype)
        pstrat = np.empty(n_g, dtype=np.int64)
        y_hi = int(y.max())
        for g in range(n_g):
            idx = ginv == g
            vals, cnt = np.unique(a[idx], return_counts=True)
            pa[g] = vals[np.argmax(cnt)]
            ys = np.unique(y[idx])
            pstrat[g] = int(ys[0]) if len(ys) == 1 else y_hi + 1
        out[0] = pa[ginv]
        strata = [np.arange(n_g)] if mode == "global" else [np.flatnonzero(pstrat == s) for s in np.unique(pstrat)]
        for i in range(1, k + 1):
            prow = pa.copy()
            for idx in strata:
                prow[idx] = pa[idx][rng.permutation(len(idx))]
            out[i] = prow[ginv]
        return out
    strata = [np.arange(n)] if mode == "global" else [np.flatnonzero(y == v) for v in np.unique(y)]
    for i in range(1, k + 1):
        row = a.copy()
        for idx in strata:
            row[idx] = a[idx][rng.permutation(len(idx))]
        out[i] = row
    return out

benjamini_hochberg

benjamini_hochberg

benjamini_hochberg(p: ndarray) -> np.ndarray

BH adjusted p-values along the last axis.

Source code in shortcut_detect/audit/conditional.py

def benjamini_hochberg(p: np.ndarray) -> np.ndarray:
    """BH adjusted p-values along the last axis."""
    p = np.asarray(p, dtype=float)
    m = p.shape[-1]
    order = np.argsort(p, axis=-1)
    ranked = np.take_along_axis(p, order, axis=-1) * m / np.arange(1, m + 1)
    ranked = np.minimum.accumulate(ranked[..., ::-1], axis=-1)[..., ::-1]
    out = np.empty_like(ranked)
    np.put_along_axis(out, order, np.minimum(ranked, 1.0), axis=-1)
    return out