{"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":[{"sourceId":116062,"databundleVersionId":14875579,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":14053785,"sourceType":"datasetVersion","datasetId":8945796},{"sourceId":14477966,"sourceType":"datasetVersion","datasetId":9247384}],"dockerImageVersionId":31260,"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-01-21T03:53:31.330568Z","iopub.execute_input":"2026-01-21T03:53:31.330775Z","iopub.status.idle":"2026-01-21T03:53:32.428971Z","shell.execute_reply.started":"2026-01-21T03:53:31.330756Z","shell.execute_reply":"2026-01-21T03:53:32.428137Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install Biopython --q\n!pip install goatools --q\n!pip install iterative-stratification --q","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T03:54:37.254868Z","iopub.execute_input":"2026-01-21T03:54:37.255727Z","iopub.status.idle":"2026-01-21T03:55:21.842254Z","shell.execute_reply.started":"2026-01-21T03:54:37.255695Z","shell.execute_reply":"2026-01-21T03:55:21.841186Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# import libraries\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader, TensorDataset\nfrom collections import defaultdict\nfrom Bio import SeqIO\nprint(\"Done\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T03:55:52.523928Z","iopub.execute_input":"2026-01-21T03:55:52.524756Z","iopub.status.idle":"2026-01-21T03:55:55.772874Z","shell.execute_reply.started":"2026-01-21T03:55:52.524718Z","shell.execute_reply":"2026-01-21T03:55:55.772125Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load protein sequences\nfrom Bio import SeqIO\nimport pandas as pd\ntrain_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\"\ntest_path = \"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\"\n\ntrain_records = list(SeqIO.parse(train_path,\"fasta\"))\ndf_train = pd.DataFrame({\n    \"protein_id\":[rec.id.split(\"|\")[1] for rec in train_records],\n    \"sequence\":[str(rec.seq) for rec in train_records],\n})\n\ntest_records = list(SeqIO.parse(test_path,\"fasta\"))\ndf_test = pd.DataFrame({\n    \"protein_id\":[rec.id.split(\" \")[0] for rec in test_records],\n    \"sequence\":[str(rec.seq) for rec in test_records],\n})\n\n# Load train_taxonomy\ntaxonomy_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_taxonomy.tsv\"\ndf_tax = pd.read_csv(taxonomy_path,sep=\"\\t\", header=None)\ndf_tax.columns = [\"protein_id\", \"taxon_id\"]\n\ntrain_terms_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\"\n\ndf_terms = pd.read_csv(train_terms_path, sep=\"\\t\")\n\ndf_terms.columns = [\"protein_id\", \"go_id\", \"aspect\"]\ndf_terms.head()\n\n# Combine everything into a unified table\ndf_train = df_train.merge(df_tax, on = \"protein_id\",how=\"left\")\n\nterms_per_protein = (\n    df_terms.groupby(\"protein_id\")[\"go_id\"]\n    .apply(list)\n    .to_dict()\n)\n\ndf_train[\"go_terms\"] = df_train[\"protein_id\"].map(terms_per_protein).fillna(\"\").apply(\n    lambda x:x if isinstance(x,list) else[]\n)\n\nprint(\"Proteins:\", len(df_train))\nprint(\"Proteins with GO terms:\", df_train['go_terms'].apply(len).gt(0).sum())\nprint(\"Unique GO terms:\", df_terms['go_id'].nunique())\nprint(\"Aspect distribution:\", df_terms['aspect'].value_counts())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T03:55:59.480342Z","iopub.execute_input":"2026-01-21T03:55:59.480998Z","iopub.status.idle":"2026-01-21T03:56:05.852665Z","shell.execute_reply.started":"2026-01-21T03:55:59.480969Z","shell.execute_reply":"2026-01-21T03:56:05.851985Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load test dataset\ntest_records = list(SeqIO.parse(test_path, \"fasta\"))\n\nprotein_ids = []\ntaxon_ids = []\nsequences = []\n\nfor rec in test_records:\n    header = rec.description.split()\n    protein_id = header[0]          # A0A0C5B5G6\n    taxon_id = header[1]            # 9606\n    \n    protein_ids.append(protein_id)\n    taxon_ids.append(taxon_id)\n    sequences.append(str(rec.seq))\n\ndf_test = pd.DataFrame({\n    \"protein_id\": protein_ids,\n    \"taxon_id\": taxon_ids,\n    \"sequence\": sequences\n})\n\ndf_tax_species=pd.read_csv(\"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset-taxon-list.tsv\", sep=\"\\t\")\ndf_tax_species.columns=[\"taxon_id\",\"species\"]\ndf_tax_species[\"taxon_id\"] = df_tax_species[\"taxon_id\"].astype(str)\ndf_test = df_test.merge(df_tax_species, on=\"taxon_id\", how=\"left\")\ndf_test.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T03:56:14.767093Z","iopub.execute_input":"2026-01-21T03:56:14.767397Z","iopub.status.idle":"2026-01-21T03:56:17.014506Z","shell.execute_reply.started":"2026-01-21T03:56:14.767372Z","shell.execute_reply":"2026-01-21T03:56:17.013859Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# GO label propagation\n\nfrom goatools.obo_parser import GODag\ngo_dag = GODag(\"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\")\n\ndef get_ancestors(go_id,go_dag):\n    if go_id not in go_dag:\n        return set()\n    term = go_dag[go_id]\n    ancestors = term.get_all_parents()\n    ancestors.add(go_id)\n    return ancestors\n\n# Propagate GO terms for one protein\n\ndef propagate_terms(go_list, go_dag):\n    \"\"\"Propagate a list of GO terms to include all ancestors.\"\"\"\n    propagated = set()\n    for go_id in go_list:\n        if go_id not in go_dag:\n            continue\n        term = go_dag[go_id]\n        propagated.add(go_id)\n        ancestors = term.get_all_parents()\n        propagated |= ancestors\n    return propagated\n\n\ndf_train[\"go_terms_propagated\"]=df_train[\"go_terms\"].apply(\n    lambda terms: propagate_terms(terms,go_dag)\n)\n\n# Split MF / BP / C\ndef split_by_namespace(go_set, go_dag):\n    mf, bp, cc = [], [], []\n    for go_id in go_set:\n        if go_id in go_dag:\n            ns = go_dag[go_id].namespace\n            if ns == \"molecular_function\":\n                mf.append(go_id)\n            elif ns == \"biological_process\":\n                bp.append(go_id)\n            elif ns == \"cellular_component\":\n                cc.append(go_id)\n    return mf, bp, cc\n\ndf_train[\"MF\"], df_train[\"BP\"], df_train[\"CC\"] = zip(*df_train[\"go_terms_propagated\"].apply(\n    lambda s: split_by_namespace(s, go_dag)\n))\n\ndf_train.head()\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T03:56:18.622985Z","iopub.execute_input":"2026-01-21T03:56:18.623760Z","iopub.status.idle":"2026-01-21T03:56:26.150804Z","shell.execute_reply.started":"2026-01-21T03:56:18.623730Z","shell.execute_reply":"2026-01-21T03:56:26.150103Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.preprocessing import MultiLabelBinarizer\nmlb_mf = MultiLabelBinarizer()\nY_mf = mlb_mf.fit_transform(df_train[\"MF\"])\n\nmlb_bp = MultiLabelBinarizer()\nY_bp = mlb_bp.fit_transform(df_train[\"BP\"])\n\nmlb_cc = MultiLabelBinarizer()\nY_cc = mlb_cc.fit_transform(df_train[\"CC\"])\n\n\nprint(\"Avg MF labels:\", Y_mf.sum(axis=1).mean())\nprint(\"Avg BP labels:\", Y_bp.sum(axis=1).mean())\nprint(\"Avg CC labels:\", Y_cc.sum(axis=1).mean())\n\nprint(\"Y_mf.shape:\",Y_mf.shape)\nprint(\"Y_bp.shape:\",Y_bp.shape)\nprint(\"Y_cc.shape:\",Y_cc.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T03:56:30.183435Z","iopub.execute_input":"2026-01-21T03:56:30.184310Z","iopub.status.idle":"2026-01-21T03:56:41.072829Z","shell.execute_reply.started":"2026-01-21T03:56:30.184283Z","shell.execute_reply":"2026-01-21T03:56:41.072167Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import joblib\n\n\njoblib.dump(mlb_mf, \"/kaggle/working/mlb_mf.pkl\")\njoblib.dump(mlb_bp, \"/kaggle/working/mlb_bp.pkl\")\njoblib.dump(mlb_cc, \"/kaggle/working/mlb_cc.pkl\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T03:56:45.047490Z","iopub.execute_input":"2026-01-21T03:56:45.047936Z","iopub.status.idle":"2026-01-21T03:56:45.061195Z","shell.execute_reply.started":"2026-01-21T03:56:45.047910Z","shell.execute_reply":"2026-01-21T03:56:45.060589Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# filter rare GO terms\n\nROOT_TERMS = {'GO:0008150', 'GO:0003674', 'GO:0005575'}\n\nmin_count = 30  # common threshold\n\nmf_counts = Y_mf.sum(axis=0)\nkeep_mf = [\n    (count >= min_count) and (term not in ROOT_TERMS)\n    for count, term in zip(mf_counts, mlb_mf.classes_)\n]\nY_mf = Y_mf[:, keep_mf]\nmf_terms = [t for t, k in zip(mlb_mf.classes_, keep_mf) if k]\n\n\nmin_count = 50  # common threshold\nbp_counts = Y_bp.sum(axis=0)\nkeep_bp = [\n    (count >= min_count) and (term not in ROOT_TERMS)\n    for count, term in zip(bp_counts, mlb_bp.classes_)\n]\nY_bp = Y_bp[:, keep_bp]\nbp_terms = [t for t, k in zip(mlb_bp.classes_, keep_bp) if k]\n\nmin_count = 20\ncc_counts = Y_cc.sum(axis=0)\nkeep_cc = [\n    (count >= min_count) and (term not in ROOT_TERMS)\n    for count, term in zip(cc_counts, mlb_cc.classes_)\n]\nY_cc = Y_cc[:, keep_cc]\ncc_terms = [t for t, k in zip(mlb_cc.classes_, keep_cc) if k]\n\n\n\nprint(\"Y_mf.shape:\",Y_mf.shape)\nprint(\"Y_bp.shape:\",Y_bp.shape)\nprint(\"Y_cc.shape:\",Y_cc.shape)\n\nnp.save(\"Y_mf.npy\", Y_mf)\nnp.save(\"Y_bp.npy\", Y_bp)\nnp.save(\"Y_cc.npy\", Y_cc)\n\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T03:56:47.859398Z","iopub.execute_input":"2026-01-21T03:56:47.860263Z","iopub.status.idle":"2026-01-21T03:57:00.946227Z","shell.execute_reply.started":"2026-01-21T03:56:47.860234Z","shell.execute_reply":"2026-01-21T03:57:00.945396Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\nprint(f\"Total proteins: {len(df_train):,}\")\nprint(f\"MF terms: {Y_mf.shape[1]}\")\nprint(f\"BP terms: {Y_bp.shape[1]}\")\nprint(f\"CC terms: {Y_cc.shape[1]}\")\nprint(f\"Total GO terms: {Y_mf.shape[1] + Y_bp.shape[1] + Y_cc.shape[1]}\")\n\nprint(f\"MF: avg {Y_mf.sum(axis=1).mean():.1f} terms/protein, \"\n      f\"avg {Y_mf.sum(axis=0).mean():.1f} proteins/term\")\nprint(f\"BP: avg {Y_bp.sum(axis=1).mean():.1f} terms/protein, \"\n      f\"avg {Y_bp.sum(axis=0).mean():.1f} proteins/term\")\nprint(f\"CC: avg {Y_cc.sum(axis=1).mean():.1f} terms/protein, \"\n      f\"avg {Y_cc.sum(axis=0).mean():.1f} proteins/term\")\n\nprint(f\"MF: {Y_mf.mean()*100:.2f}% positive labels\")\nprint(f\"BP: {Y_bp.mean()*100:.2f}% positive labels\")\nprint(f\"CC: {Y_cc.mean()*100:.2f}% positive labels\")\n# Verify no root terms remain\nassert not any(term in ROOT_TERMS for term in mf_terms), \"MF still has root!\"\nassert not any(term in ROOT_TERMS for term in bp_terms), \"BP still has root!\"\nassert not any(term in ROOT_TERMS for term in cc_terms), \"CC still has root!\"\n\nprint(\"No root terms in training labels\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T03:57:11.309366Z","iopub.execute_input":"2026-01-21T03:57:11.310125Z","iopub.status.idle":"2026-01-21T03:57:12.579398Z","shell.execute_reply.started":"2026-01-21T03:57:11.310060Z","shell.execute_reply":"2026-01-21T03:57:12.578626Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load the protT5 embeddings\ntrain_path = \"/kaggle/input/prott5-embeddings/protT5_embeddings.npy\"\ntest_path = \"/kaggle/input/test-embeddings-prot-t5-xl/protT5_test_embeddings.npy\"\n\nX_np = np.load(train_path)\nX_train_torch = torch.from_numpy(X_np).float()\n\ntest  =  np.load(test_path)\nX_test_torch = torch.from_numpy(test).float()\n\ndef load_ids(fasta):\n    return [r.id.split()[0] for r in SeqIO.parse(fasta,\"fasta\")]\n\ntrain_ids = load_ids(\"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\")\nprint(train_ids[0])\ntrain_ids = [id.split(\"|\")[1] for id in train_ids]\ntest_ids = load_ids(\"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\")\n\nprint(train_ids[0],len(train_ids))\nprint(test_ids[0],len(test_ids))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T03:57:17.333404Z","iopub.execute_input":"2026-01-21T03:57:17.334310Z","iopub.status.idle":"2026-01-21T03:57:26.919350Z","shell.execute_reply.started":"2026-01-21T03:57:17.334277Z","shell.execute_reply":"2026-01-21T03:57:26.918584Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_embedding_dict = {\n    pid: X_np[i]\n    for i, pid in enumerate(train_ids)\n}\ntest_embedding_dict = {\n    pid: test[i]\n    for i, pid in enumerate(test_ids)\n}\n\nX_train = np.vstack([train_embedding_dict[pid] for pid in df_train[\"protein_id\"]])\nassert len(df_train) == len(X_train), \"Length mismatch!\"\n\nmissing = set(df_train[\"protein_id\"]) - set(train_embedding_dict.keys())\nassert len(missing) == 0, f\"Missing embeddings: {len(missing)}\"\n\n\n\nX_test = np.vstack([test_embedding_dict[pid] for pid in df_test[\"protein_id\"]])\nnp.save(\"X_train_protT5_xl.npy\", X_train)\nnp.save(\"X_test_protT5_xl.npy\", X_test)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T03:57:33.115826Z","iopub.execute_input":"2026-01-21T03:57:33.116159Z","iopub.status.idle":"2026-01-21T03:57:34.454769Z","shell.execute_reply.started":"2026-01-21T03:57:33.116131Z","shell.execute_reply":"2026-01-21T03:57:34.454129Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"go_counts = (\n    np.asarray(Y_mf.sum(axis=1)).ravel() +\n    np.asarray(Y_bp.sum(axis=1)).ravel() +\n    np.asarray(Y_cc.sum(axis=1)).ravel()\n)\n\ngo_counts_capped = np.clip(go_counts, 0, 20)\n\nbins = pd.qcut(\n    go_counts_capped,\n    q=10,\n    duplicates=\"drop\"\n)\n\nbins.value_counts().sort_index()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T03:57:38.184310Z","iopub.execute_input":"2026-01-21T03:57:38.184980Z","iopub.status.idle":"2026-01-21T03:57:38.616753Z","shell.execute_reply.started":"2026-01-21T03:57:38.184943Z","shell.execute_reply":"2026-01-21T03:57:38.615857Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.model_selection import train_test_split\n# Split\n\nidx = np.arange(len(go_counts))\n\ntrain_idx, val_idx = train_test_split(\n    idx,\n    test_size=0.2,\n    random_state=42,\n    stratify=bins\n)\n\nX_tr = X_train[train_idx]\nX_val = X_train[val_idx]\n\nY_mf_tr, Y_mf_val = Y_mf[train_idx], Y_mf[val_idx]\nY_bp_tr, Y_bp_val = Y_bp[train_idx], Y_bp[val_idx]\nY_cc_tr, Y_cc_val = Y_cc[train_idx], Y_cc[val_idx]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T04:11:08.602971Z","iopub.execute_input":"2026-01-21T04:11:08.603634Z","iopub.status.idle":"2026-01-21T04:11:16.886328Z","shell.execute_reply.started":"2026-01-21T04:11:08.603601Z","shell.execute_reply":"2026-01-21T04:11:16.885370Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"Train bin distribution:\")\nprint(pd.qcut(go_counts[train_idx], q=10, duplicates=\"drop\").value_counts().sort_index())\n\nprint(\"\\nVal bin distribution:\")\nprint(pd.qcut(go_counts[val_idx], q=10, duplicates=\"drop\").value_counts().sort_index())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T04:11:21.459736Z","iopub.execute_input":"2026-01-21T04:11:21.460370Z","iopub.status.idle":"2026-01-21T04:11:21.481334Z","shell.execute_reply.started":"2026-01-21T04:11:21.460340Z","shell.execute_reply":"2026-01-21T04:11:21.480495Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class MF_MLP(nn.Module):\n    def __init__(self, in_dim, out_dim):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(in_dim, 2048),\n            nn.LayerNorm(2048),\n            nn.ReLU(),\n            nn.Dropout(0.4),\n            nn.Linear(2048, out_dim)\n        )\n    \n    def forward(self, x):\n        return self.net(x)\n\nclass BP_MLP(nn.Module):\n    def __init__(self, in_dim, out_dim):\n        super().__init__()\n     \n        self.net = nn.Sequential(\n            nn.Linear(in_dim, 2048),  # Increased from 512\n            nn.LayerNorm(2048),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            \n            nn.Linear(2048, 1024),  # Increased from 256\n            nn.LayerNorm(1024),\n            nn.ReLU(),\n            nn.Dropout(0.2),\n\n                              \n            nn.Linear(1024, out_dim)\n        )\n    \n    def forward(self, x):\n        return self.net(x)\n\nclass CC_MLP(nn.Module):\n    def __init__(self, in_dim, out_dim):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(in_dim, 1024),\n            nn.LayerNorm(1024),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            \n            nn.Linear(1024, 512),\n            nn.LayerNorm(512),\n            nn.ReLU(),\n            nn.Dropout(0.2),\n            \n            nn.Linear(512, out_dim)\n        )\n    \n    def forward(self, x):\n        return self.net(x)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T06:41:09.105937Z","iopub.execute_input":"2026-01-21T06:41:09.106260Z","iopub.status.idle":"2026-01-21T06:41:09.114329Z","shell.execute_reply.started":"2026-01-21T06:41:09.106228Z","shell.execute_reply":"2026-01-21T06:41:09.113581Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# Train the model\n\ndef train_aspect(\n    model_class, \n    X_train, Y_train, \n    X_val, Y_val, \n    aspect_name=\"\",\n    pos_weight_alpha=0.5,  # Smoothing factor\n    epochs=50, \n    batch_size=5000, \n    lr=1e-3, \n    patience=10\n):\n    \"\"\"\n    Train a single-aspect model.\n    \n    Args:\n        model_class: Model class (MF_MLP, BP_MLP, or CC_MLP)\n        X_train, Y_train: Training data (numpy arrays or tensors)\n        X_val, Y_val: Validation data\n        aspect_name: Name for logging (MF/BP/CC)\n        pos_weight_alpha: Smoothing for pos_weight (0.5 = moderate)\n        epochs: Maximum epochs\n        batch_size: Batch size\n        lr: Learning rate\n        patience: Early stopping patience\n    \n    Returns:\n        Trained model\n    \"\"\"\n    print(f\"\\n{'='*60}\")\n    print(f\"Training {aspect_name} Model\")\n    print(f\"{'='*60}\")\n    \n    # Convert to tensors if needed\n    if not isinstance(X_train, torch.Tensor):\n        X_train = torch.tensor(X_train, dtype=torch.float32)\n    if not isinstance(Y_train, torch.Tensor):\n        Y_train = torch.tensor(Y_train, dtype=torch.float32)\n    if not isinstance(X_val, torch.Tensor):\n        X_val = torch.tensor(X_val, dtype=torch.float32)\n    if not isinstance(Y_val, torch.Tensor):\n        Y_val = torch.tensor(Y_val, dtype=torch.float32)\n    \n    # Move validation data to GPU once (make copies to avoid modifying originals)\n    X_val_gpu = X_val.to(device)\n    Y_val_gpu = Y_val.to(device)\n    \n    # Compute pos_weight with smoothing\n    pos_freq = Y_train.mean(dim=0)  # Frequency of each class\n    pos_freq = torch.clamp(pos_freq, min=0.001, max=0.999)  # Avoid division by zero\n    \n    # Smooth pos_weight: ((1-p)/p)^alpha\n    pos_weight = ((1 - pos_freq) / pos_freq) ** pos_weight_alpha\n    pos_weight = pos_weight.to(device)\n    \n    print(f\"Class balance: {pos_freq.mean().item()*100:.2f}% positive\")\n    print(f\"Avg pos_weight: {pos_weight.mean().item():.1f}x\")\n    print(f\"Training samples: {len(X_train):,}\")\n    print(f\"Output dimensions: {Y_train.shape[1]}\")\n    \n    # DataLoader\n    dataset = TensorDataset(X_train, Y_train)\n    loader = DataLoader(\n        dataset, \n        batch_size=batch_size, \n        shuffle=True, \n        pin_memory=True,\n        num_workers=0 \n    )\n    \n    # Initialize model\n    model = model_class(X_train.shape[1], Y_train.shape[1]).to(device)\n    \n    # Optimizer & loss\n    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)\n    loss_fn = nn.BCEWithLogitsLoss(pos_weight=pos_weight)\n    \n    # Learning rate scheduler\n    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n        optimizer,\n        mode='min',\n        factor=0.5,\n        patience=3\n    )\n    \n    # Early stopping\n    best_val_loss = float(\"inf\")\n    patience_counter = 0\n    best_model_state = None\n    \n    print(f\"\\nStarting training...\")\n    \n    for epoch in range(epochs):\n    \n        model.train()\n        train_loss = 0.0\n        \n        for xb, yb in loader:\n            xb, yb = xb.to(device), yb.to(device)\n            \n            optimizer.zero_grad()\n            logits = model(xb)\n            loss = loss_fn(logits, yb)\n            loss.backward()\n            \n            # Gradient clipping\n            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n            \n            optimizer.step()\n            train_loss += loss.item() * xb.size(0)\n        \n        train_loss /= len(X_train)\n        \n        # ==================\n        # VALIDATION\n        # ==================\n        model.eval()\n        with torch.no_grad():\n            val_logits = model(X_val_gpu)\n            val_loss = loss_fn(val_logits, Y_val_gpu).item()\n        \n        # Update scheduler\n        scheduler.step(val_loss)\n        \n        # Print progress\n        print(f\"Epoch {epoch+1:3d}/{epochs}: \"\n              f\"train_loss={train_loss:.4f}, \"\n              f\"val_loss={val_loss:.4f}, \"\n              f\"lr={optimizer.param_groups[0]['lr']:.6f}\")\n        \n        # Early stopping check\n        if val_loss < best_val_loss:\n            best_val_loss = val_loss\n            patience_counter = 0\n            best_model_state = model.state_dict()\n            print(f\"  → New best validation loss: {val_loss:.4f}\")\n        else:\n            patience_counter += 1\n            if patience_counter >= patience:\n                print(f\"\\nEarly stopping triggered after {epoch+1} epochs\")\n                break\n    \n    # Load best model\n    if best_model_state is not None:\n        model.load_state_dict(best_model_state)\n        print(f\"\\n✓ Best model restored (val_loss={best_val_loss:.4f})\")\n    \n    return model\n\n\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T06:41:12.373046Z","iopub.execute_input":"2026-01-21T06:41:12.373856Z","iopub.status.idle":"2026-01-21T06:41:12.386658Z","shell.execute_reply.started":"2026-01-21T06:41:12.373828Z","shell.execute_reply":"2026-01-21T06:41:12.385872Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def compute_fmax(preds, labels, thresholds=None):\n    \"\"\"\n    CAFA-style protein-centric Fmax.\n    \n    Args:\n        preds: (n_proteins, n_terms) prediction scores\n        labels: (n_proteins, n_terms) binary labels\n        thresholds: array of thresholds to try\n    \n    Returns:\n        best_fmax, best_threshold\n    \"\"\"\n    if thresholds is None:\n        # Use finer thresholds in low range where optimal often is\n        thresholds = np.concatenate([\n            np.arange(0.01, 0.1, 0.01),   # Fine-grained for low thresholds\n            np.arange(0.1, 0.5, 0.02),    # Medium-grained\n            np.arange(0.5, 1.0, 0.05)     # Coarse for high thresholds\n        ])\n    \n    best_f = 0\n    best_t = 0\n    n_proteins = labels.shape[0]\n    \n    for t in thresholds:\n        # Binarize predictions at threshold t\n        preds_binary = (preds >= t).astype(np.int32)\n        \n        # Vectorized computation per protein\n        tp = np.sum(labels & preds_binary, axis=1)  # True positives per protein\n        fp = np.sum((1 - labels) & preds_binary, axis=1)  # False positives per protein\n        fn = np.sum(labels & (1 - preds_binary), axis=1)  # False negatives per protein\n        \n        # Recall for all proteins (can be 0)\n        recall_denom = tp + fn\n        recalls = np.zeros(n_proteins)\n        valid_recall = recall_denom > 0\n        recalls[valid_recall] = tp[valid_recall] / recall_denom[valid_recall]\n        \n        # Precision only for proteins with predictions\n        precision_denom = tp + fp\n        has_predictions = precision_denom > 0\n        \n        if not np.any(has_predictions):\n            # No predictions at this threshold, skip\n            continue\n        \n        precisions = tp[has_predictions] / precision_denom[has_predictions]\n        \n        # Average precision (only over proteins with predictions)\n        # Average recall (over ALL proteins)\n        P = np.mean(precisions)\n        R = np.mean(recalls)\n        \n        if P + R > 0:\n            F = 2 * P * R / (P + R)\n        else:\n            F = 0\n        \n        if F > best_f:\n            best_f = F\n            best_t = t\n    \n    return best_f, best_t","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T06:49:00.610750Z","iopub.execute_input":"2026-01-21T06:49:00.611404Z","iopub.status.idle":"2026-01-21T06:49:00.618556Z","shell.execute_reply.started":"2026-01-21T06:49:00.611374Z","shell.execute_reply":"2026-01-21T06:49:00.617798Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n\n# Convert validation data to tensors once (for evaluation later)\nif not isinstance(X_val, torch.Tensor):\n    X_val_tensor = torch.tensor(X_val, dtype=torch.float32)\nelse:\n    X_val_tensor = X_val\n    \n# Train MF model\nmodel_mf = train_aspect(\n    MF_MLP, \n    X_tr, Y_mf_tr, \n    X_val, Y_mf_val,\n    aspect_name=\"MF\",\n    pos_weight_alpha=0, \n    epochs=100,\n    batch_size= 1024 ,\n    lr=1e-3,\n    patience=15\n)\n    \n# Evaluate MF\nmodel_mf.eval()\nwith torch.no_grad():\n    pred_mf = torch.sigmoid(model_mf(X_val_tensor.to(device))).cpu().numpy()\n    \nfmax_mf, thresh_mf = compute_fmax(pred_mf, Y_mf_val if isinstance(Y_mf_val, np.ndarray) else Y_mf_val.numpy())\nprint(f\"\\n{'='*60}\")\nprint(f\"MF Results: Fmax={fmax_mf:.4f} @ t={thresh_mf:.3f} \")\nprint(f\"{'='*60}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T06:41:55.905728Z","iopub.execute_input":"2026-01-21T06:41:55.906341Z","iopub.status.idle":"2026-01-21T06:43:23.200839Z","shell.execute_reply.started":"2026-01-21T06:41:55.906313Z","shell.execute_reply":"2026-01-21T06:43:23.200110Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Train BP model\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nmodel_bp = train_aspect(\n    BP_MLP, \n    X_tr, Y_bp_tr,                               \n    X_val, Y_bp_val,\n    aspect_name=\"BP\",\n    pos_weight_alpha=0,\n    epochs=100,\n    batch_size=5000,\n    lr=1e-3,\n    patience=15\n)\n    \n# Evaluate BP\nmodel_bp.eval()\nwith torch.no_grad():\n    pred_bp = torch.sigmoid(model_bp(X_val_tensor.to(device))).cpu().numpy()\n    \nfmax_bp, thresh_bp = compute_fmax(pred_bp, Y_bp_val if isinstance(Y_bp_val, np.ndarray) else Y_bp_val.numpy())\nprint(f\"\\n{'='*60}\")\nprint(f\"MF Results: Fmax={fmax_bp:.4f} @ t={thresh_bp:.3f} \")\nprint(f\"{'='*60}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T06:45:14.442200Z","iopub.execute_input":"2026-01-21T06:45:14.442903Z","iopub.status.idle":"2026-01-21T06:48:36.865434Z","shell.execute_reply.started":"2026-01-21T06:45:14.442874Z","shell.execute_reply":"2026-01-21T06:48:36.864674Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n\n# Train CC model\nmodel_cc = train_aspect(\n    CC_MLP, \n    X_tr, Y_cc_tr, \n    X_val, Y_cc_val,\n    aspect_name=\"CC\",\n    pos_weight_alpha=0,\n    epochs=50,\n    batch_size=1024,\n    lr=1e-3,\n    patience=10\n)\n    \n# Evaluate CC\nmodel_cc.eval()\nwith torch.no_grad():\n    pred_cc = torch.sigmoid(model_cc(X_val_tensor.to(device))).cpu().numpy()\n    \nfmax_cc, thresh_cc = compute_fmax(pred_cc, Y_cc_val if isinstance(Y_cc_val, np.ndarray) else Y_cc_val.numpy())\nprint(f\"\\n{'='*60}\")\nprint(f\"CC Results: Fmax={fmax_cc:.4f} @ t={thresh_cc:.3f}\" )\nprint(f\"{'='*60}\")\n    \n# Overall results\navg_fmax = (fmax_mf + fmax_bp + fmax_cc) / 3\nprint(f\"\\n{'='*60}\")\nprint(f\"OVERALL RESULTS\")\nprint(f\"{'='*60}\")\nprint(f\"Average Fmax: {avg_fmax:.4f}\")\nprint(f\"MF: {fmax_mf:.4f}\")\nprint(f\"BP: {fmax_bp:.4f}\")\nprint(f\"CC: {fmax_cc:.4f}\")\n    \n# Save models\ntorch.save(model_mf.state_dict(), \"model_mf_best.pt\")\ntorch.save(model_bp.state_dict(), \"model_bp_best.pt\")\ntorch.save(model_cc.state_dict(), \"model_cc_best.pt\")\nprint(f\"\\n✓ Models saved!\")\n    ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T06:50:13.951966Z","iopub.execute_input":"2026-01-21T06:50:13.952690Z","iopub.status.idle":"2026-01-21T06:50:59.392548Z","shell.execute_reply.started":"2026-01-21T06:50:13.952658Z","shell.execute_reply":"2026-01-21T06:50:59.391958Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def compute_fmax_with_details(preds, labels, thresholds=None):\n    \"\"\"\n    Extended version that returns precision and recall at optimal threshold.\n    \"\"\"\n    if thresholds is None:\n        thresholds = np.concatenate([\n            np.arange(0.01, 0.1, 0.01),\n            np.arange(0.1, 0.5, 0.02),\n            np.arange(0.5, 1.0, 0.05)\n        ])\n    \n    best_f = 0\n    best_t = 0\n    best_p = 0\n    best_r = 0\n    n_proteins = labels.shape[0]\n    \n    for t in thresholds:\n        preds_binary = (preds >= t).astype(np.int32)\n        \n        tp = np.sum(labels & preds_binary, axis=1)\n        fp = np.sum((1 - labels) & preds_binary, axis=1)\n        fn = np.sum(labels & (1 - preds_binary), axis=1)\n        \n        recall_denom = tp + fn\n        recalls = np.zeros(n_proteins)\n        valid_recall = recall_denom > 0\n        recalls[valid_recall] = tp[valid_recall] / recall_denom[valid_recall]\n        \n        precision_denom = tp + fp\n        has_predictions = precision_denom > 0\n        \n        if not np.any(has_predictions):\n            continue\n        \n        precisions = tp[has_predictions] / precision_denom[has_predictions]\n        \n        P = np.mean(precisions)\n        R = np.mean(recalls)\n        \n        if P + R > 0:\n            F = 2 * P * R / (P + R)\n        else:\n            F = 0\n        \n        if F > best_f:\n            best_f = F\n            best_t = t\n            best_p = P\n            best_r = R\n    \n    return best_f, best_t, best_p, best_r","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T07:09:03.812818Z","iopub.execute_input":"2026-01-21T07:09:03.813139Z","iopub.status.idle":"2026-01-21T07:09:03.820576Z","shell.execute_reply.started":"2026-01-21T07:09:03.813113Z","shell.execute_reply":"2026-01-21T07:09:03.819892Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport numpy as np\nfrom tqdm import tqdm\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(f\"Using device: {device}\")\n\n# ============================================\n# LOAD TRAINED MODELS\n# ============================================\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"LOADING TRAINED MODELS\")\nprint(\"=\"*60)\n\n# Initialize models\nmf_model = MF_MLP(X_val.shape[1], Y_mf_val.shape[1]).to(device)\nbp_model = BP_MLP(X_val.shape[1], Y_bp_val.shape[1]).to(device)\ncc_model = CC_MLP(X_val.shape[1], Y_cc_val.shape[1]).to(device)\n\n# Load state dicts\ntry:\n    mf_model.load_state_dict(torch.load(\"/kaggle/working/model_mf_best.pt\", map_location=device))\n    print(\"✓ MF model loaded\")\nexcept FileNotFoundError:\n    print(\"✗ MF model not found!\")\n\ntry:\n    bp_model.load_state_dict(torch.load(\"/kaggle/working/model_bp_best.pt\", map_location=device))\n    print(\"✓ BP model loaded\")\nexcept FileNotFoundError:\n    print(\"✗ BP model not found!\")\n\ntry:\n    cc_model.load_state_dict(torch.load(\"/kaggle/working/model_cc_best.pt\", map_location=device))\n    print(\"✓ CC model loaded\")\nexcept FileNotFoundError:\n    print(\"✗ CC model not found!\")\n\n# Set to evaluation mode\nmf_model.eval()\nbp_model.eval()\ncc_model.eval()\n\nprint(\"✓ All models set to eval mode\")\n\n# ============================================\n# VALIDATION SET PREDICTIONS\n# ============================================\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"GENERATING VALIDATION PREDICTIONS\")\nprint(\"=\"*60)\n\n# Convert to tensor\nX_val_tensor = torch.tensor(X_val, dtype=torch.float32).to(device)\n\n# Generate predictions\nwith torch.no_grad():\n    print(\"Predicting MF...\")\n    mf_val_logits = mf_model(X_val_tensor)\n    mf_val_pred = torch.sigmoid(mf_val_logits).cpu().numpy()\n    \n    print(\"Predicting BP...\")\n    bp_val_logits = bp_model(X_val_tensor)\n    bp_val_pred = torch.sigmoid(bp_val_logits).cpu().numpy()\n    \n    print(\"Predicting CC...\")\n    cc_val_logits = cc_model(X_val_tensor)\n    cc_val_pred = torch.sigmoid(cc_val_logits).cpu().numpy()\n\nprint(\"\\n✓ Predictions generated\")\n\n# Verify shapes\nprint(\"\\nShape verification:\")\nprint(f\"  MF: {mf_val_pred.shape} (expected: {Y_mf_val.shape})\")\nprint(f\"  BP: {bp_val_pred.shape} (expected: {Y_bp_val.shape})\")\nprint(f\"  CC: {cc_val_pred.shape} (expected: {Y_cc_val.shape})\")\n\nassert mf_val_pred.shape == (Y_mf_val.shape if isinstance(Y_mf_val, np.ndarray) else Y_mf_val.numpy().shape), \"MF shape mismatch!\"\nassert bp_val_pred.shape == (Y_bp_val.shape if isinstance(Y_bp_val, np.ndarray) else Y_bp_val.numpy().shape), \"BP shape mismatch!\"\nassert cc_val_pred.shape == (Y_cc_val.shape if isinstance(Y_cc_val, np.ndarray) else Y_cc_val.numpy().shape), \"CC shape mismatch!\"\n\n# Save predictions\nnp.save(\"mlp_val_mf.npy\", mf_val_pred)\nnp.save(\"mlp_val_bp.npy\", bp_val_pred)\nnp.save(\"mlp_val_cc.npy\", cc_val_pred)\n\nprint(\"\\n✓ Predictions saved to .npy files\")\n\n# ============================================\n# PREDICTION STATISTICS\n# ============================================\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"PREDICTION STATISTICS\")\nprint(\"=\"*60)\n\nprint(\"\\nMF Predictions:\")\nprint(f\"  Mean: {mf_val_pred.mean():.4f}\")\nprint(f\"  Median: {np.median(mf_val_pred):.4f}\")\nprint(f\"  Std: {mf_val_pred.std():.4f}\")\nprint(f\"  Min: {mf_val_pred.min():.4f}, Max: {mf_val_pred.max():.4f}\")\nprint(f\"  % > 0.5: {(mf_val_pred > 0.5).mean()*100:.2f}%\")\n\nprint(\"\\nBP Predictions:\")\nprint(f\"  Mean: {bp_val_pred.mean():.4f}\")\nprint(f\"  Median: {np.median(bp_val_pred):.4f}\")\nprint(f\"  Std: {bp_val_pred.std():.4f}\")\nprint(f\"  Min: {bp_val_pred.min():.4f}, Max: {bp_val_pred.max():.4f}\")\nprint(f\"  % > 0.5: {(bp_val_pred > 0.5).mean()*100:.2f}%\")\n\nprint(\"\\nCC Predictions:\")\nprint(f\"  Mean: {cc_val_pred.mean():.4f}\")\nprint(f\"  Median: {np.median(cc_val_pred):.4f}\")\nprint(f\"  Std: {cc_val_pred.std():.4f}\")\nprint(f\"  Min: {cc_val_pred.min():.4f}, Max: {cc_val_pred.max():.4f}\")\nprint(f\"  % > 0.5: {(cc_val_pred > 0.5).mean()*100:.2f}%\")\n\n# ============================================\n# COMPUTE FMAX\n# ============================================\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"COMPUTING FMAX SCORES\")\nprint(\"=\"*60)\n\n# Convert labels to numpy if needed\nY_mf_val_np = Y_mf_val if isinstance(Y_mf_val, np.ndarray) else Y_mf_val.numpy()\nY_bp_val_np = Y_bp_val if isinstance(Y_bp_val, np.ndarray) else Y_bp_val.numpy()\nY_cc_val_np = Y_cc_val if isinstance(Y_cc_val, np.ndarray) else Y_cc_val.numpy()\n\n# Compute Fmax for each aspect\nprint(\"\\nComputing MF Fmax...\")\nfmax_mf, thresh_mf, prec_mf, rec_mf = compute_fmax_with_details(mf_val_pred, Y_mf_val_np)\n\nprint(\"Computing BP Fmax...\")\nfmax_bp, thresh_bp, prec_bp, rec_bp = compute_fmax_with_details(bp_val_pred, Y_bp_val_np)\n\nprint(\"Computing CC Fmax...\")\nfmax_cc, thresh_cc, prec_cc, rec_cc = compute_fmax_with_details(cc_val_pred, Y_cc_val_np)\n\n# ============================================\n# FINAL RESULTS\n# ============================================\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"FINAL VALIDATION RESULTS\")\nprint(\"=\"*60)\n\nprint(f\"\\nMF:\")\nprint(f\"  Fmax:      {fmax_mf:.4f}\")\nprint(f\"  Threshold: {thresh_mf:.3f}\")\nprint(f\"  Precision: {prec_mf:.3f}\")\nprint(f\"  Recall:    {rec_mf:.3f}\")\n\nprint(f\"\\nBP:\")\nprint(f\"  Fmax:      {fmax_bp:.4f}\")\nprint(f\"  Threshold: {thresh_bp:.3f}\")\nprint(f\"  Precision: {prec_bp:.3f}\")\nprint(f\"  Recall:    {rec_bp:.3f}\")\n\nprint(f\"\\nCC:\")\nprint(f\"  Fmax:      {fmax_cc:.4f}\")\nprint(f\"  Threshold: {thresh_cc:.3f}\")\nprint(f\"  Precision: {prec_cc:.3f}\")\nprint(f\"  Recall:    {rec_cc:.3f}\")\n\navg_fmax = (fmax_mf + fmax_bp + fmax_cc) / 3\n\nprint(f\"\\n\" + \"=\"*60)\nprint(f\"AVERAGE FMAX: {avg_fmax:.4f}\")\nprint(f\"=\"*60)\n\n# Save thresholds for test set inference\noptimal_thresholds = {\n    'MF': thresh_mf,\n    'BP': thresh_bp,\n    'CC': thresh_cc\n}\n\nimport pickle\nwith open('optimal_thresholds.pkl', 'wb') as f:\n    pickle.dump(optimal_thresholds, f)\n\nprint(\"\\n✓ Optimal thresholds saved to 'optimal_thresholds.pkl'\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T07:09:19.974298Z","iopub.execute_input":"2026-01-21T07:09:19.976378Z","iopub.status.idle":"2026-01-21T07:10:13.598102Z","shell.execute_reply.started":"2026-01-21T07:09:19.975879Z","shell.execute_reply":"2026-01-21T07:10:13.597120Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"MF:\", mf_val_pred.min(), mf_val_pred.max())\nprint(\"BP:\", bp_val_pred.min(), bp_val_pred.max())\nprint(\"CC:\", cc_val_pred.min(), cc_val_pred.max())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T07:10:31.354539Z","iopub.execute_input":"2026-01-21T07:10:31.355143Z","iopub.status.idle":"2026-01-21T07:10:31.430496Z","shell.execute_reply.started":"2026-01-21T07:10:31.355106Z","shell.execute_reply":"2026-01-21T07:10:31.429609Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nfrom goatools.obo_parser import GODag\ngo_dag = GODag(\"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T07:26:32.981362Z","iopub.execute_input":"2026-01-21T07:26:32.981694Z","iopub.status.idle":"2026-01-21T07:26:33.636414Z","shell.execute_reply.started":"2026-01-21T07:26:32.981667Z","shell.execute_reply":"2026-01-21T07:26:33.635788Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nX_test_tensor = torch.tensor(X_test, dtype=torch.float32).to(device)\n\nmf_model = MF_MLP(1024, 1062).to(device)\nbp_model = BP_MLP(1024, 3539).to(device)\ncc_model = CC_MLP(1024, 884).to(device)\n\n\nmf_model.load_state_dict(torch.load(\"/kaggle/working/model_mf_best.pt\", map_location=device))\nbp_model.load_state_dict(torch.load(\"/kaggle/working/model_bp_best.pt\", map_location=device))\ncc_model.load_state_dict(torch.load(\"/kaggle/working/model_cc_best.pt\", map_location=device))\n\nmf_model.eval()\nbp_model.eval()\ncc_model.eval()\n\nwith torch.no_grad():\n    mf_logits = mf_model(X_test_tensor)\n    bp_logits = bp_model(X_test_tensor)\n    cc_logits = cc_model(X_test_tensor)\n\nmf_pred_mlp = torch.sigmoid(mf_logits).cpu().numpy()\nbp_pred_mlp = torch.sigmoid(bp_logits).cpu().numpy()\ncc_pred_mlp = torch.sigmoid(cc_logits).cpu().numpy()\n\nnp.save(\"mf_pred_mlp.npy\", mf_pred_mlp)\nnp.save(\"bp_pred_mlp.npy\", bp_pred_mlp)\nnp.save(\"cc_pred_mlp.npy\", cc_pred_mlp)\n\nprint(mf_pred_mlp.shape)  \nprint(bp_pred_mlp.shape) \nprint(cc_pred_mlp.shape) \n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T07:17:42.084343Z","iopub.execute_input":"2026-01-21T07:17:42.085044Z","iopub.status.idle":"2026-01-21T07:18:04.399375Z","shell.execute_reply.started":"2026-01-21T07:17:42.085003Z","shell.execute_reply":"2026-01-21T07:18:04.398646Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pickle\nimport pandas as pd\nimport numpy as np\nfrom collections import defaultdict\nfrom tqdm import tqdm\n\n# ============================================\n# LOAD OPTIMAL THRESHOLDS\n# ============================================\n\nprint(\"=\"*60)\nprint(\"CREATING CAFA SUBMISSION\")\nprint(\"=\"*60)\n\nwith open('/kaggle/working/optimal_thresholds.pkl', 'rb') as f:\n    optimal_thresholds = pickle.load(f)\n\nthresh_mf = optimal_thresholds['MF']  # 0.260\nthresh_bp = optimal_thresholds['BP']  # 0.180\nthresh_cc = optimal_thresholds['CC']  # 0.240\n\nprint(f\"\\nUsing optimal thresholds:\")\nprint(f\"  MF: {thresh_mf:.3f}\")\nprint(f\"  BP: {thresh_bp:.3f}\")\nprint(f\"  CC: {thresh_cc:.3f}\")\n\n# ============================================\n# LOAD PREDICTIONS (you already generated these)\n# ============================================\n\nprint(f\"\\nLoading test predictions...\")\nmf_pred_mlp = np.load(\"/kaggle/working/mf_pred_mlp.npy\")\nbp_pred_mlp = np.load(\"/kaggle/working/bp_pred_mlp.npy\")\ncc_pred_mlp = np.load(\"/kaggle/working/cc_pred_mlp.npy\")\n\nprint(f\"Prediction shapes:\")\nprint(f\"  MF: {mf_pred_mlp.shape}\")\nprint(f\"  BP: {bp_pred_mlp.shape}\")\nprint(f\"  CC: {cc_pred_mlp.shape}\")\n\n\n# IMPORTANT: Make sure you use the FILTERED terms that match your model output dimensions!\nprint(f\"\\nGO term counts:\")\nprint(f\"  MF: {len(mf_terms)} terms\")\nprint(f\"  BP: {len(bp_terms)} terms\")\nprint(f\"  CC: {len(cc_terms)} terms\")\n\n# ============================================\n# CREATE SUBMISSION WITH PROPAGATION\n# ============================================\n\ndef create_cafa_submission_with_propagation(\n    pred_mf, pred_bp, pred_cc,\n    thresh_mf, thresh_bp, thresh_cc,\n    mf_terms, bp_terms, cc_terms,\n    protein_ids, go_dag,\n    max_terms=1500\n):\n    \"\"\"\n    Create CAFA submission with hierarchical propagation and 1500-term limit.\n    \n    Pipeline:\n    1. Apply thresholds to raw predictions\n    2. Propagate to parent terms\n    3. Limit to 1500 terms per protein\n    4. Ensure hierarchical consistency\n    \"\"\"\n    print(f\"\\nProcessing {len(protein_ids):,} proteins...\")\n    \n    submissions = []\n    \n    for i in tqdm(range(len(protein_ids)), desc=\"Creating submission\"):\n        protein_id = protein_ids[i]\n        \n        # ===========================================\n        # STEP 1: Apply thresholds\n        # ===========================================\n        protein_predictions = {}\n        \n        # MF predictions above threshold\n        for j, go_id in enumerate(mf_terms):\n            if pred_mf[i, j] >= thresh_mf:\n                protein_predictions[go_id] = float(pred_mf[i, j])\n        \n        # BP predictions above threshold\n        for j, go_id in enumerate(bp_terms):\n            if pred_bp[i, j] >= thresh_bp:\n                protein_predictions[go_id] = float(pred_bp[i, j])\n        \n        # CC predictions above threshold\n        for j, go_id in enumerate(cc_terms):\n            if pred_cc[i, j] >= thresh_cc:\n                protein_predictions[go_id] = float(pred_cc[i, j])\n        \n        # Skip if no predictions\n        if len(protein_predictions) == 0:\n            continue\n        \n        # ===========================================\n        # STEP 2: Propagate to parents\n        # ===========================================\n        propagated = {}\n        \n        for go_id, score in protein_predictions.items():\n            # Add the term itself\n            propagated[go_id] = score\n            \n            # Add all parent terms\n            if go_id in go_dag:\n                try:\n                    parents = go_dag[go_id].get_all_parents()\n                    for parent in parents:\n                        # Parent gets max score of all its children\n                        propagated[parent] = max(\n                            propagated.get(parent, 0),\n                            score * 0.95  # Slight decay for ancestors\n                        )\n                except:\n                    # Skip if term not in DAG\n                    pass\n        \n        # ===========================================\n        # STEP 3: Apply 1500-term limit\n        # ===========================================\n        if len(propagated) > max_terms:\n            # Sort by score (descending)\n            sorted_terms = sorted(propagated.items(), key=lambda x: x[1], reverse=True)\n            \n            # Keep top max_terms\n            selected = dict(sorted_terms[:max_terms])\n            \n            # ===========================================\n            # STEP 4: Ensure hierarchical consistency\n            # ===========================================\n            # If we keep a child, we MUST keep all its parents\n            must_include = set(selected.keys())\n            \n            for go_id in list(selected.keys()):\n                if go_id in go_dag:\n                    try:\n                        parents = go_dag[go_id].get_all_parents()\n                        for parent in parents:\n                            if parent not in must_include:\n                                must_include.add(parent)\n                                if parent in propagated:\n                                    selected[parent] = propagated[parent]\n                    except:\n                        pass\n            \n            # If adding parents pushed us over max_terms, remove lowest-scoring leaves\n            while len(selected) > max_terms:\n                # Find leaf terms (terms with no children in selected set)\n                leaves = []\n                \n                for go_id in selected.keys():\n                    is_leaf = True\n                    \n                    if go_id in go_dag:\n                        try:\n                            children = go_dag[go_id].children\n                            for child in children:\n                                if child.id in selected:\n                                    is_leaf = False\n                                    break\n                        except:\n                            pass\n                    \n                    if is_leaf:\n                        leaves.append(go_id)\n                \n                if not leaves:\n                    # No leaves found, just remove lowest-scoring term\n                    lowest = min(selected, key=selected.get)\n                    del selected[lowest]\n                else:\n                    # Remove lowest-scoring leaf\n                    lowest_leaf = min(leaves, key=lambda x: selected[x])\n                    del selected[lowest_leaf]\n            \n            propagated = selected\n        \n        # ===========================================\n        # STEP 5: Add to submission\n        # ===========================================\n        for go_id, score in propagated.items():\n            submissions.append({\n                'protein_id': protein_id,\n                'go_id': go_id,\n                'score': score\n            })\n    \n    return submissions\n\n\n# ============================================\n# GENERATE SUBMISSION\n# ============================================\n\nprint(\"\\nGenerating submission...\")\n\nsubmissions = create_cafa_submission_with_propagation(\n    mf_pred_mlp, bp_pred_mlp, cc_pred_mlp,\n    thresh_mf, thresh_bp, thresh_cc,\n    mf_terms, bp_terms, cc_terms,\n    df_test['protein_id'].tolist(),\n    go_dag,\n    max_terms=1500\n)\n\n# ============================================\n# CREATE DATAFRAME AND VERIFY\n# ============================================\n\nsubmission_df = pd.DataFrame(submissions)\n\nprint(f\"\\n{'='*60}\")\nprint(\"SUBMISSION STATISTICS\")\nprint(f\"{'='*60}\")\nprint(f\"Total predictions: {len(submission_df):,}\")\nprint(f\"Unique proteins: {submission_df['protein_id'].nunique():,}\")\nprint(f\"Unique GO terms: {submission_df['go_id'].nunique():,}\")\n\n# Check terms per protein\nterms_per_protein = submission_df.groupby('protein_id').size()\nprint(f\"\\nTerms per protein:\")\nprint(f\"  Mean: {terms_per_protein.mean():.1f}\")\nprint(f\"  Median: {terms_per_protein.median():.1f}\")\nprint(f\"  Min: {terms_per_protein.min()}\")\nprint(f\"  Max: {terms_per_protein.max()}\")\nprint(f\"  Std: {terms_per_protein.std():.1f}\")\n\n# Verify 1500 limit\nover_limit = (terms_per_protein > 1500).sum()\nif over_limit > 0:\n    print(f\"\\n⚠️  WARNING: {over_limit} proteins exceed 1500 term limit!\")\n    print(\"This should not happen. Check the code.\")\nelse:\n    print(f\"\\n✓ All proteins within 1500 term limit\")\n\n# Check score distribution\nprint(f\"\\nScore distribution:\")\nprint(f\"  Mean: {submission_df['score'].mean():.4f}\")\nprint(f\"  Median: {submission_df['score'].median():.4f}\")\nprint(f\"  Min: {submission_df['score'].min():.4f}\")\nprint(f\"  Max: {submission_df['score'].max():.4f}\")\n\n# Check aspect distribution\nmf_count = submission_df[submission_df['go_id'].isin(mf_terms)].shape[0]\nbp_count = submission_df[submission_df['go_id'].isin(bp_terms)].shape[0]\ncc_count = submission_df[submission_df['go_id'].isin(cc_terms)].shape[0]\n\nprint(f\"\\nPredictions by aspect:\")\nprint(f\"  MF: {mf_count:,} ({mf_count/len(submission_df)*100:.1f}%)\")\nprint(f\"  BP: {bp_count:,} ({bp_count/len(submission_df)*100:.1f}%)\")\nprint(f\"  CC: {cc_count:,} ({cc_count/len(submission_df)*100:.1f}%)\")\n\n# ============================================\n# SAVE SUBMISSION\n# ============================================\n\n# CAFA format: protein_id \\t go_id \\t score (no header)\nsubmission_df[['protein_id', 'go_id', 'score']].to_csv(\n    'submission.tsv', \n    sep='\\t', \n    index=False, \n    header=False\n)\n\nprint(f\"\\n{'='*60}\")\nprint(\"✓ Submission saved to 'submission.tsv'\")\nprint(f\"{'='*60}\")\n\n# ============================================\n# SANITY CHECKS\n# ============================================\n\nprint(f\"\\nRunning sanity checks...\")\n\n# Check 1: All proteins have at least one prediction\ntest_proteins = set(df_test['protein_id'].tolist())\nsubmission_proteins = set(submission_df['protein_id'].unique())\nmissing_proteins = test_proteins - submission_proteins\n\nif len(missing_proteins) > 0:\n    print(f\"⚠️  WARNING: {len(missing_proteins)} proteins have no predictions\")\n    print(f\"   First 10: {list(missing_proteins)[:10]}\")\nelse:\n    print(f\"✓ All test proteins have predictions\")\n\n# Check 2: All GO terms are valid\ninvalid_terms = []\nfor go_id in submission_df['go_id'].unique():\n    if go_id not in go_dag:\n        invalid_terms.append(go_id)\n\nif len(invalid_terms) > 0:\n    print(f\"⚠️  WARNING: {len(invalid_terms)} invalid GO terms found\")\n    print(f\"   First 10: {invalid_terms[:10]}\")\nelse:\n    print(f\"✓ All GO terms are valid\")\n\n# Check 3: Scores in valid range\nif (submission_df['score'] < 0).any() or (submission_df['score'] > 1).any():\n    print(f\"⚠️  WARNING: Some scores are outside [0,1] range\")\nelse:\n    print(f\"✓ All scores in valid range [0,1]\")\n\nprint(f\"\\n{'='*60}\")\nprint(\"SUBMISSION COMPLETE!\")\nprint(f\"{'='*60}\")\nprint(f\"\\nYour submission file 'submission.tsv' is ready for upload!\")\nprint(f\"Expected performance (based on validation):\")\nprint(f\"  MF Fmax: ~0.51\")\nprint(f\"  BP Fmax: ~0.35\")\nprint(f\"  CC Fmax: ~0.48\")\nprint(f\"  Average: ~0.45\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T07:26:42.121591Z","iopub.execute_input":"2026-01-21T07:26:42.122122Z","iopub.status.idle":"2026-01-21T07:34:22.375486Z","shell.execute_reply.started":"2026-01-21T07:26:42.122090Z","shell.execute_reply":"2026-01-21T07:34:22.374536Z"}},"outputs":[],"execution_count":null}]}