| import numpy as np, pandas as pd, torch |
| from rdkit import Chem, RDLogger |
| from rdkit.Chem import rdFingerprintGenerator |
| from transformers import AutoTokenizer, EsmModel |
| from scipy.sparse import coo_matrix |
| from scipy.sparse.csgraph import connected_components |
|
|
| RDLogger.DisableLog("rdApp.*") |
| T = ["log10_kcat", "log10_km"] |
| _gen = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=2048) |
|
|
| def smiles_feat(smi): |
| m = Chem.MolFromSmiles(str(smi)) |
| if m is None: |
| return np.zeros(2048, dtype=np.float32) |
| return _gen.GetCountFingerprintAsNumPy(m).astype(np.float32) |
|
|
| @torch.no_grad() |
| def embed_unique(seqs, model_dir, max_len=1000, bs=16): |
| dev = "cuda" if torch.cuda.is_available() else "cpu" |
| tok = AutoTokenizer.from_pretrained(model_dir) |
| model = EsmModel.from_pretrained(model_dir).to(dev).eval() |
| if dev == "cuda": |
| model.half() |
| uniq = sorted(set(map(str, seqs)), key=len) |
| out = {} |
| for i in range(0, len(uniq), bs): |
| b = uniq[i:i + bs] |
| enc = tok(b, return_tensors="pt", padding=True, truncation=True, max_length=max_len).to(dev) |
| h = model(**enc).last_hidden_state.float() |
| m = enc["attention_mask"].unsqueeze(-1).float() |
| pooled = ((h * m).sum(1) / m.sum(1)).cpu().numpy() |
| for s, v in zip(b, pooled): |
| out[s] = v.astype(np.float32) |
| return out |
|
|
| def unit(M): |
| return M / (np.linalg.norm(M, axis=1, keepdims=True) + 1e-8) |
|
|
| def base_X(df, emb): |
| fp = {s: smiles_feat(s) for s in df["smiles"].unique()} |
| return np.hstack([ |
| np.stack([emb[s] for s in df["sequence"]]), |
| np.stack([fp[s] for s in df["smiles"]]), |
| np.log1p(df["sequence"].str.len().values)[:, None], |
| ]).astype(np.float32) |
|
|
| def smiles_stats(ref, q): |
| g = ref.groupby("smiles")[T].agg(["mean", "median", "std", "count"]) |
| g.columns = ["_".join(c) for c in g.columns] |
| out = pd.DataFrame({"smiles": list(q)}).join(g, on="smiles").drop(columns="smiles") |
| for t in T: |
| out[t + "_mean"] = out[t + "_mean"].fillna(ref[t].mean()) |
| out[t + "_median"] = out[t + "_median"].fillna(ref[t].median()) |
| out[t + "_std"] = out[t + "_std"].fillna(0) |
| out[t + "_count"] = out[t + "_count"].fillna(0) |
| return out.values.astype(np.float32) |
|
|
| class SeqKNN: |
| def __init__(self, ref, emb, mu, k=10): |
| tab = ref.groupby("sequence")[T].mean() |
| self.Y = tab.values.astype(np.float32) |
| self.M = unit(np.stack([emb[s] for s in tab.index]) - mu) |
| self.mu, self.k = mu, min(k, len(tab)) |
|
|
| def rows(self, seqs, emb): |
| seqs = list(seqs); uq = list(dict.fromkeys(seqs)); res = {} |
| Q = unit(np.stack([emb[s] for s in uq]) - self.mu) |
| for i in range(0, len(uq), 1000): |
| S = Q[i:i + 1000] @ self.M.T |
| idx = np.argpartition(-S, self.k - 1, axis=1)[:, :self.k] |
| sims = np.take_along_axis(S, idx, 1) |
| o = np.argsort(-sims, 1) |
| idx = np.take_along_axis(idx, o, 1); sims = np.take_along_axis(sims, o, 1) |
| w = np.exp((sims - sims[:, :1]) / 0.05); w /= w.sum(1, keepdims=True) |
| Yn = self.Y[idx] |
| f = np.hstack([sims[:, :1], sims.mean(1, keepdims=True), Yn[:, 0, :], |
| Yn.mean(1), (Yn * w[:, :, None]).sum(1)]) |
| for s, v in zip(uq[i:i + 1000], f): |
| res[s] = v |
| return np.stack([res[s] for s in seqs]).astype(np.float32) |
|
|
| def label_feats(ref, q_df, emb, mu): |
| knn = SeqKNN(ref, emb, mu) |
| return np.hstack([smiles_stats(ref, q_df["smiles"]), knn.rows(q_df["sequence"], emb)]) |
|
|
| def cluster_groups(uniq, emb, mu, thr=0.9): |
| U = unit(np.stack([emb[s] for s in uniq]) - mu); n = len(U); r_, c_ = [], [] |
| for i in range(0, n, 2000): |
| r, c = np.nonzero(U[i:i + 2000] @ U.T > thr) |
| r_.append(r + i); c_.append(c) |
| r_ = np.concatenate(r_); c_ = np.concatenate(c_) |
| _, lab = connected_components(coo_matrix((np.ones(len(r_)), (r_, c_)), shape=(n, n)), directed=False) |
| return dict(zip(uniq, lab)) |
|
|
| class SubKNN: |
| """Labels of the most similar training substrates (cosine on Morgan counts).""" |
| def __init__(self, ref, k=10): |
| tab = ref.groupby("smiles")[T].mean() |
| self.Y = tab.values.astype(np.float32) |
| self.M = unit(np.stack([smiles_feat(s) for s in tab.index])) |
| self.k = min(k, len(tab)) |
|
|
| def rows(self, smiles): |
| smiles = list(smiles); uq = list(dict.fromkeys(smiles)); res = {} |
| Q = unit(np.stack([smiles_feat(s) for s in uq])) |
| for i in range(0, len(uq), 1000): |
| S = Q[i:i + 1000] @ self.M.T |
| idx = np.argpartition(-S, self.k - 1, axis=1)[:, :self.k] |
| sims = np.take_along_axis(S, idx, 1) |
| o = np.argsort(-sims, 1) |
| idx = np.take_along_axis(idx, o, 1); sims = np.take_along_axis(sims, o, 1) |
| w = np.exp((sims - sims[:, :1]) / 0.1); w /= w.sum(1, keepdims=True) |
| Yn = self.Y[idx] |
| f = np.hstack([sims[:, :1], sims.mean(1, keepdims=True), Yn[:, 0, :], |
| Yn.mean(1), (Yn * w[:, :, None]).sum(1)]) |
| for s, v in zip(uq[i:i + 1000], f): |
| res[s] = v |
| return np.stack([res[s] for s in smiles]).astype(np.float32) |
|
|
| def label_feats(ref, q_df, emb, mu): |
| knn = SeqKNN(ref, emb, mu) |
| sub = SubKNN(ref) |
| return np.hstack([smiles_stats(ref, q_df["smiles"]), |
| knn.rows(q_df["sequence"], emb), |
| sub.rows(q_df["smiles"])]) |
|
|
| class PairKNN: |
| """Training rows with the SAME substrate, weighted by enzyme similarity.""" |
| def __init__(self, ref, emb, mu): |
| self.ref = ref.reset_index(drop=True) |
| self.mu = mu |
| self.vec = unit(np.stack([emb[s] for s in self.ref["sequence"]]) - mu) |
| self.Y = self.ref[T].values.astype(np.float32) |
| self.idx = self.ref.groupby("smiles").indices |
| self.gm = self.Y.mean(0) |
|
|
| def rows(self, q_df, emb): |
| Q = unit(np.stack([emb[s] for s in q_df["sequence"]]) - self.mu) |
| out = np.zeros((len(q_df), 8), dtype=np.float32) |
| for i, smi in enumerate(q_df["smiles"].values): |
| ix = self.idx.get(smi) |
| if ix is None: |
| out[i] = [0, 0, *self.gm, *self.gm, *self.gm] |
| continue |
| sims = self.vec[ix] @ Q[i] |
| Yg = self.Y[ix] |
| j = int(np.argmax(sims)) |
| w = np.exp((sims - sims[j]) / 0.05); w /= w.sum() |
| out[i] = [1, sims[j], *Yg[j], *(Yg * w[:, None]).sum(0), *Yg.mean(0)] |
| return out |
|
|
| def label_feats(ref, q_df, emb, mu): |
| knn = SeqKNN(ref, emb, mu) |
| sub = SubKNN(ref) |
| pair = PairKNN(ref, emb, mu) |
| return np.hstack([smiles_stats(ref, q_df["smiles"]), |
| knn.rows(q_df["sequence"], emb), |
| sub.rows(q_df["smiles"]), |
| pair.rows(q_df, emb)]) |
|
|
