| """Evidence for the v2 reranker: one feature vector per (query, candidate EC). |
| |
| Hits are read from a stream, not a file: DIAMOND's alignments run to gigabytes and the task forbids |
| predict.py from writing any file over 16 MiB. Each query is reduced to its features as it arrives. |
| |
| Four groups of evidence: |
| |
| V1 the homology and hierarchy features of v1. |
| HOM per-EC alignment geometry: both coverages, how an EC's best hit compares with the best hit |
| that does not carry it, and what share of an EC's known members the query reached. |
| PROF residue-level evidence (see `profile_features`), new in v2. |
| EMB the hierarchical ESM-2 head: L4 and L3 probabilities, centroid similarity, nearest |
| neighbours in the contrastive space. |
| """ |
| import math |
| from collections import defaultdict |
|
|
| import numpy as np |
|
|
| V1_NAMES = ["vote_frac", "vote_rank", "max_bits", "max_bits_rel", "n_hits", "max_pid", "mean_pid", |
| "max_cov", "best_hit_is_ec", "best_hit_n_ecs", "q_best_bits", "q_best_pid", "q_n_hits", |
| "q_n_cands", "q_len", "l1_mass", "l2_mass", "l3_mass", "ec_prior"] |
| HOM_NAMES = ["scov_max", "len_ratio", "bits_gap", "top3_rel", "first_rank", "member_frac", |
| "hit_frac", "log_n_db", "qcov_union", "share_in_l3", "n_sib_l3"] |
| PROF_NAMES = ["vote_hk", "prof_llr", "prof_llr_l3", "prof_top", "sdp_match", "sdp_n", "prof_self", |
| "prof_n"] |
| EMB_NAMES = ["emb_cls", "emb_l3", "emb_cen", "emb_knn", "emb_cls_rank", "emb_cen_rank", |
| "emb_knn_rank", "has_emb", "emb_new"] |
| FEATURE_NAMES = V1_NAMES + HOM_NAMES + PROF_NAMES + EMB_NAMES |
| N_SCAN = len(V1_NAMES) + len(HOM_NAMES) + len(PROF_NAMES) |
|
|
| OUTFMT = ["qseqid", "sseqid", "pident", "bitscore", "length", "qlen", "slen", "qstart", "qend", |
| "sstart", "send", "qseq_gapped", "sseq_gapped"] |
|
|
| AA = "ACDEFGHIKLMNPQRSTVWY" |
| LUT = np.full(256, 20, np.int8) |
| for _i, _a in enumerate(AA): |
| LUT[ord(_a)] = _i |
| |
| BG = np.array([.0780, .0193, .0536, .0630, .0386, .0738, .0220, .0515, .0574, .0902, |
| .0224, .0449, .0520, .0426, .0513, .0712, .0584, .0644, .0133, .0322]) |
| BG = BG / BG.sum() |
|
|
|
|
| def read_hits(lines, k, keep_aln): |
| """Group DIAMOND tabular lines by query; only the top `keep_aln` hits keep their alignment.""" |
| cur, rows = None, [] |
| for line in lines: |
| p = line.rstrip("\n").split("\t") |
| if len(p) < 9: |
| continue |
| q = int(p[0]) |
| if q != cur: |
| if rows: |
| yield cur, rows |
| cur, rows = q, [] |
| if len(rows) >= k: |
| continue |
| complete = len(p) > 12 and len(p[11]) == len(p[12]) |
| aln = (p[11], p[12]) if complete and len(rows) < keep_aln else None |
| rows.append((int(p[1]), float(p[2]), float(p[3]), int(p[4]), int(p[5]), int(p[6]), |
| int(p[7]), int(p[8]), aln)) |
| if rows: |
| yield cur, rows |
|
|
|
|
| def project(rows, qlen): |
| """Alignments as a hits x query-columns matrix: -1 unaligned, 0-19 residue, 20 gap or other.""" |
| A = np.full((len(rows), qlen), -1, np.int8) |
| for h, r in enumerate(rows): |
| q = np.frombuffer(r[8][0].encode(), np.uint8) |
| s = np.frombuffer(r[8][1].encode(), np.uint8) |
| col = LUT[s[q != 45]] |
| A[h, r[6] - 1:r[6] - 1 + len(col)] = col |
| return A |
|
|
|
|
| def profile_features(rows, seq, cands, ecsets, beta=1.0, n_top=20): |
| """Residue-level evidence for each candidate EC, from a query-anchored alignment. |
| |
| The fourth EC digit is usually decided by a few specificity-determining residues, invisible to |
| any global similarity score. Every hit is projected on the query's columns, then for candidate |
| EC e the hits are split into members and non-members, each giving a profile (Henikoff-weighted |
| so redundant reference sequences do not dominate). Then we ask which profile the query fits: |
| |
| prof_llr mean log-odds of the query's residues, members vs non-members |
| prof_llr_l3 the same against same-L3 ECs only: the competitors that matter for L4 |
| prof_top the same on the 20 most discriminative columns (Jensen-Shannon divergence) |
| sdp_match of columns conserved in e and different elsewhere, the share the query matches |
| sdp_n how many such columns exist |
| prof_self log-odds of the query's residues against e's profile vs background |
| vote_hk redundancy-corrected vote share |
| |
| Returns {ec: [vote_hk, llr, llr_l3, top, sdp_match, sdp_n, self, n_members]}. |
| """ |
| A = project(rows, len(seq)) |
| H, L = A.shape |
| hh, cc = np.nonzero(A >= 0) |
| if not len(hh): |
| return {} |
| aa = A[hh, cc].astype(np.int64) |
| cell = cc * 21 + aa |
| cnt = np.bincount(cell, minlength=L * 21) |
| kinds = (cnt.reshape(L, 21) > 0).sum(1) |
| w = np.bincount(hh, 1.0 / (kinds[cc] * cnt[cell]), minlength=H) |
| w = w / np.maximum(np.bincount(hh, minlength=H), 1) |
| w = w * H / max(w.sum(), 1e-9) |
| cw = w[hh] |
| total = np.bincount(cell, cw, minlength=L * 21).reshape(L, 21)[:, :20] |
| qi = LUT[np.frombuffer(seq.encode(), np.uint8)].astype(np.int64) |
| known = np.where(qi < 20)[0] |
| wsum = w.sum() |
|
|
| def freq(c, n): |
| return (c + beta * BG) / (n + beta)[:, None] |
|
|
| out = {} |
| member = {e: np.array([e in s for s in ecsets]) for e in cands} |
| for e in cands: |
| m = member[e] |
| if not m.any(): |
| continue |
| vote_hk = float(w[m].sum() / wsum) |
| sel = m[hh] |
| ce = np.bincount(cell[sel], cw[sel], minlength=L * 21).reshape(L, 21)[:, :20] |
| co = np.maximum(total - ce, 0.0) |
| ne, no = ce.sum(1), co.sum(1) |
| pe = freq(ce, ne) |
| cols = known[ne[known] > 0] |
| f = [vote_hk, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, float(m.sum())] |
| if len(cols): |
| f[6] = float(np.mean(np.log(pe[cols, qi[cols]] / BG[qi[cols]]))) |
| both = known[(ne[known] > 0) & (no[known] > 0)] |
| if len(both): |
| po = freq(co, no) |
| llr = np.log(pe[both, qi[both]] / po[both, qi[both]]) |
| f[1] = float(llr.mean()) |
| mix = 0.5 * (pe[both] + po[both]) |
| js = 0.5 * (pe[both] * np.log(pe[both] / mix)).sum(1) \ |
| + 0.5 * (po[both] * np.log(po[both] / mix)).sum(1) |
| f[3] = float(llr[np.argsort(-js)[:n_top]].mean()) |
| sdp = (js > 0.25) & (pe[both].max(1) >= 0.6) |
| if sdp.any(): |
| f[4] = float(np.mean(pe[both[sdp]].argmax(1) == qi[both[sdp]])) |
| f[5] = float(sdp.sum()) |
| pre = e.rsplit(".", 1)[0] |
| sib = np.array([any(x != e and x.rsplit(".", 1)[0] == pre for x in s) and not m[h] |
| for h, s in enumerate(ecsets)]) |
| if sib.any(): |
| sel = sib[hh] |
| cs = np.bincount(cell[sel], cw[sel], minlength=L * 21).reshape(L, 21)[:, :20] |
| ns = cs.sum(1) |
| c2 = known[(ne[known] > 0) & (ns[known] > 0)] |
| if len(c2): |
| ps = freq(cs, ns) |
| f[2] = float(np.log(pe[c2, qi[c2]] / ps[c2, qi[c2]]).mean()) |
| out[e] = f |
| return out |
|
|
|
|
| def query_features(rows, seq, db_ecs, ec_count, ec_prior, cfg, profile=True): |
| """All homology-side features for one query: (candidates, matrix, query-level info).""" |
| rows = sorted(rows, key=lambda r: -r[2]) |
| best_bits = rows[0][2] |
| best_pid = max(r[1] for r in rows) |
| ecsets = [set(db_ecs[r[0]]) for r in rows] |
| best_ecs = ecsets[0] |
| qlen = rows[0][4] |
| power, rel_cut, max_cands = cfg["power"], cfg["rel_cut"], cfg["max_cands"] |
|
|
| agg = {} |
| for h, (r, ecs) in enumerate(zip(rows, ecsets)): |
| s, pid, bits, length, _, slen, qstart, qend, _ = r |
| if bits < rel_cut * best_bits: |
| continue |
| for ec in ecs: |
| a = agg.get(ec) |
| if a is None: |
| a = agg[ec] = {"sum": 0.0, "bits": [], "pids": [], "max_cov": 0.0, "max_scov": 0.0, |
| "first": h, "spans": [], "len_ratio": slen / max(qlen, 1)} |
| a["sum"] += bits ** power / len(ecs) |
| a["bits"].append(bits) |
| a["pids"].append(pid) |
| a["max_cov"] = max(a["max_cov"], length / max(qlen, 1)) |
| a["max_scov"] = max(a["max_scov"], length / max(slen, 1)) |
| a["spans"].append((qstart, qend)) |
| total = sum(a["sum"] for a in agg.values()) or 1.0 |
| ranked = sorted(agg, key=lambda e: -agg[e]["sum"])[:max_cands] |
|
|
| mass = [defaultdict(float) for _ in range(3)] |
| l3_ecs = defaultdict(list) |
| for ec, a in agg.items(): |
| parts = ec.split(".") |
| for lv in range(3): |
| mass[lv][".".join(parts[:lv + 1])] += a["sum"] / total |
| l3_ecs[".".join(parts[:3])].append(ec) |
|
|
| prof = {} |
| if profile and len(ranked) > 1: |
| kept = [h for h, r in enumerate(rows) if r[8] is not None] |
| prof = profile_features([rows[h] for h in kept], seq, ranked[:cfg["prof_cands"]], |
| [ecsets[h] for h in kept]) |
|
|
| q_level = [best_bits, best_pid, len(rows), len(agg), len(seq)] |
| X = np.zeros((len(ranked), N_SCAN), np.float32) |
| for j, ec in enumerate(ranked): |
| a = agg[ec] |
| pre = ec.split(".") |
| l3 = ".".join(pre[:3]) |
| v1 = [a["sum"] / total, j, a["bits"][0], a["bits"][0] / best_bits, len(a["bits"]), |
| max(a["pids"]), float(np.mean(a["pids"])), a["max_cov"], float(ec in best_ecs), |
| len(best_ecs)] + q_level + [ |
| mass[0][pre[0]], mass[1][".".join(pre[:2])], mass[2][l3], ec_prior.get(ec, 0.0)] |
| other = next((r[2] for r, s in zip(rows, ecsets) if ec not in s), 0.0) |
| n_db = ec_count.get(ec, 0) |
| spans, cover, end = sorted(a["spans"]), 0, 0 |
| for st, en in spans: |
| if en > end: |
| cover += en - max(st - 1, end) |
| end = en |
| hom = [a["max_scov"], a["len_ratio"], (a["bits"][0] - other) / best_bits, |
| float(np.mean(a["bits"][:3])) / best_bits, math.log1p(a["first"]), |
| min(len(a["bits"]) / max(n_db, 1), 1.0), len(a["bits"]) / len(rows), |
| math.log1p(n_db), cover / max(qlen, 1), |
| a["sum"] / max(sum(agg[x]["sum"] for x in l3_ecs[l3]), 1e-9), len(l3_ecs[l3])] |
| X[j] = v1 + hom + prof.get(ec, [0.0] * len(PROF_NAMES)) |
| return ranked, X, {"q_level": q_level, "best_ecs": len(best_ecs), |
| "mass": [dict(m) for m in mass], "max_cands": max_cands} |
|
|
|
|
| def assemble(scanned, n, emb, ec_prior, ec_count, seqs, extra_from_emb=5): |
| """Merge homology features with head evidence into the final (q, ec, X) arrays. |
| |
| scanned[i] is query_features' result or None; emb maps a query to {ec: (cls, l3, cen, knn)}. |
| A query missing from `emb` gets zeroed columns and has_emb = 0. ECs only the head proposes join |
| as candidates with zeroed homology columns and emb_new = 1. |
| """ |
| out_q, out_ec, out_X = [], [], [] |
| i_rank, i_prior, i_ndb = (V1_NAMES.index("vote_rank"), V1_NAMES.index("ec_prior"), |
| len(V1_NAMES) + HOM_NAMES.index("log_n_db")) |
| i_qlevel = V1_NAMES.index("q_best_bits") |
| i_best = V1_NAMES.index("best_hit_n_ecs") |
| i_mass = V1_NAMES.index("l1_mass") |
| for i in range(n): |
| res = scanned[i] |
| ev = emb.get(i, {}) |
| cands, X = (res[0], res[1]) if res else ([], np.zeros((0, N_SCAN), np.float32)) |
| known = set(cands) |
| new = [e for e in sorted(ev, key=lambda e: -(ev[e][0] + ev[e][2] + ev[e][3])) |
| if e not in known][:extra_from_emb] |
| if not cands and not new: |
| continue |
| if new: |
| Y = np.zeros((len(new), N_SCAN), np.float32) |
| Y[:, i_rank] = 25 |
| for j, e in enumerate(new): |
| Y[j, i_prior] = ec_prior.get(e, 0.0) |
| Y[j, i_ndb] = math.log1p(ec_count.get(e, 0)) |
| parts = e.split(".") |
| if res: |
| info = res[2] |
| Y[j, i_qlevel:i_qlevel + 5] = info["q_level"] |
| Y[j, i_best] = info["best_ecs"] |
| for lv in range(3): |
| Y[j, i_mass + lv] = info["mass"][lv].get(".".join(parts[:lv + 1]), 0.0) |
| else: |
| Y[j, V1_NAMES.index("q_len")] = len(seqs[i]) |
| X = np.vstack([X, Y]) if len(X) else Y |
| names = list(cands) + new |
| E = np.zeros((len(names), len(EMB_NAMES)), np.float32) |
| if ev: |
| rank = [{e: r for r, e in enumerate(sorted(ev, key=lambda e: -ev[e][k]))} |
| for k in (0, 2, 3)] |
| for j, e in enumerate(names): |
| c = ev.get(e, (0.0, 0.0, 0.0, 0.0)) |
| E[j] = [c[0], c[1], c[2], c[3], rank[0].get(e, 99), rank[1].get(e, 99), |
| rank[2].get(e, 99), 1.0, float(e in new)] |
| else: |
| E[:, 4:7] = 99 |
| out_q.extend([i] * len(names)) |
| out_ec.extend(names) |
| out_X.append(np.hstack([X, E])) |
| F = np.vstack(out_X) if out_X else np.zeros((0, len(FEATURE_NAMES)), np.float32) |
| return np.array(out_q, np.int64), np.array(out_ec), F |
|
|
