| """Shared by training (notebook) and inference (predict.py): candidate features + a numpy tree-ensemble evaluator.""" |
| import json, math |
| import numpy as np |
|
|
| FEATS = ['mm_share', 'mm_maxb_ratio', 'mm_maxb', 'mm_cnt', 'mm_rank', 'mm_maxpid', 'mm_maxqcov', 'mm_maxtcov', 'topbits', 'top_pid', 'qlen', 'nhits', |
| 'h_prob', 'h_rank', 'h_max', 'mm_l3share', 'h_l3mass', 'logprior', 'ncand_mm', |
| 'in_top1', 'top3_cnt', 'share_gap', 'share_rank', 'h_rel', 'h_gap', 'top1_nlab', 'sup_nlab', 'tb_per_len', 'mb_per_len', 'q_margin', 'q_entropy', 'q_nclose', 'bits_gap', 'pid_gap', 'cnt_frac', 'maxb_rank', 'k_share', 'k_max', 'k_cnt', 'k_top1', |
| 'n2_share', 'n2_gap', 'n2_rank', 'n_pur', 'n_top1', 'is_unseen', 'fam_nunseen', 'fam_nseen', 'fam_rank', |
| 'hs_prob', 'hs_rank_fam', 'hs_gap_fam', 'ps_sim', 'ps_rank_fam', 'ps_gap_fam', 'ps_gap_all', |
| 'mo_prob', 'mo_rank', 'mo_fam', 'mo_fam_rank', 'mo_gap_fam', 'mo_l3mass'] |
| KNN_FEATS = ['k_share', 'k_max', 'k_cnt', 'k_top1'] |
| CTX_FEATS = ['q_margin', 'q_entropy', 'q_nclose', 'bits_gap', 'pid_gap', 'cnt_frac', 'maxb_rank'] |
| UNSEEN_FEATS = ['is_unseen', 'fam_nunseen', 'fam_nseen', 'fam_rank'] |
| NB_FEATS = ['n2_share', 'n2_gap', 'n2_rank', 'n_pur', 'n_top1'] |
| HS_FEATS = ['hs_prob', 'hs_rank_fam', 'hs_gap_fam'] |
| PS_FEATS = ['ps_sim', 'ps_rank_fam', 'ps_gap_fam', 'ps_gap_all'] |
| HIER_FEATS = HS_FEATS + PS_FEATS |
| MOTIF_FEATS = ['mo_prob', 'mo_rank', 'mo_fam', 'mo_fam_rank', 'mo_gap_fam', 'mo_l3mass'] |
| HEAD_FEATS = ['h_prob', 'h_rank', 'h_max', 'h_l3mass', 'h_rel', 'h_gap'] + KNN_FEATS + HIER_FEATS |
| COLS = ['q', 'e'] + FEATS |
| NAN = float('nan') |
|
|
|
|
| class Ctx: |
| """labels: list (per target sequence) of tuples of EC indices; ecs: EC strings; prior: per-EC counts.""" |
| def __init__(self, labels, ecs, prior): |
| self.labels, self.ecs = labels, list(ecs) |
| self.prior = np.asarray(prior, float); self.logprior = np.log1p(self.prior) |
| ids = {} |
| self.l3_id = np.array([ids.setdefault('.'.join(e.split('.')[:3]), len(ids)) for e in self.ecs]); self.nl3 = len(ids) |
| self.top_ec = self.ecs[int(self.prior.argmax())] |
| self.nb_e = self.nb_s = None |
| |
| def key4(e): |
| f = e.split('.')[-1] |
| return (0, int(f), f) if f.isdigit() else (1, 0, f) |
| un = [i for i in range(len(self.ecs)) if self.prior[i] <= 0] |
| self.unseen_by_l3 = {} |
| for i in sorted(un, key=lambda i: key4(self.ecs[i])): self.unseen_by_l3.setdefault(int(self.l3_id[i]), []).append(i) |
| self.fam_members = [np.flatnonzero(self.l3_id == f) for f in range(self.nl3)] |
| self.nunseen_l3 = np.zeros(self.nl3); self.nseen_l3 = np.zeros(self.nl3) |
| for i in range(len(self.ecs)): |
| if self.prior[i] <= 0: self.nunseen_l3[self.l3_id[i]] += 1 |
| else: self.nseen_l3[self.l3_id[i]] += 1 |
| self.use_unseen = False |
| self.use_hier_cands = False |
| self.use_motif_cands = False |
|
|
|
|
| def set_nb(self, E, S): |
| """E[j], S[j]: the top EC indices among database sequence j's own neighbours and their (bits^2-weighted) shares; S==0 marks empty slots.""" |
| self.nb_e = None if E is None else np.asarray(E); self.nb_s = None if S is None else np.asarray(S, np.float64) |
|
|
|
|
| def query_rows(qi, hit, qlen, h, ctx, TOPM=50, TOPH=10, knn=None, TOPK=5, TOPN=20, TOPN2C=3, h2=None, ps=None, mot=None, TOPMO=3): |
| """hit: None or (t, bits, pident, qcov, tcov) arrays sorted by bits desc; h: None or head probabilities over all ECs.""" |
| if hit is not None: |
| t, b, pid, qc, tc = (np.asarray(a)[:TOPM] for a in hit); tb = float(b[0]); tpid = float(pid[0]) |
| else: |
| t = b = pid = qc = tc = np.empty(0); tb = 0.0; tpid = 0.0 |
| votes, maxb, cnt, rank, mpid, mqc, mtc = {}, {}, {}, {}, {}, {}, {} |
| nlab = {} |
| for r, (j, bj, pj, qj, tj) in enumerate(zip(t, b, pid, qc, tc)): |
| l = ctx.labels[int(j)] |
| for e in l: |
| votes[e] = votes.get(e, 0.0) + float(bj) ** 2 / len(l); maxb[e] = max(maxb.get(e, 0.0), float(bj)); cnt[e] = cnt.get(e, 0) + 1; rank.setdefault(e, r); nlab[e] = nlab.get(e, 0) + len(l) |
| mpid[e] = max(mpid.get(e, 0.0), float(pj)); mqc[e] = max(mqc.get(e, 0.0), float(qj)); mtc[e] = max(mtc.get(e, 0.0), float(tj)) |
| tot = sum(votes.values()) or 1.0 |
| l3v = {} |
| for e, v in votes.items(): l3v[ctx.l3_id[e]] = l3v.get(ctx.l3_id[e], 0.0) + v / tot |
| if h is not None: |
| l3m = np.bincount(ctx.l3_id, weights=h, minlength=ctx.nl3); hmax = float(h.max()) |
| cands = set(votes) | set(np.argpartition(-h, TOPH)[:TOPH].tolist()) |
| else: |
| l3m = None; hmax = NAN; cands = set(votes) |
| if knn is not None: |
| ki, ks = knn; ks = np.asarray(ks, float); kw = np.exp((ks - ks.max()) * 20.0); kv, km, kc = {}, {}, {} |
| for j, sj, wj in zip(ki, ks, kw): |
| l = ctx.labels[int(j)] |
| for e in l: kv[e] = kv.get(e, 0.0) + wj / len(l); km[e] = max(km.get(e, -1.0), float(sj)); kc[e] = kc.get(e, 0) + 1 |
| ktot = sum(kv.values()) or 1.0; ktop = float(ks.max()); cands |= set(sorted(kv, key=lambda e: -kv[e])[:TOPK]) |
| fam_rank = {} |
| if ctx.use_unseen: |
| fams = sorted(l3v, key=lambda f: -l3v[f])[:2] |
| if l3m is not None: fams += [int(np.argmax(l3m))] |
| for r_, f in enumerate(fams): |
| fam_rank.setdefault(int(f), r_) |
| u = ctx.unseen_by_l3.get(int(f)) |
| if u: cands.add(u[0]) |
| use_h2, use_ps = h2 is not None, ps is not None |
| if ctx.use_hier_cands and (use_h2 or use_ps): |
| fams2 = sorted(l3v, key=lambda f: -l3v[f])[:2] |
| if l3m is not None: fams2 += [int(np.argmax(l3m))] |
| for f in set(int(x) for x in fams2): |
| mem = ctx.fam_members[f] |
| if use_h2: cands.add(int(mem[int(np.argmax(h2[mem]))])) |
| if use_ps: |
| v = ps[mem] |
| if np.isfinite(v).any(): cands.add(int(mem[int(np.nanargmax(v))])) |
| use_mo = mot is not None |
| if use_mo: |
| mfam = np.bincount(ctx.l3_id, weights=mot, minlength=ctx.nl3) |
| if ctx.use_motif_cands: cands |= set(np.argpartition(-mot, TOPMO)[:TOPMO].tolist()) |
| use_nb = ctx.nb_e is not None |
| n2, pur_v, pur_w, nt1 = {}, {}, {}, None |
| if use_nb and len(t): |
| wsum = 0.0 |
| for r, (j, bj) in enumerate(zip(t[:TOPN], b[:TOPN])): |
| j = int(j); w = float(bj) ** 2; sj = ctx.nb_s[j] |
| if sj[0] <= 0: |
| continue |
| nd = dict(zip(ctx.nb_e[j].tolist(), sj.tolist())); wsum += w |
| if r == 0: nt1 = nd |
| for e, v in nd.items(): |
| if v > 0: n2[e] = n2.get(e, 0.0) + w * v |
| for e in ctx.labels[j]: |
| pur_v[e] = pur_v.get(e, 0.0) + w * nd.get(e, 0.0); pur_w[e] = pur_w.get(e, 0.0) + w |
| if wsum > 0: n2 = {e: v / wsum for e, v in n2.items()} |
| cands |= set(sorted(n2, key=lambda e: -n2[e])[:TOPN2C]) |
| top1 = set(ctx.labels[int(t[0])]) if len(t) else set() |
| top3 = {} |
| for j in t[:3]: |
| for e in ctx.labels[int(j)]: top3[e] = top3.get(e, 0) + 1 |
| shares = {e: votes.get(e, 0.0) / tot for e in cands}; smax = max(shares.values()) if shares else 0.0 |
| srank = {e: r for r, e in enumerate(sorted(shares, key=lambda x: -shares[x]))} |
| mbs = {e: maxb.get(e, 0.0) for e in cands}; mbmax = max(mbs.values()) if mbs else 0.0 |
| mps = {e: mpid.get(e, 0.0) for e in cands}; mpmax = max(mps.values()) if mps else 0.0 |
| brank = {e: r for r, e in enumerate(sorted(mbs, key=lambda x: -mbs[x]))} |
| sv = sorted(shares.values(), reverse=True) |
| qmargin = (sv[0] - (sv[1] if len(sv) > 1 else 0.0)) if sv else 0.0 |
| pp = np.array([v for v in sv if v > 0]); qent = float(-(pp / pp.sum() * np.log(pp / pp.sum())).sum()) if len(pp) else 0.0 |
| nclose = sum(1 for v in sv if sv[0] > 0 and v >= 0.5 * sv[0]); nt = max(len(t), 1) |
| n2c = {e: n2.get(e, 0.0) for e in cands}; n2max = max(n2c.values()) if n2c else 0.0 |
| n2rank = {e: r for r, e in enumerate(sorted(n2c, key=lambda x: -n2c[x]))} |
| mor, mofr, mofg = {}, {}, {} |
| if use_mo: |
| for e in cands: mor[e] = min(int((mot > mot[e]).sum()), 99) |
| for f in set(int(ctx.l3_id[e]) for e in cands): |
| mem = ctx.fam_members[f]; v = mot[mem]; mx = float(v.max()) |
| for e in cands: |
| if int(ctx.l3_id[e]) == f: mofr[e] = int((v > mot[e]).sum()); mofg[e] = (float(mot[e]) - mx) / max(float(mfam[f]), 1e-12) |
| hsr, hsg, psr, psg, psa = {}, {}, {}, {}, {} |
| if use_h2 or use_ps: |
| for f in set(int(ctx.l3_id[e]) for e in cands): |
| mem = ctx.fam_members[f] |
| if use_h2: |
| v = h2[mem]; mx = float(v.max()); order = {int(e): r for r, e in enumerate(mem[np.argsort(-v, kind='stable')])} |
| for e in cands: |
| if int(ctx.l3_id[e]) == f: hsr[e] = order[int(e)]; hsg[e] = float(h2[e]) - mx |
| if use_ps: |
| v = ps[mem]; fin = np.isfinite(v) |
| if fin.any(): |
| mx = float(np.nanmax(v)); vv = np.where(fin, v, -np.inf); order = {int(e): r for r, e in enumerate(mem[np.argsort(-vv, kind='stable')])} |
| for e in cands: |
| if int(ctx.l3_id[e]) == f and np.isfinite(ps[e]): psr[e] = order[int(e)]; psg[e] = float(ps[e]) - mx |
| if use_ps: |
| pc = [float(ps[e]) for e in cands if np.isfinite(ps[e])]; pcmax = max(pc) if pc else NAN |
| rows = [] |
| for e in cands: |
| rows.append((qi, e, votes.get(e, 0.0) / tot, maxb.get(e, 0.0) / tb if tb > 0 else 0.0, maxb.get(e, 0.0), cnt.get(e, 0), rank.get(e, 99), |
| mpid.get(e, 0.0), mqc.get(e, 0.0), mtc.get(e, 0.0), tb, tpid, qlen, len(t), |
| float(h[e]) if h is not None else NAN, int((h > h[e]).sum()) if h is not None else NAN, hmax, |
| l3v.get(ctx.l3_id[e], 0.0), float(l3m[ctx.l3_id[e]]) if h is not None else NAN, float(ctx.logprior[e]), len(votes), |
| int(e in top1), top3.get(e, 0), shares[e] - smax, srank[e], |
| float(h[e]) / hmax if h is not None and hmax > 0 else NAN, float(h[e]) - hmax if h is not None else NAN, |
| len(ctx.labels[int(t[0])]) if len(t) else 0, nlab.get(e, 0) / cnt[e] if cnt.get(e, 0) else 0.0, |
| tb / max(qlen, 1), maxb.get(e, 0.0) / max(qlen, 1), |
| qmargin, qent, nclose, mbs[e] - mbmax, mps[e] - mpmax, cnt.get(e, 0) / nt, brank[e], |
| kv.get(e, 0.0) / ktot if knn is not None else NAN, km.get(e, -1.0) if knn is not None else NAN, kc.get(e, 0) if knn is not None else NAN, ktop if knn is not None else NAN, |
| n2c[e] if use_nb else NAN, n2c[e] - n2max if use_nb else NAN, n2rank[e] if use_nb else NAN, |
| pur_v[e] / pur_w[e] if use_nb and pur_w.get(e, 0) > 0 else NAN, |
| (nt1.get(e, 0.0) if nt1 is not None else NAN) if use_nb else NAN, |
| int(ctx.prior[e] <= 0), float(ctx.nunseen_l3[ctx.l3_id[e]]), float(ctx.nseen_l3[ctx.l3_id[e]]), fam_rank.get(int(ctx.l3_id[e]), 9), |
| float(h2[e]) if use_h2 else NAN, hsr.get(e, NAN) if use_h2 else NAN, hsg.get(e, NAN) if use_h2 else NAN, |
| float(ps[e]) if use_ps else NAN, psr.get(e, NAN) if use_ps else NAN, psg.get(e, NAN) if use_ps else NAN, |
| (float(ps[e]) - pcmax) if use_ps and np.isfinite(ps[e]) else NAN, |
| float(mot[e]) if use_mo else NAN, mor.get(e, NAN) if use_mo else NAN, |
| float(mot[e]) / max(float(mfam[ctx.l3_id[e]]), 1e-12) if use_mo else NAN, mofr.get(e, NAN) if use_mo else NAN, |
| mofg.get(e, NAN) if use_mo else NAN, float(mfam[ctx.l3_id[e]]) if use_mo else NAN)) |
| return rows |
|
|
|
|
| |
| def forest_from_lgb(booster): |
| out = [] |
| for ti in booster.dump_model()['tree_info']: |
| feat, thr, left, right, dl, mnan, val = [], [], [], [], [], [], [] |
| def add(node): |
| k = len(feat); feat.append(-1); thr.append(0.0); left.append(-1); right.append(-1); dl.append(False); mnan.append(False); val.append(0.0) |
| if 'leaf_value' in node: val[k] = float(node['leaf_value']); return k |
| assert node['decision_type'] == '<=' and node['missing_type'] in ('None', 'NaN'), node['missing_type'] |
| feat[k] = int(node['split_feature']); thr[k] = float(node['threshold']); dl[k] = bool(node['default_left']); mnan[k] = node['missing_type'] == 'NaN' |
| left[k] = add(node['left_child']); right[k] = add(node['right_child']); return k |
| add(ti['tree_structure']) |
| out.append(dict(feat=feat, thr=thr, left=left, right=right, dl=dl, mnan=mnan, val=val)) |
| return out |
|
|
|
|
| class Forest: |
| def __init__(self, trees): |
| self.trees = [{k: np.asarray(v, dtype={'feat': np.int64, 'left': np.int64, 'right': np.int64, 'dl': bool, 'mnan': bool}.get(k, np.float64)) for k, v in t.items()} for t in trees] |
| def predict(self, X): |
| X = np.asarray(X, np.float64); n = len(X); out = np.zeros(n) |
| for tr in self.trees: |
| node = np.zeros(n, np.int64) |
| while True: |
| idx = np.where(tr['feat'][node] >= 0)[0] |
| if len(idx) == 0: break |
| nd = node[idx]; x = X[idx, tr['feat'][nd]]; isn = np.isnan(x) |
| left = np.where(isn, np.where(tr['mnan'][nd], tr['dl'][nd], 0.0 <= tr['thr'][nd]), x <= tr['thr'][nd]) |
| node[idx] = np.where(left, tr['left'][nd], tr['right'][nd]) |
| out += tr['val'][node] |
| return out |
|
|
|
|
| |
| def gelu(x): |
| from scipy.special import erf |
| return 0.5 * x * (1.0 + erf(x / math.sqrt(2.0))) |
|
|
| def softmax(z): |
| z = z - z.max(1, keepdims=True); e = np.exp(z); return e / e.sum(1, keepdims=True) |
|
|
|
|
| def family_softmax(z, l3_id, nl3): |
| """z: logits over all ECs, 1-D or (n, n_ECs). -> P(EC | its own L3 family) for every EC.""" |
| z = np.asarray(z, np.float64); one = z.ndim == 1; Z = z[None] if one else z |
| ez = np.exp(Z - Z.max(1, keepdims=True)) |
| S = np.zeros((len(Z), nl3)); np.add.at(S.T, l3_id, ez.T) |
| out = ez / np.maximum(S[:, l3_id], 1e-300) |
| return out[0] if one else out |
|
|
|
|
| |
| _KLUT = np.full(256, 20, np.int64) |
| for _i, _a in enumerate('ACDEFGHIKLMNPQRSTVWY'): _KLUT[ord(_a)] = _i |
|
|
| def kmer_ids(seq, ks=(3, 4, 5, 6), nb=2 ** 20, maxlen=1000): |
| """Hashed ids of every k-mer (k in ks) of the first maxlen residues. Identical in training and inference.""" |
| c = _KLUT[np.frombuffer(seq[:maxlen].encode('ascii', 'ignore'), np.uint8)] |
| out = [] |
| for k in ks: |
| n = len(c) - k + 1 |
| if n <= 0: continue |
| code = np.zeros(n, np.int64) |
| for j in range(k): code = code * 21 + c[j:j + n] |
| out.append((code * 2654435761 + k * 97531) % nb) |
| return np.concatenate(out).astype(np.int64) if out else np.zeros(1, np.int64) |
|
|
| def motif_probs(seqs, mw, ks=(3, 4, 5, 6), batch=1000): |
| """mw: dict(table (nb, D) float16, W (nEC, D), b (nEC)). -> (n, nEC) softmax probabilities.""" |
| T = mw['table']; W = mw['W'].astype(np.float32); b = mw['b'].astype(np.float32); nb = T.shape[0] |
| out = np.zeros((len(seqs), len(b)), np.float32) |
| for st in range(0, len(seqs), batch): |
| X = np.stack([T[kmer_ids(s, ks, nb)].astype(np.float32).mean(0) for s in seqs[st:st + batch]]) |
| out[st:st + batch] = softmax(X @ W.T + b) |
| return out |
|
|
