| """Homology-graph neural network for EC prediction (offline, 2 CPUs). |
| |
| 1. Retrieval: k-mer prefilter (5-mers + 4-mers) finds candidate training enzymes; Smith-Waterman alignment |
| (BLOSUM62) scores them; the 20 best hits form the query's homology graph. |
| 2. Neural network: a graph attention network (gnn.py) reads the 20 neighbours (alignment features + EC agreement), |
| passes messages between them with two transformer layers and assigns each neighbour a weight. |
| The EC carrying the most weight is the prediction. |
| """ |
| import argparse |
| import gc |
| import json |
| import os |
| import time |
| from collections import Counter |
| from concurrent.futures import ThreadPoolExecutor |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
| import pyopal |
| import scoring_matrices |
|
|
| from gnn import K, HomologyGNN, features |
| from inputs import prepare_inputs |
| from kmer import KmerIndex |
|
|
| ROOT = Path(__file__).parent |
| T0 = time.time() |
| BUDGET = float(os.environ.get('TIME_BUDGET', 480)) |
| N5 = 1000 |
| N4 = 300 |
| MAX_LEN_RATIO = 3 |
| THREADS = 2 |
|
|
|
|
| def log(msg): |
| print(f'[{time.time() - T0:6.1f}s] {msg}', flush=True) |
|
|
|
|
| def find(name): |
| for p in (ROOT / 'weights' / name, ROOT / name): |
| if p.exists(): |
| return p |
| raise FileNotFoundError(name) |
|
|
|
|
| def load_db(): |
| db = pd.read_parquet(find('db.parquet')) |
| 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 self_scores(seqs): |
| m = scoring_matrices.ScoringMatrix.from_name('BLOSUM62') |
| diag = np.zeros(256) |
| for i, ch in enumerate(m.alphabet): |
| diag[ord(ch)] = m[ch][i] |
| return np.array([diag[np.frombuffer(s.encode('ascii', 'replace'), np.uint8)].sum() for s in seqs], np.float32) |
|
|
|
|
| def prefix_ids(ecs): |
| """(N,4) ids of the L1..L4 prefixes of each enzyme's first EC.""" |
| pid = {} |
| return np.array([[pid.setdefault('.'.join(row[0].split('.')[:L]) + f'#{L}', len(pid)) for L in (1, 2, 3, 4)] |
| for row in ecs], np.int64) |
|
|
|
|
| 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) |
| |
| 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 = np.array([int(c) for c in dict.fromkeys(pool.tolist()) if c >= 0]) |
| if len(cand) == 0: |
| return q, None |
| scores = np.array([r.score for r in aligner.align(queries[q], database.extract(cand.tolist()))], np.float32) |
| if scores.max() <= 0: |
| return q, None |
| o = np.argsort(-scores)[:K] |
| return q, (cand[o], scores[o]) |
|
|
| NI = np.full((len(queries), K), -1, np.int64) |
| NS = np.zeros((len(queries), K), np.float32) |
| with ThreadPoolExecutor(THREADS) as pool: |
| for q, res in pool.map(align, range(len(queries))): |
| if res is not None: |
| NI[q, :len(res[0])], NS[q, :len(res[1])] = res |
| aligned = np.where(NI[:, 0] >= 0)[0] |
| log(f'aligned {len(aligned)}/{len(queries)} queries') |
|
|
| |
| net = HomologyGNN(find('gnn.npz')) |
| dself, dpref = self_scores(seqs), prefix_ids(ecs) |
| qself = self_scores(queries) |
| qlen = np.array([max(len(s), 1) for s in queries], np.float32) |
| for s in range(0, len(aligned), 1024): |
| b = aligned[s:s + 1024] |
| ni, valid = NI[b].clip(0), NI[b] >= 0 |
| f, rel = features(NS[b], qself[b], dself[ni], qlen[b], lengths[ni].astype(np.float32), valid) |
| w = net.neighbour_weights(f, rel, dpref[ni], valid) |
| for r, q in enumerate(b): |
| mass = {} |
| for j in range(K): |
| if valid[r, j]: |
| row = ecs[NI[q, j]] |
| for e in row: |
| mass[e] = mass.get(e, 0) + w[r, j] / len(row) |
| preds[q] = max(mass, key=mass.get) |
| log(f'graph network predicted {len(aligned)} 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) |
|
|
