| """Train the kcat/Km model. |
| |
| Typical use (in the Colab notebook): |
| python train.py --train ../data/train.parquet --out weights --esm --lab # compare ideas on a local holdout, then stop |
| python train.py --train ../data/train.parquet --out weights --esm --preset NAME --seeds 5 # real training |
| """ |
| import argparse |
| import json |
| import os |
| import time |
|
|
| import lightgbm as lgb |
| import numpy as np |
| import pandas as pd |
| from scipy.stats import spearmanr |
| from sklearn.model_selection import GroupKFold, GroupShuffleSplit |
|
|
| import features as F |
|
|
| ESM_NAME = "facebook/esm2_t12_35M_UR50D" |
| REG = dict(num_leaves=31, min_child_samples=60, colsample_bytree=0.15, reg_lambda=10.0) |
| PRESETS = { |
| "base": {}, |
| "clip": {"clip": True}, |
| "w05": {"weight_pow": 0.5}, |
| "reg": {"params": REG}, |
| "knn": {"knn": True}, |
| "combo": {"clip": True, "weight_pow": 0.5, "params": REG}, |
| "combo_knn": {"clip": True, "weight_pow": 0.5, "params": REG, "knn": True}, |
| "noesm": {"noesm": True}, |
| |
| "knn_w05": {"knn": True, "weight_pow": 0.5}, |
| |
| "L12": {"esm_blocks": [4]}, |
| "L8": {"esm_blocks": [2]}, |
| "L6": {"esm_blocks": [1]}, |
| "L8_12": {"esm_blocks": [2, 4]}, |
| "L6_8_12": {"esm_blocks": [1, 2, 4]}, |
| "Lall": {}, |
| |
| "res": {"res": True}, |
| "rank": {"rank": True}, |
| "res_knn_w05": {"res": True, "knn": True, "weight_pow": 0.5}, |
| "res_w05": {"res": True, "weight_pow": 0.5}, |
| "pca32": {"esm_pca": 32}, |
| "nodi": {"nodi": True}, |
| "pca32_nodi": {"esm_pca": 32, "nodi": True}, |
| "pca32_knn": {"esm_pca": 32, "knn": True}, |
| "knn_multi": {"knn": True, "ks": (1, 5, 20)}, |
| "knn_multi_w05": {"knn": True, "ks": (1, 5, 20), "weight_pow": 0.5}, |
| } |
|
|
|
|
| def log(*a): |
| print(time.strftime("%H:%M:%S"), *a, flush=True) |
|
|
|
|
| def te_oof(df, groups, smi_index, bits_u, seq_index, embn, ks=(5,), n_splits=5, want_res=False): |
| """Substrate-level (and ESM-neighbour) features for training rows, computed WITHOUT |
| looking at the row's own enzyme group, so the model learns to use them honestly.""" |
| te = np.full((len(df), 12), np.nan, dtype=np.float32) |
| kn = None if embn is None else np.full((len(df), 6 * len(ks)), np.nan, dtype=np.float32) |
| er = np.full((len(df), 3), np.nan, dtype=np.float32) if (want_res and embn is not None) else None |
| smis = df["smiles"].values |
| sidx = df["sequence"].map(seq_index).values |
| ys = df[["log10_kcat", "log10_km"]].values.astype(np.float32) |
| for tr, va in GroupKFold(n_splits).split(df, groups=groups): |
| st = F.te_stats(df.iloc[tr]) |
| rb = bits_u[[smi_index[s] for s in st.index]] |
| te[va] = F.te_full(smis[va], smi_index, bits_u, st, rb) |
| if kn is not None: |
| kn[va] = F.knn_feats(smis[va], sidx[va], embn, smis[tr], sidx[tr], embn, ys[tr], ks=ks) |
| if er is not None: |
| rsr = F.enzyme_effect_ref(smis[tr], sidx[tr], ys[tr], embn.shape[0], st) |
| er[va] = F.enzyme_effect_feats(sidx[va], embn, embn, rsr) |
| return te, kn, er |
|
|
|
|
| def mlp_fit(Ztr, Ytr, Zva=None, Yva=None, epochs=15, seed=0, hidden=(384, 128), alpha=1e-3, lr=1e-3): |
| """Small neural net (sklearn). Returns model, per-epoch validation SCC, per-epoch validation predictions.""" |
| from sklearn.neural_network import MLPRegressor |
| m = MLPRegressor(hidden_layer_sizes=hidden, alpha=alpha, learning_rate_init=lr, |
| batch_size=256, random_state=seed) |
| curve, preds = [], [] |
| for ep in range(epochs): |
| m.partial_fit(Ztr, Ytr) |
| if Zva is not None: |
| P = m.predict(Zva) |
| preds.append(P) |
| curve.append(float(np.mean([spearmanr(Yva[:, j], P[:, j]).correlation for j in (0, 1)]))) |
| return m, curve, preds |
|
|
|
|
| def nn_fit_torch(Ztr, Ytr, Zva=None, Yva=None, epochs=12, seed=0, hidden=(512, 256), |
| dropout=0.3, in_drop=0.1, wd=1e-2, lr=1e-3, bs=256): |
| """Regularised neural net trained with PyTorch (dropout + weight decay + Huber loss). |
| Returns numpy layers so prediction needs only numpy.""" |
| import torch |
| torch.manual_seed(seed) |
| dev = "cuda" if torch.cuda.is_available() else "cpu" |
| mods, prev = [torch.nn.Dropout(in_drop)], Ztr.shape[1] |
| for h in hidden: |
| mods += [torch.nn.Linear(prev, h), torch.nn.ReLU(), torch.nn.Dropout(dropout)] |
| prev = h |
| mods += [torch.nn.Linear(prev, Ytr.shape[1])] |
| net = torch.nn.Sequential(*mods).to(dev) |
| opt = torch.optim.AdamW(net.parameters(), lr=lr, weight_decay=wd) |
| Xt = torch.tensor(Ztr, dtype=torch.float32, device=dev) |
| Yt = torch.tensor(Ytr, dtype=torch.float32, device=dev) |
| Xv = None if Zva is None else torch.tensor(Zva, dtype=torch.float32, device=dev) |
| curve, preds = [], [] |
| for ep in range(epochs): |
| net.train() |
| perm = torch.randperm(len(Xt), device=dev) |
| for i in range(0, len(Xt), bs): |
| idx = perm[i:i + bs] |
| loss = torch.nn.functional.huber_loss(net(Xt[idx]), Yt[idx], delta=1.0) |
| opt.zero_grad() |
| loss.backward() |
| opt.step() |
| if Xv is not None: |
| net.eval() |
| with torch.no_grad(): |
| P = net(Xv).cpu().numpy() |
| preds.append(P) |
| curve.append(float(np.mean([spearmanr(Yva[:, j], P[:, j]).correlation for j in (0, 1)]))) |
| lin = [m for m in net if isinstance(m, torch.nn.Linear)] |
| layers = [(l.weight.detach().cpu().numpy().T.astype(np.float32), |
| l.bias.detach().cpu().numpy().astype(np.float32)) for l in lin] |
| return layers, curve, preds |
|
|
|
|
| def fit_nn(args, Ztr, Ytr, Zva=None, Yva=None, epochs=12, seed=0): |
| """-> (numpy layers, per-epoch val SCC, per-epoch val predictions)""" |
| if args.nn == "torch": |
| try: |
| return nn_fit_torch(Ztr, Ytr, Zva, Yva, epochs=epochs, seed=seed) |
| except Exception as e: |
| log(" torch net failed (", repr(e)[:150], ") -> using the simple sklearn net instead") |
| m, curve, preds = mlp_fit(Ztr, Ytr, Zva, Yva, epochs=epochs, seed=seed) |
| return mlp_layers(m), curve, preds |
|
|
|
|
| def mlp_layers(m): |
| return [(W.astype(np.float32), b.astype(np.float32)) for W, b in zip(m.coefs_, m.intercepts_)] |
|
|
|
|
| def make_model(seed, n_estimators, lr, **over): |
| p = dict(n_estimators=n_estimators, learning_rate=lr, num_leaves=63, |
| min_child_samples=20, subsample=0.8, subsample_freq=1, |
| colsample_bytree=0.25, reg_lambda=1.0, random_state=seed, |
| n_jobs=-1, verbose=-1) |
| p.update(over) |
| return lgb.LGBMRegressor(**p) |
|
|
|
|
| def row_weights(g, pw): |
| """enzymes with hundreds of mutant rows should not dominate training""" |
| if not pw: |
| return None |
| s = pd.Series(g) |
| sizes = s.map(s.value_counts()).values.astype(float) |
| w = sizes ** (-pw) |
| return w / w.mean() |
|
|
|
|
| def target(y, preset): |
| if preset.get("rank"): |
| from scipy.stats import norm, rankdata |
| return norm.ppf(rankdata(y) / (len(y) + 1.0)) |
| if preset.get("clip"): |
| lo, hi = np.quantile(y, [0.01, 0.99]) |
| return np.clip(y, lo, hi) |
| return y |
|
|
|
|
| def report(name, y, p): |
| scc = spearmanr(y, p).correlation |
| rmse = float(np.sqrt(np.mean((y - p) ** 2))) |
| log(f" {name}: SCC={scc:.4f} RMSE={rmse:.4f}") |
| return scc |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--train", required=True) |
| ap.add_argument("--out", default="weights") |
| ap.add_argument("--esm", action="store_true", help="add ESM2 protein embeddings") |
| ap.add_argument("--esm_name", default=ESM_NAME) |
| ap.add_argument("--esm_layers", default="", help="e.g. 6,9,12 = pool several ESM layers (default: last layer only)") |
| ap.add_argument("--preset", default="base", choices=list(PRESETS)) |
| ap.add_argument("--lab", action="store_true", help="compare presets on local holdout(s), then stop") |
| ap.add_argument("--lab_presets", default="", help="comma separated presets for --lab") |
| ap.add_argument("--split", default="esm", choices=["esm", "cosine"], |
| help="esm = hold out whole ESM-similarity families (harsh, closest to the real validation set)") |
| ap.add_argument("--repeats", type=int, default=1, help="how many different holdouts to average in --lab") |
| ap.add_argument("--esm_clusters", type=int, default=300) |
| ap.add_argument("--seeds", type=int, default=1) |
| ap.add_argument("--lr", type=float, default=0.03) |
| ap.add_argument("--max_trees", type=int, default=3000) |
| ap.add_argument("--no_holdout", action="store_true") |
| ap.add_argument("--external", default="", help="parquet of EXTRA training rows (smiles, sequence, log10_kcat, log10_km; one label may be NaN)") |
| ap.add_argument("--mlp", action="store_true", help="also train a neural net and blend it with LightGBM") |
| ap.add_argument("--mlp_epochs", type=int, default=12) |
| ap.add_argument("--nn", default="torch", choices=["torch", "sk"], help="neural net implementation") |
| ap.add_argument("--mlp_w", type=float, default=0.4, help="weight of the neural net in the blend") |
| ap.add_argument("--blend_with", default="base", help="LightGBM preset to blend with in --lab") |
| args = ap.parse_args() |
| os.makedirs(args.out, exist_ok=True) |
| preset = PRESETS[args.preset] |
| needs_esm = lambda p: p.get("knn") or p.get("esm_pca") or p.get("noesm") or p.get("res") |
| if (needs_esm(preset) or args.lab or args.split == "esm") and not args.esm: |
| raise SystemExit("this needs --esm") |
| if preset.get("noesm") and not args.lab: |
| raise SystemExit("'noesm' is only for --lab") |
| if preset.get("rank") and args.mlp and not args.lab: |
| raise SystemExit("'rank' cannot be blended with the neural net (different scales)") |
|
|
| df = pd.read_parquet(args.train).reset_index(drop=True) |
| df = df[["smiles", "sequence", "log10_kcat", "log10_km"]].copy() |
| df["_ext"] = False |
| if args.external: |
| ext = pd.read_parquet(args.external)[["smiles", "sequence", "log10_kcat", "log10_km"]].copy() |
| ext["_ext"] = True |
| log(f"external rows added: {len(ext)} (kcat labels: {ext.log10_kcat.notna().sum()}, km labels: {ext.log10_km.notna().sum()})") |
| df = pd.concat([df, ext], ignore_index=True) |
| is_task = ~df["_ext"].values |
| log("task rows:", int(is_task.sum()), "| all rows:", len(df), "| unique smiles:", df.smiles.nunique(), "| unique sequences:", df.sequence.nunique()) |
| usmi = df["smiles"].unique().tolist() |
| useq = df["sequence"].unique().tolist() |
| smi_index = {s: i for i, s in enumerate(usmi)} |
| seq_index = {s: i for i, s in enumerate(useq)} |
| log("computing SMILES features ...") |
| smi_mat = F.smiles_features(usmi) |
| log("computing sequence features ...") |
| seq_mat = F.seq_features(useq) |
| bits_u = (smi_mat[:, :F.N_MORGAN] > 0).astype(np.float32) |
|
|
| esm_mat, esm_mean, embn = None, None, None |
| layers = [int(x) for x in args.esm_layers.split(",") if x] or None |
| n_blocks = len(layers) if layers else 1 |
| if args.esm: |
| esm_dir = os.path.join(args.out, "esm") |
| if not os.path.exists(esm_dir): |
| from transformers import AutoTokenizer, EsmModel |
| log("downloading", args.esm_name) |
| AutoTokenizer.from_pretrained(args.esm_name).save_pretrained(esm_dir) |
| EsmModel.from_pretrained(args.esm_name).save_pretrained(esm_dir) |
| cache = "esm_cache_" + args.esm_name.split("/")[-1] + ("_L" + "-".join(map(str, layers)) if layers else "") + ".npz" |
| if os.path.exists(cache) and list(np.load(cache, allow_pickle=True)["seqs"]) == useq: |
| esm_mat = np.load(cache, allow_pickle=True)["emb"] |
| log("loaded ESM embeddings from cache") |
| else: |
| import torch |
| dev = "cuda" if torch.cuda.is_available() else "cpu" |
| log(f"computing ESM embeddings on {dev} (this is the slow part) ...") |
| esm_mat = F.esm_embed(useq, esm_dir, device=dev, |
| batch_tokens=16384 if dev == "cuda" else 4096, |
| n_threads=os.cpu_count() or 2, layers=layers) |
| np.savez(cache, seqs=np.array(useq, dtype=object), emb=esm_mat) |
| task_seqs = set(df.loc[is_task, "sequence"]) |
| seq_is_task = np.array([s in task_seqs for s in useq]) |
| esm_mean = esm_mat[seq_is_task].mean(0) |
| np.save(os.path.join(args.out, "esm_mean.npy"), esm_mean) |
| embn = F.normalize_emb(esm_mat, esm_mean) |
|
|
| log("grouping near-identical sequences (mutants) ...") |
| lab_ = F.cluster_by_cosine(useq, log=log) |
| groups = df["sequence"].map(dict(zip(useq, lab_))).values |
| df["_g"] = groups |
| if args.split == "esm": |
| from sklearn.cluster import KMeans |
| kmf = KMeans(n_clusters=args.esm_clusters, n_init=3, random_state=0).fit(embn[seq_is_task]) |
| km_lab = np.zeros(len(useq), dtype=int) |
| km_lab[seq_is_task] = kmf.labels_ |
| if (~seq_is_task).any(): |
| km_lab[~seq_is_task] = kmf.predict(embn[~seq_is_task]) |
| split_groups = df["sequence"].map(dict(zip(useq, km_lab))).values |
| log(f"harsh split: {args.esm_clusters} enzyme families by ESM similarity") |
| else: |
| split_groups = groups |
|
|
| smi_dim, seq_dim = smi_mat.shape[1], seq_mat.shape[1] |
| esm_dim = esm_mat.shape[1] if esm_mat is not None else 0 |
| pcas = {} |
|
|
| def get_pca(p): |
| n = p.get("esm_pca") |
| if not n: |
| return None |
| if n not in pcas: |
| pcas[n] = F.fit_pca(esm_mat, n) |
| return pcas[n] |
|
|
| def variant(X0, kn, p, ks_all, er=None): |
| return F.apply_variant(X0, kn, p, smi_dim, seq_dim, esm_dim, ks_all=ks_all, pca=get_pca(p), n_blocks=n_blocks, er=er) |
|
|
| best_iter = {"kcat": 800, "km": 800} |
|
|
| |
| if not args.no_holdout or args.lab: |
| if args.lab: |
| lab_names = [n for n in args.lab_presets.split(",") if n] or \ |
| ["base", "noesm", "w05", "knn", "knn_w05", "pca32", "nodi", "pca32_nodi"] |
| else: |
| lab_names = [args.preset] |
| all_ks = sorted({k for n in lab_names for k in PRESETS[n].get("ks", (5,))}) |
| want_res = any(PRESETS[n].get("res") for n in lab_names) |
|
|
| def make_split(seed): |
| task_ix = np.nonzero(is_task)[0] |
| ext_ix = np.nonzero(~is_task)[0] |
| tr_t, va_t = next(GroupShuffleSplit(1, test_size=0.2, random_state=seed).split(task_ix, groups=split_groups[task_ix])) |
| va = task_ix[va_t] |
| ext_keep = ext_ix[~np.isin(split_groups[ext_ix], np.unique(split_groups[va]))] |
| tr = np.concatenate([task_ix[tr_t], ext_keep]) |
| dtr, dva = df.iloc[tr].reset_index(drop=True), df.iloc[va].reset_index(drop=True) |
| te_tr, kn_tr, er_tr = te_oof(dtr, split_groups[tr], smi_index, bits_u, seq_index, embn, ks=tuple(all_ks), want_res=want_res) |
| st_tr = F.te_stats(dtr) |
| te_va = F.te_full(dva["smiles"].values, smi_index, bits_u, st_tr, |
| bits_u[[smi_index[s] for s in st_tr.index]]) |
| kn_va = None |
| if embn is not None: |
| kn_va = F.knn_feats(dva["smiles"].values, dva["sequence"].map(seq_index).values, embn, |
| dtr["smiles"].values, dtr["sequence"].map(seq_index).values, embn, |
| dtr[["log10_kcat", "log10_km"]].values.astype(np.float32), ks=tuple(all_ks)) |
| er_va = None |
| if want_res: |
| rsr = F.enzyme_effect_ref(dtr["smiles"].values, dtr["sequence"].map(seq_index).values, |
| dtr[["log10_kcat", "log10_km"]].values, embn.shape[0], st_tr) |
| er_va = F.enzyme_effect_feats(dva["sequence"].map(seq_index).values, embn, embn, rsr) |
| return dict(dtr=dtr, dva=dva, gtr=groups[tr], kn_tr=kn_tr, kn_va=kn_va, er_tr=er_tr, er_va=er_va, |
| Xtr0=F.assemble(dtr, usmi, smi_mat, useq, seq_mat, esm_mat, te_tr), |
| Xva0=F.assemble(dva, usmi, smi_mat, useq, seq_mat, esm_mat, te_va)) |
|
|
| def run(pname, d): |
| p = PRESETS[pname] |
| Xtr = variant(d["Xtr0"], d["kn_tr"], p, all_ks, d["er_tr"]) |
| Xva = variant(d["Xva0"], d["kn_va"], p, all_ks, d["er_va"]) |
| w = row_weights(d["gtr"], p.get("weight_pow")) |
| res, its, prd = {}, {}, {} |
| for name, col in [("kcat", "log10_kcat"), ("km", "log10_km")]: |
| m = make_model(0, args.max_trees, args.lr, **p.get("params", {})) |
| mk = ~np.isnan(d["dtr"][col].values) |
| m.fit(Xtr[mk], target(d["dtr"][col].values[mk], p), sample_weight=None if w is None else w[mk], |
| eval_set=[(Xva, d["dva"][col].values)], |
| callbacks=[lgb.early_stopping(100, verbose=False)]) |
| its[name] = max(100, int(m.best_iteration_ or args.max_trees)) |
| prd[name] = m.predict(Xva) |
| res[name] = report(f"{pname}/{name}", d["dva"][col].values, prd[name]) |
| return res, its, prd |
|
|
| if args.lab: |
| if args.mlp and args.blend_with not in lab_names: |
| lab_names.append(args.blend_with) |
| scores = {n: [] for n in lab_names} |
| lgb_pred, mlp_pred, ytrue = {}, [], [] |
| nodi = {"nodi": True} |
| for rep in range(args.repeats): |
| log(f"##### holdout {rep + 1} of {args.repeats}") |
| d = make_split(rep) |
| log("holdout features:", d["Xtr0"].shape, d["Xva0"].shape) |
| for pname in lab_names: |
| log(f"=== preset: {pname}") |
| r, _, pr = run(pname, d) |
| scores[pname].append((r["kcat"], r["km"])) |
| lgb_pred[(rep, pname)] = pr |
| if args.mlp: |
| log("=== neural net") |
| Xa = F.apply_variant(d["Xtr0"], None, nodi, smi_dim, seq_dim, esm_dim) |
| Xb = F.apply_variant(d["Xva0"], None, nodi, smi_dim, seq_dim, esm_dim) |
| Ytr = d["dtr"][["log10_kcat", "log10_km"]].values.astype(np.float32) |
| Yva = d["dva"][["log10_kcat", "log10_km"]].values.astype(np.float32) |
| both = ~np.isnan(Ytr).any(1) |
| Xa, Ytr = Xa[both], Ytr[both] |
| st = F.mlp_prep_fit(Xa) |
| ym, ys = Ytr.mean(0), Ytr.std(0) |
| _, curve, preds = fit_nn(args, F.mlp_prep(Xa, st), (Ytr - ym) / ys, F.mlp_prep(Xb, st), |
| (Yva - ym) / ys, epochs=args.mlp_epochs) |
| log(" nn SCC after each epoch:", " ".join(f"{c:.3f}" for c in curve)) |
| mlp_pred.append([p * ys + ym for p in preds]); ytrue.append(Yva) |
| log("========== LAB RESULTS (harsh local holdout, higher is better) ==========") |
| log(f"{'preset':14s} {'kcat':>8s} {'km':>8s} {'mean':>8s} {'+-':>7s}") |
| rows = [] |
| for n, v in scores.items(): |
| a = np.array(v) |
| m_ = a.mean(1) |
| rows.append((n, a[:, 0].mean(), a[:, 1].mean(), m_.mean(), m_.std())) |
| for r_ in sorted(rows, key=lambda r: -r[3]): |
| log(f"{r_[0]:14s} {r_[1]:8.4f} {r_[2]:8.4f} {r_[3]:8.4f} {r_[4]:7.4f}") |
| if args.mlp: |
| def scc2(y, p): |
| return [spearmanr(y[:, j], p[:, j]).correlation for j in (0, 1)] |
| n_ep = len(mlp_pred[0]) |
| mean_curve = [np.mean([np.mean(scc2(ytrue[r], mlp_pred[r][e])) for r in range(args.repeats)]) for e in range(n_ep)] |
| be = int(np.argmax(mean_curve)) |
| log(f"neural net alone: best after epoch {be + 1} (mean SCC {mean_curve[be]:.4f})") |
| log(f"BLEND of '{args.blend_with}' LightGBM and neural net (epoch {be + 1}):") |
| log(f"{'nn weight':>10s} {'kcat':>8s} {'km':>8s} {'mean':>8s}") |
| for w_ in (0.0, 0.2, 0.3, 0.4, 0.5, 0.6, 0.8, 1.0): |
| ks_, ms_ = [], [] |
| for r in range(args.repeats): |
| lp = lgb_pred[(r, args.blend_with)] |
| L = np.column_stack([lp["kcat"], lp["km"]]) |
| s2 = scc2(ytrue[r], (1 - w_) * L + w_ * mlp_pred[r][be]) |
| ks_.append(s2[0]); ms_.append(s2[1]) |
| log(f"{w_:10.1f} {np.mean(ks_):8.4f} {np.mean(ms_):8.4f} {(np.mean(ks_) + np.mean(ms_)) / 2:8.4f}") |
| log(f"-> to use it for real set MLP_EPOCHS = {be + 1} and MLP_W to the best weight above") |
| return |
| d = make_split(0) |
| r, best_iter, _ = run(args.preset, d) |
| log(f"LOCAL HOLDOUT mean SCC = {(r['kcat'] + r['km']) / 2:.4f} (preset {args.preset}, best iters: {best_iter})") |
|
|
| |
| log("training final models on all rows ...") |
| fks = tuple(preset.get("ks", (5,))) |
| te_all, kn_all, er_all = te_oof(df, split_groups, smi_index, bits_u, seq_index, embn, ks=fks, want_res=bool(preset.get("res"))) |
| stats = F.te_stats(df) |
| stats.to_parquet(os.path.join(args.out, "te_stats.parquet")) |
| np.save(os.path.join(args.out, "ref_bits.npy"), |
| bits_u[[smi_index[s] for s in stats.index]].astype(np.uint8)) |
| if preset.get("res"): |
| np.save(os.path.join(args.out, "ref_seq_res.npy"), |
| F.enzyme_effect_ref(df["smiles"].values, df["sequence"].map(seq_index).values, |
| df[["log10_kcat", "log10_km"]].values, embn.shape[0], stats)) |
| if preset.get("knn") or preset.get("res"): |
| np.save(os.path.join(args.out, "ref_embn.npy"), embn.astype(np.float16)) |
| if preset.get("knn"): |
| pd.DataFrame({"smiles": df["smiles"].values, "seq_id": df["sequence"].map(seq_index).values, |
| "kc": df["log10_kcat"].values, "km": df["log10_km"].values} |
| ).to_parquet(os.path.join(args.out, "ref_rows.parquet")) |
| pca = get_pca(preset) |
| if pca is not None: |
| np.savez(os.path.join(args.out, "esm_pca.npz"), mean=pca["mean"], comps=pca["comps"]) |
| X = variant(F.assemble(df, usmi, smi_mat, useq, seq_mat, esm_mat, te_all), kn_all, preset, list(fks), er_all) |
| w = row_weights(groups, preset.get("weight_pow")) |
| for name, col in [("kcat", "log10_kcat"), ("km", "log10_km")]: |
| mk = ~np.isnan(df[col].values) |
| y = target(df[col].values[mk], preset) |
| for s in range(args.seeds): |
| m = make_model(s, int(best_iter[name] * 1.15), args.lr, **preset.get("params", {})) |
| m.fit(X[mk], y, sample_weight=None if w is None else w[mk]) |
| m.booster_.save_model(os.path.join(args.out, f"{name}_{s}.txt")) |
| log(f"saved {name}_{s}.txt") |
|
|
| mlp_meta = None |
| if args.mlp: |
| log("training neural nets on all rows ...") |
| nodi = {"nodi": True} |
| Xm = F.apply_variant(F.assemble(df, usmi, smi_mat, useq, seq_mat, esm_mat, te_all), None, nodi, |
| smi_dim, seq_dim, esm_dim) |
| Y = df[["log10_kcat", "log10_km"]].values.astype(np.float32) |
| both = ~np.isnan(Y).any(1) |
| Xm, Y = Xm[both], Y[both] |
| st = F.mlp_prep_fit(Xm) |
| ym, ys = Y.mean(0), Y.std(0) |
| Zm = F.mlp_prep(Xm, st) |
| np.savez(os.path.join(args.out, "mlp_prep.npz"), mean=st["mean"], std=st["std"], ym=ym, ys=ys) |
| n_nn = max(3, args.seeds) |
| for s in range(n_nn): |
| lay, _, _ = fit_nn(args, Zm, (Y - ym) / ys, epochs=args.mlp_epochs, seed=s) |
| np.savez(os.path.join(args.out, f"mlp_{s}.npz"), n=len(lay), |
| **{f"W{i}": W for i, (W, b) in enumerate(lay)}, **{f"b{i}": b for i, (W, b) in enumerate(lay)}) |
| log(f"saved mlp_{s}.npz") |
| mlp_meta = {"w": args.mlp_w, "n": n_nn} |
|
|
| pflags = {k: (list(v) if isinstance(v, tuple) else v) for k, v in preset.items() if k != "params"} |
| meta = {"use_esm": bool(args.esm), "knn": bool(preset.get("knn")), "ks": list(fks), "preset": pflags, |
| "dims": [smi_dim, seq_dim, esm_dim], "layers": layers, "n_seeds": args.seeds, "mlp": mlp_meta, |
| "mean_kcat": float(np.nanmean(df["log10_kcat"])), "mean_km": float(np.nanmean(df["log10_km"])), |
| "external_rows": int((~is_task).sum())} |
| with open(os.path.join(args.out, "meta.json"), "w") as f: |
| json.dump(meta, f) |
| log("DONE. Weights are in", args.out) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
