Use AminoGrapes
Complete Model Implementation, Colab Reproduction Protocol & Python Code
This page contains the complete implementation of the AminoGrapes viral fitness prediction model. Every cell from Aminograpes_reproduce.ipynb is provided below with executable Python code, mathematical breakdown, and benchmark verification.
In Plain Words
How AminoGrapes Works
ESM2-650M: A 650-million parameter protein masked language model trained by Meta AI. It serves as the primary language model baseline.
AminoGrapes: A closed-form scoring method that computes how conserved an amino acid is among evolutionary homologs (from a multiple sequence alignment) multiplied by how buried that residue is in the 3D protein fold (from AlphaFold2 predicted coordinates). No parameters are trained or tuned on viral targets.
Public Model Ensemble (Step 10): Combines 5 third-party ProteinGym models selected strictly on non-viral assays. It is included only for context; AminoGrapes is not used in that ensemble.
Expected Benchmark Results31 Viral Assays
ESM2-650M (official ProteinGym)0.2687
ESM2-650M (our notebook rerun)0.2693
Alignment only (Evolutionary log-odds)0.3926
AminoGrapes (Alignment × Burial)0.4304
AminoGrapes + ESM2 blend (λ=0.7)0.3764
Ensemble (5 public models, Step 10)0.5732
AminoGrapes vs ESM2-650M: +0.1617 (95% CI [0.0775, 0.2450]), winning 25 of 31 assays.
Execution Protocol in Google Colab
1. Hardware Selection
In Colab: Runtime > Change runtime type > T4 GPU. Needs ~20 GB free disk space.
2. Core Evaluation
Steps 1 to 8 reproduce all published primary metrics (~1 to 2 hours, mostly downloads and ESM2 inference).
3. Automatic Caching
Every completed assay is cached to disk. If Colab disconnects, re-running the cell will skip finished assays.
IMPLEMENTATION & CELLS
Step-by-Step Model Execution
Below is the verbatim source code for each of the 11 steps in the notebook. You can copy individual snippets or download the full notebook above.
Step 1
Install and Set Up
Environment initialization and GPU check
Installs required dependencies (biopython, transformers), configures workspace directories, and verifies GPU availability.
Requirement: Hardware acceleration (T4 GPU or better recommended for ESM2 scoring and PyTorch operations).
python • cell_1.py
Step 1 of 11!pip -q install biopython transformers
import os, re, glob, zipfile, shutil, numpy as np, pandas as pd
from scipy.stats import spearmanr, wilcoxon
WORK = '/content/aminograpes_repro'; P = f'{WORK}/pg'; OUT = f'{WORK}/out'
os.makedirs(P, exist_ok=True); os.makedirs(OUT, exist_ok=True)
import torch; print('GPU available:', torch.cuda.is_available())
assert torch.cuda.is_available(), 'Turn on a GPU: Runtime > Change runtime type'Expected Output
GPU available: True
Step 2
Download Public Data
Fetch ProteinGym v1.3 archives (~10–20 min)
Downloads official benchmark files from the Harvard Marks Lab ProteinGym v1.3 release:
- DMS_ProteinGym_substitutions.zip (assay experimental measurements)
- DMS_msa_weights.zip & DMS_msa_files.zip (evolutionary alignments)
- zero_shot_substitutions_scores.zip (official benchmark baselines)
- ProteinGym_AF2_structures.zip (predicted 3D coordinates)
python • cell_2.py
Step 2 of 11B = 'https://marks.hms.harvard.edu/proteingym/ProteinGym_v1.3'
def fetch(zipname, flag, dest=P):
if os.path.exists(f'{P}/{flag}'): return
print('downloading', zipname, '...')
assert os.system(f'curl -fsSL -o {P}/t.zip {B}/{zipname}') == 0, 'download failed: ' + zipname
zipfile.ZipFile(f'{P}/t.zip').extractall(dest); os.remove(f'{P}/t.zip'); open(f'{P}/{flag}', 'w').close()
fetch('DMS_ProteinGym_substitutions.zip', 'subs.ok')
fetch('DMS_msa_weights.zip', 'w.ok')
fetch('DMS_msa_files.zip', 'msa.ok')
fetch('zero_shot_substitutions_scores.zip', 'zs.ok', dest=f'{P}/zs')
if not os.path.exists(f'{P}/DMS_substitutions.csv'):
os.system(f'curl -fsSL -o {P}/DMS_substitutions.csv https://raw.githubusercontent.com/OATML-Markslab/ProteinGym/main/reference_files/DMS_substitutions.csv')
if not os.path.isdir(f'{P}/af2'):
ok = False
for u in [f'{B}/ProteinGym_AF2_structures.zip', 'https://marks.hms.harvard.edu/proteingym/ProteinGym_AF2_structures.zip']:
if os.system(f'curl -fsSL -o {P}/af2.zip {u}') == 0: ok = True; break
assert ok, 'structure download failed'
zipfile.ZipFile(f'{P}/af2.zip').extractall(f'{P}/af2'); os.remove(f'{P}/af2.zip')
print('done')Expected Output
done
Step 3
Pick the 31 Viral Assays
Filter benchmark reference to viral taxon
Filters the ProteinGym reference metadata down to the 31 viral deep mutational scanning (DMS) assays, indexing the alignment files, substitution datasets, and AlphaFold2 PDB coordinates.
python • cell_3.py
Step 3 of 11ref = pd.read_csv(f'{P}/DMS_substitutions.csv')
viral = ref[ref['taxon'] == 'Virus'].reset_index(drop=True)
print('viral assays:', len(viral)); assert len(viral) == 31
# index of data files (skip the structure and official-score folders so names do not collide)
idx = {os.path.basename(p): p for p in glob.glob(f'{P}/**/*', recursive=True)
if os.path.isfile(p) and not p.startswith((f'{P}/af2', f'{P}/zs'))}
pdbs = {os.path.splitext(os.path.basename(p))[0]: p for p in glob.glob(f'{P}/af2/**/*.pdb', recursive=True)}
print(len(pdbs), 'structure files')Expected Output
viral assays: 31 217 structure files
Step 4Core Architecture
The AminoGrapes Model Code
Complete closed-form scoring engine (zero fitted parameters on viral data)
This single cell is the complete AminoGrapes model. It executes in three core stages:
build_logf: Reads the multiple sequence alignment, down-weights near-duplicate sequences (similarity cutoff θ = 0.2), and computes position-wise log-odds amino acid frequencies with pseudocount λ = 0.01.struct_weights: Parses the AlphaFold2 PDB structure, calculates Shrake-Rupley solvent accessibility normalized by theoretical maximums (Tien et al. 2013), and computes residue burial: Burial = 1 − RSA.aminograpes_scores: Scores each mutant as:Score = ∑ [ Buriali × (log fi(mutant) − log fi(wildtype)) ]
python • cell_4.py
Step 4 of 11from Bio import Align
from Bio.Align import substitution_matrices
from Bio.PDB import PDBParser, ShrakeRupley
from Bio.SeqUtils import seq1
AA = 'ACDEFGHIKLMNPQRSTVWY'; aidx = {a: i for i, a in enumerate(AA)}
LUT = np.full(256, 21, dtype=np.uint8)
for i, a in enumerate(AA): LUT[ord(a)] = i
LUT[ord('-')] = 20
LAM, THETA, MAXB = 0.01, 0.2, 8e9 # pseudocount, clustering cutoff, memory budget
MAXASA = dict(A=129, R=274, N=195, D=193, C=167, E=223, Q=225, G=104, H=224, I=197,
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
def read_a2m(path):
seqs, cur = [], []
for line in open(path):
line = line.strip()
if line.startswith('>'):
if cur: seqs.append(''.join(cur)); cur = []
elif line: cur.append(line)
if cur: seqs.append(''.join(cur))
return [re.sub('[a-z.]', '', s) for s in seqs]
@torch.no_grad()
def gpu_weights(X, chunk=512):
N, L = X.shape # gap counts as its own state
Xt = torch.from_numpy(X.astype(np.int64)).cuda()
oh = torch.zeros(N, L, 21, dtype=torch.float16, device='cuda')
oh.scatter_(2, Xt.unsqueeze(2), 1.0); oh = oh.view(N, L * 21)
cnt = torch.zeros(N, device='cuda'); thr = (1 - THETA) * L
for s in range(0, N, chunk):
cnt[s:s + chunk] = ((oh[s:s + chunk] @ oh.T) >= thr).sum(1).float()
w = (1.0 / cnt.clamp(min=1)).cpu().numpy()
del oh, Xt; torch.cuda.empty_cache(); return w
def build_logf(msa_path, w_path):
seqs = read_a2m(msa_path); q = seqs[0]; L = len(q); seqs = [s for s in seqs if len(s) == L]
X = LUT[np.frombuffer(''.join(seqs).encode(), dtype=np.uint8).reshape(len(seqs), L)]
focus = np.array([c != '-' for c in q]); X = X[:, focus]; qf = q.replace('-', '')
keep = (X < 21).all(1) & ((X == 20).mean(1) <= 0.5)
wf = np.load(w_path).astype(np.float32)
if len(wf) == keep.sum(): X, w = X[keep], wf # ProteinGym weights match
elif len(wf) == len(X): w = wf
else: # recompute weights
X = X[keep]
if len(X) * X.shape[1] * 42 > MAXB:
n = int(MAXB / (X.shape[1] * 42)); X = X[np.sort(np.random.RandomState(0).choice(len(X), n, replace=False))]
w = gpu_weights(X)
c = np.stack([w @ (X == a) for a in range(20)], 1)
f = (1 - LAM) * c / np.maximum(c.sum(1, keepdims=True), 1e-9) + LAM / 20
return np.log(f), qf
aligner = Align.PairwiseAligner(); aligner.mode = 'global'
aligner.substitution_matrix = substitution_matrices.load('BLOSUM62')
aligner.open_gap_score = -10; aligner.extend_gap_score = -0.5; aligner.end_gap_score = 0.0
def map_positions(target, other):
aln = aligner.align(target, other)[0]; tmap = -np.ones(len(target), dtype=int)
for (t0, t1), (q0, q1) in zip(aln.aligned[0], aln.aligned[1]): tmap[t0:t1] = np.arange(q0, q1)
return tmap
def rsa_from_pdb(path):
model = PDBParser(QUIET=True).get_structure('x', path)[0]
ShrakeRupley().compute(model, level='R')
chain = next(model.get_chains()); seq, rsa = [], []
for res in chain:
if res.id[0] != ' ': continue
aa = seq1(res.get_resname())
if aa in MAXASA: seq.append(aa); rsa.append(min(1.0, res.sasa / MAXASA[aa]))
return ''.join(seq), np.array(rsa)
def find_struct(r):
for key in [r['DMS_id'], r['UniProt_ID']]:
if key in pdbs: return pdbs[key]
for k, p in pdbs.items():
if len(k) >= 6 and (r['UniProt_ID'] in k or k in r['DMS_id']): return p
return None
def struct_weights(r):
"burial = 1 - RSA per target position, or None if no usable structure"
tseq = r['target_seq']; sp = find_struct(r)
if not sp: return None
sseq, srsa = rsa_from_pdb(sp); sm = map_positions(tseq, sseq); m_ = sm >= 0
if m_.mean() < 0.3: return None
if np.mean([tseq[t] == sseq[sm[t]] for t in np.where(m_)[0]]) < 0.8: return None
rsa = np.full(len(tseq), np.nan); rsa[m_] = srsa[sm[m_]]; rsa[np.isnan(rsa)] = np.nanmean(rsa)
return 1 - rsa
def aminograpes_scores(r, mutants):
"returns (plain alignment score, Aminograpes score = alignment x burial, structure_used)"
tseq = r['target_seq']
logf, qf = build_logf(idx[r['MSA_filename']], idx[r['weight_file_name']])
tmap = map_positions(tseq, qf)
w_t = struct_weights(r); used = w_t is not None
if not used: w_t = np.full(len(tseq), 0.5) # constant weight: ranking equals plain alignment
plain, afr = [], []
for x in mutants:
a = b = 0.0
for m in x.split(':'):
wt, p, mt = m[0], int(m[1:-1]) - 1, m[-1]; q = tmap[p]
if q >= 0:
lor = logf[q, aidx[mt]] - logf[q, aidx[wt]]; a += lor; b += w_t[p] * lor
plain.append(a); afr.append(b)
return plain, afr, used
print('Aminograpes model code loaded')Expected Output
Aminograpes model code loaded
Step 5
Score Viral Assays with ESM2-650M
Baseline Language Model Evaluation
Scores the 31 viral assays with Meta's ESM2-650M masked language model baseline (facebook/esm2_t33_650M_UR50D) using a 1,022-residue sliding window. All computed scores are cached to disk so interrupted sessions resume immediately.
python • cell_5.py
Step 5 of 11from transformers import AutoTokenizer, EsmForMaskedLM
name = 'facebook/esm2_t33_650M_UR50D'
tok = AutoTokenizer.from_pretrained(name)
model = EsmForMaskedLM.from_pretrained(name, torch_dtype=torch.float16).cuda().eval()
aa_id = torch.tensor([tok.convert_tokens_to_ids(a) for a in AA]).cuda()
MAXL, BS = 1022, 16
@torch.no_grad()
def esm2_scores(seq, mutants):
L = len(seq); pos = sorted({int(m[1:-1]) - 1 for x in mutants for m in x.split(':')}); lp = {}
for s in range(0, len(pos), BS):
chunk = pos[s:s + BS]; wins = []
for p in chunk:
start = 0 if L <= MAXL else max(0, min(p - MAXL // 2, L - MAXL))
wins.append((seq[start:start + MAXL], p - start))
enc = tok([w for w, _ in wins], return_tensors='pt').to('cuda'); ids = enc['input_ids'].clone()
for r_, (w, rel) in enumerate(wins): ids[r_, rel + 1] = tok.mask_token_id
logits = model(input_ids=ids, attention_mask=enc['attention_mask']).logits
for r_, (w, rel) in enumerate(wins): lp[chunk[r_]] = logits[r_, rel + 1].float().log_softmax(-1)[aa_id].cpu()
out = []
for x in mutants:
s = 0.0
for m in x.split(':'):
wt, p, mt = m[0], int(m[1:-1]) - 1, m[-1]; s += (lp[p][aidx[mt]] - lp[p][aidx[wt]]).item()
out.append(s)
return out
for _, r in viral.iterrows():
f = f"{OUT}/esm2_{r['DMS_id']}.csv"
if os.path.exists(f): continue
df = pd.read_csv(idx[r['DMS_filename']], usecols=['mutant', 'DMS_score'])
df['esm2'] = esm2_scores(r['target_seq'], list(df.mutant)); df.to_csv(f, index=False)
print(f"{r['DMS_id']:45s} esm2 rho = {spearmanr(df.esm2, df.DMS_score).correlation:.3f}")Step 6
Run AminoGrapes on 31 Viral Assays
Execute scoring and verify match against published run
Executes AminoGrapes on all 31 viral assays, computing both alignment-only log-odds and structure-weighted AminoGrapes scores.
Expected Results: Plain alignment mean = 0.3926 | AminoGrapes mean = 0.4304 | Structures used for 29 of 31 assays.
python • cell_6.py
Step 6 of 11EXPECTED_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}
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}
rows = []
for _, r in viral.iterrows():
i = r['DMS_id']; f = f'{OUT}/both_{i}.csv'
if os.path.exists(f): df = pd.read_csv(f)
else:
df = pd.read_csv(f'{OUT}/esm2_{i}.csv')
df['aminograpes_plain'], df['aminograpes'], used = aminograpes_scores(r, df.mutant); df['struct_used'] = used; df.to_csv(f, index=False)
rows.append((i, bool(df.struct_used.iloc[0]), spearmanr(df.esm2, df.DMS_score).correlation,
spearmanr(df.aminograpes_plain, df.DMS_score).correlation, spearmanr(df.aminograpes, df.DMS_score).correlation))
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}")
res = pd.DataFrame(rows, columns=['DMS_id', 'struct', 'esm2_ours', 'plain', 'aminograpes']).set_index('DMS_id')
res.to_csv(f'{OUT}/aminograpes_viral_per_assay.csv')
print(f"\nstructures used for {int(res.struct.sum())} of 31 assays")
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)")
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())
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')Expected Output
structures used for 29 of 31 assays MEAN: our ESM2 0.2693 | alignment only 0.3926 | AMINOGRAPES 0.4304 (expected 0.2693 | 0.3926 | 0.4304) largest per-assay difference from the published run: alignment 0.0004, Aminograpes 0.0004 -> MATCHES
Step 7
Compare with 95 Official ProteinGym Models
Cross-model leaderboard evaluation across all 31 viral targets
Evaluates all 95 published models from the official ProteinGym zero-shot release that have complete scores across all 31 viral assays, computing plain per-assay Spearman correlations identical to the primary benchmark.
python • cell_7.py
Step 7 of 11root = f'{P}/zs'
files = {os.path.splitext(os.path.basename(p))[0]: p for p in glob.glob(f'{root}/**/*.csv', recursive=True)}
assert all(i in files for i in viral.DMS_id), 'some official score files missing'
SKIP = {'DMS_score', 'DMS_score_bin', 'mutant', 'mutated_sequence', 'DMS_id'}
rows_ = {}
for i, f in zip(viral.DMS_id, viral.DMS_filename):
g = pd.read_csv(idx[f], usecols=['mutant', 'DMS_score']).set_index('mutant').DMS_score
d = pd.read_csv(files[i]).drop_duplicates('mutant').set_index('mutant').reindex(g.index)
r_ = {}
for c in d.columns:
if c in SKIP or not pd.api.types.is_numeric_dtype(d[c]): continue
ok = d[c].notna()
if ok.mean() >= 0.9: r_[c] = spearmanr(d[c][ok], g[ok]).correlation
rows_[i] = r_
Mo = pd.DataFrame(rows_).T.dropna(axis=1).reindex(res.index)
Mo.to_csv(f'{OUT}/official_viral_spearman_all_models.csv')
means = Mo.mean().sort_values(ascending=False)
print(len(means), 'official models with scores for all 31 viral assays\n')
print('Top 10 official models (plain mean over 31 viral assays):'); print(means.head(10).round(4).to_string())
n_above = int((means > res.aminograpes.mean()).sum())
print(f'\nOfficial models scoring above Aminograpes: {n_above} of {len(means)}')
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}')
print('Our ESM2 rerun vs official: per-assay agreement r = %.3f (expected 1.000)' % np.corrcoef(Mo.ESM2_650M, res.esm2_ours)[0, 1])
print(means[[c for c in means.index if c.startswith('ESM2')]].round(4).to_string())Step 8
Confidence Intervals & Statistical Significance
10,000 bootstrap resamples and Wilcoxon signed-rank testing
Performs bootstrap resampling over assays (10,000 iterations, fixed random seed = 0) and computes non-parametric Wilcoxon signed-rank tests for statistical hypothesis validation.
Expected: AminoGrapes − ESM2-650M = +0.1617, 95% CI [0.0775, 0.2450], 25 of 31 wins, p = 0.0008.
python • cell_8.py
Step 8 of 11def boot(rng, d):
b = np.array([rng.choice(d, len(d)).mean() for _ in range(10000)]); return np.percentile(b, [2.5, 97.5])
rng = np.random.default_rng(0)
for nm, a, b in [('Aminograpes - alignment only (burial effect)', res.aminograpes, res.plain),
('Aminograpes - official ESM2-650M', res.aminograpes, Mo.ESM2_650M),
('Aminograpes - official RSALOR', res.aminograpes, Mo.RSALOR)]:
d = (a - b).values; lo, hi = boot(rng, d)
print(f'{nm}: {d.mean():+.4f}, 95% CI [{lo:.4f}, {hi:.4f}], wins {(d > 0).sum()}/31, Wilcoxon p={wilcoxon(d).pvalue:.4f}')
d = (res.plain - Mo.ESM2_650M).values; lo, hi = boot(np.random.default_rng(0), d)
print(f'alignment only - official ESM2-650M: {d.mean():+.4f}, 95% CI [{lo:.4f}, {hi:.4f}], wins {(d > 0).sum()}/31')
print('\nAssays where official ESM2-650M beats Aminograpes:', list(res.index[(res.aminograpes - Mo.ESM2_650M) <= 0]))
print('Note: bootstrap CIs can differ in the 3rd decimal if the comparisons are run in a different order.')Expected Output
Aminograpes - alignment only (burial effect): +0.0378, 95% CI [0.0152, 0.0610], wins 22/31, Wilcoxon p=0.0012 Aminograpes - official ESM2-650M: +0.1617, 95% CI [0.0775, 0.2450], wins 25/31, Wilcoxon p=0.0008 Aminograpes - official RSALOR: -0.0493, 95% CI [-0.0931, -0.0084], wins 7/31, Wilcoxon p=0.0177 alignment only - official ESM2-650M: +0.1239, 95% CI [0.0461, 0.2014], wins 23/31
Step 9Optional
ESM2 Blend Weight Tuning on Non-Viral Assays
Unsupervised zero-shot blend tuning (Optional)
Tunes the blend weight between AminoGrapes and ESM2 on a development subset of 30 non-viral assays, then applies that frozen parameter (λ = 0.7) to the 31 viral targets.
Note: Viral assays are never used to choose the weight, maintaining zero-shot integrity. Expected viral blend: 0.3764 (+0.1077 over ESM2-650M alone).
python • cell_9.py
Step 9 of 11cand = ref[(ref['taxon'] != 'Virus') & (ref['seq_len'] <= 400)].sample(frac=1, random_state=0)
tune = []
for _, r in cand.iterrows():
if len(pd.read_csv(idx[r['DMS_filename']], usecols=['mutant'])) <= 6000: tune.append(r)
if len(tune) == 30: break
z = lambda x: (x - x.mean()) / (x.std() + 1e-9)
tdfs = []
for r in tune:
i = r['DMS_id']; f = f'{OUT}/nv_{i}.csv'
if os.path.exists(f): tdfs.append(pd.read_csv(f)); continue
w_t = struct_weights(r)
if w_t is None: print('no usable structure, excluded:', i); continue
df = pd.read_csv(idx[r['DMS_filename']], usecols=['mutant', 'DMS_score'])
df['esm2'] = esm2_scores(r['target_seq'], list(df.mutant))
logf, qf = build_logf(idx[r['MSA_filename']], idx[r['weight_file_name']]); tmap = map_positions(r['target_seq'], qf)
out_ = []
for x in df.mutant:
b = 0.0
for m in x.split(':'):
wt, p, mt = m[0], int(m[1:-1]) - 1, m[-1]; q = tmap[p]
if q >= 0: b += w_t[p] * (logf[q, aidx[mt]] - logf[q, aidx[wt]])
out_.append(b)
df['aminograpes'] = out_; df.to_csv(f, index=False); tdfs.append(df)
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}')
Wg = np.round(np.arange(0, 1.01, 0.1), 1)
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])
cn = curve(tdfs); w_star = float(Wg[int(np.argmax(cn))])
vd = {i: pd.read_csv(f'{OUT}/both_{i}.csv') for i in viral.DMS_id}
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)
print(f'\nNON-VIRAL ({len(tdfs)} assays): best weight on ESM2 = {w_star} (expected 0.7)')
print('weight on ESM2 :', ' '.join(f'{w:5.1f}' for w in Wg))
print('non-viral mean :', ' '.join(f'{c:5.3f}' for c in cn))
print('viral mean :', ' '.join(f'{c:5.3f}' for c in S.mean(axis=1)), ' <- display only')
fin = S.loc[w_star]; alo = S.loc[0.0]
print(f'\nFINAL blend (weight {w_star}): {fin.mean():.4f} (expected 0.3764)')
rng = np.random.default_rng(0)
for nm, a, b in [('blend - official ESM2-650M', fin, Mo.ESM2_650M), ('blend - Aminograpes alone', fin, alo), ('blend - official RSALOR', fin, Mo.RSALOR)]:
d = (a - b).values; lo, hi = boot(rng, d)
print(f'{nm}: {d.mean():+.4f}, 95% CI [{lo:.4f}, {hi:.4f}], wins {(d > 0).sum()}/31, p={wilcoxon(d).pvalue:.4f}')
up = viral.set_index('DMS_id').UniProt_ID
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})
print(f'sensitivity: weight chosen on other viral proteins (leave-one-protein-out): {lopo.mean():.4f} (expected about 0.4408)')Step 10Optional
Public Models Ensemble Analysis
Contextual stack of published methods (Optional)
Combines top published models (VenusREM, ProSST-4096, S3F_MSA, ProSST-1024, RSALOR) chosen strictly on 186 non-viral assays and evaluated on viral targets (ensemble mean = 0.5732).
Key finding: When testing whether adding AminoGrapes improves this 5-model stack, the chosen weight on AminoGrapes is 0.0 (AminoGrapes adds nothing new to an already saturated multi-model ensemble).
python • cell_10.py
Step 10 of 11from scipy.stats import rankdata
isviral = ref.set_index('DMS_id').taxon.eq('Virus')
SKIP2 = {'mutant', 'mutated_sequence', 'DMS_score', 'DMS_score_bin'}
data, comp = {}, {}
for i in ref.DMS_id:
if i not in files: continue
d = pd.read_csv(files[i]).dropna(subset=['DMS_score'])
if not isviral[i] and len(d) > 3000: d = d.sample(3000, random_state=0)
cols = [c for c in d.columns if c not in SKIP2 and pd.api.types.is_numeric_dtype(d[c])]
X_ = d[cols]; comp[i] = set(X_.columns[X_.notna().all()])
data[i] = (((X_ - X_.mean()) / (X_.std() + 1e-9)).fillna(0).astype('float32'), d.DMS_score.values)
VI = [i for i in data if isviral[i]]; NV = [i for i in data if not isviral[i]]
cols_all = sorted(set.intersection(*[set(Z.columns) for Z, _ in data.values()]))
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]
print(f'{len(NV)} non-viral / {len(VI)} viral assays; candidate pool = {len(pool)} models')
A_ = {i: data[i][0][pool].values for i in data}
Yr = {i: (lambda r_: (r_ - r_.mean()) / (r_.std() + 1e-12))(rankdata(data[i][1])) for i in data}
col = {m: k for k, m in enumerate(pool)}
def rho_(i, members):
s = A_[i][:, [col[m] for m in members]].sum(1); r_ = rankdata(s)
return float(np.dot((r_ - r_.mean()) / (r_.std() + 1e-12), Yr[i]) / len(r_))
mean_rho = lambda ids_, members: np.mean([rho_(i, members) for i in ids_])
sel, best = [], -1.0
while len(sel) < 5:
s_, m_ = max((mean_rho(NV, sel + [m]), m) for m in pool if m not in sel)
if s_ <= best + 0.002: break
sel.append(m_); best = s_; print(f' + {m_:22s} -> non-viral mean {s_:.4f}')
print('selected on non-viral only:', sel)
fin_e = pd.Series({i: rho_(i, sel) for i in VI}); Mv = Mo.reindex(VI)
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}')
rng = np.random.default_rng(0)
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()])]:
d = (fin_e - s_).values; lo, hi = boot(rng, d)
print(f'ensemble - {nm}: {d.mean():+.4f}, 95% CI [{lo:.4f}, {hi:.4f}], wins {(d > 0).sum()}/{len(d)}')
# does adding the Aminograpes score help the ensemble?
WA = [0, 0.5, 1, 2, 3]
def per_assay(i, afr_df):
d = pd.read_csv(files[i]).drop_duplicates('mutant').set_index('mutant')
a = afr_df.drop_duplicates('mutant').set_index('mutant')
j = d[sel + ['DMS_score']].join(a[['aminograpes']], how='inner').dropna(); base = sum(z(j[m]) for m in sel)
return [spearmanr(base + w * z(j.aminograpes), j.DMS_score).correlation for w in WA]
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)
vi = pd.DataFrame({i: per_assay(i, vd[i]) for i in VI}, index=WA)
print('\nweight on Aminograpes score :', ' '.join(f'{w:6.1f}' for w in WA))
print('non-viral mean :', ' '.join(f'{c:6.4f}' for c in nv.mean(axis=1)))
print('viral mean (display) :', ' '.join(f'{c:6.4f}' for c in vi.mean(axis=1)))
print('weight chosen on non-viral:', float(nv.mean(axis=1).idxmax()), ' (expected 0.0: Aminograpes adds nothing to the ensemble)')Step 11
Save Results Archive
Export all per-assay CSV files and summary tables
Packages all generated predictions, per-assay correlation CSVs, and model matrices into a single zip archive (aminograpes_results.zip) for local inspection and verification.
python • cell_11.py
Step 11 of 11shutil.make_archive('/content/aminograpes_results', 'zip', OUT)
print('saved /content/aminograpes_results.zip')
try:
from google.colab import files as _f; _f.download('/content/aminograpes_results.zip')
except Exception as e: print('download it from the Files panel on the left')Expected Output
saved /content/aminograpes_results.zip
Ready to reproduce the evaluation?
Download the complete Aminograpes_reproduce.ipynb notebook to run the pipeline locally or in Google Colab.