{"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":14053785,"sourceType":"datasetVersion","datasetId":8945796},{"sourceId":14239152,"sourceType":"datasetVersion","datasetId":9084547},{"sourceId":14297097,"sourceType":"datasetVersion","datasetId":9126339},{"sourceId":14394957,"sourceType":"datasetVersion","datasetId":9118404}],"dockerImageVersionId":31236,"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-07T06:09:37.691947Z","iopub.execute_input":"2026-01-07T06:09:37.692823Z","iopub.status.idle":"2026-01-07T06:09:37.729214Z","shell.execute_reply.started":"2026-01-07T06:09:37.692788Z","shell.execute_reply":"2026-01-07T06:09:37.728458Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install Biopython --q\n!pip install goatools --q","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:37.730423Z","iopub.execute_input":"2026-01-07T06:09:37.730661Z","iopub.status.idle":"2026-01-07T06:09:46.327822Z","shell.execute_reply.started":"2026-01-07T06:09:37.730629Z","shell.execute_reply":"2026-01-07T06:09:46.327020Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# import libraries\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nimport matplotlib.pyplot as plt\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader, TensorDataset\nfrom scipy.sparse import csr_matrix\nfrom collections import defaultdict\nfrom Bio import SeqIO\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:46.329090Z","iopub.execute_input":"2026-01-07T06:09:46.329456Z","iopub.status.idle":"2026-01-07T06:09:46.335356Z","shell.execute_reply.started":"2026-01-07T06:09:46.329407Z","shell.execute_reply":"2026-01-07T06:09:46.334734Z"}},"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())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:46.337040Z","iopub.execute_input":"2026-01-07T06:09:46.337509Z","iopub.status.idle":"2026-01-07T06:09:54.347938Z","shell.execute_reply.started":"2026-01-07T06:09:46.337487Z","shell.execute_reply":"2026-01-07T06:09:54.347026Z"}},"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-07T06:09:54.349084Z","iopub.execute_input":"2026-01-07T06:09:54.349341Z","iopub.status.idle":"2026-01-07T06:09:55.826257Z","shell.execute_reply.started":"2026-01-07T06:09:54.349316Z","shell.execute_reply":"2026-01-07T06:09:55.825620Z"}},"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\n# Helper - get ancestors for a GO term\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\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()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:14:26.572559Z","iopub.execute_input":"2026-01-07T06:14:26.573430Z","iopub.status.idle":"2026-01-07T06:14:33.346269Z","shell.execute_reply.started":"2026-01-07T06:14:26.573395Z","shell.execute_reply":"2026-01-07T06:14:33.345576Z"}},"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\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-07T06:15:32.376022Z","iopub.execute_input":"2026-01-07T06:15:32.376584Z","iopub.status.idle":"2026-01-07T06:15:36.836045Z","shell.execute_reply.started":"2026-01-07T06:15:32.376555Z","shell.execute_reply":"2026-01-07T06:15:36.835201Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# filter rare GO terms\n\nmin_count = 30  # common threshold\n\nmf_counts = Y_mf.sum(axis=0)\nkeep_mf = mf_counts >= min_count\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 = bp_counts >= min_count\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 = cc_counts >= min_count\nY_cc = Y_cc[:, keep_cc]\ncc_terms = [t for t, k in zip(mlb_bp.classes_, keep_bp) if k]\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\nimport joblib\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},"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\n# Helper - get ancestors for a GO term\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\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()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:15:51.129081Z","iopub.execute_input":"2026-01-07T06:15:51.129729Z","iopub.status.idle":"2026-01-07T06:15:58.119127Z","shell.execute_reply.started":"2026-01-07T06:15:51.129695Z","shell.execute_reply":"2026-01-07T06:15:58.118378Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# filter rare GO terms\n\nmin_count = 30  # common threshold\n\nmf_counts = Y_mf.sum(axis=0)\nkeep_mf = mf_counts >= min_count\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 = bp_counts >= min_count\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 = cc_counts >= min_count\nY_cc = Y_cc[:, keep_cc]\ncc_terms = [t for t, k in zip(mlb_bp.classes_, keep_bp) if k]\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\nimport joblib\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-07T06:16:03.709069Z","iopub.execute_input":"2026-01-07T06:16:03.709727Z","iopub.status.idle":"2026-01-07T06:16:31.094299Z","shell.execute_reply.started":"2026-01-07T06:16:03.709694Z","shell.execute_reply":"2026-01-07T06:16:31.093533Z"}},"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/prott5-embeddings-test-npy/test_prott5_embeddings.npy\"\n\nX_np = np.load(train_path)\n\nX_np.shape\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.897304Z","iopub.status.idle":"2026-01-07T06:09:55.897550Z","shell.execute_reply.started":"2026-01-07T06:09:55.897439Z","shell.execute_reply":"2026-01-07T06:09:55.897453Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# X_np is a 2D numpy matrix.\n1. Row i corresponds to protein i\n2. The file does not store protein IDs\n# I need load protein ID to the embeddings\n   \n","metadata":{}},{"cell_type":"code","source":"# Load the data\ntrain_path = \"/kaggle/input/prott5-embeddings/protT5_embeddings.npy\"\ntest_path = \"/kaggle/input/prott5-embeddings-test-npy/test_prott5_embeddings.npy\"\n\n\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-07T06:09:55.898619Z","iopub.status.idle":"2026-01-07T06:09:55.898987Z","shell.execute_reply.started":"2026-01-07T06:09:55.898807Z","shell.execute_reply":"2026-01-07T06:09:55.898828Z"}},"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}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.900244Z","iopub.status.idle":"2026-01-07T06:09:55.903076Z","shell.execute_reply.started":"2026-01-07T06:09:55.902867Z","shell.execute_reply":"2026-01-07T06:09:55.902902Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"X_train = np.vstack([train_embedding_dict[pid] for pid in df_train[\"protein_id\"]])\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-07T06:09:55.904034Z","iopub.status.idle":"2026-01-07T06:09:55.904441Z","shell.execute_reply.started":"2026-01-07T06:09:55.904238Z","shell.execute_reply":"2026-01-07T06:09:55.904262Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nos.listdir(\"/kaggle/working/\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.906184Z","iopub.status.idle":"2026-01-07T06:09:55.906612Z","shell.execute_reply.started":"2026-01-07T06:09:55.906429Z","shell.execute_reply":"2026-01-07T06:09:55.906451Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# MLP model training\n","metadata":{}},{"cell_type":"code","source":"# Load data for training\nimport numpy as np\nimport joblib\n\nX_train = np.load(\"/kaggle/input/cafa-6-prediction-using-prott5-embeddings-2/X_train_protT5.npy\")\nX_test  = np.load(\"/kaggle/input/cafa-6-prediction-using-prott5-embeddings-2/X_test_protT5.npy\")\n\nY_mf = np.load(\"/kaggle/input/cafa-6-prediction-using-prott5-embeddings-2/Y_mf.npy\")\nY_bp = np.load(\"/kaggle/input/cafa-6-prediction-using-prott5-embeddings-2/Y_bp.npy\")\nY_cc = np.load(\"/kaggle/input/cafa-6-prediction-using-prott5-embeddings-2/Y_cc.npy\")\n\nmlb_mf = joblib.load(\"/kaggle/input/cafa-6-prediction-using-prott5-embeddings-2/mlb_mf.pkl\")\nmlb_bp = joblib.load(\"/kaggle/input/cafa-6-prediction-using-prott5-embeddings-2/mlb_bp.pkl\")\nmlb_cc = joblib.load(\"/kaggle/input/cafa-6-prediction-using-prott5-embeddings-2/mlb_cc.pkl\")\n\nfrom sklearn.model_selection import train_test_split\n\n# Split\nX_tr, X_val, mf_tr, mf_val, bp_tr, bp_val, cc_tr, cc_val = train_test_split(\n    X_train, Y_mf, Y_bp, Y_cc, test_size=0.2, random_state=42\n)\n\n# Convert to tensors\nX_tr = torch.tensor(X_tr, dtype=torch.float32)\nX_val = torch.tensor(X_val, dtype=torch.float32)\n\nmf_tr = torch.tensor(mf_tr, dtype=torch.float32)\nbp_tr = torch.tensor(bp_tr, dtype=torch.float32)\ncc_tr = torch.tensor(cc_tr, dtype=torch.float32)\n\nmf_val = torch.tensor(mf_val, dtype=torch.float32)\nbp_val = torch.tensor(bp_val, dtype=torch.float32)\ncc_val = torch.tensor(cc_val, dtype=torch.float32)\n","metadata":{"execution":{"iopub.status.busy":"2026-01-07T06:09:55.907735Z","iopub.status.idle":"2026-01-07T06:09:55.908155Z","shell.execute_reply.started":"2026-01-07T06:09:55.907972Z","shell.execute_reply":"2026-01-07T06:09:55.907993Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from torch.utils.data import Dataset, DataLoader\nclass ProteinDataset(Dataset):\n    def __init__(self, X, mf, bp, cc):\n        self.X = X\n        self.mf = mf\n        self.bp = bp\n        self.cc = cc\n\n    def __len__(self):\n        return len(self.X)\n\n    def __getitem__(self, idx):\n        return self.X[idx], self.mf[idx], self.bp[idx], self.cc[idx]\n\ntrain_loader = DataLoader(ProteinDataset(X_tr, mf_tr, bp_tr, cc_tr),\n                          batch_size=1024, shuffle=True)\n\nval_loader = DataLoader(ProteinDataset(X_val, mf_val, bp_val, cc_val),\n                        batch_size=1024, shuffle=False)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.909246Z","iopub.status.idle":"2026-01-07T06:09:55.909760Z","shell.execute_reply.started":"2026-01-07T06:09:55.909590Z","shell.execute_reply":"2026-01-07T06:09:55.909607Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class MultiTaskMLP(nn.Module):\n    def __init__(self, input_dim, mf_dim, bp_dim, cc_dim):\n        super().__init__()\n\n        self.shared = nn.Sequential(\n            nn.Linear(input_dim, 1024),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(1024, 512),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n        )\n\n        self.mf_head = nn.Linear(512, mf_dim)\n        self.bp_head = nn.Linear(512, bp_dim)\n        self.cc_head = nn.Linear(512, cc_dim)\n\n    def forward(self, x):\n        h = self.shared(x)\n        return (\n            torch.sigmoid(self.mf_head(h)),\n            torch.sigmoid(self.bp_head(h)),\n            torch.sigmoid(self.cc_head(h)),\n        )\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\nmodel = MultiTaskMLP(\n    input_dim=1024,\n    mf_dim=Y_mf.shape[1],\n    bp_dim=Y_bp.shape[1],\n    cc_dim=Y_cc.shape[1]\n).to(device)\n\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\ncriterion = nn.BCELoss()\n\nfor epoch in range(10):\n    model.train()\n    total_loss = 0\n\n    for X, mf, bp, cc in train_loader:\n        X, mf, bp, cc = X.to(device), mf.to(device), bp.to(device), cc.to(device)\n\n        optimizer.zero_grad()\n        pred_mf, pred_bp, pred_cc = model(X)\n\n        loss = (\n            criterion(pred_mf, mf) +\n            criterion(pred_bp, bp) +\n            criterion(pred_cc, cc)\n        )\n\n        loss.backward()\n        optimizer.step()\n\n        total_loss += loss.item()\n\n    print(f\"Epoch {epoch+1}, Loss = {total_loss:.4f}\")\n\n    \n       \n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.910735Z","iopub.status.idle":"2026-01-07T06:09:55.910996Z","shell.execute_reply.started":"2026-01-07T06:09:55.910874Z","shell.execute_reply":"2026-01-07T06:09:55.910893Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Evaluation\ndef get_validation_predictions(model, val_loader, device):\n    model.eval()\n    preds_mf, preds_bp, preds_cc = [], [], []\n    trues_mf, trues_bp, trues_cc = [], [], []\n\n    with torch.no_grad():\n        for X, mf, bp, cc in val_loader:\n            X = X.to(device)\n            pmf, pbp, pcc = model(X)\n\n            preds_mf.append(pmf.cpu())\n            preds_bp.append(pbp.cpu())\n            preds_cc.append(pcc.cpu())\n\n            trues_mf.append(mf)\n            trues_bp.append(bp)\n            trues_cc.append(cc)\n\n    return (\n        torch.cat(preds_mf).numpy(),\n        torch.cat(preds_bp).numpy(),\n        torch.cat(preds_cc).numpy(),\n        torch.cat(trues_mf).numpy(),\n        torch.cat(trues_bp).numpy(),\n        torch.cat(trues_cc).numpy(),\n    )\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.911878Z","iopub.status.idle":"2026-01-07T06:09:55.912096Z","shell.execute_reply.started":"2026-01-07T06:09:55.911991Z","shell.execute_reply":"2026-01-07T06:09:55.912004Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\n\ndef compute_fmax(pred, true):\n    thresholds = np.linspace(0, 1, 101)\n    best_f, best_t = 0, 0\n\n    for t in thresholds:\n        bin_pred = (pred >= t).astype(int)\n\n        tp = (bin_pred * true).sum()\n        fp = (bin_pred * (1 - true)).sum()\n        fn = ((1 - bin_pred) * true).sum()\n\n        if tp == 0:\n            continue\n\n        precision = tp / (tp + fp)\n        recall = tp / (tp + fn)\n        f = 2 * precision * recall / (precision + recall)\n\n        if f > best_f:\n            best_f = f\n            best_t = t\n\n    return best_f, best_t\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.912897Z","iopub.status.idle":"2026-01-07T06:09:55.913116Z","shell.execute_reply.started":"2026-01-07T06:09:55.913013Z","shell.execute_reply":"2026-01-07T06:09:55.913026Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pred_mf, pred_bp, pred_cc, true_mf, true_bp, true_cc = \\\n    get_validation_predictions(model, val_loader, device)\n\nbest_f_mf, best_t_mf = compute_fmax(pred_mf, true_mf)\nbest_f_bp, best_t_bp = compute_fmax(pred_bp, true_bp)\nbest_f_cc, best_t_cc = compute_fmax(pred_cc, true_cc)\n\nprint(\"MF:\", best_f_mf, best_t_mf)\nprint(\"BP:\", best_f_bp, best_t_bp)\nprint(\"CC:\", best_f_cc, best_t_cc)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.913760Z","iopub.status.idle":"2026-01-07T06:09:55.914086Z","shell.execute_reply.started":"2026-01-07T06:09:55.913950Z","shell.execute_reply":"2026-01-07T06:09:55.913972Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Train more epochs\n\nfor epoch in range(30):\n    model.train()\n    total_loss = 0\n\n    for X, mf, bp, cc in train_loader:\n        X, mf, bp, cc = X.to(device), mf.to(device), bp.to(device), cc.to(device)\n\n        optimizer.zero_grad()\n        pred_mf, pred_bp, pred_cc = model(X)\n\n        loss = (\n            criterion(pred_mf, mf) +\n            criterion(pred_bp, bp) +\n            criterion(pred_cc, cc)\n        )\n\n        loss.backward()\n        optimizer.step()\n\n        total_loss += loss.item()\n\n    print(f\"Epoch {epoch+1}, Loss = {total_loss:.4f}\")\n    \npred_mf, pred_bp, pred_cc, true_mf, true_bp, true_cc = \\\n    get_validation_predictions(model, val_loader, device)\n\nbest_f_mf, best_t_mf = compute_fmax(pred_mf, true_mf)\nbest_f_bp, best_t_bp = compute_fmax(pred_bp, true_bp)\nbest_f_cc, best_t_cc = compute_fmax(pred_cc, true_cc)\n\nprint(\"MF:\", best_f_mf, best_t_mf)\nprint(\"BP:\", best_f_bp, best_t_bp)\nprint(\"CC:\", best_f_cc, best_t_cc)\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.915522Z","iopub.status.idle":"2026-01-07T06:09:55.915859Z","shell.execute_reply.started":"2026-01-07T06:09:55.915680Z","shell.execute_reply":"2026-01-07T06:09:55.915699Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Add batchNorm\n\nclass MultiTaskMLP(nn.Module):\n    def __init__(self, input_dim, mf_dim, bp_dim, cc_dim):\n        super().__init__()\n\n        self.shared = nn.Sequential(\n            nn.Linear(input_dim, 1024),\n            nn.BatchNorm1d(1024),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(1024, 512),\n            nn.BatchNorm1d(512),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n        )\n\n        self.mf_head = nn.Linear(512, mf_dim)\n        self.bp_head = nn.Linear(512, bp_dim)\n        self.cc_head = nn.Linear(512, cc_dim)\n\n    def forward(self, x):\n        h = self.shared(x)\n        return (\n            torch.sigmoid(self.mf_head(h)),\n            torch.sigmoid(self.bp_head(h)),\n            torch.sigmoid(self.cc_head(h)),\n        )\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\nmodel = MultiTaskMLP(\n    input_dim=1024,\n    mf_dim=Y_mf.shape[1],\n    bp_dim=Y_bp.shape[1],\n    cc_dim=Y_cc.shape[1]\n).to(device)\n\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\ncriterion = nn.BCELoss()\n\nfor epoch in range(30):\n    model.train()\n    total_loss = 0\n\n    for X, mf, bp, cc in train_loader:\n        X, mf, bp, cc = X.to(device), mf.to(device), bp.to(device), cc.to(device)\n\n        optimizer.zero_grad()\n        pred_mf, pred_bp, pred_cc = model(X)\n\n        loss = (\n            criterion(pred_mf, mf) +\n            criterion(pred_bp, bp) +\n            criterion(pred_cc, cc)\n        )\n\n        loss.backward()\n        optimizer.step()\n\n        total_loss += loss.item()\n\n    print(f\"Epoch {epoch+1}, Loss = {total_loss:.4f}\")\n    \npred_mf, pred_bp, pred_cc, true_mf, true_bp, true_cc = \\\n    get_validation_predictions(model, val_loader, device)\n\nbest_f_mf, best_t_mf = compute_fmax(pred_mf, true_mf)\nbest_f_bp, best_t_bp = compute_fmax(pred_bp, true_bp)\nbest_f_cc, best_t_cc = compute_fmax(pred_cc, true_cc)\n\nprint(\"MF:\", best_f_mf, best_t_mf)\nprint(\"BP:\", best_f_bp, best_t_bp)\nprint(\"CC:\", best_f_cc, best_t_cc)\n\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.917130Z","iopub.status.idle":"2026-01-07T06:09:55.917403Z","shell.execute_reply.started":"2026-01-07T06:09:55.917262Z","shell.execute_reply":"2026-01-07T06:09:55.917280Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# early stopping and learning rate scheduler\nimport torch\nimport numpy as np\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\ncriterion = torch.nn.BCELoss()\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n\nscheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n    optimizer,\n    mode=\"min\",\n    factor=0.5,\n    patience=2,\n)\n\nmax_epochs = 50\nearly_stop_patience = 5\nbest_val_loss = float(\"inf\")\nepochs_without_improvement = 0\n\nbest_model_state = None\n\nfor epoch in range(max_epochs):\n    # -------------------------\n    # TRAINING\n    # -------------------------\n    model.train()\n    train_loss = 0\n\n    for X, mf, bp, cc in train_loader:\n        X, mf, bp, cc = X.to(device), mf.to(device), bp.to(device), cc.to(device)\n\n        optimizer.zero_grad()\n        pmf, pbp, pcc = model(X)\n\n        loss = (\n            criterion(pmf, mf) +\n            criterion(pbp, bp) +\n            criterion(pcc, cc)\n        )\n\n        loss.backward()\n        optimizer.step()\n\n        train_loss += loss.item()\n\n    # -------------------------\n    # VALIDATION\n    # -------------------------\n    model.eval()\n    val_loss = 0\n\n    with torch.no_grad():\n        for X, mf, bp, cc in val_loader:\n            X, mf, bp, cc = X.to(device), mf.to(device), bp.to(device), cc.to(device)\n            pmf, pbp, pcc = model(X)\n\n            loss = (\n                criterion(pmf, mf) +\n                criterion(pbp, bp) +\n                criterion(pcc, cc)\n            )\n            val_loss += loss.item()\n\n    print(f\"Epoch {epoch+1}/{max_epochs} | Train Loss={train_loss:.4f} | Val Loss={val_loss:.4f}\")\n\n    # -------------------------\n    # LR SCHEDULER\n    # -------------------------\n    scheduler.step(val_loss)\n\n    # -------------------------\n    # EARLY STOPPING\n    # -------------------------\n    if val_loss < best_val_loss:\n        best_val_loss = val_loss\n        epochs_without_improvement = 0\n        best_model_state = model.state_dict()  # save best model\n        print(\"  → New best model saved\")\n    else:\n        epochs_without_improvement += 1\n        print(f\"  → No improvement for {epochs_without_improvement} epochs\")\n\n        if epochs_without_improvement >= early_stop_patience:\n            print(\"Early stopping triggered\")\n            break\n\n# Load best model\nmodel.load_state_dict(best_model_state)\nprint(\"Best model restored\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.918816Z","iopub.status.idle":"2026-01-07T06:09:55.919129Z","shell.execute_reply.started":"2026-01-07T06:09:55.918960Z","shell.execute_reply":"2026-01-07T06:09:55.919000Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pred_mf, pred_bp, pred_cc, true_mf, true_bp, true_cc = \\\n    get_validation_predictions(model, val_loader, device)\n\nmlp_val_mf = pred_mf\nmlp_val_bp = pred_bp\nmlp_val_cc = pred_cc\n\nnp.save(\"mlp_val_mf.npy\", mlp_val_mf)\nnp.save(\"mlp_val_bp.npy\", mlp_val_bp)\nnp.save(\"mlp_val_cc.npy\", mlp_val_cc)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.920420Z","iopub.status.idle":"2026-01-07T06:09:55.920753Z","shell.execute_reply.started":"2026-01-07T06:09:55.920585Z","shell.execute_reply":"2026-01-07T06:09:55.920611Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n\nbest_f_mf, best_t_mf = compute_fmax(pred_mf, true_mf)\nbest_f_bp, best_t_bp = compute_fmax(pred_bp, true_bp)\nbest_f_cc, best_t_cc = compute_fmax(pred_cc, true_cc)\n\nprint(\"MF:\", best_f_mf, best_t_mf)\nprint(\"BP:\", best_f_bp, best_t_bp)\nprint(\"CC:\", best_f_cc, best_t_cc)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.922202Z","iopub.status.idle":"2026-01-07T06:09:55.922447Z","shell.execute_reply.started":"2026-01-07T06:09:55.922335Z","shell.execute_reply":"2026-01-07T06:09:55.922349Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Add a deeper shared trunk\nclass MultiTaskMLP(nn.Module):\n    def __init__(self, input_dim, mf_dim, bp_dim, cc_dim):\n        super().__init__()\n\n        self.shared = nn.Sequential(\n            nn.Linear(input_dim, 1024),\n            nn.ReLU(),\n            nn.Dropout(0.1),\n            \n            nn.Linear(1024, 512),\n            nn.ReLU(),\n            nn.Dropout(0.1),\n            \n            nn.Linear(512, 256),\n            nn.ReLU(),\n            nn.Dropout(0.1),\n        )\n\n        self.mf_head = nn.Linear(256, mf_dim)\n        self.bp_head = nn.Linear(256, bp_dim)\n        self.cc_head = nn.Linear(256, cc_dim)\n\n    def forward(self, x):\n        h = self.shared(x)\n        return (\n            torch.sigmoid(self.mf_head(h)),\n            torch.sigmoid(self.bp_head(h)),\n            torch.sigmoid(self.cc_head(h)),\n        )\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\nmodel = MultiTaskMLP(\n    input_dim=1024,\n    mf_dim=Y_mf.shape[1],\n    bp_dim=Y_bp.shape[1],\n    cc_dim=Y_cc.shape[1]\n).to(device)\n\ncriterion = torch.nn.BCELoss()\noptimizer = torch.optim.Adam(model.parameters(), lr=5e-4)\n\n\nscheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n    optimizer,\n    mode=\"min\",\n    factor=0.5,\n    patience=4,\n)\n\nmax_epochs =150\nearly_stop_patience = 10\nbest_val_loss = float(\"inf\")\nepochs_without_improvement = 0\n\nbest_model_state = None\n\nfor epoch in range(max_epochs):\n    # -------------------------\n    # TRAINING\n    # -------------------------\n    model.train()\n    train_loss = 0\n\n    for X, mf, bp, cc in train_loader:\n        X, mf, bp, cc = X.to(device), mf.to(device), bp.to(device), cc.to(device)\n\n        optimizer = torch.optim.Adam(model.parameters(), lr=5e-4)\n\n        pmf, pbp, pcc = model(X)\n\n        loss = (\n            criterion(pmf, mf) +\n            criterion(pbp, bp) +\n            criterion(pcc, cc)\n        )\n\n        loss.backward()\n        optimizer.step()\n\n        train_loss += loss.item()\n\n    # -------------------------\n    # VALIDATION\n    # -------------------------\n    model.eval()\n    val_loss = 0\n\n    with torch.no_grad():\n        for X, mf, bp, cc in val_loader:\n            X, mf, bp, cc = X.to(device), mf.to(device), bp.to(device), cc.to(device)\n            pmf, pbp, pcc = model(X)\n\n            loss = (\n                criterion(pmf, mf) +\n                criterion(pbp, bp) +\n                criterion(pcc, cc)\n            )\n            val_loss += loss.item()\n\n    print(f\"Epoch {epoch+1}/{max_epochs} | Train Loss={train_loss:.4f} | Val Loss={val_loss:.4f}\")\n\n    # -------------------------\n    # LR SCHEDULER\n    # -------------------------\n    scheduler.step(val_loss)\n\n    # -------------------------\n    # EARLY STOPPING\n    # -------------------------\n    if val_loss < best_val_loss:\n        best_val_loss = val_loss\n        epochs_without_improvement = 0\n        best_model_state = model.state_dict()  # save best model\n        print(\"  → New best model saved\")\n    else:\n        epochs_without_improvement += 1\n        print(f\"  → No improvement for {epochs_without_improvement} epochs\")\n\n        if epochs_without_improvement >= early_stop_patience:\n            print(\"Early stopping triggered\")\n            break\n\n# Load best model\nmodel.load_state_dict(best_model_state)\nprint(\"Best model restored\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.924022Z","iopub.status.idle":"2026-01-07T06:09:55.924341Z","shell.execute_reply.started":"2026-01-07T06:09:55.924191Z","shell.execute_reply":"2026-01-07T06:09:55.924217Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pred_mf, pred_bp, pred_cc, true_mf, true_bp, true_cc = \\\n    get_validation_predictions(model, val_loader, device)\n\nbest_f_mf, best_t_mf = compute_fmax(pred_mf, true_mf)\nbest_f_bp, best_t_bp = compute_fmax(pred_bp, true_bp)\nbest_f_cc, best_t_cc = compute_fmax(pred_cc, true_cc)\n\nprint(\"MF:\", best_f_mf, best_t_mf)\nprint(\"BP:\", best_f_bp, best_t_bp)\nprint(\"CC:\", best_f_cc, best_t_cc)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.925195Z","iopub.status.idle":"2026-01-07T06:09:55.925513Z","shell.execute_reply.started":"2026-01-07T06:09:55.925330Z","shell.execute_reply":"2026-01-07T06:09:55.925354Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Final MLP model:\nclass MultiTaskMLP(nn.Module):\n    def __init__(self, input_dim, mf_dim, bp_dim, cc_dim):\n        super().__init__()\n\n        self.shared = nn.Sequential(\n            nn.Linear(input_dim, 1024),\n            nn.BatchNorm1d(1024),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(1024, 512),\n            nn.BatchNorm1d(512),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n        )\n\n        self.mf_head = nn.Linear(512, mf_dim)\n        self.bp_head = nn.Linear(512, bp_dim)\n        self.cc_head = nn.Linear(512, cc_dim)\n\n    def forward(self, x):\n        h = self.shared(x)\n        return (\n            torch.sigmoid(self.mf_head(h)),\n            torch.sigmoid(self.bp_head(h)),\n            torch.sigmoid(self.cc_head(h)),\n        )\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\nmodel = MultiTaskMLP(\n    input_dim=1024,\n    mf_dim=Y_mf.shape[1],\n    bp_dim=Y_bp.shape[1],\n    cc_dim=Y_cc.shape[1]\n).to(device)\n\n\ncriterion = torch.nn.BCELoss()\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n\nscheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n    optimizer,\n    mode=\"min\",\n    factor=0.5,\n    patience=2,\n)\n\nmax_epochs = 50\nearly_stop_patience = 5\nbest_val_loss = float(\"inf\")\nepochs_without_improvement = 0\n\nbest_model_state = None\n\nfor epoch in range(max_epochs):\n    # -------------------------\n    # TRAINING\n    # -------------------------\n    model.train()\n    train_loss = 0\n\n    for X, mf, bp, cc in train_loader:\n        X, mf, bp, cc = X.to(device), mf.to(device), bp.to(device), cc.to(device)\n\n        optimizer.zero_grad()\n        pmf, pbp, pcc = model(X)\n\n        loss = (\n            criterion(pmf, mf) +\n            criterion(pbp, bp) +\n            criterion(pcc, cc)\n        )\n\n        loss.backward()\n        optimizer.step()\n\n        train_loss += loss.item()\n\n    # -------------------------\n    # VALIDATION\n    # -------------------------\n    model.eval()\n    val_loss = 0\n\n    with torch.no_grad():\n        for X, mf, bp, cc in val_loader:\n            X, mf, bp, cc = X.to(device), mf.to(device), bp.to(device), cc.to(device)\n            pmf, pbp, pcc = model(X)\n\n            loss = (\n                criterion(pmf, mf) +\n                criterion(pbp, bp) +\n                criterion(pcc, cc)\n            )\n            val_loss += loss.item()\n\n    print(f\"Epoch {epoch+1}/{max_epochs} | Train Loss={train_loss:.4f} | Val Loss={val_loss:.4f}\")\n\n    # -------------------------\n    # LR SCHEDULER\n    # -------------------------\n    scheduler.step(val_loss)\n\n    # -------------------------\n    # EARLY STOPPING\n    # -------------------------\n    if val_loss < best_val_loss:\n        best_val_loss = val_loss\n        epochs_without_improvement = 0\n        best_model_state = model.state_dict()  # save best model\n        print(\"  → New best model saved\")\n    else:\n        epochs_without_improvement += 1\n        print(f\"  → No improvement for {epochs_without_improvement} epochs\")\n\n        if epochs_without_improvement >= early_stop_patience:\n            print(\"Early stopping triggered\")\n            break\n\n# Load best model\nmodel.load_state_dict(best_model_state)\nprint(\"Best model restored\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.926868Z","iopub.status.idle":"2026-01-07T06:09:55.927207Z","shell.execute_reply.started":"2026-01-07T06:09:55.927035Z","shell.execute_reply":"2026-01-07T06:09:55.927054Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pred_mf, pred_bp, pred_cc, true_mf, true_bp, true_cc = \\\n    get_validation_predictions(model, val_loader, device)\n\nbest_f_mf, best_t_mf = compute_fmax(pred_mf, true_mf)\nbest_f_bp, best_t_bp = compute_fmax(pred_bp, true_bp)\nbest_f_cc, best_t_cc = compute_fmax(pred_cc, true_cc)\n\nprint(\"MF:\", best_f_mf, best_t_mf)\nprint(\"BP:\", best_f_bp, best_t_bp)\nprint(\"CC:\", best_f_cc, best_t_cc)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.927943Z","iopub.status.idle":"2026-01-07T06:09:55.928159Z","shell.execute_reply.started":"2026-01-07T06:09:55.928055Z","shell.execute_reply":"2026-01-07T06:09:55.928068Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"y_pred_mf, y_pred_bp, y_pred_cc, true_mf, true_bp, true_cc = \\\n    get_validation_predictions(model, val_loader, device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.929014Z","iopub.status.idle":"2026-01-07T06:09:55.929302Z","shell.execute_reply.started":"2026-01-07T06:09:55.929146Z","shell.execute_reply":"2026-01-07T06:09:55.929160Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"go_terms_mf = list(mlb_mf.classes_)  # List of MF GO terms\ngo_terms_bp = list(mlb_bp.classes_)  # List of BP GO terms  \ngo_terms_cc = list(mlb_cc.classes_)  # List of CC GO terms\n\nprint(f\"Type of go_terms_mf: {type(go_terms_mf)}\")\nprint(f\"Length of go_terms_mf: {len(go_terms_mf)}\")\nprint(f\"Shape of y_pred_mf: {y_pred_mf.shape}\")\nprint(f\"First 5 GO terms: {go_terms_mf[:5]}\")\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.930657Z","iopub.status.idle":"2026-01-07T06:09:55.930989Z","shell.execute_reply.started":"2026-01-07T06:09:55.930836Z","shell.execute_reply":"2026-01-07T06:09:55.930859Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_active_go_terms(Y_matrix, mlb):\n    \"\"\"\n    Get only the GO terms that have at least one positive example.\n    This matches what your model actually predicts.\n    \n    Args:\n        Y_matrix: binary label matrix [n_samples, n_terms]\n        mlb: fitted MultiLabelBinarizer\n    \n    Returns:\n        tuple: (active_go_terms, column_indices)\n    \"\"\"\n    all_go_terms = list(mlb.classes_)\n    \n    # Find columns with at least one positive example\n    active_cols = np.where(Y_matrix.sum(axis=0) > 0)[0]\n    active_go_terms = [all_go_terms[i] for i in active_cols]\n    \n    print(f\"Total GO terms in MLB: {len(all_go_terms)}\")\n    print(f\"Active GO terms (with examples): {len(active_go_terms)}\")\n    \n    return active_go_terms, active_cols\n\n# For MF\ngo_terms_mf, active_cols_mf = get_active_go_terms(Y_mf, mlb_mf)\nprint(f\"MF: {len(go_terms_mf)} active terms\")\n\n# For BP\ngo_terms_bp, active_cols_bp = get_active_go_terms(Y_bp, mlb_bp)\nprint(f\"BP: {len(go_terms_bp)} active terms\")\n\n# For CC\ngo_terms_cc, active_cols_cc = get_active_go_terms(Y_cc, mlb_cc)\nprint(f\"CC: {len(go_terms_cc)} active terms\")\n\n# Verify dimensions match\nprint(f\"\\nDimension check:\")\nprint(f\"MF: go_terms={len(go_terms_mf)}, y_pred shape={y_pred_mf.shape}\")\nprint(f\"BP: go_terms={len(go_terms_bp)}, y_pred shape={y_pred_bp.shape}\")\nprint(f\"CC: go_terms={len(go_terms_cc)}, y_pred shape={y_pred_cc.shape}\")\n\n\n\n\n\ny_pred_mf = y_pred_mf[:, active_cols_mf]\ny_pred_bp = y_pred_bp[:, active_cols_bp]\ny_pred_cc = y_pred_cc[:, active_cols_cc]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.932854Z","iopub.status.idle":"2026-01-07T06:09:55.933194Z","shell.execute_reply.started":"2026-01-07T06:09:55.933021Z","shell.execute_reply":"2026-01-07T06:09:55.933042Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nfrom goatools.obo_parser import GODag\n\ndef apply_go_dag_consistency(y_pred, gos, go_dag):\n    \"\"\"\n    Enforce GO DAG consistency on prediction scores.\n    y_pred: np.ndarray [n_samples, n_terms]\n    gos: list of GO IDs aligned to y_pred columns\n    go_dag: GODag object from goatools\n    \"\"\"\n    # Convert gos to list if it's not already\n    gos = list(gos)\n    go_to_idx = {go: i for i, go in enumerate(gos)}\n    \n    # Sanity check\n    assert len(go_to_idx) == y_pred.shape[1], \\\n        f\"Mismatch: {len(go_to_idx)} GO terms but {y_pred.shape[1]} columns in y_pred\"\n    \n    y = y_pred.copy()\n    \n    for child_go in gos:\n        if child_go not in go_dag:\n            continue\n        child_idx = go_to_idx[child_go]\n        child_scores = y[:, child_idx]\n        \n        # Get all parent terms\n        for parent_go in go_dag[child_go].get_all_parents():\n            # Only propagate to parents that exist in our term list\n            if parent_go not in go_to_idx:\n                continue\n            \n            parent_idx = go_to_idx[parent_go]\n            y[:, parent_idx] = np.maximum(\n                y[:, parent_idx],\n                child_scores\n            )\n    return y\n\n\ndef propagate_ground_truth(y_true, gos, go_dag):\n    \"\"\"\n    Propagate ground truth labels up the GO DAG (add all ancestors).\n    This is critical for fair F-max evaluation!\n    \n    y_true: np.ndarray [n_samples, n_terms] - binary labels\n    gos: list of GO IDs aligned to y_true columns\n    go_dag: GODag object from goatools\n    \n    Returns: np.ndarray with ancestors added\n    \"\"\"\n    go_to_idx = {go: i for i, go in enumerate(gos)}\n    y_propagated = y_true.copy()\n    \n    for i in range(y_true.shape[0]):  # for each sample\n        for j, child_go in enumerate(gos):\n            if y_true[i, j] == 0:  # not annotated\n                continue\n            if child_go not in go_dag:\n                continue\n            \n            # Add all parent terms\n            for parent_go in go_dag[child_go].get_all_parents():\n                if parent_go in go_to_idx:\n                    parent_idx = go_to_idx[parent_go]\n                    y_propagated[i, parent_idx] = 1\n    \n    return y_propagated\n\ny_pred_mf = apply_go_dag_consistency(\n    y_pred_mf,\n    go_terms_mf,\n    go_dag\n)\n\ny_pred_bp = apply_go_dag_consistency(\n    y_pred_bp,\n    go_terms_bp,\n    go_dag\n)\n\ny_pred_cc = apply_go_dag_consistency(\n    y_pred_cc,\n    go_terms_cc,\n    go_dag\n)\n","metadata":{"execution":{"iopub.status.busy":"2026-01-07T06:09:55.934182Z","iopub.status.idle":"2026-01-07T06:09:55.934445Z","shell.execute_reply.started":"2026-01-07T06:09:55.934325Z","shell.execute_reply":"2026-01-07T06:09:55.934341Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"true_mf = propagate_ground_truth(true_mf, go_terms_mf, go_dag)\ntrue_bp = propagate_ground_truth(true_bp, go_terms_bp, go_dag)\ntrue_cc = propagate_ground_truth(true_cc, go_terms_cc, go_dag)\n\n\nbest_f_mf, best_t_mf = compute_fmax(y_pred_mf, true_mf)\nbest_f_bp, best_t_bp = compute_fmax(y_pred_bp, true_bp)\nbest_f_cc, best_t_cc = compute_fmax(y_pred_cc, true_cc)\n\nprint(\"MF:\", best_f_mf, best_t_mf)\nprint(\"BP:\", best_f_bp, best_t_bp)\nprint(\"CC:\", best_f_cc, best_t_cc)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.935275Z","iopub.status.idle":"2026-01-07T06:09:55.935486Z","shell.execute_reply.started":"2026-01-07T06:09:55.935385Z","shell.execute_reply":"2026-01-07T06:09:55.935397Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"X_full = torch.tensor(X_train, dtype=torch.float32)\n\nmf_full = torch.tensor(Y_mf, dtype=torch.float32)\nbp_full = torch.tensor(Y_bp, dtype=torch.float32)\ncc_full = torch.tensor(Y_cc, dtype=torch.float32)\n\n\ngo_terms_mf, active_cols_mf = get_active_go_terms(Y_mf, mlb_mf)\ngo_terms_bp, active_cols_bp = get_active_go_terms(Y_bp, mlb_bp)\ngo_terms_cc, active_cols_cc = get_active_go_terms(Y_cc, mlb_cc)\n\n\nmf_full = mf_full[:, active_cols_mf]\nbp_full = bp_full[:, active_cols_bp]\ncc_full = cc_full[:, active_cols_cc]\n\n\nfull_dataset = TensorDataset(X_full, mf_full, bp_full, cc_full)\nfull_loader = DataLoader(full_dataset, batch_size=1024, shuffle=True)\n\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\nmodel = MultiTaskMLP(\n    input_dim=1024,\n    mf_dim=len(active_cols_mf),\n    bp_dim=len(active_cols_bp),\n    cc_dim=len(active_cols_cc)\n).to(device)\n\n\ncriterion = torch.nn.BCELoss()\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n\n\nfor epoch in range(50):\n    model.train()\n    train_loss = 0\n\n    for X, mf, bp, cc in full_loader:\n        X, mf, bp, cc = X.to(device), mf.to(device), bp.to(device), cc.to(device)\n\n        optimizer.zero_grad()\n        pmf, pbp, pcc = model(X)\n\n        loss = (\n            criterion(pmf, mf) +\n            criterion(pbp, bp) +\n            criterion(pcc, cc)\n        )\n\n        loss.backward()\n        optimizer.step()\n\n        train_loss += loss.item()\n\n    print(f\"Epoch {epoch+1}/50 | Train Loss={train_loss:.4f}\")\n\n\n  \n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.937154Z","iopub.status.idle":"2026-01-07T06:09:55.937777Z","shell.execute_reply.started":"2026-01-07T06:09:55.937563Z","shell.execute_reply":"2026-01-07T06:09:55.937586Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"X_test_tensor = torch.tensor(X_test, dtype=torch.float32)\ntest_loader = DataLoader(X_test_tensor, batch_size=1024, shuffle=False)\n\nmodel.eval()\npred_mf, pred_bp, pred_cc = [], [], []\n\nwith torch.no_grad():\n    for X in test_loader:\n        X = X.to(device)\n        pmf, pbp, pcc = model(X)\n        pred_mf.append(pmf.cpu().numpy())\n        pred_bp.append(pbp.cpu().numpy())\n        pred_cc.append(pcc.cpu().numpy())\n\npred_mf = np.vstack(pred_mf)\npred_bp = np.vstack(pred_bp)\npred_cc = np.vstack(pred_cc)\n\nnp.save(\"mlp_pred_mf.npy\", pred_mf)\nnp.save(\"mlp_pred_bp.npy\", pred_bp)\nnp.save(\"mlp_pred_cc.npy\", pred_cc)\n\npred_mf = pred_mf[:, active_cols_mf]\npred_bp = pred_bp[:, active_cols_bp]\npred_cc = pred_cc[:, active_cols_cc]\n\npred_mf = apply_go_dag_consistency(pred_mf, go_terms_mf, go_dag)\npred_bp = apply_go_dag_consistency(pred_bp, go_terms_bp, go_dag)\npred_cc = apply_go_dag_consistency(pred_cc, go_terms_cc, go_dag)\n\n\n\n\nt_mf, t_bp, t_cc = 0.37, 0.24, 0.29\nbin_mf = (pred_mf >= t_mf)\nbin_bp = (pred_bp >= t_bp)\nbin_cc = (pred_cc >= t_cc)\n\nrows = []\ntest_protein_ids = df_test[\"protein_id\"].tolist()\n\n\nfor i, pid in enumerate(test_protein_ids):\n    # MF\n    for j, term in enumerate(go_terms_mf):\n        if bin_mf[i, j]:\n            rows.append([pid, term, float(pred_mf[i, j])])\n    # BP\n    for j, term in enumerate(go_terms_bp):\n        if bin_bp[i, j]:\n            rows.append([pid, term, float(pred_bp[i, j])])\n    # CC\n    for j, term in enumerate(go_terms_cc):\n        if bin_cc[i, j]:\n            rows.append([pid, term, float(pred_cc[i, j])])\n\ndf = pd.DataFrame(rows, columns=[\"ProteinID\", \"GO_term\", \"Score\"])\ndf.to_csv(\"final_submission.tsv\", sep=\"\\t\", index=False)\n\nprint(\"Final submission saved as final_submission.tsv\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:09:55.938808Z","iopub.status.idle":"2026-01-07T06:09:55.939101Z","shell.execute_reply.started":"2026-01-07T06:09:55.938972Z","shell.execute_reply":"2026-01-07T06:09:55.938994Z"}},"outputs":[],"execution_count":null}]}