| 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 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()) |
|
|
| df = pd.read_parquet(args.input, columns=['id', 'sequence']) |
| |
| uniq = df.drop_duplicates('sequence', keep='first').reset_index(drop=True) |
| qid = ['q%d' % i for i in range(len(uniq))] |
|
|
| with tempfile.TemporaryDirectory(dir='/tmp') as tmpdir: |
| tmp = Path(tmpdir) |
| with open(tmp / 'q.fa', 'w') as f: |
| for i, s in zip(qid, uniq['sequence']): |
| f.write('>%s\n%s\n' % (i, s)) |
| out = tmp / 'hits.tsv' |
| cmd = [str(HERE / 'bin' / 'mmseqs'), 'easy-search', str(tmp / 'q.fa'), str(HERE / 'db' / 'train.fa'), |
| str(out), str(tmp / 'work'), '-s', str(cfg['sens']), '--max-seqs', '100', '-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: |
| hits = pd.read_csv(out, sep='\t', names=['query', 'target', 'bits']) |
| else: |
| hits = pd.DataFrame(columns=['query', 'target', 'bits']) |
| sys.stderr.write('search done after %.0fs\n' % (time.perf_counter() - t0)) |
|
|
| hl = knn.hits_from_df(hits, cfg['kmax']) |
| with open(args.output, 'w') as f: |
| for q, i in zip(qid, uniq['id']): |
| ecs = knn.select_ecs(hl.get(q, []), ec_map, fallback, cfg['K'], cfg['pw'], cfg['ratio'], cfg['max_ec']) |
| f.write(json.dumps({'id': i, 'prediction': ecs}) + '\n') |
| sys.stderr.write('total %.0fs for %d unique inputs\n' % (time.perf_counter() - t0, len(uniq))) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
