"""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}, # lab only: shows whether ESM helps at all # lab results (local holdout): knn best, w05 helps a little, clip/reg hurt -> combine the winners "knn_w05": {"knn": True, "weight_pow": 0.5}, # layer experiments: run with --esm_layers 4,6,8,10,12 (blocks 0..4) "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": {}, # v9: enzyme-effect (residual) idea + rank-transformed target "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: # never let the extra model break the whole run 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) # centre on the competition's own enzymes (keeps holdouts identical with/without external data) 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 # mutant groups (used for weights) df["_g"] = groups # lets te_stats average per enzyme group 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} # ---------------- local validation run(s) ---------------- 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]))] # no external rows from held-out families 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})") # ---------------- final training on ALL training rows ---------------- 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()