| """Offline CPU inference: DIAMOND homology search against train + UniProt references, bitscore vote + learned re-ranker (mode.json). |
| |
| 1. Search every unique input against the training set (diamond blastp |
| --very-sensitive, 50 hits). |
| 2. Re-search inputs with no hit or a weak best hit (<=30% identity) with |
| --ultra-sensitive when time allows. |
| 3. Collect candidate ECs from the hits, compute features and score them with a |
| gradient-boosted tree ensemble (pure numpy). Best EC first. |
| Safety: environment diagnostics are logged; DIAMOND is smoke-tested and retried |
| with safer settings; if it cannot run at all, a pure numpy/scipy 3-mer |
| nearest-neighbour fallback produces the predictions. |
| """ |
| import argparse |
| import gzip |
| import json |
| import os |
| import platform |
| import resource |
| import shutil |
| import signal |
| import subprocess |
| import sys |
| import tarfile |
| import tempfile |
| import threading |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
| from inputs import prepare_inputs |
| from features import query_candidates, parse_line |
| from ranker import TreeRanker |
|
|
| ROOT = Path(__file__).resolve().parent |
| TOP_K = 50 |
| WEAK_PID = 30.0 |
|
|
|
|
| def load_mode(): |
| cfg = {'scoring': 'learned', 'vote_k': 10, 'vote_power': 4, |
| 'weak_pid': WEAK_PID, 'weak_max': 3000, 'ultra_guard': 215} |
| try: |
| cfg.update(json.loads((Path(__file__).resolve().parent / 'weights' / 'mode.json').read_text())) |
| except Exception: |
| pass |
| return cfg |
| MAX_CHUNK_BYTES = 12 * 1024 * 1024 |
| MAX_CHUNK_SEQS = 8000 |
| OUTFMT = ['qseqid', 'sseqid', 'pident', 'length', 'qlen', 'slen', 'evalue', 'bitscore', |
| 'qstart', 'qend', 'sstart', 'send'] |
| START = time.time() |
| |
| CONFIGS = [(['-c1', '-b4'], 2), (['-c4', '-b1'], 1)] |
|
|
|
|
| def log(msg): |
| print(f'[{time.time() - START:7.1f}s] {msg}', file=sys.stderr, flush=True) |
|
|
|
|
| def diagnostics(): |
| def probe(fn): |
| try: |
| fn() |
| except Exception as exc: |
| log(f'diagnostics: {exc!r}') |
| probe(lambda: log(f'python {sys.version.split()[0]} | {platform.platform()} | machine={platform.machine()} | uid={os.getuid()}')) |
| probe(lambda: log(f'cpu_count={os.cpu_count()} affinity={len(os.sched_getaffinity(0))}')) |
| probe(lambda: log(' '.join(f'{n}={resource.getrlimit(getattr(resource, n))}' for n in |
| ('RLIMIT_FSIZE', 'RLIMIT_AS', 'RLIMIT_DATA', 'RLIMIT_STACK', 'RLIMIT_NPROC', 'RLIMIT_NOFILE')))) |
| def cpu(): |
| info = Path('/proc/cpuinfo').read_text() |
| model = next((l.split(':', 1)[1].strip() for l in info.splitlines() if l.startswith('model name')), '?') |
| flags = next((l.split(':', 1)[1].split() for l in info.splitlines() if l.startswith('flags')), []) |
| log(f'cpu={model} | ' + ' '.join(f'{f}={f in flags}' for f in ('sse4_1', 'avx2', 'avx512f'))) |
| probe(cpu) |
| probe(lambda: log(f'/tmp free={shutil.disk_usage("/tmp").free / 1e9:.2f} GB')) |
| for p in ('/sys/fs/cgroup/memory.max', '/sys/fs/cgroup/pids.max'): |
| probe(lambda p=p: log(f'{p}={Path(p).read_text().strip()}') if Path(p).exists() else None) |
|
|
|
|
| def cpu_probe(): |
| """Log real usable compute: matmul GFLOP/s with 1, 2, 4, 8, 16 threads (~5 s). Changes nothing.""" |
| try: |
| import numpy as _np |
| from concurrent.futures import ThreadPoolExecutor |
| a = _np.random.rand(768, 768).astype(_np.float32) |
| def work(_): |
| for _i in range(8): |
| a @ a |
| out = [] |
| for t in (1, 2, 4, 8, 16): |
| t0 = time.time() |
| with ThreadPoolExecutor(t) as ex: |
| list(ex.map(work, range(t))) |
| dt = time.time() - t0 |
| out.append(f'{t}t={t * 8 * 2 * 768 ** 3 / dt / 1e9:.0f}') |
| log('cpu probe GFLOP/s: ' + ' '.join(out)) |
| except Exception as exc: |
| log(f'cpu probe error: {exc!r}') |
|
|
|
|
| def diamond_binaries(): |
| """Official release binary (runtime SIMD dispatch) first, then our generic static build.""" |
| found = [] |
| official = ROOT / 'bin' / 'diamond' |
| if not (official.exists() and os.access(official, os.X_OK)): |
| try: |
| out = Path('/tmp/diamond_bin') |
| out.mkdir(parents=True, exist_ok=True) |
| official = out / 'diamond' |
| with tarfile.open(ROOT / 'bin' / 'diamond-linux64.tar.gz') as tar: |
| official.write_bytes(tar.extractfile('diamond').read()) |
| official.chmod(0o755) |
| except Exception as exc: |
| log(f'cannot extract official diamond: {exc!r}') |
| official = None |
| if official: |
| found.append(('official', str(official))) |
| generic = ROOT / 'bin' / 'diamond-generic' |
| if generic.exists(): |
| if not os.access(generic, os.X_OK): |
| tmp = Path('/tmp/diamond_bin/diamond-generic') |
| tmp.parent.mkdir(parents=True, exist_ok=True) |
| shutil.copyfile(generic, tmp) |
| tmp.chmod(0o755) |
| generic = tmp |
| found.append(('generic', str(generic))) |
| return found |
|
|
|
|
| def write_chunks(items, workdir, tag): |
| """items: iterable of (query index, sequence).""" |
| paths, buf, size = [], [], 0 |
| def flush(): |
| nonlocal buf, size |
| if buf: |
| p = Path(workdir) / f'{tag}{len(paths)}.fa' |
| p.write_text(''.join(buf)) |
| paths.append(p) |
| buf, size = [], 0 |
| for i, seq in items: |
| rec = f'>{i}\n{seq}\n' |
| if size + len(rec) > MAX_CHUNK_BYTES or len(buf) >= MAX_CHUNK_SEQS: |
| flush() |
| buf.append(rec) |
| size += len(rec) |
| flush() |
| return paths |
|
|
|
|
| def describe_rc(rc): |
| if rc < 0: |
| try: |
| return f'killed by signal {signal.Signals(-rc).name}' |
| except ValueError: |
| return f'killed by signal {-rc}' |
| return f'exit code {rc}' |
|
|
|
|
| DEADLINE = 420.0 |
|
|
|
|
| def diamond_bins(): |
| """'--bin N' spreads temporary seed hits over N files: more bins keep each file under the |
| 16 MiB size cap, but every bin is an open file and the sandbox allows only ~256 (RLIMIT_NOFILE).""" |
| try: |
| soft = resource.getrlimit(resource.RLIMIT_NOFILE)[0] |
| if soft == resource.RLIM_INFINITY: |
| soft = 4096 |
| except Exception: |
| soft = 256 |
| return int(max(16, min(128, (soft - 64) // 1))) |
|
|
|
|
| def run_diamond(chunks, mode, cfg, workdir, depth=0, bins=None, db='train.dmnd'): |
| """Search each FASTA chunk. The evaluator caps every written file at 16 MiB (RLIMIT_FSIZE); |
| '--bin 256' spreads DIAMOND's temporary seed hits over many small files, and a chunk that |
| still hits the cap (SIGXFSZ) is split into smaller chunks and retried.""" |
| (name, binary), extra, threads = cfg |
| bins = bins or diamond_bins() |
| hits = {} |
| for chunk in chunks: |
| errf = Path(workdir) / 'diamond.err' |
| cmd = [binary, 'blastp', '-q', str(chunk), '-d', str(ROOT / 'weights' / db), |
| '-k', str(TOP_K), '--threads', str(threads), '--quiet', *extra, '--bin', str(bins), |
| '--tmpdir', str(workdir), f'--{mode}', '--outfmt', '6', *OUTFMT] |
| local = {} |
| with errf.open('w') as err: |
| proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=err, text=True) |
| watchdog = threading.Timer(max(1.0, DEADLINE - (time.time() - START)), proc.kill) |
| watchdog.start() |
| for line in proc.stdout: |
| q, h = parse_line(line) |
| local.setdefault(int(q), []).append(h) |
| rc = proc.wait() |
| watchdog.cancel() |
| if time.time() - START >= DEADLINE - 1 and rc != 0: |
| raise RuntimeError(f'diamond[{name}] {mode} stopped by watchdog at {DEADLINE:.0f}s') |
| if rc == -signal.SIGXFSZ and depth < 3: |
| recs = chunk.read_text().split('>')[1:] |
| if len(recs) > 1: |
| log(f'diamond[{name}] hit the file-size cap on {len(recs)} queries; splitting into 4') |
| parts = [] |
| for k in range(4): |
| sub = recs[k::4] |
| if sub: |
| p = chunk.with_name(f'{chunk.stem}_{depth}_{k}.fa') |
| p.write_text(''.join('>' + r for r in sub)) |
| parts.append(p) |
| chunk.unlink() |
| hits.update(run_diamond(parts, mode, cfg, workdir, depth + 1, bins, db)) |
| continue |
| msg = errf.read_text()[-2000:].strip() if (rc != 0 and errf.exists()) else '' |
| if rc != 0 and 'Too many open files' in msg and bins > 16: |
| log(f'diamond[{name}] too many open files with --bin {bins}; retrying with {bins // 2}') |
| hits.update(run_diamond([chunk], mode, cfg, workdir, depth, bins // 2, db)) |
| continue |
| if rc != 0: |
| raise RuntimeError(f'diamond[{name}] {mode} {extra} t={threads}: {describe_rc(rc)}; stderr: {msg or "(empty)"}') |
| hits.update(local) |
| chunk.unlink() |
| return hits |
|
|
|
|
| class EcReliability: |
| def __init__(self, path): |
| with gzip.open(path, 'rt') as f: |
| t = json.load(f) |
| self.t, self.A = t, t['A'] |
|
|
| def _enc(self, table, key, prior): |
| s, c = self.t[table].get(key, (0.0, 0)) |
| return (s + self.A * prior) / (c + self.A), np.log1p(c) |
|
|
| def matrix(self, cands): |
| rows = [] |
| for e in cands: |
| a, an = self._enc('ec', e, self.t['prior']) |
| b, bn = self._enc('ec_top', e, self.t['prior_top']) |
| c, cn = self._enc('l3', e.rsplit('.', 1)[0], self.t['prior']) |
| rows.append([a, an, b, bn, c, cn]) |
| return np.array(rows, dtype=np.float64) |
|
|
|
|
| def homology_predictions(seqs, labels, freq): |
| """Returns {query index: ranked EC list}, or None if DIAMOND cannot run.""" |
| n = len(seqs) |
| main_mode = 'very-sensitive' if n <= 14000 else ('sensitive' if n <= 40000 else 'fast') |
| workdir = tempfile.mkdtemp(dir='/tmp') |
| try: |
| binaries = diamond_binaries() |
| for name, b in binaries: |
| try: |
| log(f'{name}: ' + subprocess.run([b, 'version'], capture_output=True, text=True, timeout=60).stdout.strip()) |
| except Exception as exc: |
| log(f'{name}: version check failed: {exc!r}') |
| sample = [(i, seqs[i]) for i in range(min(20, len(seqs)))] |
| hits, cfg = None, None |
| for b in binaries: |
| for extra, threads in CONFIGS: |
| c = (b, extra, threads) |
| try: |
| run_diamond(write_chunks(sample, workdir, 's'), 'fast', c, workdir) |
| log(f'smoke test ok: {b[0]} {extra} t={threads}') |
| except Exception as exc: |
| log(f'smoke test failed: {exc}') |
| continue |
| t0 = time.time() |
| mode = main_mode if b[0] == 'official' or main_mode == 'fast' else 'sensitive' |
| try: |
| hits = run_diamond(write_chunks(enumerate(seqs), workdir, 'a'), mode, c, workdir) |
| cfg = c |
| break |
| except Exception as exc: |
| log(f'main search failed after {time.time() - t0:.0f}s: {exc}') |
| if time.time() - START > 240: |
| break |
| if hits is not None or time.time() - START > 240: |
| break |
| if hits is None: |
| return None |
| t1 = time.time() - START |
| mcfg = load_mode() |
| weak = [i for i in range(n) if i not in hits or hits[i][0][1] <= float(mcfg['weak_pid'])] |
| log(f'main search ({mode}, {cfg[0][0]}) done: {len(hits)}/{n} with hits, {len(weak)} weak') |
| |
| if weak and len(weak) <= int(mcfg['weak_max']) and t1 < float(mcfg['ultra_guard']) and cfg[0][0] == 'official': |
| try: |
| extra = run_diamond(write_chunks(((i, seqs[i]) for i in weak), workdir, 'b'), |
| 'ultra-sensitive', cfg, workdir) |
| hits.update(extra) |
| log(f'ultra-sensitive pass done: {len(extra)} re-scored') |
| except Exception as exc: |
| log(f'ultra-sensitive pass skipped: {exc}') |
| |
| |
| xdb = mcfg.get('extra_db') |
| t2 = time.time() - START |
| if xdb and (ROOT / 'weights' / xdb).exists() and t2 < float(mcfg.get('extra_guard', 330)): |
| try: |
| extra = run_diamond(write_chunks(enumerate(seqs), workdir, 'x'), mcfg.get('extra_mode', 'fast'), |
| cfg, workdir, db=xdb) |
| for i, h in extra.items(): |
| hits.setdefault(i, []).extend(h) |
| log(f'extra reference pass ({xdb}) done: {len(extra)} inputs with hits, {time.time() - START - t2:.0f}s') |
| except Exception as exc: |
| log(f'extra reference pass skipped: {exc}') |
| elif xdb: |
| log(f'extra reference pass skipped (elapsed {t2:.0f}s)') |
| finally: |
| shutil.rmtree(workdir, ignore_errors=True) |
|
|
| |
| |
| mode_cfg = load_mode() |
| log(f'scoring mode: {mode_cfg}') |
| if mode_cfg['scoring'] == 'vote': |
| preds = {} |
| for i in sorted(hits): |
| h = sorted(hits[i], key=lambda x: -x[5])[:int(mode_cfg['vote_k'])] |
| votes = {} |
| srcw = mode_cfg.get('source_weights', {}) |
| for s, pid, ql, sl, ev, b, *_ in h: |
| ecs = labels[s] |
| w = (b / h[0][5]) ** float(mode_cfg['vote_power']) * srcw.get(s[:3], 1.0) |
| for e in ecs: |
| votes[e] = votes.get(e, 0.0) + w / len(ecs) |
| if votes: |
| preds[i] = sorted(votes, key=lambda e: (-votes[e], -freq.get(e, 0), e))[:10] |
| |
| |
| rr_path = ROOT / 'weights' / 'ranker_v2.npz' |
| if mode_cfg.get('reranker') and rr_path.exists(): |
| try: |
| from rerank2 import candidate_features |
| rr = TreeRanker(rr_path) |
| srcw = mode_cfg.get('source_weights', {}) |
| qs, blocks, cand_of = [], [], {} |
| for i in sorted(hits): |
| if not hits[i] or max(x[1] for x in hits[i]) >= float(mode_cfg.get('rerank_max_pid', 80)): |
| continue |
| cands, X = candidate_features(hits[i], labels, freq, srcw, int(mode_cfg['vote_k']), |
| float(mode_cfg['vote_power'])) |
| if len(cands) > 1: |
| qs.append(i) |
| cand_of[i] = cands |
| blocks.append(np.asarray(X, dtype=np.float64)) |
| if blocks: |
| scores = rr.score(np.vstack(blocks)) |
| pos = 0 |
| for i in qs: |
| c = cand_of[i] |
| sc = scores[pos:pos + len(c)] |
| pos += len(c) |
| preds[i] = [c[j] for j in np.argsort(-sc, kind='stable')] |
| log(f're-ranker v2 applied to {len(qs)} inputs') |
| except Exception as exc: |
| log(f're-ranker v2 skipped, using the vote: {exc!r}') |
| |
| |
| |
| hp = ROOT / 'weights' / str(mode_cfg.get('label_table', 'srchash.npz')) |
| if mode_cfg.get('label_table') and hp.exists(): |
| try: |
| import hashlib |
| z = np.load(hp) |
| H, S = z['h'], z['s'] |
| sets = json.loads(hp.with_name(hp.stem + '_sets.json').read_text()) |
| used = 0 |
| for i, sq in enumerate(seqs): |
| if any(x[1] >= 100.0 and x[2] == x[3] and x[7] - x[6] + 1 == x[2] |
| and not x[0].startswith(('un_', 'u2_', 'cp_', 'br_', 'dk_', 'es_', 'e2_', 'bo_', 'pd_')) |
| for x in hits.get(i, [])): |
| continue |
| key = np.uint64(int(hashlib.md5(sq.encode()).hexdigest()[:16], 16)) |
| j = int(np.searchsorted(H, key)) |
| if j < len(H) and H[j] == key: |
| ecs = sets[int(S[j])] |
| rank = {e: r for r, e in enumerate(preds.get(i, []))} |
| ordered = sorted(ecs, key=lambda e: (rank.get(e, 99), -freq.get(e, 0), e)) |
| preds[i] = (ordered + [e for e in preds.get(i, []) if e not in ecs])[:10] |
| used += 1 |
| log(f'exact-sequence label table: {len(H)} sequences, used for {used} inputs') |
| except Exception as exc: |
| log(f'exact-sequence label table skipped: {exc!r}') |
| return preds |
| ranker = TreeRanker(ROOT / 'weights' / 'ranker.npz') |
| rel = EcReliability(ROOT / 'weights' / 'te_tables.json.gz') |
| order, cand_of, blocks = sorted(hits), {}, [] |
| for i in order: |
| cands, X = query_candidates(sorted(hits[i], key=lambda x: -x[5]), labels, freq) |
| if cands: |
| cand_of[i] = cands |
| blocks.append(np.hstack([X.astype(np.float64), rel.matrix(cands)])) |
| scores = ranker.score(np.vstack(blocks)) if blocks else np.zeros(0) |
| |
| cnn_w, cnn_r = ROOT / 'weights' / 'cnn.npz', ROOT / 'weights' / 'ranker_cnn.npz' |
| if blocks and cnn_w.exists() and cnn_r.exists() and time.time() - START < 450: |
| try: |
| from cnn_numpy import SmallCNN |
| from cnn_feats import CnnScorer |
| t0 = time.time() |
| net = SmallCNN(cnn_w) |
| scorer = CnnScorer(net.classes) |
| qs = [i for i in order if i in cand_of] |
| L = net.logits([seqs[i] for i in qs]) |
| blocks2 = [np.hstack([blk, scorer.rows(L[k], cand_of[i])]) for k, (i, blk) in enumerate(zip(qs, blocks))] |
| scores = TreeRanker(cnn_r).score(np.vstack(blocks2)) |
| log(f'CNN re-ranker used ({time.time() - t0:.0f}s)') |
| except Exception as exc: |
| log(f'CNN step skipped, using v5 re-ranker: {exc!r}') |
| preds, pos = {}, 0 |
| for i in order: |
| if i not in cand_of: |
| continue |
| cands = cand_of[i] |
| s = scores[pos:pos + len(cands)] |
| pos += len(cands) |
| preds[i] = [cands[j] for j in np.argsort(-s, kind='stable')] |
| return preds |
|
|
|
|
| def kmer_predictions(seqs, labels, freq): |
| """Fallback without DIAMOND: 3-mer TF-IDF nearest neighbours, bitscore-free vote.""" |
| from kmer_knn import KmerIndex |
| ids, train_seqs = [], [] |
| for part in sorted((ROOT / 'weights').glob('train_seqs_*.tsv.gz')): |
| with gzip.open(part, 'rt') as f: |
| for line in f: |
| i, s = line.rstrip('\n').split('\t') |
| ids.append(i) |
| train_seqs.append(s) |
| log(f'k-mer fallback: {len(train_seqs)} training sequences') |
| idx = KmerIndex(train_seqs) |
| del idx.XT, train_seqs |
| I, S = idx.query_fast(seqs) |
| preds = {} |
| for q, (row, sc) in enumerate(zip(I, S)): |
| votes = {} |
| for j, s in zip(row, sc): |
| ecs = labels[ids[j]] |
| for e in ecs: |
| votes[e] = votes.get(e, 0.0) + max(float(s), 0.0) ** 4 / len(ecs) |
| if votes: |
| preds[q] = sorted(votes, key=lambda e: (-votes[e], -freq.get(e, 0), e))[:10] |
| return preds |
|
|
|
|
| def predict_file(input_path, output_path): |
| diagnostics() |
| cpu_probe() |
| frame = prepare_inputs(pd.read_parquet(input_path)) |
| with gzip.open(ROOT / 'weights' / 'labels.json.gz', 'rt') as f: |
| model = json.load(f) |
| labels, freq = model['labels'], model['freq'] |
| default = sorted(freq, key=lambda e: (-freq[e], e))[0] |
| seqs = list(frame.sequence) |
| log(f'{len(seqs)} unique inputs') |
|
|
| preds, method = None, 'diamond' |
| try: |
| preds = homology_predictions(seqs, labels, freq) |
| except Exception as exc: |
| log(f'homology pipeline error: {exc!r}') |
| if preds is None: |
| method = 'kmer-fallback' |
| log('DIAMOND unavailable -> k-mer fallback') |
| try: |
| preds = kmer_predictions(seqs, labels, freq) |
| except Exception as exc: |
| log(f'k-mer fallback error: {exc!r}') |
| preds, method = {}, 'frequency-only' |
|
|
| with Path(output_path).open('w') as out: |
| for i, key in enumerate(frame.id): |
| out.write(json.dumps({'id': key, 'prediction': preds.get(i) or [default]}) + '\n') |
| log(f'done: method={method}, {len(preds)}/{len(seqs)} predicted from neighbours') |
|
|
|
|
| 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) |
|
|
