| import argparse |
| import json |
| import math |
| import os |
| import shutil |
| import stat |
| import subprocess |
| import tempfile |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
|
|
| ROOT = Path(__file__).parent |
| BETA, GATE, K_HITS = 0.7, 20, 30 |
| EXACT = 'ACDEFGHIKLMNPQRSTVWY' |
| GROUPS = ['LVIM', 'C', 'A', 'G', 'ST', 'P', 'FYW', 'EDNQ', 'KR', 'H'] |
| K, BASE = 6, 11 |
| MAXDF, MAXPOS, BAND, TOP = 2000, 60000, 32, 200 |
| KEYMUL = 2 * MAXPOS // BAND + 2 |
| LETTERS = np.frombuffer(b'ACDEFGHIKLMNPQRSTVWYX', dtype=np.uint8) |
| COLS = ['q', 't', 'bits', 'pident', 'qcov', 'tcov'] |
| NF = 29 |
|
|
| LUT = np.full(256, 20, dtype=np.int64) |
| for i, c in enumerate(EXACT): |
| LUT[ord(c)] = i |
| MAP = np.full(21, 10, dtype=np.int64) |
| for g, letters in enumerate(GROUPS): |
| for c in letters: |
| MAP[EXACT.index(c)] = g |
| LUT_NET = np.full(256, 21, dtype=np.int64) |
| for i, c in enumerate(EXACT): |
| LUT_NET[ord(c)] = i + 1 |
|
|
|
|
| def to_codes(seq): |
| return LUT[np.frombuffer(seq.encode('ascii', 'ignore'), dtype=np.uint8)] |
|
|
|
|
| def hits_of(codes): |
| codes = MAP[codes[:MAXPOS].astype(np.int64)] |
| n = len(codes) - K + 1 |
| if n <= 0: |
| return np.empty(0, dtype=np.int32), np.empty(0, dtype=np.int32) |
| ids = np.zeros(n, dtype=np.int64) |
| for j in range(K): |
| ids = ids * BASE + codes[j:j + n] |
| return ids.astype(np.int32), np.arange(n, dtype=np.int32) |
|
|
|
|
| def build_index(codes, lens): |
| starts = np.cumsum(lens) - lens |
| parts = [hits_of(codes[s:s + n]) for s, n in zip(starts, lens)] |
| sizes = np.array([len(p[0]) for p in parts]) |
| ids_all = np.concatenate([p[0] for p in parts]) |
| pos_all = np.concatenate([p[1] for p in parts]).astype(np.uint16) |
| del parts |
| order = np.argsort(ids_all, kind='stable') |
| kmer = ids_all[order] |
| del ids_all |
| pos = pos_all[order] |
| del pos_all |
| seq = np.repeat(np.arange(len(lens), dtype=np.int32), sizes)[order] |
| return {'kmer': kmer, 'pos': pos, 'seq': seq} |
|
|
|
|
| def candidates(index, codes): |
| empty = np.empty(0, dtype=np.int64) |
| ids, qpos = hits_of(codes) |
| if len(ids) == 0: |
| return empty, empty |
| lo = np.searchsorted(index['kmer'], ids, 'left') |
| cnt = np.searchsorted(index['kmer'], ids, 'right') - lo |
| keep = cnt <= MAXDF |
| lo, cnt, qpos = lo[keep], cnt[keep], qpos[keep] |
| total = cnt.sum() |
| if total == 0: |
| return empty, empty |
| offs = np.cumsum(cnt) - cnt |
| idx = np.arange(total) - np.repeat(offs, cnt) + np.repeat(lo, cnt) |
| diag = index['pos'][idx].astype(np.int64) - np.repeat(qpos, cnt).astype(np.int64) + MAXPOS |
| key = index['seq'][idx].astype(np.int64) * KEYMUL + diag // BAND |
| uk, counts = np.unique(key, return_counts=True) |
| nxt = np.minimum(np.searchsorted(uk, uk + 1), len(uk) - 1) |
| score = counts + np.where(uk[nxt] == uk + 1, counts[nxt], 0) |
| top = np.argsort(-score)[:TOP] |
| cand = uk[top] // KEYMUL |
| _, first = np.unique(cand, return_index=True) |
| first.sort() |
| return cand[first], score[top][first] |
|
|
|
|
| def vote(cand, score, lab_flat, lab_ptr, n_ec, k=5): |
| v = np.zeros(n_ec, dtype=np.float32) |
| for j, s in zip(cand[:k], score[:k]): |
| labs = lab_flat[lab_ptr[j]:lab_ptr[j + 1]] |
| v[labs] += s / len(labs) |
| total = v.sum() |
| return v / total if total > 0 else v |
|
|
|
|
| def red_votes(index, lab_flat, lab_ptr, seqs, n_ec): |
| out = np.zeros((len(seqs), n_ec), dtype=np.float32) |
| best = np.zeros(len(seqs)) |
| for r, seq in enumerate(seqs): |
| cand, s = candidates(index, to_codes(seq)) |
| out[r] = vote(cand, s, lab_flat, lab_ptr, n_ec) |
| best[r] = s[0] if len(s) else 0 |
| return out, best |
|
|
|
|
| def binary_path(): |
| flags = '' |
| try: |
| flags = Path('/proc/cpuinfo').read_text() |
| except OSError: |
| pass |
| path = ROOT.resolve() / ('mmseqs_avx2' if ' avx2' in flags else 'mmseqs_sse41') |
| try: |
| os.chmod(path, os.stat(path).st_mode | stat.S_IEXEC) |
| except OSError: |
| pass |
| return str(path) |
|
|
|
|
| def search(codes, lens, queries, threads=2, sens=4.0, max_seqs=50, timeout=480): |
| work = tempfile.mkdtemp(dir='/tmp') |
| try: |
| exe = binary_path() |
| starts = np.cumsum(lens) - lens |
| with open(os.path.join(work, 'target.fa'), 'wb') as f: |
| for i in range(len(lens)): |
| seq = LETTERS[codes[starts[i]:starts[i] + lens[i]].astype(np.int64)].tobytes() |
| f.write(b'>' + str(i).encode() + b'\n' + seq + b'\n') |
| with open(os.path.join(work, 'query.fa'), 'wb') as f: |
| for i, s in enumerate(queries): |
| f.write(b'>' + str(i).encode() + b'\n' + s.encode('ascii', 'ignore') + b'\n') |
|
|
| def run(*args): |
| subprocess.run([exe, *args, '-v', '1'], check=True, cwd=work, timeout=timeout, |
| stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) |
|
|
| run('createdb', 'target.fa', 'tDB') |
| run('createdb', 'query.fa', 'qDB') |
| run('search', 'qDB', 'tDB', 'res', 'tmp', '-s', str(sens), '--max-seqs', str(max_seqs), '-e', '10', |
| '--threads', str(threads), '--split-memory-limit', '6G') |
| run('convertalis', 'qDB', 'tDB', 'res', 'res.m8', '--format-output', 'query,target,bits,pident,qcov,tcov') |
| out = os.path.join(work, 'res.m8') |
| if os.path.getsize(out) == 0: |
| return pd.DataFrame(columns=COLS) |
| table = pd.read_csv(out, sep='\t', names=COLS) |
| return table.sort_values(['q', 'bits'], ascending=[True, False], kind='stable') |
| finally: |
| shutil.rmtree(work, ignore_errors=True) |
|
|
|
|
| def build_vocab_maps(ec_list): |
| l3 = {} |
| l2 = {} |
| l3of = np.zeros(len(ec_list), dtype=np.int32) |
| l2of = np.zeros(len(ec_list), dtype=np.int32) |
| for i, e in enumerate(ec_list): |
| p = e.split('.') |
| l3of[i] = l3.setdefault('.'.join(p[:3]), len(l3)) |
| l2of[i] = l2.setdefault('.'.join(p[:2]), len(l2)) |
| return l3of, l2of |
|
|
|
|
| def query_features(t, bits, pid, qcov, tcov, labs, prior_log, l3of, l2of, qlen, tlens): |
| prior_cnt = np.exp(prior_log) |
| K = len(t) |
| top = float(bits[0]) |
| rel = bits / top |
| r4 = rel ** 4 |
| r10 = rel ** 10 |
| cand = {} |
| tot4 = 0.0 |
| tot1 = 0.0 |
| l3s = {} |
| l2s = {} |
| for r in range(K): |
| L = labs[t[r]] |
| n = len(L) |
| w4 = r4[r] / n |
| w1 = rel[r] / n |
| tot4 += r4[r] |
| tot1 += rel[r] |
| for e in L: |
| d = cand.get(e) |
| if d is None: |
| d = [0.0] * 14 |
| d[5] = r |
| d[6] = bits[r] |
| d[7] = pid[r] |
| d[8] = n |
| d[9] = qcov[r] * tcov[r] |
| d[10] = tlens[t[r]] / max(qlen, 1) |
| cand[e] = d |
| d[0] += w1 |
| d[12] += pid[r] |
| d[13] += bits[r] |
| d[1] += w4 |
| d[2] += r10[r] / n |
| d[3] += 1 |
| if r < 5: |
| d[4] += 1 |
| if r < 10: |
| d[11] += 1 |
| for e in set(int(l3of[x]) for x in L): |
| l3s[e] = l3s.get(e, 0.0) + r4[r] / len(set(int(l3of[x]) for x in L)) |
| for e in set(int(l2of[x]) for x in L): |
| l2s[e] = l2s.get(e, 0.0) + r4[r] / len(set(int(l2of[x]) for x in L)) |
| ecs = list(cand.keys()) |
| F = np.zeros((len(ecs), NF), dtype=np.float32) |
| top_pid = float(pid[0]) |
| top_labs = set(labs[t[0]]) |
| for i, e in enumerate(ecs): |
| d = cand[e] |
| F[i, 0] = d[0] / tot1 |
| F[i, 1] = d[1] / tot4 |
| F[i, 2] = d[2] |
| F[i, 3] = d[3] |
| F[i, 4] = d[4] |
| F[i, 5] = d[5] |
| F[i, 6] = d[6] / top |
| F[i, 7] = d[7] |
| F[i, 8] = d[8] |
| F[i, 9] = d[9] |
| F[i, 10] = d[10] |
| F[i, 11] = d[11] |
| F[i, 12] = prior_log[e] |
| F[i, 13] = l3s[int(l3of[e])] / tot4 |
| F[i, 14] = l2s[int(l2of[e])] / tot4 |
| F[i, 15] = 1.0 if e in top_labs else 0.0 |
| F[i, 16] = np.log(top) |
| F[i, 17] = top_pid |
| F[i, 18] = np.log(max(qlen, 1)) |
| F[i, 19] = K |
| F[i, 20] = d[1] |
| F[i, 21] = len(top_labs) |
| F[i, 22] = bits[min(1, K - 1)] / top |
| F[i, 23] = d[6] |
| F[i, 24] = d[3] / prior_cnt[e] |
| F[i, 25] = d[12] / d[3] |
| F[i, 26] = d[13] / bits.sum() |
| F[i, 27] = np.log(d[3] / prior_cnt[e] + 1e-4) |
| F[i, 28] = pid.mean() |
| return ecs, F |
|
|
|
|
| def full_features(F): |
| v = F[:, 1] |
| v0 = F[:, 0] |
| srt = np.sort(v)[::-1] |
| ex = np.zeros((len(F), 7), dtype=np.float32) |
| ex[:, 0] = len(F) |
| ex[:, 1] = v - srt[0] |
| ex[:, 2] = v - (srt[1] if len(srt) > 1 else 0) |
| ex[:, 3] = (-v).argsort().argsort() |
| ex[:, 4] = v0 - v0.max() |
| pos = v[v > 0] |
| ex[:, 5] = -(pos * np.log(pos)).sum() |
| ex[:, 6] = F[:, 12] - F[:, 12].max() |
| return np.hstack([F, ex]) |
|
|
|
|
| class Ranker: |
| def __init__(self, z): |
| self.feat, self.thr = z['r_feat'], z['r_thr'] |
| self.left, self.right = z['r_left'], z['r_right'] |
| self.val, self.leaf = z['r_val'], z['r_leaf'] |
| self.base = float(z['r_base'][0]) |
|
|
| def score(self, X): |
| rows = np.arange(len(X)) |
| raw = np.full(len(X), self.base) |
| for i in range(len(self.feat)): |
| node = np.zeros(len(X), dtype=np.int64) |
| for _ in range(64): |
| lf = self.leaf[i][node] |
| if lf.all(): |
| break |
| f = self.feat[i][node] |
| go_left = X[rows, f] <= self.thr[i][node] |
| nxt = np.where(go_left, self.left[i][node], self.right[i][node]) |
| node = np.where(lf, node, nxt) |
| raw += self.val[i][node] |
| return raw |
|
|
|
|
| def rank_queries(hits, seqs, labs, prior_log, l3of, l2of, tlens, ranker, ec_list): |
| out = {} |
| q = hits.q.values |
| if len(q) == 0: |
| return out |
| cuts = np.concatenate([[0], np.where(q[1:] != q[:-1])[0] + 1, [len(q)]]) |
| t_all, b_all = hits.t.values, hits.bits.values.astype(float) |
| p_all, qc_all, tc_all = hits.pident.values, hits.qcov.values, hits.tcov.values |
| feats, owners = [], [] |
| for a, b in zip(cuts[:-1], cuts[1:]): |
| b = min(b, a + K_HITS) |
| qi = int(q[a]) |
| ecs, F = query_features(t_all[a:b], b_all[a:b], p_all[a:b], qc_all[a:b], tc_all[a:b], labs, prior_log, |
| l3of, l2of, len(seqs[qi]), tlens) |
| feats.append(full_features(F)) |
| owners.append((qi, ecs)) |
| scores = ranker.score(np.vstack(feats)) |
| pos = 0 |
| for qi, ecs in owners: |
| s = scores[pos:pos + len(ecs)] |
| pos += len(ecs) |
| out[qi] = ec_list[ecs[int(np.argmax(s))]] |
| return out |
|
|
|
|
| def network_scores(z, seqs): |
| import torch |
| import torch.nn.functional as F |
|
|
| torch.set_grad_enabled(False) |
| torch.set_num_threads(2) |
| P = {k[2:]: torch.from_numpy(z[k]) for k in z.files if k.startswith('p_')} |
| dilations = z['dilations'].tolist() |
| alpha = float(z['alpha'][0]) |
| bank = F.normalize(torch.from_numpy(z['bank'].astype(np.float32)), dim=1) |
| bank_flat = z['bank_flat'].astype(np.int64) |
| bank_ptr = np.concatenate([[0], np.cumsum(z['bank_len'])]) |
| n_ec = len(z['ec_list']) |
|
|
| def softmax(x): |
| e = torch.exp(x - x.max(1, keepdim=True).values) |
| return e / e.sum(1, keepdim=True) |
|
|
| def encode(seq): |
| codes = LUT_NET[np.frombuffer(seq.encode('ascii', 'ignore'), dtype=np.uint8)] |
| if len(codes) > 800: |
| start = (len(codes) - 800) // 2 |
| codes = codes[start:start + 800] |
| return codes |
|
|
| def forward(tokens): |
| m = (tokens > 0).unsqueeze(-1).to(P['embed.w'].dtype) |
| x = P['embed.w'][tokens] * m |
| for i, d in enumerate(dilations): |
| pre = f'b{i}.' |
| w = P[pre + 'dw.w'] |
| xc = x.transpose(1, 2).contiguous() |
| h = F.conv1d(xc, w, P[pre + 'dw.b'], padding=d * (w.shape[-1] // 2), dilation=d, groups=xc.shape[1]) |
| h = h.transpose(1, 2).contiguous() * m |
| mu = h.mean(-1, keepdim=True) |
| var = h.var(-1, unbiased=False, keepdim=True) |
| h = P[pre + 'ln.g'] * ((h - mu) * (1.0 / torch.sqrt(var + 1e-6))) + P[pre + 'ln.b'] |
| h = h @ P[pre + 'pw1.w'] + P[pre + 'pw1.b'] |
| h = h * 0.5 * (1.0 + torch.erf(h / math.sqrt(2.0))) |
| h = h @ P[pre + 'pw2.w'] + P[pre + 'pw2.b'] |
| x = (x + h) * m |
| n = m.sum(1).clamp(min=1.0) |
| mean = x.sum(1) / n |
| mx, idx = (x - (1.0 - m) * 1e4).max(1) |
| pooled = torch.cat([mx, mean], 1) |
| mu = pooled.mean(-1, keepdim=True) |
| var = pooled.var(-1, unbiased=False, keepdim=True) |
| z1 = P['head.ln.g'] * ((pooled - mu) * (1.0 / torch.sqrt(var + 1e-6))) + P['head.ln.b'] |
| e = z1 @ P['head.proj.w'] + P['head.proj.b'] |
| a = e * 0.5 * (1.0 + torch.erf(e / math.sqrt(2.0))) |
| return a @ P['head.out.w'] + P['head.out.b'], e |
|
|
| order = np.argsort([len(s) for s in seqs], kind='stable') |
| emb = np.zeros((len(seqs), 128), np.float32) |
| probs = np.zeros((len(seqs), n_ec), np.float32) |
| for s in range(0, len(seqs), 64): |
| ids = order[s:s + 64] |
| enc = [encode(seqs[i]) for i in ids] |
| x = np.zeros((len(enc), max(len(e) for e in enc)), dtype=np.int64) |
| for r, e in enumerate(enc): |
| x[r, :len(e)] = e |
| logits, e = forward(torch.from_numpy(x)) |
| emb[ids] = e.cpu().numpy() |
| probs[ids] = softmax(logits).cpu().numpy() |
| q = F.normalize(torch.from_numpy(emb), dim=1) |
| votes = np.zeros((len(q), n_ec), np.float32) |
| for s in range(0, len(q), 256): |
| v, idx = (q[s:s + 256] @ bank.T).topk(5, dim=1) |
| w = softmax(v / 0.05).numpy() |
| idx = idx.numpy() |
| for r in range(len(idx)): |
| for j in range(5): |
| labs = bank_flat[bank_ptr[idx[r, j]]:bank_ptr[idx[r, j] + 1]] |
| votes[s + r, labs] += w[r, j] / len(labs) |
| return alpha * votes + (1 - alpha) * probs |
|
|
|
|
| def predict_file(input_path, output_path): |
| t0 = time.time() |
| frame = pd.read_parquet(input_path)[['id', 'sequence']] |
| if frame.id.isna().any() or frame.id.duplicated().any(): |
| raise ValueError('Input IDs must be present and unique.') |
| frame = frame.drop_duplicates('sequence', keep='first').reset_index(drop=True) |
| seqs = frame.sequence.tolist() |
|
|
| z = np.load(ROOT / 'model.npz') |
| ec_list = z['ec_list'].tolist() |
| n_ec = len(ec_list) |
| codes, lens = z['codes'], z['lens'] |
| lab_flat = z['lab_flat'].astype(np.int64) |
| lab_ptr = np.concatenate([[0], np.cumsum(z['lab_len'])]) |
|
|
| chosen = {} |
| try: |
| labs = [lab_flat[lab_ptr[i]:lab_ptr[i + 1]] for i in range(len(lens))] |
| prior_log = np.log(np.maximum(np.bincount(lab_flat, minlength=n_ec), 1).astype(float)) |
| l3of, l2of = build_vocab_maps(ec_list) |
| ranker = Ranker(z) |
| hits = search(codes, lens, seqs) |
| print('search finished after', round(time.time() - t0), 's, hits:', len(hits), flush=True) |
| chosen = rank_queries(hits, seqs, labs, prior_log, l3of, l2of, lens, ranker, ec_list) |
| except Exception as error: |
| print('search or ranking failed, using the backup method for every input:', error, flush=True) |
| chosen = {} |
| rest = np.array([i for i in range(len(seqs)) if i not in chosen], dtype=np.int64) |
| print('ranked inputs:', len(chosen), ' inputs without a search hit:', len(rest), ' of', len(seqs), flush=True) |
|
|
| preds = [chosen.get(i) for i in range(len(seqs))] |
| if len(rest): |
| index = build_index(codes, lens) |
| red, best = red_votes(index, lab_flat, lab_ptr, [seqs[i] for i in rest], n_ec) |
| weak = np.where(best < GATE)[0] |
| if len(weak): |
| try: |
| net = network_scores(z, [seqs[rest[i]] for i in weak]) |
| red[weak] = BETA * red[weak] + (1 - BETA) * net |
| except Exception as error: |
| print('network step skipped:', error, flush=True) |
| for r, i in enumerate(rest): |
| preds[i] = ec_list[int(red[r].argmax())] |
|
|
| with Path(output_path).open('w') as output: |
| for key, ec in zip(frame.id, preds): |
| output.write(json.dumps({'id': key, 'prediction': [ec]}) + '\n') |
| print('total time', round(time.time() - t0), 's', flush=True) |
|
|
|
|
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--input', required=True) |
| parser.add_argument('--output', required=True) |
| args = parser.parse_args() |
| predict_file(args.input, args.output) |
|
|
