| """Shared feature code. train.py and predict.py BOTH import this file, |
| so features are guaranteed to be identical at training and inference time.""" |
| import time |
| import zlib |
| from collections import Counter |
|
|
| import numpy as np |
|
|
| AA = "ACDEFGHIKLMNPQRSTVWY" |
| _LUT = np.full(256, -1, dtype=np.int16) |
| for _i, _a in enumerate(AA): |
| _LUT[ord(_a)] = _i |
|
|
| N_MORGAN = 1024 |
|
|
|
|
| |
| def _smiles_vec_factory(): |
| from rdkit import Chem, RDLogger |
| from rdkit.Chem import (Descriptors, MACCSkeys, rdFingerprintGenerator, |
| rdMolDescriptors) |
| RDLogger.DisableLog("rdApp.*") |
| gen = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=N_MORGAN) |
| n_desc = 12 |
| dim = N_MORGAN + 167 + n_desc |
|
|
| def vec(smi): |
| out = np.zeros(dim, dtype=np.float32) |
| try: |
| mol = Chem.MolFromSmiles(smi) |
| if mol is None: |
| return out |
| out[:N_MORGAN] = gen.GetCountFingerprintAsNumPy(mol) |
| out[N_MORGAN:N_MORGAN + 167] = np.array(list(MACCSkeys.GenMACCSKeys(mol)), dtype=np.float32) |
| d = [ |
| Descriptors.MolWt(mol), Descriptors.MolLogP(mol), Descriptors.TPSA(mol), |
| Descriptors.NumHDonors(mol), Descriptors.NumHAcceptors(mol), |
| Descriptors.NumRotatableBonds(mol), rdMolDescriptors.CalcNumRings(mol), |
| rdMolDescriptors.CalcNumAromaticRings(mol), mol.GetNumHeavyAtoms(), |
| Chem.GetFormalCharge(mol), rdMolDescriptors.CalcFractionCSP3(mol), |
| len(Chem.GetMolFrags(mol)), |
| ] |
| out[N_MORGAN + 167:] = np.array(d, dtype=np.float32) |
| except Exception: |
| pass |
| return out |
|
|
| return vec, dim |
|
|
|
|
| def smiles_features(unique_smiles): |
| """-> (matrix [n_unique, dim]) in the same order as unique_smiles""" |
| vec, dim = _smiles_vec_factory() |
| return np.stack([vec(s) for s in unique_smiles]).astype(np.float32) |
|
|
|
|
| |
| def _seq_vec(s): |
| a = _LUT[np.frombuffer(s.encode("ascii", "ignore"), dtype=np.uint8)] |
| a = a[a >= 0].astype(np.int64) |
| n = len(a) |
| comp = np.bincount(a, minlength=20) / max(n, 1) |
| if n > 1: |
| di = np.bincount(a[:-1] * 20 + a[1:], minlength=400) / (n - 1) |
| else: |
| di = np.zeros(400) |
| return np.concatenate([[len(s), np.log1p(len(s))], comp, di]).astype(np.float32) |
|
|
|
|
| def seq_features(unique_seqs): |
| return np.stack([_seq_vec(s) for s in unique_seqs]).astype(np.float32) |
|
|
|
|
| |
| def esm_embed(seqs, model_dir, fallback=None, max_len=800, budget_s=None, |
| batch_tokens=4096, n_threads=2, device="cpu", log=print, layers=None): |
| """Mean-pooled ESM2 embedding for each sequence (list of str). |
| Short sequences are done first. If budget_s runs out, the remaining |
| sequences get `fallback` (the training-mean embedding) so we never time out.""" |
| import torch |
| from transformers import AutoTokenizer, EsmModel |
| torch.set_num_threads(n_threads) |
| tok = AutoTokenizer.from_pretrained(model_dir) |
| model = EsmModel.from_pretrained(model_dir).to(device).eval() |
| dim = model.config.hidden_size * (len(layers) if layers else 1) |
| out = np.full((len(seqs), dim), np.nan, dtype=np.float32) |
| order = np.argsort([len(s) for s in seqs]) |
| t0 = time.time() |
| i = 0 |
| while i < len(order): |
| if budget_s is not None and time.time() - t0 > budget_s: |
| log(f"[esm] time budget hit after {i}/{len(order)} sequences; using fallback for the rest") |
| break |
| j, maxl = i, 0 |
| while j < len(order): |
| l = min(len(seqs[order[j]]), max_len) + 2 |
| if j > i and (j - i + 1) * max(maxl, l) > batch_tokens: |
| break |
| maxl = max(maxl, l) |
| j += 1 |
| idx = order[i:j] |
| batch = [seqs[k][:max_len] for k in idx] |
| enc = tok(batch, return_tensors="pt", padding=True) |
| enc = {k: v.to(device) for k, v in enc.items()} |
| with torch.inference_mode(): |
| o = model(**enc, output_hidden_states=bool(layers)) |
| m = enc["attention_mask"].clone().float() |
| lens = enc["attention_mask"].sum(1) |
| m[:, 0] = 0 |
| m[torch.arange(len(idx), device=m.device), lens - 1] = 0 |
| hs = [o.hidden_states[l] for l in layers] if layers else [o.last_hidden_state] |
| e = torch.cat([(h * m.unsqueeze(-1)).sum(1) / m.sum(1, keepdim=True).clamp(min=1) for h in hs], dim=1) |
| out[idx] = e.cpu().numpy() |
| i = j |
| if (i // 50) != ((i - len(idx)) // 50): |
| log(f"[esm] {i}/{len(order)} done, {time.time() - t0:.0f}s") |
| missing = np.isnan(out[:, 0]) |
| if missing.any(): |
| out[missing] = fallback if fallback is not None else 0.0 |
| return out |
|
|
|
|
| |
| def cluster_sequences(seqs, k=6, mod=2, min_frac=0.3): |
| """Cheap grouping of similar sequences (mutants / close homologs). |
| Used ONLY to make an honest local validation split. Returns a group id per sequence.""" |
| order = sorted(range(len(seqs)), key=lambda i: -len(seqs[i])) |
| kmer2leaders = {} |
| leader_of = [0] * len(seqs) |
| for i in order: |
| s = seqs[i] |
| kms = {s[j:j + k] for j in range(len(s) - k + 1) |
| if zlib.crc32(s[j:j + k].encode()) % mod == 0} |
| hits = Counter() |
| for km in kms: |
| for L in kmer2leaders.get(km, ()): |
| hits[L] += 1 |
| best = None |
| if hits: |
| L, c = hits.most_common(1)[0] |
| if c >= min_frac * max(len(kms), 1): |
| best = L |
| if best is None: |
| leader_of[i] = i |
| for km in kms: |
| kmer2leaders.setdefault(km, []).append(i) |
| else: |
| leader_of[i] = best |
| return np.array(leader_of) |
|
|
|
|
| def cluster_by_cosine(seqs, thr=0.35, max_frac=0.2, log=print): |
| """Group sequences that are even remotely related (3-mer cosine similarity >= thr), |
| so a local holdout looks like the real validation set (whose enzymes are unrelated to training). |
| If one giant group appears, the threshold is raised automatically.""" |
| from sklearn.feature_extraction.text import CountVectorizer |
| from sklearn.preprocessing import normalize |
| from scipy.sparse import coo_matrix |
| from scipy.sparse.csgraph import connected_components |
| cv = CountVectorizer(analyzer="char", ngram_range=(3, 3), lowercase=False, dtype=np.float32) |
| A = normalize(cv.fit_transform(seqs)).tocsr() |
| AT = A.T.tocsr() |
| n = A.shape[0] |
| rows, cols, vals = [], [], [] |
| for i in range(0, n, 500): |
| S = (A[i:i + 500] @ AT).toarray() |
| r, c = np.nonzero(S >= 0.30) |
| keep = (r + i) < c |
| rows.append(r[keep] + i); cols.append(c[keep]); vals.append(S[r[keep], c[keep]]) |
| rows, cols, vals = np.concatenate(rows), np.concatenate(cols), np.concatenate(vals) |
| t = thr |
| while True: |
| m = vals >= t |
| g = coo_matrix((np.ones(int(m.sum())), (rows[m], cols[m])), shape=(n, n)) |
| _, lab = connected_components(g, directed=False) |
| big = np.bincount(lab).max() / n |
| if big <= max_frac or t >= 0.9: |
| break |
| t += 0.05 |
| log(f"grouping threshold used: {t:.2f} | groups: {len(set(lab))} | biggest group: {big:.1%} of sequences") |
| return lab |
|
|
|
|
| def _wmean(Y, w): |
| """weighted mean over axis 1 that ignores NaN labels. Y [b,k,2], w [b,k] -> [b,2] (NaN if no valid label)""" |
| ok = ~np.isnan(Y) |
| ww = w[:, :, None] * ok |
| den = ww.sum(1) |
| num = (np.where(ok, Y, 0.0) * ww).sum(1) |
| return np.where(den > 0, num / np.maximum(den, 1e-12), np.nan) |
|
|
|
|
| |
| def te_stats(df): |
| """per-substrate: row count, mean/std of log10_kcat and log10_km, plus ENZYME-BALANCED means |
| (average over enzyme groups instead of rows, so enzymes with hundreds of mutants do not dominate).""" |
| g = df.groupby("smiles").agg(n=("log10_kcat", "size"), |
| kc=("log10_kcat", "mean"), km=("log10_km", "mean"), |
| kc_sd=("log10_kcat", "std"), km_sd=("log10_km", "std")) |
| if "_g" in df.columns: |
| d = df.groupby(["smiles", "_g"]).agg(kc=("log10_kcat", "mean"), km=("log10_km", "mean")).reset_index() |
| bal = d.groupby("smiles").agg(kcb=("kc", "mean"), kmb=("km", "mean"), ng=("kc", "size")) |
| g = g.join(bal) |
| else: |
| g["kcb"], g["kmb"], g["ng"] = g["kc"], g["km"], g["n"] |
| return g |
|
|
|
|
| def te_apply(smiles, stats): |
| s = stats.reindex(smiles) |
| return np.column_stack([s["n"].fillna(0).values, s["kc"].values, s["km"].values, |
| s["kc_sd"].values, s["km_sd"].values, |
| s["kcb"].values, s["kmb"].values, s["ng"].fillna(0).values]).astype(np.float32) |
|
|
|
|
| def sub_neighbor_feats(q_bits, ref_bits, ref_vals, k=5): |
| """For each query substrate: how similar are the k most similar TRAINING substrates |
| (Tanimoto on Morgan bits) and what are their weighted-mean kcat / km. |
| This helps for the ~1/3 of validation substrates never seen in training.""" |
| nq = q_bits.shape[0] |
| out = np.full((nq, 4), np.nan, dtype=np.float32) |
| rs = ref_bits.sum(1) |
| k = min(k, ref_bits.shape[0] - 1) |
| for i in range(0, nq, 512): |
| qb = q_bits[i:i + 512] |
| inter = qb @ ref_bits.T |
| union = qb.sum(1)[:, None] + rs[None, :] - inter |
| sim = inter / np.maximum(union, 1) |
| idx = np.argpartition(-sim, k, axis=1)[:, :k] |
| s = np.take_along_axis(sim, idx, 1) |
| w = s ** 2 + 1e-6 |
| vals = ref_vals[idx] |
| out[i:i + 512, 0] = s.max(1) |
| out[i:i + 512, 1:3] = _wmean(vals, w) |
| out[i:i + 512, 3] = s.mean(1) |
| return out |
|
|
|
|
| def te_full(smiles_values, smi_index, bits_u, stats, ref_bits): |
| """9 substrate-level features per row: exact-substrate stats (5) + similar-substrate stats (4). |
| smi_index: dict smiles -> row of bits_u; ref_bits aligned with stats.index.""" |
| exact = te_apply(smiles_values, stats) |
| uq = list(dict.fromkeys(smiles_values)) |
| q = bits_u[[smi_index[s] for s in uq]] |
| ref_vals = stats[["kc", "km"]].values.astype(np.float32) |
| nb = sub_neighbor_feats(q, ref_bits, ref_vals) |
| pos = {s: i for i, s in enumerate(uq)} |
| nb_rows = nb[[pos[s] for s in smiles_values]] |
| return np.hstack([exact, nb_rows]).astype(np.float32) |
|
|
|
|
| |
| def assemble(df, smi_unique, smi_mat, seq_unique, seq_mat, esm_mat=None, te=None): |
| smi_idx = {s: i for i, s in enumerate(smi_unique)} |
| seq_idx = {s: i for i, s in enumerate(seq_unique)} |
| si = df["smiles"].map(smi_idx).values |
| qi = df["sequence"].map(seq_idx).values |
| parts = [smi_mat[si], seq_mat[qi]] |
| if esm_mat is not None: |
| parts.append(esm_mat[qi]) |
| if te is not None: |
| parts.append(te) |
| return np.hstack(parts).astype(np.float32) |
|
|
|
|
| |
| def _knn_one(q_smiles, q_si, q_embn, ref_smiles, ref_si, ref_embn, ref_y, k=5, temp=5.0): |
| """For each query row, look at training enzymes that ESM thinks are similar. |
| cols 0-2: among training rows with the SAME substrate: best similarity, weighted mean kcat, weighted mean km |
| cols 3-5: among ALL training enzymes (per-enzyme mean targets): best similarity, mean kcat, mean km |
| q_si / ref_si = row index of each row's sequence inside q_embn / ref_embn (centered + length-normalised).""" |
| from collections import defaultdict |
| nq = len(q_smiles) |
| out = np.full((nq, 6), np.nan, dtype=np.float32) |
| q_si = np.asarray(q_si); ref_si = np.asarray(ref_si); ref_y = np.asarray(ref_y, dtype=np.float32) |
| ref_by = defaultdict(list) |
| for j, s in enumerate(ref_smiles): |
| ref_by[s].append(j) |
| q_by = defaultdict(list) |
| for i, s in enumerate(q_smiles): |
| q_by[s].append(i) |
| for s, qi in q_by.items(): |
| rj = ref_by.get(s) |
| if not rj: |
| continue |
| rj = np.array(rj); qi = np.array(qi) |
| sim = q_embn[q_si[qi]] @ ref_embn[ref_si[rj]].T |
| kk = min(k, len(rj)) |
| idx = np.argpartition(-sim, kk - 1, axis=1)[:, :kk] |
| sv = np.take_along_axis(sim, idx, 1) |
| w = np.exp(temp * sv) |
| y = ref_y[rj][idx] |
| out[qi, 0] = sv.max(1) |
| out[qi, 1:3] = _wmean(y, w) |
| |
| n_ref = ref_embn.shape[0] |
| ok = ~np.isnan(ref_y) |
| ysum = np.zeros((n_ref, 2)); cnt = np.zeros((n_ref, 2)) |
| np.add.at(ysum, ref_si, np.where(ok, ref_y, 0.0)); np.add.at(cnt, ref_si, ok.astype(float)) |
| have = np.nonzero(cnt.sum(1))[0] |
| ref_seq_y = np.where(cnt[have] > 0, ysum[have] / np.maximum(cnt[have], 1), np.nan).astype(np.float32) |
| R = ref_embn[have] |
| uq = np.unique(q_si) |
| G = np.full((len(uq), 3), np.nan, dtype=np.float32) |
| kk = min(k, len(have)) |
| for i in range(0, len(uq), 1024): |
| sim = q_embn[uq[i:i + 1024]] @ R.T |
| idx = np.argpartition(-sim, kk - 1, axis=1)[:, :kk] |
| sv = np.take_along_axis(sim, idx, 1) |
| w = np.exp(temp * sv) |
| G[i:i + 1024, 0] = sv.max(1) |
| G[i:i + 1024, 1:3] = _wmean(ref_seq_y[idx], w) |
| pos = {int(s): i for i, s in enumerate(uq)} |
| out[:, 3:6] = G[[pos[int(s)] for s in q_si]] |
| return out |
|
|
|
|
| def normalize_emb(emb, mean): |
| c = emb - mean |
| return (c / np.maximum(np.linalg.norm(c, axis=1, keepdims=True), 1e-6)).astype(np.float32) |
|
|
|
|
| def knn_feats(q_smiles, q_si, q_embn, ref_smiles, ref_si, ref_embn, ref_y, ks=(5,), temp=5.0): |
| """6 columns per value of k (see _knn_one).""" |
| return np.hstack([_knn_one(q_smiles, q_si, q_embn, ref_smiles, ref_si, ref_embn, ref_y, k=k, temp=temp) |
| for k in ks]) |
|
|
|
|
| |
| def fit_pca(emb, n): |
| mean = emb.mean(0) |
| _, _, vt = np.linalg.svd(emb - mean, full_matrices=False) |
| return {"mean": mean.astype(np.float32), "comps": vt[:n].astype(np.float32)} |
|
|
|
|
| def apply_variant(X0, kn, p, smi_dim, seq_dim, esm_dim, ks_all=(5,), pca=None, n_blocks=1, er=None): |
| """Shared by train.py and predict.py. X0 columns = [smiles | sequence | esm | substrate-stats].""" |
| esm_lo = smi_dim + seq_dim |
| drop = [] |
| if p.get("esm_blocks") is not None and esm_dim: |
| bs = esm_dim // n_blocks |
| for b in range(n_blocks): |
| if b not in p["esm_blocks"]: |
| drop += list(range(esm_lo + b * bs, esm_lo + (b + 1) * bs)) |
| if p.get("noesm") or p.get("esm_pca"): |
| drop += list(range(esm_lo, esm_lo + esm_dim)) |
| if p.get("nodi"): |
| drop += list(range(smi_dim + 22, smi_dim + seq_dim)) |
| keep = np.setdiff1d(np.arange(X0.shape[1]), np.array(drop, dtype=int)) |
| parts = [X0[:, keep]] |
| if p.get("esm_pca"): |
| parts.append((X0[:, esm_lo:esm_lo + esm_dim] - pca["mean"]) @ pca["comps"].T) |
| if p.get("knn"): |
| cols = [] |
| for k in p.get("ks", (5,)): |
| j = list(ks_all).index(k) |
| cols.extend(range(6 * j, 6 * j + 6)) |
| parts.append(kn[:, cols]) |
| if p.get("res"): |
| parts.append(er) |
| return np.hstack(parts).astype(np.float32) |
|
|
|
|
| |
| def mlp_prep_fit(X): |
| X = X.astype(np.float32).copy() |
| X[:, :N_MORGAN] = np.log1p(np.maximum(X[:, :N_MORGAN], 0)) |
| with np.errstate(all="ignore"): |
| mean = np.nanmean(X, axis=0) |
| std = np.nanstd(X, axis=0) |
| mean = np.nan_to_num(mean); std = np.nan_to_num(std) |
| std[std < 1e-6] = 1.0 |
| return {"mean": mean.astype(np.float32), "std": std.astype(np.float32)} |
|
|
|
|
| def mlp_prep(X, st): |
| X = X.astype(np.float32).copy() |
| X[:, :N_MORGAN] = np.log1p(np.maximum(X[:, :N_MORGAN], 0)) |
| Z = (X - st["mean"]) / st["std"] |
| return np.clip(np.nan_to_num(Z, nan=0.0), -8, 8).astype(np.float32) |
|
|
|
|
| def mlp_forward(Z, layers): |
| h = Z |
| for i, (W, b) in enumerate(layers): |
| h = h @ W + b |
| if i < len(layers) - 1: |
| h = np.maximum(h, 0) |
| return h |
|
|
|
|
| |
| def enzyme_effect_ref(ref_smiles, ref_si, ref_y, n_ref, stats, m=3.0): |
| """Split each label into substrate effect + enzyme effect. Returns, for every enzyme (row of the |
| embedding matrix), its average residual = label - (shrunk) substrate mean. NaN if the enzyme is absent.""" |
| ref_y = np.asarray(ref_y, dtype=np.float64) |
| gm = np.nanmean(ref_y, axis=0) |
| s = stats.reindex(ref_smiles) |
| n = s["n"].values[:, None] |
| mu = s[["kc", "km"]].values |
| shr = (n * mu + m * gm) / (n + m) |
| res = ref_y - shr |
| ok = ~np.isnan(res) |
| rs = np.zeros((n_ref, 2)); cnt = np.zeros((n_ref, 2)) |
| np.add.at(rs, ref_si, np.where(ok, res, 0.0)); np.add.at(cnt, ref_si, ok.astype(float)) |
| return np.where(cnt > 0, rs / np.maximum(cnt, 1), np.nan).astype(np.float32) |
|
|
|
|
| def enzyme_effect_feats(q_si, q_embn, ref_embn, ref_seq_res, k=10, temp=5.0): |
| """How much faster / slower than average (for its substrate) are the training enzymes that ESM finds |
| most similar to this one? 3 columns: best similarity, weighted mean residual kcat, residual km.""" |
| q_si = np.asarray(q_si) |
| has = np.nonzero(~np.isnan(ref_seq_res).all(1))[0] |
| R = ref_embn[has]; Y = ref_seq_res[has] |
| uq = np.unique(q_si) |
| G = np.full((len(uq), 3), np.nan, dtype=np.float32) |
| kk = min(k, len(has)) |
| for i in range(0, len(uq), 1024): |
| sim = q_embn[uq[i:i + 1024]] @ R.T |
| idx = np.argpartition(-sim, kk - 1, axis=1)[:, :kk] |
| sv = np.take_along_axis(sim, idx, 1) |
| w = np.exp(temp * sv) |
| G[i:i + 1024, 0] = sv.max(1) |
| G[i:i + 1024, 1:3] = _wmean(Y[idx], w) |
| pos = {int(s): i for i, s in enumerate(uq)} |
| return G[[pos[int(s)] for s in q_si]] |
|
|
