{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":116062,"databundleVersionId":14084779,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":5549164,"sourceType":"datasetVersion","datasetId":3197305}],"dockerImageVersionId":31192,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# CAFA6 Protein Function: EDA + Baseline\n**Goal:** Show initial data exploration and a baseline model for the CAFA6 competition.  \n**Reader guide:** Short sections. Clear plots. One simple model. A valid submission.\n\n**Data sources**\n- `/kaggle/input/cafa-6-protein-function-prediction`\n- This project uses precomputed, fixed dimension protein embeddings from Kaggle to focus on model design and evaluation, not on deep network training. ESM2 and ProtBERT (public datasets on Kaggle)\n\n**What you will see**\n1) Load data and quick checks  \n2) Descriptive stats and plots  \n3) Baseline model with Logistic Regression on embeddings  \n4) Create a small, valid submission  \n5) Notes on limits and next steps\n","metadata":{}},{"cell_type":"code","source":"# Install lightweight deps\n!pip -q install obonet biopython\n\nimport os, gc, numpy as np, pandas as pd, matplotlib.pyplot as plt, seaborn as sns\nfrom pathlib import Path\nfrom collections import Counter, defaultdict\nfrom Bio import SeqIO\nimport obonet\nimport networkx as nx\n\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.linear_model import LogisticRegression\nfrom sklearn.multiclass import OneVsRestClassifier\nfrom sklearn.preprocessing import MultiLabelBinarizer\nfrom sklearn.metrics import f1_score, classification_report\nsns.set_theme()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-29T16:48:28.591872Z","iopub.execute_input":"2025-11-29T16:48:28.592162Z","iopub.status.idle":"2025-11-29T16:48:31.611740Z","shell.execute_reply.started":"2025-11-29T16:48:28.592141Z","shell.execute_reply":"2025-11-29T16:48:31.610896Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"BASE = Path(\"/kaggle/input/cafa-6-protein-function-prediction\")\nTRAIN_DIR = BASE/\"Train\"\nTEST_DIR  = BASE/\"Test\"\n\n# Embeddings (use either one; ESM2 is common and small enough)\nESM_DIR = Path(\"/kaggle/input/cafa-5-ems-2-embeddings-numpy\")\nPROT_DIR = Path(\"/kaggle/input/protbert-embeddings-for-cafa5\")\n\n\nSAMPLE_N = 3000       # number of training proteins to use for baseline\nTOP_K_TERMS = 300       # predict most frequent 300 terms overall\nVAL_SIZE = 0.2\nRANDOM_STATE = 42\nTHRESH = 0.2            # simple probability threshold for multilabel\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-29T16:52:15.522024Z","iopub.execute_input":"2025-11-29T16:52:15.522320Z","iopub.status.idle":"2025-11-29T16:52:15.527028Z","shell.execute_reply.started":"2025-11-29T16:52:15.522300Z","shell.execute_reply":"2025-11-29T16:52:15.526260Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"terms = pd.read_csv(TRAIN_DIR/\"train_terms.tsv\", sep=\"\\t\", names=[\"protein\",\"term\",\"aspect\"])\ntaxa  = pd.read_csv(TRAIN_DIR/\"train_taxonomy.tsv\", sep=\"\\t\", names=[\"protein\",\"taxon\"])\nprint(terms.shape, \"annotations\")\nprint(\"Unique proteins:\", terms[\"protein\"].nunique())\nprint(\"Unique terms:\", terms[\"term\"].nunique())\nprint(terms[\"aspect\"].value_counts())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-29T17:05:09.610965Z","iopub.execute_input":"2025-11-29T17:05:09.611706Z","iopub.status.idle":"2025-11-29T17:05:10.067841Z","shell.execute_reply.started":"2025-11-29T17:05:09.611685Z","shell.execute_reply":"2025-11-29T17:05:10.066733Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Terms per protein\ntpp = terms.groupby(\"protein\")[\"term\"].nunique()\n# Proteins per term\nppt = terms.groupby(\"term\")[\"protein\"].nunique()\n\nfig, axs = plt.subplots(1,3, figsize=(16,4))\nterms[\"aspect\"].value_counts().plot(kind=\"bar\", ax=axs[0], title=\"Annotations by aspect\")\naxs[0].set_xlabel(\"Aspect\")\n\naxs[1].hist(tpp, bins=50)\naxs[1].set_title(\"Terms per protein\")\n\naxs[2].hist(ppt, bins=50, log=True)\naxs[2].set_title(\"Proteins per term (log scale)\")\nplt.tight_layout()\nplt.show()\n\nprint(\"Mean terms/protein:\", tpp.mean())\nprint(\"Median terms/protein:\", tpp.median())\nprint(\"Mean proteins/term:\", ppt.mean())\nprint(\"Median proteins/term:\", ppt.median())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-29T17:06:17.279271Z","iopub.execute_input":"2025-11-29T17:06:17.279537Z","iopub.status.idle":"2025-11-29T17:06:18.631370Z","shell.execute_reply.started":"2025-11-29T17:06:17.279520Z","shell.execute_reply":"2025-11-29T17:06:18.630381Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_fasta = TRAIN_DIR/\"train_sequences.fasta\"\nseq_len = []\nseen = set()\nfor rec in SeqIO.parse(str(train_fasta), \"fasta\"):\n    pid = rec.id.split(\"|\")[1] if \"|\" in rec.id else rec.id\n    if pid in seen:\n        continue\n    seen.add(pid)\n    seq_len.append(len(rec.seq))\n\nplt.figure(figsize=(6,4))\nplt.hist(seq_len, bins=60)\nplt.title(\"Training sequence length\")\nplt.xlabel(\"Length\")\nplt.ylabel(\"Count\")\nplt.show()\n\nprint(\"Mean length:\", int(np.mean(seq_len)))\nprint(\"Median length:\", int(np.median(seq_len)))\nprint(\"Range:\", (int(np.min(seq_len)), int(np.max(seq_len))))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-29T17:07:03.166776Z","iopub.execute_input":"2025-11-29T17:07:03.167181Z","iopub.status.idle":"2025-11-29T17:07:04.420143Z","shell.execute_reply.started":"2025-11-29T17:07:03.167158Z","shell.execute_reply":"2025-11-29T17:07:04.419337Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"go = obonet.read_obo(TRAIN_DIR/\"go-basic.obo\")\nprint(\"GO nodes:\", len(go))\nnamespaces = pd.Series({n: go.nodes[n].get(\"namespace\",\"\") for n in go.nodes()})\nprint(namespaces.value_counts().head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-29T17:08:26.113897Z","iopub.execute_input":"2025-11-29T17:08:26.114529Z","iopub.status.idle":"2025-11-29T17:08:30.562294Z","shell.execute_reply.started":"2025-11-29T17:08:26.114506Z","shell.execute_reply":"2025-11-29T17:08:30.561377Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Use most frequent overall TOP_K_TERMS\ntop_terms = terms[\"term\"].value_counts().head(TOP_K_TERMS).index.tolist()\n\n# Build multilabel targets for proteins we have embeddings for\n# We will align to ESM2 train_ids.npy to avoid order mismatch\ntrain_ids = np.load(ESM_DIR/\"train_ids.npy\", allow_pickle=True)\ntrain_ids = train_ids.astype(str)\n\n# Group terms by protein for quick lookup\nby_protein = terms.groupby(\"protein\")[\"term\"].apply(list)\n\n# Create label matrix aligned to train_ids order\nmlb = MultiLabelBinarizer(classes=top_terms)\nlabels_list = [by_protein.get(pid, []) for pid in train_ids]\nY = mlb.fit_transform(labels_list)  # uses provided classes as order\n\nprint(\"Y shape:\", Y.shape, \"positive rate:\", Y.mean().round(4))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-29T17:09:00.896880Z","iopub.execute_input":"2025-11-29T17:09:00.897273Z","iopub.status.idle":"2025-11-29T17:09:03.220029Z","shell.execute_reply.started":"2025-11-29T17:09:00.897242Z","shell.execute_reply":"2025-11-29T17:09:03.208731Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\nfrom sklearn.model_selection import train_test_split\nimport numpy as np\n\n# Load embeddings\nX = np.load(ESM_DIR / \"train_embeddings.npy\")\n\n# Sample subset for speed\nrs = np.random.RandomState(RANDOM_STATE)\nidx = rs.choice(len(X), size=min(SAMPLE_N, len(X)), replace=False)\nXs = X[idx]\nYs = Y[idx]  # make sure Y is defined from your earlier label matrix step\n\n# Split into train and validation sets\nXtr, Xva, Ytr, Yva = train_test_split(Xs, Ys, test_size=VAL_SIZE, random_state=RANDOM_STATE)\n\nprint(\"Train shapes:\", Xtr.shape, Ytr.shape)\nprint(\"Val shapes:\", Xva.shape, Yva.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-29T17:09:51.414727Z","iopub.execute_input":"2025-11-29T17:09:51.414983Z","iopub.status.idle":"2025-11-29T17:09:51.624543Z","shell.execute_reply.started":"2025-11-29T17:09:51.414969Z","shell.execute_reply":"2025-11-29T17:09:51.623673Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nfor d, _, f in os.walk(\"/kaggle/input\"):\n    print(d)\n    for x in f:\n        print(\"   \", x)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-29T17:09:58.332000Z","iopub.execute_input":"2025-11-29T17:09:58.332268Z","iopub.status.idle":"2025-11-29T17:09:58.353282Z","shell.execute_reply.started":"2025-11-29T17:09:58.332253Z","shell.execute_reply":"2025-11-29T17:09:58.352470Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"terms = pd.read_csv(TRAIN_DIR/\"train_terms.tsv\", sep=\"\\t\", names=[\"protein\",\"term\",\"aspect\"])\ntaxa  = pd.read_csv(TRAIN_DIR/\"train_taxonomy.tsv\", sep=\"\\t\", names=[\"protein\",\"taxon\"])\nprint(terms.shape, \"annotations\")\nprint(\"Unique proteins:\", terms[\"protein\"].nunique())\nprint(\"Unique terms:\", terms[\"term\"].nunique())\nprint(terms[\"aspect\"].value_counts())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-29T17:10:03.607872Z","iopub.execute_input":"2025-11-29T17:10:03.608155Z","iopub.status.idle":"2025-11-29T17:10:03.903748Z","shell.execute_reply.started":"2025-11-29T17:10:03.608140Z","shell.execute_reply":"2025-11-29T17:10:03.903033Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Terms per protein\ntpp = terms.groupby(\"protein\")[\"term\"].nunique()\n# Proteins per term\nppt = terms.groupby(\"term\")[\"protein\"].nunique()\n\nfig, axs = plt.subplots(1,3, figsize=(16,4))\nterms[\"aspect\"].value_counts().plot(kind=\"bar\", ax=axs[0], title=\"Annotations by aspect\")\naxs[0].set_xlabel(\"Aspect\")\n\naxs[1].hist(tpp, bins=50)\naxs[1].set_title(\"Terms per protein\")\n\naxs[2].hist(ppt, bins=50, log=True)\naxs[2].set_title(\"Proteins per term (log scale)\")\nplt.tight_layout()\nplt.show()\n\nprint(\"Mean terms/protein:\", tpp.mean())\nprint(\"Median terms/protein:\", tpp.median())\nprint(\"Mean proteins/term:\", ppt.mean())\nprint(\"Median proteins/term:\", ppt.median())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-29T17:10:07.665028Z","iopub.execute_input":"2025-11-29T17:10:07.665325Z","iopub.status.idle":"2025-11-29T17:10:09.144032Z","shell.execute_reply.started":"2025-11-29T17:10:07.665305Z","shell.execute_reply":"2025-11-29T17:10:09.142999Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_fasta = TRAIN_DIR/\"train_sequences.fasta\"\nseq_len = []\nseen = set()\nfor rec in SeqIO.parse(str(train_fasta), \"fasta\"):\n    pid = rec.id.split(\"|\")[1] if \"|\" in rec.id else rec.id\n    if pid in seen:\n        continue\n    seen.add(pid)\n    seq_len.append(len(rec.seq))\n\nplt.figure(figsize=(6,4))\nplt.hist(seq_len, bins=60)\nplt.title(\"Training sequence length\")\nplt.xlabel(\"Length\")\nplt.ylabel(\"Count\")\nplt.show()\n\nprint(\"Mean length:\", int(np.mean(seq_len)))\nprint(\"Median length:\", int(np.median(seq_len)))\nprint(\"Range:\", (int(np.min(seq_len)), int(np.max(seq_len))))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-29T17:10:13.833509Z","iopub.execute_input":"2025-11-29T17:10:13.833775Z","iopub.status.idle":"2025-11-29T17:10:14.570225Z","shell.execute_reply.started":"2025-11-29T17:10:13.833761Z","shell.execute_reply":"2025-11-29T17:10:14.569296Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"go = obonet.read_obo(TRAIN_DIR/\"go-basic.obo\")\nprint(\"GO nodes:\", len(go))\nnamespaces = pd.Series({n: go.nodes[n].get(\"namespace\",\"\") for n in go.nodes()})\nprint(namespaces.value_counts().head())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-29T17:10:19.247162Z","iopub.execute_input":"2025-11-29T17:10:19.247444Z","iopub.status.idle":"2025-11-29T17:10:23.524131Z","shell.execute_reply.started":"2025-11-29T17:10:19.247426Z","shell.execute_reply":"2025-11-29T17:10:23.521923Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# How many labels to model\nTOP_K_TERMS = 300  \n\n# Compute the top-K most frequent GO terms from the training annotations\ntop_terms = terms[\"term\"].value_counts().head(TOP_K_TERMS).index.tolist()\nprint(\"Top terms:\", len(top_terms))\n\n# Build multilabel targets for proteins we have embeddings for\n# We will align to ESM2 train_ids.npy to avoid order mismatch\ntrain_ids = np.load(ESM_DIR/\"train_ids.npy\", allow_pickle=True)\ntrain_ids = train_ids.astype(str)\n\n# Group terms by protein for quick lookup\nby_protein = terms.groupby(\"protein\")[\"term\"].apply(list)\n\n# Keep only top-K terms in each protein's label list\ntop_set = set(top_terms)\nlabels_list = [[t for t in by_protein.get(pid, []) if t in top_set] for pid in train_ids]\n\n# Create label matrix aligned to train_ids order\nfrom sklearn.preprocessing import MultiLabelBinarizer\nmlb = MultiLabelBinarizer(classes=top_terms)\nY = mlb.fit_transform(labels_list)\n\nprint(\"Y shape:\", Y.shape, \"positive rate:\", Y.mean().round(4))\n\n\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-29T17:11:21.205391Z","iopub.execute_input":"2025-11-29T17:11:21.205631Z","iopub.status.idle":"2025-11-29T17:11:23.969628Z","shell.execute_reply.started":"2025-11-29T17:11:21.205616Z","shell.execute_reply":"2025-11-29T17:11:23.969031Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"X = np.load(ESM_DIR/\"train_embeddings.npy\")\nprint(\"X shape:\", X.shape)\n\nif SAMPLE_N < len(train_ids):\n    idx = np.random.RandomState(RANDOM_STATE).choice(len(train_ids), size=SAMPLE_N, replace=False)\n    Xs = X[idx]\n    Ys = Y[idx]\n    ids_s = train_ids[idx]\nelse:\n    Xs, Ys, ids_s = X, Y, train_ids\n\nprint(\"Sampled:\", Xs.shape, Ys.shape)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-29T17:11:32.716926Z","iopub.execute_input":"2025-11-29T17:11:32.717159Z","iopub.status.idle":"2025-11-29T17:11:32.884836Z","shell.execute_reply.started":"2025-11-29T17:11:32.717144Z","shell.execute_reply":"2025-11-29T17:11:32.884101Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"Xs:\", Xs.shape, \"Ys:\", Ys.shape)\nprint(\"Split:\", Xtr.shape, Ytr.shape, \"|\", Xva.shape, Yva.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-29T17:11:36.002902Z","iopub.execute_input":"2025-11-29T17:11:36.003485Z","iopub.status.idle":"2025-11-29T17:11:36.008383Z","shell.execute_reply.started":"2025-11-29T17:11:36.003469Z","shell.execute_reply":"2025-11-29T17:11:36.007494Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.pipeline import make_pipeline\nfrom sklearn.preprocessing import StandardScaler\nfrom sklearn.decomposition import PCA\nfrom sklearn.linear_model import LogisticRegression\nfrom sklearn.multiclass import OneVsRestClassifier\nfrom sklearn.metrics import f1_score\nimport numpy as np\n\nimport warnings\nfrom sklearn.exceptions import ConvergenceWarning\nwarnings.filterwarnings(\"ignore\", category=ConvergenceWarning)\n\n# faster pipeline\npipe = make_pipeline(\n    StandardScaler(),\n    PCA(n_components=128, random_state=42),   # compress 1280 -> 128\n    LogisticRegression(\n        solver=\"saga\",\n        max_iter=200,\n        tol=1e-3,\n        C=1.0,\n        class_weight=\"balanced\",\n        n_jobs=-1\n    )\n)\n\nbase_clf = OneVsRestClassifier(pipe, n_jobs=-1)\n\n# train\nimport time\nstart = time.time()\n\nbase_clf.fit(Xtr, Ytr)\n\nprint(\"Training time:\", round(time.time() - start, 2), \"seconds\")\n\nbase_clf.fit(Xtr, Ytr)\n\n# robust probs\nTHRESH = 0.2\ndef _sigmoid(z): return 1/(1+np.exp(-z))\ntry:\n    probs = base_clf.predict_proba(Xva)\nexcept Exception:\n    s = base_clf.decision_function(Xva)\n    if isinstance(s, list):\n        s = np.column_stack(s)\n    probs = _sigmoid(s)\n\nYhat = (probs >= THRESH).astype(int)\nprint(\"Micro F1:\", round(f1_score(Yva, Yhat, average=\"micro\"), 4))\nprint(\"Macro F1:\", round(f1_score(Yva, Yhat, average=\"macro\"), 4))\nprint(\"Pred density:\", round(Yhat.mean(), 4))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-29T17:11:41.757784Z","iopub.execute_input":"2025-11-29T17:11:41.758177Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np, pandas as pd\n\n# Load test embeddings and ids\ntest_ids = np.load(ESM_DIR/\"test_ids.npy\", allow_pickle=True).astype(str)\ntest_emb = np.load(ESM_DIR/\"test_embeddings.npy\")\n\n# Predict probabilities safely\ndef _sigmoid(z): return 1/(1+np.exp(-z))\ntry:\n    test_probs = base_clf.predict_proba(test_emb)\nexcept Exception:\n    s = base_clf.decision_function(test_emb)\n    if isinstance(s, list):\n        s = np.column_stack(s)\n    test_probs = _sigmoid(s)\n\n# Build a term list for writing\nterms_out = mlb.classes_.tolist()  # same order as columns\n\n# Write a valid small submission\nTHRESH = 0.2\nMAX_PER_PROTEIN = 200\nout_path = \"submission.tsv\"\n\nwith open(out_path, \"w\") as f:\n    for i, pid in enumerate(test_ids):\n        scores = test_probs[i]\n        keep = np.where(scores >= THRESH)[0]\n        if len(keep) == 0:\n            continue\n        order = np.argsort(scores[keep])[::-1][:MAX_PER_PROTEIN]\n        for j in keep[order]:\n            f.write(f\"{pid}\\t{terms_out[j]}\\t{scores[j]:.3f}\\n\")\n\n# Preview\nsub_df = pd.read_csv(out_path, sep=\"\\t\", header=None, names=[\"Id\",\"GO_term\",\"Confidence\"])\nprint(sub_df.head())\nprint(\"Rows:\", len(sub_df))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-29T17:33:11.759602Z","iopub.execute_input":"2025-11-29T17:33:11.759849Z","execution_failed":"2025-11-29T18:45:52.789Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Conclusion\nBaseline on precomputed ESM2 embeddings with logistic regression.\nValid split: 80/20 on 3k samples.\nResults: Micro F1 = 0.076, Macro F1 = 0.034 at threshold 0.2.\nWe created a competition-ready submission (3.22M rows).\nNext steps: threshold tuning, per-ontology models, class weights, ancestor propagation.","metadata":{}}]}