File size: 4,056 Bytes
66db5a1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | """Offline CPU inference: k-mer prefilter (5-mers + 4-mers) -> Smith-Waterman re-ranking -> top-20 weighted EC vote."""
import argparse
import gc
import json
import os
import time
from collections import Counter, defaultdict
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import numpy as np
import pandas as pd
import pyopal
from inputs import prepare_inputs
from kmer import KmerIndex
ROOT = Path(__file__).parent
T0 = time.time()
BUDGET = float(os.environ.get('TIME_BUDGET', 480)) # seconds; hard limit is 600
N5 = 1000 # 5-mer candidates per query
N4 = 300 # extra 4-mer candidates per query (merged, duplicates removed)
TOP_K = 20 # neighbours that vote
EXPONENT = 16 # vote weight = (score / best score) ** EXPONENT
MAX_LEN_RATIO = 3 # skip 4-mer extras longer than 3x the query (slow to align, rarely homologs)
THREADS = 2
def log(msg):
print(f'[{time.time() - T0:6.1f}s] {msg}', flush=True)
def load_db():
path = ROOT / 'weights/db.parquet'
if not path.exists():
path = ROOT / 'db.parquet'
db = pd.read_parquet(path)
vocab = set(json.loads((ROOT / 'ec_vocabulary.json').read_text()))
ecs = [[e for e in row if e in vocab] for row in db.ecs]
keep = [i for i, e in enumerate(ecs) if e]
seqs = [db.sequence.iat[i] for i in keep]
ecs = [ecs[i] for i in keep]
fallback = Counter(e for row in ecs for e in row).most_common(1)[0][0]
return seqs, ecs, fallback
def vote(cand, scores, ecs, fallback):
if len(cand) == 0 or scores.max() <= 0:
return fallback
order = np.argsort(-scores)[:TOP_K]
best = scores[order[0]]
votes = defaultdict(float)
for i in order:
weight = (max(scores[i], 0) / best) ** EXPONENT
row = ecs[cand[i]]
for e in row:
votes[e] += weight / len(row)
return max(votes, key=votes.get)
def candidates(seqs, queries, k, n):
index = KmerIndex(seqs, k=k)
idx, _ = index.topk(queries, n)
del index
gc.collect()
log(f'{k}-mer candidates ready')
return idx
def predict_file(input_path, output_path):
frame = prepare_inputs(pd.read_parquet(input_path))
queries = [s if isinstance(s, str) else '' for s in frame.sequence]
log(f'{len(queries)} unique inputs')
seqs, ecs, fallback = load_db()
c5 = candidates(seqs, queries, 5, N5)
# Default prediction: 5-mer nearest neighbour (used if time runs short).
preds = [ecs[i[0]][0] if i[0] >= 0 else fallback for i in c5]
c4 = candidates(seqs, queries, 4, N4) if time.time() - T0 < BUDGET / 3 else None
aligner = pyopal.Aligner('BLOSUM62', gap_open=11, gap_extend=1)
database = pyopal.Database(seqs)
lengths = np.array([len(s) for s in seqs])
def align(q):
if time.time() - T0 > BUDGET or not queries[q]:
return q, None
pool = c5[q]
if c4 is not None:
extra = c4[q][c4[q] >= 0]
extra = extra[lengths[extra] <= MAX_LEN_RATIO * len(queries[q])]
pool = np.concatenate([pool, extra])
cand = [int(c) for c in dict.fromkeys(pool.tolist()) if c >= 0]
if not cand:
return q, None
scores = np.array([r.score for r in aligner.align(queries[q], database.extract(cand))], float)
return q, vote(cand, scores, ecs, fallback)
done = 0
with ThreadPoolExecutor(THREADS) as pool:
for q, p in pool.map(align, range(len(queries))):
if p is not None:
preds[q] = p
done += 1
log(f'aligned {done}/{len(queries)} queries')
with Path(output_path).open('w') as out:
for key, p in zip(frame.id, preds):
out.write(json.dumps({'id': key, 'prediction': [p]}) + '\n')
log(f'wrote {output_path}')
if __name__ == '__main__':
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--input', required=True)
parser.add_argument('--output', required=True)
args = parser.parse_args()
predict_file(args.input, args.output)
|