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)])