| """Candidate-EC features from DIAMOND hits (shared by training and inference).""" |
| import math |
| import numpy as np |
|
|
| MAX_CAND = 10 |
| FEATURES = [ |
| 'v1', 'v4', 'v8', 'vrank', 'best_bits_rel', 'best_bits', 'best_pid', 'best_rank', |
| 'cnt10', 'cnt50', 'frac10', 'qcov', 'scov', 'log_ev', 'best_nec', 'l3', 'l2', 'l1', |
| 'log_freq', 'top_bits', 'top_pid', 'log_qlen', 'n_hits', 'n_cand', 'top_nec', |
| 'margin4', 'pid_gap', |
| ] |
|
|
|
|
| def query_candidates(hits, labels, freq): |
| """hits: list of (sseqid, pid, qlen, slen, evalue, bits, qstart, qend, sstart, send), |
| sorted by bitscore descending. Returns (candidate ECs, feature matrix).""" |
| if not hits: |
| return [], np.zeros((0, len(FEATURES)), dtype=np.float32) |
| top_bits = hits[0][5] |
| top_pid = hits[0][1] |
| qlen = hits[0][2] |
| stats = {} |
| tot1 = tot4 = tot8 = 0.0 |
| lvl = [dict(), dict(), dict()] |
| for r, (s, pid, ql, sl, ev, b, qs, qe, ss, se) in enumerate(hits): |
| ecs = labels[s] |
| n = len(ecs) |
| rb = b / top_bits |
| w1, w4, w8 = rb / n, rb ** 4 / n, rb ** 8 / n |
| tot1 += rb |
| tot4 += rb ** 4 |
| tot8 += rb ** 8 |
| prefixes = [set(), set(), set()] |
| for e in ecs: |
| st = stats.get(e) |
| if st is None: |
| st = stats[e] = [0.0, 0.0, 0.0, b, pid, r, 0, 0, |
| (qe - qs + 1) / ql, (se - ss + 1) / sl, ev, n] |
| st[0] += w1 |
| st[1] += w4 |
| st[2] += w8 |
| if r < 10: |
| st[6] += 1 |
| st[7] += 1 |
| p = e.split('.') |
| for L in range(3): |
| prefixes[L].add('.'.join(p[:3 - L])) |
| for L in range(3): |
| for pre in prefixes[L]: |
| lvl[L][pre] = lvl[L].get(pre, 0.0) + rb ** 4 |
| cands = sorted(stats, key=lambda e: (-stats[e][1], -freq.get(e, 0), e))[:MAX_CAND] |
| v4s = [stats[e][1] for e in cands] |
| margin = (v4s[0] - v4s[1]) / tot4 if len(v4s) > 1 else v4s[0] / tot4 |
| top_nec = len(labels[hits[0][0]]) |
| n10 = min(10, len(hits)) |
| X = np.zeros((len(cands), len(FEATURES)), dtype=np.float32) |
| for i, e in enumerate(cands): |
| st = stats[e] |
| p = e.split('.') |
| X[i] = [st[0] / tot1, st[1] / tot4, st[2] / tot8, i, st[3] / top_bits, st[3], st[4], st[5], |
| st[6], st[7], st[6] / n10, st[8], st[9], math.log10(st[10] + 1e-180), st[11], |
| lvl[0]['.'.join(p[:3])] / tot4, lvl[1]['.'.join(p[:2])] / tot4, |
| lvl[2][p[0]] / tot4, math.log1p(freq.get(e, 0)), top_bits, top_pid, |
| math.log(qlen), len(hits), len(stats), top_nec, margin, top_pid - st[4]] |
| return cands, X |
|
|
|
|
| def parse_line(line): |
| q, s, pid, length, ql, sl, ev, b, qs, qe, ss, se = line.rstrip('\n').split('\t') |
| return q, (s, float(pid), int(ql), int(sl), float(ev), float(b), int(qs), int(qe), int(ss), int(se)) |
|
|
