"""EC number prediction backend for the ProtHGT Space. Unlike GO term prediction (one flat decoder per sub-ontology), EC prediction uses a single joint model with a shared HGT backbone (plus a normalized skip connection) and four level-specific decoders, run as a top-down cascade so the predicted EC levels form a consistent hierarchy. """ import os import gzip import json import shutil from functools import lru_cache import torch from torch.nn import Linear, LayerNorm from torch_geometric.nn import HGTConv, MLP import pandas as pd import yaml from huggingface_hub import hf_hub_download from run_prothgt_app import load_name_info EC_MODELS_DIR = os.path.join("data", "models", "ec") EC_CONFIG_PATH = os.path.join(EC_MODELS_DIR, "prothgt-ec-config.yaml") EC_MODEL_PATH = os.path.join(EC_MODELS_DIR, "prothgt-ec-model.pt") EC_KG_PATH = os.path.join("data", "prothgt-ec-kg.pt") # EC class/subclass/sub-subclass (L1-L3) names from the ExPASy ENZYME nomenclature; # L4 (full EC number) names come from name_info instead. EC_CLASS_NAMES_PATH = os.path.join("data", "ec_class_names.json") LEVELS = ['l1', 'l2', 'l3', 'l4'] LEVEL_TARGET_TYPE = {'l1': 'EC_class_1', 'l2': 'EC_class_2', 'l3': 'EC_class_3', 'l4': 'EC_number'} LEVEL_LABEL = {'l1': 'L1', 'l2': 'L2', 'l3': 'L3', 'l4': 'L4'} CHILDREN_KEY = {'l2': 'l1_children', 'l3': 'l2_children', 'l4': 'l3_children'} # EC taxonomy edges used to find each node's children for the cascade. CHILD_EDGE_TYPES = { 'l1_children': ('EC_class_1', 'rev_ec_parent_1', 'EC_class_2'), 'l2_children': ('EC_class_2', 'rev_ec_parent_2', 'EC_class_3'), 'l3_children': ('EC_class_3', 'rev_ec_parent_3', 'EC_number'), } @lru_cache(maxsize=1) def load_ec_class_names(path=EC_CLASS_NAMES_PATH): """{EC class prefix (L1-L3) -> name}, e.g. '2.7.11' -> 'Protein-serine/threonine kinases'.""" with open(path, encoding='utf-8') as f: return json.load(f) # --------------------------------------------------------------------------- # Model (ported from the ProtHGT src/ec package) # --------------------------------------------------------------------------- class ProtHGT(torch.nn.Module): """HGT backbone with a normalized pre/post-GNN skip connection.""" def __init__(self, data, hidden_channels, num_heads, num_layers, mlp_hidden_layers, mlp_dropout): super().__init__() self.lin_dict = torch.nn.ModuleDict() for node_type in data.node_types: self.lin_dict[node_type] = Linear(data[node_type].x.size(1), hidden_channels) self.convs = torch.nn.ModuleList() for _ in range(num_layers): self.convs.append(HGTConv(hidden_channels, hidden_channels, data.metadata(), num_heads, group='sum')) self.initial_norm = LayerNorm(hidden_channels) self.gnn_norm = LayerNorm(hidden_channels) self.mlp = MLP(mlp_hidden_layers, dropout=mlp_dropout, norm=None) def generate_embeddings(self, x_dict, edge_index_dict): initial_dict = {nt: self.lin_dict[nt](x).relu_() for nt, x in x_dict.items()} gnn_dict = initial_dict for conv in self.convs: out = conv(gnn_dict, edge_index_dict) gnn_dict = {nt: (out[nt] if out.get(nt) is not None else prev) for nt, prev in gnn_dict.items()} return {nt: torch.cat([self.initial_norm(initial_dict[nt]), self.gnn_norm(gnn_dict[nt])], dim=-1) for nt in initial_dict} class ProtHGTHier(torch.nn.Module): """Shared ProtHGT backbone + one MLP decoder per EC level.""" def __init__(self, data, hidden_channels, num_heads, num_layers, mlp_hidden_layers, mlp_dropout): super().__init__() self.backbone = ProtHGT(data, hidden_channels, num_heads, num_layers, mlp_hidden_layers, mlp_dropout) self.backbone.mlp = torch.nn.Identity() self.decoders = torch.nn.ModuleDict({ level: MLP(mlp_hidden_layers, dropout=mlp_dropout, norm=None) for level in LEVELS }) def generate_embeddings(self, x_dict, edge_index_dict): return self.backbone.generate_embeddings(x_dict, edge_index_dict) def load_hierarchy_masks(data): """{'l1_children': {ec_class_1_idx: [ec_class_2_idx, ...]}, ...} from the EC taxonomy edges.""" from collections import defaultdict masks = {} for name, et in CHILD_EDGE_TYPES.items(): row, col = data[et].edge_index children = defaultdict(list) for r, c in zip(row.tolist(), col.tolist()): children[r].append(c) masks[name] = dict(children) return masks # --------------------------------------------------------------------------- # Loading # --------------------------------------------------------------------------- def load_ec_kg(output=EC_KG_PATH): """Download (if needed) and load the EC knowledge graph. Callers should cache the result.""" if not os.path.exists(output): cached_path = hf_hub_download( repo_id="HUBioDataLab/ProtHGT", filename="knowledge_graphs/ec/prothgt-ec-kg.pt", repo_type="dataset", ) os.makedirs(os.path.dirname(output), exist_ok=True) shutil.copy(cached_path, output) print(f"EC KG downloaded to {output}") print('Loading EC knowledge graph...') return torch.load(output) def load_ec_model(heterodata, config_path=EC_CONFIG_PATH, model_path=EC_MODEL_PATH): """Build ProtHGTHier and load its weights. Callers should cache the result.""" with open(config_path) as f: cfg = yaml.safe_load(f) h = cfg['hidden_channels'] model = ProtHGTHier( heterodata, hidden_channels=h, num_heads=cfg['num_heads'], num_layers=cfg['num_layers'], mlp_hidden_layers=[4 * h, h, 1], mlp_dropout=cfg['mlp_dropout'], ) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model.load_state_dict(torch.load(model_path, map_location=device)) model.eval() print(f'Loaded EC model weights from {model_path}') return model # --------------------------------------------------------------------------- # Cascade prediction # --------------------------------------------------------------------------- @torch.no_grad() def _cascade_protein(model, x_dict, masks, protein_idx, inv_maps, threshold, device): """Top-down cascade for one protein: L1 over all classes, each deeper level only over the children of the kept nodes above. Returns (level_key, ec_string, prob) rows.""" rows = [] kept_prev = [] for level in LEVELS: target_type = LEVEL_TARGET_TYPE[level] if level == 'l1': candidates = list(range(x_dict[target_type].size(0))) else: ckey = CHILDREN_KEY[level] candidates = sorted({c for p in kept_prev for c in masks[ckey].get(p, [])}) if not candidates: break cand = torch.tensor(candidates, dtype=torch.long, device=device) p_emb = x_dict['Protein'][protein_idx].unsqueeze(0).expand(len(candidates), -1) z = torch.cat([p_emb, x_dict[target_type][cand]], dim=-1) probs = torch.sigmoid(model.decoders[level](z).view(-1)) keep = (probs > threshold).nonzero(as_tuple=False).view(-1) kept_prev = [candidates[i] for i in keep.tolist()] for i in keep.tolist(): rows.append((level, inv_maps[target_type][candidates[i]], float(probs[i]))) return rows def generate_ec_prediction_df(protein_ids, heterodata, model, threshold=0.5, progress_callback=None): """Cascade EC predictions for the given proteins. The backbone embeddings are computed once over the whole graph; the decoders then run per protein.""" if isinstance(protein_ids, str): protein_ids = [protein_ids] # The full-graph forward over the whole KG is memory-heavy for this model # (hidden=256); fall back to CPU if a GPU runs out of memory. device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') def _embed(dev): model.to(dev) hd = heterodata.to(dev) with torch.no_grad(): return hd, model.generate_embeddings(hd.x_dict, hd.edge_index_dict) try: heterodata, x_dict = _embed(device) except RuntimeError as e: if device.type == 'cuda' and 'out of memory' in str(e).lower(): print('GPU out of memory during the full-graph forward; falling back to CPU...') torch.cuda.empty_cache() device = torch.device('cpu') heterodata, x_dict = _embed(device) else: raise masks = load_hierarchy_masks(heterodata) inv_maps = { LEVEL_TARGET_TYPE[level]: {idx: tok for tok, idx in heterodata[LEVEL_TARGET_TYPE[level]]['id_mapping'].items()} for level in LEVELS } id_mapping = heterodata['Protein']['id_mapping'] name_info = load_name_info() ec_l4_names = name_info['EC_number'] # full EC number (L4) -> enzyme name ec_class_names = load_ec_class_names() # L1-L3 class prefix -> class name rows = [] for i, pid in enumerate(protein_ids): if progress_callback is not None: progress_callback(i, len(protein_ids)) protein_name = name_info['Protein'].get(pid, pid) for level, token, prob in _cascade_protein( model, x_dict, masks, id_mapping[pid], inv_maps, threshold, device): ec_name = ec_l4_names.get(token, '') if level == 'l4' else ec_class_names.get(token, '') rows.append({ 'UniProt_ID': pid, 'Protein': protein_name, 'EC': token, 'EC_name': ec_name, 'EC_level': LEVEL_LABEL[level], 'Probability': prob, }) df = pd.DataFrame(rows, columns=['UniProt_ID', 'Protein', 'EC', 'EC_name', 'EC_level', 'Probability']) return heterodata, df