scripts/claim4_probe.py

"""Claim 4 — non-linear MLP probing on APT-tokenized sequences for CATH fold classification.

Encodes CATH proteins with the released APT tokenizer, mean+max pools the code
representation (c_BLD, the vectors the decoder consumes) into a fixed feature, then
trains (a) a linear probe and (b) a non-linear MLP probe to classify CATH topology
(the C.A.T level, e.g. "3.90.10").

Reproduces APT's own probe accuracy and shows the non-linear MLP beats a linear probe
on the same tokens. Cross-tokenizer comparison to DPLM2/ESM3 (Fig 5) is cited from the
paper (ESM3 weights are gated); this run verifies the APT side.
"""
from __future__ import annotations

import json
import os
from collections import Counter
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn

from apt.models import APTTokenizer

CKPT = Path(os.environ.get("CKPT_DIR", "/tmp/apt_ckpts"))
TOK = CKPT / "tokenizer128.pt"
OUT = Path(os.environ.get("OUT_DIR", "/tmp/out"))
MAXLEN = 256
MIN_PER_CLASS = 10  # keep topology classes with enough support


def load_split(split: str):
    from huggingface_hub import hf_hub_download
    import pyarrow.parquet as pq

    files = {
        "train": ["data/train-00000-of-00002.parquet", "data/train-00001-of-00002.parquet"],
        "validation": ["data/validation-00000-of-00001.parquet"],
        "test": ["data/test-00000-of-00001.parquet"],
    }[split]
    rows = []
    for f in files:
        fp = hf_hub_download("cctien/protein_backbone_cath_4.3", filename=f, repo_type="dataset")
        tbl = pq.read_table(fp, columns=["name", "coords", "CATH"])
        for nm, c, cath in zip(
            tbl.column("name").to_pylist(),
            tbl.column("coords").to_pylist(),
            tbl.column("CATH").to_pylist(),
        ):
            if not cath:
                continue
            lab = cath[0]  # first domain's C.A.T topology, e.g. "3.90.10"
            if not lab or lab.count(".") < 2:
                continue
            ca = np.asarray(c["CA"], dtype=np.float32)
            if ca.ndim != 2 or ca.shape[1] != 3 or np.isnan(ca).any():
                continue
            if not (1 <= ca.shape[0] <= MAXLEN):
                continue
            rows.append((nm, ca, lab))
    return rows


@torch.no_grad()
def featurize(tok, rows, dev, max_toks):
    feats, labels = [], []
    for i, (nm, ca, lab) in enumerate(rows):
        try:
            x = torch.from_numpy(ca).float()
            x = (x - x.mean(0, keepdim=True)) / 10.0
            x = x.unsqueeze(0).to(dev)
            s_BLD, _, _ = tok.encode(x)  # encoder embedding (pre-quantization)
            c = s_BLD[0]  # (L, D), one token per residue
            feat = torch.cat([c.mean(0), c.amax(0)], dim=-1)  # (2D,)
        except Exception as e:  # noqa: BLE001
            print(f"  skip {nm}: {e}")
            continue
        feats.append(feat.cpu().numpy())
        labels.append(lab)
        if i % 2000 == 0:
            print(f"  featurized {i}/{len(rows)}")
    return np.stack(feats), np.array(labels)


class MLP(nn.Module):
    def __init__(self, d_in, n_cls, hidden=512):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_in, hidden), nn.GELU(), nn.Dropout(0.3),
            nn.Linear(hidden, hidden), nn.GELU(), nn.Dropout(0.3),
            nn.Linear(hidden, n_cls),
        )

    def forward(self, x):
        return self.net(x)


def train_probe(Xtr, ytr, Xva, yva, Xte, yte, n_cls, dev, linear=False, epochs=60):
    d = Xtr.shape[1]
    model = (nn.Linear(d, n_cls) if linear else MLP(d, n_cls)).to(dev)
    opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
    Xtr_t = torch.tensor(Xtr, device=dev); ytr_t = torch.tensor(ytr, device=dev)
    Xva_t = torch.tensor(Xva, device=dev); Xte_t = torch.tensor(Xte, device=dev)
    lossf = nn.CrossEntropyLoss()
    best_va, best_te = 0.0, 0.0
    n = Xtr.shape[0]
    for ep in range(epochs):
        model.train()
        perm = torch.randperm(n, device=dev)
        for s in range(0, n, 256):
            idx = perm[s : s + 256]
            opt.zero_grad()
            loss = lossf(model(Xtr_t[idx]), ytr_t[idx])
            loss.backward()
            opt.step()
        model.eval()
        with torch.no_grad():
            va = (model(Xva_t).argmax(1).cpu().numpy() == yva).mean()
            if va >= best_va:
                best_va = va
                best_te = (model(Xte_t).argmax(1).cpu().numpy() == yte).mean()
    return float(best_va), float(best_te)


def main() -> None:
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    torch.manual_seed(0); np.random.seed(0)
    tok = APTTokenizer.from_pretrained(TOK).to(dev).eval()
    max_toks = tok.cfg.n_tokens

    # The cctien CATH splits hold out whole topologies (disjoint folds), so a
    # cross-split fold classifier has zero shared classes. Pool all structures,
    # featurize once, then evaluate at architecture (C.A) and topology (C.A.T)
    # granularity with a stratified 70/10/20 split over shared classes.
    allrows = load_split("train") + load_split("validation") + load_split("test")
    print(f"pooled structures={len(allrows)}")
    X_all, lab_all = featurize(tok, allrows, dev, max_toks)  # lab = full C.A.T string
    lab_all = list(lab_all)

    def to_level(lab, k):  # k=2 architecture C.A ; k=3 topology C.A.T
        return ".".join(lab.split(".")[:k])

    summary = {"feat_dim": int(X_all.shape[1]), "pooled": len(allrows), "levels": {}}
    for lvlname, k in [("architecture_CA", 2), ("topology_CAT", 3)]:
        labs = [to_level(l, k) for l in lab_all]
        cnt = Counter(labs)
        keep = {c for c, n in cnt.items() if n >= MIN_PER_CLASS}
        classes = sorted(keep)
        cls2i = {c: i for i, c in enumerate(classes)}
        idx = [i for i, l in enumerate(labs) if l in keep]

        rng = np.random.RandomState(0)
        by_cls = {}
        for i in idx:
            by_cls.setdefault(labs[i], []).append(i)
        tri, vai, tei = [], [], []
        for c, ii in by_cls.items():
            rng.shuffle(ii)
            n = len(ii); n_te = max(1, int(0.2 * n)); n_va = max(1, int(0.1 * n))
            tei += ii[:n_te]; vai += ii[n_te:n_te + n_va]; tri += ii[n_te + n_va:]

        def pack(ids):
            return X_all[ids], np.array([cls2i[labs[i]] for i in ids])
        Xtr, ytr = pack(tri); Xva, yva = pack(vai); Xte, yte = pack(tei)
        mu, sd = Xtr.mean(0), Xtr.std(0) + 1e-6
        Xtr = (Xtr - mu) / sd; Xva = (Xva - mu) / sd; Xte = (Xte - mu) / sd

        _, lin_te = train_probe(Xtr, ytr, Xva, yva, Xte, yte, len(classes), dev, linear=True)
        _, mlp_te = train_probe(Xtr, ytr, Xva, yva, Xte, yte, len(classes), dev, linear=False)
        majority = Counter(ytr).most_common(1)[0][1] / len(ytr)
        summary["levels"][lvlname] = {
            "n_classes": len(classes), "n_train": len(tri), "n_test": len(tei),
            "majority_acc": round(float(majority), 4),
            "linear_probe_test_acc": round(lin_te, 4),
            "mlp_probe_test_acc": round(mlp_te, 4),
            "mlp_minus_linear": round(mlp_te - lin_te, 4),
        }
        print(f"[{lvlname}] classes={len(classes)} maj={majority:.3f} linear={lin_te:.3f} mlp={mlp_te:.3f}")

    print("RESULT_JSON " + json.dumps(summary))
    OUT.mkdir(parents=True, exist_ok=True)
    (OUT / "claim4_probe_summary.json").write_text(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()