| import argparse |
| import json |
| import subprocess |
| import sys |
| import tempfile |
| import time |
| from pathlib import Path |
|
|
| import pandas as pd |
|
|
| HERE = Path(__file__).resolve().parent |
| sys.path.insert(0, str(HERE)) |
| import knn |
|
|
|
|
| def search_chunk(rows, seqs, cfg): |
| with tempfile.TemporaryDirectory(dir='/tmp') as td: |
| td = Path(td) |
| with open(td / 'q.fa', 'w') as f: |
| for i in rows: |
| f.write('>q%d\n%s\n' % (i, seqs[i])) |
| out = td / 'hits.tsv' |
| cmd = [str(HERE / 'bin' / 'mmseqs'), 'easy-search', str(td / 'q.fa'), str(HERE / 'db' / 'trainDB'), str(out), |
| str(td / 'work'), '-s', str(cfg['sens']), '--max-seqs', str(cfg['max_seqs']), '-e', '10', |
| '--split-memory-limit', cfg['split_mem'], '--threads', str(cfg['threads']), |
| '--format-output', 'query,target,bits', '-v', '1'] |
| r = subprocess.run(cmd, capture_output=True, text=True) |
| if r.returncode != 0: |
| sys.stderr.write(r.stdout[-2000:] + '\n' + r.stderr[-2000:] + '\n') |
| sys.exit('mmseqs failed with code %d' % r.returncode) |
| if out.exists() and out.stat().st_size > 0: |
| return pd.read_csv(out, sep='\t', names=['query', 'target', 'bits']) |
| return pd.DataFrame(columns=['query', 'target', 'bits']) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--input', required=True) |
| ap.add_argument('--output', required=True) |
| args = ap.parse_args() |
| t0 = time.perf_counter() |
| cfg = json.loads((HERE / 'config.json').read_text()) |
| ec_map = json.loads((HERE / 'db' / 'ec_map.json').read_text()) |
| fallback = json.loads((HERE / 'db' / 'fallback.json').read_text()) |
| freq = json.loads((HERE / 'db' / 'freq.json').read_text()) |
|
|
| df = pd.read_parquet(args.input, columns=['id', 'sequence']) |
| uniq = df.drop_duplicates('sequence', keep='first').reset_index(drop=True) |
| seqs = uniq['sequence'].tolist() |
| order = sorted(range(len(seqs)), key=lambda i: len(seqs[i])) |
| chunks = [order[i:i + cfg['chunk']] for i in range(0, len(order), cfg['chunk'])] |
|
|
| hl, times, skipped = {}, [], 0 |
| for ci, rows in enumerate(chunks): |
| elapsed = time.perf_counter() - t0 |
| if times and elapsed + 1.3 * max(times) > cfg['time_budget']: |
| skipped = sum(len(c) for c in chunks[ci:]) |
| sys.stderr.write('TIME GUARD: stopping before chunk %d; %d inputs get the fallback EC\n' % (ci, skipped)) |
| break |
| tc = time.perf_counter() |
| hits = search_chunk(rows, seqs, cfg) |
| hl.update(knn.hits_from_df(hits, cfg['kmax'])) |
| times.append(time.perf_counter() - tc) |
| sys.stderr.write('chunk %d/%d done in %.0fs (elapsed %.0fs)\n' % (ci + 1, len(chunks), times[-1], time.perf_counter() - t0)) |
|
|
| with open(args.output, 'w') as f: |
| for i, rid in enumerate(uniq['id']): |
| ecs = knn.select_ecs(hl.get('q%d' % i, []), ec_map, fallback, cfg['K'], cfg['pw'], cfg['ratio'], cfg['max_ec'], freq) |
| f.write(json.dumps({'id': rid, 'prediction': ecs}) + '\n') |
| sys.stderr.write('total %.0fs for %d unique inputs (%d without search)\n' % (time.perf_counter() - t0, len(seqs), skipped)) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
