{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":41875,"databundleVersionId":5521661,"isSourceIdPinned":false}],"dockerImageVersionId":31329,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-05-02T16:05:53.751148Z","iopub.execute_input":"2026-05-02T16:05:53.751533Z","iopub.status.idle":"2026-05-02T16:05:54.863402Z","shell.execute_reply.started":"2026-05-02T16:05:53.751505Z","shell.execute_reply":"2026-05-02T16:05:54.862558Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# ── Install all required packages ──────────────────────────────────\nimport subprocess, sys\n\ndef pip(pkg):\n    subprocess.run([sys.executable, '-m', 'pip', 'install', '-q', pkg], check=True)\n\npip('fair-esm')           # ESM-2 protein language model\npip('biopython')          # FASTA parsing\npip('goatools')           # GO OBO parser\npip('faiss-cpu')          # fast kNN (use faiss-cpu if no GPU)\npip('torch-scatter')      # PyG dependency\npip('torch-sparse')       # PyG dependency\n\n# PyTorch Geometric (match your torch/cuda version on Kaggle)\nimport torch\nTORCH = torch.__version__.split('+')[0]\nCUDA  = 'cu' + torch.version.cuda.replace('.', '') if torch.cuda.is_available() else 'cpu'\nprint(f'torch={TORCH}  cuda={CUDA}')\n\nsubprocess.run([\n    sys.executable, '-m', 'pip', 'install', '-q',\n    'torch-geometric',\n    '-f', f'https://data.pyg.org/whl/torch-{TORCH}+{CUDA}.html'\n], check=True)\n\npip('networkx')           # graph visualization\npip('matplotlib')         # plotting\npip('seaborn')            # heatmaps\npip('plotly')             # interactive visualization\npip('scikit-learn')       # MultiLabelBinarizer, metrics\n\nprint('✅ All packages installed')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-02T16:08:11.185011Z","iopub.execute_input":"2026-05-02T16:08:11.185408Z","iopub.status.idle":"2026-05-02T16:38:50.705143Z","shell.execute_reply.started":"2026-05-02T16:08:11.185377Z","shell.execute_reply":"2026-05-02T16:38:50.704273Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os, re, math, warnings\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.optim import AdamW\nfrom torch.optim.lr_scheduler import CosineAnnealingLR\nfrom torch_geometric.data import Data\nfrom torch_geometric.loader import NeighborLoader\nfrom torch_geometric.nn import SAGEConv, JumpingKnowledge\nfrom torch_geometric.utils import to_networkx\nfrom sklearn.preprocessing import MultiLabelBinarizer\nimport faiss\nfrom Bio import SeqIO\nfrom goatools.obo_parser import GODag\nimport networkx as nx\nimport matplotlib.pyplot as plt\nimport matplotlib.patches as mpatches\nimport seaborn as sns\nimport plotly.graph_objects as go\nimport plotly.express as px\nfrom plotly.subplots import make_subplots\nfrom collections import Counter\nwarnings.filterwarnings('ignore')\n\n# ── Config ────────────────────────────────────────────────────────\nCFG = dict(\n    # Kaggle dataset paths  →  /kaggle/input/<dataset-name>/\n    train_fasta  = '/kaggle/input/competitions/cafa-5-protein-function-prediction/Train/train_sequences.fasta',\n    train_terms  = '/kaggle/input/competitions/cafa-5-protein-function-prediction/Train/train_terms.tsv',\n    test_fasta   = '/kaggle/input/competitions/cafa-5-protein-function-prediction/Test (Targets)/testsuperset.fasta',\n    obo_path     = '/kaggle/input/competitions/cafa-5-protein-function-prediction/Train/go-basic.obo',\n    ia_path      = '/kaggle/input/competitions/cafa-5-protein-function-prediction/IA.txt',\n    emb_cache    = '/kaggle/working/node_features.npy',\n    test_emb_cache = '/kaggle/working/test_features.npy',\n    ontology     = 'BPO',       # Biological Process\n    hidden_dim   = 512,\n    num_layers   = 3,\n    dropout      = 0.3,\n    lr           = 3e-4,\n    weight_decay = 1e-4,\n    batch_size   = 512,\n    epochs       = 60,\n    patience     = 10,\n    max_grad_norm= 1.0,\n    hier_lambda  = 0.5,\n    graph_k      = 20,\n    graph_thresh = 0.80,\n    val_frac     = 0.1,\n    seed         = 42,\n    device       = 'cuda' if torch.cuda.is_available() else 'cpu',\n    # How many proteins to use for fast demo (set None for full run)\n    demo_limit   = 5000,\n)\n\ntorch.manual_seed(CFG['seed'])\nnp.random.seed(CFG['seed'])\nDEVICE = torch.device(CFG['device'])\nprint(f'Device: {DEVICE}')\nif torch.cuda.is_available():\n    print(f'GPU: {torch.cuda.get_device_name(0)}')\n    print(f'VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-02T16:48:18.586802Z","iopub.execute_input":"2026-05-02T16:48:18.587416Z","iopub.status.idle":"2026-05-02T16:48:18.599735Z","shell.execute_reply.started":"2026-05-02T16:48:18.587386Z","shell.execute_reply":"2026-05-02T16:48:18.599086Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ─────────────────────────────────────────────────────────────────\n# 3a. Read FASTA sequences\n# ─────────────────────────────────────────────────────────────────\ndef load_fasta(path):\n    \"\"\"Returns {protein_id: sequence} dict.\"\"\"\n    return {rec.id: str(rec.seq) for rec in SeqIO.parse(path, 'fasta')}\n\nprint('Reading train sequences...')\ntrain_seqs = load_fasta(CFG['train_fasta'])\nprint(f'  → {len(train_seqs):,} train proteins')\n\nprint('Reading test sequences...')\ntest_seqs = load_fasta(CFG['test_fasta'])\nprint(f'  → {len(test_seqs):,} test proteins')\n\n# Demo mode: subsample for fast Kaggle run\nif CFG['demo_limit']:\n    prot_ids = list(train_seqs.keys())[:CFG['demo_limit']]\n    train_seqs = {k: train_seqs[k] for k in prot_ids}\n    print(f'  [Demo mode] Using {len(train_seqs):,} proteins')\nelse:\n    prot_ids = list(train_seqs.keys())\n\n# ─────────────────────────────────────────────────────────────────\n# 3b. Read GO term labels\n# ─────────────────────────────────────────────────────────────────\nprint('\\nReading train_terms.tsv...')\ndf_terms = pd.read_csv(CFG['train_terms'], sep='\\t', header=None,\n                       names=['protein_id', 'term', 'aspect'])\nprint(f'  → {len(df_terms):,} rows')\nprint(f'  Aspects: {df_terms[\"aspect\"].value_counts().to_dict()}')\n\n# Filter to BPO (Biological Process)\ndf_bp = df_terms[df_terms['aspect'] == CFG['ontology']]\ngo_label_map = df_bp.groupby('protein_id')['term'].apply(list).to_dict()\nprint(f'  → {len(go_label_map):,} proteins have BP annotations')\n\n# Show distribution of label counts\nlabel_counts = [len(v) for v in go_label_map.values()]\nprint(f'  Labels per protein: mean={np.mean(label_counts):.1f}, '\n      f'median={np.median(label_counts):.0f}, max={max(label_counts)}')\n\n# ─────────────────────────────────────────────────────────────────\n# 3c. Read GO OBO hierarchy\n# ─────────────────────────────────────────────────────────────────\nprint('\\nParsing go-basic.obo...')\ngo_dag = GODag(CFG['obo_path'], optional_attrs=['relationship'])\nbp_terms = [t for t, node in go_dag.items()\n            if hasattr(node, 'namespace') and node.namespace == 'biological_process']\nprint(f'  → {len(bp_terms):,} BP terms in GO hierarchy')\n\n# ─────────────────────────────────────────────────────────────────\n# 3d. Read IA weights\n# ─────────────────────────────────────────────────────────────────\nprint('\\nReading IA.txt...')\nia_map = {}\nwith open(CFG['ia_path']) as f:\n    for line in f:\n        parts = line.strip().split('\\t')\n        if len(parts) == 2:\n            try:\n                ia_map[parts[0]] = float(parts[1])\n            except ValueError:\n                pass\nprint(f'  → {len(ia_map):,} terms with IA weights')\nprint(f'  IA range: [{min(ia_map.values()):.3f}, {max(ia_map.values()):.3f}]')\n\nprint('\\n✅ All files loaded')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-02T16:48:20.200036Z","iopub.execute_input":"2026-05-02T16:48:20.200350Z","iopub.status.idle":"2026-05-02T16:48:27.879646Z","shell.execute_reply.started":"2026-05-02T16:48:20.200324Z","shell.execute_reply":"2026-05-02T16:48:27.878761Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fig, axes = plt.subplots(2, 3, figsize=(18, 10))\nfig.suptitle('CAFA 5 — Data Exploration', fontsize=16, fontweight='bold')\n\n# ── Plot 1: Sequence length distribution ──\nseq_lens = [len(s) for s in train_seqs.values()]\naxes[0,0].hist(seq_lens, bins=60, color='#4A90D9', edgecolor='white', linewidth=0.4)\naxes[0,0].axvline(np.median(seq_lens), color='#E55', linewidth=2,\n                  label=f'Median={int(np.median(seq_lens))}')\naxes[0,0].set_xlabel('Sequence length (aa)')\naxes[0,0].set_ylabel('Count')\naxes[0,0].set_title('Protein sequence lengths')\naxes[0,0].legend()\n\n# ── Plot 2: GO annotations per protein ──\nann_counts = [len(go_label_map.get(p, [])) for p in prot_ids]\naxes[0,1].hist([c for c in ann_counts if c > 0], bins=50,\n               color='#6BCB77', edgecolor='white', linewidth=0.4)\naxes[0,1].axvline(np.median([c for c in ann_counts if c > 0]),\n                  color='#E55', linewidth=2)\naxes[0,1].set_xlabel('GO terms per protein')\naxes[0,1].set_ylabel('Count')\naxes[0,1].set_title('GO annotations per protein (BP)')\n\n# ── Plot 3: Annotation aspect distribution ──\naspect_counts = df_terms['aspect'].value_counts()\ncolors = ['#4A90D9', '#6BCB77', '#FF6B6B']\naxes[0,2].bar(aspect_counts.index, aspect_counts.values, color=colors)\nfor i, (asp, cnt) in enumerate(aspect_counts.items()):\n    axes[0,2].text(i, cnt + 50, f'{cnt:,}', ha='center', fontsize=10)\naxes[0,2].set_xlabel('Ontology aspect')\naxes[0,2].set_ylabel('Annotation count')\naxes[0,2].set_title('Annotations by ontology')\n\n# ── Plot 4: Top 20 most frequent GO terms ──\nterm_freq = Counter(t for terms in go_label_map.values() for t in terms)\ntop20 = term_freq.most_common(20)\nterms_plot, freqs_plot = zip(*top20)\naxes[1,0].barh(range(20), freqs_plot, color='#9B59B6')\naxes[1,0].set_yticks(range(20))\naxes[1,0].set_yticklabels([f'{t}' for t in terms_plot], fontsize=7)\naxes[1,0].set_xlabel('Frequency')\naxes[1,0].set_title('Top 20 GO terms (BP)')\naxes[1,0].invert_yaxis()\n\n# ── Plot 5: IA weight distribution ──\nia_vals = list(ia_map.values())\naxes[1,1].hist(ia_vals, bins=60, color='#FF6B6B', edgecolor='white', linewidth=0.4)\naxes[1,1].set_xlabel('Information Accretion (IA) weight')\naxes[1,1].set_ylabel('Count')\naxes[1,1].set_title('IA weight distribution')\naxes[1,1].axvline(np.mean(ia_vals), color='#333', linewidth=2,\n                  label=f'Mean={np.mean(ia_vals):.2f}')\naxes[1,1].legend()\n\n# ── Plot 6: GO hierarchy depth distribution ──\ndepths = [go_dag[t].depth for t in bp_terms if t in go_dag]\naxes[1,2].hist(depths, bins=30, color='#F39C12', edgecolor='white', linewidth=0.4)\naxes[1,2].set_xlabel('Depth in GO hierarchy')\naxes[1,2].set_ylabel('Count')\naxes[1,2].set_title('GO term depth distribution')\n\nplt.tight_layout()\nplt.savefig('/kaggle/working/eda_plots.png', dpi=150, bbox_inches='tight')\nplt.show()\nprint('✅ EDA plots saved')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-02T16:48:54.295449Z","iopub.execute_input":"2026-05-02T16:48:54.296105Z","iopub.status.idle":"2026-05-02T16:48:57.465687Z","shell.execute_reply.started":"2026-05-02T16:48:54.296076Z","shell.execute_reply":"2026-05-02T16:48:57.464789Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ── Multi-label binarizer ─────────────────────────────────────────\nprint('Building label matrix...')\nmlb = MultiLabelBinarizer()\nlabel_list = [go_label_map.get(pid, []) for pid in prot_ids]\nY = mlb.fit_transform(label_list)   # [N_train, C]\ngo_terms = list(mlb.classes_)\nn_go = len(go_terms)\nprint(f'  Label matrix shape: {Y.shape}  (proteins × GO terms)')\nprint(f'  Positive rate: {Y.mean():.4f}')\nprint(f'  GO terms in matrix: {n_go:,}')\n\n# ── Build GO hierarchy index for loss ────────────────────────────\nprint('\\nBuilding GO hierarchy index...')\nterm2idx = {t: i for i, t in enumerate(go_terms)}\nparent_idx_list, child_idx_list = [], []\n\nfor term in go_terms:\n    if term not in go_dag:\n        continue\n    node = go_dag[term]\n    parents = set(node.parents)\n    if hasattr(node, 'relationship'):\n        parents |= set(node.relationship.get('part_of', []))\n    for parent in parents:\n        pid = parent.item_id\n        if pid in term2idx:\n            child_idx_list.append(term2idx[term])\n            parent_idx_list.append(term2idx[pid])\n\nprint(f'  Parent-child pairs: {len(parent_idx_list):,}')\n\n# ── IA weight tensor ──────────────────────────────────────────────\nia_weights = torch.tensor(\n    [ia_map.get(t, 1.0) for t in go_terms], dtype=torch.float32\n)\nprint(f'  IA weights shape: {ia_weights.shape}')\nprint('✅ Labels & hierarchy ready')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-02T17:05:02.637430Z","iopub.execute_input":"2026-05-02T17:05:02.637767Z","iopub.status.idle":"2026-05-02T17:05:02.973111Z","shell.execute_reply.started":"2026-05-02T17:05:02.637741Z","shell.execute_reply":"2026-05-02T17:05:02.972268Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import esm\n\nclass AttentionPooling(nn.Module):\n    \"\"\"\n    Learned weighted pooling over residue tokens.\n    score_i = w^T · tanh(W·h_i)   →   alpha = softmax   →   output = Σ alpha_i * h_i\n    \"\"\"\n    def __init__(self, embed_dim=1280):\n        super().__init__()\n        self.proj  = nn.Linear(embed_dim, 256, bias=False)\n        self.score = nn.Linear(256, 1, bias=False)\n\n    def forward(self, x):          # x: [B, L, D]\n        a = torch.tanh(self.proj(x))               # [B, L, 256]\n        a = self.score(a).squeeze(-1)              # [B, L]\n        a = torch.softmax(a, dim=-1)               # [B, L]\n        return (a.unsqueeze(-1) * x).sum(dim=1)    # [B, D]\n\n\nclass ESM2Embedder:\n    MAX_LEN = 1022\n    OVERLAP  = 64\n\n    def __init__(self, device):\n        self.device = device\n        print('Loading ESM-2 (650M)...')\n        self.model, self.alphabet = esm.pretrained.esm2_t33_650M_UR50D()\n        self.model = self.model.eval().to(device)\n        self.batch_converter = self.alphabet.get_batch_converter()\n        self.attn_pool = AttentionPooling(1280).to(device)\n        print('  ESM-2 loaded ✅')\n\n    @torch.no_grad()\n    def embed(self, sequences: dict, batch_size=8):\n        ids  = list(sequences.keys())\n        seqs = list(sequences.values())\n        all_embs = []\n        for i in range(0, len(ids), batch_size):\n            batch_embs = [self._embed_one(s) for s in seqs[i:i+batch_size]]\n            all_embs.append(torch.stack(batch_embs))\n            if (i // batch_size) % 20 == 0:\n                print(f'  [{i}/{len(ids)}] embedded...', end='\\r')\n        print(f'  [{len(ids)}/{len(ids)}] done       ')\n        return torch.cat(all_embs, dim=0).cpu().numpy()\n\n    def _embed_one(self, seq):\n        if len(seq) <= self.MAX_LEN:\n            return self._forward_chunk(seq)\n        step = self.MAX_LEN - self.OVERLAP\n        chunks, start = [], 0\n        while start < len(seq):\n            end = min(start + self.MAX_LEN, len(seq))\n            chunks.append(self._forward_chunk(seq[start:end]))\n            if end == len(seq): break\n            start += step\n        return torch.stack(chunks).mean(0)\n\n    def _forward_chunk(self, seq):\n        _, _, tokens = self.batch_converter([('p', seq)])\n        tokens = tokens.to(self.device)\n        out = self.model(tokens, repr_layers=[33], return_contacts=False)\n        reps = out['representations'][33][0]     # [L+2, 1280]\n        cls_emb = reps[0]\n        eos_emb = reps[-1]\n        mid     = reps[1:-1]                     # [L, 1280]\n        pooled  = self.attn_pool(mid.unsqueeze(0)).squeeze(0)\n        return ((cls_emb + eos_emb + pooled) / 3.0).float()\n\nprint('ESM-2 embedder class defined ✅')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-02T17:05:35.190655Z","iopub.execute_input":"2026-05-02T17:05:35.191400Z","iopub.status.idle":"2026-05-02T17:05:35.210560Z","shell.execute_reply.started":"2026-05-02T17:05:35.191371Z","shell.execute_reply":"2026-05-02T17:05:35.209735Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def build_protein_graph(embeddings, k=20, threshold=0.80):\n    \"\"\"\n    L2-normalize → FAISS inner product search (= cosine similarity)\n    → keep edges with cosine ≥ threshold\n    Returns: edge_index [2, E], edge_weight [E]\n    \"\"\"\n    n, d = embeddings.shape\n    norms  = np.linalg.norm(embeddings, axis=1, keepdims=True) + 1e-12\n    normed = (embeddings / norms).astype(np.float32)\n\n    # FAISS index: inner product over L2-normed = cosine similarity\n    index = faiss.IndexFlatIP(d)\n    # Use GPU index if available\n    if torch.cuda.is_available():\n        res = faiss.StandardGpuResources()\n        index = faiss.index_cpu_to_gpu(res, 0, index)\n    index.add(normed)\n\n    print(f'  Searching {k} neighbors for {n:,} proteins...')\n    sims, nbrs = index.search(normed, k + 1)   # k+1: first hit is self\n\n    src, dst, wts = [], [], []\n    for i in range(n):\n        for j, sim in zip(nbrs[i], sims[i]):\n            if j == i or sim < threshold:\n                continue\n            src.append(i); dst.append(j); wts.append(float(sim))\n\n    edge_index  = torch.tensor([src, dst], dtype=torch.long)\n    edge_weight = torch.tensor(wts, dtype=torch.float32)\n    return edge_index, edge_weight\n\n\nprint('Building protein similarity graph...')\nedge_index, edge_weight = build_protein_graph(\n    embeddings, k=CFG['graph_k'], threshold=CFG['graph_thresh']\n)\n\nn_nodes = len(prot_ids)\nn_edges = edge_index.shape[1]\navg_deg = n_edges / n_nodes\n\nprint(f'\\n📊 Graph Statistics:')\nprint(f'  Nodes  : {n_nodes:,}')\nprint(f'  Edges  : {n_edges:,}')\nprint(f'  Avg degree: {avg_deg:.2f}')\nprint(f'  Edge weight range: [{edge_weight.min():.4f}, {edge_weight.max():.4f}]')\nprint('✅ Graph built')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-02T17:06:44.616490Z","iopub.execute_input":"2026-05-02T17:06:44.616905Z","iopub.status.idle":"2026-05-02T17:06:44.629118Z","shell.execute_reply.started":"2026-05-02T17:06:44.616876Z","shell.execute_reply":"2026-05-02T17:06:44.628261Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}