{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1e54a7de",
   "metadata": {},
   "source": [
    "# Aminograpes: reproduce the viral results on Google Colab\n",
    "\n",
    "This notebook reruns the whole Aminograpes evaluation on the **31 viral assays of ProteinGym v1.3** and checks your numbers against the published run.\n",
    "\n",
    "## In plain words\n",
    "- **ESM2-650M** is a protein language model made by Meta. It is the baseline Aminograpes is compared against.\n",
    "- **Aminograpes** is a simple score: how conserved a position is among related viruses (from an alignment), multiplied by how buried that position is inside the protein (from a predicted 3D structure). No training. **The Aminograpes model code is in Step 4.**\n",
    "- **The \"ensemble\" in Step 10 is not Aminograpes.** It combines other groups' public models and is only there for context. Aminograpes adds nothing to it.\n",
    "- \"Meta\" always means the company that made ESM2. \"Ensemble\" means the stack of public models.\n",
    "\n",
    "## What you should get\n",
    "| Result (mean Spearman, 31 viral assays) | Expected |\n",
    "|---|---|\n",
    "| ESM2-650M, official ProteinGym scores | 0.2687 (our rerun 0.2693) |\n",
    "| Aminograpes, alignment only | 0.3926 |\n",
    "| **Aminograpes (alignment x burial)** | **0.4304** |\n",
    "| Aminograpes + ESM2 blend (weight 0.7) | 0.3764 |\n",
    "| Ensemble of 5 public models (not Aminograpes) | 0.5732 |\n",
    "\n",
    "Aminograpes minus official ESM2-650M: **+0.1617, 95% CI [0.0775, 0.2450], better on 25 of 31 assays.**\n",
    "\n",
    "## How to run it\n",
    "1. **Runtime > Change runtime type > T4 GPU** (any GPU works). Then **Runtime > Run all**, or run the steps in order.\n",
    "2. Steps 1 to 8 are the core check (about 1 to 2 hours, mostly downloads and ESM2). Steps 9 and 10 are optional extras (another 1 to 2 hours).\n",
    "3. At the end, Step 11 zips your results so you can download them.\n",
    "\n",
    "Needs about 20 GB of free disk and a high-RAM runtime helps for the biggest alignments. If Colab disconnects, run again: finished assays are cached and skipped.\n",
    "\n",
    "Honest scope: this is an evaluation on a public benchmark, not an official ProteinGym leaderboard entry."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f151236b",
   "metadata": {},
   "source": [
    "## Step 1. Install and set up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2658cd6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip -q install biopython transformers\n",
    "import os, re, glob, zipfile, shutil, numpy as np, pandas as pd\n",
    "from scipy.stats import spearmanr, wilcoxon\n",
    "WORK = '/content/aminograpes_repro'; P = f'{WORK}/pg'; OUT = f'{WORK}/out'\n",
    "os.makedirs(P, exist_ok=True); os.makedirs(OUT, exist_ok=True)\n",
    "import torch; print('GPU available:', torch.cuda.is_available())\n",
    "assert torch.cuda.is_available(), 'Turn on a GPU: Runtime > Change runtime type'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94ce6c29",
   "metadata": {},
   "source": [
    "## Step 2. Download the public data (several GB, about 10 to 20 minutes)\n",
    "Everything comes from the official ProteinGym v1.3 release."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aea15ede",
   "metadata": {},
   "outputs": [],
   "source": [
    "B = 'https://marks.hms.harvard.edu/proteingym/ProteinGym_v1.3'\n",
    "def fetch(zipname, flag, dest=P):\n",
    "    if os.path.exists(f'{P}/{flag}'): return\n",
    "    print('downloading', zipname, '...')\n",
    "    assert os.system(f'curl -fsSL -o {P}/t.zip {B}/{zipname}') == 0, 'download failed: ' + zipname\n",
    "    zipfile.ZipFile(f'{P}/t.zip').extractall(dest); os.remove(f'{P}/t.zip'); open(f'{P}/{flag}', 'w').close()\n",
    "fetch('DMS_ProteinGym_substitutions.zip', 'subs.ok')\n",
    "fetch('DMS_msa_weights.zip', 'w.ok')\n",
    "fetch('DMS_msa_files.zip', 'msa.ok')\n",
    "fetch('zero_shot_substitutions_scores.zip', 'zs.ok', dest=f'{P}/zs')\n",
    "if not os.path.exists(f'{P}/DMS_substitutions.csv'):\n",
    "    os.system(f'curl -fsSL -o {P}/DMS_substitutions.csv https://raw.githubusercontent.com/OATML-Markslab/ProteinGym/main/reference_files/DMS_substitutions.csv')\n",
    "if not os.path.isdir(f'{P}/af2'):\n",
    "    ok = False\n",
    "    for u in [f'{B}/ProteinGym_AF2_structures.zip', 'https://marks.hms.harvard.edu/proteingym/ProteinGym_AF2_structures.zip']:\n",
    "        if os.system(f'curl -fsSL -o {P}/af2.zip {u}') == 0: ok = True; break\n",
    "    assert ok, 'structure download failed'\n",
    "    zipfile.ZipFile(f'{P}/af2.zip').extractall(f'{P}/af2'); os.remove(f'{P}/af2.zip')\n",
    "print('done')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8471d5c2",
   "metadata": {},
   "source": [
    "## Step 3. Pick the 31 viral assays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52479de0",
   "metadata": {},
   "outputs": [],
   "source": [
    "ref = pd.read_csv(f'{P}/DMS_substitutions.csv')\n",
    "viral = ref[ref['taxon'] == 'Virus'].reset_index(drop=True)\n",
    "print('viral assays:', len(viral)); assert len(viral) == 31\n",
    "# index of data files (skip the structure and official-score folders so names do not collide)\n",
    "idx = {os.path.basename(p): p for p in glob.glob(f'{P}/**/*', recursive=True)\n",
    "       if os.path.isfile(p) and not p.startswith((f'{P}/af2', f'{P}/zs'))}\n",
    "pdbs = {os.path.splitext(os.path.basename(p))[0]: p for p in glob.glob(f'{P}/af2/**/*.pdb', recursive=True)}\n",
    "print(len(pdbs), 'structure files')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acc0e61e",
   "metadata": {},
   "source": [
    "## Step 4. THE AMINOGRAPES MODEL CODE\n",
    "This one cell is the whole Aminograpes model. Read it top to bottom:\n",
    "1. `build_logf` reads the alignment, down-weights near-duplicate sequences, and gives each position a frequency for every amino acid (conservation).\n",
    "2. `struct_weights` reads the predicted structure and works out how buried each position is (burial, 1 minus relative solvent accessibility).\n",
    "3. `aminograpes_scores` scores each mutation as: burial x [ log(frequency of new amino acid) minus log(frequency of original) ], summed over the mutated positions.\n",
    "\n",
    "Fixed constants (nothing is trained or tuned): similarity cutoff 0.2, pseudocount 0.01, structure accepted if at least 30% of the protein maps and at least 80% of mapped residues match."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1bc8469",
   "metadata": {},
   "outputs": [],
   "source": [
    "from Bio import Align\n",
    "from Bio.Align import substitution_matrices\n",
    "from Bio.PDB import PDBParser, ShrakeRupley\n",
    "from Bio.SeqUtils import seq1\n",
    "\n",
    "AA = 'ACDEFGHIKLMNPQRSTVWY'; aidx = {a: i for i, a in enumerate(AA)}\n",
    "LUT = np.full(256, 21, dtype=np.uint8)\n",
    "for i, a in enumerate(AA): LUT[ord(a)] = i\n",
    "LUT[ord('-')] = 20\n",
    "LAM, THETA, MAXB = 0.01, 0.2, 8e9           # pseudocount, clustering cutoff, memory budget\n",
    "MAXASA = dict(A=129, R=274, N=195, D=193, C=167, E=223, Q=225, G=104, H=224, I=197,\n",
    "              L=201, K=236, M=224, F=240, P=159, S=155, T=172, W=285, Y=263, V=174)  # Tien et al. 2013, theoretical\n",
    "\n",
    "def read_a2m(path):\n",
    "    seqs, cur = [], []\n",
    "    for line in open(path):\n",
    "        line = line.strip()\n",
    "        if line.startswith('>'):\n",
    "            if cur: seqs.append(''.join(cur)); cur = []\n",
    "        elif line: cur.append(line)\n",
    "    if cur: seqs.append(''.join(cur))\n",
    "    return [re.sub('[a-z.]', '', s) for s in seqs]\n",
    "\n",
    "@torch.no_grad()\n",
    "def gpu_weights(X, chunk=512):\n",
    "    N, L = X.shape                                   # gap counts as its own state\n",
    "    Xt = torch.from_numpy(X.astype(np.int64)).cuda()\n",
    "    oh = torch.zeros(N, L, 21, dtype=torch.float16, device='cuda')\n",
    "    oh.scatter_(2, Xt.unsqueeze(2), 1.0); oh = oh.view(N, L * 21)\n",
    "    cnt = torch.zeros(N, device='cuda'); thr = (1 - THETA) * L\n",
    "    for s in range(0, N, chunk):\n",
    "        cnt[s:s + chunk] = ((oh[s:s + chunk] @ oh.T) >= thr).sum(1).float()\n",
    "    w = (1.0 / cnt.clamp(min=1)).cpu().numpy()\n",
    "    del oh, Xt; torch.cuda.empty_cache(); return w\n",
    "\n",
    "def build_logf(msa_path, w_path):\n",
    "    seqs = read_a2m(msa_path); q = seqs[0]; L = len(q); seqs = [s for s in seqs if len(s) == L]\n",
    "    X = LUT[np.frombuffer(''.join(seqs).encode(), dtype=np.uint8).reshape(len(seqs), L)]\n",
    "    focus = np.array([c != '-' for c in q]); X = X[:, focus]; qf = q.replace('-', '')\n",
    "    keep = (X < 21).all(1) & ((X == 20).mean(1) <= 0.5)\n",
    "    wf = np.load(w_path).astype(np.float32)\n",
    "    if len(wf) == keep.sum(): X, w = X[keep], wf              # ProteinGym weights match\n",
    "    elif len(wf) == len(X): w = wf\n",
    "    else:                                                      # recompute weights\n",
    "        X = X[keep]\n",
    "        if len(X) * X.shape[1] * 42 > MAXB:\n",
    "            n = int(MAXB / (X.shape[1] * 42)); X = X[np.sort(np.random.RandomState(0).choice(len(X), n, replace=False))]\n",
    "        w = gpu_weights(X)\n",
    "    c = np.stack([w @ (X == a) for a in range(20)], 1)\n",
    "    f = (1 - LAM) * c / np.maximum(c.sum(1, keepdims=True), 1e-9) + LAM / 20\n",
    "    return np.log(f), qf\n",
    "\n",
    "aligner = Align.PairwiseAligner(); aligner.mode = 'global'\n",
    "aligner.substitution_matrix = substitution_matrices.load('BLOSUM62')\n",
    "aligner.open_gap_score = -10; aligner.extend_gap_score = -0.5; aligner.end_gap_score = 0.0\n",
    "def map_positions(target, other):\n",
    "    aln = aligner.align(target, other)[0]; tmap = -np.ones(len(target), dtype=int)\n",
    "    for (t0, t1), (q0, q1) in zip(aln.aligned[0], aln.aligned[1]): tmap[t0:t1] = np.arange(q0, q1)\n",
    "    return tmap\n",
    "\n",
    "def rsa_from_pdb(path):\n",
    "    model = PDBParser(QUIET=True).get_structure('x', path)[0]\n",
    "    ShrakeRupley().compute(model, level='R')\n",
    "    chain = next(model.get_chains()); seq, rsa = [], []\n",
    "    for res in chain:\n",
    "        if res.id[0] != ' ': continue\n",
    "        aa = seq1(res.get_resname())\n",
    "        if aa in MAXASA: seq.append(aa); rsa.append(min(1.0, res.sasa / MAXASA[aa]))\n",
    "    return ''.join(seq), np.array(rsa)\n",
    "\n",
    "def find_struct(r):\n",
    "    for key in [r['DMS_id'], r['UniProt_ID']]:\n",
    "        if key in pdbs: return pdbs[key]\n",
    "    for k, p in pdbs.items():\n",
    "        if len(k) >= 6 and (r['UniProt_ID'] in k or k in r['DMS_id']): return p\n",
    "    return None\n",
    "\n",
    "def struct_weights(r):\n",
    "    \"burial = 1 - RSA per target position, or None if no usable structure\"\n",
    "    tseq = r['target_seq']; sp = find_struct(r)\n",
    "    if not sp: return None\n",
    "    sseq, srsa = rsa_from_pdb(sp); sm = map_positions(tseq, sseq); m_ = sm >= 0\n",
    "    if m_.mean() < 0.3: return None\n",
    "    if np.mean([tseq[t] == sseq[sm[t]] for t in np.where(m_)[0]]) < 0.8: return None\n",
    "    rsa = np.full(len(tseq), np.nan); rsa[m_] = srsa[sm[m_]]; rsa[np.isnan(rsa)] = np.nanmean(rsa)\n",
    "    return 1 - rsa\n",
    "\n",
    "def aminograpes_scores(r, mutants):\n",
    "    \"returns (plain alignment score, Aminograpes score = alignment x burial, structure_used)\"\n",
    "    tseq = r['target_seq']\n",
    "    logf, qf = build_logf(idx[r['MSA_filename']], idx[r['weight_file_name']])\n",
    "    tmap = map_positions(tseq, qf)\n",
    "    w_t = struct_weights(r); used = w_t is not None\n",
    "    if not used: w_t = np.full(len(tseq), 0.5)       # constant weight: ranking equals plain alignment\n",
    "    plain, afr = [], []\n",
    "    for x in mutants:\n",
    "        a = b = 0.0\n",
    "        for m in x.split(':'):\n",
    "            wt, p, mt = m[0], int(m[1:-1]) - 1, m[-1]; q = tmap[p]\n",
    "            if q >= 0:\n",
    "                lor = logf[q, aidx[mt]] - logf[q, aidx[wt]]; a += lor; b += w_t[p] * lor\n",
    "        plain.append(a); afr.append(b)\n",
    "    return plain, afr, used\n",
    "print('Aminograpes model code loaded')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5cfa873",
   "metadata": {},
   "source": [
    "## Step 5. Score the viral assays with ESM2-650M (the baseline)\n",
    "Uses a 1022-residue window for long proteins. Results are cached, so a reconnect does not lose them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53b8875b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, EsmForMaskedLM\n",
    "name = 'facebook/esm2_t33_650M_UR50D'\n",
    "tok = AutoTokenizer.from_pretrained(name)\n",
    "model = EsmForMaskedLM.from_pretrained(name, torch_dtype=torch.float16).cuda().eval()\n",
    "aa_id = torch.tensor([tok.convert_tokens_to_ids(a) for a in AA]).cuda()\n",
    "MAXL, BS = 1022, 16\n",
    "\n",
    "@torch.no_grad()\n",
    "def esm2_scores(seq, mutants):\n",
    "    L = len(seq); pos = sorted({int(m[1:-1]) - 1 for x in mutants for m in x.split(':')}); lp = {}\n",
    "    for s in range(0, len(pos), BS):\n",
    "        chunk = pos[s:s + BS]; wins = []\n",
    "        for p in chunk:\n",
    "            start = 0 if L <= MAXL else max(0, min(p - MAXL // 2, L - MAXL))\n",
    "            wins.append((seq[start:start + MAXL], p - start))\n",
    "        enc = tok([w for w, _ in wins], return_tensors='pt').to('cuda'); ids = enc['input_ids'].clone()\n",
    "        for r_, (w, rel) in enumerate(wins): ids[r_, rel + 1] = tok.mask_token_id\n",
    "        logits = model(input_ids=ids, attention_mask=enc['attention_mask']).logits\n",
    "        for r_, (w, rel) in enumerate(wins): lp[chunk[r_]] = logits[r_, rel + 1].float().log_softmax(-1)[aa_id].cpu()\n",
    "    out = []\n",
    "    for x in mutants:\n",
    "        s = 0.0\n",
    "        for m in x.split(':'):\n",
    "            wt, p, mt = m[0], int(m[1:-1]) - 1, m[-1]; s += (lp[p][aidx[mt]] - lp[p][aidx[wt]]).item()\n",
    "        out.append(s)\n",
    "    return out\n",
    "\n",
    "for _, r in viral.iterrows():\n",
    "    f = f\"{OUT}/esm2_{r['DMS_id']}.csv\"\n",
    "    if os.path.exists(f): continue\n",
    "    df = pd.read_csv(idx[r['DMS_filename']], usecols=['mutant', 'DMS_score'])\n",
    "    df['esm2'] = esm2_scores(r['target_seq'], list(df.mutant)); df.to_csv(f, index=False)\n",
    "    print(f\"{r['DMS_id']:45s} esm2 rho = {spearmanr(df.esm2, df.DMS_score).correlation:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77382540",
   "metadata": {},
   "source": [
    "## Step 6. Run Aminograpes on the 31 viral assays\n",
    "Expected: mean alignment-only **0.3926**, mean Aminograpes **0.4304**, structure used for **29 of 31** assays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fb5fe35",
   "metadata": {},
   "outputs": [],
   "source": [
    "EXPECTED_PLAIN = {\"A0A140D2T1_ZIKV_Sourisseau_2019\": 0.399, \"A0A192B1T2_9HIV1_Haddox_2018\": 0.471, \"A0A2Z5U3Z0_9INFA_Doud_2016\": 0.493, \"A0A2Z5U3Z0_9INFA_Wu_2014\": 0.452, \"A4D664_9INFA_Soh_2019\": 0.403, \"C6KNH7_9INFA_Lee_2018\": 0.417, \"CAPSD_AAV2S_Sinai_2021\": 0.447, \"ENV_HV1B9_DuenasDecamp_2016\": 0.319, \"ENV_HV1BR_Haddox_2016\": 0.307, \"HCP_LAMBD_Tsuboyama_2023_2L6Q\": 0.473, \"I6TAH8_I68A0_Doud_2015\": 0.355, \"NCAP_I34A1_Doud_2015\": 0.373, \"NRAM_I33A0_Jiang_2016\": 0.498, \"PA_I34A1_Wu_2015\": 0.522, \"POLG_CXB3N_Mattenberger_2021\": 0.416, \"POLG_DEN26_Suphatrakul_2023\": 0.507, \"POLG_HCVJF_Qi_2014\": 0.661, \"POLG_PESV_Tsuboyama_2023_2MXD\": 0.208, \"Q2N0S5_9HIV1_Haddox_2018\": 0.485, \"R1AB_SARS2_Flynn_2022\": 0.562, \"RCRO_LAMBD_Tsuboyama_2023_1ORC\": 0.377, \"RDRP_I33A0_Li_2023\": 0.371, \"REV_HV1H2_Fernandes_2016\": 0.239, \"RPC1_BP434_Tsuboyama_2023_1R69\": 0.726, \"RPC1_LAMBD_Li_2019_high-expression\": 0.284, \"RPC1_LAMBD_Li_2019_low-expression\": 0.243, \"SPIKE_SARS2_Starr_2020_binding\": 0.172, \"SPIKE_SARS2_Starr_2020_expression\": 0.203, \"TAT_HV1BR_Fernandes_2016\": 0.361, \"VG08_BPP22_Tsuboyama_2023_2GP8\": 0.373, \"VRPI_BPT7_Tsuboyama_2023_2WNM\": 0.055}\n",
    "EXPECTED_AMINOGRAPES = {\"A0A140D2T1_ZIKV_Sourisseau_2019\": 0.399, \"A0A192B1T2_9HIV1_Haddox_2018\": 0.5, \"A0A2Z5U3Z0_9INFA_Doud_2016\": 0.545, \"A0A2Z5U3Z0_9INFA_Wu_2014\": 0.465, \"A4D664_9INFA_Soh_2019\": 0.287, \"C6KNH7_9INFA_Lee_2018\": 0.528, \"CAPSD_AAV2S_Sinai_2021\": 0.541, \"ENV_HV1B9_DuenasDecamp_2016\": 0.309, \"ENV_HV1BR_Haddox_2016\": 0.285, \"HCP_LAMBD_Tsuboyama_2023_2L6Q\": 0.602, \"I6TAH8_I68A0_Doud_2015\": 0.236, \"NCAP_I34A1_Doud_2015\": 0.281, \"NRAM_I33A0_Jiang_2016\": 0.569, \"PA_I34A1_Wu_2015\": 0.4, \"POLG_CXB3N_Mattenberger_2021\": 0.406, \"POLG_DEN26_Suphatrakul_2023\": 0.588, \"POLG_HCVJF_Qi_2014\": 0.661, \"POLG_PESV_Tsuboyama_2023_2MXD\": 0.263, \"Q2N0S5_9HIV1_Haddox_2018\": 0.525, \"R1AB_SARS2_Flynn_2022\": 0.641, \"RCRO_LAMBD_Tsuboyama_2023_1ORC\": 0.426, \"RDRP_I33A0_Li_2023\": 0.326, \"REV_HV1H2_Fernandes_2016\": 0.321, \"RPC1_BP434_Tsuboyama_2023_1R69\": 0.748, \"RPC1_LAMBD_Li_2019_high-expression\": 0.364, \"RPC1_LAMBD_Li_2019_low-expression\": 0.315, \"SPIKE_SARS2_Starr_2020_binding\": 0.333, \"SPIKE_SARS2_Starr_2020_expression\": 0.359, \"TAT_HV1BR_Fernandes_2016\": 0.42, \"VG08_BPP22_Tsuboyama_2023_2GP8\": 0.488, \"VRPI_BPT7_Tsuboyama_2023_2WNM\": 0.213}\n",
    "rows = []\n",
    "for _, r in viral.iterrows():\n",
    "    i = r['DMS_id']; f = f'{OUT}/both_{i}.csv'\n",
    "    if os.path.exists(f): df = pd.read_csv(f)\n",
    "    else:\n",
    "        df = pd.read_csv(f'{OUT}/esm2_{i}.csv')\n",
    "        df['aminograpes_plain'], df['aminograpes'], used = aminograpes_scores(r, df.mutant); df['struct_used'] = used; df.to_csv(f, index=False)\n",
    "    rows.append((i, bool(df.struct_used.iloc[0]), spearmanr(df.esm2, df.DMS_score).correlation,\n",
    "                 spearmanr(df.aminograpes_plain, df.DMS_score).correlation, spearmanr(df.aminograpes, df.DMS_score).correlation))\n",
    "    print(f\"{i[:36]:36s} struct={'Y' if rows[-1][1] else 'N'} | ESM2 {rows[-1][2]:6.3f} | alignment {rows[-1][3]:6.3f} | AMINOGRAPES {rows[-1][4]:6.3f}\")\n",
    "res = pd.DataFrame(rows, columns=['DMS_id', 'struct', 'esm2_ours', 'plain', 'aminograpes']).set_index('DMS_id')\n",
    "res.to_csv(f'{OUT}/aminograpes_viral_per_assay.csv')\n",
    "print(f\"\\nstructures used for {int(res.struct.sum())} of 31 assays\")\n",
    "print(f\"MEAN: our ESM2 {res.esm2_ours.mean():.4f} | alignment only {res.plain.mean():.4f} | AMINOGRAPES {res.aminograpes.mean():.4f}   (expected 0.2693 | 0.3926 | 0.4304)\")\n",
    "dp = max(abs(res.plain[k] - v) for k, v in EXPECTED_PLAIN.items()); da = max(abs(res.aminograpes[k] - v) for k, v in EXPECTED_AMINOGRAPES.items())\n",
    "print(f\"largest per-assay difference from the published run: alignment {dp:.4f}, Aminograpes {da:.4f}  ->\", 'MATCHES' if max(dp, da) < 0.006 else 'CHECK: differs from published run')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab08302d",
   "metadata": {},
   "source": [
    "## Step 7. Compare with the official ProteinGym scores of other models\n",
    "This scores all 95 models that have scores for every viral assay, the same way as Aminograpes (plain mean of per-assay Spearman)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d732a96",
   "metadata": {},
   "outputs": [],
   "source": [
    "root = f'{P}/zs'\n",
    "files = {os.path.splitext(os.path.basename(p))[0]: p for p in glob.glob(f'{root}/**/*.csv', recursive=True)}\n",
    "assert all(i in files for i in viral.DMS_id), 'some official score files missing'\n",
    "SKIP = {'DMS_score', 'DMS_score_bin', 'mutant', 'mutated_sequence', 'DMS_id'}\n",
    "rows_ = {}\n",
    "for i, f in zip(viral.DMS_id, viral.DMS_filename):\n",
    "    g = pd.read_csv(idx[f], usecols=['mutant', 'DMS_score']).set_index('mutant').DMS_score\n",
    "    d = pd.read_csv(files[i]).drop_duplicates('mutant').set_index('mutant').reindex(g.index)\n",
    "    r_ = {}\n",
    "    for c in d.columns:\n",
    "        if c in SKIP or not pd.api.types.is_numeric_dtype(d[c]): continue\n",
    "        ok = d[c].notna()\n",
    "        if ok.mean() >= 0.9: r_[c] = spearmanr(d[c][ok], g[ok]).correlation\n",
    "    rows_[i] = r_\n",
    "Mo = pd.DataFrame(rows_).T.dropna(axis=1).reindex(res.index)\n",
    "Mo.to_csv(f'{OUT}/official_viral_spearman_all_models.csv')\n",
    "means = Mo.mean().sort_values(ascending=False)\n",
    "print(len(means), 'official models with scores for all 31 viral assays\\n')\n",
    "print('Top 10 official models (plain mean over 31 viral assays):'); print(means.head(10).round(4).to_string())\n",
    "n_above = int((means > res.aminograpes.mean()).sum())\n",
    "print(f'\\nOfficial models scoring above Aminograpes: {n_above} of {len(means)}')\n",
    "print(f'Official ESM2-650M mean: {Mo.ESM2_650M.mean():.4f} (expected 0.2687) | RSALOR {Mo.RSALOR.mean():.4f} (0.4797) | GEMME {Mo.GEMME.mean():.4f} | PoET {Mo.PoET.mean():.4f}')\n",
    "print('Our ESM2 rerun vs official: per-assay agreement r = %.3f (expected 1.000)' % np.corrcoef(Mo.ESM2_650M, res.esm2_ours)[0, 1])\n",
    "print(means[[c for c in means.index if c.startswith('ESM2')]].round(4).to_string())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "381ea828",
   "metadata": {},
   "source": [
    "## Step 8. Is the win real? Confidence intervals\n",
    "Bootstrap over assays (10,000 resamples, fixed seed). Expected: Aminograpes minus ESM2-650M **+0.1617, CI [0.0775, 0.2450], 25 of 31**; Aminograpes minus RSALOR about **-0.049**, 7 of 31."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8817185",
   "metadata": {},
   "outputs": [],
   "source": [
    "def boot(rng, d): \n",
    "    b = np.array([rng.choice(d, len(d)).mean() for _ in range(10000)]); return np.percentile(b, [2.5, 97.5])\n",
    "rng = np.random.default_rng(0)\n",
    "for nm, a, b in [('Aminograpes - alignment only (burial effect)', res.aminograpes, res.plain),\n",
    "                 ('Aminograpes - official ESM2-650M', res.aminograpes, Mo.ESM2_650M),\n",
    "                 ('Aminograpes - official RSALOR', res.aminograpes, Mo.RSALOR)]:\n",
    "    d = (a - b).values; lo, hi = boot(rng, d)\n",
    "    print(f'{nm}: {d.mean():+.4f}, 95% CI [{lo:.4f}, {hi:.4f}], wins {(d > 0).sum()}/31, Wilcoxon p={wilcoxon(d).pvalue:.4f}')\n",
    "d = (res.plain - Mo.ESM2_650M).values; lo, hi = boot(np.random.default_rng(0), d)\n",
    "print(f'alignment only - official ESM2-650M: {d.mean():+.4f}, 95% CI [{lo:.4f}, {hi:.4f}], wins {(d > 0).sum()}/31')\n",
    "print('\\nAssays where official ESM2-650M beats Aminograpes:', list(res.index[(res.aminograpes - Mo.ESM2_650M) <= 0]))\n",
    "print('Note: bootstrap CIs can differ in the 3rd decimal if the comparisons are run in a different order.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1079daea",
   "metadata": {},
   "source": [
    "## Step 9 (optional). The blend with ESM2: choose the weight on NON-viral assays\n",
    "This picks the ESM2 blend weight using 30 small non-viral assays, then applies it, frozen, to the viral assays. Expected best weight **0.7**, viral blend **0.3764** (+0.1077, CI [0.0781, 0.1397], 27 of 31). The viral numbers in the table are for display only and are never used to choose the weight."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf69c6e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "cand = ref[(ref['taxon'] != 'Virus') & (ref['seq_len'] <= 400)].sample(frac=1, random_state=0)\n",
    "tune = []\n",
    "for _, r in cand.iterrows():\n",
    "    if len(pd.read_csv(idx[r['DMS_filename']], usecols=['mutant'])) <= 6000: tune.append(r)\n",
    "    if len(tune) == 30: break\n",
    "z = lambda x: (x - x.mean()) / (x.std() + 1e-9)\n",
    "tdfs = []\n",
    "for r in tune:\n",
    "    i = r['DMS_id']; f = f'{OUT}/nv_{i}.csv'\n",
    "    if os.path.exists(f): tdfs.append(pd.read_csv(f)); continue\n",
    "    w_t = struct_weights(r)\n",
    "    if w_t is None: print('no usable structure, excluded:', i); continue\n",
    "    df = pd.read_csv(idx[r['DMS_filename']], usecols=['mutant', 'DMS_score'])\n",
    "    df['esm2'] = esm2_scores(r['target_seq'], list(df.mutant))\n",
    "    logf, qf = build_logf(idx[r['MSA_filename']], idx[r['weight_file_name']]); tmap = map_positions(r['target_seq'], qf)\n",
    "    out_ = []\n",
    "    for x in df.mutant:\n",
    "        b = 0.0\n",
    "        for m in x.split(':'):\n",
    "            wt, p, mt = m[0], int(m[1:-1]) - 1, m[-1]; q = tmap[p]\n",
    "            if q >= 0: b += w_t[p] * (logf[q, aidx[mt]] - logf[q, aidx[wt]])\n",
    "        out_.append(b)\n",
    "    df['aminograpes'] = out_; df.to_csv(f, index=False); tdfs.append(df)\n",
    "    print(f'{i[:40]:40s} esm2={spearmanr(df.esm2, df.DMS_score).correlation:6.3f} aminograpes={spearmanr(df.aminograpes, df.DMS_score).correlation:6.3f}')\n",
    "\n",
    "Wg = np.round(np.arange(0, 1.01, 0.1), 1)\n",
    "curve = lambda dfs: np.array([np.mean([spearmanr(w * z(d.esm2) + (1 - w) * z(d.aminograpes), d.DMS_score).correlation for d in dfs]) for w in Wg])\n",
    "cn = curve(tdfs); w_star = float(Wg[int(np.argmax(cn))])\n",
    "vd = {i: pd.read_csv(f'{OUT}/both_{i}.csv') for i in viral.DMS_id}\n",
    "S = pd.DataFrame({i: [spearmanr(w * z(d.esm2) + (1 - w) * z(d.aminograpes), d.DMS_score).correlation for w in Wg] for i, d in vd.items()}, index=Wg)\n",
    "print(f'\\nNON-VIRAL ({len(tdfs)} assays): best weight on ESM2 = {w_star}  (expected 0.7)')\n",
    "print('weight on ESM2 :', ' '.join(f'{w:5.1f}' for w in Wg))\n",
    "print('non-viral mean :', ' '.join(f'{c:5.3f}' for c in cn))\n",
    "print('viral mean     :', ' '.join(f'{c:5.3f}' for c in S.mean(axis=1)), '  <- display only')\n",
    "fin = S.loc[w_star]; alo = S.loc[0.0]\n",
    "print(f'\\nFINAL blend (weight {w_star}): {fin.mean():.4f}   (expected 0.3764)')\n",
    "rng = np.random.default_rng(0)\n",
    "for nm, a, b in [('blend - official ESM2-650M', fin, Mo.ESM2_650M), ('blend - Aminograpes alone', fin, alo), ('blend - official RSALOR', fin, Mo.RSALOR)]:\n",
    "    d = (a - b).values; lo, hi = boot(rng, d)\n",
    "    print(f'{nm}: {d.mean():+.4f}, 95% CI [{lo:.4f}, {hi:.4f}], wins {(d > 0).sum()}/31, p={wilcoxon(d).pvalue:.4f}')\n",
    "up = viral.set_index('DMS_id').UniProt_ID\n",
    "lopo = pd.Series({i: S.loc[S[[j for j in S.columns if up[j] != up[i]]].mean(axis=1).idxmax(), i] for i in S.columns})\n",
    "print(f'sensitivity: weight chosen on other viral proteins (leave-one-protein-out): {lopo.mean():.4f}  (expected about 0.4408)')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6962786",
   "metadata": {},
   "source": [
    "## Step 10 (optional). Ensemble of public models (this is NOT Aminograpes)\n",
    "This is context only. It stacks other groups' official ProteinGym models, chooses which ones using the 186 non-viral assays, and then tests on viral. Expected: ensemble **0.5732**, selected models VenusREM, ProSST-4096, S3F_MSA, ProSST-1024, RSALOR.\n",
    "Then it asks: does adding the Aminograpes score help? Expected: no (chosen weight 0)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b40eaf50",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import rankdata\n",
    "isviral = ref.set_index('DMS_id').taxon.eq('Virus')\n",
    "SKIP2 = {'mutant', 'mutated_sequence', 'DMS_score', 'DMS_score_bin'}\n",
    "data, comp = {}, {}\n",
    "for i in ref.DMS_id:\n",
    "    if i not in files: continue\n",
    "    d = pd.read_csv(files[i]).dropna(subset=['DMS_score'])\n",
    "    if not isviral[i] and len(d) > 3000: d = d.sample(3000, random_state=0)\n",
    "    cols = [c for c in d.columns if c not in SKIP2 and pd.api.types.is_numeric_dtype(d[c])]\n",
    "    X_ = d[cols]; comp[i] = set(X_.columns[X_.notna().all()])\n",
    "    data[i] = (((X_ - X_.mean()) / (X_.std() + 1e-9)).fillna(0).astype('float32'), d.DMS_score.values)\n",
    "VI = [i for i in data if isviral[i]]; NV = [i for i in data if not isviral[i]]\n",
    "cols_all = sorted(set.intersection(*[set(Z.columns) for Z, _ in data.values()]))\n",
    "pool = [m for m in cols_all if all(m in comp[i] for i in VI) and np.mean([m in comp[i] for i in NV]) >= 0.9]\n",
    "print(f'{len(NV)} non-viral / {len(VI)} viral assays; candidate pool = {len(pool)} models')\n",
    "A_ = {i: data[i][0][pool].values for i in data}\n",
    "Yr = {i: (lambda r_: (r_ - r_.mean()) / (r_.std() + 1e-12))(rankdata(data[i][1])) for i in data}\n",
    "col = {m: k for k, m in enumerate(pool)}\n",
    "def rho_(i, members):\n",
    "    s = A_[i][:, [col[m] for m in members]].sum(1); r_ = rankdata(s)\n",
    "    return float(np.dot((r_ - r_.mean()) / (r_.std() + 1e-12), Yr[i]) / len(r_))\n",
    "mean_rho = lambda ids_, members: np.mean([rho_(i, members) for i in ids_])\n",
    "sel, best = [], -1.0\n",
    "while len(sel) < 5:\n",
    "    s_, m_ = max((mean_rho(NV, sel + [m]), m) for m in pool if m not in sel)\n",
    "    if s_ <= best + 0.002: break\n",
    "    sel.append(m_); best = s_; print(f'  + {m_:22s} -> non-viral mean {s_:.4f}')\n",
    "print('selected on non-viral only:', sel)\n",
    "fin_e = pd.Series({i: rho_(i, sel) for i in VI}); Mv = Mo.reindex(VI)\n",
    "print(f'\\nVIRAL: ensemble {fin_e.mean():.4f} (expected 0.5732) | best single official {Mv.mean().max():.4f} ({Mv.mean().idxmax()}) | Aminograpes {res.aminograpes.mean():.4f}')\n",
    "rng = np.random.default_rng(0)\n",
    "for nm, s_ in [('ESM2-650M', Mv.ESM2_650M), ('RSALOR', Mv.RSALOR), ('PoET', Mv.PoET), (f'best single ({Mv.mean().idxmax()})', Mv[Mv.mean().idxmax()])]:\n",
    "    d = (fin_e - s_).values; lo, hi = boot(rng, d)\n",
    "    print(f'ensemble - {nm}: {d.mean():+.4f}, 95% CI [{lo:.4f}, {hi:.4f}], wins {(d > 0).sum()}/{len(d)}')\n",
    "\n",
    "# does adding the Aminograpes score help the ensemble?\n",
    "WA = [0, 0.5, 1, 2, 3]\n",
    "def per_assay(i, afr_df):\n",
    "    d = pd.read_csv(files[i]).drop_duplicates('mutant').set_index('mutant')\n",
    "    a = afr_df.drop_duplicates('mutant').set_index('mutant')\n",
    "    j = d[sel + ['DMS_score']].join(a[['aminograpes']], how='inner').dropna(); base = sum(z(j[m]) for m in sel)\n",
    "    return [spearmanr(base + w * z(j.aminograpes), j.DMS_score).correlation for w in WA]\n",
    "nv = pd.DataFrame({os.path.basename(p)[3:-4]: per_assay(os.path.basename(p)[3:-4], pd.read_csv(p)) for p in glob.glob(f'{OUT}/nv_*.csv')}, index=WA)\n",
    "vi = pd.DataFrame({i: per_assay(i, vd[i]) for i in VI}, index=WA)\n",
    "print('\\nweight on Aminograpes score :', ' '.join(f'{w:6.1f}' for w in WA))\n",
    "print('non-viral mean         :', ' '.join(f'{c:6.4f}' for c in nv.mean(axis=1)))\n",
    "print('viral mean (display)   :', ' '.join(f'{c:6.4f}' for c in vi.mean(axis=1)))\n",
    "print('weight chosen on non-viral:', float(nv.mean(axis=1).idxmax()), ' (expected 0.0: Aminograpes adds nothing to the ensemble)')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a8effeb",
   "metadata": {},
   "source": [
    "## Step 11. Save your results\n",
    "Zips every per-assay score file and summary table so you can download them and compare with the published run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "380fbcb7",
   "metadata": {},
   "outputs": [],
   "source": [
    "shutil.make_archive('/content/aminograpes_results', 'zip', OUT)\n",
    "print('saved /content/aminograpes_results.zip')\n",
    "try:\n",
    "    from google.colab import files as _f; _f.download('/content/aminograpes_results.zip')\n",
    "except Exception as e: print('download it from the Files panel on the left')"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}