{"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"}],"dockerImageVersionId":31260,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ============================================================\n# CAFA 6 Protein Function Prediction\n# ============================================================\n\nimport pandas as pd\nfrom tqdm import tqdm\nfrom collections import defaultdict\n\n# ------------------------------------------------------------\n# Paths\n# ------------------------------------------------------------\nDATA_ROOT = \"/kaggle/input/cafa-6-protein-function-prediction\"\n\nTRAIN_TERMS = f\"{DATA_ROOT}/Train/train_terms.tsv\"\nTRAIN_TAX  = f\"{DATA_ROOT}/Train/train_taxonomy.tsv\"\nTEST_FASTA = f\"{DATA_ROOT}/Test/testsuperset.fasta\"\nTEST_TAX   = f\"{DATA_ROOT}/Test/testsuperset-taxon-list.tsv\"\n\nOUT_FILE = \"submission.tsv\"\n\n# ------------------------------------------------------------\n# Hyperparameters\n# ------------------------------------------------------------\nTOP_K = {\n    \"F\": 30,\n    \"P\": 60,\n    \"C\": 30\n}\n\nMIN_PROTEINS_PER_TAXON = 50   # below this → fallback to global\nCHUNK_SIZE = 1000             # proteins per write chunk\n\n# ------------------------------------------------------------\n# Load training data\n# ------------------------------------------------------------\nprint(\"Loading training annotations...\")\ntrain_terms = pd.read_csv(TRAIN_TERMS, sep=\"\\t\")\ntrain_terms.columns = [\"protein_id\", \"go_id\", \"aspect\"]\n\ntrain_tax = pd.read_csv(\n    TRAIN_TAX,\n    sep=\"\\t\",\n    names=[\"protein_id\", \"taxon\"]\n)\n\ntrain = train_terms.merge(train_tax, on=\"protein_id\")\n\n# ------------------------------------------------------------\n# Load test taxonomy\n# ------------------------------------------------------------\ntest_tax = pd.read_csv(\n    TEST_TAX,\n    sep=\"\\t\",\n    names=[\"protein_id\", \"taxon\"]\n)\n\n# ------------------------------------------------------------\n# FAST FASTA ID reader\n# ------------------------------------------------------------\ndef read_fasta_ids(path):\n    with open(path) as f:\n        return [line[1:].split()[0] for line in f if line.startswith(\">\")]\n\nprint(\"Reading test FASTA...\")\ntest_ids = read_fasta_ids(TEST_FASTA)\n\ntest_df = (\n    pd.DataFrame({\"protein_id\": test_ids})\n    .merge(test_tax, on=\"protein_id\", how=\"left\")\n)\n\nprint(\"Test proteins:\", len(test_df))\n\n# ------------------------------------------------------------\n# GLOBAL GO PRIORS (fallback)\n# ------------------------------------------------------------\nprint(\"Computing global priors...\")\nglobal_stats = (\n    train\n    .groupby([\"aspect\", \"go_id\"])\n    .size()\n    .reset_index(name=\"count\")\n)\n\nglobal_stats[\"prob\"] = (\n    global_stats.groupby(\"aspect\")[\"count\"]\n    .transform(lambda x: (x / x.max()) ** 0.5)\n)\n\nglobal_top = {}\nfor aspect, k in TOP_K.items():\n    df = (\n        global_stats[global_stats[\"aspect\"] == aspect]\n        .sort_values(\"prob\", ascending=False)\n        .head(k)\n    )\n    global_top[aspect] = df[[\"go_id\", \"prob\"]].values.tolist()\n\n# ------------------------------------------------------------\n# TAXONOMY-CONDITIONED PRIORS\n# ------------------------------------------------------------\nprint(\"Computing taxonomy-conditioned priors...\")\ntaxon_counts = train.groupby(\"taxon\")[\"protein_id\"].nunique()\n\ntaxon_priors = defaultdict(dict)\n\nfor taxon, n_prot in tqdm(taxon_counts.items()):\n    if n_prot < MIN_PROTEINS_PER_TAXON:\n        continue\n\n    subset = train[train[\"taxon\"] == taxon]\n\n    stats = (\n        subset\n        .groupby([\"aspect\", \"go_id\"])\n        .size()\n        .reset_index(name=\"count\")\n    )\n\n    stats[\"prob\"] = (\n        stats.groupby(\"aspect\")[\"count\"]\n        .transform(lambda x: (x / x.max()) ** 0.5)\n    )\n\n    for aspect, k in TOP_K.items():\n        df = (\n            stats[stats[\"aspect\"] == aspect]\n            .sort_values(\"prob\", ascending=False)\n            .head(k)\n        )\n        if len(df) > 0:\n            taxon_priors[taxon][aspect] = df[[\"go_id\", \"prob\"]].values.tolist()\n\nprint(\"Taxa with learned priors:\", len(taxon_priors))\n\n# ------------------------------------------------------------\n# WRITE SUBMISSION\n# ------------------------------------------------------------\nprint(\"Writing submission...\")\nwith open(OUT_FILE, \"w\") as f_out:\n    for i in tqdm(range(0, len(test_df), CHUNK_SIZE)):\n        chunk = test_df.iloc[i:i + CHUNK_SIZE]\n\n        rows = []\n        for _, row in chunk.iterrows():\n            pid = row[\"protein_id\"]\n            taxon = row[\"taxon\"]\n\n            for aspect in [\"F\", \"P\", \"C\"]:\n                if taxon in taxon_priors and aspect in taxon_priors[taxon]:\n                    terms = taxon_priors[taxon][aspect]\n                else:\n                    terms = global_top[aspect]\n\n                for go_id, prob in terms:\n                    rows.append([pid, go_id, f\"{prob:.3f}\"])\n\n        pd.DataFrame(rows).to_csv(\n            f_out,\n            sep=\"\\t\",\n            header=False,\n            index=False\n        )\n\nprint(\"\\n Submission written:\", OUT_FILE)\n\n# Preview\nprint(pd.read_csv(OUT_FILE, sep=\"\\t\", header=None, nrows=10))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T17:29:45.835624Z","iopub.execute_input":"2026-01-26T17:29:45.835878Z"}},"outputs":[],"execution_count":null}]}