Blame for run_prothgt_ec_app.py

"""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