{"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":14477966,"sourceType":"datasetVersion","datasetId":9247384}],"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-23T21:25:47.818931Z","iopub.execute_input":"2026-01-23T21:25:47.819115Z","iopub.status.idle":"2026-01-23T21:25:48.941986Z","shell.execute_reply.started":"2026-01-23T21:25:47.819097Z","shell.execute_reply":"2026-01-23T21:25:48.941309Z"}},"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-23T21:25:48.943724Z","iopub.execute_input":"2026-01-23T21:25:48.944065Z","iopub.status.idle":"2026-01-23T21:26:31.553654Z","shell.execute_reply.started":"2026-01-23T21:25:48.944042Z","shell.execute_reply":"2026-01-23T21:26:31.552877Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nimport matplotlib.pyplot as plt\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import DataLoader, TensorDataset\nfrom scipy.sparse import csr_matrix\nfrom collections import defaultdict\nimport joblib","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T21:26:31.554880Z","iopub.execute_input":"2026-01-23T21:26:31.555151Z","iopub.status.idle":"2026-01-23T21:26:35.381297Z","shell.execute_reply.started":"2026-01-23T21:26:31.555122Z","shell.execute_reply":"2026-01-23T21:26:35.380634Z"}},"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-23T21:26:35.382109Z","iopub.execute_input":"2026-01-23T21:26:35.382569Z","iopub.status.idle":"2026-01-23T21:26:41.423747Z","shell.execute_reply.started":"2026-01-23T21:26:35.382532Z","shell.execute_reply":"2026-01-23T21:26:41.422968Z"}},"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-23T21:26:41.425142Z","iopub.execute_input":"2026-01-23T21:26:41.425417Z","iopub.status.idle":"2026-01-23T21:26:43.609974Z","shell.execute_reply.started":"2026-01-23T21:26:41.425381Z","shell.execute_reply":"2026-01-23T21:26:43.609236Z"}},"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()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T21:26:43.610891Z","iopub.execute_input":"2026-01-23T21:26:43.611150Z","iopub.status.idle":"2026-01-23T21:26:50.654401Z","shell.execute_reply.started":"2026-01-23T21:26:43.611129Z","shell.execute_reply":"2026-01-23T21:26:50.653769Z"}},"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-23T21:26:50.656401Z","iopub.execute_input":"2026-01-23T21:26:50.656694Z","iopub.status.idle":"2026-01-23T21:26:59.797358Z","shell.execute_reply.started":"2026-01-23T21:26:50.656673Z","shell.execute_reply":"2026-01-23T21:26:59.796731Z"}},"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-23T21:26:59.798261Z","iopub.execute_input":"2026-01-23T21:26:59.798689Z","iopub.status.idle":"2026-01-23T21:26:59.810632Z","shell.execute_reply.started":"2026-01-23T21:26:59.798666Z","shell.execute_reply":"2026-01-23T21:26:59.810022Z"}},"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)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T21:26:59.811423Z","iopub.execute_input":"2026-01-23T21:26:59.811668Z","iopub.status.idle":"2026-01-23T21:27:11.632327Z","shell.execute_reply.started":"2026-01-23T21:26:59.811649Z","shell.execute_reply":"2026-01-23T21:27:11.631377Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load the protT5 embeddings\ntrain_path = \"/kaggle/input/prott5-embeddings/protT5_embeddings.npy\"\ntest_path = \"/kaggle/input/test-embeddings-prot-t5-xl/protT5_test_embeddings.npy\"\n\nX_np = np.load(train_path)\nX_train_torch = torch.from_numpy(X_np).float()\n\ntest  =  np.load(test_path)\nX_test_torch = torch.from_numpy(test).float()\n\ndef load_ids(fasta):\n    return [r.id.split()[0] for r in SeqIO.parse(fasta,\"fasta\")]\n\ntrain_ids = load_ids(\"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\")\nprint(train_ids[0])\ntrain_ids = [id.split(\"|\")[1] for id in train_ids]\ntest_ids = load_ids(\"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\")\n\nprint(train_ids[0],len(train_ids))\nprint(test_ids[0],len(test_ids))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T21:27:11.633323Z","iopub.execute_input":"2026-01-23T21:27:11.633599Z","iopub.status.idle":"2026-01-23T21:27:19.708236Z","shell.execute_reply.started":"2026-01-23T21:27:11.633570Z","shell.execute_reply":"2026-01-23T21:27:19.707204Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_embedding_dict = {\n    pid: X_np[i]\n    for i, pid in enumerate(train_ids)\n}\ntest_embedding_dict = {\n    pid: test[i]\n    for i, pid in enumerate(test_ids)\n}\n\nX_train = np.vstack([train_embedding_dict[pid] for pid in df_train[\"protein_id\"]])\nassert len(df_train) == len(X_train), \"Length mismatch!\"\n\nmissing = set(df_train[\"protein_id\"]) - set(train_embedding_dict.keys())\nassert len(missing) == 0, f\"Missing embeddings: {len(missing)}\"\n\n\n\nX_test = np.vstack([test_embedding_dict[pid] for pid in df_test[\"protein_id\"]])\nnp.save(\"X_train_protT5_xl.npy\", X_train)\nnp.save(\"X_test_protT5_xl.npy\", X_test)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T21:27:19.709533Z","iopub.execute_input":"2026-01-23T21:27:19.709917Z","iopub.status.idle":"2026-01-23T21:27:20.795503Z","shell.execute_reply.started":"2026-01-23T21:27:19.709890Z","shell.execute_reply":"2026-01-23T21:27:20.794676Z"}},"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-23T21:27:20.796556Z","iopub.execute_input":"2026-01-23T21:27:20.796868Z","iopub.status.idle":"2026-01-23T21:27:29.858619Z","shell.execute_reply.started":"2026-01-23T21:27:20.796818Z","shell.execute_reply":"2026-01-23T21:27:29.857914Z"}},"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-23T21:27:29.859580Z","iopub.execute_input":"2026-01-23T21:27:29.859896Z","iopub.status.idle":"2026-01-23T21:27:37.267350Z","shell.execute_reply.started":"2026-01-23T21:27:29.859864Z","shell.execute_reply":"2026-01-23T21:27:37.266488Z"}},"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-23T21:27:37.268307Z","iopub.execute_input":"2026-01-23T21:27:37.268765Z","iopub.status.idle":"2026-01-23T21:27:37.287464Z","shell.execute_reply.started":"2026-01-23T21:27:37.268741Z","shell.execute_reply":"2026-01-23T21:27:37.286671Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Convert to tensors\nX_tr = torch.tensor(X_tr, dtype=torch.float32)\nX_val = torch.tensor(X_val, dtype=torch.float32)\n\nY_mf_tr = torch.tensor(Y_mf_tr, dtype=torch.float32)\nY_bp_tr = torch.tensor(Y_bp_tr, dtype=torch.float32)\nY_cc_tr = torch.tensor(Y_cc_tr, dtype=torch.float32)\n\nY_mf_val = torch.tensor(Y_mf_val, dtype=torch.float32)\nY_bp_val = torch.tensor(Y_bp_val, dtype=torch.float32)\nY_cc_val = torch.tensor(Y_cc_val, dtype=torch.float32)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T21:27:37.288477Z","iopub.execute_input":"2026-01-23T21:27:37.289064Z","iopub.status.idle":"2026-01-23T21:27:38.224085Z","shell.execute_reply.started":"2026-01-23T21:27:37.289042Z","shell.execute_reply":"2026-01-23T21:27:38.223435Z"}},"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, Y_mf_tr, Y_bp_tr, Y_cc_tr),\n                          batch_size=1024, shuffle=True)\n\nval_loader = DataLoader(ProteinDataset(X_val, Y_mf_val, Y_bp_val, Y_cc_val),\n                        batch_size=1024, shuffle=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T21:27:38.225096Z","iopub.execute_input":"2026-01-23T21:27:38.225401Z","iopub.status.idle":"2026-01-23T21:27:38.231219Z","shell.execute_reply.started":"2026-01-23T21:27:38.225371Z","shell.execute_reply":"2026-01-23T21:27:38.230630Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n# Define Logistic Regression\nclass LogisticRegression(nn.Module):\n    def __init__(self, input_dim, mf_dim, bp_dim, cc_dim):\n        super().__init__()\n        self.mf_head = nn.Linear(input_dim, mf_dim)\n        self.bp_head = nn.Linear(input_dim, bp_dim)\n        self.cc_head = nn.Linear(input_dim, cc_dim)\n\n    def forward(self, x):\n        pmf = self.mf_head(x)\n        pbp = self.bp_head(x)\n        pcc = self.cc_head(x)\n        return pmf, pbp, pcc\n\ndef compute_pos_weight(Y):\n    pos = Y.sum(dim=0)\n    neg = Y.shape[0] - pos\n    return neg / (pos + 1e-8)\n\npos_weight_mf = compute_pos_weight(Y_mf_tr).to(device)\npos_weight_mf = torch.sqrt(pos_weight_mf)\npos_weight_mf = torch.clamp(pos_weight_mf, max=10.0)\npos_weight_bp = compute_pos_weight(Y_bp_tr).to(device)\npos_weight_bp = torch.sqrt(pos_weight_bp)\npos_weight_bp = torch.clamp(pos_weight_bp, max=10.0)\n\npos_weight_cc = compute_pos_weight(Y_cc_tr).to(device)\npos_weight_cc = torch.sqrt(pos_weight_cc)\npos_weight_cc = torch.clamp(pos_weight_cc, max=10.0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T21:27:38.232189Z","iopub.execute_input":"2026-01-23T21:27:38.232573Z","iopub.status.idle":"2026-01-23T21:27:38.739667Z","shell.execute_reply.started":"2026-01-23T21:27:38.232542Z","shell.execute_reply":"2026-01-23T21:27:38.738899Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = LogisticRegression(\n    input_dim=1024,\n    mf_dim=Y_mf_tr.shape[1],\n    bp_dim=Y_bp_tr.shape[1],\n    cc_dim=Y_cc_tr.shape[1]\n).to(device)\n\ncriterion_mf = nn.BCEWithLogitsLoss(pos_weight=pos_weight_mf)\ncriterion_bp = nn.BCEWithLogitsLoss(pos_weight=pos_weight_bp)\ncriterion_cc = nn.BCEWithLogitsLoss(pos_weight=pos_weight_cc)\n\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-4)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T21:27:38.740647Z","iopub.execute_input":"2026-01-23T21:27:38.740936Z","iopub.status.idle":"2026-01-23T21:27:42.045368Z","shell.execute_reply.started":"2026-01-23T21:27:38.740907Z","shell.execute_reply":"2026-01-23T21:27:42.044816Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"logit_reg_weight = 1e-4  # start here, can tune\nmax_epochs = 50\nfor epoch in range(max_epochs):\n    model.train()\n    total_loss = 0.0\n\n    for X, mf, bp, cc in train_loader:\n        X  = X.to(device)\n        mf = mf.to(device)\n        bp = bp.to(device)\n        cc = cc.to(device)\n\n        optimizer.zero_grad()\n\n        pred_mf, pred_bp, pred_cc = model(X)  # RAW LOGITS\n\n        loss_mf = criterion_mf(pred_mf, mf)\n        loss_bp = criterion_bp(pred_bp, bp)\n        loss_cc = criterion_cc(pred_cc, cc)\n\n        # 🔴 LOGIT REGULARIZATION\n        logit_penalty = (\n            pred_mf.mean() +\n            pred_bp.mean() +\n            pred_cc.mean()\n        )\n\n        loss = (\n            1.0 * loss_mf +\n            0.5 * loss_bp +\n            1.0 * loss_cc +\n            logit_reg_weight * logit_penalty\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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T21:27:42.046292Z","iopub.execute_input":"2026-01-23T21:27:42.046717Z","iopub.status.idle":"2026-01-23T21:29:39.367760Z","shell.execute_reply.started":"2026-01-23T21:27:42.046685Z","shell.execute_reply":"2026-01-23T21:29:39.367013Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ProteinInferenceDataset(Dataset):\n    def __init__(self, X):\n        if isinstance(X, torch.Tensor):\n            self.X = X.detach().clone().float()\n        else:\n            self.X = torch.tensor(X, dtype=torch.float32)\n\n    def __len__(self):\n        return len(self.X)\n\n    def __getitem__(self, idx):\n        return self.X[idx]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T21:29:39.368686Z","iopub.execute_input":"2026-01-23T21:29:39.369193Z","iopub.status.idle":"2026-01-23T21:29:39.373912Z","shell.execute_reply.started":"2026-01-23T21:29:39.369168Z","shell.execute_reply":"2026-01-23T21:29:39.373211Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\nfrom torch.utils.data import DataLoader\nimport numpy as np\n\nval_loader = DataLoader(\n    ProteinInferenceDataset(X_val),\n    batch_size=1024,\n    shuffle=False\n)\n\n\nmodel.eval()\nlr_val_mf, lr_val_bp, lr_val_cc = [], [], []\n\nwith torch.no_grad():\n    for X in val_loader:\n        X = X.to(device)\n        pmf, pbp, pcc = model(X)\n\n        lr_val_mf.append(torch.sigmoid(pmf).cpu())\n        lr_val_bp.append(torch.sigmoid(pbp).cpu())\n        lr_val_cc.append(torch.sigmoid(pcc).cpu())\n\nlr_val_mf = torch.cat(lr_val_mf).numpy()\nlr_val_bp = torch.cat(lr_val_bp).numpy()\nlr_val_cc = torch.cat(lr_val_cc).numpy()\n\nnp.save(\"lr_val_mf.npy\", lr_val_mf)\nnp.save(\"lr_val_bp.npy\", lr_val_bp)\nnp.save(\"lr_val_cc.npy\", lr_val_cc)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T21:29:39.374787Z","iopub.execute_input":"2026-01-23T21:29:39.375044Z","iopub.status.idle":"2026-01-23T21:29:40.083276Z","shell.execute_reply.started":"2026-01-23T21:29:39.375024Z","shell.execute_reply":"2026-01-23T21:29:40.082433Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def compute_fmax_with_details(preds, labels, thresholds=None):\n    \"\"\"\n    CAFA-style Fmax (protein-centric).\n    Computes precision/recall per protein, then averages.\n    \"\"\"\n    if thresholds is None:\n        thresholds = np.linspace(0.01, 0.99, 99)\n\n    # 🔑 FIX: ensure binary integer labels\n    labels = labels.astype(np.int32)\n\n    best_f = 0\n    best_t = 0\n    best_p = 0\n    best_r = 0\n\n    for t in thresholds:\n        preds_binary = (preds >= t).astype(np.int32)\n\n        # Per-protein counts\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        # Per-protein precision & recall\n        precision = np.divide(\n            tp, tp + fp, out=np.zeros_like(tp, dtype=float), where=(tp + fp) > 0\n        )\n        recall = np.divide(\n            tp, tp + fn, out=np.zeros_like(tp, dtype=float), where=(tp + fn) > 0\n        )\n\n        # Average over proteins with ≥1 true label\n        valid = np.sum(labels, axis=1) > 0\n        if valid.sum() == 0:\n            continue\n\n        p = precision[valid].mean()\n        r = recall[valid].mean()\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\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T21:29:40.086036Z","iopub.execute_input":"2026-01-23T21:29:40.086290Z","iopub.status.idle":"2026-01-23T21:29:40.093697Z","shell.execute_reply.started":"2026-01-23T21:29:40.086268Z","shell.execute_reply":"2026-01-23T21:29:40.093170Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fmax_mf, thresh_mf,prec_mf,recall_mf = compute_fmax_with_details(lr_val_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={recall_mf:.3f})\")\nprint(f\"{'='*60}\")\n    \nfmax_bp, thresh_bp,prec_bp,recall_bp = compute_fmax_with_details(lr_val_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={recall_bp:.3f})\")\nprint(f\"{'='*60}\")\nfmax_cc, thresh_cc,prec_cc,recall_cc = compute_fmax_with_details(lr_val_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={recall_cc:.3f})\")\nprint(f\"{'='*60}\")\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T21:29:40.094656Z","iopub.execute_input":"2026-01-23T21:29:40.094922Z","iopub.status.idle":"2026-01-23T21:31:02.176868Z","shell.execute_reply.started":"2026-01-23T21:29:40.094894Z","shell.execute_reply":"2026-01-23T21:31:02.175979Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_dataset = ProteinInferenceDataset(X_test)\ntest_loader = DataLoader(\n    test_dataset,\n    batch_size=1024,\n    shuffle=False,\n    pin_memory=True\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T21:31:02.178476Z","iopub.execute_input":"2026-01-23T21:31:02.178760Z","iopub.status.idle":"2026-01-23T21:31:02.551440Z","shell.execute_reply.started":"2026-01-23T21:31:02.178737Z","shell.execute_reply":"2026-01-23T21:31:02.550801Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\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\n        pred_mf.append(torch.sigmoid(pmf).cpu().numpy())\n        pred_bp.append(torch.sigmoid(pbp).cpu().numpy())\n        pred_cc.append(torch.sigmoid(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(\"lr_pred_mf.npy\",pred_mf)\nnp.save(\"lr_pred_bp.npy\",pred_bp)\nnp.save(\"lr_pred_cc.npy\",pred_cc)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T21:31:02.552344Z","iopub.execute_input":"2026-01-23T21:31:02.552635Z","iopub.status.idle":"2026-01-23T21:31:31.938721Z","shell.execute_reply.started":"2026-01-23T21:31:02.552603Z","shell.execute_reply":"2026-01-23T21:31:31.938049Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(pred_mf.shape, pred_bp.shape, pred_cc.shape)\nprint(pred_mf.min(), pred_mf.max())\nprint(pred_bp.min(), pred_bp.max())\nprint(pred_cc.min(), pred_cc.max())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T21:31:31.939634Z","iopub.execute_input":"2026-01-23T21:31:31.939994Z","iopub.status.idle":"2026-01-23T21:31:32.817200Z","shell.execute_reply.started":"2026-01-23T21:31:31.939960Z","shell.execute_reply":"2026-01-23T21:31:32.816432Z"}},"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\ngo_parents = {\n    go: go_dag[go].get_all_parents()\n    for go in set(mf_terms + bp_terms + cc_terms)\n    if go in go_dag\n}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T21:31:32.818245Z","iopub.execute_input":"2026-01-23T21:31:32.818808Z","iopub.status.idle":"2026-01-23T21:31:40.735339Z","shell.execute_reply.started":"2026-01-23T21:31:32.818778Z","shell.execute_reply":"2026-01-23T21:31:40.734532Z"}},"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\nimport pandas as pd\nimport numpy as np\nfrom collections import defaultdict\nfrom tqdm import tqdm\n\n# ============================================\n# LOAD OPTIMAL THRESHOLDS\n# ============================================\n\nprint(\"=\"*60)\nprint(\"CREATING CAFA SUBMISSION\")\nprint(\"=\"*60)\n\n\nprint(f\"\\nUsing optimal thresholds:\")\nprint(f\"  MF: {thresh_mf:.3f}\")\nprint(f\"  BP: {thresh_bp:.3f}\")\nprint(f\"  CC: {thresh_cc:.3f}\")\n\n# ============================================\n# LOAD PREDICTIONS (you already generated these)\n# ============================================\n\nprint(f\"\\nLoading test predictions...\")\nmf_pred = np.load(\"lr_pred_mf.npy\", mmap_mode=\"r\")\nbp_pred = np.load(\"lr_pred_bp.npy\", mmap_mode=\"r\")\ncc_pred = np.load(\"lr_pred_cc.npy\", mmap_mode=\"r\")\n\nprint(f\"Prediction shapes:\")\nprint(f\"  MF: {mf_pred.shape}\")\nprint(f\"  BP: {bp_pred.shape}\")\nprint(f\"  CC: {cc_pred.shape}\")\n\n\n# IMPORTANT: Make sure you use the FILTERED terms that match your model output dimensions!\nprint(f\"\\nGO term counts:\")\nprint(f\"  MF: {len(mf_terms)} terms\")\nprint(f\"  BP: {len(bp_terms)} terms\")\nprint(f\"  CC: {len(cc_terms)} terms\")\n\n# ============================================\n# CREATE SUBMISSION WITH PROPAGATION\n# ============================================\n\ndef stream_cafa_submission(\n    mf_pred, bp_pred, cc_pred,\n    thresh_mf, thresh_bp, thresh_cc,\n    mf_terms, bp_terms, cc_terms,\n    protein_ids, go_parents,\n    out_path,\n    max_terms=1500,\n    batch_size=256\n):\n    with open(out_path, \"w\") as f:\n        for i in tqdm(range(0, len(protein_ids), batch_size)):\n            end = min(i + batch_size, len(protein_ids))\n\n            for k in range(i, end):\n                pid = protein_ids[k]\n                scores = {}\n\n                # ---------- threshold ----------\n                for j, go in enumerate(mf_terms):\n                    s = mf_pred[k, j]\n                    if s >= thresh_mf:\n                        scores[go] = s\n\n                for j, go in enumerate(bp_terms):\n                    s = bp_pred[k, j]\n                    if s >= thresh_bp:\n                        scores[go] = max(scores.get(go, 0), s)\n\n                for j, go in enumerate(cc_terms):\n                    s = cc_pred[k, j]\n                    if s >= thresh_cc:\n                        scores[go] = max(scores.get(go, 0), s)\n\n                if not scores:\n                    continue\n\n                # ---------- propagate ----------\n                propagated = dict(scores)\n                for go, s in scores.items():\n                    for parent in go_parents.get(go, []):\n                        propagated[parent] = max(propagated.get(parent, 0), s)\n\n                # ---------- top 1500 ----------\n                items = sorted(propagated.items(), key=lambda x: x[1], reverse=True)\n                items = items[:max_terms]\n\n                # ---------- write ----------\n                for go, s in items:\n                    f.write(f\"{pid}\\t{go}\\t{s:.6f}\\n\")\n    \n\nstream_cafa_submission(\n    mf_pred, bp_pred, cc_pred,\n    thresh_mf, thresh_bp, thresh_cc,\n    mf_terms, bp_terms, cc_terms,\n    protein_ids, go_parents,\n    \"submission.tsv\",\n    max_terms=1500,\n    batch_size=256\n)\nprint(\"submission is complete\")\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-23T22:12:56.700625Z","iopub.execute_input":"2026-01-23T22:12:56.701038Z","iopub.status.idle":"2026-01-23T22:22:36.598033Z","shell.execute_reply.started":"2026-01-23T22:12:56.701008Z","shell.execute_reply":"2026-01-23T22:22:36.597274Z"}},"outputs":[],"execution_count":null}]}