{"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,"sourceType":"competition"},{"sourceId":14477952,"sourceType":"datasetVersion","datasetId":9247377},{"sourceId":14478028,"sourceType":"datasetVersion","datasetId":9084547}],"dockerImageVersionId":31234,"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-26T04:41:38.499218Z","iopub.execute_input":"2026-01-26T04:41:38.499829Z","iopub.status.idle":"2026-01-26T04:41:38.565564Z","shell.execute_reply.started":"2026-01-26T04:41:38.499796Z","shell.execute_reply":"2026-01-26T04:41:38.564970Z"}},"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-26T03:53:03.789624Z","iopub.execute_input":"2026-01-26T03:53:03.789968Z","iopub.status.idle":"2026-01-26T03:53:46.229928Z","shell.execute_reply.started":"2026-01-26T03:53:03.789942Z","shell.execute_reply":"2026-01-26T03:53:46.229136Z"}},"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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T03:53:46.233459Z","iopub.execute_input":"2026-01-26T03:53:46.233802Z","iopub.status.idle":"2026-01-26T03:53:50.378916Z","shell.execute_reply.started":"2026-01-26T03:53:46.233765Z","shell.execute_reply":"2026-01-26T03:53:50.378167Z"}},"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-26T03:53:50.379981Z","iopub.execute_input":"2026-01-26T03:53:50.380611Z","iopub.status.idle":"2026-01-26T03:53:56.228050Z","shell.execute_reply.started":"2026-01-26T03:53:50.380586Z","shell.execute_reply":"2026-01-26T03:53:56.227201Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_train.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T03:53:56.229229Z","iopub.execute_input":"2026-01-26T03:53:56.229660Z","iopub.status.idle":"2026-01-26T03:53:56.251642Z","shell.execute_reply.started":"2026-01-26T03:53:56.229623Z","shell.execute_reply":"2026-01-26T03:53:56.251029Z"}},"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-26T03:53:56.253649Z","iopub.execute_input":"2026-01-26T03:53:56.253954Z","iopub.status.idle":"2026-01-26T03:53:58.407986Z","shell.execute_reply.started":"2026-01-26T03:53:56.253921Z","shell.execute_reply":"2026-01-26T03:53:58.407224Z"}},"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-26T03:53:58.409143Z","iopub.execute_input":"2026-01-26T03:53:58.409469Z","iopub.status.idle":"2026-01-26T03:54:05.576768Z","shell.execute_reply.started":"2026-01-26T03:53:58.409421Z","shell.execute_reply":"2026-01-26T03:54:05.576097Z"}},"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-26T03:54:05.577722Z","iopub.execute_input":"2026-01-26T03:54:05.577990Z","iopub.status.idle":"2026-01-26T03:54:14.717733Z","shell.execute_reply.started":"2026-01-26T03:54:05.577967Z","shell.execute_reply":"2026-01-26T03:54:14.717021Z"}},"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-26T03:54:14.718596Z","iopub.execute_input":"2026-01-26T03:54:14.719049Z","iopub.status.idle":"2026-01-26T03:54:14.731340Z","shell.execute_reply.started":"2026-01-26T03:54:14.719024Z","shell.execute_reply":"2026-01-26T03:54:14.730755Z"}},"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-26T03:54:14.732282Z","iopub.execute_input":"2026-01-26T03:54:14.732587Z","iopub.status.idle":"2026-01-26T03:54:26.479207Z","shell.execute_reply.started":"2026-01-26T03:54:14.732556Z","shell.execute_reply":"2026-01-26T03:54:26.478372Z"}},"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-26T03:54:26.480203Z","iopub.execute_input":"2026-01-26T03:54:26.480484Z","iopub.status.idle":"2026-01-26T03:54:27.758529Z","shell.execute_reply.started":"2026-01-26T03:54:26.480452Z","shell.execute_reply":"2026-01-26T03:54:27.757711Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load the protT5 embeddings\ntrain_path = \"/kaggle/input/train-embeddings-prott5-xl-half/train_prott5_xl_half_embeddings.npy\"\ntest_path = \"/kaggle/input/prott5-embeddings-test-npy/test_prott5_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-26T03:54:27.759875Z","iopub.execute_input":"2026-01-26T03:54:27.760237Z","iopub.status.idle":"2026-01-26T03:54:35.734638Z","shell.execute_reply.started":"2026-01-26T03:54:27.760197Z","shell.execute_reply":"2026-01-26T03:54:35.733872Z"}},"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.npy\", X_train)\nnp.save(\"X_test_protT5.npy\", X_test)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T03:54:35.735645Z","iopub.execute_input":"2026-01-26T03:54:35.736050Z","iopub.status.idle":"2026-01-26T03:54:36.800447Z","shell.execute_reply.started":"2026-01-26T03:54:35.736015Z","shell.execute_reply":"2026-01-26T03:54:36.799668Z"}},"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-26T03:54:36.801449Z","iopub.execute_input":"2026-01-26T03:54:36.801784Z","iopub.status.idle":"2026-01-26T03:54:37.914215Z","shell.execute_reply.started":"2026-01-26T03:54:36.801742Z","shell.execute_reply":"2026-01-26T03:54:37.913462Z"}},"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-26T03:54:37.915164Z","iopub.execute_input":"2026-01-26T03:54:37.915568Z","iopub.status.idle":"2026-01-26T03:54:44.970200Z","shell.execute_reply.started":"2026-01-26T03:54:37.915535Z","shell.execute_reply":"2026-01-26T03:54:44.969593Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"np.save(\"X_tr\", X_tr)\nnp.save(\"X_val\", X_val)\n\nnp.save(\"Y_mf_tr\", Y_mf_tr)\nnp.save(\"Y_mf_val\",Y_mf_val)\nnp.save(\"Y_bp_tr\", Y_bp_tr)\nnp.save(\"Y_bp_val\",Y_bp_val)\nnp.save(\"Y_cc_tr\", Y_cc_tr)\nnp.save(\"Y_cc_val\",Y_cc_val)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T03:54:44.971147Z","iopub.execute_input":"2026-01-26T03:54:44.971639Z","iopub.status.idle":"2026-01-26T03:54:48.234982Z","shell.execute_reply.started":"2026-01-26T03:54:44.971601Z","shell.execute_reply":"2026-01-26T03:54:48.234189Z"}},"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-26T03:54:48.235939Z","iopub.execute_input":"2026-01-26T03:54:48.236223Z","iopub.status.idle":"2026-01-26T03:54:48.260111Z","shell.execute_reply.started":"2026-01-26T03:54:48.236188Z","shell.execute_reply":"2026-01-26T03:54:48.259340Z"}},"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            \n                      \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            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.3),\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-26T03:54:48.261233Z","iopub.execute_input":"2026-01-26T03:54:48.261547Z","iopub.status.idle":"2026-01-26T03:54:59.564870Z","shell.execute_reply.started":"2026-01-26T03:54:48.261497Z","shell.execute_reply":"2026-01-26T03:54:59.564222Z"}},"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=4096, \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-26T03:54:59.565808Z","iopub.execute_input":"2026-01-26T03:54:59.566082Z","iopub.status.idle":"2026-01-26T03:55:01.457364Z","shell.execute_reply.started":"2026-01-26T03:54:59.566060Z","shell.execute_reply":"2026-01-26T03:55:01.456588Z"}},"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-26T03:55:01.458418Z","iopub.execute_input":"2026-01-26T03:55:01.459123Z","iopub.status.idle":"2026-01-26T03:55:02.658396Z","shell.execute_reply.started":"2026-01-26T03:55:01.459098Z","shell.execute_reply":"2026-01-26T03:55:02.657813Z"}},"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-26T03:55:02.661815Z","iopub.execute_input":"2026-01-26T03:55:02.662040Z","iopub.status.idle":"2026-01-26T03:55:02.677210Z","shell.execute_reply.started":"2026-01-26T03:55:02.662020Z","shell.execute_reply":"2026-01-26T03:55:02.676475Z"}},"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=50,\n    batch_size=1024,\n    lr=1e-3,\n    patience=10\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, prec_mf, rec_mf = compute_fmax_with_details(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} \"\n        f\"(P={prec_mf:.3f}, R={rec_mf:.3f})\")\nprint(f\"{'='*60}\")\n    \n# Train BP model\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=4096,\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, prec_bp, rec_bp = compute_fmax_with_details(pred_bp, Y_bp_val if isinstance(Y_bp_val, np.ndarray) else Y_bp_val.numpy())\nprint(f\"\\n{'='*60}\")\nprint(f\"BP Results: Fmax={fmax_bp:.4f} @ t={thresh_bp:.3f} \"\n        f\"(P={prec_bp:.3f}, R={rec_bp:.3f})\")\nprint(f\"{'='*60}\")\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, prec_cc, rec_cc = compute_fmax_with_details(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} \"\n        f\"(P={prec_cc:.3f}, R={rec_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    \n    ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T03:55:02.678005Z","iopub.execute_input":"2026-01-26T03:55:02.678268Z","iopub.status.idle":"2026-01-26T04:00:04.677610Z","shell.execute_reply.started":"2026-01-26T03:55:02.678244Z","shell.execute_reply":"2026-01-26T04:00:04.676927Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport numpy as np\nfrom tqdm import tqdm\nimport pickle\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-26T04:00:04.678522Z","iopub.execute_input":"2026-01-26T04:00:04.679014Z","iopub.status.idle":"2026-01-26T04:00:56.085539Z","shell.execute_reply.started":"2026-01-26T04:00:04.678988Z","shell.execute_reply":"2026-01-26T04:00:56.084797Z"}},"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())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T04:00:56.086470Z","iopub.execute_input":"2026-01-26T04:00:56.086795Z","iopub.status.idle":"2026-01-26T04:00:56.154011Z","shell.execute_reply.started":"2026-01-26T04:00:56.086770Z","shell.execute_reply":"2026-01-26T04:00:56.153419Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport numpy as np\nfrom tqdm import tqdm","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T04:00:56.154804Z","iopub.execute_input":"2026-01-26T04:00:56.155122Z","iopub.status.idle":"2026-01-26T04:00:56.158539Z","shell.execute_reply.started":"2026-01-26T04:00:56.155101Z","shell.execute_reply":"2026-01-26T04:00:56.157801Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================\n# STEP 1: TRAIN MULTIPLE MODELS WITH DIFFERENT SEEDS\n# ============================================\n\ndef set_seed(seed):\n    \"\"\"Set random seed for reproducibility\"\"\"\n    import random\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    if torch.cuda.is_available():\n        torch.cuda.manual_seed_all(seed)\n        torch.backends.cudnn.deterministic = True\n        torch.backends.cudnn.benchmark = False\n\n\ndef train_ensemble_models(\n    model_class,\n    X_train, Y_train,\n    X_val, Y_val,\n    aspect_name=\"\",\n    n_models=5,\n    seeds=None,\n    **training_kwargs  # All other train_aspect parameters\n):\n    \"\"\"\n    Train multiple models with different random seeds.\n    \n    Args:\n        model_class: Model class (MF_MLP, BP_MLP, CC_MLP)\n        X_train, Y_train: Training data\n        X_val, Y_val: Validation data\n        aspect_name: Name for logging\n        n_models: Number of models to train\n        seeds: List of seeds (if None, use [0, 1, 2, 3, 4])\n        **training_kwargs: All parameters for train_aspect()\n    \n    Returns:\n        List of trained models\n    \"\"\"\n    if seeds is None:\n        seeds = list(range(n_models))\n    \n    print(f\"\\n{'='*60}\")\n    print(f\"TRAINING ENSEMBLE: {n_models} {aspect_name} MODELS\")\n    print(f\"{'='*60}\")\n    \n    models = []\n    \n    for i, seed in enumerate(seeds):\n        print(f\"\\n--- Model {i+1}/{n_models} (seed={seed}) ---\")\n        \n        # Set seed\n        set_seed(seed)\n        \n        # Train model\n        model = train_aspect(\n            model_class,\n            X_train, Y_train,\n            X_val, Y_val,\n            aspect_name=f\"{aspect_name}_seed{seed}\",\n            **training_kwargs\n        )\n        \n        models.append(model)\n        \n        # Save individual model\n        torch.save(model.state_dict(), f\"model_{aspect_name.lower()}_seed{seed}.pt\")\n        print(f\"✓ Model saved: model_{aspect_name.lower()}_seed{seed}.pt\")\n    \n    print(f\"\\n✓ Ensemble training complete: {n_models} models\")\n    return models\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T04:00:56.159486Z","iopub.execute_input":"2026-01-26T04:00:56.159855Z","iopub.status.idle":"2026-01-26T04:00:56.178679Z","shell.execute_reply.started":"2026-01-26T04:00:56.159823Z","shell.execute_reply":"2026-01-26T04:00:56.178053Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================\n# STEP 2: ENSEMBLE PREDICTION STRATEGIES\n# ============================================\n\ndef ensemble_predict_average(models, X, device='cuda'):\n    \"\"\"\n    Simple averaging: Average predictions from all models.\n    \n    This is the most common and often best-performing ensemble method.\n    \"\"\"\n    X_tensor = torch.tensor(X, dtype=torch.float32).to(device)\n    \n    predictions = []\n    \n    for model in models:\n        model.eval()\n        with torch.no_grad():\n            logits = model(X_tensor)\n            preds = torch.sigmoid(logits).cpu().numpy()\n            predictions.append(preds)\n    \n    # Simple average\n    ensemble_pred = np.mean(predictions, axis=0)\n    \n    return ensemble_pred","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T04:00:56.179455Z","iopub.execute_input":"2026-01-26T04:00:56.179737Z","iopub.status.idle":"2026-01-26T04:00:56.196096Z","shell.execute_reply.started":"2026-01-26T04:00:56.179706Z","shell.execute_reply":"2026-01-26T04:00:56.195361Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def ensemble_predict_weighted(models, X, weights, device='cuda'):\n    \"\"\"\n    Weighted averaging: Weight models by their validation performance.\n    \n    Args:\n        models: List of models\n        X: Input data\n        weights: List of weights (should sum to 1.0)\n        device: Device to use\n    \n    Returns:\n        Weighted ensemble predictions\n    \"\"\"\n    X_tensor = torch.tensor(X, dtype=torch.float32).to(device)\n    \n    predictions = []\n    \n    for model in models:\n        model.eval()\n        with torch.no_grad():\n            logits = model(X_tensor)\n            preds = torch.sigmoid(logits).cpu().numpy()\n            predictions.append(preds)\n    \n    # Weighted average\n    ensemble_pred = np.zeros_like(predictions[0])\n    for pred, weight in zip(predictions, weights):\n        ensemble_pred += pred * weight\n    \n    return ensemble_pred\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T04:00:56.196926Z","iopub.execute_input":"2026-01-26T04:00:56.197184Z","iopub.status.idle":"2026-01-26T04:00:56.210031Z","shell.execute_reply.started":"2026-01-26T04:00:56.197164Z","shell.execute_reply":"2026-01-26T04:00:56.209412Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def ensemble_predict_rank_average(models, X, device='cuda'):\n    \"\"\"\n    Rank averaging: Convert predictions to ranks, then average.\n    \n    More robust to different prediction scales across models.\n    \"\"\"\n    X_tensor = torch.tensor(X, dtype=torch.float32).to(device)\n    \n    predictions = []\n    \n    for model in models:\n        model.eval()\n        with torch.no_grad():\n            logits = model(X_tensor)\n            preds = torch.sigmoid(logits).cpu().numpy()\n            predictions.append(preds)\n    \n    # Convert to ranks (per protein)\n    ranked_predictions = []\n    for pred in predictions:\n        ranked = np.zeros_like(pred)\n        for i in range(pred.shape[0]):\n            # Rank predictions for this protein\n            ranks = pred[i].argsort().argsort()\n            # Normalize to [0, 1]\n            ranked[i] = ranks / (pred.shape[1] - 1)\n        ranked_predictions.append(ranked)\n    \n    # Average ranks\n    ensemble_pred = np.mean(ranked_predictions, axis=0)\n    \n    return ensemble_pred\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T04:00:56.210856Z","iopub.execute_input":"2026-01-26T04:00:56.211127Z","iopub.status.idle":"2026-01-26T04:00:56.225268Z","shell.execute_reply.started":"2026-01-26T04:00:56.211106Z","shell.execute_reply":"2026-01-26T04:00:56.224525Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# ============================================\n# STEP 3: COMPUTE ENSEMBLE WEIGHTS FROM VALIDATION\n# ============================================\n\ndef compute_ensemble_weights(models, X_val, Y_val, device='cuda'):\n    \"\"\"\n    Compute optimal weights based on validation Fmax.\n    \n    Args:\n        models: List of models\n        X_val: Validation input\n        Y_val: Validation labels\n        device: Device\n    \n    Returns:\n        List of weights (sum to 1.0)\n    \"\"\"\n    print(\"\\nComputing ensemble weights from validation performance...\")\n    \n    fmax_scores = []\n    \n    for i, model in enumerate(models):\n        model.eval()\n        X_tensor = torch.tensor(X_val, dtype=torch.float32).to(device)\n        \n        with torch.no_grad():\n            logits = model(X_tensor)\n            preds = torch.sigmoid(logits).cpu().numpy()\n        \n        # Compute Fmax\n        fmax, _ = compute_fmax(preds, Y_val)\n        fmax_scores.append(fmax)\n        print(f\"  Model {i+1}: Fmax = {fmax:.4f}\")\n    \n    # Convert to weights (higher Fmax = higher weight)\n    # Use softmax to normalize\n    fmax_array = np.array(fmax_scores)\n    weights = np.exp(fmax_array * 10)  # Scale up differences\n    weights = weights / weights.sum()  # Normalize\n    \n    print(f\"\\nEnsemble weights:\")\n    for i, (fmax, weight) in enumerate(zip(fmax_scores, weights)):\n        print(f\"  Model {i+1}: Fmax={fmax:.4f}, Weight={weight:.3f}\")\n    \n    return weights.tolist()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T04:00:56.226160Z","iopub.execute_input":"2026-01-26T04:00:56.226375Z","iopub.status.idle":"2026-01-26T04:00:56.241363Z","shell.execute_reply.started":"2026-01-26T04:00:56.226355Z","shell.execute_reply":"2026-01-26T04:00:56.240678Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nn_models = 5  # Train 5 models per aspect\nseeds = [42, 123, 456, 789, 2024]  # Different random seeds\n\n# TRAIN ENSEMBLE FOR EACH ASPECT\nprint(\"\\n\" + \"=\"*60)\nprint(\"TRAINING MF ENSEMBLE\")\nprint(\"\\n\" + \"=\"*60)\n\nmf_models = train_ensemble_models(\n    MF_MLP,\n    X_tr, Y_mf_tr,\n    X_val, Y_mf_val,\n    aspect_name=\"MF\",\n    n_models=n_models,\n    seeds=seeds,\n    pos_weight_alpha=0,\n    epochs=50,\n    batch_size=1024,\n    lr=1e-3,\n    patience=10\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T04:00:56.242128Z","iopub.execute_input":"2026-01-26T04:00:56.242390Z","iopub.status.idle":"2026-01-26T04:04:10.189482Z","shell.execute_reply.started":"2026-01-26T04:00:56.242361Z","shell.execute_reply":"2026-01-26T04:04:10.188690Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"bp_models = train_ensemble_models(\n    BP_MLP,\n    X_tr, Y_bp_tr,\n    X_val, Y_bp_val,\n    aspect_name=\"BP\",\n    n_models=n_models,\n    seeds=seeds,\n    pos_weight_alpha=0,\n    epochs=50,\n    batch_size=1024,\n    lr=1e-3,\n    patience=10\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T04:04:10.190572Z","iopub.execute_input":"2026-01-26T04:04:10.190963Z","iopub.status.idle":"2026-01-26T04:10:16.977813Z","shell.execute_reply.started":"2026-01-26T04:04:10.190930Z","shell.execute_reply":"2026-01-26T04:10:16.977177Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"cc_models = train_ensemble_models(\n    CC_MLP,\n    X_tr, Y_cc_tr,\n    X_val, Y_cc_val,\n    aspect_name=\"CC\",\n    n_models=n_models,\n    seeds=seeds,\n    pos_weight_alpha=0,\n    epochs=50,\n    batch_size=1024,\n    lr=1e-3,\n    patience=10\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T04:10:16.978701Z","iopub.execute_input":"2026-01-26T04:10:16.978939Z","iopub.status.idle":"2026-01-26T04:13:03.752385Z","shell.execute_reply.started":"2026-01-26T04:10:16.978917Z","shell.execute_reply":"2026-01-26T04:13:03.751661Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ========================================\n# EVALUATE ENSEMBLE ON VALIDATION\n# ========================================\n    \nprint(\"\\n\" + \"=\"*60)\nprint(\"EVALUATING ENSEMBLE METHODS\")\nprint(\"=\"*60)\n    \n# Convert validation data\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# Method 1: Simple Average\nprint(\"\\n--- Method 1: Simple Average ---\")\n    \nmf_pred_avg = ensemble_predict_average(mf_models, X_val)\nbp_pred_avg = ensemble_predict_average(bp_models, X_val)\ncc_pred_avg = ensemble_predict_average(cc_models, X_val)\n    \nfmax_mf_avg, _ = compute_fmax(mf_pred_avg, Y_mf_val_np)\nfmax_bp_avg, _ = compute_fmax(bp_pred_avg, Y_bp_val_np)\nfmax_cc_avg, _ = compute_fmax(cc_pred_avg, Y_cc_val_np)\n\n# Method 2: Weighted Average\nprint(\"\\n--- Method 2: Weighted Average ---\")\n    \nweights_mf = compute_ensemble_weights(mf_models, X_val, Y_mf_val_np)\nweights_bp = compute_ensemble_weights(bp_models, X_val, Y_bp_val_np)\nweights_cc = compute_ensemble_weights(cc_models, X_val, Y_cc_val_np)\n    \nmf_pred_weighted = ensemble_predict_weighted(mf_models, X_val, weights_mf)\nbp_pred_weighted = ensemble_predict_weighted(bp_models, X_val, weights_bp)\ncc_pred_weighted = ensemble_predict_weighted(cc_models, X_val, weights_cc)\n    \nfmax_mf_weighted, _ = compute_fmax(mf_pred_weighted, Y_mf_val_np)\nfmax_bp_weighted, _ = compute_fmax(bp_pred_weighted, Y_bp_val_np)\nfmax_cc_weighted, _ = compute_fmax(cc_pred_weighted, Y_cc_val_np)\n    \nprint(f\"MF Fmax: {fmax_mf_weighted:.4f}\")\nprint(f\"BP Fmax: {fmax_bp_weighted:.4f}\")\nprint(f\"CC Fmax: {fmax_cc_weighted:.4f}\")\nprint(f\"Average: {(fmax_mf_weighted + fmax_bp_weighted + fmax_cc_weighted)/3:.4f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T04:13:03.753338Z","iopub.execute_input":"2026-01-26T04:13:03.753642Z","iopub.status.idle":"2026-01-26T04:18:54.839498Z","shell.execute_reply.started":"2026-01-26T04:13:03.753618Z","shell.execute_reply":"2026-01-26T04:18:54.838739Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ========================================\n# SAVE BEST ENSEMBLE METHOD\n# ========================================\n    \navg_fmax = (fmax_mf_avg + fmax_bp_avg + fmax_cc_avg) / 3\nweighted_fmax = (fmax_mf_weighted + fmax_bp_weighted + fmax_cc_weighted) / 3\n    \nprint(f\"\\n{'='*60}\")\nprint(\"ENSEMBLE COMPARISON\")\nprint(f\"{'='*60}\")\nprint(f\"Simple Average:  {avg_fmax:.4f}\")\nprint(f\"Weighted Average: {weighted_fmax:.4f}\")\n    \nif weighted_fmax > avg_fmax:\n    print(\"\\n✓ Weighted average performs better! Using weighted ensemble.\")\n    best_method = \"weighted\"\n    # Save weights\n    ensemble_weights = {\n        'MF': weights_mf,\n        'BP': weights_bp,\n        'CC': weights_cc\n    }\nelse:\n    print(\"\\n✓ Simple average performs better! Using simple ensemble.\")\n    best_method = \"average\"\n    ensemble_weights = None\n    \n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T04:18:54.840992Z","iopub.execute_input":"2026-01-26T04:18:54.841327Z","iopub.status.idle":"2026-01-26T04:18:54.847223Z","shell.execute_reply.started":"2026-01-26T04:18:54.841298Z","shell.execute_reply":"2026-01-26T04:18:54.846564Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Compute Fmax and thresholds\nprint(\"Computing Fmax and optimal thresholds...\")\nfmax_mf, thresh_mf, prec_mf, rec_mf = compute_fmax_with_details(mf_pred_avg, Y_mf_val_np)\nfmax_bp, thresh_bp, prec_bp, rec_bp = compute_fmax_with_details(bp_pred_avg, Y_bp_val_np)\nfmax_cc, thresh_cc, prec_cc, rec_cc = compute_fmax_with_details(cc_pred_avg, Y_cc_val_np)\n\n# Print results\nprint(f\"\\n{'='*60}\")\nprint(\"ENSEMBLE VALIDATION RESULTS\")\nprint(f\"{'='*60}\")\nprint(f\"\\nMF: Fmax={fmax_mf:.4f}, Threshold={thresh_mf:.3f}, P={prec_mf:.3f}, R={rec_mf:.3f}\")\nprint(f\"BP: Fmax={fmax_bp:.4f}, Threshold={thresh_bp:.3f}, P={prec_bp:.3f}, R={rec_bp:.3f}\")\nprint(f\"CC: Fmax={fmax_cc:.4f}, Threshold={thresh_cc:.3f}, P={prec_cc:.3f}, R={rec_cc:.3f}\")\nprint(f\"\\nAverage Fmax: {(fmax_mf + fmax_bp + fmax_cc)/3:.4f}\")\n\n# Save thresholds\nensemble_info = {\n    'thresholds': {\n        'MF': thresh_mf,\n        'BP': thresh_bp,\n        'CC': thresh_cc\n    },\n    'fmax_scores': {\n        'MF': fmax_mf,\n        'BP': fmax_bp,\n        'CC': fmax_cc,\n        'average': (fmax_mf + fmax_bp + fmax_cc) / 3\n    },\n    'n_models': n_models,\n    'seeds': seeds\n}\n\nwith open('ensemble_info.pkl', 'wb') as f:\n    pickle.dump(ensemble_info, f)\n\nprint(\"\\n✓ Saved ensemble_info.pkl\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T04:18:54.848330Z","iopub.execute_input":"2026-01-26T04:18:54.848840Z","iopub.status.idle":"2026-01-26T04:19:44.014752Z","shell.execute_reply.started":"2026-01-26T04:18:54.848818Z","shell.execute_reply":"2026-01-26T04:19:44.014093Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport pickle\nfrom tqdm import tqdm\n\n# ============================================\n# STEP 1: GENERATE ENSEMBLE TEST PREDICTIONS\n# ============================================\n\n# ============================================\n# MEMORY-EFFICIENT ENSEMBLE PREDICTION\n# ============================================\n\nprint(\"=\"*60)\nprint(\"GENERATING ENSEMBLE TEST PREDICTIONS (Memory-Efficient)\")\nprint(\"=\"*60)\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nX_test_tensor = torch.tensor(X_test, dtype=torch.float32).to(device)\n\nn_models = 5\nseeds = [42, 123, 456, 789, 2024]\n\n# --- MF Ensemble (Running Average) ---\nprint(\"\\nMF ensemble predictions...\")\ntest_pred_mf = None\n\nfor i, seed in enumerate(seeds):\n    print(f\"  Loading model {i+1}/{n_models} (seed={seed})...\")\n    \n    # Load model\n    model = MF_MLP(1024, len(mf_terms)).to(device)\n    model.load_state_dict(torch.load(f\"model_mf_seed{seed}.pt\", map_location=device))\n    model.eval()\n    \n    # Generate predictions\n    with torch.no_grad():\n        logits = model(X_test_tensor)\n        preds = torch.sigmoid(logits).cpu().numpy()\n    \n    # Running average\n    if test_pred_mf is None:\n        test_pred_mf = preds / n_models  # First model\n    else:\n        test_pred_mf += preds / n_models  # Add subsequent models\n    \n    # Delete model to free memory\n    del model\n    torch.cuda.empty_cache()  # Clear GPU cache\n    \n    print(f\"  ✓ Model {i+1} processed\")\n\nprint(f\"✓ MF ensemble complete: {test_pred_mf.shape}\")\n\n# --- BP Ensemble (Running Average) ---\nprint(\"\\nBP ensemble predictions...\")\ntest_pred_bp = None\n\nfor i, seed in enumerate(seeds):\n    print(f\"  Loading model {i+1}/{n_models} (seed={seed})...\")\n    \n    model = BP_MLP(1024, len(bp_terms)).to(device)\n    model.load_state_dict(torch.load(f\"model_bp_seed{seed}.pt\", map_location=device))\n    model.eval()\n    \n    with torch.no_grad():\n        logits = model(X_test_tensor)\n        preds = torch.sigmoid(logits).cpu().numpy()\n    \n    if test_pred_bp is None:\n        test_pred_bp = preds / n_models\n    else:\n        test_pred_bp += preds / n_models\n    \n    del model\n    torch.cuda.empty_cache()\n    \n    print(f\"  ✓ Model {i+1} processed\")\n\nprint(f\"✓ BP ensemble complete: {test_pred_bp.shape}\")\n\n# --- CC Ensemble (Running Average) ---\nprint(\"\\nCC ensemble predictions...\")\ntest_pred_cc = None\n\nfor i, seed in enumerate(seeds):\n    print(f\"  Loading model {i+1}/{n_models} (seed={seed})...\")\n    \n    model = CC_MLP(1024, len(cc_terms)).to(device)\n    model.load_state_dict(torch.load(f\"model_cc_seed{seed}.pt\", map_location=device))\n    model.eval()\n    \n    with torch.no_grad():\n        logits = model(X_test_tensor)\n        preds = torch.sigmoid(logits).cpu().numpy()\n    \n    if test_pred_cc is None:\n        test_pred_cc = preds / n_models\n    else:\n        test_pred_cc += preds / n_models\n    \n    del model\n    torch.cuda.empty_cache()\n    \n    print(f\"  ✓ Model {i+1} processed\")\n\nprint(f\"✓ CC ensemble complete: {test_pred_cc.shape}\")\n\n# Save ensemble predictions\nnp.save('ensemble_test_pred_mf.npy', test_pred_mf)\nnp.save('ensemble_test_pred_bp.npy', test_pred_bp)\nnp.save('ensemble_test_pred_cc.npy', test_pred_cc)\n\nprint(f\"\\n{'='*60}\")\nprint(\"✓ ALL ENSEMBLE PREDICTIONS SAVED\")\nprint(f\"{'='*60}\")\nprint(f\"  ensemble_test_pred_mf.npy: {test_pred_mf.shape}\")\nprint(f\"  ensemble_test_pred_bp.npy: {test_pred_bp.shape}\")\nprint(f\"  ensemble_test_pred_cc.npy: {test_pred_cc.shape}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T04:19:44.015640Z","iopub.execute_input":"2026-01-26T04:19:44.015918Z","iopub.status.idle":"2026-01-26T04:20:35.518324Z","shell.execute_reply.started":"2026-01-26T04:19:44.015894Z","shell.execute_reply":"2026-01-26T04:20:35.517531Z"}},"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-26T04:20:35.519375Z","iopub.execute_input":"2026-01-26T04:20:35.519686Z","iopub.status.idle":"2026-01-26T04:20:37.048278Z","shell.execute_reply.started":"2026-01-26T04:20:35.519661Z","shell.execute_reply":"2026-01-26T04:20:37.047678Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================\n# STEP 2: LOAD OPTIMAL THRESHOLDS\n# ============================================\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"LOADING OPTIMAL THRESHOLDS\")\nprint(\"=\"*60)\n\n# These were computed from validation set during ensemble training\n# You should have saved them like this:\n# with open('ensemble_thresholds.pkl', 'wb') as f:\n#     pickle.dump({'MF': thresh_mf, 'BP': thresh_bp, 'CC': thresh_cc}, f)\n\nwith open('/kaggle/working/ensemble_info.pkl', 'rb') as f:\n    ensemble_info = pickle.load(f)\n    optimal_thresholds = ensemble_info[\"thresholds\"]\n\nthresh_mf = optimal_thresholds['MF']\nthresh_bp = optimal_thresholds['BP']\nthresh_cc = optimal_thresholds['CC']\n\nprint(f\"\\nOptimal thresholds:\")\nprint(f\"  MF: {thresh_mf:.3f}\")\nprint(f\"  BP: {thresh_bp:.3f}\")\nprint(f\"  CC: {thresh_cc:.3f}\")\n\n# ============================================\n# STEP 3: CREATE SUBMISSION WITH THRESHOLDS + 1500 LIMIT\n# ============================================\n\ndef create_final_submission(\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    min_predictions=5\n):\n    \"\"\"\n    Create CAFA submission with:\n    1. Thresholding\n    2. Hierarchical propagation\n    3. 1500-term limit per protein\n    4. Fallback for proteins with no predictions\n    \"\"\"\n    print(f\"\\n{'='*60}\")\n    print(\"CREATING FINAL SUBMISSION\")\n    print(f\"{'='*60}\")\n    print(f\"\\nProcessing {len(protein_ids):,} proteins...\")\n    \n    submissions = []\n    proteins_with_fallback = 0\n    \n    for i in tqdm(range(len(protein_ids)), desc=\"Creating submission\"):\n        protein_id = protein_ids[i]\n        \n        # ==========================================\n        # STEP 3.1: Apply thresholds\n        # ==========================================\n        protein_predictions = {}\n        \n        # MF terms 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 terms 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 terms 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        # ==========================================\n        # FALLBACK: If no predictions, use top-N\n        # ==========================================\n        if len(protein_predictions) == 0:\n            proteins_with_fallback += 1\n            \n            # Collect all predictions\n            all_scores = []\n            for j, go_id in enumerate(mf_terms):\n                all_scores.append((go_id, float(pred_mf[i, j])))\n            for j, go_id in enumerate(bp_terms):\n                all_scores.append((go_id, float(pred_bp[i, j])))\n            for j, go_id in enumerate(cc_terms):\n                all_scores.append((go_id, float(pred_cc[i, j])))\n            \n            # Take top min_predictions\n            all_scores.sort(key=lambda x: x[1], reverse=True)\n            for go_id, score in all_scores[:min_predictions]:\n                protein_predictions[go_id] = score\n        \n        # ==========================================\n        # STEP 3.2: Hierarchical propagation\n        # ==========================================\n        propagated = {}\n        \n        for go_id, score in protein_predictions.items():\n            # Add term itself\n            propagated[go_id] = score\n            \n            # Add all parents\n            if go_id in go_dag:\n                try:\n                    parents = go_dag[go_id].get_all_parents()\n                    for parent in parents:\n                        propagated[parent] = max(\n                            propagated.get(parent, 0),\n                            score * 0.95  # Slight decay\n                        )\n                except:\n                    pass\n        \n        # ==========================================\n        # STEP 3.3: Apply 1500-term limit\n        # ==========================================\n        if len(propagated) > max_terms:\n            # Sort by score, keep top max_terms\n            sorted_terms = sorted(propagated.items(), key=lambda x: x[1], reverse=True)\n            selected = dict(sorted_terms[:max_terms])\n            \n            # Ensure hierarchical consistency\n            must_include = set(selected.keys())\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 and parent in propagated:\n                                must_include.add(parent)\n                                selected[parent] = propagated[parent]\n                    except:\n                        pass\n            \n            # Remove lowest-scoring leaves if over limit\n            while len(selected) > max_terms:\n                leaves = []\n                for go_id in selected.keys():\n                    is_leaf = True\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                    if is_leaf:\n                        leaves.append(go_id)\n                \n                if not leaves:\n                    lowest = min(selected, key=selected.get)\n                    del selected[lowest]\n                else:\n                    lowest_leaf = min(leaves, key=lambda x: selected[x])\n                    del selected[lowest_leaf]\n            \n            propagated = selected\n        \n        # ==========================================\n        # STEP 3.4: 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    if proteins_with_fallback > 0:\n        print(f\"\\n⚠️  {proteins_with_fallback} proteins used fallback predictions\")\n    \n    return submissions\n\n\n# ============================================\n# GENERATE SUBMISSION\n# ============================================\n\nsubmissions = create_final_submission(\n    test_pred_mf, test_pred_bp, test_pred_cc,\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    min_predictions=5\n)\n\n# ============================================\n# CREATE SUBMISSION DATAFRAME\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# Verify 1500 limit\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()}\")\n\nassert terms_per_protein.max() <= 1500, \"ERROR: Some proteins exceed 1500 terms!\"\nprint(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# ============================================\n# SAVE SUBMISSION FILE\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: submission.tsv\")\nprint(f\"{'='*60}\")\n\n# ============================================\n# SANITY CHECKS\n# ============================================\n\nprint(f\"\\nRunning sanity checks...\")\n\n# Check 1: Coverage\ntest_proteins = set(df_test['protein_id'].tolist())\nsubmission_proteins = set(submission_df['protein_id'].unique())\nmissing = test_proteins - submission_proteins\n\nif len(missing) > 0:\n    print(f\"⚠️  {len(missing)} proteins missing from submission\")\nelse:\n    print(f\"✓ All {len(test_proteins)} test proteins have predictions\")\n\n# Check 2: Valid GO terms\ninvalid_terms = [t for t in submission_df['go_id'].unique() if t not in go_dag]\nif len(invalid_terms) > 0:\n    print(f\"⚠️  {len(invalid_terms)} invalid GO terms\")\nelse:\n    print(f\"✓ All GO terms are valid\")\n\n# Check 3: Score range\nif (submission_df['score'] < 0).any() or (submission_df['score'] > 1).any():\n    print(f\"⚠️  Some scores outside [0, 1] range\")\nelse:\n    print(f\"✓ All scores in valid range\")\n\nprint(f\"\\n{'='*60}\")\nprint(\"ENSEMBLE SUBMISSION COMPLETE!\")\nprint(f\"{'='*60}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T04:42:25.106319Z","iopub.execute_input":"2026-01-26T04:42:25.106703Z","iopub.status.idle":"2026-01-26T04:49:25.820333Z","shell.execute_reply.started":"2026-01-26T04:42:25.106675Z","shell.execute_reply":"2026-01-26T04:49:25.819581Z"}},"outputs":[],"execution_count":null}]}