| import numpy as np |
| import pandas as pd |
|
|
| HIT_FIELDS = ['query', 'target', 'pident', 'qcov', 'tcov', 'bits', 'evalue', 'qlen', 'tlen', 'alnlen'] |
|
|
| FEATS = ['sum_bits', 'max_bits', 'n_sup', 'best_rank', 'pid_max', 'pid_best', 'qcov_max', 'tcov_max', 'nec_best', |
| 'top_bits', 'tot_bits', 'n_hits', 'top_pident', 'r_max', 'r_sum', 'r_n', 'log_top', 'ec_logfreq', |
| 'l3_support', 'l2_support', 'ec_share_l3', 'qlen', 'tlen_best', 'aln_frac_best', 'bits_per_qlen', |
| 'log_evalue_best'] |
|
|
|
|
| def prefix(ec, k): |
| return '.'.join(ec.split('.')[:k]) |
|
|
|
|
| def l3_totals(freq): |
| tot = {} |
| for e, c in freq.items(): |
| p = prefix(e, 3) |
| tot[p] = tot.get(p, 0) + c |
| return tot |
|
|
|
|
| def build_candidates(h, ec_map, freq, kmax, l3_tot): |
| """One row per (query, candidate EC) with features computed from the top-kmax alignment hits.""" |
| if len(h) == 0: |
| return pd.DataFrame(columns=['query', 'ec'] + FEATS) |
| H = h.sort_values(['query', 'bits'], ascending=[True, False], kind='mergesort') |
| H = H.groupby('query', sort=False).head(kmax).copy() |
| H['rank'] = H.groupby('query', sort=False).cumcount() |
| qs = H.groupby('query', sort=False).agg(top_bits=('bits', 'max'), tot_bits=('bits', 'sum'), |
| n_hits=('bits', 'size'), top_pident=('pident', 'first'), |
| qlen=('qlen', 'first')) |
| H['ecs'] = H['target'].map(ec_map) |
| H['n_ec'] = H['ecs'].map(len) |
| H = H.explode('ecs').rename(columns={'ecs': 'ec'}) |
| H['l3'] = H['ec'].map(lambda e: prefix(e, 3)) |
| H['l2'] = H['ec'].map(lambda e: prefix(e, 2)) |
| H['log_ev'] = np.log10(H['evalue'].astype(float) + 1e-200) |
| a = H.groupby(['query', 'ec'], sort=False).agg( |
| sum_bits=('bits', 'sum'), max_bits=('bits', 'max'), n_sup=('bits', 'size'), best_rank=('rank', 'min'), |
| pid_max=('pident', 'max'), pid_best=('pident', 'first'), qcov_max=('qcov', 'max'), tcov_max=('tcov', 'max'), |
| nec_best=('n_ec', 'first'), tlen_best=('tlen', 'first'), aln_best=('alnlen', 'first'), |
| log_evalue_best=('log_ev', 'first'), l3=('l3', 'first'), l2=('l2', 'first')).reset_index() |
| for lvl in ('l3', 'l2'): |
| s = (H.drop_duplicates(['query', 'target', lvl]).groupby(['query', lvl], sort=False)['bits'] |
| .sum().rename(lvl + '_bits').reset_index()) |
| a = a.merge(s, on=['query', lvl], how='left') |
| a = a.join(qs, on='query') |
| a['r_max'] = a.max_bits / a.top_bits |
| a['r_sum'] = a.sum_bits / a.tot_bits |
| a['r_n'] = a.n_sup / a.n_hits |
| a['log_top'] = np.log(a.top_bits) |
| a['ec_logfreq'] = np.log1p(a['ec'].map(lambda e: freq.get(e, 0))) |
| a['ec_share_l3'] = a['ec'].map(lambda e: freq.get(e, 0)) / a['l3'].map(lambda p: max(l3_tot.get(p, 0), 1)) |
| a['aln_frac_best'] = a.aln_best / a.qlen |
| a['bits_per_qlen'] = a.max_bits / a.qlen |
| a['l3_support'] = a.l3_bits / a.tot_bits |
| a['l2_support'] = a.l2_bits / a.tot_bits |
| a['tlen_best'] = a['tlen_best'].astype(float) |
| return a[['query', 'ec', 'ec_logfreq'] + [f for f in FEATS if f != 'ec_logfreq']] |
|
|
|
|
| def select_from_scores(c, rel=0.5, max_ec=3): |
| """c: DataFrame with query, ec, ec_logfreq, p. Returns {query: best-first EC list}.""" |
| out = {} |
| if len(c) == 0: |
| return out |
| c = c.sort_values(['query', 'p', 'ec_logfreq', 'ec'], ascending=[True, False, False, True], kind='mergesort') |
| for q, g in c.groupby('query', sort=False): |
| ecs = g['ec'].tolist() |
| ps = g['p'].tolist() |
| top = ps[0] |
| out[q] = [ecs[0]] + [e for e, pp in zip(ecs[1:max_ec], ps[1:max_ec]) if top > 0 and pp >= rel * top] |
| return out |
|
|
|
|
| def rerank_select(cand, model, rel=0.5, max_ec=3): |
| if len(cand) == 0: |
| return {} |
| c = cand[['query', 'ec', 'ec_logfreq']].copy() |
| c['p'] = model.predict(cand[FEATS].to_numpy(dtype=float)) |
| return select_from_scores(c, rel, max_ec) |
|
|
