Blame for predict.py

"""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   # every file we write stays under 16 MiB
MAX_CHUNK_SEQS = 8000                # bounds DIAMOND's temporary seed-hit files too
OUTFMT = ['qseqid', 'sseqid', 'pident', 'length', 'qlen', 'slen', 'evalue', 'bitscore',
          'qstart', 'qend', 'sstart', 'send']
START = time.time()
# DIAMOND settings to try, most efficient first.
CONFIGS = [(['-c1', '-b4'], 2), (['-c4', '-b1'], 1)]   # (extra args, threads)


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:  # diagnostics must never break inference
            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   # seconds after start; DIAMOND is killed past this point


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')   # v14: larger reference
    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')
        # The second pass costs about as much as the first, so run it only with a wide margin.
        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}')
        # v16: optional pass over a second, larger reference (unreviewed TrEMBL, transcript-level evidence)
        # in DIAMOND's fast mode. It only adds hits; if time is short or it fails, it is skipped.
        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)

    # Scoring mode (weights/mode.json): "learned" = gradient-boosted re-ranker (default),
    # "vote" = plain bitscore^power vote over the top-k hits.
    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', {})   # e.g. down-weight unreviewed (TrEMBL) references
            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]
        # v17: learned re-ranker over the vote's top 10 candidates, used only when the best hit is below
        # rerank_max_pid identity (it was trained and validated on official-label hard proteins there).
        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}')
        # v24: exact-sequence label table (weights/srchash.npz, built by tools/build_source_hash.py). For an input with no
        # identical train / Swiss-Prot sequence, the table's label (source priority measured on the training labels:
        # BRENDA > CatPred-DB > CataPro > DLKcat > PDB/SIFTS > UniProtKB unreviewed > ESP) is placed first.
        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                      # identical train / Swiss-Prot sequence: keep its label
                    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)
    # v7: add CNN features and use the CNN-aware re-ranker when available and time allows.
    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)