{"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":14875579,"sourceType":"competition"}],"dockerImageVersionId":31192,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ============================================================\n# CAFA6 EDA — FULL END-TO-END PIPELINE (Kaggle)\n# - Robust to missing headers (train_taxonomy.tsv, IA.tsv)\n# - Correctly handles testsuperset-taxon-list.tsv as: taxon_id -> taxon_name\n# - Exports: figures (PNG+PDF, 300dpi) + CSV tables to /kaggle/working/eda_out\n# ============================================================\n\n# %% [1] Setup\nimport os, re\nfrom pathlib import Path\nfrom collections import Counter\nfrom itertools import combinations\n\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\n\nplt.rcParams.update({\n    \"figure.dpi\": 120,\n    \"savefig.dpi\": 300,\n    \"font.size\": 10,\n    \"axes.titlesize\": 11,\n    \"axes.labelsize\": 10,\n    \"legend.fontsize\": 9,\n})\n\nDATA_DIR  = Path(\"/kaggle/input/cafa-6-protein-function-prediction\")\nTRAIN_DIR = DATA_DIR / \"Train\"\nTEST_DIR  = DATA_DIR / \"Test\"\n\nOUT_DIR = Path(\"/kaggle/working/eda_out\")\nFIG_DIR = OUT_DIR / \"figures\"\nTAB_DIR = OUT_DIR / \"tables\"\nFIG_DIR.mkdir(parents=True, exist_ok=True)\nTAB_DIR.mkdir(parents=True, exist_ok=True)\n\nprint(\"DATA_DIR:\", DATA_DIR)\nfor p in [DATA_DIR, TRAIN_DIR, TEST_DIR]:\n    print(f\"\\nListing: {p}\")\n    for f in sorted(p.glob(\"*\")):\n        print(\" -\", f.name)\n\ndef savefig(name: str):\n    png = FIG_DIR / f\"{name}.png\"\n    pdf = FIG_DIR / f\"{name}.pdf\"\n    plt.savefig(png, bbox_inches=\"tight\")\n    plt.savefig(pdf, bbox_inches=\"tight\")\n    print(\"Saved:\", png.name, \"and\", pdf.name)\n\n# %% [2] FASTA streaming utilities (no full sequence store)\ndef iter_fasta(path: Path):\n    header = None\n    seq_chunks = []\n    with open(path, \"r\") as f:\n        for line in f:\n            line = line.strip()\n            if not line:\n                continue\n            if line.startswith(\">\"):\n                if header is not None:\n                    yield header, \"\".join(seq_chunks)\n                header = line[1:].split()[0]\n                seq_chunks = []\n            else:\n                seq_chunks.append(line)\n        if header is not None:\n            yield header, \"\".join(seq_chunks)\n\ndef fasta_lengths_and_composition(path: Path):\n    ids, lens = [], []\n    aa_counter = Counter()\n    other_counter = Counter()\n    AA20 = set(\"ACDEFGHIKLMNPQRSTVWY\")\n    for pid, seq in iter_fasta(path):\n        ids.append(pid)\n        lens.append(len(seq))\n        for ch in seq:\n            if ch in AA20:\n                aa_counter[ch] += 1\n            else:\n                other_counter[ch] += 1\n    return pd.DataFrame({\"EntryID\": ids, \"seq_len\": lens}), aa_counter, other_counter\n\n# %% [3] Robust TSV readers\ndef read_two_col_tsv(path: Path, names=(\"col1\", \"col2\"), header_expected_first_tokens=None):\n    \"\"\"\n    Reads a 2-column TSV that may or may not have a header.\n    If header_expected_first_tokens is provided (set of lowercase strings),\n    it uses that to detect header from first cell.\n    \"\"\"\n    first = path.open(\"r\").readline().strip().split(\"\\t\")\n    first0 = first[0].strip().lower() if first else \"\"\n\n    header_like = False\n    if header_expected_first_tokens is not None:\n        header_like = first0 in header_expected_first_tokens\n\n    if header_like:\n        df = pd.read_csv(path, sep=\"\\t\", dtype=str)\n        df = df.iloc[:, :2].copy()\n        df.columns = list(names)\n    else:\n        df = pd.read_csv(path, sep=\"\\t\", header=None, names=list(names), dtype=str)\n        df = df.iloc[:, :2].copy()\n        df.columns = list(names)\n\n    return df\n\ndef read_train_taxonomy(path: Path):\n    # train_taxonomy.tsv: EntryID \\t taxon_id (often no header)\n    return read_two_col_tsv(\n        path,\n        names=(\"EntryID\", \"taxon_id\"),\n        header_expected_first_tokens={\"entryid\", \"entry_id\", \"entry\", \"protein\", \"protein_id\", \"id\"},\n    )\n\ndef read_ia(path: Path):\n    # IA.tsv: GO:xxxxxxx \\t float  (often no header)\n    first = path.open(\"r\").readline().strip().split(\"\\t\")\n    go_like = bool(first) and bool(re.match(r\"^GO:\\d{7}$\", first[0].strip()))\n    if go_like:\n        df = pd.read_csv(path, sep=\"\\t\", header=None, names=[\"term\", \"ia\"], dtype=str)\n    else:\n        df = pd.read_csv(path, sep=\"\\t\", dtype=str).iloc[:, :2].copy()\n        df.columns = [\"term\", \"ia\"]\n    df[\"ia\"] = pd.to_numeric(df[\"ia\"], errors=\"coerce\")\n    return df\n\ndef read_test_taxon_list(path: Path):\n    # testsuperset-taxon-list.tsv is taxon_id -> taxon_name (your log confirms)\n    df = read_two_col_tsv(\n        path,\n        names=(\"taxon_id\", \"taxon_name\"),\n        header_expected_first_tokens={\"taxon\", \"taxon_id\", \"taxid\", \"taxonomy\", \"taxonid\"},\n    )\n    return df\n\n# %% [4] Load core files\ntrain_terms = pd.read_csv(TRAIN_DIR / \"train_terms.tsv\", sep=\"\\t\", dtype=str)\ntrain_tax   = read_train_taxonomy(TRAIN_DIR / \"train_taxonomy.tsv\")\ntest_taxons = read_test_taxon_list(TEST_DIR  / \"testsuperset-taxon-list.tsv\")\nia_df       = read_ia(DATA_DIR / \"IA.tsv\")\n\nprint(\"\\ntrain_terms:\", train_terms.shape, train_terms.columns.tolist())\nprint(\"train_tax  :\", train_tax.shape, train_tax.columns.tolist(), \"\\n\", train_tax.head())\nprint(\"test_taxons:\", test_taxons.shape, test_taxons.columns.tolist(), \"\\n\", test_taxons.head())\nprint(\"IA         :\", ia_df.shape, ia_df.columns.tolist(), \"\\n\", ia_df.head())\n\n# Expected columns in train_terms\nassert {\"EntryID\", \"term\"}.issubset(set(train_terms.columns)), \"train_terms.tsv must contain EntryID and term\"\nASPECT_COL = \"aspect\" if \"aspect\" in train_terms.columns else None\n\n# %% [5] Load FASTA stats\ntrain_seq_df, train_aa, train_other = fasta_lengths_and_composition(TRAIN_DIR / \"train_sequences.fasta\")\ntest_seq_df,  test_aa,  test_other  = fasta_lengths_and_composition(TEST_DIR  / \"testsuperset.fasta\")\n\nprint(\"\\ntrain_seq_df:\", train_seq_df.shape, \"test_seq_df:\", test_seq_df.shape)\nprint(\"Non-20AA tokens in train (top):\", train_other.most_common(10))\nprint(\"Non-20AA tokens in test  (top):\", test_other.most_common(10))\n\n# %% [6] Sanity checks: ID coverage\ntrain_seq_ids = set(train_seq_df[\"EntryID\"])\nterms_ids = set(train_terms[\"EntryID\"])\ntax_ids = set(train_tax[\"EntryID\"])\n\ncoverage = {\n    \"n_train_seq_ids\": len(train_seq_ids),\n    \"n_train_terms_proteins\": len(terms_ids),\n    \"n_train_tax_proteins\": len(tax_ids),\n    \"terms_in_sequences_%\": 100.0 * len(terms_ids & train_seq_ids) / max(1, len(terms_ids)),\n    \"tax_in_sequences_%\":   100.0 * len(tax_ids & train_seq_ids) / max(1, len(tax_ids)),\n}\ncoverage_df = pd.DataFrame([coverage])\ndisplay(coverage_df)\ncoverage_df.to_csv(TAB_DIR / \"table_id_coverage.csv\", index=False)\n\n# %% [7] Labels per protein + term frequency\nlabels_per_prot = train_terms.groupby(\"EntryID\")[\"term\"].size().rename(\"n_labels\").reset_index()\nterm_freq = train_terms[\"term\"].value_counts().rename_axis(\"term\").reset_index(name=\"freq\")\n\n# Merge with seq length + taxonomy\ntrain_merged = train_seq_df.merge(labels_per_prot, on=\"EntryID\", how=\"left\")\ntrain_merged[\"n_labels\"] = train_merged[\"n_labels\"].fillna(0).astype(int)\n\ntrain_tax[\"taxon_id\"] = train_tax[\"taxon_id\"].astype(str)\ntrain_merged = train_merged.merge(train_tax, on=\"EntryID\", how=\"left\")\n\n# %% [8] Dataset summary table (IEEE Table I)\nsummary = {\n    \"n_train_proteins\": int(train_seq_df.shape[0]),\n    \"n_test_proteins\":  int(test_seq_df.shape[0]),\n    \"n_train_annotations\": int(train_terms.shape[0]),\n    \"n_unique_terms\": int(train_terms[\"term\"].nunique()),\n    \"avg_labels_per_train_protein\": float(train_merged[\"n_labels\"].mean()),\n    \"median_labels_per_train_protein\": float(train_merged[\"n_labels\"].median()),\n    \"avg_train_seq_len\": float(train_seq_df[\"seq_len\"].mean()),\n    \"median_train_seq_len\": float(train_seq_df[\"seq_len\"].median()),\n    \"avg_test_seq_len\": float(test_seq_df[\"seq_len\"].mean()),\n    \"median_test_seq_len\": float(test_seq_df[\"seq_len\"].median()),\n    \"n_train_taxa\": int(train_tax[\"taxon_id\"].nunique()),\n    \"n_test_taxa_listed\": int(test_taxons[\"taxon_id\"].nunique()),\n    \"n_taxa_overlap(train_vs_test_list)\": int(len(set(train_tax[\"taxon_id\"]) & set(test_taxons[\"taxon_id\"]))),\n}\nsummary_df = pd.DataFrame([summary])\ndisplay(summary_df)\nsummary_df.to_csv(TAB_DIR / \"table_dataset_summary.csv\", index=False)\n\n# %% [9] Plot: sequence length distribution (Train vs Test)\nplt.figure(figsize=(6.2, 4.0))\nplt.hist(train_seq_df[\"seq_len\"], bins=120, alpha=0.6, label=\"Train\")\nplt.hist(test_seq_df[\"seq_len\"],  bins=120, alpha=0.6, label=\"Test\")\nplt.yscale(\"log\")\nplt.xlabel(\"Sequence length (aa)\")\nplt.ylabel(\"Count (log scale)\")\nplt.title(\"CAFA6: Sequence length distribution (Train vs Test)\")\nplt.legend()\nsavefig(\"fig_seq_length_train_vs_test\")\nplt.show()\n\n# %% [10] Plot: labels per protein distribution\nplt.figure(figsize=(6.2, 4.0))\nplt.hist(train_merged[\"n_labels\"], bins=60)\nplt.yscale(\"log\")\nplt.xlabel(\"#GO terms per protein\")\nplt.ylabel(\"Count (log scale)\")\nplt.title(\"CAFA6: Number of GO labels per protein (Train)\")\nsavefig(\"fig_labels_per_protein\")\nplt.show()\n\n# %% [11] Plot: seq length vs #labels (sample scatter)\nsample = train_merged.sample(n=min(60000, len(train_merged)), random_state=0)\nplt.figure(figsize=(6.2, 4.0))\nplt.scatter(sample[\"seq_len\"], sample[\"n_labels\"], s=6, alpha=0.15)\nplt.xlabel(\"Sequence length (aa)\")\nplt.ylabel(\"#GO terms\")\nplt.title(\"CAFA6: Sequence length vs #labels (Train sample)\")\nsavefig(\"fig_len_vs_labels_scatter\")\nplt.show()\n\n# %% [12] Plot: GO term long-tail + Top-20\nrank_freq = term_freq.sort_values(\"freq\", ascending=False).reset_index(drop=True)\nrank_freq[\"rank\"] = np.arange(1, len(rank_freq) + 1)\n\nplt.figure(figsize=(6.2, 4.0))\nplt.plot(rank_freq[\"rank\"], rank_freq[\"freq\"])\nplt.xscale(\"log\")\nplt.yscale(\"log\")\nplt.xlabel(\"Term rank (log)\")\nplt.ylabel(\"Frequency (log)\")\nplt.title(\"CAFA6: GO term frequency long-tail (Train)\")\nsavefig(\"fig_term_rank_frequency\")\nplt.show()\n\ntopk = 20\nplt.figure(figsize=(6.2, 4.2))\ntop_terms = rank_freq.head(topk).iloc[::-1]\nplt.barh(top_terms[\"term\"].astype(str), top_terms[\"freq\"])\nplt.xlabel(\"Frequency\")\nplt.title(f\"CAFA6: Top-{topk} most frequent GO terms (Train)\")\nsavefig(\"fig_top20_terms\")\nplt.show()\n\n# %% [13] Plot: aspect distribution (if available)\nif ASPECT_COL is not None:\n    aspect_counts = train_terms[ASPECT_COL].value_counts()\n    aspect_df = aspect_counts.rename_axis(\"aspect\").reset_index(name=\"count\")\n    display(aspect_df)\n    aspect_df.to_csv(TAB_DIR / \"table_aspect_counts.csv\", index=False)\n\n    plt.figure(figsize=(6.0, 3.8))\n    plt.bar(aspect_counts.index.astype(str), aspect_counts.values)\n    plt.xlabel(\"Aspect\")\n    plt.ylabel(\"Annotations\")\n    plt.title(\"CAFA6: Annotation counts by aspect (Train)\")\n    savefig(\"fig_aspect_counts\")\n    plt.show()\n\n# %% [14] Taxonomy (Train): top taxa + seq length by taxa\ntrain_tax_counts = train_tax[\"taxon_id\"].value_counts().reset_index()\ntrain_tax_counts.columns = [\"taxon_id\", \"n_proteins\"]\ntrain_tax_counts.to_csv(TAB_DIR / \"table_train_taxon_counts.csv\", index=False)\n\nplt.figure(figsize=(6.2, 4.2))\ntop_tax = train_tax_counts.head(20).iloc[::-1]\nplt.barh(top_tax[\"taxon_id\"].astype(str), top_tax[\"n_proteins\"])\nplt.xlabel(\"Number of proteins\")\nplt.title(\"CAFA6: Top-20 taxa in Train\")\nsavefig(\"fig_top20_taxa_train\")\nplt.show()\n\ntop10_taxa = train_tax_counts.head(10)[\"taxon_id\"].tolist()\nbox_data = [train_merged.loc[train_merged[\"taxon_id\"] == t, \"seq_len\"].dropna().values for t in top10_taxa]\n\nplt.figure(figsize=(6.8, 4.2))\nplt.boxplot(box_data, labels=[str(t) for t in top10_taxa], showfliers=False)\nplt.xticks(rotation=45, ha=\"right\")\nplt.ylabel(\"Sequence length (aa)\")\nplt.title(\"CAFA6: Sequence length by taxon (Top-10 taxa, Train)\")\nsavefig(\"fig_seq_len_by_taxon_top10\")\nplt.show()\n\n# %% [15] Test taxa list (taxon_id -> taxon_name): overlap with train\ntest_taxons[\"taxon_id\"] = test_taxons[\"taxon_id\"].astype(str)\ntest_taxons.to_csv(TAB_DIR / \"table_test_taxon_list.csv\", index=False)\n\noverlap_ids = sorted(set(train_tax[\"taxon_id\"]) & set(test_taxons[\"taxon_id\"]))\noverlap_df = test_taxons[test_taxons[\"taxon_id\"].isin(overlap_ids)].sort_values(\"taxon_id\")\noverlap_df.to_csv(TAB_DIR / \"table_taxon_overlap_train_vs_test_list.csv\", index=False)\nprint(\"Taxa overlap count:\", len(overlap_ids))\ndisplay(overlap_df.head(20))\n\n# %% [16] Amino acid composition (Train)\nAA20 = list(\"ACDEFGHIKLMNPQRSTVWY\")\ntrain_total_aa = sum(train_aa.values()) if sum(train_aa.values()) > 0 else 1\naa_freqs = np.array([train_aa[a] / train_total_aa for a in AA20])\naa_df = pd.DataFrame({\"aa\": AA20, \"freq\": aa_freqs})\naa_df.to_csv(TAB_DIR / \"table_amino_acid_freq_train.csv\", index=False)\n\nplt.figure(figsize=(7.0, 3.8))\nplt.bar(AA20, aa_freqs)\nplt.xlabel(\"Amino acid\")\nplt.ylabel(\"Relative frequency\")\nplt.title(\"CAFA6: Amino acid composition (Train, global)\")\nsavefig(\"fig_amino_acid_composition_train\")\nplt.show()\n\n# %% [17] IA analysis: distribution + IA vs frequency\nterm_stats = term_freq.merge(ia_df.drop_duplicates(\"term\"), on=\"term\", how=\"left\")\nterm_stats.to_csv(TAB_DIR / \"table_term_freq_with_ia.csv\", index=False)\n\nplt.figure(figsize=(6.2, 4.0))\nplt.hist(term_stats[\"ia\"].dropna(), bins=80)\nplt.xlabel(\"Information Accretion (IA)\")\nplt.ylabel(\"Number of terms\")\nplt.title(\"CAFA6: IA distribution for observed GO terms (Train)\")\nsavefig(\"fig_ia_distribution\")\nplt.show()\n\nscatter_df = term_stats.dropna(subset=[\"ia\"]).copy()\nplt.figure(figsize=(6.2, 4.0))\nplt.scatter(np.log10(scatter_df[\"freq\"]), scatter_df[\"ia\"], s=10, alpha=0.25)\nplt.xlabel(\"log10(term frequency)\")\nplt.ylabel(\"IA\")\nplt.title(\"CAFA6: IA vs term frequency (Train)\")\nsavefig(\"fig_ia_vs_freq\")\nplt.show()\n\n# %% [18] GO ontology parsing (go-basic.obo): namespace + depth\nobo_path = TRAIN_DIR / \"go-basic.obo\"\nprint(\"Parsing OBO:\", obo_path)\n\ndef parse_go_basic_obo(path: Path):\n    terms = {}\n    cur = None\n    in_term = False\n    with open(path, \"r\") as f:\n        for line in f:\n            line = line.rstrip(\"\\n\")\n            if line == \"[Term]\":\n                in_term = True\n                cur = {\"id\": None, \"namespace\": None, \"parents\": []}\n                continue\n            if line.startswith(\"[\") and line != \"[Term]\":\n                in_term = False\n                cur = None\n                continue\n            if not in_term or cur is None:\n                continue\n            if not line:\n                if cur.get(\"id\"):\n                    terms[cur[\"id\"]] = {\"namespace\": cur[\"namespace\"], \"parents\": cur[\"parents\"]}\n                cur = {\"id\": None, \"namespace\": None, \"parents\": []}\n                continue\n            if line.startswith(\"id:\"):\n                cur[\"id\"] = line.split(\"id:\")[1].strip()\n            elif line.startswith(\"namespace:\"):\n                cur[\"namespace\"] = line.split(\"namespace:\")[1].strip()\n            elif line.startswith(\"is_a:\"):\n                parent = line.split(\"is_a:\")[1].strip().split()[0]\n                cur[\"parents\"].append(parent)\n        if in_term and cur and cur.get(\"id\"):\n            terms[cur[\"id\"]] = {\"namespace\": cur[\"namespace\"], \"parents\": cur[\"parents\"]}\n    return terms\n\ngo_terms = parse_go_basic_obo(obo_path)\nprint(\"Parsed GO terms:\", len(go_terms))\n\nROOTS = {\"GO:0008150\", \"GO:0003674\", \"GO:0005575\"}  # BP, MF, CC\ndepth_memo = {}\n\ndef go_depth(go_id: str):\n    if go_id in depth_memo:\n        return depth_memo[go_id]\n    if go_id in ROOTS:\n        depth_memo[go_id] = 0\n        return 0\n    node = go_terms.get(go_id)\n    if node is None:\n        depth_memo[go_id] = np.nan\n        return np.nan\n    parents = node.get(\"parents\", [])\n    if not parents:\n        depth_memo[go_id] = 0\n        return 0\n    ds = []\n    for p in parents:\n        d = go_depth(p)\n        if not (isinstance(d, float) and np.isnan(d)):\n            ds.append(d)\n    if not ds:\n        depth_memo[go_id] = 0\n        return 0\n    depth_memo[go_id] = 1 + max(ds)\n    return depth_memo[go_id]\n\nobserved_terms = pd.DataFrame({\"term\": term_freq[\"term\"]})\nobserved_terms[\"namespace\"] = observed_terms[\"term\"].map(lambda t: go_terms.get(t, {}).get(\"namespace\"))\nobserved_terms[\"depth\"] = observed_terms[\"term\"].map(go_depth)\nobserved_terms.to_csv(TAB_DIR / \"table_observed_terms_namespace_depth.csv\", index=False)\n\nns_counts = observed_terms[\"namespace\"].value_counts(dropna=False)\nplt.figure(figsize=(6.2, 4.0))\nplt.bar(ns_counts.index.astype(str), ns_counts.values)\nplt.xticks(rotation=30, ha=\"right\")\nplt.ylabel(\"Number of observed terms\")\nplt.title(\"CAFA6: Observed GO terms by namespace (Train)\")\nsavefig(\"fig_terms_by_namespace\")\nplt.show()\n\ndepth_vals = observed_terms[\"depth\"].dropna()\nplt.figure(figsize=(6.2, 4.0))\nplt.hist(depth_vals, bins=np.arange(0, int(depth_vals.max()) + 2) - 0.5)\nplt.xlabel(\"GO term depth (from root)\")\nplt.ylabel(\"Number of observed terms\")\nplt.title(\"CAFA6: Depth distribution of observed GO terms (Train)\")\nsavefig(\"fig_term_depth_distribution\")\nplt.show()\n\n# %% [19] Co-occurrence heatmap (Top-N terms) — memory-safe (no labels list for all proteins)\nTOPN = 30\ntop_terms_list = rank_freq.head(TOPN)[\"term\"].tolist()\ntop_set = set(top_terms_list)\n\n# per-term counts for TopN\ntop_counts = dict(zip(term_freq[\"term\"], term_freq[\"freq\"]))\nc_i = {t: int(top_counts.get(t, 0)) for t in top_terms_list}\n\n# filter train_terms to only top terms, then iterate by EntryID\ntt = train_terms.loc[train_terms[\"term\"].isin(top_set), [\"EntryID\", \"term\"]].copy()\ntt.sort_values(\"EntryID\", inplace=True)\n\npair_counts = Counter()\ncur_id = None\ncur_terms = set()\n\ndef flush_terms(s):\n    if len(s) >= 2:\n        for a, b in combinations(sorted(s), 2):\n            pair_counts[(a, b)] += 1\n\nfor row in tt.itertuples(index=False):\n    pid, term = row.EntryID, row.term\n    if cur_id is None:\n        cur_id = pid\n    if pid != cur_id:\n        flush_terms(cur_terms)\n        cur_terms = set()\n        cur_id = pid\n    cur_terms.add(term)\nflush_terms(cur_terms)\n\nidx = {t: i for i, t in enumerate(top_terms_list)}\nM = np.zeros((TOPN, TOPN), dtype=float)\nfor i in range(TOPN):\n    M[i, i] = 1.0\n\nfor (a, b), cij in pair_counts.items():\n    denom = c_i[a] + c_i[b] - cij\n    j = (cij / denom) if denom > 0 else 0.0\n    ia, ib = idx[a], idx[b]\n    M[ia, ib] = j\n    M[ib, ia] = j\n\nplt.figure(figsize=(7.2, 6.2))\nplt.imshow(M, aspect=\"auto\")\nplt.xticks(range(TOPN), top_terms_list, rotation=90, fontsize=7)\nplt.yticks(range(TOPN), top_terms_list, fontsize=7)\nplt.title(f\"CAFA6: GO term co-occurrence (Jaccard) among Top-{TOPN} terms\")\nplt.colorbar(label=\"Jaccard\")\nsavefig(\"fig_term_cooccurrence_top30_jaccard\")\nplt.show()\n\n# %% [20] Export rank-frequency table + final file list\nrank_freq.to_csv(TAB_DIR / \"table_term_rank_frequency.csv\", index=False)\n\nprint(\"\\nDONE. Outputs written to:\", OUT_DIR)\nprint(\"\\nFigures:\", len(list(FIG_DIR.glob(\"*.png\"))), \"PNGs\")\nprint(\"Tables :\", len(list(TAB_DIR.glob(\"*.csv\"))), \"CSVs\")\nprint(\"\\nExample outputs:\")\nfor p in sorted(list(FIG_DIR.glob(\"*.png\")))[:10]:\n    print(\" - figures/\", p.name)\nfor p in sorted(list(TAB_DIR.glob(\"*.csv\")))[:10]:\n    print(\" - tables /\", p.name)\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-12T17:30:21.037204Z","iopub.execute_input":"2025-12-12T17:30:21.037573Z","iopub.status.idle":"2025-12-12T17:31:15.204136Z","shell.execute_reply.started":"2025-12-12T17:30:21.037547Z","shell.execute_reply":"2025-12-12T17:31:15.203106Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# %% [ADD-ON CELL] Extra CAFA6 EDA for IEEE report (drop-in after your current pipeline)\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom pathlib import Path\nfrom collections import Counter, defaultdict\nimport math\n\n# ---- safe display fallback\ntry:\n    from IPython.display import display\nexcept Exception:\n    def display(x): \n        print(x)\n\n# ---- paths / dirs (reuse if already defined)\nDATA_DIR  = globals().get(\"DATA_DIR\", Path(\"/kaggle/input/cafa-6-protein-function-prediction\"))\nTRAIN_DIR = globals().get(\"TRAIN_DIR\", DATA_DIR / \"Train\")\nTEST_DIR  = globals().get(\"TEST_DIR\",  DATA_DIR / \"Test\")\nFIG_DIR   = globals().get(\"FIG_DIR\", Path(\"/kaggle/working/figures\"))\nFIG_DIR.mkdir(parents=True, exist_ok=True)\n\ndef _savefig(name: str):\n    # reuse your savefig() if exists; else create local one\n    if \"savefig\" in globals() and callable(globals()[\"savefig\"]):\n        globals()[\"savefig\"](name)\n        return\n    png = FIG_DIR / f\"{name}.png\"\n    pdf = FIG_DIR / f\"{name}.pdf\"\n    plt.savefig(png, bbox_inches=\"tight\", dpi=300)\n    plt.savefig(pdf, bbox_inches=\"tight\", dpi=300)\n    print(\"Saved:\", png, \"and\", pdf)\n\ndef _norm(s: str) -> str:\n    return str(s).strip().lower().replace(\" \", \"\").replace(\"-\", \"\").replace(\"_\", \"\")\n\ndef _guess_col(df: pd.DataFrame, candidates):\n    cols = list(df.columns)\n    norm_map = {_norm(c): c for c in cols}\n    for cand in candidates:\n        if _norm(cand) in norm_map:\n            return norm_map[_norm(cand)]\n    for c in cols:\n        nc = _norm(c)\n        for cand in candidates:\n            if _norm(cand) in nc:\n                return c\n    raise KeyError(f\"Missing columns; tried {candidates}; have {cols}\")\n\n# ---- 1) Load / normalize core tables if missing\nif \"train_terms\" not in globals():\n    train_terms = pd.read_csv(TRAIN_DIR / \"train_terms.tsv\", sep=\"\\t\", dtype=str)\n\nif \"train_tax\" not in globals():\n    train_tax = pd.read_csv(TRAIN_DIR / \"train_taxonomy.tsv\", sep=\"\\t\", dtype=str)\n\n# IA.tsv is headerless; protect against mis-read\nif \"ia_df\" not in globals():\n    ia_df = pd.read_csv(DATA_DIR / \"IA.tsv\", sep=\"\\t\", header=None, names=[\"term\", \"IA\"])\nelse:\n    # if columns look like first data row (e.g., ['GO:0000001','0.0']), reload correctly\n    if ia_df.shape[1] == 2 and all(str(c).startswith(\"GO:\") for c in ia_df.columns):\n        ia_df = pd.read_csv(DATA_DIR / \"IA.tsv\", sep=\"\\t\", header=None, names=[\"term\", \"IA\"])\n    else:\n        ia_df = ia_df.iloc[:, :2].copy()\n        ia_df.columns = [\"term\", \"IA\"]\n\nia_df[\"IA\"] = ia_df[\"IA\"].astype(float)\nia_df = ia_df.dropna().drop_duplicates(subset=[\"term\"]).copy()\nia_map = dict(zip(ia_df[\"term\"], ia_df[\"IA\"]))\n\n# normalize train_terms columns\nENTRY_COL  = _guess_col(train_terms, [\"EntryID\", \"entry_id\", \"entry\", \"protein\", \"protein_id\", \"id\"])\nTERM_COL   = _guess_col(train_terms, [\"term\", \"go_id\", \"go\", \"GO_ID\"])\nASPECT_COL = _guess_col(train_terms, [\"aspect\", \"ontology\", \"namespace\"])\ntrain_terms = train_terms[[ENTRY_COL, TERM_COL, ASPECT_COL]].copy()\ntrain_terms.columns = [\"EntryID\", \"term\", \"aspect\"]\n\n# normalize taxonomy columns\nTAX_ENTRY_COL = _guess_col(train_tax, [\"EntryID\", \"entry_id\", \"entry\", \"protein\", \"protein_id\", \"id\"])\nTAX_COL       = _guess_col(train_tax, [\"taxon\", \"taxon_id\", \"taxonomy\", \"taxid\"])\ntrain_tax2 = train_tax[[TAX_ENTRY_COL, TAX_COL]].copy()\ntrain_tax2.columns = [\"EntryID\", \"taxon\"]\n\n# ---- 2) FASTA header parsing to get Swiss-Prot vs TrEMBL source (sp/tr) + robust EntryID\ndef parse_fasta_token(token: str):\n    # examples: \"sp|A0A0C5B5G6|MOTSC_HUMAN\" or \"A0A0C5B5G6\"\n    parts = token.split(\"|\")\n    if len(parts) >= 2 and parts[0] in (\"sp\", \"tr\"):\n        return parts[0], parts[1]\n    return \"na\", token\n\ndef fasta_id_table(fasta_path: Path, max_records=None):\n    ids, accs, dbs, lens = [], [], [], []\n    n = 0\n    cur = None\n    seq_len = 0\n    with open(fasta_path, \"r\") as f:\n        for line in f:\n            line = line.strip()\n            if not line:\n                continue\n            if line.startswith(\">\"):\n                if cur is not None:\n                    token = cur.split()[0]\n                    db, acc = parse_fasta_token(token)\n                    ids.append(token); accs.append(acc); dbs.append(db); lens.append(seq_len)\n                    n += 1\n                    if max_records and n >= max_records:\n                        break\n                cur = line[1:]\n                seq_len = 0\n            else:\n                seq_len += len(line)\n        if (not max_records or n < max_records) and cur is not None:\n            token = cur.split()[0]\n            db, acc = parse_fasta_token(token)\n            ids.append(token); accs.append(acc); dbs.append(db); lens.append(seq_len)\n    return pd.DataFrame({\"fasta_token\": ids, \"EntryID\": accs, \"db_source\": dbs, \"seq_len\": lens})\n\n# Prefer existing train_seq_df if present; else build from fasta (only once)\nif \"train_seq_df\" in globals() and isinstance(globals()[\"train_seq_df\"], pd.DataFrame) and \"EntryID\" in globals()[\"train_seq_df\"].columns:\n    tmp = globals()[\"train_seq_df\"][[\"EntryID\"]].astype(str).head(200)[\"EntryID\"]\n    looks_like_pipe = (tmp.str.contains(r\"\\|\").mean() > 0.5)\n    if looks_like_pipe:\n        _df = globals()[\"train_seq_df\"].copy()\n        parsed = _df[\"EntryID\"].astype(str).map(lambda x: parse_fasta_token(x.split()[0]))\n        _df[\"db_source\"] = [p[0] for p in parsed]\n        _df[\"EntryID\"] = [p[1] for p in parsed]\n        if \"seq_len\" not in _df.columns:\n            # fallback: rebuild seq_len from fasta (safer than guessing)\n            _df = fasta_id_table(TRAIN_DIR / \"train_sequences.fasta\")[[\"EntryID\", \"db_source\", \"seq_len\"]]\n        else:\n            _df = _df.rename(columns={\"seq_len\": \"seq_len\"})[[\"EntryID\", \"db_source\", \"seq_len\"]]\n        seq_df = _df.drop_duplicates(\"EntryID\")\n    else:\n        _df = globals()[\"train_seq_df\"].copy()\n        if \"seq_len\" not in _df.columns:\n            seq_df = fasta_id_table(TRAIN_DIR / \"train_sequences.fasta\")[[\"EntryID\", \"db_source\", \"seq_len\"]]\n        else:\n            seq_df = _df[[\"EntryID\", \"seq_len\"]].copy()\n            seq_df[\"db_source\"] = \"na\"\nelse:\n    seq_df = fasta_id_table(TRAIN_DIR / \"train_sequences.fasta\")[[\"EntryID\", \"db_source\", \"seq_len\"]]\n\nseq_df = seq_df.drop_duplicates(\"EntryID\")\n\n# ---- 3) Protein-level difficulty stats: IA mass, labels, aspect IA mass\naspect_map = {\"F\": \"MFO\", \"P\": \"BPO\", \"C\": \"CCO\"}\ntrain_terms[\"aspect_name\"] = train_terms[\"aspect\"].map(aspect_map).fillna(train_terms[\"aspect\"])\ntrain_terms[\"IA\"] = train_terms[\"term\"].map(ia_map).fillna(0.0)\n\nprot_stats = (\n    train_terms.groupby(\"EntryID\")\n    .agg(\n        n_labels=(\"term\", \"count\"),\n        n_unique_terms=(\"term\", \"nunique\"),\n        IA_mass=(\"IA\", \"sum\"),\n    )\n    .reset_index()\n)\n\naspect_ia = (\n    train_terms.groupby([\"EntryID\", \"aspect_name\"])[\"IA\"]\n    .sum()\n    .unstack(fill_value=0.0)\n    .reset_index()\n)\n\nprot_stats = prot_stats.merge(aspect_ia, on=\"EntryID\", how=\"left\")\nfor col in [\"BPO\", \"CCO\", \"MFO\"]:\n    if col not in prot_stats.columns:\n        prot_stats[col] = 0.0\n\nprot_stats = prot_stats.merge(seq_df, on=\"EntryID\", how=\"left\")\nprot_stats = prot_stats.merge(train_tax2, on=\"EntryID\", how=\"left\")\nprot_stats[\"db_source\"] = prot_stats[\"db_source\"].fillna(\"na\")\n\nprot_stats.to_csv(FIG_DIR / \"protein_level_stats.csv\", index=False)\n\n# ---- 4) Plot: Protein IA mass distribution + IA vs length\nplt.figure(figsize=(6.2, 4.0))\nplt.hist(prot_stats[\"IA_mass\"].dropna(), bins=120)\nplt.yscale(\"log\")\nplt.xlabel(\"Protein IA mass = Σ IA(term)\")\nplt.ylabel(\"Count (log scale)\")\nplt.title(\"CAFA6 Train: Protein-level IA mass distribution\")\n_savefig(\"fig_protein_ia_mass_dist\")\nplt.show()\n\n_s = prot_stats.dropna(subset=[\"seq_len\", \"IA_mass\"]).sample(n=min(60000, len(prot_stats)), random_state=0)\nplt.figure(figsize=(6.2, 4.0))\nplt.scatter(_s[\"seq_len\"], _s[\"IA_mass\"], s=6, alpha=0.12)\nplt.xscale(\"log\")\nplt.yscale(\"log\")\nplt.xlabel(\"Sequence length (aa, log)\")\nplt.ylabel(\"IA mass (log)\")\nplt.title(\"CAFA6 Train: IA mass vs sequence length (sample)\")\n_savefig(\"fig_ia_mass_vs_len\")\nplt.show()\n\n# ---- 5) Truncation coverage for common max_len (useful for ProtBERT/ESM/ProtT5 design)\nmax_lens = [256, 512, 1024, 2048, 4096]\nL = prot_stats[\"seq_len\"].dropna().astype(float).values\ntotal_res = float(L.sum()) if L.size else 1.0\n\nrows = []\nfor m in max_lens:\n    pct_full = float((L <= m).mean() * 100.0) if L.size else 0.0\n    pct_res  = float((np.minimum(L, m).sum() / total_res) * 100.0) if L.size else 0.0\n    rows.append((m, pct_full, pct_res))\n\ncov_df = pd.DataFrame(rows, columns=[\"max_len\", \"pct_sequences_fully_covered\", \"pct_residues_covered\"])\ndisplay(cov_df)\ncov_df.to_csv(FIG_DIR / \"truncation_coverage.csv\", index=False)\n\nplt.figure(figsize=(6.2, 4.0))\nplt.plot(cov_df[\"max_len\"], cov_df[\"pct_sequences_fully_covered\"], marker=\"o\", label=\"% sequences fully covered\")\nplt.plot(cov_df[\"max_len\"], cov_df[\"pct_residues_covered\"], marker=\"o\", label=\"% residues covered\")\nplt.xscale(\"log\")\nplt.xlabel(\"Max sequence length (log)\")\nplt.ylabel(\"Coverage (%)\")\nplt.title(\"CAFA6 Train: Truncation coverage trade-off\")\nplt.legend()\n_savefig(\"fig_truncation_coverage\")\nplt.show()\n\n# ---- 6) Swiss-Prot vs TrEMBL stratification (sp/tr/na): counts + IA_mass by source\ndb_counts = prot_stats[\"db_source\"].value_counts()\nplt.figure(figsize=(6.2, 4.0))\nplt.bar(db_counts.index.astype(str), db_counts.values)\nplt.ylabel(\"# proteins\")\nplt.title(\"CAFA6 Train: FASTA source distribution (sp/tr/na)\")\n_savefig(\"fig_db_source_counts\")\nplt.show()\n\ngroups = [prot_stats.loc[prot_stats[\"db_source\"] == g, \"IA_mass\"].dropna().values for g in db_counts.index]\nplt.figure(figsize=(6.8, 4.2))\nplt.boxplot(groups, labels=[str(g) for g in db_counts.index], showfliers=False)\nplt.yscale(\"log\")\nplt.ylabel(\"IA mass (log)\")\nplt.title(\"CAFA6 Train: Protein IA mass by FASTA source\")\n_savefig(\"fig_db_source_ia_mass_box\")\nplt.show()\n\n# ---- 7) Taxonomy skew (Gini) + Taxonomy × aspect IA heatmap (top taxa)\ntax_counts = prot_stats[\"taxon\"].value_counts(dropna=True)\ntax_counts.to_csv(FIG_DIR / \"train_taxon_counts.csv\", header=[\"count\"])\n\ndef gini(arr):\n    x = np.array(arr, dtype=float)\n    x = x[x >= 0]\n    if x.size == 0:\n        return np.nan\n    x = np.sort(x)\n    n = x.size\n    cumx = np.cumsum(x)\n    return (n + 1 - 2 * np.sum(cumx) / cumx[-1]) / n\n\ngini_tax = gini(tax_counts.values)\nprint(f\"Gini(taxon protein counts) = {gini_tax:.4f}  (higher = more skewed)\")\n\nTOP_TAXA = 20\ntop_taxa = tax_counts.head(TOP_TAXA).index.tolist()\n\nheat = (\n    prot_stats[prot_stats[\"taxon\"].isin(top_taxa)]\n    .groupby(\"taxon\")[[\"BPO\", \"CCO\", \"MFO\"]]\n    .mean()\n    .reindex(top_taxa)\n)\nheat.to_csv(FIG_DIR / \"taxon_aspect_mean_IA.csv\")\n\nplt.figure(figsize=(7.0, 5.2))\nplt.imshow(heat.values, aspect=\"auto\")\nplt.yticks(range(len(heat.index)), [str(x) for x in heat.index], fontsize=8)\nplt.xticks(range(heat.shape[1]), heat.columns.tolist())\nplt.colorbar(label=\"Mean IA mass (per protein)\")\nplt.title(\"CAFA6 Train: Taxonomy × aspect (mean IA mass, top taxa)\")\n_savefig(\"fig_taxon_aspect_mean_IA_heatmap\")\nplt.show()\n\n# ---- 8) GO polyhierarchy (parents/children per observed term) using a minimal OBO parser\ndef parse_go_basic_obo_min(path: Path):\n    terms = {}\n    cur = None\n    in_term = False\n    with open(path, \"r\") as f:\n        for line in f:\n            line = line.rstrip(\"\\n\")\n            if line == \"[Term]\":\n                in_term = True\n                cur = {\"id\": None, \"namespace\": None, \"parents\": []}\n                continue\n            if line.startswith(\"[\") and line != \"[Term]\":\n                in_term = False\n                cur = None\n                continue\n            if not in_term or cur is None:\n                continue\n            if not line:\n                if cur.get(\"id\"):\n                    terms[cur[\"id\"]] = {\"namespace\": cur[\"namespace\"], \"parents\": cur[\"parents\"]}\n                cur = {\"id\": None, \"namespace\": None, \"parents\": []}\n                continue\n            if line.startswith(\"id:\"):\n                cur[\"id\"] = line.split(\"id:\")[1].strip()\n            elif line.startswith(\"namespace:\"):\n                cur[\"namespace\"] = line.split(\"namespace:\")[1].strip()\n            elif line.startswith(\"is_a:\"):\n                cur[\"parents\"].append(line.split(\"is_a:\")[1].strip().split()[0])\n        if in_term and cur and cur.get(\"id\"):\n            terms[cur[\"id\"]] = {\"namespace\": cur[\"namespace\"], \"parents\": cur[\"parents\"]}\n    return terms\n\ngo_terms = globals().get(\"go_terms\", None)\nif go_terms is None or not isinstance(go_terms, dict) or len(go_terms) == 0:\n    go_terms = parse_go_basic_obo_min(TRAIN_DIR / \"go-basic.obo\")\n\nobserved_terms = set(train_terms[\"term\"].unique())\nchildren = defaultdict(list)\nparent_ct = {}\nfor t in observed_terms:\n    ps = go_terms.get(t, {}).get(\"parents\", [])\n    parent_ct[t] = len(ps)\n    for p in ps:\n        children[p].append(t)\n\nchild_ct = {t: len(children.get(t, [])) for t in observed_terms}\n\ndeg_df = pd.DataFrame({\n    \"term\": list(observed_terms),\n    \"n_parents\": [parent_ct[t] for t in observed_terms],\n    \"n_children\": [child_ct.get(t, 0) for t in observed_terms],\n    \"IA\": [ia_map.get(t, np.nan) for t in observed_terms],\n})\ndeg_df.to_csv(FIG_DIR / \"observed_term_degrees.csv\", index=False)\n\nplt.figure(figsize=(6.2, 4.0))\nplt.hist(deg_df[\"n_parents\"], bins=np.arange(0, deg_df[\"n_parents\"].max() + 2) - 0.5)\nplt.yscale(\"log\")\nplt.xlabel(\"# parents per observed GO term\")\nplt.ylabel(\"Count (log scale)\")\nplt.title(\"CAFA6 Train: GO polyhierarchy (parents per term)\")\n_savefig(\"fig_go_parents_per_term\")\nplt.show()\n\nplt.figure(figsize=(6.2, 4.0))\nplt.hist(np.clip(deg_df[\"n_children\"], 0, 200), bins=80)\nplt.yscale(\"log\")\nplt.xlabel(\"# children per observed GO term (clipped at 200)\")\nplt.ylabel(\"Count (log scale)\")\nplt.title(\"CAFA6 Train: GO branching (children per term)\")\n_savefig(\"fig_go_children_per_term\")\nplt.show()\n\n# ---- 9) Co-occurrence with PMI (more informative than raw co-occur counts)\nTOPN = 60\nN_PROT_SAMPLE = 20000   # increase if you want (slower)\nMIN_SUPPORT = 30\n\ntop_terms = train_terms[\"term\"].value_counts().head(TOPN).index.tolist()\ntop_set = set(top_terms)\n\nprot_ids = train_terms[\"EntryID\"].drop_duplicates().values\nrng = np.random.default_rng(0)\nsample_prots = set(rng.choice(prot_ids, size=min(N_PROT_SAMPLE, len(prot_ids)), replace=False))\n\nsub = train_terms[train_terms[\"EntryID\"].isin(sample_prots)]\nlabels = sub.groupby(\"EntryID\")[\"term\"].agg(lambda x: list(set([t for t in x if t in top_set])))\n\nc_i = Counter()\nc_ij = Counter()\nfor terms in labels.values:\n    for t in terms:\n        c_i[t] += 1\n    if len(terms) >= 2:\n        terms = sorted(terms)\n        for i in range(len(terms)):\n            for j in range(i + 1, len(terms)):\n                c_ij[(terms[i], terms[j])] += 1\n\nN = len(labels)\nrows = []\nfor (a, b), cij in c_ij.items():\n    p_ab = cij / N\n    p_a = c_i[a] / N\n    p_b = c_i[b] / N\n    if p_ab > 0 and p_a > 0 and p_b > 0:\n        pmi = math.log2(p_ab / (p_a * p_b))\n        lift = p_ab / (p_a * p_b)\n    else:\n        pmi = -np.inf\n        lift = np.nan\n    rows.append((a, b, cij, c_i[a], c_i[b], pmi, lift))\n\npmi_df = pd.DataFrame(rows, columns=[\"term_a\", \"term_b\", \"cooccur\", \"count_a\", \"count_b\", \"PMI\", \"lift\"])\npmi_df = pmi_df.sort_values([\"PMI\", \"cooccur\"], ascending=[False, False])\n\ndisplay(pmi_df.head(20))\npmi_df.to_csv(FIG_DIR / \"top_term_pairs_by_PMI_all.csv\", index=False)\npmi_df.head(200).to_csv(FIG_DIR / \"top_term_pairs_by_PMI_top200.csv\", index=False)\n\ntop_pairs = pmi_df[pmi_df[\"cooccur\"] >= MIN_SUPPORT].head(20).iloc[::-1]\nplt.figure(figsize=(7.2, 4.6))\nplt.barh([f\"{a} | {b}\" for a, b in zip(top_pairs[\"term_a\"], top_pairs[\"term_b\"])], top_pairs[\"PMI\"])\nplt.xlabel(\"PMI (log2)\")\nplt.title(f\"CAFA6 Train: Top term pairs by PMI (top-{TOPN} terms, support≥{MIN_SUPPORT})\")\n_savefig(\"fig_top_term_pairs_pmi\")\nplt.show()\n\nprint(\"✅ Extra EDA finished. Outputs saved to:\", FIG_DIR)\nprint(\"Sample files:\", sorted([p.name for p in FIG_DIR.glob(\"*\")])[:30])\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T17:45:52.576149Z","iopub.execute_input":"2025-12-12T17:45:52.576518Z"}},"outputs":[],"execution_count":null}]}